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 6cfcf1c298..0f657e98b5 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 @@ -3958,13 +3958,22 @@ describe("LocalRuntimeHost", () => { createAgent: () => agent as never, }); let endedEvents = 0; + const submitted: string[] = []; + const discarded: string[] = []; manager.subscribe((event) => { if (event.type === "ended") endedEvents += 1; + if (event.type === "pending_prompt_submitted") { + submitted.push(event.payload.prompt); + } + if (event.type === "pending_prompts" && event.payload.discarded) { + discarded.push(...event.payload.discarded.map((p) => p.prompt)); + } }); + const capture = vi.fn(); await manager.startSession( normalizeStartInput({ - config: createConfig({ sessionId }), + config: createConfig({ sessionId, telemetry: { capture } as never }), interactive: true, }), ); @@ -3998,6 +4007,19 @@ describe("LocalRuntimeHost", () => { expect(sentPrompts).toHaveLength(1); expect(endedEvents).toBe(1); expect(statuses()).toEqual(["running", "idle", "running", "cancelled"]); + // The drain had already announced the prompt as submitted; the client + // must learn that it never reached the agent. + expect(submitted).toEqual(["drained before stop"]); + expect(discarded).toEqual(["drained before stop"]); + expect(capture).toHaveBeenCalledWith( + expect.objectContaining({ + event: CORE_TELEMETRY_EVENTS.SESSION.PENDING_PROMPTS_DISCARDED, + properties: expect.objectContaining({ + count: 1, + reason: "session_teardown", + }), + }), + ); }); function createOneShotStopFixture( 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 64a049a3eb..9b9db15bcf 100644 --- a/sdk/packages/core/src/runtime/host/local-runtime-host.ts +++ b/sdk/packages/core/src/runtime/host/local-runtime-host.ts @@ -107,7 +107,10 @@ import { type SessionStatus, } from "../../types/common"; import type { CoreSessionConfig } from "../../types/config"; -import type { CoreSessionEvent } from "../../types/events"; +import type { + CoreSessionEvent, + SessionPendingPrompt, +} from "../../types/events"; import type { ActiveSession, PreparedTurnInput } from "../../types/session"; import type { SessionRecord } from "../../types/sessions"; import type { RuntimeCapabilities } from "../capabilities"; @@ -121,7 +124,10 @@ import { } from "../orchestration/runtime-oauth-token-manager"; import type { RuntimeBuilder } from "../orchestration/session-runtime"; import { SessionRuntime } from "../orchestration/session-runtime-orchestrator"; -import { PendingPromptsController } from "../turn-queue/pending-prompt-service"; +import { + type PendingPromptEntry, + PendingPromptsController, +} from "../turn-queue/pending-prompt-service"; import { manifestToSessionRecord } from "./history"; import { AgentEventBridge } from "./local/agent-event-bridge"; import { @@ -281,6 +287,15 @@ function isSessionDirectedPluginEvent(name: string): boolean { ); } +interface TurnInput { + prompt: string; + mode?: SendSessionInput["mode"]; + userImages?: string[]; + userFiles?: string[]; + /** Set when the prompt was dequeued by the drain. */ + drained?: PendingPromptEntry; +} + export class LocalRuntimeHost implements RuntimeHost { public readonly runtimeAddress = undefined; public readonly pendingPrompts: PendingPromptsServiceApi; @@ -411,7 +426,8 @@ export class LocalRuntimeHost implements RuntimeHost { this.pendingPromptsController = new PendingPromptsController({ getSession: (sid) => this.sessions.get(sid), emit: (event) => this.emit(event), - send: (input) => this.runTurn(input), + send: ({ pendingPrompt, ...input }) => + this.runTurnWithOrigin(input, pendingPrompt), }); this.pendingPrompts = { steerFirst: async (input) => @@ -1229,6 +1245,18 @@ export class LocalRuntimeHost implements RuntimeHost { } async runTurn(input: SendSessionInput): Promise { + return this.runTurnWithOrigin(input); + } + + /** + * `drained` is the queue entry this turn was dequeued from, when it was. + * The drain has already announced it as submitted, so a teardown that + * stops the turn before the agent sees it must report it as discarded. + */ + private async runTurnWithOrigin( + input: SendSessionInput, + drained?: PendingPromptEntry, + ): Promise { const session = this.getSessionOrThrow(input.sessionId); // Only the input the queue drain replays carries this marker; any // other input during a queued run is fresh user input. @@ -1282,6 +1310,7 @@ export class LocalRuntimeHost implements RuntimeHost { mode: input.mode, userImages: input.userImages, userFiles: input.userFiles, + drained, }); if (!session.interactive) { await this.finalizeSingleRun(session, result.finishReason); @@ -1917,12 +1946,7 @@ export class LocalRuntimeHost implements RuntimeHost { private async executeTurn( session: ActiveSession, - input: { - prompt: string; - mode?: SendSessionInput["mode"]; - userImages?: string[]; - userFiles?: string[]; - }, + input: TurnInput, ): Promise { // The agent only knows about a run once executeAgentTurn starts it; // the flag covers the preparation before that, so teardown treats a @@ -1942,12 +1966,7 @@ export class LocalRuntimeHost implements RuntimeHost { private async prepareAndRunTurn( session: ActiveSession, - input: { - prompt: string; - mode?: SendSessionInput["mode"]; - userImages?: string[]; - userFiles?: string[]; - }, + input: TurnInput, ): Promise { // An abort that arrived between turns targeted a run that had already // ended; only aborts issued from here on belong to this turn. @@ -1964,11 +1983,7 @@ export class LocalRuntimeHost implements RuntimeHost { // A session torn down before its first turn persisted must not be // materialized on disk by that turn afterwards. if (session.shuttingDown) { - return this.abortedTurnResult(session, { - messages: session.agent.getMessages(), - usage: createInitialAccumulatedUsage(), - endedAt: new Date(), - }); + return this.skipTurnForTeardown(session, input); } await this.ensureSessionPersisted(session); // A seeded session (fork, checkpoint restore, missing-session @@ -1997,11 +2012,7 @@ export class LocalRuntimeHost implements RuntimeHost { // Teardown may have started during the awaits above. It only aborts a // run that is already in progress, so this turn has to stop itself. if (session.shuttingDown) { - return this.abortedTurnResult(session, { - messages: session.agent.getMessages(), - usage: createInitialAccumulatedUsage(), - endedAt: new Date(), - }); + return this.skipTurnForTeardown(session, input); } try { @@ -2122,6 +2133,24 @@ export class LocalRuntimeHost implements RuntimeHost { return this.abortedTurnResult(session, { messages, usage, endedAt }); } + private skipTurnForTeardown( + session: ActiveSession, + input: TurnInput, + ): AgentResult { + if (input.drained) { + this.recordDiscardedPrompts( + session, + this.pendingPromptsController.emitDiscarded(session, [input.drained]), + "session_teardown", + ); + } + return this.abortedTurnResult(session, { + messages: session.agent.getMessages(), + usage: createInitialAccumulatedUsage(), + endedAt: new Date(), + }); + } + private abortedTurnResult( session: ActiveSession, input: { @@ -2687,7 +2716,18 @@ export class LocalRuntimeHost implements RuntimeHost { session: ActiveSession, reason: "queue_abort" | "session_teardown", ): void { - const discarded = this.pendingPromptsController.discardQueue(session); + this.recordDiscardedPrompts( + session, + this.pendingPromptsController.discardQueue(session), + reason, + ); + } + + private recordDiscardedPrompts( + session: ActiveSession, + discarded: SessionPendingPrompt[], + reason: "queue_abort" | "session_teardown", + ): void { if (discarded.length === 0) return; capturePendingPromptsDiscarded(session.config.telemetry, { sessionId: session.sessionId, 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 e97d3b980f..e9f6e983b4 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 @@ -35,6 +35,8 @@ export interface PendingPromptsControllerDeps { mode?: AgentMode; userImages?: string[]; userFiles?: string[]; + /** The dequeued entry, so a turn stopped before it reaches the agent can report it as discarded. */ + pendingPrompt: PendingPromptEntry; }): Promise; } @@ -305,6 +307,16 @@ export class PendingPromptsController { return discarded; } + /** Reports entries that left the queue but were never sent to the agent. */ + emitDiscarded( + session: ActiveSession, + entries: PendingPromptEntry[], + ): SessionPendingPrompt[] { + const discarded = entries.map(snapshotPrompt); + this.emitPrompts(session, discarded); + return discarded; + } + emitPrompts( session: ActiveSession, discarded?: SessionPendingPrompt[], @@ -361,6 +373,7 @@ export class PendingPromptsController { ...(next.mode ? { mode: next.mode } : {}), userImages: next.userImages, userFiles: next.userFiles, + pendingPrompt: next, }); // A turn that resolves with an error finish ran (the prompt is in // the conversation and the error is surfaced), so the entry is not