refactor(coding-agent): simplify extension runtime architecture

- Replace per-extension closures with shared ExtensionRuntime - Split context actions: ExtensionContextActions (required) + ExtensionCommandContextActions (optional) - Rename LoadedExtension to Extension, remove setter methods - Change runner.initialize() from options object to positional params - Derive hasUI from uiContext presence (no separate param) - Add warning when extensions override built-in tools - RPC and print modes now provide full command context actions BREAKING CHANGE: Extension system types and initialization API changed. See CHANGELOG.md for migration details.
2026-04-15 20:03:05 +00:00 · 2026-01-07 23:50:18 +01:00 · 2026-01-07 23:50:18 +01:00 · cb3ac0ba9e
commit cb3ac0ba9e
parent faa26ffbf9
16 changed files with 580 additions and 736 deletions
--- a/packages/coding-agent/src/core/extensions/index.ts
+++ b/packages/coding-agent/src/core/extensions/index.ts
@ -2,7 +2,12 @@
 * Extension system for lifecycle events and custom tools.
 */

-export { discoverAndLoadExtensions, loadExtensionFromFactory, loadExtensions } from "./loader.js";
+export {
+	createExtensionRuntime,
+	discoverAndLoadExtensions,
+	loadExtensionFromFactory,
+	loadExtensions,
+} from "./loader.js";
 export type { BranchHandler, ExtensionErrorListener, NavigateTreeHandler, NewSessionHandler } from "./runner.js";
 export { ExtensionRunner } from "./runner.js";
 export type {
@ -25,17 +30,23 @@ export type {
 	EditToolResultEvent,
 	ExecOptions,
 	ExecResult,
+	Extension,
+	ExtensionActions,
 	// API
 	ExtensionAPI,
 	ExtensionCommandContext,
+	ExtensionCommandContextActions,
 	// Context
 	ExtensionContext,
+	ExtensionContextActions,
 	// Errors
 	ExtensionError,
 	ExtensionEvent,
 	ExtensionFactory,
 	ExtensionFlag,
 	ExtensionHandler,
+	// Runtime
+	ExtensionRuntime,
 	ExtensionShortcut,
 	ExtensionUIContext,
 	ExtensionUIDialogOptions,
@ -46,8 +57,6 @@ export type {
 	GrepToolResultEvent,
 	KeybindingsManager,
 	LoadExtensionsResult,
-	// Loaded Extension
-	LoadedExtension,
 	LsToolResultEvent,
 	// Message Rendering
 	MessageRenderer,
--- a/packages/coding-agent/src/core/extensions/loader.ts
+++ b/packages/coding-agent/src/core/extensions/loader.ts
@ -10,30 +10,17 @@ import { fileURLToPath } from "node:url";
 import type { KeyId } from "@mariozechner/pi-tui";
 import { createJiti } from "jiti";
 import { getAgentDir, isBunBinary } from "../../config.js";
-import { theme } from "../../modes/interactive/theme/theme.js";
 import { createEventBus, type EventBus } from "../event-bus.js";
 import type { ExecOptions } from "../exec.js";
 import { execCommand } from "../exec.js";
 import type {
-	AppendEntryHandler,
+	Extension,
 	ExtensionAPI,
 	ExtensionFactory,
-	ExtensionFlag,
-	ExtensionShortcut,
-	ExtensionUIContext,
-	GetActiveToolsHandler,
-	GetAllToolsHandler,
-	GetThinkingLevelHandler,
+	ExtensionRuntime,
 	LoadExtensionsResult,
-	LoadedExtension,
 	MessageRenderer,
 	RegisteredCommand,
-	RegisteredTool,
-	SendMessageHandler,
-	SendUserMessageHandler,
-	SetActiveToolsHandler,
-	SetModelHandler,
-	SetThinkingLevelHandler,
 	ToolDefinition,
 } from "./types.js";

@ -84,87 +71,59 @@ function resolvePath(extPath: string, cwd: string): string {
 	return path.resolve(cwd, expanded);
 }

-function createNoOpUIContext(): ExtensionUIContext {
+type HandlerFn = (...args: unknown[]) => Promise<unknown>;
+
+/**
+ * Create a runtime with throwing stubs for action methods.
+ * Runner.initialize() replaces these with real implementations.
+ */
+export function createExtensionRuntime(): ExtensionRuntime {
+	const notInitialized = () => {
+		throw new Error("Extension runtime not initialized. Action methods cannot be called during extension loading.");
+	};
+
 	return {
-		select: async () => undefined,
-		confirm: async () => false,
-		input: async () => undefined,
-		notify: () => {},
-		setStatus: () => {},
-		setWidget: () => {},
-		setFooter: () => {},
-		setHeader: () => {},
-		setTitle: () => {},
-		custom: async () => undefined as never,
-		setEditorText: () => {},
-		getEditorText: () => "",
-		editor: async () => undefined,
-		setEditorComponent: () => {},
-		get theme() {
-			return theme;
-		},
+		sendMessage: notInitialized,
+		sendUserMessage: notInitialized,
+		appendEntry: notInitialized,
+		getActiveTools: notInitialized,
+		getAllTools: notInitialized,
+		setActiveTools: notInitialized,
+		setModel: () => Promise.reject(new Error("Extension runtime not initialized")),
+		getThinkingLevel: notInitialized,
+		setThinkingLevel: notInitialized,
+		flagValues: new Map(),
 	};
 }

-type HandlerFn = (...args: unknown[]) => Promise<unknown>;
-
+/**
+ * Create the ExtensionAPI for an extension.
+ * Registration methods write to the extension object.
+ * Action methods delegate to the shared runtime.
+ */
 function createExtensionAPI(
-	handlers: Map<string, HandlerFn[]>,
-	tools: Map<string, RegisteredTool>,
+	extension: Extension,
+	runtime: ExtensionRuntime,
 	cwd: string,
-	extensionPath: string,
 	eventBus: EventBus,
-	_sharedUI: { ui: ExtensionUIContext; hasUI: boolean },
-): {
-	api: ExtensionAPI;
-	messageRenderers: Map<string, MessageRenderer>;
-	commands: Map<string, RegisteredCommand>;
-	flags: Map<string, ExtensionFlag>;
-	flagValues: Map<string, boolean | string>;
-	shortcuts: Map<KeyId, ExtensionShortcut>;
-	setSendMessageHandler: (handler: SendMessageHandler) => void;
-	setSendUserMessageHandler: (handler: SendUserMessageHandler) => void;
-	setAppendEntryHandler: (handler: AppendEntryHandler) => void;
-	setGetActiveToolsHandler: (handler: GetActiveToolsHandler) => void;
-	setGetAllToolsHandler: (handler: GetAllToolsHandler) => void;
-	setSetActiveToolsHandler: (handler: SetActiveToolsHandler) => void;
-	setSetModelHandler: (handler: SetModelHandler) => void;
-	setGetThinkingLevelHandler: (handler: GetThinkingLevelHandler) => void;
-	setSetThinkingLevelHandler: (handler: SetThinkingLevelHandler) => void;
-	setFlagValue: (name: string, value: boolean | string) => void;
-} {
-	let sendMessageHandler: SendMessageHandler = () => {};
-	let sendUserMessageHandler: SendUserMessageHandler = () => {};
-	let appendEntryHandler: AppendEntryHandler = () => {};
-	let getActiveToolsHandler: GetActiveToolsHandler = () => [];
-	let getAllToolsHandler: GetAllToolsHandler = () => [];
-	let setActiveToolsHandler: SetActiveToolsHandler = () => {};
-	let setModelHandler: SetModelHandler = async () => false;
-	let getThinkingLevelHandler: GetThinkingLevelHandler = () => "off";
-	let setThinkingLevelHandler: SetThinkingLevelHandler = () => {};
-
-	const messageRenderers = new Map<string, MessageRenderer>();
-	const commands = new Map<string, RegisteredCommand>();
-	const flags = new Map<string, ExtensionFlag>();
-	const flagValues = new Map<string, boolean | string>();
-	const shortcuts = new Map<KeyId, ExtensionShortcut>();
-
+): ExtensionAPI {
 	const api = {
+		// Registration methods - write to extension
 		on(event: string, handler: HandlerFn): void {
-			const list = handlers.get(event) ?? [];
+			const list = extension.handlers.get(event) ?? [];
 			list.push(handler);
-			handlers.set(event, list);
+			extension.handlers.set(event, list);
 		},

 		registerTool(tool: ToolDefinition): void {
-			tools.set(tool.name, {
+			extension.tools.set(tool.name, {
 				definition: tool,
-				extensionPath,
+				extensionPath: extension.path,
 			});
 		},

 		registerCommand(name: string, options: { description?: string; handler: RegisteredCommand["handler"] }): void {
-			commands.set(name, { name, ...options });
+			extension.commands.set(name, { name, ...options });
 		},

 		registerShortcut(
@ -174,37 +133,40 @@ function createExtensionAPI(
 				handler: (ctx: import("./types.js").ExtensionContext) => Promise<void> | void;
 			},
 		): void {
-			shortcuts.set(shortcut, { shortcut, extensionPath, ...options });
+			extension.shortcuts.set(shortcut, { shortcut, extensionPath: extension.path, ...options });
 		},

 		registerFlag(
 			name: string,
 			options: { description?: string; type: "boolean" | "string"; default?: boolean | string },
 		): void {
-			flags.set(name, { name, extensionPath, ...options });
+			extension.flags.set(name, { name, extensionPath: extension.path, ...options });
 			if (options.default !== undefined) {
-				flagValues.set(name, options.default);
+				runtime.flagValues.set(name, options.default);
 			}
 		},

-		getFlag(name: string): boolean | string | undefined {
-			return flagValues.get(name);
-		},
-
 		registerMessageRenderer<T>(customType: string, renderer: MessageRenderer<T>): void {
-			messageRenderers.set(customType, renderer as MessageRenderer);
+			extension.messageRenderers.set(customType, renderer as MessageRenderer);
 		},

+		// Flag access - checks extension registered it, reads from runtime
+		getFlag(name: string): boolean | string | undefined {
+			if (!extension.flags.has(name)) return undefined;
+			return runtime.flagValues.get(name);
+		},
+
+		// Action methods - delegate to shared runtime
 		sendMessage(message, options): void {
-			sendMessageHandler(message, options);
+			runtime.sendMessage(message, options);
 		},

 		sendUserMessage(content, options): void {
-			sendUserMessageHandler(content, options);
+			runtime.sendUserMessage(content, options);
 		},

 		appendEntry(customType: string, data?: unknown): void {
-			appendEntryHandler(customType, data);
+			runtime.appendEntry(customType, data);
 		},

 		exec(command: string, args: string[], options?: ExecOptions) {
@ -212,222 +174,86 @@ function createExtensionAPI(
 		},

 		getActiveTools(): string[] {
-			return getActiveToolsHandler();
+			return runtime.getActiveTools();
 		},

 		getAllTools(): string[] {
-			return getAllToolsHandler();
+			return runtime.getAllTools();
 		},

 		setActiveTools(toolNames: string[]): void {
-			setActiveToolsHandler(toolNames);
+			runtime.setActiveTools(toolNames);
 		},

 		setModel(model) {
-			return setModelHandler(model);
+			return runtime.setModel(model);
 		},

 		getThinkingLevel() {
-			return getThinkingLevelHandler();
+			return runtime.getThinkingLevel();
 		},

 		setThinkingLevel(level) {
-			setThinkingLevelHandler(level);
+			runtime.setThinkingLevel(level);
 		},

 		events: eventBus,
 	} as ExtensionAPI;

-	return {
-		api,
-		messageRenderers,
-		commands,
-		flags,
-		flagValues,
-		shortcuts,
-		setSendMessageHandler: (handler: SendMessageHandler) => {
-			sendMessageHandler = handler;
-		},
-		setSendUserMessageHandler: (handler: SendUserMessageHandler) => {
-			sendUserMessageHandler = handler;
-		},
-		setAppendEntryHandler: (handler: AppendEntryHandler) => {
-			appendEntryHandler = handler;
-		},
-		setGetActiveToolsHandler: (handler: GetActiveToolsHandler) => {
-			getActiveToolsHandler = handler;
-		},
-		setGetAllToolsHandler: (handler: GetAllToolsHandler) => {
-			getAllToolsHandler = handler;
-		},
-		setSetActiveToolsHandler: (handler: SetActiveToolsHandler) => {
-			setActiveToolsHandler = handler;
-		},
-		setSetModelHandler: (handler: SetModelHandler) => {
-			setModelHandler = handler;
-		},
-		setGetThinkingLevelHandler: (handler: GetThinkingLevelHandler) => {
-			getThinkingLevelHandler = handler;
-		},
-		setSetThinkingLevelHandler: (handler: SetThinkingLevelHandler) => {
-			setThinkingLevelHandler = handler;
-		},
-		setFlagValue: (name: string, value: boolean | string) => {
-			flagValues.set(name, value);
-		},
-	};
+	return api;
 }

-async function loadExtensionWithBun(
-	resolvedPath: string,
-	cwd: string,
-	extensionPath: string,
-	eventBus: EventBus,
-	sharedUI: { ui: ExtensionUIContext; hasUI: boolean },
-): Promise<{ extension: LoadedExtension | null; error: string | null }> {
-	try {
-		const module = await import(resolvedPath);
-		const factory = (module.default ?? module) as ExtensionFactory;
+async function loadBun(path: string) {
+	const module = await import(path);
+	const factory = (module.default ?? module) as ExtensionFactory;
+	return typeof factory !== "function" ? undefined : factory;
+}

-		if (typeof factory !== "function") {
-			return { extension: null, error: "Extension must export a default function" };
-		}
+async function loadJiti(path: string) {
+	const jiti = createJiti(import.meta.url, {
+		alias: getAliases(),
+	});

-		const handlers = new Map<string, HandlerFn[]>();
-		const tools = new Map<string, RegisteredTool>();
-		const {
-			api,
-			messageRenderers,
-			commands,
-			flags,
-			flagValues,
-			shortcuts,
-			setSendMessageHandler,
-			setSendUserMessageHandler,
-			setAppendEntryHandler,
-			setGetActiveToolsHandler,
-			setGetAllToolsHandler,
-			setSetActiveToolsHandler,
-			setSetModelHandler,
-			setGetThinkingLevelHandler,
-			setSetThinkingLevelHandler,
-			setFlagValue,
-		} = createExtensionAPI(handlers, tools, cwd, extensionPath, eventBus, sharedUI);
+	const module = await jiti.import(path, { default: true });
+	const factory = module as ExtensionFactory;
+	return typeof factory !== "function" ? undefined : factory;
+}

-		await factory(api);
-
-		return {
-			extension: {
-				path: extensionPath,
-				resolvedPath,
-				handlers,
-				tools,
-				messageRenderers,
-				commands,
-				flags,
-				flagValues,
-				shortcuts,
-				setSendMessageHandler,
-				setSendUserMessageHandler,
-				setAppendEntryHandler,
-				setGetActiveToolsHandler,
-				setGetAllToolsHandler,
-				setSetActiveToolsHandler,
-				setSetModelHandler,
-				setGetThinkingLevelHandler,
-				setSetThinkingLevelHandler,
-				setFlagValue,
-			},
-			error: null,
-		};
-	} catch (err) {
-		const message = err instanceof Error ? err.message : String(err);
-
-		if (message.includes("Cannot find module") && message.includes("@mariozechner/")) {
-			return {
-				extension: null,
-				error:
-					`${message}\n` +
-					"Note: Extensions importing from @mariozechner/* packages are not supported in the standalone binary.\n" +
-					"Please install pi via npm: npm install -g @mariozechner/pi-coding-agent",
-			};
-		}
-
-		return { extension: null, error: `Failed to load extension: ${message}` };
-	}
+/**
+ * Create an Extension object with empty collections.
+ */
+function createExtension(extensionPath: string, resolvedPath: string): Extension {
+	return {
+		path: extensionPath,
+		resolvedPath,
+		handlers: new Map(),
+		tools: new Map(),
+		messageRenderers: new Map(),
+		commands: new Map(),
+		flags: new Map(),
+		shortcuts: new Map(),
+	};
 }

 async function loadExtension(
 	extensionPath: string,
 	cwd: string,
 	eventBus: EventBus,
-	sharedUI: { ui: ExtensionUIContext; hasUI: boolean },
-): Promise<{ extension: LoadedExtension | null; error: string | null }> {
+	runtime: ExtensionRuntime,
+): Promise<{ extension: Extension | null; error: string | null }> {
 	const resolvedPath = resolvePath(extensionPath, cwd);

-	if (isBunBinary) {
-		return loadExtensionWithBun(resolvedPath, cwd, extensionPath, eventBus, sharedUI);
-	}
-
 	try {
-		const jiti = createJiti(import.meta.url, {
-			alias: getAliases(),
-		});
-
-		const module = await jiti.import(resolvedPath, { default: true });
-		const factory = module as ExtensionFactory;
-
-		if (typeof factory !== "function") {
-			return { extension: null, error: "Extension must export a default function" };
+		const factory = isBunBinary ? await loadBun(resolvedPath) : await loadJiti(resolvedPath);
+		if (!factory) {
+			return { extension: null, error: `Extension does not export a valid factory function: ${extensionPath}` };
 		}

-		const handlers = new Map<string, HandlerFn[]>();
-		const tools = new Map<string, RegisteredTool>();
-		const {
-			api,
-			messageRenderers,
-			commands,
-			flags,
-			flagValues,
-			shortcuts,
-			setSendMessageHandler,
-			setSendUserMessageHandler,
-			setAppendEntryHandler,
-			setGetActiveToolsHandler,
-			setGetAllToolsHandler,
-			setSetActiveToolsHandler,
-			setSetModelHandler,
-			setGetThinkingLevelHandler,
-			setSetThinkingLevelHandler,
-			setFlagValue,
-		} = createExtensionAPI(handlers, tools, cwd, extensionPath, eventBus, sharedUI);
-
+		const extension = createExtension(extensionPath, resolvedPath);
+		const api = createExtensionAPI(extension, runtime, cwd, eventBus);
 		await factory(api);

-		return {
-			extension: {
-				path: extensionPath,
-				resolvedPath,
-				handlers,
-				tools,
-				messageRenderers,
-				commands,
-				flags,
-				flagValues,
-				shortcuts,
-				setSendMessageHandler,
-				setSendUserMessageHandler,
-				setAppendEntryHandler,
-				setGetActiveToolsHandler,
-				setGetAllToolsHandler,
-				setSetActiveToolsHandler,
-				setSetModelHandler,
-				setGetThinkingLevelHandler,
-				setSetThinkingLevelHandler,
-				setFlagValue,
-			},
-			error: null,
-		};
+		return { extension, error: null };
 	} catch (err) {
 		const message = err instanceof Error ? err.message : String(err);
 		return { extension: null, error: `Failed to load extension: ${message}` };
@ -435,72 +261,32 @@ async function loadExtension(
 }

 /**
- * Create a LoadedExtension from an inline factory function.
+ * Create an Extension from an inline factory function.
 */
 export async function loadExtensionFromFactory(
 	factory: ExtensionFactory,
 	cwd: string,
 	eventBus: EventBus,
-	sharedUI: { ui: ExtensionUIContext; hasUI: boolean },
-	name = "<inline>",
-): Promise<LoadedExtension> {
-	const handlers = new Map<string, HandlerFn[]>();
-	const tools = new Map<string, RegisteredTool>();
-	const {
-		api,
-		messageRenderers,
-		commands,
-		flags,
-		flagValues,
-		shortcuts,
-		setSendMessageHandler,
-		setSendUserMessageHandler,
-		setAppendEntryHandler,
-		setGetActiveToolsHandler,
-		setGetAllToolsHandler,
-		setSetActiveToolsHandler,
-		setSetModelHandler,
-		setGetThinkingLevelHandler,
-		setSetThinkingLevelHandler,
-		setFlagValue,
-	} = createExtensionAPI(handlers, tools, cwd, name, eventBus, sharedUI);
-
+	runtime: ExtensionRuntime,
+	extensionPath = "<inline>",
+): Promise<Extension> {
+	const extension = createExtension(extensionPath, extensionPath);
+	const api = createExtensionAPI(extension, runtime, cwd, eventBus);
 	await factory(api);
-
-	return {
-		path: name,
-		resolvedPath: name,
-		handlers,
-		tools,
-		messageRenderers,
-		commands,
-		flags,
-		flagValues,
-		shortcuts,
-		setSendMessageHandler,
-		setSendUserMessageHandler,
-		setAppendEntryHandler,
-		setGetActiveToolsHandler,
-		setGetAllToolsHandler,
-		setSetActiveToolsHandler,
-		setSetModelHandler,
-		setGetThinkingLevelHandler,
-		setSetThinkingLevelHandler,
-		setFlagValue,
-	};
+	return extension;
 }

 /**
 * Load extensions from paths.
 */
 export async function loadExtensions(paths: string[], cwd: string, eventBus?: EventBus): Promise<LoadExtensionsResult> {
-	const extensions: LoadedExtension[] = [];
+	const extensions: Extension[] = [];
 	const errors: Array<{ path: string; error: string }> = [];
 	const resolvedEventBus = eventBus ?? createEventBus();
-	const sharedUI = { ui: createNoOpUIContext(), hasUI: false };
+	const runtime = createExtensionRuntime();

 	for (const extPath of paths) {
-		const { extension, error } = await loadExtension(extPath, cwd, resolvedEventBus, sharedUI);
+		const { extension, error } = await loadExtension(extPath, cwd, resolvedEventBus, runtime);

 		if (error) {
 			errors.push({ path: extPath, error });
@ -515,10 +301,7 @@ export async function loadExtensions(paths: string[], cwd: string, eventBus?: Ev
 	return {
 		extensions,
 		errors,
-		setUIContext(uiContext, hasUI) {
-			sharedUI.ui = uiContext;
-			sharedUI.hasUI = hasUI;
-		},
+		runtime,
 	};
 }

--- a/packages/coding-agent/src/core/extensions/runner.ts
+++ b/packages/coding-agent/src/core/extensions/runner.ts
@ -9,32 +9,27 @@ import { theme } from "../../modes/interactive/theme/theme.js";
 import type { ModelRegistry } from "../model-registry.js";
 import type { SessionManager } from "../session-manager.js";
 import type {
-	AppendEntryHandler,
 	BeforeAgentStartEvent,
 	BeforeAgentStartEventResult,
 	ContextEvent,
 	ContextEventResult,
+	Extension,
+	ExtensionActions,
 	ExtensionCommandContext,
+	ExtensionCommandContextActions,
 	ExtensionContext,
+	ExtensionContextActions,
 	ExtensionError,
 	ExtensionEvent,
 	ExtensionFlag,
+	ExtensionRuntime,
 	ExtensionShortcut,
 	ExtensionUIContext,
-	GetActiveToolsHandler,
-	GetAllToolsHandler,
-	GetThinkingLevelHandler,
-	LoadedExtension,
 	MessageRenderer,
 	RegisteredCommand,
 	RegisteredTool,
-	SendMessageHandler,
-	SendUserMessageHandler,
 	SessionBeforeCompactResult,
 	SessionBeforeTreeResult,
-	SetActiveToolsHandler,
-	SetModelHandler,
-	SetThinkingLevelHandler,
 	ToolCallEvent,
 	ToolCallEventResult,
 	ToolResultEventResult,
@ -81,9 +76,9 @@ const noOpUIContext: ExtensionUIContext = {
 };

 export class ExtensionRunner {
-	private extensions: LoadedExtension[];
+	private extensions: Extension[];
+	private runtime: ExtensionRuntime;
 	private uiContext: ExtensionUIContext;
-	private hasUI: boolean;
 	private cwd: string;
 	private sessionManager: SessionManager;
 	private modelRegistry: ModelRegistry;
@ -98,78 +93,60 @@ export class ExtensionRunner {
 	private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });

 	constructor(
-		extensions: LoadedExtension[],
+		extensions: Extension[],
+		runtime: ExtensionRuntime,
 		cwd: string,
 		sessionManager: SessionManager,
 		modelRegistry: ModelRegistry,
 	) {
 		this.extensions = extensions;
+		this.runtime = runtime;
 		this.uiContext = noOpUIContext;
-		this.hasUI = false;
 		this.cwd = cwd;
 		this.sessionManager = sessionManager;
 		this.modelRegistry = modelRegistry;
 	}

-	initialize(options: {
-		getModel: () => Model<any> | undefined;
-		sendMessageHandler: SendMessageHandler;
-		sendUserMessageHandler: SendUserMessageHandler;
-		appendEntryHandler: AppendEntryHandler;
-		getActiveToolsHandler: GetActiveToolsHandler;
-		getAllToolsHandler: GetAllToolsHandler;
-		setActiveToolsHandler: SetActiveToolsHandler;
-		setModelHandler: SetModelHandler;
-		getThinkingLevelHandler: GetThinkingLevelHandler;
-		setThinkingLevelHandler: SetThinkingLevelHandler;
-		newSessionHandler?: NewSessionHandler;
-		branchHandler?: BranchHandler;
-		navigateTreeHandler?: NavigateTreeHandler;
-		isIdle?: () => boolean;
-		waitForIdle?: () => Promise<void>;
-		abort?: () => void;
-		hasPendingMessages?: () => boolean;
-		uiContext?: ExtensionUIContext;
-		hasUI?: boolean;
-	}): void {
-		this.getModel = options.getModel;
-		this.isIdleFn = options.isIdle ?? (() => true);
-		this.waitForIdleFn = options.waitForIdle ?? (async () => {});
-		this.abortFn = options.abort ?? (() => {});
-		this.hasPendingMessagesFn = options.hasPendingMessages ?? (() => false);
+	initialize(
+		actions: ExtensionActions,
+		contextActions: ExtensionContextActions,
+		commandContextActions?: ExtensionCommandContextActions,
+		uiContext?: ExtensionUIContext,
+	): void {
+		// Copy actions into the shared runtime (all extension APIs reference this)
+		this.runtime.sendMessage = actions.sendMessage;
+		this.runtime.sendUserMessage = actions.sendUserMessage;
+		this.runtime.appendEntry = actions.appendEntry;
+		this.runtime.getActiveTools = actions.getActiveTools;
+		this.runtime.getAllTools = actions.getAllTools;
+		this.runtime.setActiveTools = actions.setActiveTools;
+		this.runtime.setModel = actions.setModel;
+		this.runtime.getThinkingLevel = actions.getThinkingLevel;
+		this.runtime.setThinkingLevel = actions.setThinkingLevel;

-		if (options.newSessionHandler) {
-			this.newSessionHandler = options.newSessionHandler;
-		}
-		if (options.branchHandler) {
-			this.branchHandler = options.branchHandler;
-		}
-		if (options.navigateTreeHandler) {
-			this.navigateTreeHandler = options.navigateTreeHandler;
+		// Context actions (required)
+		this.getModel = contextActions.getModel;
+		this.isIdleFn = contextActions.isIdle;
+		this.abortFn = contextActions.abort;
+		this.hasPendingMessagesFn = contextActions.hasPendingMessages;
+
+		// Command context actions (optional, only for interactive mode)
+		if (commandContextActions) {
+			this.waitForIdleFn = commandContextActions.waitForIdle;
+			this.newSessionHandler = commandContextActions.newSession;
+			this.branchHandler = commandContextActions.branch;
+			this.navigateTreeHandler = commandContextActions.navigateTree;
 		}

-		for (const ext of this.extensions) {
-			ext.setSendMessageHandler(options.sendMessageHandler);
-			ext.setSendUserMessageHandler(options.sendUserMessageHandler);
-			ext.setAppendEntryHandler(options.appendEntryHandler);
-			ext.setGetActiveToolsHandler(options.getActiveToolsHandler);
-			ext.setGetAllToolsHandler(options.getAllToolsHandler);
-			ext.setSetActiveToolsHandler(options.setActiveToolsHandler);
-			ext.setSetModelHandler(options.setModelHandler);
-			ext.setGetThinkingLevelHandler(options.getThinkingLevelHandler);
-			ext.setSetThinkingLevelHandler(options.setThinkingLevelHandler);
-		}
-
-		this.uiContext = options.uiContext ?? noOpUIContext;
-		this.hasUI = options.hasUI ?? false;
+		this.uiContext = uiContext ?? noOpUIContext;
 	}

-	getUIContext(): ExtensionUIContext | null {
+	getUIContext(): ExtensionUIContext {
 		return this.uiContext;
 	}

-	getHasUI(): boolean {
-		return this.hasUI;
+	hasUI(): boolean {
+		return this.uiContext !== noOpUIContext;
 	}

 	getExtensionPaths(): string[] {
@ -198,11 +175,7 @@ export class ExtensionRunner {
 	}

 	setFlagValue(name: string, value: boolean | string): void {
-		for (const ext of this.extensions) {
-			if (ext.flags.has(name)) {
-				ext.setFlagValue(name, value);
-			}
-		}
+		this.runtime.flagValues.set(name, value);
 	}

 	private static readonly RESERVED_SHORTCUTS = new Set([
@ -301,7 +274,7 @@ export class ExtensionRunner {
 	private createContext(): ExtensionContext {
 		return {
 			ui: this.uiContext,
-			hasUI: this.hasUI,
+			hasUI: this.hasUI(),
 			cwd: this.cwd,
 			sessionManager: this.sessionManager,
 			modelRegistry: this.modelRegistry,
--- a/packages/coding-agent/src/core/extensions/types.ts
+++ b/packages/coding-agent/src/core/extensions/types.ts
@ -21,7 +21,7 @@ import type { Theme } from "../../modes/interactive/theme/theme.js";
 import type { CompactionPreparation, CompactionResult } from "../compaction/index.js";
 import type { EventBus } from "../event-bus.js";
 import type { ExecOptions, ExecResult } from "../exec.js";
-import type { AppAction, KeybindingsManager } from "../keybindings.js";
+import type { KeybindingsManager } from "../keybindings.js";
 import type { CustomMessage } from "../messages.js";
 import type { ModelRegistry } from "../model-registry.js";
 import type {
@ -742,8 +742,63 @@ export type GetThinkingLevelHandler = () => ThinkingLevel;

 export type SetThinkingLevelHandler = (level: ThinkingLevel) => void;

+/**
+ * Shared state created by loader, used during registration and runtime.
+ * Contains flag values (defaults set during registration, CLI values set after).
+ */
+export interface ExtensionRuntimeState {
+	flagValues: Map<string, boolean | string>;
+}
+
+/**
+ * Action implementations for pi.* API methods.
+ * Provided to runner.initialize(), copied into the shared runtime.
+ */
+export interface ExtensionActions {
+	sendMessage: SendMessageHandler;
+	sendUserMessage: SendUserMessageHandler;
+	appendEntry: AppendEntryHandler;
+	getActiveTools: GetActiveToolsHandler;
+	getAllTools: GetAllToolsHandler;
+	setActiveTools: SetActiveToolsHandler;
+	setModel: SetModelHandler;
+	getThinkingLevel: GetThinkingLevelHandler;
+	setThinkingLevel: SetThinkingLevelHandler;
+}
+
+/**
+ * Actions for ExtensionContext (ctx.* in event handlers).
+ * Required by all modes.
+ */
+export interface ExtensionContextActions {
+	getModel: () => Model<any> | undefined;
+	isIdle: () => boolean;
+	abort: () => void;
+	hasPendingMessages: () => boolean;
+}
+
+/**
+ * Actions for ExtensionCommandContext (ctx.* in command handlers).
+ * Only needed for interactive mode where extension commands are invokable.
+ */
+export interface ExtensionCommandContextActions {
+	waitForIdle: () => Promise<void>;
+	newSession: (options?: {
+		parentSession?: string;
+		setup?: (sessionManager: SessionManager) => Promise<void>;
+	}) => Promise<{ cancelled: boolean }>;
+	branch: (entryId: string) => Promise<{ cancelled: boolean }>;
+	navigateTree: (targetId: string, options?: { summarize?: boolean }) => Promise<{ cancelled: boolean }>;
+}
+
+/**
+ * Full runtime = state + actions.
+ * Created by loader with throwing action stubs, completed by runner.initialize().
+ */
+export interface ExtensionRuntime extends ExtensionRuntimeState, ExtensionActions {}
+
 /** Loaded extension with all registered items. */
-export interface LoadedExtension {
+export interface Extension {
 	path: string;
 	resolvedPath: string;
 	handlers: Map<string, HandlerFn[]>;
@ -751,25 +806,15 @@ export interface LoadedExtension {
 	messageRenderers: Map<string, MessageRenderer>;
 	commands: Map<string, RegisteredCommand>;
 	flags: Map<string, ExtensionFlag>;
-	flagValues: Map<string, boolean | string>;
 	shortcuts: Map<KeyId, ExtensionShortcut>;
-	setSendMessageHandler: (handler: SendMessageHandler) => void;
-	setSendUserMessageHandler: (handler: SendUserMessageHandler) => void;
-	setAppendEntryHandler: (handler: AppendEntryHandler) => void;
-	setGetActiveToolsHandler: (handler: GetActiveToolsHandler) => void;
-	setGetAllToolsHandler: (handler: GetAllToolsHandler) => void;
-	setSetActiveToolsHandler: (handler: SetActiveToolsHandler) => void;
-	setSetModelHandler: (handler: SetModelHandler) => void;
-	setGetThinkingLevelHandler: (handler: GetThinkingLevelHandler) => void;
-	setSetThinkingLevelHandler: (handler: SetThinkingLevelHandler) => void;
-	setFlagValue: (name: string, value: boolean | string) => void;
 }

 /** Result of loading extensions. */
 export interface LoadExtensionsResult {
-	extensions: LoadedExtension[];
+	extensions: Extension[];
 	errors: Array<{ path: string; error: string }>;
-	setUIContext(uiContext: ExtensionUIContext, hasUI: boolean): void;
+	/** Shared runtime - actions are throwing stubs until runner.initialize() */
+	runtime: ExtensionRuntime;
 }

 // ============================================================================
--- a/packages/coding-agent/src/core/index.ts
+++ b/packages/coding-agent/src/core/index.ts
@ -26,6 +26,7 @@ export {
 	discoverAndLoadExtensions,
 	type ExecOptions,
 	type ExecResult,
+	type Extension,
 	type ExtensionAPI,
 	type ExtensionCommandContext,
 	type ExtensionContext,
@ -38,7 +39,6 @@ export {
 	type ExtensionShortcut,
 	type ExtensionUIContext,
 	type LoadExtensionsResult,
-	type LoadedExtension,
 	type MessageRenderer,
 	type RegisteredCommand,
 	type SessionBeforeBranchEvent,
--- a/packages/coding-agent/src/core/sdk.ts
+++ b/packages/coding-agent/src/core/sdk.ts
@ -28,11 +28,11 @@ import { AgentSession } from "./agent-session.js";
 import { AuthStorage } from "./auth-storage.js";
 import { createEventBus, type EventBus } from "./event-bus.js";
 import {
+	createExtensionRuntime,
 	discoverAndLoadExtensions,
 	type ExtensionFactory,
 	ExtensionRunner,
 	type LoadExtensionsResult,
-	type LoadedExtension,
 	loadExtensionFromFactory,
 	type ToolDefinition,
 	wrapRegisteredTools,
@ -106,10 +106,10 @@ export interface CreateAgentSessionOptions {
 	/** Additional extension paths to load (merged with discovery). */
 	additionalExtensionPaths?: string[];
 	/**
-	 * Pre-loaded extensions (skips file discovery).
+	 * Pre-loaded extensions result (skips file discovery).
 	 * @internal Used by CLI when extensions are loaded early to parse custom flags.
 	 */
-	preloadedExtensions?: LoadedExtension[];
+	preloadedExtensionsResult?: LoadExtensionsResult;

 	/** Shared event bus for tool/extension communication. Default: creates new bus. */
 	eventBus?: EventBus;
@ -438,20 +438,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}

 	// Load extensions (discovers from standard locations + configured paths)
 	let extensionsResult: LoadExtensionsResult;
-	if (options.preloadedExtensions !== undefined && options.preloadedExtensions.length > 0) {
+	if (options.preloadedExtensionsResult !== undefined) {
 		// Use pre-loaded extensions (from early CLI flag discovery)
-		extensionsResult = {
-			extensions: options.preloadedExtensions,
-			errors: [],
-			setUIContext: () => {},
-		};
+		extensionsResult = options.preloadedExtensionsResult;
 	} else if (options.extensions !== undefined) {
 		// User explicitly provided extensions array (even if empty) - skip discovery
-		// Inline factories from options.extensions are loaded below
+		// Create runtime for inline extensions
+		const runtime = createExtensionRuntime();
 		extensionsResult = {
 			extensions: [],
 			errors: [],
-			setUIContext: () => {},
+			runtime,
 		};
 	} else {
 		// Discover extensions, merging with additional paths
@ -465,45 +462,29 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}

 	// Load inline extensions from factories
 	if (options.extensions && options.extensions.length > 0) {
-		// Create shared UI context holder that will be set later
-		const uiHolder: { ui: any; hasUI: boolean } = {
-			ui: {
-				select: async () => undefined,
-				confirm: async () => false,
-				input: async () => undefined,
-				notify: () => {},
-				setStatus: () => {},
-				setWidget: () => {},
-				setFooter: () => {},
-				setTitle: () => {},
-				custom: async () => undefined as never,
-				setEditorText: () => {},
-				getEditorText: () => "",
-				editor: async () => undefined,
-				get theme() {
-					return {} as any;
-				},
-			},
-			hasUI: false,
-		};
 		for (let i = 0; i < options.extensions.length; i++) {
 			const factory = options.extensions[i];
-			const loaded = await loadExtensionFromFactory(factory, cwd, eventBus, uiHolder, `<inline-${i}>`);
+			const loaded = await loadExtensionFromFactory(
+				factory,
+				cwd,
+				eventBus,
+				extensionsResult.runtime,
+				`<inline-${i}>`,
+			);
 			extensionsResult.extensions.push(loaded);
 		}
-		// Extend setUIContext to update inline extensions too
-		const originalSetUIContext = extensionsResult.setUIContext;
-		extensionsResult.setUIContext = (uiContext, hasUI) => {
-			originalSetUIContext(uiContext, hasUI);
-			uiHolder.ui = uiContext;
-			uiHolder.hasUI = hasUI;
-		};
 	}

 	// Create extension runner if we have extensions
 	let extensionRunner: ExtensionRunner | undefined;
 	if (extensionsResult.extensions.length > 0) {
-		extensionRunner = new ExtensionRunner(extensionsResult.extensions, cwd, sessionManager, modelRegistry);
+		extensionRunner = new ExtensionRunner(
+			extensionsResult.extensions,
+			extensionsResult.runtime,
+			cwd,
+			sessionManager,
+			modelRegistry,
+		);
 	}

 	// Wrap extension-registered tools and SDK-provided custom tools with context getter
@ -536,7 +517,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 				return {} as any;
 			},
 		},
-		hasUI: extensionRunner?.getHasUI() ?? false,
+		hasUI: extensionRunner?.hasUI() ?? false,
 		cwd,
 		sessionManager,
 		modelRegistry,