repoburg 1.0.68 → 1.1.0

package/backend/packages/gemini-core/src/core/contentGenerator.ts

@@ -0,0 +1,136 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type {
+  CountTokensResponse,
+  GenerateContentResponse,
+  GenerateContentParameters,
+  CountTokensParameters,
+  EmbedContentResponse,
+  EmbedContentParameters,
+} from '@google/genai';
+import { GoogleGenAI } from '@google/genai';
+import { createCodeAssistContentGenerator } from '../code_assist/codeAssist';
+import type { Config } from '../config/config';
+
+import type { UserTierId } from '../code_assist/types';
+
+/**
+ * Interface abstracting the core functionalities for generating content and counting tokens.
+ */
+export interface ContentGenerator {
+  generateContent(
+    request: GenerateContentParameters,
+    userPromptId: string,
+  ): Promise<GenerateContentResponse>;
+
+  generateContentStream(
+    request: GenerateContentParameters,
+    userPromptId: string,
+  ): Promise<AsyncGenerator<GenerateContentResponse>>;
+
+  countTokens(request: CountTokensParameters): Promise<CountTokensResponse>;
+
+  embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse>;
+
+  userTier?: UserTierId;
+}
+
+export enum AuthType {
+  LOGIN_WITH_GOOGLE = 'oauth-personal',
+  USE_GEMINI = 'gemini-api-key',
+  USE_VERTEX_AI = 'vertex-ai',
+  CLOUD_SHELL = 'cloud-shell',
+}
+
+export type ContentGeneratorConfig = {
+  apiKey?: string;
+  vertexai?: boolean;
+  authType?: AuthType;
+};
+
+export function createContentGeneratorConfig(
+  config: Config,
+  authType: AuthType | undefined,
+): ContentGeneratorConfig {
+  const geminiApiKey = process.env['GEMINI_API_KEY'] || undefined;
+  const googleApiKey = process.env['GOOGLE_API_KEY'] || undefined;
+  const googleCloudProject = process.env['GOOGLE_CLOUD_PROJECT'] || undefined;
+  const googleCloudLocation = process.env['GOOGLE_CLOUD_LOCATION'] || undefined;
+
+  const contentGeneratorConfig: ContentGeneratorConfig = {
+    authType,
+  };
+
+  // If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
+  if (
+    authType === AuthType.LOGIN_WITH_GOOGLE ||
+    authType === AuthType.CLOUD_SHELL
+  ) {
+    return contentGeneratorConfig;
+  }
+
+  if (authType === AuthType.USE_GEMINI && geminiApiKey) {
+    contentGeneratorConfig.apiKey = geminiApiKey;
+    contentGeneratorConfig.vertexai = false;
+
+    return contentGeneratorConfig;
+  }
+
+  if (
+    authType === AuthType.USE_VERTEX_AI &&
+    (googleApiKey || (googleCloudProject && googleCloudLocation))
+  ) {
+    contentGeneratorConfig.apiKey = googleApiKey;
+    contentGeneratorConfig.vertexai = true;
+
+    return contentGeneratorConfig;
+  }
+
+  return contentGeneratorConfig;
+}
+
+export async function createContentGenerator(
+  config: ContentGeneratorConfig,
+  gcConfig: Config,
+  sessionId?: string,
+): Promise<ContentGenerator> {
+  const version = process.env['CLI_VERSION'] || process.version;
+  const userAgent = `GeminiCLI/${version} (${process.platform}; ${process.arch})`;
+  const baseHeaders: Record<string, string> = {
+    'User-Agent': userAgent,
+  };
+
+  if (
+    config.authType === AuthType.LOGIN_WITH_GOOGLE ||
+    config.authType === AuthType.CLOUD_SHELL
+  ) {
+    const httpOptions = { headers: baseHeaders };
+    return createCodeAssistContentGenerator(
+      httpOptions,
+      config.authType,
+      gcConfig,
+      sessionId,
+    );
+  }
+
+  if (
+    config.authType === AuthType.USE_GEMINI ||
+    config.authType === AuthType.USE_VERTEX_AI
+  ) {
+    const httpOptions = { headers: baseHeaders };
+
+    const googleGenAI = new GoogleGenAI({
+      apiKey: config.apiKey === '' ? undefined : config.apiKey,
+      vertexai: config.vertexai,
+      httpOptions,
+    });
+    return googleGenAI.models;
+  }
+  throw new Error(
+    `Error creating contentGenerator: Unsupported authType: ${config.authType}`,
+  );
+}

package/backend/packages/gemini-core/src/core/geminiChat.ts

@@ -0,0 +1,385 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// DISCLAIMER: This is a copied version of https://github.com/googleapis/js-genai/blob/main/src/chats.ts with the intention of working around a key bug
+// where function responses are not treated as "valid" responses: https://b.corp.google.com/issues/420354090
+
+import {
+  GenerateContentResponse,
+  type Content,
+  type GenerateContentConfig,
+  type SendMessageParameters,
+  type Part,
+  ApiError,
+} from '@google/genai';
+import { createUserContent } from '@google/genai';
+import { retryWithBackoff } from '../utils/retry';
+import type { Config } from '../config/config';
+import { getEffectiveModel } from '../config/models';
+
+export enum StreamEventType {
+  /** A regular content chunk from the API. */
+  CHUNK = 'chunk',
+  /** A signal that a retry is about to happen. The UI should discard any partial
+   * content from the attempt that just failed. */
+  RETRY = 'retry',
+}
+
+export type StreamEvent =
+  | { type: StreamEventType.CHUNK; value: GenerateContentResponse }
+  | { type: StreamEventType.RETRY };
+
+/**
+ * Options for retrying due to invalid content from the model.
+ */
+interface ContentRetryOptions {
+  /** Total number of attempts to make (1 initial + N retries). */
+  maxAttempts: number;
+  /** The base delay in milliseconds for linear backoff. */
+  initialDelayMs: number;
+}
+
+const INVALID_CONTENT_RETRY_OPTIONS: ContentRetryOptions = {
+  maxAttempts: 2, // 1 initial call + 1 retry
+  initialDelayMs: 500,
+};
+
+/**
+ * Returns true if the response is valid, false otherwise.
+ */
+function isValidResponse(response: GenerateContentResponse): boolean {
+  if (response.candidates === undefined || response.candidates.length === 0) {
+    return false;
+  }
+  const content = response.candidates[0]?.content;
+  if (content === undefined) {
+    return false;
+  }
+  return isValidContent(content);
+}
+
+export function isValidNonThoughtTextPart(part: Part): boolean {
+  return (
+    typeof part.text === 'string' &&
+    !part.thought &&
+    // Technically, the model should never generate parts that have text and
+    //  any of these but we don't trust them so check anyways.
+    !part.inlineData &&
+    !part.fileData
+  );
+}
+
+function isValidContent(content: Content): boolean {
+  if (content.parts === undefined || content.parts.length === 0) {
+    return false;
+  }
+  for (const part of content.parts) {
+    if (part === undefined || Object.keys(part).length === 0) {
+      return false;
+    }
+    if (!part.thought && part.text !== undefined && part.text === '') {
+      return false;
+    }
+  }
+  return true;
+}
+
+/**
+ * Validates the history contains the correct roles.
+ *
+ * @throws Error if the history does not start with a user turn.
+ * @throws Error if the history contains an invalid role.
+ */
+function validateHistory(history: Content[]) {
+  for (const content of history) {
+    if (content.role !== 'user' && content.role !== 'model') {
+      throw new Error(`Role must be user or model, but got ${content.role}.`);
+    }
+  }
+}
+
+/**
+ * Extracts the curated (valid) history from a comprehensive history.
+ *
+ * @remarks
+ * The model may sometimes generate invalid or empty contents(e.g., due to safety
+ * filters or recitation). Extracting valid turns from the history
+ * ensures that subsequent requests could be accepted by the model.
+ */
+function extractCuratedHistory(comprehensiveHistory: Content[]): Content[] {
+  if (comprehensiveHistory === undefined || comprehensiveHistory.length === 0) {
+    return [];
+  }
+  const curatedHistory: Content[] = [];
+  const length = comprehensiveHistory.length;
+  let i = 0;
+  while (i < length) {
+    if (comprehensiveHistory[i].role === 'user') {
+      curatedHistory.push(comprehensiveHistory[i]);
+      i++;
+    } else {
+      const modelOutput: Content[] = [];
+      let isValid = true;
+      while (i < length && comprehensiveHistory[i].role === 'model') {
+        modelOutput.push(comprehensiveHistory[i]);
+        if (isValid && !isValidContent(comprehensiveHistory[i])) {
+          isValid = false;
+        }
+        i++;
+      }
+      if (isValid) {
+        curatedHistory.push(...modelOutput);
+      }
+    }
+  }
+  return curatedHistory;
+}
+
+/**
+ * Custom error to signal that a stream completed with invalid content,
+ * which should trigger a retry.
+ */
+export class InvalidStreamError extends Error {
+  readonly type: 'NO_FINISH_REASON' | 'NO_RESPONSE_TEXT';
+
+  constructor(message: string, type: 'NO_FINISH_REASON' | 'NO_RESPONSE_TEXT') {
+    super(message);
+    this.name = 'InvalidStreamError';
+    this.type = type;
+  }
+}
+
+/**
+ * Chat session that enables sending messages to the model with previous
+ * conversation context.
+ *
+ * @remarks
+ * The session maintains all the turns between user and model.
+ */
+export class GeminiChat {
+  // A promise to represent the current state of the message being sent to the
+  // model.
+  private sendPromise: Promise<void> = Promise.resolve();
+
+  constructor(
+    private readonly config: Config,
+    private readonly generationConfig: GenerateContentConfig = {},
+    private history: Content[] = [],
+  ) {
+    validateHistory(history);
+  }
+
+  setSystemInstruction(sysInstr: string) {
+    this.generationConfig.systemInstruction = sysInstr;
+  }
+
+  async sendMessageStream(
+    model: string,
+    params: SendMessageParameters,
+    prompt_id: string,
+  ): Promise<AsyncGenerator<StreamEvent>> {
+    await this.sendPromise;
+
+    let streamDoneResolver: () => void;
+    const streamDonePromise = new Promise<void>((resolve) => {
+      streamDoneResolver = resolve;
+    });
+    this.sendPromise = streamDonePromise;
+
+    const userContent = createUserContent(params.message);
+
+    // Add user content to history ONCE before any attempts.
+    this.history.push(userContent);
+    const requestContents = this.getHistory(true);
+
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const self = this;
+    return (async function* () {
+      try {
+        let lastError: unknown = new Error('Request failed after all retries.');
+
+        for (
+          let attempt = 0;
+          attempt < INVALID_CONTENT_RETRY_OPTIONS.maxAttempts;
+          attempt++
+        ) {
+          try {
+            if (attempt > 0) {
+              yield { type: StreamEventType.RETRY };
+            }
+
+            const stream = await self.makeApiCallAndProcessStream(
+              model,
+              requestContents,
+              params,
+              prompt_id,
+            );
+
+            for await (const chunk of stream) {
+              yield { type: StreamEventType.CHUNK, value: chunk };
+            }
+
+            lastError = null;
+            break;
+          } catch (error) {
+            lastError = error;
+            const isContentError = error instanceof InvalidStreamError;
+
+            if (isContentError) {
+              // Check if we have more attempts left.
+              if (attempt < INVALID_CONTENT_RETRY_OPTIONS.maxAttempts - 1) {
+                await new Promise((res) =>
+                  setTimeout(
+                    res,
+                    INVALID_CONTENT_RETRY_OPTIONS.initialDelayMs *
+                      (attempt + 1),
+                  ),
+                );
+                continue;
+              }
+            }
+            break;
+          }
+        }
+
+        if (lastError) {
+          if (self.history[self.history.length - 1] === userContent) {
+            self.history.pop();
+          }
+          throw lastError;
+        }
+      } finally {
+        streamDoneResolver!();
+      }
+    })();
+  }
+
+  private async makeApiCallAndProcessStream(
+    model: string,
+    requestContents: Content[],
+    params: SendMessageParameters,
+    prompt_id: string,
+  ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    const apiCall = () => {
+      const modelToUse = getEffectiveModel(model);
+
+      return this.config.getContentGenerator().generateContentStream(
+        {
+          model: modelToUse,
+          contents: requestContents,
+          config: { ...this.generationConfig, ...params.config },
+        },
+        prompt_id,
+      );
+    };
+
+    const streamResponse = await retryWithBackoff(apiCall, {
+      shouldRetryOnError: (error: unknown) => {
+        if (error instanceof ApiError && error.message) {
+          if (error.status === 400) return false;
+          if (isSchemaDepthError(error.message)) return false;
+          if (error.status === 429) return true;
+          if (error.status >= 500 && error.status < 600) return true;
+        }
+        return false;
+      },
+      authType: this.config.getContentGeneratorConfig()?.authType,
+    });
+
+    return this.processStreamResponse(model, streamResponse);
+  }
+
+  getHistory(curated: boolean = false): Content[] {
+    const history = curated
+      ? extractCuratedHistory(this.history)
+      : this.history;
+    // Deep copy the history to avoid mutating the history outside of the
+    // chat session.
+    return structuredClone(history);
+  }
+
+  clearHistory(): void {
+    this.history = [];
+  }
+
+  addHistory(content: Content): void {
+    this.history.push(content);
+  }
+
+  setHistory(history: Content[]): void {
+    this.history = history;
+  }
+
+  private async *processStreamResponse(
+    model: string,
+    streamResponse: AsyncGenerator<GenerateContentResponse>,
+  ): AsyncGenerator<GenerateContentResponse> {
+    const modelResponseParts: Part[] = [];
+
+    let hasFinishReason = false;
+
+    for await (const chunk of streamResponse) {
+      hasFinishReason =
+        chunk?.candidates?.some((candidate) => candidate.finishReason) ?? false;
+      if (isValidResponse(chunk)) {
+        const content = chunk.candidates?.[0]?.content;
+        if (content?.parts) {
+          modelResponseParts.push(
+            ...content.parts.filter((part) => !part.thought),
+          );
+        }
+      }
+
+      yield chunk; // Yield every chunk to the UI immediately.
+    }
+
+    // String thoughts and consolidate text parts.
+    const consolidatedParts: Part[] = [];
+    for (const part of modelResponseParts) {
+      const lastPart = consolidatedParts[consolidatedParts.length - 1];
+      if (
+        lastPart?.text &&
+        isValidNonThoughtTextPart(lastPart) &&
+        isValidNonThoughtTextPart(part)
+      ) {
+        lastPart.text += part.text;
+      } else {
+        consolidatedParts.push(part);
+      }
+    }
+
+    const responseText = consolidatedParts
+      .filter((part) => part.text)
+      .map((part) => part.text)
+      .join('')
+      .trim();
+
+    if (!hasFinishReason || !responseText) {
+      if (!hasFinishReason) {
+        throw new InvalidStreamError(
+          'Model stream ended without a finish reason.',
+          'NO_FINISH_REASON',
+        );
+      } else {
+        throw new InvalidStreamError(
+          'Model stream ended with empty response text.',
+          'NO_RESPONSE_TEXT',
+        );
+      }
+    }
+
+    this.history.push({ role: 'model', parts: consolidatedParts });
+  }
+}
+
+/** Visible for Testing */
+export function isSchemaDepthError(errorMessage: string): boolean {
+  return errorMessage.includes('maximum schema depth exceeded');
+}
+
+export function isInvalidArgumentError(errorMessage: string): boolean {
+  return errorMessage.includes('Request contains an invalid argument');
+}

package/backend/packages/gemini-core/src/core/geminiRequest.ts

@@ -0,0 +1,19 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { type PartListUnion } from '@google/genai';
+import { partToString } from '../utils/partUtils';
+
+/**
+ * Represents a request to be sent to the Gemini API.
+ * For now, it's an alias to PartListUnion as the primary content.
+ * This can be expanded later to include other request parameters.
+ */
+export type GeminiCodeRequest = PartListUnion;
+
+export function partListUnionToString(value: PartListUnion): string {
+  return partToString(value, { verbose: true });
+}

package/backend/packages/gemini-core/src/core/prompts.ts

@@ -0,0 +1,109 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import path from 'node:path';
+import fs from 'node:fs';
+import os from 'node:os';
+import process from 'node:process';
+
+export const GEMINI_CONFIG_DIR = '.gemini';
+
+export function resolvePathFromEnv(envVar?: string): {
+  isSwitch: boolean;
+  value: string | null;
+  isDisabled: boolean;
+} {
+  // Handle the case where the environment variable is not set, empty, or just whitespace.
+  const trimmedEnvVar = envVar?.trim();
+  if (!trimmedEnvVar) {
+    return { isSwitch: false, value: null, isDisabled: false };
+  }
+
+  const lowerEnvVar = trimmedEnvVar.toLowerCase();
+  // Check if the input is a common boolean-like string.
+  if (['0', 'false', '1', 'true'].includes(lowerEnvVar)) {
+    // If so, identify it as a "switch" and return its value.
+    const isDisabled = ['0', 'false'].includes(lowerEnvVar);
+    return { isSwitch: true, value: lowerEnvVar, isDisabled };
+  }
+
+  // If it's not a switch, treat it as a potential file path.
+  let customPath = trimmedEnvVar;
+
+  // Safely expand the tilde (~) character to the user's home directory.
+  if (customPath.startsWith('~/') || customPath === '~') {
+    try {
+      const home = os.homedir(); // This is the call that can throw an error.
+      if (customPath === '~') {
+        customPath = home;
+      } else {
+        customPath = path.join(home, customPath.slice(2));
+      }
+    } catch (error) {
+      // If os.homedir() fails, we catch the error instead of crashing.
+      console.warn(
+        `Could not resolve home directory for path: ${trimmedEnvVar}`,
+        error,
+      );
+      // Return null to indicate the path resolution failed.
+      return { isSwitch: false, value: null, isDisabled: false };
+    }
+  }
+
+  // Return it as a non-switch with the fully resolved absolute path.
+  return {
+    isSwitch: false,
+    value: path.resolve(customPath),
+    isDisabled: false,
+  };
+}
+
+export function getCoreSystemPrompt(): string {
+  // A flag to indicate whether the system prompt override is active.
+  let systemMdEnabled = false;
+  // The default path for the system prompt file. This can be overridden.
+  let systemMdPath = path.resolve(path.join(GEMINI_CONFIG_DIR, 'system.md'));
+  // Resolve the environment variable to get either a path or a switch value.
+  const systemMdResolution = resolvePathFromEnv(
+    process.env['GEMINI_SYSTEM_MD'],
+  );
+
+  // Proceed only if the environment variable is set and is not disabled.
+  if (systemMdResolution.value && !systemMdResolution.isDisabled) {
+    systemMdEnabled = true;
+
+    // We update systemMdPath to this new custom path.
+    if (!systemMdResolution.isSwitch) {
+      systemMdPath = systemMdResolution.value;
+    }
+
+    // require file to exist when override is enabled
+    if (!fs.existsSync(systemMdPath)) {
+      throw new Error(`missing system prompt file '${systemMdPath}'`);
+    }
+  }
+  const basePrompt = systemMdEnabled
+    ? fs.readFileSync(systemMdPath, 'utf8')
+    : `You are a helpful assistant.`;
+
+  // if GEMINI_WRITE_SYSTEM_MD is set (and not 0|false), write base system prompt to file
+  const writeSystemMdResolution = resolvePathFromEnv(
+    process.env['GEMINI_WRITE_SYSTEM_MD'],
+  );
+
+  // Check if the feature is enabled. This proceeds only if the environment
+  // variable is set and is not explicitly '0' or 'false'.
+  if (writeSystemMdResolution.value && !writeSystemMdResolution.isDisabled) {
+    const writePath = writeSystemMdResolution.isSwitch
+      ? systemMdPath
+      : writeSystemMdResolution.value;
+
+    fs.mkdirSync(path.dirname(writePath), { recursive: true });
+    fs.writeFileSync(writePath, basePrompt);
+  }
+
+  return basePrompt;
+}

package/backend/packages/gemini-core/src/core/tokenLimits.ts

@@ -0,0 +1,32 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+type Model = string;
+type TokenCount = number;
+
+export const DEFAULT_TOKEN_LIMIT = 1_048_576;
+
+export function tokenLimit(model: Model): TokenCount {
+  // Add other models as they become relevant or if specified by config
+  // Pulled from https://ai.google.dev/gemini-api/docs/models
+  switch (model) {
+    case 'gemini-1.5-pro':
+      return 2_097_152;
+    case 'gemini-1.5-flash':
+    case 'gemini-2.5-pro-preview-05-06':
+    case 'gemini-2.5-pro-preview-06-05':
+    case 'gemini-2.5-pro':
+    case 'gemini-2.5-flash-preview-05-20':
+    case 'gemini-2.5-flash':
+    case 'gemini-2.5-flash-lite':
+    case 'gemini-2.0-flash':
+      return 1_048_576;
+    case 'gemini-2.0-flash-preview-image-generation':
+      return 32_000;
+    default:
+      return DEFAULT_TOKEN_LIMIT;
+  }
+}