platform/seed/extensions/tool-loader/index.ts

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': {
        // values can be a string[] from deeper YAML parsing,
        // or a comma-separated string like "single,batch" from flat YAML
        const raw = param.values;
        const values = Array.isArray(raw)
          ? raw
          : typeof raw === 'string'
            ? raw.split(',').map((v) => v.trim())
            : [];
        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})`);
    }
  }
}