diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a4d66e8d2..51571235c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/src/extensibility/extensions/loader.ts b/packages/coding-agent/src/extensibility/extensions/loader.ts index 4a641f91c..c07b82a72 100644 --- a/packages/coding-agent/src/extensibility/extensions/loader.ts +++ b/packages/coding-agent/src/extensibility/extensions/loader.ts @@ -175,10 +175,12 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { } registerTool(tool: ToolDefinition): 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(), diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 500f4bd50..273ed73a6 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/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>(); #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 { + 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 { + 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`; diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 6d46fedfb..51b345fd5 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -1467,6 +1467,9 @@ export interface RegisteredTool void; + export interface ExtensionFlag { name: string; description?: string; @@ -1589,6 +1592,7 @@ export interface Extension { label?: string; handlers: Map; tools: Map>; + toolRegistrationListeners?: Set; assistantThinkingRenderers: AssistantThinkingRenderer[]; messageRenderers: Map; commands: Map; diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index f770442c5..163ba891a 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/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(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(); + const scheduleToolRegistration = (toolName: string): Promise => { + 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("mcp-notification", { build: buildMcpNotificationBatchMessage, }); session.yieldQueue.register(LSP_LATE_DIAGNOSTIC_MESSAGE_TYPE, { - isStale: entry => entry.isStale(), build: buildLateDiagnosticsBatchMessage, }); diff --git a/packages/coding-agent/test/sdk-tool-activation.test.ts b/packages/coding-agent/test/sdk-tool-activation.test.ts index bfcc3700f..bf0392948 100644 --- a/packages/coding-agent/test/sdk-tool-activation.test.ts +++ b/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(); + const lateRegistrationExtension: ExtensionFactory = pi => { + let startupPromise: Promise | 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");