huskies/frontend/src/hooks/useChatWebSocket.ts

import * as React from "react";
import type {
	PipelineState,
	StatusEvent,
	WizardStateData,
} from "../api/client";
import { api, ChatWebSocket } from "../api/client";
import type { LogEntry } from "../components/ServerLogsPanel";
import type { Message } from "../types";
import { formatToolActivity } from "../utils/chatUtils";

const { useEffect, useRef, useState } = React;

type SetState<T> = React.Dispatch<React.SetStateAction<T>>;

interface UseChatWebSocketParams {
	setMessages: SetState<Message[]>;
	setLoading: SetState<boolean>;
	setClaudeSessionId: SetState<string | null>;
	queuedMessagesRef: React.MutableRefObject<{ id: string; text: string }[]>;
	setQueuedMessages: SetState<{ id: string; text: string }[]>;
	setPendingAutoSendBatch: SetState<string[] | null>;
}

interface ReconciliationEvent {
	id: string;
	storyId: string;
	status: string;
	message: string;
}

export interface UseChatWebSocketResult {
	wsRef: React.MutableRefObject<ChatWebSocket | null>;
	wsConnected: boolean;
	streamingContent: string;
	setStreamingContent: SetState<string>;
	streamingThinking: string;
	setStreamingThinking: SetState<string>;
	activityStatus: string | null;
	setActivityStatus: SetState<string | null>;
	permissionQueue: {
		requestId: string;
		toolName: string;
		toolInput: Record<string, unknown>;
	}[];
	setPermissionQueue: SetState<
		{
			requestId: string;
			toolName: string;
			toolInput: Record<string, unknown>;
		}[]
	>;
	pipeline: PipelineState;
	pipelineVersion: number;
	reconciliationActive: boolean;
	reconciliationEvents: ReconciliationEvent[];
	agentConfigVersion: number;
	agentStateVersion: number;
	needsOnboarding: boolean;
	setNeedsOnboarding: SetState<boolean>;
	wizardState: WizardStateData | null;
	setWizardState: SetState<WizardStateData | null>;
	sideQuestion: {
		question: string;
		response: string;
		loading: boolean;
	} | null;
	setSideQuestion: SetState<{
		question: string;
		response: string;
		loading: boolean;
	} | null>;
	serverLogs: LogEntry[];
	storyTokenCosts: Map<string, number>;
	/** Structured pipeline status events. Each entry preserves the full
	 * StatusEvent so future UI stories can render per-type icons or filters. */
	statusEvents: Array<{ receivedAt: string; event: StatusEvent }>;
}

export function useChatWebSocket({
	setMessages,
	setLoading,
	setClaudeSessionId,
	queuedMessagesRef,
	setQueuedMessages,
	setPendingAutoSendBatch,
}: UseChatWebSocketParams): UseChatWebSocketResult {
	const wsRef = useRef<ChatWebSocket | null>(null);
	const [wsConnected, setWsConnected] = useState(false);
	const [streamingContent, setStreamingContent] = useState("");
	const [streamingThinking, setStreamingThinking] = useState("");
	const [activityStatus, setActivityStatus] = useState<string | null>(null);
	const [permissionQueue, setPermissionQueue] = useState<
		{
			requestId: string;
			toolName: string;
			toolInput: Record<string, unknown>;
		}[]
	>([]);
	const [pipeline, setPipeline] = useState<PipelineState>({
		backlog: [],
		current: [],
		qa: [],
		merge: [],
		done: [],
		deterministic_merges_in_flight: [],
	});
	const [pipelineVersion, setPipelineVersion] = useState(0);
	const [reconciliationActive, setReconciliationActive] = useState(false);
	const [reconciliationEvents, setReconciliationEvents] = useState<
		ReconciliationEvent[]
	>([]);
	const reconciliationEventIdRef = useRef(0);
	const [agentConfigVersion, setAgentConfigVersion] = useState(0);
	const [agentStateVersion, setAgentStateVersion] = useState(0);
	const [needsOnboarding, setNeedsOnboarding] = useState(false);
	const [wizardState, setWizardState] = useState<WizardStateData | null>(null);
	const [sideQuestion, setSideQuestion] = useState<{
		question: string;
		response: string;
		loading: boolean;
	} | null>(null);
	const [serverLogs, setServerLogs] = useState<LogEntry[]>([]);
	const [storyTokenCosts, setStoryTokenCosts] = useState<Map<string, number>>(
		new Map(),
	);
	const [statusEvents, setStatusEvents] = useState<
		Array<{ receivedAt: string; event: StatusEvent }>
	>([]);

	useEffect(() => {
		const ws = new ChatWebSocket();
		wsRef.current = ws;

		ws.connect({
			onToken: (content) => {
				setStreamingContent((prev: string) => prev + content);
			},
			onThinkingToken: (content) => {
				setStreamingThinking((prev: string) => prev + content);
			},
			onUpdate: (history) => {
				setMessages(history);
				setStreamingContent("");
				setStreamingThinking("");
				const last = history[history.length - 1];
				if (last?.role === "assistant" && !last.tool_calls) {
					setLoading(false);
					setActivityStatus(null);
					if (queuedMessagesRef.current.length > 0) {
						const batch = queuedMessagesRef.current.map((item) => item.text);
						queuedMessagesRef.current = [];
						setQueuedMessages([]);
						setPendingAutoSendBatch(batch);
					}
				}
			},
			onSessionId: (sessionId) => {
				setClaudeSessionId(sessionId);
			},
			onError: (message) => {
				console.error("WebSocket error:", message);
				setLoading(false);
				setActivityStatus(null);
				const markdownMessage = message.replace(
					/(https?:\/\/[^\s]+)/g,
					"[$1]($1)",
				);
				setMessages((prev) => [
					...prev,
					{ role: "assistant", content: markdownMessage },
				]);
				if (queuedMessagesRef.current.length > 0) {
					const batch = queuedMessagesRef.current.map((item) => item.text);
					queuedMessagesRef.current = [];
					setQueuedMessages([]);
					setPendingAutoSendBatch(batch);
				}
			},
			onPipelineState: (state) => {
				setPipeline(state);
				setPipelineVersion((v) => v + 1);
				const allItems = [
					...state.backlog,
					...state.current,
					...state.qa,
					...state.merge,
					...state.done,
				];
				for (const item of allItems) {
					api
						.getTokenCost(item.story_id)
						.then((cost) => {
							if (cost.total_cost_usd > 0) {
								setStoryTokenCosts((prev) => {
									const next = new Map(prev);
									next.set(item.story_id, cost.total_cost_usd);
									return next;
								});
							}
						})
						.catch(() => {
							// Silently ignore — cost data may not exist yet.
						});
				}
			},
			onPermissionRequest: (requestId, toolName, toolInput) => {
				setPermissionQueue((prev) => [
					...prev,
					{ requestId, toolName, toolInput },
				]);
			},
			onActivity: (toolName) => {
				setActivityStatus(formatToolActivity(toolName));
			},
			onReconciliationProgress: (storyId, status, message) => {
				if (status === "done") {
					setReconciliationActive(false);
				} else {
					setReconciliationActive(true);
					setReconciliationEvents((prev) => {
						const id = String(reconciliationEventIdRef.current++);
						const next = [...prev, { id, storyId, status, message }];
						// Keep only the last 8 events to avoid the banner growing too tall.
						return next.slice(-8);
					});
				}
			},
			onAgentConfigChanged: () => {
				setAgentConfigVersion((v) => v + 1);
			},
			onAgentStateChanged: () => {
				setAgentStateVersion((v) => v + 1);
			},
			onOnboardingStatus: (onboarding: boolean) => {
				setNeedsOnboarding(onboarding);
			},
			onWizardState: (state: WizardStateData) => {
				setWizardState(state);
			},
			onSideQuestionToken: (content) => {
				setSideQuestion((prev) =>
					prev ? { ...prev, response: prev.response + content } : prev,
				);
			},
			onSideQuestionDone: (response) => {
				setSideQuestion((prev) =>
					prev ? { ...prev, response, loading: false } : prev,
				);
			},
			onLogEntry: (timestamp, level, message) => {
				setServerLogs((prev) => [...prev, { timestamp, level, message }]);
			},
			onStatusUpdate: (event) => {
				// Preserve the structured event and receive timestamp so future stories
				// can render per-type icons, banners, or filters without format changes.
				setStatusEvents((prev) => [
					...prev,
					{ receivedAt: new Date().toISOString(), event },
				]);
			},
			onConnected: () => {
				setWsConnected(true);
			},
		});

		return () => {
			ws.close();
			wsRef.current = null;
		};
	}, []);

	return {
		wsRef,
		wsConnected,
		streamingContent,
		setStreamingContent,
		streamingThinking,
		setStreamingThinking,
		activityStatus,
		setActivityStatus,
		permissionQueue,
		setPermissionQueue,
		pipeline,
		pipelineVersion,
		reconciliationActive,
		reconciliationEvents,
		agentConfigVersion,
		agentStateVersion,
		needsOnboarding,
		setNeedsOnboarding,
		wizardState,
		setWizardState,
		sideQuestion,
		setSideQuestion,
		serverLogs,
		storyTokenCosts,
		statusEvents,
	};
}