@oh-my-pi/pi-coding-agent 3.37.1 → 4.0.0

package/src/core/extensions/runner.ts

@@ -5,41 +5,43 @@
 import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
 import type { ImageContent, Model } from "@oh-my-pi/pi-ai";
 import type { KeyId } from "@oh-my-pi/pi-tui";
-import { theme } from "../../modes/interactive/theme/theme";
+import { type Theme, theme } from "../../modes/interactive/theme/theme";
+import { logger } from "../logger";
 import type { ModelRegistry } from "../model-registry";
 import type { SessionManager } from "../session-manager";
 import type {
-	AppendEntryHandler,
 	BeforeAgentStartEvent,
 	BeforeAgentStartEventResult,
 	ContextEvent,
 	ContextEventResult,
+	Extension,
+	ExtensionActions,
 	ExtensionCommandContext,
+	ExtensionCommandContextActions,
 	ExtensionContext,
+	ExtensionContextActions,
 	ExtensionError,
 	ExtensionEvent,
 	ExtensionFlag,
+	ExtensionRuntime,
 	ExtensionShortcut,
 	ExtensionUIContext,
-	GetActiveToolsHandler,
-	GetAllToolsHandler,
-	LoadedExtension,
 	MessageRenderer,
 	RegisteredCommand,
 	RegisteredTool,
-	SendMessageHandler,
 	SessionBeforeCompactResult,
 	SessionBeforeTreeResult,
-	SetActiveToolsHandler,
 	ToolCallEvent,
 	ToolCallEventResult,
 	ToolResultEventResult,
+	UserBashEvent,
+	UserBashEventResult,
 } from "./types";
 
 /** Combined result from all before_agent_start handlers */
 interface BeforeAgentStartCombinedResult {
 	messages?: NonNullable<BeforeAgentStartEventResult["message"]>[];
-	systemPromptAppend?: string;
+	systemPrompt?: string;
 }
 
 export type ExtensionErrorListener = (error: ExtensionError) => void;
@@ -56,27 +58,49 @@ export type NavigateTreeHandler = (
 	options?: { summarize?: boolean },
 ) => Promise<{ cancelled: boolean }>;
 
+export type ShutdownHandler = () => void;
+
+/**
+ * Helper function to emit session_shutdown event to extensions.
+ * Returns true if the event was emitted, false if there were no handlers.
+ */
+export async function emitSessionShutdownEvent(extensionRunner: ExtensionRunner | undefined): Promise<boolean> {
+	if (extensionRunner?.hasHandlers("session_shutdown")) {
+		await extensionRunner.emit({
+			type: "session_shutdown",
+		});
+		return true;
+	}
+	return false;
+}
+
 const noOpUIContext: ExtensionUIContext = {
-	select: async () => undefined,
-	confirm: async () => false,
-	input: async () => undefined,
+	select: async (_title, _options, _dialogOptions) => undefined,
+	confirm: async (_title, _message, _dialogOptions) => false,
+	input: async (_title, _placeholder, _dialogOptions) => undefined,
 	notify: () => {},
 	setStatus: () => {},
 	setWidget: () => {},
+	setFooter: () => {},
+	setHeader: () => {},
 	setTitle: () => {},
 	custom: async () => undefined as never,
 	setEditorText: () => {},
 	getEditorText: () => "",
 	editor: async () => undefined,
+	setEditorComponent: () => {},
 	get theme() {
 		return theme;
 	},
+	getAllThemes: () => [],
+	getTheme: () => undefined,
+	setTheme: (_theme: string | Theme) => ({ success: false, error: "UI not available" }),
 };
 
 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;
@@ -89,72 +113,64 @@ export class ExtensionRunner {
 	private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
 	private branchHandler: BranchHandler = async () => ({ cancelled: false });
 	private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });
+	private shutdownHandler: ShutdownHandler = () => {};
 
 	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;
-		appendEntryHandler: AppendEntryHandler;
-		getActiveToolsHandler: GetActiveToolsHandler;
-		getAllToolsHandler: GetAllToolsHandler;
-		setActiveToolsHandler: SetActiveToolsHandler;
-		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);
-
-		if (options.newSessionHandler) {
-			this.newSessionHandler = options.newSessionHandler;
-		}
-		if (options.branchHandler) {
-			this.branchHandler = options.branchHandler;
-		}
-		if (options.navigateTreeHandler) {
-			this.navigateTreeHandler = options.navigateTreeHandler;
-		}
-
-		for (const ext of this.extensions) {
-			ext.setSendMessageHandler(options.sendMessageHandler);
-			ext.setAppendEntryHandler(options.appendEntryHandler);
-			ext.setGetActiveToolsHandler(options.getActiveToolsHandler);
-			ext.setGetAllToolsHandler(options.getAllToolsHandler);
-			ext.setSetActiveToolsHandler(options.setActiveToolsHandler);
+	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;
+
+		// Context actions (required)
+		this.getModel = contextActions.getModel;
+		this.isIdleFn = contextActions.isIdle;
+		this.abortFn = contextActions.abort;
+		this.hasPendingMessagesFn = contextActions.hasPendingMessages;
+		this.shutdownHandler = contextActions.shutdown;
+
+		// 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;
 		}
 
-		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[] {
@@ -183,11 +199,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([
@@ -214,17 +226,20 @@ export class ExtensionRunner {
 				const normalizedKey = key.toLowerCase() as KeyId;
 
 				if (ExtensionRunner.RESERVED_SHORTCUTS.has(normalizedKey)) {
-					console.warn(
-						`Extension shortcut '${key}' from ${shortcut.extensionPath} conflicts with built-in shortcut. Skipping.`,
-					);
+					logger.warn("Extension shortcut conflicts with built-in shortcut", {
+						key,
+						extensionPath: shortcut.extensionPath,
+					});
 					continue;
 				}
 
 				const existing = allShortcuts.get(normalizedKey);
 				if (existing) {
-					console.warn(
-						`Extension shortcut conflict: '${key}' registered by both ${existing.extensionPath} and ${shortcut.extensionPath}. Using ${shortcut.extensionPath}.`,
-					);
+					logger.warn("Extension shortcut conflict", {
+						key,
+						extensionPath: shortcut.extensionPath,
+						existingExtensionPath: existing.extensionPath,
+					});
 				}
 				allShortcuts.set(normalizedKey, shortcut);
 			}
@@ -283,10 +298,10 @@ export class ExtensionRunner {
 		return undefined;
 	}
 
-	private createContext(): ExtensionContext {
+	createContext(): ExtensionContext {
 		return {
 			ui: this.uiContext,
-			hasUI: this.hasUI,
+			hasUI: this.hasUI(),
 			cwd: this.cwd,
 			sessionManager: this.sessionManager,
 			modelRegistry: this.modelRegistry,
@@ -294,10 +309,18 @@ export class ExtensionRunner {
 			isIdle: () => this.isIdleFn(),
 			abort: () => this.abortFn(),
 			hasPendingMessages: () => this.hasPendingMessagesFn(),
+			shutdown: () => this.shutdownHandler(),
 			hasQueuedMessages: () => this.hasPendingMessagesFn(),
 		};
 	}
 
+	/**
+	 * Request a graceful shutdown. Called by extension tools and event handlers.
+	 */
+	shutdown(): void {
+		this.shutdownHandler();
+	}
+
 	createCommandContext(): ExtensionCommandContext {
 		return {
 			...this.createContext(),
@@ -394,6 +417,35 @@ export class ExtensionRunner {
 		return result;
 	}
 
+	async emitUserBash(event: UserBashEvent): Promise<UserBashEventResult | undefined> {
+		const ctx = this.createContext();
+
+		for (const ext of this.extensions) {
+			const handlers = ext.handlers.get("user_bash");
+			if (!handlers || handlers.length === 0) continue;
+
+			for (const handler of handlers) {
+				try {
+					const handlerResult = await handler(event, ctx);
+					if (handlerResult) {
+						return handlerResult as UserBashEventResult;
+					}
+				} catch (err) {
+					const message = err instanceof Error ? err.message : String(err);
+					const stack = err instanceof Error ? err.stack : undefined;
+					this.emitError({
+						extensionPath: ext.path,
+						event: "user_bash",
+						error: message,
+						stack,
+					});
+				}
+			}
+		}
+
+		return undefined;
+	}
+
 	async emitContext(messages: AgentMessage[]): Promise<AgentMessage[]> {
 		const ctx = this.createContext();
 		let currentMessages = structuredClone(messages);
@@ -428,11 +480,13 @@ export class ExtensionRunner {
 
 	async emitBeforeAgentStart(
 		prompt: string,
-		images?: ImageContent[],
+		images: ImageContent[] | undefined,
+		systemPrompt: string,
 	): Promise<BeforeAgentStartCombinedResult | undefined> {
 		const ctx = this.createContext();
 		const messages: NonNullable<BeforeAgentStartEventResult["message"]>[] = [];
-		const systemPromptAppends: string[] = [];
+		let currentSystemPrompt = systemPrompt;
+		let systemPromptModified = false;
 
 		for (const ext of this.extensions) {
 			const handlers = ext.handlers.get("before_agent_start");
@@ -440,7 +494,12 @@ export class ExtensionRunner {
 
 			for (const handler of handlers) {
 				try {
-					const event: BeforeAgentStartEvent = { type: "before_agent_start", prompt, images };
+					const event: BeforeAgentStartEvent = {
+						type: "before_agent_start",
+						prompt,
+						images,
+						systemPrompt: currentSystemPrompt,
+					};
 					const handlerResult = await handler(event, ctx);
 
 					if (handlerResult) {
@@ -448,8 +507,9 @@ export class ExtensionRunner {
 						if (result.message) {
 							messages.push(result.message);
 						}
-						if (result.systemPromptAppend) {
-							systemPromptAppends.push(result.systemPromptAppend);
+						if (result.systemPrompt !== undefined) {
+							currentSystemPrompt = result.systemPrompt;
+							systemPromptModified = true;
 						}
 					}
 				} catch (err) {
@@ -465,10 +525,10 @@ export class ExtensionRunner {
 			}
 		}
 
-		if (messages.length > 0 || systemPromptAppends.length > 0) {
+		if (messages.length > 0 || systemPromptModified) {
 			return {
 				messages: messages.length > 0 ? messages : undefined,
-				systemPromptAppend: systemPromptAppends.length > 0 ? systemPromptAppends.join("\n\n") : undefined,
+				systemPrompt: systemPromptModified ? currentSystemPrompt : undefined,
 			};
 		}