@arcote.tech/arc-chat 0.7.9 → 0.7.11

package/src/routes/chat-stream-route.ts

@@ -7,18 +7,29 @@ export function createChatStreamRoute(config: {
   userToken: Token;
 }) {
   return route(`${config.name}ChatStream`)
-    .path(`/chat/${config.name}/stream/:streamId`)
+    .path(`/chat/${config.name}/stream/:messageId`)
     .protectBy(config.userToken, () => true)
     .handle({
-      GET: async (_ctx, req: Request, params: Record<string, string>) => {
-        // Klient po reload przekazuje `?afterSeq=N` z `partialLastSeq`
-        // odczytanego z DB. Replay buffer pomija eventy już zaaplikowane.
-        const url = new URL(req.url);
-        const afterSeqRaw = url.searchParams.get("afterSeq");
-        const afterSeq = afterSeqRaw ? Number.parseInt(afterSeqRaw, 10) || 0 : 0;
-        const stream = subscribe(params.streamId, afterSeq);
-
-        return new Response(stream, {
+      GET: async (_ctx, _req: Request, params: Record<string, string>) => {
+        const result = subscribe(params.messageId);
+        if (!result) {
+          // Brak in-memory streamu: generacja zakończona poza grace window
+          // ALBO proces zrestartował się mid-stream. Klient rozróżnia
+          // sytuacje po `isGenerating` w DB: jeśli row ma `isGenerating=false`
+          // → użyj final `blocks`, jeśli `true` → "Generation interrupted"
+          // + retry.
+          return new Response(
+            JSON.stringify({
+              error: "stream_not_found",
+              messageId: params.messageId,
+            }),
+            {
+              status: 410,
+              headers: { "Content-Type": "application/json" },
+            },
+          );
+        }
+        return new Response(result.stream, {
           headers: {
             "Content-Type": "text/event-stream",
             "Cache-Control": "no-cache",

package/src/streaming/stream-registry.ts

@@ -1,23 +1,46 @@
-import type { ChatStreamEvent } from "@arcote.tech/arc-ai";
+import type {
+  AssistantContentBlock,
+  ChatStreamEvent,
+  FinishReason,
+  TokenUsage,
+  ToolCall,
+  ToolResult,
+} from "@arcote.tech/arc-ai";
 
-// ─── ChatStreamManager — per session SSE registry with replay buffer ───
+// ─── Per-message in-memory stream registry ──────────────────────────
 //
-// Per-session state:
-//  - `streams[sessionId]` — live controllers currently subscribed
-//  - `buffers[sessionId]` — every event broadcast since the session started,
-//    so a late subscriber (e.g. after a page refresh mid-generation) gets
-//    the full prefix replayed before going live
-//  - `keepAliveIntervals[sessionId]` — heartbeat ping interval
-
-const streams = new Map<string, Set<ReadableStreamDefaultController<Uint8Array>>>();
-const buffers = new Map<string, ChatStreamEvent[]>();
-const keepAliveIntervals = new Map<string, ReturnType<typeof setInterval>>();
-const encoder = new TextEncoder();
+// Live state of an in-progress assistant message lives ONLY here, never in
+// DB. `currentBlocks` accumulates text + tool_call blocks as the LLM streams.
+// Each `publish()` mutates it AND broadcasts the SSE event to subscribers.
+// A new subscriber receives `{ type: "init", currentBlocks }` as the first
+// event — that's the snapshot of the in-memory state — and then live events.
+//
+// After `finalize()` the stream is kept for FINALIZE_GRACE_MS so a client
+// subscribing in that small window gets a `done` event instead of 410. After
+// grace, the stream is dropped — the client falls back to the final `blocks`
+// from DB via the normal aggregate query.
+//
+// Server restart / process crash → all in-memory state is lost. A client
+// observing `isGenerating: true` in DB but getting 410 from `subscribe()`
+// renders an "interrupted" state with a retry button.
+
+interface MessageStream {
+  messageId: string;
+  currentBlocks: AssistantContentBlock[];
+  toolCallsById: Map<
+    string,
+    Extract<AssistantContentBlock, { type: "tool_call" }>
+  >;
+  subscribers: Set<ReadableStreamDefaultController<Uint8Array>>;
+  keepAliveInterval?: ReturnType<typeof setInterval>;
+  finalized: boolean;
+  finalEvent?: ChatStreamEvent;
+}
 
-/** Hard cap on per-session buffer size. Each typical generation produces a
- *  few hundred chunks; 5000 is generous but bounds memory if a stream
- *  somehow runs without `endStream`. */
-const MAX_BUFFER = 5000;
+const streams = new Map<string, MessageStream>();
+const FINALIZE_GRACE_MS = 5_000;
+const KEEPALIVE_INTERVAL_MS = 5_000;
+const encoder = new TextEncoder();
 
 function encode(event: ChatStreamEvent): Uint8Array {
   return encoder.encode(`data: ${JSON.stringify(event)}\n\n`);
@@ -27,126 +50,237 @@ function encodePing(): Uint8Array {
   return encoder.encode(`: ping\n\n`);
 }
 
-export function broadcast(sessionId: string, event: ChatStreamEvent): void {
-  // Append to the replay buffer first — even if no client is currently
-  // subscribed (initial connect race) the event survives for replay.
-  let buf = buffers.get(sessionId);
-  if (!buf) {
-    buf = [];
-    buffers.set(sessionId, buf);
+function ensureKeepAlive(s: MessageStream): void {
+  if (s.keepAliveInterval) return;
+  s.keepAliveInterval = setInterval(() => {
+    if (s.subscribers.size === 0) return;
+    const ping = encodePing();
+    for (const ctrl of s.subscribers) {
+      try {
+        ctrl.enqueue(ping);
+      } catch {
+        s.subscribers.delete(ctrl);
+      }
+    }
+  }, KEEPALIVE_INTERVAL_MS);
+}
+
+/**
+ * Open an in-memory stream for `messageId`. Idempotent — safe to call when
+ * the assistant row was pre-created by a mutation (`preCreatedAssistantMessageId`).
+ *
+ * Listener calls this before the first `publish()`. From this point on
+ * `subscribe(messageId)` returns a live stream; before it, `subscribe()`
+ * returns null.
+ */
+export function startStream(messageId: string): void {
+  if (streams.has(messageId)) return;
+  streams.set(messageId, {
+    messageId,
+    currentBlocks: [],
+    toolCallsById: new Map(),
+    subscribers: new Set(),
+    finalized: false,
+  });
+}
+
+/**
+ * Event shape that the listener pushes into the stream. `messageId` is
+ * injected by the registry — caller doesn't pass it.
+ */
+export type PublishableEvent =
+  | { type: "text_delta"; textDelta: string }
+  | { type: "tool_call_pending"; toolCallId: string; toolCallName?: string }
+  | {
+      type: "tool_call_arguments_delta";
+      toolCallId: string;
+      argumentsDelta: string;
+    }
+  | {
+      type: "tool_call_arguments_complete";
+      toolCallId: string;
+      toolCallName?: string;
+      arguments: Record<string, unknown>;
+      executionCount?: number;
+    }
+  | {
+      type: "tool_call_executed";
+      toolCallId: string;
+      toolCallName: string;
+      toolResult: ToolResult;
+      executionCount?: number;
+    }
+  | {
+      type: "interactive_tool_request";
+      toolCalls: ToolCall[];
+      executionCount?: number;
+    }
+  | { type: "usage_update"; usage: TokenUsage }
+  | { type: "error"; error: string; executionCount?: number };
+
+/**
+ * Mutate the in-memory `currentBlocks` for the relevant event types, then
+ * broadcast the event to all active subscribers as SSE.
+ *
+ * No-op if the stream has been finalized or never started (race with a
+ * client disconnect / listener teardown).
+ */
+export function publish(messageId: string, event: PublishableEvent): void {
+  const s = streams.get(messageId);
+  if (!s || s.finalized) return;
+
+  switch (event.type) {
+    case "text_delta": {
+      const last = s.currentBlocks[s.currentBlocks.length - 1];
+      if (last && last.type === "text") {
+        last.text += event.textDelta;
+      } else {
+        s.currentBlocks.push({ type: "text", text: event.textDelta });
+      }
+      break;
+    }
+    case "tool_call_pending": {
+      const block: Extract<AssistantContentBlock, { type: "tool_call" }> = {
+        type: "tool_call",
+        id: event.toolCallId,
+        name: event.toolCallName ?? "",
+        arguments: {},
+      };
+      s.currentBlocks.push(block);
+      s.toolCallsById.set(event.toolCallId, block);
+      break;
+    }
+    case "tool_call_arguments_complete": {
+      const block = s.toolCallsById.get(event.toolCallId);
+      if (block) {
+        block.arguments = event.arguments;
+        if (event.toolCallName) block.name = event.toolCallName;
+      }
+      break;
+    }
+    // tool_call_arguments_delta, tool_call_executed, interactive_tool_request,
+    // usage_update, error — broadcast-only. They don't mutate `currentBlocks`:
+    //  - args_delta is incremental JSON the listener replaces via args_complete
+    //  - executed / interactive_tool_request relate to tool execution flow
+    //    (tool_result rows exist separately in DB)
+    //  - usage_update / error are transient signals to the UI
   }
-  if (buf.length < MAX_BUFFER) buf.push(event);
 
-  const controllers = streams.get(sessionId);
-  if (!controllers || controllers.size === 0) return;
-  const data = encode(event);
-  for (const controller of controllers) {
+  const payload = { ...event, messageId } as ChatStreamEvent;
+  const data = encode(payload);
+  for (const ctrl of s.subscribers) {
     try {
-      controller.enqueue(data);
+      ctrl.enqueue(data);
     } catch {
-      controllers.delete(controller);
+      s.subscribers.delete(ctrl);
     }
   }
 }
 
-export function subscribe(
-  sessionId: string,
-  afterSeq = 0,
-): ReadableStream<Uint8Array> {
-  return new ReadableStream<Uint8Array>({
-    async start(controller) {
-      // Replay buffered events with seq > afterSeq. Klient po reload czyta
-      // `partialLastSeq` z DB i przekazuje go jako afterSeq — buffer pomija
-      // już-zaaplikowane chunki (eliminuje duplikację typu "Dobrze — Dobrze —").
-      //
-      // Co 10 eventów yield (`setTimeout(16)`) — bez tego cały bufor leci
-      // w jednym chunku TCP, klient nie ma szansy zrobić reader.read() +
-      // render między burstami → React batchuje setTimeline w jeden render,
-      // streaming niewidoczny. 16ms ≈ rAF tick: 100 eventów = 160ms widocznego
-      // streamu zamiast 0ms burst.
-      const buf = buffers.get(sessionId);
-      if (buf) {
-        let count = 0;
-        for (const e of buf) {
-          if (e.seq <= afterSeq) continue;
-          try {
-            controller.enqueue(encode(e));
-          } catch {
-            return;
-          }
-          count++;
-          if (count % 10 === 0) {
-            await new Promise<void>((r) => setTimeout(r, 16));
-          }
-        }
-      }
+/**
+ * Open an SSE stream for a subscriber. Returns null if the stream has been
+ * dropped (server restart / past the FINALIZE_GRACE_MS window). The route
+ * handler maps null → HTTP 410, which signals the UI to show "Generation
+ * interrupted" + Retry.
+ *
+ * The stream starts with an `init` event carrying the current in-memory
+ * snapshot of `currentBlocks` so the subscriber sees what's already been
+ * generated. If the stream is already finalized within grace, the controller
+ * receives `init` then `done` then closes.
+ */
+export function subscribe(messageId: string): {
+  stream: ReadableStream<Uint8Array>;
+  currentBlocks: AssistantContentBlock[];
+} | null {
+  const s = streams.get(messageId);
+  if (!s) return null;
 
-      let set = streams.get(sessionId);
-      if (!set) {
-        set = new Set();
-        streams.set(sessionId, set);
-      }
-      set.add(controller);
-
-      // Start keep-alive if not running
-      if (!keepAliveIntervals.has(sessionId)) {
-        const interval = setInterval(() => {
-          const s = streams.get(sessionId);
-          if (s && s.size > 0) {
-            const ping = encodePing();
-            for (const c of s) {
-              try { c.enqueue(ping); } catch { s.delete(c); }
-            }
-          } else if (!buffers.has(sessionId)) {
-            // Stream truly inactive: no live clients AND no buffer. Stop
-            // pinging. We never proactively drop the buffer here — that
-            // happens in `endStream` so a late re-subscribe still gets the
-            // full replay.
-            cleanup(sessionId);
-          }
-        }, 5000);
-        keepAliveIntervals.set(sessionId, interval);
+  // Deep snapshot — the subscriber sees a stable `init` payload even if the
+  // listener mutates `currentBlocks` on a parallel publish.
+  const snapshot: AssistantContentBlock[] = JSON.parse(
+    JSON.stringify(s.currentBlocks),
+  );
+
+  const stream = new ReadableStream<Uint8Array>({
+    start(controller) {
+      controller.enqueue(
+        encode({
+          type: "init",
+          messageId,
+          currentBlocks: snapshot,
+        }),
+      );
+      if (s.finalized) {
+        if (s.finalEvent) controller.enqueue(encode(s.finalEvent));
+        try {
+          controller.close();
+        } catch {}
+        return;
       }
+      s.subscribers.add(controller);
+      ensureKeepAlive(s);
     },
     cancel() {
-      // One client disconnected — don't tear down session state. The buffer
-      // and other subscribers (if any) remain.
+      // Subscriber disconnected. Garbage-collection of dead controllers
+      // happens in `publish()` via the enqueue-catch path.
     },
   });
+
+  return { stream, currentBlocks: snapshot };
 }
 
-/** Called by the AI generation listener when a turn finishes (success or
- *  error). Closes all live SSE streams and drops the replay buffer. After
- *  this, a fresh `subscribe(sessionId)` returns an empty stream — the
- *  client should fall back to reading the final `blocks` from DB. */
-export function endStream(sessionId: string): void {
-  const controllers = streams.get(sessionId);
-  if (controllers) {
-    // seq dla `done` — przerzucamy ostatni z bufora lub 0 gdy pusto.
-    const buf = buffers.get(sessionId);
-    const lastSeq = buf && buf.length > 0 ? buf[buf.length - 1].seq : 0;
-    const done = encode({ type: "done", sessionId, seq: lastSeq + 1 } as ChatStreamEvent);
-    for (const controller of controllers) {
-      try {
-        controller.enqueue(done);
-        controller.close();
-      } catch {}
-    }
+/**
+ * Called by the listener after `completeAssistantTurn` returns. Broadcasts
+ * `done` to all live subscribers, closes the controllers, then drops the
+ * stream from the registry after FINALIZE_GRACE_MS.
+ */
+export function finalize(
+  messageId: string,
+  finalDetails?: {
+    usage?: TokenUsage;
+    finishReason?: FinishReason;
+    error?: string;
+    executionCount?: number;
+  },
+): void {
+  const s = streams.get(messageId);
+  if (!s || s.finalized) return;
+  s.finalized = true;
+
+  const done: ChatStreamEvent = {
+    type: "done",
+    messageId,
+    ...finalDetails,
+  };
+  s.finalEvent = done;
+
+  const data = encode(done);
+  for (const ctrl of s.subscribers) {
+    try {
+      ctrl.enqueue(data);
+      ctrl.close();
+    } catch {}
+  }
+  s.subscribers.clear();
+  if (s.keepAliveInterval) {
+    clearInterval(s.keepAliveInterval);
+    s.keepAliveInterval = undefined;
   }
-  cleanup(sessionId);
+
+  setTimeout(() => {
+    streams.delete(messageId);
+  }, FINALIZE_GRACE_MS);
 }
 
-export function hasActiveStream(sessionId: string): boolean {
-  const s = streams.get(sessionId);
-  return !!s && s.size > 0;
+export function isActive(messageId: string): boolean {
+  const s = streams.get(messageId);
+  return !!s && !s.finalized;
 }
 
-function cleanup(sessionId: string): void {
-  const interval = keepAliveIntervals.get(sessionId);
-  if (interval) {
-    clearInterval(interval);
-    keepAliveIntervals.delete(sessionId);
-  }
-  streams.delete(sessionId);
-  buffers.delete(sessionId);
+/** Test/debug — readonly snapshot of `currentBlocks` for a stream. */
+export function getCurrentBlocks(
+  messageId: string,
+): AssistantContentBlock[] | null {
+  const s = streams.get(messageId);
+  return s ? s.currentBlocks : null;
 }

package/src/tools/ask-questions.tsx

@@ -40,12 +40,15 @@ function AskQuestionsView({
   const { answerBelowLabel } = useChatLabels();
 
   useEffect(() => {
-    if (calling) {
+    // `params.questions` może być undefined podczas streamingu tool args —
+    // partial JSON nie ma jeszcze tablicy. Czekamy aż streaming skończy
+    // wypełniać argumenty.
+    if (calling && Array.isArray(params.questions)) {
       const questions: Question[] = params.questions.map((q) => ({
         id: q.id,
         label: q.label,
         description: q.description,
-        options: [...q.options],
+        options: [...(q.options ?? [])],
       }));
 
       registerInputOverride(