diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 1d6f323deb..fd28fe5550 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -956,6 +956,73 @@ test("unavailable steer interrupts once and starts one replacement turn", async } }); +test("steers a send that arrives while the previous turn is still starting", async () => { + const startHeld = deferred(); + class SlowStartSteeringSession extends SteeringTestSession { + override async startTurn(prompt: AgentPromptInput): Promise<{ turnId: string }> { + await startHeld.promise; + return await super.startTurn(prompt); + } + } + const session = new SlowStartSteeringSession({ + provider: "codex", + cwd: process.cwd(), + }); + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-steer-starting-")); + const manager = new AgentManager({ + clients: { + codex: new (class extends TestAgentClient { + override async createSession(): Promise { + return session; + } + })(), + }, + logger, + }); + let agentId: string | null = null; + try { + const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, { + workspaceId: undefined, + }); + agentId = agent.id; + const steerOptions = (clientMessageId: string) => ({ + replaceRunning: true, + activeTurnBehavior: "steer" as const, + runOptions: { clientMessageId }, + }); + await startAgentRun(manager, agent.id, "FIRST", logger, steerOptions("first-client")); + const second = startAgentRun( + manager, + agent.id, + "SECOND", + logger, + steerOptions("second-client"), + ); + + startHeld.resolve(); + await expect(second).resolves.toEqual({ disposition: "steered" }); + + expect(session.startPrompts).toEqual(["FIRST"]); + expect(session.interruptCount).toBe(0); + expect( + manager + .fetchTimeline(agent.id, { limit: 0 }) + .rows.filter((row) => row.item.type === "user_message") + .map((row) => ({ + text: row.item.type === "user_message" ? row.item.text : null, + turnId: row.turnId, + })), + ).toEqual([ + { text: "FIRST", turnId: "active-turn-1" }, + { text: "SECOND", turnId: "active-turn-1" }, + ]); + } finally { + startHeld.resolve(); + if (agentId) await manager.closeAgent(agentId).catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("orders an accepted steer before output emitted while acknowledgement is pending", async () => { const entered = deferred(); const release = deferred(); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 46c822ea2d..a89eb338b7 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -2734,7 +2734,7 @@ export class AgentManager { throw error; } if (isStaleProviderSessionError(error)) { - pendingRun.start = { status: "failed", error: error.message }; + this.runs.settleForegroundStart(pendingRun, { status: "failed", error: error.message }); agent.pendingReplacement = false; if (!agent.activeForegroundTurnId) agent.lifecycle = "idle"; this.runs.settleForegroundRun(agentId, pendingRun.token); @@ -2742,7 +2742,7 @@ export class AgentManager { } agent.pendingReplacement = false; const errorMsg = error instanceof Error ? error.message : "Failed to start turn"; - pendingRun.start = { status: "failed", error: errorMsg }; + this.runs.settleForegroundStart(pendingRun, { status: "failed", error: errorMsg }); // A terminal rejection belongs after the submitted prompt even though no provider turn exists. if (options?.clientMessageId) { this.recordSubmittedPrompt(agent, prompt, options.clientMessageId, "rejected", { @@ -2832,7 +2832,7 @@ export class AgentManager { agent.pendingReplacement = false; } const turnStartedAt = new Date(); - pendingRun.start = { status: "started", turnId }; + this.runs.settleForegroundStart(pendingRun, { status: "started", turnId }); agent.activeForegroundTurnId = turnId; this.openActiveTurn(agent, turnId, turnStartedAt); agent.lifecycle = "running"; @@ -3041,6 +3041,9 @@ export class AgentManager { prompt: AgentPromptInput, options?: AgentSteerOptions, ): Promise { + // A turn that is still starting has no id to steer yet. Replacing it would drop its + // prompt, so wait until it is accepted (or fails) and steer the turn it became. + await this.runs.waitForForegroundStart(agentId); const agent = this.requireSessionAgent(agentId); const expectedTurnId = agent.activeForegroundTurnId ?? agent.activeTurnId; if (!expectedTurnId) { diff --git a/packages/server/src/server/agent/agent-run-state.ts b/packages/server/src/server/agent/agent-run-state.ts index ad1eb0af34..f2f4ee3d16 100644 --- a/packages/server/src/server/agent/agent-run-state.ts +++ b/packages/server/src/server/agent/agent-run-state.ts @@ -10,14 +10,18 @@ export interface ForegroundTurnWaiter { resolveSettled: () => void; } +type ForegroundRunStart = + | { status: "pending" } + | { status: "started"; turnId: string } + | { status: "failed"; error: string }; + export interface PendingForegroundRun { token: string; kind: "foreground"; stagedEvents: AgentStreamEvent[]; - start: - | { status: "pending" } - | { status: "started"; turnId: string } - | { status: "failed"; error: string }; + start: ForegroundRunStart; + startSettledPromise: Promise; + resolveStartSettled: () => void; settled: boolean; settledPromise: Promise; resolveSettled: () => void; @@ -54,6 +58,23 @@ export class AgentRunState { return run?.kind === "foreground" ? run : null; } + settleForegroundStart( + run: PendingForegroundRun, + start: Exclude, + ): void { + run.start = start; + run.resolveStartSettled(); + } + + /** Resolves once the agent's foreground run has a turn id, failed to start, or ended. */ + async waitForForegroundStart(agentId: string): Promise { + const run = this.getPendingRun(agentId); + if (run?.start.status !== "pending") { + return; + } + await Promise.race([run.startSettledPromise, run.settledPromise]); + } + hasPendingRun(agentId: string): boolean { return this.getPendingRun(agentId) !== null; } @@ -277,10 +298,16 @@ export class ForegroundTurnStream { } function createPendingForegroundRun(): PendingForegroundRun { + let resolveStartSettled!: () => void; + const startSettledPromise = new Promise((resolvePromise) => { + resolveStartSettled = resolvePromise; + }); return { ...createTrackedRunState(), kind: "foreground", start: { status: "pending" }, + startSettledPromise, + resolveStartSettled, stagedEvents: [], }; }