From 8a2793977f89f31477a5f8e2c790a720b74cd3a3 Mon Sep 17 00:00:00 2001
From: stayzappy <adezapelectronics@gmail.com>
Date: Tue, 28 Jul 2026 11:16:19 +0000
Subject: [PATCH] feat: implement caching for account and balance endpoints

---
 src/client/createSorokitClient.ts | 39 ++++++++++++++++++++++++-------
 src/tests/client.test.ts          | 39 +++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/src/client/createSorokitClient.ts b/src/client/createSorokitClient.ts
index bec320a..823bc54 100644
--- a/src/client/createSorokitClient.ts
+++ b/src/client/createSorokitClient.ts
@@ -656,9 +656,16 @@ export function createSorokitClient(
           errorHandler,
           { functionName: "account.get", params: { publicKey } },
           () =>
-            withLogging(logger, "account.get", { publicKey }, () =>
-              getAccount(horizonUrl, publicKey),
-            ),
+            withLogging(logger, "account.get", { publicKey }, async () => {
+              const cacheKey = `account:get:${horizonUrl}:${publicKey}`;
+              if (cache) {
+                const cachedVal = cache.get(cacheKey);
+                if (cachedVal) return ok(cachedVal as AccountInfo);
+              }
+              const res = await getAccount(horizonUrl, publicKey);
+              if (cache && res.status === "ok") cache.set(cacheKey, res.data);
+              return res;
+            }),
         ).then(applyTx),
       getAccountsBatch: (publicKeys) =>
         withErrorHandling(
@@ -674,18 +681,32 @@ export function createSorokitClient(
           errorHandler,
           { functionName: "account.getBalances", params: { publicKey } },
           () =>
-            withLogging(logger, "account.getBalances", { publicKey }, () =>
-              getBalances(horizonUrl, publicKey),
-            ),
+            withLogging(logger, "account.getBalances", { publicKey }, async () => {
+              const cacheKey = `account:balances:${horizonUrl}:${publicKey}`;
+              if (cache) {
+                const cachedVal = cache.get(cacheKey);
+                if (cachedVal) return ok(cachedVal as AssetBalance[]);
+              }
+              const res = await getBalances(horizonUrl, publicKey);
+              if (cache && res.status === "ok") cache.set(cacheKey, res.data);
+              return res;
+            }),
         ).then(applyTx),
       getAssetBalances: (publicKey, filter) =>
         withErrorHandling(
           errorHandler,
           { functionName: "account.getAssetBalances", params: { publicKey, filter } },
           () =>
-            withLogging(logger, "account.getAssetBalances", { publicKey, filter }, () =>
-              getAssetBalances(horizonUrl, publicKey, filter),
-            ),
+            withLogging(logger, "account.getAssetBalances", { publicKey, filter }, async () => {
+              const cacheKey = `account:assetBalances:${horizonUrl}:${publicKey}:${JSON.stringify(filter ?? {})}`;
+              if (cache) {
+                const cachedVal = cache.get(cacheKey);
+                if (cachedVal) return ok(cachedVal as AssetBalance[]);
+              }
+              const res = await getAssetBalances(horizonUrl, publicKey, filter);
+              if (cache && res.status === "ok") cache.set(cacheKey, res.data);
+              return res;
+            }),
         ).then(applyTx),
       stream: (publicKey, streamConfig, signal) =>
         streamAccount(horizonUrl, publicKey, streamConfig, signal, logger),
diff --git a/src/tests/client.test.ts b/src/tests/client.test.ts
index 84d1766..6993e5d 100644
--- a/src/tests/client.test.ts
+++ b/src/tests/client.test.ts
@@ -30,6 +30,19 @@ vi.mock("../transaction/streamTransactions", async (importOriginal) => {
   };
 });
 
+const { mockGetAccount } = vi.hoisted(() => ({
+  mockGetAccount: vi.fn(),
+}));
+
+vi.mock("../account/getAccount", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../account/getAccount")>();
+  mockGetAccount.mockImplementation(actual.getAccount);
+  return {
+    ...actual,
+    getAccount: mockGetAccount,
+  };
+});
+
 describe("createSorokitClient", () => {
   it("creates a client for testnet", () => {
     const result = createSorokitClient({ network: "testnet" });
@@ -461,4 +474,30 @@ describe("createSorokitClient", () => {
       }
     }
   });
+
+  it("returns cached result and does not call Horizon if cache is warm", async () => {
+    mockGetAccount.mockClear();
+    const mockCache = {
+      get: vi.fn(),
+      set: vi.fn(),
+      invalidate: vi.fn(),
+      delete: vi.fn(),
+    };
+    const cachedAccount = { publicKey: "G...", sequence: "1", balances: [] };
+    mockCache.get.mockReturnValue(cachedAccount);
+
+    const result = createSorokitClient({ network: "testnet", cache: mockCache });
+    expect(result.status).toBe("ok");
+    if (result.status === "ok") {
+      const client = result.data;
+      const res = await client.account.get("G...");
+      expect(res.status).toBe("ok");
+      expect(mockCache.get).toHaveBeenCalledWith(expect.stringContaining("account:get:"));
+      expect(mockGetAccount).not.toHaveBeenCalled();
+      if (res.status === "ok") {
+        expect(res.data).toEqual(cachedAccount);
+      }
+    }
+    mockGetAccount.mockClear();
+  });
 });