Skip to content

Commit e7a8ad3

Browse files
amanthanviamanthanvi
authored
fix: preserve sidebar usage values on partial rate-limit updates (#496)
1 parent b131fc3 commit e7a8ad3

9 files changed

Lines changed: 434 additions & 81 deletions

‎src/features/app/hooks/useAppServerEvents.test.tsx‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ describe("useAppServerEvents", () => {
6060
onRequestUserInput: vi.fn(),
6161
onItemCompleted: vi.fn(),
6262
onAgentMessageCompleted: vi.fn(),
63+
onAccountRateLimitsUpdated: vi.fn(),
6364
onAccountUpdated: vi.fn(),
6465
onAccountLoginCompleted: vi.fn(),
6566
};
@@ -287,6 +288,36 @@ describe("useAppServerEvents", () => {
287288
text: "Done",
288289
});
289290

291+
act(() => {
292+
listener?.({
293+
workspace_id: "ws-1",
294+
message: {
295+
method: "account/rateLimits/updated",
296+
params: {
297+
rateLimits: { primary: { usedPercent: 25 } },
298+
},
299+
},
300+
});
301+
});
302+
expect(handlers.onAccountRateLimitsUpdated).toHaveBeenCalledWith("ws-1", {
303+
primary: { usedPercent: 25 },
304+
});
305+
306+
act(() => {
307+
listener?.({
308+
workspace_id: "ws-1",
309+
message: {
310+
method: "account/rateLimits/updated",
311+
params: {
312+
rate_limits: { primary: { used_percent: 30 } },
313+
},
314+
},
315+
});
316+
});
317+
expect(handlers.onAccountRateLimitsUpdated).toHaveBeenCalledWith("ws-1", {
318+
primary: { used_percent: 30 },
319+
});
320+
290321
act(() => {
291322
listener?.({
292323
workspace_id: "ws-1",

‎src/features/threads/hooks/useThreadEventHandlers.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { useCallback, useMemo } from "react";
22
import type { Dispatch, MutableRefObject } from "react";
3-
import type { AppServerEvent, DebugEntry, TurnPlan } from "@/types";
3+
import type { AppServerEvent, DebugEntry, RateLimitSnapshot, TurnPlan } from "@/types";
44
import { getAppServerRawMethod } from "@utils/appServerEvents";
55
import { useThreadApprovalEvents } from "./useThreadApprovalEvents";
66
import { useThreadItemEvents } from "./useThreadItemEvents";
@@ -12,6 +12,7 @@ type ThreadEventHandlersOptions = {
1212
activeThreadId: string | null;
1313
dispatch: Dispatch<ThreadAction>;
1414
planByThreadRef: MutableRefObject<Record<string, TurnPlan | null>>;
15+
getCurrentRateLimits?: (workspaceId: string) => RateLimitSnapshot | null;
1516
getCustomName: (workspaceId: string, threadId: string) => string | undefined;
1617
isThreadHidden: (workspaceId: string, threadId: string) => boolean;
1718
markProcessing: (threadId: string, isProcessing: boolean) => void;
@@ -46,6 +47,7 @@ export function useThreadEventHandlers({
4647
activeThreadId,
4748
dispatch,
4849
planByThreadRef,
50+
getCurrentRateLimits,
4951
getCustomName,
5052
isThreadHidden,
5153
markProcessing,
@@ -110,6 +112,7 @@ export function useThreadEventHandlers({
110112
} = useThreadTurnEvents({
111113
dispatch,
112114
planByThreadRef,
115+
getCurrentRateLimits,
113116
getCustomName,
114117
isThreadHidden,
115118
markProcessing,

‎src/features/threads/hooks/useThreadRateLimits.test.tsx‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,47 @@ describe("useThreadRateLimits", () => {
9696
});
9797
});
9898

99+
it("does not auto-refresh again when accessor callback identity changes", async () => {
100+
const dispatch = vi.fn();
101+
102+
vi.mocked(getAccountRateLimits).mockResolvedValue({
103+
result: { rate_limits: {} },
104+
});
105+
106+
const { rerender } = renderHook(
107+
({
108+
getCurrentRateLimits,
109+
}: {
110+
getCurrentRateLimits: (workspaceId: string) => null;
111+
}) =>
112+
useThreadRateLimits({
113+
activeWorkspaceId: "ws-1",
114+
activeWorkspaceConnected: true,
115+
dispatch,
116+
getCurrentRateLimits,
117+
}),
118+
{
119+
initialProps: {
120+
getCurrentRateLimits: () => null,
121+
},
122+
},
123+
);
124+
125+
await waitFor(() => {
126+
expect(getAccountRateLimits).toHaveBeenCalledTimes(1);
127+
});
128+
129+
rerender({
130+
getCurrentRateLimits: () => null,
131+
});
132+
133+
await act(async () => {
134+
await Promise.resolve();
135+
});
136+
137+
expect(getAccountRateLimits).toHaveBeenCalledTimes(1);
138+
});
139+
99140
it("reports errors via debug callback without dispatching", async () => {
100141
const dispatch = vi.fn();
101142
const onDebug = vi.fn();
@@ -123,4 +164,70 @@ describe("useThreadRateLimits", () => {
123164
}),
124165
);
125166
});
167+
168+
it("merges partial payloads with previous workspace rate limits", async () => {
169+
const dispatch = vi.fn();
170+
const previousRateLimits = {
171+
primary: {
172+
usedPercent: 42,
173+
windowDurationMins: 60,
174+
resetsAt: 12345,
175+
},
176+
secondary: {
177+
usedPercent: 70,
178+
windowDurationMins: 10080,
179+
resetsAt: 99999,
180+
},
181+
credits: {
182+
hasCredits: true,
183+
unlimited: false,
184+
balance: "5",
185+
},
186+
planType: "pro",
187+
} as const;
188+
189+
vi.mocked(getAccountRateLimits).mockResolvedValue({
190+
result: {
191+
rate_limits: {
192+
primary: { resets_at: 88888 },
193+
secondary: {},
194+
},
195+
},
196+
});
197+
198+
const { result } = renderHook(() =>
199+
useThreadRateLimits({
200+
activeWorkspaceId: "ws-1",
201+
dispatch,
202+
getCurrentRateLimits: () => previousRateLimits,
203+
}),
204+
);
205+
206+
await act(async () => {
207+
await result.current.refreshAccountRateLimits();
208+
});
209+
210+
expect(dispatch).toHaveBeenCalledWith({
211+
type: "setRateLimits",
212+
workspaceId: "ws-1",
213+
rateLimits: {
214+
primary: {
215+
usedPercent: 42,
216+
windowDurationMins: 60,
217+
resetsAt: 88888,
218+
},
219+
secondary: {
220+
usedPercent: 70,
221+
windowDurationMins: 10080,
222+
resetsAt: 99999,
223+
},
224+
credits: {
225+
hasCredits: true,
226+
unlimited: false,
227+
balance: "5",
228+
},
229+
planType: "pro",
230+
},
231+
});
232+
});
126233
});

‎src/features/threads/hooks/useThreadRateLimits.ts‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,29 @@
1-
import { useCallback, useEffect } from "react";
2-
import type { DebugEntry } from "@/types";
1+
import { useCallback, useEffect, useRef } from "react";
2+
import type { DebugEntry, RateLimitSnapshot } from "@/types";
33
import { getAccountRateLimits } from "@services/tauri";
44
import { normalizeRateLimits } from "@threads/utils/threadNormalize";
55
import type { ThreadAction } from "./useThreadsReducer";
66

77
type UseThreadRateLimitsOptions = {
88
activeWorkspaceId: string | null;
99
activeWorkspaceConnected?: boolean;
10+
getCurrentRateLimits?: (workspaceId: string) => RateLimitSnapshot | null;
1011
dispatch: React.Dispatch<ThreadAction>;
1112
onDebug?: (entry: DebugEntry) => void;
1213
};
1314

1415
export function useThreadRateLimits({
1516
activeWorkspaceId,
1617
activeWorkspaceConnected,
18+
getCurrentRateLimits,
1719
dispatch,
1820
onDebug,
1921
}: UseThreadRateLimitsOptions) {
22+
const getCurrentRateLimitsRef = useRef(getCurrentRateLimits);
23+
useEffect(() => {
24+
getCurrentRateLimitsRef.current = getCurrentRateLimits;
25+
}, [getCurrentRateLimits]);
26+
2027
const refreshAccountRateLimits = useCallback(
2128
async (workspaceId?: string) => {
2229
const targetId = workspaceId ?? activeWorkspaceId;
@@ -45,10 +52,12 @@ export function useThreadRateLimits({
4552
(response?.rateLimits as Record<string, unknown> | undefined) ??
4653
(response?.rate_limits as Record<string, unknown> | undefined);
4754
if (rateLimits) {
55+
const previousRateLimits =
56+
getCurrentRateLimitsRef.current?.(targetId) ?? null;
4857
dispatch({
4958
type: "setRateLimits",
5059
workspaceId: targetId,
51-
rateLimits: normalizeRateLimits(rateLimits),
60+
rateLimits: normalizeRateLimits(rateLimits, previousRateLimits),
5261
});
5362
}
5463
} catch (error) {

‎src/features/threads/hooks/useThreadTurnEvents.test.tsx‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// @vitest-environment jsdom
22
import { act, renderHook } from "@testing-library/react";
33
import { beforeEach, describe, expect, it, vi } from "vitest";
4-
import type { TurnPlan } from "@/types";
4+
import type { RateLimitSnapshot, TurnPlan } from "@/types";
55
import { interruptTurn } from "@services/tauri";
66
import {
77
normalizePlanUpdate,
@@ -26,6 +26,7 @@ type SetupOverrides = {
2626
pendingInterrupts?: string[];
2727
planByThread?: Record<string, TurnPlan | null>;
2828
activeTurnByThread?: Record<string, string | null>;
29+
rateLimitsByWorkspace?: Record<string, RateLimitSnapshot | null>;
2930
};
3031

3132
const makeOptions = (overrides: SetupOverrides = {}) => {
@@ -38,6 +39,9 @@ const makeOptions = (overrides: SetupOverrides = {}) => {
3839
const getActiveTurnId = vi.fn(
3940
(threadId: string) => overrides.activeTurnByThread?.[threadId] ?? null,
4041
);
42+
const getCurrentRateLimits = vi.fn(
43+
(workspaceId: string) => overrides.rateLimitsByWorkspace?.[workspaceId] ?? null,
44+
);
4145
const pushThreadErrorMessage = vi.fn();
4246
const safeMessageActivity = vi.fn();
4347
const recordThreadActivity = vi.fn();
@@ -52,6 +56,7 @@ const makeOptions = (overrides: SetupOverrides = {}) => {
5256
useThreadTurnEvents({
5357
dispatch,
5458
planByThreadRef,
59+
getCurrentRateLimits,
5560
getCustomName,
5661
isThreadHidden,
5762
markProcessing,
@@ -74,6 +79,7 @@ const makeOptions = (overrides: SetupOverrides = {}) => {
7479
markReviewing,
7580
setActiveTurnId,
7681
getActiveTurnId,
82+
getCurrentRateLimits,
7783
pushThreadErrorMessage,
7884
safeMessageActivity,
7985
recordThreadActivity,
@@ -557,7 +563,20 @@ describe("useThreadTurnEvents", () => {
557563
});
558564

559565
it("dispatches normalized rate limits updates", () => {
560-
const { result, dispatch } = makeOptions();
566+
const previousRateLimits = {
567+
primary: {
568+
usedPercent: 35,
569+
windowDurationMins: 60,
570+
resetsAt: 1_700_000_000,
571+
},
572+
secondary: null,
573+
credits: null,
574+
planType: null,
575+
} satisfies RateLimitSnapshot;
576+
577+
const { result, dispatch, getCurrentRateLimits } = makeOptions({
578+
rateLimitsByWorkspace: { "ws-1": previousRateLimits },
579+
});
561580
const normalized = { primary: { usedPercent: 10 } };
562581

563582
vi.mocked(normalizeRateLimits).mockReturnValue(normalized as never);
@@ -566,6 +585,11 @@ describe("useThreadTurnEvents", () => {
566585
result.current.onAccountRateLimitsUpdated("ws-1", { primary: {} });
567586
});
568587

588+
expect(getCurrentRateLimits).toHaveBeenCalledWith("ws-1");
589+
expect(normalizeRateLimits).toHaveBeenCalledWith(
590+
{ primary: {} },
591+
previousRateLimits,
592+
);
569593
expect(dispatch).toHaveBeenCalledWith({
570594
type: "setRateLimits",
571595
workspaceId: "ws-1",

‎src/features/threads/hooks/useThreadTurnEvents.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { useCallback, useRef } from "react";
22
import type { Dispatch, MutableRefObject } from "react";
3-
import type { TurnPlan } from "@/types";
3+
import type { RateLimitSnapshot, TurnPlan } from "@/types";
44
import { interruptTurn as interruptTurnService } from "@services/tauri";
55
import { getThreadTimestamp } from "@utils/threadItems";
66
import {
@@ -14,6 +14,7 @@ import type { ThreadAction } from "./useThreadsReducer";
1414
type UseThreadTurnEventsOptions = {
1515
dispatch: Dispatch<ThreadAction>;
1616
planByThreadRef: MutableRefObject<Record<string, TurnPlan | null>>;
17+
getCurrentRateLimits?: (workspaceId: string) => RateLimitSnapshot | null;
1718
getCustomName: (workspaceId: string, threadId: string) => string | undefined;
1819
isThreadHidden: (workspaceId: string, threadId: string) => boolean;
1920
markProcessing: (threadId: string, isProcessing: boolean) => void;
@@ -40,6 +41,7 @@ function normalizeThreadStatusType(status: Record<string, unknown>): string {
4041
export function useThreadTurnEvents({
4142
dispatch,
4243
planByThreadRef,
44+
getCurrentRateLimits,
4345
getCustomName,
4446
isThreadHidden,
4547
markProcessing,
@@ -318,13 +320,14 @@ export function useThreadTurnEvents({
318320

319321
const onAccountRateLimitsUpdated = useCallback(
320322
(workspaceId: string, rateLimits: Record<string, unknown>) => {
323+
const previousRateLimits = getCurrentRateLimits?.(workspaceId) ?? null;
321324
dispatch({
322325
type: "setRateLimits",
323326
workspaceId,
324-
rateLimits: normalizeRateLimits(rateLimits),
327+
rateLimits: normalizeRateLimits(rateLimits, previousRateLimits),
325328
});
326329
},
327-
[dispatch],
330+
[dispatch, getCurrentRateLimits],
328331
);
329332

330333
const onTurnError = useCallback(

‎src/features/threads/hooks/useThreads.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ export function useThreads({
119119
threadsByWorkspaceRef.current = state.threadsByWorkspace;
120120
activeTurnIdByThreadRef.current = state.activeTurnIdByThread;
121121
threadParentByIdRef.current = state.threadParentById;
122+
const rateLimitsByWorkspaceRef = useRef(state.rateLimitsByWorkspace);
123+
rateLimitsByWorkspaceRef.current = state.rateLimitsByWorkspace;
122124
const { approvalAllowlistRef, handleApprovalDecision, handleApprovalRemember } =
123125
useThreadApprovals({ dispatch, onDebug });
124126
const { handleUserInputSubmit } = useThreadUserInput({ dispatch });
@@ -141,9 +143,15 @@ export function useThreads({
141143
itemsByThread: state.itemsByThread,
142144
});
143145

146+
const getCurrentRateLimits = useCallback(
147+
(workspaceId: string) => rateLimitsByWorkspaceRef.current[workspaceId] ?? null,
148+
[],
149+
);
150+
144151
const { refreshAccountRateLimits } = useThreadRateLimits({
145152
activeWorkspaceId,
146153
activeWorkspaceConnected: activeWorkspace?.connected,
154+
getCurrentRateLimits,
147155
dispatch,
148156
onDebug,
149157
});
@@ -380,6 +388,7 @@ export function useThreads({
380388
activeThreadId,
381389
dispatch,
382390
planByThreadRef,
391+
getCurrentRateLimits,
383392
getCustomName,
384393
isThreadHidden,
385394
markProcessing,

0 commit comments

Comments
 (0)