Gateway: fix multi-agent sessions.usage discovery (#11523)

* Gateway: fix multi-agent sessions.usage discovery

* Gateway: resolve sessions.usage keys via sessionId
This commit is contained in:
Tak Hoffman
2026-02-07 17:40:56 -06:00
committed by GitHub
parent b8f740fb14
commit 9271fcb3d4
2 changed files with 214 additions and 12 deletions

View File

@@ -19,6 +19,7 @@ import {
loadSessionCostSummary,
loadSessionUsageTimeSeries,
discoverAllSessions,
type DiscoveredSession,
} from "../../infra/session-cost-usage.js";
import { parseAgentSessionKey } from "../../routing/session-key.js";
import {
@@ -27,7 +28,11 @@ import {
formatValidationErrors,
validateSessionsUsageParams,
} from "../protocol/index.js";
import { loadCombinedSessionStoreForGateway, loadSessionEntry } from "../session-utils.js";
import {
listAgentsForGateway,
loadCombinedSessionStoreForGateway,
loadSessionEntry,
} from "../session-utils.js";
const COST_USAGE_CACHE_TTL_MS = 30_000;
@@ -109,6 +114,27 @@ const parseDateRange = (params: {
return { startMs: defaultStartMs, endMs: todayEndMs };
};
type DiscoveredSessionWithAgent = DiscoveredSession & { agentId: string };
async function discoverAllSessionsForUsage(params: {
config: ReturnType<typeof loadConfig>;
startMs: number;
endMs: number;
}): Promise<DiscoveredSessionWithAgent[]> {
const agents = listAgentsForGateway(params.config).agents;
const results = await Promise.all(
agents.map(async (agent) => {
const sessions = await discoverAllSessions({
agentId: agent.id,
startMs: params.startMs,
endMs: params.endMs,
});
return sessions.map((session) => ({ ...session, agentId: agent.id }));
}),
);
return results.flat().toSorted((a, b) => b.mtime - a.mtime);
}
async function loadCostUsageSummaryCached(params: {
startMs: number;
endMs: number;
@@ -166,6 +192,7 @@ export const __test = {
parseDateToMs,
parseDays,
parseDateRange,
discoverAllSessionsForUsage,
loadCostUsageSummaryCached,
costUsageCache,
};
@@ -282,18 +309,37 @@ export const usageHandlers: GatewayRequestHandlers = {
// Optimization: If a specific key is requested, skip full directory scan
if (specificKey) {
// Check if it's a named session in the store
const storeEntry = store[specificKey];
let sessionId = storeEntry?.sessionId ?? specificKey;
const parsed = parseAgentSessionKey(specificKey);
const agentIdFromKey = parsed?.agentId;
const keyRest = parsed?.rest ?? specificKey;
// Prefer the store entry when available, even if the caller provides a discovered key
// (`agent:<id>:<sessionId>`) for a session that now has a canonical store key.
const storeBySessionId = new Map<string, { key: string; entry: SessionEntry }>();
for (const [key, entry] of Object.entries(store)) {
if (entry?.sessionId) {
storeBySessionId.set(entry.sessionId, { key, entry });
}
}
const storeMatch = store[specificKey]
? { key: specificKey, entry: store[specificKey] }
: null;
const storeByIdMatch = storeBySessionId.get(keyRest) ?? null;
const resolvedStoreKey = storeMatch?.key ?? storeByIdMatch?.key ?? specificKey;
const storeEntry = storeMatch?.entry ?? storeByIdMatch?.entry;
const sessionId = storeEntry?.sessionId ?? keyRest;
// Resolve the session file path
const sessionFile = resolveSessionFilePath(sessionId, storeEntry);
const sessionFile = resolveSessionFilePath(sessionId, storeEntry, {
agentId: agentIdFromKey,
});
try {
const stats = fs.statSync(sessionFile);
if (stats.isFile()) {
mergedEntries.push({
key: specificKey,
key: resolvedStoreKey,
sessionId,
sessionFile,
label: storeEntry?.label,
@@ -306,7 +352,8 @@ export const usageHandlers: GatewayRequestHandlers = {
}
} else {
// Full discovery for list view
const discoveredSessions = await discoverAllSessions({
const discoveredSessions = await discoverAllSessionsForUsage({
config,
startMs,
endMs,
});
@@ -334,7 +381,8 @@ export const usageHandlers: GatewayRequestHandlers = {
} else {
// Unnamed session - use session ID as key, no label
mergedEntries.push({
key: discovered.sessionId,
// Keep agentId in the key so the dashboard can attribute sessions and later fetch logs.
key: `agent:${discovered.agentId}:${discovered.sessionId}`,
sessionId: discovered.sessionId,
sessionFile: discovered.sessionFile,
label: undefined, // No label for unnamed sessions
@@ -711,8 +759,12 @@ export const usageHandlers: GatewayRequestHandlers = {
const { entry } = loadSessionEntry(key);
// For discovered sessions (not in store), try using key as sessionId directly
const sessionId = entry?.sessionId ?? key;
const sessionFile = entry?.sessionFile ?? resolveSessionFilePath(key);
const parsed = parseAgentSessionKey(key);
const agentId = parsed?.agentId;
const rawSessionId = parsed?.rest ?? key;
const sessionId = entry?.sessionId ?? rawSessionId;
const sessionFile =
entry?.sessionFile ?? resolveSessionFilePath(rawSessionId, entry, { agentId });
const timeseries = await loadSessionUsageTimeSeries({
sessionId,
@@ -749,8 +801,12 @@ export const usageHandlers: GatewayRequestHandlers = {
const { entry } = loadSessionEntry(key);
// For discovered sessions (not in store), try using key as sessionId directly
const sessionId = entry?.sessionId ?? key;
const sessionFile = entry?.sessionFile ?? resolveSessionFilePath(key);
const parsed = parseAgentSessionKey(key);
const agentId = parsed?.agentId;
const rawSessionId = parsed?.rest ?? key;
const sessionId = entry?.sessionId ?? rawSessionId;
const sessionFile =
entry?.sessionFile ?? resolveSessionFilePath(rawSessionId, entry, { agentId });
const { loadSessionLogs } = await import("../../infra/session-cost-usage.js");
const logs = await loadSessionLogs({