diff --git a/src/gateway/server-methods/usage.sessions-usage.test.ts b/src/gateway/server-methods/usage.sessions-usage.test.ts index 7eb48fd821..3027abe1e4 100644 --- a/src/gateway/server-methods/usage.sessions-usage.test.ts +++ b/src/gateway/server-methods/usage.sessions-usage.test.ts @@ -82,6 +82,33 @@ import { import { loadCombinedSessionStoreForGateway } from "../session-utils.js"; import { usageHandlers } from "./usage.js"; +async function runSessionsUsage(params: Record) { + const respond = vi.fn(); + await usageHandlers["sessions.usage"]({ + respond, + params, + } as unknown as Parameters<(typeof usageHandlers)["sessions.usage"]>[0]); + return respond; +} + +async function runSessionsUsageTimeseries(params: Record) { + const respond = vi.fn(); + await usageHandlers["sessions.usage.timeseries"]({ + respond, + params, + } as unknown as Parameters<(typeof usageHandlers)["sessions.usage.timeseries"]>[0]); + return respond; +} + +async function runSessionsUsageLogs(params: Record) { + const respond = vi.fn(); + await usageHandlers["sessions.usage.logs"]({ + respond, + params, + } as unknown as Parameters<(typeof usageHandlers)["sessions.usage.logs"]>[0]); + return respond; +} + describe("sessions.usage", () => { beforeEach(() => { vi.useRealTimers(); @@ -89,16 +116,11 @@ describe("sessions.usage", () => { }); it("discovers sessions across configured agents and keeps agentId in key", async () => { - const respond = vi.fn(); - - await usageHandlers["sessions.usage"]({ - respond, - params: { - startDate: "2026-02-01", - endDate: "2026-02-02", - limit: 10, - }, - } as unknown as Parameters<(typeof usageHandlers)["sessions.usage"]>[0]); + const respond = await runSessionsUsage({ + startDate: "2026-02-01", + endDate: "2026-02-02", + limit: 10, + }); expect(vi.mocked(discoverAllSessions)).toHaveBeenCalledTimes(2); expect(vi.mocked(discoverAllSessions).mock.calls[0]?.[0]?.agentId).toBe("main"); @@ -129,7 +151,6 @@ describe("sessions.usage", () => { fs.mkdirSync(agentSessionsDir, { recursive: true }); const sessionFile = path.join(agentSessionsDir, "s-opus.jsonl"); fs.writeFileSync(sessionFile, "", "utf-8"); - const respond = vi.fn(); // Swap the store mock for this test: the canonical key differs from the discovered key // but points at the same sessionId. @@ -146,15 +167,12 @@ describe("sessions.usage", () => { }); // Query via discovered key: agent:: - await usageHandlers["sessions.usage"]({ - respond, - params: { - startDate: "2026-02-01", - endDate: "2026-02-02", - key: "agent:opus:s-opus", - limit: 10, - }, - } as unknown as Parameters<(typeof usageHandlers)["sessions.usage"]>[0]); + const respond = await runSessionsUsage({ + startDate: "2026-02-01", + endDate: "2026-02-02", + key: "agent:opus:s-opus", + limit: 10, + }); expect(respond).toHaveBeenCalledTimes(1); expect(respond.mock.calls[0]?.[0]).toBe(true); @@ -172,17 +190,12 @@ describe("sessions.usage", () => { }); it("rejects traversal-style keys in specific session usage lookups", async () => { - const respond = vi.fn(); - - await usageHandlers["sessions.usage"]({ - respond, - params: { - startDate: "2026-02-01", - endDate: "2026-02-02", - key: "agent:opus:../../etc/passwd", - limit: 10, - }, - } as unknown as Parameters<(typeof usageHandlers)["sessions.usage"]>[0]); + const respond = await runSessionsUsage({ + startDate: "2026-02-01", + endDate: "2026-02-02", + key: "agent:opus:../../etc/passwd", + limit: 10, + }); expect(respond).toHaveBeenCalledTimes(1); expect(respond.mock.calls[0]?.[0]).toBe(false); @@ -191,30 +204,44 @@ describe("sessions.usage", () => { }); it("passes parsed agentId into sessions.usage.timeseries", async () => { - const respond = vi.fn(); - - await usageHandlers["sessions.usage.timeseries"]({ - respond, - params: { - key: "agent:opus:s-opus", - }, - } as unknown as Parameters<(typeof usageHandlers)["sessions.usage.timeseries"]>[0]); + await runSessionsUsageTimeseries({ + key: "agent:opus:s-opus", + }); expect(vi.mocked(loadSessionUsageTimeSeries)).toHaveBeenCalled(); expect(vi.mocked(loadSessionUsageTimeSeries).mock.calls[0]?.[0]?.agentId).toBe("opus"); }); it("passes parsed agentId into sessions.usage.logs", async () => { - const respond = vi.fn(); - - await usageHandlers["sessions.usage.logs"]({ - respond, - params: { - key: "agent:opus:s-opus", - }, - } as unknown as Parameters<(typeof usageHandlers)["sessions.usage.logs"]>[0]); + await runSessionsUsageLogs({ + key: "agent:opus:s-opus", + }); expect(vi.mocked(loadSessionLogs)).toHaveBeenCalled(); expect(vi.mocked(loadSessionLogs).mock.calls[0]?.[0]?.agentId).toBe("opus"); }); + + it("rejects traversal-style keys in timeseries/log lookups", async () => { + const timeseriesRespond = await runSessionsUsageTimeseries({ + key: "agent:opus:../../etc/passwd", + }); + expect(timeseriesRespond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + message: expect.stringContaining("Invalid session key"), + }), + ); + + const logsRespond = await runSessionsUsageLogs({ + key: "agent:opus:../../etc/passwd", + }); + expect(logsRespond).toHaveBeenCalledWith( + false, + undefined, + expect.objectContaining({ + message: expect.stringContaining("Invalid session key"), + }), + ); + }); });