| 1 | import type { WireEvent } from "./types"; |
| 2 | import type { LiveStream } from "./useController"; |
| 3 | |
| 4 | export interface StreamDeltaEntry { |
| 5 | tabId: string; |
| 6 | e: WireEvent; |
| 7 | } |
| 8 | |
| 9 | // StreamSegment is one run of consecutive same-kind deltas within a frame. |
| 10 | // Segment order is authoritative: a reasoning→text boundary completes |
| 11 | // reasoning exactly as per-delta delivery would, so kinds are never bucketed. |
| 12 | export interface StreamSegment { |
| 13 | kind: "text" | "reasoning"; |
| 14 | delta: string; |
| 15 | } |
| 16 | |
| 17 | export interface TabStreamBatch { |
| 18 | tabId: string; |
| 19 | segments: StreamSegment[]; |
| 20 | } |
| 21 | |
| 22 | // coalesceStreamDeltas groups one rAF batch by tab and merges consecutive |
| 23 | // same-kind deltas into ordered segments, so a frame dispatches one |
| 24 | // stream_batch action (one reducer pass, one live-store notification) per tab |
| 25 | // no matter how many token deltas the bridge delivered. Tabs are independent |
| 26 | // state machines, so per-tab grouping cannot reorder anything observable. |
| 27 | // Empty deltas are kept: an empty text delta still completes live reasoning. |
| 28 | export function coalesceStreamDeltas(batch: StreamDeltaEntry[]): TabStreamBatch[] { |
| 29 | const out: TabStreamBatch[] = []; |
| 30 | const byTab = new Map<string, StreamSegment[]>(); |
| 31 | for (const { tabId, e } of batch) { |
| 32 | const kind = e.kind === "reasoning" ? "reasoning" : "text"; |
| 33 | const delta = e.text ?? e.reasoning ?? ""; |
| 34 | let segments = byTab.get(tabId); |
| 35 | if (!segments) { |
| 36 | segments = []; |
| 37 | byTab.set(tabId, segments); |
| 38 | out.push({ tabId, segments }); |
| 39 | } |
| 40 | const last = segments[segments.length - 1]; |
| 41 | if (last && last.kind === kind) last.delta += delta; |
| 42 | else segments.push({ kind, delta }); |
| 43 | } |
| 44 | return out; |
| 45 | } |
| 46 | |
| 47 | export function completeLiveReasoning(live: LiveStream, now = Date.now()): LiveStream { |
| 48 | if (!live.reasoning || live.reasoningCompletedAt) { |
| 49 | return { ...live, reasoningComplete: live.reasoning !== "" || live.reasoningComplete }; |
| 50 | } |
| 51 | return { |
| 52 | ...live, |
| 53 | reasoningComplete: true, |
| 54 | reasoningCompletedAt: now, |
| 55 | }; |
| 56 | } |
| 57 | |
| 58 | // applyLiveSegments folds one frame's ordered segments into the live stream, |
| 59 | // replicating per-delta semantics: text completes reasoning first; reasoning |
| 60 | // reopens it and stamps its start on the first non-empty delta. |
| 61 | export function applyLiveSegments(base: LiveStream, segments: StreamSegment[], now: number): LiveStream { |
| 62 | let live = base; |
| 63 | for (const seg of segments) { |
| 64 | live = |
| 65 | seg.kind === "text" |
| 66 | ? { ...completeLiveReasoning(live, now), text: live.text + seg.delta } |
| 67 | : { |
| 68 | ...live, |
| 69 | reasoning: live.reasoning + seg.delta, |
| 70 | reasoningComplete: false, |
| 71 | reasoningStartedAt: live.reasoningStartedAt ?? (seg.delta ? now : undefined), |
| 72 | reasoningCompletedAt: undefined, |
| 73 | }; |
| 74 | } |
| 75 | return live; |
| 76 | } |
| 77 |