app/src/protocol.mjs

import { DOMParser } from '@xmldom/xmldom';

export const textContent = content => typeof content === 'string' ? content :
  (content ?? []).filter(c => c.type === 'text').map(c => c.text).join('\n');

export function toMiniMessages(context) {
  const messages = context.messages.filter(m => m.role !== 'assistant' || !['error', 'aborted'].includes(m.stopReason));
  return [{ role: 'system', content: context.systemPrompt }, ...messages.map(m => {
    if (m.role === 'toolResult') return { role: 'tool', content: JSON.stringify({
      name: m.toolName, isError: m.isError, output: textContent(m.content),
    }) };
    if (m.role === 'assistant') return { role: 'assistant',
      // Explicit even when empty: the upstream template otherwise mistakes
      // literal </think> inside code or tool arguments for a reasoning boundary.
      reasoning_content: m.content.filter(c => c.type === 'thinking').map(c => c.thinking).join('\n'),
      content: m.content.map(c => c.type === 'text' ? c.text : c.type === 'toolCall' ? '<tool_sep>' : '').join(''),
      tool_calls: m.content.filter(c => c.type === 'toolCall').map(c => ({
        type: 'function', function: { name: c.name, arguments: c.arguments },
      })) };
    return { role: 'user', content: textContent(m.content) };
  })];
}

// Scan function boundaries outside CDATA; regex alone breaks shell commands
// containing XML (including literal </function> inside a here-document).
function functionEnd(raw, start) {
  let at = start;
  while (at < raw.length) {
    const cdata = raw.indexOf('<![CDATA[', at);
    const end = raw.indexOf('</function>', at);
    if (end < 0) throw Error('Incomplete tool call. No command was executed. Try again with a shorter task.');
    if (cdata < 0 || cdata > end) return end + 11;
    const close = raw.indexOf(']]>', cdata + 9);
    if (close < 0) throw Error('Incomplete CDATA in tool call. No command was executed.');
    at = close + 3;
  }
}

function nextFunction(raw, start) {
  let fence = null, inline = 0;
  for (let at = start; at < raw.length;) {
    if (at === 0 || raw[at - 1] === '\n') {
      const line = raw.slice(at).split('\n')[0];
      const marker = /^ {0,3}(`{3,}|~{3,})/.exec(line)?.[1];
      if (marker) {
        if (!fence) fence = marker;
        else if (marker[0] === fence[0] && marker.length >= fence.length) fence = null;
        at += line.length + 1; continue;
      }
    }
    if (fence) { at++; continue; }
    if (raw[at] === '`') {
      const count = /^`+/.exec(raw.slice(at))[0].length;
      if (!inline) inline = count; else if (inline === count) inline = 0;
      at += count; continue;
    }
    if (!inline && /^<function\s/.test(raw.slice(at))) return at;
    at++;
  }
  return -1;
}

export function splitThinking(raw, { thinkingPrefilled = false } = {}) {
  const opening = /^\s*<think>\n?/.exec(raw);
  if (!thinkingPrefilled && !opening) return { thinking: '', answer: raw, complete: true };
  const start = opening?.[0].length ?? 0;
  const end = raw.indexOf('</think>', start);
  if (end < 0) return { thinking: raw.slice(start), answer: '', complete: false };
  return { thinking: raw.slice(start, end), answer: raw.slice(end + 8), complete: true };
}

export function parseCompletion(raw, tools, options) {
  raw = raw.replace(/<\|im_end\|>$|<\/s>$/, '');
  const reasoning = splitThinking(raw, options);
  if (!reasoning.complete) throw Error('The model used its response budget before finishing its thoughts. Try a smaller task.');
  raw = reasoning.answer;
  const blocks = reasoning.thinking.trim() ? [{ type: 'thinking', thinking: reasoning.thinking.trim() }] : [];
  let at = 0;
  while (at < raw.length) {
    const start = nextFunction(raw, at);
    if (start < 0) { if (raw.slice(at).trim()) blocks.push({ type: 'text', text: raw.slice(at).trim() }); break; }
    if (raw.slice(at, start).trim()) blocks.push({ type: 'text', text: raw.slice(at, start).trim() });
    const end = functionEnd(raw, start);
    const errors = [];
    const doc = new DOMParser({ onError: (level, message) => errors.push(message) })
      .parseFromString(raw.slice(start, end), 'application/xml');
    if (errors.length || doc.doctype || doc.documentElement.tagName !== 'function') throw Error('Malformed model tool call.');
    const fn = doc.documentElement, name = fn.getAttribute('name');
    const tool = tools.find(t => t.name === name);
    if (!tool) throw Error('Unknown model tool: ' + name);
    const args = Object.create(null);
    for (const node of Array.from(fn.childNodes)) {
      if (node.nodeType === 3 && !node.textContent.trim()) continue;
      if (node.nodeType !== 1 || node.tagName !== 'param') throw Error('Unexpected content in tool call.');
      const key = node.getAttribute('name');
      if (Object.hasOwn(args, key) || !Object.hasOwn(tool.parameters.properties, key)) throw Error('Unknown or duplicate tool parameter: ' + key);
      if (Array.from(node.childNodes).some(c => ![3, 4].includes(c.nodeType))) throw Error('Tool values containing XML must use CDATA.');
      const schema = tool.parameters.properties[key];
      args[key] = schema.type === 'string' ? node.textContent : JSON.parse(node.textContent);
    }
    blocks.push({ type: 'toolCall', id: crypto.randomUUID(), name, arguments: args });
    at = end;
  }
  return blocks;
}

// Keep complete user turns together so tool results never lose their calls.
export function fitContext(context, tokenize, maxPromptTokens) {
  const messages = [...context.messages];
  let dropped = 0;
  while (true) {
    const inputs = tokenize(toMiniMessages({ ...context, messages }));
    if (inputs.input_ids.dims[1] <= maxPromptTokens) return { inputs, dropped };
    const nextUser = messages.findIndex((m, i) => i > 0 && m.role === 'user');
    if (nextUser < 0) throw Error('This task exceeds the browser context budget. Start a new chat or use smaller files and shorter tool output.');
    messages.splice(0, nextUser); dropped += nextUser;
  }
}