mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-20 08:47:44 +00:00
952 lines
33 KiB
TypeScript
952 lines
33 KiB
TypeScript
/**
|
|
* @license
|
|
* Copyright 2025 Google LLC
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
|
import type { Mock } from 'vitest';
|
|
import type { ConfigParameters, SandboxConfig } from './config.js';
|
|
import { Config, ApprovalMode } from './config.js';
|
|
import * as path from 'node:path';
|
|
import { setGeminiMdFilename as mockSetGeminiMdFilename } from '../tools/memoryTool.js';
|
|
import {
|
|
DEFAULT_TELEMETRY_TARGET,
|
|
DEFAULT_OTLP_ENDPOINT,
|
|
QwenLogger,
|
|
} from '../telemetry/index.js';
|
|
import type { ContentGeneratorConfig } from '../core/contentGenerator.js';
|
|
import { DEFAULT_DASHSCOPE_BASE_URL } from '../core/openaiContentGenerator/constants.js';
|
|
import {
|
|
AuthType,
|
|
createContentGeneratorConfig,
|
|
} from '../core/contentGenerator.js';
|
|
import { GeminiClient } from '../core/client.js';
|
|
import { GitService } from '../services/gitService.js';
|
|
|
|
vi.mock('fs', async (importOriginal) => {
|
|
const actual = await importOriginal<typeof import('fs')>();
|
|
return {
|
|
...actual,
|
|
existsSync: vi.fn().mockReturnValue(true),
|
|
statSync: vi.fn().mockReturnValue({
|
|
isDirectory: vi.fn().mockReturnValue(true),
|
|
}),
|
|
realpathSync: vi.fn((path) => path),
|
|
};
|
|
});
|
|
|
|
import { ShellTool } from '../tools/shell.js';
|
|
import { ReadFileTool } from '../tools/read-file.js';
|
|
import { GrepTool } from '../tools/grep.js';
|
|
import { canUseRipgrep } from '../utils/ripgrepUtils.js';
|
|
import { RipGrepTool } from '../tools/ripGrep.js';
|
|
import { logRipgrepFallback } from '../telemetry/loggers.js';
|
|
import { RipgrepFallbackEvent } from '../telemetry/types.js';
|
|
import { ToolRegistry } from '../tools/tool-registry.js';
|
|
|
|
vi.mock('fs', async (importOriginal) => {
|
|
const actual = await importOriginal<typeof import('fs')>();
|
|
return {
|
|
...actual,
|
|
existsSync: vi.fn().mockReturnValue(true),
|
|
statSync: vi.fn().mockReturnValue({
|
|
isDirectory: vi.fn().mockReturnValue(true),
|
|
}),
|
|
realpathSync: vi.fn((path) => path),
|
|
};
|
|
});
|
|
|
|
// 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.discoverAllTools = 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.js');
|
|
vi.mock('../tools/ripGrep.js', () => ({
|
|
RipGrepTool: class MockRipGrepTool {},
|
|
}));
|
|
vi.mock('../utils/ripgrepUtils.js', () => ({
|
|
canUseRipgrep: vi.fn(),
|
|
}));
|
|
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',
|
|
QWEN_CONFIG_DIR: '.qwen',
|
|
}));
|
|
|
|
vi.mock('../core/contentGenerator.js');
|
|
|
|
vi.mock('../core/client.js', () => ({
|
|
GeminiClient: vi.fn().mockImplementation(() => ({
|
|
initialize: vi.fn().mockResolvedValue(undefined),
|
|
isInitialized: vi.fn().mockReturnValue(true),
|
|
stripThoughtsFromHistory: vi.fn(),
|
|
setTools: vi.fn(),
|
|
})),
|
|
}));
|
|
|
|
vi.mock('../telemetry/index.js', async (importOriginal) => {
|
|
const actual = await importOriginal<typeof import('../telemetry/index.js')>();
|
|
return {
|
|
...actual,
|
|
initializeTelemetry: vi.fn(),
|
|
uiTelemetryService: {
|
|
getLastPromptTokenCount: vi.fn(),
|
|
},
|
|
};
|
|
});
|
|
|
|
vi.mock('../telemetry/loggers.js', async (importOriginal) => {
|
|
const actual =
|
|
await importOriginal<typeof import('../telemetry/loggers.js')>();
|
|
return {
|
|
...actual,
|
|
logRipgrepFallback: vi.fn(),
|
|
};
|
|
});
|
|
|
|
vi.mock('../services/gitService.js', () => {
|
|
const GitServiceMock = vi.fn();
|
|
GitServiceMock.prototype.initialize = vi.fn();
|
|
return { GitService: GitServiceMock };
|
|
});
|
|
|
|
vi.mock('../ide/ide-client.js', () => ({
|
|
IdeClient: {
|
|
getInstance: vi.fn().mockResolvedValue({
|
|
getConnectionStatus: vi.fn(),
|
|
initialize: vi.fn(),
|
|
shutdown: vi.fn(),
|
|
}),
|
|
},
|
|
}));
|
|
|
|
import { BaseLlmClient } from '../core/baseLlmClient.js';
|
|
import { tokenLimit } from '../core/tokenLimits.js';
|
|
import { uiTelemetryService } from '../telemetry/index.js';
|
|
|
|
vi.mock('../core/baseLlmClient.js');
|
|
vi.mock('../core/tokenLimits.js', () => ({
|
|
tokenLimit: vi.fn(),
|
|
}));
|
|
|
|
describe('Server Config (config.ts)', () => {
|
|
const MODEL = 'qwen3-coder-plus';
|
|
const SANDBOX: SandboxConfig = {
|
|
command: 'docker',
|
|
image: 'qwen-code-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,
|
|
usageStatisticsEnabled: false,
|
|
};
|
|
|
|
beforeEach(() => {
|
|
// Reset mocks if necessary
|
|
vi.clearAllMocks();
|
|
vi.spyOn(QwenLogger.prototype, 'logStartSessionEvent').mockImplementation(
|
|
() => undefined,
|
|
);
|
|
});
|
|
|
|
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();
|
|
});
|
|
|
|
it('should throw an error if initialized more than once', async () => {
|
|
const config = new Config({
|
|
...baseParams,
|
|
checkpointing: false,
|
|
});
|
|
|
|
await expect(config.initialize()).resolves.toBeUndefined();
|
|
await expect(config.initialize()).rejects.toThrow(
|
|
'Config was already initialized',
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('refreshAuth', () => {
|
|
it('should refresh auth and update config', async () => {
|
|
const config = new Config(baseParams);
|
|
const authType = AuthType.USE_GEMINI;
|
|
const mockContentConfig = {
|
|
apiKey: 'test-key',
|
|
model: 'qwen3-coder-plus',
|
|
};
|
|
|
|
vi.mocked(createContentGeneratorConfig).mockReturnValue(
|
|
mockContentConfig,
|
|
);
|
|
|
|
// Set fallback mode to true to ensure it gets reset
|
|
config.setFallbackMode(true);
|
|
expect(config.isInFallbackMode()).toBe(true);
|
|
|
|
await config.refreshAuth(authType);
|
|
|
|
expect(createContentGeneratorConfig).toHaveBeenCalledWith(
|
|
config,
|
|
authType,
|
|
{
|
|
model: MODEL,
|
|
baseUrl: DEFAULT_DASHSCOPE_BASE_URL,
|
|
},
|
|
);
|
|
// Verify that contentGeneratorConfig is updated
|
|
expect(config.getContentGeneratorConfig()).toEqual(mockContentConfig);
|
|
expect(GeminiClient).toHaveBeenCalledWith(config);
|
|
// Verify that fallback mode is reset
|
|
expect(config.isInFallbackMode()).toBe(false);
|
|
});
|
|
|
|
it('should strip thoughts when switching from GenAI to Vertex', async () => {
|
|
const config = new Config(baseParams);
|
|
|
|
vi.mocked(createContentGeneratorConfig).mockImplementation(
|
|
(_: Config, authType: AuthType | undefined) =>
|
|
({ authType }) as unknown as ContentGeneratorConfig,
|
|
);
|
|
|
|
await config.refreshAuth(AuthType.USE_GEMINI);
|
|
|
|
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
|
|
|
|
expect(
|
|
config.getGeminiClient().stripThoughtsFromHistory,
|
|
).toHaveBeenCalledWith();
|
|
});
|
|
|
|
it('should not strip thoughts when switching from Vertex to GenAI', async () => {
|
|
const config = new Config(baseParams);
|
|
|
|
vi.mocked(createContentGeneratorConfig).mockImplementation(
|
|
(_: Config, authType: AuthType | undefined) =>
|
|
({ authType }) as unknown as ContentGeneratorConfig,
|
|
);
|
|
|
|
await config.refreshAuth(AuthType.USE_VERTEX_AI);
|
|
|
|
await config.refreshAuth(AuthType.USE_GEMINI);
|
|
|
|
expect(
|
|
config.getGeminiClient().stripThoughtsFromHistory,
|
|
).not.toHaveBeenCalledWith();
|
|
});
|
|
});
|
|
|
|
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('should initialize WorkspaceContext with includeDirectories', () => {
|
|
const includeDirectories = ['/path/to/dir1', '/path/to/dir2'];
|
|
const paramsWithIncludeDirs: ConfigParameters = {
|
|
...baseParams,
|
|
includeDirectories,
|
|
};
|
|
const config = new Config(paramsWithIncludeDirs);
|
|
const workspaceContext = config.getWorkspaceContext();
|
|
const directories = workspaceContext.getDirectories();
|
|
|
|
// Should include the target directory plus the included directories
|
|
expect(directories).toHaveLength(3);
|
|
expect(directories).toContain(path.resolve(baseParams.targetDir));
|
|
expect(directories).toContain('/path/to/dir1');
|
|
expect(directories).toContain('/path/to/dir2');
|
|
});
|
|
|
|
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('Config constructor should set telemetry useCollector to true when provided', () => {
|
|
const paramsWithTelemetry: ConfigParameters = {
|
|
...baseParams,
|
|
telemetry: { enabled: true, useCollector: true },
|
|
};
|
|
const config = new Config(paramsWithTelemetry);
|
|
expect(config.getTelemetryUseCollector()).toBe(true);
|
|
});
|
|
|
|
it('Config constructor should set telemetry useCollector to false when provided', () => {
|
|
const paramsWithTelemetry: ConfigParameters = {
|
|
...baseParams,
|
|
telemetry: { enabled: true, useCollector: false },
|
|
};
|
|
const config = new Config(paramsWithTelemetry);
|
|
expect(config.getTelemetryUseCollector()).toBe(false);
|
|
});
|
|
|
|
it('Config constructor should default telemetry useCollector to false if not provided', () => {
|
|
const paramsWithTelemetry: ConfigParameters = {
|
|
...baseParams,
|
|
telemetry: { enabled: true },
|
|
};
|
|
const config = new Config(paramsWithTelemetry);
|
|
expect(config.getTelemetryUseCollector()).toBe(false);
|
|
});
|
|
|
|
it('should have a getFileService method that returns FileDiscoveryService', () => {
|
|
const config = new Config(baseParams);
|
|
const fileService = config.getFileService();
|
|
expect(fileService).toBeDefined();
|
|
});
|
|
|
|
describe('Usage Statistics', () => {
|
|
it('defaults usage statistics to enabled if not specified', () => {
|
|
const config = new Config({
|
|
...baseParams,
|
|
usageStatisticsEnabled: undefined,
|
|
});
|
|
|
|
expect(config.getUsageStatisticsEnabled()).toBe(true);
|
|
});
|
|
|
|
it.each([{ enabled: true }, { enabled: false }])(
|
|
'sets usage statistics based on the provided value (enabled: $enabled)',
|
|
({ enabled }) => {
|
|
const config = new Config({
|
|
...baseParams,
|
|
usageStatisticsEnabled: enabled,
|
|
});
|
|
expect(config.getUsageStatisticsEnabled()).toBe(enabled);
|
|
},
|
|
);
|
|
|
|
it('logs the session start event', async () => {
|
|
const config = new Config({
|
|
...baseParams,
|
|
usageStatisticsEnabled: true,
|
|
});
|
|
await config.refreshAuth(AuthType.USE_GEMINI);
|
|
|
|
expect(QwenLogger.prototype.logStartSessionEvent).toHaveBeenCalledOnce();
|
|
});
|
|
});
|
|
|
|
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);
|
|
});
|
|
|
|
it('should return provided OTLP protocol', () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
telemetry: { enabled: true, otlpProtocol: 'http' },
|
|
};
|
|
const config = new Config(params);
|
|
expect(config.getTelemetryOtlpProtocol()).toBe('http');
|
|
});
|
|
|
|
it('should return default OTLP protocol if not provided', () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
telemetry: { enabled: true },
|
|
};
|
|
const config = new Config(params);
|
|
expect(config.getTelemetryOtlpProtocol()).toBe('grpc');
|
|
});
|
|
|
|
it('should return default OTLP protocol if telemetry object is not provided', () => {
|
|
const paramsWithoutTelemetry: ConfigParameters = { ...baseParams };
|
|
delete paramsWithoutTelemetry.telemetry;
|
|
const config = new Config(paramsWithoutTelemetry);
|
|
expect(config.getTelemetryOtlpProtocol()).toBe('grpc');
|
|
});
|
|
});
|
|
|
|
describe('UseRipgrep Configuration', () => {
|
|
it('should default useRipgrep to true when not provided', () => {
|
|
const config = new Config(baseParams);
|
|
expect(config.getUseRipgrep()).toBe(true);
|
|
});
|
|
|
|
it('should set useRipgrep to false when provided as false', () => {
|
|
const paramsWithRipgrep: ConfigParameters = {
|
|
...baseParams,
|
|
useRipgrep: false,
|
|
};
|
|
const config = new Config(paramsWithRipgrep);
|
|
expect(config.getUseRipgrep()).toBe(false);
|
|
});
|
|
|
|
it('should set useRipgrep to true when explicitly provided as true', () => {
|
|
const paramsWithRipgrep: ConfigParameters = {
|
|
...baseParams,
|
|
useRipgrep: true,
|
|
};
|
|
const config = new Config(paramsWithRipgrep);
|
|
expect(config.getUseRipgrep()).toBe(true);
|
|
});
|
|
|
|
it('should default useRipgrep to true when undefined', () => {
|
|
const paramsWithUndefinedRipgrep: ConfigParameters = {
|
|
...baseParams,
|
|
useRipgrep: undefined,
|
|
};
|
|
const config = new Config(paramsWithUndefinedRipgrep);
|
|
expect(config.getUseRipgrep()).toBe(true);
|
|
});
|
|
});
|
|
|
|
describe('createToolRegistry', () => {
|
|
it('should register a tool if coreTools contains an argument-specific pattern', async () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
coreTools: ['ShellTool(git status)'],
|
|
};
|
|
const config = new Config(params);
|
|
await config.initialize();
|
|
|
|
// The ToolRegistry class is mocked, so we can inspect its prototype's methods.
|
|
const registerToolMock = (
|
|
(await vi.importMock('../tools/tool-registry')) as {
|
|
ToolRegistry: { prototype: { registerTool: Mock } };
|
|
}
|
|
).ToolRegistry.prototype.registerTool;
|
|
|
|
// Check that registerTool was called for ShellTool
|
|
const wasShellToolRegistered = (registerToolMock as Mock).mock.calls.some(
|
|
(call) => call[0] instanceof vi.mocked(ShellTool),
|
|
);
|
|
expect(wasShellToolRegistered).toBe(true);
|
|
|
|
// Check that registerTool was NOT called for ReadFileTool
|
|
const wasReadFileToolRegistered = (
|
|
registerToolMock as Mock
|
|
).mock.calls.some((call) => call[0] instanceof vi.mocked(ReadFileTool));
|
|
expect(wasReadFileToolRegistered).toBe(false);
|
|
});
|
|
|
|
describe('with minified tool class names', () => {
|
|
beforeEach(() => {
|
|
Object.defineProperty(
|
|
vi.mocked(ShellTool).prototype.constructor,
|
|
'name',
|
|
{
|
|
value: '_ShellTool',
|
|
configurable: true,
|
|
},
|
|
);
|
|
});
|
|
|
|
afterEach(() => {
|
|
Object.defineProperty(
|
|
vi.mocked(ShellTool).prototype.constructor,
|
|
'name',
|
|
{
|
|
value: 'ShellTool',
|
|
},
|
|
);
|
|
});
|
|
|
|
it('should register a tool if coreTools contains the non-minified class name', async () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
coreTools: ['ShellTool'],
|
|
};
|
|
const config = new Config(params);
|
|
await config.initialize();
|
|
|
|
const registerToolMock = (
|
|
(await vi.importMock('../tools/tool-registry')) as {
|
|
ToolRegistry: { prototype: { registerTool: Mock } };
|
|
}
|
|
).ToolRegistry.prototype.registerTool;
|
|
|
|
const wasShellToolRegistered = (
|
|
registerToolMock as Mock
|
|
).mock.calls.some((call) => call[0] instanceof vi.mocked(ShellTool));
|
|
expect(wasShellToolRegistered).toBe(true);
|
|
});
|
|
|
|
it('should not register a tool if excludeTools contains the non-minified class name', async () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
coreTools: undefined, // all tools enabled by default
|
|
excludeTools: ['ShellTool'],
|
|
};
|
|
const config = new Config(params);
|
|
await config.initialize();
|
|
|
|
const registerToolMock = (
|
|
(await vi.importMock('../tools/tool-registry')) as {
|
|
ToolRegistry: { prototype: { registerTool: Mock } };
|
|
}
|
|
).ToolRegistry.prototype.registerTool;
|
|
|
|
const wasShellToolRegistered = (
|
|
registerToolMock as Mock
|
|
).mock.calls.some((call) => call[0] instanceof vi.mocked(ShellTool));
|
|
expect(wasShellToolRegistered).toBe(false);
|
|
});
|
|
|
|
it('should register a tool if coreTools contains an argument-specific pattern with the non-minified class name', async () => {
|
|
const params: ConfigParameters = {
|
|
...baseParams,
|
|
coreTools: ['ShellTool(git status)'],
|
|
};
|
|
const config = new Config(params);
|
|
await config.initialize();
|
|
|
|
const registerToolMock = (
|
|
(await vi.importMock('../tools/tool-registry')) as {
|
|
ToolRegistry: { prototype: { registerTool: Mock } };
|
|
}
|
|
).ToolRegistry.prototype.registerTool;
|
|
|
|
const wasShellToolRegistered = (
|
|
registerToolMock as Mock
|
|
).mock.calls.some((call) => call[0] instanceof vi.mocked(ShellTool));
|
|
expect(wasShellToolRegistered).toBe(true);
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('getTruncateToolOutputThreshold', () => {
|
|
beforeEach(() => {
|
|
vi.clearAllMocks();
|
|
});
|
|
|
|
it('should return the calculated threshold when it is smaller than the default', () => {
|
|
const config = new Config(baseParams);
|
|
vi.mocked(tokenLimit).mockReturnValue(32000);
|
|
vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(
|
|
1000,
|
|
);
|
|
// 4 * (32000 - 1000) = 4 * 31000 = 124000
|
|
// default is 4_000_000
|
|
expect(config.getTruncateToolOutputThreshold()).toBe(124000);
|
|
});
|
|
|
|
it('should return the default threshold when the calculated value is larger', () => {
|
|
const config = new Config(baseParams);
|
|
vi.mocked(tokenLimit).mockReturnValue(2_000_000);
|
|
vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(
|
|
500_000,
|
|
);
|
|
// 4 * (2_000_000 - 500_000) = 4 * 1_500_000 = 6_000_000
|
|
// default is 4_000_000
|
|
expect(config.getTruncateToolOutputThreshold()).toBe(4_000_000);
|
|
});
|
|
|
|
it('should use a custom truncateToolOutputThreshold if provided', () => {
|
|
const customParams = {
|
|
...baseParams,
|
|
truncateToolOutputThreshold: 50000,
|
|
};
|
|
const config = new Config(customParams);
|
|
vi.mocked(tokenLimit).mockReturnValue(8000);
|
|
vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(
|
|
2000,
|
|
);
|
|
// 4 * (8000 - 2000) = 4 * 6000 = 24000
|
|
// custom threshold is 50000
|
|
expect(config.getTruncateToolOutputThreshold()).toBe(24000);
|
|
|
|
vi.mocked(tokenLimit).mockReturnValue(32000);
|
|
vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(
|
|
1000,
|
|
);
|
|
// 4 * (32000 - 1000) = 124000
|
|
// custom threshold is 50000
|
|
expect(config.getTruncateToolOutputThreshold()).toBe(50000);
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('setApprovalMode with folder trust', () => {
|
|
const baseParams: ConfigParameters = {
|
|
sessionId: 'test',
|
|
targetDir: '.',
|
|
debugMode: false,
|
|
model: 'test-model',
|
|
cwd: '.',
|
|
};
|
|
|
|
it('should throw an error when setting YOLO mode in an untrusted folder', () => {
|
|
const config = new Config(baseParams);
|
|
vi.spyOn(config, 'isTrustedFolder').mockReturnValue(false);
|
|
expect(() => config.setApprovalMode(ApprovalMode.YOLO)).toThrow(
|
|
'Cannot enable privileged approval modes in an untrusted folder.',
|
|
);
|
|
});
|
|
|
|
it('should throw an error when setting AUTO_EDIT mode in an untrusted folder', () => {
|
|
const config = new Config(baseParams);
|
|
vi.spyOn(config, 'isTrustedFolder').mockReturnValue(false);
|
|
expect(() => config.setApprovalMode(ApprovalMode.AUTO_EDIT)).toThrow(
|
|
'Cannot enable privileged approval modes in an untrusted folder.',
|
|
);
|
|
});
|
|
|
|
it('should NOT throw an error when setting DEFAULT mode in an untrusted folder', () => {
|
|
const config = new Config(baseParams);
|
|
vi.spyOn(config, 'isTrustedFolder').mockReturnValue(false);
|
|
expect(() => config.setApprovalMode(ApprovalMode.DEFAULT)).not.toThrow();
|
|
});
|
|
|
|
it('should NOT throw an error when setting PLAN mode in an untrusted folder', () => {
|
|
const config = new Config({
|
|
sessionId: 'test',
|
|
targetDir: '.',
|
|
debugMode: false,
|
|
model: 'test-model',
|
|
cwd: '.',
|
|
trustedFolder: false, // Untrusted
|
|
});
|
|
expect(() => config.setApprovalMode(ApprovalMode.PLAN)).not.toThrow();
|
|
});
|
|
|
|
it('should NOT throw an error when setting any mode in a trusted folder', () => {
|
|
const config = new Config(baseParams);
|
|
vi.spyOn(config, 'isTrustedFolder').mockReturnValue(true);
|
|
expect(() => config.setApprovalMode(ApprovalMode.YOLO)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.AUTO_EDIT)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.DEFAULT)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.PLAN)).not.toThrow();
|
|
});
|
|
|
|
it('should NOT throw an error when setting any mode if trustedFolder is undefined', () => {
|
|
const config = new Config(baseParams);
|
|
vi.spyOn(config, 'isTrustedFolder').mockReturnValue(true); // isTrustedFolder defaults to true
|
|
expect(() => config.setApprovalMode(ApprovalMode.YOLO)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.AUTO_EDIT)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.DEFAULT)).not.toThrow();
|
|
expect(() => config.setApprovalMode(ApprovalMode.PLAN)).not.toThrow();
|
|
});
|
|
|
|
describe('registerCoreTools', () => {
|
|
beforeEach(() => {
|
|
vi.clearAllMocks();
|
|
});
|
|
|
|
it('should register RipGrepTool when useRipgrep is true and it is available', async () => {
|
|
(canUseRipgrep as Mock).mockResolvedValue(true);
|
|
const config = new Config({ ...baseParams, useRipgrep: true });
|
|
await config.initialize();
|
|
|
|
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
|
|
const wasRipGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(RipGrepTool),
|
|
);
|
|
const wasGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(GrepTool),
|
|
);
|
|
|
|
expect(wasRipGrepRegistered).toBe(true);
|
|
expect(wasGrepRegistered).toBe(false);
|
|
expect(logRipgrepFallback).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it('should register GrepTool as a fallback when useRipgrep is true but it is not available', async () => {
|
|
(canUseRipgrep as Mock).mockResolvedValue(false);
|
|
const config = new Config({ ...baseParams, useRipgrep: true });
|
|
await config.initialize();
|
|
|
|
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
|
|
const wasRipGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(RipGrepTool),
|
|
);
|
|
const wasGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(GrepTool),
|
|
);
|
|
|
|
expect(wasRipGrepRegistered).toBe(false);
|
|
expect(wasGrepRegistered).toBe(true);
|
|
expect(logRipgrepFallback).toHaveBeenCalledWith(
|
|
config,
|
|
expect.any(RipgrepFallbackEvent),
|
|
);
|
|
const event = (logRipgrepFallback as Mock).mock.calls[0][1];
|
|
expect(event.error).toBeUndefined();
|
|
});
|
|
|
|
it('should register GrepTool as a fallback when canUseRipgrep throws an error', async () => {
|
|
const error = new Error('ripGrep check failed');
|
|
(canUseRipgrep as Mock).mockRejectedValue(error);
|
|
const config = new Config({ ...baseParams, useRipgrep: true });
|
|
await config.initialize();
|
|
|
|
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
|
|
const wasRipGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(RipGrepTool),
|
|
);
|
|
const wasGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(GrepTool),
|
|
);
|
|
|
|
expect(wasRipGrepRegistered).toBe(false);
|
|
expect(wasGrepRegistered).toBe(true);
|
|
expect(logRipgrepFallback).toHaveBeenCalledWith(
|
|
config,
|
|
expect.any(RipgrepFallbackEvent),
|
|
);
|
|
const event = (logRipgrepFallback as Mock).mock.calls[0][1];
|
|
expect(event.error).toBe(String(error));
|
|
});
|
|
|
|
it('should register GrepTool when useRipgrep is false', async () => {
|
|
const config = new Config({ ...baseParams, useRipgrep: false });
|
|
await config.initialize();
|
|
|
|
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
|
|
const wasRipGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(RipGrepTool),
|
|
);
|
|
const wasGrepRegistered = calls.some(
|
|
(call) => call[0] instanceof vi.mocked(GrepTool),
|
|
);
|
|
|
|
expect(wasRipGrepRegistered).toBe(false);
|
|
expect(wasGrepRegistered).toBe(true);
|
|
expect(canUseRipgrep).not.toHaveBeenCalled();
|
|
expect(logRipgrepFallback).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
});
|
|
|
|
describe('BaseLlmClient Lifecycle', () => {
|
|
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,
|
|
usageStatisticsEnabled: false,
|
|
};
|
|
|
|
it('should throw an error if getBaseLlmClient is called before refreshAuth', () => {
|
|
const config = new Config(baseParams);
|
|
expect(() => config.getBaseLlmClient()).toThrow(
|
|
'BaseLlmClient not initialized. Ensure authentication has occurred and ContentGenerator is ready.',
|
|
);
|
|
});
|
|
|
|
it('should successfully initialize BaseLlmClient after refreshAuth is called', async () => {
|
|
const config = new Config(baseParams);
|
|
const authType = AuthType.USE_GEMINI;
|
|
const mockContentConfig = { model: 'gemini-flash', apiKey: 'test-key' };
|
|
|
|
vi.mocked(createContentGeneratorConfig).mockReturnValue(mockContentConfig);
|
|
|
|
await config.refreshAuth(authType);
|
|
|
|
// Should not throw
|
|
const llmService = config.getBaseLlmClient();
|
|
expect(llmService).toBeDefined();
|
|
expect(BaseLlmClient).toHaveBeenCalledWith(
|
|
config.getContentGenerator(),
|
|
config,
|
|
);
|
|
});
|
|
});
|