tools system - web-fetch
This commit is contained in:
@@ -0,0 +1,237 @@
|
||||
import type { ExtensionAPI } from '@mariozechner/pi-coding-agent';
|
||||
import { Type, type TSchema } from '@sinclair/typebox';
|
||||
import { readdirSync, existsSync, readFileSync } from 'node:fs';
|
||||
import { join } from 'node:path';
|
||||
|
||||
type ToolParamType = 'string' | 'number' | 'boolean' | 'enum';
|
||||
|
||||
type ToolParam = {
|
||||
type: ToolParamType;
|
||||
description: string;
|
||||
values?: string[];
|
||||
default?: unknown;
|
||||
optional?: boolean;
|
||||
};
|
||||
|
||||
type ToolMeta = {
|
||||
name: string;
|
||||
label: string;
|
||||
description: string;
|
||||
language: 'typescript' | 'bash' | 'python';
|
||||
inputs: Record<string, ToolParam>;
|
||||
};
|
||||
|
||||
function parseFrontmatter(content: string): { meta: Partial<ToolMeta>; body: string } {
|
||||
const match = content.match(/^---\n([\s\S]*?)\n---\n?([\s\S]*)$/);
|
||||
if (!match) return { meta: {}, body: content };
|
||||
|
||||
const yamlBlock = match[1]!;
|
||||
const body = match[2]!;
|
||||
const meta: Record<string, unknown> = {};
|
||||
|
||||
const lines = yamlBlock.split('\n');
|
||||
let currentKey: string | null = null;
|
||||
let currentObj: Record<string, unknown> | null = null;
|
||||
let currentSubKey: string | null = null;
|
||||
let currentSubObj: Record<string, unknown> | null = null;
|
||||
|
||||
for (const line of lines) {
|
||||
const topMatch = line.match(/^(\w[\w-]*):\s*(.*)$/);
|
||||
if (topMatch && !line.startsWith(' ')) {
|
||||
if (currentSubObj && currentSubKey && currentObj) {
|
||||
currentObj[currentSubKey] = currentSubObj;
|
||||
currentSubObj = null;
|
||||
currentSubKey = null;
|
||||
}
|
||||
if (currentObj && currentKey) {
|
||||
meta[currentKey] = currentObj;
|
||||
currentObj = null;
|
||||
currentKey = null;
|
||||
}
|
||||
const [, key, value] = topMatch;
|
||||
if (!value || value.trim() === '') {
|
||||
currentKey = key!;
|
||||
currentObj = {};
|
||||
} else {
|
||||
meta[key!] = value.trim();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const midMatch = line.match(/^ (\w[\w-]*):\s*(.*)$/);
|
||||
if (midMatch && currentObj !== null) {
|
||||
if (currentSubObj && currentSubKey) {
|
||||
currentObj[currentSubKey] = currentSubObj;
|
||||
currentSubObj = null;
|
||||
currentSubKey = null;
|
||||
}
|
||||
const [, key, value] = midMatch;
|
||||
if (!value || value.trim() === '') {
|
||||
currentSubKey = key!;
|
||||
currentSubObj = {};
|
||||
} else {
|
||||
currentObj[key!] = value.trim();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const deepMatch = line.match(/^ (\w[\w-]*):\s*(.*)$/);
|
||||
if (deepMatch && currentSubObj !== null) {
|
||||
const [, key, value] = deepMatch;
|
||||
currentSubObj[key!] = value!.trim();
|
||||
continue;
|
||||
}
|
||||
|
||||
const arrayMatch = line.match(/^ - (.+)$/);
|
||||
if (arrayMatch && currentSubObj !== null) {
|
||||
const key = Object.keys(currentSubObj).at(-1);
|
||||
if (key) {
|
||||
const arr = currentSubObj[key];
|
||||
if (Array.isArray(arr)) {
|
||||
arr.push(arrayMatch[1]!.trim());
|
||||
} else {
|
||||
currentSubObj[key] = [arrayMatch[1]!.trim()];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (currentSubObj && currentSubKey && currentObj) {
|
||||
currentObj[currentSubKey] = currentSubObj;
|
||||
}
|
||||
if (currentObj && currentKey) {
|
||||
meta[currentKey] = currentObj;
|
||||
}
|
||||
|
||||
return { meta: meta as Partial<ToolMeta>, body };
|
||||
}
|
||||
|
||||
function buildSchema(inputs: Record<string, ToolParam>): TSchema {
|
||||
const props: Record<string, TSchema> = {};
|
||||
|
||||
for (const [paramName, param] of Object.entries(inputs)) {
|
||||
let schema: TSchema;
|
||||
|
||||
switch (param.type) {
|
||||
case 'enum': {
|
||||
const values = param.values ?? [];
|
||||
schema = Type.Union(values.map((v) => Type.Literal(v)), {
|
||||
description: param.description,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case 'number':
|
||||
schema = Type.Number({ description: param.description });
|
||||
break;
|
||||
case 'boolean':
|
||||
schema = Type.Boolean({ description: param.description });
|
||||
break;
|
||||
default:
|
||||
schema = Type.String({ description: param.description });
|
||||
}
|
||||
|
||||
props[paramName] = param.optional ? Type.Optional(schema) : schema;
|
||||
}
|
||||
|
||||
return Type.Object(props);
|
||||
}
|
||||
|
||||
function discoverTools(dir: string): Array<{ toolDir: string; entryFile: string; meta: ToolMeta }> {
|
||||
if (!existsSync(dir)) return [];
|
||||
|
||||
const discovered: Array<{ toolDir: string; entryFile: string; meta: ToolMeta }> = [];
|
||||
const entries = readdirSync(dir, { withFileTypes: true });
|
||||
|
||||
for (const entry of entries) {
|
||||
if (!entry.isDirectory()) continue;
|
||||
|
||||
const toolDir = join(dir, entry.name);
|
||||
const toolMdPath = join(toolDir, 'TOOL.md');
|
||||
if (!existsSync(toolMdPath)) continue;
|
||||
|
||||
const indexTs = join(toolDir, 'index.ts');
|
||||
const indexJs = join(toolDir, 'index.js');
|
||||
const entryFile = existsSync(indexTs) ? indexTs : existsSync(indexJs) ? indexJs : null;
|
||||
if (!entryFile) {
|
||||
console.warn(`[tool-loader] Skipping ${entry.name}: no index.ts or index.js found`);
|
||||
continue;
|
||||
}
|
||||
|
||||
const content = readFileSync(toolMdPath, 'utf-8');
|
||||
const { meta } = parseFrontmatter(content);
|
||||
|
||||
if (!meta.name || !meta.description) {
|
||||
console.warn(`[tool-loader] Skipping ${entry.name}: missing name or description in TOOL.md`);
|
||||
continue;
|
||||
}
|
||||
|
||||
discovered.push({ toolDir, entryFile, meta: meta as ToolMeta });
|
||||
}
|
||||
|
||||
return discovered;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const rawDirs = process.env.PI_TOOLS_DIRS ?? '';
|
||||
const toolDirs = rawDirs.split(':').filter(Boolean);
|
||||
|
||||
if (toolDirs.length === 0) {
|
||||
console.warn('[tool-loader] PI_TOOLS_DIRS not set — no custom tools will be loaded');
|
||||
return;
|
||||
}
|
||||
|
||||
// Register synchronously in the factory function so tools appear in the system prompt.
|
||||
// Implementations are lazy-loaded on first call to avoid async import issues at startup.
|
||||
const seen = new Set<string>();
|
||||
|
||||
for (const dir of toolDirs) {
|
||||
const tools = discoverTools(dir);
|
||||
|
||||
for (const { entryFile, meta } of tools) {
|
||||
// User dirs come after global — last writer wins, so skip if already registered
|
||||
if (seen.has(meta.name)) continue;
|
||||
seen.add(meta.name);
|
||||
|
||||
const schema = buildSchema(meta.inputs ?? {});
|
||||
|
||||
// Capture entryFile in closure for lazy load
|
||||
const capturedEntry = entryFile;
|
||||
|
||||
pi.registerTool({
|
||||
name: meta.name,
|
||||
label: meta.label ?? meta.name,
|
||||
description: meta.description,
|
||||
parameters: schema,
|
||||
|
||||
async execute(toolCallId, params, signal, onUpdate, ctx) {
|
||||
// Lazy-load the implementation on first actual call.
|
||||
// Dynamic import works here because we're already in an async tool execution
|
||||
// context — jiti/Node has had time to set up its module hooks.
|
||||
let executeFn: Function | undefined;
|
||||
try {
|
||||
const mod = await import(capturedEntry);
|
||||
executeFn = mod.execute ?? mod.default?.execute;
|
||||
} catch (err) {
|
||||
return {
|
||||
content: [{ type: 'text', text: `[tool-loader] Failed to load ${meta.name}: ${String(err)}` }],
|
||||
details: { error: String(err) },
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
if (typeof executeFn !== 'function') {
|
||||
return {
|
||||
content: [{ type: 'text', text: `[tool-loader] ${meta.name}/index.ts must export an "execute" function` }],
|
||||
details: {},
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
return executeFn(toolCallId, params, signal, onUpdate, ctx);
|
||||
},
|
||||
});
|
||||
|
||||
console.log(`[tool-loader] Registered tool: ${meta.name} (${capturedEntry})`);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user