mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 18:01:22 +00:00
fix(coding-agent): support dynamic tool registration and tool prompt snippets closes #1720
This commit is contained in:
parent
ca5510158d
commit
bc2fa8d6d0
12 changed files with 285 additions and 47 deletions
|
|
@ -268,6 +268,7 @@ export class AgentSession {
|
|||
|
||||
// Tool registry for extension getTools/setTools
|
||||
private _toolRegistry: Map<string, AgentTool> = new Map();
|
||||
private _toolPromptSnippets: Map<string, string> = new Map();
|
||||
|
||||
// Base system prompt (without extension appends) - used to apply fresh appends each turn
|
||||
private _baseSystemPrompt = "";
|
||||
|
|
@ -682,8 +683,25 @@ export class AgentSession {
|
|||
return this._resourceLoader.getPrompts().prompts;
|
||||
}
|
||||
|
||||
private _normalizePromptSnippet(text: string | undefined): string | undefined {
|
||||
if (!text) return undefined;
|
||||
const oneLine = text
|
||||
.replace(/[\r\n]+/g, " ")
|
||||
.replace(/\s+/g, " ")
|
||||
.trim();
|
||||
return oneLine.length > 0 ? oneLine : undefined;
|
||||
}
|
||||
|
||||
private _rebuildSystemPrompt(toolNames: string[]): string {
|
||||
const validToolNames = toolNames.filter((name) => this._baseToolRegistry.has(name));
|
||||
const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name));
|
||||
const toolSnippets: Record<string, string> = {};
|
||||
for (const name of validToolNames) {
|
||||
const snippet = this._toolPromptSnippets.get(name);
|
||||
if (snippet) {
|
||||
toolSnippets[name] = snippet;
|
||||
}
|
||||
}
|
||||
|
||||
const loaderSystemPrompt = this._resourceLoader.getSystemPrompt();
|
||||
const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt();
|
||||
const appendSystemPrompt =
|
||||
|
|
@ -698,6 +716,7 @@ export class AgentSession {
|
|||
customPrompt: loaderSystemPrompt,
|
||||
appendSystemPrompt,
|
||||
selectedTools: validToolNames,
|
||||
toolSnippets,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -1934,6 +1953,7 @@ export class AgentSession {
|
|||
getActiveTools: () => this.getActiveToolNames(),
|
||||
getAllTools: () => this.getAllTools(),
|
||||
setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames),
|
||||
refreshTools: () => this._refreshToolRegistry(),
|
||||
getCommands,
|
||||
setModel: async (model) => {
|
||||
const key = await this.modelRegistry.getApiKey(model);
|
||||
|
|
@ -1969,6 +1989,60 @@ export class AgentSession {
|
|||
);
|
||||
}
|
||||
|
||||
private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void {
|
||||
const previousRegistryNames = new Set(this._toolRegistry.keys());
|
||||
const previousActiveToolNames = this.getActiveToolNames();
|
||||
|
||||
const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
|
||||
const allCustomTools = [
|
||||
...registeredTools,
|
||||
...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
|
||||
];
|
||||
this._toolPromptSnippets = new Map(
|
||||
allCustomTools
|
||||
.map((registeredTool) => {
|
||||
const snippet = this._normalizePromptSnippet(
|
||||
registeredTool.definition.promptSnippet ?? registeredTool.definition.description,
|
||||
);
|
||||
return snippet ? ([registeredTool.definition.name, snippet] as const) : undefined;
|
||||
})
|
||||
.filter((entry): entry is readonly [string, string] => entry !== undefined),
|
||||
);
|
||||
const wrappedExtensionTools = this._extensionRunner
|
||||
? wrapRegisteredTools(allCustomTools, this._extensionRunner)
|
||||
: [];
|
||||
|
||||
const toolRegistry = new Map(this._baseToolRegistry);
|
||||
for (const tool of wrappedExtensionTools as AgentTool[]) {
|
||||
toolRegistry.set(tool.name, tool);
|
||||
}
|
||||
|
||||
if (this._extensionRunner) {
|
||||
const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
|
||||
this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
|
||||
} else {
|
||||
this._toolRegistry = toolRegistry;
|
||||
}
|
||||
|
||||
const nextActiveToolNames = options?.activeToolNames
|
||||
? [...options.activeToolNames]
|
||||
: [...previousActiveToolNames];
|
||||
|
||||
if (options?.includeAllExtensionTools) {
|
||||
for (const tool of wrappedExtensionTools) {
|
||||
nextActiveToolNames.push(tool.name);
|
||||
}
|
||||
} else if (!options?.activeToolNames) {
|
||||
for (const toolName of this._toolRegistry.keys()) {
|
||||
if (!previousRegistryNames.has(toolName)) {
|
||||
nextActiveToolNames.push(toolName);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
this.setActiveToolsByName([...new Set(nextActiveToolNames)]);
|
||||
}
|
||||
|
||||
private _buildRuntime(options: {
|
||||
activeToolNames?: string[];
|
||||
flagValues?: Map<string, boolean | string>;
|
||||
|
|
@ -2012,52 +2086,14 @@ export class AgentSession {
|
|||
this._applyExtensionBindings(this._extensionRunner);
|
||||
}
|
||||
|
||||
const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
|
||||
const allCustomTools = [
|
||||
...registeredTools,
|
||||
...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
|
||||
];
|
||||
const wrappedExtensionTools = this._extensionRunner
|
||||
? wrapRegisteredTools(allCustomTools, this._extensionRunner)
|
||||
: [];
|
||||
|
||||
const toolRegistry = new Map(this._baseToolRegistry);
|
||||
for (const tool of wrappedExtensionTools as AgentTool[]) {
|
||||
toolRegistry.set(tool.name, tool);
|
||||
}
|
||||
|
||||
const defaultActiveToolNames = this._baseToolsOverride
|
||||
? Object.keys(this._baseToolsOverride)
|
||||
: ["read", "bash", "edit", "write"];
|
||||
const baseActiveToolNames = options.activeToolNames ?? defaultActiveToolNames;
|
||||
const activeToolNameSet = new Set<string>(baseActiveToolNames);
|
||||
if (options.includeAllExtensionTools) {
|
||||
for (const tool of wrappedExtensionTools as AgentTool[]) {
|
||||
activeToolNameSet.add(tool.name);
|
||||
}
|
||||
}
|
||||
|
||||
const extensionToolNames = new Set(wrappedExtensionTools.map((tool) => tool.name));
|
||||
const activeBaseTools = Array.from(activeToolNameSet)
|
||||
.filter((name) => this._baseToolRegistry.has(name) && !extensionToolNames.has(name))
|
||||
.map((name) => this._baseToolRegistry.get(name) as AgentTool);
|
||||
const activeExtensionTools = wrappedExtensionTools.filter((tool) => activeToolNameSet.has(tool.name));
|
||||
const activeToolsArray: AgentTool[] = [...activeBaseTools, ...activeExtensionTools];
|
||||
|
||||
if (this._extensionRunner) {
|
||||
const wrappedActiveTools = wrapToolsWithExtensions(activeToolsArray, this._extensionRunner);
|
||||
this.agent.setTools(wrappedActiveTools as AgentTool[]);
|
||||
|
||||
const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
|
||||
this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
|
||||
} else {
|
||||
this.agent.setTools(activeToolsArray);
|
||||
this._toolRegistry = toolRegistry;
|
||||
}
|
||||
|
||||
const systemPromptToolNames = Array.from(activeToolNameSet).filter((name) => this._baseToolRegistry.has(name));
|
||||
this._baseSystemPrompt = this._rebuildSystemPrompt(systemPromptToolNames);
|
||||
this.agent.setSystemPrompt(this._baseSystemPrompt);
|
||||
this._refreshToolRegistry({
|
||||
activeToolNames: baseActiveToolNames,
|
||||
includeAllExtensionTools: options.includeAllExtensionTools,
|
||||
});
|
||||
}
|
||||
|
||||
async reload(): Promise<void> {
|
||||
|
|
|
|||
|
|
@ -119,6 +119,8 @@ export function createExtensionRuntime(): ExtensionRuntime {
|
|||
getActiveTools: notInitialized,
|
||||
getAllTools: notInitialized,
|
||||
setActiveTools: notInitialized,
|
||||
// registerTool() is valid during extension load; refresh is only needed post-bind.
|
||||
refreshTools: () => {},
|
||||
getCommands: notInitialized,
|
||||
setModel: () => Promise.reject(new Error("Extension runtime not initialized")),
|
||||
getThinkingLevel: notInitialized,
|
||||
|
|
@ -162,6 +164,7 @@ function createExtensionAPI(
|
|||
definition: tool,
|
||||
extensionPath: extension.path,
|
||||
});
|
||||
runtime.refreshTools();
|
||||
},
|
||||
|
||||
registerCommand(name: string, options: Omit<RegisteredCommand, "name">): void {
|
||||
|
|
|
|||
|
|
@ -244,6 +244,7 @@ export class ExtensionRunner {
|
|||
this.runtime.getActiveTools = actions.getActiveTools;
|
||||
this.runtime.getAllTools = actions.getAllTools;
|
||||
this.runtime.setActiveTools = actions.setActiveTools;
|
||||
this.runtime.refreshTools = actions.refreshTools;
|
||||
this.runtime.getCommands = actions.getCommands;
|
||||
this.runtime.setModel = actions.setModel;
|
||||
this.runtime.getThinkingLevel = actions.getThinkingLevel;
|
||||
|
|
|
|||
|
|
@ -339,6 +339,8 @@ export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = un
|
|||
label: string;
|
||||
/** Description for LLM */
|
||||
description: string;
|
||||
/** Optional one-line snippet for the Available tools section in the default system prompt. Falls back to description when omitted. */
|
||||
promptSnippet?: string;
|
||||
/** Parameter schema (TypeBox) */
|
||||
parameters: TParams;
|
||||
|
||||
|
|
@ -1251,6 +1253,8 @@ export type GetCommandsHandler = () => SlashCommandInfo[];
|
|||
|
||||
export type SetActiveToolsHandler = (toolNames: string[]) => void;
|
||||
|
||||
export type RefreshToolsHandler = () => void;
|
||||
|
||||
export type SetModelHandler = (model: Model<any>) => Promise<boolean>;
|
||||
|
||||
export type GetThinkingLevelHandler = () => ThinkingLevel;
|
||||
|
|
@ -1291,6 +1295,7 @@ export interface ExtensionActions {
|
|||
getActiveTools: GetActiveToolsHandler;
|
||||
getAllTools: GetAllToolsHandler;
|
||||
setActiveTools: SetActiveToolsHandler;
|
||||
refreshTools: RefreshToolsHandler;
|
||||
getCommands: GetCommandsHandler;
|
||||
setModel: SetModelHandler;
|
||||
getThinkingLevel: GetThinkingLevelHandler;
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ export interface BuildSystemPromptOptions {
|
|||
customPrompt?: string;
|
||||
/** Tools to include in prompt. Default: [read, bash, edit, write] */
|
||||
selectedTools?: string[];
|
||||
/** Optional one-line tool snippets keyed by tool name. */
|
||||
toolSnippets?: Record<string, string>;
|
||||
/** Text to append to system prompt. */
|
||||
appendSystemPrompt?: string;
|
||||
/** Working directory. Default: process.cwd() */
|
||||
|
|
@ -36,6 +38,7 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin
|
|||
const {
|
||||
customPrompt,
|
||||
selectedTools,
|
||||
toolSnippets,
|
||||
appendSystemPrompt,
|
||||
cwd,
|
||||
contextFiles: providedContextFiles,
|
||||
|
|
@ -94,9 +97,18 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin
|
|||
const docsPath = getDocsPath();
|
||||
const examplesPath = getExamplesPath();
|
||||
|
||||
// Build tools list based on selected tools (only built-in tools with known descriptions)
|
||||
const tools = (selectedTools || ["read", "bash", "edit", "write"]).filter((t) => t in toolDescriptions);
|
||||
const toolsList = tools.length > 0 ? tools.map((t) => `- ${t}: ${toolDescriptions[t]}`).join("\n") : "(none)";
|
||||
// Build tools list based on selected tools.
|
||||
// Built-ins use toolDescriptions. Custom tools can provide one-line snippets.
|
||||
const tools = selectedTools || ["read", "bash", "edit", "write"];
|
||||
const toolsList =
|
||||
tools.length > 0
|
||||
? tools
|
||||
.map((name) => {
|
||||
const snippet = toolSnippets?.[name] ?? toolDescriptions[name] ?? name;
|
||||
return `- ${name}: ${snippet}`;
|
||||
})
|
||||
.join("\n")
|
||||
: "(none)";
|
||||
|
||||
// Build guidelines based on which tools are actually available
|
||||
const guidelinesList: string[] = [];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue