diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 1d6f323deb..bb521abc71 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -956,6 +956,131 @@ 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("starts its own turn when the previous turn fails to start under a steer send", async () => { + const startHeld = deferred(); + class FailingFirstStartSession extends SteeringTestSession { + override async startTurn(prompt: AgentPromptInput): Promise<{ turnId: string }> { + if (prompt === "FIRST") { + this.startPrompts.push(prompt); + await startHeld.promise; + throw new Error("first start failed"); + } + return await super.startTurn(prompt); + } + } + const session = new FailingFirstStartSession({ + provider: "codex", + cwd: process.cwd(), + }); + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-steer-failed-start-")); + 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: "turn_started" }); + await vi.waitFor(() => expect(session.startPrompts).toEqual(["FIRST", "SECOND"])); + expect(session.interruptCount).toBe(0); + } 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..ae9e424d78 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,20 +2742,24 @@ export class AgentManager { } agent.pendingReplacement = false; const errorMsg = error instanceof Error ? error.message : "Failed to start turn"; - pendingRun.start = { 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", { - messageId: options.clientMessageId, + this.runs.settleForegroundStart(pendingRun, { status: "failed", error: errorMsg }); + try { + // A terminal rejection belongs after the submitted prompt even though no provider turn exists. + if (options?.clientMessageId) { + this.recordSubmittedPrompt(agent, prompt, options.clientMessageId, "rejected", { + messageId: options.clientMessageId, + }); + } + await this.handleStreamEvent(agent, { + type: "turn_failed", + provider: agent.provider, + error: errorMsg, }); + this.finalizeForegroundTurn(agent); + } finally { + // A steer waiting on this start is released only once the run is cleared. + this.runs.settleForegroundRun(agentId, pendingRun.token); } - await this.handleStreamEvent(agent, { - type: "turn_failed", - provider: agent.provider, - error: errorMsg, - }); - this.finalizeForegroundTurn(agent); - this.runs.settleForegroundRun(agentId, pendingRun.token); throw error; } } @@ -2832,7 +2836,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 +3045,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.test.ts b/packages/server/src/server/agent/agent-run-state.test.ts new file mode 100644 index 0000000000..29e09a5f94 --- /dev/null +++ b/packages/server/src/server/agent/agent-run-state.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, test } from "vitest"; + +import { AgentRunState } from "./agent-run-state.js"; + +function track(promise: Promise): { readonly resolved: boolean } { + const state = { resolved: false }; + void (async () => { + await promise; + state.resolved = true; + })(); + return state; +} + +async function flushMicrotasks(): Promise { + await new Promise((resolve) => setImmediate(resolve)); +} + +describe("waitForForegroundStart", () => { + test("resolves once the starting turn has an id", async () => { + const runs = new AgentRunState(); + const run = runs.createPendingRun("agent"); + const wait = track(runs.waitForForegroundStart("agent")); + + await flushMicrotasks(); + expect(wait.resolved).toBe(false); + + runs.settleForegroundStart(run, { status: "started", turnId: "turn-1" }); + await flushMicrotasks(); + expect(wait.resolved).toBe(true); + }); + + test("waits for a failed start's run to be cleared", async () => { + const runs = new AgentRunState(); + const run = runs.createPendingRun("agent"); + const wait = track(runs.waitForForegroundStart("agent")); + + runs.settleForegroundStart(run, { status: "failed", error: "start failed" }); + await flushMicrotasks(); + expect(wait.resolved).toBe(false); + expect(runs.hasPendingRun("agent")).toBe(true); + + runs.settleForegroundRun("agent", run.token); + await flushMicrotasks(); + expect(wait.resolved).toBe(true); + }); + + test("resolves at once when no turn is starting", async () => { + const runs = new AgentRunState(); + await expect(runs.waitForForegroundStart("agent")).resolves.toBeUndefined(); + }); +}); diff --git a/packages/server/src/server/agent/agent-run-state.ts b/packages/server/src/server/agent/agent-run-state.ts index ad1eb0af34..0ea85f6d81 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; + turnStartedPromise: Promise; + resolveTurnStarted: () => void; settled: boolean; settledPromise: Promise; resolveSettled: () => void; @@ -54,6 +58,28 @@ export class AgentRunState { return run?.kind === "foreground" ? run : null; } + settleForegroundStart( + run: PendingForegroundRun, + start: Exclude, + ): void { + run.start = start; + if (start.status === "started") { + run.resolveTurnStarted(); + } + } + + /** + * Resolves once the agent's foreground run has a turn id or has ended. A failed start + * resolves only once its run is cleared, so callers never see a failed run still tracked. + */ + async waitForForegroundStart(agentId: string): Promise { + const run = this.getPendingRun(agentId); + if (run?.start.status !== "pending") { + return; + } + await Promise.race([run.turnStartedPromise, run.settledPromise]); + } + hasPendingRun(agentId: string): boolean { return this.getPendingRun(agentId) !== null; } @@ -277,10 +303,16 @@ export class ForegroundTurnStream { } function createPendingForegroundRun(): PendingForegroundRun { + let resolveTurnStarted!: () => void; + const turnStartedPromise = new Promise((resolvePromise) => { + resolveTurnStarted = resolvePromise; + }); return { ...createTrackedRunState(), kind: "foreground", start: { status: "pending" }, + turnStartedPromise, + resolveTurnStarted, stagedEvents: [], }; } diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index 2169a059f2..e709fd490c 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts @@ -226,6 +226,64 @@ function deferred() { } describe("Codex active-turn steering admission", () => { + test("steers an accepted turn whose native id arrives after startTurn returns", async () => { + const appServer = createFakeCodexAppServer({ + "turn/steer": () => ({ turn: { id: "native-A" } }), + }); + const session = new CodexAppServerAgentSession( + createConfig({ cwd: "/workspace/project" }), + null, + createTestLogger(), + async () => appServer.child, + ); + try { + const { turnId } = await session.startTurn("FIRST"); + const steer = session.steerActiveTurn("SECOND", { expectedTurnId: turnId }); + appServer.startsTurn({ threadId: "thread-1", turnId: "native-A" }); + + await expect(steer).resolves.toEqual({ status: "accepted" }); + expect( + appServer.requests().filter((request) => request.method === "turn/start"), + ).toHaveLength(1); + expect(appServer.requests().filter((request) => request.method === "turn/interrupt")).toEqual( + [], + ); + expect(appServer.requests().find((request) => request.method === "turn/steer")).toMatchObject( + { + params: { + threadId: "thread-1", + expectedTurnId: "native-A", + input: [{ type: "text", text: "SECOND" }], + }, + }, + ); + appServer.assertNoErrors(); + } finally { + await session.close(); + } + }); + + test("releases a waiting steer when the accepted turn ends before identification", async () => { + const appServer = createFakeCodexAppServer(); + const session = new CodexAppServerAgentSession( + createConfig({ cwd: "/workspace/project" }), + null, + createTestLogger(), + async () => appServer.child, + ); + try { + const { turnId } = await session.startTurn("FIRST"); + const steer = session.steerActiveTurn("SECOND", { expectedTurnId: turnId }); + appServer.completeTurn(); + + await expect(steer).resolves.toEqual({ status: "unavailable" }); + expect(appServer.requests().filter((request) => request.method === "turn/steer")).toEqual([]); + appServer.assertNoErrors(); + } finally { + await session.close(); + } + }); + test("a steer without the clearing contract leaves permissions open", async () => { const appServer = createFakeCodexAppServer({ "turn/steer": () => ({ turn: { id: "native-A" } }), diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index cdad017d45..5e43f582c5 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -4456,9 +4456,18 @@ export class CodexAppServerAgentSession implements AgentSession { ): Promise { const client = this.client; const threadId = this.currentThreadId; - const nativeTurnId = this.currentTurnId; + let nativeTurnId = this.currentTurnId; const foregroundTurnId = this.activeForegroundTurnId; - if (!client || !threadId || !nativeTurnId || foregroundTurnId !== options.expectedTurnId) { + if (!client || !threadId || foregroundTurnId !== options.expectedTurnId) { + return { status: "unavailable" }; + } + const pendingIdentification = this.pendingForegroundTurnIdentification; + // The start acknowledgement precedes the notification that identifies the native + // turn. Keep this steer attached to the accepted foreground turn while it arrives. + if (!nativeTurnId && pendingIdentification?.foregroundTurnId === foregroundTurnId) { + nativeTurnId = await pendingIdentification.promise; + } + if (!nativeTurnId) { return { status: "unavailable" }; } if (await this.resolveSlashCommandInvocation(prompt)) return { status: "unavailable" }; diff --git a/packages/server/src/server/daemon-e2e/codex-steer-starting.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/codex-steer-starting.real.e2e.test.ts new file mode 100644 index 0000000000..60508205f2 --- /dev/null +++ b/packages/server/src/server/daemon-e2e/codex-steer-starting.real.e2e.test.ts @@ -0,0 +1,80 @@ +import { expect, test } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import pino from "pino"; + +import { CodexAppServerAgentClient } from "../agent/providers/codex-app-server-agent.js"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import { createMessageCollector } from "../test-utils/message-collector.js"; +import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; +import { canRunRealProvider } from "./real-provider-test-config.js"; + +test("keeps two concurrent startup steer messages in one real Codex turn", async (context) => { + if (!(await canRunRealProvider("codex"))) { + context.skip(); + return; + } + const cwd = mkdtempSync(path.join(tmpdir(), "codex-steer-starting-")); + const logger = pino({ level: "silent" }); + const daemon = await createTestPaseoDaemon({ + agentClients: { codex: new CodexAppServerAgentClient(logger) }, + logger, + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + const collector = createMessageCollector(client); + const subscriptions: ReturnType[] = []; + try { + await client.connect(); + await client.fetchAgents({ subscribe: {} }); + for (let attempt = 0; attempt < 3; attempt += 1) { + const agent = await client.createAgent({ + cwd, + provider: "codex", + modeId: "full-access", + }); + const subscription = client.subscribeAgentTimeline(agent.id, () => {}); + subscriptions.push(subscription); + await subscription.ready; + await Promise.all([ + client.sendAgentMessage(agent.id, "Reply exactly: FIRST", { + messageId: `first-${attempt}`, + activeTurnBehavior: "steer", + }), + client.sendAgentMessage(agent.id, "Reply exactly: SECOND", { + messageId: `second-${attempt}`, + activeTurnBehavior: "steer", + }), + ]); + await client.waitForFinish(agent.id, 90_000); + const timeline = await client.fetchAgentTimeline(agent.id, { limit: 100 }); + const users = timeline.entries.filter((entry) => entry.item.type === "user_message"); + const events = collector.messages.flatMap((message) => + message.type === "agent_stream" && message.payload.agentId === agent.id + ? [message.payload.event] + : [], + ); + expect(users.map((entry) => entry.item.text).sort()).toEqual([ + "Reply exactly: FIRST", + "Reply exactly: SECOND", + ]); + expect(new Set(users.map((entry) => entry.turnId)).size).toBe(1); + expect(events.filter((event) => event.type === "turn_started")).toHaveLength(1); + expect( + events.filter((event) => event.type === "turn_canceled" || event.type === "turn_failed"), + ).toHaveLength(0); + expect( + timeline.entries.some( + (entry) => + entry.item.type === "assistant_message" && entry.item.text.includes("[System Error]"), + ), + ).toBe(false); + } + } finally { + await Promise.all(subscriptions.map((subscription) => subscription.release())); + collector.unsubscribe(); + await client.close(); + await daemon.close(); + rmSync(cwd, { recursive: true, force: true }); + } +}, 300_000);