frontend/src/components/SessionChat.tsx

/**
 * Per-session chat component.
 *
 * Each session renders its own SessionChat. The hook (useAgentChat) always
 * runs — processing events — but only the active session renders visible
 * UI (MessageList + ChatInput).
 */
import { useCallback, useEffect, useState } from 'react';
import { useAgentChat } from '@/hooks/useAgentChat';
import { useAgentStore } from '@/store/agentStore';
import { useSessionStore } from '@/store/sessionStore';
import MessageList from '@/components/Chat/MessageList';
import ChatInput from '@/components/Chat/ChatInput';
import ExpiredBanner from '@/components/Chat/ExpiredBanner';
import BillingBanner from '@/components/Chat/BillingBanner';
import ChatErrorBanner from '@/components/Chat/ChatErrorBanner';
import { useUserQuota } from '@/hooks/useUserQuota';
import { isPremiumPath } from '@/utils/model';
import { apiFetch } from '@/utils/api';
import { logger } from '@/utils/logger';

interface SessionChatProps {
  sessionId: string;
  isActive: boolean;
  onSessionDead: (sessionId: string) => void;
}

export default function SessionChat({ sessionId, isActive, onSessionDead }: SessionChatProps) {
  const { isConnected, isProcessing, activityStatus, updateSession } = useAgentStore();
  const { updateSessionTitle, sessions, setSessionProcessing } = useSessionStore();
  const sessionMeta = sessions.find((s) => s.id === sessionId);
  const isExpired = sessionMeta?.expired === true;
  const [chatError, setChatError] = useState<string | null>(null);

  const {
    messages,
    sendMessage,
    stop,
    status,
    undoLastTurn,
    editAndRegenerate,
    approveTools,
    refreshMessages,
  } = useAgentChat({
    sessionId,
    isActive,
    // A backgrounded session that the backend reports mid-turn still needs to
    // mount its live subscription; idle backgrounded sessions do not (that's
    // what stops app load from reactivating every historical runtime).
    isProcessing: sessionMeta?.isProcessing ?? false,
    onReady: () => logger.log(`Session ${sessionId} ready`),
    onError: (error) => {
      logger.error(`Session ${sessionId} error:`, error);
      setChatError(error);
    },
    onSessionDead,
  });

  // When this session becomes active, restore its per-session state to the
  // global flat fields. The per-session state map is kept up-to-date by
  // side-channel callbacks even while the session is in the background.
  useEffect(() => {
    if (isActive) {
      useAgentStore.getState().switchActiveSession(sessionId);
      useAgentStore.getState().setConnected(true);
    }
  }, [isActive, sessionId]);

  // Re-sync state when the browser tab regains focus (Chrome throttles
  // timers in background tabs which can stall the AI SDK's update flushing).
  // Fires for ALL sessions so background sessions also recover after sleep.
  useEffect(() => {
    const onVisible = () => {
      if (document.visibilityState === 'visible' && isActive) {
        useAgentStore.getState().switchActiveSession(sessionId);
      }
    };
    document.addEventListener('visibilitychange', onVisible);
    return () => document.removeEventListener('visibilitychange', onVisible);
  }, [isActive, sessionId]);

  // Wrap stop to show cancelled shimmer
  const handleStop = useCallback(() => {
    stop();
    updateSession(sessionId, { activityStatus: { type: 'cancelled' } });
  }, [stop, updateSession, sessionId]);

  // SDK status is the ground truth — if it's streaming/submitted, agent is busy
  const sdkBusy = status === 'streaming' || status === 'submitted';
  const busy = isProcessing || sdkBusy;
  const { quota, refresh: refreshQuota } = useUserQuota({ enabled: isActive });

  // Whether this session's premium usage is being billed to the user's own HF
  // account (past the subsidized daily allowance). Re-read after each turn,
  // since the backend flips it at submit time. Only premium-model sessions can
  // ever be user-billed, so skip the fetch for free models.
  const [premiumBilled, setPremiumBilled] = useState(false);
  const [premiumQuotaCounted, setPremiumQuotaCounted] = useState(false);
  const onPremiumModel = isPremiumPath(sessionMeta?.model ?? undefined);
  useEffect(() => {
    if (!isActive || !onPremiumModel) {
      setPremiumBilled(false);
      setPremiumQuotaCounted(false);
      return;
    }
    if (busy) return;
    let cancelled = false;
    apiFetch(`/api/session/${sessionId}`)
      .then((r) => (r.ok ? r.json() : null))
      .then((d) => {
        if (!cancelled && d) {
          setPremiumBilled(Boolean(d.premium_user_billed));
          setPremiumQuotaCounted(Boolean(d.premium_quota_counted));
        }
      })
      .catch(() => {});
    return () => {
      cancelled = true;
    };
  }, [busy, isActive, onPremiumModel, sessionId]);

  const sessionPremiumBilled = premiumBilled || Boolean(sessionMeta?.premiumUserBilled);
  const sessionPremiumQuotaCounted =
    premiumQuotaCounted || Boolean(sessionMeta?.premiumQuotaCounted);
  const premiumBillingNotice =
    sessionPremiumBilled ||
    (isActive &&
      onPremiumModel &&
      quota?.premiumRemaining === 0 &&
      !sessionPremiumQuotaCounted);

  const handleSendMessage = useCallback(
    async (text: string) => {
      if (!text.trim() || busy) return;

      setChatError(null);
      updateSession(sessionId, { isProcessing: true, activityStatus: { type: 'thinking' } });
      setSessionProcessing(sessionId, true);
      sendMessage({ text: text.trim(), metadata: { createdAt: new Date().toISOString() } });

      // Auto-title the session from the first user message
      const isFirstMessage = messages.filter((m) => m.role === 'user').length === 0;
      if (isFirstMessage) {
        apiFetch('/api/title', {
          method: 'POST',
          body: JSON.stringify({ session_id: sessionId, text: text.trim() }),
        })
          .then((res) => res.json())
          .then((data) => {
            if (data.title) updateSessionTitle(sessionId, data.title);
          })
          .catch(() => {
            const raw = text.trim();
            updateSessionTitle(sessionId, raw.length > 40 ? raw.slice(0, 40) + '\u2026' : raw);
          });
      }
    },
    [sessionId, sendMessage, messages, updateSessionTitle, busy, updateSession, setSessionProcessing],
  );

  // Don't render UI for background sessions — hooks still run
  if (!isActive) return null;

  return (
    <>
      <MessageList
        messages={messages}
        isProcessing={busy}
        sessionId={sessionId}
        approveTools={approveTools}
        onUndoLastTurn={undoLastTurn}
        onEditAndRegenerate={editAndRegenerate}
      />
      {isExpired ? (
        <ExpiredBanner sessionId={sessionId} />
      ) : (
        <>
          {premiumBillingNotice && <BillingBanner />}
          {chatError && (
            <ChatErrorBanner
              error={chatError}
              sessionId={sessionId}
              model={sessionMeta?.model}
              onDismiss={() => setChatError(null)}
            />
          )}
          <ChatInput
            sessionId={sessionId}
            initialModelPath={sessionMeta?.model}
            onSend={handleSendMessage}
            onStop={handleStop}
            onDatasetUploaded={refreshMessages}
            isProcessing={busy}
            disabled={!isConnected || activityStatus.type === 'waiting-approval'}
            quota={quota}
            refreshQuota={refreshQuota}
            placeholder={
              activityStatus.type === 'waiting-approval'
                ? 'Approve or reject pending tools first...'
                : undefined
            }
          />
        </>
      )}
    </>
  );
}