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.
77 lines
2.4 KiB
TypeScript
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,
|
|
};
|
|
}
|