src/piAgent.js

import { Agent } from "@earendil-works/pi-agent-core";
import { createAssistantMessageEventStream, Type } from "@earendil-works/pi-ai";
import { env, pipeline } from "@huggingface/transformers";

const LOCAL_MODELS = {
  qwen25coder: {
    id: "onnx-community/Qwen2.5-Coder-0.5B-Instruct",
    name: "Qwen2.5 Coder 0.5B ONNX",
    promptStyle: "qwen",
  },
  qwen: {
    id: "onnx-community/Qwen3-0.6B-ONNX",
    name: "Qwen3 0.6B ONNX",
    promptStyle: "qwen",
  },
  minicpm: {
    id: "Mike0021/MiniCPM5-1B-ONNX-Web",
    name: "MiniCPM5-1B ONNX Web",
    promptStyle: "json",
  },
};

const DEFAULT_LOCAL_MODEL_KEY = "qwen25coder";

function getLocalModel(key) {
  return LOCAL_MODELS[key] || LOCAL_MODELS[DEFAULT_LOCAL_MODEL_KEY];
}

function createLocalModelMetadata(model) {
  return {
    id: model.id,
    name: model.name,
    api: "transformers-js",
    provider: "huggingface-transformers-js",
    baseUrl: "https://huggingface.co",
    reasoning: false,
    input: ["text"],
    cost: {
      input: 0,
      output: 0,
      cacheRead: 0,
      cacheWrite: 0,
    },
    contextWindow: 131072,
    maxTokens: 8192,
  };
}

const EMPTY_USAGE = {
  input: 0,
  output: 0,
  cacheRead: 0,
  cacheWrite: 0,
  totalTokens: 0,
  cost: {
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
    total: 0,
  },
};

function textFromContent(content) {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n");
}

function now() {
  return Date.now();
}

function createMessage(content, stopReason = "stop", model = getLocalModel(DEFAULT_LOCAL_MODEL_KEY)) {
  return {
    role: "assistant",
    content,
    api: "transformers-js",
    provider: "huggingface-transformers-js",
    model: model.id,
    usage: EMPTY_USAGE,
    stopReason,
    timestamp: now(),
  };
}

function stringifyToolResult(message) {
  const text = textFromContent(message.content);
  return `${message.toolName}(${message.toolCallId}) ${message.isError ? "failed" : "succeeded"}:\n${text}`;
}

function buildPrompt(context) {
  const tools = (context.tools || []).map((tool) => ({
    name: tool.name,
    description: tool.description,
    parameters: tool.parameters,
  }));
  const transcript = context.messages
    .slice(-8)
    .map((message) => {
      if (message.role === "toolResult") return `TOOL_RESULT:\n${stringifyToolResult(message)}`;
      return `${message.role.toUpperCase()}:\n${textFromContent(message.content)}`;
    })
    .join("\n\n");

  return `You are running as a pi Agent inside a browser-only app. The sandbox is a WebContainer: it has a virtual filesystem and can spawn browser-contained Node.js processes.

Use tools by returning strict JSON only. Do not use markdown.

To call tools:
{"toolCalls":[{"tool":"write_file","args":{"path":"hello.js","content":"console.log(2 + 2)\\n"}},{"tool":"run_command","args":{"command":"node","args":["hello.js"]}}]}

To answer the user after tool results:
{"final":"Short answer that explains what happened."}

Available tools:
${JSON.stringify(tools, null, 2)}

Conversation:
${transcript}

Return JSON now.`;
}

function buildQwenMessages(context) {
  const lastUserText = [...context.messages]
    .reverse()
    .find((message) => message.role === "user");
  const taskText = textFromContent(lastUserText?.content || "");
  const transcript = context.messages
    .slice(-8)
    .map((message) => {
      if (message.role === "toolResult") return `TOOL_RESULT ${message.toolName}:\n${textFromContent(message.content)}`;
      return `${message.role.toUpperCase()}:\n${textFromContent(message.content)}`;
    })
    .join("\n\n");
  const lastMessage = context.messages[context.messages.length - 1];

  if (lastMessage?.role === "toolResult") {
    return [
      {
        role: "system",
        content: "You are Pi Web Agent. Give a concise final answer from the sandbox tool results. Do not emit JSON.",
      },
      {
        role: "user",
        content: `Conversation and tool results:\n${transcript}\n\nFinal answer:`,
      },
    ];
  }

  const loweredTask = taskText.toLowerCase();
  let example = `<plan>
<write_file path="app.js">
console.log(2 * 2);
</write_file>
<run_command command="node" args="app.js" timeoutMs="10000"></run_command>
</plan>`;
  if (/\b(npm|install|package|dependency)\b/.test(loweredTask)) {
    example = `<plan>
<write_file path="pad.mjs">
import leftPad from 'left-pad';
console.log('padded: ' + leftPad('5', 3, '0'));
</write_file>
<run_command command="npm" args="install left-pad@1.3.0" timeoutMs="120000"></run_command>
<run_command command="node" args="pad.mjs" timeoutMs="10000"></run_command>
</plan>`;
  } else if (/\b(src\/|import|export|module|multi[- ]?file)\b/.test(loweredTask)) {
    example = `<plan>
<write_file path="src/util.mjs">
export function add(a, b) {
  return a + b;
}
</write_file>
<write_file path="test.mjs">
import { add } from './src/util.mjs';
console.log('sum: ' + add(2, 3));
</write_file>
<run_command command="node" args="test.mjs" timeoutMs="10000"></run_command>
</plan>`;
  }

  return [
    {
      role: "system",
      content:
        'Return only a complete <plan>...</plan>. No markdown or prose. Inside the plan use <write_file path="...">code</write_file>, <run_command command="..." args="..." timeoutMs="..."></run_command>, <read_file path="..."></read_file>, and <list_files path="..."></list_files>. Every task that says run must end with a run_command tag. Multi-file tasks must write every requested file before run_command. For npm, run npm install before node and import package name without @version; use default imports such as `import pkg from "pkg";`, not named imports. Use exact filenames. Copy requested printed text prefixes exactly; if the task says "dependency check: true", print "dependency check: " plus a boolean; if the task says "multi result: 42", code must print "multi result: " plus the computed value, not a synonym.',
    },
    {
      role: "user",
      content: `Example output:
${example}

Conversation:
${transcript}

Complete plan:`,
    },
  ];
}

function generatedTextFromResult(result) {
  const generated = result?.[0]?.generated_text;
  if (Array.isArray(generated)) return generated.at(-1)?.content ?? "";
  return generated ?? "";
}

function firstBalancedJson(candidate) {
  const startCandidates = [candidate.indexOf("{"), candidate.indexOf("[")].filter((index) => index >= 0);
  if (startCandidates.length === 0) return null;
  const start = Math.min(...startCandidates);
  const stack = [];
  let inString = false;
  let escaping = false;

  for (let index = start; index < candidate.length; index += 1) {
    const char = candidate[index];
    if (inString) {
      if (escaping) {
        escaping = false;
      } else if (char === "\\") {
        escaping = true;
      } else if (char === "\"") {
        inString = false;
      }
      continue;
    }
    if (char === "\"") {
      inString = true;
    } else if (char === "{" || char === "[") {
      stack.push(char);
    } else if (char === "}" || char === "]") {
      const expected = char === "}" ? "{" : "[";
      if (stack.pop() !== expected) return null;
      if (stack.length === 0) return candidate.slice(start, index + 1);
    }
  }
  return null;
}

function extractJsonPayload(text) {
  const trimmed = String(text || "").trim();
  const fence = trimmed.match(/```(?:json)?\s*([\s\S]*?)```/i);
  const candidate = fence ? fence[1].trim() : trimmed;
  const jsonText = firstBalancedJson(candidate);
  if (!jsonText) return null;
  try {
    return JSON.parse(jsonText);
  } catch {
    return null;
  }
}

function decodeAttribute(value = "") {
  return String(value)
    .replace(/&quot;/g, "\"")
    .replace(/&apos;/g, "'")
    .replace(/&lt;/g, "<")
    .replace(/&gt;/g, ">")
    .replace(/&amp;/g, "&");
}

function parseAttributes(source = "") {
  const attributes = {};
  const pattern = /([A-Za-z][\w-]*)\s*=\s*"([^"]*)"/g;
  let match = pattern.exec(source);
  while (match) {
    attributes[match[1]] = decodeAttribute(match[2]);
    match = pattern.exec(source);
  }
  return attributes;
}

function trimCodeBlock(text) {
  return String(text || "").replace(/^\n/, "").replace(/\n?$/, "\n");
}

function splitCommandLine(value) {
  const parts = [];
  let current = "";
  let quote = "";
  let escaping = false;
  for (const char of String(value || "")) {
    if (escaping) {
      current += char;
      escaping = false;
      continue;
    }
    if (char === "\\") {
      escaping = true;
      continue;
    }
    if (quote) {
      if (char === quote) {
        quote = "";
      } else {
        current += char;
      }
      continue;
    }
    if (char === "\"" || char === "'") {
      quote = char;
    } else if (/\s/.test(char)) {
      if (current) {
        parts.push(current);
        current = "";
      }
    } else {
      current += char;
    }
  }
  if (current) parts.push(current);
  return parts;
}

function normalizeCommandArgs(args = {}) {
  const normalized = { ...args };
  if (typeof normalized.args === "string") normalized.args = splitCommandLine(normalized.args);
  if (!Array.isArray(normalized.args)) normalized.args = [];
  if (typeof normalized.command === "string" && normalized.command.includes(" ") && normalized.args.length === 0) {
    const [command, ...rest] = splitCommandLine(normalized.command);
    normalized.command = command;
    normalized.args = rest;
  }
  if (normalized.timeoutMs !== undefined) normalized.timeoutMs = Number(normalized.timeoutMs);
  return normalized;
}

function lastUserText(context) {
  const message = [...context.messages].reverse().find((item) => item.role === "user");
  return textFromContent(message?.content || "");
}

function packageNameFromSpecifier(specifier) {
  const value = String(specifier || "");
  const versionAt = value.lastIndexOf("@");
  return versionAt > 0 ? value.slice(0, versionAt) : value;
}

function requestedPackageSpecifiers(context) {
  const text = lastUserText(context);
  const specs = [];
  const pattern = /\b(?:install|package|dependency)\b[^\n,;]*?\s((?:@[\w.-]+\/)?[\w.-]+@[0-9][^\s,;)]*)/gi;
  let match = pattern.exec(text);
  while (match) {
    specs.push(match[1]);
    match = pattern.exec(text);
  }
  return specs;
}

function isRunCommand(call, command) {
  if (call.name !== "run_command") return false;
  const args = normalizeCommandArgs(call.arguments || {});
  return args.command === command;
}

function repairToolCalls(toolCalls, context) {
  const repaired = toolCalls.map((call) => ({ ...call, arguments: { ...(call.arguments || {}) } }));
  const userText = lastUserText(context).toLowerCase();
  const packageSpecifiers = requestedPackageSpecifiers(context);

  for (const specifier of packageSpecifiers) {
    const packageName = packageNameFromSpecifier(specifier);
    const hasInstall = repaired.some((call) => {
      if (!isRunCommand(call, "npm")) return false;
      const args = normalizeCommandArgs(call.arguments || {}).args;
      return args.includes("install") && args.includes(specifier);
    });
    if (!hasInstall) {
      const firstNodeRun = repaired.findIndex((call) => isRunCommand(call, "node"));
      const installCall = {
        type: "toolCall",
        id: `tool-${now()}-repair-install`,
        name: "run_command",
        arguments: {
          command: "npm",
          args: ["install", specifier],
          timeoutMs: 120000,
        },
      };
      repaired.splice(firstNodeRun >= 0 ? firstNodeRun : repaired.length, 0, installCall);
    }

    if (packageName === "is-number" && userText.includes("dependency check")) {
      const file = repaired.find((call) => call.name === "write_file" && String(call.arguments?.path || "").endsWith(".mjs"));
      if (file) {
        file.arguments.content = "import isNumber from 'is-number';\nconsole.log('dependency check: ' + isNumber(42));\n";
      }
    }
  }

  return repaired;
}

function parseTaggedToolCalls(text) {
  const toolCalls = [];
  const source = String(text || "");
  const writePattern = /<write_file\b([^>]*)>([\s\S]*?)<\/write_file>/gi;
  let write = writePattern.exec(source);
  while (write) {
    const attributes = parseAttributes(write[1]);
    if (attributes.path) {
      toolCalls.push({
        tool: "write_file",
        args: {
          path: attributes.path,
          content: trimCodeBlock(write[2]),
        },
      });
    }
    write = writePattern.exec(source);
  }

  const emptyTagPattern = /<(run_command|read_file|list_files)\b([^>]*)>(?:\s*<\/\1>)?/gi;
  let tag = emptyTagPattern.exec(source);
  while (tag) {
    const tool = tag[1];
    const attributes = parseAttributes(tag[2]);
    if (tool === "run_command" && attributes.command) {
      toolCalls.push({
        tool,
        args: normalizeCommandArgs({
          command: attributes.command,
          args: attributes.args || [],
          timeoutMs: attributes.timeoutMs,
        }),
      });
    } else if ((tool === "read_file" || tool === "list_files") && attributes.path) {
      toolCalls.push({ tool, args: { path: attributes.path } });
    }
    tag = emptyTagPattern.exec(source);
  }

  return toolCalls;
}

function normalizeToolCalls(payload, generated = "") {
  const rawCalls = [];
  function collect(value) {
    if (!value) return;
    if (Array.isArray(value)) {
      for (const item of value) collect(item);
      return;
    }
    if (Array.isArray(value.toolCalls)) collect(value.toolCalls);
    if (Array.isArray(value.tools)) collect(value.tools);
    if (Array.isArray(value.actions)) collect(value.actions);
    if (value.tool || value.name) rawCalls.push(value);
  }
  collect(payload);
  if (rawCalls.length === 0) rawCalls.push(...parseTaggedToolCalls(generated));
  return rawCalls
    .map((call, index) => ({
      type: "toolCall",
      id: `tool-${now()}-${index}`,
      name: String(call.tool || call.name || ""),
      arguments: call.tool === "run_command" || call.name === "run_command" ? normalizeCommandArgs(call.args || call.arguments || {}) : call.args || call.arguments || {},
    }))
    .filter((call) => call.name);
}

function normalizeFinalText(payload, fallback) {
  if (payload && typeof payload.final === "string") return payload.final;
  if (payload && typeof payload.message === "string") return payload.message;
  if (payload && typeof payload.answer === "string") return payload.answer;
  return String(fallback || "")
    .replace(/```(?:\w+)?/g, "")
    .trim() || "Done.";
}

function mockPlan(context) {
  const last = context.messages[context.messages.length - 1];
  if (last?.role === "toolResult") {
    return {
      final: summarizeToolResults(context),
    };
  }

  const userText = textFromContent(last?.content || "").toLowerCase();
  if (userText.includes("read")) {
    return {
      toolCalls: [{ tool: "read_file", args: { path: "hello.js" } }],
    };
  }
  if (userText.includes("list")) {
    return {
      toolCalls: [{ tool: "list_files", args: { path: "." } }],
    };
  }
  if (userText.includes("install") || userText.includes("dependency") || userText.includes("package")) {
    return {
      toolCalls: [
        {
          tool: "write_file",
          args: {
            path: "check-package.mjs",
            content: 'import isNumber from "is-number";\nconsole.log(`dependency check: ${isNumber(42)}`);\n',
          },
        },
        {
          tool: "run_command",
          args: {
            command: "npm",
            args: ["install", "is-number@7.0.0"],
            timeoutMs: 120000,
          },
        },
        {
          tool: "run_command",
          args: {
            command: "node",
            args: ["check-package.mjs"],
          },
        },
      ],
    };
  }
  return {
    toolCalls: [
      {
        tool: "write_file",
        args: {
          path: "hello.js",
          content: 'const value = 21 * 2;\nconsole.log(`pi sandbox result: ${value}`);\n',
        },
      },
      {
        tool: "run_command",
        args: {
          command: "node",
          args: ["hello.js"],
        },
      },
    ],
  };
}

function summarizeToolResults(context) {
  const results = context.messages.filter((message) => message.role === "toolResult").slice(-4).map(stringifyToolResult);
  return `The sandbox work completed.\n\n${results.join("\n\n")}`;
}

function emitFinal(stream, message, text = "") {
  stream.push({ type: "start", partial: { ...message, content: [{ type: "text", text: "" }] } });
  if (text) {
    const partial = { ...message, content: [{ type: "text", text }] };
    stream.push({ type: "text_start", contentIndex: 0, partial: { ...message, content: [{ type: "text", text: "" }] } });
    stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial });
    stream.push({ type: "text_end", contentIndex: 0, content: text, partial });
  }
  if (message.stopReason === "error" || message.stopReason === "aborted") {
    stream.push({ type: "error", reason: message.stopReason, error: message });
  } else {
    stream.push({ type: "done", reason: message.stopReason, message });
  }
}

export function createPiAgent({ sandbox, modelMode, device, maxTokens, temperature, onModelStatus = () => {} }) {
  env.allowLocalModels = false;
  env.allowRemoteModels = true;
  env.backends.onnx.wasm.numThreads = Math.min(2, navigator.hardwareConcurrency || 2);

  let generatorPromise = null;
  let generatorKey = "";

  async function getGenerator() {
    const model = getLocalModel(modelMode());
    const key = `${model.id}:${device()}`;
    if (!generatorPromise || generatorKey !== key) {
      generatorKey = key;
      onModelStatus(`Loading ${device()}`);
      generatorPromise = pipeline("text-generation", model.id, {
        dtype: "q4",
        device: device(),
        progress_callback: (event) => {
          if (event.status === "progress") {
            onModelStatus(`${event.file} ${Math.round(event.progress)}%`);
          } else if (event.status) {
            onModelStatus(event.status);
          }
        },
      });
    }
    return generatorPromise;
  }

  async function releaseGenerator() {
    if (!generatorPromise) return;
    const generator = await generatorPromise.catch(() => null);
    generatorPromise = null;
    generatorKey = "";
    await generator?.dispose?.();
  }

  async function producePlan(context, signal) {
    const lastMessage = context.messages[context.messages.length - 1];
    if (lastMessage?.role === "toolResult") {
      return JSON.stringify({ final: summarizeToolResults(context) });
    }

    if (modelMode() === "mock") {
      return JSON.stringify(mockPlan(context));
    }

    const model = getLocalModel(modelMode());
    const generator = await getGenerator();
    if (signal?.aborted) throw new Error("Aborted");
    const input = model.promptStyle === "qwen" ? buildQwenMessages(context) : buildPrompt(context);
    const result = await generator(input, {
      max_new_tokens: Number(maxTokens()) || 128,
      temperature: Number(temperature()) || 0,
      do_sample: Number(temperature()) > 0,
      return_full_text: false,
      tokenizer_encode_kwargs: model.promptStyle === "qwen" ? { enable_thinking: false } : undefined,
    });
    onModelStatus("Model ready");
    return generatedTextFromResult(result);
  }

  function streamFn(_model, context, options = {}) {
    const stream = createAssistantMessageEventStream();
    queueMicrotask(async () => {
      try {
        const model = getLocalModel(modelMode());
        const generated = await producePlan(context, options.signal);
        const payload = extractJsonPayload(generated);
        const lastMessage = context.messages[context.messages.length - 1];
        const forceFinal = lastMessage?.role === "toolResult";
        const toolCalls = forceFinal ? [] : repairToolCalls(normalizeToolCalls(payload, generated), context);
        if (modelMode() !== "mock") await releaseGenerator();
        if (toolCalls.length > 0) {
          const message = createMessage([{ type: "text", text: "Using sandbox tools." }, ...toolCalls], "toolUse", model);
          emitFinal(stream, message, "Using sandbox tools.");
          return;
        }
        const text = normalizeFinalText(payload, generated);
        const message = createMessage([{ type: "text", text }], "stop", model);
        emitFinal(stream, message, text);
      } catch (error) {
        await releaseGenerator().catch(() => {});
        const text = error instanceof Error ? error.message : String(error);
        const message = {
          ...createMessage([{ type: "text", text }], options.signal?.aborted ? "aborted" : "error", getLocalModel(modelMode())),
          errorMessage: text,
        };
        emitFinal(stream, message, text);
      }
    });
    return stream;
  }

  const tools = [
    {
      name: "list_files",
      label: "List files",
      description: "List files in the sandbox workspace.",
      parameters: Type.Object({
        path: Type.Optional(Type.String()),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.listFiles(args.path || ".");
        return { content: [{ type: "text", text: output }], details: { output } };
      },
    },
    {
      name: "read_file",
      label: "Read file",
      description: "Read a UTF-8 file from the sandbox workspace.",
      parameters: Type.Object({
        path: Type.String(),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.readFile(args.path);
        return { content: [{ type: "text", text: output }], details: { path: args.path, output } };
      },
    },
    {
      name: "write_file",
      label: "Write file",
      description: "Create or replace a UTF-8 file inside the sandbox workspace.",
      parameters: Type.Object({
        path: Type.String(),
        content: Type.String(),
      }),
      execute: async (_id, args) => {
        const output = await sandbox.writeFile(args.path, args.content);
        return { content: [{ type: "text", text: output }], details: { path: args.path } };
      },
    },
    {
      name: "run_command",
      label: "Run command",
      description: "Spawn a process inside the browser-only WebContainer sandbox.",
      parameters: Type.Object({
        command: Type.String(),
        args: Type.Optional(Type.Array(Type.String())),
        timeoutMs: Type.Optional(Type.Number()),
      }),
      execute: async (_id, args) => {
        const result = await sandbox.runCommand(args.command, args.args || [], args.timeoutMs || 10000);
        const text = `$ ${result.command}\nexit ${result.exitCode}\n${result.output}`;
        return {
          content: [{ type: "text", text }],
          details: result,
        };
      },
      executionMode: "sequential",
    },
  ];

  const agent = new Agent({
    initialState: {
      model: createLocalModelMetadata(getLocalModel(modelMode())),
      systemPrompt:
        "You are Pi Web Agent. Use the sandbox tools for filesystem, npm dependency, or command tasks, then give concise results.",
      tools,
    },
    streamFn,
    toolExecution: "sequential",
  });

  agent.preloadModel = async () => {
    if (modelMode() === "mock") {
      onModelStatus("Deterministic test model");
      return;
    }
    await getGenerator();
    onModelStatus("Model ready");
  };

  return agent;
}

export { DEFAULT_LOCAL_MODEL_KEY, LOCAL_MODELS };