@futurity/chat-react 0.0.2 → 0.1.0

package/src/stream-accumulator.ts

@@ -0,0 +1,448 @@
+import {
+  type MessagePart,
+  Z_DataEndSplitSchema,
+  Z_DataSplitSchema,
+  Z_DataSubagentPartSchema,
+} from "@futurity/chat-protocol";
+
+export type ProcessedPart =
+  | { type: "regular"; part: MessagePart; originalIndex: number }
+  | {
+      type: "split-group";
+      title?: string;
+      subtitle?: string;
+      parts: MessagePart[];
+      startIndex: number;
+      endIndex: number;
+      desktopSessionId?: string;
+    };
+
+// biome-ignore lint/suspicious/noExplicitAny: raw stream protocol chunks have dynamic shapes
+type Chunk = Record<string, any>;
+
+/** Narrow `unknown` to a `Chunk` whose `type` is a string. */
+export function isTypedChunk(x: unknown): x is Chunk & { type: string } {
+  return (
+    typeof x === "object" &&
+    x !== null &&
+    "type" in x &&
+    typeof x.type === "string"
+  );
+}
+
+/**
+ * Incrementally accumulates raw UI message stream protocol chunks into
+ * rendered message parts. Maintains internal state so repeated calls with a
+ * growing `rawParts` array only process newly appended chunks.
+ *
+ * Internally tracks mutable objects via `Chunk` maps (text-by-id, tool-by-id)
+ * and casts to `MessagePart` at the output boundary. The constructed shapes
+ * conform to the `MessagePart` union at runtime.
+ */
+export class StreamAccumulator {
+  readonly parts: MessagePart[] = [];
+  private textById = new Map<string, Chunk>();
+  private reasoningById = new Map<string, Chunk>();
+  private toolById = new Map<string, Chunk>();
+  private partialToolText = new Map<string, string>();
+  private processedCount = 0;
+
+  private push(part: Chunk): void {
+    this.parts.push(part as MessagePart);
+  }
+
+  /** Process any new raw chunks and return the full accumulated parts array. */
+  accumulate(rawParts: unknown[]): MessagePart[] {
+    for (let i = this.processedCount; i < rawParts.length; i++) {
+      const chunk = rawParts[i];
+      if (!isTypedChunk(chunk)) {
+        this.parts.push(chunk as MessagePart);
+        continue;
+      }
+
+      switch (chunk.type) {
+        // --- Text accumulation ---
+        case "text-start": {
+          const part = { type: "text", text: "", state: "streaming" };
+          this.textById.set(chunk.id, part);
+          this.push(part);
+          break;
+        }
+        case "text-delta": {
+          const part = this.textById.get(chunk.id);
+          if (part) part.text += chunk.delta;
+          break;
+        }
+        case "text-end": {
+          const part = this.textById.get(chunk.id);
+          if (part) part.state = "done";
+          break;
+        }
+
+        // --- Reasoning accumulation ---
+        case "reasoning-start": {
+          const part = { type: "reasoning", text: "", state: "streaming" };
+          this.reasoningById.set(chunk.id, part);
+          this.push(part);
+          break;
+        }
+        case "reasoning-delta": {
+          const part = this.reasoningById.get(chunk.id);
+          if (part) part.text += chunk.delta;
+          break;
+        }
+        case "reasoning-end": {
+          const part = this.reasoningById.get(chunk.id);
+          if (part) part.state = "done";
+          break;
+        }
+
+        // --- Tool invocation accumulation ---
+        case "tool-input-start": {
+          const typeName = chunk.dynamic
+            ? "dynamic-tool"
+            : `tool-${chunk.toolName}`;
+          const part: Chunk = {
+            type: typeName,
+            toolCallId: chunk.toolCallId,
+            toolName: chunk.toolName,
+            state: "input-streaming",
+            input: undefined,
+            providerExecuted: chunk.providerExecuted,
+            title: chunk.title,
+          };
+          this.toolById.set(chunk.toolCallId, part);
+          this.partialToolText.set(chunk.toolCallId, "");
+          this.push(part);
+          break;
+        }
+        case "tool-input-delta": {
+          const text =
+            (this.partialToolText.get(chunk.toolCallId) ?? "") +
+            (chunk.inputTextDelta ?? "");
+          this.partialToolText.set(chunk.toolCallId, text);
+          const part = this.toolById.get(chunk.toolCallId);
+          if (part) {
+            try {
+              part.input = JSON.parse(text);
+            } catch {
+              // Partial JSON — leave input as-is until complete
+            }
+          }
+          break;
+        }
+        case "tool-input-available": {
+          const typeName = chunk.dynamic
+            ? "dynamic-tool"
+            : `tool-${chunk.toolName}`;
+          let part = this.toolById.get(chunk.toolCallId);
+          if (part) {
+            part.state = "input-available";
+            part.input = chunk.input;
+            part.providerExecuted = chunk.providerExecuted;
+          } else {
+            part = {
+              type: typeName,
+              toolCallId: chunk.toolCallId,
+              toolName: chunk.toolName,
+              state: "input-available",
+              input: chunk.input,
+              providerExecuted: chunk.providerExecuted,
+              title: chunk.title,
+            };
+            this.toolById.set(chunk.toolCallId, part);
+            this.push(part);
+          }
+          break;
+        }
+        case "tool-input-error": {
+          const part = this.toolById.get(chunk.toolCallId);
+          if (part) {
+            part.state = "output-error";
+            part.errorText = chunk.errorText;
+            part.input = chunk.input;
+          }
+          break;
+        }
+        case "tool-output-available": {
+          const part = this.toolById.get(chunk.toolCallId);
+          if (part) {
+            part.state = "output-available";
+            part.output = chunk.output;
+            part.preliminary = chunk.preliminary;
+          }
+          break;
+        }
+        case "tool-output-error": {
+          const part = this.toolById.get(chunk.toolCallId);
+          if (part) {
+            part.state = "output-error";
+            part.errorText = chunk.errorText;
+          }
+          break;
+        }
+
+        // --- Pass-through parts ---
+        case "file":
+          this.push({
+            type: "file",
+            mediaType: chunk.mediaType,
+            url: chunk.url,
+            filename: chunk.filename,
+          });
+          break;
+        case "source-url":
+          this.push({
+            type: "source-url",
+            sourceId: chunk.sourceId,
+            url: chunk.url,
+            title: chunk.title,
+          });
+          break;
+        case "source-document":
+          this.push({
+            type: "source-document",
+            sourceId: chunk.sourceId,
+            mediaType: chunk.mediaType,
+            title: chunk.title,
+            filename: chunk.filename,
+          });
+          break;
+        case "start-step":
+          this.push({ type: "step-start" });
+          break;
+
+        // --- Ignored protocol chunks ---
+        case "finish-step":
+        case "start":
+        case "finish":
+        case "error":
+          break;
+
+        // --- Unknown / custom data parts → pass through ---
+        default:
+          if (chunk.type.startsWith("data-")) {
+            this.push(chunk);
+          }
+          break;
+      }
+    }
+
+    this.processedCount = rawParts.length;
+    return this.parts;
+  }
+}
+
+type PendingGroup = {
+  startIndex: number;
+  title?: string;
+  subtitle?: string;
+  desktopSessionId?: string;
+  subAgentId: string;
+  /** Raw stream protocol chunks from data-subagent-part messages (pre-accumulation). */
+  innerParts: unknown[];
+  endIndex: number | null;
+};
+
+export type PreprocessorState = {
+  messageId: string | null;
+  /** How many raw parts pass 1 has already scanned. */
+  scannedLength: number;
+  /** Indices consumed by subAgentId-based groups. */
+  claimedIndices: Set<number>;
+  /** Groups still collecting inner parts (keyed by subAgentId). */
+  openGroups: Map<string, PendingGroup>;
+  /** All groups keyed by their startIndex (for pass 2 lookup). */
+  groupsByStartIndex: Map<number, PendingGroup>;
+  /** Reverse lookup: subAgentId → group (includes closed groups). */
+  groupBySubAgentId: Map<string, PendingGroup>;
+  /** Per-group stream accumulators (keyed by group startIndex). */
+  accumulators: Map<number, StreamAccumulator>;
+};
+
+export function createPreprocessorState(messageId?: string): PreprocessorState {
+  return {
+    messageId: messageId ?? null,
+    scannedLength: 0,
+    claimedIndices: new Set(),
+    openGroups: new Map(),
+    groupsByStartIndex: new Map(),
+    groupBySubAgentId: new Map(),
+    accumulators: new Map(),
+  };
+}
+
+/**
+ * Incrementally preprocess message parts to group subagent content.
+ *
+ * Pass 1 (incremental): only scans parts appended since the last call,
+ * building group data structures without re-parsing already-seen parts.
+ *
+ * Per-group accumulation (incremental): each group's stream protocol chunks
+ * are accumulated via a stateful {@link StreamAccumulator} that only processes
+ * new chunks on each call.
+ *
+ * Pass 2 (rebuild): constructs the output array from cached state. This is
+ * cheap — O(n) Map/Set lookups with no Zod parsing for subAgentId groups.
+ *
+ * Falls back to legacy greedy scanning for old messages without subAgentId.
+ */
+export function incrementalPreprocess(
+  state: PreprocessorState,
+  parts: MessagePart[],
+): ProcessedPart[] {
+  // --- Incremental Pass 1: only scan newly appended parts ---
+  for (let i = state.scannedLength; i < parts.length; i++) {
+    const part = parts[i];
+    if (!isTypedChunk(part)) continue;
+
+    if (part.type === "data-split") {
+      const splitParse = Z_DataSplitSchema.safeParse(part);
+      if (splitParse.success && splitParse.data.data.subAgentId) {
+        const sid = splitParse.data.data.subAgentId;
+        const group: PendingGroup = {
+          startIndex: i,
+          title: splitParse.data.data.title,
+          subtitle: splitParse.data.data.subtitle,
+          desktopSessionId: splitParse.data.data.desktopSessionId,
+          subAgentId: sid,
+          innerParts: [],
+          endIndex: null,
+        };
+        state.openGroups.set(sid, group);
+        state.groupsByStartIndex.set(i, group);
+        state.groupBySubAgentId.set(sid, group);
+        state.claimedIndices.add(i);
+        continue;
+      }
+    }
+
+    if (part.type === "data-endsplit") {
+      const endSplitParse = Z_DataEndSplitSchema.safeParse(part);
+      if (endSplitParse.success && endSplitParse.data?.data?.subAgentId) {
+        const sid = endSplitParse.data.data.subAgentId;
+        const group = state.openGroups.get(sid);
+        if (group) {
+          group.endIndex = i;
+          group.desktopSessionId =
+            group.desktopSessionId ?? endSplitParse.data.data.desktopSessionId;
+          state.openGroups.delete(sid);
+        }
+        state.claimedIndices.add(i);
+        continue;
+      }
+    }
+
+    if (part.type === "data-subagent-part") {
+      const subagentParse = Z_DataSubagentPartSchema.safeParse(part);
+      if (subagentParse.success) {
+        const sid = subagentParse.data.data.subAgentId;
+        // Prefer open group, fall back to closed group (late-arriving part)
+        const group =
+          state.openGroups.get(sid) ?? state.groupBySubAgentId.get(sid);
+        if (group) {
+          group.innerParts.push(subagentParse.data.data.part);
+        }
+        state.claimedIndices.add(i);
+      }
+    }
+  }
+  state.scannedLength = parts.length;
+
+  // --- Pass 2: build output array using cached accumulators ---
+  const processed: ProcessedPart[] = [];
+
+  let i = 0;
+  while (i < parts.length) {
+    // Check if a subAgentId group starts here
+    const groupAtIndex = state.groupsByStartIndex.get(i);
+    if (groupAtIndex) {
+      // Get or create accumulator for this group
+      let acc = state.accumulators.get(groupAtIndex.startIndex);
+      if (!acc) {
+        acc = new StreamAccumulator();
+        state.accumulators.set(groupAtIndex.startIndex, acc);
+      }
+
+      processed.push({
+        type: "split-group",
+        title: groupAtIndex.title,
+        subtitle: groupAtIndex.subtitle,
+        // Incrementally accumulate — only processes newly added innerParts
+        parts: acc.accumulate(groupAtIndex.innerParts),
+        startIndex: groupAtIndex.startIndex,
+        endIndex: groupAtIndex.endIndex ?? parts.length,
+        desktopSessionId: groupAtIndex.desktopSessionId,
+      });
+      // Don't jump to endIndex — parallel subagent groups have overlapping
+      // index ranges. All inner data-subagent-part indices are claimed, so
+      // advancing one-by-one correctly skips them while still visiting other
+      // groups whose startIndex falls within this range.
+      i++;
+      continue;
+    }
+
+    // Skip any already-claimed index (subagent-part or endsplit handled above)
+    if (state.claimedIndices.has(i)) {
+      i++;
+      continue;
+    }
+
+    const part = parts[i];
+    const type: string | undefined = isTypedChunk(part) ? part.type : undefined;
+
+    // Legacy: data-split WITHOUT subAgentId → greedy first-match (backward compat)
+    if (type === "data-split") {
+      const splitParse = Z_DataSplitSchema.safeParse(part);
+      if (splitParse.success && !splitParse.data.data.subAgentId) {
+        let endIndex = parts.length;
+        let desktopSessionId = splitParse.data.data.desktopSessionId;
+        for (
+          let searchIndex = i + 1;
+          searchIndex < parts.length;
+          searchIndex++
+        ) {
+          if (state.claimedIndices.has(searchIndex)) continue;
+          const checkPart = parts[searchIndex];
+          const endParse = Z_DataEndSplitSchema.safeParse(checkPart);
+          if (endParse.success && !endParse.data?.data?.subAgentId) {
+            endIndex = searchIndex;
+            desktopSessionId =
+              desktopSessionId ?? endParse.data?.data?.desktopSessionId;
+            break;
+          }
+        }
+
+        const groupedParts = parts.slice(i + 1, endIndex).filter((_, idx) => {
+          const absIdx = i + 1 + idx;
+          return !state.claimedIndices.has(absIdx);
+        });
+
+        processed.push({
+          type: "split-group",
+          title: splitParse.data.data.title,
+          subtitle: splitParse.data.data.subtitle,
+          parts: groupedParts,
+          startIndex: i,
+          endIndex,
+          desktopSessionId,
+        });
+
+        i = endIndex < parts.length ? endIndex + 1 : endIndex;
+        continue;
+      }
+    }
+
+    // Orphaned data-endsplit without subAgentId → skip
+    if (type === "data-endsplit") {
+      i++;
+      continue;
+    }
+
+    // Regular part
+    processed.push({ type: "regular", part, originalIndex: i });
+    i++;
+  }
+
+  return processed;
+}

package/src/types.ts

@@ -8,6 +8,7 @@ import type {
 } from "@futurity/chat-protocol";
 import { type MessagePart, Z_MessagePart } from "@futurity/chat-protocol";
 import { z } from "zod";
+import type { ProcessedPart } from "./stream-accumulator";
 
 export type { MessagePart };
 
@@ -22,6 +23,8 @@ export type ChatMessage = {
   id: string;
   role: "user" | "assistant" | "system";
   parts: MessagePart[];
+  /** Pre-computed processed parts with subagent grouping and stream accumulation. */
+  processedParts: ProcessedPart[];
   createdAt?: Date;
   metadata?: MessageMetadata;
 };