ai-mcp-web-ui/frontend-v2/lib/hooks.ts

/**
 * 自定义 React Hooks - 用于与 FastAPI 后端交互
 *
 * 注意：useChat Hook 已被移除
 * 请使用官方 @json-render/react 的 useChatUI 替代：
 *   import { useChatUI } from '@json-render/react'
 */
'use client';

import { useState, useCallback, useEffect } from 'react';
import { apiClient, type UserInfo, type UserRole } from './api-client';

// ============================================================================
// 认证 Hook
// ============================================================================
export function useAuth() {
  const [isAuthenticated, setIsAuthenticated] = useState(false);
  const [username, setUsername] = useState<string | null>(null);
  const [role, setRole] = useState<UserRole | null>(null);
  const [isLoading, setIsLoading] = useState(false);

  // 从 localStorage 恢复认证状态
  useEffect(() => {
    const storedSessionId = localStorage.getItem('session_id');
    const storedUserInfo = localStorage.getItem('userInfo');
    if (storedSessionId) {
      apiClient.setSession(storedSessionId);
      setIsAuthenticated(true);
    }
    if (storedUserInfo) {
      try {
        const userInfo: UserInfo = JSON.parse(storedUserInfo);
        setUsername(userInfo.username);
        setRole(userInfo.role);
        apiClient.setUserInfo(userInfo);
      } catch (e) {
        console.error('Failed to parse stored userInfo:', e);
      }
    }
  }, []);

  // 根据角色获取 MCP URL
  const getMcpUrl = useCallback((): string => {
    return apiClient.getMcpUrl();
  }, []);

  const login = useCallback(async (email: string, password: string) => {
    setIsLoading(true);
    try {
      console.log('[useAuth] Starting login:', email);
      const response = await apiClient.login(email, password);
      console.log('[useAuth] Login response:', response);
      // 检查 session_id 而非 success 字段
      if (response.session_id) {
        console.log('[useAuth] Login successful, setting state...');
        setIsAuthenticated(true);
        const username = response.username ?? 'Unknown';
        const role = response.role ?? 'user';
        setUsername(username);
        setRole(role);
        // 持久化到 localStorage
        localStorage.setItem('session_id', response.session_id);
        localStorage.setItem('username', username);
        const userInfo: UserInfo = {
          username,
          role,
        };
        localStorage.setItem('userInfo', JSON.stringify(userInfo));
        console.log('[useAuth] State updated, userInfo:', userInfo);
        return response;
      }
      throw new Error('Login failed: No session_id returned');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const register = useCallback(async (email: string, username: string, password: string) => {
    setIsLoading(true);
    try {
      // 先注册
      const registerResponse = await apiClient.register(email, username, password);
      if (!registerResponse.success) {
        throw new Error(registerResponse.message || 'Registration failed');
      }
      // 注册成功后自动登录
      const loginResponse = await apiClient.login(email, password);
      if (loginResponse.session_id) {
        setIsAuthenticated(true);
        const loginUsername = loginResponse.username ?? username;
        const loginRole = loginResponse.role ?? 'user';
        setUsername(loginUsername);
        setRole(loginRole);
        // 持久化到 localStorage
        localStorage.setItem('session_id', loginResponse.session_id);
        localStorage.setItem('username', loginUsername);
        const userInfo: UserInfo = {
          username: loginUsername,
          role: loginRole,
        };
        localStorage.setItem('userInfo', JSON.stringify(userInfo));
        return { ...registerResponse, session: loginResponse };
      }
      throw new Error('Auto-login after registration failed');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const logout = useCallback(async () => {
    setIsLoading(true);
    try {
      await apiClient.logout();
      setIsAuthenticated(false);
      setUsername(null);
      setRole(null);
      // 清除 localStorage
      localStorage.removeItem('session_id');
      localStorage.removeItem('username');
      localStorage.removeItem('userInfo');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const checkStatus = useCallback(async () => {
    try {
      const status = await apiClient.authStatus();
      setIsAuthenticated(status.authenticated);
      setUsername(status.username || null);
      setRole((status.role as UserRole) || null);
      return status;
    } catch (err) {
      setIsAuthenticated(false);
      setUsername(null);
      setRole(null);
      return null;
    }
  }, []);

  return {
    isAuthenticated,
    username,
    role,
    isLoading,
    login,
    register,
    logout,
    checkStatus,
    getMcpUrl,
  };
}

// MCP 工具 Hook
export function useMcpTools() {
  const [tools, setTools] = useState<Array<{ name: string; description: string }>>([]);
  const [isLoading, setIsLoading] = useState(false);

  const fetchTools = useCallback(async () => {
    setIsLoading(true);
    try {
      const response = await apiClient.listMcpTools();
      setTools(response.tools);
      return response;
    } finally {
      setIsLoading(false);
    }
  }, []);

  return {
    tools,
    isLoading,
    fetchTools,
  };
}