mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-17 08:00:59 +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();
|
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(),
|
||||||
|
|
|
||||||
|
|
@ -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));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue