返回 JoyAI-Echo
useNanobotStream.ts
根目录 / echo_longvideo / Director_Agent / webui / src / hooks / useNanobotStream.ts
1 import { useCallback, useEffect, useRef, useState } from "react";
2
3 import { splitMediaByKind, wireMediaRefs } from "@/lib/media";
4 import {
5 stripAskUserFollowupBlurbs,
6 wireQuestionCards,
7 } from "@/lib/questions";
8 import { useClient } from "@/providers/ClientProvider";
9 import type { StreamError } from "@/lib/nanobot-client";
10 import { randomId } from "@/lib/utils";
11 import type {
12 InboundEvent,
13 OutboundMedia,
14 UIImage,
15 UIMessage,
16 } from "@/lib/types";
17
18 interface StreamBuffer {
19 /** ID of the assistant message currently receiving deltas. */
20 messageId: string;
21 /** Sequence of deltas accumulated in order. */
22 parts: string[];
23 }
24
25 /** Locate the empty top-of-turn wait bubble, if any. */
26 function findTurnAnchorId(list: UIMessage[]): string | null {
27 const anchor = list.find(
28 (m) =>
29 m.turnWaiting &&
30 m.isStreaming &&
31 !m.content.trim() &&
32 !m.questions?.length,
33 );
34 return anchor?.id ?? null;
35 }
36
37 /**
38 * Subscribe to a chat by ID. Returns the in-memory message list for the chat,
39 * a streaming flag, and a ``send`` function. Initial history must be seeded
40 * separately (e.g. via ``fetchSessionMessages``) since the server only replays
41 * live events.
42 */
43 /** Payload passed to ``send`` when the user attaches one or more images.
44 *
45 * ``media`` is handed to the wire client verbatim; ``preview`` powers the
46 * optimistic user bubble (blob URLs so the preview appears before the server
47 * acks the frame). Keeping the two separate lets the bubble re-use the local
48 * blob URL even after the server persists the file under a different name. */
49 export interface SendImage {
50 /** Omit for a UI-only preview whose payload is carried in message extras. */
51 media?: OutboundMedia;
52 preview: UIImage;
53 }
54
55 export function useNanobotStream(
56 chatId: string | null,
57 initialMessages: UIMessage[] = [],
58 options?: { onReplyEnd?: () => void },
59 ): {
60 messages: UIMessage[];
61 isStreaming: boolean;
62 /** ``true`` after ``stream_end`` with ``resuming: false`` (or a non-resuming reply). */
63 turnComplete: boolean;
64 send: (
65 content: string,
66 images?: SendImage[],
67 extras?: {
68 autoGenerate?: boolean;
69 duration_sec?: number;
70 reference_image_url?: string;
71 reference_image_name?: string;
72 reference_image_width?: number;
73 reference_image_height?: number;
74 },
75 ) => void;
76 answerQuestion: (messageId: string, cardId: string, value: string) => void;
77 setMessages: React.Dispatch<React.SetStateAction<UIMessage[]>>;
78 /** Latest transport-level fault raised since the last ``dismissStreamError``.
79 * ``null`` when there is nothing to show. */
80 streamError: StreamError | null;
81 /** Clear the current ``streamError`` (e.g. after the user dismisses the
82 * notification or starts a fresh action). */
83 dismissStreamError: () => void;
84 } {
85 const { client } = useClient();
86 const messagesRef = useRef<UIMessage[]>(initialMessages);
87 const [messages, setMessages] = useState<UIMessage[]>(initialMessages);
88 messagesRef.current = messages;
89 const [isStreaming, setIsStreaming] = useState(false);
90 const [turnComplete, setTurnComplete] = useState(true);
91 const turnCompleteRef = useRef(true);
92 turnCompleteRef.current = turnComplete;
93 /** Set when an intermediate ``stream_end`` arrives with ``resuming: true``. */
94 const sawResumingStreamEndRef = useRef(false);
95 /** Empty ``turnWaiting`` bubble pinned at the top of the in-flight agent turn. */
96 const turnAnchorIdRef = useRef<string | null>(null);
97 const [streamError, setStreamError] = useState<StreamError | null>(null);
98 const buffer = useRef<StreamBuffer | null>(null);
99 const onReplyEndRef = useRef(options?.onReplyEnd);
100 onReplyEndRef.current = options?.onReplyEnd;
101
102 useEffect(() => {
103 return client.onError((err) => setStreamError(err));
104 }, [client]);
105
106 const dismissStreamError = useCallback(() => setStreamError(null), []);
107
108 useEffect(() => {
109 if (!chatId) {
110 buffer.current = null;
111 setIsStreaming(false);
112 setTurnComplete(true);
113 sawResumingStreamEndRef.current = false;
114 turnAnchorIdRef.current = null;
115 return;
116 }
117
118 // Seed the target chat before onChat: offline replay is scheduled in a
119 // microtask and must not read the previous session's messagesRef.
120 // Seed messagesRef before subscribe; defer React state until after offline
121 // replay microtasks so a switch-back does not wipe replayed deltas.
122 messagesRef.current = initialMessages;
123 buffer.current = null;
124 const restoredAnchorId = findTurnAnchorId(initialMessages);
125 turnAnchorIdRef.current = restoredAnchorId;
126 const openTurn = restoredAnchorId !== null;
127 setIsStreaming(openTurn);
128 setTurnComplete(!openTurn);
129 sawResumingStreamEndRef.current = false;
130 setStreamError(null);
131
132 const stripTurnAnchor = (list: UIMessage[]): UIMessage[] => {
133 turnAnchorIdRef.current = null;
134 return list.filter(
135 (m) => !(m.turnWaiting && !m.content.trim() && !m.questions?.length),
136 );
137 };
138
139 const finalizeTurnEnd = (list: UIMessage[]): UIMessage[] => {
140 return stripAskUserFollowupBlurbs(
141 stripTurnAnchor(
142 list
143 .map((m) => (m.isStreaming ? { ...m, isStreaming: false } : m))
144 .filter(
145 (m) =>
146 !(
147 m.role === "assistant" &&
148 !m.content.trim() &&
149 !m.isStreaming &&
150 !m.questions?.length
151 ),
152 ),
153 ),
154 );
155 };
156
157 const commitMessages = (
158 updater: UIMessage[] | ((prev: UIMessage[]) => UIMessage[]),
159 ) => {
160 const next =
161 typeof updater === "function" ? updater(messagesRef.current) : updater;
162 messagesRef.current = next;
163 setMessages(next);
164 };
165
166 /** Bind or create the single top-of-turn wait bubble (TypingDots). */
167 const bindWaitBuffer = (): string => {
168 if (buffer.current) return buffer.current.messageId;
169 const anchorId = turnAnchorIdRef.current;
170 if (anchorId) {
171 const anchor = messagesRef.current.find((m) => m.id === anchorId);
172 if (anchor?.isStreaming && !anchor.content.trim()) {
173 buffer.current = { messageId: anchor.id, parts: [] };
174 setIsStreaming(true);
175 return anchor.id;
176 }
177 }
178 const last = messagesRef.current.at(-1);
179 if (
180 last?.role === "assistant" &&
181 last.isStreaming &&
182 !last.content.trim()
183 ) {
184 if (last.turnWaiting) {
185 turnAnchorIdRef.current = last.id;
186 }
187 buffer.current = { messageId: last.id, parts: [] };
188 setIsStreaming(true);
189 return last.id;
190 }
191 const id = randomId();
192 turnAnchorIdRef.current = id;
193 const placeholder: UIMessage = {
194 id,
195 role: "assistant",
196 content: "",
197 isStreaming: true,
198 turnWaiting: true,
199 createdAt: Date.now(),
200 };
201 buffer.current = { messageId: id, parts: [] };
202 setIsStreaming(true);
203 const next = [...messagesRef.current, placeholder];
204 messagesRef.current = next;
205 commitMessages(next);
206 return id;
207 };
208
209 /** Tail segment for streamed deltas — never reuses the top turn anchor. */
210 const beginDeltaStream = (): string => {
211 if (buffer.current) return buffer.current.messageId;
212 const anchorId = turnAnchorIdRef.current;
213 const last = messagesRef.current.at(-1);
214 if (
215 last?.role === "assistant" &&
216 last.isStreaming &&
217 !last.content.trim() &&
218 last.id !== anchorId &&
219 !last.turnWaiting
220 ) {
221 buffer.current = { messageId: last.id, parts: [] };
222 setIsStreaming(true);
223 return last.id;
224 }
225 const id = randomId();
226 const placeholder: UIMessage = {
227 id,
228 role: "assistant",
229 content: "",
230 isStreaming: true,
231 createdAt: Date.now(),
232 };
233 buffer.current = { messageId: id, parts: [] };
234 setIsStreaming(true);
235 const next = [...messagesRef.current, placeholder];
236 messagesRef.current = next;
237 commitMessages(next);
238 return id;
239 };
240
241 const beginStreamingTurn = bindWaitBuffer;
242
243 const handle = (ev: InboundEvent) => {
244 if (ev.event === "delta") {
245 const anchorId = turnAnchorIdRef.current;
246 const bufferTargetsAnchor = (): boolean => {
247 const id = buffer.current?.messageId;
248 if (!id) return false;
249 if (anchorId && id === anchorId) return true;
250 return (
251 messagesRef.current.find((m) => m.id === id)?.turnWaiting === true
252 );
253 };
254 if (!buffer.current || bufferTargetsAnchor()) {
255 if (buffer.current && bufferTargetsAnchor()) {
256 buffer.current = null;
257 }
258 beginDeltaStream();
259 }
260 buffer.current!.parts.push(ev.text);
261 const combined = buffer.current!.parts.join("");
262 const targetId = buffer.current!.messageId;
263 commitMessages(
264 messagesRef.current.map((m) =>
265 m.id === targetId ? { ...m, content: combined } : m,
266 ),
267 );
268 return;
269 }
270
271 if (ev.event === "stream_end") {
272 if (ev.resuming === true) {
273 sawResumingStreamEndRef.current = true;
274 setTurnComplete(false);
275 } else if (ev.resuming === false) {
276 sawResumingStreamEndRef.current = false;
277 setTurnComplete(true);
278 }
279
280 if (!buffer.current) {
281 if (ev.resuming === false) {
282 setIsStreaming(false);
283 commitMessages((prev) => finalizeTurnEnd(prev));
284 onReplyEndRef.current?.();
285 return;
286 }
287 // Director may lead with stream_end on replay; keep the wait UI alive.
288 beginStreamingTurn();
289 return;
290 }
291 const finalId = buffer.current.messageId;
292 const combined = buffer.current.parts.join("");
293 const hadContent = combined.trim().length > 0;
294 if (!hadContent) {
295 if (ev.resuming === false) {
296 buffer.current = null;
297 setIsStreaming(false);
298 commitMessages((prev) => finalizeTurnEnd(prev));
299 onReplyEndRef.current?.();
300 }
301 return;
302 }
303
304 if (ev.resuming === true) {
305 buffer.current = null;
306 setIsStreaming(true);
307 commitMessages((prev) =>
308 prev.map((m) =>
309 m.id === finalId
310 ? { ...m, content: combined, isStreaming: false }
311 : m,
312 ),
313 );
314 return;
315 }
316
317 buffer.current = null;
318 setIsStreaming(false);
319 if (ev.resuming === undefined) {
320 sawResumingStreamEndRef.current = false;
321 setTurnComplete(true);
322 }
323 commitMessages((prev) => {
324 let next = prev.map((m) =>
325 m.id === finalId
326 ? { ...m, content: combined, isStreaming: false }
327 : m,
328 );
329 if (ev.resuming === false) {
330 next = stripTurnAnchor(next);
331 }
332 return next;
333 });
334 onReplyEndRef.current?.();
335 return;
336 }
337
338 if (ev.event === "message") {
339 // Tool hints / progress are not surfaced in the thread UI; the send
340 // placeholder progress bar covers the wait instead.
341 if (ev.kind === "tool_hint" || ev.kind === "progress") {
342 return;
343 }
344
345 // A complete (non-streamed) assistant message. If a stream was in
346 // flight, drop the placeholder so we don't render the text twice.
347 const midTurn = sawResumingStreamEndRef.current;
348 const activeId = buffer.current?.messageId;
349 const anchorId = turnAnchorIdRef.current;
350 buffer.current = null;
351 if (!midTurn) {
352 setIsStreaming(false);
353 setTurnComplete(true);
354 turnAnchorIdRef.current = null;
355 } else {
356 setIsStreaming(true);
357 }
358 commitMessages((prev) => {
359 const shouldRemoveActive =
360 !!activeId && (!midTurn || activeId !== anchorId);
361 const base = shouldRemoveActive
362 ? prev.filter((m) => m.id !== activeId)
363 : prev;
364 const media = splitMediaByKind(
365 wireMediaRefs(ev.media_urls, ev.media),
366 );
367 const questions =
368 ev.questions && ev.questions.length > 0
369 ? wireQuestionCards(ev.questions)
370 : undefined;
371
372 const attachCards = (list: UIMessage[]): UIMessage[] => {
373 if (!questions?.length) {
374 const assistantMsg: UIMessage = {
375 id: randomId(),
376 role: "assistant",
377 content: ev.text,
378 createdAt: Date.now(),
379 ...(media.images ? { images: media.images } : {}),
380 ...(media.videos ? { videos: media.videos } : {}),
381 };
382 return [...list, assistantMsg];
383 }
384
385 const cardMessage = (): UIMessage => ({
386 id: randomId(),
387 role: "assistant",
388 content: "",
389 createdAt: Date.now(),
390 ...(media.images ? { images: media.images } : {}),
391 ...(media.videos ? { videos: media.videos } : {}),
392 questions,
393 ...(ev.question_batch_id
394 ? { questionBatchId: ev.question_batch_id }
395 : {}),
396 });
397
398 const last = list.at(-1);
399 if (last?.role === "assistant") {
400 if (
401 last.questions?.length &&
402 ev.question_batch_id &&
403 last.questionBatchId === ev.question_batch_id
404 ) {
405 return list;
406 }
407 if (
408 last.questions?.length &&
409 ev.question_batch_id &&
410 last.questionBatchId !== ev.question_batch_id
411 ) {
412 return stripAskUserFollowupBlurbs([...list, cardMessage()]);
413 }
414 if (
415 last.turnWaiting ||
416 (last.isStreaming &&
417 !last.content.trim() &&
418 !last.questions?.length)
419 ) {
420 return stripAskUserFollowupBlurbs([
421 ...list.slice(0, -1),
422 {
423 ...last,
424 content: "",
425 isStreaming: false,
426 turnWaiting: undefined,
427 questions,
428 ...(ev.question_batch_id
429 ? { questionBatchId: ev.question_batch_id }
430 : {}),
431 },
432 ]);
433 }
434 return stripAskUserFollowupBlurbs([
435 ...list.slice(0, -1),
436 {
437 ...last,
438 content: "",
439 isStreaming: false,
440 turnWaiting: undefined,
441 questions,
442 ...(ev.question_batch_id
443 ? { questionBatchId: ev.question_batch_id }
444 : {}),
445 },
446 ]);
447 }
448
449 return stripAskUserFollowupBlurbs([...list, cardMessage()]);
450 };
451
452 const resolvedAnchorId = anchorId ?? findTurnAnchorId(base);
453 if (midTurn && resolvedAnchorId) {
454 turnAnchorIdRef.current = resolvedAnchorId;
455 const anchorIdx = base.findIndex((m) => m.id === resolvedAnchorId);
456 const next = [...base];
457 if (anchorIdx >= 0) {
458 const prefix = next.slice(0, anchorIdx + 1);
459 const suffix = next.slice(anchorIdx + 1);
460 return [...prefix, ...attachCards(suffix)];
461 }
462 return attachCards(next);
463 }
464
465 return attachCards(stripTurnAnchor(base));
466 });
467 if (!midTurn) {
468 onReplyEndRef.current?.();
469 }
470 return;
471 }
472
473 if (ev.event === "question_answer_ok") {
474 commitMessages((prev) =>
475 prev.map((m) => {
476 if (m.questionBatchId !== ev.question_batch_id || !m.questions) {
477 return m;
478 }
479 return {
480 ...m,
481 questions: m.questions.map((q) =>
482 q.id === ev.card_id ? { ...q, answered: ev.value } : q,
483 ),
484 };
485 }),
486 );
487 return;
488 }
489 // ``attached`` / ``error`` frames aren't actionable here; the client
490 // shell handles them separately.
491 };
492
493 const pendingReplay =
494 typeof client.hasOfflineBuffer === "function"
495 ? client.hasOfflineBuffer(chatId)
496 : false;
497 const unsub = client.onChat(chatId, handle);
498 if (pendingReplay && typeof client.replayOfflineNow === "function") {
499 client.replayOfflineNow(chatId);
500 }
501 setMessages(messagesRef.current);
502 return () => {
503 unsub();
504 buffer.current = null;
505 };
506 // Only re-subscribe on chat switch; history HTTP refresh is merged via ThreadShell.
507 // eslint-disable-next-line react-hooks/exhaustive-deps
508 }, [chatId, client]);
509
510 const send = useCallback(
511 (
512 content: string,
513 images?: SendImage[],
514 extras?: {
515 autoGenerate?: boolean;
516 duration_sec?: number;
517 reference_image_url?: string;
518 reference_image_name?: string;
519 reference_image_width?: number;
520 reference_image_height?: number;
521 },
522 ) => {
523 if (!chatId) return;
524 const hasImages = !!images && images.length > 0;
525 // Text is optional when images are attached — the agent will still see
526 // the image blocks via ``media`` paths.
527 if (!hasImages && !content.trim()) return;
528
529 const pendingId = randomId();
530 turnAnchorIdRef.current = pendingId;
531 buffer.current = null;
532 sawResumingStreamEndRef.current = false;
533 setTurnComplete(false);
534 setIsStreaming(true);
535 const previews = hasImages ? images!.map((i) => i.preview) : undefined;
536 setMessages((prev) => [
537 ...prev,
538 {
539 id: randomId(),
540 role: "user",
541 content,
542 createdAt: Date.now(),
543 ...(previews ? { images: previews } : {}),
544 },
545 {
546 id: pendingId,
547 role: "assistant",
548 content: "",
549 isStreaming: true,
550 turnWaiting: true,
551 createdAt: Date.now(),
552 },
553 ]);
554 const wireMedia = hasImages
555 ? images!
556 .map((i) => i.media)
557 .filter((media): media is OutboundMedia => media !== undefined)
558 : undefined;
559 const mergedExtras = {
560 ...(extras ?? {}),
561 };
562 client.sendMessage(
563 chatId,
564 content,
565 wireMedia && wireMedia.length > 0 ? wireMedia : undefined,
566 Object.keys(mergedExtras).length > 0
567 ? (mergedExtras as {
568 temperature?: number;
569 top_p?: number;
570 top_k?: number;
571 autoGenerate?: boolean;
572 duration_sec?: number;
573 reference_image_url?: string;
574 reference_image_name?: string;
575 reference_image_width?: number;
576 reference_image_height?: number;
577 })
578 : undefined,
579 );
580 },
581 [chatId, client],
582 );
583
584 const answerQuestion = useCallback(
585 (messageId: string, cardId: string, value: string) => {
586 if (!chatId) return;
587 if (!turnCompleteRef.current) return;
588 const trimmed = value.trim();
589 if (!trimmed) return;
590
591 const target = messagesRef.current.find((m) => m.id === messageId);
592 const batchId = target?.questionBatchId;
593 const pendingId = randomId();
594 turnAnchorIdRef.current = pendingId;
595 buffer.current = null;
596 sawResumingStreamEndRef.current = false;
597 setTurnComplete(false);
598 setIsStreaming(true);
599
600 setMessages((prev) => {
601 const withAnswer = prev.map((m) => {
602 if (m.id !== messageId || !m.questions) return m;
603 return {
604 ...m,
605 questions: m.questions.map((q) =>
606 q.id === cardId ? { ...q, answered: trimmed } : q,
607 ),
608 };
609 });
610 return [
611 ...withAnswer,
612 {
613 id: randomId(),
614 role: "user",
615 content: trimmed,
616 createdAt: Date.now(),
617 },
618 {
619 id: pendingId,
620 role: "assistant",
621 content: "",
622 isStreaming: true,
623 turnWaiting: true,
624 createdAt: Date.now(),
625 },
626 ];
627 });
628
629 if (batchId) {
630 client.answerQuestion(chatId, batchId, cardId, trimmed);
631 }
632 client.sendMessage(chatId, trimmed);
633 },
634 [chatId, client],
635 );
636
637 return {
638 messages,
639 isStreaming,
640 turnComplete,
641 send,
642 answerQuestion,
643 setMessages,
644 streamError,
645 dismissStreamError,
646 };
647 }
648
648 lines TYPESCRIPT