diff --git a/src/gateway/server-methods/usage.test.ts b/src/gateway/server-methods/usage.test.ts index 54b6a2a69dc..e259bb3cca2 100644 --- a/src/gateway/server-methods/usage.test.ts +++ b/src/gateway/server-methods/usage.test.ts @@ -29,10 +29,22 @@ vi.mock("../../infra/session-cost-usage.js", async () => { missingCostEntries: 0, }, })), + discoverAllSessions: vi.fn(async () => []), }; }); -import { loadCostUsageSummaryFromCache } from "../../infra/session-cost-usage.js"; +vi.mock("../session-utils.js", async () => { + const actual = await vi.importActual("../session-utils.js"); + return { + ...actual, + loadCombinedSessionStoreForGateway: vi.fn(() => ({ storePath: "(multiple)", store: {} })), + }; +}); + +import { + discoverAllSessions, + loadCostUsageSummaryFromCache, +} from "../../infra/session-cost-usage.js"; import { testApi, usageHandlers } from "./usage.js"; describe("gateway usage helpers", () => { @@ -576,4 +588,77 @@ describe("gateway usage helpers", () => { undefined, ); }); + + it("rejects an all-agent usage load when one agent task fails", async () => { + const failure = new Error("agent usage load failed"); + vi.mocked(loadCostUsageSummaryFromCache) + .mockResolvedValueOnce(costSummary({ totalTokens: 1, totalCost: 0 })) + .mockRejectedValueOnce(failure); + + const respond = vi.fn(); + const request = usageHandlers["usage.cost"]({ + respond, + params: { startDate: "2026-02-01", endDate: "2026-02-02", agentScope: "all" }, + context: { + getRuntimeConfig: () => ({ + agents: { list: [{ id: "main" }, { id: "broken" }] }, + }), + }, + } as unknown as Parameters<(typeof usageHandlers)["usage.cost"]>[0]); + + await expect(request).rejects.toBe(failure); + expect(respond).not.toHaveBeenCalled(); + }); + + it("bounds sessions.usage all-agent session discovery", async () => { + const agentCount = 13; + const concurrencyLimit = 12; + let releaseLoads!: () => void; + const loadsReleased = new Promise((resolve) => { + releaseLoads = resolve; + }); + let resolveFirstBatchStarted!: () => void; + const firstBatchStarted = new Promise((resolve) => { + resolveFirstBatchStarted = resolve; + }); + let started = 0; + let inFlight = 0; + let peakInFlight = 0; + + vi.mocked(discoverAllSessions).mockImplementation(async () => { + started += 1; + inFlight += 1; + peakInFlight = Math.max(peakInFlight, inFlight); + if (started === concurrencyLimit) { + resolveFirstBatchStarted(); + } + await loadsReleased; + inFlight -= 1; + return []; + }); + + const respond = vi.fn(); + const request = usageHandlers["sessions.usage"]({ + respond, + params: { startDate: "2026-02-01", endDate: "2026-02-02", agentScope: "all" }, + context: { + getRuntimeConfig: () => ({ + agents: { + list: Array.from({ length: agentCount }, (_, i) => ({ id: `agent-${i}` })), + }, + }), + }, + } as unknown as Parameters<(typeof usageHandlers)["sessions.usage"]>[0]); + + await firstBatchStarted; + const startedBeforeRelease = started; + const peakBeforeRelease = peakInFlight; + releaseLoads(); + await request; + + expect(startedBeforeRelease).toBe(concurrencyLimit); + expect(peakBeforeRelease).toBe(concurrencyLimit); + expect(vi.mocked(discoverAllSessions)).toHaveBeenCalledTimes(agentCount); + expect(respond).toHaveBeenCalledWith(true, expect.any(Object), undefined); + }); }); diff --git a/src/gateway/server-methods/usage.ts b/src/gateway/server-methods/usage.ts index b7ccd2b83c7..03052c36573 100644 --- a/src/gateway/server-methods/usage.ts +++ b/src/gateway/server-methods/usage.ts @@ -66,6 +66,20 @@ const COST_USAGE_CACHE_TTL_MS = 30_000; const COST_USAGE_CACHE_MAX = 256; const USAGE_AGENT_LOAD_CONCURRENCY = 12; +async function runUsageAgentTasks(tasks: Array<() => Promise>): Promise { + const result = await runTasksWithConcurrency({ + tasks, + limit: USAGE_AGENT_LOAD_CONCURRENCY, + errorMode: "stop", + }); + // These fan-outs historically rejected as one unit. Never return partial + // per-agent usage; successful results retain their input order. + if (result.hasError) { + throw result.firstError; + } + return result.results; +} + type DateRange = { startMs: number; endMs: number }; // Keep validation and parsed timestamps in one result so handlers cannot forward // an invalid or backwards window to the usage loaders. @@ -477,8 +491,8 @@ async function discoverAllSessionsForUsage(params: { const agents = requestedAgentId ? [{ id: normalizeAgentId(requestedAgentId) }] : listAgentsForGateway(params.config).agents; - const discovered = await Promise.all( - agents.map(async (agent) => { + const discovered = await runUsageAgentTasks( + agents.map((agent) => async () => { const agentId = normalizeAgentId(agent.id); const sessions = await discoverAllSessions({ agentId, @@ -864,8 +878,8 @@ async function loadAllAgentCostUsageSummary(params: { const agentIds = listAgentsForGateway(params.config).agents.map((agent) => normalizeAgentId(agent.id), ); - const agentLoadResult = await runTasksWithConcurrency({ - tasks: agentIds.map( + const summaries = await runUsageAgentTasks( + agentIds.map( (agentId) => () => loadCostUsageSummaryFromCache({ startMs: params.startMs, @@ -877,13 +891,7 @@ async function loadAllAgentCostUsageSummary(params: { refreshMode: "background", }), ), - limit: USAGE_AGENT_LOAD_CONCURRENCY, - errorMode: "stop", - }); - if (agentLoadResult.hasError) { - throw agentLoadResult.firstError; - } - const summaries = agentLoadResult.results; + ); const dailyByDate = new Map(); const totals = createEmptyCostUsageTotals(); let cacheStatus: UsageCacheStatus | undefined; @@ -1275,8 +1283,8 @@ export const usageHandlers: GatewayRequestHandlers = { } } - const agentLoadResult = await runTasksWithConcurrency({ - tasks: Array.from(sessionsByAgent.entries()).map(([agentId, agentSessions]) => async () => ({ + const agentLoads = await runUsageAgentTasks( + Array.from(sessionsByAgent.entries()).map(([agentId, agentSessions]) => async () => ({ agentSessions, loaded: await loadSessionCostSummariesFromCache({ sessions: agentSessions, @@ -1287,13 +1295,8 @@ export const usageHandlers: GatewayRequestHandlers = { dailyUtcOffsetMinutes, }), })), - limit: USAGE_AGENT_LOAD_CONCURRENCY, - errorMode: "stop", - }); - if (agentLoadResult.hasError) { - throw agentLoadResult.firstError; - } - for (const { agentSessions, loaded } of agentLoadResult.results) { + ); + for (const { agentSessions, loaded } of agentLoads) { cacheStatus = mergeUsageCacheStatus(cacheStatus, loaded.cacheStatus); for (const [index, summary] of loaded.summaries.entries()) { if (!summary) {