mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 04:01:56 +00:00
feat(coding-agent): add tool override support via extensions
- Add setActiveTools() to ExtensionAPI for dynamic tool management - Extensions can now override, wrap, or disable built-in tools - Add tool-override.ts example demonstrating the pattern - Update documentation for tool override capabilities
This commit is contained in:
parent
7a2c19cdf0
commit
e3dd4f21d1
10 changed files with 211 additions and 61 deletions
|
|
@ -4,7 +4,7 @@
|
|||
|
||||
### Added
|
||||
|
||||
- `--no-tools` flag to disable all built-in tools, allowing extension-only tool setups ([#555](https://github.com/badlogic/pi-mono/issues/555))
|
||||
- `--no-tools` flag to disable all built-in tools, allowing extension-only tool setups ([#557](https://github.com/badlogic/pi-mono/pull/557) by [@cv](https://github.com/cv))
|
||||
|
||||
### Fixed
|
||||
|
||||
|
|
|
|||
|
|
@ -1122,6 +1122,7 @@ pi [options] [@files...] [messages...]
|
|||
| `--continue`, `-c` | Continue most recent session |
|
||||
| `--resume`, `-r` | Select session to resume |
|
||||
| `--models <patterns>` | Comma-separated patterns for Ctrl+P cycling. Supports glob patterns (e.g., `anthropic/*`, `*sonnet*:high`) and fuzzy matching (e.g., `sonnet,haiku:low`) |
|
||||
| `--no-tools` | Disable all built-in tools (use with `-e` for extension-only setups) |
|
||||
| `--tools <tools>` | Comma-separated tool list (default: `read,bash,edit,write`) |
|
||||
| `--thinking <level>` | Thinking level: `off`, `minimal`, `low`, `medium`, `high` |
|
||||
| `--extension <path>`, `-e` | Load an extension file (can be used multiple times) |
|
||||
|
|
|
|||
|
|
@ -946,6 +946,17 @@ pi.registerTool({
|
|||
|
||||
Extensions can override built-in tools (`read`, `bash`, `edit`, `write`, `grep`, `find`, `ls`) by registering a tool with the same name. Interactive mode displays a warning when this happens.
|
||||
|
||||
Use `--no-tools` to start without built-in tools, then add back specific ones with `--tools`:
|
||||
```bash
|
||||
# No built-in tools, only extension tools
|
||||
pi --no-tools -e ./my-extension.ts
|
||||
|
||||
# No built-in read, use extension's read, keep other built-ins
|
||||
pi -e ./tool-override.ts
|
||||
```
|
||||
|
||||
See [examples/extensions/tool-override.ts](../examples/extensions/tool-override.ts) for a complete example that overrides `read` with logging and access control.
|
||||
|
||||
**Your implementation must match the exact result shape**, including the `details` type. The UI and session logic depend on these shapes for rendering and state tracking.
|
||||
|
||||
Built-in tool implementations:
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/
|
|||
| `todo.ts` | Todo list tool + `/todos` command with custom rendering and state persistence |
|
||||
| `hello.ts` | Minimal custom tool example |
|
||||
| `question.ts` | Demonstrates `ctx.ui.select()` for asking the user questions |
|
||||
| `tool-override.ts` | Override built-in tools (e.g., add logging/access control to `read`) |
|
||||
| `subagent/` | Delegate tasks to specialized subagents with isolated context windows |
|
||||
|
||||
### Commands & UI
|
||||
|
|
|
|||
178
packages/coding-agent/examples/extensions/tool-override.ts
Normal file
178
packages/coding-agent/examples/extensions/tool-override.ts
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
/**
|
||||
* Tool Override Example - Demonstrates overriding built-in tools
|
||||
*
|
||||
* Extensions can register tools with the same name as built-in tools to replace them.
|
||||
* This is useful for:
|
||||
* - Adding logging or auditing to tool calls
|
||||
* - Implementing access control or sandboxing
|
||||
* - Routing tool calls to remote systems (e.g., pi-ssh-remote)
|
||||
* - Modifying tool behavior for specific workflows
|
||||
*
|
||||
* This example overrides the `read` tool to:
|
||||
* 1. Log all file access to a log file
|
||||
* 2. Block access to sensitive paths (e.g., .env files)
|
||||
* 3. Delegate to the original read implementation for allowed files
|
||||
*
|
||||
* Usage:
|
||||
* pi --no-tools --tools bash,edit,write -e ./tool-override.ts
|
||||
*
|
||||
* The --no-tools flag disables all built-in tools, then --tools adds back the ones
|
||||
* we want (excluding read, which we're overriding). The extension provides our
|
||||
* custom read implementation.
|
||||
*
|
||||
* Alternatively, without --no-tools the extension's read tool will override the
|
||||
* built-in read tool automatically.
|
||||
*/
|
||||
|
||||
import type { TextContent } from "@mariozechner/pi-ai";
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import { Text } from "@mariozechner/pi-tui";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { appendFileSync, constants, readFileSync } from "fs";
|
||||
import { access, readFile } from "fs/promises";
|
||||
import { homedir } from "os";
|
||||
import { join, resolve } from "path";
|
||||
|
||||
const LOG_FILE = join(homedir(), ".pi", "agent", "read-access.log");
|
||||
|
||||
// Paths that are blocked from reading
|
||||
const BLOCKED_PATTERNS = [
|
||||
/\.env$/,
|
||||
/\.env\..+$/,
|
||||
/secrets?\.(json|yaml|yml|toml)$/i,
|
||||
/credentials?\.(json|yaml|yml|toml)$/i,
|
||||
/\/\.ssh\//,
|
||||
/\/\.aws\//,
|
||||
/\/\.gnupg\//,
|
||||
];
|
||||
|
||||
function isBlockedPath(path: string): boolean {
|
||||
return BLOCKED_PATTERNS.some((pattern) => pattern.test(path));
|
||||
}
|
||||
|
||||
function logAccess(path: string, allowed: boolean, reason?: string) {
|
||||
const timestamp = new Date().toISOString();
|
||||
const status = allowed ? "ALLOWED" : "BLOCKED";
|
||||
const msg = reason ? ` (${reason})` : "";
|
||||
const line = `[${timestamp}] ${status}: ${path}${msg}\n`;
|
||||
|
||||
try {
|
||||
appendFileSync(LOG_FILE, line);
|
||||
} catch {
|
||||
// Ignore logging errors
|
||||
}
|
||||
}
|
||||
|
||||
const readSchema = Type.Object({
|
||||
path: Type.String({ description: "Path to the file to read (relative or absolute)" }),
|
||||
offset: Type.Optional(Type.Number({ description: "Line number to start reading from (1-indexed)" })),
|
||||
limit: Type.Optional(Type.Number({ description: "Maximum number of lines to read" })),
|
||||
});
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
pi.registerTool({
|
||||
name: "read", // Same name as built-in - this will override it
|
||||
label: "read (audited)",
|
||||
description:
|
||||
"Read the contents of a file with access logging. Some sensitive paths (.env, secrets, credentials) are blocked.",
|
||||
parameters: readSchema,
|
||||
|
||||
async execute(_toolCallId, params, _onUpdate, ctx) {
|
||||
const { path, offset, limit } = params;
|
||||
const absolutePath = resolve(ctx.cwd, path);
|
||||
|
||||
// Check if path is blocked
|
||||
if (isBlockedPath(absolutePath)) {
|
||||
logAccess(absolutePath, false, "matches blocked pattern");
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: `Access denied: "${path}" matches a blocked pattern (sensitive file). This tool blocks access to .env files, secrets, credentials, and SSH/AWS/GPG directories.`,
|
||||
},
|
||||
],
|
||||
details: { blocked: true },
|
||||
};
|
||||
}
|
||||
|
||||
// Log allowed access
|
||||
logAccess(absolutePath, true);
|
||||
|
||||
// Perform the actual read (simplified implementation)
|
||||
try {
|
||||
await access(absolutePath, constants.R_OK);
|
||||
const content = await readFile(absolutePath, "utf-8");
|
||||
const lines = content.split("\n");
|
||||
|
||||
// Apply offset and limit
|
||||
const startLine = offset ? Math.max(0, offset - 1) : 0;
|
||||
const endLine = limit ? startLine + limit : lines.length;
|
||||
const selectedLines = lines.slice(startLine, endLine);
|
||||
|
||||
// Basic truncation (50KB limit)
|
||||
let text = selectedLines.join("\n");
|
||||
const maxBytes = 50 * 1024;
|
||||
if (Buffer.byteLength(text, "utf-8") > maxBytes) {
|
||||
text = `${text.slice(0, maxBytes)}\n\n[Output truncated at 50KB]`;
|
||||
}
|
||||
|
||||
return {
|
||||
content: [{ type: "text", text }] as TextContent[],
|
||||
details: { lines: lines.length },
|
||||
};
|
||||
} catch (error: any) {
|
||||
return {
|
||||
content: [{ type: "text", text: `Error reading file: ${error.message}` }] as TextContent[],
|
||||
details: { error: true },
|
||||
};
|
||||
}
|
||||
},
|
||||
|
||||
// Custom rendering to show it's the audited version
|
||||
renderCall(args, theme) {
|
||||
return new Text(theme.fg("toolTitle", theme.bold("read ")) + theme.fg("accent", args.path), 0, 0);
|
||||
},
|
||||
|
||||
renderResult(result, { expanded }, theme) {
|
||||
const content = result.content[0];
|
||||
if (content?.type === "text" && content.text.startsWith("Access denied")) {
|
||||
return new Text(theme.fg("error", "Access denied (sensitive file)"), 0, 0);
|
||||
}
|
||||
if (content?.type === "text" && content.text.startsWith("Error")) {
|
||||
return new Text(theme.fg("error", content.text), 0, 0);
|
||||
}
|
||||
|
||||
// Show preview of content
|
||||
if (content?.type === "text") {
|
||||
const lines = content.text.split("\n");
|
||||
const preview = lines.slice(0, expanded ? 10 : 3);
|
||||
let text = theme.fg("success", `Read ${lines.length} lines`);
|
||||
if (expanded) {
|
||||
for (const line of preview) {
|
||||
text += `\n${theme.fg("dim", line.slice(0, 100))}`;
|
||||
}
|
||||
if (lines.length > 10) {
|
||||
text += `\n${theme.fg("muted", `... ${lines.length - 10} more lines`)}`;
|
||||
}
|
||||
}
|
||||
return new Text(text, 0, 0);
|
||||
}
|
||||
|
||||
return new Text(theme.fg("dim", "Read complete"), 0, 0);
|
||||
},
|
||||
});
|
||||
|
||||
// Also register a command to view the access log
|
||||
pi.registerCommand("read-log", {
|
||||
description: "View the file access log",
|
||||
handler: async (_args, ctx) => {
|
||||
try {
|
||||
const log = readFileSync(LOG_FILE, "utf-8");
|
||||
const lines = log.trim().split("\n").slice(-20); // Last 20 entries
|
||||
ctx.ui.notify(`Recent file access:\n${lines.join("\n")}`, "info");
|
||||
} catch {
|
||||
ctx.ui.notify("No access log found", "info");
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
|
|
@ -90,24 +90,18 @@ export function parseArgs(args: string[], extensionFlags?: Map<string, { type: "
|
|||
} else if (arg === "--no-tools") {
|
||||
result.noTools = true;
|
||||
} else if (arg === "--tools" && i + 1 < args.length) {
|
||||
const toolsArg = args[++i];
|
||||
// Handle empty string as no tools (e.g., --tools '')
|
||||
if (toolsArg === "") {
|
||||
result.tools = [];
|
||||
} else {
|
||||
const toolNames = toolsArg.split(",").map((s) => s.trim());
|
||||
const validTools: ToolName[] = [];
|
||||
for (const name of toolNames) {
|
||||
if (name in allTools) {
|
||||
validTools.push(name as ToolName);
|
||||
} else {
|
||||
console.error(
|
||||
chalk.yellow(`Warning: Unknown tool "${name}". Valid tools: ${Object.keys(allTools).join(", ")}`),
|
||||
);
|
||||
}
|
||||
const toolNames = args[++i].split(",").map((s) => s.trim());
|
||||
const validTools: ToolName[] = [];
|
||||
for (const name of toolNames) {
|
||||
if (name in allTools) {
|
||||
validTools.push(name as ToolName);
|
||||
} else {
|
||||
console.error(
|
||||
chalk.yellow(`Warning: Unknown tool "${name}". Valid tools: ${Object.keys(allTools).join(", ")}`),
|
||||
);
|
||||
}
|
||||
result.tools = validTools;
|
||||
}
|
||||
result.tools = validTools;
|
||||
} else if (arg === "--thinking" && i + 1 < args.length) {
|
||||
const level = args[++i];
|
||||
if (isValidThinkingLevel(level)) {
|
||||
|
|
@ -188,7 +182,7 @@ ${chalk.bold("Options:")}
|
|||
Available: read, bash, edit, write, grep, find, ls
|
||||
--thinking <level> Set thinking level: off, minimal, low, medium, high, xhigh
|
||||
--extension, -e <path> Load an extension file (can be used multiple times)
|
||||
--no-extensions Disable extensions discovery and loading
|
||||
--no-extensions Disable extension discovery (explicit -e paths still work)
|
||||
--no-skills Disable skills discovery and loading
|
||||
--skills <patterns> Comma-separated glob patterns to filter skills (e.g., git-*,docker)
|
||||
--export <file> Export session file to HTML and exit
|
||||
|
|
|
|||
|
|
@ -514,7 +514,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
|||
}
|
||||
|
||||
// Initially active tools = active built-in + extension tools
|
||||
let activeToolsArray: Tool[] = [...initialActiveBuiltInTools, ...wrappedExtensionTools];
|
||||
// Extension tools can override built-in tools with the same name
|
||||
const extensionToolNames = new Set(wrappedExtensionTools.map((t) => t.name));
|
||||
const nonOverriddenBuiltInTools = initialActiveBuiltInTools.filter((t) => !extensionToolNames.has(t.name));
|
||||
let activeToolsArray: Tool[] = [...nonOverriddenBuiltInTools, ...wrappedExtensionTools];
|
||||
time("combineTools");
|
||||
|
||||
// Wrap tools with extensions if available
|
||||
|
|
|
|||
|
|
@ -222,12 +222,6 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin
|
|||
const hasLs = tools.includes("ls");
|
||||
const hasRead = tools.includes("read");
|
||||
|
||||
// Read-only mode notice (only if we have some read-only tools but no write tools)
|
||||
// Skip this if there are no built-in tools at all (extensions may provide write capabilities)
|
||||
if (tools.length > 0 && !hasBash && !hasEdit && !hasWrite) {
|
||||
guidelinesList.push("You are in READ-ONLY mode - you cannot modify files or execute arbitrary commands");
|
||||
}
|
||||
|
||||
// Bash without edit/write = read-only bash mode
|
||||
if (hasBash && !hasEdit && !hasWrite) {
|
||||
guidelinesList.push(
|
||||
|
|
@ -266,9 +260,7 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin
|
|||
|
||||
// Always include these
|
||||
guidelinesList.push("Be concise in your responses");
|
||||
if (tools.length > 0) {
|
||||
guidelinesList.push("Show file paths clearly when working with files");
|
||||
}
|
||||
guidelinesList.push("Show file paths clearly when working with files");
|
||||
|
||||
const guidelines = guidelinesList.map((g) => `- ${g}`).join("\n");
|
||||
|
||||
|
|
|
|||
|
|
@ -181,11 +181,6 @@ describe("parseArgs", () => {
|
|||
expect(result.noTools).toBe(true);
|
||||
expect(result.tools).toEqual(["read", "bash"]);
|
||||
});
|
||||
|
||||
test("parses --tools with empty string", () => {
|
||||
const result = parseArgs(["--tools", ""]);
|
||||
expect(result.tools).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("messages and file args", () => {
|
||||
|
|
|
|||
|
|
@ -3,18 +3,6 @@ import { buildSystemPrompt } from "../src/core/system-prompt.js";
|
|||
|
||||
describe("buildSystemPrompt", () => {
|
||||
describe("empty tools", () => {
|
||||
test("does not show READ-ONLY mode when no built-in tools", () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
selectedTools: [],
|
||||
contextFiles: [],
|
||||
skills: [],
|
||||
});
|
||||
|
||||
// Should not mention READ-ONLY mode when there are no tools
|
||||
// (extensions may provide write capabilities)
|
||||
expect(prompt).not.toContain("READ-ONLY mode");
|
||||
});
|
||||
|
||||
test("shows (none) for empty tools list", () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
selectedTools: [],
|
||||
|
|
@ -25,37 +13,24 @@ describe("buildSystemPrompt", () => {
|
|||
expect(prompt).toContain("Available tools:\n(none)");
|
||||
});
|
||||
|
||||
test("does not show file paths guideline when no tools", () => {
|
||||
test("shows file paths guideline even with no tools", () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
selectedTools: [],
|
||||
contextFiles: [],
|
||||
skills: [],
|
||||
});
|
||||
|
||||
expect(prompt).not.toContain("Show file paths clearly");
|
||||
});
|
||||
});
|
||||
|
||||
describe("read-only tools", () => {
|
||||
test("shows READ-ONLY mode when only read tools available", () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
selectedTools: ["read", "grep", "find", "ls"],
|
||||
contextFiles: [],
|
||||
skills: [],
|
||||
});
|
||||
|
||||
expect(prompt).toContain("READ-ONLY mode");
|
||||
expect(prompt).toContain("Show file paths clearly");
|
||||
});
|
||||
});
|
||||
|
||||
describe("default tools", () => {
|
||||
test("does not show READ-ONLY mode with default tools", () => {
|
||||
test("includes all default tools", () => {
|
||||
const prompt = buildSystemPrompt({
|
||||
contextFiles: [],
|
||||
skills: [],
|
||||
});
|
||||
|
||||
expect(prompt).not.toContain("READ-ONLY mode");
|
||||
expect(prompt).toContain("- read:");
|
||||
expect(prompt).toContain("- bash:");
|
||||
expect(prompt).toContain("- edit:");
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue