getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-18 02:03:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(coding-agent): support dynamic tool registration and tool prompt snippets closes #1720

Browse source
This commit is contained in:
Mario Zechner 2026-03-02 22:32:07 +01:00
parent ca5510158d
commit bc2fa8d6d0
12 changed files with 285 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

122
packages/coding-agent/src/core/agent-session.ts
View file

@ -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> {