src/worker/toolParser.ts

// Phân tích đầu ra của MiniCPM5: khối <think>…</think> và lời gọi tool dạng XML:
//
//   <function name="bash"><param name="command"><![CDATA[ls -la]]></param></function>
//
// Bản gốc dùng DOMParser (xmldom). Worker không có DOMParser nên ở đây dùng một
// parser XML nhỏ, chặt chẽ, chỉ hiểu đúng cú pháp trên.

import type { AgentMessage, AssistantContent, TextContent, ToolCall, ToolDef } from '../shared/types';

// ───────────────────────── Thinking ─────────────────────────

export interface ThinkingSplit {
  thinking: string;
  answer: string;
  complete: boolean;
}

/**
 * Tách phần suy nghĩ và phần trả lời. Với `thinkingPrefilled`, chat template đã
 * mở sẵn thẻ <think>, nên đầu ra bắt đầu ngay trong phần suy nghĩ.
 */
export function splitThinking(text: string, opts: { thinkingPrefilled?: boolean } = {}): ThinkingSplit {
  const prefilled = opts.thinkingPrefilled ?? false;
  const open = /^\s*<think>\n?/.exec(text);
  if (!prefilled && !open) return { thinking: '', answer: text, complete: true };
  const start = open ? open[0].length : 0;
  const end = text.indexOf('</think>', start);
  if (end < 0) return { thinking: text.slice(start), answer: '', complete: false };
  return { thinking: text.slice(start, end), answer: text.slice(end + '</think>'.length), complete: true };
}

// ───────────────────────── Vị trí lời gọi tool ─────────────────────────

/** Tìm `<function ` đầu tiên nằm ngoài code fence và ngoài inline code. */
export function findFunctionStart(text: string, from: number): number {
  let fence: string | null = null;
  let inlineTicks = 0;
  for (let i = from; i < text.length; ) {
    if (i === 0 || text[i - 1] === '\n') {
      const nl = text.indexOf('\n', i);
      const line = nl < 0 ? text.slice(i) : text.slice(i, nl);
      const marker = /^ {0,3}(`{3,}|~{3,})/.exec(line)?.[1];
      if (marker) {
        if (fence) {
          if (marker[0] === fence[0] && marker.length >= fence.length) fence = null;
        } else {
          fence = marker;
        }
        i += line.length + 1;
        continue;
      }
    }
    if (fence) {
      i++;
      continue;
    }
    if (text[i] === '`') {
      let n = 0;
      while (text[i + n] === '`') n++;
      if (inlineTicks) {
        if (inlineTicks === n) inlineTicks = 0;
      } else {
        inlineTicks = n;
      }
      i += n;
      continue;
    }
    if (!inlineTicks && text.startsWith('<function', i) && /\s/.test(text[i + 9] ?? '')) return i;
    i++;
  }
  return -1;
}

/** Tìm vị trí ngay sau `</function>` tương ứng, bỏ qua `</function>` nằm trong CDATA. */
export function findFunctionEnd(text: string, from: number): number {
  let n = from;
  while (n < text.length) {
    const cdata = text.indexOf('<![CDATA[', n);
    const close = text.indexOf('</function>', n);
    if (close < 0) throw new Error('Incomplete tool call. No command was executed. Try again with a shorter task.');
    if (cdata < 0 || cdata > close) return close + '</function>'.length;
    const cdataEnd = text.indexOf(']]>', cdata + 9);
    if (cdataEnd < 0) throw new Error('Incomplete CDATA in tool call. No command was executed.');
    n = cdataEnd + 3;
  }
  throw new Error('Incomplete tool call. No command was executed. Try again with a shorter task.');
}

// ───────────────────────── Parser XML tối giản ─────────────────────────

const MALFORMED = 'Malformed model tool call.';

function decodeEntities(s: string): string {
  return s.replace(/&(#x[0-9a-fA-F]+|#[0-9]+|lt|gt|amp|quot|apos);/g, (_m, ref: string) => {
    switch (ref) {
      case 'lt':
        return '<';
      case 'gt':
        return '>';
      case 'amp':
        return '&';
      case 'quot':
        return '"';
      case 'apos':
        return "'";
      default: {
        const code = ref[1] === 'x' ? parseInt(ref.slice(2), 16) : parseInt(ref.slice(1), 10);
        if (!Number.isFinite(code) || code > 0x10ffff) throw new Error(MALFORMED);
        return String.fromCodePoint(code);
      }
    }
  });
}

interface OpenTag {
  attrs: Record<string, string>;
  selfClosing: boolean;
  end: number;
}

/** Đọc thẻ mở bắt đầu tại `pos` (text[pos] === '<'), ví dụ `<param name="x">`. */
function readOpenTag(xml: string, pos: number, tag: string): OpenTag {
  if (!xml.startsWith('<' + tag, pos)) throw new Error(MALFORMED);
  let i = pos + 1 + tag.length;
  const attrs: Record<string, string> = Object.create(null);
  for (;;) {
    while (/\s/.test(xml[i] ?? '')) i++;
    if (xml[i] === '>') return { attrs, selfClosing: false, end: i + 1 };
    if (xml[i] === '/' && xml[i + 1] === '>') return { attrs, selfClosing: true, end: i + 2 };
    const m = /^([A-Za-z_][\w.-]*)\s*=\s*("([^"<]*)"|'([^'<]*)')/.exec(xml.slice(i));
    if (!m) throw new Error(MALFORMED);
    const name = m[1];
    if (name in attrs) throw new Error(MALFORMED);
    attrs[name] = decodeEntities(m[3] ?? m[4] ?? '');
    i += m[0].length;
  }
}

export interface ParsedCall {
  name: string;
  params: Array<{ name: string; text: string }>;
}

/** Phân tích đúng một phần tử `<function name="…"><param name="…">…</param>…</function>`. */
export function parseFunctionXml(xml: string): ParsedCall {
  if (!xml.startsWith('<function')) throw new Error(MALFORMED);
  const head = readOpenTag(xml, 0, 'function');
  const name = head.attrs.name;
  if (name === undefined) throw new Error(MALFORMED);
  const params: ParsedCall['params'] = [];
  let i = head.end;

  if (head.selfClosing) {
    if (i !== xml.length) throw new Error(MALFORMED);
    return { name, params };
  }

  for (;;) {
    // Khoảng trắng giữa các phần tử được phép; chữ khác thì không.
    const ws = /^\s*/.exec(xml.slice(i))![0];
    i += ws.length;
    if (xml.startsWith('</function>', i)) {
      i += '</function>'.length;
      if (i !== xml.length) throw new Error(MALFORMED);
      return { name, params };
    }
    if (xml[i] !== '<') throw new Error('Unexpected content in tool call.');
    if (!xml.startsWith('<param', i) || !/[\s>/]/.test(xml[i + 6] ?? '')) {
      // Có thể là phần tử lạ hoặc XML hỏng.
      if (/^<[A-Za-z]/.test(xml.slice(i))) throw new Error('Unexpected content in tool call.');
      throw new Error(MALFORMED);
    }
    const open = readOpenTag(xml, i, 'param');
    const pname = open.attrs.name;
    if (pname === undefined) throw new Error(MALFORMED);
    i = open.end;
    let text = '';
    if (!open.selfClosing) {
      for (;;) {
        if (i >= xml.length) throw new Error(MALFORMED);
        if (xml.startsWith('</param>', i)) {
          i += '</param>'.length;
          break;
        }
        if (xml.startsWith('<![CDATA[', i)) {
          const end = xml.indexOf(']]>', i + 9);
          if (end < 0) throw new Error(MALFORMED);
          text += xml.slice(i + 9, end);
          i = end + 3;
          continue;
        }
        if (xml[i] === '<') throw new Error('Tool values containing XML must use CDATA.');
        const next = xml.indexOf('<', i);
        if (next < 0) throw new Error(MALFORMED);
        text += decodeEntities(xml.slice(i, next));
        i = next;
      }
    }
    params.push({ name: pname, text });
  }
}

// ───────────────────────── Đầu ra mô hình → nội dung assistant ─────────────────────────

/**
 * Chuyển đầu ra thô của mô hình thành danh sách nội dung: thinking, text và toolCall.
 * Ném lỗi (không thực thi tool nào) nếu lời gọi tool sai cú pháp.
 */
export function parseAssistantOutput(
  raw: string,
  tools: ToolDef[],
  opts: { thinkingPrefilled?: boolean } = {},
): AssistantContent[] {
  const text = raw.replace(/<\|im_end\|>$|<\/s>$/, '');
  const split = splitThinking(text, opts);
  if (!split.complete) {
    throw new Error('The model used its response budget before finishing its thoughts. Try a smaller task.');
  }
  const answer = split.answer;
  const out: AssistantContent[] = [];
  if (split.thinking.trim()) out.push({ type: 'thinking', thinking: split.thinking.trim() });

  let pos = 0;
  while (pos < answer.length) {
    const start = findFunctionStart(answer, pos);
    if (start < 0) {
      const rest = answer.slice(pos).trim();
      if (rest) out.push({ type: 'text', text: rest });
      break;
    }
    const before = answer.slice(pos, start).trim();
    if (before) out.push({ type: 'text', text: before });

    const end = findFunctionEnd(answer, start);
    const parsed = parseFunctionXml(answer.slice(start, end));
    const tool = tools.find((t) => t.name === parsed.name);
    if (!tool) throw new Error('Unknown model tool: ' + parsed.name);

    const args: Record<string, unknown> = Object.create(null);
    for (const p of parsed.params) {
      if (Object.hasOwn(args, p.name) || !Object.hasOwn(tool.parameters.properties, p.name)) {
        throw new Error('Unknown or duplicate tool parameter: ' + p.name);
      }
      args[p.name] = tool.parameters.properties[p.name].type === 'string' ? p.text : JSON.parse(p.text);
    }
    const call: ToolCall = { type: 'toolCall', id: crypto.randomUUID(), name: parsed.name, arguments: args };
    out.push(call);
    pos = end;
  }
  return out;
}

// ───────────────────────── Lịch sử chat → chat template ─────────────────────────

const textOf = (content: string | TextContent[] | undefined): string =>
  typeof content === 'string' ? content : (content ?? []).filter((c) => c.type === 'text').map((c) => c.text).join('\n');

export interface TemplateMessage {
  role: 'system' | 'user' | 'assistant' | 'tool';
  content: string;
  reasoning_content?: string;
  tool_calls?: Array<{ type: 'function'; function: { name: string; arguments: Record<string, unknown> } }>;
}

/** Bỏ các lượt assistant lỗi/bị dừng rồi dựng danh sách message cho chat template. */
export function toTemplateMessages(systemPrompt: string, messages: AgentMessage[]): TemplateMessage[] {
  const kept = messages.filter((m) => m.role !== 'assistant' || !['error', 'aborted'].includes(m.stopReason ?? ''));
  return [
    { role: 'system', content: systemPrompt },
    ...kept.map((m): TemplateMessage => {
      if (m.role === 'toolResult') {
        return {
          role: 'tool',
          content: JSON.stringify({ name: m.toolName, isError: m.isError, output: textOf(m.content) }),
        };
      }
      if (m.role === 'assistant') {
        return {
          role: 'assistant',
          reasoning_content: m.content
            .filter((c): c is Extract<AssistantContent, { type: 'thinking' }> => 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 is ToolCall => c.type === 'toolCall')
            .map((c) => ({ type: 'function', function: { name: c.name, arguments: c.arguments } })),
        };
      }
      return { role: 'user', content: m.content };
    }),
  ];
}