src/lib/utils/tool-parser-htmlparser2.ts

export interface ToolCall {
  name: string;
  args: Record<string, unknown>;
  rawArgs: string;
  startIndex?: number;
  endIndex?: number;
  complete: boolean;
}

export interface ParseResult {
  safeText: string;
  toolCalls: ToolCall[];
  pendingToolCall: boolean;
  errors: string[];
}

export class ToolParser {
  private buffer = "";
  private textBuffer = "";
  private lastReturnedTextLength = 0;
  private toolCalls: ToolCall[] = [];
  private errors: string[];
  private processedUpTo = 0;

  constructor() {
    this.errors = [];
  }

  reset(): void {
    this.buffer = "";
    this.textBuffer = "";
    this.lastReturnedTextLength = 0;
    this.toolCalls = [];
    this.errors = [];
    this.processedUpTo = 0;
  }

  process(text: string): ParseResult {
    this.buffer += text;

    // Use regex to find complete tool tags
    // This avoids htmlparser2 interpreting HTML inside JSON strings
    const toolRegex = /<tool\s+name="([^"]+)">([^]*?)<\/tool>/g;
    toolRegex.lastIndex = 0;

    let match;
    let lastMatchEnd = this.processedUpTo;

    while ((match = toolRegex.exec(this.buffer)) !== null) {
      // Only process new matches
      if (match.index < this.processedUpTo) continue;

      // Add text before the tool tag to textBuffer
      if (match.index > lastMatchEnd) {
        this.textBuffer += this.buffer.slice(lastMatchEnd, match.index);
      }

      const toolName = match[1];
      const rawContent = match[2];

      // Parse the JSON content
      let args: Record<string, unknown> = {};
      try {
        const trimmedContent = rawContent.trim();
        if (trimmedContent) {
          args = JSON.parse(trimmedContent);
        }
      } catch (e) {
        this.errors.push(`Failed to parse tool args for ${toolName}: ${e}`);
        args = {};
      }

      this.toolCalls.push({
        name: toolName,
        args,
        rawArgs: rawContent,
        startIndex: match.index,
        endIndex: match.index + match[0].length,
        complete: true,
      });

      lastMatchEnd = match.index + match[0].length;
      this.processedUpTo = lastMatchEnd;
    }

    // Check for incomplete tool tag at the end
    const remainingText = this.buffer.slice(this.processedUpTo);
    const incompleteToolRegex = /<tool\s+name="[^"]*"?\s*>?[^]*$/;
    const incompleteMatch = incompleteToolRegex.test(remainingText);

    if (!incompleteMatch && remainingText) {
      // No pending tool tag, add remaining text
      this.textBuffer += remainingText;
      this.processedUpTo = this.buffer.length;
    }

    // Return only new text since last call
    const newText = this.textBuffer.slice(this.lastReturnedTextLength);
    this.lastReturnedTextLength = this.textBuffer.length;

    return {
      safeText: newText,
      toolCalls: [...this.toolCalls],
      pendingToolCall: incompleteMatch,
      errors: [...this.errors],
    };
  }

  finalize(): ParseResult {
    // Process any remaining text
    const remainingText = this.buffer.slice(this.processedUpTo);
    if (remainingText && !/<tool\s/.test(remainingText)) {
      this.textBuffer += remainingText;
    }

    // Return any remaining text
    const newText = this.textBuffer.slice(this.lastReturnedTextLength);
    this.lastReturnedTextLength = this.textBuffer.length;

    return {
      safeText: newText,
      toolCalls: [...this.toolCalls],
      pendingToolCall: false,
      errors: [...this.errors],
    };
  }

  getAllText(): string {
    return this.textBuffer;
  }
}

export class StreamingToolParser {
  private parser: ToolParser;
  private lastToolCount = 0;
  private safeTextCallback?: (text: string) => void;
  private toolCallCallback?: (tool: ToolCall) => void;
  private errorCallback?: (error: string) => void;

  constructor(options?: {
    onSafeText?: (text: string) => void;
    onToolCall?: (tool: ToolCall) => void;
    onError?: (error: string) => void;
  }) {
    this.parser = new ToolParser();
    this.safeTextCallback = options?.onSafeText;
    this.toolCallCallback = options?.onToolCall;
    this.errorCallback = options?.onError;
  }

  write(chunk: string): ParseResult {
    const result = this.parser.process(chunk);

    // Handle callbacks
    if (result.safeText && this.safeTextCallback) {
      this.safeTextCallback(result.safeText);
    }

    if (result.toolCalls.length > this.lastToolCount) {
      const newTools = result.toolCalls.slice(this.lastToolCount);
      for (const tool of newTools) {
        if (tool.complete && this.toolCallCallback) {
          this.toolCallCallback(tool);
        }
      }
      this.lastToolCount = result.toolCalls.length;
    }

    if (result.errors.length > 0 && this.errorCallback) {
      for (const error of result.errors) {
        this.errorCallback(error);
      }
    }

    return result;
  }

  end(): ParseResult {
    return this.parser.finalize();
  }

  reset(): void {
    this.parser.reset();
    this.lastToolCount = 0;
  }
}

export function extractToolCalls(text: string): ToolCall[] {
  const parser = new ToolParser();
  parser.process(text);
  const result = parser.finalize();

  if (result.errors.length > 0) {
    console.warn("Tool parsing errors:", result.errors);
  }

  return result.toolCalls.filter((tc) => tc.complete);
}

export function filterToolCalls(text: string): string {
  const parser = new ToolParser();
  parser.process(text);
  parser.finalize();
  return parser.getAllText();
}