metona-ai-desktop/electron/harness/adapters/shared/sse-stream.ts

/**
 * SSE 流式解析工具
 *
 * 解析 OpenAI 兼容的 Server-Sent Events (SSE) 流式响应，
 * 产出 MetonaStreamEvent。DeepSeek、Agnes AI 和 MiMo 共享此工具。
 *
 * SSE 格式：data: {json}\n\n
 * 结束标记：data: [DONE]
 */

import { nanoid } from 'nanoid';
import type { MetonaStreamEvent, MetonaTokenUsage } from '../../types';
import { MetonaStreamEventType } from '../../types';

/**
 * L-4 修复: 提取 flushToolCallBuffer 辅助函数，消除 [DONE] 分支和 finish_reason='tool_calls' 分支的重复代码
 *
 * 遍历工具调用缓冲区，对每个缓冲的工具调用：
 * 1. JSON.parse argsBuffer（失败则跳过）
 * 2. yield 一个 TOOL_CALL_COMPLETE 事件
 * 3. 清空缓冲区
 *
 * @param toolCallsBuffer - 工具调用缓冲区（index → { name, argsBuffer }）
 * @param requestId - 请求 ID
 * @param sessionId - 会话 ID
 * @param iteration - 当前迭代轮次
 * @param seqRef - seq 计数器引用（递增）
 * @yields MetonaStreamEvent
 */
function* flushToolCallBuffer(
  toolCallsBuffer: Map<number, { name: string; argsBuffer: string }>,
  requestId: string,
  sessionId: string,
  iteration: number,
  seqRef: { seq: number },
): Generator<MetonaStreamEvent> {
  for (const [, buf] of toolCallsBuffer) {
    try {
      yield {
        type: MetonaStreamEventType.TOOL_CALL_COMPLETE,
        requestId,
        sessionId,
        iteration,
        seq: seqRef.seq++,
        timestamp: Date.now(),
        toolCall: {
          id: `tc_${nanoid(8)}`,
          name: buf.name,
          args: buf.argsBuffer ? JSON.parse(buf.argsBuffer) : {},
          iteration,
          timestamp: Date.now(),
        },
      };
    } catch {
      // JSON 解析失败，跳过该工具调用
    }
  }
  toolCallsBuffer.clear();
}

/**
 * 解析 OpenAI 兼容 SSE 流式响应
 *
 * @param responseBody - fetch Response.body (ReadableStream<Uint8Array>)
 * @param requestId - 对应的请求 ID
 * @param sessionId - 会话 ID
 * @param iteration - 当前迭代轮次
 * @yields MetonaStreamEvent
 */
export async function* parseSSEStream(
  responseBody: ReadableStream<Uint8Array>,
  requestId: string,
  sessionId: string,
  iteration: number,
): AsyncGenerator<MetonaStreamEvent> {
  const reader = responseBody.getReader();
  const decoder = new TextDecoder();
  const seqRef = { seq: 0 };
  let buffer = '';

  // 工具调用缓冲区：index → { name, argsBuffer }
  const toolCallsBuffer = new Map<number, { name: string; argsBuffer: string }>();

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split('\n');
    buffer = lines.pop() ?? '';

    for (const line of lines) {
      const trimmed = line.trim();
      if (!trimmed || !trimmed.startsWith('data: ')) continue;
      const data = trimmed.slice(6);

      // 流结束
      if (data === '[DONE]') {
        // L-4 修复: 使用 flushToolCallBuffer 替代重复的遍历代码
        yield* flushToolCallBuffer(toolCallsBuffer, requestId, sessionId, iteration, seqRef);

        yield {
          type: MetonaStreamEventType.DONE,
          requestId,
          sessionId,
          iteration,
          seq: seqRef.seq++,
          timestamp: Date.now(),
        };
        return;
      }

      try {
        const chunk = JSON.parse(data);
        const delta = chunk.choices?.[0]?.delta;

        // 文本内容增量
        if (delta?.content) {
          yield {
            type: MetonaStreamEventType.TEXT_DELTA,
            requestId,
            sessionId,
            iteration,
            seq: seqRef.seq++,
            timestamp: Date.now(),
            delta: delta.content,
          };
        }

        // 推理内容增量（Thinking 模式）
        if (delta?.reasoning_content) {
          yield {
            type: MetonaStreamEventType.REASONING_DELTA,
            requestId,
            sessionId,
            iteration,
            seq: seqRef.seq++,
            timestamp: Date.now(),
            delta: delta.reasoning_content,
          };
        }

        // 工具调用增量 — 缓冲拼接
        if (delta?.tool_calls) {
          for (const tc of delta.tool_calls) {
            const idx = tc.index ?? 0;
            if (!toolCallsBuffer.has(idx)) {
              toolCallsBuffer.set(idx, { name: tc.function?.name ?? '', argsBuffer: '' });
            }
            const buf = toolCallsBuffer.get(idx)!;
            if (tc.function?.name) buf.name = tc.function.name;
            if (tc.function?.arguments) buf.argsBuffer += tc.function.arguments;

            yield {
              type: MetonaStreamEventType.TOOL_CALL_DELTA,
              requestId,
              sessionId,
              iteration,
              seq: seqRef.seq++,
              timestamp: Date.now(),
              toolCallDelta: {
                index: idx,
                name: tc.function?.name,
                argsDelta: tc.function?.arguments,
              },
            };
          }
        }

        // Token 使用统计 / finish_reason
        if (chunk.usage) {
          const usage: MetonaTokenUsage = {
            inputTokens: chunk.usage.prompt_tokens ?? 0,
            outputTokens: chunk.usage.completion_tokens ?? 0,
            totalTokens: chunk.usage.total_tokens ?? 0,
            reasoningTokens: chunk.usage.completion_tokens_details?.reasoning_tokens,
            // DeepSeek: prompt_cache_hit_tokens / prompt_cache_miss_tokens
            // MiMo: prompt_tokens_details.cached_tokens
            cacheHitTokens: chunk.usage.prompt_cache_hit_tokens
              ?? chunk.usage.prompt_tokens_details?.cached_tokens,
            cacheMissTokens: chunk.usage.prompt_cache_miss_tokens,
          };

          yield {
            type: MetonaStreamEventType.USAGE,
            requestId,
            sessionId,
            iteration,
            seq: seqRef.seq++,
            timestamp: Date.now(),
            usage,
          };
        }

        // 非 [DONE] 但 finish_reason 为 tool_calls 时提前 flush 缓冲区
        const finishReason = chunk.choices?.[0]?.finish_reason as string | undefined;
        if (finishReason === 'tool_calls') {
          // L-4 修复: 使用 flushToolCallBuffer 替代重复的遍历代码
          yield* flushToolCallBuffer(toolCallsBuffer, requestId, sessionId, iteration, seqRef);
        }
      } catch {
        // 跳过解析失败的行
      }
    }
  }
}

/**
 * 解析 OpenAI 兼容的非流式 JSON 响应 → MetonaResponse
 */
export function parseOpenAICompatibleResponse(
  data: Record<string, unknown>,
  requestId: string,
  provider: string,
  defaultModel: string,
): {
  content: string;
  reasoningContent?: string;
  toolCalls?: Array<{ id: string; name: string; args: Record<string, unknown>; iteration: number; timestamp: number }>;
  finishReason: string;
  usage: MetonaTokenUsage;
} {
  const choice = (data.choices as Array<Record<string, unknown>>)?.[0];
  const message = choice?.message as Record<string, unknown> | undefined;
  const usage = data.usage as Record<string, unknown> | undefined;
  const rawToolCalls = message?.tool_calls as Array<Record<string, unknown>> | undefined;

  return {
    content: (message?.content as string) ?? '',
    reasoningContent: message?.reasoning_content as string | undefined,
    toolCalls: rawToolCalls?.map((tc) => {
      const fn = tc.function as Record<string, unknown>;
      let args: Record<string, unknown> = {};
      const rawArgs = fn?.arguments;
      if (typeof rawArgs === 'string') {
        try {
          args = JSON.parse(rawArgs);
        } catch {
          args = {};
        }
      } else if (rawArgs && typeof rawArgs === 'object') {
        args = rawArgs as Record<string, unknown>;
      }
      return {
        id: tc.id as string,
        name: fn.name as string,
        args,
        iteration: 0,
        timestamp: Date.now(),
      };
    }),
    finishReason: mapOpenAIFinishReason(choice?.finish_reason as string),
    usage: {
      inputTokens: (usage?.prompt_tokens as number) ?? 0,
      outputTokens: (usage?.completion_tokens as number) ?? 0,
      totalTokens: (usage?.total_tokens as number) ?? 0,
      reasoningTokens: (usage?.completion_tokens_details as Record<string, unknown>)?.reasoning_tokens as number | undefined,
      // DeepSeek: prompt_cache_hit_tokens / MiMo: prompt_tokens_details.cached_tokens
      cacheHitTokens: (usage?.prompt_cache_hit_tokens as number | undefined)
        ?? (usage?.prompt_tokens_details as Record<string, unknown> | undefined)?.cached_tokens as number | undefined,
      cacheMissTokens: usage?.prompt_cache_miss_tokens as number | undefined,
    },
  };
}

function mapOpenAIFinishReason(reason: string): string {
  switch (reason) {
    case 'stop': return 'stop';
    case 'length': return 'length';
    case 'tool_calls': return 'tool_calls';
    case 'content_filter': return 'content_filter';
    // MiMo 特有：检测到复读截断
    case 'repetition_truncation': return 'stop';
    default: return 'stop';
  }
}