use-ai-stream.ts

use-ai-stream.tsapps/web/src/hooks/use-ai-stream.tsraw ← hooks15,347 B← Back to the app
1"use client";
2 
3/**
4 * TASK-0122 — Ask AI panel state (the streaming conversation hook).
5 *
6 * Owns, for ONE workspace: the conversation list, the active conversation's
7 * messages, the in-flight SSE stream, and the error banner state. The panel
8 * component is a thin view over this hook.
9 *
10 * Streaming model (mirrors docs/api/ai.md):
11 *   - `send(content)` → ensure a conversation (seeded with the entry point's
12 *     context) → optimistic user bubble → `sendAiMessage` frames append to a
13 *     live assistant bubble (`delta`), collect tool calls, then `done`
14 *     finalizes the rows (the server persisted them — the ids come back).
15 *   - `stop()` aborts the fetch; the server persists the partial
16 *     (`metadata.truncated`) and we refresh the rows on the next open.
17 *   - `retry()` re-sends the LAST user message with `retryOf` (no duplicate
18 *     user row) after a mid-stream / pre-stream error.
19 */
20import { useCallback, useEffect, useRef, useState } from "react";
21import {
22  aiApi,
23  sendAiMessage,
24  AiPreStreamError,
25  type AiContextRef,
26  type AiConversationSummary,
27  type AiMessage,
28  type AiSseFrame
29} from "@/lib/ai";
30 
31export interface AiBannerError {
32  code: string;
33  message: string;
34  retryable: boolean;
35  retryAfterSeconds?: number;
36}
37 
38interface AiTurn {
39  userMessageId?: string;
40  content: string;
41  toolCalls: Array<{ id: string; name: string; arguments: Record<string, unknown> }>;
42  usage?: { tokensIn: number; tokensOut: number };
43  done?: boolean;
44  assistantMessageId?: string;
45  truncated?: boolean;
46  error?: string;
47}
48 
49/** Optimistic row ids (never collide with server cuids). */
50let tmpSeq = 0;
51const tmpId = (prefix: string) => `${prefix}-tmp-${Date.now().toString(36)}-${(tmpSeq += 1)}`;
52 
53export function useAiStream(workspaceId: string | null, context: AiContextRef | null) {
54  const [conversations, setConversations] = useState<AiConversationSummary[]>([]);
55  const [convError, setConvError] = useState<string | null>(null);
56  const [activeId, setActiveId] = useState<string | null>(null);
57  const [messages, setMessages] = useState<AiMessage[]>([]);
58  const [messagesLoading, setMessagesLoading] = useState(false);
59  const [streaming, setStreaming] = useState(false);
60  const [banner, setBanner] = useState<AiBannerError | null>(null);
61  /** The in-flight assistant bubble (empty = none). */
62  const [live, setLive] = useState<AiTurn | null>(null);
63  /** The optimistic user bubble for the in-flight turn. */
64  const [pendingUser, setPendingUser] = useState<{ id: string; content: string } | null>(null);
65 
66  /**
67   * Synchronous source of truth for the in-flight turn. `setLive` is only the
68   * render trigger; every transition is computed from `liveRef` and the ref
69   * is updated in the SAME synchronous block (before the React state commit).
70   *
71   * Why a ref instead of reading the `live` state: the finalize runs right
72   * after `sendAiMessage` resolves — before React has re-rendered, so the
73   * `live` closure is stale. And the ref must NOT be updated inside a
74   * `setLive` *updater*: React runs updaters during render (asynchronously,
75   * twice under StrictMode), so a finalize that reads a ref written in an
76   * updater would see the pre-`done` value. That bug made Retry send the
77   * optimistic `tmpId("u")` as `retryOf` (server 404: "Message not found")
78   * because the server's real `userMessageId` (from the `done` frame) had
79   * never landed in the ref.
80   */
81  const liveRef = useRef<AiTurn | null>(null);
82 
83  const abortRef = useRef<AbortController | null>(null);
84  /** The last turn's text (+ its server row id once known) — the Retry
85   *  button re-sends THIS, whether or not the user row was persisted
86   *  (a pre-stream 429 on a fresh conversation has no row to retryOf). */
87  const lastSentRef = useRef<{ content: string; userMessageId?: string } | null>(null);
88  // Refs so `send`/`retry` (stable callbacks) see fresh context without
89  // re-creating on every keystroke of an unrelated input.
90  const contextRef = useRef(context);
91  contextRef.current = context;
92 
93  /* --------------------------------------------------------- conv list */
94 
95  const refreshConversations = useCallback(async () => {
96    if (!workspaceId) return;
97    try {
98      const list = await aiApi.listConversations(workspaceId);
99      setConversations(list);
100      setConvError(null);
101    } catch {
102      setConvError("Could not load conversations.");
103    }
104  }, [workspaceId]);
105 
106  useEffect(() => {
107    setConversations([]);
108    setActiveId(null);
109    setMessages([]);
110    setBanner(null);
111    if (workspaceId) void refreshConversations();
112  }, [workspaceId, refreshConversations]);
113 
114  /* -------------------------------------------------------- active conv */
115 
116  const openConversation = useCallback(
117    async (id: string) => {
118      if (streaming) stop();
119      setActiveId(id);
120      setMessages([]);
121      setBanner(null);
122      setMessagesLoading(true);
123      try {
124        const rows = await aiApi.listMessages(id);
125        setMessages(rows);
126      } catch {
127        setBanner({ code: "ERROR", message: "Could not load this conversation.", retryable: false });
128      } finally {
129        setMessagesLoading(false);
130      }
131    },
132    // eslint-disable-next-line react-hooks/exhaustive-deps
133    [streaming]
134  );
135 
136  const startNew = useCallback(() => {
137    if (streaming) stop();
138    setActiveId(null);
139    setMessages([]);
140    setBanner(null);
141    // eslint-disable-next-line react-hooks/exhaustive-deps
142  }, [streaming]);
143 
144  const removeConversation = useCallback(
145    async (id: string) => {
146      try {
147        await aiApi.deleteConversation(id);
148      } catch {
149        /* the row simply stays if the delete failed — non-fatal */
150      }
151      setConversations((list) => list.filter((c) => c.id !== id));
152      if (activeId === id) startNew();
153    },
154    [activeId, startNew]
155  );
156 
157  /** Inline rename (Enter/blur) — the server PATCHes the title. */
158  const renameConversation = useCallback(
159    async (id: string, title: string) => {
160      const t = title.trim();
161      if (!t) return;
162      try {
163        const updated = await aiApi.renameConversation(id, t);
164        setConversations((list) => list.map((c) => (c.id === id ? { ...c, title: updated.title ?? t } : c)));
165      } catch {
166        /* keep the old title; the row is untouched */
167      }
168    },
169    []
170  );
171 
172  /* -------------------------------------------------------------- send */
173 
174  const applyFrame = useCallback((frame: AiSseFrame) => {
175    // Transition from the ref (the synchronous truth) and commit BOTH the ref
176    // and the render state in the same synchronous block. `setLive(next)` is
177    // a plain value (not an updater), so StrictMode's double-invocation does
178    // not apply and there is no async-timing hazard.
179    const t = liveRef.current;
180    if (t) {
181      let next = t;
182      if (frame.type === "delta") next = { ...t, content: t.content + frame.text };
183      else if (frame.type === "tool_call")
184        next = { ...t, toolCalls: [...t.toolCalls, { id: frame.id, name: frame.name, arguments: frame.arguments }] };
185      else if (frame.type === "usage") next = { ...t, usage: { tokensIn: frame.tokensIn, tokensOut: frame.tokensOut } };
186      else if (frame.type === "done")
187        next = {
188          ...t,
189          done: true,
190          assistantMessageId: frame.assistantMessageId,
191          userMessageId: frame.userMessageId ?? t.userMessageId,
192          truncated: frame.reason === "aborted",
193          ...(frame.reason === "aborted" ? { error: "aborted" } : {})
194        };
195      if (next !== t) {
196        liveRef.current = next;
197        setLive(next);
198      }
199    }
200    if (frame.type === "error") {
201      setBanner({
202        code: frame.code,
203        message: frame.message,
204        retryable: frame.retryable === true,
205        ...(typeof frame.retryAfterSeconds === "number" ? { retryAfterSeconds: frame.retryAfterSeconds } : {})
206      });
207    }
208  }, []);
209 
210  /** Ensure the turn has a conversation (create on the fly, seeded context). */
211  const ensureConversation = useCallback(async (): Promise<string | null> => {
212    if (activeId) return activeId;
213    if (!workspaceId) return null;
214    const created = await aiApi.createConversation(workspaceId, {
215      ...(contextRef.current && Object.keys(contextRef.current).length > 0 ? { context: contextRef.current } : {})
216    });
217    setActiveId(created.id);
218    setConversations((list) => [created, ...list]);
219    return created.id;
220  }, [activeId, workspaceId]);
221 
222  const send = useCallback(
223    async (content: string) => {
224      const text = content.trim();
225      if (!text || streaming || !workspaceId) return;
226      setBanner(null);
227      const convId = await ensureConversation();
228      if (!convId) return;
229      const userRow = { id: tmpId("u"), content: text };
230      lastSentRef.current = { content: text };
231      setPendingUser(userRow);
232      liveRef.current = { userMessageId: userRow.id, content: "", toolCalls: [] };
233      setLive(liveRef.current);
234      setStreaming(true);
235      const ac = new AbortController();
236      abortRef.current = ac;
237      try {
238        const reason = await sendAiMessage(
239          { workspaceId, conversationId: convId, content: text, context: contextRef.current ?? undefined, signal: ac.signal },
240          applyFrame
241        );
242        // Finalize: `liveRef` is the synchronous truth (every transition is
243        // computed from the ref in `applyFrame`, so it already carries the
244        // server's real `userMessageId` from the `done` frame). Reading it
245        // here — not the stale `live` closure, and not a ref written inside a
246        // `setLive` updater (React runs updaters asynchronously during
247        // render) — is what lets Retry send the REAL row id as `retryOf`
248        // instead of the optimistic tmpId (which 404s server-side).
249        const t = liveRef.current;
250        if (t) {
251          if (t.userMessageId) lastSentRef.current = { content: text, userMessageId: t.userMessageId };
252          setMessages((rows) => [
253            ...rows,
254            {
255              id: t.userMessageId ?? userRow.id,
256              role: "user",
257              content: text,
258              createdAt: new Date().toISOString()
259            },
260            ...(t.content
261              ? [
262                  {
263                    id: t.assistantMessageId ?? tmpId("a"),
264                    role: "assistant" as const,
265                    content: t.content,
266                    model: null,
267                    tokensIn: t.usage?.tokensIn ?? null,
268                    tokensOut: t.usage?.tokensOut ?? null,
269                    createdAt: new Date().toISOString(),
270                    metadata: { feature: "chat", ...(t.truncated ? { truncated: true } : {}), ...(t.error ? { error: t.error } : {}) }
271                  }
272                ]
273              : [])
274          ]);
275        }
276        liveRef.current = null;
277        setLive(null);
278        void reason;
279      } catch (err) {
280        if (err instanceof AiPreStreamError) {
281          setBanner({
282            code: err.code,
283            message: err.message,
284            retryable: err.code === "RATE_LIMITED" || err.status >= 500,
285            ...(err.retryAfterSeconds !== undefined ? { retryAfterSeconds: err.retryAfterSeconds } : {})
286          });
287          // pre-stream: the user row was NOT persisted — re-queue the text
288          // on Retry (no retryOf — the server has no row to re-send).
289        } else if (!(err instanceof Error && err.name === "AbortError")) {
290          setBanner({ code: "ERROR", message: "Network error — the AI reply could not be completed.", retryable: true });
291        }
292      } finally {
293        setPendingUser(null);
294        setStreaming(false);
295        abortRef.current = null;
296        void refreshConversations(); // auto-title / message counts changed
297      }
298    },
299    [workspaceId, streaming, ensureConversation, applyFrame, refreshConversations]
300  );
301 
302  const stop = useCallback(() => {
303    abortRef.current?.abort();
304  }, []);
305 
306  /**
307   * Re-send the last turn after an error (the Retry button).
308   *
309   * `lastSentRef` carries the turn's text + the persisted user-row id (set
310   * when the `done` frame arrived). When the row exists (a mid-stream error)
311   * we send `retryOf` so the server re-runs WITHOUT duplicating the user
312   * message; when it does not (a pre-stream 402/429 on a fresh conversation)
313   * we send a plain new message — same text, new row.
314   */
315  const retry = useCallback(async () => {
316    if (streaming || !workspaceId || !activeId) return;
317    const last = lastSentRef.current;
318    if (!last) return;
319    setBanner(null);
320    liveRef.current = { userMessageId: last.userMessageId, content: "", toolCalls: [] };
321    setLive(liveRef.current);
322    setStreaming(true);
323    const ac = new AbortController();
324    abortRef.current = ac;
325    try {
326      await sendAiMessage(
327        {
328          workspaceId,
329          conversationId: activeId,
330          content: last.content,
331          context: contextRef.current ?? undefined,
332          ...(last.userMessageId ? { retryOf: last.userMessageId } : {}),
333          signal: ac.signal
334        },
335        applyFrame
336      );
337      // Finalize outside the updater (same StrictMode reason as send). The
338      // user row already lives in `messages` (the original send added it, or
339      // it came from the server) — retry appends ONLY the new assistant row;
340      // with `retryOf` the server did not create a second user row either.
341      const t = liveRef.current;
342      if (t && t.content && !ac.signal.aborted) {
343        if (t.userMessageId) lastSentRef.current = { content: last.content, userMessageId: t.userMessageId };
344        setMessages((rows) => [
345          ...rows,
346          {
347            id: t.assistantMessageId ?? tmpId("a"),
348            role: "assistant",
349            content: t.content,
350            model: null,
351            tokensIn: t.usage?.tokensIn ?? null,
352            tokensOut: t.usage?.tokensOut ?? null,
353            createdAt: new Date().toISOString(),
354            metadata: { feature: "chat", ...(t.truncated ? { truncated: true } : {}) }
355          }
356        ]);
357      }
358      liveRef.current = null;
359      setLive(null);
360    } catch (err) {
361      if (err instanceof AiPreStreamError) {
362        setBanner({
363          code: err.code,
364          message: err.message,
365          retryable: err.code === "RATE_LIMITED" || err.status >= 500,
366          ...(err.retryAfterSeconds !== undefined ? { retryAfterSeconds: err.retryAfterSeconds } : {})
367        });
368      } else if (!(err instanceof Error && err.name === "AbortError")) {
369        setBanner({ code: "ERROR", message: "Network error — the AI reply could not be completed.", retryable: true });
370      }
371    } finally {
372      setStreaming(false);
373      abortRef.current = null;
374      void refreshConversations();
375    }
376  }, [streaming, workspaceId, activeId, applyFrame, refreshConversations]);
377 
378  // Aborting on unmount (panel close) → the server persists the partial.
379  useEffect(() => () => abortRef.current?.abort(), []);
380 
381  return {
382    conversations,
383    convError,
384    activeId,
385    messages,
386    messagesLoading,
387    streaming,
388    banner,
389    live,
390    pendingUser,
391    send,
392    stop,
393    retry,
394    openConversation,
395    startNew,
396    removeConversation,
397    renameConversation,
398    refreshConversations
399  };
400}
401