Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 74 additions & 0 deletions packages/core/src/tools/mcp-client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1792,6 +1792,80 @@ describe('mcp-client', () => {
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBeUndefined();
});

it('should include extension settings with defined values in environment', async () => {
const mockedTransport = vi
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
.mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

await createTransport(
'test-server',
{
command: 'test-command',
extension: {
name: 'test-ext',
resolvedSettings: [
{
envVar: 'GEMINI_CLI_EXT_VAR',
value: 'defined-value',
sensitive: false,
name: 'ext-setting',
},
],
version: '',
isActive: false,
path: '',
contextFiles: [],
id: '',
},
},
false,
EMPTY_CONFIG,
);

const callArgs = mockedTransport.mock.calls[0][0];
expect(callArgs.env).toBeDefined();
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('defined-value');
});

it('should resolve environment variables in mcpServerConfig.env using extension settings', async () => {
const mockedTransport = vi
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
.mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

await createTransport(
'test-server',
{
command: 'test-command',
env: {
RESOLVED_VAR: '$GEMINI_CLI_EXT_VAR',
},
extension: {
name: 'test-ext',
resolvedSettings: [
{
envVar: 'GEMINI_CLI_EXT_VAR',
value: 'ext-value',
sensitive: false,
name: 'ext-setting',
},
],
version: '',
isActive: false,
path: '',
contextFiles: [],
id: '',
},
},
false,
EMPTY_CONFIG,
);

const callArgs = mockedTransport.mock.calls[0][0];
expect(callArgs.env).toBeDefined();
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')
Expand Down
59 changes: 53 additions & 6 deletions packages/core/src/tools/mcp-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,11 @@ import {
ProgressNotificationSchema,
} from '@modelcontextprotocol/sdk/types.js';
import { parse } from 'shell-quote';
import type { Config, MCPServerConfig } from '../config/config.js';
import type {
Config,
MCPServerConfig,
GeminiCLIExtension,
} from '../config/config.js';
import { AuthProviderType } from '../config/config.js';
import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
import { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js';
Expand Down Expand Up @@ -778,15 +782,25 @@ async function handleAutomaticOAuth(
*
* @param mcpServerConfig The MCP server configuration
* @param headers Additional headers
* @param sanitizationConfig Configuration for environment sanitization
*/
function createTransportRequestInit(
mcpServerConfig: MCPServerConfig,
headers: Record<string, string>,
sanitizationConfig: EnvironmentSanitizationConfig,
): RequestInit {
const extensionEnv = getExtensionEnvironment(mcpServerConfig.extension);
const expansionEnv = { ...process.env, ...extensionEnv };
Comment thread
chrstnb marked this conversation as resolved.

const sanitizedEnv = sanitizeEnvironment(expansionEnv, {
...sanitizationConfig,
enableEnvironmentVariableRedaction: true,
});

const expandedHeaders: Record<string, string> = {};
if (mcpServerConfig.headers) {
for (const [key, value] of Object.entries(mcpServerConfig.headers)) {
expandedHeaders[key] = expandEnvVars(value, process.env);
expandedHeaders[key] = expandEnvVars(value, sanitizedEnv);
}
}

Expand Down Expand Up @@ -826,12 +840,14 @@ function createAuthProvider(
* @param mcpServerName The name of the MCP server
* @param mcpServerConfig The MCP server configuration
* @param accessToken The OAuth access token
* @param sanitizationConfig Configuration for environment sanitization
* @returns The transport with OAuth token, or null if creation fails
*/
async function createTransportWithOAuth(
mcpServerName: string,
mcpServerConfig: MCPServerConfig,
accessToken: string,
sanitizationConfig: EnvironmentSanitizationConfig,
): Promise<StreamableHTTPClientTransport | SSEClientTransport | null> {
try {
const headers: Record<string, string> = {
Expand All @@ -840,7 +856,11 @@ async function createTransportWithOAuth(
const transportOptions:
| StreamableHTTPClientTransportOptions
| SSEClientTransportOptions = {
requestInit: createTransportRequestInit(mcpServerConfig, headers),
requestInit: createTransportRequestInit(
mcpServerConfig,
headers,
sanitizationConfig,
),
};

return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
Expand Down Expand Up @@ -1435,13 +1455,15 @@ async function showAuthRequiredMessage(serverName: string): Promise<never> {
* @param config The MCP server configuration
* @param accessToken The OAuth access token to use
* @param httpReturned404 Whether the HTTP transport returned 404 (indicating SSE-only server)
* @param sanitizationConfig Configuration for environment sanitization
*/
async function retryWithOAuth(
client: Client,
serverName: string,
config: MCPServerConfig,
accessToken: string,
httpReturned404: boolean,
sanitizationConfig: EnvironmentSanitizationConfig,
): Promise<void> {
if (httpReturned404) {
// HTTP returned 404, only try SSE
Expand All @@ -1462,6 +1484,7 @@ async function retryWithOAuth(
serverName,
config,
accessToken,
sanitizationConfig,
);
if (!httpTransport) {
throw new Error(
Expand Down Expand Up @@ -1741,6 +1764,7 @@ export async function connectToMcpServer(
mcpServerConfig,
accessToken,
httpReturned404,
sanitizationConfig,
);
return mcpClient;
} else {
Expand Down Expand Up @@ -1813,6 +1837,7 @@ export async function connectToMcpServer(
mcpServerName,
mcpServerConfig,
accessToken,
sanitizationConfig,
);
if (!oauthTransport) {
throw new Error(
Expand Down Expand Up @@ -1960,23 +1985,31 @@ export async function createTransport(
const transportOptions:
| StreamableHTTPClientTransportOptions
| SSEClientTransportOptions = {
requestInit: createTransportRequestInit(mcpServerConfig, headers),
requestInit: createTransportRequestInit(
mcpServerConfig,
headers,
sanitizationConfig,
),
authProvider,
};

return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
}

if (mcpServerConfig.command) {
const extensionEnv = getExtensionEnvironment(mcpServerConfig.extension);
const expansionEnv = { ...process.env, ...extensionEnv };
Comment thread
chrstnb marked this conversation as resolved.

// 1. Sanitize the base process environment to prevent unintended leaks of system-wide secrets.
const sanitizedEnv = sanitizeEnvironment(process.env, {
const sanitizedEnv = sanitizeEnvironment(expansionEnv, {
...sanitizationConfig,
enableEnvironmentVariableRedaction: true,
});

const finalEnv: Record<string, string> = {
[GEMINI_CLI_IDENTIFICATION_ENV_VAR]:
GEMINI_CLI_IDENTIFICATION_ENV_VAR_VALUE,
...extensionEnv,
};
for (const [key, value] of Object.entries(sanitizedEnv)) {
if (value !== undefined) {
Expand All @@ -1987,7 +2020,7 @@ export async function createTransport(
// Expand and merge explicit environment variables from the MCP configuration.
if (mcpServerConfig.env) {
for (const [key, value] of Object.entries(mcpServerConfig.env)) {
finalEnv[key] = expandEnvVars(value, process.env);
finalEnv[key] = expandEnvVars(value, expansionEnv);
}
}

Expand Down Expand Up @@ -2045,6 +2078,20 @@ interface NamedTool {
name?: string;
}

function getExtensionEnvironment(
extension?: GeminiCLIExtension,
): Record<string, string> {
const env: Record<string, string> = {};
if (extension?.resolvedSettings) {
for (const setting of extension.resolvedSettings) {
if (setting.value !== undefined) {
env[setting.envVar] = setting.value;
}
}
}
return env;
}

/** Visible for testing */
export function isEnabled(
funcDecl: NamedTool,
Expand Down
Loading