Settle reloads on dispose and fix first-turn screenshot ownership

Builds on the reentrant-reload coalescing already in place:

- reload() returns an outcome (reloaded/coalesced/disposed). When a /reload
  coalesces onto an in-flight reload the reentrancy guard returns before the
  changes are applied, so the command now reports "a reload is already in
  progress" instead of a false success.
- dispose() awaits an in-flight queued reload before tearing down, so the
  print/interactive/action cleanup `finally` blocks can't dispose the runner and
  close the browser mid-reload. Teardown is split into disposeNow() so a shutdown
  raised from inside reload() (an extension calling ctx.shutdown() during
  session_shutdown) tears down without awaiting its own reload — which would
  deadlock.
- The first-turn screenshot is decided once at startup, before extensions load,
  and threaded via skipInitialScreenshot. An extension that sends a user message
  during session_start no longer flips the live prior-turn check and leaves the
  user's real first prompt without the browser frame; a startedUp guard keeps the
  host from consuming the screenshot during startup.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: rgarcia

packages/cli/src/action/harness-runner.ts

@@ -139,24 +139,14 @@ export async function runAction(
 }
 
 async function maybeInitialScreenshot(opts: HarnessRunOptions): Promise<ImageContent[] | undefined> {
+	// `skipInitialScreenshot` is decided once at startup (before extensions load),
+	// so an extension's startup message can't suppress the user's first-turn frame.
 	if (opts.skipInitialScreenshot) return undefined;
-	const hasPriorTurn = await sessionHasPriorTurn(opts.session);
-	if (hasPriorTurn) return undefined;
 	const png = await captureScreenshot(opts.browserHandle.client, opts.browserHandle.browser.session_id);
 	if (!png) return undefined;
 	return [{ type: "image", data: png.toString("base64"), mimeType: "image/png" }];
 }
 
-async function sessionHasPriorTurn(session: Session): Promise<boolean> {
-	const entries = await session.getBranch();
-	for (const entry of entries) {
-		if (entry.type === "message" && (entry.message.role === "user" || entry.message.role === "assistant")) {
-			return true;
-		}
-	}
-	return false;
-}
-
 function textFromAssistant(message: AssistantMessage): string {
 	const parts: string[] = [];
 	for (const block of message.content) {

packages/cli/src/cli-harness.ts

@@ -351,6 +351,13 @@ interface HarnessRuntime {
 	harness: ReturnType<typeof buildCuaHarness>;
 	provider: string;
 	modelRef: CuaModelRef;
+	/**
+	 * Whether the first prompt should skip the initial browser screenshot.
+	 * Decided once here, before extensions load, so an extension that sends a
+	 * user message during startup can't flip the live prior-turn check and leave
+	 * the user's real first prompt without the screenshot.
+	 */
+	skipInitialScreenshot: boolean;
 	/** Loaded pi-extension host. Undefined with --no-extensions or an untrusted project + no global extensions. */
 	host?: HarnessExtensionHost;
 }
@@ -432,6 +439,11 @@ async function setupHarnessRuntime(
 		const png = await captureScreenshot(handle.client, handle.browser.session_id);
 		return png ? [{ type: "image", data: png.toString("base64"), mimeType: "image/png" }] : undefined;
 	};
+	// Decide the first-turn screenshot before extensions load: a resumed session
+	// already has turns, and capturing it now (rather than re-reading the live
+	// transcript at prompt time) keeps an extension's startup sendUserMessage from
+	// flipping the check and leaving the user's real first prompt without it.
+	const skipInitialScreenshot = resolved?.resumed === true || (await sessionHasPriorTurn(session));
 	// A throwing extension load (e.g. a tool name colliding with a base tool)
 	// must not leak the already-provisioned browser handle: the caller's finally
 	// only runs once this returns, so close the handle here before rethrowing.
@@ -459,10 +471,20 @@ async function setupHarnessRuntime(
 		harness,
 		provider,
 		modelRef: auth.modelRef,
+		skipInitialScreenshot,
 		host,
 	};
 }
 
+async function sessionHasPriorTurn(session: Session): Promise<boolean> {
+	const entries = await session.getBranch();
+	return entries.some(
+		(entry) =>
+			entry.type === "message" &&
+			(entry.message.role === "user" || entry.message.role === "assistant"),
+	);
+}
+
 function hasExplicitSessionFlag(flags: HarnessCliFlags): boolean {
 	return (
 		!!flags.sessionRef ||
@@ -515,7 +537,7 @@ export async function runPrintCommand(prompt: string, flags: HarnessCliFlags): P
 			provider: runtime.provider,
 			prompt,
 			skills: runtime.skills,
-			skipInitialScreenshot: runtime.resolved?.resumed === true,
+			skipInitialScreenshot: runtime.skipInitialScreenshot,
 			verbose: flags.verbose,
 			jsonlMode,
 			jsonlIncludeDeltas: flags.jsonlIncludeDeltas,
@@ -560,7 +582,7 @@ export async function runInteractiveCommand(
 			debugTui: flags.debugTui,
 			resumed: runtime.resolved?.resumed === true,
 			transcriptPath: runtime.resolved?.transcriptPath,
-			skipInitialScreenshot: runtime.resolved?.resumed === true,
+			skipInitialScreenshot: runtime.skipInitialScreenshot,
 			host: runtime.host,
 		});
 	} finally {
@@ -596,7 +618,7 @@ export async function runActionCommand(
 			harness: runtime.harness,
 			browserHandle: runtime.handle,
 			session: runtime.session,
-			skipInitialScreenshot: runtime.resolved?.resumed === true,
+			skipInitialScreenshot: runtime.skipInitialScreenshot,
 		}, screenshotOut);
 		return emitCompact(res);
 	} finally {

packages/cli/src/extensions/host.ts

@@ -60,6 +60,16 @@ interface WriteExtensionDetails {
 	hostLoadErrors: Array<{ path: string; error: string }>;
 }
 
+/**
+ * Result of a {@link HarnessExtensionHost.reload} call, so callers like the
+ * `/reload` command can report honestly instead of assuming success:
+ * - `reloaded`  — extensions were re-discovered and re-applied.
+ * - `coalesced` — a reload was already in flight; this request was latched onto
+ *   it rather than run concurrently, so nothing new was applied yet.
+ * - `disposed`  — the host was (or became) torn down, so no reload happened.
+ */
+export type ReloadOutcome = "reloaded" | "coalesced" | "disposed";
+
 const WRITE_EXTENSION_TOOL_NAME = "write_extension";
 
 const WRITE_EXTENSION_DESCRIPTION = [
@@ -146,13 +156,21 @@ export class HarnessExtensionHost {
 	/** Set by `write_extension`; drained into a single reload at the next idle boundary. */
 	private reloadRequested = false;
 	/** The in-flight queued reload, so a drain can await one already running. */
-	private pendingReload: Promise<void> | undefined;
+	private pendingReload: Promise<ReloadOutcome> | undefined;
 	/** Sticky shutdown request raised by `ctx.shutdown()` or owner disposal. */
 	private shutdownRequested = false;
 	/** Guards `dispose` so `ctx.shutdown()` and an owner call don't double-tear-down. */
 	private disposed = false;
+	/** Guards the actual teardown (`disposeNow`); `disposed` is set before the await. */
+	private teardownDone = false;
 	/** True once `load()` has run; guards against double-load and load-after-dispose. */
 	private loaded = false;
+	/**
+	 * False until `load()` finishes emitting the startup `session_start`. While
+	 * false, an extension-initiated `sendUserMessage` does not consume the
+	 * first-turn screenshot, so it can't pre-empt the user's real first prompt.
+	 */
+	private startedUp = false;
 	private sessionName: string | undefined;
 
 	/** Load errors surfaced from the last discover; non-fatal. */
@@ -182,7 +200,9 @@ export class HarnessExtensionHost {
 			getSignal: () => undefined,
 			shutdown: () => this.requestShutdown(),
 		});
-		this.commandActions = makeExtensionCommandContextActions(this.harness, () => this.reload());
+		this.commandActions = makeExtensionCommandContextActions(this.harness, async () => {
+			await this.reload();
+		});
 		this.hostTools = options.selfExtend ? [this.makeWriteExtensionTool()] : [];
 	}
 
@@ -203,6 +223,10 @@ export class HarnessExtensionHost {
 		await this.reapplyTools();
 		this.installBridge();
 		await this.runner?.emit({ type: "session_start", reason: "startup" });
+		// Startup is over: from here an extension sendUserMessage may carry the
+		// first-turn screenshot (it can no longer steal it from the user's first
+		// prompt, which the CLI captured before extensions loaded).
+		this.startedUp = true;
 		// An extension that calls ctx.shutdown() during session_start disposes via
 		// requestShutdown; honor it so load doesn't resolve a torn-down host as ready.
 		if (this.shutdownRequested) await this.dispose();
@@ -215,14 +239,16 @@ export class HarnessExtensionHost {
 	 * the bridge, then emit `session_start`. No extension cache is cleared because
 	 * the loader imports each extension fresh from disk.
 	 */
-	async reload(): Promise<void> {
-		if (this.disposed) return;
+	async reload(): Promise<ReloadOutcome> {
+		if (this.disposed) return "disposed";
 		// Reentrancy guard: a reload triggered (e.g. via ctx.reload()) while one is
 		// in flight must not run concurrently and double-tear-down the bridge. Re-arm
-		// the latch so the in-flight reload's drain picks up the newer request.
+		// the latch so the in-flight reload's loop picks up the newer request, and
+		// report `coalesced` so a caller (e.g. the /reload command) doesn't claim a
+		// completed reload it didn't perform.
 		if (this.reloading) {
 			this.reloadRequested = true;
-			return;
+			return "coalesced";
 		}
 		this.reloading = true;
 		try {
@@ -238,19 +264,19 @@ export class HarnessExtensionHost {
 				// scheduled off-stack at agent_end), so this resolves promptly and
 				// cannot deadlock on an awaited-in-loop reload.
 				await this.harness.waitForIdle();
-				if (this.disposed) return;
+				if (this.disposed) return "disposed";
 				const flags = this.runner?.getFlagValues() ?? new Map<string, boolean | string>();
 				await this.runner?.emit({ type: "session_shutdown", reason: "reload" });
-				if (await this.disposeIfShutdownRequested()) return;
+				if (await this.disposeIfShutdownRequested()) return "disposed";
 				this.teardownBridge?.();
 				this.teardownBridge = undefined;
 				try {
 					await this.buildRunner();
-					if (await this.disposeIfShutdownRequested()) return;
+					if (await this.disposeIfShutdownRequested()) return "disposed";
 					for (const [name, value] of flags) this.runner?.setFlagValue(name, value);
 
 					await this.reapplyTools();
-					if (await this.disposeIfShutdownRequested()) return;
+					if (await this.disposeIfShutdownRequested()) return "disposed";
 					this.installBridge();
 					await this.runner?.emit({ type: "session_start", reason: "reload" });
 				} catch (error) {
@@ -265,7 +291,13 @@ export class HarnessExtensionHost {
 			this.reloading = false;
 		}
 		// Honor a shutdown requested during the final emit, after `reloading` cleared.
-		if (this.shutdownRequested) await this.dispose();
+		// Still inside reload() (pendingReload may point at us), so tear down via
+		// disposeNow rather than dispose to avoid awaiting our own reload.
+		if (this.shutdownRequested) {
+			await this.disposeNow();
+			return "disposed";
+		}
+		return "reloaded";
 	}
 
 	/**
@@ -280,6 +312,29 @@ export class HarnessExtensionHost {
 		if (this.disposed) return;
 		this.shutdownRequested = true;
 		this.disposed = true;
+		// A write_extension reload scheduled off-stack at agent_end may still be in
+		// flight (print/interactive/action cleanup runs in `finally`). Setting
+		// `disposed` stops new reloads and makes the in-flight one bail at its next
+		// await boundary; await it so teardown — and the caller closing the browser —
+		// doesn't race a live reload. Shutdowns raised from inside reload() use
+		// `disposeNow` directly, since awaiting the running reload from within its own
+		// call stack would deadlock.
+		const inFlight = this.pendingReload;
+		if (inFlight) await inFlight.catch(() => {});
+		await this.disposeNow();
+	}
+
+	/**
+	 * The actual teardown, split from `dispose` so reload()'s own shutdown paths
+	 * can run it without awaiting the in-flight reload (which is their call stack).
+	 * Idempotent via `teardownDone` — `disposed` is set before `dispose` awaits, so
+	 * it can't double as the teardown guard.
+	 */
+	private async disposeNow(): Promise<void> {
+		if (this.teardownDone) return;
+		this.teardownDone = true;
+		this.shutdownRequested = true;
+		this.disposed = true;
 		this.teardownBridge?.();
 		this.teardownBridge = undefined;
 		await this.runner?.emit({ type: "session_shutdown", reason: "quit" });
@@ -423,12 +478,14 @@ export class HarnessExtensionHost {
 	 * handler — reloading there would swap the runner out from under the in-flight
 	 * loop and the listener dispatching the event. The `reloading` guard keeps an
 	 * in-flight reload from being re-entered; a write during a reload re-arms the
-	 * latch for the next boundary. `disposed` makes this a no-op during teardown.
+	 * latch, which reload()'s own loop drains before it resolves. `disposed` makes
+	 * this a no-op during teardown.
 	 */
 	async drainPendingReload(): Promise<void> {
 		// Await a reload already running (the bridge fires this fire-and-forget, so
 		// a caller that awaits the drain — e.g. a test asserting the new tool is
-		// live — must observe that reload settle).
+		// live — must observe that reload settle). reload() drains any request
+		// latched mid-reload via its own loop, so a single pass suffices here.
 		if (this.pendingReload) {
 			await this.pendingReload;
 			return;
@@ -564,7 +621,9 @@ export class HarnessExtensionHost {
 	/** Honor a shutdown latched during reload. Returns true if the host disposed. */
 	private async disposeIfShutdownRequested(): Promise<boolean> {
 		if (!this.shutdownRequested && !this.disposed) return false;
-		await this.dispose();
+		// disposeNow, not dispose: this runs inside reload(), whose promise is the
+		// `pendingReload` dispose() would await — awaiting it here would deadlock.
+		await this.disposeNow();
 		return true;
 	}
 
@@ -590,6 +649,9 @@ export class HarnessExtensionHost {
 
 	private async maybeInitialScreenshot(): Promise<ImageContent[] | undefined> {
 		if (!this.initialScreenshot) return undefined;
+		// During startup the user's first prompt owns the first-turn screenshot; an
+		// extension message here must not consume it (see `startedUp`).
+		if (!this.startedUp) return undefined;
 		if (await sessionHasPriorTurn(this.session)) return undefined;
 		return this.initialScreenshot();
 	}

packages/cli/src/print.ts

@@ -93,9 +93,9 @@ export async function runPrint(opts: RunPrintOptions): Promise<number> {
 }
 
 async function maybeInitialScreenshot(opts: RunPrintOptions): Promise<ImageContent[] | undefined> {
+	// `skipInitialScreenshot` is decided once at startup (before extensions load),
+	// so an extension's startup message can't suppress the user's first-turn frame.
 	if (opts.skipInitialScreenshot) return undefined;
-	const hasPriorTurn = await sessionHasPriorTurn(opts.session);
-	if (hasPriorTurn) return undefined;
 	const png = await captureScreenshot(opts.browserHandle.client, opts.browserHandle.browser.session_id);
 	if (!png) return undefined;
 	return [
@@ -106,13 +106,3 @@ async function maybeInitialScreenshot(opts: RunPrintOptions): Promise<ImageConte
 		},
 	];
 }
-
-async function sessionHasPriorTurn(session: Session): Promise<boolean> {
-	const entries = await session.getBranch();
-	for (const entry of entries) {
-		if (entry.type === "message" && (entry.message.role === "user" || entry.message.role === "assistant")) {
-			return true;
-		}
-	}
-	return false;
-}

packages/cli/src/tui/main.ts

@@ -450,23 +450,14 @@ async function maybeInitialScreenshot(
 	firstPromptSent: boolean,
 ): Promise<ImageContent[] | undefined> {
 	if (firstPromptSent) return undefined;
+	// `skipInitialScreenshot` is decided once at startup (before extensions load),
+	// so an extension's startup message can't suppress the user's first-turn frame.
 	if (opts.skipInitialScreenshot) return undefined;
-	if (await sessionHasPriorTurn(opts.session)) return undefined;
 	const png = await captureScreenshot(opts.browserHandle.client, opts.browserHandle.browser.session_id);
 	if (!png) return undefined;
 	return [{ type: "image", data: png.toString("base64"), mimeType: "image/png" }];
 }
 
-async function sessionHasPriorTurn(session: Session): Promise<boolean> {
-	const entries = await session.getBranch();
-	for (const entry of entries) {
-		if (entry.type === "message" && (entry.message.role === "user" || entry.message.role === "assistant")) {
-			return true;
-		}
-	}
-	return false;
-}
-
 async function applyModelCommand(
 	opts: InteractiveOptions,
 	footer: TelemetryFooter,
@@ -540,11 +531,15 @@ export async function applyReloadCommand(opts: InteractiveOptions, messages: Mes
 		// reload() emits no harness event, so this helper is the only source of
 		// feedback; surface loadErrors so a broken edited extension isn't silently
 		// dropped with its tool missing.
-		await opts.host.reload();
-		if (opts.host.isDisposed()) {
+		const outcome = await opts.host.reload();
+		if (outcome === "disposed" || opts.host.isDisposed()) {
 			// An extension calling ctx.shutdown() during the reload tears the host
 			// down; don't claim a successful reload.
 			messages.addNotice("session is shutting down; extensions were not reloaded");
+		} else if (outcome === "coalesced") {
+			// Another reload was already in flight (e.g. a self-extend reload); this
+			// request was latched onto it, so nothing new has been applied yet.
+			messages.addNotice("a reload is already in progress");
 		} else if (opts.host.loadErrors.length > 0) {
 			for (const { path, error } of opts.host.loadErrors) messages.addError(`${path}: ${error}`);
 		} else {