From 6225ec30082cda96a2467183925e55f8b25b3484 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 2 May 2026 15:34:47 +0800 Subject: [PATCH] refactor(server): extract foreground run state from agent-manager (#666) Centralize pending-run, waiter, and finalized-turn bookkeeping in a new ForegroundRunState module so streamAgent, replaceAgentRun, cancelAgentRun, reload/close, and stream dispatch share one coordinator instead of hand-rolling the lifecycle in five places. --- .../server/src/server/agent/agent-manager.ts | 203 +++------------ .../src/server/agent/foreground-run-state.ts | 232 ++++++++++++++++++ 2 files changed, 268 insertions(+), 167 deletions(-) create mode 100644 packages/server/src/server/agent/foreground-run-state.ts diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 105e63a98..88a2788ea 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -48,6 +48,7 @@ import { AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS, AgentStreamCoalescer, } from "./agent-stream-coalescer.js"; +import { ForegroundRunState, type ForegroundTurnWaiter } from "./foreground-run-state.js"; import { getAgentProviderDefinition } from "./provider-manifest.js"; const RELOAD_SESSION_CLOSE_TIMEOUT_MS = 3_000; @@ -177,22 +178,6 @@ interface HandleStreamEventOptions { canonicalUserMessagesById?: ReadonlyMap; } -interface ForegroundTurnWaiter { - turnId: string; - callback: (event: AgentStreamEvent) => void; - settled: boolean; - settledPromise: Promise; - resolveSettled: () => void; -} - -interface PendingForegroundRun { - token: string; - started: boolean; - settled: boolean; - settledPromise: Promise; - resolveSettled: () => void; -} - interface ManagedAgentBase { id: string; provider: AgentProvider; @@ -399,7 +384,7 @@ export class AgentManager { private readonly timelineStore = new InMemoryAgentTimelineStore(); private readonly agentsAwaitingInitialSnapshotPersist = new Set(); private readonly sessionEventTails = new Map>(); - private readonly pendingForegroundRuns = new Map(); + private readonly foregroundRuns = new ForegroundRunState(); private readonly subscribers = new Set(); private readonly idFactory: () => string; private readonly registry?: AgentStorage; @@ -524,7 +509,7 @@ export class AgentManager { return ( agent.lifecycle === "running" || Boolean(agent.activeForegroundTurnId) || - this.hasPendingForegroundRun(agentId) + this.foregroundRuns.hasPendingRun(agentId) ); } @@ -848,11 +833,7 @@ export class AgentManager { existing.unsubscribeSession(); existing.unsubscribeSession = null; } - for (const waiter of existing.foregroundTurnWaiters) { - this.settleForegroundTurnWaiter(waiter); - } - existing.foregroundTurnWaiters.clear(); - this.settlePendingForegroundRun(agentId); + this.foregroundRuns.clearAgent(agentId, existing); await this.closeReloadedSession(existing.session, agentId); // Preserve existing labels and timeline during reload. @@ -1291,18 +1272,18 @@ export class AgentManager { agentId, lifecycle: existingAgent.lifecycle, activeForegroundTurnId: existingAgent.activeForegroundTurnId, - hasPendingForegroundRun: this.hasPendingForegroundRun(agentId), + hasPendingForegroundRun: this.foregroundRuns.hasPendingRun(agentId), promptType: typeof prompt === "string" ? "string" : "structured", hasRunOptions: Boolean(options), }, "streamAgent: requested", ); - if (existingAgent.activeForegroundTurnId || this.hasPendingForegroundRun(agentId)) { + if (existingAgent.activeForegroundTurnId || this.foregroundRuns.hasPendingRun(agentId)) { this.logger.trace( { agentId, lifecycle: existingAgent.lifecycle, - hasPendingForegroundRun: this.hasPendingForegroundRun(agentId), + hasPendingForegroundRun: this.foregroundRuns.hasPendingRun(agentId), }, "streamAgent: rejected because a foreground run is already in flight", ); @@ -1313,12 +1294,11 @@ export class AgentManager { agent.pendingReplacement = false; agent.lastError = undefined; - const pendingRun = this.createPendingForegroundRun(); - this.pendingForegroundRuns.set(agentId, pendingRun); + const pendingRun = this.foregroundRuns.createPendingRun(agentId); const streamForwarder = async function* streamForwarder(this: AgentManager) { let turnId: string; - let waiter: ForegroundTurnWaiter | null = null; + let turnStream: ReturnType | null = null; try { const result = await agent.session.startTurn(prompt, options); turnId = result.turnId; @@ -1330,7 +1310,7 @@ export class AgentManager { error: errorMsg, }); this.finalizeForegroundTurn(agent); - this.settlePendingForegroundRun(agentId, pendingRun.token); + this.foregroundRuns.settlePendingRun(agentId, pendingRun.token); throw error; } @@ -1348,55 +1328,18 @@ export class AgentManager { "streamAgent: started", ); - // Create a pushable queue for this foreground turn - const queue: AgentStreamEvent[] = []; - let queueResolve: (() => void) | null = null; - let done = false; - let resolveSettled!: () => void; - const settledPromise = new Promise((resolvePromise) => { - resolveSettled = resolvePromise; - }); - - waiter = { - turnId, - settled: false, - settledPromise, - resolveSettled, - callback: (event: AgentStreamEvent) => { - queue.push(event); - if (queueResolve) { - queueResolve(); - queueResolve = null; - } - }, - }; - agent.foregroundTurnWaiters.add(waiter); + turnStream = this.foregroundRuns.createTurnStream(turnId); + this.foregroundRuns.addWaiter(agent, turnStream.waiter); try { - while (!done) { - while (queue.length > 0) { - const event = queue.shift()!; - yield event; - if (isTurnTerminalEvent(event)) { - done = true; - break; - } - } - if (!done && queue.length === 0) { - if (waiter.settled) { - break; - } - await new Promise((resolvePromise) => { - queueResolve = resolvePromise; - }); - } + for await (const event of turnStream.events(isTurnTerminalEvent)) { + yield event; } } finally { - if (waiter) { - agent.foregroundTurnWaiters.delete(waiter); - this.settleForegroundTurnWaiter(waiter); + if (turnStream) { + this.foregroundRuns.deleteWaiter(agent, turnStream.waiter); } - this.settlePendingForegroundRun(agentId, pendingRun.token); + this.foregroundRuns.settlePendingRun(agentId, pendingRun.token); if (!agent.activeForegroundTurnId) { await this.refreshRuntimeInfo(agent); } @@ -1409,7 +1352,7 @@ export class AgentManager { private finalizeForegroundTurn(agent: ActiveManagedAgent, turnId?: string): void { const mutableAgent = agent as ActiveManagedAgent; if (turnId) { - this.rememberFinalizedForegroundTurn(mutableAgent, turnId); + this.foregroundRuns.rememberFinalizedTurn(mutableAgent, turnId); } mutableAgent.activeForegroundTurnId = null; const terminalError = mutableAgent.lastError; @@ -1455,7 +1398,7 @@ export class AgentManager { if ( snapshot.lifecycle !== "running" && !snapshot.activeForegroundTurnId && - !this.hasPendingForegroundRun(agentId) + !this.foregroundRuns.hasPendingRun(agentId) ) { return this.streamAgent(agentId, prompt, options); } @@ -1495,7 +1438,7 @@ export class AgentManager { throw new Error(`Agent ${agentId} not found`); } - const pendingRun = this.getPendingForegroundRun(agentId); + const pendingRun = this.foregroundRuns.getPendingRun(agentId); if ((snapshot.lifecycle === "running" || pendingRun?.started) && !snapshot.pendingReplacement) { return; } @@ -1559,7 +1502,7 @@ export class AgentManager { return true; } - const currentPendingRun = this.getPendingForegroundRun(agentId); + const currentPendingRun = this.foregroundRuns.getPendingRun(agentId); if ( (current.lifecycle === "running" || currentPendingRun?.started) && !current.pendingReplacement @@ -1632,7 +1575,7 @@ export class AgentManager { async cancelAgentRun(agentId: string): Promise { const agent = this.requireSessionAgent(agentId); - const pendingRun = this.getPendingForegroundRun(agentId); + const pendingRun = this.foregroundRuns.getPendingRun(agentId); const foregroundTurnId = agent.activeForegroundTurnId; const hasForegroundTurn = Boolean(foregroundTurnId); const isAutonomousRunning = agent.lifecycle === "running" && !hasForegroundTurn && !pendingRun; @@ -1701,7 +1644,7 @@ export class AgentManager { }); // The synthetic event unblocks the streamForwarder generator, whose finally // block settles the pending foreground run asynchronously. Wait for it. - const staleRun = this.getPendingForegroundRun(agentId); + const staleRun = this.foregroundRuns.getPendingRun(agentId); if (staleRun && !staleRun.settled) { await staleRun.settledPromise; } @@ -1876,7 +1819,7 @@ export class AgentManager { throw new Error(`Agent ${agentId} not found`); } - const pendingForegroundRun = this.getPendingForegroundRun(agentId); + const pendingForegroundRun = this.foregroundRuns.getPendingRun(agentId); const hasForegroundTurn = Boolean(snapshot.activeForegroundTurnId) || Boolean(pendingForegroundRun); @@ -2198,17 +2141,13 @@ export class AgentManager { agent.unsubscribeSession(); agent.unsubscribeSession = null; } - for (const waiter of agent.foregroundTurnWaiters) { - waiter.callback({ - type: "turn_canceled", - provider: agent.provider, - reason: cancelReason, - turnId: waiter.turnId, - }); - this.settleForegroundTurnWaiter(waiter); - } - agent.foregroundTurnWaiters.clear(); - this.settlePendingForegroundRun(agent.id); + this.foregroundRuns.cancelWaiters(agent, (turnId) => ({ + type: "turn_canceled", + provider: agent.provider, + reason: cancelReason, + turnId, + })); + this.foregroundRuns.settlePendingRun(agent.id); return { ...agent, lifecycle: "closed", @@ -2267,12 +2206,7 @@ export class AgentManager { event: AgentStreamEvent, ): Promise { const turnId = (event as { turnId?: string }).turnId; - const matchingWaiters = - turnId == null - ? [] - : Array.from(agent.foregroundTurnWaiters).filter( - (waiter) => waiter.turnId === turnId && !waiter.settled, - ); + const matchingWaiters = this.foregroundRuns.getMatchingWaiters(agent, turnId); const shouldNotifyWaiters = await this.handleStreamEvent(agent, event); @@ -2280,70 +2214,9 @@ export class AgentManager { return; } - for (const waiter of matchingWaiters) { - waiter.callback(event); - if (isTurnTerminalEvent(event)) { - this.settleForegroundTurnWaiter(waiter); - } - } - } - - private settleForegroundTurnWaiter(waiter: ForegroundTurnWaiter): void { - if (waiter.settled) { - return; - } - waiter.settled = true; - waiter.resolveSettled(); - } - - private rememberFinalizedForegroundTurn(agent: ActiveManagedAgent, turnId: string): void { - agent.finalizedForegroundTurnIds.add(turnId); - if (agent.finalizedForegroundTurnIds.size <= 50) { - return; - } - const oldest = agent.finalizedForegroundTurnIds.values().next().value; - if (oldest) { - agent.finalizedForegroundTurnIds.delete(oldest); - } - } - - private createPendingForegroundRun(): PendingForegroundRun { - let resolveSettled!: () => void; - const settledPromise = new Promise((resolvePromise) => { - resolveSettled = resolvePromise; + this.foregroundRuns.notifyWaiters(matchingWaiters, event, { + terminal: isTurnTerminalEvent(event), }); - return { - token: randomUUID(), - started: false, - settled: false, - settledPromise, - resolveSettled, - }; - } - - private getPendingForegroundRun(agentId: string): PendingForegroundRun | null { - return this.pendingForegroundRuns.get(agentId) ?? null; - } - - private hasPendingForegroundRun(agentId: string): boolean { - return this.pendingForegroundRuns.has(agentId); - } - - private settlePendingForegroundRun(agentId: string, token?: string): void { - const pendingRun = this.pendingForegroundRuns.get(agentId); - if (!pendingRun) { - return; - } - if (token && pendingRun.token !== token) { - return; - } - - this.pendingForegroundRuns.delete(agentId); - if (pendingRun.settled) { - return; - } - pendingRun.settled = true; - pendingRun.resolveSettled(); } private async resolveInitialPersistedTitle( @@ -2466,11 +2339,7 @@ export class AgentManager { return; } - for (const waiter of agent.foregroundTurnWaiters) { - if (waiter.turnId === turnId && !waiter.settled) { - waiter.callback(event); - } - } + this.foregroundRuns.notifyAgentWaiters(agent, event); } private async handleStreamEvent( @@ -2483,7 +2352,7 @@ export class AgentManager { if ( eventTurnId && isTurnTerminalEvent(event) && - agent.finalizedForegroundTurnIds.has(eventTurnId) + this.foregroundRuns.hasFinalizedTurn(agent, eventTurnId) ) { return false; } diff --git a/packages/server/src/server/agent/foreground-run-state.ts b/packages/server/src/server/agent/foreground-run-state.ts new file mode 100644 index 000000000..44d33c045 --- /dev/null +++ b/packages/server/src/server/agent/foreground-run-state.ts @@ -0,0 +1,232 @@ +import { randomUUID } from "node:crypto"; + +import type { AgentStreamEvent } from "./agent-sdk-types.js"; + +export interface ForegroundTurnWaiter { + turnId: string; + callback: (event: AgentStreamEvent) => void; + settled: boolean; + settledPromise: Promise; + resolveSettled: () => void; +} + +export interface PendingForegroundRun { + token: string; + started: boolean; + settled: boolean; + settledPromise: Promise; + resolveSettled: () => void; +} + +export interface ForegroundRunAgentState { + foregroundTurnWaiters: Set; + finalizedForegroundTurnIds: Set; +} + +export class ForegroundRunState { + private readonly pendingRuns = new Map(); + + createPendingRun(agentId: string): PendingForegroundRun { + const pendingRun = createPendingForegroundRun(); + this.pendingRuns.set(agentId, pendingRun); + return pendingRun; + } + + getPendingRun(agentId: string): PendingForegroundRun | null { + return this.pendingRuns.get(agentId) ?? null; + } + + hasPendingRun(agentId: string): boolean { + return this.pendingRuns.has(agentId); + } + + settlePendingRun(agentId: string, token?: string): void { + const pendingRun = this.pendingRuns.get(agentId); + if (!pendingRun) { + return; + } + if (token && pendingRun.token !== token) { + return; + } + + this.pendingRuns.delete(agentId); + settlePendingForegroundRun(pendingRun); + } + + createTurnStream(turnId: string): ForegroundTurnStream { + return new ForegroundTurnStream(turnId); + } + + addWaiter(agent: ForegroundRunAgentState, waiter: ForegroundTurnWaiter): void { + agent.foregroundTurnWaiters.add(waiter); + } + + deleteWaiter(agent: ForegroundRunAgentState, waiter: ForegroundTurnWaiter): void { + agent.foregroundTurnWaiters.delete(waiter); + this.settleWaiter(waiter); + } + + settleWaiter(waiter: ForegroundTurnWaiter): void { + if (waiter.settled) { + return; + } + waiter.settled = true; + waiter.resolveSettled(); + } + + getMatchingWaiters( + agent: ForegroundRunAgentState, + turnId: string | undefined, + ): ForegroundTurnWaiter[] { + if (turnId == null) { + return []; + } + + return Array.from(agent.foregroundTurnWaiters).filter( + (waiter) => waiter.turnId === turnId && !waiter.settled, + ); + } + + notifyWaiters( + waiters: Iterable, + event: AgentStreamEvent, + options: { terminal: boolean }, + ): void { + for (const waiter of waiters) { + waiter.callback(event); + if (options.terminal) { + this.settleWaiter(waiter); + } + } + } + + notifyAgentWaiters( + agent: ForegroundRunAgentState, + event: AgentStreamEvent, + options?: { terminal?: boolean }, + ): void { + const waiters = this.getMatchingWaiters(agent, eventTurnId(event)); + this.notifyWaiters(waiters, event, { terminal: options?.terminal ?? false }); + } + + cancelWaiters( + agent: ForegroundRunAgentState, + createEvent: (turnId: string) => AgentStreamEvent, + ): void { + for (const waiter of agent.foregroundTurnWaiters) { + waiter.callback(createEvent(waiter.turnId)); + this.settleWaiter(waiter); + } + agent.foregroundTurnWaiters.clear(); + } + + clearAgent(agentId: string, agent: ForegroundRunAgentState): void { + for (const waiter of agent.foregroundTurnWaiters) { + this.settleWaiter(waiter); + } + agent.foregroundTurnWaiters.clear(); + this.settlePendingRun(agentId); + } + + rememberFinalizedTurn(agent: ForegroundRunAgentState, turnId: string): void { + agent.finalizedForegroundTurnIds.add(turnId); + if (agent.finalizedForegroundTurnIds.size <= 50) { + return; + } + + const oldest = agent.finalizedForegroundTurnIds.values().next().value; + if (oldest) { + agent.finalizedForegroundTurnIds.delete(oldest); + } + } + + hasFinalizedTurn(agent: ForegroundRunAgentState, turnId: string): boolean { + return agent.finalizedForegroundTurnIds.has(turnId); + } +} + +export class ForegroundTurnStream { + private readonly queue: AgentStreamEvent[] = []; + private queueResolve: (() => void) | null = null; + + readonly waiter: ForegroundTurnWaiter; + + constructor(turnId: string) { + let resolveSettled!: () => void; + const settledPromise = new Promise((resolvePromise) => { + resolveSettled = resolvePromise; + }); + + this.waiter = { + turnId, + settled: false, + settledPromise, + resolveSettled, + callback: (event) => { + this.queue.push(event); + this.wake(); + }, + }; + } + + async *events( + isTerminalEvent: (event: AgentStreamEvent) => boolean, + ): AsyncGenerator { + let done = false; + while (!done) { + while (this.queue.length > 0) { + const event = this.queue.shift()!; + yield event; + if (isTerminalEvent(event)) { + done = true; + break; + } + } + + if (!done && this.queue.length === 0) { + if (this.waiter.settled) { + break; + } + await new Promise((resolvePromise) => { + this.queueResolve = resolvePromise; + }); + } + } + } + + private wake(): void { + if (!this.queueResolve) { + return; + } + + this.queueResolve(); + this.queueResolve = null; + } +} + +function createPendingForegroundRun(): PendingForegroundRun { + let resolveSettled!: () => void; + const settledPromise = new Promise((resolvePromise) => { + resolveSettled = resolvePromise; + }); + return { + token: randomUUID(), + started: false, + settled: false, + settledPromise, + resolveSettled, + }; +} + +function settlePendingForegroundRun(pendingRun: PendingForegroundRun): void { + if (pendingRun.settled) { + return; + } + + pendingRun.settled = true; + pendingRun.resolveSettled(); +} + +function eventTurnId(event: AgentStreamEvent): string | undefined { + return (event as { turnId?: string }).turnId; +}