fix(extensions): register lifecycle tools

packages/coding-agent/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed Pi-compatible extensions registering tools during asynchronous session startup being omitted from the live model tool registry.
+
 ## [17.2.12] - 2026-08-08
 
 ### Fixed

packages/coding-agent/src/extensibility/extensions/loader.ts

@@ -175,10 +175,12 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 	}
 
 	registerTool<TParams extends TSchema = TSchema, TDetails = unknown>(tool: ToolDefinition<TParams, TDetails>): void {
-		this.extension.tools.set(tool.name, {
+		const registered = {
 			definition: tool,
 			extensionPath: this.extension.path,
-		});
+		};
+		this.extension.tools.set(tool.name, registered);
+		for (const listener of this.extension.toolRegistrationListeners ?? []) listener(tool.name);
 	}
 
 	registerCommand(
@@ -316,6 +318,7 @@ function createExtension(extensionPath: string, resolvedPath: string): Extension
 		resolvedPath,
 		handlers: new Map(),
 		tools: new Map(),
+		toolRegistrationListeners: new Set(),
 		assistantThinkingRenderers: [],
 		messageRenderers: new Map(),
 		commands: new Map(),

packages/coding-agent/src/extensibility/extensions/runner.ts

@@ -62,6 +62,7 @@ import type {
 	SessionStopEventResult,
 	ToolCallEvent,
 	ToolCallEventResult,
+	ToolRegistrationListener,
 	ToolResultEvent,
 	ToolResultEventResult,
 	UserBashEvent,
@@ -352,6 +353,7 @@ export class ExtensionRunner {
 	#shutdownHandler: ShutdownHandler = () => {};
 	#getMemoryFn?: () => MemoryRuntimeContext | undefined;
 	#commandDiagnostics: Array<{ type: string; message: string; path: string }> = [];
+	#pendingToolRegistrations = new Set<Promise<void>>();
 	#initialized = false;
 	/**
 	 * Buffer for `credential_disabled` events received via {@link emitCredentialDisabled}
@@ -674,6 +676,51 @@ export class ExtensionRunner {
 		return tools;
 	}
 
+	/** Get the effective registered tool for a name using normal last-extension-wins precedence. */
+	getRegisteredTool(name: string): RegisteredTool | undefined {
+		for (let index = this.extensions.length - 1; index >= 0; index -= 1) {
+			const tool = this.extensions[index]?.tools.get(name);
+			if (tool) return tool;
+		}
+		return undefined;
+	}
+
+	/**
+	 * Observe tools registered after extension factories have loaded. Listener
+	 * promises are drained before the lifecycle handler that registered them
+	 * completes, keeping the model tool snapshot and system prompt coherent.
+	 */
+	onToolRegistered(listener: (toolName: string) => void | Promise<void>): () => void {
+		const wrapped: ToolRegistrationListener = toolName => {
+			try {
+				const pending = listener(toolName);
+				if (!pending) return;
+				this.#pendingToolRegistrations.add(pending);
+				void pending.then(
+					() => this.#pendingToolRegistrations.delete(pending),
+					() => this.#pendingToolRegistrations.delete(pending),
+				);
+			} catch (error) {
+				const pending = Promise.reject(error);
+				this.#pendingToolRegistrations.add(pending);
+				void pending.catch(() => this.#pendingToolRegistrations.delete(pending));
+			}
+		};
+		for (const extension of this.extensions) {
+			extension.toolRegistrationListeners ??= new Set();
+			extension.toolRegistrationListeners.add(wrapped);
+		}
+		return () => {
+			for (const extension of this.extensions) extension.toolRegistrationListeners?.delete(wrapped);
+		};
+	}
+
+	async #flushToolRegistrations(): Promise<void> {
+		while (this.#pendingToolRegistrations.size > 0) {
+			await Promise.all(this.#pendingToolRegistrations);
+		}
+	}
+
 	/**
 	 * Aggregate the registered CLI flags across a set of extensions (last write
 	 * wins on name collision). Static so callers that need the flag set before a
@@ -940,6 +987,7 @@ export class ExtensionRunner {
 				timeoutMs,
 				signal,
 			);
+			await this.#flushToolRegistrations();
 			if (handlerResult === EXTENSION_HANDLER_ABORTED) return undefined;
 			if (handlerResult === EXTENSION_HANDLER_TIMEOUT) {
 				const error = `handler timed out after ${timeoutMs}ms`;

packages/coding-agent/src/extensibility/extensions/types.ts

@@ -1467,6 +1467,9 @@ export interface RegisteredTool<TParams extends TSchema = TSchema, TDetails = un
 	extensionPath: string;
 }
 
+/** Internal observer invoked when an already-loaded extension registers or replaces a tool. */
+export type ToolRegistrationListener = (toolName: string) => void;
+
 export interface ExtensionFlag {
 	name: string;
 	description?: string;
@@ -1589,6 +1592,7 @@ export interface Extension {
 	label?: string;
 	handlers: Map<string, HandlerFn[]>;
 	tools: Map<string, RegisteredTool<any, any>>;
+	toolRegistrationListeners?: Set<ToolRegistrationListener>;
 	assistantThinkingRenderers: AssistantThinkingRenderer[];
 	messageRenderers: Map<string, MessageRenderer>;
 	commands: Map<string, RegisteredCommand>;

packages/coding-agent/src/sdk.ts

@@ -91,6 +91,7 @@ import {
 	type LoadExtensionsResult,
 	loadExtensionFromFactory,
 	loadExtensions,
+	type RegisteredTool,
 	type ToolDefinition,
 	wrapRegisteredTools,
 } from "./extensibility/extensions";
@@ -2584,6 +2585,7 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
 		const nativeToolsByName = new Map<string, Tool>(toolSession.xdev?.tools ?? undefined);
 
 		const registeredTools = restrictToolNames ? [] : extensionRunner.getAllRegisteredTools();
+		const initialRegisteredToolsByName = new Map(registeredTools.map(tool => [tool.definition.name, tool] as const));
 		const sdkCustomTools =
 			restrictToolNames && options.allowRestrictedCustomTools !== true
 				? []
@@ -3433,11 +3435,47 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
 			titleSystemPrompt: options.titleSystemPrompt,
 		});
 		hasSession = true;
+		// Extension factories normally register tools before session construction,
+		// but Pi-compatible extensions may discover them asynchronously from a
+		// session_start handler. Install those late registrations into the live
+		// registry and serialize activation so no update can overwrite a sibling.
+		let dynamicToolRegistrationChain = Promise.resolve();
+		const scheduledToolRegistrations = new WeakSet<RegisteredTool>();
+		const scheduleToolRegistration = (toolName: string): Promise<void> => {
+			const registered = extensionRunner.getRegisteredTool(toolName);
+			if (!registered) return dynamicToolRegistrationChain;
+			if (scheduledToolRegistrations.has(registered)) return dynamicToolRegistrationChain;
+			scheduledToolRegistrations.add(registered);
+
+			const [wrapped] = wrapRegisteredTools([registered], extensionRunner);
+			if (!wrapped) return dynamicToolRegistrationChain;
+			const name = registered.definition.name;
+			toolRegistry.set(name, new ExtensionToolWrapper(wrapToolWithMetaNotice(wrapped), extensionRunner));
+			builtInRegistryToolNames.delete(name);
+
+			const activation = dynamicToolRegistrationChain.then(async () => {
+				const enabled = session.getEnabledToolNames();
+				const alreadyEnabled = enabled.includes(name);
+				if (!alreadyEnabled && registered.definition.defaultInactive) return;
+				await session.setActiveToolsByName(alreadyEnabled ? enabled : [...enabled, name]);
+			});
+			dynamicToolRegistrationChain = activation.catch(() => {});
+			return activation;
+		};
+		const unsubscribeToolRegistrations = extensionRunner.onToolRegistered(scheduleToolRegistration);
+		disposeCallbacks.add(unsubscribeToolRegistrations);
+
+		// Close the construction race: a background registration can land after
+		// the initial snapshot but before the live listener above is attached.
+		for (const registered of extensionRunner.getAllRegisteredTools()) {
+			if (initialRegisteredToolsByName.get(registered.definition.name) !== registered) {
+				await scheduleToolRegistration(registered.definition.name);
+			}
+		}
 		session.yieldQueue.register<McpNotificationEntry>("mcp-notification", {
 			build: buildMcpNotificationBatchMessage,
 		});
 		session.yieldQueue.register<DeferredDiagnosticsEntry>(LSP_LATE_DIAGNOSTIC_MESSAGE_TYPE, {
-			isStale: entry => entry.isStale(),
 			build: buildLateDiagnosticsBatchMessage,
 		});
 

packages/coding-agent/test/sdk-tool-activation.test.ts

@@ -141,6 +141,64 @@ describe("createAgentSession defaultInactive tool activation", () => {
 		}
 	});
 
+	it("publishes tools from lazy session startup before the input lifecycle completes", async () => {
+		const tempDir = makeTempDir();
+		const startupGate = Promise.withResolvers<void>();
+		const lateRegistrationExtension: ExtensionFactory = pi => {
+			let startupPromise: Promise<void> | undefined;
+			pi.on("session_start", () => {
+				startupPromise = (async () => {
+					await startupGate.promise;
+					pi.registerTool({
+						name: "late_active_tool",
+						label: "Late Active Tool",
+						description: "Registered after asynchronous session initialization.",
+						parameters: type({}),
+						async execute() {
+							return { content: [{ type: "text", text: "late active" }] };
+						},
+					});
+					pi.registerTool({
+						name: "late_inactive_tool",
+						label: "Late Inactive Tool",
+						description: "Registered late but left disabled by default.",
+						parameters: type({}),
+						defaultInactive: true,
+						async execute() {
+							return { content: [{ type: "text", text: "late inactive" }] };
+						},
+					});
+				})();
+			});
+			pi.on("input", async () => {
+				await startupPromise;
+			});
+		};
+
+		const { session } = await createAgentSession({
+			...baseOptions(tempDir),
+			extensions: [lateRegistrationExtension],
+		});
+
+		try {
+			expect(session.getAllToolNames()).not.toContain("late_active_tool");
+			const runner = session.extensionRunner;
+			if (!runner) throw new Error("expected extension runner");
+			await runner.emit({ type: "session_start" });
+			expect(session.getAllToolNames()).not.toContain("late_active_tool");
+			startupGate.resolve();
+			await runner.emitInput("probe", undefined, "interactive");
+
+			expect(session.getAllToolNames()).toEqual(expect.arrayContaining(["late_active_tool", "late_inactive_tool"]));
+			expect(session.getEnabledToolNames()).toContain("late_active_tool");
+			expect(session.getEnabledToolNames()).not.toContain("late_inactive_tool");
+			expect(session.systemPrompt.join("\n")).toContain("late_active_tool");
+			expect(session.systemPrompt.join("\n")).not.toContain("late_inactive_tool");
+		} finally {
+			await session.dispose();
+		}
+	});
+
 	it("forwards built-in and external xd:// devices to Cursor provider contexts", async () => {
 		const tempDir = makeTempDir();
 		const cursorModel = getBundledModel("cursor", "composer-1.5");