diff --git a/src/gateway/server-methods/agent-job.ts b/src/gateway/server-methods/agent-job.ts index f08de1f2594d..bb6b2334c517 100644 --- a/src/gateway/server-methods/agent-job.ts +++ b/src/gateway/server-methods/agent-job.ts @@ -10,6 +10,7 @@ import { setSafeTimeout } from "../../utils/timer-delay.js"; import type { AgentWaitTerminalSnapshot } from "./agent-wait-dedupe.js"; const AGENT_RUN_CACHE_TTL_MS = 10 * 60_000; +const AGENT_RUN_CACHE_MAX_ENTRIES = 5_000; /** * Embedded runs can emit transient lifecycle `error` events while auth/model * failover is still in progress. Give errors a short grace window so a @@ -63,6 +64,28 @@ function recordAgentRunSnapshot(entry: AgentRunSnapshot) { return; } agentRunCache.set(entry.runId, entry); + // Time-based prune only fires on the TTL window; under high run fan-out a + // burst can add far more entries than the window reclaims. Cap with a FIFO + // drop so the cache cannot grow without bound between prunes. + enforceAgentRunCacheMaxEntries(); +} + +function enforceAgentRunCacheMaxEntries() { + if (agentRunCache.size <= AGENT_RUN_CACHE_MAX_ENTRIES) { + return; + } + const toRemove = agentRunCache.size - AGENT_RUN_CACHE_MAX_ENTRIES; + let removed = 0; + for (const runId of agentRunCache.keys()) { + if (removed >= toRemove) { + break; + } + if ((agentRunWaiterCounts.get(runId) ?? 0) > 0) { + continue; + } + agentRunCache.delete(runId); + removed += 1; + } } function shouldPreserveTerminalSnapshot( @@ -494,5 +517,12 @@ export const testing = { resetWaiters(): void { agentRunWaiterCounts.clear(); }, + getAgentRunCacheSize(): number { + return agentRunCache.size; + }, + resetAgentRunCache(): void { + agentRunCache.clear(); + }, + agentRunCacheMaxEntries: AGENT_RUN_CACHE_MAX_ENTRIES, }; export { testing as __testing }; diff --git a/src/gateway/server-methods/server-methods.test.ts b/src/gateway/server-methods/server-methods.test.ts index f9f51c18ac3c..aa4e22dd7763 100644 --- a/src/gateway/server-methods/server-methods.test.ts +++ b/src/gateway/server-methods/server-methods.test.ts @@ -37,7 +37,7 @@ import { } from "../chat-display-projection.js"; import { sanitizeChatSendMessageInput } from "../chat-input-sanitize.js"; import { ExecApprovalManager } from "../exec-approval-manager.js"; -import { waitForAgentJob } from "./agent-job.js"; +import { __testing as agentJobTesting, waitForAgentJob } from "./agent-job.js"; import { injectTimestamp, timestampOptsFromConfig } from "./agent-timestamp.js"; import { normalizeRpcAttachmentsToChatAttachments } from "./attachment-normalize.js"; import { createExecApprovalHandlers } from "./exec-approval.js"; @@ -592,6 +592,74 @@ describe("waitForAgentJob", () => { vi.useRealTimers(); } }); + + it("caps agentRunCache at AGENT_RUN_CACHE_MAX_ENTRIES via FIFO drop", () => { + agentJobTesting.resetAgentRunCache(); + const max = agentJobTesting.agentRunCacheMaxEntries; + const overflow = 25; + const prefix = `cap-${Date.now()}-${Math.random().toString(36).slice(2)}`; + for (let i = 0; i < max + overflow; i++) { + emitAgentEvent({ + runId: `${prefix}-${i}`, + stream: "lifecycle", + data: { phase: "end", startedAt: i, endedAt: i + 1 }, + }); + } + expect(agentJobTesting.getAgentRunCacheSize()).toBe(max); + agentJobTesting.resetAgentRunCache(); + }); + + it("does not evict cached terminal snapshots with active fresh waiters", async () => { + agentJobTesting.resetAgentRunCache(); + const max = agentJobTesting.agentRunCacheMaxEntries; + const prefix = `cap-waiter-${Date.now()}-${Math.random().toString(36).slice(2)}`; + const waitedRunId = `${prefix}-waited`; + emitAgentEvent({ + runId: waitedRunId, + stream: "lifecycle", + data: { phase: "end", startedAt: 1_000, endedAt: 1_100 }, + }); + const waitPromise = waitForAgentJob({ + runId: waitedRunId, + timeoutMs: 5_000, + ignoreCachedSnapshot: true, + }); + + for (let i = 0; i < max + 25; i++) { + emitAgentEvent({ + runId: `${prefix}-${i}`, + stream: "lifecycle", + data: { phase: "end", startedAt: i, endedAt: i + 1 }, + }); + } + const cached = await waitForAgentJob({ runId: waitedRunId, timeoutMs: 0 }); + expectRecordFields(cached, { + status: "ok", + startedAt: 1_000, + endedAt: 1_100, + }); + expect(agentJobTesting.getAgentRunCacheSize()).toBe(max); + + emitAgentEvent({ + runId: waitedRunId, + stream: "lifecycle", + data: { phase: "end", startedAt: 10_000, endedAt: 10_100 }, + }); + + const waited = await waitPromise; + expectRecordFields(waited, { + status: "ok", + startedAt: 10_000, + endedAt: 10_100, + }); + emitAgentEvent({ + runId: `${prefix}-after-waiter`, + stream: "lifecycle", + data: { phase: "end", startedAt: 20_000, endedAt: 20_100 }, + }); + expect(agentJobTesting.getAgentRunCacheSize()).toBe(max); + agentJobTesting.resetAgentRunCache(); + }); }); describe("augmentChatHistoryWithCanvasBlocks", () => {