fix(core): report a dequeued prompt that teardown stops before the agent sees it

The drain announces a prompt as submitted the moment it leaves the
queue, then sends it. If teardown starts while that turn is still
preparing, the pre-run guard stops it without it ever reaching the
agent, and discardQueue cannot report it because it is no longer in
the queue: clients that offer discarded text back missed exactly that
prompt, and the discard count was one short. The drain now hands the
dequeued entry to the turn, and the guard reports it as discarded on
the same pending_prompts event and telemetry as a queue discard.
This commit is contained in:
Mikołaj Kondratek
2026-10-23 12:04:53 +02:00
parent 5b20862b4e
commit 2353b7c4cc
3 changed files with 102 additions and 27 deletions
@@ -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(
@@ -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<AgentResult | undefined> {
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<AgentResult | undefined> {
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<AgentResult> {
// 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<AgentResult> {
// 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,
@@ -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<unknown>;
}
@@ -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