(
- (resolve, reject) => {
- let serverPort: number;
-
- const server = http.createServer(
- async (req: http.IncomingMessage, res: http.ServerResponse) => {
- try {
- const url = new URL(req.url!, `http://localhost:${serverPort}`);
-
- if (url.pathname !== REDIRECT_PATH) {
- res.writeHead(404);
- res.end('Not found');
- return;
- }
-
- const code = url.searchParams.get('code');
- const state = url.searchParams.get('state');
- const error = url.searchParams.get('error');
-
- if (error) {
- res.writeHead(HTTP_OK, { 'Content-Type': 'text/html' });
- res.end(`
-
-
- Authentication Failed
- Error: ${error.replace(//g, '>')}
- ${(url.searchParams.get('error_description') || '').replace(//g, '>')}
- You can close this window.
-
-
- `);
- server.close();
- reject(new Error(`OAuth error: ${error}`));
- return;
- }
-
- if (!code || !state) {
- res.writeHead(400);
- res.end('Missing code or state parameter');
- return;
- }
-
- if (state !== expectedState) {
- res.writeHead(400);
- res.end('Invalid state parameter');
- server.close();
- reject(new Error('State mismatch - possible CSRF attack'));
- return;
- }
-
- // Send success response to browser
- res.writeHead(HTTP_OK, { 'Content-Type': 'text/html' });
- res.end(`
-
-
- Authentication Successful!
- You can close this window and return to Gemini CLI.
-
-
-
- `);
-
- server.close();
- resolve({ code, state });
- } catch (error) {
- server.close();
- reject(error);
- }
- },
- );
-
- server.on('error', (error) => {
- portReject(error);
- reject(error);
- });
-
- // Determine which port to use (env var, argument, or OS-assigned)
- let listenPort = 0; // Default to OS-assigned port
-
- const portStr = process.env['OAUTH_CALLBACK_PORT'];
- if (portStr) {
- const envPort = parseInt(portStr, 10);
- if (isNaN(envPort) || envPort <= 0 || envPort > 65535) {
- const error = new Error(
- `Invalid value for OAUTH_CALLBACK_PORT: "${portStr}"`,
- );
- portReject(error);
- reject(error);
- return;
- }
- listenPort = envPort;
- } else if (port !== undefined) {
- listenPort = port;
- }
-
- server.listen(listenPort, () => {
- // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- const address = server.address() as net.AddressInfo;
- serverPort = address.port;
- debugLogger.log(
- `OAuth callback server listening on port ${serverPort}`,
- );
- portResolve(serverPort); // Resolve port promise immediately
- });
-
- // Timeout after 5 minutes
- setTimeout(
- () => {
- server.close();
- reject(new Error('OAuth callback timeout'));
- },
- 5 * 60 * 1000,
- );
- },
- );
-
- return { port: portPromise, response: responsePromise };
- }
-
- /**
- * Extract the port number from a URL string if available and valid.
- *
- * @param urlString The URL string to parse
- * @returns The port number or undefined if not found or invalid
- */
- private getPortFromUrl(urlString?: string): number | undefined {
- if (!urlString) {
- return undefined;
- }
-
+ private buildResourceParam(mcpServerUrl?: string): string | undefined {
+ if (!mcpServerUrl) return undefined;
try {
- const url = new URL(urlString);
- if (url.port) {
- const parsedPort = parseInt(url.port, 10);
- if (!isNaN(parsedPort) && parsedPort > 0 && parsedPort <= 65535) {
- return parsedPort;
- }
- }
- } catch {
- // Ignore invalid URL
- }
-
- return undefined;
- }
-
- /**
- * Build the authorization URL for the OAuth flow.
-
- *
- * @param config OAuth configuration
- * @param pkceParams PKCE parameters
- * @param redirectPort The port to use for the redirect URI
- * @param mcpServerUrl The MCP server URL to use as the resource parameter
- * @returns The authorization URL
- */
- private buildAuthorizationUrl(
- config: MCPOAuthConfig,
- pkceParams: PKCEParams,
- redirectPort: number,
- mcpServerUrl?: string,
- ): string {
- const redirectUri =
- config.redirectUri || `http://localhost:${redirectPort}${REDIRECT_PATH}`;
-
- const params = new URLSearchParams({
- client_id: config.clientId!,
- response_type: 'code',
- redirect_uri: redirectUri,
- state: pkceParams.state,
- code_challenge: pkceParams.codeChallenge,
- code_challenge_method: 'S256',
- });
-
- if (config.scopes && config.scopes.length > 0) {
- params.append('scope', config.scopes.join(' '));
- }
-
- if (config.audiences && config.audiences.length > 0) {
- params.append('audience', config.audiences.join(' '));
- }
-
- // Add resource parameter for MCP OAuth spec compliance
- // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth)
- if (mcpServerUrl) {
- try {
- params.append(
- 'resource',
- OAuthUtils.buildResourceParameter(mcpServerUrl),
- );
- } catch (error) {
- debugLogger.warn(
- `Could not add resource parameter: ${getErrorMessage(error)}`,
- );
- }
- }
-
- const url = new URL(config.authorizationUrl!);
- params.forEach((value, key) => {
- url.searchParams.append(key, value);
- });
- return url.toString();
- }
-
- /**
- * Exchange authorization code for tokens.
- *
- * @param config OAuth configuration
- * @param code Authorization code
- * @param codeVerifier PKCE code verifier
- * @param redirectPort The port to use for the redirect URI
- * @param mcpServerUrl The MCP server URL to use as the resource parameter
- * @returns The token response
- */
- private async exchangeCodeForToken(
- config: MCPOAuthConfig,
- code: string,
- codeVerifier: string,
- redirectPort: number,
- mcpServerUrl?: string,
- ): Promise {
- const redirectUri =
- config.redirectUri || `http://localhost:${redirectPort}${REDIRECT_PATH}`;
-
- const params = new URLSearchParams({
- grant_type: 'authorization_code',
- code,
- redirect_uri: redirectUri,
- code_verifier: codeVerifier,
- client_id: config.clientId!,
- });
-
- if (config.clientSecret) {
- params.append('client_secret', config.clientSecret);
- }
-
- if (config.audiences && config.audiences.length > 0) {
- params.append('audience', config.audiences.join(' '));
- }
-
- // Add resource parameter for MCP OAuth spec compliance
- // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth)
- if (mcpServerUrl) {
- const resourceUrl = mcpServerUrl;
- try {
- params.append(
- 'resource',
- OAuthUtils.buildResourceParameter(resourceUrl),
- );
- } catch (error) {
- debugLogger.warn(
- `Could not add resource parameter: ${getErrorMessage(error)}`,
- );
- }
- }
-
- const response = await fetch(config.tokenUrl!, {
- method: 'POST',
- headers: {
- 'Content-Type': 'application/x-www-form-urlencoded',
- Accept: 'application/json, application/x-www-form-urlencoded',
- },
- body: params.toString(),
- });
-
- const responseText = await response.text();
- const contentType = response.headers.get('content-type') || '';
-
- if (!response.ok) {
- // Try to parse error from form-urlencoded response
- let errorMessage: string | null = null;
- try {
- const errorParams = new URLSearchParams(responseText);
- const error = errorParams.get('error');
- const errorDescription = errorParams.get('error_description');
- if (error) {
- errorMessage = `Token exchange failed: ${error} - ${errorDescription || 'No description'}`;
- }
- } catch {
- // Fall back to raw error
- }
- throw new Error(
- errorMessage ||
- `Token exchange failed: ${response.status} - ${responseText}`,
- );
- }
-
- // Log unexpected content types for debugging
- if (
- !contentType.includes('application/json') &&
- !contentType.includes('application/x-www-form-urlencoded')
- ) {
+ return OAuthUtils.buildResourceParameter(mcpServerUrl);
+ } catch (error) {
debugLogger.warn(
- `Token endpoint returned unexpected content-type: ${contentType}. ` +
- `Expected application/json or application/x-www-form-urlencoded. ` +
- `Will attempt to parse response.`,
+ `Could not add resource parameter: ${getErrorMessage(error)}`,
);
- }
-
- // Try to parse as JSON first, fall back to form-urlencoded
- try {
- // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- return JSON.parse(responseText) as OAuthTokenResponse;
- } catch {
- // Parse form-urlencoded response
- const tokenParams = new URLSearchParams(responseText);
- const accessToken = tokenParams.get('access_token');
- const tokenType = tokenParams.get('token_type') || 'Bearer';
- const expiresIn = tokenParams.get('expires_in');
- const refreshToken = tokenParams.get('refresh_token');
- const scope = tokenParams.get('scope');
-
- if (!accessToken) {
- // Check for error in response
- const error = tokenParams.get('error');
- const errorDescription = tokenParams.get('error_description');
- throw new Error(
- `Token exchange failed: ${error || 'no_access_token'} - ${errorDescription || responseText}`,
- );
- }
-
- return {
- access_token: accessToken,
- token_type: tokenType,
- expires_in: expiresIn ? parseInt(expiresIn, 10) : undefined,
- refresh_token: refreshToken || undefined,
- scope: scope || undefined,
- } as OAuthTokenResponse;
+ return undefined;
}
}
@@ -626,112 +253,21 @@ export class MCPOAuthProvider {
tokenUrl: string,
mcpServerUrl?: string,
): Promise {
- const params = new URLSearchParams({
- grant_type: 'refresh_token',
- refresh_token: refreshToken,
- client_id: config.clientId!,
- });
-
- if (config.clientSecret) {
- params.append('client_secret', config.clientSecret);
+ if (!config.clientId) {
+ throw new Error('Missing required clientId for token refresh');
}
- if (config.scopes && config.scopes.length > 0) {
- params.append('scope', config.scopes.join(' '));
- }
-
- if (config.audiences && config.audiences.length > 0) {
- params.append('audience', config.audiences.join(' '));
- }
-
- // Add resource parameter for MCP OAuth spec compliance
- // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth)
- if (mcpServerUrl) {
- try {
- params.append(
- 'resource',
- OAuthUtils.buildResourceParameter(mcpServerUrl),
- );
- } catch (error) {
- debugLogger.warn(
- `Could not add resource parameter: ${getErrorMessage(error)}`,
- );
- }
- }
-
- const response = await fetch(tokenUrl, {
- method: 'POST',
- headers: {
- 'Content-Type': 'application/x-www-form-urlencoded',
- Accept: 'application/json, application/x-www-form-urlencoded',
+ return refreshAccessTokenShared(
+ {
+ clientId: config.clientId,
+ clientSecret: config.clientSecret,
+ scopes: config.scopes,
+ audiences: config.audiences,
},
- body: params.toString(),
- });
-
- const responseText = await response.text();
- const contentType = response.headers.get('content-type') || '';
-
- if (!response.ok) {
- // Try to parse error from form-urlencoded response
- let errorMessage: string | null = null;
- try {
- const errorParams = new URLSearchParams(responseText);
- const error = errorParams.get('error');
- const errorDescription = errorParams.get('error_description');
- if (error) {
- errorMessage = `Token refresh failed: ${error} - ${errorDescription || 'No description'}`;
- }
- } catch {
- // Fall back to raw error
- }
- throw new Error(
- errorMessage ||
- `Token refresh failed: ${response.status} - ${responseText}`,
- );
- }
-
- // Log unexpected content types for debugging
- if (
- !contentType.includes('application/json') &&
- !contentType.includes('application/x-www-form-urlencoded')
- ) {
- debugLogger.warn(
- `Token refresh endpoint returned unexpected content-type: ${contentType}. ` +
- `Expected application/json or application/x-www-form-urlencoded. ` +
- `Will attempt to parse response.`,
- );
- }
-
- // Try to parse as JSON first, fall back to form-urlencoded
- try {
- // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- return JSON.parse(responseText) as OAuthTokenResponse;
- } catch {
- // Parse form-urlencoded response
- const tokenParams = new URLSearchParams(responseText);
- const accessToken = tokenParams.get('access_token');
- const tokenType = tokenParams.get('token_type') || 'Bearer';
- const expiresIn = tokenParams.get('expires_in');
- const refreshToken = tokenParams.get('refresh_token');
- const scope = tokenParams.get('scope');
-
- if (!accessToken) {
- // Check for error in response
- const error = tokenParams.get('error');
- const errorDescription = tokenParams.get('error_description');
- throw new Error(
- `Token refresh failed: ${error || 'unknown_error'} - ${errorDescription || responseText}`,
- );
- }
-
- return {
- access_token: accessToken,
- token_type: tokenType,
- expires_in: expiresIn ? parseInt(expiresIn, 10) : undefined,
- refresh_token: refreshToken || undefined,
- scope: scope || undefined,
- } as OAuthTokenResponse;
- }
+ refreshToken,
+ tokenUrl,
+ this.buildResourceParam(mcpServerUrl),
+ );
}
/**
@@ -830,17 +366,14 @@ export class MCPOAuthProvider {
}
// Generate PKCE parameters
- const pkceParams = this.generatePKCEParams();
+ const pkceParams = generatePKCEParams();
// Determine preferred port from redirectUri if available
- const preferredPort = this.getPortFromUrl(config.redirectUri);
+ const preferredPort = getPortFromUrl(config.redirectUri);
// Start callback server first to allocate port
// This ensures we only create one server and eliminates race conditions
- const callbackServer = this.startCallbackServer(
- pkceParams.state,
- preferredPort,
- );
+ const callbackServer = startCallbackServer(pkceParams.state, preferredPort);
// Wait for server to start and get the allocated port
// We need this port for client registration and auth URL building
@@ -892,12 +425,24 @@ export class MCPOAuthProvider {
);
}
+ // Build flow config for shared utilities
+ const flowConfig: OAuthFlowConfig = {
+ clientId: config.clientId,
+ clientSecret: config.clientSecret,
+ authorizationUrl: config.authorizationUrl,
+ tokenUrl: config.tokenUrl,
+ scopes: config.scopes,
+ audiences: config.audiences,
+ redirectUri: config.redirectUri,
+ };
+
// Build authorization URL
- const authUrl = this.buildAuthorizationUrl(
- config,
+ const resource = this.buildResourceParam(mcpServerUrl);
+ const authUrl = buildAuthorizationUrl(
+ flowConfig,
pkceParams,
redirectPort,
- mcpServerUrl,
+ resource,
);
const userConsent = await getConsentForOauth(
@@ -933,12 +478,12 @@ ${authUrl}
);
// Exchange code for tokens
- const tokenResponse = await this.exchangeCodeForToken(
- config,
+ const tokenResponse = await exchangeCodeForToken(
+ flowConfig,
code,
pkceParams.codeVerifier,
redirectPort,
- mcpServerUrl,
+ resource,
);
// Convert to our token format
diff --git a/packages/core/src/mcp/oauth-utils.test.ts b/packages/core/src/mcp/oauth-utils.test.ts
index c318261aef..f27ee7727b 100644
--- a/packages/core/src/mcp/oauth-utils.test.ts
+++ b/packages/core/src/mcp/oauth-utils.test.ts
@@ -5,11 +5,11 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import type {
- OAuthAuthorizationServerMetadata,
- OAuthProtectedResourceMetadata,
+import {
+ OAuthUtils,
+ type OAuthAuthorizationServerMetadata,
+ type OAuthProtectedResourceMetadata,
} from './oauth-utils.js';
-import { OAuthUtils } from './oauth-utils.js';
// Mock fetch globally
const mockFetch = vi.fn();
diff --git a/packages/core/src/mcp/token-storage/hybrid-token-storage.ts b/packages/core/src/mcp/token-storage/hybrid-token-storage.ts
index 3bda6050e6..20560ba30e 100644
--- a/packages/core/src/mcp/token-storage/hybrid-token-storage.ts
+++ b/packages/core/src/mcp/token-storage/hybrid-token-storage.ts
@@ -6,8 +6,11 @@
import { BaseTokenStorage } from './base-token-storage.js';
import { FileTokenStorage } from './file-token-storage.js';
-import type { TokenStorage, OAuthCredentials } from './types.js';
-import { TokenStorageType } from './types.js';
+import {
+ TokenStorageType,
+ type TokenStorage,
+ type OAuthCredentials,
+} from './types.js';
import { coreEvents } from '../../utils/events.js';
import { TokenStorageInitializationEvent } from '../../telemetry/types.js';
diff --git a/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts b/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts
index 8b402ff7cd..2192abbc45 100644
--- a/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts
+++ b/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts
@@ -5,54 +5,41 @@
*/
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
-import type { KeychainTokenStorage } from './keychain-token-storage.js';
+import { KeychainTokenStorage } from './keychain-token-storage.js';
import type { OAuthCredentials } from './types.js';
+import { KeychainService } from '../../services/keychainService.js';
import { coreEvents } from '../../utils/events.js';
-
-// Hoist the mock to be available in the vi.mock factory
-const mockKeytar = vi.hoisted(() => ({
- getPassword: vi.fn(),
- setPassword: vi.fn(),
- deletePassword: vi.fn(),
- findCredentials: vi.fn(),
-}));
-
-const mockServiceName = 'service-name';
-const mockCryptoRandomBytesString = 'random-string';
-
-// Mock the dynamic import of 'keytar'
-vi.mock('keytar', () => ({
- default: mockKeytar,
-}));
-
-vi.mock('node:crypto', async (importOriginal) => {
- const actual = await importOriginal();
- return {
- ...actual,
- randomBytes: vi.fn(() => ({
- toString: vi.fn(() => mockCryptoRandomBytesString),
- })),
- };
-});
-
-vi.mock('../../utils/events.js', () => ({
- coreEvents: {
- emitFeedback: vi.fn(),
- emitTelemetryKeychainAvailability: vi.fn(),
- },
-}));
+import { KEYCHAIN_TEST_PREFIX } from '../../services/keychainTypes.js';
describe('KeychainTokenStorage', () => {
let storage: KeychainTokenStorage;
+ const mockServiceName = 'service-name';
+ let storageState: Map;
- beforeEach(async () => {
- vi.resetAllMocks();
- // Reset the internal state of the keychain-token-storage module
- vi.resetModules();
- const { KeychainTokenStorage } = await import(
- './keychain-token-storage.js'
- );
+ beforeEach(() => {
+ vi.clearAllMocks();
storage = new KeychainTokenStorage(mockServiceName);
+ storageState = new Map();
+
+ // Use stateful spies to verify logic behaviorally
+ vi.spyOn(KeychainService.prototype, 'getPassword').mockImplementation(
+ async (account) => storageState.get(account) ?? null,
+ );
+ vi.spyOn(KeychainService.prototype, 'setPassword').mockImplementation(
+ async (account, value) => {
+ storageState.set(account, value);
+ },
+ );
+ vi.spyOn(KeychainService.prototype, 'deletePassword').mockImplementation(
+ async (account) => storageState.delete(account),
+ );
+ vi.spyOn(KeychainService.prototype, 'findCredentials').mockImplementation(
+ async () =>
+ Array.from(storageState.entries()).map(([account, password]) => ({
+ account,
+ password,
+ })),
+ );
});
afterEach(() => {
@@ -70,375 +57,149 @@ describe('KeychainTokenStorage', () => {
updatedAt: Date.now(),
} as OAuthCredentials;
- describe('checkKeychainAvailability', () => {
- it('should return true if keytar is available and functional', async () => {
- mockKeytar.setPassword.mockResolvedValue(undefined);
- mockKeytar.getPassword.mockResolvedValue('test');
- mockKeytar.deletePassword.mockResolvedValue(true);
-
- const isAvailable = await storage.checkKeychainAvailability();
- expect(isAvailable).toBe(true);
- expect(mockKeytar.setPassword).toHaveBeenCalledWith(
- mockServiceName,
- `__keychain_test__${mockCryptoRandomBytesString}`,
- 'test',
- );
- expect(mockKeytar.getPassword).toHaveBeenCalledWith(
- mockServiceName,
- `__keychain_test__${mockCryptoRandomBytesString}`,
- );
- expect(mockKeytar.deletePassword).toHaveBeenCalledWith(
- mockServiceName,
- `__keychain_test__${mockCryptoRandomBytesString}`,
- );
- });
-
- it('should return false if keytar fails to set password', async () => {
- const error = new Error('write error');
- mockKeytar.setPassword.mockRejectedValue(error);
- const isAvailable = await storage.checkKeychainAvailability();
- expect(isAvailable).toBe(false);
- });
-
- it('should return false if retrieved password does not match', async () => {
- mockKeytar.setPassword.mockResolvedValue(undefined);
- mockKeytar.getPassword.mockResolvedValue('wrong-password');
- mockKeytar.deletePassword.mockResolvedValue(true);
- const isAvailable = await storage.checkKeychainAvailability();
- expect(isAvailable).toBe(false);
- });
-
- it('should cache the availability result', async () => {
- mockKeytar.setPassword.mockResolvedValue(undefined);
- mockKeytar.getPassword.mockResolvedValue('test');
- mockKeytar.deletePassword.mockResolvedValue(true);
-
- await storage.checkKeychainAvailability();
- await storage.checkKeychainAvailability();
-
- expect(mockKeytar.setPassword).toHaveBeenCalledTimes(1);
- });
- });
-
- describe('with keychain unavailable', () => {
- beforeEach(async () => {
- // Force keychain to be unavailable
- mockKeytar.setPassword.mockRejectedValue(new Error('keychain error'));
- await storage.checkKeychainAvailability();
- });
-
- it('getCredentials should throw', async () => {
- await expect(storage.getCredentials('server')).rejects.toThrow(
- 'Keychain is not available',
- );
- });
-
- it('setCredentials should throw', async () => {
- await expect(storage.setCredentials(validCredentials)).rejects.toThrow(
- 'Keychain is not available',
- );
- });
-
- it('deleteCredentials should throw', async () => {
- await expect(storage.deleteCredentials('server')).rejects.toThrow(
- 'Keychain is not available',
- );
- });
-
- it('listServers should throw', async () => {
- await expect(storage.listServers()).rejects.toThrow(
- 'Keychain is not available',
- );
- });
-
- it('getAllCredentials should throw', async () => {
- await expect(storage.getAllCredentials()).rejects.toThrow(
- 'Keychain is not available',
- );
- });
- });
-
describe('with keychain available', () => {
- beforeEach(async () => {
- mockKeytar.setPassword.mockResolvedValue(undefined);
- mockKeytar.getPassword.mockResolvedValue('test');
- mockKeytar.deletePassword.mockResolvedValue(true);
- await storage.checkKeychainAvailability();
- // Reset mocks after availability check
- vi.resetAllMocks();
+ beforeEach(() => {
+ vi.spyOn(KeychainService.prototype, 'isAvailable').mockResolvedValue(
+ true,
+ );
});
- describe('getCredentials', () => {
- it('should return null if no credentials are found', async () => {
- mockKeytar.getPassword.mockResolvedValue(null);
- const result = await storage.getCredentials('test-server');
- expect(result).toBeNull();
- expect(mockKeytar.getPassword).toHaveBeenCalledWith(
- mockServiceName,
- 'test-server',
- );
- });
+ it('should store and retrieve credentials correctly', async () => {
+ await storage.setCredentials(validCredentials);
+ const retrieved = await storage.getCredentials('test-server');
- it('should return credentials if found and not expired', async () => {
- mockKeytar.getPassword.mockResolvedValue(
- JSON.stringify(validCredentials),
- );
- const result = await storage.getCredentials('test-server');
- expect(result).toEqual(validCredentials);
- });
-
- it('should return null if credentials have expired', async () => {
- const expiredCreds = {
- ...validCredentials,
- token: { ...validCredentials.token, expiresAt: Date.now() - 1000 },
- };
- mockKeytar.getPassword.mockResolvedValue(JSON.stringify(expiredCreds));
- const result = await storage.getCredentials('test-server');
- expect(result).toBeNull();
- });
-
- it('should throw if stored data is corrupted JSON', async () => {
- mockKeytar.getPassword.mockResolvedValue('not-json');
- await expect(storage.getCredentials('test-server')).rejects.toThrow(
- 'Failed to parse stored credentials for test-server',
- );
- });
+ expect(retrieved?.token.accessToken).toBe('access-token');
+ expect(retrieved?.serverName).toBe('test-server');
});
- describe('setCredentials', () => {
- it('should save credentials to keychain', async () => {
- vi.useFakeTimers();
- mockKeytar.setPassword.mockResolvedValue(undefined);
- await storage.setCredentials(validCredentials);
- expect(mockKeytar.setPassword).toHaveBeenCalledWith(
- mockServiceName,
- 'test-server',
- JSON.stringify({ ...validCredentials, updatedAt: Date.now() }),
- );
- });
+ it('should return null if no credentials are found or they are expired', async () => {
+ expect(await storage.getCredentials('missing')).toBeNull();
- it('should throw if saving to keychain fails', async () => {
- mockKeytar.setPassword.mockRejectedValue(
- new Error('keychain write error'),
- );
- await expect(storage.setCredentials(validCredentials)).rejects.toThrow(
- 'keychain write error',
- );
- });
+ const expiredCreds = {
+ ...validCredentials,
+ token: { ...validCredentials.token, expiresAt: Date.now() - 1000 },
+ };
+ await storage.setCredentials(expiredCreds);
+ expect(await storage.getCredentials('test-server')).toBeNull();
});
- describe('deleteCredentials', () => {
- it('should delete credentials from keychain', async () => {
- mockKeytar.deletePassword.mockResolvedValue(true);
- await storage.deleteCredentials('test-server');
- expect(mockKeytar.deletePassword).toHaveBeenCalledWith(
- mockServiceName,
- 'test-server',
- );
- });
-
- it('should throw if no credentials were found to delete', async () => {
- mockKeytar.deletePassword.mockResolvedValue(false);
- await expect(storage.deleteCredentials('test-server')).rejects.toThrow(
- 'No credentials found for test-server',
- );
- });
-
- it('should throw if deleting from keychain fails', async () => {
- mockKeytar.deletePassword.mockRejectedValue(
- new Error('keychain delete error'),
- );
- await expect(storage.deleteCredentials('test-server')).rejects.toThrow(
- 'keychain delete error',
- );
- });
+ it('should throw if stored data is corrupted JSON', async () => {
+ storageState.set('bad-server', 'not-json');
+ await expect(storage.getCredentials('bad-server')).rejects.toThrow(
+ /Failed to parse/,
+ );
});
- describe('listServers', () => {
- it('should return a list of server names', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: 'server1', password: '' },
- { account: 'server2', password: '' },
- ]);
- const result = await storage.listServers();
- expect(result).toEqual(['server1', 'server2']);
+ it('should list servers and filter internal keys', async () => {
+ await storage.setCredentials(validCredentials);
+ await storage.setCredentials({
+ ...validCredentials,
+ serverName: 'server2',
});
+ storageState.set(`${KEYCHAIN_TEST_PREFIX}internal`, '...');
+ storageState.set('__secret__key', '...');
- it('should not include internal test keys in the server list', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: 'server1', password: '' },
- {
- account: `__keychain_test__${mockCryptoRandomBytesString}`,
- password: '',
- },
- { account: 'server2', password: '' },
- ]);
- const result = await storage.listServers();
- expect(result).toEqual(['server1', 'server2']);
- });
-
- it('should return an empty array on error', async () => {
- const error = new Error('find error');
- mockKeytar.findCredentials.mockRejectedValue(error);
- const result = await storage.listServers();
- expect(result).toEqual([]);
- expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
- 'error',
- 'Failed to list servers from keychain',
- error,
- );
- });
+ const servers = await storage.listServers();
+ expect(servers).toEqual(['test-server', 'server2']);
});
- describe('getAllCredentials', () => {
- it('should return a map of all valid credentials and emit feedback for invalid ones', async () => {
- const creds2 = {
- ...validCredentials,
- serverName: 'server2',
- };
- const expiredCreds = {
- ...validCredentials,
- serverName: 'expired-server',
- token: { ...validCredentials.token, expiresAt: Date.now() - 1000 },
- };
- const structurallyInvalidCreds = {
- serverName: 'invalid-server',
- };
+ it('should handle getAllCredentials with individual parse errors', async () => {
+ await storage.setCredentials(validCredentials);
+ storageState.set('bad', 'not-json');
+ const emitFeedbackSpy = vi.spyOn(coreEvents, 'emitFeedback');
- mockKeytar.findCredentials.mockResolvedValue([
- {
- account: 'test-server',
- password: JSON.stringify(validCredentials),
- },
- { account: 'server2', password: JSON.stringify(creds2) },
- {
- account: 'expired-server',
- password: JSON.stringify(expiredCreds),
- },
- { account: 'bad-server', password: 'not-json' },
- {
- account: 'invalid-server',
- password: JSON.stringify(structurallyInvalidCreds),
- },
- ]);
-
- const result = await storage.getAllCredentials();
- expect(result.size).toBe(2);
- expect(result.get('test-server')).toEqual(validCredentials);
- expect(result.get('server2')).toEqual(creds2);
- expect(result.has('expired-server')).toBe(false);
- expect(result.has('bad-server')).toBe(false);
- expect(result.has('invalid-server')).toBe(false);
-
- expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
- 'error',
- 'Failed to parse credentials for bad-server',
- expect.any(SyntaxError),
- );
- expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
- 'error',
- 'Failed to parse credentials for invalid-server',
- expect.any(Error),
- );
- });
-
- it('should emit feedback and return empty map if findCredentials fails', async () => {
- const error = new Error('find all error');
- mockKeytar.findCredentials.mockRejectedValue(error);
-
- const result = await storage.getAllCredentials();
- expect(result.size).toBe(0);
- expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
- 'error',
- 'Failed to get all credentials from keychain',
- error,
- );
- });
+ const result = await storage.getAllCredentials();
+ expect(result.size).toBe(1);
+ expect(emitFeedbackSpy).toHaveBeenCalled();
});
- describe('clearAll', () => {
- it('should delete all credentials for the service', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: 'server1', password: '' },
- { account: 'server2', password: '' },
- ]);
- mockKeytar.deletePassword.mockResolvedValue(true);
+ it('should aggregate errors in clearAll', async () => {
+ storageState.set('s1', '...');
+ storageState.set('s2', '...');
- await storage.clearAll();
+ // Aggregating a system error (rejection)
+ vi.spyOn(KeychainService.prototype, 'deletePassword')
+ .mockResolvedValueOnce(true)
+ .mockRejectedValueOnce(new Error('system fail'));
- expect(mockKeytar.deletePassword).toHaveBeenCalledTimes(2);
- expect(mockKeytar.deletePassword).toHaveBeenCalledWith(
- mockServiceName,
- 'server1',
- );
- expect(mockKeytar.deletePassword).toHaveBeenCalledWith(
- mockServiceName,
- 'server2',
- );
- });
+ await expect(storage.clearAll()).rejects.toThrow(
+ /Failed to clear some credentials: system fail/,
+ );
- it('should throw an aggregated error if deletions fail', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: 'server1', password: '' },
- { account: 'server2', password: '' },
- ]);
- mockKeytar.deletePassword
- .mockResolvedValueOnce(true)
- .mockRejectedValueOnce(new Error('delete failed'));
+ // Aggregating a 'not found' error (returns false)
+ vi.spyOn(KeychainService.prototype, 'deletePassword')
+ .mockResolvedValueOnce(true)
+ .mockResolvedValueOnce(false);
- await expect(storage.clearAll()).rejects.toThrow(
- 'Failed to clear some credentials: delete failed',
- );
- });
+ await expect(storage.clearAll()).rejects.toThrow(
+ /Failed to clear some credentials: No credentials found/,
+ );
});
- describe('Secrets', () => {
- it('should set and get a secret', async () => {
- mockKeytar.setPassword.mockResolvedValue(undefined);
- mockKeytar.getPassword.mockResolvedValue('secret-value');
+ it('should manage secrets with prefix independently', async () => {
+ await storage.setSecret('key1', 'val1');
+ await storage.setCredentials(validCredentials);
- await storage.setSecret('secret-key', 'secret-value');
- const value = await storage.getSecret('secret-key');
+ expect(await storage.getSecret('key1')).toBe('val1');
+ expect(await storage.listSecrets()).toEqual(['key1']);
+ expect(await storage.listServers()).not.toContain('key1');
+ });
+ });
- expect(mockKeytar.setPassword).toHaveBeenCalledWith(
- mockServiceName,
- '__secret__secret-key',
- 'secret-value',
- );
- expect(mockKeytar.getPassword).toHaveBeenCalledWith(
- mockServiceName,
- '__secret__secret-key',
- );
- expect(value).toBe('secret-value');
- });
+ describe('unavailability handling', () => {
+ beforeEach(() => {
+ vi.spyOn(KeychainService.prototype, 'isAvailable').mockResolvedValue(
+ false,
+ );
+ vi.spyOn(KeychainService.prototype, 'getPassword').mockRejectedValue(
+ new Error('Keychain is not available'),
+ );
+ vi.spyOn(KeychainService.prototype, 'setPassword').mockRejectedValue(
+ new Error('Keychain is not available'),
+ );
+ vi.spyOn(KeychainService.prototype, 'deletePassword').mockRejectedValue(
+ new Error('Keychain is not available'),
+ );
+ vi.spyOn(KeychainService.prototype, 'findCredentials').mockRejectedValue(
+ new Error('Keychain is not available'),
+ );
+ });
- it('should delete a secret', async () => {
- mockKeytar.deletePassword.mockResolvedValue(true);
- await storage.deleteSecret('secret-key');
- expect(mockKeytar.deletePassword).toHaveBeenCalledWith(
- mockServiceName,
- '__secret__secret-key',
- );
- });
+ it.each([
+ { method: 'getCredentials', args: ['s'] },
+ { method: 'setCredentials', args: [validCredentials] },
+ { method: 'deleteCredentials', args: ['s'] },
+ { method: 'clearAll', args: [] },
+ ])(
+ '$method should propagate unavailability error',
+ async ({ method, args }) => {
+ await expect(
+ (
+ storage as unknown as Record<
+ string,
+ (...args: unknown[]) => Promise
+ >
+ )[method](...args),
+ ).rejects.toThrow('Keychain is not available');
+ },
+ );
- it('should list secrets', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: '__secret__secret1', password: '' },
- { account: '__secret__secret2', password: '' },
- { account: 'server1', password: '' },
- ]);
- const secrets = await storage.listSecrets();
- expect(secrets).toEqual(['secret1', 'secret2']);
- });
-
- it('should not list secrets in listServers', async () => {
- mockKeytar.findCredentials.mockResolvedValue([
- { account: '__secret__secret1', password: '' },
- { account: 'server1', password: '' },
- ]);
- const servers = await storage.listServers();
- expect(servers).toEqual(['server1']);
- });
+ it.each([
+ { method: 'listServers' },
+ { method: 'getAllCredentials' },
+ { method: 'listSecrets' },
+ ])('$method should emit feedback and return empty', async ({ method }) => {
+ const emitFeedbackSpy = vi.spyOn(coreEvents, 'emitFeedback');
+ expect(
+ await (storage as unknown as Record Promise>)[
+ method
+ ](),
+ ).toEqual(method === 'getAllCredentials' ? new Map() : []);
+ expect(emitFeedbackSpy).toHaveBeenCalledWith(
+ 'error',
+ expect.any(String),
+ expect.any(Error),
+ );
});
});
});
diff --git a/packages/core/src/mcp/token-storage/keychain-token-storage.ts b/packages/core/src/mcp/token-storage/keychain-token-storage.ts
index 4be0d082e5..d0b4990279 100644
--- a/packages/core/src/mcp/token-storage/keychain-token-storage.ts
+++ b/packages/core/src/mcp/token-storage/keychain-token-storage.ts
@@ -4,70 +4,30 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import * as crypto from 'node:crypto';
import { BaseTokenStorage } from './base-token-storage.js';
import type { OAuthCredentials, SecretStorage } from './types.js';
import { coreEvents } from '../../utils/events.js';
-import { KeychainAvailabilityEvent } from '../../telemetry/types.js';
-
-interface Keytar {
- getPassword(service: string, account: string): Promise;
- setPassword(
- service: string,
- account: string,
- password: string,
- ): Promise;
- deletePassword(service: string, account: string): Promise;
- findCredentials(
- service: string,
- ): Promise>;
-}
-
-const KEYCHAIN_TEST_PREFIX = '__keychain_test__';
-const SECRET_PREFIX = '__secret__';
+import { KeychainService } from '../../services/keychainService.js';
+import {
+ KEYCHAIN_TEST_PREFIX,
+ SECRET_PREFIX,
+} from '../../services/keychainTypes.js';
export class KeychainTokenStorage
extends BaseTokenStorage
implements SecretStorage
{
- private keychainAvailable: boolean | null = null;
- private keytarModule: Keytar | null = null;
- private keytarLoadAttempted = false;
+ private readonly keychainService: KeychainService;
- async getKeytar(): Promise {
- // If we've already tried loading (successfully or not), return the result
- if (this.keytarLoadAttempted) {
- return this.keytarModule;
- }
-
- this.keytarLoadAttempted = true;
-
- try {
- // Try to import keytar without any timeout - let the OS handle it
- const moduleName = 'keytar';
- // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
- const module = await import(moduleName);
- // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
- this.keytarModule = module.default || module;
- } catch (_) {
- //Keytar is optional so we shouldn't raise an error of log anything.
- }
- return this.keytarModule;
+ constructor(serviceName: string) {
+ super(serviceName);
+ this.keychainService = new KeychainService(serviceName);
}
async getCredentials(serverName: string): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
-
try {
const sanitizedName = this.sanitizeServerName(serverName);
- const data = await keytar.getPassword(this.serviceName, sanitizedName);
+ const data = await this.keychainService.getPassword(sanitizedName);
if (!data) {
return null;
@@ -90,15 +50,6 @@ export class KeychainTokenStorage
}
async setCredentials(credentials: OAuthCredentials): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
-
this.validateCredentials(credentials);
const sanitizedName = this.sanitizeServerName(credentials.serverName);
@@ -108,24 +59,12 @@ export class KeychainTokenStorage
};
const data = JSON.stringify(updatedCredentials);
- await keytar.setPassword(this.serviceName, sanitizedName, data);
+ await this.keychainService.setPassword(sanitizedName, data);
}
async deleteCredentials(serverName: string): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
-
const sanitizedName = this.sanitizeServerName(serverName);
- const deleted = await keytar.deletePassword(
- this.serviceName,
- sanitizedName,
- );
+ const deleted = await this.keychainService.deletePassword(sanitizedName);
if (!deleted) {
throw new Error(`No credentials found for ${serverName}`);
@@ -133,17 +72,8 @@ export class KeychainTokenStorage
}
async listServers(): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
-
try {
- const credentials = await keytar.findCredentials(this.serviceName);
+ const credentials = await this.keychainService.findCredentials();
return credentials
.filter(
(cred) =>
@@ -162,20 +92,9 @@ export class KeychainTokenStorage
}
async getAllCredentials(): Promise> {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
-
const result = new Map();
try {
- const credentials = (
- await keytar.findCredentials(this.serviceName)
- ).filter(
+ const credentials = (await this.keychainService.findCredentials()).filter(
(c) =>
!c.account.startsWith(KEYCHAIN_TEST_PREFIX) &&
!c.account.startsWith(SECRET_PREFIX),
@@ -208,119 +127,48 @@ export class KeychainTokenStorage
}
async clearAll(): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
-
- const servers = this.keytarModule
- ? await this.keytarModule
- .findCredentials(this.serviceName)
- .then((creds) => creds.map((c) => c.account))
- .catch((error: Error) => {
- throw new Error(
- `Failed to list servers for clearing: ${error.message}`,
- );
- })
- : [];
- const errors: Error[] = [];
-
- for (const server of servers) {
- try {
- await this.deleteCredentials(server);
- } catch (error) {
- // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- errors.push(error as Error);
- }
- }
-
- if (errors.length > 0) {
- throw new Error(
- `Failed to clear some credentials: ${errors.map((e) => e.message).join(', ')}`,
- );
- }
- }
-
- // Checks whether or not a set-get-delete cycle with the keychain works.
- // Returns false if any operation fails.
- async checkKeychainAvailability(): Promise {
- if (this.keychainAvailable !== null) {
- return this.keychainAvailable;
- }
-
try {
- const keytar = await this.getKeytar();
- if (!keytar) {
- this.keychainAvailable = false;
- return false;
+ const credentials = await this.keychainService.findCredentials();
+ const errors: Error[] = [];
+
+ for (const cred of credentials) {
+ try {
+ await this.deleteCredentials(cred.account);
+ } catch (error) {
+ // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+ errors.push(error as Error);
+ }
}
- const testAccount = `${KEYCHAIN_TEST_PREFIX}${crypto.randomBytes(8).toString('hex')}`;
- const testPassword = 'test';
-
- await keytar.setPassword(this.serviceName, testAccount, testPassword);
- const retrieved = await keytar.getPassword(this.serviceName, testAccount);
- const deleted = await keytar.deletePassword(
- this.serviceName,
- testAccount,
+ if (errors.length > 0) {
+ throw new Error(
+ `Failed to clear some credentials: ${errors.map((e) => e.message).join(', ')}`,
+ );
+ }
+ } catch (error) {
+ coreEvents.emitFeedback(
+ 'error',
+ 'Failed to clear credentials from keychain',
+ error,
);
-
- const success = deleted && retrieved === testPassword;
- this.keychainAvailable = success;
-
- coreEvents.emitTelemetryKeychainAvailability(
- new KeychainAvailabilityEvent(success),
- );
-
- return success;
- } catch (_error) {
- this.keychainAvailable = false;
-
- // Do not log the raw error message to avoid potential PII leaks
- // (e.g. from OS-level error messages containing file paths)
- coreEvents.emitTelemetryKeychainAvailability(
- new KeychainAvailabilityEvent(false),
- );
-
- return false;
+ throw error;
}
}
async isAvailable(): Promise {
- return this.checkKeychainAvailability();
+ return this.keychainService.isAvailable();
}
async setSecret(key: string, value: string): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
- await keytar.setPassword(this.serviceName, `${SECRET_PREFIX}${key}`, value);
+ await this.keychainService.setPassword(`${SECRET_PREFIX}${key}`, value);
}
async getSecret(key: string): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
- return keytar.getPassword(this.serviceName, `${SECRET_PREFIX}${key}`);
+ return this.keychainService.getPassword(`${SECRET_PREFIX}${key}`);
}
async deleteSecret(key: string): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
- const deleted = await keytar.deletePassword(
- this.serviceName,
+ const deleted = await this.keychainService.deletePassword(
`${SECRET_PREFIX}${key}`,
);
if (!deleted) {
@@ -329,15 +177,8 @@ export class KeychainTokenStorage
}
async listSecrets(): Promise {
- if (!(await this.checkKeychainAvailability())) {
- throw new Error('Keychain is not available');
- }
- const keytar = await this.getKeytar();
- if (!keytar) {
- throw new Error('Keytar module not available');
- }
try {
- const credentials = await keytar.findCredentials(this.serviceName);
+ const credentials = await this.keychainService.findCredentials();
return credentials
.filter((cred) => cred.account.startsWith(SECRET_PREFIX))
.map((cred) => cred.account.substring(SECRET_PREFIX.length));
diff --git a/packages/core/src/output/stream-json-formatter.test.ts b/packages/core/src/output/stream-json-formatter.test.ts
index 69dbaac23b..c911a9dbc2 100644
--- a/packages/core/src/output/stream-json-formatter.test.ts
+++ b/packages/core/src/output/stream-json-formatter.test.ts
@@ -6,14 +6,14 @@
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { StreamJsonFormatter } from './stream-json-formatter.js';
-import { JsonStreamEventType } from './types.js';
-import type {
- InitEvent,
- MessageEvent,
- ToolUseEvent,
- ToolResultEvent,
- ErrorEvent,
- ResultEvent,
+import {
+ JsonStreamEventType,
+ type InitEvent,
+ type MessageEvent,
+ type ToolUseEvent,
+ type ToolResultEvent,
+ type ErrorEvent,
+ type ResultEvent,
} from './types.js';
import type { SessionMetrics } from '../telemetry/uiTelemetry.js';
import { ToolCallDecision } from '../telemetry/tool-call-decision.js';
diff --git a/packages/core/src/policy/config.test.ts b/packages/core/src/policy/config.test.ts
index 3ded361084..9cee9997b2 100644
--- a/packages/core/src/policy/config.test.ts
+++ b/packages/core/src/policy/config.test.ts
@@ -8,8 +8,12 @@ import { describe, it, expect, vi, afterEach, beforeEach } from 'vitest';
import nodePath from 'node:path';
-import type { PolicySettings } from './types.js';
-import { ApprovalMode, PolicyDecision, InProcessCheckerType } from './types.js';
+import {
+ ApprovalMode,
+ PolicyDecision,
+ InProcessCheckerType,
+ type PolicySettings,
+} from './types.js';
import { isDirectorySecure } from '../utils/security.js';
vi.unmock('../config/storage.js');
@@ -202,8 +206,7 @@ describe('createPolicyEngineConfig', () => {
'/tmp/mock/default/policies',
);
const rule = config.rules?.find(
- (r) =>
- r.toolName === 'my-server__*' && r.decision === PolicyDecision.ALLOW,
+ (r) => r.mcpName === 'my-server' && r.decision === PolicyDecision.ALLOW,
);
expect(rule).toBeDefined();
expect(rule?.priority).toBe(4.1); // MCP allowed server
@@ -220,8 +223,7 @@ describe('createPolicyEngineConfig', () => {
'/tmp/mock/default/policies',
);
const rule = config.rules?.find(
- (r) =>
- r.toolName === 'my-server__*' && r.decision === PolicyDecision.DENY,
+ (r) => r.mcpName === 'my-server' && r.decision === PolicyDecision.DENY,
);
expect(rule).toBeDefined();
expect(rule?.priority).toBe(4.9); // MCP excluded server
@@ -247,8 +249,7 @@ describe('createPolicyEngineConfig', () => {
const trustedRule = config.rules?.find(
(r) =>
- r.toolName === 'trusted-server__*' &&
- r.decision === PolicyDecision.ALLOW,
+ r.mcpName === 'trusted-server' && r.decision === PolicyDecision.ALLOW,
);
expect(trustedRule).toBeDefined();
expect(trustedRule?.priority).toBe(4.2); // MCP trusted server
@@ -256,8 +257,7 @@ describe('createPolicyEngineConfig', () => {
// Untrusted server should not have an allow rule
const untrustedRule = config.rules?.find(
(r) =>
- r.toolName === 'untrusted-server__*' &&
- r.decision === PolicyDecision.ALLOW,
+ r.mcpName === 'untrusted-server' && r.decision === PolicyDecision.ALLOW,
);
expect(untrustedRule).toBeUndefined();
});
@@ -284,8 +284,7 @@ describe('createPolicyEngineConfig', () => {
// Check allowed server
const allowedRule = config.rules?.find(
(r) =>
- r.toolName === 'allowed-server__*' &&
- r.decision === PolicyDecision.ALLOW,
+ r.mcpName === 'allowed-server' && r.decision === PolicyDecision.ALLOW,
);
expect(allowedRule).toBeDefined();
expect(allowedRule?.priority).toBe(4.1); // MCP allowed server
@@ -293,8 +292,7 @@ describe('createPolicyEngineConfig', () => {
// Check trusted server
const trustedRule = config.rules?.find(
(r) =>
- r.toolName === 'trusted-server__*' &&
- r.decision === PolicyDecision.ALLOW,
+ r.mcpName === 'trusted-server' && r.decision === PolicyDecision.ALLOW,
);
expect(trustedRule).toBeDefined();
expect(trustedRule?.priority).toBe(4.2); // MCP trusted server
@@ -302,8 +300,7 @@ describe('createPolicyEngineConfig', () => {
// Check excluded server
const excludedRule = config.rules?.find(
(r) =>
- r.toolName === 'excluded-server__*' &&
- r.decision === PolicyDecision.DENY,
+ r.mcpName === 'excluded-server' && r.decision === PolicyDecision.DENY,
);
expect(excludedRule).toBeDefined();
expect(excludedRule?.priority).toBe(4.9); // MCP excluded server
@@ -368,7 +365,7 @@ describe('createPolicyEngineConfig', () => {
const { createPolicyEngineConfig } = await import('./config.js');
const settings: PolicySettings = {
mcp: { excluded: ['my-server'] },
- tools: { allowed: ['my-server__specific-tool'] },
+ tools: { allowed: ['mcp_my-server_specific-tool'] },
};
const config = await createPolicyEngineConfig(
settings,
@@ -377,12 +374,11 @@ describe('createPolicyEngineConfig', () => {
);
const serverDenyRule = config.rules?.find(
- (r) =>
- r.toolName === 'my-server__*' && r.decision === PolicyDecision.DENY,
+ (r) => r.mcpName === 'my-server' && r.decision === PolicyDecision.DENY,
);
const toolAllowRule = config.rules?.find(
(r) =>
- r.toolName === 'my-server__specific-tool' &&
+ r.toolName === 'mcp_my-server_specific-tool' &&
r.decision === PolicyDecision.ALLOW,
);
@@ -404,7 +400,7 @@ describe('createPolicyEngineConfig', () => {
trust: true,
},
},
- tools: { exclude: ['my-server__dangerous-tool'] },
+ tools: { exclude: ['mcp_my-server_dangerous-tool'] },
};
const config = await createPolicyEngineConfig(
settings,
@@ -413,12 +409,11 @@ describe('createPolicyEngineConfig', () => {
);
const serverAllowRule = config.rules?.find(
- (r) =>
- r.toolName === 'my-server__*' && r.decision === PolicyDecision.ALLOW,
+ (r) => r.mcpName === 'my-server' && r.decision === PolicyDecision.ALLOW,
);
const toolDenyRule = config.rules?.find(
(r) =>
- r.toolName === 'my-server__dangerous-tool' &&
+ r.toolName === 'mcp_my-server_dangerous-tool' &&
r.decision === PolicyDecision.DENY,
);
@@ -432,8 +427,8 @@ describe('createPolicyEngineConfig', () => {
it('should handle complex priority scenarios correctly', async () => {
const settings: PolicySettings = {
tools: {
- allowed: ['my-server__tool1', 'other-tool'], // Priority 4.3
- exclude: ['my-server__tool2', 'glob'], // Priority 4.4
+ allowed: ['mcp_trusted-server_tool1', 'other-tool'], // Priority 4.3
+ exclude: ['mcp_trusted-server_tool2', 'glob'], // Priority 4.4
},
mcp: {
allowed: ['allowed-server'], // Priority 4.1
@@ -564,13 +559,12 @@ describe('createPolicyEngineConfig', () => {
// Neither server should have an allow rule
const noTrustRule = config.rules?.find(
(r) =>
- r.toolName === 'no-trust-property__*' &&
+ r.mcpName === 'no-trust-property' &&
r.decision === PolicyDecision.ALLOW,
);
const explicitFalseRule = config.rules?.find(
(r) =>
- r.toolName === 'explicit-false__*' &&
- r.decision === PolicyDecision.ALLOW,
+ r.mcpName === 'explicit-false' && r.decision === PolicyDecision.ALLOW,
);
expect(noTrustRule).toBeUndefined();
diff --git a/packages/core/src/policy/config.ts b/packages/core/src/policy/config.ts
index a1e337436e..8437cb9845 100644
--- a/packages/core/src/policy/config.ts
+++ b/packages/core/src/policy/config.ts
@@ -31,6 +31,7 @@ import { debugLogger } from '../utils/debugLogger.js';
import { SHELL_TOOL_NAMES } from '../utils/shell-utils.js';
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
import { isNodeError } from '../utils/errors.js';
+import { MCP_TOOL_PREFIX } from '../tools/mcp-tool.js';
import { isDirectorySecure } from '../utils/security.js';
@@ -342,7 +343,11 @@ export async function createPolicyEngineConfig(
if (settings.mcp?.excluded) {
for (const serverName of settings.mcp.excluded) {
rules.push({
- toolName: `${serverName}__*`,
+ toolName:
+ serverName === '*'
+ ? `${MCP_TOOL_PREFIX}*`
+ : `${MCP_TOOL_PREFIX}${serverName}_*`,
+ mcpName: serverName,
decision: PolicyDecision.DENY,
priority: MCP_EXCLUDED_PRIORITY,
source: 'Settings (MCP Excluded)',
@@ -423,9 +428,10 @@ export async function createPolicyEngineConfig(
)) {
if (serverConfig.trust) {
// Trust all tools from this MCP server
- // Using pattern matching for MCP tool names which are formatted as "serverName__toolName"
+ // Using explicit mcpName metadata and FQN mcp_{serverName}_*
rules.push({
- toolName: `${serverName}__*`,
+ toolName: `${MCP_TOOL_PREFIX}${serverName}_*`,
+ mcpName: serverName,
decision: PolicyDecision.ALLOW,
priority: TRUSTED_MCP_SERVER_PRIORITY,
source: 'Settings (MCP Trusted)',
@@ -439,7 +445,11 @@ export async function createPolicyEngineConfig(
if (settings.mcp?.allowed) {
for (const serverName of settings.mcp.allowed) {
rules.push({
- toolName: `${serverName}__*`,
+ toolName:
+ serverName === '*'
+ ? `${MCP_TOOL_PREFIX}*`
+ : `${MCP_TOOL_PREFIX}${serverName}_*`,
+ mcpName: serverName,
decision: PolicyDecision.ALLOW,
priority: ALLOWED_MCP_SERVER_PRIORITY,
source: 'Settings (MCP Allowed)',
diff --git a/packages/core/src/policy/policies/read-only.toml b/packages/core/src/policy/policies/read-only.toml
index c9c96923e7..ad996864b2 100644
--- a/packages/core/src/policy/policies/read-only.toml
+++ b/packages/core/src/policy/policies/read-only.toml
@@ -51,3 +51,8 @@ priority = 50
toolName = "google_web_search"
decision = "allow"
priority = 50
+
+[[rule]]
+toolName = ["codebase_investigator", "cli_help"]
+decision = "allow"
+priority = 50
\ No newline at end of file
diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts
index b8e6968af9..baf475701c 100644
--- a/packages/core/src/policy/policy-engine.test.ts
+++ b/packages/core/src/policy/policy-engine.test.ts
@@ -150,7 +150,8 @@ describe('PolicyEngine', () => {
it('should match unqualified tool names with qualified rules when serverName is provided', async () => {
const rules: PolicyRule[] = [
{
- toolName: 'my-server__tool',
+ toolName: 'mcp_my-server_tool',
+ mcpName: 'my-server',
decision: PolicyDecision.ALLOW,
},
];
@@ -159,23 +160,9 @@ describe('PolicyEngine', () => {
// Match with qualified name (standard)
expect(
- (await engine.check({ name: 'my-server__tool' }, 'my-server')).decision,
+ (await engine.check({ name: 'mcp_my-server_tool' }, 'my-server'))
+ .decision,
).toBe(PolicyDecision.ALLOW);
-
- // Match with unqualified name + serverName (the fix)
- expect((await engine.check({ name: 'tool' }, 'my-server')).decision).toBe(
- PolicyDecision.ALLOW,
- );
-
- // Should NOT match with unqualified name but NO serverName
- expect((await engine.check({ name: 'tool' }, undefined)).decision).toBe(
- PolicyDecision.ASK_USER,
- );
-
- // Should NOT match with unqualified name but WRONG serverName
- expect(
- (await engine.check({ name: 'tool' }, 'wrong-server')).decision,
- ).toBe(PolicyDecision.ASK_USER);
});
it('should match by args pattern', async () => {
@@ -476,61 +463,41 @@ describe('PolicyEngine', () => {
(await engine.check({ name: 'read_file' }, undefined)).decision,
).toBe(PolicyDecision.ALLOW);
expect(
- (await engine.check({ name: 'my-server__tool' }, 'my-server')).decision,
+ (await engine.check({ name: 'mcp_my-server_tool' }, 'my-server'))
+ .decision,
).toBe(PolicyDecision.ALLOW);
});
- it('should match any MCP tool when toolName is *__*', async () => {
+ it('should match any MCP tool when toolName is mcp_*', async () => {
engine = new PolicyEngine({
rules: [
- { toolName: '*__*', decision: PolicyDecision.ALLOW, priority: 10 },
+ { toolName: 'mcp_*', decision: PolicyDecision.ALLOW, priority: 10 },
],
defaultDecision: PolicyDecision.DENY,
});
- expect((await engine.check({ name: 'mcp__tool' }, 'mcp')).decision).toBe(
- PolicyDecision.ALLOW,
- );
expect(
- (await engine.check({ name: 'other__tool' }, 'other')).decision,
+ (await engine.check({ name: 'mcp_mcp_tool' }, 'mcp')).decision,
+ ).toBe(PolicyDecision.ALLOW);
+ expect(
+ (await engine.check({ name: 'mcp_other_tool' }, 'other')).decision,
).toBe(PolicyDecision.ALLOW);
expect(
(await engine.check({ name: 'read_file' }, undefined)).decision,
).toBe(PolicyDecision.DENY);
});
- it('should match specific tool across all servers when using *__tool', async () => {
- engine = new PolicyEngine({
- rules: [
- {
- toolName: '*__search',
- decision: PolicyDecision.ALLOW,
- priority: 10,
- },
- ],
- defaultDecision: PolicyDecision.DENY,
- });
-
- expect((await engine.check({ name: 'ws__search' }, 'ws')).decision).toBe(
- PolicyDecision.ALLOW,
- );
- expect((await engine.check({ name: 'gh__search' }, 'gh')).decision).toBe(
- PolicyDecision.ALLOW,
- );
- expect((await engine.check({ name: 'gh__list' }, 'gh')).decision).toBe(
- PolicyDecision.DENY,
- );
- });
-
it('should match MCP server wildcard patterns', async () => {
const rules: PolicyRule[] = [
{
- toolName: 'my-server__*',
+ toolName: 'mcp_my-server_*',
+ mcpName: 'my-server',
decision: PolicyDecision.ALLOW,
priority: 10,
},
{
- toolName: 'blocked-server__*',
+ toolName: 'mcp_blocked-server_*',
+ mcpName: 'blocked-server',
decision: PolicyDecision.DENY,
priority: 20,
},
@@ -540,19 +507,23 @@ describe('PolicyEngine', () => {
// Should match my-server tools
expect(
- (await engine.check({ name: 'my-server__tool1' }, 'my-server'))
+ (await engine.check({ name: 'mcp_my-server_tool1' }, 'my-server'))
.decision,
).toBe(PolicyDecision.ALLOW);
expect(
- (await engine.check({ name: 'my-server__another_tool' }, 'my-server'))
- .decision,
+ (
+ await engine.check(
+ { name: 'mcp_my-server_another_tool' },
+ 'my-server',
+ )
+ ).decision,
).toBe(PolicyDecision.ALLOW);
// Should match blocked-server tools
expect(
(
await engine.check(
- { name: 'blocked-server__tool1' },
+ { name: 'mcp_blocked-server_tool1' },
'blocked-server',
)
).decision,
@@ -560,7 +531,7 @@ describe('PolicyEngine', () => {
expect(
(
await engine.check(
- { name: 'blocked-server__dangerous' },
+ { name: 'mcp_blocked-server_dangerous' },
'blocked-server',
)
).decision,
@@ -568,7 +539,7 @@ describe('PolicyEngine', () => {
// Should not match other patterns
expect(
- (await engine.check({ name: 'other-server__tool' }, 'other-server'))
+ (await engine.check({ name: 'mcp_other-server_tool' }, 'other-server'))
.decision,
).toBe(PolicyDecision.ASK_USER);
expect(
@@ -582,12 +553,14 @@ describe('PolicyEngine', () => {
it('should prioritize specific tool rules over server wildcards', async () => {
const rules: PolicyRule[] = [
{
- toolName: 'my-server__*',
+ toolName: 'mcp_my-server_*',
+ mcpName: 'my-server',
decision: PolicyDecision.ALLOW,
priority: 10,
},
{
- toolName: 'my-server__dangerous-tool',
+ toolName: 'mcp_my-server_dangerous-tool',
+ mcpName: 'my-server',
decision: PolicyDecision.DENY,
priority: 20,
},
@@ -597,33 +570,38 @@ describe('PolicyEngine', () => {
// Specific tool deny should override server allow
expect(
- (await engine.check({ name: 'my-server__dangerous-tool' }, 'my-server'))
- .decision,
+ (
+ await engine.check(
+ { name: 'mcp_my-server_dangerous-tool' },
+ 'my-server',
+ )
+ ).decision,
).toBe(PolicyDecision.DENY);
expect(
- (await engine.check({ name: 'my-server__safe-tool' }, 'my-server'))
+ (await engine.check({ name: 'mcp_my-server_safe-tool' }, 'my-server'))
.decision,
).toBe(PolicyDecision.ALLOW);
});
it('should NOT match spoofed server names when using wildcards', async () => {
- // Vulnerability: A rule for 'prefix__*' matches 'prefix__suffix__tool'
- // effectively allowing a server named 'prefix__suffix' to spoof 'prefix'.
+ // Vulnerability: A rule for 'mcp_prefix_*' matches 'mcp_prefix__suffix_tool'
+ // effectively allowing a server named 'mcp_prefix_suffix' to spoof 'prefix'.
const rules: PolicyRule[] = [
{
- toolName: 'safe_server__*',
+ toolName: 'mcp_safe_server_*',
+ mcpName: 'safe_server',
decision: PolicyDecision.ALLOW,
},
];
engine = new PolicyEngine({ rules });
- // A tool from a different server 'safe_server__malicious'
- const spoofedToolCall = { name: 'safe_server__malicious__tool' };
+ // A tool from a different server 'mcp_safe_server_malicious'
+ const spoofedToolCall = { name: 'mcp_mcp_safe_server_malicious_tool' };
// CURRENT BEHAVIOR (FIXED): Matches because it starts with 'safe_server__' BUT serverName doesn't match 'safe_server'
// We expect this to FAIL matching the ALLOW rule, thus falling back to default (ASK_USER)
expect(
- (await engine.check(spoofedToolCall, 'safe_server__malicious'))
+ (await engine.check(spoofedToolCall, 'mcp_safe_server_malicious'))
.decision,
).toBe(PolicyDecision.ASK_USER);
});
@@ -631,14 +609,15 @@ describe('PolicyEngine', () => {
it('should verify tool name prefix even if serverName matches', async () => {
const rules: PolicyRule[] = [
{
- toolName: 'safe_server__*',
+ toolName: 'mcp_safe_server_*',
+ mcpName: 'safe_server',
decision: PolicyDecision.ALLOW,
},
];
engine = new PolicyEngine({ rules });
// serverName matches, but tool name does not start with prefix
- const invalidToolCall = { name: 'other_server__tool' };
+ const invalidToolCall = { name: 'mcp_other_server_tool' };
expect(
(await engine.check(invalidToolCall, 'safe_server')).decision,
).toBe(PolicyDecision.ASK_USER);
@@ -647,13 +626,14 @@ describe('PolicyEngine', () => {
it('should allow when both serverName and tool name prefix match', async () => {
const rules: PolicyRule[] = [
{
- toolName: 'safe_server__*',
+ toolName: 'mcp_safe_server_*',
+ mcpName: 'safe_server',
decision: PolicyDecision.ALLOW,
},
];
engine = new PolicyEngine({ rules });
- const validToolCall = { name: 'safe_server__tool' };
+ const validToolCall = { name: 'mcp_safe_server_tool' };
expect((await engine.check(validToolCall, 'safe_server')).decision).toBe(
PolicyDecision.ALLOW,
);
@@ -2007,11 +1987,16 @@ describe('PolicyEngine', () => {
it('should support wildcard patterns for checkers', async () => {
const rules: PolicyRule[] = [
- { toolName: 'server__tool', decision: PolicyDecision.ALLOW },
+ {
+ toolName: 'mcp_server_tool',
+ mcpName: 'server',
+ decision: PolicyDecision.ALLOW,
+ },
];
const wildcardChecker: SafetyCheckerRule = {
checker: { type: 'external', name: 'wildcard' },
- toolName: 'server__*',
+ toolName: 'mcp_server_*',
+ mcpName: 'server',
};
engine = new PolicyEngine(
@@ -2023,7 +2008,7 @@ describe('PolicyEngine', () => {
decision: SafetyCheckDecision.ALLOW,
});
- await engine.check({ name: 'server__tool' }, 'server');
+ await engine.check({ name: 'mcp_server_tool' }, 'server');
expect(mockCheckerRunner.runChecker).toHaveBeenCalledWith(
expect.anything(),
@@ -2137,6 +2122,8 @@ describe('PolicyEngine', () => {
rules: PolicyRule[];
approvalMode?: ApprovalMode;
nonInteractive?: boolean;
+ allToolNames?: string[];
+ metadata?: Map>;
expected: string[];
}
@@ -2144,11 +2131,13 @@ describe('PolicyEngine', () => {
{
name: 'should return empty set when no rules provided',
rules: [],
+ allToolNames: ['tool1'],
expected: [],
},
{
name: 'should apply rules without explicit modes to all modes',
rules: [{ toolName: 'tool1', decision: PolicyDecision.DENY }],
+ allToolNames: ['tool1', 'tool2'],
expected: ['tool1'],
},
{
@@ -2168,6 +2157,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['tool1'],
expected: [],
},
{
@@ -2184,6 +2174,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['tool1', 'tool2', 'tool3'],
expected: ['tool1'],
},
{
@@ -2202,6 +2193,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['tool1'],
expected: ['tool1'],
},
{
@@ -2220,6 +2212,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['tool1'],
expected: [],
},
{
@@ -2232,7 +2225,8 @@ describe('PolicyEngine', () => {
},
],
nonInteractive: true,
- expected: [],
+ allToolNames: ['tool1'],
+ expected: ['tool1'],
},
{
name: 'should ignore rules with argsPattern',
@@ -2244,6 +2238,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['tool1'],
expected: [],
},
{
@@ -2256,6 +2251,7 @@ describe('PolicyEngine', () => {
},
],
approvalMode: ApprovalMode.PLAN,
+ allToolNames: ['tool1'],
expected: ['tool1'],
},
{
@@ -2268,6 +2264,7 @@ describe('PolicyEngine', () => {
},
],
approvalMode: ApprovalMode.DEFAULT,
+ allToolNames: ['tool1'],
expected: [],
},
{
@@ -2286,36 +2283,55 @@ describe('PolicyEngine', () => {
},
],
approvalMode: ApprovalMode.YOLO,
+ allToolNames: ['dangerous-tool', 'safe-tool'],
expected: [],
},
{
name: 'should respect server wildcard DENY',
rules: [
{
- toolName: 'server__*',
+ toolName: 'mcp_server_*',
+ mcpName: 'server',
decision: PolicyDecision.DENY,
modes: [ApprovalMode.DEFAULT],
},
],
- expected: ['server__*'],
+ allToolNames: [
+ 'mcp_server_tool1',
+ 'mcp_server_tool2',
+ 'mcp_other_tool',
+ ],
+ metadata: new Map([
+ ['mcp_server_tool1', { _serverName: 'server' }],
+ ['mcp_server_tool2', { _serverName: 'server' }],
+ ['mcp_other_tool', { _serverName: 'other' }],
+ ]),
+ expected: ['mcp_server_tool1', 'mcp_server_tool2'],
},
{
name: 'should expand server wildcard for specific tools if already processed',
rules: [
{
- toolName: 'server__*',
+ toolName: 'mcp_server_*',
+ mcpName: 'server',
decision: PolicyDecision.DENY,
priority: 100,
modes: [ApprovalMode.DEFAULT],
},
{
- toolName: 'server__tool1',
- decision: PolicyDecision.DENY,
+ toolName: 'mcp_server_tool1',
+ mcpName: 'server',
+ decision: PolicyDecision.DENY, // redundant but tests ordering
priority: 10,
modes: [ApprovalMode.DEFAULT],
},
],
- expected: ['server__*', 'server__tool1'],
+ allToolNames: ['mcp_server_tool1', 'mcp_server_tool2'],
+ metadata: new Map([
+ ['mcp_server_tool1', { _serverName: 'server' }],
+ ['mcp_server_tool2', { _serverName: 'server' }],
+ ]),
+ expected: ['mcp_server_tool1', 'mcp_server_tool2'],
},
{
name: 'should exclude run_shell_command but NOT write_file in simulated Plan Mode',
@@ -2342,24 +2358,29 @@ describe('PolicyEngine', () => {
priority: 10,
},
],
- expected: ['run_shell_command'],
+ allToolNames: ['write_file', 'run_shell_command', 'read_file'],
+ expected: ['run_shell_command', 'read_file'],
},
{
name: 'should NOT exclude tool if covered by a higher priority wildcard ALLOW',
rules: [
{
- toolName: 'server__*',
+ toolName: 'mcp_server_*',
+ mcpName: 'server',
decision: PolicyDecision.ALLOW,
priority: 100,
modes: [ApprovalMode.DEFAULT],
},
{
- toolName: 'server__tool1',
+ toolName: 'mcp_server_tool1',
+ mcpName: 'server',
decision: PolicyDecision.DENY,
priority: 10,
modes: [ApprovalMode.DEFAULT],
},
],
+ allToolNames: ['mcp_server_tool1'],
+ metadata: new Map([['mcp_server_tool1', { _serverName: 'server' }]]),
expected: [],
},
{
@@ -2371,41 +2392,63 @@ describe('PolicyEngine', () => {
priority: 10,
},
],
- expected: ['*'],
+ allToolNames: ['toolA', 'toolB', 'mcp_server_toolC'],
+ expected: ['toolA', 'toolB', 'mcp_server_toolC'], // all tools denied by *
},
{
name: 'should handle MCP category wildcard *__* in getExcludedTools',
rules: [
{
- toolName: '*__*',
+ toolName: 'mcp_*',
decision: PolicyDecision.DENY,
priority: 10,
},
],
- expected: ['*__*'],
+ allToolNames: ['localTool', 'mcp_myserver_mytool'],
+ metadata: new Map([
+ ['mcp_myserver_mytool', { _serverName: 'myserver' }],
+ ]),
+ expected: ['mcp_myserver_mytool'],
},
{
- name: 'should handle tool wildcard *__search in getExcludedTools',
+ name: 'should handle tool wildcard mcp_server_* in getExcludedTools',
rules: [
{
- toolName: '*__search',
+ toolName: 'mcp_server_*',
decision: PolicyDecision.DENY,
priority: 10,
},
],
- expected: ['*__search'],
+ allToolNames: [
+ 'localTool',
+ 'mcp_server_search',
+ 'mcp_otherserver_read',
+ ],
+ metadata: new Map([
+ ['mcp_server_search', { _serverName: 'server' }],
+ ['mcp_otherserver_read', { _serverName: 'otherserver' }],
+ ]),
+ expected: ['mcp_server_search'],
},
];
it.each(testCases)(
'$name',
- ({ rules, approvalMode, nonInteractive, expected }) => {
+ ({
+ rules,
+ approvalMode,
+ nonInteractive,
+ allToolNames,
+ metadata,
+ expected,
+ }) => {
engine = new PolicyEngine({
rules,
approvalMode: approvalMode ?? ApprovalMode.DEFAULT,
nonInteractive: nonInteractive ?? false,
});
- const excluded = engine.getExcludedTools();
+ const toolsSet = allToolNames ? new Set(allToolNames) : undefined;
+ const excluded = engine.getExcludedTools(metadata, toolsSet);
expect(Array.from(excluded).sort()).toEqual(expected.sort());
},
);
@@ -2420,7 +2463,10 @@ describe('PolicyEngine', () => {
},
],
});
- const excluded = engine.getExcludedTools();
+ const excluded = engine.getExcludedTools(
+ undefined,
+ new Set(['dangerous_tool']),
+ );
expect(Array.from(excluded)).toEqual([]);
});
@@ -2438,7 +2484,10 @@ describe('PolicyEngine', () => {
['dangerous_tool', { destructiveHint: true }],
['safe_tool', { readOnlyHint: true }],
]);
- const excluded = engine.getExcludedTools(metadata);
+ const excluded = engine.getExcludedTools(
+ metadata,
+ new Set(['dangerous_tool', 'safe_tool']),
+ );
expect(Array.from(excluded)).toEqual(['dangerous_tool']);
});
@@ -2455,7 +2504,10 @@ describe('PolicyEngine', () => {
const metadata = new Map>([
['safe_tool', { readOnlyHint: true }],
]);
- const excluded = engine.getExcludedTools(metadata);
+ const excluded = engine.getExcludedTools(
+ metadata,
+ new Set(['safe_tool']),
+ );
expect(Array.from(excluded)).toEqual([]);
});
@@ -2463,7 +2515,8 @@ describe('PolicyEngine', () => {
engine = new PolicyEngine({
rules: [
{
- toolName: 'server__*',
+ toolName: 'mcp_server_*',
+ mcpName: 'server',
toolAnnotations: { destructiveHint: true },
decision: PolicyDecision.DENY,
priority: 10,
@@ -2471,12 +2524,25 @@ describe('PolicyEngine', () => {
],
});
const metadata = new Map>([
- ['server__dangerous_tool', { destructiveHint: true }],
- ['other__dangerous_tool', { destructiveHint: true }],
- ['server__safe_tool', { readOnlyHint: true }],
+ [
+ 'mcp_server_dangerous_tool',
+ { destructiveHint: true, _serverName: 'server' },
+ ],
+ [
+ 'mcp_other_dangerous_tool',
+ { destructiveHint: true, _serverName: 'other' },
+ ],
+ ['mcp_server_safe_tool', { readOnlyHint: true, _serverName: 'server' }],
]);
- const excluded = engine.getExcludedTools(metadata);
- expect(Array.from(excluded)).toEqual(['server__dangerous_tool']);
+ const excluded = engine.getExcludedTools(
+ metadata,
+ new Set([
+ 'mcp_server_dangerous_tool',
+ 'mcp_other_dangerous_tool',
+ 'mcp_server_safe_tool',
+ ]),
+ );
+ expect(Array.from(excluded)).toEqual(['mcp_server_dangerous_tool']);
});
it('should exclude unprocessed tools from allToolNames when global DENY is active', () => {
@@ -2493,8 +2559,8 @@ describe('PolicyEngine', () => {
priority: 70,
},
{
- // Simulates plan.toml: mcpName="*" → toolName="*__*"
- toolName: '*__*',
+ // Simulates plan.toml: mcpName="*" → toolName="mcp_*"
+ toolName: 'mcp_*',
toolAnnotations: { readOnlyHint: true },
decision: PolicyDecision.ASK_USER,
priority: 70,
@@ -2505,36 +2571,42 @@ describe('PolicyEngine', () => {
},
],
});
- // MCP tools are registered with unqualified names in ToolRegistry
+ // MCP tools are registered with qualified names in ToolRegistry
const allToolNames = new Set([
'glob',
'read_file',
'shell',
'web_fetch',
- 'read_mcp_tool',
- 'write_mcp_tool',
+ 'mcp_my-server_read_mcp_tool',
+ 'mcp_my-server_write_mcp_tool',
]);
// buildToolMetadata() includes _serverName for MCP tools
const toolMetadata = new Map>([
- ['read_mcp_tool', { readOnlyHint: true, _serverName: 'my-server' }],
- ['write_mcp_tool', { readOnlyHint: false, _serverName: 'my-server' }],
+ [
+ 'mcp_my-server_read_mcp_tool',
+ { readOnlyHint: true, _serverName: 'my-server' },
+ ],
+ [
+ 'mcp_my-server_write_mcp_tool',
+ { readOnlyHint: false, _serverName: 'my-server' },
+ ],
]);
const excluded = engine.getExcludedTools(toolMetadata, allToolNames);
expect(excluded.has('shell')).toBe(true);
expect(excluded.has('web_fetch')).toBe(true);
// Non-read-only MCP tool excluded by catch-all DENY
- expect(excluded.has('write_mcp_tool')).toBe(true);
+ expect(excluded.has('mcp_my-server_write_mcp_tool')).toBe(true);
expect(excluded.has('glob')).toBe(false);
expect(excluded.has('read_file')).toBe(false);
// Read-only MCP tool allowed by annotation rule
- expect(excluded.has('read_mcp_tool')).toBe(false);
+ expect(excluded.has('mcp_my-server_read_mcp_tool')).toBe(false);
});
- it('should match already-qualified MCP tool names without _serverName', () => {
+ it('should match MCP wildcard rules when explicitly mapped with _serverName', () => {
engine = new PolicyEngine({
rules: [
{
- toolName: '*__*',
+ toolName: 'mcp_*',
toolAnnotations: { readOnlyHint: true },
decision: PolicyDecision.ASK_USER,
priority: 70,
@@ -2547,17 +2619,23 @@ describe('PolicyEngine', () => {
});
// Tool registered with qualified name (collision case)
const allToolNames = new Set([
- 'myserver__read_tool',
- 'myserver__write_tool',
+ 'mcp_myserver_read_tool',
+ 'mcp_myserver_write_tool',
]);
const toolMetadata = new Map>([
- ['myserver__read_tool', { readOnlyHint: true }],
- ['myserver__write_tool', { readOnlyHint: false }],
+ [
+ 'mcp_myserver_read_tool',
+ { readOnlyHint: true, _serverName: 'myserver' },
+ ],
+ [
+ 'mcp_myserver_write_tool',
+ { readOnlyHint: false, _serverName: 'myserver' },
+ ],
]);
const excluded = engine.getExcludedTools(toolMetadata, allToolNames);
- // Qualified name already contains __, matched directly without _serverName
- expect(excluded.has('myserver__read_tool')).toBe(false);
- expect(excluded.has('myserver__write_tool')).toBe(true);
+ // Qualified name matched using explicit _serverName
+ expect(excluded.has('mcp_myserver_read_tool')).toBe(false);
+ expect(excluded.has('mcp_myserver_write_tool')).toBe(true);
});
it('should not exclude unprocessed tools when allToolNames is not provided (backward compat)', () => {
@@ -2648,7 +2726,7 @@ describe('PolicyEngine', () => {
modes: [ApprovalMode.PLAN],
},
{
- toolName: '*__*',
+ toolName: 'mcp_*',
toolAnnotations: { readOnlyHint: true },
decision: PolicyDecision.ASK_USER,
priority: 70,
@@ -2679,13 +2757,19 @@ describe('PolicyEngine', () => {
'write_todos',
'memory',
'save_memory',
- 'read_tool',
- 'write_tool',
+ 'mcp_mcp-server_read_tool',
+ 'mcp_mcp-server_write_tool',
]);
// buildToolMetadata() includes _serverName for MCP tools
const toolMetadata = new Map>([
- ['read_tool', { readOnlyHint: true, _serverName: 'mcp-server' }],
- ['write_tool', { readOnlyHint: false, _serverName: 'mcp-server' }],
+ [
+ 'mcp_mcp-server_read_tool',
+ { readOnlyHint: true, _serverName: 'mcp-server' },
+ ],
+ [
+ 'mcp_mcp-server_write_tool',
+ { readOnlyHint: false, _serverName: 'mcp-server' },
+ ],
]);
const excluded = engine.getExcludedTools(toolMetadata, allToolNames);
// These should be excluded (caught by catch-all DENY)
@@ -2698,7 +2782,7 @@ describe('PolicyEngine', () => {
expect(excluded.has('write_file')).toBe(true);
expect(excluded.has('replace')).toBe(true);
// Non-read-only MCP tool excluded by catch-all DENY
- expect(excluded.has('write_tool')).toBe(true);
+ expect(excluded.has('mcp_mcp-server_write_tool')).toBe(true);
// These should NOT be excluded (explicitly allowed)
expect(excluded.has('glob')).toBe(false);
expect(excluded.has('grep_search')).toBe(false);
@@ -2710,7 +2794,7 @@ describe('PolicyEngine', () => {
expect(excluded.has('exit_plan_mode')).toBe(false);
expect(excluded.has('save_memory')).toBe(false);
// Read-only MCP tool allowed by annotation rule (matched via _serverName)
- expect(excluded.has('read_tool')).toBe(false);
+ expect(excluded.has('mcp_mcp-server_read_tool')).toBe(false);
});
});
@@ -3023,7 +3107,6 @@ describe('PolicyEngine', () => {
expect(checkers[0].checker.name).toBe('c2');
});
});
-
describe('Tool Annotations', () => {
it('should match tools by semantic annotations', async () => {
engine = new PolicyEngine({
@@ -3058,13 +3141,13 @@ describe('PolicyEngine', () => {
engine = new PolicyEngine({
rules: [
{
- toolName: '*__*',
+ toolName: 'mcp_*',
toolAnnotations: { experimental: true },
decision: PolicyDecision.DENY,
priority: 20,
},
{
- toolName: '*__*',
+ toolName: 'mcp_*',
decision: PolicyDecision.ALLOW,
priority: 10,
},
@@ -3073,21 +3156,20 @@ describe('PolicyEngine', () => {
expect(
(
- await engine.check({ name: 'mcp__test' }, 'mcp', {
+ await engine.check({ name: 'mcp_mcp_test' }, 'mcp', {
experimental: true,
})
).decision,
).toBe(PolicyDecision.DENY);
expect(
(
- await engine.check({ name: 'mcp__stable' }, 'mcp', {
+ await engine.check({ name: 'mcp_mcp_stable' }, 'mcp', {
experimental: false,
})
).decision,
).toBe(PolicyDecision.ALLOW);
});
});
-
describe('hook checkers', () => {
it('should add and retrieve hook checkers in priority order', () => {
engine.addHookChecker({
diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts
index 03087716ff..0d6a043da0 100644
--- a/packages/core/src/policy/policy-engine.ts
+++ b/packages/core/src/policy/policy-engine.ts
@@ -25,6 +25,11 @@ import {
hasRedirection,
} from '../utils/shell-utils.js';
import { getToolAliases } from '../tools/tool-names.js';
+import {
+ MCP_TOOL_PREFIX,
+ isMcpToolAnnotation,
+ parseMcpToolName,
+} from '../tools/mcp-tool.js';
function isWildcardPattern(name: string): boolean {
return name === '*' || name.includes('*');
@@ -32,7 +37,7 @@ function isWildcardPattern(name: string): boolean {
/**
* Checks if a tool call matches a wildcard pattern.
- * Supports global (*) and composite (server__*, *__tool, *__*) patterns.
+ * Supports global (*) and the explicit MCP (*mcp_serverName_**) format.
*/
function matchesWildcard(
pattern: string,
@@ -43,59 +48,25 @@ function matchesWildcard(
return true;
}
- if (pattern.includes('__')) {
- return matchesCompositePattern(pattern, toolName, serverName);
+ if (pattern === `${MCP_TOOL_PREFIX}*`) {
+ return serverName !== undefined;
}
+ if (pattern.startsWith(MCP_TOOL_PREFIX) && pattern.endsWith('_*')) {
+ const expectedServerName = pattern.slice(MCP_TOOL_PREFIX.length, -2);
+ // 1. Must be an MCP tool call (has serverName)
+ // 2. Server name must match
+ // 3. Tool name must be properly qualified by that server
+ if (serverName === undefined || serverName !== expectedServerName) {
+ return false;
+ }
+ return toolName.startsWith(`${MCP_TOOL_PREFIX}${expectedServerName}_`);
+ }
+
+ // Not a recognized wildcard pattern, fallback to exact match just in case
return toolName === pattern;
}
-/**
- * Matches composite patterns like "server__*", "*__tool", or "*__*".
- */
-function matchesCompositePattern(
- pattern: string,
- toolName: string,
- serverName: string | undefined,
-): boolean {
- const parts = pattern.split('__');
- if (parts.length !== 2) return false;
- const [patternServer, patternTool] = parts;
-
- // 1. Identify the tool's components
- const { actualServer, actualTool } = getToolMetadata(toolName, serverName);
-
- // 2. Composite patterns require a server context
- if (actualServer === undefined) {
- return false;
- }
-
- // 3. Robustness: if serverName is provided, toolName MUST be qualified by it.
- // This prevents "malicious-server" from spoofing "trusted-server" by naming itself "trusted-server__malicious".
- if (serverName !== undefined && !toolName.startsWith(serverName + '__')) {
- return false;
- }
-
- // 4. Match components
- const serverMatch = patternServer === '*' || patternServer === actualServer;
- const toolMatch = patternTool === '*' || patternTool === actualTool;
-
- return serverMatch && toolMatch;
-}
-
-/**
- * Extracts the server and unqualified tool name from a tool call context.
- */
-function getToolMetadata(toolName: string, serverName: string | undefined) {
- const sepIndex = toolName.indexOf('__');
- const isQualified = sepIndex !== -1;
- return {
- actualServer:
- serverName ?? (isQualified ? toolName.substring(0, sepIndex) : undefined),
- actualTool: isQualified ? toolName.substring(sepIndex + 2) : toolName,
- };
-}
-
function ruleMatches(
rule: PolicyRule | SafetyCheckerRule,
toolCall: FunctionCall,
@@ -111,9 +82,20 @@ function ruleMatches(
}
}
+ // Strictly enforce mcpName identity if the rule dictates it
+ if (rule.mcpName) {
+ if (rule.mcpName === '*') {
+ // Rule requires it to be ANY MCP tool
+ if (serverName === undefined) return false;
+ } else {
+ // Rule requires it to be a specific MCP server
+ if (serverName !== rule.mcpName) return false;
+ }
+ }
+
// Check tool name if specified
if (rule.toolName) {
- // Support wildcard patterns: "serverName__*" matches "serverName__anyTool"
+ // Support wildcard patterns: "mcp_serverName_*" matches "mcp_serverName_anyTool"
if (rule.toolName === '*') {
// Match all tools
} else if (isWildcardPattern(rule.toolName)) {
@@ -371,6 +353,22 @@ export class PolicyEngine {
serverName: string | undefined,
toolAnnotations?: Record,
): Promise {
+ // Case 1: Metadata injection is the primary and safest way to identify an MCP server.
+ // If we have explicit `_serverName` metadata (usually injected by tool-registry for active tools), use it.
+ if (!serverName && isMcpToolAnnotation(toolAnnotations)) {
+ serverName = toolAnnotations._serverName;
+ }
+
+ // Case 2: Fallback for static FQN strings (e.g. from TOML policies or allowed/excluded settings strings).
+ // These strings don't have active metadata objects associated with them during policy generation,
+ // so we must extract the server name from the qualified `mcp_{server}_{tool}` format.
+ if (!serverName && toolCall.name) {
+ const parsed = parseMcpToolName(toolCall.name);
+ if (parsed.serverName) {
+ serverName = parsed.serverName;
+ }
+ }
+
let stringifiedArgs: string | undefined;
// Compute stringified args once before the loop
if (
@@ -404,20 +402,12 @@ export class PolicyEngine {
let matchedRule: PolicyRule | undefined;
let decision: PolicyDecision | undefined;
- // For tools with a server name, we want to try matching both the
- // original name and the fully qualified name (server__tool).
// We also want to check legacy aliases for the tool name.
const toolNamesToTry = toolCall.name ? getToolAliases(toolCall.name) : [];
const toolCallsToTry: FunctionCall[] = [];
for (const name of toolNamesToTry) {
toolCallsToTry.push({ ...toolCall, name });
- if (serverName && !name.includes('__')) {
- toolCallsToTry.push({
- ...toolCall,
- name: `${serverName}__${name}`,
- });
- }
}
for (const rule of this.rules) {
@@ -654,145 +644,69 @@ export class PolicyEngine {
allToolNames?: Set,
): Set {
const excludedTools = new Set();
- const processedTools = new Set();
- let globalVerdict: PolicyDecision | undefined;
- for (const rule of this.rules) {
- if (rule.argsPattern) {
- if (rule.toolName && rule.decision !== PolicyDecision.DENY) {
- processedTools.add(rule.toolName);
- }
- continue;
- }
-
- // Check if rule applies to current approval mode
- if (rule.modes && rule.modes.length > 0) {
- if (!rule.modes.includes(this.approvalMode)) {
- continue;
- }
- }
-
- // Handle annotation-based rules
- if (rule.toolAnnotations) {
- if (!toolMetadata) {
- // Without metadata, we can't evaluate annotation rules — skip (conservative fallback)
- continue;
- }
- // Iterate over all known tools and check if their annotations match this rule
- for (const [toolName, annotations] of toolMetadata) {
- if (processedTools.has(toolName)) {
- continue;
- }
- // Check if annotations match the rule's toolAnnotations (partial match)
- let annotationsMatch = true;
- for (const [key, value] of Object.entries(rule.toolAnnotations)) {
- if (annotations[key] !== value) {
- annotationsMatch = false;
- break;
- }
- }
- if (!annotationsMatch) {
- continue;
- }
- // Check if the tool name matches the rule's toolName pattern (if any)
- if (rule.toolName) {
- if (isWildcardPattern(rule.toolName)) {
- // For composite patterns (e.g. "*__*"), construct a qualified
- // name from metadata so matchesWildcard can resolve it.
- const rawServerName = annotations['_serverName'];
- const serverName =
- typeof rawServerName === 'string' ? rawServerName : undefined;
- const qualifiedName =
- serverName && !toolName.includes('__')
- ? `${serverName}__${toolName}`
- : toolName;
- if (!matchesWildcard(rule.toolName, qualifiedName, undefined)) {
- continue;
- }
- } else if (toolName !== rule.toolName) {
- continue;
- }
- }
- // Determine decision considering global verdict
- let decision: PolicyDecision;
- if (globalVerdict !== undefined) {
- decision = globalVerdict;
- } else {
- decision = rule.decision;
- }
- if (decision === PolicyDecision.DENY) {
- excludedTools.add(toolName);
- }
- processedTools.add(toolName);
- }
- continue;
- }
-
- // Handle Global Rules
- if (!rule.toolName) {
- if (globalVerdict === undefined) {
- globalVerdict = rule.decision;
- if (globalVerdict !== PolicyDecision.DENY) {
- // Global ALLOW/ASK found.
- // Since rules are sorted by priority, this overrides any lower-priority rules.
- // We can stop processing because nothing else will be excluded.
- break;
- }
- // If Global DENY, we continue to find specific tools to add to excluded set
- }
- continue;
- }
-
- const toolName = rule.toolName;
-
- // Check if already processed (exact match)
- if (processedTools.has(toolName)) {
- continue;
- }
-
- // Check if covered by a processed wildcard
- let coveredByWildcard = false;
- for (const processed of processedTools) {
- if (
- isWildcardPattern(processed) &&
- matchesWildcard(processed, toolName, undefined)
- ) {
- // It's covered by a higher-priority wildcard rule.
- // If that wildcard rule resulted in exclusion, this tool should also be excluded.
- if (excludedTools.has(processed)) {
- excludedTools.add(toolName);
- }
- coveredByWildcard = true;
- break;
- }
- }
- if (coveredByWildcard) {
- continue;
- }
-
- processedTools.add(toolName);
-
- // Determine decision
- let decision: PolicyDecision;
- if (globalVerdict !== undefined) {
- decision = globalVerdict;
- } else {
- decision = rule.decision;
- }
-
- if (decision === PolicyDecision.DENY) {
- excludedTools.add(toolName);
- }
+ if (!allToolNames) {
+ return excludedTools;
}
- // If there's a global DENY and we know all tool names, exclude any tool
- // that wasn't explicitly allowed by a higher-priority rule.
- if (globalVerdict === PolicyDecision.DENY && allToolNames) {
- for (const name of allToolNames) {
- if (!processedTools.has(name)) {
- excludedTools.add(name);
+ for (const toolName of allToolNames) {
+ const annotations = toolMetadata?.get(toolName);
+ const serverName = isMcpToolAnnotation(annotations)
+ ? annotations._serverName
+ : undefined;
+
+ let staticallyExcluded = false;
+ let matchFound = false;
+
+ // Evaluate rules in priority order (they are already sorted in constructor)
+ for (const rule of this.rules) {
+ // Create a copy of the rule without argsPattern to see if it targets the tool
+ // regardless of the runtime arguments it might receive.
+ const ruleWithoutArgs: PolicyRule = { ...rule, argsPattern: undefined };
+ const toolCall: FunctionCall = { name: toolName, args: {} };
+
+ const appliesToTool = ruleMatches(
+ ruleWithoutArgs,
+ toolCall,
+ undefined, // stringifiedArgs
+ serverName,
+ this.approvalMode,
+ annotations,
+ );
+
+ if (appliesToTool) {
+ if (rule.argsPattern) {
+ // Exclusions only apply statically before arguments are known.
+ if (rule.decision !== PolicyDecision.DENY) {
+ // Conditionally allowed/asked based on args. Therefore NOT statically excluded.
+ staticallyExcluded = false;
+ matchFound = true;
+ break;
+ }
+ // If it's conditionally DENIED based on args, it means it's not unconditionally denied.
+ // We must keep evaluating lower priority rules to see the default/unconditional state.
+ continue;
+ } else {
+ // Unconditional rule for this tool
+ const decision = this.applyNonInteractiveMode(rule.decision);
+ staticallyExcluded = decision === PolicyDecision.DENY;
+ matchFound = true;
+ break;
+ }
}
}
+
+ if (!matchFound) {
+ // Fallback to default decision if no rule matches
+ const defaultDec = this.applyNonInteractiveMode(this.defaultDecision);
+ if (defaultDec === PolicyDecision.DENY) {
+ staticallyExcluded = true;
+ }
+ }
+
+ if (staticallyExcluded) {
+ excludedTools.add(toolName);
+ }
}
return excludedTools;
diff --git a/packages/core/src/policy/toml-loader.test.ts b/packages/core/src/policy/toml-loader.test.ts
index a65248cfea..959f09ba80 100644
--- a/packages/core/src/policy/toml-loader.test.ts
+++ b/packages/core/src/policy/toml-loader.test.ts
@@ -17,8 +17,8 @@ import { fileURLToPath } from 'node:url';
import {
loadPoliciesFromToml,
validateMcpPolicyToolNames,
+ type PolicyLoadResult,
} from './toml-loader.js';
-import type { PolicyLoadResult } from './toml-loader.js';
import { PolicyEngine } from './policy-engine.js';
const __filename = fileURLToPath(import.meta.url);
@@ -129,7 +129,7 @@ priority = 10
`);
expect(result.rules).toHaveLength(1);
- expect(result.rules[0].toolName).toBe('*__*');
+ expect(result.rules[0].toolName).toBe('mcp_*');
expect(result.rules[0].decision).toBe(PolicyDecision.ASK_USER);
expect(result.errors).toHaveLength(0);
});
@@ -144,7 +144,7 @@ priority = 10
`);
expect(result.rules).toHaveLength(1);
- expect(result.rules[0].toolName).toBe('*__search');
+ expect(result.rules[0].toolName).toBe('mcp_*_search');
expect(result.errors).toHaveLength(0);
});
@@ -215,8 +215,12 @@ priority = 100
`);
expect(result.rules).toHaveLength(2);
- expect(result.rules[0].toolName).toBe('google-workspace__calendar.list');
- expect(result.rules[1].toolName).toBe('google-workspace__calendar.get');
+ expect(result.rules[0].toolName).toBe(
+ 'mcp_google-workspace_calendar.list',
+ );
+ expect(result.rules[1].toolName).toBe(
+ 'mcp_google-workspace_calendar.get',
+ );
expect(result.errors).toHaveLength(0);
});
@@ -678,12 +682,12 @@ priority = 100
it('should not warn for MCP format tool names', async () => {
const result = await runLoadPoliciesFromToml(`
[[rule]]
-toolName = "my-server__my-tool"
+toolName = "mcp_my-server_my-tool"
decision = "allow"
priority = 100
[[rule]]
-toolName = "my-server__*"
+toolName = "mcp_my-server_*"
decision = "allow"
priority = 100
`);
@@ -822,7 +826,7 @@ priority = 100
annotationRule,
'Should have loaded a rule with toolAnnotations',
).toBeDefined();
- expect(annotationRule!.toolName).toBe('*__*');
+ expect(annotationRule!.toolName).toBe('mcp_*');
expect(annotationRule!.toolAnnotations).toEqual({
readOnlyHint: true,
});
@@ -863,7 +867,7 @@ priority = 100
// 4. MCP tool WITHOUT annotations should be DENIED
const denyResult = await engine.check(
- { name: 'github__create_issue' },
+ { name: 'mcp_github_create_issue' },
'github',
undefined,
);
@@ -874,7 +878,7 @@ priority = 100
// 5. MCP tool with readOnlyHint=false should also be DENIED
const denyResult2 = await engine.check(
- { name: 'github__delete_issue' },
+ { name: 'mcp_github_delete_issue' },
'github',
{ readOnlyHint: false },
);
@@ -883,9 +887,9 @@ priority = 100
'MCP tool with readOnlyHint=false should be DENIED in Plan Mode',
).toBe(PolicyDecision.DENY);
- // 6. Test with qualified tool name format (server__tool) but no separate serverName
+ // 6. Test with qualified tool name format (mcp_server_tool) but no separate serverName
const qualifiedResult = await engine.check(
- { name: 'github__list_repos' },
+ { name: 'mcp_github_list_repos' },
undefined,
{ readOnlyHint: true },
);
@@ -990,7 +994,8 @@ priority = 100
['people.getMe', 'calendar.list', 'calendar.get'],
[
{
- toolName: 'google-workspace__people.getxMe',
+ toolName: 'mcp_google-workspace_people.getxMe',
+ mcpName: 'google-workspace',
source: 'User: workspace.toml',
},
],
@@ -1007,8 +1012,14 @@ priority = 100
'google-workspace',
['people.getMe', 'calendar.list'],
[
- { toolName: 'google-workspace__people.getMe' },
- { toolName: 'google-workspace__calendar.list' },
+ {
+ toolName: 'mcp_google-workspace_people.getMe',
+ mcpName: 'google-workspace',
+ },
+ {
+ toolName: 'mcp_google-workspace_calendar.list',
+ mcpName: 'google-workspace',
+ },
],
);
@@ -1019,7 +1030,7 @@ priority = 100
const warnings = validateMcpPolicyToolNames(
'my-server',
['tool1', 'tool2'],
- [{ toolName: 'my-server__*' }],
+ [{ toolName: 'mcp_my-server_*', mcpName: 'my-server' }],
);
expect(warnings).toHaveLength(0);
@@ -1029,7 +1040,7 @@ priority = 100
const warnings = validateMcpPolicyToolNames(
'server-a',
['tool1'],
- [{ toolName: 'server-b__toolx' }],
+ [{ toolName: 'mcp_server-b_toolx', mcpName: 'server-b' }],
);
expect(warnings).toHaveLength(0);
@@ -1039,7 +1050,12 @@ priority = 100
const warnings = validateMcpPolicyToolNames(
'my-server',
['tool1', 'tool2'],
- [{ toolName: 'my-server__completely_different_name' }],
+ [
+ {
+ toolName: 'mcp_my-server_completely_different_name',
+ mcpName: 'my-server',
+ },
+ ],
);
expect(warnings).toHaveLength(0);
@@ -1059,7 +1075,13 @@ priority = 100
const warnings = validateMcpPolicyToolNames(
'my-server',
['tool1'],
- [{ toolName: 'my-server__tol1', source: 'User: custom.toml' }],
+ [
+ {
+ toolName: 'mcp_my-server_tol1',
+ mcpName: 'my-server',
+ source: 'User: custom.toml',
+ },
+ ],
);
expect(warnings).toHaveLength(1);
diff --git a/packages/core/src/policy/toml-loader.ts b/packages/core/src/policy/toml-loader.ts
index d2a24aa100..c91930a21d 100644
--- a/packages/core/src/policy/toml-loader.ts
+++ b/packages/core/src/policy/toml-loader.ts
@@ -24,6 +24,7 @@ import path from 'node:path';
import toml from '@iarna/toml';
import { z, type ZodError } from 'zod';
import { isNodeError } from '../utils/errors.js';
+import { MCP_TOOL_PREFIX, formatMcpToolName } from '../tools/mcp-tool.js';
/**
* Maximum Levenshtein distance to consider a name a likely typo of a built-in tool.
@@ -262,11 +263,15 @@ function validateShellCommandSyntax(
* tool name, or null if valid or not close to any built-in name.
*/
function validateToolName(name: string, ruleIndex: number): string | null {
+ if (name.includes('__')) {
+ return `Rule #${ruleIndex + 1}: The "__" syntax for MCP tools is strictly deprecated. Please use the 'mcpName = "..."' property or the 'mcp_server_tool' format instead.`;
+ }
+
// A name that looks like an MCP tool (e.g., "re__ad") could be a typo of a
// built-in tool ("read_file"). We should let such names fall through to the
// Levenshtein distance check below. Non-MCP-like names that are valid can
// be safely skipped.
- if (isValidToolName(name, { allowWildcards: true }) && !name.includes('__')) {
+ if (isValidToolName(name, { allowWildcards: true })) {
return null;
}
@@ -402,8 +407,8 @@ export async function loadPoliciesFromToml(
// Validate tool names in rules
for (let i = 0; i < tomlRules.length; i++) {
const rule = tomlRules[i];
- // Skip MCP-scoped rules — MCP tool names are server-defined and dynamic
- if (rule.mcpName) continue;
+ // We no longer skip MCP-scoped rules because we need to specifically
+ // warn users if they use deprecated "__" syntax for MCP tool names
const toolNames: string[] = rule.toolName
? Array.isArray(rule.toolName)
@@ -447,18 +452,19 @@ export async function loadPoliciesFromToml(
// Create a policy rule for each tool name
return toolNames.map((toolName) => {
- // Transform mcpName field to composite toolName format
- let effectiveToolName: string | undefined;
- if (rule.mcpName && toolName) {
- effectiveToolName = `${rule.mcpName}__${toolName}`;
- } else if (rule.mcpName) {
- effectiveToolName = `${rule.mcpName}__*`;
- } else {
- effectiveToolName = toolName;
+ let effectiveToolName: string | undefined = toolName;
+ const mcpName = rule.mcpName;
+
+ if (mcpName) {
+ effectiveToolName = formatMcpToolName(
+ mcpName,
+ effectiveToolName,
+ );
}
const policyRule: PolicyRule = {
toolName: effectiveToolName,
+ mcpName: rule.mcpName,
decision: rule.decision,
priority: transformPriority(rule.priority, tier),
modes: rule.modes,
@@ -563,15 +569,16 @@ export async function loadPoliciesFromToml(
return toolNames.map((toolName) => {
let effectiveToolName: string | undefined;
if (checker.mcpName && toolName) {
- effectiveToolName = `${checker.mcpName}__${toolName}`;
+ effectiveToolName = `${MCP_TOOL_PREFIX}${checker.mcpName}_${toolName}`;
} else if (checker.mcpName) {
- effectiveToolName = `${checker.mcpName}__*`;
+ effectiveToolName = `${MCP_TOOL_PREFIX}${checker.mcpName}_*`;
} else {
effectiveToolName = toolName;
}
const safetyCheckerRule: SafetyCheckerRule = {
toolName: effectiveToolName,
+ mcpName: checker.mcpName,
priority: transformPriority(checker.priority, tier),
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
checker: checker.checker as SafetyCheckerConfig,
@@ -655,16 +662,28 @@ export async function loadPoliciesFromToml(
export function validateMcpPolicyToolNames(
serverName: string,
discoveredToolNames: string[],
- policyRules: ReadonlyArray<{ toolName?: string; source?: string }>,
+ policyRules: ReadonlyArray<{
+ toolName?: string;
+ mcpName?: string;
+ source?: string;
+ }>,
): string[] {
- const prefix = `${serverName}__`;
+ const prefix = `${MCP_TOOL_PREFIX}${serverName}_`;
const warnings: string[] = [];
for (const rule of policyRules) {
if (!rule.toolName) continue;
- if (!rule.toolName.startsWith(prefix)) continue;
- const toolPart = rule.toolName.slice(prefix.length);
+ let toolPart: string | undefined;
+
+ // The toolName is typically transformed into an FQN if mcpName was used.
+ if (rule.mcpName === serverName && rule.toolName.startsWith(prefix)) {
+ toolPart = rule.toolName.slice(prefix.length);
+ } else if (rule.toolName.startsWith(prefix)) {
+ toolPart = rule.toolName.slice(prefix.length);
+ } else {
+ continue;
+ }
// Skip wildcards
if (toolPart === '*') continue;
diff --git a/packages/core/src/policy/types.ts b/packages/core/src/policy/types.ts
index 18c621c176..f59821b093 100644
--- a/packages/core/src/policy/types.ts
+++ b/packages/core/src/policy/types.ts
@@ -110,6 +110,13 @@ export interface PolicyRule {
*/
toolName?: string;
+ /**
+ * Identifies the MCP server this rule applies to.
+ * Enables precise rule matching against `serverName` metadata instead
+ * of parsing composite string names.
+ */
+ mcpName?: string;
+
/**
* Pattern to match against tool arguments.
* Can be used for more fine-grained control.
@@ -166,6 +173,11 @@ export interface SafetyCheckerRule {
*/
toolName?: string;
+ /**
+ * Identifies the MCP server this rule applies to.
+ */
+ mcpName?: string;
+
/**
* Pattern to match against tool arguments.
* Can be used for more fine-grained control.
diff --git a/packages/core/src/prompts/promptProvider.test.ts b/packages/core/src/prompts/promptProvider.test.ts
index b74f159e4f..2d96dee7ef 100644
--- a/packages/core/src/prompts/promptProvider.test.ts
+++ b/packages/core/src/prompts/promptProvider.test.ts
@@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { PromptProvider } from './promptProvider.js';
import type { Config } from '../config/config.js';
import {
@@ -35,6 +35,9 @@ describe('PromptProvider', () => {
beforeEach(() => {
vi.resetAllMocks();
+ vi.stubEnv('GEMINI_SYSTEM_MD', '');
+ vi.stubEnv('GEMINI_WRITE_SYSTEM_MD', '');
+
mockConfig = {
getToolRegistry: vi.fn().mockReturnValue({
getAllToolNames: vi.fn().mockReturnValue([]),
@@ -56,9 +59,14 @@ describe('PromptProvider', () => {
}),
getApprovedPlanPath: vi.fn().mockReturnValue(undefined),
getApprovalMode: vi.fn(),
+ isTrackerEnabled: vi.fn().mockReturnValue(false),
} as unknown as Config;
});
+ afterEach(() => {
+ vi.unstubAllEnvs();
+ });
+
it('should handle multiple context filenames in the system prompt', () => {
vi.mocked(getAllGeminiMdFilenames).mockReturnValue([
DEFAULT_CONTEXT_FILENAME,
@@ -153,7 +161,7 @@ describe('PromptProvider', () => {
const provider = new PromptProvider();
const prompt = provider.getCoreSystemPrompt(mockConfig);
- expect(prompt).toContain('`mcp_read` (my-mcp-server)');
+ expect(prompt).toContain('`mcp_my-mcp-server_mcp_read` (my-mcp-server)');
});
it('should include write constraint message in plan mode prompt', () => {
@@ -176,4 +184,41 @@ describe('PromptProvider', () => {
expect(prompt).toContain('/tmp/project-temp/plans/');
});
});
+
+ describe('getCompressionPrompt', () => {
+ it('should include plan preservation instructions when an approved plan path is provided', () => {
+ const planPath = '/path/to/plan.md';
+ (
+ mockConfig.getApprovedPlanPath as ReturnType
+ ).mockReturnValue(planPath);
+
+ const provider = new PromptProvider();
+ const prompt = provider.getCompressionPrompt(mockConfig);
+
+ expect(prompt).toContain('### APPROVED PLAN PRESERVATION');
+ expect(prompt).toContain(planPath);
+
+ // Verify it's BEFORE the structure example
+ const structureMarker = 'The structure MUST be as follows:';
+ const planPreservationMarker = '### APPROVED PLAN PRESERVATION';
+
+ const structureIndex = prompt.indexOf(structureMarker);
+ const planPreservationIndex = prompt.indexOf(planPreservationMarker);
+
+ expect(planPreservationIndex).toBeGreaterThan(-1);
+ expect(structureIndex).toBeGreaterThan(-1);
+ expect(planPreservationIndex).toBeLessThan(structureIndex);
+ });
+
+ it('should NOT include plan preservation instructions when no approved plan path is provided', () => {
+ (
+ mockConfig.getApprovedPlanPath as ReturnType
+ ).mockReturnValue(undefined);
+
+ const provider = new PromptProvider();
+ const prompt = provider.getCompressionPrompt(mockConfig);
+
+ expect(prompt).not.toContain('### APPROVED PLAN PRESERVATION');
+ });
+ });
});
diff --git a/packages/core/src/prompts/promptProvider.ts b/packages/core/src/prompts/promptProvider.ts
index 9b8759c2af..01dbd8d4d4 100644
--- a/packages/core/src/prompts/promptProvider.ts
+++ b/packages/core/src/prompts/promptProvider.ts
@@ -159,6 +159,7 @@ export class PromptProvider {
approvedPlan: approvedPlanPath
? { path: approvedPlanPath }
: undefined,
+ taskTracker: config.isTrackerEnabled(),
}),
!isPlanMode,
),
@@ -168,9 +169,11 @@ export class PromptProvider {
planModeToolsList,
plansDir: config.storage.getPlansDir(),
approvedPlanPath: config.getApprovedPlanPath(),
+ taskTracker: config.isTrackerEnabled(),
}),
isPlanMode,
),
+ taskTracker: config.isTrackerEnabled(),
operationalGuidelines: this.withSection(
'operationalGuidelines',
() => ({
@@ -231,7 +234,7 @@ export class PromptProvider {
);
const isModernModel = supportsModernFeatures(desiredModel);
const activeSnippets = isModernModel ? snippets : legacySnippets;
- return activeSnippets.getCompressionPrompt();
+ return activeSnippets.getCompressionPrompt(config.getApprovedPlanPath());
}
private withSection(
diff --git a/packages/core/src/prompts/snippets.legacy.ts b/packages/core/src/prompts/snippets.legacy.ts
index 3671490089..227b06be45 100644
--- a/packages/core/src/prompts/snippets.legacy.ts
+++ b/packages/core/src/prompts/snippets.legacy.ts
@@ -15,6 +15,7 @@ import {
GREP_TOOL_NAME,
MEMORY_TOOL_NAME,
READ_FILE_TOOL_NAME,
+ SHELL_PARAM_IS_BACKGROUND,
SHELL_TOOL_NAME,
WRITE_FILE_TOOL_NAME,
WRITE_TODOS_TOOL_NAME,
@@ -599,12 +600,12 @@ function toolUsageInteractive(
interactiveShellEnabled: boolean,
): string {
if (interactive) {
- const ctrlF = interactiveShellEnabled
- ? ' If you choose to execute an interactive command consider letting the user know they can press `ctrl + f` to focus into the shell to provide input.'
+ const focusHint = interactiveShellEnabled
+ ? ' If you choose to execute an interactive command consider letting the user know they can press `tab` to focus into the shell to provide input.'
: '';
return `
-- **Background Processes:** To run a command in the background, set the \`is_background\` parameter to true. If unsure, ask the user.
-- **Interactive Commands:** Always prefer non-interactive commands (e.g., using 'run once' or 'CI' flags for test runners to avoid persistent watch modes or 'git --no-pager') unless a persistent process is specifically required; however, some commands are only interactive and expect user input during their execution (e.g. ssh, vim).${ctrlF}`;
+ - **Background Processes:** To run a command in the background, set the \`${SHELL_PARAM_IS_BACKGROUND}\` parameter to true.
+ - **Interactive Commands:** Always prefer non-interactive commands (e.g., using 'run once' or 'CI' flags for test runners to avoid persistent watch modes or 'git --no-pager') unless a persistent process is specifically required; however, some commands are only interactive and expect user input during their execution (e.g. ssh, vim).${focusHint}`;
}
return `
- **Background Processes:** To run a command in the background, set the \`is_background\` parameter to true.
diff --git a/packages/core/src/prompts/snippets.ts b/packages/core/src/prompts/snippets.ts
index bebd3c9146..3b3334f96b 100644
--- a/packages/core/src/prompts/snippets.ts
+++ b/packages/core/src/prompts/snippets.ts
@@ -27,6 +27,9 @@ import {
READ_FILE_PARAM_END_LINE,
SHELL_PARAM_IS_BACKGROUND,
EDIT_PARAM_OLD_STRING,
+ TRACKER_CREATE_TASK_TOOL_NAME,
+ TRACKER_LIST_TASKS_TOOL_NAME,
+ TRACKER_UPDATE_TASK_TOOL_NAME,
} from '../tools/tool-names.js';
import type { HierarchicalMemory } from '../config/memory.js';
import { DEFAULT_CONTEXT_FILENAME } from '../tools/memoryTool.js';
@@ -41,6 +44,7 @@ export interface SystemPromptOptions {
hookContext?: boolean;
primaryWorkflows?: PrimaryWorkflowsOptions;
planningWorkflow?: PlanningWorkflowOptions;
+ taskTracker?: boolean;
operationalGuidelines?: OperationalGuidelinesOptions;
sandbox?: SandboxMode;
interactiveYoloMode?: boolean;
@@ -66,6 +70,7 @@ export interface PrimaryWorkflowsOptions {
enableGrep: boolean;
enableGlob: boolean;
approvedPlan?: { path: string };
+ taskTracker?: boolean;
}
export interface OperationalGuidelinesOptions {
@@ -83,6 +88,7 @@ export interface PlanningWorkflowOptions {
planModeToolsList: string;
plansDir: string;
approvedPlanPath?: string;
+ taskTracker?: boolean;
}
export interface AgentSkillOptions {
@@ -120,6 +126,8 @@ ${
: renderPrimaryWorkflows(options.primaryWorkflows)
}
+${options.taskTracker ? renderTaskTracker() : ''}
+
${renderOperationalGuidelines(options.operationalGuidelines)}
${renderInteractiveYoloMode(options.interactiveYoloMode)}
@@ -243,6 +251,8 @@ Operate as a **strategic orchestrator**. Your own context window is your most pr
When you delegate, the sub-agent's entire execution is consolidated into a single summary in your history, keeping your main loop lean.
+**Concurrency Safety and Mandate:** You should NEVER run multiple subagents in a single turn if their abilities mutate the same files or resources. This is to prevent race conditions and ensure that the workspace is in a consistent state. Only run multiple subagents in parallel when their tasks are independent (e.g., multiple concurrent research or read-only tasks) or if parallel execution is explicitly requested by the user.
+
**High-Impact Delegation Candidates:**
- **Repetitive Batch Tasks:** Tasks involving more than 3 files or repeated steps (e.g., "Add license headers to all files in src/", "Fix all lint errors in the project").
- **High-Volume Output:** Commands or tools expected to return large amounts of data (e.g., verbose builds, exhaustive file searches).
@@ -464,6 +474,24 @@ ${trimmed}
return `\n---\n\n\n${sections.join('\n')}\n`;
}
+export function renderTaskTracker(): string {
+ const trackerCreate = formatToolName(TRACKER_CREATE_TASK_TOOL_NAME);
+ const trackerList = formatToolName(TRACKER_LIST_TASKS_TOOL_NAME);
+ const trackerUpdate = formatToolName(TRACKER_UPDATE_TASK_TOOL_NAME);
+
+ return `
+# TASK MANAGEMENT PROTOCOL
+You are operating with a persistent file-based task tracking system located at \`.tracker/tasks/\`. You must adhere to the following rules:
+
+1. **NO IN-MEMORY LISTS**: Do not maintain a mental list of tasks or write markdown checkboxes in the chat. Use the provided tools (${trackerCreate}, ${trackerList}, ${trackerUpdate}) for all state management.
+2. **IMMEDIATE DECOMPOSITION**: Upon receiving a task, evaluate its functional complexity and scope. If the request involves more than a single atomic modification, or necessitates research before execution, you MUST immediately decompose it into discrete entries using ${trackerCreate}.
+3. **IGNORE FORMATTING BIAS**: Trigger the protocol based on the **objective complexity** of the goal, regardless of whether the user provided a structured list or a single block of text/paragraph. "Paragraph-style" goals that imply multiple actions are multi-step projects and MUST be tracked.
+4. **PLAN MODE INTEGRATION**: If an approved plan exists, you MUST use the ${trackerCreate} tool to decompose it into discrete tasks before writing any code. Maintain a bidirectional understanding between the plan document and the task graph.
+5. **VERIFICATION**: Before marking a task as complete, verify the work is actually done (e.g., run the test, check the file existence).
+6. **STATE OVER CHAT**: If the user says "I think we finished that," but the tool says it is 'pending', trust the tool--or verify explicitly before updating.
+7. **DEPENDENCY MANAGEMENT**: Respect task topology. Never attempt to execute a task if its dependencies are not marked as 'closed'. If you are blocked, focus only on the leaf nodes of the task graph.`.trim();
+}
+
export function renderPlanningWorkflow(
options?: PlanningWorkflowOptions,
): string {
@@ -580,6 +608,10 @@ function workflowStepResearch(options: PrimaryWorkflowsOptions): string {
}
function workflowStepStrategy(options: PrimaryWorkflowsOptions): string {
+ if (options.approvedPlan && options.taskTracker) {
+ return `2. **Strategy:** An approved plan is available for this task. Treat this file as your single source of truth and invoke the task tracker tool to create tasks for this plan. You MUST read this file before proceeding. If you discover new requirements or need to change the approach, confirm with the user and update this plan file to reflect the updated design decisions or discovered requirements. Make sure to update the tracker task list based on this updated plan. Once all implementation and verification steps are finished, provide a **final summary** of the work completed against the plan and offer clear **next steps** to the user (e.g., 'Open a pull request').`;
+ }
+
if (options.approvedPlan) {
return `2. **Strategy:** An approved plan is available for this task. Treat this file as your single source of truth. You MUST read this file before proceeding. If you discover new requirements or need to change the approach, confirm with the user and update this plan file to reflect the updated design decisions or discovered requirements. Once all implementation and verification steps are finished, provide a **final summary** of the work completed against the plan and offer clear **next steps** to the user (e.g., 'Open a pull request').`;
}
@@ -665,12 +697,12 @@ function toolUsageInteractive(
interactiveShellEnabled: boolean,
): string {
if (interactive) {
- const ctrlF = interactiveShellEnabled
- ? ' If you choose to execute an interactive command consider letting the user know they can press `ctrl + f` to focus into the shell to provide input.'
+ const focusHint = interactiveShellEnabled
+ ? ' If you choose to execute an interactive command consider letting the user know they can press `tab` to focus into the shell to provide input.'
: '';
return `
- **Background Processes:** To run a command in the background, set the \`${SHELL_PARAM_IS_BACKGROUND}\` parameter to true. If unsure, ask the user.
-- **Interactive Commands:** Always prefer non-interactive commands (e.g., using 'run once' or 'CI' flags for test runners to avoid persistent watch modes or 'git --no-pager') unless a persistent process is specifically required; however, some commands are only interactive and expect user input during their execution (e.g. ssh, vim).${ctrlF}`;
+- **Interactive Commands:** Always prefer non-interactive commands (e.g., using 'run once' or 'CI' flags for test runners to avoid persistent watch modes or 'git --no-pager') unless a persistent process is specifically required; however, some commands are only interactive and expect user input during their execution (e.g. ssh, vim).${focusHint}`;
}
return `
- **Background Processes:** To run a command in the background, set the \`${SHELL_PARAM_IS_BACKGROUND}\` parameter to true.
@@ -702,7 +734,17 @@ function formatToolName(name: string): string {
/**
* Provides the system prompt for history compression.
*/
-export function getCompressionPrompt(): string {
+export function getCompressionPrompt(approvedPlanPath?: string): string {
+ const planPreservation = approvedPlanPath
+ ? `
+
+### APPROVED PLAN PRESERVATION
+An approved implementation plan exists at ${approvedPlanPath}. You MUST preserve the following in your snapshot:
+- The plan's file path in
+- Completion status of each plan step in (mark as [DONE], [IN PROGRESS], or [TODO])
+- Any user feedback or modifications to the plan in `
+ : '';
+
return `
You are a specialized system component responsible for distilling chat history into a structured XML .
@@ -718,7 +760,7 @@ When the conversation history grows too large, you will be invoked to distill th
First, you will think through the entire history in a private . Review the user's overall goal, the agent's actions, tool outputs, file modifications, and any unresolved questions. Identify every piece of information for future actions.
-After your reasoning is complete, generate the final XML object. Be incredibly dense with information. Omit any irrelevant conversational filler.
+After your reasoning is complete, generate the final XML object. Be incredibly dense with information. Omit any irrelevant conversational filler.${planPreservation}
The structure MUST be as follows:
diff --git a/packages/core/src/resources/resource-registry.ts b/packages/core/src/resources/resource-registry.ts
index 1c2c754504..ce30456df5 100644
--- a/packages/core/src/resources/resource-registry.ts
+++ b/packages/core/src/resources/resource-registry.ts
@@ -69,4 +69,17 @@ export class ResourceRegistry {
clear(): void {
this.resources.clear();
}
+
+ /**
+ * Returns an array of resources registered from a specific MCP server.
+ */
+ getResourcesByServer(serverName: string): MCPResource[] {
+ const serverResources: MCPResource[] = [];
+ for (const resource of this.resources.values()) {
+ if (resource.serverName === serverName) {
+ serverResources.push(resource);
+ }
+ }
+ return serverResources.sort((a, b) => a.uri.localeCompare(b.uri));
+ }
}
diff --git a/packages/core/src/routing/strategies/gemmaClassifierStrategy.test.ts b/packages/core/src/routing/strategies/gemmaClassifierStrategy.test.ts
index 9425208fd7..967a185eaf 100644
--- a/packages/core/src/routing/strategies/gemmaClassifierStrategy.test.ts
+++ b/packages/core/src/routing/strategies/gemmaClassifierStrategy.test.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { Mock } from 'vitest';
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { GemmaClassifierStrategy } from './gemmaClassifierStrategy.js';
import type { RoutingContext } from '../routingStrategy.js';
import type { Config } from '../../config/config.js';
diff --git a/packages/core/src/safety/built-in.test.ts b/packages/core/src/safety/built-in.test.ts
index d940929009..ecfc8e6bd5 100644
--- a/packages/core/src/safety/built-in.test.ts
+++ b/packages/core/src/safety/built-in.test.ts
@@ -9,8 +9,7 @@ import * as fs from 'node:fs/promises';
import * as os from 'node:os';
import * as path from 'node:path';
import { AllowedPathChecker } from './built-in.js';
-import type { SafetyCheckInput } from './protocol.js';
-import { SafetyCheckDecision } from './protocol.js';
+import { SafetyCheckDecision, type SafetyCheckInput } from './protocol.js';
import type { FunctionCall } from '@google/genai';
describe('AllowedPathChecker', () => {
diff --git a/packages/core/src/safety/built-in.ts b/packages/core/src/safety/built-in.ts
index 540af36290..aae8c8ee53 100644
--- a/packages/core/src/safety/built-in.ts
+++ b/packages/core/src/safety/built-in.ts
@@ -6,8 +6,11 @@
import * as path from 'node:path';
import * as fs from 'node:fs';
-import type { SafetyCheckInput, SafetyCheckResult } from './protocol.js';
-import { SafetyCheckDecision } from './protocol.js';
+import {
+ SafetyCheckDecision,
+ type SafetyCheckInput,
+ type SafetyCheckResult,
+} from './protocol.js';
import type { AllowedPathConfig } from '../policy/types.js';
/**
diff --git a/packages/core/src/safety/checker-runner.test.ts b/packages/core/src/safety/checker-runner.test.ts
index cd3c0e18ba..6358541ecf 100644
--- a/packages/core/src/safety/checker-runner.test.ts
+++ b/packages/core/src/safety/checker-runner.test.ts
@@ -13,8 +13,7 @@ import {
type InProcessCheckerConfig,
InProcessCheckerType,
} from '../policy/types.js';
-import type { SafetyCheckResult } from './protocol.js';
-import { SafetyCheckDecision } from './protocol.js';
+import { SafetyCheckDecision, type SafetyCheckResult } from './protocol.js';
import type { Config } from '../config/config.js';
// Mock dependencies
diff --git a/packages/core/src/safety/checker-runner.ts b/packages/core/src/safety/checker-runner.ts
index a46c3e6dbd..c0ed57aa20 100644
--- a/packages/core/src/safety/checker-runner.ts
+++ b/packages/core/src/safety/checker-runner.ts
@@ -11,8 +11,11 @@ import type {
InProcessCheckerConfig,
ExternalCheckerConfig,
} from '../policy/types.js';
-import type { SafetyCheckInput, SafetyCheckResult } from './protocol.js';
-import { SafetyCheckDecision } from './protocol.js';
+import {
+ SafetyCheckDecision,
+ type SafetyCheckInput,
+ type SafetyCheckResult,
+} from './protocol.js';
import type { CheckerRegistry } from './registry.js';
import type { ContextBuilder } from './context-builder.js';
import { z } from 'zod';
diff --git a/packages/core/src/safety/conseca/conseca.test.ts b/packages/core/src/safety/conseca/conseca.test.ts
index 8d871777de..2ad9ef3295 100644
--- a/packages/core/src/safety/conseca/conseca.test.ts
+++ b/packages/core/src/safety/conseca/conseca.test.ts
@@ -6,8 +6,7 @@
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { ConsecaSafetyChecker } from './conseca.js';
-import { SafetyCheckDecision } from '../protocol.js';
-import type { SafetyCheckInput } from '../protocol.js';
+import { SafetyCheckDecision, type SafetyCheckInput } from '../protocol.js';
import {
logConsecaPolicyGeneration,
logConsecaVerdict,
diff --git a/packages/core/src/safety/conseca/conseca.ts b/packages/core/src/safety/conseca/conseca.ts
index 4d837bbc47..3964911796 100644
--- a/packages/core/src/safety/conseca/conseca.ts
+++ b/packages/core/src/safety/conseca/conseca.ts
@@ -5,8 +5,11 @@
*/
import type { InProcessChecker } from '../built-in.js';
-import type { SafetyCheckInput, SafetyCheckResult } from '../protocol.js';
-import { SafetyCheckDecision } from '../protocol.js';
+import {
+ SafetyCheckDecision,
+ type SafetyCheckInput,
+ type SafetyCheckResult,
+} from '../protocol.js';
import {
logConsecaPolicyGeneration,
diff --git a/packages/core/src/scheduler/confirmation.test.ts b/packages/core/src/scheduler/confirmation.test.ts
index e9e55e807d..abd07ba86e 100644
--- a/packages/core/src/scheduler/confirmation.test.ts
+++ b/packages/core/src/scheduler/confirmation.test.ts
@@ -28,8 +28,11 @@ import {
} from '../tools/tools.js';
import type { SchedulerStateManager } from './state-manager.js';
import type { ToolModificationHandler } from './tool-modifier.js';
-import type { ValidatingToolCall, WaitingToolCall } from './types.js';
-import { ROOT_SCHEDULER_ID } from './types.js';
+import {
+ ROOT_SCHEDULER_ID,
+ type ValidatingToolCall,
+ type WaitingToolCall,
+} from './types.js';
import type { Config } from '../config/config.js';
import { type EditorType } from '../utils/editor.js';
import { randomUUID } from 'node:crypto';
diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts
index be79b7c62d..05f5b08a2f 100644
--- a/packages/core/src/scheduler/policy.test.ts
+++ b/packages/core/src/scheduler/policy.test.ts
@@ -25,16 +25,16 @@ import {
type ToolExecuteConfirmationDetails,
type AnyToolInvocation,
} from '../tools/tools.js';
-import type {
- ValidatingToolCall,
- ToolCallRequestInfo,
- CompletedToolCall,
+import {
+ ROOT_SCHEDULER_ID,
+ type ValidatingToolCall,
+ type ToolCallRequestInfo,
+ type CompletedToolCall,
} from './types.js';
import type { PolicyEngine } from '../policy/policy-engine.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { CoreToolScheduler } from '../core/coreToolScheduler.js';
import { Scheduler } from './scheduler.js';
-import { ROOT_SCHEDULER_ID } from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
import type { ToolRegistry } from '../tools/tool-registry.js';
diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts
index b2c1adade0..ee5438c319 100644
--- a/packages/core/src/scheduler/scheduler.test.ts
+++ b/packages/core/src/scheduler/scheduler.test.ts
@@ -75,19 +75,20 @@ import {
type AnyDeclarativeTool,
type AnyToolInvocation,
} from '../tools/tools.js';
-import type {
- ToolCallRequestInfo,
- ValidatingToolCall,
- SuccessfulToolCall,
- ErroredToolCall,
- CancelledToolCall,
- CompletedToolCall,
- ToolCallResponseInfo,
- ExecutingToolCall,
- Status,
- ToolCall,
+import {
+ CoreToolCallStatus,
+ ROOT_SCHEDULER_ID,
+ type ToolCallRequestInfo,
+ type ValidatingToolCall,
+ type SuccessfulToolCall,
+ type ErroredToolCall,
+ type CancelledToolCall,
+ type CompletedToolCall,
+ type ToolCallResponseInfo,
+ type ExecutingToolCall,
+ type Status,
+ type ToolCall,
} from './types.js';
-import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
import * as ToolUtils from '../utils/tool-utils.js';
@@ -946,7 +947,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
CoreToolCallStatus.Cancelled,
- 'Operation cancelled',
+ { callId: 'call-1', responseParts: [] },
);
});
diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts
index 58e4586887..38e001ea90 100644
--- a/packages/core/src/scheduler/scheduler.ts
+++ b/packages/core/src/scheduler/scheduler.ts
@@ -24,8 +24,7 @@ import {
type ScheduledToolCall,
} from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
-import type { ApprovalMode } from '../policy/types.js';
-import { PolicyDecision } from '../policy/types.js';
+import { PolicyDecision, type ApprovalMode } from '../policy/types.js';
import {
ToolConfirmationOutcome,
type AnyDeclarativeTool,
@@ -741,7 +740,7 @@ export class Scheduler {
this.state.updateStatus(
callId,
CoreToolCallStatus.Cancelled,
- 'Operation cancelled',
+ result.response,
);
} else {
this.state.updateStatus(
diff --git a/packages/core/src/scheduler/scheduler_parallel.test.ts b/packages/core/src/scheduler/scheduler_parallel.test.ts
index 9633784323..56e6e26243 100644
--- a/packages/core/src/scheduler/scheduler_parallel.test.ts
+++ b/packages/core/src/scheduler/scheduler_parallel.test.ts
@@ -72,14 +72,14 @@ import {
type AnyToolInvocation,
Kind,
} from '../tools/tools.js';
-import type {
- ToolCallRequestInfo,
- CompletedToolCall,
- SuccessfulToolCall,
- Status,
- ToolCall,
+import {
+ ROOT_SCHEDULER_ID,
+ type ToolCallRequestInfo,
+ type CompletedToolCall,
+ type SuccessfulToolCall,
+ type Status,
+ type ToolCall,
} from './types.js';
-import { ROOT_SCHEDULER_ID } from './types.js';
import { GeminiCliOperation } from '../telemetry/constants.js';
import type { EditorType } from '../utils/editor.js';
diff --git a/packages/core/src/scheduler/state-manager.test.ts b/packages/core/src/scheduler/state-manager.test.ts
index b27e51de8f..dd5071c5bf 100644
--- a/packages/core/src/scheduler/state-manager.test.ts
+++ b/packages/core/src/scheduler/state-manager.test.ts
@@ -6,17 +6,18 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SchedulerStateManager } from './state-manager.js';
-import type {
- ValidatingToolCall,
- WaitingToolCall,
- SuccessfulToolCall,
- ErroredToolCall,
- CancelledToolCall,
- ExecutingToolCall,
- ToolCallRequestInfo,
- ToolCallResponseInfo,
+import {
+ CoreToolCallStatus,
+ ROOT_SCHEDULER_ID,
+ type ValidatingToolCall,
+ type WaitingToolCall,
+ type SuccessfulToolCall,
+ type ErroredToolCall,
+ type CancelledToolCall,
+ type ExecutingToolCall,
+ type ToolCallRequestInfo,
+ type ToolCallResponseInfo,
} from './types.js';
-import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import {
ToolConfirmationOutcome,
type AnyDeclarativeTool,
diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts
index b14b492e4b..428b7f87a8 100644
--- a/packages/core/src/scheduler/state-manager.ts
+++ b/packages/core/src/scheduler/state-manager.ts
@@ -4,20 +4,21 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type {
- ToolCall,
- Status,
- WaitingToolCall,
- CompletedToolCall,
- SuccessfulToolCall,
- ErroredToolCall,
- CancelledToolCall,
- ScheduledToolCall,
- ValidatingToolCall,
- ExecutingToolCall,
- ToolCallResponseInfo,
+import {
+ CoreToolCallStatus,
+ ROOT_SCHEDULER_ID,
+ type ToolCall,
+ type Status,
+ type WaitingToolCall,
+ type CompletedToolCall,
+ type SuccessfulToolCall,
+ type ErroredToolCall,
+ type CancelledToolCall,
+ type ScheduledToolCall,
+ type ValidatingToolCall,
+ type ExecutingToolCall,
+ type ToolCallResponseInfo,
} from './types.js';
-import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import type {
ToolConfirmationOutcome,
ToolResultDisplay,
@@ -30,6 +31,7 @@ import {
MessageBusType,
type SerializableConfirmationDetails,
} from '../confirmation-bus/types.js';
+import { isToolCallResponseInfo } from '../utils/tool-utils.js';
/**
* Handler for terminal tool calls.
@@ -127,7 +129,7 @@ export class SchedulerStateManager {
updateStatus(
callId: string,
status: CoreToolCallStatus.Cancelled,
- data: string,
+ data: string | ToolCallResponseInfo,
): void;
updateStatus(
callId: string,
@@ -264,7 +266,7 @@ export class SchedulerStateManager {
): ToolCall {
switch (newStatus) {
case CoreToolCallStatus.Success: {
- if (!this.isToolCallResponseInfo(auxiliaryData)) {
+ if (!isToolCallResponseInfo(auxiliaryData)) {
throw new Error(
`Invalid data for 'success' transition (callId: ${call.request.callId})`,
);
@@ -272,7 +274,7 @@ export class SchedulerStateManager {
return this.toSuccess(call, auxiliaryData);
}
case CoreToolCallStatus.Error: {
- if (!this.isToolCallResponseInfo(auxiliaryData)) {
+ if (!isToolCallResponseInfo(auxiliaryData)) {
throw new Error(
`Invalid data for 'error' transition (callId: ${call.request.callId})`,
);
@@ -290,9 +292,12 @@ export class SchedulerStateManager {
case CoreToolCallStatus.Scheduled:
return this.toScheduled(call);
case CoreToolCallStatus.Cancelled: {
- if (typeof auxiliaryData !== 'string') {
+ if (
+ typeof auxiliaryData !== 'string' &&
+ !isToolCallResponseInfo(auxiliaryData)
+ ) {
throw new Error(
- `Invalid reason (string) for 'cancelled' transition (callId: ${call.request.callId})`,
+ `Invalid reason (string) or response for 'cancelled' transition (callId: ${call.request.callId})`,
);
}
return this.toCancelled(call, auxiliaryData);
@@ -317,15 +322,6 @@ export class SchedulerStateManager {
}
}
- private isToolCallResponseInfo(data: unknown): data is ToolCallResponseInfo {
- return (
- typeof data === 'object' &&
- data !== null &&
- 'callId' in data &&
- 'responseParts' in data
- );
- }
-
private isExecutingToolCallPatch(
data: unknown,
): data is Partial {
@@ -451,7 +447,10 @@ export class SchedulerStateManager {
};
}
- private toCancelled(call: ToolCall, reason: string): CancelledToolCall {
+ private toCancelled(
+ call: ToolCall,
+ reason: string | ToolCallResponseInfo,
+ ): CancelledToolCall {
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Cancelled);
const startTime = 'startTime' in call ? call.startTime : undefined;
@@ -478,6 +477,31 @@ export class SchedulerStateManager {
}
}
+ // Capture any existing live output so it isn't lost when forcing cancellation.
+ let existingOutput: ToolResultDisplay | undefined = undefined;
+ if (call.status === CoreToolCallStatus.Executing && call.liveOutput) {
+ existingOutput = call.liveOutput;
+ }
+
+ if (isToolCallResponseInfo(reason)) {
+ const finalResponse = { ...reason };
+ if (!finalResponse.resultDisplay) {
+ finalResponse.resultDisplay = resultDisplay ?? existingOutput;
+ }
+
+ return {
+ request: call.request,
+ tool: call.tool,
+ invocation: call.invocation,
+ status: CoreToolCallStatus.Cancelled,
+ response: finalResponse,
+ durationMs: startTime ? Date.now() - startTime : undefined,
+ outcome: call.outcome,
+ schedulerId: call.schedulerId,
+ approvalMode: call.approvalMode,
+ };
+ }
+
const errorMessage = `[Operation Cancelled] Reason: ${reason}`;
return {
request: call.request,
@@ -495,7 +519,7 @@ export class SchedulerStateManager {
},
},
],
- resultDisplay,
+ resultDisplay: resultDisplay ?? existingOutput,
error: undefined,
errorType: undefined,
contentLength: errorMessage.length,
diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts
index 3049d13928..1211a75c22 100644
--- a/packages/core/src/scheduler/tool-executor.test.ts
+++ b/packages/core/src/scheduler/tool-executor.test.ts
@@ -13,8 +13,7 @@ import {
} from '../index.js';
import { makeFakeConfig } from '../test-utils/config.js';
import { MockTool } from '../test-utils/mock-tool.js';
-import type { ScheduledToolCall } from './types.js';
-import { CoreToolCallStatus } from './types.js';
+import { CoreToolCallStatus, type ScheduledToolCall } from './types.js';
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import type { CallableTool } from '@google/genai';
@@ -608,4 +607,113 @@ describe('ToolExecutor', () => {
}),
);
});
+
+ it('should return cancelled result with partial output when signal is aborted', async () => {
+ const mockTool = new MockTool({
+ name: 'slowTool',
+ });
+ const invocation = mockTool.build({});
+
+ const partialOutput = 'Some partial output before cancellation';
+ vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation(
+ async () => ({
+ llmContent: partialOutput,
+ returnDisplay: `[Cancelled] ${partialOutput}`,
+ }),
+ );
+
+ const scheduledCall: ScheduledToolCall = {
+ status: CoreToolCallStatus.Scheduled,
+ request: {
+ callId: 'call-cancel-partial',
+ name: 'slowTool',
+ args: {},
+ isClientInitiated: false,
+ prompt_id: 'prompt-cancel',
+ },
+ tool: mockTool,
+ invocation: invocation as unknown as AnyToolInvocation,
+ startTime: Date.now(),
+ };
+
+ const controller = new AbortController();
+ controller.abort();
+
+ const result = await executor.execute({
+ call: scheduledCall,
+ signal: controller.signal,
+ onUpdateToolCall: vi.fn(),
+ });
+
+ expect(result.status).toBe(CoreToolCallStatus.Cancelled);
+ if (result.status === CoreToolCallStatus.Cancelled) {
+ const response = result.response.responseParts[0]?.functionResponse
+ ?.response as Record;
+ expect(response).toEqual({
+ error: '[Operation Cancelled] User cancelled tool execution.',
+ output: partialOutput,
+ });
+ expect(result.response.resultDisplay).toBe(
+ `[Cancelled] ${partialOutput}`,
+ );
+ }
+ });
+
+ it('should truncate large shell output even on cancellation', async () => {
+ // 1. Setup Config for Truncation
+ vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
+ vi.spyOn(config.storage, 'getProjectTempDir').mockReturnValue('/tmp');
+
+ const mockTool = new MockTool({ name: SHELL_TOOL_NAME });
+ const invocation = mockTool.build({});
+ const longOutput = 'This is a very long output that should be truncated.';
+
+ // 2. Mock execution returning long content
+ vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
+ llmContent: longOutput,
+ returnDisplay: longOutput,
+ });
+
+ const scheduledCall: ScheduledToolCall = {
+ status: CoreToolCallStatus.Scheduled,
+ request: {
+ callId: 'call-trunc-cancel',
+ name: SHELL_TOOL_NAME,
+ args: { command: 'echo long' },
+ isClientInitiated: false,
+ prompt_id: 'prompt-trunc-cancel',
+ },
+ tool: mockTool,
+ invocation: invocation as unknown as AnyToolInvocation,
+ startTime: Date.now(),
+ };
+
+ // 3. Abort immediately
+ const controller = new AbortController();
+ controller.abort();
+
+ // 4. Execute
+ const result = await executor.execute({
+ call: scheduledCall,
+ signal: controller.signal,
+ onUpdateToolCall: vi.fn(),
+ });
+
+ // 5. Verify Truncation Logic was applied in cancelled path
+ expect(fileUtils.saveTruncatedToolOutput).toHaveBeenCalledWith(
+ longOutput,
+ SHELL_TOOL_NAME,
+ 'call-trunc-cancel',
+ expect.any(String),
+ 'test-session-id',
+ );
+
+ expect(result.status).toBe(CoreToolCallStatus.Cancelled);
+ if (result.status === CoreToolCallStatus.Cancelled) {
+ const response = result.response.responseParts[0]?.functionResponse
+ ?.response as Record;
+ expect(response['output']).toBe('TruncatedContent...');
+ expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
+ }
+ });
});
diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts
index c293af0ac4..0521d211a7 100644
--- a/packages/core/src/scheduler/tool-executor.ts
+++ b/packages/core/src/scheduler/tool-executor.ts
@@ -6,23 +6,21 @@
import fsPromises from 'node:fs/promises';
import { debugLogger } from '../utils/debugLogger.js';
-import type {
- ToolCallRequestInfo,
- ToolCallResponseInfo,
- ToolResult,
- Config,
- ToolResultDisplay,
- ToolLiveOutput,
-} from '../index.js';
import {
ToolErrorType,
ToolOutputTruncatedEvent,
logToolOutputTruncated,
runInDevTraceSpan,
+ type ToolCallRequestInfo,
+ type ToolCallResponseInfo,
+ type ToolResult,
+ type ToolResultDisplay,
+ type Config,
+ type ToolLiveOutput,
} from '../index.js';
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
-import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { ShellToolInvocation } from '../tools/shell.js';
+import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
import {
saveTruncatedToolOutput,
@@ -30,15 +28,16 @@ import {
moveToolOutputToFile,
} from '../utils/fileUtils.js';
import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js';
-import type {
- CompletedToolCall,
- ToolCall,
- ExecutingToolCall,
- ErroredToolCall,
- SuccessfulToolCall,
- CancelledToolCall,
+import {
+ CoreToolCallStatus,
+ type CompletedToolCall,
+ type ToolCall,
+ type ExecutingToolCall,
+ type ErroredToolCall,
+ type SuccessfulToolCall,
+ type CancelledToolCall,
} from './types.js';
-import { CoreToolCallStatus } from './types.js';
+import type { PartListUnion, Part } from '@google/genai';
import {
GeminiCliOperation,
GEN_AI_TOOL_CALL_ID,
@@ -145,10 +144,10 @@ export class ToolExecutor {
);
});
}
- completedToolCall = this.createCancelledResult(
+ completedToolCall = await this.createCancelledResult(
call,
'User cancelled tool execution.',
- toolResult.returnDisplay,
+ toolResult,
);
} else if (toolResult.error === undefined) {
completedToolCall = await this.createSuccessResult(
@@ -186,7 +185,7 @@ export class ToolExecutor {
executionError.message.includes('Operation cancelled by user'));
if (signal.aborted || isAbortError) {
- completedToolCall = this.createCancelledResult(
+ completedToolCall = await this.createCancelledResult(
call,
'User cancelled tool execution.',
);
@@ -209,56 +208,13 @@ export class ToolExecutor {
);
}
- private createCancelledResult(
+ private async truncateOutputIfNeeded(
call: ToolCall,
- reason: string,
- resultDisplay?: ToolResultDisplay,
- ): CancelledToolCall {
- const errorMessage = `[Operation Cancelled] ${reason}`;
- const startTime = 'startTime' in call ? call.startTime : undefined;
-
- if (!('tool' in call) || !('invocation' in call)) {
- // This should effectively never happen in execution phase, but we handle
- // it safely
- throw new Error('Cancelled tool call missing tool/invocation references');
- }
-
- return {
- status: CoreToolCallStatus.Cancelled,
- request: call.request,
- response: {
- callId: call.request.callId,
- responseParts: [
- {
- functionResponse: {
- id: call.request.callId,
- name: call.request.name,
- response: { error: errorMessage },
- },
- },
- ],
- resultDisplay,
- error: undefined,
- errorType: undefined,
- contentLength: errorMessage.length,
- },
- tool: call.tool,
- invocation: call.invocation,
- durationMs: startTime ? Date.now() - startTime : undefined,
- startTime,
- endTime: Date.now(),
- outcome: call.outcome,
- };
- }
-
- private async createSuccessResult(
- call: ToolCall,
- toolResult: ToolResult,
- ): Promise {
- let content = toolResult.llmContent;
- let outputFile: string | undefined;
- const toolName = call.request.originalRequestName || call.request.name;
+ content: PartListUnion,
+ ): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> {
+ const toolName = call.request.name;
const callId = call.request.callId;
+ let outputFile: string | undefined;
if (toolResult.fullOutputFilePath) {
const threshold = this.config.getTruncateToolOutputThreshold();
@@ -300,17 +256,23 @@ export class ToolExecutor {
this.config.getSessionId(),
);
outputFile = savedPath;
- content = formatTruncatedToolOutput(content, outputFile, threshold);
+ const truncatedContent = formatTruncatedToolOutput(
+ content,
+ outputFile,
+ threshold,
+ );
logToolOutputTruncated(
this.config,
new ToolOutputTruncatedEvent(call.request.prompt_id, {
toolName,
originalContentLength,
- truncatedContentLength: content.length,
+ truncatedContentLength: truncatedContent.length,
threshold,
}),
);
+
+ return { truncatedContent, outputFile };
}
} else if (
Array.isArray(content) &&
@@ -338,7 +300,12 @@ export class ToolExecutor {
outputFile,
threshold,
);
- content[0] = { ...firstPart, text: truncatedText };
+
+ // We need to return a NEW array to avoid mutating the original toolResult if it matters,
+ // though here we are creating the response so it's probably fine to mutate or return new.
+ const truncatedContent: Part[] = [
+ { ...firstPart, text: truncatedText },
+ ];
logToolOutputTruncated(
this.config,
@@ -349,10 +316,95 @@ export class ToolExecutor {
threshold,
}),
);
+
+ return { truncatedContent, outputFile };
}
}
}
+ return { truncatedContent: content, outputFile };
+ }
+
+ private async createCancelledResult(
+ call: ToolCall,
+ reason: string,
+ toolResult?: ToolResult,
+ ): Promise {
+ const errorMessage = `[Operation Cancelled] ${reason}`;
+ const startTime = 'startTime' in call ? call.startTime : undefined;
+
+ if (!('tool' in call) || !('invocation' in call)) {
+ // This should effectively never happen in execution phase, but we handle
+ // it safely
+ throw new Error('Cancelled tool call missing tool/invocation references');
+ }
+
+ let responseParts: Part[] = [];
+ let outputFile: string | undefined;
+
+ if (toolResult?.llmContent) {
+ // Attempt to truncate and save output if we have content, even in cancellation case
+ // This is to handle cases where the tool may have produced output before cancellation
+ const { truncatedContent: output, outputFile: truncatedOutputFile } =
+ await this.truncateOutputIfNeeded(call, toolResult?.llmContent);
+
+ outputFile = truncatedOutputFile;
+ responseParts = convertToFunctionResponse(
+ call.request.name,
+ call.request.callId,
+ output,
+ this.config.getActiveModel(),
+ );
+
+ // Inject the cancellation error into the response object
+ const mainPart = responseParts[0];
+ if (mainPart?.functionResponse?.response) {
+ const respObj = mainPart.functionResponse.response;
+ respObj['error'] = errorMessage;
+ }
+ } else {
+ responseParts = [
+ {
+ functionResponse: {
+ id: call.request.callId,
+ name: call.request.name,
+ response: { error: errorMessage },
+ },
+ },
+ ];
+ }
+
+ return {
+ status: CoreToolCallStatus.Cancelled,
+ request: call.request,
+ response: {
+ callId: call.request.callId,
+ responseParts,
+ resultDisplay: toolResult?.returnDisplay,
+ error: undefined,
+ errorType: undefined,
+ outputFile,
+ contentLength: JSON.stringify(responseParts).length,
+ },
+ tool: call.tool,
+ invocation: call.invocation,
+ durationMs: startTime ? Date.now() - startTime : undefined,
+ startTime,
+ endTime: Date.now(),
+ outcome: call.outcome,
+ };
+ }
+
+ private async createSuccessResult(
+ call: ToolCall,
+ toolResult: ToolResult,
+ ): Promise {
+ const { truncatedContent: content, outputFile } =
+ await this.truncateOutputIfNeeded(call, toolResult.llmContent);
+
+ const toolName = call.request.originalRequestName || call.request.name;
+ const callId = call.request.callId;
+
const response = convertToFunctionResponse(
toolName,
callId,
diff --git a/packages/core/src/scheduler/tool-modifier.test.ts b/packages/core/src/scheduler/tool-modifier.test.ts
index 35ff2cd79c..98be4098c4 100644
--- a/packages/core/src/scheduler/tool-modifier.test.ts
+++ b/packages/core/src/scheduler/tool-modifier.test.ts
@@ -4,11 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { ToolModificationHandler } from './tool-modifier.js';
-import type { WaitingToolCall, ToolCallRequestInfo } from './types.js';
-import { CoreToolCallStatus } from './types.js';
+import {
+ CoreToolCallStatus,
+ type WaitingToolCall,
+ type ToolCallRequestInfo,
+} from './types.js';
import * as modifiableToolModule from '../tools/modifiable-tool.js';
+import type { ModifyContext } from '../tools/modifiable-tool.js';
import * as Diff from 'diff';
import { MockModifiableTool, MockTool } from '../test-utils/mock-tool.js';
import type {
@@ -16,8 +20,6 @@ import type {
ToolInvocation,
ToolConfirmationPayload,
} from '../tools/tools.js';
-import type { ModifyContext } from '../tools/modifiable-tool.js';
-import type { Mock } from 'vitest';
// Mock the modules that export functions we need to control
vi.mock('diff', () => ({
diff --git a/packages/core/src/services/FolderTrustDiscoveryService.ts b/packages/core/src/services/FolderTrustDiscoveryService.ts
index e81273af22..bdf5d76297 100644
--- a/packages/core/src/services/FolderTrustDiscoveryService.ts
+++ b/packages/core/src/services/FolderTrustDiscoveryService.ts
@@ -132,7 +132,11 @@ export class FolderTrustDiscoveryService {
for (const event of Object.values(hooksConfig)) {
if (!Array.isArray(event)) continue;
for (const hook of event) {
- if (this.isRecord(hook) && typeof hook['command'] === 'string') {
+ if (
+ this.isRecord(hook) &&
+ // eslint-disable-next-line no-restricted-syntax
+ typeof hook['command'] === 'string'
+ ) {
hooks.add(hook['command']);
}
}
diff --git a/packages/core/src/services/chatCompressionService.test.ts b/packages/core/src/services/chatCompressionService.test.ts
index 4ddd38e25c..7ae9549a25 100644
--- a/packages/core/src/services/chatCompressionService.test.ts
+++ b/packages/core/src/services/chatCompressionService.test.ts
@@ -10,14 +10,15 @@ import {
findCompressSplitPoint,
modelStringToModelConfigAlias,
} from './chatCompressionService.js';
-import type { Content, GenerateContentResponse } from '@google/genai';
+import type { Content, GenerateContentResponse, Part } from '@google/genai';
import { CompressionStatus } from '../core/turn.js';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
import type { GeminiChat } from '../core/geminiChat.js';
import type { Config } from '../config/config.js';
import * as fileUtils from '../utils/fileUtils.js';
-import { TOOL_OUTPUTS_DIR } from '../utils/fileUtils.js';
import { getInitialChatHistory } from '../utils/environmentContext.js';
+
+const { TOOL_OUTPUTS_DIR } = fileUtils;
import * as tokenCalculation from '../utils/tokenCalculation.js';
import { tokenLimit } from '../core/tokenLimits.js';
import os from 'node:os';
@@ -188,6 +189,7 @@ describe('ChatCompressionService', () => {
storage: {
getProjectTempDir: vi.fn().mockReturnValue(testTempDir),
},
+ getApprovedPlanPath: vi.fn().mockReturnValue('/path/to/plan.md'),
} as unknown as Config;
vi.mocked(getInitialChatHistory).mockImplementation(
@@ -354,6 +356,63 @@ describe('ChatCompressionService', () => {
);
});
+ it('should include the approved plan path in the system instruction', async () => {
+ const planPath = '/custom/plan/path.md';
+ vi.mocked(mockConfig.getApprovedPlanPath).mockReturnValue(planPath);
+ vi.mocked(mockConfig.getActiveModel).mockReturnValue(
+ 'gemini-3.1-pro-preview',
+ );
+
+ const history: Content[] = [
+ { role: 'user', parts: [{ text: 'msg1' }] },
+ { role: 'model', parts: [{ text: 'msg2' }] },
+ ];
+ vi.mocked(mockChat.getHistory).mockReturnValue(history);
+ vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000);
+
+ await service.compress(
+ mockChat,
+ mockPromptId,
+ false,
+ mockModel,
+ mockConfig,
+ false,
+ );
+
+ const firstCallText = (
+ vi.mocked(mockConfig.getBaseLlmClient().generateContent).mock.calls[0][0]
+ .systemInstruction as Part
+ ).text;
+ expect(firstCallText).toContain('### APPROVED PLAN PRESERVATION');
+ expect(firstCallText).toContain(planPath);
+ });
+
+ it('should not include the approved plan section if no approved plan path exists', async () => {
+ vi.mocked(mockConfig.getApprovedPlanPath).mockReturnValue(undefined);
+
+ const history: Content[] = [
+ { role: 'user', parts: [{ text: 'msg1' }] },
+ { role: 'model', parts: [{ text: 'msg2' }] },
+ ];
+ vi.mocked(mockChat.getHistory).mockReturnValue(history);
+ vi.mocked(mockChat.getLastPromptTokenCount).mockReturnValue(600000);
+
+ await service.compress(
+ mockChat,
+ mockPromptId,
+ false,
+ mockModel,
+ mockConfig,
+ false,
+ );
+
+ const firstCallText = (
+ vi.mocked(mockConfig.getBaseLlmClient().generateContent).mock.calls[0][0]
+ .systemInstruction as Part
+ ).text;
+ expect(firstCallText).not.toContain('### APPROVED PLAN PRESERVATION');
+ });
+
it('should force compress even if under threshold', async () => {
const history: Content[] = [
{ role: 'user', parts: [{ text: 'msg1' }] },
diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts
index 5303a1a82a..8dceb18f4b 100644
--- a/packages/core/src/services/chatCompressionService.ts
+++ b/packages/core/src/services/chatCompressionService.ts
@@ -156,11 +156,13 @@ async function truncateHistoryToBudget(
} else if (responseObj && typeof responseObj === 'object') {
if (
'output' in responseObj &&
+ // eslint-disable-next-line no-restricted-syntax
typeof responseObj['output'] === 'string'
) {
contentStr = responseObj['output'];
} else if (
'content' in responseObj &&
+ // eslint-disable-next-line no-restricted-syntax
typeof responseObj['content'] === 'string'
) {
contentStr = responseObj['content'];
diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts
index 50a363a1db..2b8e8f1977 100644
--- a/packages/core/src/services/chatRecordingService.test.ts
+++ b/packages/core/src/services/chatRecordingService.test.ts
@@ -8,14 +8,14 @@ import { expect, it, describe, vi, beforeEach, afterEach } from 'vitest';
import fs from 'node:fs';
import path from 'node:path';
import os from 'node:os';
-import type {
- ConversationRecord,
- ToolCallRecord,
- MessageRecord,
+import {
+ ChatRecordingService,
+ type ConversationRecord,
+ type ToolCallRecord,
+ type MessageRecord,
} from './chatRecordingService.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import type { Content, Part } from '@google/genai';
-import { ChatRecordingService } from './chatRecordingService.js';
import type { Config } from '../config/config.js';
import { getProjectHash } from '../utils/paths.js';
@@ -245,6 +245,97 @@ describe('ChatRecordingService', () => {
tool: 0,
});
});
+
+ it('should not write to disk when queuing tokens (no last gemini message)', () => {
+ const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+
+ // Clear spy call count after initialize writes the initial file
+ writeFileSyncSpy.mockClear();
+
+ // No gemini message recorded yet, so tokens should only be queued
+ chatRecordingService.recordMessageTokens({
+ promptTokenCount: 5,
+ candidatesTokenCount: 10,
+ totalTokenCount: 15,
+ cachedContentTokenCount: 0,
+ });
+
+ // writeFileSync should NOT have been called since we only queued
+ expect(writeFileSyncSpy).not.toHaveBeenCalled();
+
+ // @ts-expect-error private property
+ expect(chatRecordingService.queuedTokens).toEqual({
+ input: 5,
+ output: 10,
+ total: 15,
+ cached: 0,
+ thoughts: 0,
+ tool: 0,
+ });
+
+ writeFileSyncSpy.mockRestore();
+ });
+
+ it('should not write to disk when queuing tokens (last message already has tokens)', () => {
+ chatRecordingService.recordMessage({
+ type: 'gemini',
+ content: 'Response',
+ model: 'gemini-pro',
+ });
+
+ // First recordMessageTokens updates the message and writes to disk
+ chatRecordingService.recordMessageTokens({
+ promptTokenCount: 1,
+ candidatesTokenCount: 1,
+ totalTokenCount: 2,
+ cachedContentTokenCount: 0,
+ });
+
+ const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+ writeFileSyncSpy.mockClear();
+
+ // Second call should only queue, NOT write to disk
+ chatRecordingService.recordMessageTokens({
+ promptTokenCount: 2,
+ candidatesTokenCount: 2,
+ totalTokenCount: 4,
+ cachedContentTokenCount: 0,
+ });
+
+ expect(writeFileSyncSpy).not.toHaveBeenCalled();
+ writeFileSyncSpy.mockRestore();
+ });
+
+ it('should use in-memory cache and not re-read from disk on subsequent operations', () => {
+ chatRecordingService.recordMessage({
+ type: 'gemini',
+ content: 'Response',
+ model: 'gemini-pro',
+ });
+
+ const readFileSyncSpy = vi.spyOn(fs, 'readFileSync');
+ readFileSyncSpy.mockClear();
+
+ // These operations should all use the in-memory cache
+ chatRecordingService.recordMessageTokens({
+ promptTokenCount: 1,
+ candidatesTokenCount: 1,
+ totalTokenCount: 2,
+ cachedContentTokenCount: 0,
+ });
+
+ chatRecordingService.recordMessage({
+ type: 'gemini',
+ content: 'Another response',
+ model: 'gemini-pro',
+ });
+
+ chatRecordingService.saveSummary('Test summary');
+
+ // readFileSync should NOT have been called since we use the in-memory cache
+ expect(readFileSyncSpy).not.toHaveBeenCalled();
+ readFileSyncSpy.mockRestore();
+ });
});
describe('recordToolCalls', () => {
@@ -769,6 +860,39 @@ describe('ChatRecordingService', () => {
expect(result[0].text).toBe('Prefix metadata or text');
expect(result[1].functionResponse!.id).toBe(callId);
});
+
+ it('should not write to disk when no tool calls match', () => {
+ chatRecordingService.recordMessage({
+ type: 'gemini',
+ content: 'Response with no tool calls',
+ model: 'gemini-pro',
+ });
+
+ const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+ writeFileSyncSpy.mockClear();
+
+ // History with a tool call ID that doesn't exist in the conversation
+ const history: Content[] = [
+ {
+ role: 'user',
+ parts: [
+ {
+ functionResponse: {
+ name: 'read_file',
+ id: 'nonexistent-call-id',
+ response: { output: 'some content' },
+ },
+ },
+ ],
+ },
+ ];
+
+ chatRecordingService.updateMessagesFromHistory(history);
+
+ // No tool calls matched, so writeFileSync should NOT have been called
+ expect(writeFileSyncSpy).not.toHaveBeenCalled();
+ writeFileSyncSpy.mockRestore();
+ });
});
describe('ENOENT (missing directory) handling', () => {
diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts
index 1748ccbe20..cd8d1e53c1 100644
--- a/packages/core/src/services/chatRecordingService.ts
+++ b/packages/core/src/services/chatRecordingService.ts
@@ -128,6 +128,7 @@ export interface ResumedSessionData {
export class ChatRecordingService {
private conversationFile: string | null = null;
private cachedLastConvData: string | null = null;
+ private cachedConversation: ConversationRecord | null = null;
private sessionId: string;
private projectHash: string;
private kind?: 'main' | 'subagent';
@@ -167,6 +168,7 @@ export class ChatRecordingService {
// Clear any cached data to force fresh reads
this.cachedLastConvData = null;
+ this.cachedConversation = null;
} else {
// Create new session
const chatsDir = path.join(
@@ -308,17 +310,19 @@ export class ChatRecordingService {
tool: respUsageMetadata.toolUsePromptTokenCount ?? 0,
total: respUsageMetadata.totalTokenCount ?? 0,
};
- this.updateConversation((conversation) => {
- const lastMsg = this.getLastMessage(conversation);
- // If the last message already has token info, it's because this new token info is for a
- // new message that hasn't been recorded yet.
- if (lastMsg && lastMsg.type === 'gemini' && !lastMsg.tokens) {
- lastMsg.tokens = tokens;
- this.queuedTokens = null;
- } else {
- this.queuedTokens = tokens;
- }
- });
+ const conversation = this.readConversation();
+ const lastMsg = this.getLastMessage(conversation);
+ // If the last message already has token info, it's because this new token info is for a
+ // new message that hasn't been recorded yet.
+ if (lastMsg && lastMsg.type === 'gemini' && !lastMsg.tokens) {
+ lastMsg.tokens = tokens;
+ this.queuedTokens = null;
+ this.writeConversation(conversation);
+ } else {
+ // Only queue tokens in memory; no disk I/O needed since the
+ // conversation record itself hasn't changed.
+ this.queuedTokens = tokens;
+ }
} catch (error) {
debugLogger.error(
'Error updating message tokens in chat history.',
@@ -427,12 +431,32 @@ export class ChatRecordingService {
/**
* Loads up the conversation record from disk.
+ *
+ * NOTE: The returned object is the live in-memory cache reference.
+ * Any mutations to it will be visible to all subsequent reads.
+ * Callers that mutate the result MUST call writeConversation() to
+ * persist the changes to disk.
*/
private readConversation(): ConversationRecord {
+ if (this.cachedConversation) {
+ return this.cachedConversation;
+ }
try {
this.cachedLastConvData = fs.readFileSync(this.conversationFile!, 'utf8');
- // eslint-disable-next-line @typescript-eslint/no-unsafe-return
- return JSON.parse(this.cachedLastConvData);
+ // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+ this.cachedConversation = JSON.parse(this.cachedLastConvData);
+ if (!this.cachedConversation) {
+ // File is corrupt or contains "null". Fallback to an empty conversation.
+ this.cachedConversation = {
+ sessionId: this.sessionId,
+ projectHash: this.projectHash,
+ startTime: new Date().toISOString(),
+ lastUpdated: new Date().toISOString(),
+ messages: [],
+ kind: this.kind,
+ };
+ }
+ return this.cachedConversation;
} catch (error) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
@@ -441,7 +465,7 @@ export class ChatRecordingService {
}
// Placeholder empty conversation if file doesn't exist.
- return {
+ this.cachedConversation = {
sessionId: this.sessionId,
projectHash: this.projectHash,
startTime: new Date().toISOString(),
@@ -449,6 +473,7 @@ export class ChatRecordingService {
messages: [],
kind: this.kind,
};
+ return this.cachedConversation;
}
}
@@ -464,15 +489,19 @@ export class ChatRecordingService {
// Don't write the file yet until there's at least one message.
if (conversation.messages.length === 0 && !allowEmpty) return;
- // Only write the file if this change would change the file.
- if (this.cachedLastConvData !== JSON.stringify(conversation, null, 2)) {
- conversation.lastUpdated = new Date().toISOString();
- const newContent = JSON.stringify(conversation, null, 2);
- this.cachedLastConvData = newContent;
- // Ensure directory exists before writing (handles cases where temp dir was cleaned)
- fs.mkdirSync(path.dirname(this.conversationFile), { recursive: true });
- fs.writeFileSync(this.conversationFile, newContent);
- }
+ const newContent = JSON.stringify(conversation, null, 2);
+ // Skip the disk write if nothing actually changed (e.g.
+ // updateMessagesFromHistory found no matching tool calls to update).
+ // Compare before updating lastUpdated so the timestamp doesn't
+ // cause a false diff.
+ if (this.cachedLastConvData === newContent) return;
+ this.cachedConversation = conversation;
+ conversation.lastUpdated = new Date().toISOString();
+ const contentToWrite = JSON.stringify(conversation, null, 2);
+ this.cachedLastConvData = contentToWrite;
+ // Ensure directory exists before writing (handles cases where temp dir was cleaned)
+ fs.mkdirSync(path.dirname(this.conversationFile), { recursive: true });
+ fs.writeFileSync(this.conversationFile, contentToWrite);
} catch (error) {
// Handle disk full (ENOSPC) gracefully - disable recording but allow conversation to continue
if (
@@ -482,6 +511,7 @@ export class ChatRecordingService {
(error as NodeJS.ErrnoException).code === 'ENOSPC'
) {
this.conversationFile = null;
+ this.cachedConversation = null;
debugLogger.warn(ENOSPC_WARNING_MESSAGE);
return; // Don't throw - allow the conversation to continue
}
diff --git a/packages/core/src/services/contextManager.test.ts b/packages/core/src/services/contextManager.test.ts
index 668a54fb56..945c9263f6 100644
--- a/packages/core/src/services/contextManager.test.ts
+++ b/packages/core/src/services/contextManager.test.ts
@@ -21,6 +21,7 @@ vi.mock('../utils/memoryDiscovery.js', async (importOriginal) => {
getEnvironmentMemoryPaths: vi.fn(),
readGeminiMdFiles: vi.fn(),
loadJitSubdirectoryMemory: vi.fn(),
+ deduplicatePathsByFileIdentity: vi.fn(),
concatenateInstructions: vi
.fn()
.mockImplementation(actual.concatenateInstructions),
@@ -33,7 +34,6 @@ describe('ContextManager', () => {
beforeEach(() => {
mockConfig = {
- getDebugMode: vi.fn().mockReturnValue(false),
getWorkingDir: vi.fn().mockReturnValue('/app'),
getImportFormat: vi.fn().mockReturnValue('tree'),
getWorkspaceContext: vi.fn().mockReturnValue({
@@ -52,6 +52,13 @@ describe('ContextManager', () => {
vi.clearAllMocks();
vi.spyOn(coreEvents, 'emit');
vi.mocked(memoryDiscovery.getExtensionMemoryPaths).mockReturnValue([]);
+ // default mock: deduplication returns paths as-is (no deduplication)
+ vi.mocked(
+ memoryDiscovery.deduplicatePathsByFileIdentity,
+ ).mockImplementation(async (paths: string[]) => ({
+ paths,
+ identityMap: new Map(),
+ }));
});
describe('refresh', () => {
@@ -74,13 +81,11 @@ describe('ContextManager', () => {
await contextManager.refresh();
expect(memoryDiscovery.getGlobalMemoryPaths).toHaveBeenCalled();
- expect(memoryDiscovery.getEnvironmentMemoryPaths).toHaveBeenCalledWith(
- ['/app'],
- false,
- );
+ expect(memoryDiscovery.getEnvironmentMemoryPaths).toHaveBeenCalledWith([
+ '/app',
+ ]);
expect(memoryDiscovery.readGeminiMdFiles).toHaveBeenCalledWith(
expect.arrayContaining([...globalPaths, ...envPaths]),
- false,
'tree',
);
@@ -128,6 +133,50 @@ describe('ContextManager', () => {
expect(contextManager.getEnvironmentMemory()).toBe('');
expect(contextManager.getGlobalMemory()).toContain('Global Content');
});
+
+ it('should deduplicate files by file identity in case-insensitive filesystems', async () => {
+ const globalPaths = ['/home/user/.gemini/GEMINI.md'];
+ const envPaths = ['/app/gemini.md', '/app/GEMINI.md'];
+
+ vi.mocked(memoryDiscovery.getGlobalMemoryPaths).mockResolvedValue(
+ globalPaths,
+ );
+ vi.mocked(memoryDiscovery.getEnvironmentMemoryPaths).mockResolvedValue(
+ envPaths,
+ );
+
+ // mock deduplication to return deduplicated paths (simulating same file)
+ vi.mocked(
+ memoryDiscovery.deduplicatePathsByFileIdentity,
+ ).mockResolvedValue({
+ paths: ['/home/user/.gemini/GEMINI.md', '/app/gemini.md'],
+ identityMap: new Map(),
+ });
+
+ vi.mocked(memoryDiscovery.readGeminiMdFiles).mockResolvedValue([
+ { filePath: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
+ { filePath: '/app/gemini.md', content: 'Project Content' },
+ ]);
+
+ await contextManager.refresh();
+
+ expect(
+ memoryDiscovery.deduplicatePathsByFileIdentity,
+ ).toHaveBeenCalledWith(
+ expect.arrayContaining([
+ '/home/user/.gemini/GEMINI.md',
+ '/app/gemini.md',
+ '/app/GEMINI.md',
+ ]),
+ );
+ expect(memoryDiscovery.readGeminiMdFiles).toHaveBeenCalledWith(
+ ['/home/user/.gemini/GEMINI.md', '/app/gemini.md'],
+ 'tree',
+ );
+ expect(contextManager.getEnvironmentMemory()).toContain(
+ 'Project Content',
+ );
+ });
});
describe('discoverContext', () => {
@@ -147,7 +196,7 @@ describe('ContextManager', () => {
'/app/src/file.ts',
['/app'],
expect.any(Set),
- false,
+ expect.any(Set),
);
expect(result).toMatch(/--- Context from: src[\\/]GEMINI\.md ---/);
expect(result).toContain('Src Content');
diff --git a/packages/core/src/services/contextManager.ts b/packages/core/src/services/contextManager.ts
index 1a33e24693..cec7c89ef9 100644
--- a/packages/core/src/services/contextManager.ts
+++ b/packages/core/src/services/contextManager.ts
@@ -13,12 +13,14 @@ import {
readGeminiMdFiles,
categorizeAndConcatenate,
type GeminiFileContent,
+ deduplicatePathsByFileIdentity,
} from '../utils/memoryDiscovery.js';
import type { Config } from '../config/config.js';
import { coreEvents, CoreEvent } from '../utils/events.js';
export class ContextManager {
private readonly loadedPaths: Set = new Set();
+ private readonly loadedFileIdentities: Set = new Set();
private readonly config: Config;
private globalMemory: string = '';
private extensionMemory: string = '';
@@ -33,49 +35,61 @@ export class ContextManager {
*/
async refresh(): Promise {
this.loadedPaths.clear();
- const debugMode = this.config.getDebugMode();
+ this.loadedFileIdentities.clear();
- const paths = await this.discoverMemoryPaths(debugMode);
- const contentsMap = await this.loadMemoryContents(paths, debugMode);
+ const paths = await this.discoverMemoryPaths();
+ const contentsMap = await this.loadMemoryContents(paths);
this.categorizeMemoryContents(paths, contentsMap);
this.emitMemoryChanged();
}
- private async discoverMemoryPaths(debugMode: boolean) {
+ private async discoverMemoryPaths() {
const [global, extension, project] = await Promise.all([
- getGlobalMemoryPaths(debugMode),
+ getGlobalMemoryPaths(),
Promise.resolve(
getExtensionMemoryPaths(this.config.getExtensionLoader()),
),
this.config.isTrustedFolder()
- ? getEnvironmentMemoryPaths(
- [...this.config.getWorkspaceContext().getDirectories()],
- debugMode,
- )
+ ? getEnvironmentMemoryPaths([
+ ...this.config.getWorkspaceContext().getDirectories(),
+ ])
: Promise.resolve([]),
]);
return { global, extension, project };
}
- private async loadMemoryContents(
- paths: { global: string[]; extension: string[]; project: string[] },
- debugMode: boolean,
- ) {
- const allPaths = Array.from(
+ private async loadMemoryContents(paths: {
+ global: string[];
+ extension: string[];
+ project: string[];
+ }) {
+ const allPathsStringDeduped = Array.from(
new Set([...paths.global, ...paths.extension, ...paths.project]),
);
+ // deduplicate by file identity to handle case-insensitive filesystems
+ const { paths: allPaths, identityMap: pathIdentityMap } =
+ await deduplicatePathsByFileIdentity(allPathsStringDeduped);
+
const allContents = await readGeminiMdFiles(
allPaths,
- debugMode,
this.config.getImportFormat(),
);
- this.markAsLoaded(
- allContents.filter((c) => c.content !== null).map((c) => c.filePath),
- );
+ const loadedFilePaths = allContents
+ .filter((c) => c.content !== null)
+ .map((c) => c.filePath);
+ this.markAsLoaded(loadedFilePaths);
+
+ // Cache file identities for performance optimization
+ for (const filePath of loadedFilePaths) {
+ const identity = pathIdentityMap.get(filePath);
+ if (identity) {
+ this.loadedFileIdentities.add(identity);
+ }
+ }
return new Map(allContents.map((c) => [c.filePath, c]));
}
@@ -123,14 +137,22 @@ export class ContextManager {
accessedPath,
trustedRoots,
this.loadedPaths,
- this.config.getDebugMode(),
+ this.loadedFileIdentities,
);
if (result.files.length === 0) {
return '';
}
- this.markAsLoaded(result.files.map((f) => f.path));
+ const newFilePaths = result.files.map((f) => f.path);
+ this.markAsLoaded(newFilePaths);
+
+ // Cache identities for newly loaded files
+ if (result.fileIdentities) {
+ for (const identity of result.fileIdentities) {
+ this.loadedFileIdentities.add(identity);
+ }
+ }
return concatenateInstructions(
result.files.map((f) => ({ filePath: f.path, content: f.content })),
this.config.getWorkingDir(),
diff --git a/packages/core/src/services/environmentSanitization.test.ts b/packages/core/src/services/environmentSanitization.test.ts
index cc26d7547d..63bb6ca5a5 100644
--- a/packages/core/src/services/environmentSanitization.test.ts
+++ b/packages/core/src/services/environmentSanitization.test.ts
@@ -32,6 +32,29 @@ describe('sanitizeEnvironment', () => {
expect(sanitized).toEqual(env);
});
+ it('should allow TERM and COLORTERM environment variables', () => {
+ const env = {
+ TERM: 'xterm-256color',
+ COLORTERM: 'truecolor',
+ };
+ const sanitized = sanitizeEnvironment(env, EMPTY_OPTIONS);
+ expect(sanitized).toEqual(env);
+ });
+
+ it('should preserve TERM and COLORTERM even in strict sanitization mode', () => {
+ const env = {
+ GITHUB_SHA: 'abc123',
+ TERM: 'xterm-256color',
+ COLORTERM: 'truecolor',
+ SOME_OTHER_VAR: 'value',
+ };
+ const sanitized = sanitizeEnvironment(env, EMPTY_OPTIONS);
+ expect(sanitized).toEqual({
+ TERM: 'xterm-256color',
+ COLORTERM: 'truecolor',
+ });
+ });
+
it('should allow variables prefixed with GEMINI_CLI_', () => {
const env = {
GEMINI_CLI_FOO: 'bar',
@@ -206,6 +229,48 @@ describe('sanitizeEnvironment', () => {
});
});
+ describe('value-first security: secret values must be caught even for allowed variable names', () => {
+ it('should redact ALWAYS_ALLOWED variables whose values contain a GitHub token', () => {
+ const env = {
+ HOME: 'ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
+ PATH: '/usr/bin',
+ };
+ const sanitized = sanitizeEnvironment(env, EMPTY_OPTIONS);
+ expect(sanitized).toEqual({ PATH: '/usr/bin' });
+ });
+
+ it('should redact ALWAYS_ALLOWED variables whose values contain a certificate', () => {
+ const env = {
+ SHELL:
+ '-----BEGIN RSA PRIVATE KEY-----\nMIIE...\n-----END RSA PRIVATE KEY-----',
+ USER: 'alice',
+ };
+ const sanitized = sanitizeEnvironment(env, EMPTY_OPTIONS);
+ expect(sanitized).toEqual({ USER: 'alice' });
+ });
+
+ it('should redact user-allowlisted variables whose values contain a secret', () => {
+ const env = {
+ MY_SAFE_VAR: 'ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
+ OTHER: 'fine',
+ };
+ const sanitized = sanitizeEnvironment(env, {
+ allowedEnvironmentVariables: ['MY_SAFE_VAR'],
+ blockedEnvironmentVariables: [],
+ enableEnvironmentVariableRedaction: true,
+ });
+ expect(sanitized).toEqual({ OTHER: 'fine' });
+ });
+
+ it('should NOT redact GEMINI_CLI_ variables even if their value looks like a secret (fully trusted)', () => {
+ const env = {
+ GEMINI_CLI_INTERNAL: 'ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
+ };
+ const sanitized = sanitizeEnvironment(env, EMPTY_OPTIONS);
+ expect(sanitized).toEqual(env);
+ });
+ });
+
it('should ensure all names in the sets are capitalized', () => {
for (const name of ALWAYS_ALLOWED_ENVIRONMENT_VARIABLES) {
expect(name).toBe(name.toUpperCase());
diff --git a/packages/core/src/services/environmentSanitization.ts b/packages/core/src/services/environmentSanitization.ts
index dc9c92484d..9d35249a8e 100644
--- a/packages/core/src/services/environmentSanitization.ts
+++ b/packages/core/src/services/environmentSanitization.ts
@@ -14,11 +14,9 @@ export function sanitizeEnvironment(
processEnv: NodeJS.ProcessEnv,
config: EnvironmentSanitizationConfig,
): NodeJS.ProcessEnv {
- // Enable strict sanitization in GitHub actions.
const isStrictSanitization =
!!processEnv['GITHUB_SHA'] || processEnv['SURFACE'] === 'Github';
- // Always sanitize when in GitHub actions.
if (!config.enableEnvironmentVariableRedaction && !isStrictSanitization) {
return { ...processEnv };
}
@@ -71,6 +69,10 @@ export const ALWAYS_ALLOWED_ENVIRONMENT_VARIABLES: ReadonlySet =
'TMPDIR',
'USER',
'LOGNAME',
+ // Terminal capability variables (needed by editors like vim/emacs and
+ // interactive commands like top)
+ 'TERM',
+ 'COLORTERM',
// GitHub Action-related variables
'ADDITIONAL_CONTEXT',
'AVAILABLE_LABELS',
@@ -148,7 +150,18 @@ function shouldRedactEnvironmentVariable(
key = key.toUpperCase();
value = value?.toUpperCase();
- // User overrides take precedence.
+ if (key.startsWith('GEMINI_CLI_')) {
+ return false;
+ }
+
+ if (value) {
+ for (const pattern of NEVER_ALLOWED_VALUE_PATTERNS) {
+ if (pattern.test(value)) {
+ return true;
+ }
+ }
+ }
+
if (allowedSet?.has(key)) {
return false;
}
@@ -156,20 +169,14 @@ function shouldRedactEnvironmentVariable(
return true;
}
- // These are never redacted.
- if (
- ALWAYS_ALLOWED_ENVIRONMENT_VARIABLES.has(key) ||
- key.startsWith('GEMINI_CLI_')
- ) {
+ if (ALWAYS_ALLOWED_ENVIRONMENT_VARIABLES.has(key)) {
return false;
}
- // These are always redacted.
if (NEVER_ALLOWED_ENVIRONMENT_VARIABLES.has(key)) {
return true;
}
- // If in strict mode (e.g. GitHub Action), and not explicitly allowed, redact it.
if (isStrictSanitization) {
return true;
}
@@ -180,14 +187,5 @@ function shouldRedactEnvironmentVariable(
}
}
- // Redact if the value looks like a key/cert.
- if (value) {
- for (const pattern of NEVER_ALLOWED_VALUE_PATTERNS) {
- if (pattern.test(value)) {
- return true;
- }
- }
- }
-
return false;
}
diff --git a/packages/core/src/services/fileDiscoveryService.ts b/packages/core/src/services/fileDiscoveryService.ts
index 44a28c1ff2..d816c42e31 100644
--- a/packages/core/src/services/fileDiscoveryService.ts
+++ b/packages/core/src/services/fileDiscoveryService.ts
@@ -4,10 +4,14 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { GitIgnoreFilter } from '../utils/gitIgnoreParser.js';
-import type { IgnoreFileFilter } from '../utils/ignoreFileParser.js';
-import { GitIgnoreParser } from '../utils/gitIgnoreParser.js';
-import { IgnoreFileParser } from '../utils/ignoreFileParser.js';
+import {
+ GitIgnoreParser,
+ type GitIgnoreFilter,
+} from '../utils/gitIgnoreParser.js';
+import {
+ IgnoreFileParser,
+ type IgnoreFileFilter,
+} from '../utils/ignoreFileParser.js';
import { isGitRepository } from '../utils/gitUtils.js';
import { GEMINI_IGNORE_FILE_NAME } from '../config/constants.js';
import fs from 'node:fs';
diff --git a/packages/core/src/services/gitService.ts b/packages/core/src/services/gitService.ts
index 2caad248ff..5409b1a526 100644
--- a/packages/core/src/services/gitService.ts
+++ b/packages/core/src/services/gitService.ts
@@ -8,8 +8,7 @@ import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import { isNodeError } from '../utils/errors.js';
import { spawnAsync } from '../utils/shell-utils.js';
-import type { SimpleGit } from 'simple-git';
-import { simpleGit, CheckRepoActions } from 'simple-git';
+import { simpleGit, CheckRepoActions, type SimpleGit } from 'simple-git';
import type { Storage } from '../config/storage.js';
import { debugLogger } from '../utils/debugLogger.js';
diff --git a/packages/core/src/services/keychainService.test.ts b/packages/core/src/services/keychainService.test.ts
new file mode 100644
index 0000000000..4ab59a5369
--- /dev/null
+++ b/packages/core/src/services/keychainService.test.ts
@@ -0,0 +1,183 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
+import { KeychainService } from './keychainService.js';
+import { coreEvents } from '../utils/events.js';
+import { debugLogger } from '../utils/debugLogger.js';
+
+type MockKeychain = {
+ getPassword: Mock | undefined;
+ setPassword: Mock | undefined;
+ deletePassword: Mock | undefined;
+ findCredentials: Mock | undefined;
+};
+
+const mockKeytar: MockKeychain = {
+ getPassword: vi.fn(),
+ setPassword: vi.fn(),
+ deletePassword: vi.fn(),
+ findCredentials: vi.fn(),
+};
+
+vi.mock('keytar', () => ({ default: mockKeytar }));
+
+vi.mock('../utils/events.js', () => ({
+ coreEvents: { emitTelemetryKeychainAvailability: vi.fn() },
+}));
+
+vi.mock('../utils/debugLogger.js', () => ({
+ debugLogger: { log: vi.fn() },
+}));
+
+describe('KeychainService', () => {
+ let service: KeychainService;
+ const SERVICE_NAME = 'test-service';
+ let passwords: Record = {};
+
+ beforeEach(() => {
+ vi.clearAllMocks();
+ service = new KeychainService(SERVICE_NAME);
+ passwords = {};
+
+ // Stateful mock implementation to verify behavioral correctness
+ mockKeytar.setPassword?.mockImplementation((_svc, acc, val) => {
+ passwords[acc] = val;
+ return Promise.resolve();
+ });
+ mockKeytar.getPassword?.mockImplementation((_svc, acc) =>
+ Promise.resolve(passwords[acc] ?? null),
+ );
+ mockKeytar.deletePassword?.mockImplementation((_svc, acc) => {
+ const exists = !!passwords[acc];
+ delete passwords[acc];
+ return Promise.resolve(exists);
+ });
+ mockKeytar.findCredentials?.mockImplementation(() =>
+ Promise.resolve(
+ Object.entries(passwords).map(([account, password]) => ({
+ account,
+ password,
+ })),
+ ),
+ );
+ });
+
+ describe('isAvailable', () => {
+ it('should return true and emit telemetry on successful functional test', async () => {
+ const available = await service.isAvailable();
+
+ expect(available).toBe(true);
+ expect(mockKeytar.setPassword).toHaveBeenCalled();
+ expect(coreEvents.emitTelemetryKeychainAvailability).toHaveBeenCalledWith(
+ expect.objectContaining({ available: true }),
+ );
+ });
+
+ it('should return false, log error, and emit telemetry on failed functional test', async () => {
+ mockKeytar.setPassword?.mockRejectedValue(new Error('locked'));
+
+ const available = await service.isAvailable();
+
+ expect(available).toBe(false);
+ expect(debugLogger.log).toHaveBeenCalledWith(
+ expect.stringContaining('encountered an error'),
+ 'locked',
+ );
+ expect(coreEvents.emitTelemetryKeychainAvailability).toHaveBeenCalledWith(
+ expect.objectContaining({ available: false }),
+ );
+ });
+
+ it('should return false, log validation error, and emit telemetry on module load failure', async () => {
+ const originalMock = mockKeytar.getPassword;
+ mockKeytar.getPassword = undefined; // Break schema
+
+ const available = await service.isAvailable();
+
+ expect(available).toBe(false);
+ expect(debugLogger.log).toHaveBeenCalledWith(
+ expect.stringContaining('failed structural validation'),
+ expect.objectContaining({ getPassword: expect.any(Array) }),
+ );
+ expect(coreEvents.emitTelemetryKeychainAvailability).toHaveBeenCalledWith(
+ expect.objectContaining({ available: false }),
+ );
+
+ mockKeytar.getPassword = originalMock;
+ });
+
+ it('should log failure if functional test cycle returns false', async () => {
+ mockKeytar.getPassword?.mockResolvedValue('wrong-password');
+
+ const available = await service.isAvailable();
+
+ expect(available).toBe(false);
+ expect(debugLogger.log).toHaveBeenCalledWith(
+ expect.stringContaining('functional verification failed'),
+ );
+ });
+
+ it('should cache the result and handle concurrent initialization attempts once', async () => {
+ await Promise.all([
+ service.isAvailable(),
+ service.isAvailable(),
+ service.isAvailable(),
+ ]);
+
+ expect(mockKeytar.setPassword).toHaveBeenCalledTimes(1);
+ });
+ });
+
+ describe('Password Operations', () => {
+ beforeEach(async () => {
+ await service.isAvailable();
+ vi.clearAllMocks();
+ });
+
+ it('should store, retrieve, and delete passwords correctly', async () => {
+ await service.setPassword('acc1', 'secret1');
+ await service.setPassword('acc2', 'secret2');
+
+ expect(await service.getPassword('acc1')).toBe('secret1');
+ expect(await service.getPassword('acc2')).toBe('secret2');
+
+ const creds = await service.findCredentials();
+ expect(creds).toHaveLength(2);
+ expect(creds).toContainEqual({ account: 'acc1', password: 'secret1' });
+
+ expect(await service.deletePassword('acc1')).toBe(true);
+ expect(await service.getPassword('acc1')).toBeNull();
+ expect(await service.findCredentials()).toHaveLength(1);
+ });
+
+ it('getPassword should return null if key is missing', async () => {
+ expect(await service.getPassword('missing')).toBeNull();
+ });
+ });
+
+ describe('When Unavailable', () => {
+ beforeEach(() => {
+ mockKeytar.setPassword?.mockRejectedValue(new Error('Unavailable'));
+ });
+
+ it.each([
+ { method: 'getPassword', args: ['acc'] },
+ { method: 'setPassword', args: ['acc', 'val'] },
+ { method: 'deletePassword', args: ['acc'] },
+ { method: 'findCredentials', args: [] },
+ ])('$method should throw a consistent error', async ({ method, args }) => {
+ await expect(
+ (
+ service as unknown as Record<
+ string,
+ (...args: unknown[]) => Promise
+ >
+ )[method](...args),
+ ).rejects.toThrow('Keychain is not available');
+ });
+ });
+});
diff --git a/packages/core/src/services/keychainService.ts b/packages/core/src/services/keychainService.ts
new file mode 100644
index 0000000000..ed28218c11
--- /dev/null
+++ b/packages/core/src/services/keychainService.ts
@@ -0,0 +1,147 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as crypto from 'node:crypto';
+import { coreEvents } from '../utils/events.js';
+import { KeychainAvailabilityEvent } from '../telemetry/types.js';
+import { debugLogger } from '../utils/debugLogger.js';
+import {
+ type Keychain,
+ KeychainSchema,
+ KEYCHAIN_TEST_PREFIX,
+} from './keychainTypes.js';
+
+/**
+ * Service for interacting with OS-level secure storage (e.g. keytar).
+ */
+export class KeychainService {
+ // Track an ongoing initialization attempt to avoid race conditions.
+ private initializationPromise?: Promise;
+
+ /**
+ * @param serviceName Unique identifier for the app in the OS keychain.
+ */
+ constructor(private readonly serviceName: string) {}
+
+ async isAvailable(): Promise {
+ return (await this.getKeychain()) !== null;
+ }
+
+ /**
+ * Retrieves a secret for the given account.
+ * @throws Error if the keychain is unavailable.
+ */
+ async getPassword(account: string): Promise {
+ const keychain = await this.getKeychainOrThrow();
+ return keychain.getPassword(this.serviceName, account);
+ }
+
+ /**
+ * Securely stores a secret.
+ * @throws Error if the keychain is unavailable.
+ */
+ async setPassword(account: string, value: string): Promise {
+ const keychain = await this.getKeychainOrThrow();
+ await keychain.setPassword(this.serviceName, account, value);
+ }
+
+ /**
+ * Removes a secret from the keychain.
+ * @returns true if the secret was deleted, false otherwise.
+ * @throws Error if the keychain is unavailable.
+ */
+ async deletePassword(account: string): Promise {
+ const keychain = await this.getKeychainOrThrow();
+ return keychain.deletePassword(this.serviceName, account);
+ }
+
+ /**
+ * Lists all account/secret pairs stored under this service.
+ * @throws Error if the keychain is unavailable.
+ */
+ async findCredentials(): Promise<
+ Array<{ account: string; password: string }>
+ > {
+ const keychain = await this.getKeychainOrThrow();
+ return keychain.findCredentials(this.serviceName);
+ }
+
+ private async getKeychainOrThrow(): Promise {
+ const keychain = await this.getKeychain();
+ if (!keychain) {
+ throw new Error('Keychain is not available');
+ }
+ return keychain;
+ }
+
+ private getKeychain(): Promise {
+ return (this.initializationPromise ??= this.initializeKeychain());
+ }
+
+ // High-level orchestration of the loading and testing cycle.
+ private async initializeKeychain(): Promise {
+ let resultKeychain: Keychain | null = null;
+
+ try {
+ const keychainModule = await this.loadKeychainModule();
+ if (keychainModule) {
+ if (await this.isKeychainFunctional(keychainModule)) {
+ resultKeychain = keychainModule;
+ } else {
+ debugLogger.log('Keychain functional verification failed');
+ }
+ }
+ } catch (error) {
+ // Avoid logging full error objects to prevent PII exposure.
+ const message = error instanceof Error ? error.message : String(error);
+ debugLogger.log('Keychain initialization encountered an error:', message);
+ }
+
+ coreEvents.emitTelemetryKeychainAvailability(
+ new KeychainAvailabilityEvent(resultKeychain !== null),
+ );
+
+ return resultKeychain;
+ }
+
+ // Low-level dynamic loading and structural validation.
+ private async loadKeychainModule(): Promise {
+ const moduleName = 'keytar';
+ const module: unknown = await import(moduleName);
+ const potential = (this.isRecord(module) && module['default']) || module;
+
+ const result = KeychainSchema.safeParse(potential);
+ if (result.success) {
+ // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+ return potential as Keychain;
+ }
+
+ debugLogger.log(
+ 'Keychain module failed structural validation:',
+ result.error.flatten().fieldErrors,
+ );
+ return null;
+ }
+
+ private isRecord(obj: unknown): obj is Record {
+ return typeof obj === 'object' && obj !== null;
+ }
+
+ // Performs a set-get-delete cycle to verify keychain functionality.
+ private async isKeychainFunctional(keychain: Keychain): Promise {
+ const testAccount = `${KEYCHAIN_TEST_PREFIX}${crypto.randomBytes(8).toString('hex')}`;
+ const testPassword = 'test';
+
+ await keychain.setPassword(this.serviceName, testAccount, testPassword);
+ const retrieved = await keychain.getPassword(this.serviceName, testAccount);
+ const deleted = await keychain.deletePassword(
+ this.serviceName,
+ testAccount,
+ );
+
+ return deleted && retrieved === testPassword;
+ }
+}
diff --git a/packages/core/src/services/keychainTypes.ts b/packages/core/src/services/keychainTypes.ts
new file mode 100644
index 0000000000..a643f763e4
--- /dev/null
+++ b/packages/core/src/services/keychainTypes.ts
@@ -0,0 +1,38 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { z } from 'zod';
+
+/**
+ * Interface for OS-level secure storage operations.
+ * Note: Method names must match the underlying library (e.g. keytar)
+ * to support correct dynamic loading and schema validation.
+ */
+export interface Keychain {
+ getPassword(service: string, account: string): Promise;
+ setPassword(
+ service: string,
+ account: string,
+ password: string,
+ ): Promise;
+ deletePassword(service: string, account: string): Promise;
+ findCredentials(
+ service: string,
+ ): Promise>;
+}
+
+/**
+ * Zod schema to validate that a module satisfies the Keychain interface.
+ */
+export const KeychainSchema = z.object({
+ getPassword: z.function(),
+ setPassword: z.function(),
+ deletePassword: z.function(),
+ findCredentials: z.function(),
+});
+
+export const KEYCHAIN_TEST_PREFIX = '__keychain_test__';
+export const SECRET_PREFIX = '__secret__';
diff --git a/packages/core/src/services/loopDetectionService.test.ts b/packages/core/src/services/loopDetectionService.test.ts
index 840c9ae18e..4695cd7bbf 100644
--- a/packages/core/src/services/loopDetectionService.test.ts
+++ b/packages/core/src/services/loopDetectionService.test.ts
@@ -9,12 +9,12 @@ import type { Content } from '@google/genai';
import type { Config } from '../config/config.js';
import type { GeminiClient } from '../core/client.js';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
-import type {
- ServerGeminiContentEvent,
- ServerGeminiStreamEvent,
- ServerGeminiToolCallRequestEvent,
+import {
+ GeminiEventType,
+ type ServerGeminiContentEvent,
+ type ServerGeminiStreamEvent,
+ type ServerGeminiToolCallRequestEvent,
} from '../core/turn.js';
-import { GeminiEventType } from '../core/turn.js';
import * as loggers from '../telemetry/loggers.js';
import { LoopType } from '../telemetry/types.js';
import { LoopDetectionService } from './loopDetectionService.js';
@@ -79,7 +79,7 @@ describe('LoopDetectionService', () => {
it(`should not detect a loop for fewer than TOOL_CALL_LOOP_THRESHOLD identical calls`, () => {
const event = createToolCallRequestEvent('testTool', { param: 'value' });
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
- expect(service.addAndCheck(event)).toBe(false);
+ expect(service.addAndCheck(event).count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -89,7 +89,7 @@ describe('LoopDetectionService', () => {
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
service.addAndCheck(event);
}
- expect(service.addAndCheck(event)).toBe(true);
+ expect(service.addAndCheck(event).count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -98,7 +98,7 @@ describe('LoopDetectionService', () => {
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
service.addAndCheck(event);
}
- expect(service.addAndCheck(event)).toBe(true);
+ expect(service.addAndCheck(event).count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -114,9 +114,9 @@ describe('LoopDetectionService', () => {
});
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 2; i++) {
- expect(service.addAndCheck(event1)).toBe(false);
- expect(service.addAndCheck(event2)).toBe(false);
- expect(service.addAndCheck(event3)).toBe(false);
+ expect(service.addAndCheck(event1).count).toBe(0);
+ expect(service.addAndCheck(event2).count).toBe(0);
+ expect(service.addAndCheck(event3).count).toBe(0);
}
});
@@ -130,14 +130,14 @@ describe('LoopDetectionService', () => {
// Send events just below the threshold
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
- expect(service.addAndCheck(toolCallEvent)).toBe(false);
+ expect(service.addAndCheck(toolCallEvent).count).toBe(0);
}
// Send a different event type
- expect(service.addAndCheck(otherEvent)).toBe(false);
+ expect(service.addAndCheck(otherEvent).count).toBe(0);
// Send the tool call event again, which should now trigger the loop
- expect(service.addAndCheck(toolCallEvent)).toBe(true);
+ expect(service.addAndCheck(toolCallEvent).count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -146,7 +146,7 @@ describe('LoopDetectionService', () => {
expect(loggers.logLoopDetectionDisabled).toHaveBeenCalledTimes(1);
const event = createToolCallRequestEvent('testTool', { param: 'value' });
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
- expect(service.addAndCheck(event)).toBe(false);
+ expect(service.addAndCheck(event).count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -156,19 +156,19 @@ describe('LoopDetectionService', () => {
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
service.addAndCheck(event);
}
- expect(service.addAndCheck(event)).toBe(true);
+ expect(service.addAndCheck(event).count).toBe(1);
service.disableForSession();
- // Should now return false even though a loop was previously detected
- expect(service.addAndCheck(event)).toBe(false);
+ // Should now return 0 even though a loop was previously detected
+ expect(service.addAndCheck(event).count).toBe(0);
});
it('should skip loop detection if disabled in config', () => {
vi.spyOn(mockConfig, 'getDisableLoopDetection').mockReturnValue(true);
const event = createToolCallRequestEvent('testTool', { param: 'value' });
for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD + 2; i++) {
- expect(service.addAndCheck(event)).toBe(false);
+ expect(service.addAndCheck(event).count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -192,8 +192,8 @@ describe('LoopDetectionService', () => {
service.reset('');
for (let i = 0; i < 1000; i++) {
const content = generateRandomString(10);
- const isLoop = service.addAndCheck(createContentEvent(content));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(content));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -202,17 +202,17 @@ describe('LoopDetectionService', () => {
service.reset('');
const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- isLoop = service.addAndCheck(createContentEvent(repeatedContent));
+ result = service.addAndCheck(createContentEvent(repeatedContent));
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
it('should not detect a loop for a list with a long shared prefix', () => {
service.reset('');
- let isLoop = false;
+ let result = { count: 0 };
const longPrefix =
'projects/my-google-cloud-project-12345/locations/us-central1/services/';
@@ -223,9 +223,9 @@ describe('LoopDetectionService', () => {
// Simulate receiving the list in a single large chunk or a few chunks
// This is the specific case where the issue occurs, as list boundaries might not reset tracking properly
- isLoop = service.addAndCheck(createContentEvent(listContent));
+ result = service.addAndCheck(createContentEvent(listContent));
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -234,12 +234,12 @@ describe('LoopDetectionService', () => {
const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);
const fillerContent = generateRandomString(500);
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- isLoop = service.addAndCheck(createContentEvent(fillerContent));
+ result = service.addAndCheck(createContentEvent(repeatedContent));
+ result = service.addAndCheck(createContentEvent(fillerContent));
}
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -248,12 +248,12 @@ describe('LoopDetectionService', () => {
const longPattern = createRepetitiveContent(1, 150);
expect(longPattern.length).toBe(150);
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
- isLoop = service.addAndCheck(createContentEvent(longPattern));
- if (isLoop) break;
+ result = service.addAndCheck(createContentEvent(longPattern));
+ if (result.count > 0) break;
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -266,13 +266,13 @@ describe('LoopDetectionService', () => {
I will wait for the user's next command.
`;
- let isLoop = false;
+ let result = { count: 0 };
// Loop enough times to trigger the threshold
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- isLoop = service.addAndCheck(createContentEvent(userPattern));
- if (isLoop) break;
+ result = service.addAndCheck(createContentEvent(userPattern));
+ if (result.count > 0) break;
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -281,12 +281,12 @@ describe('LoopDetectionService', () => {
const userPattern =
'I have added all the requested logs and verified the test file. I will now mark the task as complete.\n ';
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- isLoop = service.addAndCheck(createContentEvent(userPattern));
- if (isLoop) break;
+ result = service.addAndCheck(createContentEvent(userPattern));
+ if (result.count > 0) break;
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -294,14 +294,14 @@ describe('LoopDetectionService', () => {
service.reset('');
const alternatingPattern = 'Thinking... Done. ';
- let isLoop = false;
+ let result = { count: 0 };
// Needs more iterations because the pattern is short relative to chunk size,
// so it takes a few slides of the window to find the exact alignment.
for (let i = 0; i < CONTENT_LOOP_THRESHOLD * 3; i++) {
- isLoop = service.addAndCheck(createContentEvent(alternatingPattern));
- if (isLoop) break;
+ result = service.addAndCheck(createContentEvent(alternatingPattern));
+ if (result.count > 0) break;
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -310,12 +310,12 @@ describe('LoopDetectionService', () => {
const thoughtPattern =
'I need to check the file. The file does not exist. I will create the file. ';
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- isLoop = service.addAndCheck(createContentEvent(thoughtPattern));
- if (isLoop) break;
+ result = service.addAndCheck(createContentEvent(thoughtPattern));
+ if (result.count > 0) break;
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
});
@@ -328,12 +328,12 @@ describe('LoopDetectionService', () => {
service.addAndCheck(createContentEvent('```\n'));
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
- const isLoop = service.addAndCheck(createContentEvent('\n```'));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent('\n```'));
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -349,15 +349,15 @@ describe('LoopDetectionService', () => {
// Now transition into a code block - this should prevent loop detection
// even though we were already close to the threshold
const codeBlockStart = '```javascript\n';
- const isLoop = service.addAndCheck(createContentEvent(codeBlockStart));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(codeBlockStart));
+ expect(result.count).toBe(0);
// Continue adding repetitive content inside the code block - should not trigger loop
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- const isLoopInside = service.addAndCheck(
+ const resultInside = service.addAndCheck(
createContentEvent(repeatedContent),
);
- expect(isLoopInside).toBe(false);
+ expect(resultInside.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -372,8 +372,8 @@ describe('LoopDetectionService', () => {
// Verify we are now inside a code block and any content should be ignored for loop detection
const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -388,25 +388,25 @@ describe('LoopDetectionService', () => {
// Enter code block (1 fence) - should stop tracking
const enterResult = service.addAndCheck(createContentEvent('```\n'));
- expect(enterResult).toBe(false);
+ expect(enterResult.count).toBe(0);
// Inside code block - should not track loops
for (let i = 0; i < 5; i++) {
const insideResult = service.addAndCheck(
createContentEvent(repeatedContent),
);
- expect(insideResult).toBe(false);
+ expect(insideResult.count).toBe(0);
}
// Exit code block (2nd fence) - should reset tracking but still return false
const exitResult = service.addAndCheck(createContentEvent('```\n'));
- expect(exitResult).toBe(false);
+ expect(exitResult.count).toBe(0);
// Enter code block again (3rd fence) - should stop tracking again
const reenterResult = service.addAndCheck(
createContentEvent('```python\n'),
);
- expect(reenterResult).toBe(false);
+ expect(reenterResult.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -419,11 +419,11 @@ describe('LoopDetectionService', () => {
service.addAndCheck(createContentEvent('\nsome code\n'));
service.addAndCheck(createContentEvent('```'));
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- isLoop = service.addAndCheck(createContentEvent(repeatedContent));
+ result = service.addAndCheck(createContentEvent(repeatedContent));
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -431,9 +431,9 @@ describe('LoopDetectionService', () => {
service.reset('');
service.addAndCheck(createContentEvent('```\ncode1\n```'));
service.addAndCheck(createContentEvent('\nsome text\n'));
- const isLoop = service.addAndCheck(createContentEvent('```\ncode2\n```'));
+ const result = service.addAndCheck(createContentEvent('```\ncode2\n```'));
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -445,12 +445,12 @@ describe('LoopDetectionService', () => {
service.addAndCheck(createContentEvent('\ncode1\n'));
service.addAndCheck(createContentEvent('```'));
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- isLoop = service.addAndCheck(createContentEvent(repeatedContent));
+ result = service.addAndCheck(createContentEvent(repeatedContent));
}
- expect(isLoop).toBe(true);
+ expect(result.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
});
@@ -462,12 +462,12 @@ describe('LoopDetectionService', () => {
service.addAndCheck(createContentEvent('```\n'));
for (let i = 0; i < 20; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatingTokens));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatingTokens));
+ expect(result.count).toBe(0);
}
- const isLoop = service.addAndCheck(createContentEvent('\n```'));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent('\n```'));
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -484,10 +484,10 @@ describe('LoopDetectionService', () => {
// We are now in a code block, so loop detection should be off.
// Let's add the repeated content again, it should not trigger a loop.
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
- isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -505,8 +505,8 @@ describe('LoopDetectionService', () => {
// Add more repeated content after table - should not trigger loop
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -525,8 +525,8 @@ describe('LoopDetectionService', () => {
// Add more repeated content after list - should not trigger loop
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -545,8 +545,8 @@ describe('LoopDetectionService', () => {
// Add more repeated content after heading - should not trigger loop
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -565,8 +565,8 @@ describe('LoopDetectionService', () => {
// Add more repeated content after blockquote - should not trigger loop
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(createContentEvent(repeatedContent));
- expect(isLoop).toBe(false);
+ const result = service.addAndCheck(createContentEvent(repeatedContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
@@ -601,10 +601,10 @@ describe('LoopDetectionService', () => {
CONTENT_CHUNK_SIZE,
);
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(
+ const result = service.addAndCheck(
createContentEvent(newRepeatedContent),
);
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
}
});
@@ -638,10 +638,10 @@ describe('LoopDetectionService', () => {
CONTENT_CHUNK_SIZE,
);
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(
+ const result = service.addAndCheck(
createContentEvent(newRepeatedContent),
);
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
}
});
@@ -677,10 +677,10 @@ describe('LoopDetectionService', () => {
CONTENT_CHUNK_SIZE,
);
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
- const isLoop = service.addAndCheck(
+ const result = service.addAndCheck(
createContentEvent(newRepeatedContent),
);
- expect(isLoop).toBe(false);
+ expect(result.count).toBe(0);
}
});
@@ -691,7 +691,7 @@ describe('LoopDetectionService', () => {
describe('Edge Cases', () => {
it('should handle empty content', () => {
const event = createContentEvent('');
- expect(service.addAndCheck(event)).toBe(false);
+ expect(service.addAndCheck(event).count).toBe(0);
});
});
@@ -699,10 +699,10 @@ describe('LoopDetectionService', () => {
it('should not detect a loop for repeating divider-like content', () => {
service.reset('');
const dividerContent = '-'.repeat(CONTENT_CHUNK_SIZE);
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- isLoop = service.addAndCheck(createContentEvent(dividerContent));
- expect(isLoop).toBe(false);
+ result = service.addAndCheck(createContentEvent(dividerContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
@@ -710,15 +710,52 @@ describe('LoopDetectionService', () => {
it('should not detect a loop for repeating complex box-drawing dividers', () => {
service.reset('');
const dividerContent = '╭─'.repeat(CONTENT_CHUNK_SIZE / 2);
- let isLoop = false;
+ let result = { count: 0 };
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
- isLoop = service.addAndCheck(createContentEvent(dividerContent));
- expect(isLoop).toBe(false);
+ result = service.addAndCheck(createContentEvent(dividerContent));
+ expect(result.count).toBe(0);
}
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
});
+ describe('Strike Management', () => {
+ it('should increment strike count for repeated detections', () => {
+ const event = createToolCallRequestEvent('testTool', { param: 'value' });
+
+ // First strike
+ for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
+ service.addAndCheck(event);
+ }
+ expect(service.addAndCheck(event).count).toBe(1);
+
+ // Recovery simulated by caller calling clearDetection()
+ service.clearDetection();
+
+ // Second strike
+ expect(service.addAndCheck(event).count).toBe(2);
+ });
+
+ it('should allow recovery turn to proceed after clearDetection', () => {
+ const event = createToolCallRequestEvent('testTool', { param: 'value' });
+
+ // Trigger loop
+ for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
+ service.addAndCheck(event);
+ }
+ expect(service.addAndCheck(event).count).toBe(1);
+
+ // Caller clears detection to allow recovery
+ service.clearDetection();
+
+ // Subsequent call in the same turn (or next turn before it repeats) should be 0
+ // In reality, addAndCheck is called per event.
+ // If the model sends a NEW event, it should not immediately trigger.
+ const newEvent = createContentEvent('Recovery text');
+ expect(service.addAndCheck(newEvent).count).toBe(0);
+ });
+ });
+
describe('Reset Functionality', () => {
it('tool call should reset content count', () => {
const contentEvent = createContentEvent('Some content.');
@@ -732,19 +769,19 @@ describe('LoopDetectionService', () => {
service.addAndCheck(toolEvent);
// Should start fresh
- expect(service.addAndCheck(createContentEvent('Fresh content.'))).toBe(
- false,
- );
+ expect(
+ service.addAndCheck(createContentEvent('Fresh content.')).count,
+ ).toBe(0);
});
});
describe('General Behavior', () => {
- it('should return false for unhandled event types', () => {
+ it('should return 0 count for unhandled event types', () => {
const otherEvent = {
type: 'unhandled_event',
} as unknown as ServerGeminiStreamEvent;
- expect(service.addAndCheck(otherEvent)).toBe(false);
- expect(service.addAndCheck(otherEvent)).toBe(false);
+ expect(service.addAndCheck(otherEvent).count).toBe(0);
+ expect(service.addAndCheck(otherEvent).count).toBe(0);
});
});
});
@@ -805,16 +842,16 @@ describe('LoopDetectionService LLM Checks', () => {
}
};
- it('should not trigger LLM check before LLM_CHECK_AFTER_TURNS', async () => {
- await advanceTurns(39);
+ it('should not trigger LLM check before LLM_CHECK_AFTER_TURNS (30)', async () => {
+ await advanceTurns(29);
expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
});
- it('should trigger LLM check on the 40th turn', async () => {
+ it('should trigger LLM check on the 30th turn', async () => {
mockBaseLlmClient.generateJson = vi
.fn()
.mockResolvedValue({ unproductive_state_confidence: 0.1 });
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(
expect.objectContaining({
@@ -828,12 +865,12 @@ describe('LoopDetectionService LLM Checks', () => {
});
it('should detect a cognitive loop when confidence is high', async () => {
- // First check at turn 40
+ // First check at turn 30
mockBaseLlmClient.generateJson = vi.fn().mockResolvedValue({
unproductive_state_confidence: 0.85,
unproductive_state_analysis: 'Repetitive actions',
});
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(
expect.objectContaining({
@@ -842,16 +879,16 @@ describe('LoopDetectionService LLM Checks', () => {
);
// The confidence of 0.85 will result in a low interval.
- // The interval will be: 7 + (15 - 7) * (1 - 0.85) = 7 + 8 * 0.15 = 8.2 -> rounded to 8
- await advanceTurns(7); // advance to turn 47
+ // The interval will be: 5 + (15 - 5) * (1 - 0.85) = 5 + 10 * 0.15 = 6.5 -> rounded to 7
+ await advanceTurns(6); // advance to turn 36
mockBaseLlmClient.generateJson = vi.fn().mockResolvedValue({
unproductive_state_confidence: 0.95,
unproductive_state_analysis: 'Repetitive actions',
});
- const finalResult = await service.turnStarted(abortController.signal); // This is turn 48
+ const finalResult = await service.turnStarted(abortController.signal); // This is turn 37
- expect(finalResult).toBe(true);
+ expect(finalResult.count).toBe(1);
expect(loggers.logLoopDetected).toHaveBeenCalledWith(
mockConfig,
expect.objectContaining({
@@ -867,25 +904,25 @@ describe('LoopDetectionService LLM Checks', () => {
unproductive_state_confidence: 0.5,
unproductive_state_analysis: 'Looks okay',
});
- await advanceTurns(40);
+ await advanceTurns(30);
const result = await service.turnStarted(abortController.signal);
- expect(result).toBe(false);
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
it('should adjust the check interval based on confidence', async () => {
// Confidence is 0.0, so interval should be MAX_LLM_CHECK_INTERVAL (15)
- // Interval = 7 + (15 - 7) * (1 - 0.0) = 15
+ // Interval = 5 + (15 - 5) * (1 - 0.0) = 15
mockBaseLlmClient.generateJson = vi
.fn()
.mockResolvedValue({ unproductive_state_confidence: 0.0 });
- await advanceTurns(40); // First check at turn 40
+ await advanceTurns(30); // First check at turn 30
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
- await advanceTurns(14); // Advance to turn 54
+ await advanceTurns(14); // Advance to turn 44
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
- await service.turnStarted(abortController.signal); // Turn 55
+ await service.turnStarted(abortController.signal); // Turn 45
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
});
@@ -893,18 +930,18 @@ describe('LoopDetectionService LLM Checks', () => {
mockBaseLlmClient.generateJson = vi
.fn()
.mockRejectedValue(new Error('API error'));
- await advanceTurns(40);
+ await advanceTurns(30);
const result = await service.turnStarted(abortController.signal);
- expect(result).toBe(false);
+ expect(result.count).toBe(0);
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
});
it('should not trigger LLM check when disabled for session', async () => {
service.disableForSession();
expect(loggers.logLoopDetectionDisabled).toHaveBeenCalledTimes(1);
- await advanceTurns(40);
+ await advanceTurns(30);
const result = await service.turnStarted(abortController.signal);
- expect(result).toBe(false);
+ expect(result.count).toBe(0);
expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
});
@@ -925,7 +962,7 @@ describe('LoopDetectionService LLM Checks', () => {
.fn()
.mockResolvedValue({ unproductive_state_confidence: 0.1 });
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
@@ -950,7 +987,7 @@ describe('LoopDetectionService LLM Checks', () => {
unproductive_state_analysis: 'Main says loop',
});
- await advanceTurns(40);
+ await advanceTurns(30);
// It should have called generateJson twice
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
@@ -990,7 +1027,7 @@ describe('LoopDetectionService LLM Checks', () => {
unproductive_state_analysis: 'Main says no loop',
});
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
expect(mockBaseLlmClient.generateJson).toHaveBeenNthCalledWith(
@@ -1010,12 +1047,12 @@ describe('LoopDetectionService LLM Checks', () => {
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
// But should have updated the interval based on the main model's confidence (0.89)
- // Interval = 7 + (15-7) * (1 - 0.89) = 7 + 8 * 0.11 = 7 + 0.88 = 7.88 -> 8
+ // Interval = 5 + (15-5) * (1 - 0.89) = 5 + 10 * 0.11 = 5 + 1.1 = 6.1 -> 6
- // Advance by 7 turns
- await advanceTurns(7);
+ // Advance by 5 turns
+ await advanceTurns(5);
- // Next turn (48) should trigger another check
+ // Next turn (36) should trigger another check
await service.turnStarted(abortController.signal);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(3);
});
@@ -1033,7 +1070,7 @@ describe('LoopDetectionService LLM Checks', () => {
unproductive_state_analysis: 'Flash says loop',
});
- await advanceTurns(40);
+ await advanceTurns(30);
// It should have called generateJson only once
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
@@ -1047,8 +1084,6 @@ describe('LoopDetectionService LLM Checks', () => {
expect(loggers.logLoopDetected).toHaveBeenCalledWith(
mockConfig,
expect.objectContaining({
- 'event.name': 'loop_detected',
- loop_type: LoopType.LLM_DETECTED_LOOP,
confirmed_by_model: 'gemini-2.5-flash',
}),
);
@@ -1061,7 +1096,7 @@ describe('LoopDetectionService LLM Checks', () => {
.fn()
.mockResolvedValue({ unproductive_state_confidence: 0.1 });
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
@@ -1091,7 +1126,7 @@ describe('LoopDetectionService LLM Checks', () => {
.fn()
.mockResolvedValue({ unproductive_state_confidence: 0.1 });
- await advanceTurns(40);
+ await advanceTurns(30);
expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
diff --git a/packages/core/src/services/loopDetectionService.ts b/packages/core/src/services/loopDetectionService.ts
index 67207915c1..9bc8b406f8 100644
--- a/packages/core/src/services/loopDetectionService.ts
+++ b/packages/core/src/services/loopDetectionService.ts
@@ -6,8 +6,7 @@
import type { Content } from '@google/genai';
import { createHash } from 'node:crypto';
-import type { ServerGeminiStreamEvent } from '../core/turn.js';
-import { GeminiEventType } from '../core/turn.js';
+import { GeminiEventType, type ServerGeminiStreamEvent } from '../core/turn.js';
import {
logLoopDetected,
logLoopDetectionDisabled,
@@ -40,7 +39,7 @@ const LLM_LOOP_CHECK_HISTORY_COUNT = 20;
/**
* The number of turns that must pass in a single prompt before the LLM-based loop check is activated.
*/
-const LLM_CHECK_AFTER_TURNS = 40;
+const LLM_CHECK_AFTER_TURNS = 30;
/**
* The default interval, in number of turns, at which the LLM-based loop check is performed.
@@ -52,7 +51,7 @@ const DEFAULT_LLM_CHECK_INTERVAL = 10;
* The minimum interval for LLM-based loop checks.
* This is used when the confidence of a loop is high, to check more frequently.
*/
-const MIN_LLM_CHECK_INTERVAL = 7;
+const MIN_LLM_CHECK_INTERVAL = 5;
/**
* The maximum interval for LLM-based loop checks.
@@ -118,6 +117,15 @@ const LOOP_DETECTION_SCHEMA: Record = {
required: ['unproductive_state_analysis', 'unproductive_state_confidence'],
};
+/**
+ * Result of a loop detection check.
+ */
+export interface LoopDetectionResult {
+ count: number;
+ type?: LoopType;
+ detail?: string;
+ confirmedByModel?: string;
+}
/**
* Service for detecting and preventing infinite loops in AI responses.
* Monitors tool call repetitions and content sentence repetitions.
@@ -136,8 +144,11 @@ export class LoopDetectionService {
private contentStats = new Map();
private lastContentIndex = 0;
private loopDetected = false;
+ private detectedCount = 0;
+ private lastLoopDetail?: string;
private inCodeBlock = false;
+ private lastLoopType?: LoopType;
// LLM loop track tracking
private turnsInCurrentPrompt = 0;
private llmCheckInterval = DEFAULT_LLM_CHECK_INTERVAL;
@@ -170,31 +181,68 @@ export class LoopDetectionService {
/**
* Processes a stream event and checks for loop conditions.
* @param event - The stream event to process
- * @returns true if a loop is detected, false otherwise
+ * @returns A LoopDetectionResult
*/
- addAndCheck(event: ServerGeminiStreamEvent): boolean {
+ addAndCheck(event: ServerGeminiStreamEvent): LoopDetectionResult {
if (this.disabledForSession || this.config.getDisableLoopDetection()) {
- return false;
+ return { count: 0 };
+ }
+ if (this.loopDetected) {
+ return {
+ count: this.detectedCount,
+ type: this.lastLoopType,
+ detail: this.lastLoopDetail,
+ };
}
- if (this.loopDetected) {
- return this.loopDetected;
- }
+ let isLoop = false;
+ let detail: string | undefined;
switch (event.type) {
case GeminiEventType.ToolCallRequest:
// content chanting only happens in one single stream, reset if there
// is a tool call in between
this.resetContentTracking();
- this.loopDetected = this.checkToolCallLoop(event.value);
+ isLoop = this.checkToolCallLoop(event.value);
+ if (isLoop) {
+ detail = `Repeated tool call: ${event.value.name} with arguments ${JSON.stringify(event.value.args)}`;
+ }
break;
case GeminiEventType.Content:
- this.loopDetected = this.checkContentLoop(event.value);
+ isLoop = this.checkContentLoop(event.value);
+ if (isLoop) {
+ detail = `Repeating content detected: "${this.streamContentHistory.substring(Math.max(0, this.lastContentIndex - 20), this.lastContentIndex + CONTENT_CHUNK_SIZE).trim()}..."`;
+ }
break;
default:
break;
}
- return this.loopDetected;
+
+ if (isLoop) {
+ this.loopDetected = true;
+ this.detectedCount++;
+ this.lastLoopDetail = detail;
+ this.lastLoopType =
+ event.type === GeminiEventType.ToolCallRequest
+ ? LoopType.CONSECUTIVE_IDENTICAL_TOOL_CALLS
+ : LoopType.CONTENT_CHANTING_LOOP;
+
+ logLoopDetected(
+ this.config,
+ new LoopDetectedEvent(
+ this.lastLoopType,
+ this.promptId,
+ this.detectedCount,
+ ),
+ );
+ }
+ return isLoop
+ ? {
+ count: this.detectedCount,
+ type: this.lastLoopType,
+ detail: this.lastLoopDetail,
+ }
+ : { count: 0 };
}
/**
@@ -205,12 +253,20 @@ export class LoopDetectionService {
* is performed periodically based on the `llmCheckInterval`.
*
* @param signal - An AbortSignal to allow for cancellation of the asynchronous LLM check.
- * @returns A promise that resolves to `true` if a loop is detected, and `false` otherwise.
+ * @returns A promise that resolves to a LoopDetectionResult.
*/
- async turnStarted(signal: AbortSignal) {
+ async turnStarted(signal: AbortSignal): Promise {
if (this.disabledForSession || this.config.getDisableLoopDetection()) {
- return false;
+ return { count: 0 };
}
+ if (this.loopDetected) {
+ return {
+ count: this.detectedCount,
+ type: this.lastLoopType,
+ detail: this.lastLoopDetail,
+ };
+ }
+
this.turnsInCurrentPrompt++;
if (
@@ -218,10 +274,35 @@ export class LoopDetectionService {
this.turnsInCurrentPrompt - this.lastCheckTurn >= this.llmCheckInterval
) {
this.lastCheckTurn = this.turnsInCurrentPrompt;
- return this.checkForLoopWithLLM(signal);
- }
+ const { isLoop, analysis, confirmedByModel } =
+ await this.checkForLoopWithLLM(signal);
+ if (isLoop) {
+ this.loopDetected = true;
+ this.detectedCount++;
+ this.lastLoopDetail = analysis;
+ this.lastLoopType = LoopType.LLM_DETECTED_LOOP;
- return false;
+ logLoopDetected(
+ this.config,
+ new LoopDetectedEvent(
+ this.lastLoopType,
+ this.promptId,
+ this.detectedCount,
+ confirmedByModel,
+ analysis,
+ LLM_CONFIDENCE_THRESHOLD,
+ ),
+ );
+
+ return {
+ count: this.detectedCount,
+ type: this.lastLoopType,
+ detail: this.lastLoopDetail,
+ confirmedByModel,
+ };
+ }
+ }
+ return { count: 0 };
}
private checkToolCallLoop(toolCall: { name: string; args: object }): boolean {
@@ -233,13 +314,6 @@ export class LoopDetectionService {
this.toolCallRepetitionCount = 1;
}
if (this.toolCallRepetitionCount >= TOOL_CALL_LOOP_THRESHOLD) {
- logLoopDetected(
- this.config,
- new LoopDetectedEvent(
- LoopType.CONSECUTIVE_IDENTICAL_TOOL_CALLS,
- this.promptId,
- ),
- );
return true;
}
return false;
@@ -346,13 +420,6 @@ export class LoopDetectionService {
const chunkHash = createHash('sha256').update(currentChunk).digest('hex');
if (this.isLoopDetectedForChunk(currentChunk, chunkHash)) {
- logLoopDetected(
- this.config,
- new LoopDetectedEvent(
- LoopType.CHANTING_IDENTICAL_SENTENCES,
- this.promptId,
- ),
- );
return true;
}
@@ -446,28 +513,29 @@ export class LoopDetectionService {
return originalChunk === currentChunk;
}
- private trimRecentHistory(recentHistory: Content[]): Content[] {
+ private trimRecentHistory(history: Content[]): Content[] {
// A function response must be preceded by a function call.
// Continuously removes dangling function calls from the end of the history
// until the last turn is not a function call.
- while (
- recentHistory.length > 0 &&
- isFunctionCall(recentHistory[recentHistory.length - 1])
- ) {
- recentHistory.pop();
+ while (history.length > 0 && isFunctionCall(history[history.length - 1])) {
+ history.pop();
}
// A function response should follow a function call.
// Continuously removes leading function responses from the beginning of history
// until the first turn is not a function response.
- while (recentHistory.length > 0 && isFunctionResponse(recentHistory[0])) {
- recentHistory.shift();
+ while (history.length > 0 && isFunctionResponse(history[0])) {
+ history.shift();
}
- return recentHistory;
+ return history;
}
- private async checkForLoopWithLLM(signal: AbortSignal) {
+ private async checkForLoopWithLLM(signal: AbortSignal): Promise<{
+ isLoop: boolean;
+ analysis?: string;
+ confirmedByModel?: string;
+ }> {
const recentHistory = this.config
.getGeminiClient()
.getHistory()
@@ -507,13 +575,19 @@ export class LoopDetectionService {
);
if (!flashResult) {
- return false;
+ return { isLoop: false };
}
- // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- const flashConfidence = flashResult[
- 'unproductive_state_confidence'
- ] as number;
+ const flashConfidence =
+ // eslint-disable-next-line no-restricted-syntax
+ typeof flashResult['unproductive_state_confidence'] === 'number'
+ ? flashResult['unproductive_state_confidence']
+ : 0;
+ const flashAnalysis =
+ // eslint-disable-next-line no-restricted-syntax
+ typeof flashResult['unproductive_state_analysis'] === 'string'
+ ? flashResult['unproductive_state_analysis']
+ : '';
const doubleCheckModelName =
this.config.modelConfigService.getResolvedConfig({
@@ -531,7 +605,7 @@ export class LoopDetectionService {
),
);
this.updateCheckInterval(flashConfidence);
- return false;
+ return { isLoop: false };
}
const availability = this.config.getModelAvailabilityService();
@@ -540,8 +614,11 @@ export class LoopDetectionService {
const flashModelName = this.config.modelConfigService.getResolvedConfig({
model: 'loop-detection',
}).model;
- this.handleConfirmedLoop(flashResult, flashModelName);
- return true;
+ return {
+ isLoop: true,
+ analysis: flashAnalysis,
+ confirmedByModel: flashModelName,
+ };
}
// Double check with configured model
@@ -551,10 +628,18 @@ export class LoopDetectionService {
signal,
);
- const mainModelConfidence = mainModelResult
- ? // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
- (mainModelResult['unproductive_state_confidence'] as number)
- : 0;
+ const mainModelConfidence =
+ mainModelResult &&
+ // eslint-disable-next-line no-restricted-syntax
+ typeof mainModelResult['unproductive_state_confidence'] === 'number'
+ ? mainModelResult['unproductive_state_confidence']
+ : 0;
+ const mainModelAnalysis =
+ mainModelResult &&
+ // eslint-disable-next-line no-restricted-syntax
+ typeof mainModelResult['unproductive_state_analysis'] === 'string'
+ ? mainModelResult['unproductive_state_analysis']
+ : undefined;
logLlmLoopCheck(
this.config,
@@ -568,14 +653,17 @@ export class LoopDetectionService {
if (mainModelResult) {
if (mainModelConfidence >= LLM_CONFIDENCE_THRESHOLD) {
- this.handleConfirmedLoop(mainModelResult, doubleCheckModelName);
- return true;
+ return {
+ isLoop: true,
+ analysis: mainModelAnalysis,
+ confirmedByModel: doubleCheckModelName,
+ };
} else {
this.updateCheckInterval(mainModelConfidence);
}
}
- return false;
+ return { isLoop: false };
}
private async queryLoopDetectionModel(
@@ -597,37 +685,22 @@ export class LoopDetectionService {
if (
result &&
+ // eslint-disable-next-line no-restricted-syntax
typeof result['unproductive_state_confidence'] === 'number'
) {
return result;
}
return null;
- } catch (e) {
- this.config.getDebugMode() ? debugLogger.warn(e) : debugLogger.debug(e);
+ } catch (error) {
+ if (this.config.getDebugMode()) {
+ debugLogger.warn(
+ `Error querying loop detection model (${model}): ${String(error)}`,
+ );
+ }
return null;
}
}
- private handleConfirmedLoop(
- result: Record,
- modelName: string,
- ): void {
- if (
- typeof result['unproductive_state_analysis'] === 'string' &&
- result['unproductive_state_analysis']
- ) {
- debugLogger.warn(result['unproductive_state_analysis']);
- }
- logLoopDetected(
- this.config,
- new LoopDetectedEvent(
- LoopType.LLM_DETECTED_LOOP,
- this.promptId,
- modelName,
- ),
- );
- }
-
private updateCheckInterval(unproductive_state_confidence: number): void {
this.llmCheckInterval = Math.round(
MIN_LLM_CHECK_INTERVAL +
@@ -646,6 +719,17 @@ export class LoopDetectionService {
this.resetContentTracking();
this.resetLlmCheckTracking();
this.loopDetected = false;
+ this.detectedCount = 0;
+ this.lastLoopDetail = undefined;
+ this.lastLoopType = undefined;
+ }
+
+ /**
+ * Resets the loop detected flag to allow a recovery turn to proceed.
+ * This preserves the detectedCount so that the next detection will be count 2.
+ */
+ clearDetection(): void {
+ this.loopDetected = false;
}
private resetToolCallCount(): void {
diff --git a/packages/core/src/services/modelConfig.integration.test.ts b/packages/core/src/services/modelConfig.integration.test.ts
index 2ed2cb47af..09723b95ea 100644
--- a/packages/core/src/services/modelConfig.integration.test.ts
+++ b/packages/core/src/services/modelConfig.integration.test.ts
@@ -5,8 +5,10 @@
*/
import { describe, it, expect } from 'vitest';
-import { ModelConfigService } from './modelConfigService.js';
-import type { ModelConfigServiceConfig } from './modelConfigService.js';
+import {
+ ModelConfigService,
+ type ModelConfigServiceConfig,
+} from './modelConfigService.js';
// This test suite is designed to validate the end-to-end logic of the
// ModelConfigService with a complex, realistic configuration.
diff --git a/packages/core/src/services/modelConfigService.test.ts b/packages/core/src/services/modelConfigService.test.ts
index 767cb2ecfd..2bc69bbfe2 100644
--- a/packages/core/src/services/modelConfigService.test.ts
+++ b/packages/core/src/services/modelConfigService.test.ts
@@ -5,11 +5,11 @@
*/
import { describe, it, expect } from 'vitest';
-import type {
- ModelConfigAlias,
- ModelConfigServiceConfig,
+import {
+ ModelConfigService,
+ type ModelConfigAlias,
+ type ModelConfigServiceConfig,
} from './modelConfigService.js';
-import { ModelConfigService } from './modelConfigService.js';
describe('ModelConfigService', () => {
it('should resolve a basic alias to its model and settings', () => {
diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts
index 1285f2ac94..d4c8a694c3 100644
--- a/packages/core/src/services/shellExecutionService.test.ts
+++ b/packages/core/src/services/shellExecutionService.test.ts
@@ -16,11 +16,11 @@ import {
import EventEmitter from 'node:events';
import type { Readable } from 'node:stream';
import { type ChildProcess } from 'node:child_process';
-import type {
- ShellOutputEvent,
- ShellExecutionConfig,
+import {
+ ShellExecutionService,
+ type ShellOutputEvent,
+ type ShellExecutionConfig,
} from './shellExecutionService.js';
-import { ShellExecutionService } from './shellExecutionService.js';
import type { AnsiOutput, AnsiToken } from '../utils/terminalSerializer.js';
// Hoisted Mocks
diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts
index 768b4abb41..7914cc328f 100644
--- a/packages/core/src/services/shellExecutionService.ts
+++ b/packages/core/src/services/shellExecutionService.ts
@@ -5,8 +5,7 @@
*/
import stripAnsi from 'strip-ansi';
-import type { PtyImplementation } from '../utils/getPty.js';
-import { getPty } from '../utils/getPty.js';
+import { getPty, type PtyImplementation } from '../utils/getPty.js';
import { spawn as cpSpawn, type ChildProcess } from 'node:child_process';
import { TextDecoder } from 'node:util';
import os from 'node:os';
diff --git a/packages/core/src/services/trackerService.ts b/packages/core/src/services/trackerService.ts
index 3203b759e1..06e890175f 100644
--- a/packages/core/src/services/trackerService.ts
+++ b/packages/core/src/services/trackerService.ts
@@ -50,6 +50,15 @@ export class TrackerService {
id,
};
+ if (task.parentId) {
+ const parentList = await this.listTasks();
+ if (!parentList.find((t) => t.id === task.parentId)) {
+ throw new Error(`Parent task with ID ${task.parentId} not found.`);
+ }
+ }
+
+ TrackerTaskSchema.parse(task);
+
await this.saveTask(task);
return task;
}
@@ -70,7 +79,8 @@ export class TrackerService {
error &&
typeof error === 'object' &&
'code' in error &&
- error.code === 'ENOENT'
+ // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+ (error as NodeJS.ErrnoException).code === 'ENOENT'
) {
return null;
}
@@ -130,26 +140,48 @@ export class TrackerService {
id: string,
updates: Partial,
): Promise {
- const task = await this.getTask(id);
+ const isClosing = updates.status === TaskStatus.CLOSED;
+ const changingDependencies = updates.dependencies !== undefined;
+
+ let taskMap: Map | undefined;
+
+ if (isClosing || changingDependencies) {
+ const allTasks = await this.listTasks();
+ taskMap = new Map(allTasks.map((t) => [t.id, t]));
+ }
+
+ const task = taskMap ? taskMap.get(id) : await this.getTask(id);
+
if (!task) {
throw new Error(`Task with ID ${id} not found.`);
}
- const updatedTask = { ...task, ...updates };
+ const updatedTask = { ...task, ...updates, id: task.id };
- // Validate status transition if closing
- if (
- updatedTask.status === TaskStatus.CLOSED &&
- task.status !== TaskStatus.CLOSED
- ) {
- await this.validateCanClose(updatedTask);
+ if (updatedTask.parentId) {
+ const parentExists = taskMap
+ ? taskMap.has(updatedTask.parentId)
+ : !!(await this.getTask(updatedTask.parentId));
+ if (!parentExists) {
+ throw new Error(
+ `Parent task with ID ${updatedTask.parentId} not found.`,
+ );
+ }
}
- // Validate circular dependencies if dependencies changed
- if (updates.dependencies) {
- await this.validateNoCircularDependencies(updatedTask);
+ if (taskMap) {
+ if (isClosing && task.status !== TaskStatus.CLOSED) {
+ this.validateCanClose(updatedTask, taskMap);
+ }
+
+ if (changingDependencies) {
+ taskMap.set(updatedTask.id, updatedTask);
+ this.validateNoCircularDependencies(updatedTask, taskMap);
+ }
}
+ TrackerTaskSchema.parse(updatedTask);
+
await this.saveTask(updatedTask);
return updatedTask;
}
@@ -165,9 +197,12 @@ export class TrackerService {
/**
* Validates that a task can be closed (all dependencies must be closed).
*/
- private async validateCanClose(task: TrackerTask): Promise {
+ private validateCanClose(
+ task: TrackerTask,
+ taskMap: Map,
+ ): void {
for (const depId of task.dependencies) {
- const dep = await this.getTask(depId);
+ const dep = taskMap.get(depId);
if (!dep) {
throw new Error(`Dependency ${depId} not found for task ${task.id}.`);
}
@@ -182,16 +217,10 @@ export class TrackerService {
/**
* Validates that there are no circular dependencies.
*/
- private async validateNoCircularDependencies(
+ private validateNoCircularDependencies(
task: TrackerTask,
- ): Promise {
- const allTasks = await this.listTasks();
- const taskMap = new Map(
- allTasks.map((t) => [t.id, t]),
- );
- // Ensure the current (possibly unsaved) task state is used
- taskMap.set(task.id, task);
-
+ taskMap: Map,
+ ): void {
const visited = new Set();
const stack = new Set();
@@ -209,10 +238,11 @@ export class TrackerService {
stack.add(currentId);
const currentTask = taskMap.get(currentId);
- if (currentTask) {
- for (const depId of currentTask.dependencies) {
- check(depId);
- }
+ if (!currentTask) {
+ throw new Error(`Dependency ${currentId} not found.`);
+ }
+ for (const depId of currentTask.dependencies) {
+ check(depId);
}
stack.delete(currentId);
diff --git a/packages/core/src/telemetry/activity-monitor.test.ts b/packages/core/src/telemetry/activity-monitor.test.ts
index 8d20daa301..68dbe9a1c2 100644
--- a/packages/core/src/telemetry/activity-monitor.test.ts
+++ b/packages/core/src/telemetry/activity-monitor.test.ts
@@ -13,9 +13,9 @@ import {
recordGlobalActivity,
startGlobalActivityMonitoring,
stopGlobalActivityMonitoring,
+ type ActivityEvent,
} from './activity-monitor.js';
import { ActivityType } from './activity-types.js';
-import type { ActivityEvent } from './activity-monitor.js';
import type { Config } from '../config/config.js';
import { debugLogger } from '../utils/debugLogger.js';
diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
index b8148bac62..195c5544bf 100644
--- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
+++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
@@ -14,10 +14,17 @@ import {
afterAll,
beforeEach,
} from 'vitest';
-import type { LogEvent, LogEventEntry } from './clearcut-logger.js';
-import { ClearcutLogger, EventNames, TEST_ONLY } from './clearcut-logger.js';
-import type { ContentGeneratorConfig } from '../../core/contentGenerator.js';
-import { AuthType } from '../../core/contentGenerator.js';
+import {
+ ClearcutLogger,
+ EventNames,
+ TEST_ONLY,
+ type LogEvent,
+ type LogEventEntry,
+} from './clearcut-logger.js';
+import {
+ AuthType,
+ type ContentGeneratorConfig,
+} from '../../core/contentGenerator.js';
import type { SuccessfulToolCall } from '../../core/coreToolScheduler.js';
import type { ConfigParameters } from '../../config/config.js';
import { EventMetadataKey } from './event-metadata-key.js';
@@ -42,8 +49,7 @@ import { GIT_COMMIT_INFO, CLI_VERSION } from '../../generated/git-commit.js';
import { UserAccountManager } from '../../utils/userAccountManager.js';
import { InstallationManager } from '../../utils/installationManager.js';
-import si from 'systeminformation';
-import type { Systeminformation } from 'systeminformation';
+import si, { type Systeminformation } from 'systeminformation';
import * as os from 'node:os';
interface CustomMatchers {
diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
index 8646a3f6d4..51c5ab382f 100644
--- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
+++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
@@ -49,6 +49,7 @@ import type {
ToolOutputMaskingEvent,
KeychainAvailabilityEvent,
TokenStorageInitializationEvent,
+ StartupStatsEvent,
} from '../types.js';
import { EventMetadataKey } from './event-metadata-key.js';
import type { Config } from '../../config/config.js';
@@ -117,6 +118,7 @@ export enum EventNames {
TOKEN_STORAGE_INITIALIZATION = 'token_storage_initialization',
CONSECA_POLICY_GENERATION = 'conseca_policy_generation',
CONSECA_VERDICT = 'conseca_verdict',
+ STARTUP_STATS = 'startup_stats',
}
export interface LogResponse {
@@ -1691,6 +1693,30 @@ export class ClearcutLogger {
this.flushIfNeeded();
}
+ logStartupStatsEvent(event: StartupStatsEvent): void {
+ const data: EventValue[] = [
+ {
+ gemini_cli_key: EventMetadataKey.GEMINI_CLI_STARTUP_PHASES,
+ value: JSON.stringify(event.phases),
+ },
+ {
+ gemini_cli_key: EventMetadataKey.GEMINI_CLI_STARTUP_OS_PLATFORM,
+ value: event.os_platform,
+ },
+ {
+ gemini_cli_key: EventMetadataKey.GEMINI_CLI_STARTUP_OS_RELEASE,
+ value: event.os_release,
+ },
+ {
+ gemini_cli_key: EventMetadataKey.GEMINI_CLI_STARTUP_IS_DOCKER,
+ value: JSON.stringify(event.is_docker),
+ },
+ ];
+
+ this.enqueueLogEvent(this.createLogEvent(EventNames.STARTUP_STATS, data));
+ this.flushIfNeeded();
+ }
+
/**
* Adds default fields to data, and returns a new data array. This fields
* should exist on all log events.
diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
index d799ca1caf..43bfa3278d 100644
--- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
+++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
@@ -7,7 +7,7 @@
// Defines valid event metadata keys for Clearcut logging.
export enum EventMetadataKey {
// Deleted enums: 24
- // Next ID: 172
+ // Next ID: 176
GEMINI_CLI_KEY_UNKNOWN = 0,
@@ -54,6 +54,22 @@ export enum EventMetadataKey {
// Logs the output format of the session.
GEMINI_CLI_START_SESSION_OUTPUT_FORMAT = 94,
+ // ==========================================================================
+ // Startup Stats Event Keys
+ // ==========================================================================
+
+ // Logs the array of startup phases.
+ GEMINI_CLI_STARTUP_PHASES = 172,
+
+ // Logs the OS platform for startup stats.
+ GEMINI_CLI_STARTUP_OS_PLATFORM = 173,
+
+ // Logs the OS release for startup stats.
+ GEMINI_CLI_STARTUP_OS_RELEASE = 174,
+
+ // Logs whether the CLI is running in docker for startup stats.
+ GEMINI_CLI_STARTUP_IS_DOCKER = 175,
+
// ==========================================================================
// User Prompt Event Keys
// ===========================================================================
diff --git a/packages/core/src/telemetry/conseca-logger.ts b/packages/core/src/telemetry/conseca-logger.ts
index 41f1ac3d15..ad88d092ee 100644
--- a/packages/core/src/telemetry/conseca-logger.ts
+++ b/packages/core/src/telemetry/conseca-logger.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { LogRecord } from '@opentelemetry/api-logs';
-import { logs } from '@opentelemetry/api-logs';
+import { logs, type LogRecord } from '@opentelemetry/api-logs';
import type { Config } from '../config/config.js';
import { SERVICE_NAME } from './constants.js';
import { isTelemetrySdkInitialized } from './sdk.js';
diff --git a/packages/core/src/telemetry/file-exporters.test.ts b/packages/core/src/telemetry/file-exporters.test.ts
index 80a2ccafad..4b4f688ab8 100644
--- a/packages/core/src/telemetry/file-exporters.test.ts
+++ b/packages/core/src/telemetry/file-exporters.test.ts
@@ -13,8 +13,10 @@ import {
import { ExportResultCode } from '@opentelemetry/core';
import type { ReadableSpan } from '@opentelemetry/sdk-trace-base';
import type { ReadableLogRecord } from '@opentelemetry/sdk-logs';
-import type { ResourceMetrics } from '@opentelemetry/sdk-metrics';
-import { AggregationTemporality } from '@opentelemetry/sdk-metrics';
+import {
+ AggregationTemporality,
+ type ResourceMetrics,
+} from '@opentelemetry/sdk-metrics';
import * as fs from 'node:fs';
function createMockWriteStream(): {
diff --git a/packages/core/src/telemetry/file-exporters.ts b/packages/core/src/telemetry/file-exporters.ts
index def6e91f44..9f8d7f51c1 100644
--- a/packages/core/src/telemetry/file-exporters.ts
+++ b/packages/core/src/telemetry/file-exporters.ts
@@ -5,18 +5,17 @@
*/
import * as fs from 'node:fs';
-import type { ExportResult } from '@opentelemetry/core';
-import { ExportResultCode } from '@opentelemetry/core';
+import { ExportResultCode, type ExportResult } from '@opentelemetry/core';
import type { ReadableSpan, SpanExporter } from '@opentelemetry/sdk-trace-base';
import type {
ReadableLogRecord,
LogRecordExporter,
} from '@opentelemetry/sdk-logs';
-import type {
- ResourceMetrics,
- PushMetricExporter,
+import {
+ AggregationTemporality,
+ type ResourceMetrics,
+ type PushMetricExporter,
} from '@opentelemetry/sdk-metrics';
-import { AggregationTemporality } from '@opentelemetry/sdk-metrics';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
class FileExporter {
diff --git a/packages/core/src/telemetry/gcp-exporters.ts b/packages/core/src/telemetry/gcp-exporters.ts
index c7429383eb..3bf1781b87 100644
--- a/packages/core/src/telemetry/gcp-exporters.ts
+++ b/packages/core/src/telemetry/gcp-exporters.ts
@@ -7,10 +7,12 @@
import { type JWTInput } from 'google-auth-library';
import { TraceExporter } from '@google-cloud/opentelemetry-cloud-trace-exporter';
import { MetricExporter } from '@google-cloud/opentelemetry-cloud-monitoring-exporter';
-import { Logging } from '@google-cloud/logging';
-import type { Log } from '@google-cloud/logging';
-import { hrTimeToMilliseconds, ExportResultCode } from '@opentelemetry/core';
-import type { ExportResult } from '@opentelemetry/core';
+import { Logging, type Log } from '@google-cloud/logging';
+import {
+ hrTimeToMilliseconds,
+ ExportResultCode,
+ type ExportResult,
+} from '@opentelemetry/core';
import type {
ReadableLogRecord,
LogRecordExporter,
diff --git a/packages/core/src/telemetry/loggers.test.circular.ts b/packages/core/src/telemetry/loggers.test.circular.ts
index d6b6ea86ce..119c661e86 100644
--- a/packages/core/src/telemetry/loggers.test.circular.ts
+++ b/packages/core/src/telemetry/loggers.test.circular.ts
@@ -14,10 +14,10 @@ import { ToolCallEvent } from './types.js';
import type { Config } from '../config/config.js';
import type { CompletedToolCall } from '../core/coreToolScheduler.js';
import {
+ CoreToolCallStatus,
type ToolCallRequestInfo,
type ToolCallResponseInfo,
} from '../scheduler/types.js';
-import { CoreToolCallStatus } from '../scheduler/types.js';
import { MockTool } from '../test-utils/mock-tool.js';
describe('Circular Reference Handling', () => {
diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts
index de2f94c8d7..a3c757f5a7 100644
--- a/packages/core/src/telemetry/loggers.test.ts
+++ b/packages/core/src/telemetry/loggers.test.ts
@@ -4,14 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type {
- AnyDeclarativeTool,
- AnyToolInvocation,
- CompletedToolCall,
- ContentGeneratorConfig,
- ErroredToolCall,
- MessageBus,
-} from '../index.js';
import {
CoreToolCallStatus,
AuthType,
@@ -20,6 +12,12 @@ import {
ToolConfirmationOutcome,
ToolErrorType,
ToolRegistry,
+ type AnyDeclarativeTool,
+ type AnyToolInvocation,
+ type CompletedToolCall,
+ type ContentGeneratorConfig,
+ type ErroredToolCall,
+ type MessageBus,
} from '../index.js';
import { OutputFormat } from '../output/types.js';
import { logs } from '@opentelemetry/api-logs';
@@ -35,6 +33,7 @@ import {
logFlashFallback,
logChatCompression,
logMalformedJsonResponse,
+ logInvalidChunk,
logFileOperation,
logRipgrepFallback,
logToolOutputTruncated,
@@ -70,6 +69,7 @@ import {
EVENT_AGENT_START,
EVENT_AGENT_FINISH,
EVENT_WEB_FETCH_FALLBACK_ATTEMPT,
+ EVENT_INVALID_CHUNK,
ApiErrorEvent,
ApiRequestEvent,
ApiResponseEvent,
@@ -79,6 +79,7 @@ import {
FlashFallbackEvent,
RipgrepFallbackEvent,
MalformedJsonResponseEvent,
+ InvalidChunkEvent,
makeChatCompressionEvent,
FileOperationEvent,
ToolOutputTruncatedEvent,
@@ -1738,6 +1739,39 @@ describe('loggers', () => {
});
});
+ describe('logInvalidChunk', () => {
+ beforeEach(() => {
+ vi.spyOn(ClearcutLogger.prototype, 'logInvalidChunkEvent');
+ vi.spyOn(metrics, 'recordInvalidChunk');
+ });
+
+ it('logs the event to Clearcut and OTEL', () => {
+ const mockConfig = makeFakeConfig();
+ const event = new InvalidChunkEvent('Unexpected token');
+
+ logInvalidChunk(mockConfig, event);
+
+ expect(
+ ClearcutLogger.prototype.logInvalidChunkEvent,
+ ).toHaveBeenCalledWith(event);
+
+ expect(mockLogger.emit).toHaveBeenCalledWith({
+ body: 'Invalid chunk received from stream.',
+ attributes: {
+ 'session.id': 'test-session-id',
+ 'user.email': 'test-user@example.com',
+ 'installation.id': 'test-installation-id',
+ 'event.name': EVENT_INVALID_CHUNK,
+ 'event.timestamp': '2025-01-01T00:00:00.000Z',
+ interactive: false,
+ 'error.message': 'Unexpected token',
+ },
+ });
+
+ expect(metrics.recordInvalidChunk).toHaveBeenCalledWith(mockConfig);
+ });
+ });
+
describe('logFileOperation', () => {
const mockConfig = {
getSessionId: () => 'test-session-id',
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index e96db38596..393519d3ec 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { LogRecord } from '@opentelemetry/api-logs';
-import { logs } from '@opentelemetry/api-logs';
+import { logs, type LogRecord } from '@opentelemetry/api-logs';
import type { Config } from '../config/config.js';
import { SERVICE_NAME } from './constants.js';
import {
@@ -13,51 +12,50 @@ import {
EVENT_API_RESPONSE,
EVENT_TOOL_CALL,
EVENT_REWIND,
-} from './types.js';
-import type {
- ApiErrorEvent,
- ApiRequestEvent,
- ApiResponseEvent,
- FileOperationEvent,
- IdeConnectionEvent,
- StartSessionEvent,
- ToolCallEvent,
- UserPromptEvent,
- FlashFallbackEvent,
- NextSpeakerCheckEvent,
- LoopDetectedEvent,
- LoopDetectionDisabledEvent,
- SlashCommandEvent,
- RewindEvent,
- ConversationFinishedEvent,
- ChatCompressionEvent,
- MalformedJsonResponseEvent,
- ContentRetryEvent,
- ContentRetryFailureEvent,
- RipgrepFallbackEvent,
- ToolOutputTruncatedEvent,
- ModelRoutingEvent,
- ExtensionDisableEvent,
- ExtensionEnableEvent,
- ExtensionUninstallEvent,
- ExtensionInstallEvent,
- ModelSlashCommandEvent,
- EditStrategyEvent,
- EditCorrectionEvent,
- AgentStartEvent,
- AgentFinishEvent,
- RecoveryAttemptEvent,
- WebFetchFallbackAttemptEvent,
- ExtensionUpdateEvent,
- ApprovalModeSwitchEvent,
- ApprovalModeDurationEvent,
- HookCallEvent,
- StartupStatsEvent,
- LlmLoopCheckEvent,
- PlanExecutionEvent,
- ToolOutputMaskingEvent,
- KeychainAvailabilityEvent,
- TokenStorageInitializationEvent,
+ type ApiErrorEvent,
+ type ApiRequestEvent,
+ type ApiResponseEvent,
+ type FileOperationEvent,
+ type IdeConnectionEvent,
+ type StartSessionEvent,
+ type ToolCallEvent,
+ type UserPromptEvent,
+ type FlashFallbackEvent,
+ type NextSpeakerCheckEvent,
+ type LoopDetectedEvent,
+ type LoopDetectionDisabledEvent,
+ type SlashCommandEvent,
+ type RewindEvent,
+ type ConversationFinishedEvent,
+ type ChatCompressionEvent,
+ type MalformedJsonResponseEvent,
+ type InvalidChunkEvent,
+ type ContentRetryEvent,
+ type ContentRetryFailureEvent,
+ type RipgrepFallbackEvent,
+ type ToolOutputTruncatedEvent,
+ type ModelRoutingEvent,
+ type ExtensionDisableEvent,
+ type ExtensionEnableEvent,
+ type ExtensionUninstallEvent,
+ type ExtensionInstallEvent,
+ type ModelSlashCommandEvent,
+ type EditStrategyEvent,
+ type EditCorrectionEvent,
+ type AgentStartEvent,
+ type AgentFinishEvent,
+ type RecoveryAttemptEvent,
+ type WebFetchFallbackAttemptEvent,
+ type ExtensionUpdateEvent,
+ type ApprovalModeSwitchEvent,
+ type ApprovalModeDurationEvent,
+ type HookCallEvent,
+ type StartupStatsEvent,
+ type LlmLoopCheckEvent,
+ type PlanExecutionEvent,
+ type ToolOutputMaskingEvent,
+ type KeychainAvailabilityEvent,
+ type TokenStorageInitializationEvent,
} from './types.js';
import {
recordApiErrorMetrics,
@@ -78,10 +76,10 @@ import {
recordPlanExecution,
recordKeychainAvailability,
recordTokenStorageInitialization,
+ recordInvalidChunk,
} from './metrics.js';
import { bufferTelemetryEvent } from './sdk.js';
-import type { UiEvent } from './uiTelemetry.js';
-import { uiTelemetryService } from './uiTelemetry.js';
+import { uiTelemetryService, type UiEvent } from './uiTelemetry.js';
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { BillingTelemetryEvent } from './billingEvents.js';
@@ -471,6 +469,22 @@ export function logMalformedJsonResponse(
});
}
+export function logInvalidChunk(
+ config: Config,
+ event: InvalidChunkEvent,
+): void {
+ ClearcutLogger.getInstance(config)?.logInvalidChunkEvent(event);
+ bufferTelemetryEvent(() => {
+ const logger = logs.getLogger(SERVICE_NAME);
+ const logRecord: LogRecord = {
+ body: event.toLogBody(),
+ attributes: event.toOpenTelemetryAttributes(config),
+ };
+ logger.emit(logRecord);
+ recordInvalidChunk(config);
+ });
+}
+
export function logContentRetry(
config: Config,
event: ContentRetryEvent,
@@ -777,6 +791,7 @@ export function logStartupStats(
config: Config,
event: StartupStatsEvent,
): void {
+ ClearcutLogger.getInstance(config)?.logStartupStatsEvent(event);
bufferTelemetryEvent(() => {
// Wait for experiments to load before emitting so we capture experimentIds
void config
diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts
index d0254ec678..3b8ae1ea0c 100644
--- a/packages/core/src/telemetry/metrics.test.ts
+++ b/packages/core/src/telemetry/metrics.test.ts
@@ -105,6 +105,7 @@ describe('Telemetry Metrics', () => {
let recordPlanExecutionModule: typeof import('./metrics.js').recordPlanExecution;
let recordKeychainAvailabilityModule: typeof import('./metrics.js').recordKeychainAvailability;
let recordTokenStorageInitializationModule: typeof import('./metrics.js').recordTokenStorageInitialization;
+ let recordInvalidChunkModule: typeof import('./metrics.js').recordInvalidChunk;
beforeEach(async () => {
vi.resetModules();
@@ -154,6 +155,7 @@ describe('Telemetry Metrics', () => {
metricsJsModule.recordKeychainAvailability;
recordTokenStorageInitializationModule =
metricsJsModule.recordTokenStorageInitialization;
+ recordInvalidChunkModule = metricsJsModule.recordInvalidChunk;
const otelApiModule = await import('@opentelemetry/api');
@@ -1555,5 +1557,27 @@ describe('Telemetry Metrics', () => {
});
});
});
+
+ describe('recordInvalidChunk', () => {
+ it('should not record metrics if not initialized', () => {
+ const config = makeFakeConfig({});
+ recordInvalidChunkModule(config);
+ expect(mockCounterAddFn).not.toHaveBeenCalled();
+ });
+
+ it('should record invalid chunk when initialized', () => {
+ const config = makeFakeConfig({});
+ initializeMetricsModule(config);
+ mockCounterAddFn.mockClear();
+
+ recordInvalidChunkModule(config);
+
+ expect(mockCounterAddFn).toHaveBeenCalledWith(1, {
+ 'session.id': 'test-session-id',
+ 'installation.id': 'test-installation-id',
+ 'user.email': 'test@example.com',
+ });
+ });
+ });
});
});
diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts
index 598158af07..70b188f517 100644
--- a/packages/core/src/telemetry/metrics.ts
+++ b/packages/core/src/telemetry/metrics.ts
@@ -4,8 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { Attributes, Meter, Counter, Histogram } from '@opentelemetry/api';
-import { diag, metrics, ValueType } from '@opentelemetry/api';
+import {
+ diag,
+ metrics,
+ ValueType,
+ type Attributes,
+ type Meter,
+ type Counter,
+ type Histogram,
+} from '@opentelemetry/api';
import { SERVICE_NAME } from './constants.js';
import type { Config } from '../config/config.js';
import type {
diff --git a/packages/core/src/telemetry/sdk.test.ts b/packages/core/src/telemetry/sdk.test.ts
index e6d3f26ae7..9636212e2c 100644
--- a/packages/core/src/telemetry/sdk.test.ts
+++ b/packages/core/src/telemetry/sdk.test.ts
@@ -113,13 +113,13 @@ describe('Telemetry SDK', () => {
await initializeTelemetry(mockConfig);
expect(OTLPTraceExporterHttp).toHaveBeenCalledWith({
- url: 'http://localhost:4318/',
+ url: 'http://localhost:4318/v1/traces',
});
expect(OTLPLogExporterHttp).toHaveBeenCalledWith({
- url: 'http://localhost:4318/',
+ url: 'http://localhost:4318/v1/logs',
});
expect(OTLPMetricExporterHttp).toHaveBeenCalledWith({
- url: 'http://localhost:4318/',
+ url: 'http://localhost:4318/v1/metrics',
});
expect(NodeSDK.prototype.start).toHaveBeenCalled();
});
@@ -141,7 +141,7 @@ describe('Telemetry SDK', () => {
);
await initializeTelemetry(mockConfig);
expect(OTLPTraceExporterHttp).toHaveBeenCalledWith(
- expect.objectContaining({ url: 'https://my-collector.com/' }),
+ expect.objectContaining({ url: 'https://my-collector.com/v1/traces' }),
);
});
diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts
index 66b89523db..3752d3e40f 100644
--- a/packages/core/src/telemetry/sdk.ts
+++ b/packages/core/src/telemetry/sdk.ts
@@ -275,15 +275,21 @@ export async function initializeTelemetry(
});
} else if (useOtlp) {
if (otlpProtocol === 'http') {
+ const buildUrl = (path: string) => {
+ const url = new URL(parsedEndpoint);
+ // Join the existing pathname with the new path, handling trailing slashes.
+ url.pathname = [url.pathname.replace(/\/$/, ''), path].join('/');
+ return url.href;
+ };
spanExporter = new OTLPTraceExporterHttp({
- url: parsedEndpoint,
+ url: buildUrl('v1/traces'),
});
logExporter = new OTLPLogExporterHttp({
- url: parsedEndpoint,
+ url: buildUrl('v1/logs'),
});
metricReader = new PeriodicExportingMetricReader({
exporter: new OTLPMetricExporterHttp({
- url: parsedEndpoint,
+ url: buildUrl('v1/metrics'),
}),
exportIntervalMillis: 10000,
});
diff --git a/packages/core/src/telemetry/semantic.ts b/packages/core/src/telemetry/semantic.ts
index c05d110e9f..6dae06d381 100644
--- a/packages/core/src/telemetry/semantic.ts
+++ b/packages/core/src/telemetry/semantic.ts
@@ -11,13 +11,13 @@
* @see https://github.com/open-telemetry/semantic-conventions/blob/8b4f210f43136e57c1f6f47292eb6d38e3bf30bb/docs/gen-ai/gen-ai-events.md
*/
-import { FinishReason } from '@google/genai';
-import type {
- Candidate,
- Content,
- ContentUnion,
- Part,
- PartUnion,
+import {
+ FinishReason,
+ type Candidate,
+ type Content,
+ type ContentUnion,
+ type Part,
+ type PartUnion,
} from '@google/genai';
import { truncateString } from '../utils/textUtils.js';
@@ -63,6 +63,7 @@ function getStringReferences(parts: AnyPart[]): StringReference[] {
});
}
} else if (part instanceof GenericPart) {
+ // eslint-disable-next-line no-restricted-syntax
if (part.type === 'executableCode' && typeof part['code'] === 'string') {
refs.push({
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
@@ -73,6 +74,7 @@ function getStringReferences(parts: AnyPart[]): StringReference[] {
});
} else if (
part.type === 'codeExecutionResult' &&
+ // eslint-disable-next-line no-restricted-syntax
typeof part['output'] === 'string'
) {
refs.push({
diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts
index a4b3cfb4c9..43317f8baa 100644
--- a/packages/core/src/telemetry/types.ts
+++ b/packages/core/src/telemetry/types.ts
@@ -31,13 +31,13 @@ import type { AgentTerminateMode } from '../agents/types.js';
import { getCommonAttributes } from './telemetryAttributes.js';
import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
-import type { OTelFinishReason } from './semantic.js';
import {
toInputMessages,
toOutputMessages,
toFinishReasons,
toOutputType,
toSystemInstruction,
+ type OTelFinishReason,
} from './semantic.js';
import { sanitizeHookName } from './sanitize.js';
import { getFileDiffFromResultDisplay } from '../utils/fileDiffUtils.js';
@@ -790,25 +790,36 @@ export enum LoopType {
CONSECUTIVE_IDENTICAL_TOOL_CALLS = 'consecutive_identical_tool_calls',
CHANTING_IDENTICAL_SENTENCES = 'chanting_identical_sentences',
LLM_DETECTED_LOOP = 'llm_detected_loop',
+ // Aliases for tests/internal use
+ TOOL_CALL_LOOP = CONSECUTIVE_IDENTICAL_TOOL_CALLS,
+ CONTENT_CHANTING_LOOP = CHANTING_IDENTICAL_SENTENCES,
}
-
export class LoopDetectedEvent implements BaseTelemetryEvent {
'event.name': 'loop_detected';
'event.timestamp': string;
loop_type: LoopType;
prompt_id: string;
+ count: number;
confirmed_by_model?: string;
+ analysis?: string;
+ confidence?: number;
constructor(
loop_type: LoopType,
prompt_id: string,
+ count: number,
confirmed_by_model?: string,
+ analysis?: string,
+ confidence?: number,
) {
this['event.name'] = 'loop_detected';
this['event.timestamp'] = new Date().toISOString();
this.loop_type = loop_type;
this.prompt_id = prompt_id;
+ this.count = count;
this.confirmed_by_model = confirmed_by_model;
+ this.analysis = analysis;
+ this.confidence = confidence;
}
toOpenTelemetryAttributes(config: Config): LogAttributes {
@@ -818,17 +829,28 @@ export class LoopDetectedEvent implements BaseTelemetryEvent {
'event.timestamp': this['event.timestamp'],
loop_type: this.loop_type,
prompt_id: this.prompt_id,
+ count: this.count,
};
if (this.confirmed_by_model) {
attributes['confirmed_by_model'] = this.confirmed_by_model;
}
+ if (this.analysis) {
+ attributes['analysis'] = this.analysis;
+ }
+
+ if (this.confidence !== undefined) {
+ attributes['confidence'] = this.confidence;
+ }
+
return attributes;
}
toLogBody(): string {
- return `Loop detected. Type: ${this.loop_type}.${this.confirmed_by_model ? ` Confirmed by: ${this.confirmed_by_model}` : ''}`;
+ const status =
+ this.count === 1 ? 'Attempting recovery' : 'Terminating session';
+ return `Loop detected (Strike ${this.count}: ${status}). Type: ${this.loop_type}.${this.confirmed_by_model ? ` Confirmed by: ${this.confirmed_by_model}` : ''}`;
}
}
diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts
index d1a3b1a9a6..f78f0801af 100644
--- a/packages/core/src/telemetry/uiTelemetry.test.ts
+++ b/packages/core/src/telemetry/uiTelemetry.test.ts
@@ -7,12 +7,13 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { UiTelemetryService } from './uiTelemetry.js';
import { ToolCallDecision } from './tool-call-decision.js';
-import type { ApiErrorEvent, ApiResponseEvent } from './types.js';
import {
ToolCallEvent,
EVENT_API_ERROR,
EVENT_API_RESPONSE,
EVENT_TOOL_CALL,
+ type ApiErrorEvent,
+ type ApiResponseEvent,
} from './types.js';
import type {
CompletedToolCall,
diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts
index 669b6a8c68..36953c02c1 100644
--- a/packages/core/src/telemetry/uiTelemetry.ts
+++ b/packages/core/src/telemetry/uiTelemetry.ts
@@ -9,15 +9,13 @@ import {
EVENT_API_ERROR,
EVENT_API_RESPONSE,
EVENT_TOOL_CALL,
+ type ApiErrorEvent,
+ type ApiResponseEvent,
+ type ToolCallEvent,
+ type LlmRole,
} from './types.js';
import { ToolCallDecision } from './tool-call-decision.js';
-import type {
- ApiErrorEvent,
- ApiResponseEvent,
- ToolCallEvent,
- LlmRole,
-} from './types.js';
export type UiEvent =
| (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE })
diff --git a/packages/core/src/test-utils/config.ts b/packages/core/src/test-utils/config.ts
index 880599d9b9..5d896752f9 100644
--- a/packages/core/src/test-utils/config.ts
+++ b/packages/core/src/test-utils/config.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { ConfigParameters } from '../config/config.js';
-import { Config } from '../config/config.js';
+import { Config, type ConfigParameters } from '../config/config.js';
/**
* Default parameters used for {@link FAKE_CONFIG}
diff --git a/packages/core/src/test-utils/mock-tool.ts b/packages/core/src/test-utils/mock-tool.ts
index 4fa536d2db..5f89a506cd 100644
--- a/packages/core/src/test-utils/mock-tool.ts
+++ b/packages/core/src/test-utils/mock-tool.ts
@@ -8,15 +8,13 @@ import type {
ModifiableDeclarativeTool,
ModifyContext,
} from '../tools/modifiable-tool.js';
-import type {
- ToolCallConfirmationDetails,
- ToolInvocation,
- ToolResult,
-} from '../tools/tools.js';
import {
BaseDeclarativeTool,
BaseToolInvocation,
Kind,
+ type ToolCallConfirmationDetails,
+ type ToolInvocation,
+ type ToolResult,
} from '../tools/tools.js';
import { createMockMessageBus } from './mock-message-bus.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
diff --git a/packages/core/src/tools/activate-skill.ts b/packages/core/src/tools/activate-skill.ts
index cf6a33f3e6..21ee2e98c6 100644
--- a/packages/core/src/tools/activate-skill.ts
+++ b/packages/core/src/tools/activate-skill.ts
@@ -7,13 +7,15 @@
import * as path from 'node:path';
import { getFolderStructure } from '../utils/getFolderStructure.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
-import type {
- ToolResult,
- ToolCallConfirmationDetails,
- ToolInvocation,
- ToolConfirmationOutcome,
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolResult,
+ type ToolCallConfirmationDetails,
+ type ToolInvocation,
+ type ToolConfirmationOutcome,
} from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
import type { Config } from '../config/config.js';
import { ACTIVATE_SKILL_TOOL_NAME } from './tool-names.js';
import { ToolErrorType } from './tool-error.js';
diff --git a/packages/core/src/tools/definitions/trackerTools.ts b/packages/core/src/tools/definitions/trackerTools.ts
new file mode 100644
index 0000000000..e136d90d04
--- /dev/null
+++ b/packages/core/src/tools/definitions/trackerTools.ts
@@ -0,0 +1,161 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { ToolDefinition } from './types.js';
+import {
+ TRACKER_CREATE_TASK_TOOL_NAME,
+ TRACKER_UPDATE_TASK_TOOL_NAME,
+ TRACKER_GET_TASK_TOOL_NAME,
+ TRACKER_LIST_TASKS_TOOL_NAME,
+ TRACKER_ADD_DEPENDENCY_TOOL_NAME,
+ TRACKER_VISUALIZE_TOOL_NAME,
+} from '../tool-names.js';
+
+export const TRACKER_CREATE_TASK_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_CREATE_TASK_TOOL_NAME,
+ description: 'Creates a new task in the tracker.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {
+ title: {
+ type: 'string',
+ description: 'Short title of the task.',
+ },
+ description: {
+ type: 'string',
+ description: 'Detailed description of the task.',
+ },
+ type: {
+ type: 'string',
+ enum: ['epic', 'task', 'bug'],
+ description: 'Type of the task.',
+ },
+ parentId: {
+ type: 'string',
+ description: 'Optional ID of the parent task.',
+ },
+ dependencies: {
+ type: 'array',
+ items: { type: 'string' },
+ description: 'Optional list of task IDs that this task depends on.',
+ },
+ },
+ required: ['title', 'description', 'type'],
+ },
+ },
+};
+
+export const TRACKER_UPDATE_TASK_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_UPDATE_TASK_TOOL_NAME,
+ description: 'Updates an existing task in the tracker.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {
+ id: {
+ type: 'string',
+ description: 'The 6-character hex ID of the task to update.',
+ },
+ title: {
+ type: 'string',
+ description: 'New title for the task.',
+ },
+ description: {
+ type: 'string',
+ description: 'New description for the task.',
+ },
+ status: {
+ type: 'string',
+ enum: ['open', 'in_progress', 'blocked', 'closed'],
+ description: 'New status for the task.',
+ },
+ dependencies: {
+ type: 'array',
+ items: { type: 'string' },
+ description: 'New list of dependency IDs.',
+ },
+ },
+ required: ['id'],
+ },
+ },
+};
+
+export const TRACKER_GET_TASK_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_GET_TASK_TOOL_NAME,
+ description: 'Retrieves details for a specific task.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {
+ id: {
+ type: 'string',
+ description: 'The 6-character hex ID of the task.',
+ },
+ },
+ required: ['id'],
+ },
+ },
+};
+
+export const TRACKER_LIST_TASKS_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_LIST_TASKS_TOOL_NAME,
+ description:
+ 'Lists tasks in the tracker, optionally filtered by status, type, or parent.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {
+ status: {
+ type: 'string',
+ enum: ['open', 'in_progress', 'blocked', 'closed'],
+ description: 'Filter by status.',
+ },
+ type: {
+ type: 'string',
+ enum: ['epic', 'task', 'bug'],
+ description: 'Filter by type.',
+ },
+ parentId: {
+ type: 'string',
+ description: 'Filter by parent task ID.',
+ },
+ },
+ },
+ },
+};
+
+export const TRACKER_ADD_DEPENDENCY_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_ADD_DEPENDENCY_TOOL_NAME,
+ description: 'Adds a dependency between two tasks.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {
+ taskId: {
+ type: 'string',
+ description: 'The ID of the task that has a dependency.',
+ },
+ dependencyId: {
+ type: 'string',
+ description: 'The ID of the task that is being depended upon.',
+ },
+ },
+ required: ['taskId', 'dependencyId'],
+ },
+ },
+};
+
+export const TRACKER_VISUALIZE_DEFINITION: ToolDefinition = {
+ base: {
+ name: TRACKER_VISUALIZE_TOOL_NAME,
+ description: 'Renders an ASCII tree visualization of the task graph.',
+ parametersJsonSchema: {
+ type: 'object',
+ properties: {},
+ },
+ },
+};
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index a7169e99f2..214875c574 100644
--- a/packages/core/src/tools/edit.ts
+++ b/packages/core/src/tools/edit.ts
@@ -413,6 +413,20 @@ export interface EditToolParams {
ai_proposed_content?: string;
}
+export function isEditToolParams(args: unknown): args is EditToolParams {
+ if (typeof args !== 'object' || args === null) {
+ return false;
+ }
+ return (
+ 'file_path' in args &&
+ typeof args.file_path === 'string' &&
+ 'old_string' in args &&
+ typeof args.old_string === 'string' &&
+ 'new_string' in args &&
+ typeof args.new_string === 'string'
+ );
+}
+
interface CalculatedEdit {
currentContent: string | null;
newContent: string;
diff --git a/packages/core/src/tools/glob.test.ts b/packages/core/src/tools/glob.test.ts
index 2aa4d52c7e..f3390f5d3c 100644
--- a/packages/core/src/tools/glob.test.ts
+++ b/packages/core/src/tools/glob.test.ts
@@ -4,8 +4,12 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { GlobToolParams, GlobPath } from './glob.js';
-import { GlobTool, sortFileEntries } from './glob.js';
+import {
+ GlobTool,
+ sortFileEntries,
+ type GlobToolParams,
+ type GlobPath,
+} from './glob.js';
import { partListUnionToString } from '../core/geminiRequest.js';
import path from 'node:path';
import { isSubpath } from '../utils/paths.js';
diff --git a/packages/core/src/tools/glob.ts b/packages/core/src/tools/glob.ts
index 78b445e762..c2f3c4ab54 100644
--- a/packages/core/src/tools/glob.ts
+++ b/packages/core/src/tools/glob.ts
@@ -8,8 +8,13 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js';
import fs from 'node:fs';
import path from 'node:path';
import { glob, escape } from 'glob';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { shortenPath, makeRelative } from '../utils/paths.js';
import { type Config } from '../config/config.js';
import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js';
diff --git a/packages/core/src/tools/grep-utils.ts b/packages/core/src/tools/grep-utils.ts
index 6dd2cdc83e..2191588301 100644
--- a/packages/core/src/tools/grep-utils.ts
+++ b/packages/core/src/tools/grep-utils.ts
@@ -6,6 +6,7 @@
import fsPromises from 'node:fs/promises';
import { debugLogger } from '../utils/debugLogger.js';
+import { MAX_LINE_LENGTH_TEXT_FILE } from '../utils/constants.js';
/**
* Result object for a single grep match
@@ -198,7 +199,14 @@ export async function formatGrepResults(
// If isContext is undefined, assume it's a match (false)
const separator = match.isContext ? '-' : ':';
// trimEnd to avoid double newlines if line has them, but we want to preserve indentation
- llmContent += `L${match.lineNumber}${separator} ${match.line.trimEnd()}\n`;
+ let lineContent = match.line.trimEnd();
+ const graphemes = Array.from(lineContent);
+ if (graphemes.length > MAX_LINE_LENGTH_TEXT_FILE) {
+ lineContent =
+ graphemes.slice(0, MAX_LINE_LENGTH_TEXT_FILE).join('') +
+ '... [truncated]';
+ }
+ llmContent += `L${match.lineNumber}${separator} ${lineContent}\n`;
});
llmContent += '---\n';
}
diff --git a/packages/core/src/tools/grep.test.ts b/packages/core/src/tools/grep.test.ts
index 6f98b0f2fc..1f0a8ee98f 100644
--- a/packages/core/src/tools/grep.test.ts
+++ b/packages/core/src/tools/grep.test.ts
@@ -5,8 +5,7 @@
*/
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
-import type { GrepToolParams } from './grep.js';
-import { GrepTool } from './grep.js';
+import { GrepTool, type GrepToolParams } from './grep.js';
import type { ToolResult } from './tools.js';
import path from 'node:path';
import { isSubpath } from '../utils/paths.js';
@@ -563,6 +562,22 @@ describe('GrepTool', () => {
// Verify context after
expect(result.llmContent).toContain('L60- Line 60');
});
+
+ it('should truncate excessively long lines', async () => {
+ const longString = 'a'.repeat(3000);
+ await fs.writeFile(
+ path.join(tempRootDir, 'longline.txt'),
+ `Target match ${longString}`,
+ );
+
+ const params: GrepToolParams = { pattern: 'Target match' };
+ const invocation = grepTool.build(params);
+ const result = await invocation.execute(abortSignal);
+
+ // MAX_LINE_LENGTH_TEXT_FILE is 2000. It should be truncated.
+ expect(result.llmContent).toContain('... [truncated]');
+ expect(result.llmContent).not.toContain(longString);
+ });
});
describe('getDescription', () => {
diff --git a/packages/core/src/tools/grep.ts b/packages/core/src/tools/grep.ts
index 3d74521513..c7e676951a 100644
--- a/packages/core/src/tools/grep.ts
+++ b/packages/core/src/tools/grep.ts
@@ -10,13 +10,18 @@ import fsPromises from 'node:fs/promises';
import path from 'node:path';
import { spawn } from 'node:child_process';
import { globStream } from 'glob';
-import type { ToolInvocation, ToolResult } from './tools.js';
import { execStreaming } from '../utils/shell-utils.js';
import {
DEFAULT_TOTAL_MAX_MATCHES,
DEFAULT_SEARCH_TIMEOUT_MS,
} from './constants.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import { getErrorMessage, isNodeError } from '../utils/errors.js';
import { isGitRepository } from '../utils/gitUtils.js';
diff --git a/packages/core/src/tools/ls.ts b/packages/core/src/tools/ls.ts
index b98dfb9e38..9456f8ffc9 100644
--- a/packages/core/src/tools/ls.ts
+++ b/packages/core/src/tools/ls.ts
@@ -7,8 +7,13 @@
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import fs from 'node:fs/promises';
import path from 'node:path';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import type { Config } from '../config/config.js';
import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js';
diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts
index 96d7abf55c..43ea9715bc 100644
--- a/packages/core/src/tools/mcp-client-manager.ts
+++ b/packages/core/src/tools/mcp-client-manager.ts
@@ -173,7 +173,7 @@ export class McpClientManager {
return Promise.resolve();
}),
);
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
/**
@@ -193,7 +193,7 @@ export class McpClientManager {
}),
),
);
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
/**
@@ -251,7 +251,7 @@ export class McpClientManager {
if (!skipRefresh) {
// This is required to update the content generator configuration with the
// new tool configuration and system instructions.
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
}
}
@@ -321,7 +321,7 @@ export class McpClientManager {
this.cliConfig.getDebugMode(),
this.clientVersion,
async () => {
- debugLogger.log('Tools changed, updating Gemini context...');
+ debugLogger.log(`🔔 Refreshing context for server '${name}'...`);
await this.scheduleMcpContextRefresh();
},
);
@@ -431,7 +431,7 @@ export class McpClientManager {
this.eventEmitter?.emit('mcp-client-update', this.clients);
}
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
/**
@@ -451,7 +451,7 @@ export class McpClientManager {
},
),
);
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
/**
@@ -463,7 +463,7 @@ export class McpClientManager {
throw new Error(`No MCP server registered with the name "${name}"`);
}
await this.maybeDiscoverMcpServer(name, config);
- await this.cliConfig.refreshMcpContext();
+ await this.scheduleMcpContextRefresh();
}
/**
@@ -517,21 +517,51 @@ export class McpClientManager {
return instructions.join('\n\n');
}
+ private isRefreshingMcpContext: boolean = false;
+ private pendingMcpContextRefresh: boolean = false;
+
private async scheduleMcpContextRefresh(): Promise {
+ this.pendingMcpContextRefresh = true;
+
+ if (this.isRefreshingMcpContext) {
+ debugLogger.log(
+ 'MCP context refresh already in progress, queuing trailing execution.',
+ );
+ return this.pendingRefreshPromise ?? Promise.resolve();
+ }
+
if (this.pendingRefreshPromise) {
+ debugLogger.log(
+ 'MCP context refresh already scheduled, coalescing with existing request.',
+ );
return this.pendingRefreshPromise;
}
+ debugLogger.log('Scheduling MCP context refresh...');
this.pendingRefreshPromise = (async () => {
- // Debounce to coalesce multiple rapid updates
- await new Promise((resolve) => setTimeout(resolve, 300));
+ this.isRefreshingMcpContext = true;
try {
- await this.cliConfig.refreshMcpContext();
+ do {
+ this.pendingMcpContextRefresh = false;
+ debugLogger.log('Executing MCP context refresh...');
+ await this.cliConfig.refreshMcpContext();
+ debugLogger.log('MCP context refresh complete.');
+
+ // If more refresh requests came in during the execution, wait a bit
+ // to coalesce them before the next iteration.
+ if (this.pendingMcpContextRefresh) {
+ debugLogger.log(
+ 'Coalescing burst refresh requests (300ms delay)...',
+ );
+ await new Promise((resolve) => setTimeout(resolve, 300));
+ }
+ } while (this.pendingMcpContextRefresh);
} catch (error) {
debugLogger.error(
`Error refreshing MCP context: ${getErrorMessage(error)}`,
);
} finally {
+ this.isRefreshingMcpContext = false;
this.pendingRefreshPromise = null;
}
})();
diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts
index 126fb7ce68..8612a838ca 100644
--- a/packages/core/src/tools/mcp-client.test.ts
+++ b/packages/core/src/tools/mcp-client.test.ts
@@ -22,6 +22,7 @@ import {
PromptListChangedNotificationSchema,
ResourceListChangedNotificationSchema,
ToolListChangedNotificationSchema,
+ ProgressNotificationSchema,
} from '@modelcontextprotocol/sdk/types.js';
import type { DiscoveredMCPTool } from './mcp-tool.js';
@@ -102,6 +103,7 @@ describe('mcp-client', () => {
afterEach(() => {
vi.restoreAllMocks();
+ vi.useRealTimers();
});
describe('McpClient', () => {
@@ -140,13 +142,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -221,13 +226,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -328,13 +336,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -388,13 +399,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -701,13 +715,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -778,13 +795,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -864,13 +884,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -950,13 +973,16 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
+ getPromptsByServer: vi.fn().mockReturnValue([]),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
@@ -1086,6 +1112,7 @@ describe('mcp-client', () => {
setNotificationHandler: vi.fn(),
listTools: vi.fn().mockResolvedValue({ tools: [] }),
listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+ listResources: vi.fn().mockResolvedValue({ resources: [] }),
request: vi.fn().mockResolvedValue({}),
};
@@ -1096,12 +1123,27 @@ describe('mcp-client', () => {
{} as SdkClientStdioLib.StdioClientTransport,
);
+ const mockedToolRegistry = {
+ registerTool: vi.fn(),
+ sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
+ getMessageBus: vi.fn().mockReturnValue(undefined),
+ } as unknown as ToolRegistry;
+
const client = new McpClient(
'test-server',
{ command: 'test-command' },
- {} as ToolRegistry,
- {} as PromptRegistry,
- {} as ResourceRegistry,
+ mockedToolRegistry,
+ {
+ getPromptsByServer: vi.fn().mockReturnValue([]),
+ registerPrompt: vi.fn(),
+ } as unknown as PromptRegistry,
+ {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
+ registerResource: vi.fn(),
+ removeResourcesByServer: vi.fn(),
+ setResourcesForServer: vi.fn(),
+ } as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1136,9 +1178,21 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
- {} as ToolRegistry,
- {} as PromptRegistry,
- {} as ResourceRegistry,
+ {
+ getToolsByServer: vi.fn().mockReturnValue([]),
+ registerTool: vi.fn(),
+ sortTools: vi.fn(),
+ } as unknown as ToolRegistry,
+ {
+ getPromptsByServer: vi.fn().mockReturnValue([]),
+ registerPrompt: vi.fn(),
+ } as unknown as PromptRegistry,
+ {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
+ registerResource: vi.fn(),
+ removeResourcesByServer: vi.fn(),
+ setResourcesForServer: vi.fn(),
+ } as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1147,7 +1201,62 @@ describe('mcp-client', () => {
await client.connect();
- expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
+ // Should be called for ProgressNotificationSchema, even if no other capabilities
+ expect(mockedClient.setNotificationHandler).toHaveBeenCalled();
+ const progressCall = mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ProgressNotificationSchema,
+ );
+ expect(progressCall).toBeDefined();
+ });
+
+ it('should set up notification handler even if listChanged is false (robustness)', async () => {
+ // Setup mocks
+ const mockedClient = {
+ connect: vi.fn(),
+ getServerCapabilities: vi
+ .fn()
+ .mockReturnValue({ tools: { listChanged: false } }),
+ setNotificationHandler: vi.fn(),
+ request: vi.fn().mockResolvedValue({}),
+ registerCapabilities: vi.fn().mockResolvedValue({}),
+ setRequestHandler: vi.fn().mockResolvedValue({}),
+ };
+
+ vi.mocked(ClientLib.Client).mockReturnValue(
+ mockedClient as unknown as ClientLib.Client,
+ );
+
+ const client = new McpClient(
+ 'test-server',
+ { command: 'test-command' },
+ {
+ getToolsByServer: vi.fn().mockReturnValue([]),
+ registerTool: vi.fn(),
+ sortTools: vi.fn(),
+ } as unknown as ToolRegistry,
+ {
+ getPromptsByServer: vi.fn().mockReturnValue([]),
+ registerPrompt: vi.fn(),
+ } as unknown as PromptRegistry,
+ {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
+ registerResource: vi.fn(),
+ removeResourcesByServer: vi.fn(),
+ setResourcesForServer: vi.fn(),
+ } as unknown as ResourceRegistry,
+ workspaceContext,
+ MOCK_CONTEXT,
+ false,
+ '0.0.1',
+ );
+
+ await client.connect();
+
+ const toolUpdateCall =
+ mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ expect(toolUpdateCall).toBeDefined();
});
it('should refresh tools and notify manager when notification is received', async () => {
@@ -1167,6 +1276,7 @@ describe('mcp-client', () => {
],
}),
listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+ listResources: vi.fn().mockResolvedValue({ resources: [] }),
request: vi.fn().mockResolvedValue({}),
registerCapabilities: vi.fn().mockResolvedValue({}),
setRequestHandler: vi.fn().mockResolvedValue({}),
@@ -1183,31 +1293,38 @@ describe('mcp-client', () => {
removeMcpToolsByServer: vi.fn(),
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
- const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+ const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
- // Initialize client with onToolsUpdated callback
+ // Initialize client with onContextUpdated callback
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
- {} as ResourceRegistry,
+ {
+ removeMcpResourcesByServer: vi.fn(),
+ registerResource: vi.fn(),
+ } as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
- onToolsUpdatedSpy,
+ onContextUpdatedSpy,
);
// 1. Connect (sets up listener)
await client.connect();
- // 2. Extract the callback passed to setNotificationHandler
- const notificationCallback =
- mockedClient.setNotificationHandler.mock.calls[0][1];
+ // 2. Extract the callback passed to setNotificationHandler for tools
+ const toolUpdateCall =
+ mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const notificationCallback = toolUpdateCall![1];
// 3. Trigger the notification manually
await notificationCallback();
@@ -1225,7 +1342,7 @@ describe('mcp-client', () => {
expect(mockedToolRegistry.registerTool).toHaveBeenCalled();
// It should notify the manager
- expect(onToolsUpdatedSpy).toHaveBeenCalled();
+ expect(onContextUpdatedSpy).toHaveBeenCalled();
// It should emit feedback event
expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
@@ -1259,6 +1376,7 @@ describe('mcp-client', () => {
const mockedToolRegistry = {
removeMcpToolsByServer: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
@@ -1276,8 +1394,11 @@ describe('mcp-client', () => {
await client.connect();
- const notificationCallback =
- mockedClient.setNotificationHandler.mock.calls[0][1];
+ const toolUpdateCall =
+ mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const notificationCallback = toolUpdateCall![1];
// Trigger notification - should fail internally but catch the error
await notificationCallback();
@@ -1328,10 +1449,11 @@ describe('mcp-client', () => {
removeMcpToolsByServer: vi.fn(),
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
- const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+ const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
const clientA = new McpClient(
'server-A',
@@ -1343,7 +1465,7 @@ describe('mcp-client', () => {
MOCK_CONTEXT,
false,
'0.0.1',
- onToolsUpdatedSpy,
+ onContextUpdatedSpy,
);
const clientB = new McpClient(
@@ -1356,14 +1478,23 @@ describe('mcp-client', () => {
MOCK_CONTEXT,
false,
'0.0.1',
- onToolsUpdatedSpy,
+ onContextUpdatedSpy,
);
await clientA.connect();
await clientB.connect();
- const handlerA = mockClientA.setNotificationHandler.mock.calls[0][1];
- const handlerB = mockClientB.setNotificationHandler.mock.calls[0][1];
+ const toolUpdateCallA =
+ mockClientA.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const handlerA = toolUpdateCallA![1];
+
+ const toolUpdateCallB =
+ mockClientB.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const handlerB = toolUpdateCallB![1];
// Trigger burst updates simultaneously
await Promise.all([handlerA(), handlerB()]);
@@ -1383,12 +1514,11 @@ describe('mcp-client', () => {
expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
// Verify the update callback was triggered for both
- expect(onToolsUpdatedSpy).toHaveBeenCalledTimes(2);
+ expect(onContextUpdatedSpy).toHaveBeenCalledTimes(2);
});
it('should abort discovery and log error if timeout is exceeded during refresh', async () => {
vi.useFakeTimers();
-
const mockedClient = {
connect: vi.fn(),
getServerCapabilities: vi
@@ -1412,6 +1542,7 @@ describe('mcp-client', () => {
}),
),
listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+ listResources: vi.fn().mockResolvedValue({ resources: [] }),
request: vi.fn().mockResolvedValue({}),
registerCapabilities: vi.fn().mockResolvedValue({}),
setRequestHandler: vi.fn().mockResolvedValue({}),
@@ -1428,16 +1559,26 @@ describe('mcp-client', () => {
removeMcpToolsByServer: vi.fn(),
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const client = new McpClient(
'test-server',
- // Set a short timeout
- { command: 'test-command', timeout: 100 },
+ // Set a very short timeout
+ { command: 'test-command', timeout: 50 },
mockedToolRegistry,
- {} as PromptRegistry,
- {} as ResourceRegistry,
+ {
+ getPromptsByServer: vi.fn().mockReturnValue([]),
+ registerPrompt: vi.fn(),
+ removePromptsByServer: vi.fn(),
+ } as unknown as PromptRegistry,
+ {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
+ registerResource: vi.fn(),
+ removeResourcesByServer: vi.fn(),
+ setResourcesForServer: vi.fn(),
+ } as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1446,13 +1587,16 @@ describe('mcp-client', () => {
await client.connect();
- const notificationCallback =
- mockedClient.setNotificationHandler.mock.calls[0][1];
+ const toolUpdateCall =
+ mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const notificationCallback = toolUpdateCall![1];
const refreshPromise = notificationCallback();
- vi.advanceTimersByTime(150);
-
+ // Advance timers to trigger the timeout (11 minutes to cover even the default timeout)
+ await vi.advanceTimersByTimeAsync(11 * 60 * 1000);
await refreshPromise;
expect(mockedClient.listTools).toHaveBeenCalledWith(
@@ -1463,8 +1607,6 @@ describe('mcp-client', () => {
);
expect(mockedToolRegistry.registerTool).not.toHaveBeenCalled();
-
- vi.useRealTimers();
});
it('should pass abort signal to onToolsUpdated callback', async () => {
@@ -1492,35 +1634,51 @@ describe('mcp-client', () => {
removeMcpToolsByServer: vi.fn(),
registerTool: vi.fn(),
sortTools: vi.fn(),
+ getToolsByServer: vi.fn().mockReturnValue([]),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
- const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+ const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
- {} as PromptRegistry,
- {} as ResourceRegistry,
+ {
+ getPromptsByServer: vi.fn().mockReturnValue([]),
+ registerPrompt: vi.fn(),
+ removePromptsByServer: vi.fn(),
+ } as unknown as PromptRegistry,
+ {
+ getResourcesByServer: vi.fn().mockReturnValue([]),
+ registerResource: vi.fn(),
+ removeResourcesByServer: vi.fn(),
+ setResourcesForServer: vi.fn(),
+ } as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
- onToolsUpdatedSpy,
+ onContextUpdatedSpy,
);
await client.connect();
- const notificationCallback =
- mockedClient.setNotificationHandler.mock.calls[0][1];
+ const toolUpdateCall =
+ mockedClient.setNotificationHandler.mock.calls.find(
+ (call) => call[0] === ToolListChangedNotificationSchema,
+ );
+ const notificationCallback = toolUpdateCall![1];
- await notificationCallback();
+ vi.useFakeTimers();
+ const refreshPromise = notificationCallback();
+ await vi.advanceTimersByTimeAsync(500);
+ await refreshPromise;
- expect(onToolsUpdatedSpy).toHaveBeenCalledWith(expect.any(AbortSignal));
+ expect(onContextUpdatedSpy).toHaveBeenCalledWith(expect.any(AbortSignal));
// Verify the signal passed was not aborted (happy path)
- const signal = onToolsUpdatedSpy.mock.calls[0][0];
+ const signal = onContextUpdatedSpy.mock.calls[0][0];
expect(signal.aborted).toBe(false);
});
});
@@ -1895,7 +2053,6 @@ describe('mcp-client', () => {
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('ext-value');
expect(callArgs.env!['RESOLVED_VAR']).toBe('ext-value');
});
-
it('should expand environment variables in mcpServerConfig.env and not redact them', async () => {
const mockedTransport = vi
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index 24f93052bf..af55facaa3 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/packages/core/src/tools/mcp-client.ts
@@ -11,19 +11,16 @@ import type {
JsonSchemaType,
JsonSchemaValidator,
} from '@modelcontextprotocol/sdk/validation/types.js';
-import type { SSEClientTransportOptions } from '@modelcontextprotocol/sdk/client/sse.js';
-import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
+import {
+ SSEClientTransport,
+ type SSEClientTransportOptions,
+} from '@modelcontextprotocol/sdk/client/sse.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
-import type { StreamableHTTPClientTransportOptions } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import {
+ StreamableHTTPClientTransport,
+ type StreamableHTTPClientTransportOptions,
+} from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
-import type {
- GetPromptResult,
- Prompt,
- ReadResourceResult,
- Resource,
- Tool as McpTool,
-} from '@modelcontextprotocol/sdk/types.js';
import {
ListResourcesResultSchema,
ListRootsRequestSchema,
@@ -32,14 +29,19 @@ import {
ToolListChangedNotificationSchema,
PromptListChangedNotificationSchema,
ProgressNotificationSchema,
+ type GetPromptResult,
+ type Prompt,
+ type ReadResourceResult,
+ type Resource,
+ type Tool as McpTool,
} from '@modelcontextprotocol/sdk/types.js';
import { parse } from 'shell-quote';
-import type {
- Config,
- MCPServerConfig,
- GeminiCLIExtension,
+import {
+ AuthProviderType,
+ type Config,
+ type MCPServerConfig,
+ type GeminiCLIExtension,
} from '../config/config.js';
-import { AuthProviderType } from '../config/config.js';
import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
import { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js';
import { DiscoveredMCPTool } from './mcp-tool.js';
@@ -68,7 +70,10 @@ import type { ToolRegistry } from './tool-registry.js';
import { debugLogger } from '../utils/debugLogger.js';
import { type MessageBus } from '../confirmation-bus/message-bus.js';
import { coreEvents } from '../utils/events.js';
-import type { ResourceRegistry } from '../resources/resource-registry.js';
+import {
+ type ResourceRegistry,
+ type MCPResource,
+} from '../resources/resource-registry.js';
import { validateMcpPolicyToolNames } from '../policy/toml-loader.js';
import {
sanitizeEnvironment,
@@ -154,7 +159,7 @@ export class McpClient implements McpProgressReporter {
private readonly cliConfig: McpContext,
private readonly debugMode: boolean,
private readonly clientVersion: string,
- private readonly onToolsUpdated?: (signal?: AbortSignal) => Promise,
+ private readonly onContextUpdated?: (signal?: AbortSignal) => Promise,
) {}
/**
@@ -350,10 +355,21 @@ export class McpClient implements McpProgressReporter {
const capabilities = this.client.getServerCapabilities();
- if (capabilities?.tools?.listChanged) {
- debugLogger.log(
- `Server '${this.serverName}' supports tool updates. Listening for changes...`,
- );
+ debugLogger.log(
+ `Registering notification handlers for server '${this.serverName}'. Capabilities:`,
+ capabilities,
+ );
+
+ if (capabilities?.tools) {
+ if (capabilities.tools.listChanged) {
+ debugLogger.log(
+ `Server '${this.serverName}' supports tool updates. Listening for changes...`,
+ );
+ } else {
+ debugLogger.log(
+ `Server '${this.serverName}' has tools but did not declare 'listChanged' capability. Listening anyway for robustness...`,
+ );
+ }
this.client.setNotificationHandler(
ToolListChangedNotificationSchema,
@@ -366,10 +382,16 @@ export class McpClient implements McpProgressReporter {
);
}
- if (capabilities?.resources?.listChanged) {
- debugLogger.log(
- `Server '${this.serverName}' supports resource updates. Listening for changes...`,
- );
+ if (capabilities?.resources) {
+ if (capabilities.resources.listChanged) {
+ debugLogger.log(
+ `Server '${this.serverName}' supports resource updates. Listening for changes...`,
+ );
+ } else {
+ debugLogger.log(
+ `Server '${this.serverName}' has resources but did not declare 'listChanged' capability. Listening anyway for robustness...`,
+ );
+ }
this.client.setNotificationHandler(
ResourceListChangedNotificationSchema,
@@ -382,10 +404,16 @@ export class McpClient implements McpProgressReporter {
);
}
- if (capabilities?.prompts?.listChanged) {
- debugLogger.log(
- `Server '${this.serverName}' supports prompt updates. Listening for changes...`,
- );
+ if (capabilities?.prompts) {
+ if (capabilities.prompts.listChanged) {
+ debugLogger.log(
+ `Server '${this.serverName}' supports prompt updates. Listening for changes...`,
+ );
+ } else {
+ debugLogger.log(
+ `Server '${this.serverName}' has prompts but did not declare 'listChanged' capability. Listening anyway for robustness...`,
+ );
+ }
this.client.setNotificationHandler(
PromptListChangedNotificationSchema,
@@ -449,6 +477,25 @@ export class McpClient implements McpProgressReporter {
let newResources;
try {
newResources = await this.discoverResources();
+
+ // Verification Retry: If no resources are found or resources didn't change,
+ // wait briefly and try one more time. Some servers notify before they're fully ready.
+ const currentResources =
+ this.resourceRegistry.getResourcesByServer(this.serverName) || [];
+ const resourceMatch =
+ newResources.length === currentResources.length &&
+ newResources.every((nr: Resource) =>
+ currentResources.some((cr: MCPResource) => cr.uri === nr.uri),
+ );
+
+ if (resourceMatch && !this.pendingResourceRefresh) {
+ debugLogger.log(
+ `No resource changes detected for '${this.serverName}'. Retrying once in 500ms...`,
+ );
+ const retryDelay = 500;
+ await new Promise((resolve) => setTimeout(resolve, retryDelay));
+ newResources = await this.discoverResources();
+ }
} catch (err) {
debugLogger.error(
`Resource discovery failed during refresh: ${getErrorMessage(err)}`,
@@ -459,6 +506,10 @@ export class McpClient implements McpProgressReporter {
this.updateResourceRegistry(newResources);
+ if (this.onContextUpdated) {
+ await this.onContextUpdated(abortController.signal);
+ }
+
clearTimeout(timeoutId);
this.cliConfig.emitMcpDiagnostic(
@@ -474,7 +525,6 @@ export class McpClient implements McpProgressReporter {
);
} finally {
this.isRefreshingResources = false;
- this.pendingResourceRefresh = false;
}
}
@@ -517,9 +567,31 @@ export class McpClient implements McpProgressReporter {
const timeoutId = setTimeout(() => abortController.abort(), timeoutMs);
try {
- const newPrompts = await this.fetchPrompts({
+ let newPrompts = await this.fetchPrompts({
signal: abortController.signal,
});
+
+ // Verification Retry: If no prompts are found or prompts didn't change,
+ // wait briefly and try one more time. Some servers notify before they're fully ready.
+ const currentPrompts =
+ this.promptRegistry.getPromptsByServer(this.serverName) || [];
+ const promptsMatch =
+ newPrompts.length === currentPrompts.length &&
+ newPrompts.every((np) =>
+ currentPrompts.some((cp) => cp.name === np.name),
+ );
+
+ if (promptsMatch && !this.pendingPromptRefresh) {
+ debugLogger.log(
+ `No prompt changes detected for '${this.serverName}'. Retrying once in 500ms...`,
+ );
+ const retryDelay = 500;
+ await new Promise((resolve) => setTimeout(resolve, retryDelay));
+ newPrompts = await this.fetchPrompts({
+ signal: abortController.signal,
+ });
+ }
+
this.promptRegistry.removePromptsByServer(this.serverName);
for (const prompt of newPrompts) {
this.promptRegistry.registerPrompt(prompt);
@@ -532,6 +604,10 @@ export class McpClient implements McpProgressReporter {
break;
}
+ if (this.onContextUpdated) {
+ await this.onContextUpdated(abortController.signal);
+ }
+
clearTimeout(timeoutId);
this.cliConfig.emitMcpDiagnostic(
@@ -547,7 +623,6 @@ export class McpClient implements McpProgressReporter {
);
} finally {
this.isRefreshingPrompts = false;
- this.pendingPromptRefresh = false;
}
}
@@ -592,6 +667,38 @@ export class McpClient implements McpProgressReporter {
newTools = await this.discoverTools(this.cliConfig, {
signal: abortController.signal,
});
+ debugLogger.log(
+ `Refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
+ );
+
+ // Verification Retry (Option 3): If no tools are found or tools didn't change,
+ // wait briefly and try one more time. Some servers notify before they're fully ready.
+ const currentTools =
+ this.toolRegistry.getToolsByServer(this.serverName) || [];
+ const toolNamesMatch =
+ newTools.length === currentTools.length &&
+ newTools.every((nt) =>
+ currentTools.some(
+ (ct) =>
+ ct.name === nt.name ||
+ (ct instanceof DiscoveredMCPTool &&
+ ct.serverToolName === nt.serverToolName),
+ ),
+ );
+
+ if (toolNamesMatch && !this.pendingToolRefresh) {
+ debugLogger.log(
+ `No tool changes detected for '${this.serverName}'. Retrying once in 500ms...`,
+ );
+ const retryDelay = 500;
+ await new Promise((resolve) => setTimeout(resolve, retryDelay));
+ newTools = await this.discoverTools(this.cliConfig, {
+ signal: abortController.signal,
+ });
+ debugLogger.log(
+ `Retry refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
+ );
+ }
} catch (err) {
debugLogger.error(
`Discovery failed during refresh: ${getErrorMessage(err)}`,
@@ -607,8 +714,8 @@ export class McpClient implements McpProgressReporter {
}
this.toolRegistry.sortTools();
- if (this.onToolsUpdated) {
- await this.onToolsUpdated(abortController.signal);
+ if (this.onContextUpdated) {
+ await this.onContextUpdated(abortController.signal);
}
clearTimeout(timeoutId);
@@ -626,7 +733,6 @@ export class McpClient implements McpProgressReporter {
);
} finally {
this.isRefreshingTools = false;
- this.pendingToolRefresh = false;
}
}
}
diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts
index 4cdad89827..1d9e2a2f25 100644
--- a/packages/core/src/tools/mcp-tool.test.ts
+++ b/packages/core/src/tools/mcp-tool.test.ts
@@ -5,12 +5,22 @@
*/
/* eslint-disable @typescript-eslint/no-explicit-any */
-import type { Mocked } from 'vitest';
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import {
+ describe,
+ it,
+ expect,
+ vi,
+ beforeEach,
+ afterEach,
+ type Mocked,
+} from 'vitest';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
-import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
-import type { ToolResult } from './tools.js';
-import { ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome
+import {
+ DiscoveredMCPTool,
+ generateValidName,
+ formatMcpToolName,
+} from './mcp-tool.js'; // Added getStringifiedResultForDisplay
+import { ToolConfirmationOutcome, type ToolResult } from './tools.js';
import type { CallableTool, Part } from '@google/genai';
import { ToolErrorType } from './tool-error.js';
import {
@@ -43,23 +53,23 @@ const createSdkResponse = (
describe('generateValidName', () => {
it('should return a valid name for a simple function', () => {
- expect(generateValidName('myFunction')).toBe('myFunction');
+ expect(generateValidName('myFunction')).toBe('mcp_myFunction');
});
it('should replace invalid characters with underscores', () => {
expect(generateValidName('invalid-name with spaces')).toBe(
- 'invalid-name_with_spaces',
+ 'mcp_invalid-name_with_spaces',
);
});
it('should truncate long names', () => {
expect(generateValidName('x'.repeat(80))).toBe(
- 'xxxxxxxxxxxxxxxxxxxxxxxxxxxx___xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
+ 'mcp_xxxxxxxxxxxxxxxxxxxxxxxxxx...xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
);
});
it('should handle names with only invalid characters', () => {
- expect(generateValidName('!@#$%^&*()')).toBe('__________');
+ expect(generateValidName('!@#$%^&*()')).toBe('mcp___________');
});
it.each([
@@ -74,6 +84,30 @@ describe('generateValidName', () => {
);
});
+describe('formatMcpToolName', () => {
+ it('should format a fully qualified name', () => {
+ expect(formatMcpToolName('github', 'list_repos')).toBe(
+ 'mcp_github_list_repos',
+ );
+ });
+
+ it('should handle global wildcards', () => {
+ expect(formatMcpToolName('*')).toBe('mcp_*');
+ });
+
+ it('should handle tool-level wildcards', () => {
+ expect(formatMcpToolName('github', '*')).toBe('mcp_github_*');
+ });
+
+ it('should handle undefined toolName as a tool-level wildcard', () => {
+ expect(formatMcpToolName('github')).toBe('mcp_github_*');
+ });
+
+ it('should format explicitly global wildcard with specific tool', () => {
+ expect(formatMcpToolName('*', 'list_repos')).toBe('mcp_*_list_repos');
+ });
+});
+
describe('DiscoveredMCPTool', () => {
const serverName = 'mock-mcp-server';
const serverToolName = 'actual-server-tool-name';
@@ -110,8 +144,10 @@ describe('DiscoveredMCPTool', () => {
describe('constructor', () => {
it('should set properties correctly', () => {
- expect(tool.name).toBe(serverToolName);
- expect(tool.schema.name).toBe(serverToolName);
+ expect(tool.name).toBe('mcp_mock-mcp-server_actual-server-tool-name');
+ expect(tool.schema.name).toBe(
+ 'mcp_mock-mcp-server_actual-server-tool-name',
+ );
expect(tool.schema.description).toBe(baseDescription);
expect(tool.schema.parameters).toBeUndefined();
expect(tool.schema.parametersJsonSchema).toEqual(inputSchema);
@@ -933,3 +969,86 @@ describe('DiscoveredMCPTool', () => {
});
});
});
+
+describe('MCP Tool Naming Regression Fixes', () => {
+ describe('generateValidName', () => {
+ it('should replace spaces with underscores', () => {
+ expect(generateValidName('My Tool')).toBe('mcp_My_Tool');
+ });
+
+ it('should allow colons', () => {
+ expect(generateValidName('namespace:tool')).toBe('mcp_namespace:tool');
+ });
+
+ it('should ensure name starts with a letter or underscore', () => {
+ expect(generateValidName('valid_tool_name')).toBe('mcp_valid_tool_name');
+ expect(generateValidName('alsoValid-123.name')).toBe(
+ 'mcp_alsoValid-123.name',
+ );
+ expect(generateValidName('another:valid:name')).toBe(
+ 'mcp_another:valid:name',
+ );
+ });
+
+ it('should handle very long names by truncating in the middle', () => {
+ const longName = 'a'.repeat(40) + '__' + 'b'.repeat(40);
+ const result = generateValidName(longName);
+ expect(result.length).toBeLessThanOrEqual(63);
+ expect(result).toMatch(/^mcp_a{26}\.\.\.b{30}$/);
+ });
+
+ it('should handle very long names starting with a digit', () => {
+ const longName = '1' + 'a'.repeat(80);
+ const result = generateValidName(longName);
+ expect(result.length).toBeLessThanOrEqual(63);
+ expect(result.startsWith('mcp_1')).toBe(true);
+ });
+ });
+
+ describe('DiscoveredMCPTool qualified names', () => {
+ it('should generate a valid qualified name even with spaces in server name', () => {
+ const tool = new DiscoveredMCPTool(
+ {} as any,
+ 'My Server',
+ 'my-tool',
+ 'desc',
+ {},
+ {} as any,
+ );
+
+ const qn = tool.getFullyQualifiedName();
+ expect(qn).toBe('mcp_My_Server_my-tool');
+ });
+
+ it('should handle long server and tool names in qualified name', () => {
+ const serverName = 'a'.repeat(40);
+ const toolName = 'b'.repeat(40);
+ const tool = new DiscoveredMCPTool(
+ {} as any,
+ serverName,
+ toolName,
+ 'desc',
+ {},
+ {} as any,
+ );
+
+ const qn = tool.getFullyQualifiedName();
+ expect(qn.length).toBeLessThanOrEqual(63);
+ expect(qn).toContain('...');
+ });
+
+ it('should handle server names starting with digits', () => {
+ const tool = new DiscoveredMCPTool(
+ {} as any,
+ '123-server',
+ 'tool',
+ 'desc',
+ {},
+ {} as any,
+ );
+
+ const qn = tool.getFullyQualifiedName();
+ expect(qn).toBe('mcp_123-server_tool');
+ });
+ });
+});
diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts
index 3d492457f2..f67d1f9bea 100644
--- a/packages/core/src/tools/mcp-tool.ts
+++ b/packages/core/src/tools/mcp-tool.ts
@@ -5,18 +5,17 @@
*/
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
-import type {
- ToolCallConfirmationDetails,
- ToolInvocation,
- ToolMcpConfirmationDetails,
- ToolResult,
- PolicyUpdateOptions,
-} from './tools.js';
+import { debugLogger } from '../utils/debugLogger.js';
import {
BaseDeclarativeTool,
BaseToolInvocation,
Kind,
ToolConfirmationOutcome,
+ type ToolCallConfirmationDetails,
+ type ToolInvocation,
+ type ToolMcpConfirmationDetails,
+ type ToolResult,
+ type PolicyUpdateOptions,
} from './tools.js';
import type { CallableTool, FunctionCall, Part } from '@google/genai';
import { ToolErrorType } from './tool-error.js';
@@ -25,18 +24,92 @@ import type { McpContext } from './mcp-client.js';
/**
* The separator used to qualify MCP tool names with their server prefix.
- * e.g. "server_name__tool_name"
+ * e.g. "mcp_server_name_tool_name"
*/
-export const MCP_QUALIFIED_NAME_SEPARATOR = '__';
+export const MCP_QUALIFIED_NAME_SEPARATOR = '_';
/**
- * Returns true if `name` matches the MCP qualified name format: "server__tool",
- * i.e. exactly two non-empty parts separated by the MCP_QUALIFIED_NAME_SEPARATOR.
+ * The strict prefix that all MCP tools must start with.
+ */
+export const MCP_TOOL_PREFIX = 'mcp_';
+
+/**
+ * Returns true if `name` matches the MCP qualified name format: "mcp_server_tool",
+ * i.e. starts with the "mcp_" prefix.
*/
export function isMcpToolName(name: string): boolean {
- if (!name.includes(MCP_QUALIFIED_NAME_SEPARATOR)) return false;
- const parts = name.split(MCP_QUALIFIED_NAME_SEPARATOR);
- return parts.length === 2 && parts[0].length > 0 && parts[1].length > 0;
+ return name.startsWith(MCP_TOOL_PREFIX);
+}
+
+/**
+ * Extracts the server name and tool name from a fully qualified MCP tool name.
+ * Expected format: `mcp_{server_name}_{tool_name}`
+ * @param name The fully qualified tool name.
+ * @returns An object containing the extracted `serverName` and `toolName`, or
+ * `undefined` properties if the name doesn't match the expected format.
+ */
+export function parseMcpToolName(name: string): {
+ serverName?: string;
+ toolName?: string;
+} {
+ if (!isMcpToolName(name)) {
+ return {};
+ }
+ // Remove the prefix
+ const withoutPrefix = name.slice(MCP_TOOL_PREFIX.length);
+ // The first segment is the server name, the rest is the tool name
+ const match = withoutPrefix.match(/^([^_]+)_(.+)$/);
+ if (match) {
+ return {
+ serverName: match[1],
+ toolName: match[2],
+ };
+ }
+ return {};
+}
+
+/**
+ * Assembles a fully qualified MCP tool name (or wildcard pattern) from its server and tool components.
+ *
+ * @param serverName The backend MCP server name (can be '*' for global wildcards).
+ * @param toolName The name of the tool (can be undefined or '*' for tool-level wildcards).
+ * @returns The fully qualified name (e.g., `mcp_server_tool`, `mcp_*`, `mcp_server_*`).
+ */
+export function formatMcpToolName(
+ serverName: string,
+ toolName?: string,
+): string {
+ if (serverName === '*' && !toolName) {
+ return `${MCP_TOOL_PREFIX}*`;
+ } else if (serverName === '*') {
+ return `${MCP_TOOL_PREFIX}*_${toolName}`;
+ } else if (!toolName) {
+ return `${MCP_TOOL_PREFIX}${serverName}_*`;
+ } else {
+ return `${MCP_TOOL_PREFIX}${serverName}_${toolName}`;
+ }
+}
+
+/**
+ * Interface representing metadata annotations specific to an MCP tool.
+ * Ensures strongly-typed access to server-level properties.
+ */
+export interface McpToolAnnotation extends Record {
+ _serverName: string;
+}
+
+/**
+ * Type guard to check if tool annotations implement McpToolAnnotation.
+ */
+export function isMcpToolAnnotation(
+ annotation: unknown,
+): annotation is McpToolAnnotation {
+ return (
+ typeof annotation === 'object' &&
+ annotation !== null &&
+ // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, no-restricted-syntax
+ typeof (annotation as Record)['_serverName'] === 'string'
+ );
}
type ToolParams = Record;
@@ -96,14 +169,17 @@ export class DiscoveredMCPToolInvocation extends BaseToolInvocation<
) {
// Use composite format for policy checks: serverName__toolName
// This enables server wildcards (e.g., "google-workspace__*")
- // while still allowing specific tool rules
+ // while still allowing specific tool rules.
+ // We use the same sanitized names as the registry to ensure policy matches.
super(
params,
messageBus,
- `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
+ generateValidName(
+ `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
+ ),
displayName,
- serverName,
+ generateValidName(serverName),
toolAnnotationsData,
);
}
@@ -273,7 +349,10 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
private readonly _toolAnnotations?: Record,
) {
super(
- nameOverride ?? generateValidName(serverToolName),
+ nameOverride ??
+ generateValidName(
+ `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
+ ),
`${serverToolName} (${serverName} MCP Server)`,
description,
Kind.Other,
@@ -301,11 +380,15 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
}
getFullyQualifiedPrefix(): string {
- return `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`;
+ return generateValidName(
+ `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`,
+ );
}
getFullyQualifiedName(): string {
- return `${this.getFullyQualifiedPrefix()}${generateValidName(this.serverToolName)}`;
+ return generateValidName(
+ `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${this.serverToolName}`,
+ );
}
asFullyQualifiedTool(): DiscoveredMCPTool {
@@ -482,16 +565,36 @@ function getStringifiedResultForDisplay(rawResponse: Part[]): string {
return displayParts.join('\n');
}
+/**
+ * Maximum length for a function name in the Gemini API.
+ * @see https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/function-calling#functiondeclaration
+ */
+const MAX_FUNCTION_NAME_LENGTH = 64;
+
/** Visible for testing */
export function generateValidName(name: string) {
- // Replace invalid characters (based on 400 error message from Gemini API) with underscores
- let validToolname = name.replace(/[^a-zA-Z0-9_.-]/g, '_');
+ // Enforce the mcp_ prefix for all generated MCP tool names
+ let validToolname = name.startsWith('mcp_') ? name : `mcp_${name}`;
- // If longer than 63 characters, replace middle with '___'
- // (Gemini API says max length 64, but actual limit seems to be 63)
- if (validToolname.length > 63) {
- validToolname =
- validToolname.slice(0, 28) + '___' + validToolname.slice(-32);
+ // Replace invalid characters with underscores to conform to Gemini API:
+ // ^[a-zA-Z_][a-zA-Z0-9_\-.:]{0,63}$
+ validToolname = validToolname.replace(/[^a-zA-Z0-9_\-.:]/g, '_');
+
+ // Ensure it starts with a letter or underscore
+ if (/^[^a-zA-Z_]/.test(validToolname)) {
+ validToolname = `_${validToolname}`;
}
+
+ // If longer than the API limit, replace middle with '...'
+ // Note: We use 63 instead of 64 to be safe, as some environments have off-by-one behaviors.
+ const safeLimit = MAX_FUNCTION_NAME_LENGTH - 1;
+ if (validToolname.length > safeLimit) {
+ debugLogger.warn(
+ `Truncating MCP tool name "${validToolname}" to fit within the 64 character limit. This tool may require user approval.`,
+ );
+ validToolname =
+ validToolname.slice(0, 30) + '...' + validToolname.slice(-30);
+ }
+
return validToolname;
}
diff --git a/packages/core/src/tools/memoryTool.test.ts b/packages/core/src/tools/memoryTool.test.ts
index 12cb8baa2e..4b0aa1b616 100644
--- a/packages/core/src/tools/memoryTool.test.ts
+++ b/packages/core/src/tools/memoryTool.test.ts
@@ -4,8 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { Mock } from 'vitest';
-import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
+import {
+ vi,
+ describe,
+ it,
+ expect,
+ beforeEach,
+ afterEach,
+ type Mock,
+} from 'vitest';
import {
MemoryTool,
setGeminiMdFilename,
diff --git a/packages/core/src/tools/memoryTool.ts b/packages/core/src/tools/memoryTool.ts
index 33cb9483e1..68a0942a53 100644
--- a/packages/core/src/tools/memoryTool.ts
+++ b/packages/core/src/tools/memoryTool.ts
@@ -4,12 +4,13 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { ToolEditConfirmationDetails, ToolResult } from './tools.js';
import {
BaseDeclarativeTool,
BaseToolInvocation,
Kind,
ToolConfirmationOutcome,
+ type ToolEditConfirmationDetails,
+ type ToolResult,
} from './tools.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
diff --git a/packages/core/src/tools/modifiable-tool.test.ts b/packages/core/src/tools/modifiable-tool.test.ts
index 4f34b20b57..6ff9126478 100644
--- a/packages/core/src/tools/modifiable-tool.test.ts
+++ b/packages/core/src/tools/modifiable-tool.test.ts
@@ -5,13 +5,11 @@
*/
import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
-import type {
- ModifyContext,
- ModifiableDeclarativeTool,
-} from './modifiable-tool.js';
import {
modifyWithEditor,
isModifiableDeclarativeTool,
+ type ModifyContext,
+ type ModifiableDeclarativeTool,
} from './modifiable-tool.js';
import { DEFAULT_GUI_EDITOR } from '../utils/editor.js';
import fs from 'node:fs';
diff --git a/packages/core/src/tools/modifiable-tool.ts b/packages/core/src/tools/modifiable-tool.ts
index 328158bb78..69abeacb82 100644
--- a/packages/core/src/tools/modifiable-tool.ts
+++ b/packages/core/src/tools/modifiable-tool.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { EditorType } from '../utils/editor.js';
-import { openDiff } from '../utils/editor.js';
+import { openDiff, type EditorType } from '../utils/editor.js';
import os from 'node:os';
import path from 'node:path';
import fs from 'node:fs';
diff --git a/packages/core/src/tools/read-file.test.ts b/packages/core/src/tools/read-file.test.ts
index 8f79bffe17..6b82a152a6 100644
--- a/packages/core/src/tools/read-file.test.ts
+++ b/packages/core/src/tools/read-file.test.ts
@@ -5,8 +5,7 @@
*/
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
-import type { ReadFileToolParams } from './read-file.js';
-import { ReadFileTool } from './read-file.js';
+import { ReadFileTool, type ReadFileToolParams } from './read-file.js';
import { ToolErrorType } from './tool-error.js';
import path from 'node:path';
import { isSubpath } from '../utils/paths.js';
diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts
index 170cccf905..0f044a4998 100644
--- a/packages/core/src/tools/read-file.ts
+++ b/packages/core/src/tools/read-file.ts
@@ -7,8 +7,14 @@
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import path from 'node:path';
import { makeRelative, shortenPath } from '../utils/paths.js';
-import type { ToolInvocation, ToolLocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolLocation,
+ type ToolResult,
+} from './tools.js';
import { ToolErrorType } from './tool-error.js';
import type { PartUnion } from '@google/genai';
diff --git a/packages/core/src/tools/read-many-files.test.ts b/packages/core/src/tools/read-many-files.test.ts
index f340424a35..875ccf0bd5 100644
--- a/packages/core/src/tools/read-many-files.test.ts
+++ b/packages/core/src/tools/read-many-files.test.ts
@@ -4,8 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
-import type { Mock } from 'vitest';
+import {
+ vi,
+ describe,
+ it,
+ expect,
+ beforeEach,
+ afterEach,
+ type Mock,
+} from 'vitest';
import { mockControl } from '../__mocks__/fs/promises.js';
import { ReadManyFilesTool } from './read-many-files.js';
import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
diff --git a/packages/core/src/tools/read-many-files.ts b/packages/core/src/tools/read-many-files.ts
index 0a5d68a6ba..c9c4e230e6 100644
--- a/packages/core/src/tools/read-many-files.ts
+++ b/packages/core/src/tools/read-many-files.ts
@@ -5,18 +5,23 @@
*/
import type { MessageBus } from '../confirmation-bus/message-bus.js';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { getErrorMessage } from '../utils/errors.js';
import * as fsPromises from 'node:fs/promises';
import * as path from 'node:path';
import { glob, escape } from 'glob';
-import type { ProcessedFileReadResult } from '../utils/fileUtils.js';
import {
detectFileType,
processSingleFileContent,
DEFAULT_ENCODING,
getSpecificMimeType,
+ type ProcessedFileReadResult,
} from '../utils/fileUtils.js';
import type { PartListUnion } from '@google/genai';
import {
diff --git a/packages/core/src/tools/ripGrep.test.ts b/packages/core/src/tools/ripGrep.test.ts
index 0eaf5c0b68..a1b155fb7a 100644
--- a/packages/core/src/tools/ripGrep.test.ts
+++ b/packages/core/src/tools/ripGrep.test.ts
@@ -13,8 +13,12 @@ import {
afterAll,
vi,
} from 'vitest';
-import type { RipGrepToolParams } from './ripGrep.js';
-import { canUseRipgrep, RipGrepTool, ensureRgPath } from './ripGrep.js';
+import {
+ canUseRipgrep,
+ RipGrepTool,
+ ensureRgPath,
+ type RipGrepToolParams,
+} from './ripGrep.js';
import path from 'node:path';
import { isSubpath } from '../utils/paths.js';
import fs from 'node:fs/promises';
@@ -23,8 +27,7 @@ import type { Config } from '../config/config.js';
import { Storage } from '../config/storage.js';
import { GEMINI_IGNORE_FILE_NAME } from '../config/constants.js';
import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
-import type { ChildProcess } from 'node:child_process';
-import { spawn } from 'node:child_process';
+import { spawn, type ChildProcess } from 'node:child_process';
import { PassThrough, Readable } from 'node:stream';
import EventEmitter from 'node:events';
import { downloadRipGrep } from '@joshua.litt/get-ripgrep';
@@ -2025,6 +2028,32 @@ describe('RipGrepTool', () => {
expect(result.llmContent).not.toContain('fileB.txt');
expect(result.llmContent).toContain('Copyright 2025 Google LLC');
});
+
+ it('should truncate excessively long lines', async () => {
+ const longString = 'a'.repeat(3000);
+ mockSpawn.mockImplementation(
+ createMockSpawn({
+ outputData:
+ JSON.stringify({
+ type: 'match',
+ data: {
+ path: { text: 'longline.txt' },
+ line_number: 1,
+ lines: { text: `Target match ${longString}\n` },
+ },
+ }) + '\n',
+ exitCode: 0,
+ }),
+ );
+
+ const params: RipGrepToolParams = { pattern: 'Target match', context: 0 };
+ const invocation = grepTool.build(params);
+ const result = await invocation.execute(abortSignal);
+
+ // MAX_LINE_LENGTH_TEXT_FILE is 2000. It should be truncated.
+ expect(result.llmContent).toContain('... [truncated]');
+ expect(result.llmContent).not.toContain(longString);
+ });
});
});
diff --git a/packages/core/src/tools/ripGrep.ts b/packages/core/src/tools/ripGrep.ts
index ac65cf6362..18a1b0c133 100644
--- a/packages/core/src/tools/ripGrep.ts
+++ b/packages/core/src/tools/ripGrep.ts
@@ -9,8 +9,13 @@ import fs from 'node:fs';
import fsPromises from 'node:fs/promises';
import path from 'node:path';
import { downloadRipGrep } from '@joshua.litt/get-ripgrep';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { ToolErrorType } from './tool-error.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import { getErrorMessage, isNodeError } from '../utils/errors.js';
@@ -49,7 +54,23 @@ async function resolveExistingRgPath(): Promise {
}
let ripgrepAcquisitionPromise: Promise | null = null;
-
+/**
+ * Ensures a ripgrep binary is available.
+ *
+ * NOTE:
+ * - The Gemini CLI currently prefers a managed ripgrep binary downloaded
+ * into its global bin directory.
+ * - Even if ripgrep is available on the system PATH, it is intentionally
+ * not used at this time.
+ *
+ * Preference for system-installed ripgrep is blocked on:
+ * - checksum verification of external binaries
+ * - internalization of the get-ripgrep dependency
+ *
+ * See:
+ * - feat(core): Prefer rg in system path (#11847)
+ * - Move get-ripgrep to third_party (#12099)
+ */
async function ensureRipgrepAvailable(): Promise {
const existingPath = await resolveExistingRgPath();
if (existingPath) {
diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts
index 7c47053afb..656abed8ff 100644
--- a/packages/core/src/tools/shell.test.ts
+++ b/packages/core/src/tools/shell.test.ts
@@ -51,7 +51,6 @@ import {
} from '../services/shellExecutionService.js';
import * as fs from 'node:fs';
import * as os from 'node:os';
-import { EOL } from 'node:os';
import * as path from 'node:path';
import { isSubpath } from '../utils/paths.js';
import * as crypto from 'node:crypto';
@@ -265,7 +264,7 @@ describe('ShellTool', () => {
// Simulate pgrep output file creation by the shell command
const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
- fs.writeFileSync(tmpFile, `54321${EOL}54322${EOL}`);
+ fs.writeFileSync(tmpFile, `54321${os.EOL}54322${os.EOL}`);
const result = await promise;
diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts
index 55942b8cc2..80ab20e19c 100644
--- a/packages/core/src/tools/shell.ts
+++ b/packages/core/src/tools/shell.ts
@@ -7,7 +7,7 @@
import fsPromises from 'node:fs/promises';
import fs from 'node:fs';
import path from 'node:path';
-import os, { EOL } from 'node:os';
+import os from 'node:os';
import crypto from 'node:crypto';
import type { Config } from '../config/config.js';
import { debugLogger } from '../utils/debugLogger.js';
@@ -27,11 +27,11 @@ import {
import { getErrorMessage } from '../utils/errors.js';
import { summarizeToolOutput } from '../utils/summarizer.js';
-import type {
- ShellExecutionConfig,
- ShellOutputEvent,
+import {
+ ShellExecutionService,
+ type ShellExecutionConfig,
+ type ShellOutputEvent,
} from '../services/shellExecutionService.js';
-import { ShellExecutionService } from '../services/shellExecutionService.js';
import { formatBytes } from '../utils/formatters.js';
import type { AnsiOutput } from '../utils/terminalSerializer.js';
import {
@@ -328,7 +328,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
if (tempFileExists) {
const pgrepContent = await fsPromises.readFile(tempFilePath, 'utf8');
- const pgrepLines = pgrepContent.split(EOL).filter(Boolean);
+ const pgrepLines = pgrepContent.split(os.EOL).filter(Boolean);
for (const line of pgrepLines) {
if (!/^\d+$/.test(line)) {
debugLogger.error(`pgrep: ${line}`);
@@ -407,16 +407,17 @@ export class ShellToolInvocation extends BaseToolInvocation<
} else {
if (this.params.is_background || result.backgrounded) {
returnDisplayMessage = `Command moved to background (PID: ${result.pid}). Output hidden. Press Ctrl+B to view.`;
+ } else if (result.aborted) {
+ const cancelMsg = timeoutMessage || 'Command cancelled by user.';
+ if (result.output.trim()) {
+ returnDisplayMessage = `${cancelMsg}\n\nOutput before cancellation:\n${result.output}`;
+ } else {
+ returnDisplayMessage = cancelMsg;
+ }
} else if (result.output.trim()) {
returnDisplayMessage = result.output;
} else {
- if (result.aborted) {
- if (timeoutMessage) {
- returnDisplayMessage = timeoutMessage;
- } else {
- returnDisplayMessage = 'Command cancelled by user.';
- }
- } else if (result.signal) {
+ if (result.signal) {
returnDisplayMessage = `Command terminated by signal: ${result.signal}`;
} else if (result.error) {
returnDisplayMessage = `Command failed: ${getErrorMessage(
diff --git a/packages/core/src/tools/tool-names.test.ts b/packages/core/src/tools/tool-names.test.ts
index 344ff48376..8ff871986f 100644
--- a/packages/core/src/tools/tool-names.test.ts
+++ b/packages/core/src/tools/tool-names.test.ts
@@ -58,6 +58,8 @@ describe('tool-names', () => {
it('should validate MCP tool names (server__tool)', () => {
expect(isValidToolName('server__tool')).toBe(true);
expect(isValidToolName('my-server__my-tool')).toBe(true);
+ expect(isValidToolName('my.server__my:tool')).toBe(true);
+ expect(isValidToolName('my-server...truncated__tool')).toBe(true);
});
it('should validate legacy tool aliases', async () => {
diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts
index a2e8061fc6..fcdcbd6df6 100644
--- a/packages/core/src/tools/tool-names.ts
+++ b/packages/core/src/tools/tool-names.ts
@@ -160,6 +160,12 @@ export const EDIT_DISPLAY_NAME = 'Edit';
export const ASK_USER_DISPLAY_NAME = 'Ask User';
export const READ_FILE_DISPLAY_NAME = 'ReadFile';
export const GLOB_DISPLAY_NAME = 'FindFiles';
+export const TRACKER_CREATE_TASK_TOOL_NAME = 'tracker_create_task';
+export const TRACKER_UPDATE_TASK_TOOL_NAME = 'tracker_update_task';
+export const TRACKER_GET_TASK_TOOL_NAME = 'tracker_get_task';
+export const TRACKER_LIST_TASKS_TOOL_NAME = 'tracker_list_tasks';
+export const TRACKER_ADD_DEPENDENCY_TOOL_NAME = 'tracker_add_dependency';
+export const TRACKER_VISUALIZE_TOOL_NAME = 'tracker_visualize';
/**
* Mapping of legacy tool names to their current names.
@@ -213,11 +219,32 @@ export const ALL_BUILTIN_TOOL_NAMES = [
MEMORY_TOOL_NAME,
ACTIVATE_SKILL_TOOL_NAME,
ASK_USER_TOOL_NAME,
+ TRACKER_CREATE_TASK_TOOL_NAME,
+ TRACKER_UPDATE_TASK_TOOL_NAME,
+ TRACKER_GET_TASK_TOOL_NAME,
+ TRACKER_LIST_TASKS_TOOL_NAME,
+ TRACKER_ADD_DEPENDENCY_TOOL_NAME,
+ TRACKER_VISUALIZE_TOOL_NAME,
GET_INTERNAL_DOCS_TOOL_NAME,
ENTER_PLAN_MODE_TOOL_NAME,
EXIT_PLAN_MODE_TOOL_NAME,
] as const;
+/**
+ * Read-only tools available in Plan Mode.
+ * This list is used to dynamically generate the Plan Mode prompt,
+ * filtered by what tools are actually enabled in the current configuration.
+ */
+export const PLAN_MODE_TOOLS = [
+ GLOB_TOOL_NAME,
+ GREP_TOOL_NAME,
+ READ_FILE_TOOL_NAME,
+ LS_TOOL_NAME,
+ WEB_SEARCH_TOOL_NAME,
+ ASK_USER_TOOL_NAME,
+ ACTIVATE_SKILL_TOOL_NAME,
+] as const;
+
/**
* Validates if a tool name is syntactically valid.
* Checks against built-in tools, discovered tools, and MCP naming conventions.
@@ -260,8 +287,10 @@ export function isValidToolName(
return !!options.allowWildcards;
}
- // Basic slug validation for server and tool names
- const slugRegex = /^[a-z0-9-_]+$/i;
+ // Basic slug validation for server and tool names.
+ // We allow dots (.) and colons (:) as they are valid in function names and
+ // used for truncation markers.
+ const slugRegex = /^[a-z0-9_.:-]+$/i;
return slugRegex.test(server) && slugRegex.test(tool);
}
diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts
index d44c133705..ea560865e6 100644
--- a/packages/core/src/tools/tool-registry.test.ts
+++ b/packages/core/src/tools/tool-registry.test.ts
@@ -5,17 +5,27 @@
*/
/* eslint-disable @typescript-eslint/no-explicit-any */
-import type { Mocked, MockInstance } from 'vitest';
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import type { ConfigParameters } from '../config/config.js';
-import { Config } from '../config/config.js';
+import {
+ describe,
+ it,
+ expect,
+ vi,
+ beforeEach,
+ afterEach,
+ type Mocked,
+ type MockInstance,
+} from 'vitest';
+import { Config, type ConfigParameters } from '../config/config.js';
import { ApprovalMode } from '../policy/types.js';
import { ToolRegistry, DiscoveredTool } from './tool-registry.js';
import { DISCOVERED_TOOL_PREFIX } from './tool-names.js';
-import { DiscoveredMCPTool, MCP_QUALIFIED_NAME_SEPARATOR } from './mcp-tool.js';
-import type { FunctionDeclaration, CallableTool } from '@google/genai';
-import { mcpToTool } from '@google/genai';
+import { DiscoveredMCPTool } from './mcp-tool.js';
+import {
+ mcpToTool,
+ type FunctionDeclaration,
+ type CallableTool,
+} from '@google/genai';
import { spawn } from 'node:child_process';
import fs from 'node:fs';
@@ -307,7 +317,7 @@ describe('ToolRegistry', () => {
{
name: 'should match qualified MCP tool names when qualified or unqualified',
tools: [mcpTool, mcpTool.asFullyQualifiedTool()],
- excludedTools: [`${mcpTool.getFullyQualifiedPrefix()}${mcpTool.name}`],
+ excludedTools: [mcpTool.name],
},
{
name: 'should match class names',
@@ -395,7 +405,7 @@ describe('ToolRegistry', () => {
// Assert that the returned array contains all tool names, with MCP qualified
expect(toolNames).toContain('c-tool');
expect(toolNames).toContain('a-tool');
- expect(toolNames).toContain('my-server__my-tool');
+ expect(toolNames).toContain('mcp_my-server_my-tool');
expect(toolNames).toHaveLength(3);
});
@@ -409,7 +419,7 @@ describe('ToolRegistry', () => {
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
const toolNames = toolRegistry.getAllToolNames();
- expect(toolNames).toEqual([`${serverName}__${toolName}`]);
+ expect(toolNames).toEqual([`mcp_${serverName}_${toolName}`]);
});
});
@@ -439,7 +449,11 @@ describe('ToolRegistry', () => {
// Assert that the array has the correct tools and is sorted by name
expect(toolsFromServer1).toHaveLength(3);
- expect(toolNames).toEqual(['apple-tool', 'banana-tool', 'zebra-tool']);
+ expect(toolNames).toEqual([
+ 'mcp_mcp-server-uno_apple-tool',
+ 'mcp_mcp-server-uno_banana-tool',
+ 'mcp_mcp-server-uno_zebra-tool',
+ ]);
// Assert that all returned tools are indeed from the correct server
for (const tool of toolsFromServer1) {
@@ -481,8 +495,8 @@ describe('ToolRegistry', () => {
'builtin-1',
'builtin-2',
DISCOVERED_TOOL_PREFIX + 'discovered-1',
- 'apple-server__mcp-apple',
- 'zebra-server__mcp-zebra',
+ 'mcp_apple-server_mcp-apple',
+ 'mcp_zebra-server_mcp-zebra',
]);
});
});
@@ -598,25 +612,20 @@ describe('ToolRegistry', () => {
});
describe('getTool', () => {
- it('should retrieve an MCP tool by its fully qualified name even if registered with simple name', () => {
+ it('should retrieve an MCP tool by its fully qualified name', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
- // Register tool (will be registered as 'my-tool' since no conflict)
+ // Register tool
toolRegistry.registerTool(mcpTool);
- // Verify it is available as 'my-tool'
- expect(toolRegistry.getTool('my-tool')).toBeDefined();
- expect(toolRegistry.getTool('my-tool')?.name).toBe('my-tool');
-
- // Verify it is available as 'my-server__my-tool'
- const fullyQualifiedName = `${serverName}__${toolName}`;
+ // Verify it is available as 'mcp_my-server_my-tool'
+ const fullyQualifiedName = `mcp_${serverName}_${toolName}`;
const retrievedTool = toolRegistry.getTool(fullyQualifiedName);
expect(retrievedTool).toBeDefined();
- // The returned tool object is the same, so its name property is still 'my-tool'
- expect(retrievedTool?.name).toBe('my-tool');
+ expect(retrievedTool?.name).toBe(fullyQualifiedName);
});
it('should retrieve an MCP tool by its fully qualified name when tool name has special characters', () => {
@@ -626,19 +635,15 @@ describe('ToolRegistry', () => {
const validToolName = 'my_tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
- // Register tool (will be registered as sanitized name)
+ // Register tool
toolRegistry.registerTool(mcpTool);
- // Verify it is available as sanitized name
- expect(toolRegistry.getTool(validToolName)).toBeDefined();
- expect(toolRegistry.getTool(validToolName)?.name).toBe(validToolName);
-
- // Verify it is available as 'my-server__my_tool'
- const fullyQualifiedName = `${serverName}__${validToolName}`;
+ // Verify it is available as 'mcp_my-server_my_tool'
+ const fullyQualifiedName = `mcp_${serverName}_${validToolName}`;
const retrievedTool = toolRegistry.getTool(fullyQualifiedName);
expect(retrievedTool).toBeDefined();
- expect(retrievedTool?.name).toBe(validToolName);
+ expect(retrievedTool?.name).toBe(fullyQualifiedName);
});
it('should resolve qualified names in getFunctionDeclarationsFiltered', () => {
@@ -648,13 +653,13 @@ describe('ToolRegistry', () => {
toolRegistry.registerTool(mcpTool);
- const fullyQualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`;
+ const fullyQualifiedName = `mcp_${serverName}_${toolName}`;
const declarations = toolRegistry.getFunctionDeclarationsFiltered([
fullyQualifiedName,
]);
expect(declarations).toHaveLength(1);
- expect(declarations[0].name).toBe(toolName);
+ expect(declarations[0].name).toBe(fullyQualifiedName);
});
it('should retrieve a tool using its legacy alias', async () => {
@@ -685,7 +690,7 @@ describe('ToolRegistry', () => {
const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
- expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
+ expect(declarations[0].name).toBe(`mcp_${serverName}_${toolName}`);
});
it('should deduplicate MCP tools in declarations', () => {
@@ -699,7 +704,7 @@ describe('ToolRegistry', () => {
const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
- expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
+ expect(declarations[0].name).toBe(`mcp_${serverName}_${toolName}`);
});
});
@@ -752,7 +757,7 @@ describe('ToolRegistry', () => {
const allTools = toolRegistry.getAllTools();
const toolNames = allTools.map((t) => t.name);
- expect(toolNames).toContain('read-only-tool');
+ expect(toolNames).toContain('mcp_test-server_read-only-tool');
});
it('should exclude non-read-only MCP tools when denied by policy in plan mode', () => {
@@ -769,7 +774,7 @@ describe('ToolRegistry', () => {
const allTools = toolRegistry.getAllTools();
const toolNames = allTools.map((t) => t.name);
- expect(toolNames).not.toContain('write-mcp-tool');
+ expect(toolNames).not.toContain('mcp_test-server_write-mcp-tool');
});
});
diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts
index e7fd7a6a66..69695877c2 100644
--- a/packages/core/src/tools/tool-registry.ts
+++ b/packages/core/src/tools/tool-registry.ts
@@ -5,12 +5,14 @@
*/
import type { FunctionDeclaration } from '@google/genai';
-import type {
- AnyDeclarativeTool,
- ToolResult,
- ToolInvocation,
+import {
+ Kind,
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ type AnyDeclarativeTool,
+ type ToolResult,
+ type ToolInvocation,
} from './tools.js';
-import { Kind, BaseDeclarativeTool, BaseToolInvocation } from './tools.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../policy/types.js';
import { spawn } from 'node:child_process';
@@ -443,13 +445,15 @@ export class ToolRegistry {
private buildToolMetadata(): Map> {
const toolMetadata = new Map>();
for (const [name, tool] of this.allKnownTools) {
- if (tool.toolAnnotations) {
- const metadata: Record = { ...tool.toolAnnotations };
- // Include server name so the policy engine can resolve composite
- // wildcard patterns (e.g. "*__*") against unqualified tool names.
- if (tool instanceof DiscoveredMCPTool) {
- metadata['_serverName'] = tool.serverName;
- }
+ const metadata: Record = tool.toolAnnotations
+ ? { ...tool.toolAnnotations }
+ : {};
+ // Include server name so the policy engine can resolve composite
+ // wildcard patterns (e.g. "*__*") against unqualified tool names.
+ if (tool instanceof DiscoveredMCPTool) {
+ metadata['_serverName'] = tool.serverName;
+ }
+ if (Object.keys(metadata).length > 0) {
toolMetadata.set(name, metadata);
}
}
diff --git a/packages/core/src/tools/tools.test.ts b/packages/core/src/tools/tools.test.ts
index 41edf9f21d..edbc487160 100644
--- a/packages/core/src/tools/tools.test.ts
+++ b/packages/core/src/tools/tools.test.ts
@@ -5,8 +5,13 @@
*/
import { describe, it, expect, vi } from 'vitest';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { DeclarativeTool, hasCycleInSchema, Kind } from './tools.js';
+import {
+ DeclarativeTool,
+ hasCycleInSchema,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { ToolErrorType } from './tool-error.js';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
import { ReadFileTool } from './read-file.js';
diff --git a/packages/core/src/tools/trackerTools.test.ts b/packages/core/src/tools/trackerTools.test.ts
new file mode 100644
index 0000000000..ec0bd0e889
--- /dev/null
+++ b/packages/core/src/tools/trackerTools.test.ts
@@ -0,0 +1,145 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { Config } from '../config/config.js';
+import { MessageBus } from '../confirmation-bus/message-bus.js';
+import type { PolicyEngine } from '../policy/policy-engine.js';
+import {
+ TrackerCreateTaskTool,
+ TrackerListTasksTool,
+ TrackerUpdateTaskTool,
+ TrackerVisualizeTool,
+ TrackerAddDependencyTool,
+} from './trackerTools.js';
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import * as os from 'node:os';
+
+import { TaskStatus, TaskType } from '../services/trackerTypes.js';
+
+describe('Tracker Tools Integration', () => {
+ let tempDir: string;
+ let config: Config;
+ let messageBus: MessageBus;
+
+ beforeEach(async () => {
+ tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'tracker-tools-test-'));
+ config = new Config({
+ sessionId: 'test-session',
+ targetDir: tempDir,
+ cwd: tempDir,
+ model: 'gemini-3-flash',
+ debugMode: false,
+ });
+ messageBus = new MessageBus(null as unknown as PolicyEngine, false);
+ });
+
+ afterEach(async () => {
+ await fs.rm(tempDir, { recursive: true, force: true });
+ });
+
+ const getSignal = () => new AbortController().signal;
+
+ it('creates and lists tasks', async () => {
+ const createTool = new TrackerCreateTaskTool(config, messageBus);
+ const createResult = await createTool.buildAndExecute(
+ {
+ title: 'Test Task',
+ description: 'Test Description',
+ type: TaskType.TASK,
+ },
+ getSignal(),
+ );
+
+ expect(createResult.llmContent).toContain('Created task');
+
+ const listTool = new TrackerListTasksTool(config, messageBus);
+ const listResult = await listTool.buildAndExecute({}, getSignal());
+ expect(listResult.llmContent).toContain('Test Task');
+ expect(listResult.llmContent).toContain(`(${TaskStatus.OPEN})`);
+ });
+
+ it('updates task status', async () => {
+ const createTool = new TrackerCreateTaskTool(config, messageBus);
+ await createTool.buildAndExecute(
+ {
+ title: 'Update Me',
+ description: '...',
+ type: TaskType.TASK,
+ },
+ getSignal(),
+ );
+
+ const tasks = await config.getTrackerService().listTasks();
+ const taskId = tasks[0].id;
+
+ const updateTool = new TrackerUpdateTaskTool(config, messageBus);
+ const updateResult = await updateTool.buildAndExecute(
+ {
+ id: taskId,
+ status: TaskStatus.IN_PROGRESS,
+ },
+ getSignal(),
+ );
+
+ expect(updateResult.llmContent).toContain(
+ `Status: ${TaskStatus.IN_PROGRESS}`,
+ );
+
+ const task = await config.getTrackerService().getTask(taskId);
+ expect(task?.status).toBe(TaskStatus.IN_PROGRESS);
+ });
+
+ it('adds dependencies and visualizes the graph', async () => {
+ const createTool = new TrackerCreateTaskTool(config, messageBus);
+
+ // Create Parent
+ await createTool.buildAndExecute(
+ {
+ title: 'Parent Task',
+ description: '...',
+ type: TaskType.TASK,
+ },
+ getSignal(),
+ );
+
+ // Create Child
+ await createTool.buildAndExecute(
+ {
+ title: 'Child Task',
+ description: '...',
+ type: TaskType.TASK,
+ },
+ getSignal(),
+ );
+
+ const tasks = await config.getTrackerService().listTasks();
+ const parentId = tasks.find((t) => t.title === 'Parent Task')!.id;
+ const childId = tasks.find((t) => t.title === 'Child Task')!.id;
+
+ // Add Dependency
+ const addDepTool = new TrackerAddDependencyTool(config, messageBus);
+ await addDepTool.buildAndExecute(
+ {
+ taskId: parentId,
+ dependencyId: childId,
+ },
+ getSignal(),
+ );
+
+ const updatedParent = await config.getTrackerService().getTask(parentId);
+ expect(updatedParent?.dependencies).toContain(childId);
+
+ // Visualize
+ const vizTool = new TrackerVisualizeTool(config, messageBus);
+ const vizResult = await vizTool.buildAndExecute({}, getSignal());
+
+ expect(vizResult.llmContent).toContain('Parent Task');
+ expect(vizResult.llmContent).toContain('Child Task');
+ expect(vizResult.llmContent).toContain(childId);
+ });
+});
diff --git a/packages/core/src/tools/trackerTools.ts b/packages/core/src/tools/trackerTools.ts
new file mode 100644
index 0000000000..03ee3c3a97
--- /dev/null
+++ b/packages/core/src/tools/trackerTools.ts
@@ -0,0 +1,606 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { Config } from '../config/config.js';
+import type { MessageBus } from '../confirmation-bus/message-bus.js';
+import {
+ TRACKER_ADD_DEPENDENCY_DEFINITION,
+ TRACKER_CREATE_TASK_DEFINITION,
+ TRACKER_GET_TASK_DEFINITION,
+ TRACKER_LIST_TASKS_DEFINITION,
+ TRACKER_UPDATE_TASK_DEFINITION,
+ TRACKER_VISUALIZE_DEFINITION,
+} from './definitions/trackerTools.js';
+import { resolveToolDeclaration } from './definitions/resolver.js';
+import {
+ TRACKER_ADD_DEPENDENCY_TOOL_NAME,
+ TRACKER_CREATE_TASK_TOOL_NAME,
+ TRACKER_GET_TASK_TOOL_NAME,
+ TRACKER_LIST_TASKS_TOOL_NAME,
+ TRACKER_UPDATE_TASK_TOOL_NAME,
+ TRACKER_VISUALIZE_TOOL_NAME,
+} from './tool-names.js';
+import type { ToolResult } from './tools.js';
+import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import { ToolErrorType } from './tool-error.js';
+import type { TrackerTask, TaskType } from '../services/trackerTypes.js';
+import { TaskStatus } from '../services/trackerTypes.js';
+
+// --- tracker_create_task ---
+
+interface CreateTaskParams {
+ title: string;
+ description: string;
+ type: TaskType;
+ parentId?: string;
+ dependencies?: string[];
+}
+
+class TrackerCreateTaskInvocation extends BaseToolInvocation<
+ CreateTaskParams,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: CreateTaskParams,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return `Creating task: ${this.params.title}`;
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ try {
+ const task = await this.service.createTask({
+ title: this.params.title,
+ description: this.params.description,
+ type: this.params.type,
+ status: TaskStatus.OPEN,
+ parentId: this.params.parentId,
+ dependencies: this.params.dependencies ?? [],
+ });
+ return {
+ llmContent: `Created task ${task.id}: ${task.title}`,
+ returnDisplay: `Created task ${task.id}.`,
+ };
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ return {
+ llmContent: `Error creating task: ${errorMessage}`,
+ returnDisplay: 'Failed to create task.',
+ error: {
+ message: errorMessage,
+ type: ToolErrorType.EXECUTION_FAILED,
+ },
+ };
+ }
+ }
+}
+
+export class TrackerCreateTaskTool extends BaseDeclarativeTool<
+ CreateTaskParams,
+ ToolResult
+> {
+ static readonly Name = TRACKER_CREATE_TASK_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerCreateTaskTool.Name,
+ 'Create Task',
+ TRACKER_CREATE_TASK_DEFINITION.base.description!,
+ Kind.Edit,
+ TRACKER_CREATE_TASK_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(params: CreateTaskParams, messageBus: MessageBus) {
+ return new TrackerCreateTaskInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_CREATE_TASK_DEFINITION, modelId);
+ }
+}
+
+// --- tracker_update_task ---
+
+interface UpdateTaskParams {
+ id: string;
+ title?: string;
+ description?: string;
+ status?: TaskStatus;
+ dependencies?: string[];
+}
+
+class TrackerUpdateTaskInvocation extends BaseToolInvocation<
+ UpdateTaskParams,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: UpdateTaskParams,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return `Updating task ${this.params.id}`;
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ const { id, ...updates } = this.params;
+ try {
+ const task = await this.service.updateTask(id, updates);
+ return {
+ llmContent: `Updated task ${task.id}. Status: ${task.status}`,
+ returnDisplay: `Updated task ${task.id}.`,
+ };
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ return {
+ llmContent: `Error updating task: ${errorMessage}`,
+ returnDisplay: 'Failed to update task.',
+ error: {
+ message: errorMessage,
+ type: ToolErrorType.EXECUTION_FAILED,
+ },
+ };
+ }
+ }
+}
+
+export class TrackerUpdateTaskTool extends BaseDeclarativeTool<
+ UpdateTaskParams,
+ ToolResult
+> {
+ static readonly Name = TRACKER_UPDATE_TASK_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerUpdateTaskTool.Name,
+ 'Update Task',
+ TRACKER_UPDATE_TASK_DEFINITION.base.description!,
+ Kind.Edit,
+ TRACKER_UPDATE_TASK_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(params: UpdateTaskParams, messageBus: MessageBus) {
+ return new TrackerUpdateTaskInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_UPDATE_TASK_DEFINITION, modelId);
+ }
+}
+
+// --- tracker_get_task ---
+
+interface GetTaskParams {
+ id: string;
+}
+
+class TrackerGetTaskInvocation extends BaseToolInvocation<
+ GetTaskParams,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: GetTaskParams,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return `Retrieving task ${this.params.id}`;
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ const task = await this.service.getTask(this.params.id);
+ if (!task) {
+ return {
+ llmContent: `Task ${this.params.id} not found.`,
+ returnDisplay: 'Task not found.',
+ };
+ }
+ return {
+ llmContent: JSON.stringify(task, null, 2),
+ returnDisplay: `Retrieved task ${task.id}.`,
+ };
+ }
+}
+
+export class TrackerGetTaskTool extends BaseDeclarativeTool<
+ GetTaskParams,
+ ToolResult
+> {
+ static readonly Name = TRACKER_GET_TASK_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerGetTaskTool.Name,
+ 'Get Task',
+ TRACKER_GET_TASK_DEFINITION.base.description!,
+ Kind.Read,
+ TRACKER_GET_TASK_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(params: GetTaskParams, messageBus: MessageBus) {
+ return new TrackerGetTaskInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_GET_TASK_DEFINITION, modelId);
+ }
+}
+
+// --- tracker_list_tasks ---
+
+interface ListTasksParams {
+ status?: TaskStatus;
+ type?: TaskType;
+ parentId?: string;
+}
+
+class TrackerListTasksInvocation extends BaseToolInvocation<
+ ListTasksParams,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: ListTasksParams,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return 'Listing tasks.';
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ let tasks = await this.service.listTasks();
+ if (this.params.status) {
+ tasks = tasks.filter((t) => t.status === this.params.status);
+ }
+ if (this.params.type) {
+ tasks = tasks.filter((t) => t.type === this.params.type);
+ }
+ if (this.params.parentId) {
+ tasks = tasks.filter((t) => t.parentId === this.params.parentId);
+ }
+
+ if (tasks.length === 0) {
+ return {
+ llmContent: 'No tasks found matching the criteria.',
+ returnDisplay: 'No matching tasks.',
+ };
+ }
+
+ const content = tasks
+ .map((t) => `- [${t.id}] ${t.title} (${t.status})`)
+ .join('\n');
+ return {
+ llmContent: content,
+ returnDisplay: `Listed ${tasks.length} tasks.`,
+ };
+ }
+}
+
+export class TrackerListTasksTool extends BaseDeclarativeTool<
+ ListTasksParams,
+ ToolResult
+> {
+ static readonly Name = TRACKER_LIST_TASKS_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerListTasksTool.Name,
+ 'List Tasks',
+ TRACKER_LIST_TASKS_DEFINITION.base.description!,
+ Kind.Search,
+ TRACKER_LIST_TASKS_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(params: ListTasksParams, messageBus: MessageBus) {
+ return new TrackerListTasksInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_LIST_TASKS_DEFINITION, modelId);
+ }
+}
+
+// --- tracker_add_dependency ---
+
+interface AddDependencyParams {
+ taskId: string;
+ dependencyId: string;
+}
+
+class TrackerAddDependencyInvocation extends BaseToolInvocation<
+ AddDependencyParams,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: AddDependencyParams,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return `Adding dependency: ${this.params.taskId} depends on ${this.params.dependencyId}`;
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ if (this.params.taskId === this.params.dependencyId) {
+ return {
+ llmContent: `Error: Task ${this.params.taskId} cannot depend on itself.`,
+ returnDisplay: 'Self-referential dependency rejected.',
+ error: {
+ message: 'Task cannot depend on itself',
+ type: ToolErrorType.EXECUTION_FAILED,
+ },
+ };
+ }
+
+ const [task, dep] = await Promise.all([
+ this.service.getTask(this.params.taskId),
+ this.service.getTask(this.params.dependencyId),
+ ]);
+
+ if (!task) {
+ return {
+ llmContent: `Task ${this.params.taskId} not found.`,
+ returnDisplay: 'Task not found.',
+ };
+ }
+ if (!dep) {
+ return {
+ llmContent: `Dependency task ${this.params.dependencyId} not found.`,
+ returnDisplay: 'Dependency not found.',
+ };
+ }
+
+ const newDeps = Array.from(
+ new Set([...task.dependencies, this.params.dependencyId]),
+ );
+ try {
+ await this.service.updateTask(task.id, { dependencies: newDeps });
+ return {
+ llmContent: `Linked ${task.id} -> ${dep.id}.`,
+ returnDisplay: 'Dependency added.',
+ };
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ return {
+ llmContent: `Error adding dependency: ${errorMessage}`,
+ returnDisplay: 'Failed to add dependency.',
+ error: {
+ message: errorMessage,
+ type: ToolErrorType.EXECUTION_FAILED,
+ },
+ };
+ }
+ }
+}
+
+export class TrackerAddDependencyTool extends BaseDeclarativeTool<
+ AddDependencyParams,
+ ToolResult
+> {
+ static readonly Name = TRACKER_ADD_DEPENDENCY_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerAddDependencyTool.Name,
+ 'Add Dependency',
+ TRACKER_ADD_DEPENDENCY_DEFINITION.base.description!,
+ Kind.Edit,
+ TRACKER_ADD_DEPENDENCY_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(
+ params: AddDependencyParams,
+ messageBus: MessageBus,
+ ) {
+ return new TrackerAddDependencyInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_ADD_DEPENDENCY_DEFINITION, modelId);
+ }
+}
+
+// --- tracker_visualize ---
+
+class TrackerVisualizeInvocation extends BaseToolInvocation<
+ Record,
+ ToolResult
+> {
+ constructor(
+ private readonly config: Config,
+ params: Record,
+ messageBus: MessageBus,
+ toolName: string,
+ ) {
+ super(params, messageBus, toolName);
+ }
+
+ private get service() {
+ return this.config.getTrackerService();
+ }
+ getDescription(): string {
+ return 'Visualizing the task graph.';
+ }
+
+ override async execute(_signal: AbortSignal): Promise {
+ const tasks = await this.service.listTasks();
+ if (tasks.length === 0) {
+ return {
+ llmContent: 'No tasks to visualize.',
+ returnDisplay: 'Empty tracker.',
+ };
+ }
+
+ const statusEmojis: Record = {
+ open: '⭕',
+ in_progress: '🚧',
+ blocked: '🚫',
+ closed: '✅',
+ };
+
+ const typeLabels: Record = {
+ epic: '[EPIC]',
+ task: '[TASK]',
+ bug: '[BUG]',
+ };
+
+ const childrenMap = new Map();
+ const roots: TrackerTask[] = [];
+
+ for (const task of tasks) {
+ if (task.parentId) {
+ if (!childrenMap.has(task.parentId)) {
+ childrenMap.set(task.parentId, []);
+ }
+ childrenMap.get(task.parentId)!.push(task);
+ } else {
+ roots.push(task);
+ }
+ }
+
+ let output = 'Task Tracker Graph:\n';
+
+ const renderTask = (
+ task: TrackerTask,
+ depth: number,
+ visited: Set,
+ ) => {
+ if (visited.has(task.id)) {
+ output += `${' '.repeat(depth)}[CYCLE DETECTED: ${task.id}]\n`;
+ return;
+ }
+ visited.add(task.id);
+
+ const indent = ' '.repeat(depth);
+ output += `${indent}${statusEmojis[task.status]} ${task.id} ${typeLabels[task.type]} ${task.title}\n`;
+ if (task.dependencies.length > 0) {
+ output += `${indent} └─ Depends on: ${task.dependencies.join(', ')}\n`;
+ }
+ const children = childrenMap.get(task.id) ?? [];
+ for (const child of children) {
+ renderTask(child, depth + 1, new Set(visited));
+ }
+ };
+
+ for (const root of roots) {
+ renderTask(root, 0, new Set());
+ }
+
+ return {
+ llmContent: output,
+ returnDisplay: output,
+ };
+ }
+}
+
+export class TrackerVisualizeTool extends BaseDeclarativeTool<
+ Record,
+ ToolResult
+> {
+ static readonly Name = TRACKER_VISUALIZE_TOOL_NAME;
+ constructor(
+ private config: Config,
+ messageBus: MessageBus,
+ ) {
+ super(
+ TrackerVisualizeTool.Name,
+ 'Visualize Tracker',
+ TRACKER_VISUALIZE_DEFINITION.base.description!,
+ Kind.Read,
+ TRACKER_VISUALIZE_DEFINITION.base.parametersJsonSchema,
+ messageBus,
+ );
+ }
+ protected createInvocation(
+ params: Record,
+ messageBus: MessageBus,
+ ) {
+ return new TrackerVisualizeInvocation(
+ this.config,
+ params,
+ messageBus,
+ this.name,
+ );
+ }
+ override getSchema(modelId?: string) {
+ return resolveToolDeclaration(TRACKER_VISUALIZE_DEFINITION, modelId);
+ }
+}
diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts
index 55d2474c1c..3170227188 100644
--- a/packages/core/src/tools/web-fetch.ts
+++ b/packages/core/src/tools/web-fetch.ts
@@ -4,13 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type {
- ToolCallConfirmationDetails,
- ToolInvocation,
- ToolResult,
- ToolConfirmationOutcome,
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolCallConfirmationDetails,
+ type ToolInvocation,
+ type ToolResult,
+ type ToolConfirmationOutcome,
} from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { ToolErrorType } from './tool-error.js';
import { getErrorMessage } from '../utils/errors.js';
diff --git a/packages/core/src/tools/web-search.test.ts b/packages/core/src/tools/web-search.test.ts
index 3812a54879..bd07ce0dea 100644
--- a/packages/core/src/tools/web-search.test.ts
+++ b/packages/core/src/tools/web-search.test.ts
@@ -4,10 +4,16 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { Mock } from 'vitest';
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import type { WebSearchToolParams } from './web-search.js';
-import { WebSearchTool } from './web-search.js';
+import {
+ describe,
+ it,
+ expect,
+ vi,
+ beforeEach,
+ afterEach,
+ type Mock,
+} from 'vitest';
+import { WebSearchTool, type WebSearchToolParams } from './web-search.js';
import type { Config } from '../config/config.js';
import { GeminiClient } from '../core/client.js';
import { ToolErrorType } from './tool-error.js';
diff --git a/packages/core/src/tools/web-search.ts b/packages/core/src/tools/web-search.ts
index a5ac9937b8..2756599b28 100644
--- a/packages/core/src/tools/web-search.ts
+++ b/packages/core/src/tools/web-search.ts
@@ -7,8 +7,13 @@
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { WEB_SEARCH_TOOL_NAME } from './tool-names.js';
import type { GroundingMetadata } from '@google/genai';
-import type { ToolInvocation, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type ToolResult,
+} from './tools.js';
import { ToolErrorType } from './tool-error.js';
import { getErrorMessage } from '../utils/errors.js';
diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts
index 0b978f14f9..e90937bd7d 100644
--- a/packages/core/src/tools/write-file.test.ts
+++ b/packages/core/src/tools/write-file.test.ts
@@ -13,16 +13,19 @@ import {
vi,
type Mocked,
} from 'vitest';
-import type { WriteFileToolParams } from './write-file.js';
-import { getCorrectedFileContent, WriteFileTool } from './write-file.js';
+import {
+ getCorrectedFileContent,
+ WriteFileTool,
+ type WriteFileToolParams,
+} from './write-file.js';
import { ToolErrorType } from './tool-error.js';
-import type {
- FileDiff,
- ToolEditConfirmationDetails,
- ToolInvocation,
- ToolResult,
+import {
+ ToolConfirmationOutcome,
+ type FileDiff,
+ type ToolEditConfirmationDetails,
+ type ToolInvocation,
+ type ToolResult,
} from './tools.js';
-import { ToolConfirmationOutcome } from './tools.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../policy/types.js';
import type { ToolRegistry } from './tool-registry.js';
@@ -34,8 +37,7 @@ import { GeminiClient } from '../core/client.js';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
import { ensureCorrectFileContent } from '../utils/editCorrector.js';
import { StandardFileSystemService } from '../services/fileSystemService.js';
-import type { DiffUpdateResult } from '../ide/ide-client.js';
-import { IdeClient } from '../ide/ide-client.js';
+import { IdeClient, type DiffUpdateResult } from '../ide/ide-client.js';
import { WorkspaceContext } from '../utils/workspaceContext.js';
import {
createMockMessageBus,
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
index 1c8a230001..8ec660b661 100644
--- a/packages/core/src/tools/write-file.ts
+++ b/packages/core/src/tools/write-file.ts
@@ -13,16 +13,18 @@ import { WRITE_FILE_TOOL_NAME, WRITE_FILE_DISPLAY_NAME } from './tool-names.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../policy/types.js';
-import type {
- FileDiff,
- ToolCallConfirmationDetails,
- ToolEditConfirmationDetails,
- ToolInvocation,
- ToolLocation,
- ToolResult,
- ToolConfirmationOutcome,
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type FileDiff,
+ type ToolCallConfirmationDetails,
+ type ToolEditConfirmationDetails,
+ type ToolInvocation,
+ type ToolLocation,
+ type ToolResult,
+ type ToolConfirmationOutcome,
} from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
import { ToolErrorType } from './tool-error.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import { getErrorMessage, isNodeError } from '../utils/errors.js';
@@ -72,6 +74,20 @@ export interface WriteFileToolParams {
ai_proposed_content?: string;
}
+export function isWriteFileToolParams(
+ args: unknown,
+): args is WriteFileToolParams {
+ if (typeof args !== 'object' || args === null) {
+ return false;
+ }
+ return (
+ 'file_path' in args &&
+ typeof args.file_path === 'string' &&
+ 'content' in args &&
+ typeof args.content === 'string'
+ );
+}
+
interface GetCorrectedFileContentResult {
originalContent: string;
correctedContent: string;
diff --git a/packages/core/src/tools/write-todos.ts b/packages/core/src/tools/write-todos.ts
index 5eb42c73f4..dd7ab780e6 100644
--- a/packages/core/src/tools/write-todos.ts
+++ b/packages/core/src/tools/write-todos.ts
@@ -4,8 +4,14 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { ToolInvocation, Todo, ToolResult } from './tools.js';
-import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js';
+import {
+ BaseDeclarativeTool,
+ BaseToolInvocation,
+ Kind,
+ type ToolInvocation,
+ type Todo,
+ type ToolResult,
+} from './tools.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { WRITE_TODOS_TOOL_NAME } from './tool-names.js';
import { WRITE_TODOS_DEFINITION } from './definitions/coreTools.js';
diff --git a/packages/core/src/utils/apiConversionUtils.test.ts b/packages/core/src/utils/apiConversionUtils.test.ts
index 615bcb1de8..fa907ca2e6 100644
--- a/packages/core/src/utils/apiConversionUtils.test.ts
+++ b/packages/core/src/utils/apiConversionUtils.test.ts
@@ -6,11 +6,11 @@
import { describe, it, expect } from 'vitest';
import { convertToRestPayload } from './apiConversionUtils.js';
-import type { GenerateContentParameters } from '@google/genai';
import {
FunctionCallingConfigMode,
HarmCategory,
HarmBlockThreshold,
+ type GenerateContentParameters,
} from '@google/genai';
describe('apiConversionUtils', () => {
diff --git a/packages/core/src/utils/authConsent.test.ts b/packages/core/src/utils/authConsent.test.ts
index 7fc05b2a03..2eccbd39c8 100644
--- a/packages/core/src/utils/authConsent.test.ts
+++ b/packages/core/src/utils/authConsent.test.ts
@@ -4,8 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-import type { Mock } from 'vitest';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import readline from 'node:readline';
import process from 'node:process';
import { coreEvents } from './events.js';
diff --git a/packages/core/src/utils/cache.test.ts b/packages/core/src/utils/cache.test.ts
new file mode 100644
index 0000000000..249b63fe25
--- /dev/null
+++ b/packages/core/src/utils/cache.test.ts
@@ -0,0 +1,198 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { createCache } from './cache.js';
+
+describe('CacheService', () => {
+ beforeEach(() => {
+ vi.useFakeTimers();
+ });
+
+ afterEach(() => {
+ vi.useRealTimers();
+ });
+
+ describe('Basic operations', () => {
+ it('should store and retrieve values by default (Map)', () => {
+ const cache = createCache({ storage: 'map' });
+ cache.set('key', 'value');
+ expect(cache.get('key')).toBe('value');
+ });
+
+ it('should return undefined for missing keys', () => {
+ const cache = createCache({ storage: 'map' });
+ expect(cache.get('missing')).toBeUndefined();
+ });
+
+ it('should delete entries', () => {
+ const cache = createCache({ storage: 'map' });
+ cache.set('key', 'value');
+ cache.delete('key');
+ expect(cache.get('key')).toBeUndefined();
+ });
+
+ it('should clear all entries (Map)', () => {
+ const cache = createCache({ storage: 'map' });
+ cache.set('k1', 'v1');
+ cache.set('k2', 'v2');
+ cache.clear();
+ expect(cache.get('k1')).toBeUndefined();
+ expect(cache.get('k2')).toBeUndefined();
+ });
+
+ it('should throw on clear() for WeakMap', () => {
+ const cache = createCache({ storage: 'weakmap' });
+ expect(() => cache.clear()).toThrow(
+ 'clear() is not supported on WeakMap storage',
+ );
+ });
+ });
+
+ describe('TTL and Expiration', () => {
+ it('should expire entries based on defaultTtl', () => {
+ const cache = createCache({
+ storage: 'map',
+ defaultTtl: 1000,
+ });
+ cache.set('key', 'value');
+
+ vi.advanceTimersByTime(500);
+ expect(cache.get('key')).toBe('value');
+
+ vi.advanceTimersByTime(600); // Total 1100
+ expect(cache.get('key')).toBeUndefined();
+ });
+
+ it('should expire entries based on specific ttl override', () => {
+ const cache = createCache({
+ storage: 'map',
+ defaultTtl: 5000,
+ });
+ cache.set('key', 'value', 1000);
+
+ vi.advanceTimersByTime(1100);
+ expect(cache.get('key')).toBeUndefined();
+ });
+
+ it('should not expire if ttl is undefined', () => {
+ const cache = createCache({ storage: 'map' });
+ cache.set('key', 'value');
+
+ vi.advanceTimersByTime(100000);
+ expect(cache.get('key')).toBe('value');
+ });
+ });
+
+ describe('getOrCreate', () => {
+ it('should return existing value if not expired', () => {
+ const cache = createCache({ storage: 'map' });
+ cache.set('key', 'old');
+ const creator = vi.fn().mockReturnValue('new');
+
+ const result = cache.getOrCreate('key', creator);
+ expect(result).toBe('old');
+ expect(creator).not.toHaveBeenCalled();
+ });
+
+ it('should create and store value if missing', () => {
+ const cache = createCache({ storage: 'map' });
+ const creator = vi.fn().mockReturnValue('new');
+
+ const result = cache.getOrCreate('key', creator);
+ expect(result).toBe('new');
+ expect(creator).toHaveBeenCalled();
+ expect(cache.get('key')).toBe('new');
+ });
+
+ it('should recreate value if expired', () => {
+ const cache = createCache({
+ storage: 'map',
+ defaultTtl: 1000,
+ });
+ cache.set('key', 'old');
+ vi.advanceTimersByTime(1100);
+
+ const creator = vi.fn().mockReturnValue('new');
+ const result = cache.getOrCreate('key', creator);
+ expect(result).toBe('new');
+ expect(creator).toHaveBeenCalled();
+ });
+ });
+
+ describe('Promise Support', () => {
+ beforeEach(() => {
+ vi.useRealTimers();
+ });
+
+ it('should remove failed promises from cache by default', async () => {
+ const cache = createCache>({ storage: 'map' });
+ const promise = Promise.reject(new Error('fail'));
+
+ // We need to catch it to avoid unhandled rejection in test
+ promise.catch(() => {});
+
+ cache.set('key', promise);
+ expect(cache.get('key')).toBe(promise);
+
+ // Wait for promise to settle
+ await new Promise((resolve) => setImmediate(resolve));
+
+ expect(cache.get('key')).toBeUndefined();
+ });
+
+ it('should NOT remove failed promises if deleteOnPromiseFailure is false', async () => {
+ const cache = createCache>({
+ storage: 'map',
+ deleteOnPromiseFailure: false,
+ });
+ const promise = Promise.reject(new Error('fail'));
+ promise.catch(() => {});
+
+ cache.set('key', promise);
+
+ await new Promise((resolve) => setImmediate(resolve));
+
+ expect(cache.get('key')).toBe(promise);
+ });
+
+ it('should only delete the specific failed entry', async () => {
+ const cache = createCache>({ storage: 'map' });
+
+ const failPromise = Promise.reject(new Error('fail'));
+ failPromise.catch(() => {});
+
+ cache.set('key', failPromise);
+
+ // Overwrite with a new success promise before failure settles
+ const successPromise = Promise.resolve('ok');
+ cache.set('key', successPromise);
+
+ await new Promise((resolve) => setImmediate(resolve));
+
+ // Should still be successPromise
+ expect(cache.get('key')).toBe(successPromise);
+ });
+ });
+
+ describe('WeakMap Storage', () => {
+ it('should work with object keys explicitly', () => {
+ const cache = createCache({ storage: 'weakmap' });
+ const key = { id: 1 };
+ cache.set(key, 'value');
+ expect(cache.get(key)).toBe('value');
+ });
+
+ it('should default to Map for objects', () => {
+ const cache = createCache();
+ const key = { id: 1 };
+ cache.set(key, 'value');
+ expect(cache.get(key)).toBe('value');
+ // clear() should NOT throw because default is Map
+ expect(() => cache.clear()).not.toThrow();
+ });
+ });
+});
diff --git a/packages/core/src/utils/cache.ts b/packages/core/src/utils/cache.ts
new file mode 100644
index 0000000000..948d9f637c
--- /dev/null
+++ b/packages/core/src/utils/cache.ts
@@ -0,0 +1,151 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export interface CacheEntry {
+ value: V;
+ timestamp: number;
+ ttl?: number;
+}
+
+export interface CacheOptions {
+ /**
+ * Default Time To Live in milliseconds.
+ */
+ defaultTtl?: number;
+
+ /**
+ * If true, and V is a Promise, the entry will be removed from the cache
+ * if the promise rejects.
+ */
+ deleteOnPromiseFailure?: boolean;
+
+ /**
+ * The underlying storage mechanism.
+ * Use 'weakmap' (default) for object keys to allow garbage collection.
+ * Use 'map' if you need to use strings as keys or need the clear() method.
+ */
+ storage?: 'map' | 'weakmap';
+}
+
+/**
+ * A generic caching service with TTL support.
+ */
+export class CacheService {
+ private readonly storage:
+ | Map>
+ | WeakMap>;
+ private readonly defaultTtl?: number;
+ private readonly deleteOnPromiseFailure: boolean;
+
+ constructor(options: CacheOptions = {}) {
+ // Default to map for safety unless weakmap is explicitly requested.
+ this.storage =
+ options.storage === 'weakmap'
+ ? new WeakMap>()
+ : new Map>();
+ this.defaultTtl = options.defaultTtl;
+ this.deleteOnPromiseFailure = options.deleteOnPromiseFailure ?? true;
+ }
+
+ /**
+ * Retrieves a value from the cache. Returns undefined if missing or expired.
+ */
+ get(key: K): V | undefined {
+ // We have to cast to Map or WeakMap specifically to call get()
+ // but since they have the same signature for object keys, we can
+ // safely cast to 'any' internally for the dispatch.
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion
+ const entry = (this.storage as any).get(key) as CacheEntry | undefined;
+ if (!entry) {
+ return undefined;
+ }
+
+ const ttl = entry.ttl ?? this.defaultTtl;
+ if (ttl !== undefined && Date.now() - entry.timestamp > ttl) {
+ this.delete(key);
+ return undefined;
+ }
+
+ return entry.value;
+ }
+
+ /**
+ * Stores a value in the cache.
+ */
+ set(key: K, value: V, ttl?: number): void {
+ const entry: CacheEntry = {
+ value,
+ timestamp: Date.now(),
+ ttl,
+ };
+
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion
+ (this.storage as any).set(key, entry);
+
+ if (this.deleteOnPromiseFailure && value instanceof Promise) {
+ value.catch(() => {
+ // Only delete if this exact entry is still in the cache
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion
+ if ((this.storage as any).get(key) === entry) {
+ this.delete(key);
+ }
+ });
+ }
+ }
+
+ /**
+ * Helper to retrieve a value or create it if missing/expired.
+ */
+ getOrCreate(key: K, creator: () => V, ttl?: number): V {
+ let value = this.get(key);
+ if (value === undefined) {
+ value = creator();
+ this.set(key, value, ttl);
+ }
+ return value;
+ }
+
+ /**
+ * Removes an entry from the cache.
+ */
+ delete(key: K): void {
+ if (this.storage instanceof Map) {
+ this.storage.delete(key);
+ } else {
+ // WeakMap.delete returns a boolean, we can ignore it.
+ // Cast to any to bypass the WeakKey constraint since we've already
+ // confirmed the storage type.
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion
+ (this.storage as any).delete(key);
+ }
+ }
+
+ /**
+ * Clears all entries. Only supported if using Map storage.
+ */
+ clear(): void {
+ if (this.storage instanceof Map) {
+ this.storage.clear();
+ } else {
+ throw new Error('clear() is not supported on WeakMap storage');
+ }
+ }
+}
+
+/**
+ * Factory function to create a new cache.
+ */
+export function createCache(
+ options: CacheOptions & { storage: 'map' },
+): CacheService;
+export function createCache(
+ options?: CacheOptions,
+): CacheService;
+export function createCache(
+ options: CacheOptions = {},
+): CacheService {
+ return new CacheService(options);
+}
diff --git a/packages/core/src/utils/editCorrector.test.ts b/packages/core/src/utils/editCorrector.test.ts
index 533b49b9e4..f9620d74b5 100644
--- a/packages/core/src/utils/editCorrector.test.ts
+++ b/packages/core/src/utils/editCorrector.test.ts
@@ -5,8 +5,7 @@
*/
/* eslint-disable @typescript-eslint/no-explicit-any */
-import type { Mocked } from 'vitest';
-import { vi, describe, it, expect, beforeEach } from 'vitest';
+import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
// MOCKS
diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts
index f8ff81b97e..2c58bad98f 100644
--- a/packages/core/src/utils/editCorrector.ts
+++ b/packages/core/src/utils/editCorrector.ts
@@ -112,6 +112,7 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr
if (
result &&
+ // eslint-disable-next-line no-restricted-syntax
typeof result['corrected_string_escaping'] === 'string' &&
result['corrected_string_escaping'].length > 0
) {
diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts
index 159dde2a6d..47c42c93ba 100644
--- a/packages/core/src/utils/events.ts
+++ b/packages/core/src/utils/events.ts
@@ -149,6 +149,10 @@ export interface SlashCommandConflict {
renamedTo: string;
loserExtensionName?: string;
winnerExtensionName?: string;
+ loserMcpServerName?: string;
+ winnerMcpServerName?: string;
+ loserKind?: string;
+ winnerKind?: string;
}
export interface SlashCommandConflictsPayload {
@@ -232,6 +236,7 @@ type EventBacklogItem = {
export class CoreEventEmitter extends EventEmitter {
private _eventBacklog: EventBacklogItem[] = [];
+ private _backlogHead = 0;
private static readonly MAX_BACKLOG_SIZE = 10000;
constructor() {
@@ -243,8 +248,17 @@ export class CoreEventEmitter extends EventEmitter {
...args: CoreEvents[K]
): void {
if (this.listenerCount(event) === 0) {
- if (this._eventBacklog.length >= CoreEventEmitter.MAX_BACKLOG_SIZE) {
- this._eventBacklog.shift();
+ const backlogSize = this._eventBacklog.length - this._backlogHead;
+ if (backlogSize >= CoreEventEmitter.MAX_BACKLOG_SIZE) {
+ // Evict oldest entry. Use a head pointer instead of shift() to avoid
+ // O(n) array reindexing on every eviction at capacity.
+ (this._eventBacklog as unknown[])[this._backlogHead] = undefined;
+ this._backlogHead++;
+ // Compact once dead entries exceed half capacity to bound memory
+ if (this._backlogHead >= CoreEventEmitter.MAX_BACKLOG_SIZE / 2) {
+ this._eventBacklog = this._eventBacklog.slice(this._backlogHead);
+ this._backlogHead = 0;
+ }
}
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this._eventBacklog.push({ event, args } as EventBacklogItem);
@@ -391,9 +405,13 @@ export class CoreEventEmitter extends EventEmitter {
* subscribes.
*/
drainBacklogs(): void {
- const backlog = [...this._eventBacklog];
- this._eventBacklog.length = 0; // Clear in-place
- for (const item of backlog) {
+ const backlog = this._eventBacklog;
+ const head = this._backlogHead;
+ this._eventBacklog = [];
+ this._backlogHead = 0;
+ for (let i = head; i < backlog.length; i++) {
+ const item = backlog[i];
+ if (item === undefined) continue;
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
(this.emit as (event: keyof CoreEvents, ...args: unknown[]) => boolean)(
item.event,
diff --git a/packages/core/src/utils/fileUtils.test.ts b/packages/core/src/utils/fileUtils.test.ts
index de668db3ad..dcbf22c5a7 100644
--- a/packages/core/src/utils/fileUtils.test.ts
+++ b/packages/core/src/utils/fileUtils.test.ts
@@ -14,8 +14,8 @@ import {
type Mock,
} from 'vitest';
-import * as actualNodeFs from 'node:fs'; // For setup/teardown
import fs from 'node:fs';
+import * as actualNodeFs from 'node:fs'; // For setup/teardown
import fsPromises from 'node:fs/promises';
import path from 'node:path';
import os from 'node:os';
diff --git a/packages/core/src/utils/filesearch/crawler.test.ts b/packages/core/src/utils/filesearch/crawler.test.ts
index 192c0274b8..5cdeb79fdb 100644
--- a/packages/core/src/utils/filesearch/crawler.test.ts
+++ b/packages/core/src/utils/filesearch/crawler.test.ts
@@ -10,8 +10,7 @@ import * as path from 'node:path';
import * as cache from './crawlCache.js';
import { crawl } from './crawler.js';
import { createTmpDir, cleanupTmpDir } from '@google/gemini-cli-test-utils';
-import type { Ignore } from './ignore.js';
-import { loadIgnoreRules } from './ignore.js';
+import { loadIgnoreRules, type Ignore } from './ignore.js';
import { GEMINI_IGNORE_FILE_NAME } from '../../config/constants.js';
import { FileDiscoveryService } from '../../services/fileDiscoveryService.js';
diff --git a/packages/core/src/utils/filesearch/fileSearch.test.ts b/packages/core/src/utils/filesearch/fileSearch.test.ts
index 3c2506cb13..1c001eeead 100644
--- a/packages/core/src/utils/filesearch/fileSearch.test.ts
+++ b/packages/core/src/utils/filesearch/fileSearch.test.ts
@@ -421,6 +421,47 @@ describe('FileSearch', () => {
);
});
+ it('should prioritize filenames closer to the end of the path and shorter paths', async () => {
+ tmpDir = await createTmpDir({
+ src: {
+ 'hooks.ts': '',
+ hooks: {
+ 'index.ts': '',
+ },
+ utils: {
+ 'hooks.tsx': '',
+ },
+ 'hooks-dev': {
+ 'test.ts': '',
+ },
+ },
+ });
+
+ const fileSearch = FileSearchFactory.create({
+ projectRoot: tmpDir,
+ fileDiscoveryService: new FileDiscoveryService(tmpDir, {
+ respectGitIgnore: false,
+ respectGeminiIgnore: false,
+ }),
+ ignoreDirs: [],
+ cache: false,
+ cacheTtl: 0,
+ enableRecursiveFileSearch: true,
+ enableFuzzySearch: true,
+ });
+
+ await fileSearch.initialize();
+ const results = await fileSearch.search('hooks');
+
+ // The order should prioritize matches closer to the end and shorter strings.
+ // FZF matches right-to-left.
+ expect(results[0]).toBe('src/hooks/');
+ expect(results[1]).toBe('src/hooks.ts');
+ expect(results[2]).toBe('src/utils/hooks.tsx');
+ expect(results[3]).toBe('src/hooks-dev/');
+ expect(results[4]).toBe('src/hooks/index.ts');
+ expect(results[5]).toBe('src/hooks-dev/test.ts');
+ });
it('should return empty array when no matches are found', async () => {
tmpDir = await createTmpDir({
src: ['file1.js'],
diff --git a/packages/core/src/utils/filesearch/fileSearch.ts b/packages/core/src/utils/filesearch/fileSearch.ts
index 97560f7070..e3f608e508 100644
--- a/packages/core/src/utils/filesearch/fileSearch.ts
+++ b/packages/core/src/utils/filesearch/fileSearch.ts
@@ -6,15 +6,51 @@
import path from 'node:path';
import picomatch from 'picomatch';
-import type { Ignore } from './ignore.js';
-import { loadIgnoreRules } from './ignore.js';
+import { loadIgnoreRules, type Ignore } from './ignore.js';
import { ResultCache } from './result-cache.js';
import { crawl } from './crawler.js';
-import type { FzfResultItem } from 'fzf';
-import { AsyncFzf } from 'fzf';
+import { AsyncFzf, type FzfResultItem } from 'fzf';
import { unescapePath } from '../paths.js';
import type { FileDiscoveryService } from '../../services/fileDiscoveryService.js';
+// Tiebreaker: Prefers shorter paths.
+const byLengthAsc = (a: { item: string }, b: { item: string }) =>
+ a.item.length - b.item.length;
+
+// Tiebreaker: Prefers matches at the start of the filename (basename prefix).
+const byBasenamePrefix = (
+ a: { item: string; positions: Set },
+ b: { item: string; positions: Set },
+) => {
+ const getBasenameStart = (p: string) => {
+ const trimmed = p.endsWith('/') ? p.slice(0, -1) : p;
+ return Math.max(trimmed.lastIndexOf('/'), trimmed.lastIndexOf('\\')) + 1;
+ };
+ const aDiff = Math.min(...a.positions) - getBasenameStart(a.item);
+ const bDiff = Math.min(...b.positions) - getBasenameStart(b.item);
+
+ const aIsFilenameMatch = aDiff >= 0;
+ const bIsFilenameMatch = bDiff >= 0;
+
+ if (aIsFilenameMatch && !bIsFilenameMatch) return -1;
+ if (!aIsFilenameMatch && bIsFilenameMatch) return 1;
+ if (aIsFilenameMatch && bIsFilenameMatch) return aDiff - bDiff;
+
+ return 0; // Both are directory matches, let subsequent tiebreakers decide.
+};
+
+// Tiebreaker: Prefers matches closer to the end of the path.
+const byMatchPosFromEnd = (
+ a: { item: string; positions: Set },
+ b: { item: string; positions: Set },
+) => {
+ const maxPosA = Math.max(-1, ...a.positions);
+ const maxPosB = Math.max(-1, ...b.positions);
+ const distA = a.item.length - maxPosA;
+ const distB = b.item.length - maxPosB;
+ return distA - distB;
+};
+
export interface FileSearchOptions {
projectRoot: string;
ignoreDirs: string[];
@@ -194,6 +230,8 @@ class RecursiveFileSearch implements FileSearch {
// files, because the v2 algorithm is just too slow in those cases.
this.fzf = new AsyncFzf(this.allFiles, {
fuzzy: this.allFiles.length > 20000 ? 'v1' : 'v2',
+ forward: false,
+ tiebreakers: [byBasenamePrefix, byMatchPosFromEnd, byLengthAsc],
});
}
}
diff --git a/packages/core/src/utils/generateContentResponseUtilities.test.ts b/packages/core/src/utils/generateContentResponseUtilities.test.ts
index 0562f91888..179144964e 100644
--- a/packages/core/src/utils/generateContentResponseUtilities.test.ts
+++ b/packages/core/src/utils/generateContentResponseUtilities.test.ts
@@ -16,14 +16,14 @@ import {
getCitations,
convertToFunctionResponse,
} from './generateContentResponseUtilities.js';
-import type {
- GenerateContentResponse,
- Part,
- SafetyRating,
- CitationMetadata,
- PartListUnion,
+import {
+ FinishReason,
+ type GenerateContentResponse,
+ type Part,
+ type SafetyRating,
+ type CitationMetadata,
+ type PartListUnion,
} from '@google/genai';
-import { FinishReason } from '@google/genai';
import {
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL,
diff --git a/packages/core/src/utils/getFolderStructure.ts b/packages/core/src/utils/getFolderStructure.ts
index 8f871e1283..6e1814cd90 100644
--- a/packages/core/src/utils/getFolderStructure.ts
+++ b/packages/core/src/utils/getFolderStructure.ts
@@ -12,8 +12,10 @@ import type {
FileDiscoveryService,
FilterFilesOptions,
} from '../services/fileDiscoveryService.js';
-import type { FileFilteringOptions } from '../config/constants.js';
-import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js';
+import {
+ DEFAULT_FILE_FILTERING_OPTIONS,
+ type FileFilteringOptions,
+} from '../config/constants.js';
import { debugLogger } from './debugLogger.js';
const MAX_ITEMS = 200;
diff --git a/packages/core/src/utils/googleErrors.test.ts b/packages/core/src/utils/googleErrors.test.ts
index 46a6aa7b7a..5a21e72801 100644
--- a/packages/core/src/utils/googleErrors.test.ts
+++ b/packages/core/src/utils/googleErrors.test.ts
@@ -5,8 +5,7 @@
*/
import { describe, it, expect } from 'vitest';
-import { parseGoogleApiError } from './googleErrors.js';
-import type { QuotaFailure } from './googleErrors.js';
+import { parseGoogleApiError, type QuotaFailure } from './googleErrors.js';
describe('parseGoogleApiError', () => {
it('should return null for non-gaxios errors', () => {
@@ -362,4 +361,88 @@ describe('parseGoogleApiError', () => {
),
).toBe(true);
});
+
+ it('should parse a gaxios error with SSE-corrupted JSON containing stray commas', () => {
+ // This reproduces the exact corruption pattern observed in production where
+ // SSE serialization injects a stray comma on a newline before "metadata".
+ const corruptedJson = JSON.stringify([
+ {
+ error: {
+ code: 429,
+ message:
+ 'You have exhausted your capacity on this model. Your quota will reset after 19h14m47s.',
+ details: [
+ {
+ '@type': 'type.googleapis.com/google.rpc.ErrorInfo',
+ reason: 'QUOTA_EXHAUSTED',
+ domain: 'cloudcode-pa.googleapis.com',
+ metadata: {
+ uiMessage: 'true',
+ model: 'gemini-3-flash-preview',
+ },
+ },
+ {
+ '@type': 'type.googleapis.com/google.rpc.RetryInfo',
+ retryDelay: '68940s',
+ },
+ ],
+ },
+ },
+ ]).replace(
+ '"domain": "cloudcode-pa.googleapis.com",',
+ '"domain": "cloudcode-pa.googleapis.com",\n , ',
+ );
+
+ // Test via message path (fromApiError)
+ const mockError = {
+ message: corruptedJson,
+ code: 429,
+ status: 429,
+ };
+
+ const parsed = parseGoogleApiError(mockError);
+ expect(parsed).not.toBeNull();
+ expect(parsed?.code).toBe(429);
+ expect(parsed?.message).toContain('You have exhausted your capacity');
+ expect(parsed?.details).toHaveLength(2);
+ expect(
+ parsed?.details.some(
+ (d) => d['@type'] === 'type.googleapis.com/google.rpc.ErrorInfo',
+ ),
+ ).toBe(true);
+ });
+
+ it('should parse a gaxios error with SSE-corrupted JSON in response.data', () => {
+ const corruptedJson = JSON.stringify([
+ {
+ error: {
+ code: 429,
+ message: 'Quota exceeded',
+ details: [
+ {
+ '@type': 'type.googleapis.com/google.rpc.ErrorInfo',
+ reason: 'QUOTA_EXHAUSTED',
+ domain: 'cloudcode-pa.googleapis.com',
+ metadata: { model: 'gemini-3-flash-preview' },
+ },
+ ],
+ },
+ },
+ ]).replace(
+ '"domain": "cloudcode-pa.googleapis.com",',
+ '"domain": "cloudcode-pa.googleapis.com",\n, ',
+ );
+
+ const mockError = {
+ response: {
+ status: 429,
+ data: corruptedJson,
+ },
+ };
+
+ const parsed = parseGoogleApiError(mockError);
+ expect(parsed).not.toBeNull();
+ expect(parsed?.code).toBe(429);
+ expect(parsed?.message).toBe('Quota exceeded');
+ });
});
diff --git a/packages/core/src/utils/googleErrors.ts b/packages/core/src/utils/googleErrors.ts
index f7f972f568..4439d55de5 100644
--- a/packages/core/src/utils/googleErrors.ts
+++ b/packages/core/src/utils/googleErrors.ts
@@ -9,6 +9,26 @@
* This file contains types and functions for parsing structured Google API errors.
*/
+/**
+ * Sanitize a JSON string before parsing to handle known SSE stream corruption.
+ * SSE stream parsing can inject stray commas — the observed pattern is a comma
+ * at the end of one line followed by a stray comma on the next line, e.g.:
+ * `"domain": "cloudcode-pa.googleapis.com",\n , "metadata": {`
+ * This collapses duplicate commas (possibly separated by whitespace/newlines)
+ * into a single comma, preserving the whitespace.
+ */
+function sanitizeJsonString(jsonStr: string): string {
+ // Match a comma, optional whitespace/newlines, then another comma.
+ // Replace with just a comma + the captured whitespace.
+ // Loop to handle cases like `,,,` which would otherwise become `,,` on a single pass.
+ let prev: string;
+ do {
+ prev = jsonStr;
+ jsonStr = jsonStr.replace(/,(\s*),/g, ',$1');
+ } while (jsonStr !== prev);
+ return jsonStr;
+}
+
/**
* Based on google/rpc/error_details.proto
*/
@@ -138,7 +158,7 @@ export function parseGoogleApiError(error: unknown): GoogleApiError | null {
// If error is a string, try to parse it.
if (typeof errorObj === 'string') {
try {
- errorObj = JSON.parse(errorObj);
+ errorObj = JSON.parse(sanitizeJsonString(errorObj));
} catch (_) {
// Not a JSON string, can't parse.
return null;
@@ -168,7 +188,9 @@ export function parseGoogleApiError(error: unknown): GoogleApiError | null {
try {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
const parsedMessage = JSON.parse(
- currentError.message.replace(/\u00A0/g, '').replace(/\n/g, ' '),
+ sanitizeJsonString(
+ currentError.message.replace(/\u00A0/g, '').replace(/\n/g, ' '),
+ ),
);
if (parsedMessage.error) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
@@ -209,6 +231,7 @@ export function parseGoogleApiError(error: unknown): GoogleApiError | null {
}
// Basic structural check before casting.
// Since the proto definitions are loose, we primarily rely on @type presence.
+ // eslint-disable-next-line no-restricted-syntax
if (typeof detailObj['@type'] === 'string') {
// We can just cast it; the consumer will have to switch on @type
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
@@ -260,7 +283,7 @@ function fromGaxiosError(errorObj: object): ErrorShape | undefined {
if (typeof data === 'string') {
try {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
- data = JSON.parse(data);
+ data = JSON.parse(sanitizeJsonString(data));
} catch (_) {
// Not a JSON string, can't parse.
}
@@ -310,7 +333,7 @@ function fromApiError(errorObj: object): ErrorShape | undefined {
if (typeof data === 'string') {
try {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
- data = JSON.parse(data);
+ data = JSON.parse(sanitizeJsonString(data));
} catch (_) {
// Not a JSON string, can't parse.
// Try one more fallback: look for the first '{' and last '}'
@@ -320,7 +343,9 @@ function fromApiError(errorObj: object): ErrorShape | undefined {
if (firstBrace !== -1 && lastBrace !== -1 && lastBrace > firstBrace) {
try {
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
- data = JSON.parse(data.substring(firstBrace, lastBrace + 1));
+ data = JSON.parse(
+ sanitizeJsonString(data.substring(firstBrace, lastBrace + 1)),
+ );
} catch (__) {
// Still failed
}
diff --git a/packages/core/src/utils/googleQuotaErrors.test.ts b/packages/core/src/utils/googleQuotaErrors.test.ts
index 67c9ae5d3a..cd09e53511 100644
--- a/packages/core/src/utils/googleQuotaErrors.test.ts
+++ b/packages/core/src/utils/googleQuotaErrors.test.ts
@@ -669,4 +669,53 @@ describe('classifyGoogleError', () => {
expect(result).toBe(originalError);
expect(result).not.toBeInstanceOf(ValidationRequiredError);
});
+
+ it('should return TerminalQuotaError for Cloud Code QUOTA_EXHAUSTED with SSE-corrupted domain', () => {
+ // SSE serialization can inject a trailing comma into the domain string.
+ // This test verifies that the domain sanitization handles this case.
+ const apiError: GoogleApiError = {
+ code: 429,
+ message:
+ 'You have exhausted your capacity on this model. Your quota will reset after 19h14m47s.',
+ details: [
+ {
+ '@type': 'type.googleapis.com/google.rpc.ErrorInfo',
+ reason: 'QUOTA_EXHAUSTED',
+ domain: 'cloudcode-pa.googleapis.com,',
+ metadata: {
+ uiMessage: 'true',
+ model: 'gemini-3-flash-preview',
+ },
+ },
+ {
+ '@type': 'type.googleapis.com/google.rpc.RetryInfo',
+ retryDelay: '68940s',
+ },
+ ],
+ };
+ vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
+ const result = classifyGoogleError(new Error());
+ expect(result).toBeInstanceOf(TerminalQuotaError);
+ });
+
+ it('should return ValidationRequiredError with SSE-corrupted domain', () => {
+ const apiError: GoogleApiError = {
+ code: 403,
+ message: 'Forbidden.',
+ details: [
+ {
+ '@type': 'type.googleapis.com/google.rpc.ErrorInfo',
+ reason: 'VALIDATION_REQUIRED',
+ domain: 'cloudcode-pa.googleapis.com,',
+ metadata: {
+ validationUrl: 'https://example.com/validate',
+ validationDescription: 'Please validate',
+ },
+ },
+ ],
+ };
+ vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
+ const result = classifyGoogleError(new Error());
+ expect(result).toBeInstanceOf(ValidationRequiredError);
+ });
});
diff --git a/packages/core/src/utils/googleQuotaErrors.ts b/packages/core/src/utils/googleQuotaErrors.ts
index e9955493bd..fac291f36e 100644
--- a/packages/core/src/utils/googleQuotaErrors.ts
+++ b/packages/core/src/utils/googleQuotaErrors.ts
@@ -4,14 +4,14 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type {
- ErrorInfo,
- GoogleApiError,
- Help,
- QuotaFailure,
- RetryInfo,
+import {
+ parseGoogleApiError,
+ type ErrorInfo,
+ type GoogleApiError,
+ type Help,
+ type QuotaFailure,
+ type RetryInfo,
} from './googleErrors.js';
-import { parseGoogleApiError } from './googleErrors.js';
import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
/**
@@ -109,6 +109,16 @@ const CLOUDCODE_DOMAINS = [
'autopush-cloudcode-pa.googleapis.com',
];
+/**
+ * Checks if the given domain belongs to a Cloud Code API endpoint.
+ * Sanitizes stray characters that SSE stream parsing can inject into the
+ * domain string before comparing.
+ */
+function isCloudCodeDomain(domain: string): boolean {
+ const sanitized = domain.replace(/[^a-zA-Z0-9.-]/g, '');
+ return CLOUDCODE_DOMAINS.includes(sanitized);
+}
+
/**
* Checks if a 403 error requires user validation and extracts validation details.
*
@@ -129,7 +139,7 @@ function classifyValidationRequiredError(
if (
!errorInfo.domain ||
- !CLOUDCODE_DOMAINS.includes(errorInfo.domain) ||
+ !isCloudCodeDomain(errorInfo.domain) ||
errorInfo.reason !== 'VALIDATION_REQUIRED'
) {
return null;
@@ -313,12 +323,7 @@ export function classifyGoogleError(error: unknown): unknown {
// New Cloud Code API quota handling
if (errorInfo.domain) {
- const validDomains = [
- 'cloudcode-pa.googleapis.com',
- 'staging-cloudcode-pa.googleapis.com',
- 'autopush-cloudcode-pa.googleapis.com',
- ];
- if (validDomains.includes(errorInfo.domain)) {
+ if (isCloudCodeDomain(errorInfo.domain)) {
if (errorInfo.reason === 'RATE_LIMIT_EXCEEDED') {
return new RetryableQuotaError(
`${googleApiError.message}`,
diff --git a/packages/core/src/utils/installationManager.test.ts b/packages/core/src/utils/installationManager.test.ts
index 1cc7f69926..a5251697c2 100644
--- a/packages/core/src/utils/installationManager.test.ts
+++ b/packages/core/src/utils/installationManager.test.ts
@@ -4,8 +4,15 @@
* SPDX-License-Identifier: Apache-2.0
*/
-import type { Mock } from 'vitest';
-import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
+import {
+ vi,
+ describe,
+ it,
+ expect,
+ beforeEach,
+ afterEach,
+ type Mock,
+} from 'vitest';
import { InstallationManager } from './installationManager.js';
import * as fs from 'node:fs';
import * as os from 'node:os';
diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts
index 3df110d678..c2b865dad1 100644
--- a/packages/core/src/utils/memoryDiscovery.test.ts
+++ b/packages/core/src/utils/memoryDiscovery.test.ts
@@ -20,10 +20,9 @@ import {
setGeminiMdFilename,
DEFAULT_CONTEXT_FILENAME,
} from '../tools/memoryTool.js';
-import { flattenMemory } from '../config/memory.js';
+import { flattenMemory, type HierarchicalMemory } from '../config/memory.js';
import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
import { GEMINI_DIR, normalizePath, homedir as pathsHomedir } from './paths.js';
-import type { HierarchicalMemory } from '../config/memory.js';
function flattenResult(result: {
memoryContent: HierarchicalMemory;
@@ -40,7 +39,6 @@ import { Config, type GeminiCLIExtension } from '../config/config.js';
import { Storage } from '../config/storage.js';
import { SimpleExtensionLoader } from './extensionLoader.js';
import { CoreEvent, coreEvents } from './events.js';
-import { debugLogger } from './debugLogger.js';
vi.mock('os', async (importOriginal) => {
const actualOs = await importOriginal();
@@ -130,7 +128,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
false, // untrusted
@@ -167,7 +164,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
false, // untrusted
@@ -185,7 +181,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -209,7 +204,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -242,7 +236,6 @@ default context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -276,7 +269,6 @@ custom context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -314,7 +306,6 @@ cwd context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -349,7 +340,6 @@ Subdir custom memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -384,7 +374,6 @@ Src directory memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -431,7 +420,6 @@ Subdir memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -488,7 +476,6 @@ Subdir memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -513,10 +500,6 @@ My code memory
});
it('should respect the maxDirs parameter during downward scan', async () => {
- const consoleDebugSpy = vi
- .spyOn(debugLogger, 'debug')
- .mockImplementation(() => {});
-
// Create directories in parallel for better performance
const dirPromises = Array.from({ length: 2 }, (_, i) =>
createEmptyDir(path.join(cwd, `deep_dir_${i}`)),
@@ -527,7 +510,6 @@ My code memory
await loadServerHierarchicalMemory(
cwd,
[],
- true,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -540,18 +522,13 @@ My code memory
1, // maxDirs
);
- expect(consoleDebugSpy).toHaveBeenCalledWith(
- expect.stringContaining('[DEBUG] [BfsFileSearch]'),
- expect.stringContaining('Scanning [1/1]:'),
- );
-
- consoleDebugSpy.mockRestore();
+ // Note: bfsFileSearch debug logging is no longer controlled via debugMode parameter
+ // The test verifies maxDirs is respected by checking the result, not debug logs
const result = flattenResult(
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -575,7 +552,6 @@ My code memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([
{
@@ -610,7 +586,6 @@ Extension memory content
await loadServerHierarchicalMemory(
cwd,
[includedDir],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -648,7 +623,6 @@ included directory memory
await loadServerHierarchicalMemory(
cwd,
createdFiles.map((f) => path.dirname(f)),
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -686,7 +660,6 @@ included directory memory
await loadServerHierarchicalMemory(
parentDir,
[childDir, parentDir], // Deliberately include duplicates
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -865,6 +838,173 @@ included directory memory
});
});
+ describe('case-insensitive filesystem deduplication', () => {
+ it('should deduplicate files that point to the same inode (same physical file)', async () => {
+ const geminiFile = await createTestFile(
+ path.join(projectRoot, 'gemini.md'),
+ 'Project root memory',
+ );
+
+ // create hard link to simulate case-insensitive filesystem behavior
+ const geminiFileLink = path.join(projectRoot, 'GEMINI.md');
+ try {
+ await fsPromises.link(geminiFile, geminiFileLink);
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ if (
+ errorMessage.includes('cross-device') ||
+ errorMessage.includes('EXDEV') ||
+ errorMessage.includes('EEXIST')
+ ) {
+ return;
+ }
+ throw error;
+ }
+
+ const stats1 = await fsPromises.lstat(geminiFile);
+ const stats2 = await fsPromises.lstat(geminiFileLink);
+ expect(stats1.ino).toBe(stats2.ino);
+ expect(stats1.dev).toBe(stats2.dev);
+
+ setGeminiMdFilename(['GEMINI.md', 'gemini.md']);
+
+ const result = flattenResult(
+ await loadServerHierarchicalMemory(
+ cwd,
+ [],
+ new FileDiscoveryService(projectRoot),
+ new SimpleExtensionLoader([]),
+ DEFAULT_FOLDER_TRUST,
+ ),
+ );
+
+ expect(result.fileCount).toBe(1);
+ expect(result.filePaths).toHaveLength(1);
+ expect(result.memoryContent).toContain('Project root memory');
+ const contentMatches = result.memoryContent.match(/Project root memory/g);
+ expect(contentMatches).toHaveLength(1);
+
+ try {
+ await fsPromises.unlink(geminiFileLink);
+ } catch {
+ // ignore cleanup errors
+ }
+ });
+
+ it('should handle case where files have different inodes (different files)', async () => {
+ const geminiFileLower = await createTestFile(
+ path.join(projectRoot, 'gemini.md'),
+ 'Lowercase file content',
+ );
+ const geminiFileUpper = await createTestFile(
+ path.join(projectRoot, 'GEMINI.md'),
+ 'Uppercase file content',
+ );
+
+ const stats1 = await fsPromises.lstat(geminiFileLower);
+ const stats2 = await fsPromises.lstat(geminiFileUpper);
+
+ if (stats1.ino !== stats2.ino || stats1.dev !== stats2.dev) {
+ setGeminiMdFilename(['GEMINI.md', 'gemini.md']);
+
+ const result = flattenResult(
+ await loadServerHierarchicalMemory(
+ cwd,
+ [],
+ new FileDiscoveryService(projectRoot),
+ new SimpleExtensionLoader([]),
+ DEFAULT_FOLDER_TRUST,
+ ),
+ );
+
+ expect(result.fileCount).toBe(2);
+ expect(result.filePaths).toHaveLength(2);
+ expect(result.memoryContent).toContain('Lowercase file content');
+ expect(result.memoryContent).toContain('Uppercase file content');
+ }
+ });
+
+ it("should handle files that cannot be stat'd (missing files)", async () => {
+ await createTestFile(
+ path.join(projectRoot, 'gemini.md'),
+ 'Valid file content',
+ );
+
+ setGeminiMdFilename(['gemini.md', 'missing.md']);
+
+ const result = flattenResult(
+ await loadServerHierarchicalMemory(
+ cwd,
+ [],
+ new FileDiscoveryService(projectRoot),
+ new SimpleExtensionLoader([]),
+ DEFAULT_FOLDER_TRUST,
+ ),
+ );
+
+ expect(result.fileCount).toBe(1);
+ expect(result.memoryContent).toContain('Valid file content');
+ });
+
+ it('should deduplicate multiple paths pointing to same file (3+ duplicates)', async () => {
+ const geminiFile = await createTestFile(
+ path.join(projectRoot, 'gemini.md'),
+ 'Project root memory',
+ );
+
+ const link1 = path.join(projectRoot, 'GEMINI.md');
+ const link2 = path.join(projectRoot, 'Gemini.md');
+
+ try {
+ await fsPromises.link(geminiFile, link1);
+ await fsPromises.link(geminiFile, link2);
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ if (
+ errorMessage.includes('cross-device') ||
+ errorMessage.includes('EXDEV') ||
+ errorMessage.includes('EEXIST')
+ ) {
+ return;
+ }
+ throw error;
+ }
+
+ const stats1 = await fsPromises.lstat(geminiFile);
+ const stats2 = await fsPromises.lstat(link1);
+ const stats3 = await fsPromises.lstat(link2);
+ expect(stats1.ino).toBe(stats2.ino);
+ expect(stats1.ino).toBe(stats3.ino);
+
+ setGeminiMdFilename(['gemini.md', 'GEMINI.md', 'Gemini.md']);
+
+ const result = flattenResult(
+ await loadServerHierarchicalMemory(
+ cwd,
+ [],
+ new FileDiscoveryService(projectRoot),
+ new SimpleExtensionLoader([]),
+ DEFAULT_FOLDER_TRUST,
+ ),
+ );
+
+ expect(result.fileCount).toBe(1);
+ expect(result.filePaths).toHaveLength(1);
+ expect(result.memoryContent).toContain('Project root memory');
+ const contentMatches = result.memoryContent.match(/Project root memory/g);
+ expect(contentMatches).toHaveLength(1);
+
+ try {
+ await fsPromises.unlink(link1);
+ await fsPromises.unlink(link2);
+ } catch {
+ // ignore cleanup errors
+ }
+ });
+ });
+
describe('loadJitSubdirectoryMemory', () => {
it('should load JIT memory when target is inside a trusted root', async () => {
const rootDir = await createEmptyDir(path.join(testRootDir, 'jit_root'));
@@ -938,6 +1078,57 @@ included directory memory
expect(result.files[0].content).toBe('Subdir content');
});
+ it('should deduplicate files in JIT memory loading (same inode)', async () => {
+ const rootDir = await createEmptyDir(path.join(testRootDir, 'jit_root'));
+ const subDir = await createEmptyDir(path.join(rootDir, 'subdir'));
+ const targetFile = path.join(subDir, 'target.txt');
+
+ const geminiFile = await createTestFile(
+ path.join(subDir, 'gemini.md'),
+ 'JIT memory content',
+ );
+
+ const geminiFileLink = path.join(subDir, 'GEMINI.md');
+ try {
+ await fsPromises.link(geminiFile, geminiFileLink);
+ } catch (error) {
+ const errorMessage =
+ error instanceof Error ? error.message : String(error);
+ if (
+ errorMessage.includes('cross-device') ||
+ errorMessage.includes('EXDEV') ||
+ errorMessage.includes('EEXIST')
+ ) {
+ return;
+ }
+ throw error;
+ }
+
+ const stats1 = await fsPromises.lstat(geminiFile);
+ const stats2 = await fsPromises.lstat(geminiFileLink);
+ expect(stats1.ino).toBe(stats2.ino);
+
+ setGeminiMdFilename(['gemini.md', 'GEMINI.md']);
+
+ const result = await loadJitSubdirectoryMemory(
+ targetFile,
+ [rootDir],
+ new Set(),
+ );
+
+ expect(result.files).toHaveLength(1);
+ expect(result.files[0].content).toBe('JIT memory content');
+ const contentMatches =
+ result.files[0].content.match(/JIT memory content/g);
+ expect(contentMatches).toHaveLength(1);
+
+ try {
+ await fsPromises.unlink(geminiFileLink);
+ } catch {
+ // ignore cleanup errors
+ }
+ });
+
it('should use the deepest trusted root when multiple nested roots exist', async () => {
const outerRoot = await createEmptyDir(path.join(testRootDir, 'outer'));
const innerRoot = await createEmptyDir(path.join(outerRoot, 'inner'));
@@ -982,7 +1173,6 @@ included directory memory
config.shouldLoadMemoryFromIncludeDirectories()
? config.getWorkspaceContext().getDirectories()
: [],
- config.getDebugMode(),
config.getFileService(),
config.getExtensionLoader(),
config.isTrustedFolder(),
@@ -1027,7 +1217,6 @@ included directory memory
const mockConfig = {
getWorkingDir: vi.fn().mockReturnValue(cwd),
shouldLoadMemoryFromIncludeDirectories: vi.fn().mockReturnValue(false),
- getDebugMode: vi.fn().mockReturnValue(false),
getFileService: vi
.fn()
.mockReturnValue(new FileDiscoveryService(projectRoot)),
diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts
index c35d009e1d..2d7de3327c 100644
--- a/packages/core/src/utils/memoryDiscovery.ts
+++ b/packages/core/src/utils/memoryDiscovery.ts
@@ -11,14 +11,17 @@ import { bfsFileSearch } from './bfsFileSearch.js';
import { getAllGeminiMdFilenames } from '../tools/memoryTool.js';
import type { FileDiscoveryService } from '../services/fileDiscoveryService.js';
import { processImports } from './memoryImportProcessor.js';
-import type { FileFilteringOptions } from '../config/constants.js';
-import { DEFAULT_MEMORY_FILE_FILTERING_OPTIONS } from '../config/constants.js';
+import {
+ DEFAULT_MEMORY_FILE_FILTERING_OPTIONS,
+ type FileFilteringOptions,
+} from '../config/constants.js';
import { GEMINI_DIR, homedir, normalizePath } from './paths.js';
import type { ExtensionLoader } from './extensionLoader.js';
import { debugLogger } from './debugLogger.js';
import type { Config } from '../config/config.js';
import type { HierarchicalMemory } from '../config/memory.js';
import { CoreEvent, coreEvents } from './events.js';
+import { getErrorMessage } from './errors.js';
// Simple console logger, similar to the one previously in CLI's config.ts
// TODO: Integrate with a more robust server-side logger if available/appropriate.
@@ -39,6 +42,110 @@ export interface GeminiFileContent {
content: string | null;
}
+/**
+ * Deduplicates file paths by file identity (device + inode) rather than string path.
+ * This is necessary on case-insensitive filesystems where different case variants
+ * of the same filename resolve to the same physical file but have different path strings.
+ *
+ * @param filePaths Array of file paths to deduplicate
+ * @returns Object containing deduplicated file paths and a map of path to identity key
+ */
+export async function deduplicatePathsByFileIdentity(
+ filePaths: string[],
+): Promise<{
+ paths: string[];
+ identityMap: Map;
+}> {
+ if (filePaths.length === 0) {
+ return {
+ paths: [],
+ identityMap: new Map(),
+ };
+ }
+
+ // first deduplicate by string path to avoid redundant stat calls
+ const uniqueFilePaths = Array.from(new Set(filePaths));
+
+ const fileIdentityMap = new Map