@machina.ai/cell-cli-core 1.19.4-rc3 → 1.22.5-rc1

package/dist/src/config/config.js

@@ -8,6 +8,7 @@ import { inspect } from 'node:util';
 import process from 'node:process';
 import { AuthType, createContentGenerator, createContentGeneratorConfig, } from '../core/contentGenerator.js';
 import { PromptRegistry } from '../prompts/prompt-registry.js';
+import { ResourceRegistry } from '../resources/resource-registry.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
 import { LSTool } from '../tools/ls.js';
 import { ReadFileTool } from '../tools/read-file.js';
@@ -28,17 +29,19 @@ import { GitService } from '../services/gitService.js';
 import { initializeTelemetry, DEFAULT_TELEMETRY_TARGET, DEFAULT_OTLP_ENDPOINT, uiTelemetryService, } from '../telemetry/index.js';
 import { coreEvents } from '../utils/events.js';
 import { tokenLimit } from '../core/tokenLimits.js';
-import { DEFAULT_GEMINI_EMBEDDING_MODEL, DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_THINKING_MODE, } from './models.js';
+import { DEFAULT_GEMINI_EMBEDDING_MODEL, DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL_AUTO, DEFAULT_THINKING_MODE, isPreviewModel, PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_MODEL_AUTO, } from './models.js';
 import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
 import { ideContextStore } from '../ide/ideContext.js';
 import { WriteTodosTool } from '../tools/write-todos.js';
 import { StandardFileSystemService } from '../services/fileSystemService.js';
 import { logRipgrepFallback } from '../telemetry/loggers.js';
 import { RipgrepFallbackEvent } from '../telemetry/types.js';
+import { ModelAvailabilityService } from '../availability/modelAvailabilityService.js';
 import { ModelRouterService } from '../routing/modelRouterService.js';
 import { OutputFormat } from '../output/types.js';
 import { ModelConfigService } from '../services/modelConfigService.js';
 import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
+import { ContextManager } from '../services/contextManager.js';
 import { WorkspaceContext } from '../utils/workspaceContext.js';
 import { Storage } from './storage.js';
 import { FileExclusions } from '../utils/ignorePatterns.js';
@@ -48,10 +51,12 @@ import { HookSystem } from '../hooks/index.js';
 import { getCodeAssistServer } from '../code_assist/codeAssist.js';
 import { AgentRegistry } from '../agents/registry.js';
 import { setGlobalProxy } from '../utils/fetch.js';
-import { SubagentToolWrapper } from '../agents/subagent-tool-wrapper.js';
+import { DelegateToAgentTool } from '../agents/delegate-to-agent-tool.js';
+import { DELEGATE_TO_AGENT_TOOL_NAME } from '../tools/tool-names.js';
 import { getExperiments } from '../code_assist/experiments/experiments.js';
 import { ExperimentFlags } from '../code_assist/experiments/flagNames.js';
 import { debugLogger } from '../utils/debugLogger.js';
+import { startupProfiler } from '../telemetry/startupProfiler.js';
 import { ApprovalMode } from '../policy/types.js';
 export { ApprovalMode };
 import { DEFAULT_FILE_FILTERING_OPTIONS, DEFAULT_MEMORY_FILE_FILTERING_OPTIONS, } from './constants.js';
@@ -70,6 +75,7 @@ export class MCPServerConfig {
     httpUrl;
     headers;
     tcp;
+    type;
     timeout;
     trust;
     description;
@@ -89,6 +95,12 @@ export class MCPServerConfig {
     httpUrl, headers, 
     // For websocket transport
     tcp, 
+    // Transport type (optional, for use with 'url' field)
+    // When set to 'http', uses StreamableHTTPClientTransport
+    // When set to 'sse', uses SSEClientTransport
+    // When omitted, auto-detects transport type
+    // Note: 'httpUrl' is deprecated in favor of 'url' + 'type'
+    type, 
     // Common
     timeout, trust, 
     // Metadata
@@ -109,6 +121,7 @@ export class MCPServerConfig {
         this.httpUrl = httpUrl;
         this.headers = headers;
         this.tcp = tcp;
+        this.type = type;
         this.timeout = timeout;
         this.trust = trust;
         this.description = description;
@@ -133,6 +146,7 @@ export class Config {
     allowedMcpServers;
     blockedMcpServers;
     promptRegistry;
+    resourceRegistry;
     agentRegistry;
     sessionId;
     fileSystemService;
@@ -155,7 +169,6 @@ export class Config {
     userMemory;
     geminiMdFileCount;
     geminiMdFilePaths;
-    approvalMode;
     showMemoryUsage;
     accessibility;
     telemetrySettings;
@@ -163,6 +176,7 @@ export class Config {
     geminiClient;
     baseLlmClient;
     modelRouterService;
+    modelAvailabilityService;
     fileFiltering;
     fileDiscoveryService = null;
     gitService = undefined;
@@ -176,10 +190,12 @@ export class Config {
     bugCommand;
     model;
     previewFeatures;
+    hasAccessToPreviewModel = false;
     noBrowser;
     folderTrust;
     ideMode;
     inFallbackMode = false;
+    _activeModel;
     maxSessionTurns;
     listSessions;
     deleteSession;
@@ -221,18 +237,23 @@ export class Config {
     continueOnFailedApiCall;
     retryFetchErrors;
     enableShellOutputEfficiency;
+    shellToolInactivityTimeout;
     fakeResponses;
     recordResponses;
     disableYoloMode;
     pendingIncludeDirectories;
     enableHooks;
     hooks;
+    disabledHooks;
     experiments;
     experimentsPromise;
     hookSystem;
     previewModelFallbackMode = false;
     previewModelBypassMode = false;
     enableModelAvailabilityService;
+    enableAgents;
+    experimentalJitContext;
+    contextManager;
     constructor(params) {
         this.sessionId = params.sessionId;
         this.embeddingModel =
@@ -257,7 +278,6 @@ export class Config {
         this.userMemory = params.userMemory ?? '';
         this.geminiMdFileCount = params.geminiMdFileCount ?? 0;
         this.geminiMdFilePaths = params.geminiMdFilePaths ?? [];
-        this.approvalMode = params.approvalMode ?? ApprovalMode.DEFAULT;
         this.showMemoryUsage = params.showMemoryUsage ?? false;
         this.accessibility = params.accessibility ?? {};
         this.telemetrySettings = {
@@ -268,6 +288,7 @@ export class Config {
             logPrompts: params.telemetry?.logPrompts ?? true,
             outfile: params.telemetry?.outfile,
             useCollector: params.telemetry?.useCollector,
+            useCliAuth: params.telemetry?.useCliAuth,
         };
         this.usageStatisticsEnabled = params.usageStatisticsEnabled ?? true;
         this.fileFiltering = {
@@ -288,8 +309,11 @@ export class Config {
         this.fileDiscoveryService = params.fileDiscoveryService ?? null;
         this.bugCommand = params.bugCommand;
         this.model = params.model;
-        this.enableModelAvailabilityService =
-            params.enableModelAvailabilityService ?? false;
+        this._activeModel = params.model;
+        this.enableModelAvailabilityService = true;
+        this.enableAgents = params.enableAgents ?? false;
+        this.experimentalJitContext = params.experimentalJitContext ?? false;
+        this.modelAvailabilityService = new ModelAvailabilityService();
         this.previewFeatures = params.previewFeatures ?? undefined;
         this.maxSessionTurns = params.maxSessionTurns ?? -1;
         this.experimentalZedIntegration =
@@ -328,8 +352,15 @@ export class Config {
             params.truncateToolOutputLines ?? DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES;
         this.enableToolOutputTruncation = params.enableToolOutputTruncation ?? true;
         this.useSmartEdit = params.useSmartEdit ?? true;
-        this.useWriteTodos = params.useWriteTodos ?? true;
+        // // TODO(joshualitt): Re-evaluate the todo tool for 3 family.
+        this.useWriteTodos = isPreviewModel(this.model)
+            ? false
+            : (params.useWriteTodos ?? true);
         this.enableHooks = params.enableHooks ?? false;
+        this.disabledHooks =
+            (params.hooks && 'disabled' in params.hooks
+                ? params.hooks.disabled
+                : undefined) ?? [];
         // Enable MessageBus integration if:
         // 1. Explicitly enabled via setting, OR
         // 2. Hooks are enabled and hooks are configured
@@ -344,11 +375,13 @@ export class Config {
             maxTimeMinutes: params.codebaseInvestigatorSettings?.maxTimeMinutes ?? 3,
             thinkingBudget: params.codebaseInvestigatorSettings?.thinkingBudget ??
                 DEFAULT_THINKING_MODE,
-            model: params.codebaseInvestigatorSettings?.model ?? DEFAULT_GEMINI_MODEL,
+            model: params.codebaseInvestigatorSettings?.model,
         };
         this.continueOnFailedApiCall = params.continueOnFailedApiCall ?? true;
         this.enableShellOutputEfficiency =
             params.enableShellOutputEfficiency ?? true;
+        this.shellToolInactivityTimeout =
+            (params.shellToolInactivityTimeout ?? 300) * 1000; // 5 minutes
         this.extensionManagement = params.extensionManagement ?? true;
         this.enableExtensionReloading = params.enableExtensionReloading ?? false;
         this.storage = new Storage(this.targetDir);
@@ -357,7 +390,10 @@ export class Config {
         this.enablePromptCompletion = params.enablePromptCompletion ?? false;
         this.fileExclusions = new FileExclusions(this);
         this.eventEmitter = params.eventEmitter;
-        this.policyEngine = new PolicyEngine(params.policyEngineConfig);
+        this.policyEngine = new PolicyEngine({
+            ...params.policyEngineConfig,
+            approvalMode: params.approvalMode ?? params.policyEngineConfig?.approvalMode,
+        });
         this.messageBus = new MessageBus(this.policyEngine, this.debugMode);
         this.outputSettings = {
             format: params.output?.format ?? OutputFormat.TEXT,
@@ -370,6 +406,7 @@ export class Config {
             setGeminiMdFilename(params.contextFileName);
         }
         if (this.telemetrySettings.enabled) {
+            // eslint-disable-next-line @typescript-eslint/no-floating-promises
             initializeTelemetry(this);
         }
         const proxy = this.getProxy();
@@ -391,11 +428,19 @@ export class Config {
         // TODO(12593): Fix the settings loading logic to properly merge defaults and
         // remove this hack.
         let modelConfigServiceConfig = params.modelConfigServiceConfig;
-        if (modelConfigServiceConfig && !modelConfigServiceConfig.aliases) {
-            modelConfigServiceConfig = {
-                ...modelConfigServiceConfig,
-                aliases: DEFAULT_MODEL_CONFIGS.aliases,
-            };
+        if (modelConfigServiceConfig) {
+            if (!modelConfigServiceConfig.aliases) {
+                modelConfigServiceConfig = {
+                    ...modelConfigServiceConfig,
+                    aliases: DEFAULT_MODEL_CONFIGS.aliases,
+                };
+            }
+            if (!modelConfigServiceConfig.overrides) {
+                modelConfigServiceConfig = {
+                    ...modelConfigServiceConfig,
+                    overrides: DEFAULT_MODEL_CONFIGS.overrides,
+                };
+            }
         }
         this.modelConfigService = new ModelConfigService(modelConfigServiceConfig ?? DEFAULT_MODEL_CONFIGS);
     }
@@ -408,30 +453,40 @@ export class Config {
         }
         this.initialized = true;
         // Initialize centralized FileDiscoveryService
+        const discoverToolsHandle = startupProfiler.start('discover_tools');
         this.getFileService();
         if (this.getCheckpointingEnabled()) {
             await this.getGitService();
         }
         this.promptRegistry = new PromptRegistry();
+        this.resourceRegistry = new ResourceRegistry();
         this.agentRegistry = new AgentRegistry(this);
         await this.agentRegistry.initialize();
         this.toolRegistry = await this.createToolRegistry();
+        discoverToolsHandle?.end();
         this.mcpClientManager = new McpClientManager(this.toolRegistry, this, this.eventEmitter);
+        const initMcpHandle = startupProfiler.start('initialize_mcp_clients');
         await Promise.all([
             await this.mcpClientManager.startConfiguredMcpServers(),
             await this.getExtensionLoader().start(this),
         ]);
+        initMcpHandle?.end();
         // Initialize hook system if enabled
         if (this.enableHooks) {
             this.hookSystem = new HookSystem(this);
             await this.hookSystem.initialize();
         }
+        if (this.experimentalJitContext) {
+            this.contextManager = new ContextManager(this);
+        }
         await this.geminiClient.initialize();
     }
     getContentGenerator() {
         return this.contentGenerator;
     }
     async refreshAuth(authMethod, displayMessage) {
+        // Reset availability service when switching auth
+        this.modelAvailabilityService.reset();
         // Vertex and Genai have incompatible encryption and sending history with
         // thoughtSignature from Genai to Vertex will fail, we need to strip them
         if (this.contentGeneratorConfig?.authType === AuthType.USE_GEMINI &&
@@ -439,6 +494,8 @@ export class Config {
             // Restore the conversation history to the new client
             this.geminiClient.stripThoughtsFromHistory();
         }
+        // Reset availability status when switching auth (e.g. from limited key to OAuth)
+        this.modelAvailabilityService.reset();
         const newContentGeneratorConfig = await createContentGeneratorConfig(this, authMethod, {
             clientId: this.keycloakClientId,
             deviceEndpoint: this.keycloakDeviceEndpoint,
@@ -449,14 +506,16 @@ export class Config {
         this.contentGeneratorConfig = newContentGeneratorConfig;
         // Initialize BaseLlmClient now that the ContentGenerator is available
         this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this);
-        const previewFeatures = this.getPreviewFeatures();
         const codeAssistServer = getCodeAssistServer(this);
         if (codeAssistServer) {
+            if (codeAssistServer.projectId) {
+                await this.refreshUserQuota();
+            }
             this.experimentsPromise = getExperiments(codeAssistServer)
                 .then((experiments) => {
                 this.setExperiments(experiments);
                 // If preview features have not been set and the user authenticated through Google, we enable preview based on remote config only if it's true
-                if (previewFeatures === undefined) {
+                if (this.getPreviewFeatures() === undefined) {
                     const remotePreviewFeatures = experiments.flags[ExperimentFlags.ENABLE_PREVIEW]?.boolValue;
                     if (remotePreviewFeatures === true) {
                         this.setPreviewFeatures(remotePreviewFeatures);
@@ -471,8 +530,19 @@ export class Config {
             this.experiments = undefined;
             this.experimentsPromise = undefined;
         }
+        const authType = this.contentGeneratorConfig.authType;
+        if (authType === AuthType.USE_GEMINI ||
+            authType === AuthType.USE_VERTEX_AI) {
+            this.setHasAccessToPreviewModel(true);
+        }
         // Reset the session flag since we're explicitly changing auth and using default model
         this.inFallbackMode = false;
+        // Update model if user no longer has access to the preview model
+        if (!this.getPreviewFeatures() &&
+            !this.hasAccessToPreviewModel &&
+            isPreviewModel(this.model)) {
+            this.setModel(DEFAULT_GEMINI_MODEL_AUTO);
+        }
     }
     async getExperimentsAsync() {
         if (this.experiments) {
@@ -526,9 +596,23 @@ export class Config {
     setModel(newModel) {
         if (this.model !== newModel || this.inFallbackMode) {
             this.model = newModel;
+            // When the user explicitly sets a model, that becomes the active model.
+            this._activeModel = newModel;
             coreEvents.emitModelChanged(newModel);
         }
         this.setFallbackMode(false);
+        this.modelAvailabilityService.reset();
+    }
+    getActiveModel() {
+        return this._activeModel ?? this.model;
+    }
+    setActiveModel(model) {
+        if (this._activeModel !== model) {
+            this._activeModel = model;
+        }
+    }
+    resetTurn() {
+        this.modelAvailabilityService.resetTurn();
     }
     isInFallbackMode() {
         return this.inFallbackMode;
@@ -595,6 +679,9 @@ export class Config {
     getPromptRegistry() {
         return this.promptRegistry;
     }
+    getResourceRegistry() {
+        return this.resourceRegistry;
+    }
     getDebugMode() {
         return this.debugMode;
     }
@@ -605,7 +692,44 @@ export class Config {
         return this.previewFeatures;
     }
     setPreviewFeatures(previewFeatures) {
+        // No change in state, no action needed
+        if (this.previewFeatures === previewFeatures) {
+            return;
+        }
         this.previewFeatures = previewFeatures;
+        const currentModel = this.getModel();
+        // Case 1: Disabling preview features while on a preview model
+        if (!previewFeatures && isPreviewModel(currentModel)) {
+            this.setModel(DEFAULT_GEMINI_MODEL_AUTO);
+        }
+        // Case 2: Enabling preview features while on the default auto model
+        else if (previewFeatures && currentModel === DEFAULT_GEMINI_MODEL_AUTO) {
+            this.setModel(PREVIEW_GEMINI_MODEL_AUTO);
+        }
+    }
+    getHasAccessToPreviewModel() {
+        return this.hasAccessToPreviewModel;
+    }
+    setHasAccessToPreviewModel(hasAccess) {
+        this.hasAccessToPreviewModel = hasAccess;
+    }
+    async refreshUserQuota() {
+        const codeAssistServer = getCodeAssistServer(this);
+        if (!codeAssistServer || !codeAssistServer.projectId) {
+            return undefined;
+        }
+        try {
+            const quota = await codeAssistServer.retrieveUserQuota({
+                project: codeAssistServer.projectId,
+            });
+            const hasAccess = quota.buckets?.some((b) => b.modelId === PREVIEW_GEMINI_MODEL) ?? false;
+            this.setHasAccessToPreviewModel(hasAccess);
+            return quota;
+        }
+        catch (e) {
+            debugLogger.debug('Failed to retrieve user quota', e);
+            return undefined;
+        }
     }
     getCoreTools() {
         return this.coreTools;
@@ -666,6 +790,18 @@ export class Config {
     setUserMemory(newUserMemory) {
         this.userMemory = newUserMemory;
     }
+    getGlobalMemory() {
+        return this.contextManager?.getGlobalMemory() ?? '';
+    }
+    getEnvironmentMemory() {
+        return this.contextManager?.getEnvironmentMemory() ?? '';
+    }
+    getContextManager() {
+        return this.contextManager;
+    }
+    isJitContextEnabled() {
+        return this.experimentalJitContext;
+    }
     getGeminiMdFileCount() {
         return this.geminiMdFileCount;
     }
@@ -679,13 +815,13 @@ export class Config {
         this.geminiMdFilePaths = paths;
     }
     getApprovalMode() {
-        return this.approvalMode;
+        return this.policyEngine.getApprovalMode();
     }
     setApprovalMode(mode) {
         if (!this.isTrustedFolder() && mode !== ApprovalMode.DEFAULT) {
             throw new Error('Cannot enable privileged approval modes in an untrusted folder.');
         }
-        this.approvalMode = mode;
+        this.policyEngine.setApprovalMode(mode);
     }
     isYoloModeDisabled() {
         return this.disableYoloMode || !this.isTrustedFolder();
@@ -723,6 +859,9 @@ export class Config {
     getTelemetryUseCollector() {
         return this.telemetrySettings.useCollector ?? false;
     }
+    getTelemetryUseCliAuth() {
+        return this.telemetrySettings.useCliAuth ?? false;
+    }
     getGeminiClient() {
         return this.geminiClient;
     }
@@ -739,6 +878,9 @@ export class Config {
     getModelRouterService() {
         return this.modelRouterService;
     }
+    getModelAvailabilityService() {
+        return this.modelAvailabilityService;
+    }
     getEnableRecursiveFileSearch() {
         return this.fileFiltering.enableRecursiveFileSearch;
     }
@@ -836,6 +978,9 @@ export class Config {
     isModelAvailabilityServiceEnabled() {
         return this.enableModelAvailabilityService;
     }
+    isAgentsEnabled() {
+        return this.enableAgents;
+    }
     getNoBrowser() {
         return this.noBrowser;
     }
@@ -953,6 +1098,9 @@ export class Config {
     getEnableShellOutputEfficiency() {
         return this.enableShellOutputEfficiency;
     }
+    getShellToolInactivityTimeout() {
+        return this.shellToolInactivityTimeout;
+    }
     getShellExecutionConfig() {
         return this.shellExecutionConfig;
     }
@@ -1088,17 +1236,16 @@ export class Config {
             registerCoreTool(WriteTodosTool, this);
         }
         // Register Subagents as Tools
-        if (this.getCodebaseInvestigatorSettings().enabled) {
-            const definition = this.agentRegistry.getDefinition('codebase_investigator');
-            if (definition) {
-                // We must respect the main allowed/exclude lists for agents too.
-                const allowedTools = this.getAllowedTools();
-                const isAllowed = !allowedTools || allowedTools.includes(definition.name);
-                if (isAllowed) {
-                    const messageBusEnabled = this.getEnableMessageBusIntegration();
-                    const wrapper = new SubagentToolWrapper(definition, this, messageBusEnabled ? this.getMessageBus() : undefined);
-                    registry.registerTool(wrapper);
-                }
+        // Register DelegateToAgentTool if agents are enabled
+        if (this.isAgentsEnabled() ||
+            this.getCodebaseInvestigatorSettings().enabled) {
+            // Check if the delegate tool itself is allowed (if allowedTools is set)
+            const allowedTools = this.getAllowedTools();
+            const isAllowed = !allowedTools || allowedTools.includes(DELEGATE_TO_AGENT_TOOL_NAME);
+            if (isAllowed) {
+                const messageBusEnabled = this.getEnableMessageBusIntegration();
+                const delegateTool = new DelegateToAgentTool(this.agentRegistry, this, messageBusEnabled ? this.getMessageBus() : undefined);
+                registry.registerTool(delegateTool);
             }
         }
         await registry.discoverAllTools();
@@ -1117,6 +1264,12 @@ export class Config {
     getHooks() {
         return this.hooks;
     }
+    /**
+     * Get disabled hooks list
+     */
+    getDisabledHooks() {
+        return this.disabledHooks;
+    }
     /**
      * Get experiments configuration
      */