bmad-elsabro 1.0.0

package/renderer/hooks/useTokenTracking.ts

@@ -0,0 +1,464 @@
+/**
+ * useTokenTracking Hook
+ *
+ * Custom hook for managing token usage tracking in the renderer process.
+ * Provides real-time token statistics, mode switching, and session management.
+ *
+ * Features:
+ * - Real-time token usage tracking
+ * - Three efficiency modes (Ultra-Save, Balanced, Power)
+ * - Session management
+ * - Budget alerts
+ * - Cache statistics
+ * - Event-based updates
+ *
+ * @module renderer/hooks/useTokenTracking
+ */
+
+import { useState, useEffect, useCallback, useRef } from 'react';
+
+/**
+ * Token efficiency mode type
+ */
+export type TokenMode = 'ultra-save' | 'balanced' | 'power';
+
+/**
+ * Token mode configuration
+ */
+export interface TokenModeConfig {
+  mode: TokenMode;
+  displayName: string;
+  description: string;
+  icon: string;
+  color: string;
+  maxContextTokens: number;
+  enableCaching: boolean;
+  cacheExpiration: number | null;
+  maxExamples: number;
+  enableStreaming: boolean;
+  requestConciseResponses: boolean;
+  temperature: number;
+  maxOutputTokens: number;
+  enableAutoSummarization: boolean;
+  savingsMultiplier: number;
+}
+
+/**
+ * Token session information
+ */
+export interface TokenSession {
+  sessionId: string;
+  startTime: number;
+  endTime?: number;
+  totalTokens: number;
+  mode: TokenMode;
+  workflowId?: string;
+  metadata?: Record<string, any>;
+}
+
+/**
+ * Token usage event
+ */
+export interface TokenUsageEvent {
+  timestamp: number;
+  tokensUsed: number;
+  mode: TokenMode;
+  sessionId: string;
+  operationType?: string;
+  workflowId?: string;
+  agentName?: string;
+  model?: string;
+  inputTokens?: number;
+  outputTokens?: number;
+  costEstimate?: number;
+}
+
+/**
+ * Budget alert event
+ */
+export interface BudgetAlertEvent {
+  budgetName: string;
+  limitTokens: number;
+  usedTokens: number;
+  remainingTokens: number;
+  percentageUsed: number;
+  thresholdReached: boolean;
+}
+
+/**
+ * Cache statistics
+ */
+export interface CacheStats {
+  total_entries: number;
+  total_tokens_saved: number;
+  total_hits: number;
+}
+
+/**
+ * Hook return type
+ */
+export interface UseTokenTrackingReturn {
+  /**
+   * Current token mode
+   */
+  mode: TokenMode;
+
+  /**
+   * Current mode configuration
+   */
+  modeConfig: TokenModeConfig | null;
+
+  /**
+   * Current session
+   */
+  session: TokenSession | null;
+
+  /**
+   * Total tokens used in current session
+   */
+  sessionTokens: number;
+
+  /**
+   * Cache statistics
+   */
+  cacheStats: CacheStats | null;
+
+  /**
+   * Latest budget alert
+   */
+  latestAlert: BudgetAlertEvent | null;
+
+  /**
+   * Whether tracker is loading
+   */
+  isLoading: boolean;
+
+  /**
+   * Switch token mode
+   */
+  setMode: (mode: TokenMode) => Promise<void>;
+
+  /**
+   * Start a new tracking session
+   */
+  startSession: (workflowId?: string, metadata?: Record<string, any>) => Promise<string | null>;
+
+  /**
+   * End current session
+   */
+  endSession: () => Promise<TokenSession | null>;
+
+  /**
+   * Track token usage
+   */
+  trackUsage: (data: {
+    tokensUsed: number;
+    operationType?: string;
+    workflowId?: string;
+    agentName?: string;
+    model?: string;
+    inputTokens?: number;
+    outputTokens?: number;
+    costEstimate?: number;
+  }) => Promise<void>;
+
+  /**
+   * Refresh cache statistics
+   */
+  refreshCacheStats: () => Promise<void>;
+}
+
+/**
+ * Check if window.electron API is available
+ */
+const isElectronAvailable = (): boolean => {
+  return typeof window !== 'undefined' && 'electron' in window;
+};
+
+/**
+ * Custom hook for token tracking
+ *
+ * @example
+ * ```tsx
+ * const {
+ *   mode,
+ *   sessionTokens,
+ *   setMode,
+ *   startSession,
+ *   trackUsage,
+ * } = useTokenTracking();
+ *
+ * // Switch to ultra-save mode
+ * await setMode('ultra-save');
+ *
+ * // Start a session
+ * await startSession('workflow-123');
+ *
+ * // Track usage
+ * await trackUsage({ tokensUsed: 150, operationType: 'chat' });
+ * ```
+ */
+export const useTokenTracking = (): UseTokenTrackingReturn => {
+  const [mode, setModeState] = useState<TokenMode>('balanced');
+  const [modeConfig, setModeConfig] = useState<TokenModeConfig | null>(null);
+  const [session, setSession] = useState<TokenSession | null>(null);
+  const [sessionTokens, setSessionTokens] = useState<number>(0);
+  const [cacheStats, setCacheStats] = useState<CacheStats | null>(null);
+  const [latestAlert, setLatestAlert] = useState<BudgetAlertEvent | null>(null);
+  const [isLoading, setIsLoading] = useState<boolean>(true);
+  const eventListenersRegistered = useRef<boolean>(false);
+
+  /**
+   * Initialize token tracking
+   */
+  useEffect(() => {
+    const initializeTracking = async () => {
+      if (!isElectronAvailable()) {
+        console.warn('[useTokenTracking] Electron API not available');
+        setIsLoading(false);
+        return;
+      }
+
+      try {
+        // Get current mode
+        const modeResult = await (window as any).electron.invoke('token:get-current-mode');
+        if (modeResult.success) {
+          setModeState(modeResult.mode);
+          setModeConfig(modeResult.config);
+        }
+
+        // Get current session
+        const sessionResult = await (window as any).electron.invoke('token:get-session');
+        if (sessionResult.success && sessionResult.session) {
+          setSession(sessionResult.session);
+          setSessionTokens(sessionResult.session.totalTokens || 0);
+        }
+
+        // Get cache stats
+        const statsResult = await (window as any).electron.invoke('token:get-cache-stats');
+        if (statsResult.success) {
+          setCacheStats(statsResult.stats);
+        }
+
+        setIsLoading(false);
+      } catch (error) {
+        console.error('[useTokenTracking] Initialization error:', error);
+        setIsLoading(false);
+      }
+    };
+
+    initializeTracking();
+  }, []);
+
+  /**
+   * Set up event listeners for real-time updates
+   */
+  useEffect(() => {
+    if (!isElectronAvailable() || eventListenersRegistered.current) {
+      return;
+    }
+
+    const electron = (window as any).electron;
+
+    // Listen for token usage events
+    electron.receive('token:tokens-used', (event: TokenUsageEvent) => {
+      setSessionTokens(prev => prev + event.tokensUsed);
+
+      // Update session if it matches
+      if (session && event.sessionId === session.sessionId) {
+        setSession(prev => prev ? { ...prev, totalTokens: prev.totalTokens + event.tokensUsed } : null);
+      }
+    });
+
+    // Listen for mode changes
+    electron.receive('token:mode-changed', (data: { newMode: TokenMode; config: TokenModeConfig }) => {
+      setModeState(data.newMode);
+      setModeConfig(data.config);
+    });
+
+    // Listen for budget alerts
+    electron.receive('token:budget-alert', (alert: BudgetAlertEvent) => {
+      setLatestAlert(alert);
+    });
+
+    // Listen for budget exceeded
+    electron.receive('token:budget-exceeded', (alert: BudgetAlertEvent) => {
+      setLatestAlert(alert);
+    });
+
+    // Listen for session events
+    electron.receive('token:session-started', (newSession: TokenSession) => {
+      setSession(newSession);
+      setSessionTokens(0);
+    });
+
+    electron.receive('token:session-ended', (endedSession: TokenSession) => {
+      if (session && endedSession.sessionId === session.sessionId) {
+        setSession(null);
+        setSessionTokens(0);
+      }
+    });
+
+    // Listen for cache hits
+    electron.receive('token:cache-hit', () => {
+      // Refresh cache stats when there's a cache hit
+      refreshCacheStats();
+    });
+
+    eventListenersRegistered.current = true;
+
+    // Cleanup function
+    return () => {
+      eventListenersRegistered.current = false;
+    };
+  }, [session]);
+
+  /**
+   * Switch token mode
+   */
+  const setMode = useCallback(async (newMode: TokenMode): Promise<void> => {
+    if (!isElectronAvailable()) {
+      console.warn('[useTokenTracking] Electron API not available');
+      return;
+    }
+
+    try {
+      const result = await (window as any).electron.invoke('token:set-mode', newMode);
+
+      if (result.success) {
+        setModeState(result.mode);
+        setModeConfig(result.config);
+      } else {
+        console.error('[useTokenTracking] Failed to set mode:', result.error);
+      }
+    } catch (error) {
+      console.error('[useTokenTracking] Error setting mode:', error);
+    }
+  }, []);
+
+  /**
+   * Start a new session
+   */
+  const startSession = useCallback(async (
+    workflowId?: string,
+    metadata?: Record<string, any>
+  ): Promise<string | null> => {
+    if (!isElectronAvailable()) {
+      console.warn('[useTokenTracking] Electron API not available');
+      return null;
+    }
+
+    try {
+      const result = await (window as any).electron.invoke('token:start-session', {
+        workflowId,
+        metadata,
+      });
+
+      if (result.success) {
+        setSession(result.session);
+        setSessionTokens(0);
+        return result.sessionId;
+      } else {
+        console.error('[useTokenTracking] Failed to start session:', result.error);
+        return null;
+      }
+    } catch (error) {
+      console.error('[useTokenTracking] Error starting session:', error);
+      return null;
+    }
+  }, []);
+
+  /**
+   * End current session
+   */
+  const endSession = useCallback(async (): Promise<TokenSession | null> => {
+    if (!isElectronAvailable()) {
+      console.warn('[useTokenTracking] Electron API not available');
+      return null;
+    }
+
+    try {
+      const result = await (window as any).electron.invoke('token:end-session');
+
+      if (result.success) {
+        setSession(null);
+        setSessionTokens(0);
+        return result.session;
+      } else {
+        console.error('[useTokenTracking] Failed to end session:', result.error);
+        return null;
+      }
+    } catch (error) {
+      console.error('[useTokenTracking] Error ending session:', error);
+      return null;
+    }
+  }, []);
+
+  /**
+   * Track token usage
+   */
+  const trackUsage = useCallback(async (data: {
+    tokensUsed: number;
+    operationType?: string;
+    workflowId?: string;
+    agentName?: string;
+    model?: string;
+    inputTokens?: number;
+    outputTokens?: number;
+    costEstimate?: number;
+  }): Promise<void> => {
+    if (!isElectronAvailable()) {
+      console.warn('[useTokenTracking] Electron API not available');
+      return;
+    }
+
+    try {
+      const result = await (window as any).electron.invoke('token:track-usage', data);
+
+      if (result.success && result.session) {
+        setSession(result.session);
+        setSessionTokens(result.session.totalTokens || 0);
+      } else if (!result.success) {
+        console.error('[useTokenTracking] Failed to track usage:', result.error);
+      }
+    } catch (error) {
+      console.error('[useTokenTracking] Error tracking usage:', error);
+    }
+  }, []);
+
+  /**
+   * Refresh cache statistics
+   */
+  const refreshCacheStats = useCallback(async (): Promise<void> => {
+    if (!isElectronAvailable()) {
+      return;
+    }
+
+    try {
+      const result = await (window as any).electron.invoke('token:get-cache-stats');
+
+      if (result.success) {
+        setCacheStats(result.stats);
+      }
+    } catch (error) {
+      console.error('[useTokenTracking] Error refreshing cache stats:', error);
+    }
+  }, []);
+
+  return {
+    mode,
+    modeConfig,
+    session,
+    sessionTokens,
+    cacheStats,
+    latestAlert,
+    isLoading,
+    setMode,
+    startSession,
+    endSession,
+    trackUsage,
+    refreshCacheStats,
+  };
+};
+
+export default useTokenTracking;