import { useState, useCallback, useRef, useEffect, useMemo } from "react";
import { TASK_SUBSCRIBE_ENDPOINT, getTaskOutput } from "../lib/api";
import type { TaskOutputEvent } from "./useTaskSubscription";
export interface MultiTaskOutputEntry extends TaskOutputEvent {
/** Label for the task (e.g. step name or "Orchestrator") */
taskLabel: string;
/** Timestamp when the entry was received */
receivedAt: number;
/** Whether this entry was backfilled from historical data (not live-streamed) */
isBackfill?: boolean;
}
interface UseMultiTaskSubscriptionOptions {
/** Map of taskId -> label */
taskMap: Map<string, string>;
/** Whether to actively subscribe */
enabled?: boolean;
/** Max entries to keep in buffer */
maxEntries?: number;
}
export function useMultiTaskSubscription(options: UseMultiTaskSubscriptionOptions) {
const { taskMap, enabled = true, maxEntries = 2000 } = options;
const [connected, setConnected] = useState(false);
const [entries, setEntries] = useState<MultiTaskOutputEntry[]>([]);
const wsRef = useRef<WebSocket | null>(null);
const reconnectTimeoutRef = useRef<number | null>(null);
const subscribedTasksRef = useRef<Set<string>>(new Set());
const backfilledTasksRef = useRef<Set<string>>(new Set());
const taskMapRef = useRef(taskMap);
const enabledRef = useRef(enabled);
/** Track which task IDs have already been backfilled to avoid re-fetching */
const backfilledTasksRef = useRef<Set<string>>(new Set());
// Keep refs in sync
useEffect(() => {
taskMapRef.current = taskMap;
}, [taskMap]);
useEffect(() => {
enabledRef.current = enabled;
}, [enabled]);
/** Max number of historical events to backfill per task */
const MAX_BACKFILL_PER_TASK = 200;
/**
* Convert a TaskEvent (from the REST API) into a MultiTaskOutputEntry.
* Only converts events with event_type === 'output'.
*/
const convertTaskEventToEntry = useCallback(
(event: TaskEvent): MultiTaskOutputEntry | null => {
if (event.eventType !== "output") return null;
const data = event.eventData;
if (!data) return null;
return {
taskId: event.taskId,
messageType: (data.messageType as string) || "system",
content: (data.content as string) || "",
toolName: data.toolName as string | undefined,
toolInput: data.toolInput as Record<string, unknown> | undefined,
isError: data.isError as boolean | undefined,
costUsd: data.costUsd as number | undefined,
durationMs: data.durationMs as number | undefined,
isPartial: false,
taskLabel:
taskMapRef.current.get(event.taskId) || event.taskId,
receivedAt: new Date(event.createdAt).getTime(),
isBackfill: true,
};
},
[]
);
/**
* Backfill historical log entries for a task from the REST API.
* Only fetches once per task ID (tracked in backfilledTasksRef).
*/
const backfillTask = useCallback(
async (taskId: string) => {
if (backfilledTasksRef.current.has(taskId)) return;
backfilledTasksRef.current.add(taskId);
try {
const response = await listTaskEvents(taskId);
const events = response.events;
// The API returns events in DESC order; reverse to get chronological ASC
const chronologicalEvents = [...events].reverse();
// Filter to output events and convert, limiting to MAX_BACKFILL_PER_TASK
const backfillEntries: MultiTaskOutputEntry[] = [];
for (const event of chronologicalEvents) {
const entry = convertTaskEventToEntry(event);
if (entry) {
backfillEntries.push(entry);
if (backfillEntries.length >= MAX_BACKFILL_PER_TASK) break;
}
}
if (backfillEntries.length === 0) return;
// Prepend historical entries before any existing live entries for this task,
// maintaining overall chronological order across all tasks
setEntries((prev) => {
// Merge backfill entries with existing entries, maintaining chronological order
const merged = [...backfillEntries, ...prev];
// Sort by receivedAt to ensure proper chronological ordering
merged.sort((a, b) => a.receivedAt - b.receivedAt);
// Trim to maxEntries
if (merged.length > maxEntries) {
return merged.slice(merged.length - maxEntries);
}
return merged;
});
} catch (e) {
console.error(`Failed to backfill task events for ${taskId}:`, e);
// Remove from backfilled set so it can be retried
backfilledTasksRef.current.delete(taskId);
}
},
[convertTaskEventToEntry, maxEntries]
);
// Derive task IDs from the map, stabilized to avoid unnecessary effect triggers
const taskIdsKey = useMemo(() => Array.from(taskMap.keys()).sort().join(","), [taskMap]);
const taskIds = useMemo(() => Array.from(taskMap.keys()), [taskIdsKey]); // eslint-disable-line react-hooks/exhaustive-deps
const subscribeToTask = useCallback((ws: WebSocket, taskId: string) => {
if (ws.readyState === WebSocket.OPEN) {
ws.send(JSON.stringify({ type: "subscribeOutput", taskId }));
subscribedTasksRef.current.add(taskId);
}
}, []);
const unsubscribeFromTask = useCallback((ws: WebSocket, taskId: string) => {
if (ws.readyState === WebSocket.OPEN) {
ws.send(JSON.stringify({ type: "unsubscribeOutput", taskId }));
subscribedTasksRef.current.delete(taskId);
}
}, []);
const backfillTask = useCallback(
async (taskId: string, label: string) => {
if (backfilledTasksRef.current.has(taskId)) return;
backfilledTasksRef.current.add(taskId);
try {
const response = await getTaskOutput(taskId);
if (response.entries.length === 0) return;
const historicalEntries: MultiTaskOutputEntry[] = response.entries.map(
(entry) => ({
taskId: entry.taskId,
messageType: entry.messageType,
content: entry.content,
toolName: entry.toolName,
toolInput: entry.toolInput,
isError: entry.isError,
costUsd: entry.costUsd,
durationMs: entry.durationMs,
isPartial: false,
taskLabel: label,
receivedAt: new Date(entry.createdAt || Date.now()).getTime(),
})
);
setEntries((prev) => {
// De-duplicate by checking if content+taskId+messageType already exists
const existingKeys = new Set(
prev.map(
(e) =>
`${e.taskId}:${e.messageType}:${e.content.slice(0, 100)}`
)
);
const newHistorical = historicalEntries.filter(
(e) =>
!existingKeys.has(
`${e.taskId}:${e.messageType}:${e.content.slice(0, 100)}`
)
);
const combined = [...newHistorical, ...prev];
if (combined.length > maxEntries) {
return combined.slice(combined.length - maxEntries);
}
return combined;
});
} catch (e) {
console.error(`Failed to backfill task ${taskId}:`, e);
}
},
[maxEntries]
);
const connect = useCallback(() => {
const currentState = wsRef.current?.readyState;
if (currentState === WebSocket.OPEN || currentState === WebSocket.CONNECTING) {
return;
}
if (wsRef.current && currentState === WebSocket.CLOSING) {
wsRef.current = null;
}
try {
const ws = new WebSocket(TASK_SUBSCRIBE_ENDPOINT);
wsRef.current = ws;
ws.onopen = () => {
setConnected(true);
// Re-subscribe to all tasks
for (const taskId of subscribedTasksRef.current) {
ws.send(JSON.stringify({ type: "subscribeOutput", taskId }));
}
};
ws.onmessage = (event) => {
try {
const message = JSON.parse(event.data);
if (message.type === "taskOutput") {
const label = taskMapRef.current.get(message.taskId) || message.taskId;
const entry: MultiTaskOutputEntry = {
taskId: message.taskId,
messageType: message.messageType,
content: message.content,
toolName: message.toolName,
toolInput: message.toolInput,
isError: message.isError,
costUsd: message.costUsd,
durationMs: message.durationMs,
isPartial: message.isPartial,
taskLabel: label,
receivedAt: Date.now(),
};
setEntries((prev) => {
const next = [...prev, entry];
if (next.length > maxEntries) {
return next.slice(next.length - maxEntries);
}
return next;
});
}
} catch (e) {
console.error("Failed to parse multi-task subscription message:", e);
}
};
ws.onerror = () => {
console.error("Multi-task WebSocket connection error");
};
ws.onclose = () => {
setConnected(false);
wsRef.current = null;
// Reconnect if we still have subscriptions
if (subscribedTasksRef.current.size > 0 && enabledRef.current) {
reconnectTimeoutRef.current = window.setTimeout(() => {
connect();
}, 3000);
}
};
} catch (e) {
console.error("Failed to connect multi-task subscription:", e);
}
}, [maxEntries]);
// Manage subscriptions when task IDs change
useEffect(() => {
if (!enabled || taskIds.length === 0) {
// Close connection if no tasks
if (wsRef.current) {
subscribedTasksRef.current.clear();
wsRef.current.close();
wsRef.current = null;
}
return;
}
const newTaskIds = new Set(taskIds);
const ws = wsRef.current;
if (!ws || ws.readyState !== WebSocket.OPEN) {
// Set desired subscriptions and connect
subscribedTasksRef.current = newTaskIds;
connect();
// Backfill all initial tasks
for (const taskId of newTaskIds) {
const label = taskMapRef.current.get(taskId) || taskId;
backfillTask(taskId, label);
}
return;
}
// Unsubscribe from removed tasks
for (const existingId of subscribedTasksRef.current) {
if (!newTaskIds.has(existingId)) {
unsubscribeFromTask(ws, existingId);
}
}
// Subscribe to new tasks and backfill their history
for (const newId of newTaskIds) {
if (!subscribedTasksRef.current.has(newId)) {
subscribeToTask(ws, newId);
const label = taskMapRef.current.get(newId) || newId;
backfillTask(newId, label);
}
}
}, [taskIds, enabled, connect, subscribeToTask, unsubscribeFromTask, backfillTask]);
// Cleanup on unmount
useEffect(() => {
return () => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
}
if (wsRef.current) {
wsRef.current.close();
}
};
}, []);
const clearEntries = useCallback(() => {
setEntries([]);
backfilledTasksRef.current.clear();
}, []);
return {
connected,
entries,
clearEntries,
};
}
