@voyager-0x/agent-mcp/src/streamBuffer.ts

Voyager MCP Agent - A powerful Model Context Protocol agent

/**
 * 智能流式消息缓存器
 * 用于处理流式输出中的XML标签，确保完整标签生成后再发送
 */
export class StreamBuffer {
  private buffer: string = "";
  private isBuffering: boolean = false;

  /**
   * 处理流式消息片段
   * @param deltaMessage 消息片段
   * @returns 可以立即发送的内容，如果需要缓存则返回空字符串
   */
  processChunk(deltaMessage: string): string {
    // 将新内容添加到缓存
    this.buffer += deltaMessage;

    return this.extractAndSend();
  }

  /**
   * 提取可以发送的内容
   * @returns 可以发送的内容
   */
  private extractAndSend(): string {
    let result = "";
    let workingBuffer = this.buffer;

    while (workingBuffer.length > 0) {
      // 查找 '<' 的位置
      const ltIndex = workingBuffer.indexOf("<");

      if (ltIndex === -1) {
        // 没有 '<'，全部内容都可以发送
        result += workingBuffer;
        workingBuffer = "";
        this.isBuffering = false;
        break;
      }

      // 有 '<'，先发送 '<' 之前的内容
      if (ltIndex > 0) {
        result += workingBuffer.substring(0, ltIndex);
        workingBuffer = workingBuffer.substring(ltIndex);
      }

      // 现在 workingBuffer 以 '<' 开头
      // 检查是否是完整的目标标签
      const completeTagEnd = this.findCompleteTagEnd(workingBuffer);

      if (completeTagEnd !== -1) {
        // 找到完整标签，发送它
        const completeTag = workingBuffer.substring(0, completeTagEnd + 1);
        result += completeTag;
        workingBuffer = workingBuffer.substring(completeTagEnd + 1);
        this.isBuffering = false;
        continue;
      }

      // 没有完整标签，检查是否可能是目标标签的开始
      if (this.isPotentialTargetTag(workingBuffer)) {
        // 可能是目标标签，需要缓存等待更多内容
        this.isBuffering = true;
        break;
      } else {
        // 不是目标标签，发送第一个字符并继续
        result += workingBuffer.charAt(0);
        workingBuffer = workingBuffer.substring(1);
        this.isBuffering = false;
      }
    }

    // 更新缓存
    this.buffer = workingBuffer;
    return result;
  }

  /**
   * 查找完整标签的结束位置
   * @param content 从 '<' 开始的内容
   * @returns 结束标签的最后一个字符位置，如果没有找到返回-1
   */
  private findCompleteTagEnd(content: string): number {
    // 检查工具调用标签 - 必须有开始和结束标签
    if (content.startsWith("<use_mcp_tool>")) {
      const endTag = "</use_mcp_tool>";
      const endIndex = content.indexOf(endTag);
      if (endIndex !== -1) {
        // 确保结束标签完整
        const endTagEnd = endIndex + endTag.length;
        if (content.length >= endTagEnd) {
          return endTagEnd - 1;
        }
      }
    }

    // 检查资源读取标签 - 必须有开始和结束标签
    if (content.startsWith("<use_mcp_resource>")) {
      const endTag = "</use_mcp_resource>";
      const endIndex = content.indexOf(endTag);
      if (endIndex !== -1) {
        // 确保结束标签完整
        const endTagEnd = endIndex + endTag.length;
        if (content.length >= endTagEnd) {
          return endTagEnd - 1;
        }
      }
    }

    return -1;
  }

  /**
   * 检查内容是否可能是目标标签的开始
   * @param content 从 '<' 开始的内容
   * @returns 是否可能是目标标签
   */
  private isPotentialTargetTag(content: string): boolean {
    const targetTags = ["<use_mcp_tool>", "<use_mcp_resource>"];

    for (const tag of targetTags) {
      // 检查内容是否是目标标签的前缀
      if (content.length <= tag.length && tag.startsWith(content)) {
        return true;
      }

      // 如果已经匹配了开始标签，但还没有完整的结束标签，也要继续缓存
      if (content.startsWith(tag)) {
        const endTag =
          tag === "<use_mcp_tool>" ? "</use_mcp_tool>" : "</use_mcp_resource>";
        const endIndex = content.indexOf(endTag);
        if (endIndex === -1) {
          // 有开始标签但没有结束标签，继续缓存
          return true;
        } else {
          // 检查结束标签是否完整
          const endTagEnd = endIndex + endTag.length;
          if (content.length < endTagEnd) {
            // 结束标签不完整，继续缓存
            return true;
          }
        }
      }
    }

    return false;
  }

  /**
   * 重置缓存器
   */
  reset(): void {
    this.buffer = "";
    this.isBuffering = false;
  }

  /**
   * 强制刷新缓存（在会话结束时调用）
   * @returns 缓存中的所有内容
   */
  flush(): string {
    const result = this.buffer;
    this.reset();
    return result;
  }

  /**
   * 获取当前是否正在缓存
   */
  isCurrentlyBuffering(): boolean {
    return this.isBuffering;
  }

  /**
   * 获取当前缓存内容（用于调试）
   */
  getBufferContent(): string {
    return this.buffer;
  }
}