diff --git a/sdk/packages/core/src/runtime/host/local-runtime-host.test.ts b/sdk/packages/core/src/runtime/host/local-runtime-host.test.ts index 0f657e98b5..3f7e8239b1 100644 --- a/sdk/packages/core/src/runtime/host/local-runtime-host.test.ts +++ b/sdk/packages/core/src/runtime/host/local-runtime-host.test.ts @@ -4215,9 +4215,16 @@ describe("LocalRuntimeHost", () => { options: { holdRowCreation?: boolean; holdSecondRunningWrite?: boolean; + holdFirstRunningWrite?: boolean; + resumedStatus?: "completed"; } = {}, ) { - const manifest = createManifest(sessionId); + const manifest = { + ...createManifest(sessionId), + ...(options.resumedStatus + ? { status: options.resumedStatus, interactive: true } + : {}), + }; let releaseHold: (() => void) | undefined; let markHeld: (() => void) | undefined; const held = new Promise((resolve) => { @@ -4240,10 +4247,16 @@ describe("LocalRuntimeHost", () => { }; }), persistSessionMessages: vi.fn(), + ...(options.resumedStatus + ? { readSessionManifest: vi.fn().mockReturnValue(manifest) } + : {}), updateSessionStatus: vi.fn(async (_sessionId: string, status: string) => { if (status === "running") { runningWrites += 1; - if (options.holdSecondRunningWrite && runningWrites === 2) { + if ( + (options.holdSecondRunningWrite && runningWrites === 2) || + (options.holdFirstRunningWrite && runningWrites === 1) + ) { await hold(); } } @@ -4275,8 +4288,12 @@ describe("LocalRuntimeHost", () => { createAgent: () => agent as never, }); const endReasons: string[] = []; + const submitted: string[] = []; manager.subscribe((event) => { if (event.type === "ended") endReasons.push(event.payload.reason); + if (event.type === "pending_prompt_submitted") { + submitted.push(event.payload.prompt); + } }); const statuses = () => sessionService.updateSessionStatus.mock.calls.map((call) => call[1]); @@ -4290,6 +4307,7 @@ describe("LocalRuntimeHost", () => { held, release: () => releaseHold?.(), endReasons, + submitted, statuses, telemetryEvents, }; @@ -4328,6 +4346,51 @@ describe("LocalRuntimeHost", () => { expect(f.agent.continue).not.toHaveBeenCalled(); }); + for (const delivery of [undefined, "queue" as const]) { + it(`queues a prompt sent while another turn is still preparing (delivery: ${delivery ?? "default"}) and runs it afterwards`, async () => { + const sessionId = `sess-send-during-prep-${delivery ?? "default"}`; + const f = createPreparingTurnFixture(sessionId, { + holdFirstRunningWrite: true, + }); + await f.manager.startSession( + normalizeStartInput({ + config: createConfig({ sessionId }), + interactive: true, + }), + ); + // Turn A is parked in its pre-run status write: the agent has no run + // yet, so a second send must not race it to the agent. + const first = f.manager.runTurn({ sessionId, prompt: "first" }); + await f.held; + const second = await f.manager.runTurn({ + sessionId, + prompt: "second", + ...(delivery ? { delivery } : {}), + }); + expect(second).toBeUndefined(); + expect( + (await f.manager.pendingPrompts.list({ sessionId })).map( + (prompt) => prompt.prompt, + ), + ).toEqual(["second"]); + expect(f.agent.run).not.toHaveBeenCalled(); + expect(f.agent.continue).not.toHaveBeenCalled(); + // The drain must leave it alone until A settles: no submitted + // announcement, and no bouncing between drain and the queue. + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(f.submitted).toEqual([]); + + f.release(); + await expect(first).resolves.toMatchObject({ finishReason: "completed" }); + await vi.waitFor(async () => { + expect(f.agent.continue).toHaveBeenCalledTimes(1); + expect(await f.manager.pendingPrompts.list({ sessionId })).toEqual([]); + }); + expect(f.agent.run).toHaveBeenCalledTimes(1); + expect(f.submitted).toEqual(["second"]); + }); + } + it("records cancelled when the host is disposed while a session's next turn is still preparing", async () => { const sessionId = "sess-dispose-while-next-turn-prepares"; const f = createPreparingTurnFixture(sessionId, { diff --git a/sdk/packages/core/src/runtime/host/local-runtime-host.ts b/sdk/packages/core/src/runtime/host/local-runtime-host.ts index 9b9db15bcf..9f1e150a2d 100644 --- a/sdk/packages/core/src/runtime/host/local-runtime-host.ts +++ b/sdk/packages/core/src/runtime/host/local-runtime-host.ts @@ -1064,7 +1064,7 @@ export class LocalRuntimeHost implements RuntimeHost { (startsWithoutTurn ? "idle" : "running"), aborting: false, shuttingDown: false, - turnInFlight: false, + turnsInFlight: 0, interactive: input.interactive === true, persistedMessages: initialMessages, compactionState: initialCompactionState, @@ -1267,7 +1267,9 @@ export class LocalRuntimeHost implements RuntimeHost { if (session.shuttingDown) { throw new Error(`Session ${input.sessionId} is shutting down`); } - const canStartRun = session.agent.canStartRun(); + // A turn still preparing counts as busy, so a second send during that + // window is queued behind it instead of racing it to the agent. + const canStartRun = !this.hasActiveTurn(session); const delivery = input.delivery ?? (session.interactive && !canStartRun ? ("queue" as const) : undefined); @@ -1951,17 +1953,17 @@ export class LocalRuntimeHost implements RuntimeHost { // The agent only knows about a run once executeAgentTurn starts it; // the flag covers the preparation before that, so teardown treats a // turn in that window as active instead of as an idle session. - session.turnInFlight = true; + session.turnsInFlight += 1; try { return await this.prepareAndRunTurn(session, input); } finally { - session.turnInFlight = false; + session.turnsInFlight -= 1; } } /** True while a turn is being prepared or the agent has a run in progress. */ private hasActiveTurn(session: ActiveSession): boolean { - return session.turnInFlight || !session.agent.canStartRun(); + return session.turnsInFlight > 0 || !session.agent.canStartRun(); } private async prepareAndRunTurn( diff --git a/sdk/packages/core/src/runtime/turn-queue/pending-prompt-service.ts b/sdk/packages/core/src/runtime/turn-queue/pending-prompt-service.ts index e9f6e983b4..df4d6a7910 100644 --- a/sdk/packages/core/src/runtime/turn-queue/pending-prompt-service.ts +++ b/sdk/packages/core/src/runtime/turn-queue/pending-prompt-service.ts @@ -337,6 +337,7 @@ export class PendingPromptsController { session.shuttingDown || session.aborting || session.drainingPendingPrompts || + session.turnsInFlight > 0 || !session.agent.canStartRun() ) { return; @@ -356,7 +357,9 @@ export class PendingPromptsController { ) { return; } - if (!session.agent.canStartRun()) { + // A turn that is still preparing has no run the agent knows about yet; + // draining into it would start a second turn. + if (session.turnsInFlight > 0 || !session.agent.canStartRun()) { return; } const { entry: next } = this.service.shiftNext(session); diff --git a/sdk/packages/core/src/types/session.ts b/sdk/packages/core/src/types/session.ts index 99b46e5053..e960c5df34 100644 --- a/sdk/packages/core/src/types/session.ts +++ b/sdk/packages/core/src/types/session.ts @@ -29,8 +29,8 @@ export type ActiveSession = { lastTurnEndedAt?: number; /** Set when teardown starts; the host drops the session once it finishes. */ shuttingDown: boolean; - /** True from a turn's first await until it settles, including the pre-run work before the agent has a run. */ - turnInFlight: boolean; + /** Turns between their first await and settling, including the pre-run work before the agent has a run. */ + turnsInFlight: number; /** The session row being created for the first turn, so teardown can land its status on it. */ persistInFlight?: Promise; interactive: boolean;