Initialize MCP tools once at start up instead of every time we auth. (#3483)

This commit is contained in:
Tommaso Sciortino 2025-07-07 15:01:59 -07:00 committed by GitHub
parent aa10ccba71
commit 357546a2aa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 76 additions and 99 deletions

View File

@ -115,15 +115,7 @@ export async function main() {
setMaxSizedBoxDebugging(config.getDebugMode());
// Initialize centralized FileDiscoveryService
config.getFileService();
if (config.getCheckpointingEnabled()) {
try {
await config.getGitService();
} catch {
// For now swallow the error, later log it.
}
}
await config.initialize();
if (settings.merged.theme) {
if (!themeManager.setActiveTheme(settings.merged.theme)) {
@ -133,12 +125,11 @@ export async function main() {
}
}
const memoryArgs = settings.merged.autoConfigureMaxOldSpaceSize
? getNodeMemoryArgs(config)
: [];
// hop into sandbox if we are outside and sandboxing is enabled
if (!process.env.SANDBOX) {
const memoryArgs = settings.merged.autoConfigureMaxOldSpaceSize
? getNodeMemoryArgs(config)
: [];
const sandboxConfig = config.getSandbox();
if (sandboxConfig) {
if (settings.merged.selectedAuthType) {

View File

@ -80,6 +80,7 @@ interface VertexGenerateContentResponse {
promptFeedback?: GenerateContentResponsePromptFeedback;
usageMetadata?: GenerateContentResponseUsageMetadata;
}
export interface CaCountTokenRequest {
request: VertexCountTokenRequest;
}

View File

@ -232,32 +232,30 @@ export class Config {
}
}
async initialize(): Promise<void> {
// Initialize centralized FileDiscoveryService
this.getFileService();
if (this.getCheckpointingEnabled()) {
try {
await this.getGitService();
} catch {
// For now swallow the error, later log it.
}
}
this.toolRegistry = await this.createToolRegistry();
}
async refreshAuth(authMethod: AuthType) {
// Always use the original default model when switching auth methods
// This ensures users don't stay on Flash after switching between auth types
// and allows API key users to get proper fallback behavior from getEffectiveModel
const modelToUse = this.model; // Use the original default model
// Temporarily clear contentGeneratorConfig to prevent getModel() from returning
// the previous session's model (which might be Flash)
this.contentGeneratorConfig = undefined!;
const contentConfig = await createContentGeneratorConfig(
modelToUse,
this.contentGeneratorConfig = await createContentGeneratorConfig(
this.model,
authMethod,
this,
);
const gc = new GeminiClient(this);
this.geminiClient = gc;
this.toolRegistry = await createToolRegistry(this);
await gc.initialize(contentConfig);
this.contentGeneratorConfig = contentConfig;
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;
// Note: In the future, we may want to reset any cached state when switching auth methods
}
getSessionId(): string {
@ -469,58 +467,59 @@ export class Config {
return { memoryContent, fileCount };
}
}
export function createToolRegistry(config: Config): Promise<ToolRegistry> {
const registry = new ToolRegistry(config);
const targetDir = config.getTargetDir();
async createToolRegistry(): Promise<ToolRegistry> {
const registry = new ToolRegistry(this);
const targetDir = this.getTargetDir();
// 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 = config.getCoreTools();
const excludeTools = config.getExcludeTools();
// 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}(`),
);
}
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 (
excludeTools?.includes(className) ||
excludeTools?.includes(toolName)
) {
isEnabled = false;
}
if (isEnabled) {
registry.registerTool(new ToolClass(...args));
}
};
if (isEnabled) {
registry.registerTool(new ToolClass(...args));
}
};
registerCoreTool(LSTool, targetDir, this);
registerCoreTool(ReadFileTool, targetDir, this);
registerCoreTool(GrepTool, targetDir);
registerCoreTool(GlobTool, targetDir, this);
registerCoreTool(EditTool, this);
registerCoreTool(WriteFileTool, this);
registerCoreTool(WebFetchTool, this);
registerCoreTool(ReadManyFilesTool, targetDir, this);
registerCoreTool(ShellTool, this);
registerCoreTool(MemoryTool);
registerCoreTool(WebSearchTool, this);
registerCoreTool(LSTool, targetDir, config);
registerCoreTool(ReadFileTool, targetDir, config);
registerCoreTool(GrepTool, targetDir);
registerCoreTool(GlobTool, targetDir, config);
registerCoreTool(EditTool, config);
registerCoreTool(WriteFileTool, config);
registerCoreTool(WebFetchTool, config);
registerCoreTool(ReadManyFilesTool, targetDir, config);
registerCoreTool(ShellTool, config);
registerCoreTool(MemoryTool);
registerCoreTool(WebSearchTool, config);
return (async () => {
await registry.discoverTools();
return registry;
})();
}
}
// Export model constants for use in CLI
export { DEFAULT_GEMINI_FLASH_MODEL };

View File

@ -50,7 +50,6 @@ export type ContentGeneratorConfig = {
export async function createContentGeneratorConfig(
model: string | undefined,
authType: AuthType | undefined,
config?: { getModel?: () => string },
): Promise<ContentGeneratorConfig> {
const geminiApiKey = process.env.GEMINI_API_KEY;
const googleApiKey = process.env.GOOGLE_API_KEY;
@ -58,7 +57,7 @@ export async function createContentGeneratorConfig(
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION;
// Use runtime model from config if available, otherwise fallback to parameter or default
const effectiveModel = config?.getModel?.() || model || DEFAULT_GEMINI_MODEL;
const effectiveModel = model || DEFAULT_GEMINI_MODEL;
const contentGeneratorConfig: ContentGeneratorConfig = {
model: effectiveModel,

View File

@ -38,21 +38,19 @@ describe('EditTool', () => {
let tempDir: string;
let rootDir: string;
let mockConfig: Config;
let geminiClient: any;
beforeEach(() => {
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-'));
rootDir = path.join(tempDir, 'root');
fs.mkdirSync(rootDir);
// The client instance that EditTool will use
const mockClientInstanceWithGenerateJson = {
geminiClient = {
generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
};
mockConfig = {
getGeminiClient: vi
.fn()
.mockReturnValue(mockClientInstanceWithGenerateJson),
getGeminiClient: vi.fn().mockReturnValue(geminiClient),
getTargetDir: () => rootDir,
getApprovalMode: vi.fn(),
setApprovalMode: vi.fn(),
@ -339,7 +337,7 @@ describe('EditTool', () => {
mockCalled = true;
expect(content).toBe(originalContent);
expect(p).toBe(params);
expect(client).toBe((tool as any).client);
expect(client).toBe(geminiClient);
return {
params: {
file_path: filePath,

View File

@ -18,7 +18,6 @@ import {
import { SchemaValidator } from '../utils/schemaValidator.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import { isNodeError } from '../utils/errors.js';
import { GeminiClient } from '../core/client.js';
import { Config, ApprovalMode } from '../config/config.js';
import { ensureCorrectEdit } from '../utils/editCorrector.js';
import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js';
@ -72,15 +71,13 @@ export class EditTool
implements ModifiableTool<EditToolParams>
{
static readonly Name = 'replace';
private readonly config: Config;
private readonly rootDirectory: string;
private readonly client: GeminiClient;
/**
* Creates a new instance of the EditLogic
* @param rootDirectory Root directory to ground this tool in.
*/
constructor(config: Config) {
constructor(private readonly config: Config) {
super(
EditTool.Name,
'Edit',
@ -123,9 +120,7 @@ Expectation for required parameters:
type: 'object',
},
);
this.config = config;
this.rootDirectory = path.resolve(this.config.getTargetDir());
this.client = config.getGeminiClient();
}
/**
@ -239,7 +234,7 @@ Expectation for required parameters:
params.file_path,
currentContent,
params,
this.client,
this.config.getGeminiClient(),
abortSignal,
);
finalOldString = correctedEdit.params.old_string;

View File

@ -154,8 +154,6 @@ export class ToolRegistry {
for (const tool of this.tools.values()) {
if (tool instanceof DiscoveredTool || tool instanceof DiscoveredMCPTool) {
this.tools.delete(tool.name);
} else {
// Keep manually registered tools
}
}

View File

@ -23,7 +23,6 @@ import {
ensureCorrectEdit,
ensureCorrectFileContent,
} from '../utils/editCorrector.js';
import { GeminiClient } from '../core/client.js';
import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js';
import { ModifiableTool, ModifyContext } from './modifiable-tool.js';
import { getSpecificMimeType } from '../utils/fileUtils.js';
@ -67,7 +66,6 @@ export class WriteFileTool
implements ModifiableTool<WriteFileToolParams>
{
static readonly Name: string = 'write_file';
private readonly client: GeminiClient;
constructor(private readonly config: Config) {
super(
@ -92,8 +90,6 @@ export class WriteFileTool
type: 'object',
},
);
this.client = this.config.getGeminiClient();
}
/**
@ -374,7 +370,7 @@ export class WriteFileTool
new_string: proposedContent,
file_path: filePath,
},
this.client,
this.config.getGeminiClient(),
abortSignal,
);
correctedContent = correctedParams.new_string;
@ -382,7 +378,7 @@ export class WriteFileTool
// This implies new file (ENOENT)
correctedContent = await ensureCorrectFileContent(
proposedContent,
this.client,
this.config.getGeminiClient(),
abortSignal,
);
}