1
import type { AiUsage } from "@/shared/ai";
2
import type { AiFrame } from "@/worker/lib/ai/types";
3
 
4
export interface ProviderSseOptions {
5
  extractChunkText: (payload: unknown) => string | null;
6
  extractUsage?: (payload: unknown) => AiUsage | null;
7
  errorLabel: string;
8
  /** Fires for every parsed payload (before extraction). Used for debug shape sampling. */
9
  onPayload?: (payload: unknown, index: number) => void;
10
}
11
 
12
export async function* parseProviderSseFrames(
13
  source: ReadableStream<Uint8Array>,
14
  opts: ProviderSseOptions,
15
): AsyncGenerator<AiFrame> {
16
  const decoder = new TextDecoder();
17
  let buffer = "";
18
  let usage: AiUsage | undefined;
19
  let payloadIndex = 0;
20
  const reader = source.getReader();
21
 
22
  try {
23
    while (true) {
24
      const { value, done } = await reader.read();
25
      if (done) break;
26
      buffer += decoder.decode(value, { stream: true });
27
 
28
      let delimiterIndex = buffer.indexOf("\n\n");
29
      while (delimiterIndex !== -1) {
30
        const frame = buffer.slice(0, delimiterIndex);
31
        buffer = buffer.slice(delimiterIndex + 2);
32
        const raw = extractDataField(frame);
33
        if (raw !== null) {
34
          if (raw === "[DONE]") {
35
            if (usage) yield { type: "usage", usage };
36
            return;
37
          }
38
          const parsed = safeJsonParse(raw);
39
          if (parsed !== undefined) {
40
            opts.onPayload?.(parsed, payloadIndex++);
41
            if (opts.extractUsage) {
42
              const captured = opts.extractUsage(parsed);
43
              if (captured) usage = captured;
44
            }
45
            const text = opts.extractChunkText(parsed);
46
            if (text) yield { type: "chunk", text };
47
          }
48
        }
49
        delimiterIndex = buffer.indexOf("\n\n");
50
      }
51
    }
52
    if (usage) yield { type: "usage", usage };
53
  } catch {
54
    // Stream parser errors can carry upstream-derived text; emit only the safe
55
    // label so nothing leaks into client-facing SSE error frames.
56
    yield { type: "error", message: opts.errorLabel, code: "ai_chat_failed" };
57
  } finally {
58
    reader.releaseLock();
59
  }
60
}
61
 
62
export function parseProviderTokenUsage(raw: unknown): AiUsage | null {
63
  if (typeof raw !== "object" || raw === null) return null;
64
  const u = raw as { prompt_tokens?: unknown; completion_tokens?: unknown; total_tokens?: unknown };
65
  const usage: AiUsage = {};
66
  if (typeof u.prompt_tokens === "number") usage.promptTokens = u.prompt_tokens;
67
  if (typeof u.completion_tokens === "number") usage.completionTokens = u.completion_tokens;
68
  if (typeof u.total_tokens === "number") usage.totalTokens = u.total_tokens;
69
  return Object.keys(usage).length > 0 ? usage : null;
70
}
71
 
72
function extractDataField(frame: string): string | null {
73
  for (const line of frame.split("\n")) {
74
    if (line.startsWith("data: ")) return line.slice(6).trim();
75
  }
76
  return null;
77
}
78
 
79
function safeJsonParse(raw: string): unknown {
80
  try {
81
    return JSON.parse(raw);
82
  } catch {
83
    return undefined;
84
  }
85
}