mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[F083] clear stale usage snapshot when /ws/system leaves connected
The hook synced wsState into the store but never dropped usageData when the stream dropped, so on reconnect wsState flipped to "connected" before any fresh USAGE_SNAPSHOT arrived and UsageOverviewPanel rendered the prior session's totals/cost as if they were live. Clear usageData whenever state leaves "connected" so the panel falls back to the polling summary until a new snapshot lands. Connected->connected is a no-op clear skip.
This commit is contained in:
@@ -0,0 +1,99 @@
|
|||||||
|
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||||
|
import { render } from "@testing-library/react";
|
||||||
|
import type { ConnectionState } from "@/lib/websocket/connection";
|
||||||
|
|
||||||
|
// `useWebSocket` is the seam: drive its `state` + `lastMessage` from the test
|
||||||
|
// by mutating a shared ref and re-rendering. The hook under test only reads
|
||||||
|
// `state` + `lastMessage` off the return, so a plain mutable object suffices.
|
||||||
|
const wsRef = {
|
||||||
|
state: "disconnected" as ConnectionState,
|
||||||
|
lastMessage: null as {
|
||||||
|
type: string;
|
||||||
|
totals?: { input_tokens?: number; output_tokens?: number };
|
||||||
|
cost_estimate?: number;
|
||||||
|
period?: string;
|
||||||
|
timestamp?: string;
|
||||||
|
} | null,
|
||||||
|
};
|
||||||
|
|
||||||
|
vi.mock("@/hooks/use-websocket", () => ({
|
||||||
|
useWebSocket: () => wsRef,
|
||||||
|
}));
|
||||||
|
|
||||||
|
import { useRateLimitWebSocket } from "../use-rate-limit-websocket";
|
||||||
|
import { useUsageStore } from "@/store/usage-store";
|
||||||
|
|
||||||
|
// Harness: a component that mounts the hook under test. Re-rendering it (via
|
||||||
|
// the `rerender` returned by `render()`) after mutating `wsRef` makes the
|
||||||
|
// hook's `state`/`lastMessage` effects re-run with the new values.
|
||||||
|
function Harness() {
|
||||||
|
useRateLimitWebSocket();
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
function snapshot(n: number) {
|
||||||
|
return {
|
||||||
|
type: "USAGE_SNAPSHOT",
|
||||||
|
totals: { input_tokens: n, output_tokens: n * 2 },
|
||||||
|
cost_estimate: n / 100,
|
||||||
|
period: "live",
|
||||||
|
timestamp: new Date(n).toISOString(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
describe("useRateLimitWebSocket — stale-snapshot-on-reconnect (F083)", () => {
|
||||||
|
beforeEach(() => {
|
||||||
|
wsRef.state = "disconnected";
|
||||||
|
wsRef.lastMessage = null;
|
||||||
|
useUsageStore.setState({ usageData: null, wsState: "disconnected" });
|
||||||
|
});
|
||||||
|
afterEach(() => {
|
||||||
|
vi.clearAllMocks();
|
||||||
|
});
|
||||||
|
|
||||||
|
it("clears the stale snapshot on disconnect so a reconnect can't show it as live", () => {
|
||||||
|
// 1. Connect + receive a snapshot → store holds live data.
|
||||||
|
wsRef.state = "connected";
|
||||||
|
wsRef.lastMessage = snapshot(100);
|
||||||
|
const { rerender } = render(<Harness />);
|
||||||
|
expect(useUsageStore.getState().usageData).not.toBeNull();
|
||||||
|
expect(useUsageStore.getState().usageData?.tokens_input).toBe(100);
|
||||||
|
|
||||||
|
// 2. Stream drops → the stale snapshot MUST be cleared. Before the fix the
|
||||||
|
// hook only synced wsState and left usageData untouched.
|
||||||
|
wsRef.state = "reconnecting";
|
||||||
|
wsRef.lastMessage = null;
|
||||||
|
rerender(<Harness />);
|
||||||
|
expect(useUsageStore.getState().usageData).toBeNull();
|
||||||
|
expect(useUsageStore.getState().wsState).toBe("reconnecting");
|
||||||
|
|
||||||
|
// 3. Reconnect (state flips to connected) BEFORE any new snapshot arrives.
|
||||||
|
// This is the regression: the panel would render the prior session's
|
||||||
|
// totals/cost as "live". usageData must still be null so the panel falls
|
||||||
|
// back to the polling summary until a fresh USAGE_SNAPSHOT lands.
|
||||||
|
wsRef.state = "connected";
|
||||||
|
rerender(<Harness />);
|
||||||
|
expect(useUsageStore.getState().wsState).toBe("connected");
|
||||||
|
expect(useUsageStore.getState().usageData).toBeNull();
|
||||||
|
|
||||||
|
// 4. A fresh snapshot arrives on the new connection → live data restored.
|
||||||
|
wsRef.lastMessage = snapshot(500);
|
||||||
|
rerender(<Harness />);
|
||||||
|
expect(useUsageStore.getState().usageData?.tokens_input).toBe(500);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("does not clear live data while the stream stays connected (no false drop)", () => {
|
||||||
|
// Regression guard: clearing is gated on state !== "connected", so a
|
||||||
|
// connected → connected re-render (e.g. an unrelated parent re-render) must
|
||||||
|
// NOT wipe a valid live snapshot.
|
||||||
|
wsRef.state = "connected";
|
||||||
|
wsRef.lastMessage = snapshot(250);
|
||||||
|
const { rerender } = render(<Harness />);
|
||||||
|
expect(useUsageStore.getState().usageData?.tokens_input).toBe(250);
|
||||||
|
|
||||||
|
rerender(<Harness />);
|
||||||
|
rerender(<Harness />);
|
||||||
|
expect(useUsageStore.getState().usageData?.tokens_input).toBe(250);
|
||||||
|
expect(useUsageStore.getState().wsState).toBe("connected");
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -64,7 +64,17 @@ export function useRateLimitWebSocket(
|
|||||||
// Sync WS connection state into useUsageStore for cross-component visibility.
|
// Sync WS connection state into useUsageStore for cross-component visibility.
|
||||||
// This is the ONLY place wsState is written; no second useWebSocket call is needed.
|
// This is the ONLY place wsState is written; no second useWebSocket call is needed.
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
useUsageStore.getState().setWsState(state);
|
const store = useUsageStore.getState();
|
||||||
|
store.setWsState(state);
|
||||||
|
// Drop any held snapshot the moment the stream is no longer connected.
|
||||||
|
// Otherwise a reconnect flips wsState back to "connected" before a fresh
|
||||||
|
// USAGE_SNAPSHOT arrives, and UsageOverviewPanel would render the prior
|
||||||
|
// session's totals/cost as if they were live. Clearing here keeps `live`
|
||||||
|
// null until a new snapshot lands, so the panel falls back to the polling
|
||||||
|
// summary during the gap. (Connected → connected is a no-op clear skip.)
|
||||||
|
if (state !== "connected") {
|
||||||
|
store.clearUsageData();
|
||||||
|
}
|
||||||
}, [state]);
|
}, [state]);
|
||||||
|
|
||||||
// Fire onReconnect when state transitions from reconnecting → connected
|
// Fire onReconnect when state transitions from reconnecting → connected
|
||||||
|
|||||||
Reference in New Issue
Block a user