stable-diffusion.cpp-rest/webui/lib/api.ts

// API client for stable-diffusion REST API

// Type for server config injected by the server
declare global {
  interface Window {
    __SERVER_CONFIG__?: {
      apiUrl: string;
      apiBasePath: string;
      host: string;
      port: number;
      authMethod: 'none' | 'unix' | 'jwt';
      authEnabled: boolean;
    };
  }
}

// Get configuration from server-injected config or fallback to environment/defaults
// This function is called at runtime to ensure __SERVER_CONFIG__ is available
function getApiConfig() {
  if (typeof window !== 'undefined' && window.__SERVER_CONFIG__) {
    return {
      apiUrl: window.__SERVER_CONFIG__.apiUrl,
      apiBase: window.__SERVER_CONFIG__.apiBasePath,
    };
  }

  // Fallback for development mode - use current window location
  if (typeof window !== 'undefined') {
    const protocol = window.location.protocol;
    const host = window.location.hostname;
    const port = window.location.port;
    return {
      apiUrl: `${protocol}//${host}:${port}`,
      apiBase: '/api',
    };
  }

  // Server-side fallback
  return {
    apiUrl: process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8081',
    apiBase: process.env.NEXT_PUBLIC_API_BASE_PATH || '/api',
  };
}

export interface GenerationRequest {
  model?: string;
  prompt: string;
  negative_prompt?: string;
  width?: number;
  height?: number;
  steps?: number;
  cfg_scale?: number;
  seed?: string;
  sampling_method?: string;
  scheduler?: string;
  batch_count?: number;
  clip_skip?: number;
  strength?: number;
  control_strength?: number;
}

export interface JobInfo {
  id?: string;
  request_id?: string;
  status: 'pending' | 'processing' | 'completed' | 'failed' | 'cancelled' | 'queued';
  progress?: number;
  result?: {
    images: string[];
  };
  outputs?: Array<{
    filename: string;
    url: string;
    path: string;
  }>;
  error?: string;
  created_at?: string;
  updated_at?: string;
  message?: string;
  queue_position?: number;
}

export interface ModelInfo {
  id?: string;
  name: string;
  path?: string;
  type: string;
  size?: number;
  file_size?: number;
  file_size_mb?: number;
  sha256?: string | null;
  sha256_short?: string | null;
  loaded?: boolean;
}

export interface QueueStatus {
  active_generations: number;
  jobs: JobInfo[];
  running: boolean;
  size: number;
}

export interface HealthStatus {
  status: 'ok' | 'error' | 'degraded';
  message: string;
  timestamp: string;
  uptime?: number;
  version?: string;
}

class ApiClient {
  // Get base URL dynamically at runtime to ensure server config is loaded
  private getBaseUrl(): string {
    const { apiUrl, apiBase } = getApiConfig();
    return `${apiUrl}${apiBase}`;
  }

  private async request<T>(
    endpoint: string,
    options: RequestInit = {}
  ): Promise<T> {
    const url = `${this.getBaseUrl()}${endpoint}`;

    // Get authentication method from server config
    const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
      ? window.__SERVER_CONFIG__.authMethod
      : 'jwt';

    // Add auth token or Unix user header based on auth method
    const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
    const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      ...options.headers as Record<string, string>,
    };

    if (authMethod === 'unix' && unixUser) {
      // For Unix auth, send the username in X-Unix-User header
      headers['X-Unix-User'] = unixUser;
    } else if (token) {
      // For JWT auth, send the token in Authorization header
      headers['Authorization'] = `Bearer ${token}`;
    }

    const response = await fetch(url, {
      ...options,
      headers,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({
        error: { message: response.statusText },
      }));

      // Handle nested error structure: { error: { message: "..." } }
      const errorMessage =
        errorData.error?.message ||
        errorData.message ||
        errorData.error ||
        'API request failed';

      throw new Error(errorMessage);
    }

    return response.json();
  }

  // Enhanced health check with multiple endpoints and better error handling
  async checkHealth(): Promise<HealthStatus> {
    const endpoints = ['/queue/status', '/health', '/status', '/'];
    
    for (const endpoint of endpoints) {
      try {
        const response = await fetch(`${this.getBaseUrl()}${endpoint}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json',
          },
          // Add timeout to prevent hanging requests
          signal: AbortSignal.timeout(5000),
        });

        if (response.ok) {
          const data = await response.json();
          
          // For queue status, consider it healthy if it returns valid structure
          if (endpoint === '/queue/status' && data.queue) {
            return {
              status: 'ok',
              message: 'API is running and queue is accessible',
              timestamp: new Date().toISOString(),
            };
          }
          
          // For other health endpoints
          const healthStatus: HealthStatus = {
            status: 'ok',
            message: 'API is running',
            timestamp: new Date().toISOString(),
            uptime: data.uptime,
            version: data.version || data.build || data.git_version,
          };

          // Handle different response formats
          if (data.status) {
            if (data.status === 'healthy' || data.status === 'running' || data.status === 'ok') {
              healthStatus.status = 'ok';
              healthStatus.message = data.message || 'API is running';
            } else if (data.status === 'degraded') {
              healthStatus.status = 'degraded';
              healthStatus.message = data.message || 'API is running in degraded mode';
            } else {
              healthStatus.status = 'error';
              healthStatus.message = data.message || 'API status is unknown';
            }
          }

          return healthStatus;
        }
      } catch (error) {
        // Continue to next endpoint if this one fails
        console.warn(`Health check failed for endpoint ${endpoint}:`, error);
        continue;
      }
    }

    // If all endpoints fail
    throw new Error('All health check endpoints are unavailable');
  }

  // Alternative simple connectivity check
  async checkConnectivity(): Promise<boolean> {
    try {
      const response = await fetch(`${this.getBaseUrl()}`, {
        method: 'HEAD',
        signal: AbortSignal.timeout(3000),
      });
      return response.ok || response.status < 500; // Accept any non-server-error response
    } catch (error) {
      return false;
    }
  }

  // Generation endpoints
  async generateImage(params: GenerationRequest): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/text2img', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  async text2img(params: GenerationRequest): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/text2img', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  async img2img(params: GenerationRequest & { image: string }): Promise<JobInfo> {
    // Convert frontend field name to backend field name
    const backendParams = {
      ...params,
      init_image: params.image,
      image: undefined, // Remove the frontend field
    };
    return this.request<JobInfo>('/generate/img2img', {
      method: 'POST',
      body: JSON.stringify(backendParams),
    });
  }

  async inpainting(params: GenerationRequest & { source_image: string; mask_image: string }): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/inpainting', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  // Job management
  async getJobStatus(jobId: string): Promise<JobInfo> {
    return this.request<JobInfo>(`/queue/job/${jobId}`);
  }

  // Get authenticated image URL with cache-busting
  getImageUrl(jobId: string, filename: string): string {
    const { apiUrl, apiBase } = getApiConfig();
    const baseUrl = `${apiUrl}${apiBase}`;

    // Add cache-busting timestamp
    const timestamp = Date.now();
    const url = `${baseUrl}/queue/job/${jobId}/output/${filename}?t=${timestamp}`;

    return url;
  }

  // Download image with authentication
  async downloadImage(jobId: string, filename: string): Promise<Blob> {
    const url = this.getImageUrl(jobId, filename);

    // Get authentication method from server config
    const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
      ? window.__SERVER_CONFIG__.authMethod
      : 'jwt';

    // Add auth token or Unix user header based on auth method
    const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
    const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

    const headers: Record<string, string> = {};

    if (authMethod === 'unix' && unixUser) {
      // For Unix auth, send the username in X-Unix-User header
      headers['X-Unix-User'] = unixUser;
    } else if (token) {
      // For JWT auth, send the token in Authorization header
      headers['Authorization'] = `Bearer ${token}`;
    }

    const response = await fetch(url, {
      headers,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({
        error: { message: response.statusText },
      }));

      // Handle nested error structure: { error: { message: "..." } }
      const errorMessage =
        errorData.error?.message ||
        errorData.message ||
        errorData.error ||
        'Failed to download image';

      throw new Error(errorMessage);
    }

    return response.blob();
  }

  async cancelJob(jobId: string): Promise<void> {
    return this.request<void>('/queue/cancel', {
      method: 'POST',
      body: JSON.stringify({ job_id: jobId }),
    });
  }

  async getQueueStatus(): Promise<QueueStatus> {
    const response = await this.request<{ queue: QueueStatus }>('/queue/status');
    return response.queue;
  }

  async clearQueue(): Promise<void> {
    return this.request<void>('/queue/clear', {
      method: 'POST',
    });
  }

  // Model management
  async getModels(type?: string, loaded?: boolean): Promise<ModelInfo[]> {
    let endpoint = '/models';
    const params = [];
    if (type && type !== 'loaded') params.push(`type=${type}`);
    if (type === 'loaded' || loaded) params.push('loaded=true');
    // Request a high limit to get all models (default is 50)
    params.push('limit=1000');
    if (params.length > 0) endpoint += '?' + params.join('&');

    const response = await this.request<{ models: ModelInfo[] }>(endpoint);
    // Add id field based on sha256_short or name, and normalize size field
    return response.models.map(model => ({
      ...model,
      id: model.sha256_short || model.name,
      size: model.file_size || model.size,
      path: model.path || model.name,
    }));
  }

  async getModelInfo(modelId: string): Promise<ModelInfo> {
    return this.request<ModelInfo>(`/models/${modelId}`);
  }

  async loadModel(modelId: string): Promise<void> {
    return this.request<void>(`/models/${modelId}/load`, {
      method: 'POST',
    });
  }

  async unloadModel(modelId: string): Promise<void> {
    return this.request<void>(`/models/${modelId}/unload`, {
      method: 'POST',
    });
  }

  async scanModels(): Promise<void> {
    return this.request<void>('/models/refresh', {
      method: 'POST',
    });
  }

  async convertModel(modelName: string, quantizationType: string, outputPath?: string): Promise<{ request_id: string; message: string }> {
    return this.request<{ request_id: string; message: string }>('/models/convert', {
      method: 'POST',
      body: JSON.stringify({
        model_name: modelName,
        quantization_type: quantizationType,
        output_path: outputPath,
      }),
    });
  }

  // System endpoints
  async getHealth(): Promise<{ status: string }> {
    return this.request<{ status: string }>('/health');
  }

  async getStatus(): Promise<any> {
    return this.request<any>('/status');
  }

  async getSystemInfo(): Promise<any> {
    return this.request<any>('/system');
  }

  async restartServer(): Promise<{ message: string }> {
    return this.request<{ message: string }>('/system/restart', {
      method: 'POST',
      body: JSON.stringify({}),
    });
  }

  // Configuration endpoints
  async getSamplers(): Promise<Array<{ name: string; description: string; recommended_steps: number }>> {
    const response = await this.request<{ samplers: Array<{ name: string; description: string; recommended_steps: number }> }>('/samplers');
    return response.samplers;
  }

  async getSchedulers(): Promise<Array<{ name: string; description: string }>> {
    const response = await this.request<{ schedulers: Array<{ name: string; description: string }> }>('/schedulers');
    return response.schedulers;
  }
}

// Generic API request function for authentication
export async function apiRequest(
  endpoint: string,
  options: RequestInit = {}
): Promise<Response> {
  const { apiUrl, apiBase } = getApiConfig();
  const url = `${apiUrl}${apiBase}${endpoint}`;

  // Get authentication method from server config
  const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
    ? window.__SERVER_CONFIG__.authMethod
    : 'jwt';

  // Add auth token or Unix user header based on auth method
  const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
  const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...options.headers as Record<string, string>,
  };

  if (authMethod === 'unix' && unixUser) {
    // For Unix auth, send the username in X-Unix-User header
    headers['X-Unix-User'] = unixUser;
  } else if (token) {
    // For JWT auth, send the token in Authorization header
    headers['Authorization'] = `Bearer ${token}`;
  }

  return fetch(url, {
    ...options,
    headers,
  });
}

// Authentication API endpoints
export const authApi = {
  async login(username: string, password?: string) {
    // Get authentication method from server config
    const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
      ? window.__SERVER_CONFIG__.authMethod
      : 'jwt';

    // For both Unix and JWT auth, send username and password
    // The server will handle whether password is required based on PAM availability
    const response = await apiRequest('/auth/login', {
      method: 'POST',
      body: JSON.stringify({ username, password }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Login failed' }));
      throw new Error(error.message || 'Login failed');
    }

    return response.json();
  },

  async validateToken(token: string) {
    const response = await apiRequest('/auth/validate', {
      headers: { 'Authorization': `Bearer ${token}` },
    });

    if (!response.ok) {
      throw new Error('Token validation failed');
    }

    return response.json();
  },

  async refreshToken() {
    const response = await apiRequest('/auth/refresh', {
      method: 'POST',
    });

    if (!response.ok) {
      throw new Error('Token refresh failed');
    }

    return response.json();
  },

  async logout() {
    await apiRequest('/auth/logout', {
      method: 'POST',
    });
  },

  async getCurrentUser() {
    const response = await apiRequest('/auth/me');

    if (!response.ok) {
      throw new Error('Failed to get current user');
    }

    return response.json();
  },

  async changePassword(oldPassword: string, newPassword: string) {
    const response = await apiRequest('/auth/change-password', {
      method: 'POST',
      body: JSON.stringify({ old_password: oldPassword, new_password: newPassword }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Password change failed' }));
      throw new Error(error.message || 'Password change failed');
    }

    return response.json();
  },

  // API Key management
  async getApiKeys() {
    const response = await apiRequest('/auth/api-keys');

    if (!response.ok) {
      throw new Error('Failed to get API keys');
    }

    return response.json();
  },

  async createApiKey(name: string, scopes?: string[]) {
    const response = await apiRequest('/auth/api-keys', {
      method: 'POST',
      body: JSON.stringify({ name, scopes }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to create API key' }));
      throw new Error(error.message || 'Failed to create API key');
    }

    return response.json();
  },

  async deleteApiKey(keyId: string) {
    const response = await apiRequest(`/auth/api-keys/${keyId}`, {
      method: 'DELETE',
    });

    if (!response.ok) {
      throw new Error('Failed to delete API key');
    }

    return response.json();
  },

  // User management (admin only)
  async getUsers() {
    const response = await apiRequest('/auth/users');

    if (!response.ok) {
      throw new Error('Failed to get users');
    }

    return response.json();
  },

  async createUser(userData: { username: string; email?: string; password: string; role?: string }) {
    const response = await apiRequest('/auth/users', {
      method: 'POST',
      body: JSON.stringify(userData),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to create user' }));
      throw new Error(error.message || 'Failed to create user');
    }

    return response.json();
  },

  async updateUser(userId: string, userData: { email?: string; role?: string; active?: boolean }) {
    const response = await apiRequest(`/auth/users/${userId}`, {
      method: 'PUT',
      body: JSON.stringify(userData),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to update user' }));
      throw new Error(error.message || 'Failed to update user');
    }

    return response.json();
  },

  async deleteUser(userId: string) {
    const response = await apiRequest(`/auth/users/${userId}`, {
      method: 'DELETE',
    });

    if (!response.ok) {
      throw new Error('Failed to delete user');
    }

    return response.json();
  }
};

export const apiClient = new ApiClient();