fix(gateway): bound sessions.usage all-agent session discovery concurrency (#102013)

* fix(gateway): bound sessions.usage all-agent session discovery concurrency

* refactor(gateway): centralize usage agent tasks

---------

Co-authored-by: Peter Steinberger <steipete@gmail.com>

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<typeof import("../session-utils.js")>("../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<void>((resolve) => {
+      releaseLoads = resolve;
+    });
+    let resolveFirstBatchStarted!: () => void;
+    const firstBatchStarted = new Promise<void>((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);
+  });
 });

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<T>(tasks: Array<() => Promise<T>>): Promise<T[]> {
+  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<string, CostUsageTotals & { date: string }>();
   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) {