refactor(coding-agent): unify tool and event handler context creation

Tools now use ExtensionRunner.createContext() instead of a separate
inline context factory. This ensures tools and event handlers share
the same context, fixing ctx.shutdown() and other context methods.

- Made ExtensionRunner.createContext() public
- Changed wrapRegisteredTools to accept ExtensionRunner instead of getContext callback
- Create ExtensionRunner when SDK custom tools are present (not just extensions)
- Removed redundant inline context factory from sdk.ts
This commit is contained in:
Mario Zechner 2026-01-08 03:45:39 +01:00
parent cf0466d96c
commit b1fb910625
3 changed files with 23 additions and 49 deletions

View file

@ -296,7 +296,11 @@ export class ExtensionRunner {
this.shutdownHandler(); this.shutdownHandler();
} }
private createContext(): ExtensionContext { /**
* Create an ExtensionContext for use in event handlers and tool execution.
* Context values are resolved at call time, so changes via initialize() are reflected.
*/
createContext(): ExtensionContext {
return { return {
ui: this.uiContext, ui: this.uiContext,
hasUI: this.hasUI(), hasUI: this.hasUI(),

View file

@ -4,12 +4,13 @@
import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core"; import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
import type { ExtensionRunner } from "./runner.js"; import type { ExtensionRunner } from "./runner.js";
import type { ExtensionContext, RegisteredTool, ToolCallEventResult, ToolResultEventResult } from "./types.js"; import type { RegisteredTool, ToolCallEventResult, ToolResultEventResult } from "./types.js";
/** /**
* Wrap a RegisteredTool into an AgentTool. * Wrap a RegisteredTool into an AgentTool.
* Uses the runner's createContext() for consistent context across tools and event handlers.
*/ */
export function wrapRegisteredTool(registeredTool: RegisteredTool, getContext: () => ExtensionContext): AgentTool { export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: ExtensionRunner): AgentTool {
const { definition } = registeredTool; const { definition } = registeredTool;
return { return {
name: definition.name, name: definition.name,
@ -17,18 +18,16 @@ export function wrapRegisteredTool(registeredTool: RegisteredTool, getContext: (
description: definition.description, description: definition.description,
parameters: definition.parameters, parameters: definition.parameters,
execute: (toolCallId, params, signal, onUpdate) => execute: (toolCallId, params, signal, onUpdate) =>
definition.execute(toolCallId, params, onUpdate, getContext(), signal), definition.execute(toolCallId, params, onUpdate, runner.createContext(), signal),
}; };
} }
/** /**
* Wrap all registered tools into AgentTools. * Wrap all registered tools into AgentTools.
* Uses the runner's createContext() for consistent context across tools and event handlers.
*/ */
export function wrapRegisteredTools( export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
registeredTools: RegisteredTool[], return registeredTools.map((rt) => wrapRegisteredTool(rt, runner));
getContext: () => ExtensionContext,
): AgentTool[] {
return registeredTools.map((rt) => wrapRegisteredTool(rt, getContext));
} }
/** /**

View file

@ -475,9 +475,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
} }
} }
// Create extension runner if we have extensions // Create extension runner if we have extensions or SDK custom tools
// The runner provides consistent context for tool execution (shutdown, abort, etc.)
let extensionRunner: ExtensionRunner | undefined; let extensionRunner: ExtensionRunner | undefined;
if (extensionsResult.extensions.length > 0) { const hasExtensions = extensionsResult.extensions.length > 0;
const hasCustomTools = options.customTools && options.customTools.length > 0;
if (hasExtensions || hasCustomTools) {
extensionRunner = new ExtensionRunner( extensionRunner = new ExtensionRunner(
extensionsResult.extensions, extensionsResult.extensions,
extensionsResult.runtime, extensionsResult.runtime,
@ -487,50 +490,18 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
); );
} }
// Wrap extension-registered tools and SDK-provided custom tools with context getter // Wrap extension-registered tools and SDK-provided custom tools
// (agent/session assigned below, accessed at execute time) // Tools use runner.createContext() for consistent context with event handlers
let agent: Agent; let agent: Agent;
let session: AgentSession;
const registeredTools = extensionRunner?.getAllRegisteredTools() ?? []; const registeredTools = extensionRunner?.getAllRegisteredTools() ?? [];
// Combine extension-registered tools with SDK-provided custom tools // Combine extension-registered tools with SDK-provided custom tools
const allCustomTools = [ const allCustomTools = [
...registeredTools, ...registeredTools,
...(options.customTools?.map((def) => ({ definition: def, extensionPath: "<sdk>" })) ?? []), ...(options.customTools?.map((def) => ({ definition: def, extensionPath: "<sdk>" })) ?? []),
]; ];
const wrappedExtensionTools = wrapRegisteredTools(allCustomTools, () => ({
ui: extensionRunner?.getUIContext() ?? { // Wrap tools using runner's context (ensures shutdown, abort, etc. work correctly)
select: async () => undefined, const wrappedExtensionTools = extensionRunner ? wrapRegisteredTools(allCustomTools, extensionRunner) : [];
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 {} as any;
},
},
hasUI: extensionRunner?.hasUI() ?? false,
cwd,
sessionManager,
modelRegistry,
model: agent.state.model,
isIdle: () => !session.isStreaming,
hasPendingMessages: () => session.pendingMessageCount > 0,
abort: () => {
session.abort();
},
shutdown: () => {
extensionRunner?.shutdown();
},
}));
// Create tool registry mapping name -> tool (for extension getTools/setTools) // Create tool registry mapping name -> tool (for extension getTools/setTools)
// Registry contains ALL built-in tools so extensions can enable any of them // Registry contains ALL built-in tools so extensions can enable any of them
@ -673,7 +644,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
sessionManager.appendThinkingLevelChange(thinkingLevel); sessionManager.appendThinkingLevelChange(thinkingLevel);
} }
session = new AgentSession({ const session = new AgentSession({
agent, agent,
sessionManager, sessionManager,
settingsManager, settingsManager,