mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 17:27:54 +00:00
pre-release commit
This commit is contained in:
koalazf.99
352
packages/core/src/config/config.test.ts
Normal file
352
packages/core/src/config/config.test.ts
Normal file
@@ -0,0 +1,352 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, Mock } from 'vitest';
|
||||
import { Config, ConfigParameters, SandboxConfig } from './config.js';
|
||||
import * as path from 'path';
|
||||
import { setGeminiMdFilename as mockSetGeminiMdFilename } from '../tools/memoryTool.js';
|
||||
import {
|
||||
DEFAULT_TELEMETRY_TARGET,
|
||||
DEFAULT_OTLP_ENDPOINT,
|
||||
} from '../telemetry/index.js';
|
||||
import {
|
||||
AuthType,
|
||||
createContentGeneratorConfig,
|
||||
} from '../core/contentGenerator.js';
|
||||
import { GeminiClient } from '../core/client.js';
|
||||
import { GitService } from '../services/gitService.js';
|
||||
import { loadServerHierarchicalMemory } from '../utils/memoryDiscovery.js';
|
||||
|
||||
// Mock dependencies that might be called during Config construction or createServerConfig
|
||||
vi.mock('../tools/tool-registry', () => {
|
||||
const ToolRegistryMock = vi.fn();
|
||||
ToolRegistryMock.prototype.registerTool = vi.fn();
|
||||
ToolRegistryMock.prototype.discoverTools = vi.fn();
|
||||
ToolRegistryMock.prototype.getAllTools = vi.fn(() => []); // Mock methods if needed
|
||||
ToolRegistryMock.prototype.getTool = vi.fn();
|
||||
ToolRegistryMock.prototype.getFunctionDeclarations = vi.fn(() => []);
|
||||
return { ToolRegistry: ToolRegistryMock };
|
||||
});
|
||||
|
||||
vi.mock('../utils/memoryDiscovery.js', () => ({
|
||||
loadServerHierarchicalMemory: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock individual tools if their constructors are complex or have side effects
|
||||
vi.mock('../tools/ls');
|
||||
vi.mock('../tools/read-file');
|
||||
vi.mock('../tools/grep');
|
||||
vi.mock('../tools/glob');
|
||||
vi.mock('../tools/edit');
|
||||
vi.mock('../tools/shell');
|
||||
vi.mock('../tools/write-file');
|
||||
vi.mock('../tools/web-fetch');
|
||||
vi.mock('../tools/read-many-files');
|
||||
vi.mock('../tools/memoryTool', () => ({
|
||||
MemoryTool: vi.fn(),
|
||||
setGeminiMdFilename: vi.fn(),
|
||||
getCurrentGeminiMdFilename: vi.fn(() => 'QWEN.md'), // Mock the original filename
|
||||
DEFAULT_CONTEXT_FILENAME: 'QWEN.md',
|
||||
GEMINI_CONFIG_DIR: '.qwen',
|
||||
}));
|
||||
|
||||
vi.mock('../core/contentGenerator.js', async (importOriginal) => {
|
||||
const actual =
|
||||
await importOriginal<typeof import('../core/contentGenerator.js')>();
|
||||
return {
|
||||
...actual,
|
||||
createContentGeneratorConfig: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock('../core/client.js', () => ({
|
||||
GeminiClient: vi.fn().mockImplementation(() => ({
|
||||
initialize: vi.fn().mockResolvedValue(undefined),
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock('../telemetry/index.js', async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import('../telemetry/index.js')>();
|
||||
return {
|
||||
...actual,
|
||||
initializeTelemetry: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock('../services/gitService.js', () => {
|
||||
const GitServiceMock = vi.fn();
|
||||
GitServiceMock.prototype.initialize = vi.fn();
|
||||
return { GitService: GitServiceMock };
|
||||
});
|
||||
|
||||
describe('Server Config (config.ts)', () => {
|
||||
const MODEL = 'gemini-pro';
|
||||
const SANDBOX: SandboxConfig = {
|
||||
command: 'docker',
|
||||
image: 'gemini-cli-sandbox',
|
||||
};
|
||||
const TARGET_DIR = '/path/to/target';
|
||||
const DEBUG_MODE = false;
|
||||
const QUESTION = 'test question';
|
||||
const FULL_CONTEXT = false;
|
||||
const USER_MEMORY = 'Test User Memory';
|
||||
const TELEMETRY_SETTINGS = { enabled: false };
|
||||
const EMBEDDING_MODEL = 'gemini-embedding';
|
||||
const SESSION_ID = 'test-session-id';
|
||||
const baseParams: ConfigParameters = {
|
||||
cwd: '/tmp',
|
||||
embeddingModel: EMBEDDING_MODEL,
|
||||
sandbox: SANDBOX,
|
||||
targetDir: TARGET_DIR,
|
||||
debugMode: DEBUG_MODE,
|
||||
question: QUESTION,
|
||||
fullContext: FULL_CONTEXT,
|
||||
userMemory: USER_MEMORY,
|
||||
telemetry: TELEMETRY_SETTINGS,
|
||||
sessionId: SESSION_ID,
|
||||
model: MODEL,
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset mocks if necessary
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('initialize', () => {
|
||||
it('should throw an error if checkpointing is enabled and GitService fails', async () => {
|
||||
const gitError = new Error('Git is not installed');
|
||||
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError);
|
||||
|
||||
const config = new Config({
|
||||
...baseParams,
|
||||
checkpointing: true,
|
||||
});
|
||||
|
||||
await expect(config.initialize()).rejects.toThrow(gitError);
|
||||
});
|
||||
|
||||
it('should not throw an error if checkpointing is disabled and GitService fails', async () => {
|
||||
const gitError = new Error('Git is not installed');
|
||||
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError);
|
||||
|
||||
const config = new Config({
|
||||
...baseParams,
|
||||
checkpointing: false,
|
||||
});
|
||||
|
||||
await expect(config.initialize()).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('refreshAuth', () => {
|
||||
it('should refresh auth and update config', async () => {
|
||||
const config = new Config(baseParams);
|
||||
const authType = AuthType.USE_GEMINI;
|
||||
const newModel = 'gemini-flash';
|
||||
const mockContentConfig = {
|
||||
model: newModel,
|
||||
apiKey: 'test-key',
|
||||
};
|
||||
|
||||
(createContentGeneratorConfig as Mock).mockResolvedValue(
|
||||
mockContentConfig,
|
||||
);
|
||||
|
||||
await config.refreshAuth(authType);
|
||||
|
||||
expect(createContentGeneratorConfig).toHaveBeenCalledWith(
|
||||
MODEL, // Should be called with the original model 'gemini-pro'
|
||||
authType,
|
||||
);
|
||||
// Verify that contentGeneratorConfig is updated with the new model
|
||||
expect(config.getContentGeneratorConfig()).toEqual(mockContentConfig);
|
||||
expect(config.getContentGeneratorConfig().model).toBe(newModel);
|
||||
expect(config.getModel()).toBe(newModel); // getModel() should return the updated model
|
||||
expect(GeminiClient).toHaveBeenCalledWith(config);
|
||||
});
|
||||
});
|
||||
|
||||
it('Config constructor should store userMemory correctly', () => {
|
||||
const config = new Config(baseParams);
|
||||
|
||||
expect(config.getUserMemory()).toBe(USER_MEMORY);
|
||||
// Verify other getters if needed
|
||||
expect(config.getTargetDir()).toBe(path.resolve(TARGET_DIR)); // Check resolved path
|
||||
});
|
||||
|
||||
it('Config constructor should default userMemory to empty string if not provided', () => {
|
||||
const paramsWithoutMemory: ConfigParameters = { ...baseParams };
|
||||
delete paramsWithoutMemory.userMemory;
|
||||
const config = new Config(paramsWithoutMemory);
|
||||
|
||||
expect(config.getUserMemory()).toBe('');
|
||||
});
|
||||
|
||||
it('Config constructor should call setGeminiMdFilename with contextFileName if provided', () => {
|
||||
const contextFileName = 'CUSTOM_AGENTS.md';
|
||||
const paramsWithContextFile: ConfigParameters = {
|
||||
...baseParams,
|
||||
contextFileName,
|
||||
};
|
||||
new Config(paramsWithContextFile);
|
||||
expect(mockSetGeminiMdFilename).toHaveBeenCalledWith(contextFileName);
|
||||
});
|
||||
|
||||
it('Config constructor should not call setGeminiMdFilename if contextFileName is not provided', () => {
|
||||
new Config(baseParams); // baseParams does not have contextFileName
|
||||
expect(mockSetGeminiMdFilename).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should set default file filtering settings when not provided', () => {
|
||||
const config = new Config(baseParams);
|
||||
expect(config.getFileFilteringRespectGitIgnore()).toBe(true);
|
||||
});
|
||||
|
||||
it('should set custom file filtering settings when provided', () => {
|
||||
const paramsWithFileFiltering: ConfigParameters = {
|
||||
...baseParams,
|
||||
fileFiltering: {
|
||||
respectGitIgnore: false,
|
||||
},
|
||||
};
|
||||
const config = new Config(paramsWithFileFiltering);
|
||||
expect(config.getFileFilteringRespectGitIgnore()).toBe(false);
|
||||
});
|
||||
|
||||
it('Config constructor should set telemetry to true when provided as true', () => {
|
||||
const paramsWithTelemetry: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true },
|
||||
};
|
||||
const config = new Config(paramsWithTelemetry);
|
||||
expect(config.getTelemetryEnabled()).toBe(true);
|
||||
});
|
||||
|
||||
it('Config constructor should set telemetry to false when provided as false', () => {
|
||||
const paramsWithTelemetry: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: false },
|
||||
};
|
||||
const config = new Config(paramsWithTelemetry);
|
||||
expect(config.getTelemetryEnabled()).toBe(false);
|
||||
});
|
||||
|
||||
it('Config constructor should default telemetry to default value if not provided', () => {
|
||||
const paramsWithoutTelemetry: ConfigParameters = { ...baseParams };
|
||||
delete paramsWithoutTelemetry.telemetry;
|
||||
const config = new Config(paramsWithoutTelemetry);
|
||||
expect(config.getTelemetryEnabled()).toBe(TELEMETRY_SETTINGS.enabled);
|
||||
});
|
||||
|
||||
it('should have a getFileService method that returns FileDiscoveryService', () => {
|
||||
const config = new Config(baseParams);
|
||||
const fileService = config.getFileService();
|
||||
expect(fileService).toBeDefined();
|
||||
});
|
||||
|
||||
describe('Telemetry Settings', () => {
|
||||
it('should return default telemetry target if not provided', () => {
|
||||
const params: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true },
|
||||
};
|
||||
const config = new Config(params);
|
||||
expect(config.getTelemetryTarget()).toBe(DEFAULT_TELEMETRY_TARGET);
|
||||
});
|
||||
|
||||
it('should return provided OTLP endpoint', () => {
|
||||
const endpoint = 'http://custom.otel.collector:4317';
|
||||
const params: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true, otlpEndpoint: endpoint },
|
||||
};
|
||||
const config = new Config(params);
|
||||
expect(config.getTelemetryOtlpEndpoint()).toBe(endpoint);
|
||||
});
|
||||
|
||||
it('should return default OTLP endpoint if not provided', () => {
|
||||
const params: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true },
|
||||
};
|
||||
const config = new Config(params);
|
||||
expect(config.getTelemetryOtlpEndpoint()).toBe(DEFAULT_OTLP_ENDPOINT);
|
||||
});
|
||||
|
||||
it('should return provided logPrompts setting', () => {
|
||||
const params: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true, logPrompts: false },
|
||||
};
|
||||
const config = new Config(params);
|
||||
expect(config.getTelemetryLogPromptsEnabled()).toBe(false);
|
||||
});
|
||||
|
||||
it('should return default logPrompts setting (true) if not provided', () => {
|
||||
const params: ConfigParameters = {
|
||||
...baseParams,
|
||||
telemetry: { enabled: true },
|
||||
};
|
||||
const config = new Config(params);
|
||||
expect(config.getTelemetryLogPromptsEnabled()).toBe(true);
|
||||
});
|
||||
|
||||
it('should return default logPrompts setting (true) if telemetry object is not provided', () => {
|
||||
const paramsWithoutTelemetry: ConfigParameters = { ...baseParams };
|
||||
delete paramsWithoutTelemetry.telemetry;
|
||||
const config = new Config(paramsWithoutTelemetry);
|
||||
expect(config.getTelemetryLogPromptsEnabled()).toBe(true);
|
||||
});
|
||||
|
||||
it('should return default telemetry target if telemetry object is not provided', () => {
|
||||
const paramsWithoutTelemetry: ConfigParameters = { ...baseParams };
|
||||
delete paramsWithoutTelemetry.telemetry;
|
||||
const config = new Config(paramsWithoutTelemetry);
|
||||
expect(config.getTelemetryTarget()).toBe(DEFAULT_TELEMETRY_TARGET);
|
||||
});
|
||||
|
||||
it('should return default OTLP endpoint if telemetry object is not provided', () => {
|
||||
const paramsWithoutTelemetry: ConfigParameters = { ...baseParams };
|
||||
delete paramsWithoutTelemetry.telemetry;
|
||||
const config = new Config(paramsWithoutTelemetry);
|
||||
expect(config.getTelemetryOtlpEndpoint()).toBe(DEFAULT_OTLP_ENDPOINT);
|
||||
});
|
||||
});
|
||||
|
||||
describe('refreshMemory', () => {
|
||||
it('should update memory and file count on successful refresh', async () => {
|
||||
const config = new Config(baseParams);
|
||||
const mockMemoryData = {
|
||||
memoryContent: 'new memory content',
|
||||
fileCount: 5,
|
||||
};
|
||||
|
||||
(loadServerHierarchicalMemory as Mock).mockResolvedValue(mockMemoryData);
|
||||
|
||||
const result = await config.refreshMemory();
|
||||
|
||||
expect(loadServerHierarchicalMemory).toHaveBeenCalledWith(
|
||||
config.getWorkingDir(),
|
||||
config.getDebugMode(),
|
||||
config.getFileService(),
|
||||
config.getExtensionContextFilePaths(),
|
||||
);
|
||||
|
||||
expect(config.getUserMemory()).toBe(mockMemoryData.memoryContent);
|
||||
expect(config.getGeminiMdFileCount()).toBe(mockMemoryData.fileCount);
|
||||
expect(result).toEqual(mockMemoryData);
|
||||
});
|
||||
|
||||
it('should propagate errors from loadServerHierarchicalMemory', async () => {
|
||||
const config = new Config(baseParams);
|
||||
const testError = new Error('Failed to load memory');
|
||||
|
||||
(loadServerHierarchicalMemory as Mock).mockRejectedValue(testError);
|
||||
|
||||
await expect(config.refreshMemory()).rejects.toThrow(testError);
|
||||
});
|
||||
});
|
||||
});
|
||||
611
packages/core/src/config/config.ts
Normal file
611
packages/core/src/config/config.ts
Normal file
@@ -0,0 +1,611 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import * as path from 'node:path';
|
||||
import process from 'node:process';
|
||||
import {
|
||||
AuthType,
|
||||
ContentGeneratorConfig,
|
||||
createContentGeneratorConfig,
|
||||
} from '../core/contentGenerator.js';
|
||||
import { UserTierId } from '../code_assist/types.js';
|
||||
import { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import { LSTool } from '../tools/ls.js';
|
||||
import { ReadFileTool } from '../tools/read-file.js';
|
||||
import { GrepTool } from '../tools/grep.js';
|
||||
import { GlobTool } from '../tools/glob.js';
|
||||
import { EditTool } from '../tools/edit.js';
|
||||
import { ShellTool } from '../tools/shell.js';
|
||||
import { WriteFileTool } from '../tools/write-file.js';
|
||||
import { WebFetchTool } from '../tools/web-fetch.js';
|
||||
import { ReadManyFilesTool } from '../tools/read-many-files.js';
|
||||
import {
|
||||
MemoryTool,
|
||||
setGeminiMdFilename,
|
||||
GEMINI_CONFIG_DIR as GEMINI_DIR,
|
||||
} from '../tools/memoryTool.js';
|
||||
import { WebSearchTool } from '../tools/web-search.js';
|
||||
import { GeminiClient } from '../core/client.js';
|
||||
import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
|
||||
import { GitService } from '../services/gitService.js';
|
||||
import { loadServerHierarchicalMemory } from '../utils/memoryDiscovery.js';
|
||||
import { getProjectTempDir } from '../utils/paths.js';
|
||||
import {
|
||||
initializeTelemetry,
|
||||
DEFAULT_TELEMETRY_TARGET,
|
||||
DEFAULT_OTLP_ENDPOINT,
|
||||
TelemetryTarget,
|
||||
StartSessionEvent,
|
||||
} from '../telemetry/index.js';
|
||||
import {
|
||||
DEFAULT_GEMINI_EMBEDDING_MODEL,
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
} from './models.js';
|
||||
import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
|
||||
|
||||
export enum ApprovalMode {
|
||||
DEFAULT = 'default',
|
||||
AUTO_EDIT = 'autoEdit',
|
||||
YOLO = 'yolo',
|
||||
}
|
||||
|
||||
export interface AccessibilitySettings {
|
||||
disableLoadingPhrases?: boolean;
|
||||
}
|
||||
|
||||
export interface BugCommandSettings {
|
||||
urlTemplate: string;
|
||||
}
|
||||
|
||||
export interface TelemetrySettings {
|
||||
enabled?: boolean;
|
||||
target?: TelemetryTarget;
|
||||
otlpEndpoint?: string;
|
||||
logPrompts?: boolean;
|
||||
}
|
||||
|
||||
export interface ActiveExtension {
|
||||
name: string;
|
||||
version: string;
|
||||
}
|
||||
|
||||
export class MCPServerConfig {
|
||||
constructor(
|
||||
// For stdio transport
|
||||
readonly command?: string,
|
||||
readonly args?: string[],
|
||||
readonly env?: Record<string, string>,
|
||||
readonly cwd?: string,
|
||||
// For sse transport
|
||||
readonly url?: string,
|
||||
// For streamable http transport
|
||||
readonly httpUrl?: string,
|
||||
readonly headers?: Record<string, string>,
|
||||
// For websocket transport
|
||||
readonly tcp?: string,
|
||||
// Common
|
||||
readonly timeout?: number,
|
||||
readonly trust?: boolean,
|
||||
// Metadata
|
||||
readonly description?: string,
|
||||
readonly includeTools?: string[],
|
||||
readonly excludeTools?: string[],
|
||||
) {}
|
||||
}
|
||||
|
||||
export interface SandboxConfig {
|
||||
command: 'docker' | 'podman' | 'sandbox-exec';
|
||||
image: string;
|
||||
}
|
||||
|
||||
export type FlashFallbackHandler = (
|
||||
currentModel: string,
|
||||
fallbackModel: string,
|
||||
error?: unknown,
|
||||
) => Promise<boolean | string | null>;
|
||||
|
||||
export interface ConfigParameters {
|
||||
sessionId: string;
|
||||
embeddingModel?: string;
|
||||
sandbox?: SandboxConfig;
|
||||
targetDir: string;
|
||||
debugMode: boolean;
|
||||
question?: string;
|
||||
fullContext?: boolean;
|
||||
coreTools?: string[];
|
||||
excludeTools?: string[];
|
||||
toolDiscoveryCommand?: string;
|
||||
toolCallCommand?: string;
|
||||
mcpServerCommand?: string;
|
||||
mcpServers?: Record<string, MCPServerConfig>;
|
||||
userMemory?: string;
|
||||
geminiMdFileCount?: number;
|
||||
approvalMode?: ApprovalMode;
|
||||
showMemoryUsage?: boolean;
|
||||
contextFileName?: string | string[];
|
||||
accessibility?: AccessibilitySettings;
|
||||
telemetry?: TelemetrySettings;
|
||||
usageStatisticsEnabled?: boolean;
|
||||
fileFiltering?: {
|
||||
respectGitIgnore?: boolean;
|
||||
enableRecursiveFileSearch?: boolean;
|
||||
};
|
||||
checkpointing?: boolean;
|
||||
proxy?: string;
|
||||
cwd: string;
|
||||
fileDiscoveryService?: FileDiscoveryService;
|
||||
bugCommand?: BugCommandSettings;
|
||||
model: string;
|
||||
extensionContextFilePaths?: string[];
|
||||
maxSessionTurns?: number;
|
||||
listExtensions?: boolean;
|
||||
activeExtensions?: ActiveExtension[];
|
||||
noBrowser?: boolean;
|
||||
ideMode?: boolean;
|
||||
enableOpenAILogging?: boolean;
|
||||
sampling_params?: {
|
||||
top_p?: number;
|
||||
top_k?: number;
|
||||
repetition_penalty?: number;
|
||||
presence_penalty?: number;
|
||||
frequency_penalty?: number;
|
||||
temperature?: number;
|
||||
max_tokens?: number;
|
||||
};
|
||||
}
|
||||
|
||||
export class Config {
|
||||
private toolRegistry!: ToolRegistry;
|
||||
private readonly sessionId: string;
|
||||
private contentGeneratorConfig!: ContentGeneratorConfig;
|
||||
private readonly embeddingModel: string;
|
||||
private readonly sandbox: SandboxConfig | undefined;
|
||||
private readonly targetDir: string;
|
||||
private readonly debugMode: boolean;
|
||||
private readonly question: string | undefined;
|
||||
private readonly fullContext: boolean;
|
||||
private readonly coreTools: string[] | undefined;
|
||||
private readonly excludeTools: string[] | undefined;
|
||||
private readonly toolDiscoveryCommand: string | undefined;
|
||||
private readonly toolCallCommand: string | undefined;
|
||||
private readonly mcpServerCommand: string | undefined;
|
||||
private readonly mcpServers: Record<string, MCPServerConfig> | undefined;
|
||||
private userMemory: string;
|
||||
private geminiMdFileCount: number;
|
||||
private approvalMode: ApprovalMode;
|
||||
private readonly showMemoryUsage: boolean;
|
||||
private readonly accessibility: AccessibilitySettings;
|
||||
private readonly telemetrySettings: TelemetrySettings;
|
||||
private readonly usageStatisticsEnabled: boolean;
|
||||
private geminiClient!: GeminiClient;
|
||||
private readonly fileFiltering: {
|
||||
respectGitIgnore: boolean;
|
||||
enableRecursiveFileSearch: boolean;
|
||||
};
|
||||
private fileDiscoveryService: FileDiscoveryService | null = null;
|
||||
private gitService: GitService | undefined = undefined;
|
||||
private readonly checkpointing: boolean;
|
||||
private readonly proxy: string | undefined;
|
||||
private readonly cwd: string;
|
||||
private readonly bugCommand: BugCommandSettings | undefined;
|
||||
private readonly model: string;
|
||||
private readonly extensionContextFilePaths: string[];
|
||||
private readonly noBrowser: boolean;
|
||||
private readonly ideMode: boolean;
|
||||
private readonly enableOpenAILogging: boolean;
|
||||
private readonly sampling_params?: {
|
||||
top_p?: number;
|
||||
top_k?: number;
|
||||
repetition_penalty?: number;
|
||||
presence_penalty?: number;
|
||||
frequency_penalty?: number;
|
||||
temperature?: number;
|
||||
max_tokens?: number;
|
||||
};
|
||||
private modelSwitchedDuringSession: boolean = false;
|
||||
private readonly maxSessionTurns: number;
|
||||
private readonly listExtensions: boolean;
|
||||
private readonly _activeExtensions: ActiveExtension[];
|
||||
flashFallbackHandler?: FlashFallbackHandler;
|
||||
private quotaErrorOccurred: boolean = false;
|
||||
|
||||
constructor(params: ConfigParameters) {
|
||||
this.sessionId = params.sessionId;
|
||||
this.embeddingModel =
|
||||
params.embeddingModel ?? DEFAULT_GEMINI_EMBEDDING_MODEL;
|
||||
this.sandbox = params.sandbox;
|
||||
this.targetDir = path.resolve(params.targetDir);
|
||||
this.debugMode = params.debugMode;
|
||||
this.question = params.question;
|
||||
this.fullContext = params.fullContext ?? false;
|
||||
this.coreTools = params.coreTools;
|
||||
this.excludeTools = params.excludeTools;
|
||||
this.toolDiscoveryCommand = params.toolDiscoveryCommand;
|
||||
this.toolCallCommand = params.toolCallCommand;
|
||||
this.mcpServerCommand = params.mcpServerCommand;
|
||||
this.mcpServers = params.mcpServers;
|
||||
this.userMemory = params.userMemory ?? '';
|
||||
this.geminiMdFileCount = params.geminiMdFileCount ?? 0;
|
||||
this.approvalMode = params.approvalMode ?? ApprovalMode.DEFAULT;
|
||||
this.showMemoryUsage = params.showMemoryUsage ?? false;
|
||||
this.accessibility = params.accessibility ?? {};
|
||||
this.telemetrySettings = {
|
||||
enabled: params.telemetry?.enabled ?? false,
|
||||
target: params.telemetry?.target ?? DEFAULT_TELEMETRY_TARGET,
|
||||
otlpEndpoint: params.telemetry?.otlpEndpoint ?? DEFAULT_OTLP_ENDPOINT,
|
||||
logPrompts: params.telemetry?.logPrompts ?? true,
|
||||
};
|
||||
this.usageStatisticsEnabled = params.usageStatisticsEnabled ?? true;
|
||||
|
||||
this.fileFiltering = {
|
||||
respectGitIgnore: params.fileFiltering?.respectGitIgnore ?? true,
|
||||
enableRecursiveFileSearch:
|
||||
params.fileFiltering?.enableRecursiveFileSearch ?? true,
|
||||
};
|
||||
this.checkpointing = params.checkpointing ?? false;
|
||||
this.proxy = params.proxy;
|
||||
this.cwd = params.cwd ?? process.cwd();
|
||||
this.fileDiscoveryService = params.fileDiscoveryService ?? null;
|
||||
this.bugCommand = params.bugCommand;
|
||||
this.model = params.model;
|
||||
this.extensionContextFilePaths = params.extensionContextFilePaths ?? [];
|
||||
this.maxSessionTurns = params.maxSessionTurns ?? -1;
|
||||
this.listExtensions = params.listExtensions ?? false;
|
||||
this._activeExtensions = params.activeExtensions ?? [];
|
||||
this.noBrowser = params.noBrowser ?? false;
|
||||
this.ideMode = params.ideMode ?? false;
|
||||
this.enableOpenAILogging = params.enableOpenAILogging ?? false;
|
||||
this.sampling_params = params.sampling_params;
|
||||
|
||||
if (params.contextFileName) {
|
||||
setGeminiMdFilename(params.contextFileName);
|
||||
}
|
||||
|
||||
if (this.telemetrySettings.enabled) {
|
||||
initializeTelemetry(this);
|
||||
}
|
||||
|
||||
if (this.getUsageStatisticsEnabled()) {
|
||||
ClearcutLogger.getInstance(this)?.logStartSessionEvent(
|
||||
new StartSessionEvent(this),
|
||||
);
|
||||
} else {
|
||||
console.log('Data collection is disabled.');
|
||||
}
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
// Initialize centralized FileDiscoveryService
|
||||
this.getFileService();
|
||||
if (this.getCheckpointingEnabled()) {
|
||||
await this.getGitService();
|
||||
}
|
||||
this.toolRegistry = await this.createToolRegistry();
|
||||
}
|
||||
|
||||
async refreshAuth(authMethod: AuthType) {
|
||||
this.contentGeneratorConfig = await createContentGeneratorConfig(
|
||||
this.model,
|
||||
authMethod,
|
||||
);
|
||||
this.contentGeneratorConfig.enableOpenAILogging = this.enableOpenAILogging;
|
||||
|
||||
// Set sampling parameters from config if available
|
||||
if (this.sampling_params) {
|
||||
this.contentGeneratorConfig.samplingParams = this.sampling_params;
|
||||
}
|
||||
|
||||
this.geminiClient = new GeminiClient(this);
|
||||
await this.geminiClient.initialize(this.contentGeneratorConfig);
|
||||
|
||||
// Reset the session flag since we're explicitly changing auth and using default model
|
||||
this.modelSwitchedDuringSession = false;
|
||||
}
|
||||
|
||||
getSessionId(): string {
|
||||
return this.sessionId;
|
||||
}
|
||||
|
||||
getContentGeneratorConfig(): ContentGeneratorConfig {
|
||||
return this.contentGeneratorConfig;
|
||||
}
|
||||
|
||||
getModel(): string {
|
||||
return this.contentGeneratorConfig?.model || this.model;
|
||||
}
|
||||
|
||||
setModel(newModel: string): void {
|
||||
if (this.contentGeneratorConfig) {
|
||||
this.contentGeneratorConfig.model = newModel;
|
||||
this.modelSwitchedDuringSession = true;
|
||||
}
|
||||
}
|
||||
|
||||
isModelSwitchedDuringSession(): boolean {
|
||||
return this.modelSwitchedDuringSession;
|
||||
}
|
||||
|
||||
resetModelToDefault(): void {
|
||||
if (this.contentGeneratorConfig) {
|
||||
this.contentGeneratorConfig.model = this.model; // Reset to the original default model
|
||||
this.modelSwitchedDuringSession = false;
|
||||
}
|
||||
}
|
||||
|
||||
setFlashFallbackHandler(handler: FlashFallbackHandler): void {
|
||||
this.flashFallbackHandler = handler;
|
||||
}
|
||||
|
||||
getMaxSessionTurns(): number {
|
||||
return this.maxSessionTurns;
|
||||
}
|
||||
|
||||
setQuotaErrorOccurred(value: boolean): void {
|
||||
this.quotaErrorOccurred = value;
|
||||
}
|
||||
|
||||
getQuotaErrorOccurred(): boolean {
|
||||
return this.quotaErrorOccurred;
|
||||
}
|
||||
|
||||
async getUserTier(): Promise<UserTierId | undefined> {
|
||||
if (!this.geminiClient) {
|
||||
return undefined;
|
||||
}
|
||||
const generator = this.geminiClient.getContentGenerator();
|
||||
return await generator.getTier?.();
|
||||
}
|
||||
|
||||
getEmbeddingModel(): string {
|
||||
return this.embeddingModel;
|
||||
}
|
||||
|
||||
getSandbox(): SandboxConfig | undefined {
|
||||
return this.sandbox;
|
||||
}
|
||||
|
||||
getTargetDir(): string {
|
||||
return this.targetDir;
|
||||
}
|
||||
|
||||
getProjectRoot(): string {
|
||||
return this.targetDir;
|
||||
}
|
||||
|
||||
getToolRegistry(): Promise<ToolRegistry> {
|
||||
return Promise.resolve(this.toolRegistry);
|
||||
}
|
||||
|
||||
getDebugMode(): boolean {
|
||||
return this.debugMode;
|
||||
}
|
||||
getQuestion(): string | undefined {
|
||||
return this.question;
|
||||
}
|
||||
|
||||
getFullContext(): boolean {
|
||||
return this.fullContext;
|
||||
}
|
||||
|
||||
getCoreTools(): string[] | undefined {
|
||||
return this.coreTools;
|
||||
}
|
||||
|
||||
getExcludeTools(): string[] | undefined {
|
||||
return this.excludeTools;
|
||||
}
|
||||
|
||||
getToolDiscoveryCommand(): string | undefined {
|
||||
return this.toolDiscoveryCommand;
|
||||
}
|
||||
|
||||
getToolCallCommand(): string | undefined {
|
||||
return this.toolCallCommand;
|
||||
}
|
||||
|
||||
getMcpServerCommand(): string | undefined {
|
||||
return this.mcpServerCommand;
|
||||
}
|
||||
|
||||
getMcpServers(): Record<string, MCPServerConfig> | undefined {
|
||||
return this.mcpServers;
|
||||
}
|
||||
|
||||
getUserMemory(): string {
|
||||
return this.userMemory;
|
||||
}
|
||||
|
||||
setUserMemory(newUserMemory: string): void {
|
||||
this.userMemory = newUserMemory;
|
||||
}
|
||||
|
||||
getGeminiMdFileCount(): number {
|
||||
return this.geminiMdFileCount;
|
||||
}
|
||||
|
||||
setGeminiMdFileCount(count: number): void {
|
||||
this.geminiMdFileCount = count;
|
||||
}
|
||||
|
||||
getApprovalMode(): ApprovalMode {
|
||||
return this.approvalMode;
|
||||
}
|
||||
|
||||
setApprovalMode(mode: ApprovalMode): void {
|
||||
this.approvalMode = mode;
|
||||
}
|
||||
|
||||
getShowMemoryUsage(): boolean {
|
||||
return this.showMemoryUsage;
|
||||
}
|
||||
|
||||
getAccessibility(): AccessibilitySettings {
|
||||
return this.accessibility;
|
||||
}
|
||||
|
||||
getTelemetryEnabled(): boolean {
|
||||
return this.telemetrySettings.enabled ?? false;
|
||||
}
|
||||
|
||||
getTelemetryLogPromptsEnabled(): boolean {
|
||||
return this.telemetrySettings.logPrompts ?? true;
|
||||
}
|
||||
|
||||
getTelemetryOtlpEndpoint(): string {
|
||||
return this.telemetrySettings.otlpEndpoint ?? DEFAULT_OTLP_ENDPOINT;
|
||||
}
|
||||
|
||||
getTelemetryTarget(): TelemetryTarget {
|
||||
return this.telemetrySettings.target ?? DEFAULT_TELEMETRY_TARGET;
|
||||
}
|
||||
|
||||
getGeminiClient(): GeminiClient {
|
||||
return this.geminiClient;
|
||||
}
|
||||
|
||||
getGeminiDir(): string {
|
||||
return path.join(this.targetDir, GEMINI_DIR);
|
||||
}
|
||||
|
||||
getProjectTempDir(): string {
|
||||
return getProjectTempDir(this.getProjectRoot());
|
||||
}
|
||||
|
||||
getEnableRecursiveFileSearch(): boolean {
|
||||
return this.fileFiltering.enableRecursiveFileSearch;
|
||||
}
|
||||
|
||||
getFileFilteringRespectGitIgnore(): boolean {
|
||||
return this.fileFiltering.respectGitIgnore;
|
||||
}
|
||||
|
||||
getCheckpointingEnabled(): boolean {
|
||||
return this.checkpointing;
|
||||
}
|
||||
|
||||
getProxy(): string | undefined {
|
||||
return this.proxy;
|
||||
}
|
||||
|
||||
getWorkingDir(): string {
|
||||
return this.cwd;
|
||||
}
|
||||
|
||||
getBugCommand(): BugCommandSettings | undefined {
|
||||
return this.bugCommand;
|
||||
}
|
||||
|
||||
getFileService(): FileDiscoveryService {
|
||||
if (!this.fileDiscoveryService) {
|
||||
this.fileDiscoveryService = new FileDiscoveryService(this.targetDir);
|
||||
}
|
||||
return this.fileDiscoveryService;
|
||||
}
|
||||
|
||||
getUsageStatisticsEnabled(): boolean {
|
||||
return this.usageStatisticsEnabled;
|
||||
}
|
||||
|
||||
getExtensionContextFilePaths(): string[] {
|
||||
return this.extensionContextFilePaths;
|
||||
}
|
||||
|
||||
getListExtensions(): boolean {
|
||||
return this.listExtensions;
|
||||
}
|
||||
|
||||
getActiveExtensions(): ActiveExtension[] {
|
||||
return this._activeExtensions;
|
||||
}
|
||||
|
||||
getNoBrowser(): boolean {
|
||||
return this.noBrowser;
|
||||
}
|
||||
|
||||
getIdeMode(): boolean {
|
||||
return this.ideMode;
|
||||
}
|
||||
|
||||
async getGitService(): Promise<GitService> {
|
||||
if (!this.gitService) {
|
||||
this.gitService = new GitService(this.targetDir);
|
||||
await this.gitService.initialize();
|
||||
}
|
||||
return this.gitService;
|
||||
}
|
||||
|
||||
getEnableOpenAILogging(): boolean {
|
||||
return this.enableOpenAILogging;
|
||||
}
|
||||
|
||||
async refreshMemory(): Promise<{ memoryContent: string; fileCount: number }> {
|
||||
const { memoryContent, fileCount } = await loadServerHierarchicalMemory(
|
||||
this.getWorkingDir(),
|
||||
this.getDebugMode(),
|
||||
this.getFileService(),
|
||||
this.getExtensionContextFilePaths(),
|
||||
);
|
||||
|
||||
this.setUserMemory(memoryContent);
|
||||
this.setGeminiMdFileCount(fileCount);
|
||||
|
||||
return { memoryContent, fileCount };
|
||||
}
|
||||
|
||||
async createToolRegistry(): Promise<ToolRegistry> {
|
||||
const registry = new ToolRegistry(this);
|
||||
|
||||
// helper to create & register core tools that are enabled
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
const registerCoreTool = (ToolClass: any, ...args: unknown[]) => {
|
||||
const className = ToolClass.name;
|
||||
const toolName = ToolClass.Name || className;
|
||||
const coreTools = this.getCoreTools();
|
||||
const excludeTools = this.getExcludeTools();
|
||||
|
||||
let isEnabled = false;
|
||||
if (coreTools === undefined) {
|
||||
isEnabled = true;
|
||||
} else {
|
||||
isEnabled = coreTools.some(
|
||||
(tool) =>
|
||||
tool === className ||
|
||||
tool === toolName ||
|
||||
tool.startsWith(`${className}(`) ||
|
||||
tool.startsWith(`${toolName}(`),
|
||||
);
|
||||
}
|
||||
|
||||
if (
|
||||
excludeTools?.includes(className) ||
|
||||
excludeTools?.includes(toolName)
|
||||
) {
|
||||
isEnabled = false;
|
||||
}
|
||||
|
||||
if (isEnabled) {
|
||||
registry.registerTool(new ToolClass(...args));
|
||||
}
|
||||
};
|
||||
|
||||
registerCoreTool(LSTool, this);
|
||||
registerCoreTool(ReadFileTool, this);
|
||||
registerCoreTool(GrepTool, this);
|
||||
registerCoreTool(GlobTool, this);
|
||||
registerCoreTool(EditTool, this);
|
||||
registerCoreTool(WriteFileTool, this);
|
||||
registerCoreTool(WebFetchTool, this);
|
||||
registerCoreTool(ReadManyFilesTool, this);
|
||||
registerCoreTool(ShellTool, this);
|
||||
registerCoreTool(MemoryTool);
|
||||
registerCoreTool(WebSearchTool, this);
|
||||
|
||||
await registry.discoverTools();
|
||||
return registry;
|
||||
}
|
||||
}
|
||||
// Export model constants for use in CLI
|
||||
export { DEFAULT_GEMINI_FLASH_MODEL };
|
||||
139
packages/core/src/config/flashFallback.test.ts
Normal file
139
packages/core/src/config/flashFallback.test.ts
Normal file
@@ -0,0 +1,139 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { Config } from './config.js';
|
||||
import { DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL } from './models.js';
|
||||
|
||||
describe('Flash Model Fallback Configuration', () => {
|
||||
let config: Config;
|
||||
|
||||
beforeEach(() => {
|
||||
config = new Config({
|
||||
sessionId: 'test-session',
|
||||
targetDir: '/test',
|
||||
debugMode: false,
|
||||
cwd: '/test',
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
});
|
||||
|
||||
// Initialize contentGeneratorConfig for testing
|
||||
(
|
||||
config as unknown as { contentGeneratorConfig: unknown }
|
||||
).contentGeneratorConfig = {
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
authType: 'oauth-personal',
|
||||
};
|
||||
});
|
||||
|
||||
describe('setModel', () => {
|
||||
it('should update the model and mark as switched during session', () => {
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle multiple model switches during session', () => {
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
|
||||
config.setModel('gemini-1.5-pro');
|
||||
expect(config.getModel()).toBe('gemini-1.5-pro');
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
|
||||
it('should only mark as switched if contentGeneratorConfig exists', () => {
|
||||
// Create config without initializing contentGeneratorConfig
|
||||
const newConfig = new Config({
|
||||
sessionId: 'test-session-2',
|
||||
targetDir: '/test',
|
||||
debugMode: false,
|
||||
cwd: '/test',
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
});
|
||||
|
||||
// Should not crash when contentGeneratorConfig is undefined
|
||||
newConfig.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(newConfig.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getModel', () => {
|
||||
it('should return contentGeneratorConfig model if available', () => {
|
||||
// Simulate initialized content generator config
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
});
|
||||
|
||||
it('should fallback to initial model if contentGeneratorConfig is not available', () => {
|
||||
// Test with fresh config where contentGeneratorConfig might not be set
|
||||
const newConfig = new Config({
|
||||
sessionId: 'test-session-2',
|
||||
targetDir: '/test',
|
||||
debugMode: false,
|
||||
cwd: '/test',
|
||||
model: 'custom-model',
|
||||
});
|
||||
|
||||
expect(newConfig.getModel()).toBe('custom-model');
|
||||
});
|
||||
});
|
||||
|
||||
describe('isModelSwitchedDuringSession', () => {
|
||||
it('should start as false for new session', () => {
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
|
||||
it('should remain false if no model switch occurs', () => {
|
||||
// Perform other operations that don't involve model switching
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
|
||||
it('should persist switched state throughout session', () => {
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
|
||||
// Should remain true even after getting model
|
||||
config.getModel();
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('resetModelToDefault', () => {
|
||||
it('should reset model to default and clear session switch flag', () => {
|
||||
// Switch to Flash first
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
|
||||
// Reset to default
|
||||
config.resetModelToDefault();
|
||||
|
||||
// Should be back to default with flag cleared
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle case where contentGeneratorConfig is not initialized', () => {
|
||||
// Create config without initializing contentGeneratorConfig
|
||||
const newConfig = new Config({
|
||||
sessionId: 'test-session-2',
|
||||
targetDir: '/test',
|
||||
debugMode: false,
|
||||
cwd: '/test',
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
});
|
||||
|
||||
// Should not crash when contentGeneratorConfig is undefined
|
||||
expect(() => newConfig.resetModelToDefault()).not.toThrow();
|
||||
expect(newConfig.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
9
packages/core/src/config/models.ts
Normal file
9
packages/core/src/config/models.ts
Normal file
@@ -0,0 +1,9 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
export const DEFAULT_GEMINI_MODEL = 'qwen3-coder-max';
|
||||
export const DEFAULT_GEMINI_FLASH_MODEL = 'gemini-2.5-flash';
|
||||
export const DEFAULT_GEMINI_EMBEDDING_MODEL = 'gemini-embedding-001';
|
||||
Reference in New Issue
Block a user