@assistkick/create 1.9.0 → 1.11.0

package/templates/assistkick-product-system/packages/frontend/src/hooks/use_chat_stream.ts

@@ -0,0 +1,826 @@
+/**
+ * useChatStream — React hook for real-time Chat v2 message streaming.
+ *
+ * Manages a WebSocket connection to /api/chat, sends user messages,
+ * and processes stream-json events as they arrive from the backend.
+ * Text tokens are accumulated in real time for a streaming typing effect.
+ *
+ * Also handles permission requests from the backend: when Claude CLI needs
+ * permission to use a tool, the backend sends a permission_request message.
+ * The hook exposes the pending request and a method to respond.
+ *
+ * Supports message queuing: when streaming is active, consumers can register
+ * an onStreamEnd callback to auto-drain queued messages. The sendMergedMessages
+ * method adds each message as a separate user bubble but sends the merged
+ * text as a single CLI invocation.
+ *
+ * Usage:
+ *   const { connect, disconnect, sendMessage, sendMergedMessages, stopStream,
+ *           setOnStreamEnd, messages, streaming, connected, error,
+ *           cancelledUserMessage, clearCancelledUserMessage,
+ *           clearMessages, pendingPermission, respondToPermission } = useChatStream();
+ */
+
+import {useRef, useState, useCallback} from 'react';
+import type {PermissionDecision} from '../components/PermissionDialog';
+import {apiClient} from '../api/client';
+import {extractContextUsage} from '../lib/context_usage_helpers';
+import {markOrCreateStreamingAssistant} from '../lib/chat_message_helpers';
+import {useChatSessionStore} from '../stores/useChatSessionStore';
+
+// --- Server → Client message types ---
+
+interface StreamStartMessage {
+    type: 'stream_start';
+    claudeSessionId: string;
+}
+
+interface StreamEventMessage {
+    type: 'stream_event';
+    event: Record<string, unknown>;
+}
+
+interface StreamEndMessage {
+    type: 'stream_end';
+    exitCode: number | null;
+    claudeSessionId: string;
+    error?: string;
+}
+
+interface StreamCancelledMessage {
+    type: 'stream_cancelled';
+    claudeSessionId: string;
+}
+
+interface StreamResumedMessage {
+    type: 'stream_resumed';
+    claudeSessionId: string;
+}
+
+interface StreamErrorMessage {
+    type: 'stream_error';
+    error: string;
+}
+
+interface PermissionRequestMessage {
+    type: 'permission_request';
+    requestId: string;
+    toolName: string;
+    input: Record<string, unknown>;
+}
+
+interface SystemPromptMessage {
+    type: 'system_prompt';
+    prompt: string;
+}
+
+interface SessionRenamedMessage {
+    type: 'session_renamed';
+    sessionId: string;
+    name: string;
+}
+
+interface ErrorMessage {
+    type: 'error';
+    message: string;
+}
+
+type ChatServerMessage =
+    | StreamStartMessage
+    | StreamEventMessage
+    | StreamEndMessage
+    | StreamCancelledMessage
+    | StreamResumedMessage
+    | StreamErrorMessage
+    | PermissionRequestMessage
+    | SystemPromptMessage
+    | SessionRenamedMessage
+    | ErrorMessage;
+
+// --- Client → Server message types ---
+
+interface SendMessagePayload {
+    type: 'send_message';
+    message: string;
+    projectId: string;
+    /** The chat session ID (csess_xxx) for message persistence. */
+    sessionId?: string;
+    claudeSessionId: string;
+    isNewSession: boolean;
+    permissionMode?: 'skip' | 'allowed_tools';
+    allowedTools?: string[];
+    /** Relative workspace paths of attached files to include in the prompt */
+    attachments?: string[];
+    /** Attachment content blocks for DB persistence alongside the user message */
+    attachmentBlocks?: Array<{ type: 'image' | 'file'; path: string; name: string; mimeType: string }>;
+}
+
+interface PermissionResponsePayload {
+    type: 'permission_response';
+    requestId: string;
+    projectId: string;
+    decision: PermissionDecision;
+}
+
+interface StopStreamPayload {
+    type: 'stop_stream';
+    claudeSessionId?: string;
+}
+
+// --- Chat message model ---
+
+export interface TextBlock {
+    type: 'text';
+    text: string;
+}
+
+export interface ToolUseBlock {
+    type: 'tool_use';
+    id: string;
+    name: string;
+    input: Record<string, unknown>;
+}
+
+export interface ToolResultBlock {
+    type: 'tool_result';
+    toolUseId: string;
+    content: string;
+    isError: boolean;
+}
+
+export interface ImageBlock {
+    type: 'image';
+    path: string;
+    name: string;
+    mimeType: string;
+}
+
+export interface FileBlock {
+    type: 'file';
+    path: string;
+    name: string;
+    mimeType: string;
+}
+
+export type ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock | ImageBlock | FileBlock;
+
+export interface ChatMessage {
+    id: string;
+    role: 'user' | 'assistant';
+    content: ContentBlock[];
+    isStreaming: boolean;
+}
+
+export interface PendingPermission {
+    requestId: string;
+    toolName: string;
+    input: Record<string, unknown>;
+}
+
+export interface ContextUsage {
+    inputTokens: number;
+    outputTokens: number;
+    cacheCreationTokens: number;
+    cacheReadTokens: number;
+    totalTokens: number;
+    /** Context window limit from the model (null until result event provides it). */
+    contextWindow: number | null;
+}
+
+export interface UseChatStreamOptions {
+    projectId: string;
+    /** The chat session ID (csess_xxx) for message persistence. */
+    sessionId?: string;
+    claudeSessionId: string;
+    isNewSession: boolean;
+    permissionMode?: 'skip' | 'allowed_tools';
+    allowedTools?: string[];
+    /** Relative workspace paths of attached files to include in the prompt */
+    attachments?: string[];
+    /** Image/file content blocks for persisting attachment metadata in the user message */
+    attachmentBlocks?: (ImageBlock | FileBlock)[];
+}
+
+export interface UseChatStreamReturn {
+    connect: () => void;
+    disconnect: () => void;
+    sendMessage: (message: string, options: UseChatStreamOptions) => void;
+    sendMergedMessages: (texts: string[], options: UseChatStreamOptions) => void;
+    stopStream: (claudeSessionId?: string) => void;
+    setOnStreamEnd: (cb: (() => void) | null) => void;
+    /** Load persisted messages from the database for a session. Returns the number of messages loaded. */
+    loadMessages: (sessionId: string) => Promise<number>;
+    /** Attempt to re-attach to an in-flight stream after page reload. */
+    attachSession: (claudeSessionId: string) => void;
+    messages: ChatMessage[];
+    streaming: boolean;
+    connected: boolean;
+    error: string | null;
+    /** Text from the user message that was removed after a stop/cancel, for prepopulating the input field */
+    cancelledUserMessage: string | null;
+    clearCancelledUserMessage: () => void;
+    clearMessages: () => void;
+    /** Restore a previously cached set of messages (e.g. when switching back to a session). */
+    restoreMessages: (msgs: ChatMessage[]) => void;
+    pendingPermission: PendingPermission | null;
+    respondToPermission: (requestId: string, projectId: string, decision: PermissionDecision) => void;
+    /** Current context window usage (updated from stream events) */
+    contextUsage: ContextUsage | null;
+    /** Initialize context usage from persisted session data */
+    initContextUsage: (usage: ContextUsage | null) => void;
+    /** The composed system prompt for the current session */
+    systemPrompt: string | null;
+    /** Initialize system prompt from persisted session data */
+    initSystemPrompt: (prompt: string | null) => void;
+}
+
+let messageIdCounter = 0;
+const nextMessageId = (): string => `msg_${++messageIdCounter}_${Date.now()}`;
+
+/**
+ * Build the WebSocket URL for the chat endpoint.
+ * Uses wss:// for HTTPS, ws:// for HTTP, relative to the current host.
+ */
+const buildWsUrl = (): string => {
+    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+    return `${protocol}//${window.location.host}/api/chat`;
+};
+
+/**
+ * Extract content blocks from a stream-json assistant event.
+ * Returns text and tool_use blocks found in the message content array.
+ */
+const extractContentBlocks = (event: Record<string, unknown>): ContentBlock[] => {
+    const blocks: ContentBlock[] = [];
+    const message = event.message as Record<string, unknown> | undefined;
+    if (!message) return blocks;
+
+    const content = message.content as Array<Record<string, unknown>> | undefined;
+    if (!Array.isArray(content)) return blocks;
+
+    for (const block of content) {
+        if (block.type === 'text' && typeof block.text === 'string') {
+            blocks.push({type: 'text', text: block.text});
+        } else if (block.type === 'tool_use' && typeof block.name === 'string') {
+            blocks.push({
+                type: 'tool_use',
+                id: (block.id as string) || '',
+                name: block.name,
+                input: (block.input as Record<string, unknown>) || {},
+            });
+        } else if (block.type === 'tool_result' && typeof block.tool_use_id === 'string') {
+            // tool_result blocks carry the output of a tool execution.
+            // content can be a string or an array of {type:'text', text:...} blocks.
+            let resultText = '';
+            if (typeof block.content === 'string') {
+                resultText = block.content;
+            } else if (Array.isArray(block.content)) {
+                resultText = (block.content as Array<Record<string, unknown>>)
+                    .filter(c => c.type === 'text' && typeof c.text === 'string')
+                    .map(c => c.text as string)
+                    .join('\n');
+            }
+            blocks.push({
+                type: 'tool_result',
+                toolUseId: block.tool_use_id as string,
+                content: resultText,
+                isError: block.is_error === true,
+            });
+        }
+    }
+
+    return blocks;
+};
+
+/**
+ * Normalize a raw content block from the database (snake_case keys from Claude API)
+ * into the frontend's camelCase ContentBlock format.
+ */
+const normalizeContentBlock = (block: Record<string, unknown>): ContentBlock | null => {
+    if (block.type === 'text' && typeof block.text === 'string') {
+        return { type: 'text', text: block.text };
+    }
+    if (block.type === 'tool_use' && typeof block.name === 'string') {
+        return {
+            type: 'tool_use',
+            id: (block.id as string) || '',
+            name: block.name as string,
+            input: (block.input as Record<string, unknown>) || {},
+        };
+    }
+    if (block.type === 'tool_result') {
+        const toolUseId = (block.toolUseId ?? block.tool_use_id) as string | undefined;
+        if (!toolUseId) return null;
+        let resultText = '';
+        if (typeof block.content === 'string') {
+            resultText = block.content;
+        } else if (Array.isArray(block.content)) {
+            resultText = (block.content as Array<Record<string, unknown>>)
+                .filter(c => c.type === 'text' && typeof c.text === 'string')
+                .map(c => c.text as string)
+                .join('\n');
+        }
+        return {
+            type: 'tool_result',
+            toolUseId,
+            content: resultText,
+            isError: block.isError === true || block.is_error === true,
+        };
+    }
+    if (block.type === 'image' && typeof block.path === 'string') {
+        return {
+            type: 'image',
+            path: block.path as string,
+            name: (block.name as string) || '',
+            mimeType: (block.mimeType as string) || 'image/png',
+        };
+    }
+    if (block.type === 'file' && typeof block.path === 'string') {
+        return {
+            type: 'file',
+            path: block.path as string,
+            name: (block.name as string) || '',
+            mimeType: (block.mimeType as string) || 'application/octet-stream',
+        };
+    }
+    return null;
+};
+
+export const useChatStream = (): UseChatStreamReturn => {
+    const wsRef = useRef<WebSocket | null>(null);
+    const streamEndCallbackRef = useRef<(() => void) | null>(null);
+    const [connected, setConnected] = useState(false);
+    const [streaming, setStreaming] = useState(false);
+    const [messages, setMessages] = useState<ChatMessage[]>([]);
+    const [error, setError] = useState<string | null>(null);
+    const [pendingPermission, setPendingPermission] = useState<PendingPermission | null>(null);
+    const [cancelledUserMessage, setCancelledUserMessage] = useState<string | null>(null);
+    const [contextUsage, setContextUsage] = useState<ContextUsage | null>(null);
+    const [systemPrompt, setSystemPrompt] = useState<string | null>(null);
+
+    const connect = useCallback(() => {
+        if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) return;
+
+        const url = buildWsUrl();
+        const ws = new WebSocket(url);
+        wsRef.current = ws;
+
+        ws.onopen = () => {
+            setConnected(true);
+            setError(null);
+        };
+
+        ws.onclose = () => {
+            setConnected(false);
+            setStreaming(false);
+            wsRef.current = null;
+        };
+
+        ws.onerror = () => {
+            setError('WebSocket connection error');
+            setConnected(false);
+            setStreaming(false);
+        };
+
+        ws.onmessage = (event: MessageEvent) => {
+            let msg: ChatServerMessage;
+            try {
+                msg = JSON.parse(event.data);
+            } catch {
+                return;
+            }
+
+            switch (msg.type) {
+                case 'stream_start':
+                    setStreaming(true);
+                    setError(null);
+                    // Create an empty assistant message that will accumulate content
+                    setMessages(prev => [
+                        ...prev,
+                        {id: nextMessageId(), role: 'assistant', content: [], isStreaming: true},
+                    ]);
+                    break;
+
+                case 'stream_event':
+                    handleStreamEvent(msg.event);
+                    break;
+
+                case 'stream_end':
+                    setStreaming(false);
+                    // Surface error details from non-zero exit codes
+                    if (msg.error) {
+                        setError(msg.error);
+                    }
+                    // Mark the last assistant message as no longer streaming
+                    setMessages(prev => {
+                        const updated = [...prev];
+                        for (let i = updated.length - 1; i >= 0; i--) {
+                            if (updated[i].role === 'assistant' && updated[i].isStreaming) {
+                                updated[i] = {...updated[i], isStreaming: false};
+                                break;
+                            }
+                        }
+                        return updated;
+                    });
+                    // Notify consumers (e.g. message queue) that streaming has ended
+                    streamEndCallbackRef.current?.();
+                    break;
+
+                case 'stream_cancelled':
+                    setStreaming(false);
+                    // Discard the partial assistant response and the triggering user message.
+                    // Prepopulate the cancelled user message text for the input field.
+                    setMessages(prev => {
+                        const updated = [...prev];
+
+                        // Remove the last streaming assistant message (partial response)
+                        for (let i = updated.length - 1; i >= 0; i--) {
+                            if (updated[i].role === 'assistant' && updated[i].isStreaming) {
+                                updated.splice(i, 1);
+                                break;
+                            }
+                        }
+
+                        // Remove the last user message and capture its text
+                        for (let i = updated.length - 1; i >= 0; i--) {
+                            if (updated[i].role === 'user') {
+                                const userContent = updated[i].content;
+                                const textBlock = userContent.find(b => b.type === 'text') as TextBlock | undefined;
+                                if (textBlock) {
+                                    setCancelledUserMessage(textBlock.text);
+                                }
+                                updated.splice(i, 1);
+                                break;
+                            }
+                        }
+
+                        return updated;
+                    });
+                    break;
+
+                case 'stream_resumed':
+                    // Re-attached to an in-flight stream after page reload or session switch.
+                    // Mark the last assistant message as streaming so new events update it.
+                    // If no assistant message exists (e.g. messages not yet persisted to DB),
+                    // create an empty streaming assistant message for catch-up content.
+                    setStreaming(true);
+                    setError(null);
+                    setMessages(prev => markOrCreateStreamingAssistant(prev, nextMessageId()));
+                    break;
+
+                case 'stream_error':
+                    setStreaming(false);
+                    setError(msg.error);
+                    // Mark streaming message as done
+                    setMessages(prev => {
+                        const updated = [...prev];
+                        for (let i = updated.length - 1; i >= 0; i--) {
+                            if (updated[i].role === 'assistant' && updated[i].isStreaming) {
+                                updated[i] = {...updated[i], isStreaming: false};
+                                break;
+                            }
+                        }
+                        return updated;
+                    });
+                    break;
+
+                case 'permission_request':
+                    setPendingPermission({
+                        requestId: msg.requestId,
+                        toolName: msg.toolName,
+                        input: msg.input,
+                    });
+                    break;
+
+                case 'system_prompt':
+                    setSystemPrompt(msg.prompt);
+                    break;
+
+                case 'session_renamed':
+                    // Update the session name in the sidebar store in real time
+                    useChatSessionStore.getState().updateSessionName(msg.sessionId, msg.name);
+                    break;
+
+                case 'error':
+                    setError(msg.message);
+                    break;
+            }
+        };
+    }, []);
+
+    const handleStreamEvent = (event: Record<string, unknown>) => {
+        // Extract per-call context usage from assistant events only.
+        // The result event's usage is CUMULATIVE across all turns — not the context window size.
+        if (event.type === 'assistant') {
+            const message = event.message as Record<string, unknown> | undefined;
+            const usage = message?.usage as Record<string, unknown> | undefined;
+            if (usage) {
+                setContextUsage(prev => ({
+                    ...extractContextUsage(usage),
+                    contextWindow: prev?.contextWindow ?? null,
+                }));
+            }
+
+            const blocks = extractContentBlocks(event);
+            if (blocks.length === 0) return;
+
+            setMessages(prev => {
+                const updated = [...prev];
+                // Find the last streaming assistant message
+                for (let i = updated.length - 1; i >= 0; i--) {
+                    if (updated[i].role === 'assistant' && updated[i].isStreaming) {
+                        // Each assistant event is cumulative within its turn, but a new
+                        // turn (after tool use) starts fresh. Detect turn boundaries by
+                        // comparing tool_use IDs: if existing content has tool_use blocks
+                        // not present in the new blocks, we've moved to a new turn.
+                        const existing = updated[i].content;
+
+                        const existingToolUseIds = new Set(
+                            existing
+                                .filter((b): b is ToolUseBlock => b.type === 'tool_use')
+                                .map(b => b.id),
+                        );
+                        const newToolUseIds = new Set(
+                            blocks
+                                .filter((b): b is ToolUseBlock => b.type === 'tool_use')
+                                .map(b => b.id),
+                        );
+
+                        const isNewTurn = existingToolUseIds.size > 0 &&
+                            ![...existingToolUseIds].some(id => newToolUseIds.has(id));
+
+                        if (isNewTurn) {
+                            // New turn: preserve all existing content, append new blocks
+                            updated[i] = {...updated[i], content: [...existing, ...blocks]};
+                        } else {
+                            // Same turn: replace with cumulative snapshot, preserving tool_result blocks
+                            const existingToolResults = existing.filter(
+                                (b): b is ToolResultBlock => b.type === 'tool_result',
+                            );
+                            updated[i] = {...updated[i], content: [...blocks, ...existingToolResults]};
+                        }
+                        break;
+                    }
+                }
+                return updated;
+            });
+            return;
+        }
+
+        // Process tool result events — merge tool_result blocks into the last
+        // streaming assistant message so ToolUseCard can show output.
+        if (event.type === 'tool_result' || event.type === 'human' || event.type === 'user') {
+            const blocks = extractContentBlocks(event);
+            const resultBlocks = blocks.filter(
+                (b): b is ToolResultBlock => b.type === 'tool_result',
+            );
+            if (resultBlocks.length === 0) return;
+
+            setMessages(prev => {
+                const updated = [...prev];
+                for (let i = updated.length - 1; i >= 0; i--) {
+                    if (updated[i].role === 'assistant' && updated[i].isStreaming) {
+                        // Append tool_result blocks, deduplicating by toolUseId
+                        const existingToolResultIds = new Set(
+                            updated[i].content
+                                .filter((b): b is ToolResultBlock => b.type === 'tool_result')
+                                .map(b => b.toolUseId),
+                        );
+                        const newResults = resultBlocks.filter(
+                            b => !existingToolResultIds.has(b.toolUseId),
+                        );
+                        if (newResults.length > 0) {
+                            updated[i] = {
+                                ...updated[i],
+                                content: [...updated[i].content, ...newResults],
+                            };
+                        }
+                        break;
+                    }
+                }
+                return updated;
+            });
+        }
+
+        // Extract contextWindow from result event's modelUsage
+        if (event.type === 'result') {
+            const modelUsage = event.modelUsage as Record<string, Record<string, unknown>> | undefined;
+            if (modelUsage) {
+                for (const info of Object.values(modelUsage)) {
+                    if (typeof info.contextWindow === 'number') {
+                        setContextUsage(prev => prev ? { ...prev, contextWindow: info.contextWindow as number } : prev);
+                        break;
+                    }
+                }
+            }
+        }
+    };
+
+    const disconnect = useCallback(() => {
+        if (wsRef.current) {
+            wsRef.current.close();
+            wsRef.current = null;
+        }
+        setConnected(false);
+        setStreaming(false);
+    }, []);
+
+    const sendMessage = useCallback((message: string, options: UseChatStreamOptions) => {
+        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+            setError('Not connected');
+            return;
+        }
+
+        // Build user message content blocks: text + any attachment blocks
+        const userContent: ContentBlock[] = [{type: 'text', text: message}];
+        if (options.attachmentBlocks && options.attachmentBlocks.length > 0) {
+            userContent.push(...options.attachmentBlocks);
+        }
+
+        // Add user message to the list
+        setMessages(prev => [
+            ...prev,
+            {id: nextMessageId(), role: 'user', content: userContent, isStreaming: false},
+        ]);
+
+        const payload: SendMessagePayload = {
+            type: 'send_message',
+            message,
+            projectId: options.projectId,
+            sessionId: options.sessionId,
+            claudeSessionId: options.claudeSessionId,
+            isNewSession: options.isNewSession,
+            permissionMode: options.permissionMode,
+            allowedTools: options.allowedTools,
+            attachments: options.attachments,
+            attachmentBlocks: options.attachmentBlocks,
+        };
+
+        wsRef.current.send(JSON.stringify(payload));
+    }, []);
+
+    const respondToPermission = useCallback((requestId: string, projectId: string, decision: PermissionDecision) => {
+        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+            setError('Not connected');
+            return;
+        }
+
+        const payload: PermissionResponsePayload = {
+            type: 'permission_response',
+            requestId,
+            projectId,
+            decision,
+        };
+
+        wsRef.current.send(JSON.stringify(payload));
+        setPendingPermission(null);
+    }, []);
+
+    /**
+     * Send multiple queued messages as a single CLI invocation.
+     * Each text appears as a separate user bubble in the chat, but they are
+     * concatenated into one prompt for the 'claude -p --resume' call.
+     */
+    const sendMergedMessages = useCallback((texts: string[], options: UseChatStreamOptions) => {
+        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+            setError('Not connected');
+            return;
+        }
+
+        if (texts.length === 0) return;
+
+        // Add each queued message as a separate user bubble
+        setMessages(prev => [
+            ...prev,
+            ...texts.map(text => ({
+                id: nextMessageId(),
+                role: 'user' as const,
+                content: [{type: 'text' as const, text}],
+                isStreaming: false,
+            })),
+        ]);
+
+        // Merge all texts into a single prompt
+        const merged = texts.join('\n\n');
+
+        const payload: SendMessagePayload = {
+            type: 'send_message',
+            message: merged,
+            projectId: options.projectId,
+            sessionId: options.sessionId,
+            claudeSessionId: options.claudeSessionId,
+            isNewSession: false, // queued messages always resume an existing session
+            permissionMode: options.permissionMode,
+            allowedTools: options.allowedTools,
+        };
+
+        wsRef.current.send(JSON.stringify(payload));
+    }, []);
+
+    const stopStream = useCallback((claudeSessionId?: string) => {
+        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
+
+        const payload: StopStreamPayload = {type: 'stop_stream', claudeSessionId};
+        wsRef.current.send(JSON.stringify(payload));
+    }, []);
+
+    const clearCancelledUserMessage = useCallback(() => {
+        setCancelledUserMessage(null);
+    }, []);
+
+    const setOnStreamEnd = useCallback((cb: (() => void) | null) => {
+        streamEndCallbackRef.current = cb;
+    }, []);
+
+    const clearMessages = useCallback(() => {
+        setMessages([]);
+        setContextUsage(null);
+    }, []);
+
+    const restoreMessages = useCallback((msgs: ChatMessage[]) => {
+        setMessages(msgs);
+    }, []);
+
+    const initContextUsage = useCallback((usage: ContextUsage | null) => {
+        setContextUsage(usage);
+    }, []);
+
+    const initSystemPrompt = useCallback((prompt: string | null) => {
+        setSystemPrompt(prompt);
+    }, []);
+
+    /**
+     * Attempt to re-attach to an in-flight stream after page reload.
+     * If the backend has an active CLI process for this session, it will send
+     * stream_resumed + catch-up events. If not, nothing happens and the loaded
+     * DB messages are the final state.
+     */
+    const attachSession = useCallback((claudeSessionId: string) => {
+        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
+
+        wsRef.current.send(JSON.stringify({
+            type: 'attach_session',
+            claudeSessionId,
+        }));
+    }, []);
+
+    /**
+     * Load persisted messages from the database for a session.
+     * Replaces the current in-memory messages with the stored history.
+     */
+    const loadMessages = useCallback(async (sessionId: string): Promise<number> => {
+        try {
+            const data = await apiClient.listChatMessages(sessionId);
+            const rows = data.messages as Array<{
+                id: string;
+                role: string;
+                content: string;
+                orderIndex: number;
+                createdAt: string;
+            }>;
+            const loaded: ChatMessage[] = rows.map((row) => {
+                const rawBlocks = JSON.parse(row.content) as Array<Record<string, unknown>>;
+                const content = rawBlocks
+                    .map(normalizeContentBlock)
+                    .filter((b): b is ContentBlock => b !== null);
+                return {
+                    id: row.id,
+                    role: row.role as 'user' | 'assistant',
+                    content,
+                    isStreaming: false,
+                };
+            });
+            setMessages(loaded);
+            return loaded.length;
+        } catch {
+            // If loading fails (e.g. no messages yet), start with empty messages
+            setMessages([]);
+            return 0;
+        }
+    }, []);
+
+    return {
+        connect,
+        disconnect,
+        sendMessage,
+        sendMergedMessages,
+        stopStream,
+        setOnStreamEnd,
+        loadMessages,
+        attachSession,
+        messages,
+        streaming,
+        connected,
+        error,
+        cancelledUserMessage,
+        clearCancelledUserMessage,
+        clearMessages,
+        restoreMessages,
+        pendingPermission,
+        respondToPermission,
+        contextUsage,
+        initContextUsage,
+        systemPrompt,
+        initSystemPrompt,
+    };
+};