Files
zui/frontend/src/hooks/useRunStream.ts
2026-03-29 01:53:15 +01:00

111 lines
3.2 KiB
TypeScript

/**
* SSE subscriber hook for graph run execution.
* Connects to GET /api/runs/:id/stream and updates runStore.
*/
import { useEffect, useRef, useCallback } from 'react'
import { useRunStore } from '@/lib/graph/runStore'
import { dispatchCanvasCommand } from '@/app/canvas/canvasStore'
export function useRunStream() {
const eventSourceRef = useRef<EventSource | null>(null)
const { startRun, setRunStatus, setNodeStatus, appendChunk } = useRunStore()
const disconnect = useCallback(() => {
if (eventSourceRef.current) {
eventSourceRef.current.close()
eventSourceRef.current = null
}
}, [])
const connectToRun = useCallback(
(runId: string) => {
disconnect()
startRun(runId)
const es = new EventSource(`/api/runs/${runId}/stream`)
eventSourceRef.current = es
es.addEventListener('run/started', (e) => {
const data = JSON.parse(e.data)
setRunStatus('running')
// Initialize all nodes as pending so overlays appear immediately
const nodeIds = data.nodeIds as string[] | undefined
if (nodeIds) {
for (const nodeId of nodeIds) {
setNodeStatus(nodeId, { status: 'pending' })
}
}
})
es.addEventListener('run/completed', () => {
setRunStatus('completed')
es.close()
})
es.addEventListener('run/failed', (e) => {
const data = JSON.parse(e.data)
setRunStatus('failed', data.error)
es.close()
})
es.addEventListener('step/started', (e) => {
const data = JSON.parse(e.data)
setNodeStatus(data.nodeId, { status: 'running' })
// Fire trail animation along edges leading to this node
dispatchCanvasCommand({ type: 'path/addTrigger', payload: data.nodeId })
})
es.addEventListener('step/completed', (e) => {
const data = JSON.parse(e.data)
setNodeStatus(data.nodeId, { status: 'completed', output: data.output })
})
es.addEventListener('step/failed', (e) => {
const data = JSON.parse(e.data)
setNodeStatus(data.nodeId, { status: 'failed', error: data.error })
})
es.addEventListener('step/chunk', (e) => {
const data = JSON.parse(e.data)
appendChunk(data.nodeId, data.chunk)
})
es.onerror = () => {
setRunStatus('failed', 'Connection lost')
es.close()
}
},
[disconnect, startRun, setRunStatus, setNodeStatus, appendChunk]
)
// Cleanup on unmount
useEffect(() => disconnect, [disconnect])
return { connectToRun, disconnect }
}
/**
* Trigger a new run: POST the graph, then connect SSE.
*/
export async function createAndStreamRun(
recollectionId: string,
graph: { nodes: unknown[]; edges: unknown[] },
connectToRun: (runId: string) => void,
sceneId?: string
): Promise<void> {
const res = await fetch('/api/runs', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ recollectionId, sceneId, graph }),
})
if (!res.ok) {
const err = await res.json().catch(() => ({ error: 'Failed to create run' }))
throw new Error(err.error ?? 'Failed to create run')
}
const { id } = await res.json()
connectToRun(id)
}