konnosDashboard
use-ai-stream.ts150 lines · main
1'use client';
2
3import { useCallback, useRef, useState } from 'react';
4
5/**
6 * Hook for consuming the briven AI SSE streaming endpoints. The api
7 * emits text/event-stream with three event types:
8 *
9 * event: token data: <chunk> (one per token)
10 * event: done data: {} (terminal)
11 * event: error data: <message> (terminal, indicates failure)
12 *
13 * The hook returns:
14 * - text: the accumulated tokens so far
15 * - status: 'idle' | 'streaming' | 'done' | 'error' | 'not_configured'
16 * - error: the error message (when status==='error' or 'not_configured')
17 * - start: call to kick off a stream
18 * - reset: wipe state without making a request
19 *
20 * Browser EventSource doesn't support POST + bodies, so we POST with
21 * fetch and read the response body as a UTF-8 stream, parsing SSE
22 * frames manually.
23 */
24
25export type StreamStatus = 'idle' | 'streaming' | 'done' | 'error' | 'not_configured';
26
27interface UseAiStreamReturn {
28 text: string;
29 status: StreamStatus;
30 error: string | null;
31 start: (url: string, body: unknown) => Promise<void>;
32 reset: () => void;
33}
34
35export function useAiStream(): UseAiStreamReturn {
36 const [text, setText] = useState('');
37 const [status, setStatus] = useState<StreamStatus>('idle');
38 const [error, setError] = useState<string | null>(null);
39 const abortRef = useRef<AbortController | null>(null);
40
41 const reset = useCallback(() => {
42 abortRef.current?.abort();
43 abortRef.current = null;
44 setText('');
45 setStatus('idle');
46 setError(null);
47 }, []);
48
49 const start = useCallback(async (url: string, body: unknown) => {
50 abortRef.current?.abort();
51 const ac = new AbortController();
52 abortRef.current = ac;
53 setText('');
54 setStatus('streaming');
55 setError(null);
56
57 let res: Response;
58 try {
59 res = await fetch(url, {
60 method: 'POST',
61 headers: { 'content-type': 'application/json' },
62 credentials: 'include',
63 body: JSON.stringify(body),
64 signal: ac.signal,
65 });
66 } catch (err) {
67 if (ac.signal.aborted) return;
68 setStatus('error');
69 setError(err instanceof Error ? err.message : 'network error');
70 return;
71 }
72
73 if (res.status === 503) {
74 const errBody = (await res.json().catch(() => null)) as { message?: string } | null;
75 setStatus('not_configured');
76 setError(errBody?.message ?? 'AI features are disabled on this deployment');
77 return;
78 }
79 if (!res.ok || !res.body) {
80 const errText = await res.text().catch(() => '');
81 setStatus('error');
82 setError(errText || `request failed (${res.status})`);
83 return;
84 }
85
86 const reader = res.body.getReader();
87 const decoder = new TextDecoder();
88 let buffer = '';
89 let accumulated = '';
90 try {
91 while (true) {
92 const { done, value } = await reader.read();
93 if (done) break;
94 buffer += decoder.decode(value, { stream: true });
95 // SSE frames are separated by blank lines.
96 let idx: number;
97 while ((idx = buffer.indexOf('\n\n')) >= 0) {
98 const frame = buffer.slice(0, idx);
99 buffer = buffer.slice(idx + 2);
100 const parsed = parseFrame(frame);
101 if (parsed.event === 'token' && parsed.data) {
102 accumulated += unescapeJson(parsed.data);
103 setText(accumulated);
104 } else if (parsed.event === 'done') {
105 setStatus('done');
106 return;
107 } else if (parsed.event === 'error') {
108 setStatus('error');
109 setError(parsed.data || 'stream error');
110 return;
111 }
112 }
113 }
114 // Body ended without a `done` event — treat as success if we got
115 // any text, otherwise as an error.
116 setStatus(accumulated.length > 0 ? 'done' : 'error');
117 if (accumulated.length === 0) setError('stream ended unexpectedly');
118 } catch (err) {
119 if (ac.signal.aborted) return;
120 setStatus('error');
121 setError(err instanceof Error ? err.message : 'stream read error');
122 }
123 }, []);
124
125 return { text, status, error, start, reset };
126}
127
128interface ParsedFrame {
129 event: string;
130 data: string;
131}
132
133function parseFrame(frame: string): ParsedFrame {
134 let event = 'message';
135 const dataLines: string[] = [];
136 for (const line of frame.split('\n')) {
137 if (line.startsWith('event:')) event = line.slice('event:'.length).trim();
138 else if (line.startsWith('data:')) dataLines.push(line.slice('data:'.length).trim());
139 }
140 return { event, data: dataLines.join('\n') };
141}
142
143/** The api's SSE writer escapes data with JSON.stringify(s).slice(1,-1). Reverse. */
144function unescapeJson(s: string): string {
145 try {
146 return JSON.parse(`"${s}"`) as string;
147 } catch {
148 return s;
149 }
150}