Skip to content

Commit f773a97

Browse files
authored
Fix extension MCP server env var loading (#20374)
1 parent e9ef491 commit f773a97

2 files changed

Lines changed: 127 additions & 6 deletions

File tree

packages/core/src/tools/mcp-client.test.ts

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1792,6 +1792,80 @@ describe('mcp-client', () => {
17921792
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBeUndefined();
17931793
});
17941794

1795+
it('should include extension settings with defined values in environment', async () => {
1796+
const mockedTransport = vi
1797+
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
1798+
.mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);
1799+
1800+
await createTransport(
1801+
'test-server',
1802+
{
1803+
command: 'test-command',
1804+
extension: {
1805+
name: 'test-ext',
1806+
resolvedSettings: [
1807+
{
1808+
envVar: 'GEMINI_CLI_EXT_VAR',
1809+
value: 'defined-value',
1810+
sensitive: false,
1811+
name: 'ext-setting',
1812+
},
1813+
],
1814+
version: '',
1815+
isActive: false,
1816+
path: '',
1817+
contextFiles: [],
1818+
id: '',
1819+
},
1820+
},
1821+
false,
1822+
EMPTY_CONFIG,
1823+
);
1824+
1825+
const callArgs = mockedTransport.mock.calls[0][0];
1826+
expect(callArgs.env).toBeDefined();
1827+
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('defined-value');
1828+
});
1829+
1830+
it('should resolve environment variables in mcpServerConfig.env using extension settings', async () => {
1831+
const mockedTransport = vi
1832+
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
1833+
.mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);
1834+
1835+
await createTransport(
1836+
'test-server',
1837+
{
1838+
command: 'test-command',
1839+
env: {
1840+
RESOLVED_VAR: '$GEMINI_CLI_EXT_VAR',
1841+
},
1842+
extension: {
1843+
name: 'test-ext',
1844+
resolvedSettings: [
1845+
{
1846+
envVar: 'GEMINI_CLI_EXT_VAR',
1847+
value: 'ext-value',
1848+
sensitive: false,
1849+
name: 'ext-setting',
1850+
},
1851+
],
1852+
version: '',
1853+
isActive: false,
1854+
path: '',
1855+
contextFiles: [],
1856+
id: '',
1857+
},
1858+
},
1859+
false,
1860+
EMPTY_CONFIG,
1861+
);
1862+
1863+
const callArgs = mockedTransport.mock.calls[0][0];
1864+
expect(callArgs.env).toBeDefined();
1865+
expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('ext-value');
1866+
expect(callArgs.env!['RESOLVED_VAR']).toBe('ext-value');
1867+
});
1868+
17951869
it('should expand environment variables in mcpServerConfig.env and not redact them', async () => {
17961870
const mockedTransport = vi
17971871
.spyOn(SdkClientStdioLib, 'StdioClientTransport')

packages/core/src/tools/mcp-client.ts

Lines changed: 53 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@ import {
3434
ProgressNotificationSchema,
3535
} from '@modelcontextprotocol/sdk/types.js';
3636
import { parse } from 'shell-quote';
37-
import type { Config, MCPServerConfig } from '../config/config.js';
37+
import type {
38+
Config,
39+
MCPServerConfig,
40+
GeminiCLIExtension,
41+
} from '../config/config.js';
3842
import { AuthProviderType } from '../config/config.js';
3943
import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
4044
import { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js';
@@ -778,15 +782,25 @@ async function handleAutomaticOAuth(
778782
*
779783
* @param mcpServerConfig The MCP server configuration
780784
* @param headers Additional headers
785+
* @param sanitizationConfig Configuration for environment sanitization
781786
*/
782787
function createTransportRequestInit(
783788
mcpServerConfig: MCPServerConfig,
784789
headers: Record<string, string>,
790+
sanitizationConfig: EnvironmentSanitizationConfig,
785791
): RequestInit {
792+
const extensionEnv = getExtensionEnvironment(mcpServerConfig.extension);
793+
const expansionEnv = { ...process.env, ...extensionEnv };
794+
795+
const sanitizedEnv = sanitizeEnvironment(expansionEnv, {
796+
...sanitizationConfig,
797+
enableEnvironmentVariableRedaction: true,
798+
});
799+
786800
const expandedHeaders: Record<string, string> = {};
787801
if (mcpServerConfig.headers) {
788802
for (const [key, value] of Object.entries(mcpServerConfig.headers)) {
789-
expandedHeaders[key] = expandEnvVars(value, process.env);
803+
expandedHeaders[key] = expandEnvVars(value, sanitizedEnv);
790804
}
791805
}
792806

@@ -826,12 +840,14 @@ function createAuthProvider(
826840
* @param mcpServerName The name of the MCP server
827841
* @param mcpServerConfig The MCP server configuration
828842
* @param accessToken The OAuth access token
843+
* @param sanitizationConfig Configuration for environment sanitization
829844
* @returns The transport with OAuth token, or null if creation fails
830845
*/
831846
async function createTransportWithOAuth(
832847
mcpServerName: string,
833848
mcpServerConfig: MCPServerConfig,
834849
accessToken: string,
850+
sanitizationConfig: EnvironmentSanitizationConfig,
835851
): Promise<StreamableHTTPClientTransport | SSEClientTransport | null> {
836852
try {
837853
const headers: Record<string, string> = {
@@ -840,7 +856,11 @@ async function createTransportWithOAuth(
840856
const transportOptions:
841857
| StreamableHTTPClientTransportOptions
842858
| SSEClientTransportOptions = {
843-
requestInit: createTransportRequestInit(mcpServerConfig, headers),
859+
requestInit: createTransportRequestInit(
860+
mcpServerConfig,
861+
headers,
862+
sanitizationConfig,
863+
),
844864
};
845865

846866
return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
@@ -1435,13 +1455,15 @@ async function showAuthRequiredMessage(serverName: string): Promise<never> {
14351455
* @param config The MCP server configuration
14361456
* @param accessToken The OAuth access token to use
14371457
* @param httpReturned404 Whether the HTTP transport returned 404 (indicating SSE-only server)
1458+
* @param sanitizationConfig Configuration for environment sanitization
14381459
*/
14391460
async function retryWithOAuth(
14401461
client: Client,
14411462
serverName: string,
14421463
config: MCPServerConfig,
14431464
accessToken: string,
14441465
httpReturned404: boolean,
1466+
sanitizationConfig: EnvironmentSanitizationConfig,
14451467
): Promise<void> {
14461468
if (httpReturned404) {
14471469
// HTTP returned 404, only try SSE
@@ -1462,6 +1484,7 @@ async function retryWithOAuth(
14621484
serverName,
14631485
config,
14641486
accessToken,
1487+
sanitizationConfig,
14651488
);
14661489
if (!httpTransport) {
14671490
throw new Error(
@@ -1741,6 +1764,7 @@ export async function connectToMcpServer(
17411764
mcpServerConfig,
17421765
accessToken,
17431766
httpReturned404,
1767+
sanitizationConfig,
17441768
);
17451769
return mcpClient;
17461770
} else {
@@ -1813,6 +1837,7 @@ export async function connectToMcpServer(
18131837
mcpServerName,
18141838
mcpServerConfig,
18151839
accessToken,
1840+
sanitizationConfig,
18161841
);
18171842
if (!oauthTransport) {
18181843
throw new Error(
@@ -1960,23 +1985,31 @@ export async function createTransport(
19601985
const transportOptions:
19611986
| StreamableHTTPClientTransportOptions
19621987
| SSEClientTransportOptions = {
1963-
requestInit: createTransportRequestInit(mcpServerConfig, headers),
1988+
requestInit: createTransportRequestInit(
1989+
mcpServerConfig,
1990+
headers,
1991+
sanitizationConfig,
1992+
),
19641993
authProvider,
19651994
};
19661995

19671996
return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
19681997
}
19691998

19701999
if (mcpServerConfig.command) {
2000+
const extensionEnv = getExtensionEnvironment(mcpServerConfig.extension);
2001+
const expansionEnv = { ...process.env, ...extensionEnv };
2002+
19712003
// 1. Sanitize the base process environment to prevent unintended leaks of system-wide secrets.
1972-
const sanitizedEnv = sanitizeEnvironment(process.env, {
2004+
const sanitizedEnv = sanitizeEnvironment(expansionEnv, {
19732005
...sanitizationConfig,
19742006
enableEnvironmentVariableRedaction: true,
19752007
});
19762008

19772009
const finalEnv: Record<string, string> = {
19782010
[GEMINI_CLI_IDENTIFICATION_ENV_VAR]:
19792011
GEMINI_CLI_IDENTIFICATION_ENV_VAR_VALUE,
2012+
...extensionEnv,
19802013
};
19812014
for (const [key, value] of Object.entries(sanitizedEnv)) {
19822015
if (value !== undefined) {
@@ -1987,7 +2020,7 @@ export async function createTransport(
19872020
// Expand and merge explicit environment variables from the MCP configuration.
19882021
if (mcpServerConfig.env) {
19892022
for (const [key, value] of Object.entries(mcpServerConfig.env)) {
1990-
finalEnv[key] = expandEnvVars(value, process.env);
2023+
finalEnv[key] = expandEnvVars(value, expansionEnv);
19912024
}
19922025
}
19932026

@@ -2045,6 +2078,20 @@ interface NamedTool {
20452078
name?: string;
20462079
}
20472080

2081+
function getExtensionEnvironment(
2082+
extension?: GeminiCLIExtension,
2083+
): Record<string, string> {
2084+
const env: Record<string, string> = {};
2085+
if (extension?.resolvedSettings) {
2086+
for (const setting of extension.resolvedSettings) {
2087+
if (setting.value !== undefined) {
2088+
env[setting.envVar] = setting.value;
2089+
}
2090+
}
2091+
}
2092+
return env;
2093+
}
2094+
20482095
/** Visible for testing */
20492096
export function isEnabled(
20502097
funcDecl: NamedTool,

0 commit comments

Comments
 (0)