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
8 changed files with 76 additions and 99 deletions

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 };