fix: forward idempotency key on balance creation and document create race (EN-1205)

pkg/api/handler_balances_create.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"net/http"
 
+	"github.com/formancehq/go-libs/v5/pkg/transport/api"
 	wallet "github.com/formancehq/wallets/pkg"
 	"github.com/go-chi/render"
 )
@@ -17,7 +18,7 @@ func (m *MainHandler) createBalanceHandler(w http.ResponseWriter, r *http.Reques
 		}
 	}
 
-	balance, err := m.manager.CreateBalance(r.Context(), data)
+	balance, err := m.manager.CreateBalance(r.Context(), api.IdempotencyKeyFromRequest(r), data)
 	if err != nil {
 		switch {
 		case errors.Is(err, wallet.ErrInvalidBalanceName):

pkg/api/handler_balances_create_test.go

@@ -58,6 +58,135 @@ var balanceCreateTestCases = []balanceCreateTestCase{
 	},
 }
 
+func TestBalancesCreateForwardsIdempotencyKey(t *testing.T) {
+	t.Parallel()
+
+	const idempotencyKey = "create-balance-key-1"
+	walletID := uuid.NewString()
+
+	var forwardedKey string
+	testEnv := newTestEnv(
+		WithAddMetadataToAccount(func(ctx context.Context, ledger, account, ik string, metadata map[string]string) error {
+			forwardedKey = ik
+			return nil
+		}),
+		WithGetAccount(func(ctx context.Context, ledger, account string) (*wallet.AccountWithVolumesAndBalances, error) {
+			return nil, wallet.ErrAccountNotFound
+		}),
+	)
+
+	req := newRequest(t, http.MethodPost, "/wallets/"+walletID+"/balances", wallet.CreateBalance{Name: uuid.NewString()})
+	req.Header.Set("Idempotency-Key", idempotencyKey)
+	rec := httptest.NewRecorder()
+	testEnv.Router().ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusCreated, rec.Result().StatusCode)
+	require.Equal(t, idempotencyKey, forwardedKey)
+}
+
+func TestBalancesCreateIdempotentReplay(t *testing.T) {
+	t.Parallel()
+
+	const idempotencyKey = "create-balance-key-replay"
+	walletID := uuid.NewString()
+	const balanceName = "savings"
+
+	var (
+		created         bool
+		addCalls        int
+		appliedMetadata map[string]string
+	)
+	testEnv := newTestEnv(
+		WithGetAccount(func(ctx context.Context, ledger, account string) (*wallet.AccountWithVolumesAndBalances, error) {
+			if created {
+				return &wallet.AccountWithVolumesAndBalances{
+					Account: wallet.Account{
+						Address:  account,
+						Metadata: metadataWithExpectingTypesAfterUnmarshalling(appliedMetadata),
+					},
+				}, nil
+			}
+			return nil, wallet.ErrAccountNotFound
+		}),
+		WithAddMetadataToAccount(func(ctx context.Context, ledger, account, ik string, md map[string]string) error {
+			addCalls++
+			appliedMetadata = md
+			created = true
+			return nil
+		}),
+	)
+
+	create := func() *httptest.ResponseRecorder {
+		req := newRequest(t, http.MethodPost, "/wallets/"+walletID+"/balances", wallet.CreateBalance{Name: balanceName})
+		req.Header.Set("Idempotency-Key", idempotencyKey)
+		rec := httptest.NewRecorder()
+		testEnv.Router().ServeHTTP(rec, req)
+		return rec
+	}
+
+	first := create()
+	require.Equal(t, http.StatusCreated, first.Result().StatusCode)
+
+	// The retry under the same key replays the existing balance (201) instead
+	// of failing with 400 ALREADY_EXISTS, and does not re-write metadata.
+	second := create()
+	require.Equal(t, http.StatusCreated, second.Result().StatusCode)
+	require.Equal(t, 1, addCalls)
+
+	b1, b2 := &wallet.Balance{}, &wallet.Balance{}
+	readResponse(t, first, b1)
+	readResponse(t, second, b2)
+	require.Equal(t, balanceName, b1.Name)
+	require.Equal(t, b1.Name, b2.Name)
+}
+
+func TestBalancesCreateWithDifferentIdempotencyKeyDoesNotReplay(t *testing.T) {
+	t.Parallel()
+
+	walletID := uuid.NewString()
+	const balanceName = "savings"
+
+	var (
+		created         bool
+		addCalls        int
+		appliedMetadata map[string]string
+	)
+	testEnv := newTestEnv(
+		WithGetAccount(func(ctx context.Context, ledger, account string) (*wallet.AccountWithVolumesAndBalances, error) {
+			if created {
+				return &wallet.AccountWithVolumesAndBalances{
+					Account: wallet.Account{
+						Address:  account,
+						Metadata: metadataWithExpectingTypesAfterUnmarshalling(appliedMetadata),
+					},
+				}, nil
+			}
+			return nil, wallet.ErrAccountNotFound
+		}),
+		WithAddMetadataToAccount(func(ctx context.Context, ledger, account, ik string, md map[string]string) error {
+			addCalls++
+			appliedMetadata = md
+			created = true
+			return nil
+		}),
+	)
+
+	create := func(idempotencyKey string) *httptest.ResponseRecorder {
+		req := newRequest(t, http.MethodPost, "/wallets/"+walletID+"/balances", wallet.CreateBalance{Name: balanceName})
+		req.Header.Set("Idempotency-Key", idempotencyKey)
+		rec := httptest.NewRecorder()
+		testEnv.Router().ServeHTTP(rec, req)
+		return rec
+	}
+
+	first := create("create-balance-key-1")
+	require.Equal(t, http.StatusCreated, first.Result().StatusCode)
+
+	second := create("create-balance-key-2")
+	require.Equal(t, http.StatusBadRequest, second.Result().StatusCode)
+	require.Equal(t, 1, addCalls)
+}
+
 func TestBalancesCreate(t *testing.T) {
 	t.Parallel()
 

pkg/manager.go

@@ -2,6 +2,8 @@ package wallet
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"math/big"
 	"slices"
@@ -555,7 +557,16 @@ func (m *Manager) GetHold(ctx context.Context, id string) (*ExpandedDebitHold, e
 	return Ptr(ExpandedDebitHoldFromLedgerAccount(*account)), nil
 }
 
-func (m *Manager) CreateBalance(ctx context.Context, data *CreateBalance) (*Balance, error) {
+// CreateBalance creates a named balance for a wallet.
+//
+// This is a check-then-act on account metadata: the existence check and the
+// metadata write are not a single atomic operation. The Idempotency-Key makes
+// retries of a previously recorded create safe, but two genuinely concurrent
+// first-time creations of the same balance can both pass the existence check;
+// the writes are then last-write-wins on priority/expiresAt for that single
+// account. A fully atomic create would require a conditional metadata write on
+// the ledger side, which the API does not currently expose.
+func (m *Manager) CreateBalance(ctx context.Context, ik string, data *CreateBalance) (*Balance, error) {
 	if err := data.Validate(); err != nil {
 		return nil, err
 	}
@@ -565,6 +576,9 @@ func (m *Manager) CreateBalance(ctx context.Context, data *CreateBalance) (*Bala
 	case err == nil:
 		if ret.Metadata != nil &&
 			ret.Metadata[MetadataKeyWalletBalance] == TrueValue {
+			if ik != "" && ret.Metadata[MetadataKeyBalanceIdempotencyKey] == hashIdempotencyKey(ik) {
+				return Ptr(BalanceFromAccount(*ret)), nil
+			}
 			return nil, ErrBalanceAlreadyExists
 		}
 	default:
@@ -573,20 +587,29 @@ func (m *Manager) CreateBalance(ctx context.Context, data *CreateBalance) (*Bala
 
 	balance := NewBalance(data.Name, data.ExpiresAt)
 	balance.Priority = data.Priority
+	balanceMetadata := balance.LedgerMetadata(data.WalletID)
+	if ik != "" {
+		balanceMetadata[MetadataKeyBalanceIdempotencyKey] = hashIdempotencyKey(ik)
+	}
 
 	if err := m.client.AddMetadataToAccount(
 		ctx,
 		m.ledgerName,
 		m.chart.GetBalanceAccount(data.WalletID, balance.Name),
-		"",
-		balance.LedgerMetadata(data.WalletID),
+		ik,
+		balanceMetadata,
 	); err != nil {
 		return nil, errors.Wrap(err, "adding metadata to account")
 	}
 
 	return &balance, nil
 }
 
+func hashIdempotencyKey(ik string) string {
+	hash := sha256.Sum256([]byte(ik))
+	return hex.EncodeToString(hash[:])
+}
+
 func (m *Manager) GetBalance(ctx context.Context, walletID string, balanceName string) (*ExpandedBalance, error) {
 	account, err := m.client.GetAccount(ctx, m.ledgerName, m.chart.GetBalanceAccount(walletID, balanceName))
 	if err != nil {

pkg/metadata.go

@@ -21,6 +21,7 @@ const (
 	MetadataKeyBalanceName           = "wallets/balances/name"
 	MetadataKeyBalanceExpiresAt      = "wallets/balances/expiresAt"
 	MetadataKeyBalancePriority       = "wallets/balances/priority"
+	MetadataKeyBalanceIdempotencyKey = "wallets/balances/idempotencyKeyHash"
 	MetadataKeyWalletBalance         = "wallets/balances"
 	MetadataKeyCreatedAt             = "wallets/createdAt"