claudecodeui/server/modules/providers/list/gemini/gemini.provider.ts

import sessionManager from '@/sessionManager.js';
import { getGeminiCliSessionMessages } from '@/projects.js';
import { AbstractProvider } from '@/modules/providers/shared/base/abstract.provider.js';
import { GeminiProviderAuth } from '@/modules/providers/list/gemini/gemini-auth.provider.js';
import { GeminiMcpProvider } from '@/modules/providers/list/gemini/gemini-mcp.provider.js';
import type { IProviderAuth } from '@/shared/interfaces.js';
import type { FetchHistoryOptions, FetchHistoryResult, NormalizedMessage } from '@/shared/types.js';
import { createNormalizedMessage, generateMessageId, readObjectRecord } from '@/shared/utils.js';

const PROVIDER = 'gemini';

type RawProviderMessage = Record<string, any>;

function readRawProviderMessage(raw: unknown): RawProviderMessage | null {
  return readObjectRecord(raw) as RawProviderMessage | null;
}

export class GeminiProvider extends AbstractProvider {
  readonly mcp = new GeminiMcpProvider();
  readonly auth: IProviderAuth = new GeminiProviderAuth();

  constructor() {
    super('gemini');
  }

  /**
   * Normalizes live Gemini stream-json events into the shared message shape.
   *
   * Gemini history uses a different session file shape, so fetchHistory handles
   * that separately after loading raw persisted messages.
   */
  normalizeMessage(rawMessage: unknown, sessionId: string | null): NormalizedMessage[] {
    const raw = readRawProviderMessage(rawMessage);
    if (!raw) {
      return [];
    }

    const ts = raw.timestamp || new Date().toISOString();
    const baseId = raw.uuid || generateMessageId('gemini');

    if (raw.type === 'message' && raw.role === 'assistant') {
      const content = raw.content || '';
      const messages: NormalizedMessage[] = [];
      if (content) {
        messages.push(createNormalizedMessage({
          id: baseId,
          sessionId,
          timestamp: ts,
          provider: PROVIDER,
          kind: 'stream_delta',
          content,
        }));
      }
      if (raw.delta !== true) {
        messages.push(createNormalizedMessage({
          sessionId,
          timestamp: ts,
          provider: PROVIDER,
          kind: 'stream_end',
        }));
      }
      return messages;
    }

    if (raw.type === 'tool_use') {
      return [createNormalizedMessage({
        id: baseId,
        sessionId,
        timestamp: ts,
        provider: PROVIDER,
        kind: 'tool_use',
        toolName: raw.tool_name,
        toolInput: raw.parameters || {},
        toolId: raw.tool_id || baseId,
      })];
    }

    if (raw.type === 'tool_result') {
      return [createNormalizedMessage({
        id: baseId,
        sessionId,
        timestamp: ts,
        provider: PROVIDER,
        kind: 'tool_result',
        toolId: raw.tool_id || '',
        content: raw.output === undefined ? '' : String(raw.output),
        isError: raw.status === 'error',
      })];
    }

    if (raw.type === 'result') {
      const messages = [createNormalizedMessage({
        sessionId,
        timestamp: ts,
        provider: PROVIDER,
        kind: 'stream_end',
      })];
      if (raw.stats?.total_tokens) {
        messages.push(createNormalizedMessage({
          sessionId,
          timestamp: ts,
          provider: PROVIDER,
          kind: 'status',
          text: 'Complete',
          tokens: raw.stats.total_tokens,
          canInterrupt: false,
        }));
      }
      return messages;
    }

    if (raw.type === 'error') {
      return [createNormalizedMessage({
        id: baseId,
        sessionId,
        timestamp: ts,
        provider: PROVIDER,
        kind: 'error',
        content: raw.error || raw.message || 'Unknown Gemini streaming error',
      })];
    }

    return [];
  }

  /**
   * Loads Gemini history from the in-memory session manager first, then falls
   * back to Gemini CLI session files on disk.
   */
  async fetchHistory(
    sessionId: string,
    _options: FetchHistoryOptions = {},
  ): Promise<FetchHistoryResult> {
    let rawMessages: RawProviderMessage[];
    try {
      rawMessages = sessionManager.getSessionMessages(sessionId) as RawProviderMessage[];

      if (rawMessages.length === 0) {
        rawMessages = await getGeminiCliSessionMessages(sessionId) as RawProviderMessage[];
      }
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      console.warn(`[GeminiProvider] Failed to load session ${sessionId}:`, message);
      return { messages: [], total: 0, hasMore: false, offset: 0, limit: null };
    }

    const normalized: NormalizedMessage[] = [];
    for (let i = 0; i < rawMessages.length; i++) {
      const raw = rawMessages[i];
      const ts = raw.timestamp || new Date().toISOString();
      const baseId = raw.uuid || generateMessageId('gemini');

      const role = raw.message?.role || raw.role;
      const content = raw.message?.content || raw.content;

      if (!role || !content) {
        continue;
      }

      const normalizedRole = role === 'user' ? 'user' : 'assistant';

      if (Array.isArray(content)) {
        for (let partIdx = 0; partIdx < content.length; partIdx++) {
          const part = content[partIdx];
          if (part.type === 'text' && part.text) {
            normalized.push(createNormalizedMessage({
              id: `${baseId}_${partIdx}`,
              sessionId,
              timestamp: ts,
              provider: PROVIDER,
              kind: 'text',
              role: normalizedRole,
              content: part.text,
            }));
          } else if (part.type === 'tool_use') {
            normalized.push(createNormalizedMessage({
              id: `${baseId}_${partIdx}`,
              sessionId,
              timestamp: ts,
              provider: PROVIDER,
              kind: 'tool_use',
              toolName: part.name,
              toolInput: part.input,
              toolId: part.id || generateMessageId('gemini_tool'),
            }));
          } else if (part.type === 'tool_result') {
            normalized.push(createNormalizedMessage({
              id: `${baseId}_${partIdx}`,
              sessionId,
              timestamp: ts,
              provider: PROVIDER,
              kind: 'tool_result',
              toolId: part.tool_use_id || '',
              content: part.content === undefined ? '' : String(part.content),
              isError: Boolean(part.is_error),
            }));
          }
        }
      } else if (typeof content === 'string' && content.trim()) {
        normalized.push(createNormalizedMessage({
          id: baseId,
          sessionId,
          timestamp: ts,
          provider: PROVIDER,
          kind: 'text',
          role: normalizedRole,
          content,
        }));
      }
    }

    const toolResultMap = new Map<string, NormalizedMessage>();
    for (const msg of normalized) {
      if (msg.kind === 'tool_result' && msg.toolId) {
        toolResultMap.set(msg.toolId, msg);
      }
    }
    for (const msg of normalized) {
      if (msg.kind === 'tool_use' && msg.toolId && toolResultMap.has(msg.toolId)) {
        const toolResult = toolResultMap.get(msg.toolId);
        if (toolResult) {
          msg.toolResult = { content: toolResult.content, isError: toolResult.isError };
        }
      }
    }

    return {
      messages: normalized,
      total: normalized.length,
      hasMore: false,
      offset: 0,
      limit: null,
    };
  }
}