Feat: Add z.ai usage tracking

apps/ui/src/components/views/settings-view/api-keys/hooks/use-api-key-management.ts

@@ -1,7 +1,11 @@
 // @ts-nocheck - API key management state with validation and persistence
 import { useState, useEffect } from 'react';
+import { useQueryClient } from '@tanstack/react-query';
 import { createLogger } from '@automaker/utils/logger';
 import { useAppStore } from '@/store/app-store';
+import { useSetupStore, type ZaiAuthMethod } from '@/store/setup-store';
+import { getHttpApiClient } from '@/lib/http-api-client';
+import { queryKeys } from '@/lib/query-keys';
 
 const logger = createLogger('ApiKeyManagement');
 import { getElectronAPI } from '@/lib/electron';
@@ -16,6 +20,7 @@ interface ApiKeyStatus {
   hasAnthropicKey: boolean;
   hasGoogleKey: boolean;
   hasOpenaiKey: boolean;
+  hasZaiKey: boolean;
 }
 
 /**
@@ -24,16 +29,20 @@ interface ApiKeyStatus {
  */
 export function useApiKeyManagement() {
   const { apiKeys, setApiKeys } = useAppStore();
+  const { setZaiAuthStatus } = useSetupStore();
+  const queryClient = useQueryClient();
 
   // API key values
   const [anthropicKey, setAnthropicKey] = useState(apiKeys.anthropic);
   const [googleKey, setGoogleKey] = useState(apiKeys.google);
   const [openaiKey, setOpenaiKey] = useState(apiKeys.openai);
+  const [zaiKey, setZaiKey] = useState(apiKeys.zai);
 
   // Visibility toggles
   const [showAnthropicKey, setShowAnthropicKey] = useState(false);
   const [showGoogleKey, setShowGoogleKey] = useState(false);
   const [showOpenaiKey, setShowOpenaiKey] = useState(false);
+  const [showZaiKey, setShowZaiKey] = useState(false);
 
   // Test connection states
   const [testingConnection, setTestingConnection] = useState(false);
@@ -42,6 +51,8 @@ export function useApiKeyManagement() {
   const [geminiTestResult, setGeminiTestResult] = useState<TestResult | null>(null);
   const [testingOpenaiConnection, setTestingOpenaiConnection] = useState(false);
   const [openaiTestResult, setOpenaiTestResult] = useState<TestResult | null>(null);
+  const [testingZaiConnection, setTestingZaiConnection] = useState(false);
+  const [zaiTestResult, setZaiTestResult] = useState<TestResult | null>(null);
 
   // API key status from environment
   const [apiKeyStatus, setApiKeyStatus] = useState<ApiKeyStatus | null>(null);
@@ -54,6 +65,7 @@ export function useApiKeyManagement() {
     setAnthropicKey(apiKeys.anthropic);
     setGoogleKey(apiKeys.google);
     setOpenaiKey(apiKeys.openai);
+    setZaiKey(apiKeys.zai);
   }, [apiKeys]);
 
   // Check API key status from environment on mount
@@ -68,6 +80,7 @@ export function useApiKeyManagement() {
               hasAnthropicKey: status.hasAnthropicKey,
               hasGoogleKey: status.hasGoogleKey,
               hasOpenaiKey: status.hasOpenaiKey,
+              hasZaiKey: status.hasZaiKey || false,
             });
           }
         } catch (error) {
@@ -173,13 +186,89 @@ export function useApiKeyManagement() {
     }
   };
 
+  // Test z.ai connection
+  const handleTestZaiConnection = async () => {
+    setTestingZaiConnection(true);
+    setZaiTestResult(null);
+
+    // Validate input first
+    if (!zaiKey || zaiKey.trim().length === 0) {
+      setZaiTestResult({
+        success: false,
+        message: 'Please enter an API key to test.',
+      });
+      setTestingZaiConnection(false);
+      return;
+    }
+
+    try {
+      const api = getElectronAPI();
+      // Use the verify endpoint to test the key without storing it
+      const response = await api.zai?.verify(zaiKey);
+
+      if (response?.success && response?.authenticated) {
+        setZaiTestResult({
+          success: true,
+          message: response.message || 'Connection successful! z.ai API responded.',
+        });
+      } else {
+        setZaiTestResult({
+          success: false,
+          message: response?.error || 'Failed to connect to z.ai API.',
+        });
+      }
+    } catch {
+      setZaiTestResult({
+        success: false,
+        message: 'Network error. Please check your connection.',
+      });
+    } finally {
+      setTestingZaiConnection(false);
+    }
+  };
+
   // Save API keys
-  const handleSave = () => {
+  const handleSave = async () => {
     setApiKeys({
       anthropic: anthropicKey,
       google: googleKey,
       openai: openaiKey,
+      zai: zaiKey,
     });
+
+    // Configure z.ai service on the server with the new key
+    if (zaiKey && zaiKey.trim().length > 0) {
+      try {
+        const api = getHttpApiClient();
+        const result = await api.zai.configure(zaiKey.trim());
+
+        if (result.success || result.isAvailable) {
+          // Update z.ai auth status in the store
+          setZaiAuthStatus({
+            authenticated: true,
+            method: 'api_key' as ZaiAuthMethod,
+            hasApiKey: true,
+            hasEnvApiKey: false,
+          });
+          // Invalidate the z.ai usage query so it refetches with the new key
+          await queryClient.invalidateQueries({ queryKey: queryKeys.usage.zai() });
+          logger.info('z.ai API key configured successfully');
+        }
+      } catch (error) {
+        logger.error('Failed to configure z.ai API key:', error);
+      }
+    } else {
+      // Clear z.ai auth status if key is removed
+      setZaiAuthStatus({
+        authenticated: false,
+        method: 'none' as ZaiAuthMethod,
+        hasApiKey: false,
+        hasEnvApiKey: false,
+      });
+      // Invalidate the query to clear any cached data
+      await queryClient.invalidateQueries({ queryKey: queryKeys.usage.zai() });
+    }
+
     setSaved(true);
     setTimeout(() => setSaved(false), 2000);
   };
@@ -214,6 +303,15 @@ export function useApiKeyManagement() {
       onTest: handleTestOpenaiConnection,
       result: openaiTestResult,
     },
+    zai: {
+      value: zaiKey,
+      setValue: setZaiKey,
+      show: showZaiKey,
+      setShow: setShowZaiKey,
+      testing: testingZaiConnection,
+      onTest: handleTestZaiConnection,
+      result: zaiTestResult,
+    },
   };
 
   return {