openpencil/apps/web/server/utils/gemini-client.ts

import { spawn } from 'node:child_process'
import { resolveGeminiCli } from './resolve-gemini-cli'

type ThinkingMode = 'adaptive' | 'disabled' | 'enabled'
type ThinkingEffort = 'low' | 'medium' | 'high' | 'max'

export interface GeminiExecOptions {
  model?: string
  systemPrompt?: string
  thinkingMode?: ThinkingMode
  thinkingBudgetTokens?: number
  effort?: ThinkingEffort
  timeoutMs?: number
}

interface GeminiCliResult {
  text?: string
  error?: string
}

const DEFAULT_GEMINI_TIMEOUT_MS = 15 * 60 * 1000

/**
 * Allowlist-based env filter for Gemini CLI subprocess.
 * Passes through safe system vars and Google/Gemini-specific prefixes.
 */
const GEMINI_ENV_ALLOWLIST = new Set([
  'PATH', 'HOME', 'TERM', 'LANG', 'SHELL', 'TMPDIR',
  // Windows-essential
  'SYSTEMROOT', 'COMSPEC', 'USERPROFILE', 'APPDATA', 'LOCALAPPDATA',
  'PATHEXT', 'SYSTEMDRIVE', 'TEMP', 'TMP', 'HOMEDRIVE', 'HOMEPATH',
])

function filterGeminiEnv(
  env: Record<string, string | undefined>,
): Record<string, string | undefined> {
  const result: Record<string, string | undefined> = {}
  for (const [k, v] of Object.entries(env)) {
    if (
      GEMINI_ENV_ALLOWLIST.has(k)
      || k.startsWith('GOOGLE_')
      || k.startsWith('GEMINI_')
      || k.startsWith('GCLOUD_')
    ) {
      result[k] = v
    }
  }
  return result
}

/**
 * Run Gemini CLI in non-interactive mode with JSON output.
 * Passes prompt via stdin to avoid command-line length limits.
 * The CLI handles its own authentication (OAuth or API key).
 */
export async function runGeminiExec(
  userPrompt: string,
  options: GeminiExecOptions = {},
): Promise<GeminiCliResult> {
  const binPath = resolveGeminiCli()
  if (!binPath) {
    return { error: 'Gemini CLI not found. Install it first.' }
  }

  const prompt = buildPrompt(options.systemPrompt, userPrompt)

  const args = [
    '-o', 'json',
    '--approval-mode', 'plan',
  ]

  if (options.model) {
    args.push('-m', options.model)
  }

  // Use -p with a minimal marker; full prompt piped via stdin.
  // Gemini CLI appends -p value after stdin content.
  args.push('-p', ' ')

  try {
    const result = await executeGeminiCommand(
      binPath,
      args,
      options.timeoutMs ?? DEFAULT_GEMINI_TIMEOUT_MS,
      prompt,
    )
    return result
  } catch (error) {
    return { error: error instanceof Error ? error.message : 'Gemini execution failed' }
  }
}

/**
 * Stream Gemini CLI output in real-time using `stream-json` format.
 * Passes prompt via stdin. Yields text deltas as they arrive.
 */
export function streamGeminiExec(
  userPrompt: string,
  options: GeminiExecOptions = {},
): {
  stream: AsyncGenerator<{ type: 'text' | 'error' | 'done'; content: string }>
  kill: () => void
} {
  const binPath = resolveGeminiCli()
  if (!binPath) {
    return {
      stream: (async function* () {
        yield { type: 'error' as const, content: 'Gemini CLI not found.' }
      })(),
      kill: () => {},
    }
  }

  const prompt = buildPrompt(options.systemPrompt, userPrompt)

  const args = [
    '-o', 'stream-json',
    '--approval-mode', 'plan',
  ]

  if (options.model) {
    args.push('-m', options.model)
  }

  // Use -p with minimal marker; full prompt piped via stdin.
  args.push('-p', ' ')

  const child = spawn(binPath, args, {
    env: filterGeminiEnv(process.env as Record<string, string | undefined>),
    stdio: ['pipe', 'pipe', 'pipe'],
    ...(process.platform === 'win32' && { shell: true }),
  })

  // Pipe prompt via stdin
  if (child.stdin) {
    child.stdin.write(prompt)
    child.stdin.end()
  }

  const timeoutMs = options.timeoutMs ?? DEFAULT_GEMINI_TIMEOUT_MS
  const timer = setTimeout(() => {
    child.kill('SIGTERM')
  }, timeoutMs)

  async function* generateStream(): AsyncGenerator<{ type: 'text' | 'error' | 'done'; content: string }> {
    let buffer = ''

    child.stderr?.on('data', () => { /* discard stderr */ })

    try {
      for await (const chunk of child.stdout!) {
        buffer += chunk.toString('utf-8')
        let idx = buffer.indexOf('\n')
        while (idx >= 0) {
          const line = buffer.slice(0, idx).trim()
          buffer = buffer.slice(idx + 1)
          if (line) {
            const event = parseStreamJsonLine(line)
            if (event) yield event
          }
          idx = buffer.indexOf('\n')
        }
      }

      // Flush remaining buffer
      const tail = buffer.trim()
      if (tail) {
        const event = parseStreamJsonLine(tail)
        if (event) yield event
      }

      yield { type: 'done', content: '' }
    } catch (error) {
      const msg = error instanceof Error ? error.message : 'Stream error'
      yield { type: 'error', content: msg }
    } finally {
      clearTimeout(timer)
    }
  }

  return {
    stream: generateStream(),
    kill: () => {
      clearTimeout(timer)
      child.kill('SIGTERM')
    },
  }
}

function buildPrompt(systemPrompt: string | undefined, userPrompt: string): string {
  const userText = userPrompt.trim()
  if (!systemPrompt?.trim()) return userText

  return [
    'You are a design generation assistant. Follow the guidelines below to produce the requested output.',
    '',
    '--- GUIDELINES ---',
    systemPrompt.trim(),
    '',
    '--- TASK ---',
    userText,
  ].join('\n')
}

async function executeGeminiCommand(
  binPath: string,
  args: string[],
  timeoutMs: number,
  stdinText?: string,
): Promise<GeminiCliResult> {
  return await new Promise((resolve, reject) => {
    const child = spawn(binPath, args, {
      env: filterGeminiEnv(process.env as Record<string, string | undefined>),
      stdio: [stdinText ? 'pipe' : 'ignore', 'pipe', 'pipe'],
      ...(process.platform === 'win32' && { shell: true }),
    })

    // Pipe prompt via stdin
    if (stdinText && child.stdin) {
      child.stdin.write(stdinText)
      child.stdin.end()
    }

    let stdoutBuffer = ''
    let stderrBuffer = ''

    const timer = setTimeout(() => {
      child.kill('SIGTERM')
      reject(new Error(`Gemini request timed out after ${Math.round(timeoutMs / 1000)}s.`))
    }, timeoutMs)

    child.stdout!.on('data', (chunk: Buffer) => {
      stdoutBuffer += chunk.toString('utf-8')
    })

    child.stderr!.on('data', (chunk: Buffer) => {
      stderrBuffer += chunk.toString('utf-8')
    })

    child.on('error', (err) => {
      clearTimeout(timer)
      reject(err)
    })

    child.on('close', (code) => {
      clearTimeout(timer)

      // Parse JSON output — Gemini CLI always outputs a JSON object at the end of stdout.
      // Error text / stack traces may appear before it.
      const parsed = parseGeminiJsonOutput(stdoutBuffer)

      if (parsed) {
        if (parsed.response) {
          resolve({ text: parsed.response })
          return
        }
        if (parsed.errorMessage) {
          resolve({ error: friendlyGeminiApiError(parsed.errorMessage) })
          return
        }
      }

      if (code !== 0) {
        // Extract meaningful error from stderr or stdout
        const errorMsg = extractGeminiError(stdoutBuffer, stderrBuffer)
        resolve({ error: errorMsg || `Gemini exited with code ${code ?? 'unknown'}.` })
        return
      }

      const raw = stdoutBuffer.trim()
      resolve(raw ? { text: raw } : { error: 'Gemini returned no output.' })
    })
  })
}

/**
 * Parse Gemini CLI JSON output.
 * The CLI may print error text before the final JSON object.
 * We search from the END of stdout for the last valid JSON block.
 */
function parseGeminiJsonOutput(raw: string): { response?: string; errorMessage?: string } | null {
  const trimmed = raw.trim()
  if (!trimmed) return null

  // Search backwards for the last top-level JSON object (starts with `{` at line beginning)
  const lines = trimmed.split('\n')
  for (let i = lines.length - 1; i >= 0; i--) {
    const line = lines[i].trim()
    if (!line.startsWith('{')) continue

    // Try to parse from this line to the end
    const candidate = lines.slice(i).join('\n').trim()
    try {
      const data = JSON.parse(candidate) as Record<string, unknown>
      // Must have session_id to be a valid Gemini CLI response
      if (!data.session_id && !data.response && !data.error) continue

      const response = typeof data.response === 'string' ? data.response : undefined

      // error can be a string or an object { type, message, code }
      let errorMessage: string | undefined
      if (data.error) {
        if (typeof data.error === 'string') {
          errorMessage = data.error
        } else if (typeof data.error === 'object' && data.error !== null) {
          const errObj = data.error as Record<string, unknown>
          errorMessage = typeof errObj.message === 'string' ? errObj.message : JSON.stringify(data.error)
        }
      }

      return { response, errorMessage }
    } catch { /* not valid JSON from this point */ }
  }

  return null
}

/** Extract a human-readable error from Gemini CLI stdout/stderr */
function extractGeminiError(stdout: string, stderr: string): string | null {
  // Look for quota errors
  const quotaMatch = stdout.match(/quota will reset after (\S+)/i)
    || stderr.match(/quota will reset after (\S+)/i)
  if (quotaMatch) {
    return `Gemini quota exhausted. Resets after ${quotaMatch[1]}.`
  }

  // Look for TerminalQuotaError or other named errors
  const namedError = stdout.match(/(Terminal\w+Error|ApiError|AuthError):\s*(.+)/m)
  if (namedError) {
    return namedError[2].trim()
  }

  // Stderr fallback
  const stderrTrimmed = stderr.trim()
  if (stderrTrimmed) return stderrTrimmed

  return null
}

/** Map raw Gemini API errors to user-friendly messages */
function friendlyGeminiApiError(raw: string): string {
  if (/quota|exhausted|429|capacity/i.test(raw)) {
    const resetMatch = raw.match(/reset after (\S+)/i)
    return resetMatch
      ? `Gemini quota exhausted. Resets after ${resetMatch[1]}.`
      : 'Gemini quota exhausted. Please wait and try again.'
  }
  if (/401|unauthenticated|auth/i.test(raw)) {
    return 'Gemini auth expired. Run "gemini" in your terminal to re-authenticate.'
  }
  if (/\[object Object\]/.test(raw)) {
    return 'Gemini API error. Check your quota or try a different model.'
  }
  return raw
}

function parseStreamJsonLine(
  line: string,
): { type: 'text' | 'error' | 'done'; content: string } | null {
  // Skip non-JSON lines (e.g. "Loaded cached credentials.")
  if (!line.startsWith('{')) return null

  let parsed: Record<string, unknown>
  try {
    parsed = JSON.parse(line) as Record<string, unknown>
  } catch {
    return null
  }

  const type = typeof parsed.type === 'string' ? parsed.type : ''

  if (type === 'message' && parsed.role === 'assistant') {
    const content = typeof parsed.content === 'string' ? parsed.content : ''
    if (content) return { type: 'text', content }
  }

  if (type === 'result') {
    // Check for error in result event
    if (parsed.status === 'error' && parsed.error) {
      const errObj = parsed.error as Record<string, unknown>
      const msg = typeof errObj.message === 'string' ? errObj.message : 'Unknown error'
      return { type: 'error', content: friendlyGeminiApiError(msg) }
    }
    return null
  }

  if (type === 'error') {
    const content = typeof parsed.message === 'string' ? parsed.message : 'Unknown error'
    return { type: 'error', content: friendlyGeminiApiError(content) }
  }

  return null
}