Something went wrong. Try again.
ATProto-based AI Assistant Chat interface.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102import { useCallback, useEffect, useRef, useState } from 'react'
export type StreamedReply = { text: string; reasoning: string }
const FINISHED = new Set(['complete', 'error', 'cancelled'])const POLL_START_MS = 2_000const POLL_MAX_MS = 30_000
/** Follow a reply's stream, then refresh it from the store when it ends, polling if the stream drops. */export function useReplyStream( skey: string, refresh: () => Promise<void>, isPending: (rkey: string) => boolean,) { const [streams, setStreams] = useState<Record<string, StreamedReply>>({}) const sources = useRef(new Map<string, EventSource>()) const ended = useRef(new Set<string>()) const timers = useRef(new Set<ReturnType<typeof setTimeout>>()) // Keep the latest isPending in a ref for the polling timers to read. const pending = useRef(isPending) pending.current = isPending
/** Refresh until the reply is no longer pending, backing off between tries. */ const poll = useCallback( (replyRkey: string, delay = POLL_START_MS) => { const timer = setTimeout(() => { timers.current.delete(timer) refresh().then( () => pending.current(replyRkey) && poll(replyRkey, Math.min(delay * 2, POLL_MAX_MS)), (err: unknown) => { console.warn('Polling a reply failed', err) poll(replyRkey, Math.min(delay * 2, POLL_MAX_MS)) }, ) }, delay) timers.current.add(timer) }, [refresh], )
const follow = useCallback( (replyRkey: string) => { if (sources.current.has(replyRkey) || ended.current.has(replyRkey)) return const source = new EventSource(`/api/conversations/${skey}/messages/${replyRkey}/stream`) sources.current.set(replyRkey, source) const kinds: string[] = [] const parts: string[] = [] const finish = () => { source.close() sources.current.delete(replyRkey) ended.current.add(replyRkey) setStreams(({ [replyRkey]: _done, ...rest }) => rest) } const collect = (reasoning: boolean) => parts.filter((text, i) => text && (kinds[i] === 'reasoningPart') === reasoning).join('\n\n') source.addEventListener('part-start', (event) => { const { index, partType } = JSON.parse((event as MessageEvent).data) as { index: number partType: string } kinds[index] = partType }) source.addEventListener('delta', (event) => { const { index, text } = JSON.parse((event as MessageEvent).data) as { index: number text: string } parts[index] = (parts[index] ?? '') + text setStreams((current) => ({ ...current, [replyRkey]: { text: collect(false), reasoning: collect(true), }, })) }) source.addEventListener('status', (event) => { finish() const { status } = JSON.parse((event as MessageEvent).data) as { status: string } // Other statuses mean this server isn't running the reply, so wait for it elsewhere. if (!FINISHED.has(status)) return poll(replyRkey) refresh().catch((err: unknown) => console.warn('Refreshing a finished reply failed', err)) }) source.onerror = () => { finish() poll(replyRkey) } }, [skey, refresh, poll], )
useEffect( () => () => { for (const source of sources.current.values()) source.close() for (const timer of timers.current) clearTimeout(timer) }, [], )
return { streams, follow }}