@amalgm/chat 0.1.2 → 0.2.0

package/host/adapters/pi.js

@@ -0,0 +1,519 @@
+// Pi --mode rpc JSONL-over-stdio adapter.
+//
+// Laws come from ../../dist: event constructors via `createEventConstructors`,
+// the normalizer via `createPiNormalizer` (`normalizePiEvent`,
+// `piSessionStatsUsage`).
+
+import { spawn, spawnSync } from 'node:child_process';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { createEventConstructors } from '../../dist/events/index.js';
+import { createPiNormalizer } from '../../dist/normalizers/pi.js';
+import { allowed, requestPermission } from '../permissions.js';
+import { recordNativeEvent } from '../recorder.js';
+import { relayedMcpServers } from '../tooling/mcp-bundle.js';
+import { executableExists, findOnPath } from '../tooling/native-binaries.js';
+import { prepareHarnessRuntime } from '../tooling/runtime-home.js';
+import { composeSystemPrompt } from '../tooling/system-prompt.js';
+import {
+  COMMON_REASONING_SETTINGS,
+  nativeInputCapabilities,
+  UnsupportedExecutionInputError,
+} from './input-capabilities.js';
+import { piPrompt } from './prompt.js';
+
+const events = createEventConstructors({ clockMs: () => Date.now() });
+const { done, errorEvent } = events;
+const { normalizePiEvent, piSessionStatsUsage } = createPiNormalizer(events);
+
+const GATEWAY_PROVIDER = 'vercel-ai-gateway';
+
+// Pi's RPC protocol is JSONL over stdio, but NOT JSON-RPC: commands are
+// {id?, type, ...}, responses are {id?, type:'response', command, success, data},
+// and events are un-id'd JSON lines. Framing is LF-only by spec.
+class PiRpc {
+  constructor({ binary, args, cwd, env }) {
+    this.id = 1;
+    this.pending = new Map();
+    this.eventHandlers = new Set();
+    this.extensionUiHandler = null;
+    this.buffer = '';
+    this.closed = false;
+    this.closeError = null;
+    this.child = spawn(binary, args, {
+      cwd,
+      env,
+      stdio: ['pipe', 'pipe', 'pipe'],
+      detached: process.platform !== 'win32',
+      windowsHide: true,
+    });
+    this.child.stdout.setEncoding('utf8');
+    this.child.stdout.on('data', (chunk) => this.onData(chunk));
+    this.child.stderr.setEncoding('utf8');
+    this.child.stderr.on('data', (chunk) => {
+      const text = String(chunk || '').trim();
+      if (text && process.env.AMALGM_CHAT_DEBUG) console.warn('[Pi]', text.slice(0, 1000));
+    });
+    this.child.on('error', (err) => this.failAll(err));
+    this.child.on('exit', (code, signal) => {
+      this.failAll(new Error(`pi rpc agent exited (${code ?? signal ?? 'unknown'})`));
+    });
+    this.child.stdin.on('error', (err) => this.failAll(err));
+  }
+
+  failAll(err) {
+    this.closed = true;
+    this.closeError = err;
+    for (const pending of this.pending.values()) pending.reject(err);
+    this.pending.clear();
+  }
+
+  onData(chunk) {
+    this.buffer += chunk;
+    const lines = this.buffer.split('\n');
+    this.buffer = lines.pop() || '';
+    for (const rawLine of lines) {
+      const line = rawLine.endsWith('\r') ? rawLine.slice(0, -1) : rawLine;
+      if (!line.trim()) continue;
+      let msg;
+      try { msg = JSON.parse(line); } catch { continue; }
+      if (process.env.AMALGM_CHAT_DEBUG_PI === '1') {
+        console.log('[Pi:recv]', JSON.stringify(msg).slice(0, 4000));
+      }
+      if (msg.type === 'response') {
+        const pending = msg.id !== undefined ? this.pending.get(msg.id) : null;
+        if (pending) {
+          this.pending.delete(msg.id);
+          if (msg.success === false) {
+            pending.reject(new Error(msg.error || `pi ${msg.command} failed`));
+          } else {
+            pending.resolve(msg.data ?? {});
+          }
+        }
+        continue;
+      }
+      if (msg.type === 'extension_ui_request') {
+        this.handleExtensionUiRequest(msg);
+        continue;
+      }
+      for (const handler of this.eventHandlers) handler(msg);
+    }
+  }
+
+  async handleExtensionUiRequest(message) {
+    if (!this.extensionUiHandler) return;
+    try {
+      const response = await this.extensionUiHandler(message);
+      if (response) this.write({ type: 'extension_ui_response', id: message.id, ...response });
+    } catch {
+      this.write({ type: 'extension_ui_response', id: message.id, cancelled: true });
+    }
+  }
+
+  write(value) {
+    if (this.closed || this.child.exitCode !== null || this.child.signalCode !== null) return;
+    try {
+      this.child.stdin.write(`${JSON.stringify(value)}\n`, (error) => {
+        if (error) this.failAll(error);
+      });
+    } catch (error) {
+      this.failAll(error);
+    }
+  }
+
+  request(command, params = {}, timeoutMs = 180000) {
+    if (this.closeError) return Promise.reject(this.closeError);
+    if (this.closed || this.child.exitCode !== null || this.child.signalCode !== null) {
+      return Promise.reject(new Error('pi rpc agent is not running'));
+    }
+    const id = `req-${this.id++}`;
+    return new Promise((resolve, reject) => {
+      const timer = timeoutMs > 0
+        ? setTimeout(() => {
+            this.pending.delete(id);
+            reject(new Error(`pi ${command} timed out`));
+          }, timeoutMs)
+        : null;
+      this.pending.set(id, {
+        resolve: (value) => { if (timer) clearTimeout(timer); resolve(value); },
+        reject: (err) => { if (timer) clearTimeout(timer); reject(err); },
+      });
+      try {
+        this.child.stdin.write(`${JSON.stringify({ id, type: command, ...params })}\n`, (err) => {
+          if (err) this.failAll(err);
+        });
+      } catch (err) {
+        this.failAll(err);
+      }
+    });
+  }
+
+  onEvent(handler) {
+    this.eventHandlers.add(handler);
+    return () => this.eventHandlers.delete(handler);
+  }
+
+  onExtensionUi(handler) {
+    this.extensionUiHandler = handler;
+  }
+
+  close() {
+    try { this.child.stdin.end(); } catch {}
+    try {
+      if (process.platform !== 'win32' && this.child.pid) process.kill(-this.child.pid, 'SIGTERM');
+      else this.child.kill('SIGTERM');
+    } catch {}
+  }
+}
+
+function resolveBinary() {
+  const binaryName = process.platform === 'win32' ? 'pi.cmd' : 'pi';
+  const homes = [...new Set([process.env.AMALGM_NATIVE_HOME, os.homedir()].filter(Boolean))];
+  const candidates = [
+    process.env.PI_BINARY,
+    ...homes.map((home) => path.join(home, '.npm-global', 'bin', binaryName)),
+    '/opt/homebrew/bin/pi',
+    '/usr/local/bin/pi',
+    findOnPath('pi'),
+  ].filter(Boolean);
+  const failures = [];
+  for (const candidate of candidates) {
+    if (!executableExists(candidate)) {
+      failures.push(`${candidate} (missing or not executable)`);
+      continue;
+    }
+    const res = spawnSync(candidate, ['--version'], { stdio: ['ignore', 'pipe', 'ignore'], encoding: 'utf8' });
+    if (res.status === 0) return candidate;
+    failures.push(`${candidate} (--version failed: ${res.error?.message || res.status || res.signal || 'unknown'})`);
+  }
+  throw new Error(`Pi CLI is not available. Checked: ${failures.join(', ') || 'no candidates'}. Install it with: npm install -g @mariozechner/pi-coding-agent`);
+}
+
+// Under amalgm auth every gateway request must flow through the local egress
+// proxy. Pi's provider baseUrl is overridden by a tiny extension; the URL rides
+// in on an env var so the extension file itself is static.
+function ensureEgressExtension(runtimeHome) {
+  const file = path.join(runtimeHome, '.pi', 'amalgm-egress-extension.js');
+  fs.mkdirSync(path.dirname(file), { recursive: true });
+  const source = [
+    '// Generated by @amalgm/chat: routes the Vercel AI Gateway provider',
+    '// through the per-session amalgm egress proxy for billing and audit.',
+    'export default function (pi) {',
+    '  const baseUrl = process.env.AMALGM_PI_GATEWAY_BASE_URL;',
+    `  if (baseUrl) pi.registerProvider(${JSON.stringify(GATEWAY_PROVIDER)}, { baseUrl });`,
+    '}',
+    '',
+  ].join('\n');
+  fs.writeFileSync(file, source, { mode: 0o600 });
+  return file;
+}
+
+function piProvider(contract) {
+  if (contract.authMethod === 'amalgm') return GATEWAY_PROVIDER;
+  const provider = String(contract.providerId || '').trim();
+  if (provider === 'vercel' || provider === 'ai-gateway' || provider === 'ai_gateway') return GATEWAY_PROVIDER;
+  return provider;
+}
+
+function piModel(contract, provider) {
+  const model = String(contract.providerModelId || contract.cliModel || '').trim();
+  if (provider === GATEWAY_PROVIDER && contract.authMethod === 'amalgm' && !model.includes('/')) {
+    return `${contract.providerId}/${model}`;
+  }
+  if (provider !== GATEWAY_PROVIDER && model.startsWith(`${provider}/`)) return model.slice(provider.length + 1);
+  return model;
+}
+
+const PI_KEY_ENV = Object.freeze({
+  anthropic: 'ANTHROPIC_API_KEY',
+  openai: 'OPENAI_API_KEY',
+  'openai-codex': 'OPENAI_API_KEY',
+  google: 'GEMINI_API_KEY',
+  openrouter: 'OPENROUTER_API_KEY',
+  'azure-openai-responses': 'AZURE_OPENAI_API_KEY',
+  [GATEWAY_PROVIDER]: 'AI_GATEWAY_API_KEY',
+});
+
+function applyPiByokEnvironment(env, contract, provider) {
+  if (contract.authMethod !== 'byok') return;
+  const variable = PI_KEY_ENV[provider];
+  if (!variable) {
+    throw new UnsupportedExecutionInputError(
+      'pi',
+      'auth + provider',
+      `${provider} has no declared BYOK credential environment mapping`,
+    );
+  }
+  for (const key of Object.values(PI_KEY_ENV)) delete env[key];
+  env[variable] = contract.auth.tokenRef || '';
+}
+
+function piMcpExtensionSource(configFile, permissionMode) {
+  const clientModule = import.meta.resolve('@modelcontextprotocol/sdk/client/index.js');
+  const stdioModule = import.meta.resolve('@modelcontextprotocol/sdk/client/stdio.js');
+  const httpModule = import.meta.resolve('@modelcontextprotocol/sdk/client/streamableHttp.js');
+  const sseModule = import.meta.resolve('@modelcontextprotocol/sdk/client/sse.js');
+  return [
+    '// Generated by @amalgm/chat. Pi has no native MCP or permission UI.',
+    "import { readFileSync } from 'node:fs';",
+    `import { Client } from ${JSON.stringify(clientModule)};`,
+    `import { StdioClientTransport } from ${JSON.stringify(stdioModule)};`,
+    `import { StreamableHTTPClientTransport } from ${JSON.stringify(httpModule)};`,
+    `import { SSEClientTransport } from ${JSON.stringify(sseModule)};`,
+    `const servers = JSON.parse(readFileSync(${JSON.stringify(configFile)}, 'utf8'));`,
+    'const clients = [];',
+    'const headers = (entries = []) => Object.fromEntries(entries.map(({ name, value }) => [name, value]));',
+    'async function connect(server) {',
+    "  const client = new Client({ name: 'amalgm-chat-pi', version: '0.1.2' });",
+    '  const transport = server.type === "stdio"',
+    '    ? new StdioClientTransport({ command: server.command, args: server.args, env: { ...process.env, ...headers(server.env) } })',
+    '    : server.type === "sse"',
+    '      ? new SSEClientTransport(new URL(server.url), { requestInit: { headers: headers(server.headers) } })',
+    '      : new StreamableHTTPClientTransport(new URL(server.url), { requestInit: { headers: headers(server.headers) } });',
+    '  await client.connect(transport);',
+    '  clients.push(client);',
+    '  return client;',
+    '}',
+    'function resultContent(result) {',
+    '  return (result.content || []).map((part) => {',
+    '    if (part.type === "text" || part.type === "image") return part;',
+    '    return { type: "text", text: JSON.stringify(part) };',
+    '  });',
+    '}',
+    'export default function (pi) {',
+    ...(permissionMode === 'ask' ? [
+      '  const alwaysAllowed = new Set();',
+      '  pi.on("tool_call", async (event, context) => {',
+      '    if (alwaysAllowed.has(event.toolName)) return;',
+      '    const choice = await context.ui.select(`Allow ${event.toolName}?`, ["Allow once", "Always allow", "Deny"]);',
+      '    if (choice === "Always allow") alwaysAllowed.add(event.toolName);',
+      '    if (choice !== "Allow once" && choice !== "Always allow") return { block: true, reason: "Denied by the user" };',
+      '  });',
+    ] : []),
+    '  pi.on("session_start", async () => {',
+    '    for (const server of servers) {',
+    '      const client = await connect(server);',
+    '      const listed = await client.listTools();',
+    '      for (const tool of listed.tools) {',
+    '        const exposedName = `${server.name}__${tool.name}`.replace(/[^A-Za-z0-9_-]/g, "_");',
+    '        pi.registerTool({',
+    '          name: exposedName,',
+    '          label: tool.title || tool.name,',
+    '          description: tool.description || `${server.name} MCP tool ${tool.name}`,',
+    '          parameters: tool.inputSchema || { type: "object", additionalProperties: true },',
+    '          async execute(_id, parameters) {',
+    '            const result = await client.callTool({ name: tool.name, arguments: parameters });',
+    '            return { content: resultContent(result), details: result, isError: result.isError === true };',
+    '          },',
+    '        });',
+    '      }',
+    '    }',
+    '  });',
+    '  pi.on("session_shutdown", async () => {',
+    '    await Promise.all(clients.map((client) => client.close().catch(() => {})));',
+    '  });',
+    '}',
+    '',
+  ].join('\n');
+}
+
+function ensureAmalgmExtension(runtimeHome, contract) {
+  const directory = path.join(runtimeHome, '.pi', 'amalgm');
+  const configFile = path.join(directory, 'mcp-servers.json');
+  const extensionFile = path.join(directory, 'chat-extension.js');
+  fs.mkdirSync(directory, { recursive: true });
+  fs.writeFileSync(configFile, `${JSON.stringify(relayedMcpServers(contract), null, 2)}\n`, { mode: 0o600 });
+  fs.writeFileSync(extensionFile, piMcpExtensionSource(configFile, contract.permissionMode), { mode: 0o600 });
+  return extensionFile;
+}
+
+function thinkingLevelFor(contract) {
+  const effort = String(contract.reasoningEffort || '').toLowerCase();
+  if (!effort) return null;
+  if (['off', 'minimal', 'low', 'medium', 'high', 'xhigh'].includes(effort)) return effort;
+  throw new UnsupportedExecutionInputError(
+    'pi',
+    'model.settings.reasoning',
+    `unsupported reasoning level ${JSON.stringify(contract.reasoningEffort)}`,
+  );
+}
+
+function buildArgs(contract, runtimeHome) {
+  const provider = piProvider(contract);
+  const args = ['--mode', 'rpc', '--provider', provider, '--model', piModel(contract, provider)];
+  const thinking = thinkingLevelFor(contract);
+  if (thinking) args.push('--thinking', thinking);
+  const systemPrompt = composeSystemPrompt(contract);
+  if (systemPrompt) args.push('--system-prompt', systemPrompt);
+  if (contract.providerSessionId) args.push('--session', contract.providerSessionId);
+  if (contract.authMethod === 'amalgm') args.push('--extension', ensureEgressExtension(runtimeHome));
+  const hasSelectedTools = Array.isArray(contract.mcpServers) && contract.mcpServers.length > 0;
+  if (contract.permissionMode === 'ask' || hasSelectedTools) {
+    args.push('--extension', ensureAmalgmExtension(runtimeHome, contract));
+  }
+  return args;
+}
+
+class PiAdapter {
+  constructor(options = {}) {
+    this.permissionPort = options.requestPermission || null;
+    this.inputCapabilities = nativeInputCapabilities({
+      authModes: ['platform', 'subscription', 'byok'],
+      modelSettings: COMMON_REASONING_SETTINGS,
+      tools: true,
+    });
+  }
+
+  validatePrepared(contract) {
+    const provider = piProvider(contract);
+    applyPiByokEnvironment({}, contract, provider);
+  }
+
+  async create(contract) {
+    const runtime = prepareHarnessRuntime(contract);
+    const binary = resolveBinary();
+    const env = { ...runtime.env };
+    const provider = piProvider(contract);
+    applyPiByokEnvironment(env, contract, provider);
+    if (contract.authMethod === 'amalgm' && contract.auth.baseUrl) {
+      env.AMALGM_PI_GATEWAY_BASE_URL = contract.auth.baseUrl;
+    }
+    const client = new PiRpc({
+      binary,
+      args: buildArgs(contract, runtime.runtimeHome),
+      cwd: contract.cwd,
+      env,
+    });
+    client.onExtensionUi(async (request) => {
+      if (!['select', 'confirm', 'input', 'editor'].includes(request.method)) return null;
+      const decision = await requestPermission(this.permissionPort, {
+        adapterId: 'pi',
+        providerSessionId: contract.providerSessionId,
+        kind: 'tool',
+        title: request.title || 'Pi requests permission',
+        description: request.message || null,
+        input: request,
+      });
+      if (request.method === 'select') {
+        const wanted = decision.outcome === 'allow_always'
+          ? 'Always allow'
+          : decision.outcome === 'allow_once'
+            ? 'Allow once'
+            : 'Deny';
+        return request.options?.includes(wanted) ? { value: wanted } : { cancelled: true };
+      }
+      if (request.method === 'confirm') return { confirmed: allowed(decision) };
+      return allowed(decision) && typeof decision.updatedInput?.value === 'string'
+        ? { value: decision.updatedInput.value }
+        : { cancelled: true };
+    });
+    // Sessions persist as files under the managed home; the file path is the
+    // durable resume handle (pi's --session flag takes a path).
+    const stats = await client.request('get_session_stats', {}, 60000);
+    return {
+      sessionId: contract.sessionId,
+      providerSessionId: stats.sessionFile || contract.providerSessionId || null,
+      client,
+      cancelRequested: false,
+      activePrompt: null,
+    };
+  }
+
+  async resume(session, contract) {
+    if (session.client && !session.client.closed) return session;
+    return this.create({ ...contract, providerSessionId: session.providerSessionId || contract.providerSessionId });
+  }
+
+  async stop(session) {
+    session.cancelRequested = true;
+    if (!session.activePrompt) return;
+    await session.client.request('abort', {}, 15000).catch(() => {});
+  }
+
+  async destroy(session) {
+    session.client.close();
+  }
+
+  async *prompt(session, input, contract) {
+    const queue = [];
+    let complete = false;
+    let wake = null;
+    const wakeLoop = () => { if (wake) { wake(); wake = null; } };
+    // Unlike the other adapters, an `error` event does NOT complete the turn:
+    // pi still emits agent_end afterwards, and finish() owns the done event.
+    const push = (e) => {
+      if (e.type === 'done') complete = true;
+      queue.push(e);
+      wakeLoop();
+    };
+    const state = { providerSessionId: session.providerSessionId, tools: new Set() };
+    let sawError = false;
+    session.cancelRequested = false;
+    session.activePrompt = state;
+    const finish = async () => {
+      // Attach the post-turn context snapshot before done so the context circle
+      // lands on the final number even if per-message events raced.
+      const stats = await session.client.request('get_session_stats', {}, 30000).catch(() => null);
+      recordNativeEvent('pi.rpc.session_stats', stats, {
+        providerSessionId: session.providerSessionId,
+        sessionId: contract.sessionId,
+        assistantMessageId: contract.assistantMessageId,
+      });
+      const snapshot = stats ? piSessionStatsUsage(stats, { providerSessionId: session.providerSessionId }) : null;
+      if (snapshot) push(snapshot);
+      push(done({
+        providerSessionId: session.providerSessionId,
+        stopReason: session.cancelRequested ? 'cancelled' : sawError ? 'error' : 'end_turn',
+      }));
+    };
+    const off = session.client.onEvent((event) => {
+      recordNativeEvent('pi.rpc.event', event, {
+        providerSessionId: session.providerSessionId,
+        sessionId: contract.sessionId,
+        assistantMessageId: contract.assistantMessageId,
+        eventType: event?.type,
+      });
+      if (event?.type === 'agent_end') {
+        finish().catch(() => push(done({ providerSessionId: session.providerSessionId, stopReason: 'end_turn' })));
+        return;
+      }
+      for (const e of normalizePiEvent(event, state)) {
+        if (e.type === 'error') sawError = true;
+        push(e);
+      }
+    });
+    session.client.request('prompt', piPrompt(input), 0).catch((err) => {
+      if (session.cancelRequested) {
+        push(done({ providerSessionId: session.providerSessionId, stopReason: 'cancelled' }));
+        return;
+      }
+      push(errorEvent(err.message, { providerSessionId: session.providerSessionId }));
+      push(done({ providerSessionId: session.providerSessionId, stopReason: 'error' }));
+    });
+    try {
+      while (!complete || queue.length) {
+        if (!queue.length) await new Promise((resolve) => { wake = resolve; });
+        else yield queue.shift();
+      }
+    } finally {
+      off();
+      session.activePrompt = null;
+      session.cancelRequested = false;
+    }
+  }
+}
+
+const __private = {
+  PiRpc,
+  applyPiByokEnvironment,
+  buildArgs,
+  ensureAmalgmExtension,
+  ensureEgressExtension,
+  piModel,
+  piProvider,
+  resolveBinary,
+  thinkingLevelFor,
+};
+
+export { PiAdapter, __private };