返回 JoyAI-Echo
nanobot-client.ts
根目录 / echo_longvideo / Director_Agent / webui / src / lib / nanobot-client.ts
1 import type {
2 ConnectionStatus,
3 InboundEvent,
4 MemoryAssetUpload,
5 ShotMemoryAssetCreate,
6 MemorySlotReference,
7 Outbound,
8 OutboundMedia,
9 StoryProfile,
10 WorkplaceData,
11 } from "./types";
12
13 /** WebSocket readyState constants, referenced by value to stay portable
14 * across runtimes that don't expose a global ``WebSocket`` (tests, SSR). */
15 const WS_OPEN = 1;
16 const WS_CLOSING = 2;
17 /** Workplace save actions can carry large payloads; allow ample server round-trip. */
18 const WORKPLACE_ACTION_TIMEOUT_MS = 60_000;
19 /** Cap offline delta frames per chat (Director can stream heavily). */
20 const MAX_OFFLINE_DELTAS = 500;
21
22 type Unsubscribe = () => void;
23 type EventHandler = (ev: InboundEvent) => void;
24 type StatusHandler = (status: ConnectionStatus) => void;
25 type PeHandler = (chatId: string, active: string) => void;
26
27 /** Structured connection-level errors surfaced to the UI.
28 *
29 * These are *not* InboundEvent errors from the server application layer —
30 * those arrive as ``{event: "error"}`` messages via ``onChat``. These are
31 * transport-level or protocol-level faults the UI should make visible so
32 * the user understands *why* their action failed (as opposed to silently
33 * reconnecting under the hood).
34 */
35 export type StreamError =
36 /** Server rejected the inbound frame as too large (WS close code 1009).
37 * Typically means the user attached images whose base64 size exceeded
38 * ``maxMessageBytes`` on the server. */
39 | { kind: "message_too_big" };
40
41 type ErrorHandler = (error: StreamError) => void;
42
43 interface PendingNewChat {
44 resolve: (chatId: string) => void;
45 reject: (err: Error) => void;
46 timer: ReturnType<typeof setTimeout>;
47 }
48
49 type PendingWorkplaceAction = {
50 resolve: (value: { work_id: string; workplace: WorkplaceData }) => void;
51 reject: (err: Error) => void;
52 timer: ReturnType<typeof setTimeout>;
53 };
54
55 export interface NanobotClientOptions {
56 url: string;
57 reconnect?: boolean;
58 /** Called when a connection drops so the app can refresh its token. */
59 onReauth?: () => Promise<string | null>;
60 /** Inject a custom WebSocket factory (used by unit tests). */
61 socketFactory?: (url: string) => WebSocket;
62 /** Delay-cap for reconnect backoff (ms). */
63 maxBackoffMs?: number;
64 }
65
66 /**
67 * Singleton WebSocket client that multiplexes chat streams.
68 *
69 * One socket carries many chat_ids: the server tags every outbound event with
70 * ``chat_id``, and this class fans those events out to handlers registered
71 * per chat. Reconnects are transparent and re-attach every known chat_id.
72 */
73 export class NanobotClient {
74 private socket: WebSocket | null = null;
75 private statusHandlers = new Set<StatusHandler>();
76 private errorHandlers = new Set<ErrorHandler>();
77 /** Server-global PE-set change subscribers (event: "pe_updated"). */
78 private peHandlers = new Set<PeHandler>();
79 // chat_id -> handlers listening on it
80 private chatHandlers = new Map<string, Set<EventHandler>>();
81 // chat_ids we've attached to since connect; re-attached after reconnects
82 private knownChats = new Set<string>();
83 private pendingNewChat: PendingNewChat | null = null;
84 /** request_id → Promise for workplace_save_* responses (workplace_action_ok/error). */
85 private pendingWorkplaceActions = new Map<string, PendingWorkplaceAction>();
86 // Frames queued while the socket is not yet OPEN
87 private sendQueue: Outbound[] = [];
88 private reconnectAttempts = 0;
89 private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
90 private readonly shouldReconnect: boolean;
91 private readonly maxBackoffMs: number;
92 private readonly socketFactory: (url: string) => WebSocket;
93 private currentUrl: string;
94 private status_: ConnectionStatus = "idle";
95 private readyChatId: string | null = null;
96 // Set by ``close()`` so the onclose handler knows the drop was intentional
97 // and must not schedule a reconnect or flip status back to "reconnecting".
98 private intentionallyClosed = false;
99 /** Events received while no UI handler is subscribed (e.g. user switched chats). */
100 private offlineBuffers = new Map<string, InboundEvent[]>();
101 /** Coalesce multiple onChat calls in one React commit into a single replay. */
102 private replayScheduled = new Set<string>();
103
104 constructor(private options: NanobotClientOptions) {
105 this.shouldReconnect = options.reconnect ?? true;
106 this.maxBackoffMs = options.maxBackoffMs ?? 15_000;
107 this.socketFactory =
108 options.socketFactory ?? ((url) => new WebSocket(url));
109 this.currentUrl = options.url;
110 }
111
112 get status(): ConnectionStatus {
113 return this.status_;
114 }
115
116 get defaultChatId(): string | null {
117 return this.readyChatId;
118 }
119
120 /** Swap the URL (e.g. after fetching a fresh token) then reconnect. */
121 updateUrl(url: string): void {
122 this.currentUrl = url;
123 }
124
125 onStatus(handler: StatusHandler): Unsubscribe {
126 this.statusHandlers.add(handler);
127 handler(this.status_);
128 return () => {
129 this.statusHandlers.delete(handler);
130 };
131 }
132
133 /** Subscribe to transport-level faults (see :type:`StreamError`). */
134 onError(handler: ErrorHandler): Unsubscribe {
135 this.errorHandlers.add(handler);
136 return () => {
137 this.errorHandlers.delete(handler);
138 };
139 }
140
141 /** Subscribe to per-chat PE-set changes (``pe_updated`` / ``attached.active_pe``). */
142 onPe(handler: PeHandler): Unsubscribe {
143 this.peHandlers.add(handler);
144 return () => {
145 this.peHandlers.delete(handler);
146 };
147 }
148
149 /** Ask the server to bind the PE set for a given chat's session. */
150 setPe(name: string, chatId: string): void {
151 this.queueSend({ type: "set_pe", name, chat_id: chatId });
152 }
153
154 /** Whether events are queued for subscribe-time replay (user was away). */
155 hasOfflineBuffer(chatId: string): boolean {
156 const buf = this.offlineBuffers.get(chatId);
157 return !!buf && buf.length > 0;
158 }
159
160 /** Replay buffered events synchronously (after ``onChat`` registers a handler). */
161 replayOfflineNow(chatId: string): void {
162 this.flushOfflineBuffer(chatId);
163 }
164
165 /** Subscribe to events for a given chat_id. Auto-attaches on the next open. */
166 onChat(
167 chatId: string,
168 handler: EventHandler,
169 ): Unsubscribe {
170 let handlers = this.chatHandlers.get(chatId);
171 if (!handlers) {
172 handlers = new Set();
173 this.chatHandlers.set(chatId, handlers);
174 }
175 handlers.add(handler);
176 this.attach(chatId);
177 const pending = this.offlineBuffers.get(chatId);
178 if (pending && pending.length > 0) {
179 this.scheduleOfflineReplay(chatId);
180 }
181 return () => {
182 const current = this.chatHandlers.get(chatId);
183 if (!current) return;
184 current.delete(handler);
185 if (current.size === 0) this.chatHandlers.delete(chatId);
186 };
187 }
188
189 connect(): void {
190 if (this.socket && this.socket.readyState < WS_CLOSING) return;
191 this.intentionallyClosed = false;
192 this.setStatus("connecting");
193 const sock = this.socketFactory(this.currentUrl);
194 this.socket = sock;
195 sock.onopen = () => this.handleOpen();
196 sock.onmessage = (ev) => this.handleMessage(ev);
197 sock.onerror = () => this.setStatus("error");
198 sock.onclose = (ev) => this.handleClose(ev);
199 }
200
201 close(): void {
202 this.intentionallyClosed = true;
203 if (this.reconnectTimer) {
204 clearTimeout(this.reconnectTimer);
205 this.reconnectTimer = null;
206 }
207 const sock = this.socket;
208 this.socket = null;
209 try {
210 sock?.close();
211 } catch {
212 // ignore
213 }
214 this.offlineBuffers.clear();
215 this.replayScheduled.clear();
216 this.setStatus("closed");
217 }
218
219 /** Ask the server to provision a new chat_id; resolves with the assigned id. */
220 newChat(
221 timeoutMs: number = 5_000,
222 options?: { autoGenerate?: boolean },
223 ): Promise<string> {
224 if (this.pendingNewChat) {
225 return Promise.reject(new Error("newChat already in flight"));
226 }
227 return new Promise<string>((resolve, reject) => {
228 const timer = setTimeout(() => {
229 this.pendingNewChat = null;
230 reject(new Error("newChat timed out"));
231 }, timeoutMs);
232 this.pendingNewChat = {
233 resolve,
234 reject,
235 timer,
236 };
237 this.queueSend({
238 type: "new_chat",
239 ...(options?.autoGenerate ? { autoGenerate: true } : {}),
240 });
241 });
242 }
243
244 attach(chatId: string): void {
245 this.knownChats.add(chatId);
246 if (this.socket?.readyState === WS_OPEN) {
247 this.queueSend(this.buildAttachFrame(chatId));
248 }
249 }
250
251 sendMessage(
252 chatId: string,
253 content: string,
254 media?: OutboundMedia[],
255 extras?: {
256 temperature?: number;
257 top_p?: number;
258 top_k?: number;
259 autoGenerate?: boolean;
260 duration_sec?: number;
261 reference_image_url?: string;
262 reference_image_name?: string;
263 reference_image_width?: number;
264 reference_image_height?: number;
265 },
266 ): void {
267 this.knownChats.add(chatId);
268 const base =
269 media && media.length > 0
270 ? { type: "message" as const, chat_id: chatId, content, media }
271 : { type: "message" as const, chat_id: chatId, content };
272 const frame: Outbound = extras ? { ...base, ...extras } : base;
273 this.queueSend(frame);
274 }
275
276 answerQuestion(
277 chatId: string,
278 questionBatchId: string,
279 cardId: string,
280 value: string,
281 ): void {
282 this.knownChats.add(chatId);
283 this.queueSend({
284 type: "answer_question",
285 chat_id: chatId,
286 question_batch_id: questionBatchId,
287 card_id: cardId,
288 value,
289 });
290 }
291
292 /** Persist story markdown over WS (avoids HTTP header size limits on large scripts). */
293 saveWorkplaceStory(
294 chatId: string,
295 storyMd: string,
296 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
297 const trimmed = storyMd.trim();
298 if (!trimmed) {
299 return Promise.reject(new Error("story_md cannot be empty"));
300 }
301 const { requestId, promise } = this.registerWorkplaceAction();
302 this.attach(chatId);
303 this.queueSend({
304 type: "workplace_save_story",
305 chat_id: chatId,
306 request_id: requestId,
307 story_md: trimmed,
308 });
309 return promise;
310 }
311
312 /** Persist story profile (shot plan) over WS. */
313 saveWorkplaceStoryProfile(
314 chatId: string,
315 storyProfile: StoryProfile,
316 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
317 const { requestId, promise } = this.registerWorkplaceAction();
318 this.attach(chatId);
319 this.queueSend({
320 type: "workplace_save_story_profile",
321 chat_id: chatId,
322 request_id: requestId,
323 story_profile: storyProfile,
324 });
325 return promise;
326 }
327
328 /** Persist a first-frame image over WS so its data URL isn't put in HTTP headers. */
329 saveWorkplaceReferenceImage(
330 chatId: string,
331 image: {
332 url: string;
333 name?: string;
334 width?: number;
335 height?: number;
336 },
337 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
338 const { requestId, promise } = this.registerWorkplaceAction();
339 this.attach(chatId);
340 this.queueSend({
341 type: "workplace_save_reference_image",
342 chat_id: chatId,
343 request_id: requestId,
344 image,
345 });
346 return promise;
347 }
348
349 /** Add or update a local Memory Workspace asset over WebSocket. */
350 saveWorkplaceMemoryAsset(
351 chatId: string,
352 asset: MemoryAssetUpload,
353 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
354 const { requestId, promise } = this.registerWorkplaceAction();
355 this.attach(chatId);
356 this.queueSend({
357 type: "workplace_save_memory_asset",
358 chat_id: chatId,
359 request_id: requestId,
360 asset,
361 });
362 return promise;
363 }
364
365 /** Extract a reusable frame/audio clip from one generated shot. */
366 createWorkplaceShotMemoryAsset(
367 chatId: string,
368 shotId: number,
369 asset: ShotMemoryAssetCreate,
370 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
371 const { requestId, promise } = this.registerWorkplaceAction();
372 this.attach(chatId);
373 this.queueSend({
374 type: "workplace_create_shot_memory_asset",
375 chat_id: chatId,
376 request_id: requestId,
377 shot_id: shotId,
378 asset,
379 });
380 return promise;
381 }
382
383 /** Remove a locally uploaded asset from the Memory Workspace. */
384 deleteWorkplaceMemoryAsset(
385 chatId: string,
386 assetId: string,
387 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
388 const { requestId, promise } = this.registerWorkplaceAction();
389 this.attach(chatId);
390 this.queueSend({
391 type: "workplace_delete_memory_asset",
392 chat_id: chatId,
393 request_id: requestId,
394 asset_id: assetId,
395 });
396 return promise;
397 }
398
399 /** Apply an ordered Memory Workspace selection to one shot. */
400 saveWorkplaceShotMemorySlots(
401 chatId: string,
402 shotId: number,
403 slots: MemorySlotReference[],
404 ): Promise<{ work_id: string; workplace: WorkplaceData }> {
405 const { requestId, promise } = this.registerWorkplaceAction();
406 this.attach(chatId);
407 this.queueSend({
408 type: "workplace_save_shot_memory_slots",
409 chat_id: chatId,
410 request_id: requestId,
411 shot_id: shotId,
412 slots,
413 });
414 return promise;
415 }
416
417 // -- internals ---------------------------------------------------------
418
419 private buildAttachFrame(chatId: string): Outbound {
420 return { type: "attach", chat_id: chatId };
421 }
422
423 private setStatus(status: ConnectionStatus): void {
424 if (this.status_ === status) return;
425 this.status_ = status;
426 for (const handler of this.statusHandlers) handler(status);
427 }
428
429 private handleOpen(): void {
430 this.setStatus("open");
431 this.reconnectAttempts = 0;
432 // Re-attach every known chat_id so deliveries continue routing after a drop.
433 for (const chatId of this.knownChats) {
434 this.rawSend(this.buildAttachFrame(chatId));
435 }
436 // Flush anything queued during reconnect.
437 const queued = this.sendQueue.splice(0);
438 for (const frame of queued) this.rawSend(frame);
439 }
440
441 private handleMessage(ev: MessageEvent): void {
442 let parsed: InboundEvent;
443 try {
444 parsed = JSON.parse(typeof ev.data === "string" ? ev.data : "") as InboundEvent;
445 } catch {
446 return;
447 }
448
449 if (parsed.event === "ready") {
450 this.readyChatId = parsed.chat_id;
451 this.knownChats.add(parsed.chat_id);
452 return;
453 }
454
455 if (parsed.event === "attached") {
456 this.knownChats.add(parsed.chat_id);
457 if (this.pendingNewChat) {
458 clearTimeout(this.pendingNewChat.timer);
459 this.pendingNewChat.resolve(parsed.chat_id);
460 this.pendingNewChat = null;
461 }
462 if (parsed.active_pe) {
463 for (const handler of this.peHandlers) {
464 handler(parsed.chat_id, parsed.active_pe);
465 }
466 }
467 this.dispatch(parsed.chat_id, parsed);
468 return;
469 }
470
471 if (parsed.event === "workplace_action_ok") {
472 const pending = this.pendingWorkplaceActions.get(parsed.request_id);
473 if (pending) {
474 clearTimeout(pending.timer);
475 this.pendingWorkplaceActions.delete(parsed.request_id);
476 pending.resolve({
477 work_id: parsed.work_id,
478 workplace: parsed.workplace,
479 });
480 }
481 return;
482 }
483
484 if (parsed.event === "workplace_action_error") {
485 const pending = this.pendingWorkplaceActions.get(parsed.request_id);
486 if (pending) {
487 clearTimeout(pending.timer);
488 this.pendingWorkplaceActions.delete(parsed.request_id);
489 pending.reject(new Error(parsed.detail || "workplace action failed"));
490 }
491 return;
492 }
493
494 if (parsed.event === "pe_updated") {
495 for (const handler of this.peHandlers) handler(parsed.chat_id, parsed.active);
496 return;
497 }
498
499 const chatId = (parsed as { chat_id?: string }).chat_id;
500 if (chatId) this.dispatch(chatId, parsed);
501 }
502
503 private dispatch(chatId: string, ev: InboundEvent): void {
504 const handlers = this.chatHandlers.get(chatId);
505 if (handlers && handlers.size > 0) {
506 for (const h of handlers) h(ev);
507 return;
508 }
509 if (!this.isReplayable(ev)) return;
510 this.appendToOfflineBuffer(chatId, ev);
511 }
512
513 /** Whether an inbound frame should be retained for subscribe-time replay. */
514 private isReplayable(ev: InboundEvent): boolean {
515 if (ev.event === "delta" || ev.event === "stream_end") return true;
516 if (ev.event === "question_answer_ok" || ev.event === "workplace_updated") {
517 return true;
518 }
519 if (ev.event === "message") {
520 return ev.kind !== "tool_hint" && ev.kind !== "progress";
521 }
522 return false;
523 }
524
525 private isTurnComplete(ev: InboundEvent): boolean {
526 return (
527 ev.event === "message" &&
528 ev.kind !== "tool_hint" &&
529 ev.kind !== "progress"
530 );
531 }
532
533 private appendToOfflineBuffer(chatId: string, ev: InboundEvent): void {
534 let buf = this.offlineBuffers.get(chatId);
535 if (!buf) {
536 buf = [];
537 this.offlineBuffers.set(chatId, buf);
538 }
539
540 if (ev.event === "delta") {
541 const last = buf.at(-1);
542 if (last && this.isTurnComplete(last)) {
543 // ask_user mid-turn: stream_end(resuming:true) → message(questions) → deltas.
544 // Do not drop the prefix when the next segment is still the same agent turn.
545 const prev = buf.at(-2);
546 const midTurnInFlight =
547 prev?.event === "stream_end" &&
548 (prev as Extract<InboundEvent, { event: "stream_end" }>).resuming ===
549 true;
550 if (!midTurnInFlight) {
551 buf.length = 0;
552 }
553 }
554 const deltaCount = buf.filter((entry) => entry.event === "delta").length;
555 if (deltaCount >= MAX_OFFLINE_DELTAS) {
556 this.coalesceOfflineDeltas(chatId, buf);
557 }
558 }
559
560 buf.push(ev);
561 }
562
563 /** Merge buffered deltas into one frame so long Director streams stay bounded. */
564 private coalesceOfflineDeltas(chatId: string, buf: InboundEvent[]): void {
565 let mergedText = "";
566 let index = 0;
567 while (index < buf.length && buf[index]?.event === "delta") {
568 const entry = buf[index] as Extract<InboundEvent, { event: "delta" }>;
569 mergedText += entry.text;
570 index += 1;
571 }
572 const tail = buf.slice(index);
573 buf.length = 0;
574 if (mergedText) {
575 buf.push({ event: "delta", chat_id: chatId, text: mergedText });
576 }
577 buf.push(...tail);
578 }
579
580 private scheduleOfflineReplay(chatId: string): void {
581 if (this.replayScheduled.has(chatId)) return;
582 this.replayScheduled.add(chatId);
583 queueMicrotask(() => {
584 this.replayScheduled.delete(chatId);
585 this.flushOfflineBuffer(chatId);
586 });
587 }
588
589 /** Deliver buffered frames to every current subscriber, then drop the buffer. */
590 private flushOfflineBuffer(chatId: string): void {
591 const buf = this.offlineBuffers.get(chatId);
592 if (!buf || buf.length === 0) return;
593 const handlers = this.chatHandlers.get(chatId);
594 if (!handlers || handlers.size === 0) return;
595
596 const events = buf.splice(0, buf.length);
597 this.offlineBuffers.delete(chatId);
598 for (const ev of events) {
599 for (const h of handlers) {
600 h(ev);
601 }
602 }
603 }
604
605 private handleClose(event?: { code?: number }): void {
606 this.socket = null;
607 if (this.pendingNewChat) {
608 clearTimeout(this.pendingNewChat.timer);
609 this.pendingNewChat.reject(new Error("socket closed"));
610 this.pendingNewChat = null;
611 }
612 this.rejectAllPendingWorkplaceActions("socket closed");
613 // Surface structured reasons *before* reconnect logic so the UI can
614 // display the error even while the client transparently reconnects.
615 // Browsers populate ``CloseEvent.code`` with the wire-level close code;
616 // 1009 = Message Too Big (server's max frame guard).
617 if (event?.code === 1009) {
618 this.emitError({ kind: "message_too_big" });
619 }
620 if (this.intentionallyClosed || !this.shouldReconnect) {
621 this.setStatus("closed");
622 return;
623 }
624 this.scheduleReconnect();
625 }
626
627 private emitError(error: StreamError): void {
628 // Isolate subscribers so a throwing handler cannot abort the surrounding
629 // ``handleClose`` flow (which still owes us a reconnect decision + status
630 // update). We deliberately swallow here: error reporting is best-effort
631 // and must never be allowed to compound the failure it's reporting.
632 for (const handler of this.errorHandlers) {
633 try {
634 handler(error);
635 } catch {
636 // best-effort: subscriber fault must not stall transport bookkeeping
637 }
638 }
639 }
640
641 private scheduleReconnect(): void {
642 this.setStatus("reconnecting");
643 const attempt = this.reconnectAttempts++;
644 // Exponential backoff: 0.5s, 1s, 2s, 4s, capped.
645 const delay = Math.min(500 * 2 ** attempt, this.maxBackoffMs);
646 this.reconnectTimer = setTimeout(async () => {
647 this.reconnectTimer = null;
648 if (this.options.onReauth) {
649 try {
650 const refreshed = await this.options.onReauth();
651 if (refreshed) this.currentUrl = refreshed;
652 } catch {
653 // fall through to retry with current URL
654 }
655 }
656 this.connect();
657 }, delay);
658 }
659
660 private registerWorkplaceAction(
661 timeoutMs: number = WORKPLACE_ACTION_TIMEOUT_MS,
662 ): {
663 requestId: string;
664 promise: Promise<{ work_id: string; workplace: WorkplaceData }>;
665 } {
666 const requestId = crypto.randomUUID();
667 const promise = new Promise<{ work_id: string; workplace: WorkplaceData }>(
668 (resolve, reject) => {
669 const timer = setTimeout(() => {
670 this.pendingWorkplaceActions.delete(requestId);
671 reject(new Error("workplace action timed out"));
672 }, timeoutMs);
673 this.pendingWorkplaceActions.set(requestId, { resolve, reject, timer });
674 },
675 );
676 return { requestId, promise };
677 }
678
679 private rejectAllPendingWorkplaceActions(reason: string): void {
680 for (const [, pending] of this.pendingWorkplaceActions) {
681 clearTimeout(pending.timer);
682 pending.reject(new Error(reason));
683 }
684 this.pendingWorkplaceActions.clear();
685 }
686
687 private queueSend(frame: Outbound): void {
688 if (this.socket?.readyState === WS_OPEN) {
689 this.rawSend(frame);
690 } else {
691 this.sendQueue.push(frame);
692 }
693 }
694
695 private rawSend(frame: Outbound): void {
696 if (!this.socket) return;
697 try {
698 this.socket.send(JSON.stringify(frame));
699 } catch {
700 // Send failure will materialize as a close; queue the frame for retry.
701 this.sendQueue.push(frame);
702 }
703 }
704 }
705
705 lines TYPESCRIPT