mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 20:03:05 +00:00
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:
parent
cf0466d96c
commit
b1fb910625
3 changed files with 23 additions and 49 deletions
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue