gitdataai/src/hooks/use-ai-streaming.ts
ZhenYi e9d5407c66 feat(room): add AI streaming and message hooks for frontend
Add use-ai-streaming hook for SSE-based AI response streaming,
use-room-messages for real-time message updates, and wire up
room context + ws protocol changes.
2026-04-30 19:15:42 +08:00

77 lines
2.4 KiB
TypeScript

import { useCallback, useRef, useState } from 'react';
import type { RoomWsClient } from '@/lib/room-ws-client';
export interface AiStreamChunk {
type: string;
content: string;
seq?: number;
}
export interface ActiveAiStream {
message_id: string;
display_name: string;
}
/**
* Hook managing AI streaming state: streaming chunks, active stream indicator,
* and stream cancellation. Separated from the main room context to reduce
* the God component size (~1583 lines → ~300).
*/
export function useAiStreaming(clientRef: React.MutableRefObject<RoomWsClient | null>) {
const [streamingChunks, setStreamingChunks] = useState<Map<string, AiStreamChunk[]>>(new Map());
const [activeAiStream, setActiveAiStream] = useState<ActiveAiStream | null>(null);
// Ref to latest chunks so done handler reads current state (setState is async)
const chunksRef = useRef<Map<string, AiStreamChunk[]>>(new Map());
const clearStreamingState = useCallback((msgId: string) => {
setStreamingChunks(prev => { prev.delete(msgId); return new Map(prev); });
chunksRef.current.delete(msgId);
setActiveAiStream(null);
}, []);
const insertChunk = useCallback((
messageId: string,
chunkType: string | undefined,
content: string,
seq: number | undefined,
) => {
setStreamingChunks(prev => {
const next = new Map(prev);
const existing: AiStreamChunk[] = next.get(messageId) ?? [];
const s = seq ?? existing.length;
const newChunk: AiStreamChunk = { type: chunkType ?? 'answer', content, seq: s };
const insertIdx = existing.findIndex(c => c.seq != null && c.seq > s);
next.set(messageId,
insertIdx === -1
? [...existing, newChunk]
: [...existing.slice(0, insertIdx), newChunk, ...existing.slice(insertIdx)]
);
chunksRef.current = new Map(next);
return next;
});
}, []);
const getOrderedChunks = useCallback((msgId: string): AiStreamChunk[] => {
return chunksRef.current.get(msgId) ?? [];
}, []);
const cancelAiStream = useCallback(async () => {
const client = clientRef.current;
if (!client) return false;
const roomId = client.getSubscribedRooms().values().next().value;
if (!roomId) return false;
return client.cancelAiStream(roomId as string);
}, [clientRef]);
return {
streamingChunks,
activeAiStream,
setActiveAiStream,
clearStreamingState,
insertChunk,
getOrderedChunks,
cancelAiStream,
chunksRef,
};
}