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.

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<string, string>;
 }
 
-interface ForegroundTurnWaiter {
-  turnId: string;
-  callback: (event: AgentStreamEvent) => void;
-  settled: boolean;
-  settledPromise: Promise<void>;
-  resolveSettled: () => void;
-}
-
-interface PendingForegroundRun {
-  token: string;
-  started: boolean;
-  settled: boolean;
-  settledPromise: Promise<void>;
-  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<string>();
   private readonly sessionEventTails = new Map<string, Promise<void>>();
-  private readonly pendingForegroundRuns = new Map<string, PendingForegroundRun>();
+  private readonly foregroundRuns = new ForegroundRunState();
   private readonly subscribers = new Set<SubscriptionRecord>();
   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.foregroundRuns.clearAgent(agentId, existing);
-      this.settleForegroundTurnWaiter(waiter);
-    }
-    existing.foregroundTurnWaiters.clear();
-    this.settlePendingForegroundRun(agentId);
     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();
+    const pendingRun = this.foregroundRuns.createPendingRun(agentId);
-    this.pendingForegroundRuns.set(agentId, pendingRun);
 
     const streamForwarder = async function* streamForwarder(this: AgentManager) {
       let turnId: string;
-      let waiter: ForegroundTurnWaiter | null = null;
+      let turnStream: ReturnType<ForegroundRunState["createTurnStream"]> | 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
+      turnStream = this.foregroundRuns.createTurnStream(turnId);
-      const queue: AgentStreamEvent[] = [];
+      this.foregroundRuns.addWaiter(agent, turnStream.waiter);
-      let queueResolve: (() => void) | null = null;
-      let done = false;
-      let resolveSettled!: () => void;
-      const settledPromise = new Promise<void>((resolvePromise) => {
-        resolveSettled = resolvePromise;
-      });
-
-      waiter = {
-        turnId,
-        settled: false,
-        settledPromise,
-        resolveSettled,
-        callback: (event: AgentStreamEvent) => {
-          queue.push(event);
-          if (queueResolve) {
-            queueResolve();
-            queueResolve = null;
-          }
-        },
-      };
-      agent.foregroundTurnWaiters.add(waiter);
 
       try {
-        while (!done) {
+        for await (const event of turnStream.events(isTurnTerminalEvent)) {
-          while (queue.length > 0) {
+          yield event;
-            const event = queue.shift()!;
-            yield event;
-            if (isTurnTerminalEvent(event)) {
-              done = true;
-              break;
-            }
-          }
-          if (!done && queue.length === 0) {
-            if (waiter.settled) {
-              break;
-            }
-            await new Promise<void>((resolvePromise) => {
-              queueResolve = resolvePromise;
-            });
-          }
         }
       } finally {
-        if (waiter) {
+        if (turnStream) {
-          agent.foregroundTurnWaiters.delete(waiter);
+          this.foregroundRuns.deleteWaiter(agent, turnStream.waiter);
-          this.settleForegroundTurnWaiter(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<boolean> {
     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) {
+    this.foregroundRuns.cancelWaiters(agent, (turnId) => ({
-      waiter.callback({
+      type: "turn_canceled",
-        type: "turn_canceled",
+      provider: agent.provider,
-        provider: agent.provider,
+      reason: cancelReason,
-        reason: cancelReason,
+      turnId,
-        turnId: waiter.turnId,
+    }));
-      });
+    this.foregroundRuns.settlePendingRun(agent.id);
-      this.settleForegroundTurnWaiter(waiter);
-    }
-    agent.foregroundTurnWaiters.clear();
-    this.settlePendingForegroundRun(agent.id);
     return {
       ...agent,
       lifecycle: "closed",
@@ -2267,12 +2206,7 @@ export class AgentManager {
     event: AgentStreamEvent,
   ): Promise<void> {
     const turnId = (event as { turnId?: string }).turnId;
-    const matchingWaiters =
+    const matchingWaiters = this.foregroundRuns.getMatchingWaiters(agent, turnId);
-      turnId == null
-        ? []
-        : Array.from(agent.foregroundTurnWaiters).filter(
-            (waiter) => waiter.turnId === turnId && !waiter.settled,
-          );
 
     const shouldNotifyWaiters = await this.handleStreamEvent(agent, event);
 
@@ -2280,70 +2214,9 @@ export class AgentManager {
       return;
     }
 
-    for (const waiter of matchingWaiters) {
+    this.foregroundRuns.notifyWaiters(matchingWaiters, event, {
-      waiter.callback(event);
+      terminal: isTurnTerminalEvent(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<void>((resolvePromise) => {
-      resolveSettled = resolvePromise;
     });
-    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) {
+    this.foregroundRuns.notifyAgentWaiters(agent, event);
-      if (waiter.turnId === turnId && !waiter.settled) {
-        waiter.callback(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;
     }

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<void>;
+  resolveSettled: () => void;
+}
+
+export interface PendingForegroundRun {
+  token: string;
+  started: boolean;
+  settled: boolean;
+  settledPromise: Promise<void>;
+  resolveSettled: () => void;
+}
+
+export interface ForegroundRunAgentState {
+  foregroundTurnWaiters: Set<ForegroundTurnWaiter>;
+  finalizedForegroundTurnIds: Set<string>;
+}
+
+export class ForegroundRunState {
+  private readonly pendingRuns = new Map<string, PendingForegroundRun>();
+
+  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<ForegroundTurnWaiter>,
+    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<void>((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<AgentStreamEvent> {
+    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<void>((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<void>((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;
+}