VERTEX_AI_SEARCH_WIDGET_COMPONENT.md

Overview

The VertexAISearchWidget component is a sophisticated integration with Google’s Vertex AI Search service that provides users with the ability to search through knowledge bases and document collections. It handles complex authentication flows, token management, and third-party script integration to deliver a seamless search experience.

Purpose

AI-Powered Search: Integrate Google Vertex AI Search capabilities into the application
Authentication Management: Handle complex token-based authentication with auto-refresh
Third-Party Integration: Safely load and manage external Google search widgets
User Experience: Provide intuitive search interface with proper loading and error states

Key Features

Authentication Integration

Firebase Authentication: Integrates with Firebase auth for user verification
Google Token Management: Manages Google Cloud access tokens with automatic refresh
Token Persistence: Maintains tokens across component re-renders and config changes
Authentication Error Handling: Graceful handling of authentication failures

Third-Party Script Management

Dynamic Script Loading: Loads Google’s search widget script on-demand
Custom Element Registration: Handles custom web component integration
Script Error Handling: Manages script loading failures and retries

Configuration Management: Supports dynamic config ID changes
Widget Lifecycle: Properly initializes and manages widget state
Token Application: Applies authentication tokens to the search widget
Trigger Integration: Connects custom trigger buttons to the search widget

State Management

Complex State Coordination: Manages multiple interdependent states
Reference Management: Uses refs for direct DOM manipulation and state persistence
Effect Coordination: Coordinates multiple useEffect hooks for proper initialization

Component Interface

interface VertexSearchWidgetProps {
  configId: string;           // Vertex AI Search configuration ID
  buttonText?: string | null; // Custom button text (default: "Search Knowledge Base")
  className?: string;         // Additional CSS classes
  iconOnly?: boolean;         // Show only icon without text
}

// Custom element declaration for TypeScript
declare global {
  namespace JSX {
    interface IntrinsicElements {
      'gen-search-widget': React.DetailedHTMLProps<React.HTMLAttributes<HTMLElement> & {
        configId?: string;
        location?: string;
        triggerId?: string;
        authToken?: string;
        ref?: React.Ref<HTMLElement>;
      }, HTMLElement>;
    }
  }
}

// Extended widget element interface
interface VertexSearchWidgetElement extends HTMLElement {
  authToken?: string;
}

Core Functionality

1. Authentication Token Management

// Token fetching with Firebase authentication
const fetchSearchToken = async (user: User) => {
  // Prevent multiple concurrent requests
  if (tokenRequestInProgress.current) {
    debugLog('Token request already in progress, skipping');
    return;
  }
  
  tokenRequestInProgress.current = true;
  debugLog(`Fetching token for user: ${user.email}`);
  
  try {
    setIsLoading(true);
    setError(null);
    
    // Get Firebase ID token
    const idToken = await user.getIdToken();
    
    // Call API endpoint
    const response = await fetch('/api/aisearch', {
      method: 'GET',
      headers: {
        'Authorization': `Bearer ${idToken}`
      }
    });
    
    if (!response.ok) {
      const errorText = await response.text();
      console.error(`Vertex AI Search authentication error: ${response.status} ${errorText}`);
      setError('search-disabled');
      return;
    }
    
    const data = await response.json();
    
    // Store token regardless of widget status
    if (data.accessToken) {
      latestToken.current = data.accessToken;
      
      // Apply token if widget exists
      if (widgetRef.current) {
        widgetRef.current.authToken = data.accessToken;
      }
      
      // Set up refresh timer
      setupTokenRefresh(user, data.expiresIn);
    }
  } catch (error) {
    console.error(`Vertex AI Search token fetch error: ${error}`);
    setError('search-disabled');
  } finally {
    setIsLoading(false);
    tokenRequestInProgress.current = false;
  }
};

2. Automatic Token Refresh

// Set up automatic token refresh before expiration
const setupTokenRefresh = (user: User, expiresIn: number) => {
  if (tokenRefreshTimer) {
    clearTimeout(tokenRefreshTimer);
  }
  
  // Refresh 5 minutes before expiration
  const refreshTime = (expiresIn || 3600) * 1000 - 5 * 60 * 1000;
  const refreshTimer = setTimeout(() => {
    debugLog('Token refresh timer triggered');
    tokenRequestInProgress.current = false; // Reset flag before refreshing
    fetchSearchToken(user);
  }, refreshTime);
  
  setTokenRefreshTimer(refreshTimer);
};

3. Dynamic Script Loading

// Load Google's search widget script
useEffect(() => {
  if (!scriptLoaded.current) {
    debugLog('Loading widget script...');
    const script = document.createElement('script');
    script.src = 'https://cloud.google.com/ai/gen-app-builder/client?hl=en_US';
    script.async = true;
    script.onload = () => {
      debugLog('Widget script loaded successfully');
      scriptLoaded.current = true;
    };
    script.onerror = (error) => {
      debugLog(`Error loading widget script: ${error}`);
    };
    document.body.appendChild(script);
    
    return () => {
      if (document.body.contains(script)) {
        document.body.removeChild(script);
      }
    };
  }
}, []);

4. Configuration Change Handling

// Handle dynamic config ID changes
useEffect(() => {
  // Check if configId has changed
  if (configId !== previousConfigId.current) {
    debugLog(`ConfigId changed from ${previousConfigId.current} to ${configId}`);
    previousConfigId.current = configId;
    
    // Reset widget by recreating it
    if (widgetRef.current) {
      debugLog('Resetting widget due to configId change');
      
      // Remove configId and save a reference to the element
      const widgetElement = widgetRef.current;
      widgetElement.removeAttribute('configId');
      
      // Apply token if available
      if (latestToken.current) {
        debugLog('Reapplying token to reset widget');
        widgetElement.authToken = latestToken.current;
      }
      
      // Short delay before adding the new configId
      setTimeout(() => {
        if (widgetElement) {
          debugLog(`Setting new configId: ${configId}`);
          widgetElement.setAttribute('configId', configId);
        }
      }, 100);
    }
  }
}, [configId]);

5. Firebase Authentication Integration

// Monitor Firebase authentication state
useEffect(() => {
  const unsubscribe = FirebaseService.onAuthStateChange((user) => {
    debugLog(`Auth state changed, user: ${user ? `${user.email} (logged in)` : 'not logged in'}`);
    setIsAuthenticated(!!user);
    setCurrentUser(user);
    
    // Reset setup flag on auth change
    hasSetupWidget.current = false;
    
    // Try to immediately set up widget if user is available
    if (user) {
      debugLog('User authenticated, scheduling immediate token fetch');
      // Force token fetch on auth change
      setTimeout(() => {
        if (!tokenRequestInProgress.current) {
          debugLog('Delayed token fetch after auth change');
          fetchSearchToken(user);
        }
      }, 100); // Short delay to ensure component updates complete
    }
  });
  
  return () => {
    debugLog('Cleaning up auth state monitor');
    unsubscribe();
  };
}, []);

// Complex widget reference management with token application
<gen-search-widget
  ref={(el: VertexSearchWidgetElement | null) => {
    debugLog(`Widget ref callback: element ${el ? 'exists' : 'is null'}`);
    if (el !== widgetRef.current) {
      widgetRef.current = el;
      debugLog('Updated widget ref');
      
      // Apply token if we already have one
      if (el && latestToken.current) {
        debugLog('Applying token to newly referenced widget');
        
        // Remove the configId first, then set the token, then re-add the configId
        // This forces the widget to start fresh with the token already set
        const currentConfigId = el.getAttribute('configId');
        el.removeAttribute('configId');
        
        // Apply token
        el.authToken = latestToken.current;
        
        // Wait a moment for the token to be processed, then re-add configId
        setTimeout(() => {
          if (el && currentConfigId) {
            debugLog('Re-adding configId after token was set');
            el.setAttribute('configId', currentConfigId);
          }
        }, 100);
      }
    }
  }}
  configId={configId}
  location="eu"
  triggerId="searchWidgetTrigger">
</gen-search-widget>

State Management

Component State Variables

// Authentication state
const [isAuthenticated, setIsAuthenticated] = useState(false);
const [currentUser, setCurrentUser] = useState<User | null>(null);

// Loading and error states
const [isLoading, setIsLoading] = useState(false);
const [error, setError] = useState<string | null>(null);

// Timer management
const [tokenRefreshTimer, setTokenRefreshTimer] = useState<NodeJS.Timeout | null>(null);

// Refs for direct DOM manipulation and state persistence
const searchButtonRef = useRef<HTMLButtonElement | null>(null);
const widgetRef = useRef<VertexSearchWidgetElement | null>(null);
const scriptLoaded = useRef(false);
const tokenRequestInProgress = useRef(false);
const latestToken = useRef<string | null>(null);
const hasSetupWidget = useRef(false);
const previousConfigId = useRef<string>(configId);

Complex Effect Coordination

// Effect for initialization and setup
useEffect(() => {
  debugLog(`Setup effect running: authenticated=${isAuthenticated}, scriptLoaded=${scriptLoaded.current}, hasSetupWidget=${hasSetupWidget.current}`);
  
  // Check if we can initialize
  if (isAuthenticated && currentUser && scriptLoaded.current && !hasSetupWidget.current) {
    debugLog('All conditions met for widget setup');
    hasSetupWidget.current = true;
    
    // Force fetch token
    debugLog('Triggering initial token fetch');
    fetchSearchToken(currentUser);
  }
}, [isAuthenticated, currentUser, scriptLoaded.current]);

// Effect to apply token when widget becomes available
useEffect(() => {
  debugLog(`Widget ref effect: widgetRef=${!!widgetRef.current}, latestToken=${!!latestToken.current}`);
  if (widgetRef.current && latestToken.current) {
    debugLog('Applying stored token to widget from ref effect');
    widgetRef.current.authToken = latestToken.current;
  }
}, [widgetRef.current]);

UI States and Rendering

Authentication Required State

// If user is not authenticated, show login prompt
if (!isAuthenticated) {
  return (
    <div className="flex flex-col items-center justify-center p-2 border border-gray-200 rounded-lg">
      <p className="mb-2 text-center text-sm">Sign in to access search</p>
      <Button 
        variant="outline"
        size="sm"
        onClick={() => FirebaseService.loginWithGoogle()}
        className="w-full"
      >
        Sign in with Google
      </Button>
    </div>
  );
}

Loading State

// Show loading state
if (isLoading && !latestToken.current) {
  return (
    <div className="flex items-center justify-center p-2">
      <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-gray-900"></div>
      <span className="ml-2 text-sm">Loading search...</span>
    </div>
  );
}

Error State

// Show disabled state on error
if (error === 'search-disabled') {
  return (
    <Button
      variant="ghost"
      size="icon"
      disabled={true}
      className={`h-10 w-10 flex items-center justify-center border rounded-lg opacity-50 ${className}`}
      title="Search unavailable - authentication error"
    >
      <Search className="h-4 w-4" />
    </Button>
  );
}

// Main widget interface
return (
  <div className="search-widget-container">
    {/* Search button that triggers the widget */}
    <Button
      ref={searchButtonRef}
      id="searchWidgetTrigger"
      variant="ghost"
      size="icon"
      className={`h-10 w-10 flex items-center justify-center border rounded-lg ${className}`}
    >
      <Search className="h-4 w-4" />
      {!iconOnly && buttonText && <span className="ml-2">{buttonText}</span>}
    </Button>
    
    {/* The actual widget element */}
    {isAuthenticated && (
      <gen-search-widget
        ref={widgetRefCallback}
        configId={configId}
        location="eu"
        triggerId="searchWidgetTrigger">
      </gen-search-widget>
    )}
  </div>
);

Integration Examples

Basic Usage

// Simple search widget integration
function DocumentSearchPanel() {
  return (
    <div className="search-panel">
      <h3>Search Knowledge Base</h3>
      <VertexSearchWidget
        configId="your-vertex-ai-config-id"
        buttonText="Search Documents"
      />
    </div>
  );
}

Icon-Only Integration

// Compact icon-only version for toolbars
function ToolbarSearch() {
  return (
    <div className="toolbar">
      <VertexSearchWidget
        configId="your-config-id"
        iconOnly={true}
        className="toolbar-search"
      />
    </div>
  );
}

Dynamic Configuration

// Dynamic config based on user context
function ContextualSearch({ userRole, department }: { userRole: string, department: string }) {
  const getConfigId = (role: string, dept: string) => {
    // Return different config IDs based on user context
    if (role === 'admin') return 'admin-search-config';
    if (dept === 'legal') return 'legal-docs-config';
    return 'general-search-config';
  };

  return (
    <VertexSearchWidget
      configId={getConfigId(userRole, department)}
      buttonText={`Search ${department} Resources`}
    />
  );
}

Integration with Tool Context

// Integration with tool configuration system
function ConfigurableSearch() {
  const { toolConfigs } = useToolContext();
  const searchConfig = toolConfigs['vertex_search'] || {};
  
  if (!searchConfig.configId) {
    return (
      <div className="p-4 text-center text-gray-500">
        Search not configured. Please set up Vertex AI Search in admin settings.
      </div>
    );
  }

  return (
    <VertexSearchWidget
      configId={searchConfig.configId}
      buttonText={searchConfig.buttonText || "Search Knowledge Base"}
    />
  );
}

API Integration

Backend Authentication Endpoint

// Expected API endpoint structure
// GET /api/aisearch
// Headers: Authorization: Bearer <firebase-id-token>
// Response:
{
  "accessToken": "google-cloud-access-token",
  "expiresIn": 3600, // seconds
  "tokenType": "Bearer"
}

Error Handling

// API error responses
// 401 Unauthorized - Firebase token invalid
// 403 Forbidden - User not authorized for Vertex AI
// 500 Internal Server Error - Google Cloud authentication failed

// Local development setup error
console.error(`To fix locally, run: gcloud auth application-default login`);

Debugging and Logging

Debug Logging System

// Comprehensive debug logging
const debugLog = (message: string) => {
  console.log(`[VertexSearch] ${message}`);
};

// Usage throughout component
debugLog('ConfigId changed from ${previousConfigId.current} to ${configId}');
debugLog('All conditions met for widget setup');
debugLog('Applying stored token to widget from ref effect');

Common Debug Scenarios

Script Loading Issues: Check browser console for script load errors
Authentication Failures: Verify Firebase authentication and Google Cloud credentials
Token Refresh Problems: Monitor timer setup and refresh cycles
Widget Initialization: Track widget ref callbacks and token application
Config Changes: Verify proper widget reset on configuration updates

Performance Considerations

Script Caching: Google’s widget script is loaded once and cached
Token Persistence: Tokens are stored in refs to survive re-renders
Debounced Initialization: Prevents multiple simultaneous token requests
Efficient Re-renders: Uses refs for values that don’t need to trigger re-renders
Timer Management: Proper cleanup of timeout timers on unmount

Security Considerations

Token Management: Access tokens are handled securely and auto-refreshed
Authentication Validation: Proper Firebase ID token validation
HTTPS Required: Google Cloud integration requires HTTPS in production
Cross-Origin Security: Widget script loaded from trusted Google domain
Token Expiration: Automatic refresh prevents expired token usage

Error Handling

Authentication Errors

Firebase Auth Failures: Graceful handling with login prompts
Google Cloud Auth Issues: Disable search with helpful error messages
Token Refresh Failures: Retry logic with exponential backoff

Script Loading Errors

Network Issues: Retry mechanism for script loading
CORS Problems: Proper error messages for development environments
Script Failures: Fallback UI when widget script fails to load

Element Creation Failures: Proper error boundaries around widget elements
Token Application Issues: Retry logic for token application
Configuration Errors: Validation of config IDs and parameters

Browser Compatibility

Supported Browsers

Chrome: Full support (recommended)
Firefox: Full support with modern versions
Safari: Requires recent versions for custom elements
Edge: Chromium-based versions supported

Polyfills Required

Custom Elements: For older browser support
Promise: For token management in legacy browsers
Fetch: For API communication in older browsers

VERTEX_AI_SEARCH_WIDGET_COMPONENT.md

Overview

Purpose

Key Features

Authentication Integration

Third-Party Script Management

Search Widget Integration

State Management

Component Interface

Core Functionality

1. Authentication Token Management

2. Automatic Token Refresh

3. Dynamic Script Loading

4. Configuration Change Handling

5. Firebase Authentication Integration

6. Widget Reference Management

State Management

Component State Variables

Complex Effect Coordination

UI States and Rendering

Authentication Required State

Loading State

Error State

Active Search Widget

Integration Examples

Basic Usage

Icon-Only Integration

Dynamic Configuration

Integration with Tool Context

API Integration

Backend Authentication Endpoint

Error Handling

Debugging and Logging

Debug Logging System

Common Debug Scenarios

Performance Considerations

Security Considerations

Error Handling

Authentication Errors

Script Loading Errors

Widget Integration Errors

Browser Compatibility

Supported Browsers

Polyfills Required