Skip to content

Commit 99ae4d8

Browse files
cocosheng-gcocosheng-g
authored
fix(core): generalize MCP compliance fix for tool results (google-gemini#27045)
1 parent aebdba6 commit 99ae4d8

5 files changed

Lines changed: 275 additions & 191 deletions

File tree

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

Lines changed: 49 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import {
4040
discoverPrompts,
4141
type McpContext,
4242
} from './mcp-client.js';
43+
import { McpComplianceTransport } from './mcp-compliance-transport.js';
4344
import type { ToolRegistry } from './tool-registry.js';
4445
import type { ResourceRegistry } from '../resources/resource-registry.js';
4546
import * as fs from 'node:fs';
@@ -71,6 +72,9 @@ const MOCK_CONTEXT_DEFAULT = {
7172

7273
let MOCK_CONTEXT: McpContext = MOCK_CONTEXT_DEFAULT;
7374

75+
const unwrap = (t: any) =>
76+
t instanceof McpComplianceTransport ? t.transport : t;
77+
7478
vi.mock('@modelcontextprotocol/sdk/client/stdio.js');
7579
vi.mock('@modelcontextprotocol/sdk/client/index.js');
7680
vi.mock('@google/genai');
@@ -1937,7 +1941,7 @@ describe('mcp-client', () => {
19371941
MOCK_CONTEXT,
19381942
);
19391943

1940-
const testableTransport = transport as unknown as {
1944+
const testableTransport = unwrap(transport) as unknown as {
19411945
_authProvider?: {
19421946
tokens: () => Promise<{ access_token: string } | undefined>;
19431947
};
@@ -1983,7 +1987,7 @@ describe('mcp-client', () => {
19831987
MOCK_CONTEXT,
19841988
);
19851989

1986-
const testableTransport = transport as unknown as {
1990+
const testableTransport = unwrap(transport) as unknown as {
19871991
_authProvider?: {
19881992
tokens: () => Promise<
19891993
{ access_token: string; expires_in?: number } | undefined
@@ -2032,7 +2036,7 @@ describe('mcp-client', () => {
20322036
MOCK_CONTEXT,
20332037
);
20342038

2035-
const testableTransport = transport as unknown as {
2039+
const testableTransport = unwrap(transport) as unknown as {
20362040
_authProvider?: {
20372041
tokens: () => Promise<{ access_token: string } | undefined>;
20382042
};
@@ -2075,7 +2079,7 @@ describe('mcp-client', () => {
20752079
MOCK_CONTEXT,
20762080
);
20772081

2078-
const testableTransport = transport as unknown as {
2082+
const testableTransport = unwrap(transport) as unknown as {
20792083
_authProvider?: {
20802084
tokens: () => Promise<{ access_token: string } | undefined>;
20812085
};
@@ -2128,7 +2132,7 @@ describe('mcp-client', () => {
21282132
MOCK_CONTEXT,
21292133
);
21302134

2131-
const testableTransport = transport as unknown as {
2135+
const testableTransport = unwrap(transport) as unknown as {
21322136
_authProvider?: {
21332137
tokens: () => Promise<
21342138
{ access_token: string; expires_in?: number } | undefined
@@ -2167,7 +2171,7 @@ describe('mcp-client', () => {
21672171
);
21682172

21692173
const wrappedFetch = (
2170-
transport as unknown as {
2174+
unwrap(transport) as unknown as {
21712175
_fetch: (
21722176
url: URL | string,
21732177
init?: RequestInit,
@@ -2209,7 +2213,7 @@ describe('mcp-client', () => {
22092213

22102214
// For SSEClientTransport, the fetch is private or passed to the SDK.
22112215
// We can check if it creates the transport successfully.
2212-
expect(transport).toBeInstanceOf(SSEClientTransport);
2216+
expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
22132217
} finally {
22142218
vi.unstubAllEnvs();
22152219
vi.unstubAllGlobals();
@@ -2227,8 +2231,8 @@ describe('mcp-client', () => {
22272231
false,
22282232
MOCK_CONTEXT,
22292233
);
2230-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2231-
expect(transport).toMatchObject({
2234+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2235+
expect(unwrap(transport)).toMatchObject({
22322236
_url: new URL('http://test-server'),
22332237
_requestInit: { headers: {} },
22342238
});
@@ -2245,8 +2249,8 @@ describe('mcp-client', () => {
22452249
MOCK_CONTEXT,
22462250
);
22472251

2248-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2249-
expect(transport).toMatchObject({
2252+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2253+
expect(unwrap(transport)).toMatchObject({
22502254
_url: new URL('http://test-server'),
22512255
_requestInit: {
22522256
headers: { Authorization: 'derp' },
@@ -2265,8 +2269,8 @@ describe('mcp-client', () => {
22652269
MOCK_CONTEXT,
22662270
);
22672271

2268-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2269-
expect(transport).toMatchObject({
2272+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2273+
expect(unwrap(transport)).toMatchObject({
22702274
_url: new URL('http://test-server'),
22712275
_requestInit: { headers: {} },
22722276
});
@@ -2283,8 +2287,8 @@ describe('mcp-client', () => {
22832287
MOCK_CONTEXT,
22842288
);
22852289

2286-
expect(transport).toBeInstanceOf(SSEClientTransport);
2287-
expect(transport).toMatchObject({
2290+
expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
2291+
expect(unwrap(transport)).toMatchObject({
22882292
_url: new URL('http://test-server'),
22892293
_requestInit: { headers: {} },
22902294
});
@@ -2300,8 +2304,8 @@ describe('mcp-client', () => {
23002304
MOCK_CONTEXT,
23012305
);
23022306

2303-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2304-
expect(transport).toMatchObject({
2307+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2308+
expect(unwrap(transport)).toMatchObject({
23052309
_url: new URL('http://test-server'),
23062310
_requestInit: { headers: {} },
23072311
});
@@ -2319,8 +2323,8 @@ describe('mcp-client', () => {
23192323
MOCK_CONTEXT,
23202324
);
23212325

2322-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2323-
expect(transport).toMatchObject({
2326+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2327+
expect(unwrap(transport)).toMatchObject({
23242328
_url: new URL('http://test-server'),
23252329
_requestInit: {
23262330
headers: { Authorization: 'Bearer token' },
@@ -2340,8 +2344,8 @@ describe('mcp-client', () => {
23402344
MOCK_CONTEXT,
23412345
);
23422346

2343-
expect(transport).toBeInstanceOf(SSEClientTransport);
2344-
expect(transport).toMatchObject({
2347+
expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
2348+
expect(unwrap(transport)).toMatchObject({
23452349
_url: new URL('http://test-server'),
23462350
_requestInit: {
23472351
headers: { 'X-API-Key': 'key123' },
@@ -2361,8 +2365,8 @@ describe('mcp-client', () => {
23612365
);
23622366

23632367
// httpUrl should take priority and create HTTP transport
2364-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2365-
expect(transport).toMatchObject({
2368+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2369+
expect(unwrap(transport)).toMatchObject({
23662370
_url: new URL('http://test-server-http'),
23672371
_requestInit: { headers: {} },
23682372
});
@@ -2587,8 +2591,10 @@ describe('mcp-client', () => {
25872591
MOCK_CONTEXT,
25882592
);
25892593

2590-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2591-
const testableTransport = transport as unknown as TestableTransport;
2594+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2595+
const testableTransport = unwrap(
2596+
transport,
2597+
) as unknown as TestableTransport;
25922598
const authProvider = testableTransport._authProvider;
25932599
expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
25942600
const googUserProject =
@@ -2618,9 +2624,11 @@ describe('mcp-client', () => {
26182624
MOCK_CONTEXT,
26192625
);
26202626

2621-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2627+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
26222628
expect(mockGetRequestHeaders).toHaveBeenCalled();
2623-
const testableTransport = transport as unknown as TestableTransport;
2629+
const testableTransport = unwrap(
2630+
transport,
2631+
) as unknown as TestableTransport;
26242632
const headers = testableTransport._requestInit?.headers;
26252633
expect(headers?.['X-Goog-User-Project']).toBe('provider-project');
26262634
});
@@ -2650,8 +2658,10 @@ describe('mcp-client', () => {
26502658
MOCK_CONTEXT,
26512659
);
26522660

2653-
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
2654-
const testableTransport = transport as unknown as TestableTransport;
2661+
expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
2662+
const testableTransport = unwrap(
2663+
transport,
2664+
) as unknown as TestableTransport;
26552665
const headers = testableTransport._requestInit?.headers;
26562666
expect(headers?.['X-Goog-User-Project']).toBe('provider-project');
26572667
});
@@ -2671,8 +2681,10 @@ describe('mcp-client', () => {
26712681
MOCK_CONTEXT,
26722682
);
26732683

2674-
expect(transport).toBeInstanceOf(SSEClientTransport);
2675-
const testableTransport = transport as unknown as TestableTransport;
2684+
expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
2685+
const testableTransport = unwrap(
2686+
transport,
2687+
) as unknown as TestableTransport;
26762688
const authProvider = testableTransport._authProvider;
26772689
expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
26782690
});
@@ -2842,7 +2854,7 @@ describe('connectToMcpServer with OAuth', () => {
28422854
let capturedTransport: TestableTransport | undefined;
28432855
vi.mocked(mockedClient.connect).mockImplementationOnce(
28442856
async (transport) => {
2845-
capturedTransport = transport as unknown as TestableTransport;
2857+
capturedTransport = unwrap(transport) as unknown as TestableTransport;
28462858
return Promise.resolve();
28472859
},
28482860
);
@@ -2860,8 +2872,8 @@ describe('connectToMcpServer with OAuth', () => {
28602872
expect(mockedClient.connect).toHaveBeenCalledTimes(2);
28612873
expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
28622874

2863-
const authHeader = (capturedTransport as TestableTransport)._requestInit
2864-
?.headers?.['Authorization'];
2875+
const authHeader = (unwrap(capturedTransport) as TestableTransport)
2876+
._requestInit?.headers?.['Authorization'];
28652877
expect(authHeader).toBe('Bearer test-access-token');
28662878
});
28672879

@@ -2887,7 +2899,7 @@ describe('connectToMcpServer with OAuth', () => {
28872899
let capturedTransport: TestableTransport | undefined;
28882900
vi.mocked(mockedClient.connect).mockImplementationOnce(
28892901
async (transport) => {
2890-
capturedTransport = transport as unknown as TestableTransport;
2902+
capturedTransport = unwrap(transport) as unknown as TestableTransport;
28912903
return Promise.resolve();
28922904
},
28932905
);
@@ -2906,8 +2918,8 @@ describe('connectToMcpServer with OAuth', () => {
29062918
expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
29072919
expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(serverUrl);
29082920

2909-
const authHeader = (capturedTransport as TestableTransport)._requestInit
2910-
?.headers?.['Authorization'];
2921+
const authHeader = (unwrap(capturedTransport) as TestableTransport)
2922+
._requestInit?.headers?.['Authorization'];
29112923
expect(authHeader).toBe('Bearer test-access-token-from-discovery');
29122924
});
29132925

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

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ import {
4949
import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
5050
import { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js';
5151
import { DiscoveredMCPTool } from './mcp-tool.js';
52-
import { XcodeMcpBridgeFixTransport } from './xcode-mcp-fix-transport.js';
52+
import { McpComplianceTransport } from './mcp-compliance-transport.js';
5353

5454
import type { CallableTool, FunctionCall, Part, Tool } from '@google/genai';
5555
import { basename } from 'node:path';
@@ -1055,7 +1055,7 @@ async function createTransportWithOAuth(
10551055
mcpServerConfig: MCPServerConfig,
10561056
accessToken: string,
10571057
cliConfig: McpContext,
1058-
): Promise<StreamableHTTPClientTransport | SSEClientTransport | null> {
1058+
): Promise<Transport | null> {
10591059
try {
10601060
const headers: Record<string, string> = {
10611061
Authorization: `Bearer ${accessToken}`,
@@ -1070,7 +1070,12 @@ async function createTransportWithOAuth(
10701070
),
10711071
};
10721072

1073-
return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
1073+
const transport = createUrlTransport(
1074+
mcpServerName,
1075+
mcpServerConfig,
1076+
transportOptions,
1077+
);
1078+
return transport ? new McpComplianceTransport(transport) : null;
10741079
} catch (error) {
10751080
cliConfig.emitMcpDiagnostic(
10761081
'error',
@@ -2318,7 +2323,9 @@ export async function createTransport(
23182323
authProvider,
23192324
};
23202325

2321-
return createUrlTransport(mcpServerName, mcpServerConfig, transportOptions);
2326+
return new McpComplianceTransport(
2327+
createUrlTransport(mcpServerName, mcpServerConfig, transportOptions),
2328+
);
23222329
}
23232330

23242331
if (mcpServerConfig.command) {
@@ -2354,32 +2361,23 @@ export async function createTransport(
23542361
}
23552362
}
23562363

2357-
let transport: Transport = new StdioClientTransport({
2358-
command: mcpServerConfig.command,
2359-
args: mcpServerConfig.args || [],
2360-
env: finalEnv,
2361-
cwd: mcpServerConfig.cwd,
2362-
stderr: 'pipe',
2363-
});
2364-
2365-
// Fix for Xcode 26.3 mcpbridge non-compliant responses
2366-
// It returns JSON in `content` instead of `structuredContent`
2367-
if (
2368-
mcpServerConfig.command === 'xcrun' &&
2369-
mcpServerConfig.args?.includes('mcpbridge')
2370-
) {
2371-
transport = new XcodeMcpBridgeFixTransport(transport);
2372-
}
2364+
const transport: Transport = new McpComplianceTransport(
2365+
new StdioClientTransport({
2366+
command: mcpServerConfig.command,
2367+
args: mcpServerConfig.args || [],
2368+
env: finalEnv,
2369+
cwd: mcpServerConfig.cwd,
2370+
stderr: 'pipe',
2371+
}),
2372+
);
23732373

23742374
if (debugMode) {
2375-
// The `XcodeMcpBridgeFixTransport` wrapper hides the underlying `StdioClientTransport`,
2375+
// The `McpComplianceTransport` wrapper hides the underlying `StdioClientTransport`,
23762376
// which exposes `stderr` for debug logging. We need to unwrap it to attach the listener.
23772377

2378-
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
23792378
const underlyingTransport =
2380-
transport instanceof XcodeMcpBridgeFixTransport
2381-
? // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion
2382-
(transport as any).transport
2379+
transport instanceof McpComplianceTransport
2380+
? transport.transport
23832381
: transport;
23842382

23852383
if (

0 commit comments

Comments
 (0)