@gajae-code/coding-agent 0.8.2 → 0.9.1

package/src/extensibility/extensions/runner.ts

@@ -111,6 +111,8 @@ type RunnerEmitResult<TEvent extends RunnerEmitEvent> = TEvent extends { type: "
 				: TEvent extends { type: "session.compacting" }
 					? SessionCompactingResult | undefined
 					: undefined;
+type Handler = Extension["handlers"] extends Map<string, Array<infer T>> ? T : never;
+type IndexedHandler = { ext: Extension; handler: Handler };
 
 export type NewSessionHandler = (options?: {
 	parentSession?: string;
@@ -173,6 +175,8 @@ const noOpUIContext: ExtensionUIContext = {
 export class ExtensionRunner {
 	#uiContext: ExtensionUIContext;
 	#errorListeners: Set<ExtensionErrorListener> = new Set();
+	#handlersByEvent: Map<string, IndexedHandler[]> = new Map();
+
 	#getModel: () => Model | undefined = () => undefined;
 	#isIdleFn: () => boolean = () => true;
 	#waitForIdleFn: () => Promise<void> = async () => {};
@@ -208,6 +212,24 @@ export class ExtensionRunner {
 		private readonly sessionMetadata?: ExtensionContext["sessionMetadata"],
 	) {
 		this.#uiContext = noOpUIContext;
+		this.#handlersByEvent = ExtensionRunner.#indexHandlers(extensions);
+	}
+
+	static #indexHandlers(extensions: Extension[]): Map<string, IndexedHandler[]> {
+		const handlersByEvent = new Map<string, IndexedHandler[]>();
+		for (const ext of extensions) {
+			for (const [eventType, handlers] of ext.handlers) {
+				let indexedHandlers = handlersByEvent.get(eventType);
+				if (!indexedHandlers) {
+					indexedHandlers = [];
+					handlersByEvent.set(eventType, indexedHandlers);
+				}
+				for (const handler of handlers) {
+					indexedHandlers.push({ ext, handler });
+				}
+			}
+		}
+		return handlersByEvent;
 	}
 
 	initialize(
@@ -396,13 +418,7 @@ export class ExtensionRunner {
 	}
 
 	hasHandlers(eventType: string): boolean {
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get(eventType);
-			if (handlers && handlers.length > 0) {
-				return true;
-			}
-		}
-		return false;
+		return (this.#handlersByEvent.get(eventType)?.length ?? 0) > 0;
 	}
 
 	getMessageRenderer(customType: string): MessageRenderer | undefined {
@@ -511,11 +527,17 @@ export class ExtensionRunner {
 		ext: Extension,
 		timeoutMs: number,
 	): Promise<TResult | undefined> {
+		let timeout: ReturnType<typeof setTimeout> | undefined;
 		try {
-			const handlerResult = await Promise.race([
-				Promise.resolve(handler(event, ctx)),
-				Bun.sleep(timeoutMs).then(() => EXTENSION_HANDLER_TIMEOUT),
-			]);
+			const timeoutPromise = new Promise<typeof EXTENSION_HANDLER_TIMEOUT>(resolve => {
+				timeout = setTimeout(() => resolve(EXTENSION_HANDLER_TIMEOUT), timeoutMs);
+			});
+			const handlerResult = await Promise.race([Promise.resolve(handler(event, ctx)), timeoutPromise]);
+			if (timeout !== undefined) {
+				clearTimeout(timeout);
+				timeout = undefined;
+			}
+
 			if (handlerResult === EXTENSION_HANDLER_TIMEOUT) {
 				const error = `handler timed out after ${timeoutMs}ms`;
 				logger.warn("Extension handler timed out", {
@@ -541,73 +563,67 @@ export class ExtensionRunner {
 				stack,
 			});
 			return undefined;
+		} finally {
+			if (timeout !== undefined) {
+				clearTimeout(timeout);
+			}
 		}
 	}
 
 	async emit<TEvent extends RunnerEmitEvent>(event: TEvent): Promise<RunnerEmitResult<TEvent>> {
+		const handlers = this.#handlersByEvent.get(event.type) ?? [];
+		if (handlers.length === 0) return undefined as RunnerEmitResult<TEvent>;
+
 		const ctx = this.createContext();
 		let result: SessionBeforeEventResult | SessionCompactingResult | undefined;
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get(event.type);
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
-
-				if (this.#isSessionBeforeEvent(event) && handlerResult) {
-					result = handlerResult as SessionBeforeEventResult;
-					if (result.cancel) {
-						return result as RunnerEmitResult<TEvent>;
-					}
-				}
+		for (const { ext, handler } of handlers) {
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
 
-				if (event.type === "session.compacting" && handlerResult) {
-					result = handlerResult as SessionCompactingResult;
+			if (this.#isSessionBeforeEvent(event) && handlerResult) {
+				result = handlerResult as SessionBeforeEventResult;
+				if (result.cancel) {
+					return result as RunnerEmitResult<TEvent>;
 				}
 			}
+
+			if (event.type === "session.compacting" && handlerResult) {
+				result = handlerResult as SessionCompactingResult;
+			}
 		}
 
 		return result as RunnerEmitResult<TEvent>;
 	}
 
 	async emitToolResult(event: ToolResultEvent): Promise<ToolResultEventResult | undefined> {
+		const handlers = this.#handlersByEvent.get("tool_result") ?? [];
+		if (handlers.length === 0) return undefined;
+
 		const ctx = this.createContext();
 		const currentEvent: ToolResultEvent = { ...event };
 		let modified = false;
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("tool_result");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const handlerResult = (await this.#runHandlerWithTimeout(
-					handler,
-					currentEvent,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				)) as ToolResultEventResult | undefined;
-				if (!handlerResult) continue;
-
-				if (handlerResult.content !== undefined) {
-					currentEvent.content = handlerResult.content;
-					modified = true;
-				}
-				if (handlerResult.details !== undefined) {
-					currentEvent.details = handlerResult.details;
-					modified = true;
-				}
-				if (handlerResult.isError !== undefined) {
-					currentEvent.isError = handlerResult.isError;
-					modified = true;
-				}
+		for (const { ext, handler } of handlers) {
+			const handlerResult = (await this.#runHandlerWithTimeout(
+				handler,
+				currentEvent,
+				ctx,
+				ext,
+				extensionHandlerTimeoutMs,
+			)) as ToolResultEventResult | undefined;
+			if (!handlerResult) continue;
+
+			if (handlerResult.content !== undefined) {
+				currentEvent.content = handlerResult.content;
+				modified = true;
+			}
+			if (handlerResult.details !== undefined) {
+				currentEvent.details = handlerResult.details;
+				modified = true;
+			}
+			if (handlerResult.isError !== undefined) {
+				currentEvent.isError = handlerResult.isError;
+				modified = true;
 			}
 		}
 
@@ -621,34 +637,32 @@ export class ExtensionRunner {
 	}
 
 	async emitToolCall(event: ToolCallEvent): Promise<ToolCallEventResult | undefined> {
+		const handlers = this.#handlersByEvent.get("tool_call") ?? [];
+		if (handlers.length === 0) return undefined;
+
 		const ctx = this.createContext();
 		let result: ToolCallEventResult | undefined;
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("tool_call");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				try {
-					const handlerResult = await handler(event, ctx);
-
-					if (handlerResult) {
-						result = handlerResult as ToolCallEventResult;
-						if (result.block) {
-							return result;
-						}
+		for (const { ext, handler } of handlers) {
+			try {
+				const handlerResult = await handler(event, ctx);
+
+				if (handlerResult) {
+					result = handlerResult as ToolCallEventResult;
+					if (result.block) {
+						return result;
 					}
-				} 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: "tool_call",
-						error: message,
-						stack,
-					});
-					return { block: true, reason: `Extension ${ext.path} failed: ${message}` };
 				}
+			} 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: "tool_call",
+					error: message,
+					stack,
+				});
+				return { block: true, reason: `Extension ${ext.path} failed: ${message}` };
 			}
 		}
 
@@ -667,23 +681,15 @@ export class ExtensionRunner {
 		event: UserBashEvent | UserPythonEvent,
 		eventName: "user_bash" | "user_python",
 	): Promise<R | undefined> {
+		const handlers = this.#handlersByEvent.get(eventName) ?? [];
+		if (handlers.length === 0) return undefined;
+
 		const ctx = this.createContext();
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get(eventName);
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
-				if (handlerResult) {
-					return handlerResult as R;
-				}
+		for (const { ext, handler } of handlers) {
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
+			if (handlerResult) {
+				return handlerResult as R;
 			}
 		}
 
@@ -698,35 +704,26 @@ export class ExtensionRunner {
 		promptPaths: Array<{ path: string; extensionPath: string }>;
 		themePaths: Array<{ path: string; extensionPath: string }>;
 	}> {
+		const handlers = this.#handlersByEvent.get("resources_discover") ?? [];
+		if (handlers.length === 0) return { skillPaths: [], promptPaths: [], themePaths: [] };
 		const ctx = this.createContext();
 		const skillPaths: Array<{ path: string; extensionPath: string }> = [];
 		const promptPaths: Array<{ path: string; extensionPath: string }> = [];
 		const themePaths: Array<{ path: string; extensionPath: string }> = [];
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("resources_discover");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const event: ResourcesDiscoverEvent = { type: "resources_discover", cwd, reason };
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
-				const result = handlerResult as ResourcesDiscoverResult | undefined;
-
-				if (result?.skillPaths?.length) {
-					skillPaths.push(...result.skillPaths.map(path => ({ path, extensionPath: ext.path })));
-				}
-				if (result?.promptPaths?.length) {
-					promptPaths.push(...result.promptPaths.map(path => ({ path, extensionPath: ext.path })));
-				}
-				if (result?.themePaths?.length) {
-					themePaths.push(...result.themePaths.map(path => ({ path, extensionPath: ext.path })));
-				}
+		for (const { ext, handler } of handlers) {
+			const event: ResourcesDiscoverEvent = { type: "resources_discover", cwd, reason };
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
+			const result = handlerResult as ResourcesDiscoverResult | undefined;
+
+			if (result?.skillPaths?.length) {
+				skillPaths.push(...result.skillPaths.map(path => ({ path, extensionPath: ext.path })));
+			}
+			if (result?.promptPaths?.length) {
+				promptPaths.push(...result.promptPaths.map(path => ({ path, extensionPath: ext.path })));
+			}
+			if (result?.themePaths?.length) {
+				themePaths.push(...result.themePaths.map(path => ({ path, extensionPath: ext.path })));
 			}
 		}
 
@@ -739,39 +736,33 @@ export class ExtensionRunner {
 		images: ImageContent[] | undefined,
 		source: "interactive" | "rpc" | "extension",
 	): Promise<InputEventResult> {
+		const handlers = this.#handlersByEvent.get("input") ?? [];
+		if (handlers.length === 0) return {};
+
 		const ctx = this.createContext();
 		let currentText = text;
 		let currentImages = images;
 
-		for (const ext of this.extensions) {
-			for (const handler of ext.handlers.get("input") ?? []) {
-				const event: InputEvent = { type: "input", text: currentText, images: currentImages, source };
-				const result = (await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs)) as
-					| InputEventResult
-					| undefined;
-				if (result?.handled) return result;
-				if (result?.text !== undefined) {
-					currentText = result.text;
-					currentImages = result.images ?? currentImages;
-				}
+		for (const { ext, handler } of handlers) {
+			const event: InputEvent = { type: "input", text: currentText, images: currentImages, source };
+			const result = (await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs)) as
+				| InputEventResult
+				| undefined;
+			if (result?.handled) return result;
+			if (result?.text !== undefined) {
+				currentText = result.text;
+				currentImages = result.images ?? currentImages;
 			}
 		}
+
 		return currentText !== text || currentImages !== images ? { text: currentText, images: currentImages } : {};
 	}
 
 	async emitContext(messages: AgentMessage[]): Promise<AgentMessage[]> {
-		const ctx = this.createContext();
-
-		// Check if any extensions actually have context handlers before cloning
-		let hasContextHandlers = false;
-		for (const ext of this.extensions) {
-			if (ext.handlers.get("context")?.length) {
-				hasContextHandlers = true;
-				break;
-			}
-		}
-		if (!hasContextHandlers) return messages;
+		const handlers = this.#handlersByEvent.get("context") ?? [];
+		if (handlers.length === 0) return messages;
 
+		const ctx = this.createContext();
 		let currentMessages: AgentMessage[];
 		try {
 			currentMessages = structuredClone(messages);
@@ -782,23 +773,12 @@ export class ExtensionRunner {
 			currentMessages = [...messages];
 		}
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("context");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const event: ContextEvent = { type: "context", messages: currentMessages };
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
-
-				if (handlerResult && (handlerResult as ContextEventResult).messages) {
-					currentMessages = (handlerResult as ContextEventResult).messages!;
-				}
+		for (const { ext, handler } of handlers) {
+			const event: ContextEvent = { type: "context", messages: currentMessages };
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
+
+			if (handlerResult && (handlerResult as ContextEventResult).messages) {
+				currentMessages = (handlerResult as ContextEventResult).messages!;
 			}
 		}
 
@@ -806,28 +786,20 @@ export class ExtensionRunner {
 	}
 
 	async emitBeforeProviderRequest(payload: unknown): Promise<BeforeProviderRequestEventResult> {
+		const handlers = this.#handlersByEvent.get("before_provider_request") ?? [];
+		if (handlers.length === 0) return payload;
+
 		const ctx = this.createContext();
 		let currentPayload = payload;
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("before_provider_request");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const event: BeforeProviderRequestEvent = {
-					type: "before_provider_request",
-					payload: currentPayload,
-				};
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
-				if (handlerResult !== undefined) {
-					currentPayload = handlerResult;
-				}
+		for (const { ext, handler } of handlers) {
+			const event: BeforeProviderRequestEvent = {
+				type: "before_provider_request",
+				payload: currentPayload,
+			};
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
+			if (handlerResult !== undefined) {
+				currentPayload = handlerResult;
 			}
 		}
 
@@ -835,22 +807,20 @@ export class ExtensionRunner {
 	}
 
 	async emitAfterProviderResponse(response: ProviderResponseMetadata, _model?: Model): Promise<void> {
+		const handlers = this.#handlersByEvent.get("after_provider_response") ?? [];
+		if (handlers.length === 0) return;
+
 		const ctx = this.createContext();
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("after_provider_response");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const event: AfterProviderResponseEvent = {
-					type: "after_provider_response",
-					status: response.status,
-					headers: response.headers,
-					requestId: response.requestId,
-					metadata: response.metadata,
-				};
-				await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
-			}
+		for (const { ext, handler } of handlers) {
+			const event: AfterProviderResponseEvent = {
+				type: "after_provider_response",
+				status: response.status,
+				headers: response.headers,
+				requestId: response.requestId,
+				metadata: response.metadata,
+			};
+			await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
 		}
 	}
 
@@ -859,39 +829,31 @@ export class ExtensionRunner {
 		images: ImageContent[] | undefined,
 		systemPrompt: string[],
 	): Promise<BeforeAgentStartCombinedResult | undefined> {
+		const handlers = this.#handlersByEvent.get("before_agent_start") ?? [];
+		if (handlers.length === 0) return undefined;
+
 		const ctx = this.createContext();
 		const messages: NonNullable<BeforeAgentStartEventResult["message"]>[] = [];
 		let currentSystemPrompt = systemPrompt;
 		let systemPromptModified = false;
 
-		for (const ext of this.extensions) {
-			const handlers = ext.handlers.get("before_agent_start");
-			if (!handlers || handlers.length === 0) continue;
-
-			for (const handler of handlers) {
-				const event: BeforeAgentStartEvent = {
-					type: "before_agent_start",
-					prompt,
-					images,
-					systemPrompt: currentSystemPrompt,
-				};
-				const handlerResult = await this.#runHandlerWithTimeout(
-					handler,
-					event,
-					ctx,
-					ext,
-					extensionHandlerTimeoutMs,
-				);
+		for (const { ext, handler } of handlers) {
+			const event: BeforeAgentStartEvent = {
+				type: "before_agent_start",
+				prompt,
+				images,
+				systemPrompt: currentSystemPrompt,
+			};
+			const handlerResult = await this.#runHandlerWithTimeout(handler, event, ctx, ext, extensionHandlerTimeoutMs);
 
-				if (handlerResult) {
-					const result = handlerResult as BeforeAgentStartEventResult;
-					if (result.message) {
-						messages.push(result.message);
-					}
-					if (result.systemPrompt !== undefined) {
-						currentSystemPrompt = result.systemPrompt;
-						systemPromptModified = true;
-					}
+			if (handlerResult) {
+				const result = handlerResult as BeforeAgentStartEventResult;
+				if (result.message) {
+					messages.push(result.message);
+				}
+				if (result.systemPrompt !== undefined) {
+					currentSystemPrompt = result.systemPrompt;
+					systemPromptModified = true;
 				}
 			}
 		}