fix(coding-agent): settle extension command turns before returning

packages/coding-agent/docs/extensions.md

@@ -1440,6 +1440,8 @@ pi.sendMessage({
 
 Send a user message to the agent. Unlike `sendMessage()` which sends custom messages, this sends an actual user message that appears as if typed by the user. Always triggers a turn.
 
+The call is fire-and-forget: it returns before the turn finishes. Inside a command handler, a turn started this way is still awaited before `prompt()` resolves, so a host that disposes right after a prompt (print and JSON mode) does not tear the session down mid-turn. Await the call yourself when you need its result before continuing. The same applies to `sendMessage(..., { triggerTurn: true })`.
+
 ```typescript
 // Simple text message
 pi.sendUserMessage("What is 2+2?");

packages/coding-agent/docs/sdk.md

@@ -198,6 +198,8 @@ interface PromptOptions {
 
 It fires before `prompt()` resolves. `prompt()` still resolves only after the full accepted run finishes, including retries. Failures after acceptance are reported through the normal event and message stream, not through `preflightResult(false)`.
 
+A slash command that hands work to the agent (for example `/plan <task>`) is accepted as soon as its handler returns, so `preflightResult(true)` fires immediately, while the turn it started is still running. `prompt()` still resolves only after that turn settles, so a caller may dispose the session once it returns.
+
 The `prompt()` method handles prompt templates, extension commands, and message sending:
 
 ```typescript

packages/coding-agent/src/core/agent-session.ts

@@ -406,6 +406,22 @@ export class AgentSession {
 	private _extensionShutdownHandler?: ShutdownHandler;
 	private _extensionErrorListener?: ExtensionErrorListener;
 	private _extensionErrorUnsubscriber?: () => void;
+	/**
+	 * Open extension command handlers and the turns each one started.
+	 *
+	 * `pi.sendUserMessage` and `pi.sendMessage({ triggerTurn })` are
+	 * fire-and-forget bindings, so a command like `/plan <task>` returns while its
+	 * turn is still running. `prompt()` must settle that turn before reporting
+	 * completion: a host that tears down afterwards (print/JSON mode) would
+	 * otherwise dispose the session mid-turn and every in-flight extension would
+	 * observe a stale context.
+	 *
+	 * A stack, not a single slot: two handlers can overlap (one parked on a
+	 * dialog while another is dispatched), and they need not finish in the order
+	 * they started. Every open tracker is charged, so a turn can be over-awaited
+	 * by an unrelated command but never lost — losing it is the bug being fixed.
+	 */
+	private _openCommandTurns: Promise<unknown>[][] = [];
 
 	private _modelRuntime: ModelRuntime;
 	private _agentDir?: string;
@@ -1810,10 +1826,11 @@ export class AgentSession {
 			// Handle extension commands first (execute immediately, even during streaming)
 			// Extension commands manage their own LLM interaction via pi.sendMessage()
 			if (expandPromptTemplates && text.startsWith("/")) {
-				const handled = await this._tryExecuteExtensionCommand(text);
+				// A command is accepted as soon as its handler returns for prompt
+				// purposes, even though the turn it started may still be running.
+				const handled = await this._tryExecuteExtensionCommand(text, () => preflightResult?.(true));
 				if (handled) {
 					// Extension command executed, no prompt to send
-					preflightResult?.(true);
 					return;
 				}
 			}
@@ -1970,32 +1987,52 @@ export class AgentSession {
 	/**
 	 * Try to execute an extension command. Returns true if command was found and executed.
 	 */
-	private async _tryExecuteExtensionCommand(text: string): Promise<boolean> {
+	private async _tryExecuteExtensionCommand(text: string, onAccepted?: () => void): Promise<boolean> {
 		// Parse command name and args
 		const spaceIndex = text.indexOf(" ");
 		const commandName = spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex);
 		const args = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1);
 
-		if (await this._handleLiveContextCommand(commandName, args)) return true;
+		if (await this._handleLiveContextCommand(commandName, args)) {
+			onAccepted?.();
+			return true;
+		}
 
 		const command = this._extensionRunner.getCommand(commandName);
 		if (!command) return false;
 
 		// Get command context from extension runner (includes session control methods)
 		const ctx = this._extensionRunner.createCommandContext();
 
+		const turns: Promise<unknown>[] = [];
+		this._openCommandTurns.push(turns);
 		try {
 			await command.handler(args, ctx);
-			return true;
 		} catch (err) {
 			// Emit error via extension runner
 			this._extensionRunner.emitError({
 				extensionPath: `command:${commandName}`,
 				event: "command",
 				error: err instanceof Error ? err.message : String(err),
 			});
-			return true;
+		} finally {
+			const index = this._openCommandTurns.indexOf(turns);
+			if (index !== -1) this._openCommandTurns.splice(index, 1);
 		}
+
+		// The handler may have handed work to the agent and returned. Acceptance is
+		// known now — a caller gating on it (RPC) must not wait for the turn — but
+		// completion is not, so settle before letting the caller tear the session
+		// down. Each binding already reports its own failure; settling must not
+		// rethrow or mask the command's result.
+		onAccepted?.();
+		if (turns.length > 0) await Promise.allSettled(turns);
+		return true;
+	}
+
+	/** Record a turn started by any open command handler, so its prompt can wait. */
+	private _trackDeferredCommandTurn(turn: Promise<unknown>): void {
+		for (const tracker of this._openCommandTurns) tracker.push(turn);
 	}
 
 	/**
@@ -3284,7 +3321,9 @@ export class AgentSession {
 		runner.bindCore(
 			{
 				sendMessage: (message, options) => {
-					this.sendCustomMessage(message, options).catch((err) => {
+					const turn = this.sendCustomMessage(message, options);
+					this._trackDeferredCommandTurn(turn);
+					turn.catch((err) => {
 						runner.emitError({
 							extensionPath: "<runtime>",
 							event: "send_message",
@@ -3293,7 +3332,9 @@ export class AgentSession {
 					});
 				},
 				sendUserMessage: (content, options) => {
-					this.sendUserMessage(content, options).catch((err) => {
+					const turn = this.sendUserMessage(content, options);
+					this._trackDeferredCommandTurn(turn);
+					turn.catch((err) => {
 						runner.emitError({
 							extensionPath: "<runtime>",
 							event: "send_user_message",

packages/coding-agent/test/suite/regressions/command-deferred-turn.test.ts

@@ -0,0 +1,236 @@
+import { fauxAssistantMessage } from "@step-harness/providers";
+import { expect, it, vi } from "vitest";
+import { createHarness } from "../harness.ts";
+
+/**
+ * `pi.sendUserMessage` and `pi.sendMessage({ triggerTurn })` are fire-and-forget
+ * bindings: an extension command handler that hands work to the agent returns
+ * before that turn finishes. `prompt()` used to report the command as complete
+ * at that point, so a host that disposes right after `prompt()` — print and JSON
+ * mode both do — tore the session down mid-turn. Every in-flight extension then
+ * failed with "stale after session replacement or reload", and the task never ran.
+ *
+ * The command is still *accepted* as soon as its handler returns: the acceptance
+ * signal must not wait for the turn, or an RPC client gating on it stalls.
+ */
+
+it("waits for a turn an extension command started before reporting completion", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("deferred-task", {
+					// `async` without awaiting: the handler returns while the turn runs,
+					// which is exactly how `/plan <task>` hands the task to the session.
+					handler: async (_args, _ctx) => {
+						pi.sendUserMessage("Do the deferred task.");
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("deferred task done")]);
+
+		await harness.session.prompt("/deferred-task");
+
+		// The command's own prompt must not resolve until its turn has settled.
+		const assistant = harness.session.messages.filter((message) => message.role === "assistant");
+		expect(assistant).toHaveLength(1);
+		expect(JSON.stringify(assistant[0]?.content)).toContain("deferred task done");
+		expect(harness.getPendingResponseCount()).toBe(0);
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("waits for a custom message that triggers a turn", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("deferred-custom", {
+					handler: async (_args, _ctx) => {
+						pi.sendMessage(
+							{ customType: "deferred", content: "Do the deferred task.", display: false },
+							{ triggerTurn: true },
+						);
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("custom task done")]);
+
+		await harness.session.prompt("/deferred-custom");
+
+		const assistant = harness.session.messages.filter((message) => message.role === "assistant");
+		expect(assistant).toHaveLength(1);
+		expect(JSON.stringify(assistant[0]?.content)).toContain("custom task done");
+		expect(harness.getPendingResponseCount()).toBe(0);
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("accepts the command before the turn it started has finished", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("deferred-ack", {
+					handler: async (_args, _ctx) => {
+						pi.sendUserMessage("Do the deferred task.");
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("done")]);
+
+		let acceptedAtAssistantCount: number | undefined;
+		await harness.session.prompt("/deferred-ack", {
+			preflightResult: (accepted) => {
+				if (!accepted) return;
+				// RPC answers the client from this callback, so it must fire before the
+				// deferred turn finishes rather than after it.
+				acceptedAtAssistantCount = harness.session.messages.filter(
+					(message) => message.role === "assistant",
+				).length;
+			},
+		});
+
+		expect(acceptedAtAssistantCount).toBe(0);
+		// The turn did run, and only settled by the time the prompt resolved.
+		expect(harness.session.messages.filter((message) => message.role === "assistant")).toHaveLength(1);
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("does not delay a command that starts no turn", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("noop", { handler: async (_args, _ctx) => {} });
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("must not be requested")]);
+
+		let accepted = false;
+		const markAccepted = (succeeded: boolean): void => {
+			if (succeeded) accepted = true;
+		};
+		await harness.session.prompt("/noop", { preflightResult: markAccepted });
+
+		expect(accepted).toBe(true);
+		expect(harness.session.messages.filter((message) => message.role === "assistant")).toHaveLength(0);
+		expect(harness.getPendingResponseCount()).toBe(1);
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("keeps a failed deferred turn from failing the command itself", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("deferred-failure", {
+					handler: async (_args, _ctx) => {
+						void pi.sendUserMessage("Do the deferred task.");
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		const errors: string[] = [];
+		await harness.session.bindExtensions({ onError: (error) => errors.push(String(error.error)) });
+		// A preflight refusal (no model, compaction in progress) rejects this call.
+		vi.spyOn(harness.session, "sendUserMessage").mockRejectedValue(new Error("preflight refused"));
+
+		// The command itself was handled, so waiting on a rejected turn must not
+		// turn into a rejected prompt.
+		await expect(harness.session.prompt("/deferred-failure")).resolves.toBeUndefined();
+		// The failure is still reported through the ordinary extension error path.
+		expect(errors.join("|")).toContain("preflight refused");
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("still waits when two command handlers overlap and finish out of order", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				// `/slow` parks before deferring; `/quick` defers immediately but
+				// settles later, so the two trackers overlap and unwind out of order.
+				pi.registerCommand("slow", {
+					handler: async (_args, _ctx) => {
+						await new Promise((resolve) => setTimeout(resolve, 20));
+						pi.sendUserMessage("slow task");
+					},
+				});
+				pi.registerCommand("quick", {
+					handler: async (_args, _ctx) => {
+						await new Promise((resolve) => setTimeout(resolve, 40));
+						pi.sendUserMessage("quick task");
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("slow done"), () => fauxAssistantMessage("quick done")]);
+
+		const slow = harness.session.prompt("/slow");
+		const quick = harness.session.prompt("/quick");
+		await Promise.all([slow, quick]);
+
+		// Both turns ran; neither escaped its command's wait.
+		expect(harness.session.messages.filter((message) => message.role === "assistant")).toHaveLength(2);
+		expect(harness.getPendingResponseCount()).toBe(0);
+	} finally {
+		harness.cleanup();
+	}
+});
+
+it("waits for a nested command's deferred turn", async () => {
+	const harness = await createHarness({
+		extensionFactories: [
+			(pi) => {
+				pi.registerCommand("inner", {
+					handler: async (_args, _ctx) => {
+						pi.sendUserMessage("inner task");
+					},
+				});
+				pi.registerCommand("outer", {
+					handler: async (_args, ctx) => {
+						// A nested command gets its own tracker and settles before the
+						// outer handler continues.
+						await harness.session.prompt("/inner");
+						void ctx;
+					},
+				});
+			},
+		],
+	});
+
+	try {
+		harness.setResponses([() => fauxAssistantMessage("inner done")]);
+
+		await harness.session.prompt("/outer");
+
+		expect(harness.session.messages.filter((message) => message.role === "assistant")).toHaveLength(1);
+		expect(harness.getPendingResponseCount()).toBe(0);
+	} finally {
+		harness.cleanup();
+	}
+});