From d06e329a781fa9f2aeca715cbb7f6bff410e1509 Mon Sep 17 00:00:00 2001
From: murashit <murashit@proton.me>
Date: Thu, 4 Jun 2026 12:10:27 +0900
Subject: [PATCH] Preserve published rate limits on sparse refresh

---
 .../chat/chat-app-server-controller.ts        | 12 ++++
 .../chat/chat-view-controller-assembly.ts     |  2 +-
 .../chat/chat-app-server-controller.test.ts   | 67 +++++++++++++++++++
 3 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/src/features/chat/chat-app-server-controller.ts b/src/features/chat/chat-app-server-controller.ts
index df90b795..1a691036 100644
--- a/src/features/chat/chat-app-server-controller.ts
+++ b/src/features/chat/chat-app-server-controller.ts
@@ -171,6 +171,18 @@ export class ChatAppServerController {
     this.dispatch({ type: "thread/list-applied", rateLimit: rateLimit.data, appServerDiagnostics: diagnostics });
   }
 
+  async refreshPublishedRateLimits(): Promise<void> {
+    const rateLimit = await this.loadRateLimit();
+    const diagnostics = cloneAppServerDiagnostics(this.state.appServerDiagnostics);
+    diagnostics.probes["account/rateLimits/read"] = rateLimit.probe;
+    if (rateLimit.probe.status === "ok") {
+      this.dispatch({ type: "thread/list-applied", rateLimit: rateLimit.data, appServerDiagnostics: diagnostics });
+      this.publishAppServerMetadataSnapshot();
+      return;
+    }
+    this.dispatch({ type: "thread/list-applied", appServerDiagnostics: diagnostics });
+  }
+
   async loadRateLimit(): Promise<{ data: RateLimitSnapshot | null; probe: AppServerDiagnostics["probes"]["account/rateLimits/read"] }> {
     const client = this.host.currentClient();
     if (!client) {
diff --git a/src/features/chat/chat-view-controller-assembly.ts b/src/features/chat/chat-view-controller-assembly.ts
index 7c6c460b..3a91b943 100644
--- a/src/features/chat/chat-view-controller-assembly.ts
+++ b/src/features/chat/chat-view-controller-assembly.ts
@@ -355,7 +355,7 @@ export function createChatViewControllerAssembly(host: ChatViewControllerAssembl
       void host.refreshThreads();
     },
     refreshRateLimits: () => {
-      void appServer.refreshRateLimits();
+      void appServer.refreshPublishedRateLimits();
     },
     refreshSkills: (forceReload) => void host.refreshSkills(forceReload),
     publishAppServerMetadata: host.publishAppServerMetadataSnapshot,
diff --git a/tests/features/chat/chat-app-server-controller.test.ts b/tests/features/chat/chat-app-server-controller.test.ts
index afd95486..902012c5 100644
--- a/tests/features/chat/chat-app-server-controller.test.ts
+++ b/tests/features/chat/chat-app-server-controller.test.ts
@@ -133,6 +133,59 @@ describe("ChatAppServerController", () => {
     });
   });
 
+  it("publishes refreshed rate limits from sparse update notifications", async () => {
+    const state = createChatState();
+    const stateStore = createChatStateStore(state);
+    const rateLimit = rateLimitFixture({ primary: { usedPercent: 64, windowDurationMins: 300, resetsAt: null } });
+    const publishAppServerMetadata = vi.fn();
+    const client = {
+      readAccountRateLimits: vi.fn().mockResolvedValue({ rateLimits: rateLimit, rateLimitsByLimitId: null }),
+    } as unknown as AppServerClient;
+    const controller = new ChatAppServerController({
+      stateStore,
+      vaultPath: "/vault",
+      currentClient: () => client,
+      runtimeSnapshot: () => ({}) as never,
+      forceMessagesToBottom: () => undefined,
+      publishThreadList: () => undefined,
+      publishAppServerMetadata,
+    });
+
+    await controller.refreshPublishedRateLimits();
+
+    expect(stateStore.getState().rateLimit).toMatchObject({ primary: { usedPercent: 64 } });
+    expect(publishAppServerMetadata).toHaveBeenCalledWith(expect.objectContaining({ rateLimit }));
+  });
+
+  it("keeps the previous rate limit snapshot when sparse update refresh fails", async () => {
+    const state = createChatState();
+    const previousRateLimit = rateLimitFixture({
+      limitName: "Codex",
+      primary: { usedPercent: 42, windowDurationMins: 300, resetsAt: null },
+    });
+    state.rateLimit = previousRateLimit;
+    const stateStore = createChatStateStore(state);
+    const publishAppServerMetadata = vi.fn();
+    const client = {
+      readAccountRateLimits: vi.fn().mockRejectedValue(new Error("offline")),
+    } as unknown as AppServerClient;
+    const controller = new ChatAppServerController({
+      stateStore,
+      vaultPath: "/vault",
+      currentClient: () => client,
+      runtimeSnapshot: () => ({}) as never,
+      forceMessagesToBottom: () => undefined,
+      publishThreadList: () => undefined,
+      publishAppServerMetadata,
+    });
+
+    await controller.refreshPublishedRateLimits();
+
+    expect(stateStore.getState().rateLimit).toBe(previousRateLimit);
+    expect(stateStore.getState().appServerDiagnostics.probes["account/rateLimits/read"]).toMatchObject({ status: "failed" });
+    expect(publishAppServerMetadata).not.toHaveBeenCalled();
+  });
+
   it("loads MCP status lines with cached startup diagnostics", async () => {
     const state = createChatState();
     state.activeThreadId = "thread-1";
@@ -188,6 +241,20 @@ function threadFixture(id: string, overrides: Partial<Thread> = {}): Thread {
   };
 }
 
+function rateLimitFixture(overrides: Partial<RateLimitSnapshot> = {}): RateLimitSnapshot {
+  return {
+    limitId: "codex",
+    limitName: "Codex",
+    primary: null,
+    secondary: null,
+    credits: null,
+    individualLimit: null,
+    planType: null,
+    rateLimitReachedType: null,
+    ...overrides,
+  };
+}
+
 function mcpServerStatus(): McpServerStatus {
   return {
     name: "github",