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();
}
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 {
ui: this.uiContext,
hasUI: this.hasUI(),

View file

@ -4,12 +4,13 @@
import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
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.
* 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;
return {
name: definition.name,
@ -17,18 +18,16 @@ export function wrapRegisteredTool(registeredTool: RegisteredTool, getContext: (
description: definition.description,
parameters: definition.parameters,
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.
* Uses the runner's createContext() for consistent context across tools and event handlers.
*/
export function wrapRegisteredTools(
registeredTools: RegisteredTool[],
getContext: () => ExtensionContext,
): AgentTool[] {
return registeredTools.map((rt) => wrapRegisteredTool(rt, getContext));
export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
return registeredTools.map((rt) => wrapRegisteredTool(rt, runner));
}
/**

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;
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(
extensionsResult.extensions,
extensionsResult.runtime,
@ -487,50 +490,18 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
);
}
// Wrap extension-registered tools and SDK-provided custom tools with context getter
// (agent/session assigned below, accessed at execute time)
// Wrap extension-registered tools and SDK-provided custom tools
// Tools use runner.createContext() for consistent context with event handlers
let agent: Agent;
let session: AgentSession;
const registeredTools = extensionRunner?.getAllRegisteredTools() ?? [];
// Combine extension-registered tools with SDK-provided custom tools
const allCustomTools = [
...registeredTools,
...(options.customTools?.map((def) => ({ definition: def, extensionPath: "<sdk>" })) ?? []),
];
const wrappedExtensionTools = wrapRegisteredTools(allCustomTools, () => ({
ui: extensionRunner?.getUIContext() ?? {
select: async () => undefined,
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();
},
}));
// Wrap tools using runner's context (ensures shutdown, abort, etc. work correctly)
const wrappedExtensionTools = extensionRunner ? wrapRegisteredTools(allCustomTools, extensionRunner) : [];
// Create tool registry mapping name -> tool (for extension getTools/setTools)
// 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);
}
session = new AgentSession({
const session = new AgentSession({
agent,
sessionManager,
settingsManager,