Polish streaming thinking traces

web/src/lib/chat-stream-presentation.ts

@@ -0,0 +1,139 @@
+import type { Message, ToolCallEvent } from "@/lib/api";
+
+export const STREAMING_ASSISTANT_METADATA_KIND = "streaming_assistant";
+
+export type StreamingAssistantMetadata = {
+  kind: typeof STREAMING_ASSISTANT_METADATA_KIND;
+  ephemeralTrace: string | null;
+  traceRevision: number;
+};
+
+export type StreamingAttemptState = {
+  segmentText: string;
+  initiatedToolCallIds: string[];
+};
+
+export type StreamingToolCallTransition = {
+  attempt: StreamingAttemptState;
+  startedNewToolCall: boolean;
+  traceText: string | null;
+};
+
+export function createStreamingAssistantMetadata(): StreamingAssistantMetadata {
+  return {
+    kind: STREAMING_ASSISTANT_METADATA_KIND,
+    ephemeralTrace: null,
+    traceRevision: 0,
+  };
+}
+
+export function asStreamingAssistantMetadata(value: unknown): StreamingAssistantMetadata | null {
+  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
+  const record = value as Record<string, unknown>;
+  if (record.kind !== STREAMING_ASSISTANT_METADATA_KIND) return null;
+
+  return {
+    kind: STREAMING_ASSISTANT_METADATA_KIND,
+    ephemeralTrace: typeof record.ephemeralTrace === "string" ? record.ephemeralTrace : null,
+    traceRevision:
+      typeof record.traceRevision === "number" && Number.isFinite(record.traceRevision)
+        ? Math.max(0, Math.floor(record.traceRevision))
+        : 0,
+  };
+}
+
+export function createStreamingAttemptState(): StreamingAttemptState {
+  return {
+    segmentText: "",
+    initiatedToolCallIds: [],
+  };
+}
+
+export function appendStreamingDelta(attempt: StreamingAttemptState, delta: string): StreamingAttemptState {
+  if (!delta) return attempt;
+  return {
+    ...attempt,
+    segmentText: attempt.segmentText + delta,
+  };
+}
+
+export function transitionStreamingToolCall(
+  attempt: StreamingAttemptState,
+  event: Pick<ToolCallEvent, "toolCallId" | "status">
+): StreamingToolCallTransition {
+  if (event.status !== "initiated" || attempt.initiatedToolCallIds.includes(event.toolCallId)) {
+    return { attempt, startedNewToolCall: false, traceText: null };
+  }
+
+  const traceText = attempt.segmentText.trim();
+  return {
+    attempt: {
+      segmentText: "",
+      initiatedToolCallIds: attempt.initiatedToolCallIds.concat(event.toolCallId),
+    },
+    startedNewToolCall: true,
+    traceText: traceText || null,
+  };
+}
+
+function updateStreamingAssistant(
+  messages: Message[],
+  assistantMessagePrefix: string,
+  update: (message: Message, metadata: StreamingAssistantMetadata) => Message
+) {
+  let didUpdate = false;
+  const nextMessages = messages.map((message, index, all) => {
+    const isTarget = index === all.length - 1 && message.id.startsWith(assistantMessagePrefix);
+    if (!isTarget) return message;
+    didUpdate = true;
+    return update(message, asStreamingAssistantMetadata(message.metadata) ?? createStreamingAssistantMetadata());
+  });
+  return didUpdate ? nextMessages : messages;
+}
+
+export function setStreamingAssistantSegment(messages: Message[], assistantMessagePrefix: string, segmentText: string) {
+  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
+    ...message,
+    content: segmentText,
+    metadata: {
+      ...metadata,
+      ephemeralTrace: null,
+    } satisfies StreamingAssistantMetadata,
+  }));
+}
+
+export function promoteStreamingAssistantTrace(messages: Message[], assistantMessagePrefix: string, traceText: string) {
+  const normalizedTrace = traceText.trim();
+  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
+    ...message,
+    content: "",
+    metadata: normalizedTrace
+      ? ({
+          ...metadata,
+          ephemeralTrace: normalizedTrace,
+          traceRevision: metadata.traceRevision + 1,
+        } satisfies StreamingAssistantMetadata)
+      : metadata,
+  }));
+}
+
+export function finalizeStreamingAssistant(messages: Message[], assistantMessagePrefix: string, finalText: string) {
+  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
+    ...message,
+    content: finalText,
+    metadata: {
+      ...metadata,
+      ephemeralTrace: null,
+    } satisfies StreamingAssistantMetadata,
+  }));
+}
+
+export function clearStreamingAssistantTrace(messages: Message[], assistantMessagePrefix: string) {
+  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
+    ...message,
+    metadata: {
+      ...metadata,
+      ephemeralTrace: null,
+    } satisfies StreamingAssistantMetadata,
+  }));
+}