mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 06:02:42 +00:00
Rework custom tools API with CustomToolContext
- CustomAgentTool renamed to CustomTool
- ToolAPI renamed to CustomToolAPI
- ToolContext renamed to CustomToolContext
- ToolSessionEvent renamed to CustomToolSessionEvent
- Added CustomToolContext parameter to execute() and onSession()
- CustomToolFactory now returns CustomTool<any, any> for type compatibility
- dispose() replaced with onSession({ reason: 'shutdown' })
- Added wrapCustomTool() to convert CustomTool to AgentTool
- Session exposes setToolUIContext() instead of leaking internals
- Fix ToolExecutionComponent to sync with toolOutputExpanded state
- Update all custom tool examples for new API
This commit is contained in:
Mario Zechner
parent
b123df5fab
commit
568150f18b
27 changed files with 336 additions and 289 deletions
|
|
@ -11,7 +11,7 @@ import {
|
|||
type TUI,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import stripAnsi from "strip-ansi";
|
||||
import type { CustomAgentTool } from "../../../core/custom-tools/types.js";
|
||||
import type { CustomTool } from "../../../core/custom-tools/types.js";
|
||||
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize } from "../../../core/tools/truncate.js";
|
||||
import { getLanguageFromPath, highlightCode, theme } from "../theme/theme.js";
|
||||
import { renderDiff } from "./diff.js";
|
||||
|
|
@ -55,7 +55,7 @@ export class ToolExecutionComponent extends Container {
|
|||
private expanded = false;
|
||||
private showImages: boolean;
|
||||
private isPartial = true;
|
||||
private customTool?: CustomAgentTool;
|
||||
private customTool?: CustomTool;
|
||||
private ui: TUI;
|
||||
private result?: {
|
||||
content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
|
||||
|
|
@ -67,7 +67,7 @@ export class ToolExecutionComponent extends Container {
|
|||
toolName: string,
|
||||
args: any,
|
||||
options: ToolExecutionOptions = {},
|
||||
customTool: CustomAgentTool | undefined,
|
||||
customTool: CustomTool | undefined,
|
||||
ui: TUI,
|
||||
) {
|
||||
super();
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ import {
|
|||
import { exec, spawnSync } from "child_process";
|
||||
import { APP_NAME, getAuthPath, getDebugLogPath } from "../../config.js";
|
||||
import type { AgentSession, AgentSessionEvent } from "../../core/agent-session.js";
|
||||
import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "../../core/custom-tools/index.js";
|
||||
import type { CustomToolSessionEvent, LoadedCustomTool } from "../../core/custom-tools/index.js";
|
||||
import type { HookUIContext } from "../../core/hooks/index.js";
|
||||
import { createCompactionSummaryMessage } from "../../core/messages.js";
|
||||
import { type SessionContext, SessionManager } from "../../core/session-manager.js";
|
||||
|
|
@ -350,19 +350,20 @@ export class InteractiveMode {
|
|||
this.chatContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
// Load session entries if any
|
||||
const entries = this.session.sessionManager.getEntries();
|
||||
|
||||
// Set TUI-based UI context for custom tools
|
||||
const uiContext = this.createHookUIContext();
|
||||
// Create and set hook & tool UI context
|
||||
const uiContext: HookUIContext = {
|
||||
select: (title, options) => this.showHookSelector(title, options),
|
||||
confirm: (title, message) => this.showHookConfirm(title, message),
|
||||
input: (title, placeholder) => this.showHookInput(title, placeholder),
|
||||
notify: (message, type) => this.showHookNotify(message, type),
|
||||
custom: (component) => this.showHookCustom(component),
|
||||
};
|
||||
this.setToolUIContext(uiContext, true);
|
||||
|
||||
// Notify custom tools of session start
|
||||
await this.emitToolSessionEvent({
|
||||
entries,
|
||||
sessionFile: this.session.sessionFile,
|
||||
previousSessionFile: undefined,
|
||||
await this.emitCustomToolSessionEvent({
|
||||
reason: "start",
|
||||
previousSessionFile: undefined,
|
||||
});
|
||||
|
||||
const hookRunner = this.session.hookRunner;
|
||||
|
|
@ -370,34 +371,35 @@ export class InteractiveMode {
|
|||
return; // No hooks loaded
|
||||
}
|
||||
|
||||
// Set UI context on hook runner
|
||||
hookRunner.setUIContext(uiContext, true);
|
||||
hookRunner.initialize({
|
||||
getModel: () => this.session.model,
|
||||
sendMessageHandler: (message, triggerTurn) => {
|
||||
const wasStreaming = this.session.isStreaming;
|
||||
this.session
|
||||
.sendHookMessage(message, triggerTurn)
|
||||
.then(() => {
|
||||
// For non-streaming cases with display=true, update UI
|
||||
// (streaming cases update via message_end event)
|
||||
if (!wasStreaming && message.display) {
|
||||
this.rebuildChatFromMessages();
|
||||
}
|
||||
})
|
||||
.catch((err) => {
|
||||
this.showError(`Hook sendMessage failed: ${err instanceof Error ? err.message : String(err)}`);
|
||||
});
|
||||
},
|
||||
appendEntryHandler: (customType, data) => {
|
||||
this.sessionManager.appendCustomEntry(customType, data);
|
||||
},
|
||||
uiContext,
|
||||
hasUI: true,
|
||||
});
|
||||
|
||||
// Subscribe to hook errors
|
||||
hookRunner.onError((error) => {
|
||||
this.showHookError(error.hookPath, error.error);
|
||||
});
|
||||
|
||||
// Set up handlers for pi.sendMessage() and pi.appendEntry()
|
||||
hookRunner.setSendMessageHandler((message, triggerTurn) => {
|
||||
const wasStreaming = this.session.isStreaming;
|
||||
this.session
|
||||
.sendHookMessage(message, triggerTurn)
|
||||
.then(() => {
|
||||
// For non-streaming cases with display=true, update UI
|
||||
// (streaming cases update via message_end event)
|
||||
if (!wasStreaming && message.display) {
|
||||
this.rebuildChatFromMessages();
|
||||
}
|
||||
})
|
||||
.catch((err) => {
|
||||
this.showError(`Hook sendMessage failed: ${err instanceof Error ? err.message : String(err)}`);
|
||||
});
|
||||
});
|
||||
hookRunner.setAppendEntryHandler((customType, data) => {
|
||||
this.sessionManager.appendCustomEntry(customType, data);
|
||||
});
|
||||
|
||||
// Show loaded hooks
|
||||
const hookPaths = hookRunner.getHookPaths();
|
||||
if (hookPaths.length > 0) {
|
||||
|
|
@ -415,11 +417,15 @@ export class InteractiveMode {
|
|||
/**
|
||||
* Emit session event to all custom tools.
|
||||
*/
|
||||
private async emitToolSessionEvent(event: ToolSessionEvent): Promise<void> {
|
||||
private async emitCustomToolSessionEvent(event: CustomToolSessionEvent): Promise<void> {
|
||||
for (const { tool } of this.customTools.values()) {
|
||||
if (tool.onSession) {
|
||||
try {
|
||||
await tool.onSession(event);
|
||||
await tool.onSession(event, {
|
||||
sessionManager: this.session.sessionManager,
|
||||
modelRegistry: this.session.modelRegistry,
|
||||
model: this.session.model,
|
||||
});
|
||||
} catch (err) {
|
||||
this.showToolError(tool.name, err instanceof Error ? err.message : String(err));
|
||||
}
|
||||
|
|
@ -436,19 +442,6 @@ export class InteractiveMode {
|
|||
this.ui.requestRender();
|
||||
}
|
||||
|
||||
/**
|
||||
* Create the UI context for hooks.
|
||||
*/
|
||||
private createHookUIContext(): HookUIContext {
|
||||
return {
|
||||
select: (title, options) => this.showHookSelector(title, options),
|
||||
confirm: (title, message) => this.showHookConfirm(title, message),
|
||||
input: (title, placeholder) => this.showHookInput(title, placeholder),
|
||||
notify: (message, type) => this.showHookNotify(message, type),
|
||||
custom: (component) => this.showHookCustom(component),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Show a selector for hooks.
|
||||
*/
|
||||
|
|
@ -861,6 +854,7 @@ export class InteractiveMode {
|
|||
this.customTools.get(content.name)?.tool,
|
||||
this.ui,
|
||||
);
|
||||
component.setExpanded(this.toolOutputExpanded);
|
||||
this.chatContainer.addChild(component);
|
||||
this.pendingTools.set(content.id, component);
|
||||
} else {
|
||||
|
|
@ -909,6 +903,7 @@ export class InteractiveMode {
|
|||
this.customTools.get(event.toolName)?.tool,
|
||||
this.ui,
|
||||
);
|
||||
component.setExpanded(this.toolOutputExpanded);
|
||||
this.chatContainer.addChild(component);
|
||||
this.pendingTools.set(event.toolCallId, component);
|
||||
this.ui.requestRender();
|
||||
|
|
@ -1158,6 +1153,7 @@ export class InteractiveMode {
|
|||
this.customTools.get(content.name)?.tool,
|
||||
this.ui,
|
||||
);
|
||||
component.setExpanded(this.toolOutputExpanded);
|
||||
this.chatContainer.addChild(component);
|
||||
|
||||
if (message.stopReason === "aborted" || message.stopReason === "error") {
|
||||
|
|
@ -1251,7 +1247,7 @@ export class InteractiveMode {
|
|||
}
|
||||
|
||||
// Emit shutdown event to custom tools
|
||||
await this.session.emitToolSessionEvent("shutdown");
|
||||
await this.session.emitCustomToolSessionEvent("shutdown");
|
||||
|
||||
this.stop();
|
||||
process.exit(0);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue