nextcraft/apps/web/hooks/use-chat-stream.ts

'use client';

import { useCallback, useEffect, useRef, useState } from 'react';
import { parseSseEvents } from '../lib/sse';

const AI_SERVICE_URL =
  process.env.NEXT_PUBLIC_AI_SERVICE_URL ?? 'http://localhost:8420';

export type AgentName = 'coach' | 'tutor' | 'lab' | 'assessor' | 'proctor' | 'mentor';

export interface StreamMessage {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  agent?: AgentName;
  streaming?: boolean;
  suggestedActions?: string[];
}

interface StreamState {
  messages: StreamMessage[];
  isStreaming: boolean;
  error: string | null;
  model: string | null;
}

interface ChatStreamEvent {
  type: 'meta' | 'delta' | 'done' | 'error';
  [key: string]: unknown;
}

function decodeEvent(raw: string): ChatStreamEvent | '[DONE]' | null {
  if (raw === '[DONE]') return '[DONE]';
  try {
    const parsed = JSON.parse(raw);
    if (typeof parsed?.type === 'string') return parsed as ChatStreamEvent;
    // OpenAI-shaped chunks (id/choices) are not used by our envelope;
    // ignore anything without a type.
    return null;
  } catch {
    return null;
  }
}

export function useChatStream(agent: AgentName) {
  const [state, setState] = useState<StreamState>({
    messages: [],
    isStreaming: false,
    error: null,
    model: null,
  });
  const abortRef = useRef<AbortController | null>(null);
  // Agent-scoped sessions (A-007/D-019): switching agents starts a NEW
  // session per agent — no persona bleed across switcher flips.
  const sessionsRef = useRef<Partial<Record<AgentName, string>>>({});
  if (!sessionsRef.current[agent]) {
    const uuid =
      typeof crypto !== 'undefined' && 'randomUUID' in crypto
        ? crypto.randomUUID()
        : String(Date.now());
    sessionsRef.current[agent] = `${agent}-${uuid}`;
  }

  // Idempotent abort + cleanup on unmount or agent switch (Strict Mode safe)
  const abort = useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
  }, []);

  useEffect(() => {
    return () => {
      abortRef.current?.abort();
      abortRef.current = null;
    };
  }, []);

  const send = useCallback(
    async (text: string) => {
      const trimmed = text.trim();
      if (!trimmed || abortRef.current) return;

      const userMessage: StreamMessage = {
        id: `user-${Date.now()}`,
        role: 'user',
        content: trimmed,
      };
      const assistantId = `assistant-${Date.now()}`;

      setState((prev) => ({
        ...prev,
        messages: [...prev.messages, userMessage],
        isStreaming: true,
        error: null,
      }));

      const controller = new AbortController();
      abortRef.current = controller;

      try {
        const response = await fetch(`${AI_SERVICE_URL}/v1/chat/stream`, {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({
            agent,
            session_id: sessionsRef.current[agent],
            messages: [{ role: 'user', content: trimmed }],
          }),
          signal: controller.signal,
        });

        if (!response.ok || !response.body) {
          throw new Error(`AI service unavailable (${response.status})`);
        }

        setState((prev) => ({
          ...prev,
          messages: [
            ...prev.messages,
            { id: assistantId, role: 'assistant', content: '', agent, streaming: true },
          ],
        }));

        const reader = response.body.getReader();
        const decoder = new TextDecoder();
        let buffer = '';
        let done = false;

        while (!done) {
          const { value, done: readerDone } = await reader.read();
          if (readerDone) break;
          buffer += decoder.decode(value, { stream: true });

          const { events, rest } = parseSseEvents(buffer);
          buffer = rest;

          for (const raw of events) {
            const event = decodeEvent(raw);
            if (event === null) continue;
            if (event === '[DONE]') {
              done = true;
              setState((prev) => ({
                ...prev,
                isStreaming: false,
                messages: prev.messages.map((m) =>
                  m.id === assistantId ? { ...m, streaming: false } : m,
                ),
              }));
              break;
            }
            if (event.type === 'meta') {
              setState((prev) => ({ ...prev, model: (event.model as string) ?? null }));
            } else if (event.type === 'delta') {
              const content = event.content as string;
              setState((prev) => ({
                ...prev,
                messages: prev.messages.map((m) =>
                  m.id === assistantId ? { ...m, content: m.content + content } : m,
                ),
              }));
            } else if (event.type === 'error') {
              setState((prev) => ({
                ...prev,
                error: (event.message as string) ?? 'stream error',
              }));
            }
          }
        }

        setState((prev) => ({
          ...prev,
          isStreaming: false,
          messages: prev.messages.map((m) =>
            m.id === assistantId ? { ...m, streaming: false } : m,
          ),
        }));
      } catch (err) {
        const aborted = err instanceof DOMException && err.name === 'AbortError';
        if (!aborted) {
          setState((prev) => ({
            ...prev,
            isStreaming: false,
            error: err instanceof Error ? err.message : 'connection failed',
            messages: prev.messages.map((m) =>
              m.id === assistantId ? { ...m, streaming: false } : m,
            ),
          }));
        } else {
          setState((prev) => ({
            ...prev,
            isStreaming: false,
            messages: prev.messages.map((m) =>
              m.id === assistantId ? { ...m, streaming: false } : m,
            ),
          }));
        }
      } finally {
        abortRef.current = null;
      }
    },
    [agent],
  );

  const retry = useCallback(() => {
    setState((prev) => ({ ...prev, error: null }));
    const lastUser = [...state.messages].reverse().find((m) => m.role === 'user');
    if (lastUser) void send(lastUser.content);
  }, [send, state.messages]);

  return {
    messages: state.messages,
    isStreaming: state.isStreaming,
    error: state.error,
    model: state.model,
    send,
    retry,
    abort,
  };
}