use-db-ai-stream.ts

use-db-ai-stream.tsapps/web/src/hooks/use-db-ai-stream.tsraw ← hooks4,321 B← Back to the app
1"use client";
2 
3/**
4 * TASK-0124 — the shared state hook for the two streaming database-AI flows
5 * (generate + extract). Both are the same shape — stream SSE → terminal
6 * `done` frame carrying the result (`databaseId` / `recordIds` /
7 * `ignoredColumns`) — so one hook serves both dialogs.
8 *
9 * Error model (mirrors TASK-0122/0123 semantics):
10 *   - pre-stream failure (403/402/429/413/400) → `AiPreStreamError` with the
11 *     canonical `code` (the dialog renders the banner);
12 *   - in-band `error` frame after the stream started (invalid structure after
13 *     the retry-once, provider outage) → `reason: "error"` + the `done` frame
14 *     (the dialog shows the generic banner).
15 *
16 * Cancellation: `abort()` (dialog close / Stop) aborts the fetch; the stream
17 * resolves with `reason: "aborted"` and the server persists a truncated
18 * assistant row (no write is in flight — the write only happens after the
19 * validated tool call).
20 */
21import { useCallback, useEffect, useRef, useState } from "react";
22import { AiPreStreamError } from "@/lib/ai";
23import { sendDbExtract, sendDbGenerate, type DbAiResult } from "@/lib/ai-database";
24 
25export type DbAiStreamPhase = "idle" | "streaming" | "done" | "error";
26 
27export interface DbAiStreamError {
28  code: string;
29  message: string;
30  retryable?: boolean;
31  /** True when the failure is the monthly-token budget (the budget banner). */
32  budget?: boolean;
33}
34 
35export interface DbAiStreamState {
36  phase: DbAiStreamPhase;
37  /** The terminal `done` frame (present when phase === "done" and reason ok). */
38  result: DbAiResult | null;
39  error: DbAiStreamError | null;
40}
41 
42export interface UseDbAiStreamParams {
43  workspaceId: string;
44  databaseId?: string;
45  signal?: AbortSignal;
46}
47 
48export function useDbAiStream({ workspaceId }: UseDbAiStreamParams) {
49  const [phase, setPhase] = useState<DbAiStreamPhase>("idle");
50  const [result, setResult] = useState<DbAiResult | null>(null);
51  const [error, setError] = useState<DbAiStreamError | null>(null);
52  const abortRef = useRef<AbortController | null>(null);
53 
54  const run = useCallback(
55    async (
56      call: (signal: AbortSignal) => Promise<DbAiResult>
57    ): Promise<DbAiResult | null> => {
58      setPhase("streaming");
59      setResult(null);
60      setError(null);
61      const controller = new AbortController();
62      abortRef.current = controller;
63      try {
64        const res = await call(controller.signal);
65        if (res.reason === "aborted") {
66          setPhase("idle");
67          return null;
68        }
69        if (res.reason !== "ok" || !res.done) {
70          setPhase("error");
71          setResult(res); // carry the terminal frame (its `errorFrame`) for the banner
72          const ef = res.errorFrame;
73          setError(
74            ef
75              ? { code: ef.code, message: ef.message }
76              : { code: "AI_STREAM_ERROR", message: "The stream ended unexpectedly. Nothing was created." }
77          );
78          return null;
79        }
80        setResult(res);
81        setPhase("done");
82        return res;
83      } catch (err) {
84        if (err instanceof AiPreStreamError) {
85          setPhase("error");
86          setError({
87            code: err.code,
88            message: err.message,
89            retryable: err.retryAfterSeconds != null,
90            budget: err.code === "AI_BUDGET_EXCEEDED" || err.code === "AI_BUDGET_MONTHLY_EXCEEDED"
91          });
92          return null;
93        }
94        setPhase("error");
95        setError({ code: "NETWORK_ERROR", message: "Could not reach the server." });
96        return null;
97      } finally {
98        abortRef.current = null;
99      }
100    },
101    []
102  );
103 
104  const runGenerate = useCallback(
105    (prompt: string, parentId?: string) =>
106      run((signal) => sendDbGenerate({ workspaceId, prompt, ...(parentId ? { parentId } : {}), signal })),
107    [run, workspaceId]
108  );
109 
110  const runExtract = useCallback(
111    (databaseId: string, text: string) =>
112      run((signal) => sendDbExtract({ workspaceId, databaseId, text, signal })),
113    [run, workspaceId]
114  );
115 
116  const abort = useCallback(() => {
117    abortRef.current?.abort();
118    abortRef.current = null;
119  }, []);
120 
121  // Abort on unmount (the dialog unmounts when closed).
122  useEffect(() => () => abortRef.current?.abort(), []);
123 
124  return { phase, result, error, runGenerate, runExtract, abort };
125}
126