(
- (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/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/policy/config.test.ts b/packages/core/src/policy/config.test.ts
index f1cb8d0788..9cee9997b2 100644
--- a/packages/core/src/policy/config.test.ts
+++ b/packages/core/src/policy/config.test.ts
@@ -206,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
@@ -224,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
@@ -251,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
@@ -260,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();
});
@@ -288,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
@@ -297,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
@@ -306,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
@@ -372,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,
@@ -381,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,
);
@@ -408,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,
@@ -417,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,
);
@@ -436,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
@@ -568,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 72ffa9ebfb..959f09ba80 100644
--- a/packages/core/src/policy/toml-loader.test.ts
+++ b/packages/core/src/policy/toml-loader.test.ts
@@ -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/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts
index 05f5b08a2f..9320893bd6 100644
--- a/packages/core/src/scheduler/policy.test.ts
+++ b/packages/core/src/scheduler/policy.test.ts
@@ -49,6 +49,9 @@ describe('policy.ts', () => {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
+
const toolCall = {
request: { name: 'test-tool', args: {} },
tool: { name: 'test-tool' },
@@ -72,6 +75,9 @@ describe('policy.ts', () => {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
+
const mcpTool = Object.create(DiscoveredMCPTool.prototype);
mcpTool.serverName = 'my-server';
mcpTool._toolAnnotations = { readOnlyHint: true };
@@ -99,6 +105,9 @@ describe('policy.ts', () => {
isInteractive: vi.fn().mockReturnValue(false),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
+
const toolCall = {
request: { name: 'test-tool', args: {} },
tool: { name: 'test-tool' },
@@ -118,6 +127,9 @@ describe('policy.ts', () => {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
+
const toolCall = {
request: { name: 'test-tool', args: {} },
tool: { name: 'test-tool' },
@@ -137,6 +149,9 @@ describe('policy.ts', () => {
isInteractive: vi.fn().mockReturnValue(true),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
+
const toolCall = {
request: { name: 'test-tool', args: {} },
tool: { name: 'test-tool' },
@@ -152,6 +167,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -175,6 +193,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -200,6 +221,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -225,6 +249,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -256,6 +283,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -290,6 +320,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -308,6 +341,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -325,6 +361,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -344,6 +383,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -378,6 +420,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -410,6 +455,9 @@ describe('policy.ts', () => {
const mockConfig = {
setApprovalMode: vi.fn(),
} as unknown as Mocked;
+
+ (mockConfig as unknown as { config: Config }).config =
+ mockConfig as Config;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked;
@@ -447,6 +495,8 @@ describe('policy.ts', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
} as unknown as Config;
+ (mockConfig as unknown as { config: Config }).config = mockConfig;
+
const { errorMessage, errorType } = getPolicyDenialError(mockConfig);
expect(errorMessage).toBe('Tool execution denied by policy.');
@@ -457,6 +507,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
} as unknown as Config;
+
+ (mockConfig as unknown as { config: Config }).config = mockConfig;
const rule = {
decision: PolicyDecision.DENY,
denyMessage: 'Custom Deny',
@@ -517,7 +569,8 @@ describe('Plan Mode Denial Consistency', () => {
mockConfig = {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
- getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
+ toolRegistry: mockToolRegistry,
+ getToolRegistry: () => mockToolRegistry,
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
isInteractive: vi.fn().mockReturnValue(true),
getEnableHooks: vi.fn().mockReturnValue(false),
@@ -525,6 +578,7 @@ describe('Plan Mode Denial Consistency', () => {
setApprovalMode: vi.fn(),
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config = mockConfig as Config;
});
afterEach(() => {
diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts
index ee5438c319..4d40101140 100644
--- a/packages/core/src/scheduler/scheduler.test.ts
+++ b/packages/core/src/scheduler/scheduler.test.ts
@@ -169,13 +169,15 @@ describe('Scheduler (Orchestrator)', () => {
mockConfig = {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
- getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
+ toolRegistry: mockToolRegistry,
isInteractive: vi.fn().mockReturnValue(true),
getEnableHooks: vi.fn().mockReturnValue(true),
setApprovalMode: vi.fn(),
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config = mockConfig as Config;
+
mockMessageBus = {
publish: vi.fn(),
subscribe: vi.fn(),
@@ -1320,6 +1322,8 @@ describe('Scheduler MCP Progress', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config = mockConfig as Config;
+
mockMessageBus = {
publish: vi.fn(),
subscribe: vi.fn(),
diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts
index 38e001ea90..613e23b2d6 100644
--- a/packages/core/src/scheduler/scheduler.ts
+++ b/packages/core/src/scheduler/scheduler.ts
@@ -5,6 +5,7 @@
*/
import type { Config } from '../config/config.js';
+import type { AgentLoopContext } from '../config/agent-loop-context.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { SchedulerStateManager } from './state-manager.js';
import { resolveConfirmation } from './confirmation.js';
@@ -57,7 +58,7 @@ interface SchedulerQueueItem {
export interface SchedulerOptions {
config: Config;
- messageBus: MessageBus;
+ messageBus?: MessageBus;
getPreferredEditor: () => EditorType | undefined;
schedulerId: string;
parentCallId?: string;
@@ -97,6 +98,7 @@ export class Scheduler {
private readonly executor: ToolExecutor;
private readonly modifier: ToolModificationHandler;
private readonly config: Config;
+ private readonly context: AgentLoopContext;
private readonly messageBus: MessageBus;
private readonly getPreferredEditor: () => EditorType | undefined;
private readonly schedulerId: string;
@@ -109,7 +111,8 @@ export class Scheduler {
constructor(options: SchedulerOptions) {
this.config = options.config;
- this.messageBus = options.messageBus;
+ this.context = options.config;
+ this.messageBus = options.messageBus ?? this.context.messageBus;
this.getPreferredEditor = options.getPreferredEditor;
this.schedulerId = options.schedulerId;
this.parentCallId = options.parentCallId;
@@ -119,7 +122,7 @@ export class Scheduler {
this.schedulerId,
(call) => logToolCall(this.config, new ToolCallEvent(call)),
);
- this.executor = new ToolExecutor(this.config);
+ this.executor = new ToolExecutor(this.config, this.context);
this.modifier = new ToolModificationHandler();
this.setupMessageBusListener(this.messageBus);
@@ -294,7 +297,7 @@ export class Scheduler {
const currentApprovalMode = this.config.getApprovalMode();
try {
- const toolRegistry = this.config.getToolRegistry();
+ const toolRegistry = this.context.toolRegistry;
const newCalls: ToolCall[] = requests.map((request) => {
const enrichedRequest: ToolCallRequestInfo = {
...request,
@@ -697,7 +700,7 @@ export class Scheduler {
const originalRequestName =
result.request.originalRequestName || result.request.name;
- const newTool = this.config.getToolRegistry().getTool(tailRequest.name);
+ const newTool = this.context.toolRegistry.getTool(tailRequest.name);
const newRequest: ToolCallRequestInfo = {
callId: originalCallId,
@@ -713,7 +716,7 @@ export class Scheduler {
// Enqueue an errored tool call
const errorCall = this._createToolNotFoundErroredToolCall(
newRequest,
- this.config.getToolRegistry().getAllToolNames(),
+ this.context.toolRegistry.getAllToolNames(),
);
this.state.replaceActiveCallWithTailCall(callId, errorCall);
} else {
diff --git a/packages/core/src/scheduler/scheduler_parallel.test.ts b/packages/core/src/scheduler/scheduler_parallel.test.ts
index 56e6e26243..5342d3ac20 100644
--- a/packages/core/src/scheduler/scheduler_parallel.test.ts
+++ b/packages/core/src/scheduler/scheduler_parallel.test.ts
@@ -211,13 +211,15 @@ describe('Scheduler Parallel Execution', () => {
mockConfig = {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
- getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
+ toolRegistry: mockToolRegistry,
isInteractive: vi.fn().mockReturnValue(true),
getEnableHooks: vi.fn().mockReturnValue(true),
setApprovalMode: vi.fn(),
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
} as unknown as Mocked;
+ (mockConfig as unknown as { config: Config }).config = mockConfig as Config;
+
mockMessageBus = {
publish: vi.fn(),
subscribe: vi.fn(),
diff --git a/packages/core/src/scheduler/scheduler_waiting_callback.test.ts b/packages/core/src/scheduler/scheduler_waiting_callback.test.ts
index e878a80669..03b754fc86 100644
--- a/packages/core/src/scheduler/scheduler_waiting_callback.test.ts
+++ b/packages/core/src/scheduler/scheduler_waiting_callback.test.ts
@@ -36,7 +36,7 @@ describe('Scheduler waiting callback', () => {
mockTool = new MockTool({ name: 'test_tool' });
toolRegistry = new ToolRegistry(mockConfig, messageBus);
- vi.spyOn(mockConfig, 'getToolRegistry').mockReturnValue(toolRegistry);
+ vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(toolRegistry);
toolRegistry.registerTool(mockTool);
vi.mocked(checkPolicy).mockResolvedValue({
diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts
index 005f3004d6..428b7f87a8 100644
--- a/packages/core/src/scheduler/state-manager.ts
+++ b/packages/core/src/scheduler/state-manager.ts
@@ -477,13 +477,24 @@ 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: reason,
+ response: finalResponse,
durationMs: startTime ? Date.now() - startTime : undefined,
outcome: call.outcome,
schedulerId: call.schedulerId,
@@ -508,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 e744738341..a193c8ae69 100644
--- a/packages/core/src/scheduler/tool-executor.test.ts
+++ b/packages/core/src/scheduler/tool-executor.test.ts
@@ -64,7 +64,7 @@ describe('ToolExecutor', () => {
beforeEach(() => {
// Use the standard fake config factory
config = makeFakeConfig();
- executor = new ToolExecutor(config);
+ executor = new ToolExecutor(config, config);
// Reset mocks
vi.resetAllMocks();
diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts
index 8269f1fc41..e5491630d2 100644
--- a/packages/core/src/scheduler/tool-executor.ts
+++ b/packages/core/src/scheduler/tool-executor.ts
@@ -13,6 +13,7 @@ import {
type ToolCallResponseInfo,
type ToolResult,
type Config,
+ type AgentLoopContext,
type ToolLiveOutput,
} from '../index.js';
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
@@ -49,7 +50,10 @@ export interface ToolExecutionContext {
}
export class ToolExecutor {
- constructor(private readonly config: Config) {}
+ constructor(
+ private readonly config: Config,
+ private readonly context: AgentLoopContext,
+ ) {}
async execute(context: ToolExecutionContext): Promise {
const { call, signal, outputUpdateHandler, onUpdateToolCall } = context;
@@ -202,7 +206,7 @@ export class ToolExecutor {
toolName,
callId,
this.config.storage.getProjectTempDir(),
- this.config.getSessionId(),
+ this.context.promptId,
);
outputFile = savedPath;
const truncatedContent = formatTruncatedToolOutput(
@@ -241,7 +245,7 @@ export class ToolExecutor {
toolName,
callId,
this.config.storage.getProjectTempDir(),
- this.config.getSessionId(),
+ this.context.promptId,
);
outputFile = savedPath;
const truncatedText = formatTruncatedToolOutput(
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 2911119a25..7ae9549a25 100644
--- a/packages/core/src/services/chatCompressionService.test.ts
+++ b/packages/core/src/services/chatCompressionService.test.ts
@@ -10,7 +10,7 @@ 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';
@@ -189,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(
@@ -355,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 5aaa0a2538..2b8e8f1977 100644
--- a/packages/core/src/services/chatRecordingService.test.ts
+++ b/packages/core/src/services/chatRecordingService.test.ts
@@ -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/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.ts b/packages/core/src/services/loopDetectionService.ts
index e87de721c6..9bc8b406f8 100644
--- a/packages/core/src/services/loopDetectionService.ts
+++ b/packages/core/src/services/loopDetectionService.ts
@@ -579,10 +579,12 @@ export class LoopDetectionService {
}
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']
: '';
@@ -628,11 +630,13 @@ export class LoopDetectionService {
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;
@@ -681,6 +685,7 @@ export class LoopDetectionService {
if (
result &&
+ // eslint-disable-next-line no-restricted-syntax
typeof result['unproductive_state_confidence'] === 'number'
) {
return result;
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/loggers.ts b/packages/core/src/telemetry/loggers.ts
index 4c3ed55321..393519d3ec 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -791,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/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 cb38502c91..6dae06d381 100644
--- a/packages/core/src/telemetry/semantic.ts
+++ b/packages/core/src/telemetry/semantic.ts
@@ -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/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 508ae7775b..1f0a8ee98f 100644
--- a/packages/core/src/tools/grep.test.ts
+++ b/packages/core/src/tools/grep.test.ts
@@ -562,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/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts
index 0f7b58c39a..8612a838ca 100644
--- a/packages/core/src/tools/mcp-client.test.ts
+++ b/packages/core/src/tools/mcp-client.test.ts
@@ -2053,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-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts
index fc4a8d299a..1d9e2a2f25 100644
--- a/packages/core/src/tools/mcp-tool.test.ts
+++ b/packages/core/src/tools/mcp-tool.test.ts
@@ -15,7 +15,11 @@ import {
type Mocked,
} from 'vitest';
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
-import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
+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';
@@ -49,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(
- 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx...xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
+ 'mcp_xxxxxxxxxxxxxxxxxxxxxxxxxx...xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
);
});
it('should handle names with only invalid characters', () => {
- expect(generateValidName('!@#$%^&*()')).toBe('__________');
+ expect(generateValidName('!@#$%^&*()')).toBe('mcp___________');
});
it.each([
@@ -80,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';
@@ -116,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);
@@ -943,31 +973,35 @@ describe('DiscoveredMCPTool', () => {
describe('MCP Tool Naming Regression Fixes', () => {
describe('generateValidName', () => {
it('should replace spaces with underscores', () => {
- expect(generateValidName('My Tool')).toBe('My_Tool');
+ expect(generateValidName('My Tool')).toBe('mcp_My_Tool');
});
it('should allow colons', () => {
- expect(generateValidName('namespace:tool')).toBe('namespace:tool');
+ expect(generateValidName('namespace:tool')).toBe('mcp_namespace:tool');
});
it('should ensure name starts with a letter or underscore', () => {
- expect(generateValidName('123-tool')).toBe('_123-tool');
- expect(generateValidName('-tool')).toBe('_-tool');
- expect(generateValidName('.tool')).toBe('_.tool');
+ 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(/^a{30}\.\.\.b{30}$/);
+ 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('_1')).toBe(true);
+ expect(result.startsWith('mcp_1')).toBe(true);
});
});
@@ -983,7 +1017,7 @@ describe('MCP Tool Naming Regression Fixes', () => {
);
const qn = tool.getFullyQualifiedName();
- expect(qn).toBe('My_Server__my-tool');
+ expect(qn).toBe('mcp_My_Server_my-tool');
});
it('should handle long server and tool names in qualified name', () => {
@@ -1014,7 +1048,7 @@ describe('MCP Tool Naming Regression Fixes', () => {
);
const qn = tool.getFullyQualifiedName();
- expect(qn).toBe('_123-server__tool');
+ 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 2c52c72573..f67d1f9bea 100644
--- a/packages/core/src/tools/mcp-tool.ts
+++ b/packages/core/src/tools/mcp-tool.ts
@@ -5,6 +5,7 @@
*/
import { safeJsonStringify } from '../utils/safeJsonStringify.js';
+import { debugLogger } from '../utils/debugLogger.js';
import {
BaseDeclarativeTool,
BaseToolInvocation,
@@ -23,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;
@@ -274,7 +349,10 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
private readonly _toolAnnotations?: Record,
) {
super(
- generateValidName(nameOverride ?? serverToolName),
+ nameOverride ??
+ generateValidName(
+ `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
+ ),
`${serverToolName} (${serverName} MCP Server)`,
description,
Kind.Other,
@@ -302,7 +380,9 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
}
getFullyQualifiedPrefix(): string {
- return `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`;
+ return generateValidName(
+ `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`,
+ );
}
getFullyQualifiedName(): string {
@@ -493,8 +573,12 @@ 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}`;
+
+ // 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)) {
@@ -505,6 +589,9 @@ export function generateValidName(name: string) {
// 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);
}
diff --git a/packages/core/src/tools/ripGrep.test.ts b/packages/core/src/tools/ripGrep.test.ts
index 265bb8e53c..a1b155fb7a 100644
--- a/packages/core/src/tools/ripGrep.test.ts
+++ b/packages/core/src/tools/ripGrep.test.ts
@@ -2028,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 000b4f0071..18a1b0c133 100644
--- a/packages/core/src/tools/ripGrep.ts
+++ b/packages/core/src/tools/ripGrep.ts
@@ -54,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/tool-names.ts b/packages/core/src/tools/tool-names.ts
index c539532fd1..fcdcbd6df6 100644
--- a/packages/core/src/tools/tool-names.ts
+++ b/packages/core/src/tools/tool-names.ts
@@ -154,19 +154,18 @@ export const LS_TOOL_NAME_LEGACY = 'list_directory'; // Just to be safe if anyth
export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]);
-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';
-
// Tool Display Names
export const WRITE_FILE_DISPLAY_NAME = 'WriteFile';
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.
diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts
index eab05294d0..ea560865e6 100644
--- a/packages/core/src/tools/tool-registry.test.ts
+++ b/packages/core/src/tools/tool-registry.test.ts
@@ -20,7 +20,7 @@ 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 { DiscoveredMCPTool } from './mcp-tool.js';
import {
mcpToTool,
type FunctionDeclaration,
@@ -317,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',
@@ -405,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);
});
@@ -419,7 +419,7 @@ describe('ToolRegistry', () => {
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
const toolNames = toolRegistry.getAllToolNames();
- expect(toolNames).toEqual([`${serverName}__${toolName}`]);
+ expect(toolNames).toEqual([`mcp_${serverName}_${toolName}`]);
});
});
@@ -449,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) {
@@ -491,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',
]);
});
});
@@ -608,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', () => {
@@ -636,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', () => {
@@ -658,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 () => {
@@ -695,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', () => {
@@ -709,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}`);
});
});
@@ -762,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', () => {
@@ -779,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 bdd8c7d403..69695877c2 100644
--- a/packages/core/src/tools/tool-registry.ts
+++ b/packages/core/src/tools/tool-registry.ts
@@ -445,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/trackerTools.ts b/packages/core/src/tools/trackerTools.ts
index 2b9b301c53..03ee3c3a97 100644
--- a/packages/core/src/tools/trackerTools.ts
+++ b/packages/core/src/tools/trackerTools.ts
@@ -566,7 +566,7 @@ class TrackerVisualizeInvocation extends BaseToolInvocation<
return {
llmContent: output,
- returnDisplay: 'Graph rendered.',
+ returnDisplay: output,
};
}
}
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.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/googleErrors.test.ts b/packages/core/src/utils/googleErrors.test.ts
index 6e11d01f31..5a21e72801 100644
--- a/packages/core/src/utils/googleErrors.test.ts
+++ b/packages/core/src/utils/googleErrors.test.ts
@@ -361,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 d0c251e839..fac291f36e 100644
--- a/packages/core/src/utils/googleQuotaErrors.ts
+++ b/packages/core/src/utils/googleQuotaErrors.ts
@@ -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/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts
index a23b7660ff..c2b865dad1 100644
--- a/packages/core/src/utils/memoryDiscovery.test.ts
+++ b/packages/core/src/utils/memoryDiscovery.test.ts
@@ -39,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();
@@ -129,7 +128,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
false, // untrusted
@@ -166,7 +164,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
false, // untrusted
@@ -184,7 +181,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -208,7 +204,6 @@ describe('memoryDiscovery', () => {
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -241,7 +236,6 @@ default context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -275,7 +269,6 @@ custom context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -313,7 +306,6 @@ cwd context content
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -348,7 +340,6 @@ Subdir custom memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -383,7 +374,6 @@ Src directory memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -430,7 +420,6 @@ Subdir memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -487,7 +476,6 @@ Subdir memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -512,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}`)),
@@ -526,7 +510,6 @@ My code memory
await loadServerHierarchicalMemory(
cwd,
[],
- true,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -539,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,
@@ -574,7 +552,6 @@ My code memory
await loadServerHierarchicalMemory(
cwd,
[],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([
{
@@ -609,7 +586,6 @@ Extension memory content
await loadServerHierarchicalMemory(
cwd,
[includedDir],
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -647,7 +623,6 @@ included directory memory
await loadServerHierarchicalMemory(
cwd,
createdFiles.map((f) => path.dirname(f)),
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -685,7 +660,6 @@ included directory memory
await loadServerHierarchicalMemory(
parentDir,
[childDir, parentDir], // Deliberately include duplicates
- false,
new FileDiscoveryService(projectRoot),
new SimpleExtensionLoader([]),
DEFAULT_FOLDER_TRUST,
@@ -864,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'));
@@ -937,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'));
@@ -981,7 +1173,6 @@ included directory memory
config.shouldLoadMemoryFromIncludeDirectories()
? config.getWorkspaceContext().getDirectories()
: [],
- config.getDebugMode(),
config.getFileService(),
config.getExtensionLoader(),
config.isTrustedFolder(),
@@ -1026,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 677c571bec..2d7de3327c 100644
--- a/packages/core/src/utils/memoryDiscovery.ts
+++ b/packages/core/src/utils/memoryDiscovery.ts
@@ -21,6 +21,7 @@ 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.
@@ -41,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();
+ const deduplicatedPaths: string[] = [];
+
+ const CONCURRENT_LIMIT = 20;
+ const results: Array<{
+ path: string;
+ dev: bigint | number | null;
+ ino: bigint | number | null;
+ }> = [];
+
+ for (let i = 0; i < uniqueFilePaths.length; i += CONCURRENT_LIMIT) {
+ const batch = uniqueFilePaths.slice(i, i + CONCURRENT_LIMIT);
+ const batchPromises = batch.map(async (filePath) => {
+ try {
+ // use stat() instead of lstat() to follow symlinks and get target file identity
+ const stats = await fs.stat(filePath);
+ return {
+ path: filePath,
+ dev: stats.dev,
+ ino: stats.ino,
+ };
+ } catch (error: unknown) {
+ const message = error instanceof Error ? error.message : String(error);
+ logger.debug(
+ `could not stat file for deduplication: ${filePath}. error: ${message}`,
+ );
+ return {
+ path: filePath,
+ dev: null,
+ ino: null,
+ };
+ }
+ });
+
+ const batchResults = await Promise.allSettled(batchPromises);
+ for (const result of batchResults) {
+ if (result.status === 'fulfilled') {
+ results.push(result.value);
+ } else {
+ const message = getErrorMessage(result.reason);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] unexpected error during deduplication stat:',
+ message,
+ );
+ }
+ }
+ }
+
+ const pathToIdentityMap = new Map();
+ for (const { path, dev, ino } of results) {
+ if (dev !== null && ino !== null) {
+ const identityKey = `${dev.toString()}:${ino.toString()}`;
+ pathToIdentityMap.set(path, identityKey);
+ if (!fileIdentityMap.has(identityKey)) {
+ fileIdentityMap.set(identityKey, path);
+ deduplicatedPaths.push(path);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] deduplication: keeping',
+ path,
+ `(dev: ${dev}, ino: ${ino})`,
+ );
+ } else {
+ const existingPath = fileIdentityMap.get(identityKey);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] deduplication: skipping',
+ path,
+ `(same file as ${existingPath})`,
+ );
+ }
+ } else {
+ deduplicatedPaths.push(path);
+ }
+ }
+
+ return {
+ paths: deduplicatedPaths,
+ identityMap: pathToIdentityMap,
+ };
+}
+
async function findProjectRoot(startDir: string): Promise {
let currentDir = normalizePath(startDir);
while (true) {
@@ -91,7 +196,6 @@ async function getGeminiMdFilePathsInternal(
currentWorkingDirectory: string,
includeDirectoriesToReadGemini: readonly string[],
userHomePath: string,
- debugMode: boolean,
fileService: FileDiscoveryService,
folderTrust: boolean,
fileFilteringOptions: FileFilteringOptions,
@@ -114,7 +218,6 @@ async function getGeminiMdFilePathsInternal(
getGeminiMdFilePathsInternalForEachDir(
dir,
userHomePath,
- debugMode,
fileService,
folderTrust,
fileFilteringOptions,
@@ -146,7 +249,6 @@ async function getGeminiMdFilePathsInternal(
async function getGeminiMdFilePathsInternalForEachDir(
dir: string,
userHomePath: string,
- debugMode: boolean,
fileService: FileDiscoveryService,
folderTrust: boolean,
fileFilteringOptions: FileFilteringOptions,
@@ -167,10 +269,11 @@ async function getGeminiMdFilePathsInternalForEachDir(
try {
await fs.access(globalMemoryPath, fsSync.constants.R_OK);
globalPaths.add(globalMemoryPath);
- if (debugMode)
- logger.debug(
- `Found readable global ${geminiMdFilename}: ${globalMemoryPath}`,
- );
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Found readable global',
+ geminiMdFilename + ':',
+ globalMemoryPath,
+ );
} catch {
// It's okay if it's not found.
}
@@ -179,14 +282,18 @@ async function getGeminiMdFilePathsInternalForEachDir(
// if a valid currentWorkingDirectory is provided.
if (dir && folderTrust) {
const resolvedCwd = normalizePath(dir);
- if (debugMode)
- logger.debug(
- `Searching for ${geminiMdFilename} starting from CWD: ${resolvedCwd}`,
- );
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Searching for',
+ geminiMdFilename,
+ 'starting from CWD:',
+ resolvedCwd,
+ );
const projectRoot = await findProjectRoot(resolvedCwd);
- if (debugMode)
- logger.debug(`Determined project root: ${projectRoot ?? 'None'}`);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Determined project root:',
+ projectRoot ?? 'None',
+ );
const upwardPaths: string[] = [];
let currentDir = resolvedCwd;
@@ -230,7 +337,6 @@ async function getGeminiMdFilePathsInternalForEachDir(
const downwardPaths = await bfsFileSearch(resolvedCwd, {
fileName: geminiMdFilename,
maxDirs,
- debug: debugMode,
fileService,
fileFilteringOptions: mergedOptions,
});
@@ -249,7 +355,6 @@ async function getGeminiMdFilePathsInternalForEachDir(
export async function readGeminiMdFiles(
filePaths: string[],
- debugMode: boolean,
importFormat: 'flat' | 'tree' = 'tree',
): Promise {
// Process files in parallel with concurrency limit to prevent EMFILE errors
@@ -267,15 +372,16 @@ export async function readGeminiMdFiles(
const processedResult = await processImports(
content,
path.dirname(filePath),
- debugMode,
+ false,
undefined,
undefined,
importFormat,
);
- if (debugMode)
- logger.debug(
- `Successfully read and processed imports: ${filePath} (Length: ${processedResult.content.length})`,
- );
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Successfully read and processed imports:',
+ filePath,
+ `(Length: ${processedResult.content.length})`,
+ );
return { filePath, content: processedResult.content };
} catch (error: unknown) {
@@ -288,7 +394,10 @@ export async function readGeminiMdFiles(
`Warning: Could not read ${getAllGeminiMdFilenames()} file at ${filePath}. Error: ${message}`,
);
}
- if (debugMode) logger.debug(`Failed to read: ${filePath}`);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Failed to read:',
+ filePath,
+ );
return { filePath, content: null }; // Still include it with null content
}
},
@@ -337,11 +446,10 @@ export function concatenateInstructions(
export interface MemoryLoadResult {
files: Array<{ path: string; content: string }>;
+ fileIdentities?: string[];
}
-export async function getGlobalMemoryPaths(
- debugMode: boolean = false,
-): Promise {
+export async function getGlobalMemoryPaths(): Promise {
const userHome = homedir();
const geminiMdFilenames = getAllGeminiMdFilenames();
@@ -349,9 +457,10 @@ export async function getGlobalMemoryPaths(
const globalPath = normalizePath(path.join(userHome, GEMINI_DIR, filename));
try {
await fs.access(globalPath, fsSync.constants.R_OK);
- if (debugMode) {
- logger.debug(`Found global memory file: ${globalPath}`);
- }
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Found global memory file:',
+ globalPath,
+ );
return globalPath;
} catch {
return null;
@@ -377,19 +486,18 @@ export function getExtensionMemoryPaths(
export async function getEnvironmentMemoryPaths(
trustedRoots: string[],
- debugMode: boolean = false,
): Promise {
const allPaths = new Set();
// Trusted Roots Upward Traversal (Parallelized)
const traversalPromises = trustedRoots.map(async (root) => {
const resolvedRoot = normalizePath(root);
- if (debugMode) {
- logger.debug(
- `Loading environment memory for trusted root: ${resolvedRoot} (Stopping exactly here)`,
- );
- }
- return findUpwardGeminiFiles(resolvedRoot, resolvedRoot, debugMode);
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Loading environment memory for trusted root:',
+ resolvedRoot,
+ '(Stopping exactly here)',
+ );
+ return findUpwardGeminiFiles(resolvedRoot, resolvedRoot);
});
const pathArrays = await Promise.all(traversalPromises);
@@ -427,7 +535,6 @@ export function categorizeAndConcatenate(
async function findUpwardGeminiFiles(
startDir: string,
stopDir: string,
- debugMode: boolean,
): Promise {
const upwardPaths: string[] = [];
let currentDir = normalizePath(startDir);
@@ -435,11 +542,12 @@ async function findUpwardGeminiFiles(
const geminiMdFilenames = getAllGeminiMdFilenames();
const globalGeminiDir = normalizePath(path.join(homedir(), GEMINI_DIR));
- if (debugMode) {
- logger.debug(
- `Starting upward search from ${currentDir} stopping at ${resolvedStopDir}`,
- );
- }
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Starting upward search from',
+ currentDir,
+ 'stopping at',
+ resolvedStopDir,
+ );
while (true) {
if (currentDir === globalGeminiDir) {
@@ -485,7 +593,6 @@ export interface LoadServerHierarchicalMemoryResponse {
export async function loadServerHierarchicalMemory(
currentWorkingDirectory: string,
includeDirectoriesToReadGemini: readonly string[],
- debugMode: boolean,
fileService: FileDiscoveryService,
extensionLoader: ExtensionLoader,
folderTrust: boolean,
@@ -504,10 +611,11 @@ export async function loadServerHierarchicalMemory(
// function to signal that it should skip the workspace search.
currentWorkingDirectory = isHomeDirectory ? '' : currentWorkingDirectory;
- if (debugMode)
- logger.debug(
- `Loading server hierarchical memory for CWD: ${currentWorkingDirectory} (importFormat: ${importFormat})`,
- );
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Loading server hierarchical memory for CWD:',
+ currentWorkingDirectory,
+ `(importFormat: ${importFormat})`,
+ );
// For the server, homedir() refers to the server process's home.
// This is consistent with how MemoryTool already finds the global path.
@@ -519,7 +627,6 @@ export async function loadServerHierarchicalMemory(
currentWorkingDirectory,
includeDirectoriesToReadGemini,
userHomePath,
- debugMode,
fileService,
folderTrust,
fileFilteringOptions || DEFAULT_MEMORY_FILE_FILTERING_OPTIONS,
@@ -528,7 +635,7 @@ export async function loadServerHierarchicalMemory(
Promise.resolve(getExtensionMemoryPaths(extensionLoader)),
]);
- const allFilePaths = Array.from(
+ const allFilePathsStringDeduped = Array.from(
new Set([
...discoveryResult.global,
...discoveryResult.project,
@@ -536,9 +643,26 @@ export async function loadServerHierarchicalMemory(
]),
);
+ if (allFilePathsStringDeduped.length === 0) {
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] No GEMINI.md files found in hierarchy of the workspace.',
+ );
+ return {
+ memoryContent: { global: '', extension: '', project: '' },
+ fileCount: 0,
+ filePaths: [],
+ };
+ }
+
+ // deduplicate by file identity to handle case-insensitive filesystems
+ const { paths: allFilePaths } = await deduplicatePathsByFileIdentity(
+ allFilePathsStringDeduped,
+ );
+
if (allFilePaths.length === 0) {
- if (debugMode)
- logger.debug('No GEMINI.md files found in hierarchy of the workspace.');
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] No unique GEMINI.md files found after deduplication by file identity.',
+ );
return {
memoryContent: { global: '', extension: '', project: '' },
fileCount: 0,
@@ -547,11 +671,7 @@ export async function loadServerHierarchicalMemory(
}
// 2. GATHER: Read all files in parallel
- const allContents = await readGeminiMdFiles(
- allFilePaths,
- debugMode,
- importFormat,
- );
+ const allContents = await readGeminiMdFiles(allFilePaths, importFormat);
const contentsMap = new Map(allContents.map((c) => [c.filePath, c]));
// 3. CATEGORIZE: Back into Global, Project, Extension
@@ -584,7 +704,6 @@ export async function refreshServerHierarchicalMemory(config: Config) {
config.shouldLoadMemoryFromIncludeDirectories()
? config.getWorkspaceContext().getDirectories()
: [],
- config.getDebugMode(),
config.getFileService(),
config.getExtensionLoader(),
config.isTrustedFolder(),
@@ -611,7 +730,7 @@ export async function loadJitSubdirectoryMemory(
targetPath: string,
trustedRoots: string[],
alreadyLoadedPaths: Set,
- debugMode: boolean = false,
+ alreadyLoadedIdentities?: Set,
): Promise {
const resolvedTarget = normalizePath(targetPath);
let bestRoot: string | null = null;
@@ -634,39 +753,86 @@ export async function loadJitSubdirectoryMemory(
}
if (!bestRoot) {
- if (debugMode) {
- logger.debug(
- `JIT memory skipped: ${resolvedTarget} is not in any trusted root.`,
- );
- }
- return { files: [] };
- }
-
- if (debugMode) {
- logger.debug(
- `Loading JIT memory for ${resolvedTarget} (Trusted root: ${bestRoot})`,
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] JIT memory skipped:',
+ resolvedTarget,
+ 'is not in any trusted root.',
);
+ return { files: [], fileIdentities: [] };
}
- // Traverse from target up to the trusted root
- const potentialPaths = await findUpwardGeminiFiles(
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Loading JIT memory for',
resolvedTarget,
- bestRoot,
- debugMode,
+ `(Trusted root: ${bestRoot})`,
);
- // Filter out already loaded paths
- const newPaths = potentialPaths.filter((p) => !alreadyLoadedPaths.has(p));
+ // Traverse from target up to the trusted root
+ const potentialPaths = await findUpwardGeminiFiles(resolvedTarget, bestRoot);
+
+ if (potentialPaths.length === 0) {
+ return { files: [], fileIdentities: [] };
+ }
+
+ // deduplicate by file identity to handle case-insensitive filesystems
+ // this deduplicates within the current batch
+ const { paths: deduplicatedNewPaths, identityMap: newPathsIdentityMap } =
+ await deduplicatePathsByFileIdentity(potentialPaths);
+
+ // Use cached file identities if provided, otherwise build from paths
+ // This avoids redundant fs.stat() calls on already loaded files
+ const cachedIdentities = alreadyLoadedIdentities ?? new Set();
+ if (!alreadyLoadedIdentities && alreadyLoadedPaths.size > 0) {
+ const CONCURRENT_LIMIT = 20;
+ const alreadyLoadedArray = Array.from(alreadyLoadedPaths);
+
+ for (let i = 0; i < alreadyLoadedArray.length; i += CONCURRENT_LIMIT) {
+ const batch = alreadyLoadedArray.slice(i, i + CONCURRENT_LIMIT);
+ const batchPromises = batch.map(async (filePath) => {
+ try {
+ const stats = await fs.stat(filePath);
+ const identityKey = `${stats.dev.toString()}:${stats.ino.toString()}`;
+ cachedIdentities.add(identityKey);
+ } catch {
+ // ignore errors - if we can't stat it, we can't deduplicate by identity
+ }
+ });
+ // Await each batch to properly limit concurrency and prevent EMFILE errors
+ await Promise.allSettled(batchPromises);
+ }
+ }
+
+ // filter out paths that match already loaded files by identity
+ // reuse the identities from deduplicatePathsByFileIdentity to avoid redundant stat calls
+ const newPaths: string[] = [];
+ const newFileIdentities: string[] = [];
+ for (const filePath of deduplicatedNewPaths) {
+ const identityKey = newPathsIdentityMap.get(filePath);
+ if (identityKey && cachedIdentities.has(identityKey)) {
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] jit memory: skipping',
+ filePath,
+ '(already loaded with different case)',
+ );
+ continue;
+ }
+ // if we don't have an identity (stat failed), include it to be safe
+ newPaths.push(filePath);
+ if (identityKey) {
+ newFileIdentities.push(identityKey);
+ }
+ }
if (newPaths.length === 0) {
- return { files: [] };
+ return { files: [], fileIdentities: [] };
}
- if (debugMode) {
- logger.debug(`Found new JIT memory files: ${JSON.stringify(newPaths)}`);
- }
+ debugLogger.debug(
+ '[DEBUG] [MemoryDiscovery] Found new JIT memory files:',
+ JSON.stringify(newPaths),
+ );
- const contents = await readGeminiMdFiles(newPaths, debugMode, 'tree');
+ const contents = await readGeminiMdFiles(newPaths, 'tree');
return {
files: contents
@@ -676,5 +842,6 @@ export async function loadJitSubdirectoryMemory(
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
content: item.content as string,
})),
+ fileIdentities: newFileIdentities,
};
}
diff --git a/packages/core/src/utils/nextSpeakerChecker.ts b/packages/core/src/utils/nextSpeakerChecker.ts
index a5ce286feb..207306a2e8 100644
--- a/packages/core/src/utils/nextSpeakerChecker.ts
+++ b/packages/core/src/utils/nextSpeakerChecker.ts
@@ -87,7 +87,6 @@ export async function checkNextSpeaker(
lastComprehensiveMessage.parts &&
lastComprehensiveMessage.parts.length === 0
) {
- lastComprehensiveMessage.parts.push({ text: '' });
return {
reasoning:
'The last message was a filler model message with no content (nothing for user to act on), model should speak next.',
diff --git a/packages/core/src/utils/oauth-flow.test.ts b/packages/core/src/utils/oauth-flow.test.ts
new file mode 100644
index 0000000000..dee919c249
--- /dev/null
+++ b/packages/core/src/utils/oauth-flow.test.ts
@@ -0,0 +1,635 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
+import type {
+ OAuthFlowConfig,
+ OAuthRefreshConfig,
+ PKCEParams,
+} from './oauth-flow.js';
+import {
+ generatePKCEParams,
+ getPortFromUrl,
+ buildAuthorizationUrl,
+ startCallbackServer,
+ exchangeCodeForToken,
+ refreshAccessToken,
+ REDIRECT_PATH,
+} from './oauth-flow.js';
+
+// Save real fetch for startCallbackServer tests (which hit a real local server)
+const realFetch = global.fetch;
+
+// Mock fetch globally for token exchange/refresh tests
+const mockFetch = vi.fn();
+global.fetch = mockFetch;
+
+/**
+ * Helper to create a mock Response object.
+ */
+function createMockResponse(
+ body: string,
+ options: { status?: number; contentType?: string } = {},
+): Response {
+ const { status = 200, contentType = 'application/json' } = options;
+ return {
+ ok: status >= 200 && status < 300,
+ status,
+ text: () => Promise.resolve(body),
+ headers: new Headers({ 'content-type': contentType }),
+ } as Response;
+}
+
+const baseConfig: OAuthFlowConfig = {
+ clientId: 'test-client-id',
+ authorizationUrl: 'https://auth.example.com/authorize',
+ tokenUrl: 'https://auth.example.com/token',
+};
+
+const basePkceParams: PKCEParams = {
+ codeVerifier: 'test-verifier',
+ codeChallenge: 'test-challenge',
+ state: 'test-state',
+};
+
+describe('oauth-flow', () => {
+ beforeEach(() => {
+ vi.stubEnv('OAUTH_CALLBACK_PORT', '');
+ mockFetch.mockReset();
+ });
+
+ afterEach(() => {
+ vi.unstubAllEnvs();
+ vi.restoreAllMocks();
+ });
+
+ describe('generatePKCEParams', () => {
+ it('should return codeVerifier, codeChallenge, and state', () => {
+ const params = generatePKCEParams();
+ expect(params).toHaveProperty('codeVerifier');
+ expect(params).toHaveProperty('codeChallenge');
+ expect(params).toHaveProperty('state');
+ });
+
+ it('should generate a code verifier of at least 43 characters', () => {
+ const params = generatePKCEParams();
+ expect(params.codeVerifier.length).toBeGreaterThanOrEqual(43);
+ });
+
+ it('should generate unique values on each call', () => {
+ const params1 = generatePKCEParams();
+ const params2 = generatePKCEParams();
+ expect(params1.codeVerifier).not.toBe(params2.codeVerifier);
+ expect(params1.state).not.toBe(params2.state);
+ });
+
+ it('should generate base64url-encoded values', () => {
+ const params = generatePKCEParams();
+ const base64urlRegex = /^[A-Za-z0-9_-]+$/;
+ expect(params.codeVerifier).toMatch(base64urlRegex);
+ expect(params.codeChallenge).toMatch(base64urlRegex);
+ expect(params.state).toMatch(base64urlRegex);
+ });
+ });
+
+ describe('getPortFromUrl', () => {
+ it('should return undefined for undefined input', () => {
+ expect(getPortFromUrl(undefined)).toBeUndefined();
+ });
+
+ it('should return undefined for empty string', () => {
+ expect(getPortFromUrl('')).toBeUndefined();
+ });
+
+ it('should return undefined for invalid URL', () => {
+ expect(getPortFromUrl('not-a-url')).toBeUndefined();
+ });
+
+ it('should return the port number from a URL with an explicit port', () => {
+ expect(getPortFromUrl('http://localhost:8080/callback')).toBe(8080);
+ });
+
+ it('should return undefined for a URL without an explicit port', () => {
+ expect(getPortFromUrl('https://example.com/callback')).toBeUndefined();
+ });
+
+ it('should return port for edge case port 1', () => {
+ expect(getPortFromUrl('http://localhost:1')).toBe(1);
+ });
+
+ it('should return port for edge case port 65535', () => {
+ expect(getPortFromUrl('http://localhost:65535')).toBe(65535);
+ });
+ });
+
+ describe('buildAuthorizationUrl', () => {
+ it('should build a valid authorization URL with required parameters', () => {
+ const url = buildAuthorizationUrl(baseConfig, basePkceParams, 3000);
+ const parsed = new URL(url);
+
+ expect(parsed.origin).toBe('https://auth.example.com');
+ expect(parsed.pathname).toBe('/authorize');
+ expect(parsed.searchParams.get('client_id')).toBe('test-client-id');
+ expect(parsed.searchParams.get('response_type')).toBe('code');
+ expect(parsed.searchParams.get('state')).toBe('test-state');
+ expect(parsed.searchParams.get('code_challenge')).toBe('test-challenge');
+ expect(parsed.searchParams.get('code_challenge_method')).toBe('S256');
+ });
+
+ it('should use the default redirect URI based on port', () => {
+ const url = buildAuthorizationUrl(baseConfig, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.get('redirect_uri')).toBe(
+ `http://localhost:3000${REDIRECT_PATH}`,
+ );
+ });
+
+ it('should use a custom redirectUri from config when provided', () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ redirectUri: 'https://custom.example.com/callback',
+ };
+ const url = buildAuthorizationUrl(config, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.get('redirect_uri')).toBe(
+ 'https://custom.example.com/callback',
+ );
+ });
+
+ it('should include scopes when provided', () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ scopes: ['read', 'write'],
+ };
+ const url = buildAuthorizationUrl(config, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.get('scope')).toBe('read write');
+ });
+
+ it('should not include scope param when scopes array is empty', () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ scopes: [],
+ };
+ const url = buildAuthorizationUrl(config, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.has('scope')).toBe(false);
+ });
+
+ it('should include audiences when provided', () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ audiences: ['https://api.example.com'],
+ };
+ const url = buildAuthorizationUrl(config, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.get('audience')).toBe(
+ 'https://api.example.com',
+ );
+ });
+
+ it('should include resource parameter when provided', () => {
+ const url = buildAuthorizationUrl(
+ baseConfig,
+ basePkceParams,
+ 3000,
+ 'https://mcp.example.com',
+ );
+ const parsed = new URL(url);
+ expect(parsed.searchParams.get('resource')).toBe(
+ 'https://mcp.example.com',
+ );
+ });
+
+ it('should not include resource parameter when not provided', () => {
+ const url = buildAuthorizationUrl(baseConfig, basePkceParams, 3000);
+ const parsed = new URL(url);
+ expect(parsed.searchParams.has('resource')).toBe(false);
+ });
+ });
+
+ describe('startCallbackServer', () => {
+ it('should start a server and resolve port', async () => {
+ const server = startCallbackServer('test-state');
+ const port = await server.port;
+ expect(port).toBeGreaterThan(0);
+
+ // Make a successful callback request to close the server
+ const res = await realFetch(
+ `http://localhost:${port}${REDIRECT_PATH}?code=abc&state=test-state`,
+ );
+ expect(res.status).toBe(200);
+ await server.response;
+ });
+
+ it('should resolve response with code and state on valid callback', async () => {
+ const server = startCallbackServer('my-state');
+ const port = await server.port;
+
+ await realFetch(
+ `http://localhost:${port}${REDIRECT_PATH}?code=auth-code-123&state=my-state`,
+ );
+
+ const response = await server.response;
+ expect(response.code).toBe('auth-code-123');
+ expect(response.state).toBe('my-state');
+ });
+
+ it('should reject on state mismatch', async () => {
+ const server = startCallbackServer('expected-state');
+ const port = await server.port;
+
+ // Attach rejection handler BEFORE triggering the callback to prevent
+ // unhandled rejection race with Vitest's detection.
+ const responseResult = server.response.then(
+ () => new Error('Expected rejection'),
+ (e: Error) => e,
+ );
+
+ await realFetch(
+ `http://localhost:${port}${REDIRECT_PATH}?code=abc&state=wrong-state`,
+ ).catch(() => {
+ // Connection may be reset by server closing — expected
+ });
+
+ const error = await responseResult;
+ expect(error.message).toContain('State mismatch - possible CSRF attack');
+ });
+
+ it('should reject on OAuth error in callback', async () => {
+ const server = startCallbackServer('test-state');
+ const port = await server.port;
+
+ // Attach rejection handler BEFORE triggering the callback
+ const responseResult = server.response.then(
+ () => new Error('Expected rejection'),
+ (e: Error) => e,
+ );
+
+ await realFetch(
+ `http://localhost:${port}${REDIRECT_PATH}?error=access_denied&error_description=User+denied`,
+ ).catch(() => {
+ // Connection may be reset by server closing — expected
+ });
+
+ const error = await responseResult;
+ expect(error.message).toContain('OAuth error: access_denied');
+ });
+
+ it('should return 404 for non-callback paths', async () => {
+ const server = startCallbackServer('test-state');
+ const port = await server.port;
+
+ const res = await realFetch(`http://localhost:${port}/other-path`);
+ expect(res.status).toBe(404);
+
+ // Clean up: send valid callback to close the server
+ await realFetch(
+ `http://localhost:${port}${REDIRECT_PATH}?code=abc&state=test-state`,
+ );
+ await server.response;
+ });
+
+ it('should reject when OAUTH_CALLBACK_PORT env var is invalid', async () => {
+ vi.stubEnv('OAUTH_CALLBACK_PORT', 'not-a-number');
+
+ const server = startCallbackServer('test-state');
+
+ await expect(server.port).rejects.toThrow(
+ 'Invalid value for OAUTH_CALLBACK_PORT',
+ );
+ await expect(server.response).rejects.toThrow(
+ 'Invalid value for OAUTH_CALLBACK_PORT',
+ );
+ });
+ });
+
+ describe('exchangeCodeForToken', () => {
+ it('should exchange code for token with JSON response', async () => {
+ const tokenResponse = {
+ access_token: 'test-access-token',
+ token_type: 'Bearer',
+ expires_in: 3600,
+ refresh_token: 'test-refresh-token',
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(JSON.stringify(tokenResponse)),
+ );
+
+ const result = await exchangeCodeForToken(
+ baseConfig,
+ 'auth-code',
+ 'code-verifier',
+ 3000,
+ );
+
+ expect(result.access_token).toBe('test-access-token');
+ expect(result.token_type).toBe('Bearer');
+ expect(result.expires_in).toBe(3600);
+ expect(result.refresh_token).toBe('test-refresh-token');
+ });
+
+ it('should send correct parameters in the request body', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await exchangeCodeForToken(baseConfig, 'my-code', 'my-verifier', 4000);
+
+ expect(mockFetch).toHaveBeenCalledTimes(1);
+ const [url, options] = mockFetch.mock.calls[0] as [string, RequestInit];
+ expect(url).toBe('https://auth.example.com/token');
+ const body = new URLSearchParams(options.body as string);
+ expect(body.get('grant_type')).toBe('authorization_code');
+ expect(body.get('code')).toBe('my-code');
+ expect(body.get('code_verifier')).toBe('my-verifier');
+ expect(body.get('client_id')).toBe('test-client-id');
+ expect(body.get('redirect_uri')).toBe(
+ `http://localhost:4000${REDIRECT_PATH}`,
+ );
+ });
+
+ it('should include client_secret when provided', async () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ clientSecret: 'my-secret',
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await exchangeCodeForToken(config, 'code', 'verifier', 3000);
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('client_secret')).toBe('my-secret');
+ });
+
+ it('should include resource parameter when provided', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await exchangeCodeForToken(
+ baseConfig,
+ 'code',
+ 'verifier',
+ 3000,
+ 'https://mcp.example.com',
+ );
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('resource')).toBe('https://mcp.example.com');
+ });
+
+ it('should handle form-urlencoded token response', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ 'access_token=form-token&token_type=Bearer&expires_in=7200',
+ { contentType: 'application/x-www-form-urlencoded' },
+ ),
+ );
+
+ const result = await exchangeCodeForToken(
+ baseConfig,
+ 'code',
+ 'verifier',
+ 3000,
+ );
+
+ expect(result.access_token).toBe('form-token');
+ expect(result.token_type).toBe('Bearer');
+ expect(result.expires_in).toBe(7200);
+ });
+
+ it('should throw on non-ok response', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse('Bad request', { status: 400 }),
+ );
+
+ await expect(
+ exchangeCodeForToken(baseConfig, 'code', 'verifier', 3000),
+ ).rejects.toThrow('Token exchange failed');
+ });
+
+ it('should throw on non-ok response with form-urlencoded error', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ 'error=invalid_grant&error_description=Code+expired',
+ {
+ status: 400,
+ contentType: 'application/x-www-form-urlencoded',
+ },
+ ),
+ );
+
+ await expect(
+ exchangeCodeForToken(baseConfig, 'code', 'verifier', 3000),
+ ).rejects.toThrow('invalid_grant');
+ });
+
+ it('should throw when JSON response has no access_token and form-urlencoded fallback also fails', async () => {
+ // JSON that parses but has no access_token — falls through to form-urlencoded
+ // which also has no access_token
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(JSON.stringify({ error: 'server_error' })),
+ );
+
+ await expect(
+ exchangeCodeForToken(baseConfig, 'code', 'verifier', 3000),
+ ).rejects.toThrow('Token exchange failed');
+ });
+
+ it('should use custom redirectUri from config', async () => {
+ const config: OAuthFlowConfig = {
+ ...baseConfig,
+ redirectUri: 'https://custom.example.com/cb',
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await exchangeCodeForToken(config, 'code', 'verifier', 3000);
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('redirect_uri')).toBe('https://custom.example.com/cb');
+ });
+
+ it('should default token_type to Bearer when missing from JSON response', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(JSON.stringify({ access_token: 'tok' })),
+ );
+
+ const result = await exchangeCodeForToken(
+ baseConfig,
+ 'code',
+ 'verifier',
+ 3000,
+ );
+ expect(result.token_type).toBe('Bearer');
+ });
+ });
+
+ describe('refreshAccessToken', () => {
+ const refreshConfig: OAuthRefreshConfig = {
+ clientId: 'test-client-id',
+ };
+
+ it('should refresh a token with JSON response', async () => {
+ const tokenResponse = {
+ access_token: 'new-access-token',
+ token_type: 'Bearer',
+ expires_in: 3600,
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(JSON.stringify(tokenResponse)),
+ );
+
+ const result = await refreshAccessToken(
+ refreshConfig,
+ 'old-refresh-token',
+ 'https://auth.example.com/token',
+ );
+
+ expect(result.access_token).toBe('new-access-token');
+ expect(result.expires_in).toBe(3600);
+ });
+
+ it('should send correct parameters in the request body', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await refreshAccessToken(
+ refreshConfig,
+ 'my-refresh-token',
+ 'https://auth.example.com/token',
+ );
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('grant_type')).toBe('refresh_token');
+ expect(body.get('refresh_token')).toBe('my-refresh-token');
+ expect(body.get('client_id')).toBe('test-client-id');
+ });
+
+ it('should include client_secret when provided', async () => {
+ const config: OAuthRefreshConfig = {
+ ...refreshConfig,
+ clientSecret: 'secret',
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await refreshAccessToken(
+ config,
+ 'refresh-token',
+ 'https://auth.example.com/token',
+ );
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('client_secret')).toBe('secret');
+ });
+
+ it('should include scopes and audiences when provided', async () => {
+ const config: OAuthRefreshConfig = {
+ ...refreshConfig,
+ scopes: ['read', 'write'],
+ audiences: ['https://api.example.com'],
+ };
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await refreshAccessToken(
+ config,
+ 'refresh-token',
+ 'https://auth.example.com/token',
+ );
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('scope')).toBe('read write');
+ expect(body.get('audience')).toBe('https://api.example.com');
+ });
+
+ it('should include resource parameter when provided', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ JSON.stringify({ access_token: 'tok', token_type: 'Bearer' }),
+ ),
+ );
+
+ await refreshAccessToken(
+ refreshConfig,
+ 'refresh-token',
+ 'https://auth.example.com/token',
+ 'https://mcp.example.com',
+ );
+
+ const body = new URLSearchParams(
+ (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+ );
+ expect(body.get('resource')).toBe('https://mcp.example.com');
+ });
+
+ it('should throw on non-ok response', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse('Unauthorized', { status: 401 }),
+ );
+
+ await expect(
+ refreshAccessToken(
+ refreshConfig,
+ 'bad-token',
+ 'https://auth.example.com/token',
+ ),
+ ).rejects.toThrow('Token refresh failed');
+ });
+
+ it('should handle form-urlencoded token response', async () => {
+ mockFetch.mockResolvedValueOnce(
+ createMockResponse(
+ 'access_token=refreshed-token&token_type=Bearer&expires_in=1800',
+ { contentType: 'application/x-www-form-urlencoded' },
+ ),
+ );
+
+ const result = await refreshAccessToken(
+ refreshConfig,
+ 'refresh-token',
+ 'https://auth.example.com/token',
+ );
+
+ expect(result.access_token).toBe('refreshed-token');
+ expect(result.expires_in).toBe(1800);
+ });
+ });
+});
diff --git a/packages/core/src/utils/oauth-flow.ts b/packages/core/src/utils/oauth-flow.ts
new file mode 100644
index 0000000000..e13fd37837
--- /dev/null
+++ b/packages/core/src/utils/oauth-flow.ts
@@ -0,0 +1,520 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Shared OAuth 2.0 Authorization Code flow primitives with PKCE support.
+ *
+ * These utilities are protocol-agnostic and can be used by both MCP OAuth
+ * and A2A OAuth authentication providers.
+ */
+
+import * as http from 'node:http';
+import * as crypto from 'node:crypto';
+import type * as net from 'node:net';
+import { URL } from 'node:url';
+import { debugLogger } from './debugLogger.js';
+
+/**
+ * Configuration for an OAuth 2.0 Authorization Code flow.
+ * Contains only the fields needed by the shared flow utilities.
+ */
+export interface OAuthFlowConfig {
+ clientId: string;
+ clientSecret?: string;
+ authorizationUrl: string;
+ tokenUrl: string;
+ scopes?: string[];
+ audiences?: string[];
+ redirectUri?: string;
+}
+
+/**
+ * Configuration subset needed for token refresh operations.
+ */
+export type OAuthRefreshConfig = Pick<
+ OAuthFlowConfig,
+ 'clientId' | 'clientSecret' | 'scopes' | 'audiences'
+>;
+
+/**
+ * PKCE (Proof Key for Code Exchange) parameters.
+ */
+export interface PKCEParams {
+ codeVerifier: string;
+ codeChallenge: string;
+ state: string;
+}
+
+/**
+ * OAuth authorization response from the callback server.
+ */
+export interface OAuthAuthorizationResponse {
+ code: string;
+ state: string;
+}
+
+/**
+ * OAuth token response from the authorization server.
+ */
+export interface OAuthTokenResponse {
+ access_token: string;
+ token_type: string;
+ expires_in?: number;
+ refresh_token?: string;
+ scope?: string;
+}
+
+/** The path the local callback server listens on. */
+export const REDIRECT_PATH = '/oauth/callback';
+
+const HTTP_OK = 200;
+
+/**
+ * Generate PKCE parameters for OAuth flow.
+ *
+ * @returns PKCE parameters including code verifier, challenge, and state
+ */
+export function generatePKCEParams(): PKCEParams {
+ // Generate code verifier (43-128 characters)
+ // using 64 bytes results in ~86 characters, safely above the minimum of 43
+ const codeVerifier = crypto.randomBytes(64).toString('base64url');
+
+ // Generate code challenge using SHA256
+ const codeChallenge = crypto
+ .createHash('sha256')
+ .update(codeVerifier)
+ .digest('base64url');
+
+ // Generate state for CSRF protection
+ const state = crypto.randomBytes(16).toString('base64url');
+
+ return { codeVerifier, codeChallenge, state };
+}
+
+/**
+ * Start a local HTTP server to handle OAuth callback.
+ * The server will listen on the specified port (or port 0 for OS assignment).
+ *
+ * @param expectedState The state parameter to validate
+ * @param port Optional preferred port to listen on
+ * @returns Object containing the port (available immediately) and a promise for the auth response
+ */
+export function startCallbackServer(
+ expectedState: string,
+ port?: number,
+): {
+ port: Promise;
+ response: Promise;
+} {
+ let portResolve: (port: number) => void;
+ let portReject: (error: Error) => void;
+ const portPromise = new Promise((resolve, reject) => {
+ portResolve = resolve;
+ portReject = reject;
+ });
+
+ const responsePromise = new Promise(
+ (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
+ */
+export function getPortFromUrl(urlString?: string): number | undefined {
+ if (!urlString) {
+ 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 flow configuration
+ * @param pkceParams PKCE parameters
+ * @param redirectPort The port to use for the redirect URI
+ * @param resource Optional resource parameter value (RFC 8707)
+ * @returns The authorization URL
+ */
+export function buildAuthorizationUrl(
+ config: OAuthFlowConfig,
+ pkceParams: PKCEParams,
+ redirectPort: number,
+ resource?: 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(' '));
+ }
+
+ if (resource) {
+ params.append('resource', resource);
+ }
+
+ const url = new URL(config.authorizationUrl);
+ params.forEach((value, key) => {
+ url.searchParams.append(key, value);
+ });
+ return url.toString();
+}
+
+/**
+ * Parse a token endpoint response, handling both JSON and form-urlencoded formats.
+ *
+ * @param response The HTTP response from the token endpoint
+ * @param operationName Human-readable operation name for error messages (e.g., "Token exchange", "Token refresh")
+ * @param defaultErrorCode Default error code when access_token is missing (e.g., "no_access_token", "unknown_error")
+ * @returns The parsed token response
+ */
+async function parseTokenEndpointResponse(
+ response: Response,
+ operationName: string,
+ defaultErrorCode: string,
+): Promise {
+ 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 = `${operationName} failed: ${error} - ${errorDescription || 'No description'}`;
+ }
+ } catch {
+ // Fall back to raw error
+ }
+ throw new Error(
+ errorMessage ||
+ `${operationName} failed: ${response.status} - ${responseText}`,
+ );
+ }
+
+ // Log unexpected content types for debugging
+ if (
+ !contentType.includes('application/json') &&
+ !contentType.includes('application/x-www-form-urlencoded')
+ ) {
+ debugLogger.warn(
+ `${operationName} 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 {
+ const data: unknown = JSON.parse(responseText);
+ if (
+ data &&
+ typeof data === 'object' &&
+ 'access_token' in data &&
+ // eslint-disable-next-line no-restricted-syntax
+ typeof (data as Record)['access_token'] === 'string'
+ ) {
+ const obj = data as Record;
+ const result: OAuthTokenResponse = {
+ access_token: String(obj['access_token']),
+ token_type:
+ // eslint-disable-next-line no-restricted-syntax
+ typeof obj['token_type'] === 'string' ? obj['token_type'] : 'Bearer',
+ expires_in:
+ // eslint-disable-next-line no-restricted-syntax
+ typeof obj['expires_in'] === 'number' ? obj['expires_in'] : undefined,
+ refresh_token:
+ // eslint-disable-next-line no-restricted-syntax
+ typeof obj['refresh_token'] === 'string'
+ ? obj['refresh_token']
+ : undefined,
+ // eslint-disable-next-line no-restricted-syntax
+ scope: typeof obj['scope'] === 'string' ? obj['scope'] : undefined,
+ };
+ return result;
+ }
+ // JSON parsed but doesn't look like a token response — fall through
+ } catch {
+ // Not JSON — fall through to form-urlencoded parsing
+ }
+
+ // 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(
+ `${operationName} failed: ${error || defaultErrorCode} - ${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;
+}
+
+/**
+ * Exchange an authorization code for tokens.
+ *
+ * @param config OAuth flow configuration
+ * @param code Authorization code
+ * @param codeVerifier PKCE code verifier
+ * @param redirectPort The port to use for the redirect URI
+ * @param resource Optional resource parameter value (RFC 8707)
+ * @returns The token response
+ */
+export async function exchangeCodeForToken(
+ config: OAuthFlowConfig,
+ code: string,
+ codeVerifier: string,
+ redirectPort: number,
+ resource?: 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(' '));
+ }
+
+ if (resource) {
+ params.append('resource', resource);
+ }
+
+ 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(),
+ });
+
+ return parseTokenEndpointResponse(
+ response,
+ 'Token exchange',
+ 'no_access_token',
+ );
+}
+
+/**
+ * Refresh an access token using a refresh token.
+ *
+ * @param config OAuth configuration subset needed for refresh
+ * @param refreshToken The refresh token
+ * @param tokenUrl The token endpoint URL
+ * @param resource Optional resource parameter value (RFC 8707)
+ * @returns The new token response
+ */
+export async function refreshAccessToken(
+ config: OAuthRefreshConfig,
+ refreshToken: string,
+ tokenUrl: string,
+ resource?: 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.scopes && config.scopes.length > 0) {
+ params.append('scope', config.scopes.join(' '));
+ }
+
+ if (config.audiences && config.audiences.length > 0) {
+ params.append('audience', config.audiences.join(' '));
+ }
+
+ if (resource) {
+ params.append('resource', resource);
+ }
+
+ const response = await fetch(tokenUrl, {
+ method: 'POST',
+ headers: {
+ 'Content-Type': 'application/x-www-form-urlencoded',
+ Accept: 'application/json, application/x-www-form-urlencoded',
+ },
+ body: params.toString(),
+ });
+
+ return parseTokenEndpointResponse(response, 'Token refresh', 'unknown_error');
+}
diff --git a/packages/core/src/utils/partUtils.test.ts b/packages/core/src/utils/partUtils.test.ts
index d530107257..5a8130c97c 100644
--- a/packages/core/src/utils/partUtils.test.ts
+++ b/packages/core/src/utils/partUtils.test.ts
@@ -123,7 +123,19 @@ describe('partUtils', () => {
it('should return descriptive string for inlineData part', () => {
const part = { inlineData: { mimeType: 'image/png', data: '' } } as Part;
- expect(partToString(part, verboseOptions)).toBe('');
+ expect(partToString(part, verboseOptions)).toBe(
+ '[Image: image/png, 0.0 KB]',
+ );
+ });
+
+ it('should show size for inlineData with non-empty base64 data', () => {
+ // 4 base64 chars → ceil(4*3/4) = 3 bytes → 3/1024 ≈ 0.0 KB
+ const part = {
+ inlineData: { mimeType: 'audio/mp3', data: 'AAAA' },
+ } as Part;
+ expect(partToString(part, verboseOptions)).toBe(
+ '[Audio: audio/mp3, 0.0 KB]',
+ );
});
it('should return an empty string for an unknown part type', () => {
@@ -142,7 +154,7 @@ describe('partUtils', () => {
],
];
expect(partToString(parts as Part, verboseOptions)).toBe(
- 'start middle[Function Call: func1] end