@tonyclaw/agent-inspector 2.1.12 → 2.1.14

package/src/proxy/formats/openai/stream.ts

@@ -1,16 +1,163 @@
 import type { CapturedLog, JsonValue } from "../../schemas";
 import { trustAsJsonValue } from "../jsonSchema";
-import { OpenAISSERawChunkSchema } from "./schemas";
+import {
+  OpenAISSERawChunkSchema,
+  OpenAIResponsesResponseSchema,
+  OpenAIResponsesSSEventSchema,
+  type OpenAIResponsesResponse,
+} from "./schemas";
 
-/**
- * Extract OpenAI streaming response and return normalized JSON.
- * Optionally collects SSE chunks into log.streamingChunks for debugging/display.
- */
-export function extractOpenAIStream(
-  raw: string,
+type ParsedSseEvent = {
+  event: string | null;
+  data: string;
+};
+
+type CollectedChunk = {
+  index: number;
+  timestamp: number;
+  type: string;
+  data: JsonValue;
+};
+
+type ResponseFunctionCall = {
+  id?: string;
+  call_id?: string;
+  type: "function_call";
+  name?: string;
+  arguments: string;
+  status?: string;
+};
+
+const MAX_CHUNKS = 1000;
+
+function readProperty(value: unknown, key: string): unknown {
+  if (value === null || typeof value !== "object" || Array.isArray(value)) return undefined;
+  return Object.getOwnPropertyDescriptor(value, key)?.value;
+}
+
+function readStringProperty(value: unknown, key: string): string | null {
+  const property = readProperty(value, key);
+  return typeof property === "string" ? property : null;
+}
+
+function readNumberProperty(value: unknown, key: string): number | null {
+  const property = readProperty(value, key);
+  return typeof property === "number" && Number.isFinite(property) ? property : null;
+}
+
+function readObjectProperty(value: unknown, key: string): unknown | null {
+  const property = readProperty(value, key);
+  return property !== null && typeof property === "object" && !Array.isArray(property)
+    ? property
+    : null;
+}
+
+function readArrayProperty(value: unknown, key: string): unknown[] {
+  const property = readProperty(value, key);
+  return Array.isArray(property) ? property : [];
+}
+
+function parseSseEvents(raw: string): ParsedSseEvent[] {
+  const events: ParsedSseEvent[] = [];
+  let event: string | null = null;
+  let dataLines: string[] = [];
+
+  const commit = (): void => {
+    if (event === null && dataLines.length === 0) return;
+    events.push({ event, data: dataLines.join("\n") });
+    event = null;
+    dataLines = [];
+  };
+
+  for (const rawLine of raw.split("\n")) {
+    const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine;
+    if (line === "") {
+      commit();
+      continue;
+    }
+    if (line.startsWith(":")) continue;
+
+    const separatorIndex = line.indexOf(":");
+    const field = separatorIndex === -1 ? line : line.slice(0, separatorIndex);
+    const rawValue = separatorIndex === -1 ? "" : line.slice(separatorIndex + 1);
+    const value = rawValue.startsWith(" ") ? rawValue.slice(1) : rawValue;
+
+    switch (field) {
+      case "event":
+        event = value;
+        break;
+      case "data":
+        if (
+          event === null &&
+          dataLines.length > 0 &&
+          (value.startsWith("{") || value === "[DONE]")
+        ) {
+          commit();
+        }
+        dataLines.push(value);
+        break;
+      default:
+        break;
+    }
+  }
+
+  commit();
+  return events;
+}
+
+function parseJson(data: string): unknown | null {
+  try {
+    return JSON.parse(data);
+  } catch {
+    return null;
+  }
+}
+
+function eventTypeFor(event: ParsedSseEvent, parsed: unknown): string {
+  if (event.event !== null && event.event !== "") return event.event;
+  const typed = OpenAIResponsesSSEventSchema.safeParse(parsed);
+  return typed.success ? typed.data.type : "";
+}
+
+function isResponsesEvent(event: ParsedSseEvent): boolean {
+  if (event.data === "[DONE]") return false;
+  const parsed = parseJson(event.data);
+  const type = eventTypeFor(event, parsed);
+  return type.startsWith("response.") || type === "error";
+}
+
+function collectChunk(
+  chunks: CollectedChunk[],
+  chunkIndex: number,
+  streamStartMs: number,
+  type: string,
+  parsed: unknown,
+): void {
+  if (chunks.length >= MAX_CHUNKS) return;
+  chunks.push({
+    index: chunkIndex,
+    timestamp: Date.now() - streamStartMs,
+    type,
+    data: trustAsJsonValue(parsed),
+  });
+}
+
+function setErrorFromPayload(log: CapturedLog, payload: unknown): void {
+  const errorObject = readObjectProperty(payload, "error");
+  const message =
+    readStringProperty(errorObject, "message") ??
+    readStringProperty(payload, "message") ??
+    readStringProperty(payload, "error");
+  if (message !== null && message.length > 0) {
+    log.error = message;
+  }
+}
+
+function extractChatCompletionsStream(
+  events: ParsedSseEvent[],
   log: CapturedLog,
-  fallbackModel?: string,
-  collectChunks: boolean = true,
+  fallbackModel: string | undefined,
+  collectChunks: boolean,
 ): string {
   let id = "";
   let model = "";
@@ -27,145 +174,89 @@ export function extractOpenAIStream(
     type: "function";
     function: { name: string; arguments: string };
   };
-  // Map keyed by `index` so delta accumulation is O(1) per delta instead of
-  // O(N) (the previous Array#find). Map preserves insertion order, so the
-  // final `[...toolCalls.values()]` keeps tool calls in the order they were
-  // first seen.
   const toolCalls: Map<number, ToolCallEntry> = new Map();
 
-  // Chunk collection state
-  const MAX_CHUNKS = 1000;
   let chunkIndex = 0;
   let streamStartMs = 0;
-  const chunks: Array<{ index: number; timestamp: number; type: string; data: JsonValue }> = [];
-
-  for (const line of raw.split("\n")) {
-    const trimmedLine = line.trim();
-    if (!trimmedLine.startsWith("data: ")) continue;
-    const dataStr = trimmedLine.slice(6);
-    if (dataStr === "[DONE]") break;
-
-    try {
-      const parsed: unknown = JSON.parse(dataStr);
-      // Check for error object in SSE stream before trying to parse as chunk
-      if (typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)) {
-        const errorDesc = Object.getOwnPropertyDescriptor(parsed, "error");
-        if (
-          errorDesc !== undefined &&
-          typeof errorDesc.value === "object" &&
-          errorDesc.value !== null
-        ) {
-          const msgDesc = Object.getOwnPropertyDescriptor(errorDesc.value, "message");
-          if (msgDesc !== undefined && typeof msgDesc.value === "string") {
-            log.error = msgDesc.value;
-          }
-          continue;
-        }
-      }
-      const chunkResult = OpenAISSERawChunkSchema.safeParse(parsed);
-      if (!chunkResult.success) continue;
-      const chunk = chunkResult.data;
-
-      // Track stream start for timestamps
-      if (chunkIndex === 0) streamStartMs = Date.now();
-
-      // Collect chunks if enabled
-      if (collectChunks === true && chunks.length < MAX_CHUNKS) {
-        // `chunk` has already been validated by OpenAISSERawChunkSchema above;
-        // skipping the recursive JsonValue walk avoids a second full-tree zod
-        // parse on every SSE event. See trustAsJsonValue for the rationale.
-        chunks.push({
-          index: chunkIndex,
-          timestamp: Date.now() - streamStartMs,
-          type: "chat.completion.chunk",
-          data: trustAsJsonValue(chunk),
-        });
-      }
-      chunkIndex++;
+  const chunks: CollectedChunk[] = [];
 
-      if (!started) {
-        id = chunk.id;
-        model = chunk.model;
-        started = true;
-      }
+  for (const event of events) {
+    if (event.data === "[DONE]") break;
 
-      if (!usageCaptured && chunk.usage !== undefined && chunk.usage !== null) {
-        promptTokens = chunk.usage.prompt_tokens ?? 0;
-        completionTokens = chunk.usage.completion_tokens ?? 0;
-        log.inputTokens = promptTokens;
-        // Extract cached_tokens from raw parsed object (passthrough field in usage)
-        if (typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)) {
-          const usageDesc = Object.getOwnPropertyDescriptor(parsed, "usage");
-          if (
-            usageDesc !== undefined &&
-            typeof usageDesc.value === "object" &&
-            usageDesc.value !== null
-          ) {
-            const detailsDesc = Object.getOwnPropertyDescriptor(
-              usageDesc.value,
-              "prompt_tokens_details",
-            );
-            if (
-              detailsDesc !== undefined &&
-              typeof detailsDesc.value === "object" &&
-              detailsDesc.value !== null
-            ) {
-              const cacheDesc = Object.getOwnPropertyDescriptor(detailsDesc.value, "cached_tokens");
-              if (cacheDesc !== undefined && typeof cacheDesc.value === "number") {
-                log.cacheReadInputTokens = cacheDesc.value;
-              }
-            }
-          }
-        }
-        usageCaptured = true;
+    const parsed = parseJson(event.data);
+    if (parsed === null) continue;
+    setErrorFromPayload(log, parsed);
+
+    const chunkResult = OpenAISSERawChunkSchema.safeParse(parsed);
+    if (!chunkResult.success) continue;
+    const chunk = chunkResult.data;
+
+    if (chunkIndex === 0) streamStartMs = Date.now();
+    if (collectChunks) {
+      collectChunk(chunks, chunkIndex, streamStartMs, "chat.completion.chunk", chunk);
+    }
+    chunkIndex++;
+
+    if (!started) {
+      id = chunk.id;
+      model = chunk.model;
+      started = true;
+    }
+
+    if (!usageCaptured && chunk.usage !== undefined && chunk.usage !== null) {
+      promptTokens = chunk.usage.prompt_tokens ?? 0;
+      completionTokens = chunk.usage.completion_tokens ?? 0;
+      log.inputTokens = promptTokens;
+
+      const usage = readObjectProperty(parsed, "usage");
+      const details = readObjectProperty(usage, "prompt_tokens_details");
+      const cachedTokens = readNumberProperty(details, "cached_tokens");
+      if (cachedTokens !== null) {
+        log.cacheReadInputTokens = cachedTokens;
       }
+      usageCaptured = true;
+    }
 
-      for (const choice of chunk.choices) {
-        const delta = choice.delta;
-        if (delta.content !== undefined && delta.content !== null) {
-          completionContent += delta.content;
-        }
-        if (delta.reasoning_content !== undefined && delta.reasoning_content !== null) {
-          reasoningContent += delta.reasoning_content;
-        }
-        if (delta.thinking !== undefined && delta.thinking !== null) {
-          reasoningContent += delta.thinking;
-        }
-        // MiniMax uses 'think' field for thinking content - check via bracket notation
-        const thinkValue = delta.think;
-        if (thinkValue !== undefined && thinkValue !== null) {
-          reasoningContent += thinkValue;
-        }
-        if (choice.finish_reason !== undefined && choice.finish_reason !== null) {
-          finishReason = choice.finish_reason;
-        }
-        // Accumulate tool_calls
-        if (delta.tool_calls !== undefined && delta.tool_calls !== null) {
-          for (const tc of delta.tool_calls) {
-            // Find or create tool call entry by index
-            let existing = toolCalls.get(tc.index);
-            if (existing === undefined) {
-              existing = {
-                index: tc.index,
-                type: "function",
-                function: { name: "", arguments: "" },
-              };
-              toolCalls.set(tc.index, existing);
-            }
-            if (tc.id !== undefined) existing.id = tc.id;
-            if (tc.function?.name !== undefined) existing.function.name += tc.function.name;
-            if (tc.function?.arguments !== undefined)
-              existing.function.arguments += tc.function.arguments;
+    for (const choice of chunk.choices) {
+      const delta = choice.delta;
+      if (delta.content !== undefined && delta.content !== null) {
+        completionContent += delta.content;
+      }
+      if (delta.reasoning_content !== undefined && delta.reasoning_content !== null) {
+        reasoningContent += delta.reasoning_content;
+      }
+      if (delta.thinking !== undefined && delta.thinking !== null) {
+        reasoningContent += delta.thinking;
+      }
+      const thinkValue = delta.think;
+      if (thinkValue !== undefined && thinkValue !== null) {
+        reasoningContent += thinkValue;
+      }
+      if (choice.finish_reason !== undefined && choice.finish_reason !== null) {
+        finishReason = choice.finish_reason;
+      }
+      if (delta.tool_calls !== undefined && delta.tool_calls !== null) {
+        for (const tc of delta.tool_calls) {
+          let existing = toolCalls.get(tc.index);
+          if (existing === undefined) {
+            existing = {
+              index: tc.index,
+              type: "function",
+              function: { name: "", arguments: "" },
+            };
+            toolCalls.set(tc.index, existing);
+          }
+          if (tc.id !== undefined) existing.id = tc.id;
+          if (tc.function?.name !== undefined) existing.function.name += tc.function.name;
+          if (tc.function?.arguments !== undefined) {
+            existing.function.arguments += tc.function.arguments;
           }
         }
       }
-    } catch {
-      // non-JSON SSE line, skip
     }
   }
 
-  // Store collected chunks on the log
-  if (collectChunks === true) {
+  if (collectChunks) {
     log.streamingChunks = {
       chunks,
       truncated: chunkIndex > MAX_CHUNKS,
@@ -182,7 +273,7 @@ export function extractOpenAIStream(
     role: "assistant",
     content: completionContent,
   };
-  if (reasoningContent) message.reasoning_content = reasoningContent;
+  if (reasoningContent.length > 0) message.reasoning_content = reasoningContent;
   if (toolCalls.size > 0) message.tool_calls = [...toolCalls.values()];
 
   return JSON.stringify({
@@ -204,3 +295,306 @@ export function extractOpenAIStream(
     },
   });
 }
+
+function parseResponsesResponse(value: unknown): OpenAIResponsesResponse | null {
+  const fromResponseProperty = OpenAIResponsesResponseSchema.safeParse(
+    readProperty(value, "response"),
+  );
+  if (fromResponseProperty.success) return fromResponseProperty.data;
+  const direct = OpenAIResponsesResponseSchema.safeParse(value);
+  return direct.success ? direct.data : null;
+}
+
+function responseTextFromOutput(response: OpenAIResponsesResponse): string {
+  if (response.output_text !== undefined && response.output_text.length > 0) {
+    return response.output_text;
+  }
+
+  const parts: string[] = [];
+  for (const outputItem of response.output ?? []) {
+    const content = readArrayProperty(outputItem, "content");
+    for (const contentPart of content) {
+      const type = readStringProperty(contentPart, "type");
+      const text = readStringProperty(contentPart, "text");
+      if ((type === "output_text" || type === "text") && text !== null && text.length > 0) {
+        parts.push(text);
+      }
+    }
+  }
+  return parts.join("\n");
+}
+
+function updateUsageFromResponse(
+  response: OpenAIResponsesResponse,
+  log: CapturedLog,
+): { inputTokens: number | null; outputTokens: number | null } {
+  const usage = response.usage;
+  const inputTokens = usage?.input_tokens ?? null;
+  const outputTokens = usage?.output_tokens ?? null;
+  if (inputTokens !== null) log.inputTokens = inputTokens;
+  if (outputTokens !== null) log.outputTokens = outputTokens;
+  const cachedTokens = usage?.input_tokens_details?.cached_tokens ?? null;
+  if (cachedTokens !== null) log.cacheReadInputTokens = cachedTokens;
+  return { inputTokens, outputTokens };
+}
+
+function functionCallKey(payload: unknown, fallbackIndex: number): string {
+  const itemId = readStringProperty(payload, "item_id");
+  if (itemId !== null && itemId.length > 0) return itemId;
+  const outputIndex = readNumberProperty(payload, "output_index");
+  if (outputIndex !== null) return String(outputIndex);
+  return `function-${String(fallbackIndex)}`;
+}
+
+function ensureFunctionCall(
+  calls: Map<string, ResponseFunctionCall>,
+  key: string,
+): ResponseFunctionCall {
+  const existing = calls.get(key);
+  if (existing !== undefined) return existing;
+  const created: ResponseFunctionCall = {
+    type: "function_call",
+    arguments: "",
+  };
+  calls.set(key, created);
+  return created;
+}
+
+function updateFunctionCallFromItem(
+  calls: Map<string, ResponseFunctionCall>,
+  key: string,
+  item: unknown,
+): void {
+  if (readStringProperty(item, "type") !== "function_call") return;
+  const call = ensureFunctionCall(calls, key);
+  const id = readStringProperty(item, "id");
+  const callId = readStringProperty(item, "call_id");
+  const name = readStringProperty(item, "name");
+  const args = readStringProperty(item, "arguments");
+  const status = readStringProperty(item, "status");
+  if (id !== null) call.id = id;
+  if (callId !== null) call.call_id = callId;
+  if (name !== null) call.name = name;
+  if (args !== null) call.arguments = args;
+  if (status !== null) call.status = status;
+}
+
+function outputFromCollected(
+  outputText: string,
+  calls: Map<string, ResponseFunctionCall>,
+): unknown[] {
+  const output: unknown[] = [];
+  if (outputText.length > 0) {
+    output.push({
+      type: "message",
+      role: "assistant",
+      content: [{ type: "output_text", text: outputText, annotations: [] }],
+    });
+  }
+  for (const call of calls.values()) {
+    output.push(call);
+  }
+  return output;
+}
+
+function responseUsage(
+  inputTokens: number | null,
+  outputTokens: number | null,
+  outputText: string,
+): Record<string, unknown> {
+  const finalOutputTokens = outputTokens ?? Math.ceil(outputText.length / 4);
+  return {
+    input_tokens: inputTokens,
+    output_tokens: finalOutputTokens,
+    total_tokens: inputTokens === null ? finalOutputTokens : inputTokens + finalOutputTokens,
+  };
+}
+
+function normalizeResponsesResponse(input: {
+  completedResponse: OpenAIResponsesResponse | null;
+  responseId: string;
+  model: string;
+  status: string | null;
+  outputText: string;
+  calls: Map<string, ResponseFunctionCall>;
+  inputTokens: number | null;
+  outputTokens: number | null;
+}): Record<string, unknown> {
+  const completedText =
+    input.completedResponse === null ? "" : responseTextFromOutput(input.completedResponse);
+  const outputText = completedText.length > 0 ? completedText : input.outputText;
+  const completedOutput = input.completedResponse?.output;
+  const output =
+    completedOutput !== undefined && completedOutput.length > 0
+      ? completedOutput
+      : outputFromCollected(outputText, input.calls);
+  const usage =
+    input.completedResponse?.usage ??
+    responseUsage(input.inputTokens, input.outputTokens, outputText);
+
+  return {
+    ...(input.completedResponse ?? {}),
+    id: input.completedResponse?.id ?? input.responseId,
+    object: "response",
+    created_at: input.completedResponse?.created_at ?? Math.floor(Date.now() / 1000),
+    status: input.completedResponse?.status ?? input.status ?? "completed",
+    model: input.completedResponse?.model ?? input.model,
+    output,
+    output_text: outputText,
+    usage,
+  };
+}
+
+function extractResponsesStream(
+  events: ParsedSseEvent[],
+  log: CapturedLog,
+  fallbackModel: string | undefined,
+  collectChunks: boolean,
+): string {
+  let completedResponse: OpenAIResponsesResponse | null = null;
+  let responseId = "";
+  let model = fallbackModel ?? "";
+  let status: string | null = null;
+  let outputText = "";
+  let inputTokens: number | null = null;
+  let outputTokens: number | null = null;
+  let chunkIndex = 0;
+  let streamStartMs = 0;
+  const chunks: CollectedChunk[] = [];
+  const calls: Map<string, ResponseFunctionCall> = new Map();
+
+  for (const event of events) {
+    if (event.data === "[DONE]") break;
+    const parsed = parseJson(event.data);
+    if (parsed === null) continue;
+    const type = eventTypeFor(event, parsed);
+
+    if (chunkIndex === 0) streamStartMs = Date.now();
+    if (collectChunks) {
+      collectChunk(
+        chunks,
+        chunkIndex,
+        streamStartMs,
+        type === "" ? "response.event" : type,
+        parsed,
+      );
+    }
+    chunkIndex++;
+
+    setErrorFromPayload(log, parsed);
+    if (type === "error") continue;
+
+    const response = parseResponsesResponse(parsed);
+    if (response !== null) {
+      completedResponse = response;
+      responseId = response.id;
+      model = response.model ?? model;
+      status = response.status ?? status;
+      const usage = updateUsageFromResponse(response, log);
+      inputTokens = usage.inputTokens ?? inputTokens;
+      outputTokens = usage.outputTokens ?? outputTokens;
+      setErrorFromPayload(log, response);
+    }
+
+    const delta = readStringProperty(parsed, "delta");
+    switch (type) {
+      case "response.output_text.delta":
+        if (delta !== null) outputText += delta;
+        break;
+      case "response.output_text.done": {
+        const text = readStringProperty(parsed, "text");
+        if (outputText.length === 0 && text !== null) outputText = text;
+        break;
+      }
+      case "response.function_call_arguments.delta": {
+        if (delta !== null) {
+          const call = ensureFunctionCall(calls, functionCallKey(parsed, calls.size));
+          call.arguments += delta;
+        }
+        break;
+      }
+      case "response.function_call_arguments.done": {
+        const args = readStringProperty(parsed, "arguments");
+        if (args !== null) {
+          const call = ensureFunctionCall(calls, functionCallKey(parsed, calls.size));
+          call.arguments = args;
+        }
+        break;
+      }
+      case "response.output_item.added":
+      case "response.output_item.done": {
+        const item = readObjectProperty(parsed, "item");
+        updateFunctionCallFromItem(calls, functionCallKey(parsed, calls.size), item);
+        break;
+      }
+      case "response.failed":
+      case "response.incomplete":
+        status = type.slice("response.".length);
+        break;
+      case "response.completed":
+        status = "completed";
+        break;
+      case "response.created":
+      case "response.in_progress":
+      case "response.content_part.added":
+      case "response.content_part.done":
+      case "":
+        break;
+      default:
+        break;
+    }
+  }
+
+  if (collectChunks) {
+    log.streamingChunks = {
+      chunks,
+      truncated: chunkIndex > MAX_CHUNKS,
+    };
+  }
+
+  const normalized = normalizeResponsesResponse({
+    completedResponse,
+    responseId,
+    model,
+    status,
+    outputText,
+    calls,
+    inputTokens,
+    outputTokens,
+  });
+  const normalizedResponse = OpenAIResponsesResponseSchema.safeParse(normalized);
+  if (normalizedResponse.success) {
+    const finalUsage = updateUsageFromResponse(normalizedResponse.data, log);
+    if (finalUsage.outputTokens === null) {
+      const finalText = responseTextFromOutput(normalizedResponse.data);
+      const estimatedOutputTokens = Math.ceil(finalText.length / 4);
+      log.outputTokens = estimatedOutputTokens;
+    }
+    return JSON.stringify(normalizedResponse.data);
+  }
+
+  if (log.outputTokens === null) {
+    const finalText = readStringProperty(normalized, "output_text") ?? outputText;
+    const estimatedOutputTokens = Math.ceil(finalText.length / 4);
+    log.outputTokens = estimatedOutputTokens;
+  }
+
+  return JSON.stringify(normalized);
+}
+
+/**
+ * Extract OpenAI streaming response and return normalized JSON.
+ * Optionally collects SSE chunks into log.streamingChunks for debugging/display.
+ */
+export function extractOpenAIStream(
+  raw: string,
+  log: CapturedLog,
+  fallbackModel?: string,
+  collectChunks: boolean = true,
+): string {
+  const events = parseSseEvents(raw);
+  if (events.some(isResponsesEvent)) {
+    return extractResponsesStream(events, log, fallbackModel, collectChunks);
+  }
+  return extractChatCompletionsStream(events, log, fallbackModel, collectChunks);
+}