Correctness fixes: maxFee, 504 retries, key cache

@stellar/typescript-wallet-sdk-km/src/keyManager.ts

@@ -466,8 +466,13 @@ export class KeyManager {
   }
 
   private _writeIndexCache(id: string, key: Key | undefined) {
-    if (this.shouldCache && key) {
+    if (!this.shouldCache) {
+      return;
+    }
+    if (key) {
       this.keyCache[id] = key;
+    } else {
+      delete this.keyCache[id];
     }
   }
 }

@stellar/typescript-wallet-sdk-km/test/keyManager.test.ts

@@ -192,6 +192,64 @@ describe("KeyManager", () => {
     }
   });
 
+  test("removeKey clears the in-memory cache", async () => {
+    const testStore = new MemoryKeyStore();
+    const testKeyManager = new KeyManager({
+      keyStore: testStore,
+      shouldCache: true,
+    });
+    const network = Networks.TESTNET;
+
+    testKeyManager.registerEncrypter(IdentityEncrypter);
+
+    const keypair = Keypair.master(network);
+    const password = "test";
+
+    const keyMetadata = await testKeyManager.storeKey({
+      key: {
+        type: KeyType.plaintextKey,
+        publicKey: keypair.publicKey(),
+        privateKey: keypair.secret(),
+        network,
+      },
+      password,
+      encrypterName: "IdentityEncrypter",
+    });
+
+    // Sanity check: with the cached key, signing succeeds.
+    const source = new Account(
+      "GBBM6BKZPEHWYO3E3YKREDPQXMS4VK35YLNU7NFBRI26RAN7GI5POFBB",
+      "0",
+    );
+    const transaction = new TransactionBuilder(source, {
+      fee: "100",
+      networkPassphrase: network,
+    })
+      .addOperation(Operation.inflation({}))
+      .setTimeout(TimeoutInfinite)
+      .build();
+
+    await testKeyManager.signTransaction({
+      transaction,
+      id: keyMetadata.id,
+      password,
+    });
+
+    // Remove the key, then attempt to sign again. If the cache is still
+    // serving the deleted key, signing succeeds and this test fails.
+    await testKeyManager.removeKey(keyMetadata.id);
+
+    await expect(
+      testKeyManager.signTransaction({
+        transaction,
+        id: keyMetadata.id,
+        password,
+      }),
+    ).rejects.toThrow(
+      `Couldn't sign the transaction: no key with id '${keyMetadata.id}' found.`,
+    );
+  });
+
   test("Sign transactions", async () => {
     // set up the manager
     const testStore = new MemoryKeyStore();
@@ -1114,9 +1172,7 @@ describe("fetchAuthToken", () => {
 
       expect("This test failed: transaction didn't cause error").toBe(null);
     } catch (e) {
-      expect(e.toString()).toMatch(
-        `Transaction not signed by server`,
-      );
+      expect(e.toString()).toMatch(`Transaction not signed by server`);
     }
   });
 

@stellar/typescript-wallet-sdk/src/walletSdk/Horizon/Stellar.ts

@@ -24,6 +24,13 @@ import {
 import { getResultCode } from "../Utils/getResultCode";
 import { SigningKeypair } from "./Account";
 
+const SUBMIT_504_MAX_RETRIES = 5;
+const SUBMIT_504_BASE_DELAY_MS = 1000;
+const SUBMIT_504_MAX_DELAY_MS = 30000;
+
+const sleep = (ms: number) =>
+  new Promise<void>((resolve) => setTimeout(resolve, ms));
+
 /**
  * Interaction with the Stellar Network.
  * Do not create this object directly, use the Wallet class.
@@ -118,30 +125,58 @@ export class Stellar {
   }
 
   /**
-   * Submits a signed transaction to the server. If the submission fails with status
-   * 504 indicating a timeout error, it will automatically retry.
+   * Submits a signed transaction to Horizon.
+   *
+   * On HTTP 504 (timeout), retries with exponential backoff up to
+   * {@link SUBMIT_504_MAX_RETRIES} times. Any non-504 error is rethrown
+   * immediately without retrying.
+   *
    * @param {Transaction|FeeBumpTransaction} signedTransaction - The signed transaction to submit.
-   * @returns {boolean} `true` if the transaction was successfully submitted.
-   * @throws {TransactionSubmitFailedError} If the transaction submission fails.
+   * @returns {boolean} `true` if Horizon confirmed the submission as successful.
+   * @throws {TransactionSubmitFailedError} If Horizon responded with a non-successful
+   *   submission result (the transaction reached Horizon but was rejected).
+   * @throws The underlying 504 error if every retry attempt timed out. In this
+   *   case the transaction's on-chain status is **indeterminate** — Horizon may
+   *   have ingested it on the final attempt without responding in time. Callers
+   *   should poll the transaction hash to determine the actual outcome rather
+   *   than resubmit blindly; resubmitting a signed transaction with the same
+   *   sequence number will fail with `tx_bad_seq` once the original lands.
    */
   async submitTransaction(
     signedTransaction: Transaction | FeeBumpTransaction,
   ): Promise<boolean> {
-    try {
-      const response = await this.server.submitTransaction(signedTransaction);
-      if (!response.successful) {
-        throw new TransactionSubmitFailedError(response);
-      }
-      return true;
-    } catch (e) {
-      if (e.response.status === 504) {
-        // in case of 504, keep retrying this tx until submission succeeds or we get a different error
+    let lastError: unknown;
+    for (let attempt = 0; attempt <= SUBMIT_504_MAX_RETRIES; attempt++) {
+      try {
+        const response = await this.server.submitTransaction(signedTransaction);
+        if (!response.successful) {
+          throw new TransactionSubmitFailedError(response);
+        }
+        return true;
+      } catch (e) {
+        if (e?.response?.status !== 504) {
+          throw e;
+        }
+        lastError = e;
+        if (attempt === SUBMIT_504_MAX_RETRIES) {
+          break;
+        }
         // https://developers.stellar.org/api/errors/http-status-codes/horizon-specific/timeout
         // https://developers.stellar.org/docs/encyclopedia/error-handling#timeouts
-        return await this.submitTransaction(signedTransaction);
+        // Equal-jitter backoff: each attempt waits at least half the capped
+        // exponential delay (a deterministic, progressive floor) plus a random
+        // amount up to the other half. Keeps the schedule predictable while
+        // smoothing correlated retries across many clients during a Horizon
+        // outage. Total sleep stays within SUBMIT_504_MAX_DELAY_MS.
+        const cappedDelay = Math.min(
+          SUBMIT_504_BASE_DELAY_MS * 2 ** attempt,
+          SUBMIT_504_MAX_DELAY_MS,
+        );
+        const half = cappedDelay / 2;
+        await sleep(half + Math.random() * half);
       }
-      throw e;
     }
+    throw lastError;
   }
 
   /**
@@ -211,6 +246,7 @@ export class Stellar {
           signerFunction,
           baseFee: newFee,
           memo,
+          maxFee,
         });
       }
       throw e;

@stellar/typescript-wallet-sdk/test/stellar.test.ts

@@ -1,5 +1,12 @@
-import { Keypair, Memo, MemoText, Horizon } from "@stellar/stellar-sdk";
+import {
+  Keypair,
+  Memo,
+  MemoText,
+  Horizon,
+  Transaction,
+} from "@stellar/stellar-sdk";
 import axios from "axios";
+import sinon from "sinon";
 
 import { Stellar, Wallet } from "../src";
 import {
@@ -14,6 +21,7 @@ import {
 } from "../src/walletSdk/Asset";
 import { TransactionStatus, WithdrawTransaction } from "../src/walletSdk/Types";
 import {
+  TransactionSubmitWithFeeIncreaseFailedError,
   WithdrawalTxMissingDestinationError,
   WithdrawalTxMissingMemoError,
   WithdrawalTxNotPendingUserTransferStartError,
@@ -158,6 +166,51 @@ describe("Stellar", () => {
     expect(txn.fee).toBe("200");
   });
 
+  it("should enforce maxFee across multiple submitWithFeeIncrease retries", async () => {
+    const txTooLateRejection = {
+      response: {
+        status: 400,
+        statusText: "Bad Request",
+        data: {
+          extras: {
+            result_codes: { transaction: "tx_too_late" },
+          },
+        },
+      },
+    };
+
+    // Always reject with tx_too_late so the fee climbs on every retry. With the
+    // fix the maxFee cap is honored across all retries and submitWithFeeIncrease
+    // eventually throws; without it the cap is dropped after the first retry and
+    // the recursion never stops. Restore the spy in a finally block: a leaked
+    // mock on submitTransaction would turn the next test's real submission into
+    // a no-op.
+    const submitStub = jest
+      .spyOn(stellar, "submitTransaction")
+      .mockRejectedValue(txTooLateRejection);
+
+    const buildingFunction = (builder) =>
+      builder.transfer(kp.publicKey, new NativeAssetId(), "2");
+
+    try {
+      await expect(
+        stellar.submitWithFeeIncrease({
+          sourceAddress: kp,
+          timeout: 180,
+          baseFeeIncrease: 100,
+          buildingFunction,
+          baseFee: 100,
+          maxFee: 250,
+        }),
+      ).rejects.toThrow(TransactionSubmitWithFeeIncreaseFailedError);
+      // The cap is enforced only after several retries, proving it survives
+      // past the first retry (the bug being fixed).
+      expect(submitStub.mock.calls.length).toBeGreaterThan(1);
+    } finally {
+      submitStub.mockRestore();
+    }
+  });
+
   it("should add and remove asset support", async () => {
     const asset = new IssuedAssetId(
       "USDC",
@@ -202,6 +255,74 @@ describe("Stellar", () => {
     expect(fee).toBeTruthy();
   });
 
+  describe("submitTransaction 504 handling", () => {
+    let clock: sinon.SinonFakeTimers;
+    let randomStub: sinon.SinonStub;
+
+    beforeEach(() => {
+      clock = sinon.useFakeTimers();
+      // Zero jitter for deterministic test timing.
+      randomStub = sinon.stub(Math, "random").returns(0);
+    });
+
+    afterEach(() => {
+      clock.restore();
+      randomStub.restore();
+      jest.restoreAllMocks();
+    });
+
+    const make504 = () => ({ response: { status: 504 } });
+
+    it("retries once on 504 then returns on success", async () => {
+      const serverStub = jest
+        .spyOn(stellar.server, "submitTransaction")
+        .mockRejectedValueOnce(make504())
+        .mockResolvedValueOnce({ successful: true } as any);
+
+      const tx = {} as Transaction;
+      const promise = stellar.submitTransaction(tx);
+
+      // Equal-jitter sleep at attempt 0 with Math.random stubbed to 0 is
+      // cappedDelay/2 = 1000/2 = 500ms.
+      await clock.tickAsync(500);
+
+      await expect(promise).resolves.toBe(true);
+      expect(serverStub).toHaveBeenCalledTimes(2);
+    });
+
+    it("rethrows immediately on a non-504 error", async () => {
+      const serverStub = jest
+        .spyOn(stellar.server, "submitTransaction")
+        .mockRejectedValueOnce({ response: { status: 500 } });
+
+      await expect(
+        stellar.submitTransaction({} as Transaction),
+      ).rejects.toMatchObject({ response: { status: 500 } });
+      expect(serverStub).toHaveBeenCalledTimes(1);
+    });
+
+    it("gives up after the bounded number of retries and rethrows the last 504", async () => {
+      // SUBMIT_504_MAX_RETRIES = 5 → 1 initial + 5 retries = 6 total attempts.
+      const serverStub = jest
+        .spyOn(stellar.server, "submitTransaction")
+        .mockRejectedValue(make504());
+
+      const promise = stellar.submitTransaction({} as Transaction);
+      // Surface the rejection eagerly so an unhandled rejection doesn't crash
+      // the test runner mid-tick.
+      promise.catch(() => undefined);
+
+      // Sum of equal-jitter sleeps (cappedDelay/2) with Math.random stubbed
+      // to 0: 500 + 1000 + 2000 + 4000 + 8000 = 15500ms.
+      await clock.tickAsync(15500);
+
+      await expect(promise).rejects.toMatchObject({
+        response: { status: 504 },
+      });
+      expect(serverStub).toHaveBeenCalledTimes(6);
+    });
+  });
+
   describe("TransactionBuilder/transferWithdrawalTransaction", () => {
     it("should transfer withdrawal transaction", async () => {
       const memoExamples = [