@oh-my-pi/pi-coding-agent 15.7.5 → 15.7.6

package/src/eval/heartbeat.ts

@@ -0,0 +1,74 @@
+/**
+ * Keepalive for in-flight host-side eval bridge calls.
+ *
+ * The eval watchdog ({@link ../tools/eval IdleTimeout}) caps a cell's `timeout`
+ * as a wall-clock budget on the cell's *own* work, but pauses that budget while
+ * a host-side `agent()`/`parallel()` (via `runSubprocess`) or `llm()` (a single
+ * completion) call is in flight. Those calls are the only thing that re-arms the
+ * watchdog — and they can run for long stretches with **no** status of their own
+ * (a subagent's time-to-first-token on a reasoning model, a long quiet nested
+ * tool, or the entire body of a oneshot `llm()` call). Without a keepalive the
+ * watchdog would mistake that delegated work for the cell stalling and abort it
+ * mid-flight, killing the subagent.
+ *
+ * {@link withBridgeHeartbeat} bridges that gap by emitting a synthetic
+ * {@link EVAL_HEARTBEAT_OP} status event immediately when the call begins and
+ * then on a fixed cadence until it settles. The event rides the same
+ * `emitStatus → onStatus` channel both runtimes already forward, so it re-arms
+ * the watchdog without any new plumbing. The heartbeat is the *sole* signal that
+ * extends the budget: consumers MUST treat it as a pure keepalive — bump the
+ * watchdog and drop it (never persist or render it) — see the executor display
+ * sinks and the eval tool's `onStatus` handler. Every other status event
+ * (compute helpers, `log()`/`phase()`, tool results) counts against the budget.
+ */
+import type { JsStatusEvent } from "./js/shared/types";
+
+/**
+ * Synthetic status op emitted purely to keep the eval idle watchdog alive while
+ * a host-side bridge call is in flight. Carries no payload.
+ */
+export const EVAL_HEARTBEAT_OP = "heartbeat";
+
+/**
+ * Heartbeat cadence. Comfortably below the default 30s idle budget (and the
+ * larger budgets long fanouts run under), so a working bridge call always bumps
+ * the watchdog before it expires, while a genuine stall is still bounded once
+ * the call settles and the heartbeat stops.
+ */
+const HEARTBEAT_INTERVAL_MS = 5_000;
+
+let heartbeatIntervalMs = HEARTBEAT_INTERVAL_MS;
+
+/**
+ * Test seam: override the heartbeat cadence so integration tests can exercise
+ * the keepalive within a sub-second idle budget. Pass no value to restore the
+ * production default.
+ */
+export function setBridgeHeartbeatIntervalMs(ms?: number): void {
+	heartbeatIntervalMs = ms === undefined ? HEARTBEAT_INTERVAL_MS : Math.max(1, Math.floor(ms));
+}
+
+/**
+ * Run {@link operation}, pumping {@link EVAL_HEARTBEAT_OP} status events through
+ * {@link emitStatus} — one immediately, then on a fixed cadence — until it
+ * settles. The immediate beat pauses the watchdog the instant the call begins,
+ * so a bridge call that starts close to the budget edge (after the cell already
+ * spent most of it computing) is not aborted before the first interval tick. A
+ * no-op wrapper when no `emitStatus` sink is wired (the heartbeat would reach
+ * nobody).
+ */
+export async function withBridgeHeartbeat<T>(
+	emitStatus: ((event: JsStatusEvent) => void) | undefined,
+	operation: () => Promise<T>,
+): Promise<T> {
+	if (!emitStatus) return operation();
+	emitStatus({ op: EVAL_HEARTBEAT_OP });
+	const timer = setInterval(() => emitStatus({ op: EVAL_HEARTBEAT_OP }), heartbeatIntervalMs);
+	// Never keep the event loop alive for the heartbeat alone.
+	timer.unref?.();
+	try {
+		return await operation();
+	} finally {
+		clearInterval(timer);
+	}
+}

package/src/eval/js/executor.ts

@@ -1,6 +1,7 @@
 import { DEFAULT_MAX_BYTES, OutputSink } from "../../session/streaming-output";
 import type { ToolSession } from "../../tools";
 import { resolveOutputMaxColumns, resolveOutputSinkHeadBytes } from "../../tools/output-meta";
+import { EVAL_HEARTBEAT_OP } from "../heartbeat";
 import { executeInVmContext, type JsDisplayOutput } from "./context-manager";
 import type { JsStatusEvent } from "./shared/types";
 
@@ -59,11 +60,10 @@ function isTimeoutReason(reason: unknown): boolean {
 	);
 }
 
-function formatJsTimeoutAnnotation(timeoutMs: number | undefined, idle: boolean): string {
-	const suffix = idle ? " of inactivity" : "";
+function formatJsTimeoutAnnotation(timeoutMs: number | undefined): string {
 	if (timeoutMs === undefined) return "Command timed out";
 	const secs = Math.max(1, Math.round(timeoutMs / 1000));
-	return `Command timed out after ${secs} seconds${suffix}`;
+	return `Command timed out after ${secs} seconds`;
 }
 
 export async function executeJs(code: string, options: JsExecutorOptions): Promise<JsResult> {
@@ -85,10 +85,9 @@ export async function executeJs(code: string, options: JsExecutorOptions): Promi
 		options.signal && timeoutSignal
 			? AbortSignal.any([options.signal, timeoutSignal])
 			: (options.signal ?? timeoutSignal);
-	// Idle mode: the eval tool drives cancellation via an idle-aware `signal` and
-	// passes only an inactivity budget. Use it for worker cold-start headroom and
-	// timeout-annotation text; never derive a competing fixed timer from it.
-	const idleMode = legacyTimeoutMs === undefined && options.idleTimeoutMs !== undefined;
+	// The eval tool drives cancellation via an idle-aware `signal` and passes only
+	// an inactivity budget; use it solely as worker cold-start headroom and never
+	// derive a competing fixed timer from it.
 	const acquireBudgetMs = legacyTimeoutMs ?? options.idleTimeoutMs;
 
 	try {
@@ -105,8 +104,13 @@ export async function executeJs(code: string, options: JsExecutorOptions): Promi
 				signal,
 				onText: chunk => outputSink.push(chunk),
 				onDisplay: output => {
+					if (output.type === "status") {
+						// Heartbeats are pure idle-watchdog keepalives: forward them so
+						// the eval tool re-arms its timer, but never store or render them.
+						options.onStatus?.(output.event);
+						if (output.event.op === EVAL_HEARTBEAT_OP) return;
+					}
 					displayOutputs.push(output);
-					if (output.type === "status") options.onStatus?.(output.event);
 				},
 			},
 		});
@@ -127,7 +131,7 @@ export async function executeJs(code: string, options: JsExecutorOptions): Promi
 		if (signal?.aborted || isAbortError(error)) {
 			const timedOut = Boolean(timeoutSignal?.aborted) || isTimeoutReason(options.signal?.reason);
 			if (timedOut) {
-				outputSink.push(formatJsTimeoutAnnotation(legacyTimeoutMs ?? options.idleTimeoutMs, idleMode));
+				outputSink.push(formatJsTimeoutAnnotation(legacyTimeoutMs ?? options.idleTimeoutMs));
 			}
 			const summary = await outputSink.dump();
 			return {

package/src/eval/llm-bridge.ts

@@ -18,6 +18,7 @@ import { extractTextContent, extractToolCall, parseJsonPayload } from "../commit
 import { expandRoleAlias, formatModelString, resolveModelFromString } from "../config/model-resolver";
 import type { ToolSession } from "../tools";
 import { ToolError } from "../tools/tool-errors";
+import { withBridgeHeartbeat } from "./heartbeat";
 import type { JsStatusEvent } from "./js/shared/types";
 
 /** Synthetic bridge name reserved for the `llm()` helper across both runtimes. */
@@ -131,20 +132,25 @@ export async function runEvalLlm(args: unknown, options: EvalLlmBridgeOptions):
 
 	const telemetry = resolveTelemetry(options.session.getTelemetry?.(), options.session.getSessionId?.() ?? undefined);
 
-	const response = await instrumentedCompleteSimple(
-		model,
-		{
-			systemPrompt: system ? [system] : undefined,
-			messages: [{ role: "user", content: [{ type: "text", text: prompt }], timestamp: Date.now() }],
-			tools,
-		},
-		{
-			apiKey,
-			signal: options.signal,
-			reasoning: reasoningForTier(tier, model),
-			toolChoice: schema ? { type: "tool", name: STRUCTURED_TOOL_NAME } : undefined,
-		},
-		{ telemetry, oneshotKind: "eval_llm" },
+	// A oneshot completion emits no status until it returns, so pump a heartbeat
+	// while it runs to keep the eval idle watchdog armed across a slow (e.g.
+	// reasoning-tier) request that would otherwise look like a stalled cell.
+	const response = await withBridgeHeartbeat(options.emitStatus, () =>
+		instrumentedCompleteSimple(
+			model,
+			{
+				systemPrompt: system ? [system] : undefined,
+				messages: [{ role: "user", content: [{ type: "text", text: prompt }], timestamp: Date.now() }],
+				tools,
+			},
+			{
+				apiKey,
+				signal: options.signal,
+				reasoning: reasoningForTier(tier, model),
+				toolChoice: schema ? { type: "tool", name: STRUCTURED_TOOL_NAME } : undefined,
+			},
+			{ telemetry, oneshotKind: "eval_llm" },
+		),
 	);
 
 	if (response.stopReason === "error") {

package/src/eval/py/executor.ts

@@ -5,6 +5,7 @@ import { Settings } from "../../config/settings";
 import { OutputSink } from "../../session/streaming-output";
 import type { ToolSession } from "../../tools";
 import { resolveOutputMaxColumns, resolveOutputSinkHeadBytes } from "../../tools/output-meta";
+import { EVAL_HEARTBEAT_OP } from "../heartbeat";
 import type { JsStatusEvent } from "../js/shared/types";
 import {
 	checkPythonKernelAvailability,
@@ -231,20 +232,18 @@ async function waitForPromiseWithCancellation<T>(
 // Result formatting
 // ---------------------------------------------------------------------------
 
-function formatTimeoutAnnotation(timeoutMs?: number, idle = false): string | undefined {
-	const suffix = idle ? " of inactivity" : "";
+function formatTimeoutAnnotation(timeoutMs?: number): string | undefined {
 	if (timeoutMs === undefined) return "Command timed out";
 	const secs = Math.max(1, Math.round(timeoutMs / 1000));
-	return `Command timed out after ${secs} seconds${suffix}`;
+	return `Command timed out after ${secs} seconds`;
 }
 
-function formatKernelTimeoutAnnotation(timeoutMs: number | undefined, kernelKilled: boolean, idle = false): string {
+function formatKernelTimeoutAnnotation(timeoutMs: number | undefined, kernelKilled: boolean): string {
 	const secs = timeoutMs === undefined ? undefined : Math.max(1, Math.round(timeoutMs / 1000));
-	const suffix = idle ? " of inactivity" : "";
 	if (kernelKilled) {
-		return `eval cell timed out${suffix} and the kernel was unresponsive to interrupt; the kernel has been killed and will be recreated on the next call.`;
+		return "eval cell timed out and the kernel was unresponsive to interrupt; the kernel has been killed and will be recreated on the next call.";
 	}
-	const duration = secs === undefined ? "the configured timeout" : `${secs}s${suffix}`;
+	const duration = secs === undefined ? "the configured timeout" : `${secs}s`;
 	return `eval cell timed out after ${duration}; kernel interrupted but remains running. Reset the kernel via { reset: true } if state appears corrupted.`;
 }
 
@@ -488,16 +487,17 @@ async function executeWithKernel(
 	const displayOutputs: KernelDisplayOutput[] = [];
 	const deadlineMs = getExecutionDeadlineMs(options);
 	let executionTimeoutMs: number | undefined;
-	// Idle mode: the caller (eval tool) drives cancellation via an idle-aware
-	// signal and passes no wall-clock deadline, so annotate timeouts with the
-	// configured inactivity budget rather than a remaining-deadline figure.
-	const idleMode = deadlineMs === undefined && options?.idleTimeoutMs !== undefined;
 
 	// Collect every display output and, for status events, stream them live so
 	// long-running bridge helpers (e.g. `agent()`) surface progress mid-cell.
 	const collectDisplay = (output: KernelDisplayOutput) => {
+		if (output.type === "status") {
+			// Heartbeats are pure idle-watchdog keepalives: forward them so the
+			// eval tool re-arms its timer, but never store or render them.
+			options?.onStatus?.(output.event);
+			if (output.event.op === EVAL_HEARTBEAT_OP) return;
+		}
 		displayOutputs.push(output);
-		if (output.type === "status") options?.onStatus?.(output.event);
 	};
 	const emitStatus = options?.emitStatus ?? ((event: JsStatusEvent) => collectDisplay({ type: "status", event }));
 	const runId = `py-${crypto.randomUUID()}`;
@@ -524,11 +524,7 @@ async function executeWithKernel(
 
 		if (result.cancelled) {
 			const annotation = result.timedOut
-				? formatKernelTimeoutAnnotation(
-						executionTimeoutMs ?? options?.idleTimeoutMs,
-						result.kernelKilled ?? false,
-						idleMode,
-					)
+				? formatKernelTimeoutAnnotation(executionTimeoutMs ?? options?.idleTimeoutMs, result.kernelKilled ?? false)
 				: undefined;
 			return {
 				exitCode: undefined,
@@ -566,7 +562,7 @@ async function executeWithKernel(
 				displayOutputs,
 				stdinRequested: false,
 				...(await sink.dump(
-					timedOut ? formatTimeoutAnnotation(executionTimeoutMs ?? options?.idleTimeoutMs, idleMode) : undefined,
+					timedOut ? formatTimeoutAnnotation(executionTimeoutMs ?? options?.idleTimeoutMs) : undefined,
 				)),
 			};
 		}

package/src/exec/bash-executor.ts

@@ -52,6 +52,27 @@ export interface BashResult {
 
 const shellSessions = new Map<string, Shell>();
 const brokenShellSessions = new Set<string>();
+const shellSessionQuarantines = new Map<string, Promise<unknown>>();
+
+function quarantineShellSession(
+	sessionKey: string,
+	runPromise: Promise<ShellRunResult>,
+	abortCleanupPromise: Promise<void> | undefined,
+): void {
+	brokenShellSessions.add(sessionKey);
+	const cleanup = abortCleanupPromise
+		? Promise.allSettled([runPromise, abortCleanupPromise])
+		: Promise.allSettled([runPromise]);
+	shellSessionQuarantines.set(sessionKey, cleanup);
+	void cleanup
+		.finally(() => {
+			if (shellSessionQuarantines.get(sessionKey) === cleanup) {
+				shellSessionQuarantines.delete(sessionKey);
+				brokenShellSessions.delete(sessionKey);
+			}
+		})
+		.catch(() => undefined);
+}
 
 async function resolveShellCwd(cwd: string | undefined): Promise<string | undefined> {
 	if (!cwd) return undefined;
@@ -134,13 +155,13 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 	}
 	const userSignal = options?.signal;
 	const runAbortController = new AbortController();
+	let abortCleanupPromise: Promise<void> | undefined;
 	const abortCurrentExecution = () => {
 		if (!runAbortController.signal.aborted) {
 			runAbortController.abort();
 		}
-		if (shellSession) {
-			// Native abort is async; fire-and-forget because the caller races the command separately.
-			void shellSession.abort();
+		if (shellSession && !abortCleanupPromise) {
+			abortCleanupPromise = shellSession.abort().catch(() => undefined);
 		}
 	};
 	const abortDeferred = Promise.withResolvers<"abort">();
@@ -209,8 +230,7 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 			acceptingChunks = false;
 			if (shellSession) {
 				resetSession = true;
-				brokenShellSessions.add(sessionKey);
-				void runPromise.finally(() => brokenShellSessions.delete(sessionKey)).catch(() => undefined);
+				quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
 			} else {
 				void runPromise.catch(() => undefined);
 			}
@@ -235,6 +255,9 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 				? `Command timed out after ${Math.round(options.timeout / 1000)} seconds`
 				: "Command timed out";
 			resetSession = true;
+			if (shellSession) {
+				quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
+			}
 			return {
 				exitCode: undefined,
 				cancelled: true,
@@ -245,6 +268,9 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 		// Handle cancellation
 		if (winner.result.cancelled) {
 			resetSession = true;
+			if (shellSession) {
+				quarantineShellSession(sessionKey, runPromise, abortCleanupPromise);
+			}
 			return {
 				exitCode: undefined,
 				cancelled: true,

package/src/extensibility/extensions/loader.ts

@@ -5,7 +5,8 @@ import type * as fs1 from "node:fs";
 import * as fs from "node:fs/promises";
 import * as path from "node:path";
 import type { ThinkingLevel } from "@oh-my-pi/pi-agent-core";
-import type { ImageContent, Model, TextContent } from "@oh-my-pi/pi-ai";
+import type { ImageContent, Model, TextContent, TSchema } from "@oh-my-pi/pi-ai";
+import * as PiCodingAgent from "@oh-my-pi/pi-coding-agent";
 import type { KeyId } from "@oh-my-pi/pi-tui";
 import { hasFsCode, isEacces, isEnoent, logger } from "@oh-my-pi/pi-utils";
 import * as Zod from "zod/v4";
@@ -22,6 +23,7 @@ import * as TypeBox from "../typebox";
 
 import { resolvePath } from "../utils";
 import type {
+	AssistantThinkingRenderer,
 	Extension,
 	ExtensionAPI,
 	ExtensionContext,
@@ -29,6 +31,7 @@ import type {
 	ExtensionRuntime as IExtensionRuntime,
 	LoadExtensionsResult,
 	MessageRenderer,
+	ProviderConfig,
 	RegisteredCommand,
 	ToolDefinition,
 } from "./types";
@@ -55,8 +58,7 @@ export class ExtensionRuntimeNotInitializedError extends Error {
  */
 export class ExtensionRuntime implements IExtensionRuntime {
 	flagValues = new Map<string, boolean | string>();
-	pendingProviderRegistrations: Array<{ name: string; config: import("./types").ProviderConfig; sourceId: string }> =
-		[];
+	pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; sourceId: string }> = [];
 
 	sendMessage(): void {
 		throw new ExtensionRuntimeNotInitializedError();
@@ -123,12 +125,12 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 	readonly flagValues = new Map<string, boolean | string>();
 	readonly pendingProviderRegistrations: Array<{
 		name: string;
-		config: import("./types").ProviderConfig;
+		config: ProviderConfig;
 		sourceId: string;
 	}> = [];
 
 	constructor(
-		public readonly pi: typeof import("@oh-my-pi/pi-coding-agent"),
+		public readonly pi: typeof PiCodingAgent,
 		private readonly extension: Extension,
 		private readonly runtime: IExtensionRuntime,
 		private readonly cwd: string,
@@ -141,10 +143,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 		this.extension.handlers.set(event, list);
 	}
 
-	registerTool<
-		TParams extends import("@oh-my-pi/pi-ai").TSchema = import("@oh-my-pi/pi-ai").TSchema,
-		TDetails = unknown,
-	>(tool: ToolDefinition<TParams, TDetails>): void {
+	registerTool<TParams extends TSchema = TSchema, TDetails = unknown>(tool: ToolDefinition<TParams, TDetails>): void {
 		this.extension.tools.set(tool.name, {
 			definition: tool,
 			extensionPath: this.extension.path,
@@ -190,6 +189,10 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 		this.extension.messageRenderers.set(customType, renderer as MessageRenderer);
 	}
 
+	registerAssistantThinkingRenderer(renderer: AssistantThinkingRenderer): void {
+		this.extension.assistantThinkingRenderers.push(renderer);
+	}
+
 	getFlag(name: string): boolean | string | undefined {
 		if (!this.extension.flags.has(name)) return undefined;
 		return this.runtime.flagValues.get(name);
@@ -253,7 +256,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 		return this.runtime.setSessionName(name);
 	}
 
-	registerProvider(name: string, config: import("./types").ProviderConfig): void {
+	registerProvider(name: string, config: ProviderConfig): void {
 		this.runtime.pendingProviderRegistrations.push({ name, config, sourceId: this.extension.path });
 	}
 }
@@ -267,6 +270,7 @@ function createExtension(extensionPath: string, resolvedPath: string): Extension
 		resolvedPath,
 		handlers: new Map(),
 		tools: new Map(),
+		assistantThinkingRenderers: [],
 		messageRenderers: new Map(),
 		commands: new Map(),
 		flags: new Map(),
@@ -293,13 +297,7 @@ async function loadExtension(
 		}
 
 		const extension = createExtension(extensionPath, resolvedPath);
-		const api = new ConcreteExtensionAPI(
-			await import("@oh-my-pi/pi-coding-agent"),
-			extension,
-			runtime,
-			cwd,
-			eventBus,
-		);
+		const api = new ConcreteExtensionAPI(PiCodingAgent, extension, runtime, cwd, eventBus);
 		await factory(api);
 
 		return { extension, error: null };
@@ -320,7 +318,7 @@ export async function loadExtensionFromFactory(
 	name = "<inline>",
 ): Promise<Extension> {
 	const extension = createExtension(name, name);
-	const api = new ConcreteExtensionAPI(await import("@oh-my-pi/pi-coding-agent"), extension, runtime, cwd, eventBus);
+	const api = new ConcreteExtensionAPI(PiCodingAgent, extension, runtime, cwd, eventBus);
 	await factory(api);
 	return extension;
 }

package/src/extensibility/extensions/runner.ts

@@ -10,6 +10,7 @@ import { type Theme, theme } from "../../modes/theme/theme";
 import type { SessionManager } from "../../session/session-manager";
 import type {
 	AfterProviderResponseEvent,
+	AssistantThinkingRenderer,
 	BeforeAgentStartEvent,
 	BeforeAgentStartEventResult,
 	BeforeProviderRequestEvent,
@@ -343,22 +344,22 @@ export class ExtensionRunner {
 		this.runtime.flagValues.set(name, value);
 	}
 
-	static readonly #RESERVED_SHORTCUTS = new Set([
-		"ctrl+c",
-		"ctrl+d",
-		"ctrl+z",
-		"ctrl+k",
-		"ctrl+p",
-		"ctrl+l",
-		"ctrl+o",
-		"ctrl+t",
-		"ctrl+g",
-		"shift+tab",
-		"shift+ctrl+p",
-		"alt+enter",
-		"escape",
-		"enter",
-	]);
+	static readonly #RESERVED_SHORTCUTS: Record<string, true> = {
+		"ctrl+c": true,
+		"ctrl+d": true,
+		"ctrl+z": true,
+		"ctrl+k": true,
+		"ctrl+p": true,
+		"ctrl+l": true,
+		"ctrl+o": true,
+		"ctrl+t": true,
+		"ctrl+g": true,
+		"shift+tab": true,
+		"shift+ctrl+p": true,
+		"alt+enter": true,
+		escape: true,
+		enter: true,
+	};
 
 	getShortcuts(): Map<KeyId, ExtensionShortcut> {
 		const allShortcuts = new Map<KeyId, ExtensionShortcut>();
@@ -366,7 +367,7 @@ export class ExtensionRunner {
 			for (const [key, shortcut] of ext.shortcuts) {
 				const normalizedKey = key.toLowerCase() as KeyId;
 
-				if (ExtensionRunner.#RESERVED_SHORTCUTS.has(normalizedKey)) {
+				if (ExtensionRunner.#RESERVED_SHORTCUTS[normalizedKey]) {
 					logger.warn("Extension shortcut conflicts with built-in shortcut", {
 						key,
 						extensionPath: shortcut.extensionPath,
@@ -419,6 +420,10 @@ export class ExtensionRunner {
 		return undefined;
 	}
 
+	getAssistantThinkingRenderers(): AssistantThinkingRenderer[] {
+		return this.extensions.flatMap(ext => ext.assistantThinkingRenderers);
+	}
+
 	getRegisteredCommands(reserved?: Set<string>): RegisteredCommand[] {
 		this.#commandDiagnostics = [];
 

package/src/extensibility/extensions/types.ts

@@ -25,6 +25,8 @@ import type {
 import type { OAuthCredentials, OAuthLoginCallbacks } from "@oh-my-pi/pi-ai/utils/oauth/types";
 import type * as piCodingAgent from "@oh-my-pi/pi-coding-agent";
 import type { AutocompleteItem, Component, EditorTheme, KeyId, TUI } from "@oh-my-pi/pi-tui";
+import type { logger as PiLogger } from "@oh-my-pi/pi-utils";
+import type * as Zod from "zod/v4";
 import type { KeybindingsManager } from "../../config/keybindings";
 import type { ModelRegistry } from "../../config/model-registry";
 import type { EditToolDetails } from "../../edit";
@@ -81,6 +83,7 @@ import type {
 	TurnStartEvent,
 } from "../shared-events";
 import type { SlashCommandInfo } from "../slash-commands";
+import type * as TypeBox from "../typebox";
 
 export type { AppKeybinding, KeybindingsManager } from "../../config/keybindings";
 export type { ExecOptions, ExecResult } from "../../exec/exec";
@@ -90,6 +93,17 @@ export type { AgentToolResult, AgentToolUpdateCallback };
 // UI Context
 // ============================================================================
 
+export interface ExtensionUISelectOption {
+	label: string;
+	description?: string;
+}
+
+export type ExtensionUISelectItem = string | ExtensionUISelectOption;
+
+export function getExtensionUISelectOptionLabel(option: ExtensionUISelectItem): string {
+	return typeof option === "string" ? option : option.label;
+}
+
 /**
  * UI dialog options for extensions.
  */
@@ -135,8 +149,12 @@ export type ExtensionWidgetContent = string[] | ExtensionUiComponentFactory | un
 // and may be invoked from event handlers that have already taken the agent
 // loop's lock — hooks intentionally cannot.
 export interface ExtensionUIContext {
-	/** Show a selector and return the user's choice. */
-	select(title: string, options: string[], dialogOptions?: ExtensionUIDialogOptions): Promise<string | undefined>;
+	/** Show a selector and return the selected label, even when an option also includes a description. */
+	select(
+		title: string,
+		options: ExtensionUISelectItem[],
+		dialogOptions?: ExtensionUIDialogOptions,
+	): Promise<string | undefined>;
 
 	/** Show a confirmation dialog. */
 	confirm(title: string, message: string, dialogOptions?: ExtensionUIDialogOptions): Promise<boolean>;
@@ -799,6 +817,18 @@ export type MessageRenderer<T = unknown> = (
 	theme: Theme,
 ) => Component | undefined;
 
+export interface AssistantThinkingRenderContext {
+	contentIndex: number;
+	thinkingIndex: number;
+	text: string;
+	requestRender(): void;
+}
+
+export type AssistantThinkingRenderer = (
+	context: AssistantThinkingRenderContext,
+	theme: Theme,
+) => Component | undefined;
+
 // ============================================================================
 // Command Registration
 // ============================================================================
@@ -830,13 +860,13 @@ export interface ExtensionAPI {
 	// =========================================================================
 
 	/** File logger for error/warning/debug messages */
-	logger: typeof import("@oh-my-pi/pi-utils").logger;
+	logger: typeof PiLogger;
 
 	/** Injected zod-backed typebox shim for legacy `Type.Object(...)` parameter authoring. */
-	typebox: typeof import("../typebox");
+	typebox: typeof TypeBox;
 
 	/** Injected zod module for Zod-authored extension tools (canonical going forward). */
-	zod: typeof import("zod/v4");
+	zod: typeof Zod;
 
 	/** Injected pi-coding-agent exports for accessing SDK utilities */
 	pi: typeof piCodingAgent;
@@ -950,6 +980,9 @@ export interface ExtensionAPI {
 	/** Register a custom renderer for CustomMessageEntry. */
 	registerMessageRenderer<T = unknown>(customType: string, renderer: MessageRenderer<T>): void;
 
+	/** Register a renderer for assistant thinking blocks. Rendered after the original thinking text. */
+	registerAssistantThinkingRenderer(renderer: AssistantThinkingRenderer): void;
+
 	// =========================================================================
 	// Actions
 	// =========================================================================
@@ -1232,6 +1265,7 @@ export interface Extension {
 	label?: string;
 	handlers: Map<string, HandlerFn[]>;
 	tools: Map<string, RegisteredTool<any, any>>;
+	assistantThinkingRenderers: AssistantThinkingRenderer[];
 	messageRenderers: Map<string, MessageRenderer>;
 	commands: Map<string, RegisteredCommand>;
 	flags: Map<string, ExtensionFlag>;