@machina.ai/cell-cli-core 1.41.1-rc2 → 1.45.1-rc2

package/dist/src/tools/mcp-client.test.js

@@ -18,9 +18,11 @@ import { ErrorCode, McpError, PromptListChangedNotificationSchema, ResourceListC
 import { GoogleAuth } from 'google-auth-library';
 import { WorkspaceContext } from '../utils/workspaceContext.js';
 import { connectToMcpServer, createTransport, hasNetworkTransport, isEnabled, McpClient, populateMcpServerCommand, discoverPrompts, } from './mcp-client.js';
+import { McpComplianceTransport } from './mcp-compliance-transport.js';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
+import { cleanupTmpDir } from '@google/gemini-cli-test-utils';
 import { coreEvents } from '../utils/events.js';
 const EMPTY_CONFIG = {
     enableEnvironmentVariableRedaction: true,
@@ -34,22 +36,27 @@ const MOCK_CONTEXT_DEFAULT = {
     isTrustedFolder: vi.fn().mockReturnValue(true),
 };
 let MOCK_CONTEXT = MOCK_CONTEXT_DEFAULT;
-const { mockAuthProviderKeycloak } = vi.hoisted(() => ({
-    mockAuthProviderKeycloak: vi.fn().mockImplementation(() => ({
-        getClient: vi.fn().mockResolvedValue({
-            getAccessToken: vi.fn().mockResolvedValue({ token: 'mock-token' }),
-        }),
-    })),
-}));
+const unwrap = (t) => t instanceof McpComplianceTransport ? t.transport : t;
 vi.mock('@modelcontextprotocol/sdk/client/stdio.js');
 vi.mock('@modelcontextprotocol/sdk/client/index.js');
 vi.mock('@google/genai');
+vi.mock('undici', () => ({
+    EnvHttpProxyAgent: vi.fn(),
+    fetch: vi.fn(),
+    setGlobalDispatcher: vi.fn(),
+    Agent: vi.fn(),
+}));
 vi.mock('../mcp/oauth-provider.js');
 vi.mock('../mcp/oauth-token-storage.js');
 vi.mock('../mcp/oauth-utils.js');
-vi.mock('@machina.ai/auth-keycloak', () => ({
-    AuthProviderKeycloak: mockAuthProviderKeycloak,
-}));
+vi.mock('@machina.ai/auth-keycloak', () => {
+    const mockAuth = vi.fn().mockImplementation(() => ({
+        getToken: vi.fn().mockResolvedValue('mock-token'),
+    }));
+    return {
+        AuthProviderKeycloak: mockAuth,
+    };
+});
 vi.stubGlobal('fetch', vi.fn());
 vi.mock('../utils/events.js', () => ({
     coreEvents: {
@@ -72,9 +79,11 @@ describe('mcp-client', () => {
         testWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-agent-test-'));
         workspaceContext = new WorkspaceContext(testWorkspace);
     });
-    afterEach(() => {
-        vi.restoreAllMocks();
+    afterEach(async () => {
         vi.useRealTimers();
+        await cleanupTmpDir(testWorkspace);
+        workspaceContext = null;
+        vi.restoreAllMocks();
     });
     describe('McpClient', () => {
         it('should discover tools', async () => {
@@ -613,6 +622,89 @@ describe('mcp-client', () => {
                 },
             });
         });
+        it('should transform nullable array schemas and preserve properties during discovery', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'nullableTool',
+                            description: 'Tool with nullable array',
+                            inputSchema: {
+                                type: 'object',
+                                properties: {
+                                    tags: {
+                                        type: ['array', 'null'],
+                                        items: { type: 'string' },
+                                    },
+                                },
+                                $defs: {
+                                    SomeType: { type: 'string' },
+                                },
+                            },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({
+                    prompts: [],
+                }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getToolsByServer: vi.fn().mockReturnValue([]),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                getPromptsByServer: vi.fn().mockReturnValue([]),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                getResourcesByServer: vi.fn().mockReturnValue([]),
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
+            const client = new McpClient('test-server', {
+                command: 'test-command',
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            await client.connect();
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+            const registeredTool = vi.mocked(mockedToolRegistry.registerTool).mock
+                .calls[0][0];
+            expect(registeredTool.schema.parametersJsonSchema).toEqual({
+                type: 'object',
+                properties: {
+                    tags: {
+                        type: 'array',
+                        nullable: true,
+                        items: { type: 'string' },
+                    },
+                    wait_for_previous: {
+                        type: 'boolean',
+                        description: 'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
+                    },
+                },
+                $defs: {
+                    SomeType: { type: 'string' },
+                },
+            });
+        });
         it('should discover resources when a server only exposes resources', async () => {
             const mockedClient = {
                 connect: vi.fn(),
@@ -1291,29 +1383,222 @@ describe('mcp-client', () => {
         });
     });
     describe('createTransport', () => {
+        it('should create an HTTP transport that respects NO_PROXY', async () => {
+            const { createTransport } = await import('./mcp-client.js');
+            const { EnvHttpProxyAgent } = await import('undici');
+            const noProxyValue = 'localhost,127.0.0.1';
+            vi.stubEnv('NO_PROXY', noProxyValue);
+            await createTransport('test-server', {
+                url: 'http://test-server',
+                type: 'http',
+            }, false, MOCK_CONTEXT);
+            expect(EnvHttpProxyAgent).toHaveBeenCalledWith(expect.objectContaining({
+                noProxy: noProxyValue,
+            }));
+        });
         describe('should connect via httpUrl', () => {
-            it('without headers', async () => {
+            it('uses MCP SDK authProvider token() path for oauth-enabled servers', async () => {
+                const mockGetValidTokenWithMetadata = vi.fn().mockResolvedValue({
+                    accessToken: 'fresh-token',
+                    tokenType: 'Bearer',
+                    expiresAt: Date.now() + 10 * 60 * 1000,
+                });
+                vi.mocked(MCPOAuthProvider).mockReturnValue({
+                    getValidTokenWithMetadata: mockGetValidTokenWithMetadata,
+                });
+                vi.mocked(MCPOAuthTokenStorage).mockReturnValue({
+                    getCredentials: vi.fn().mockResolvedValue({
+                        clientId: 'cid',
+                        token: {
+                            accessToken: 'fresh-token',
+                            tokenType: 'Bearer',
+                            expiresAt: Date.now() + 10 * 60 * 1000,
+                        },
+                    }),
+                });
                 const transport = await createTransport('test-server', {
-                    httpUrl: 'http://test-server',
+                    url: 'http://test-server',
+                    type: 'http',
+                    oauth: { enabled: true },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
-                    _url: new URL('http://test-server'),
-                    _requestInit: { headers: {} },
+                const testableTransport = unwrap(transport);
+                expect(testableTransport._authProvider).toBeDefined();
+                const tokens = await testableTransport._authProvider.tokens();
+                expect(tokens?.access_token).toBe('fresh-token');
+            });
+            it('uses storage-backed expiry instead of long fallback cache for dynamic authProvider', async () => {
+                const now = Date.now();
+                const soonExpiry = now + 10 * 60 * 1000; // 10 minutes
+                const mockGetValidTokenWithMetadata = vi.fn().mockResolvedValue({
+                    accessToken: 'fresh-token',
+                    tokenType: 'Bearer',
+                    expiresAt: soonExpiry,
+                });
+                const mockGetCredentials = vi.fn().mockImplementation(async () => ({
+                    clientId: 'cid',
+                    token: {
+                        accessToken: 'fresh-token',
+                        tokenType: 'Bearer',
+                        expiresAt: soonExpiry,
+                    },
+                }));
+                vi.mocked(MCPOAuthProvider).mockReturnValue({
+                    getValidTokenWithMetadata: mockGetValidTokenWithMetadata,
+                });
+                vi.mocked(MCPOAuthTokenStorage).mockReturnValue({
+                    getCredentials: mockGetCredentials,
                 });
+                const transport = await createTransport('test-server', {
+                    url: 'http://test-server',
+                    type: 'http',
+                }, false, MOCK_CONTEXT);
+                const testableTransport = unwrap(transport);
+                expect(testableTransport._authProvider).toBeDefined();
+                const tokens = await testableTransport._authProvider.tokens();
+                expect(tokens?.access_token).toBe('fresh-token');
+                expect(tokens?.expires_in).toBeDefined();
+                expect((tokens?.expires_in ?? 0) <= 10 * 60).toBe(true);
+                expect(mockGetValidTokenWithMetadata).toHaveBeenCalledTimes(1);
+                expect(mockGetCredentials).toHaveBeenCalledTimes(1);
             });
-            it('with headers', async () => {
+            it('uses dynamic authProvider when stored OAuth token exists', async () => {
+                const mockGetValidTokenWithMetadata = vi.fn().mockResolvedValue({
+                    accessToken: 'stored-fresh-token',
+                    tokenType: 'Bearer',
+                    expiresAt: Date.now() + 10 * 60 * 1000,
+                });
+                vi.mocked(MCPOAuthProvider).mockReturnValue({
+                    getValidTokenWithMetadata: mockGetValidTokenWithMetadata,
+                });
+                vi.mocked(MCPOAuthTokenStorage).mockReturnValue({
+                    getCredentials: vi.fn().mockResolvedValue({
+                        clientId: 'cid',
+                        token: {
+                            accessToken: 'stored-fresh-token',
+                            tokenType: 'Bearer',
+                            expiresAt: Date.now() + 10 * 60 * 1000,
+                        },
+                    }),
+                });
                 const transport = await createTransport('test-server', {
-                    httpUrl: 'http://test-server',
-                    headers: { Authorization: 'derp' },
+                    url: 'http://test-server',
+                    type: 'http',
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
-                    _url: new URL('http://test-server'),
-                    _requestInit: {
-                        headers: { Authorization: 'derp' },
+                const testableTransport = unwrap(transport);
+                expect(testableTransport._authProvider).toBeDefined();
+                const tokens = await testableTransport._authProvider.tokens();
+                expect(tokens?.access_token).toBe('stored-fresh-token');
+            });
+            it('caches OAuth tokens in dynamic authProvider and avoids repeated lookups', async () => {
+                const mockGetValidTokenWithMetadata = vi.fn().mockResolvedValue({
+                    accessToken: 'cached-token',
+                    tokenType: 'Bearer',
+                    expiresAt: Date.now() + 10 * 60 * 1000,
+                });
+                const mockGetCredentials = vi.fn().mockResolvedValue({
+                    clientId: 'cid',
+                    token: {
+                        accessToken: 'cached-token',
+                        tokenType: 'Bearer',
+                        expiresAt: Date.now() + 10 * 60 * 1000,
+                    },
+                });
+                vi.mocked(MCPOAuthProvider).mockReturnValue({
+                    getValidTokenWithMetadata: mockGetValidTokenWithMetadata,
+                });
+                vi.mocked(MCPOAuthTokenStorage).mockReturnValue({
+                    getCredentials: mockGetCredentials,
+                });
+                const transport = await createTransport('test-server', {
+                    url: 'http://test-server',
+                    type: 'http',
+                }, false, MOCK_CONTEXT);
+                const testableTransport = unwrap(transport);
+                expect(testableTransport._authProvider).toBeDefined();
+                const t1 = await testableTransport._authProvider.tokens();
+                const t2 = await testableTransport._authProvider.tokens();
+                expect(t1?.access_token).toBe('cached-token');
+                expect(t2?.access_token).toBe('cached-token');
+                // one call from createTransport fallback detection + one call in first tokens();
+                // second tokens() should come from in-memory cache
+                expect(mockGetCredentials).toHaveBeenCalledTimes(1);
+                expect(mockGetValidTokenWithMetadata).toHaveBeenCalledTimes(1);
+            });
+            it('does not long-cache token when metadata has no expiresAt', async () => {
+                const mockGetValidTokenWithMetadata = vi.fn().mockResolvedValue({
+                    accessToken: 'no-exp-token',
+                    tokenType: 'Bearer',
+                    // expiresAt intentionally omitted
+                });
+                const mockGetCredentials = vi.fn().mockResolvedValue({
+                    clientId: 'cid',
+                    token: {
+                        accessToken: 'no-exp-token',
+                        tokenType: 'Bearer',
+                        // expiresAt intentionally omitted
                     },
                 });
+                vi.mocked(MCPOAuthProvider).mockReturnValue({
+                    getValidTokenWithMetadata: mockGetValidTokenWithMetadata,
+                });
+                vi.mocked(MCPOAuthTokenStorage).mockReturnValue({
+                    getCredentials: mockGetCredentials,
+                });
+                const transport = await createTransport('test-server', {
+                    url: 'http://test-server',
+                    type: 'http',
+                }, false, MOCK_CONTEXT);
+                const testableTransport = unwrap(transport);
+                expect(testableTransport._authProvider).toBeDefined();
+                const t1 = await testableTransport._authProvider.tokens();
+                const t2 = await testableTransport._authProvider.tokens();
+                expect(t1?.access_token).toBe('no-exp-token');
+                expect(t2?.access_token).toBe('no-exp-token');
+                expect(t1?.expires_in).toBeUndefined();
+                expect(t2?.expires_in).toBeUndefined();
+                // no-expiry tokens should not be long-cached in memory
+                expect(mockGetValidTokenWithMetadata).toHaveBeenCalledTimes(2);
+            });
+            it('wraps fetch to convert GET 404 to 405 for POST-only servers (e.g. n8n)', async () => {
+                const mockFetch = vi
+                    .fn()
+                    .mockResolvedValue(new Response(null, { status: 404, statusText: 'Not Found' }));
+                vi.stubGlobal('fetch', mockFetch);
+                try {
+                    const transport = await createTransport('test-server', { httpUrl: 'http://test-server' }, false, MOCK_CONTEXT);
+                    const wrappedFetch = unwrap(transport)._fetch;
+                    // GET 404 → 405: server doesn't support optional SSE GET stream
+                    const getRes = await wrappedFetch('http://test-server', {
+                        method: 'GET',
+                    });
+                    expect(getRes.status).toBe(405);
+                    expect(getRes.statusText).toBe('Method Not Allowed');
+                    // POST 404 → unchanged: real "not found" errors must still propagate
+                    const postRes = await wrappedFetch('http://test-server', {
+                        method: 'POST',
+                    });
+                    expect(postRes.status).toBe(404);
+                }
+                finally {
+                    vi.unstubAllGlobals();
+                }
+            });
+            it('respects NO_PROXY for network transports', async () => {
+                const mockFetch = vi
+                    .fn()
+                    .mockResolvedValue(new Response('OK', { status: 200 }));
+                vi.stubGlobal('fetch', mockFetch);
+                vi.stubEnv('NO_PROXY', 'localhost');
+                try {
+                    const transport = await createTransport('test-server', { url: 'http://localhost/sse', type: 'sse' }, false, MOCK_CONTEXT);
+                    // For SSEClientTransport, the fetch is private or passed to the SDK.
+                    // We can check if it creates the transport successfully.
+                    expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
+                }
+                finally {
+                    vi.unstubAllEnvs();
+                    vi.unstubAllGlobals();
+                }
             });
         });
         describe('should connect via url', () => {
@@ -1321,8 +1606,8 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: { headers: {} },
                 });
@@ -1332,8 +1617,8 @@ describe('mcp-client', () => {
                     url: 'http://test-server',
                     headers: { Authorization: 'derp' },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: {
                         headers: { Authorization: 'derp' },
@@ -1345,8 +1630,8 @@ describe('mcp-client', () => {
                     url: 'http://test-server',
                     type: 'http',
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: { headers: {} },
                 });
@@ -1356,8 +1641,8 @@ describe('mcp-client', () => {
                     url: 'http://test-server',
                     type: 'sse',
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(SSEClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: { headers: {} },
                 });
@@ -1366,8 +1651,8 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: { headers: {} },
                 });
@@ -1378,8 +1663,8 @@ describe('mcp-client', () => {
                     type: 'http',
                     headers: { Authorization: 'Bearer token' },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: {
                         headers: { Authorization: 'Bearer token' },
@@ -1392,8 +1677,8 @@ describe('mcp-client', () => {
                     type: 'sse',
                     headers: { 'X-API-Key': 'key123' },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(SSEClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server'),
                     _requestInit: {
                         headers: { 'X-API-Key': 'key123' },
@@ -1406,8 +1691,8 @@ describe('mcp-client', () => {
                     url: 'http://test-server-url',
                 }, false, MOCK_CONTEXT);
                 // httpUrl should take priority and create HTTP transport
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toMatchObject({
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toMatchObject({
                     _url: new URL('http://test-server-http'),
                     _requestInit: { headers: {} },
                 });
@@ -1576,8 +1861,8 @@ describe('mcp-client', () => {
                         'X-Goog-User-Project': 'myproject',
                     },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                const testableTransport = transport;
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                const testableTransport = unwrap(transport);
                 const authProvider = testableTransport._authProvider;
                 expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
                 const googUserProject = testableTransport._requestInit?.headers?.['X-Goog-User-Project'];
@@ -1595,9 +1880,9 @@ describe('mcp-client', () => {
                         scopes: ['scope1'],
                     },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
                 expect(mockGetRequestHeaders).toHaveBeenCalled();
-                const testableTransport = transport;
+                const testableTransport = unwrap(transport);
                 const headers = testableTransport._requestInit?.headers;
                 expect(headers?.['X-Goog-User-Project']).toBe('provider-project');
             });
@@ -1616,8 +1901,8 @@ describe('mcp-client', () => {
                         'X-Goog-User-Project': 'config-project',
                     },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                const testableTransport = transport;
+                expect(unwrap(transport)).toBeInstanceOf(StreamableHTTPClientTransport);
+                const testableTransport = unwrap(transport);
                 const headers = testableTransport._requestInit?.headers;
                 expect(headers?.['X-Goog-User-Project']).toBe('provider-project');
             });
@@ -1630,8 +1915,8 @@ describe('mcp-client', () => {
                         scopes: ['scope1'],
                     },
                 }, false, MOCK_CONTEXT);
-                expect(transport).toBeInstanceOf(SSEClientTransport);
-                const testableTransport = transport;
+                expect(unwrap(transport)).toBeInstanceOf(SSEClientTransport);
+                const testableTransport = unwrap(transport);
                 const authProvider = testableTransport._authProvider;
                 expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
             });
@@ -1739,7 +2024,10 @@ describe('connectToMcpServer with OAuth', () => {
         };
         vi.mocked(MCPOAuthProvider).mockReturnValue(mockAuthProvider);
     });
-    afterEach(() => {
+    afterEach(async () => {
+        vi.useRealTimers();
+        await cleanupTmpDir(testWorkspace);
+        workspaceContext = null;
         vi.clearAllMocks();
     });
     it('should handle automatic OAuth flow on 401 with www-authenticate header', async () => {
@@ -1756,15 +2044,15 @@ describe('connectToMcpServer with OAuth', () => {
         // We need this to be typed to dig into its private state.
         let capturedTransport;
         vi.mocked(mockedClient.connect).mockImplementationOnce(async (transport) => {
-            capturedTransport = transport;
+            capturedTransport = unwrap(transport);
             return Promise.resolve();
         });
         const client = await connectToMcpServer('0.0.1', 'test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, MOCK_CONTEXT);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
-        const authHeader = capturedTransport._requestInit
-            ?.headers?.['Authorization'];
+        const authHeader = unwrap(capturedTransport)
+            ._requestInit?.headers?.['Authorization'];
         expect(authHeader).toBe('Bearer test-access-token');
     });
     it('should discover oauth config if not in www-authenticate header', async () => {
@@ -1781,7 +2069,7 @@ describe('connectToMcpServer with OAuth', () => {
         // We need this to be typed to dig into its private state.
         let capturedTransport;
         vi.mocked(mockedClient.connect).mockImplementationOnce(async (transport) => {
-            capturedTransport = transport;
+            capturedTransport = unwrap(transport);
             return Promise.resolve();
         });
         const client = await connectToMcpServer('0.0.1', 'test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, MOCK_CONTEXT);
@@ -1789,8 +2077,8 @@ describe('connectToMcpServer with OAuth', () => {
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
         expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(serverUrl);
-        const authHeader = capturedTransport._requestInit
-            ?.headers?.['Authorization'];
+        const authHeader = unwrap(capturedTransport)
+            ._requestInit?.headers?.['Authorization'];
         expect(authHeader).toBe('Bearer test-access-token-from-discovery');
     });
     it('should use discoverOAuthFromWWWAuthenticate when it succeeds and skip discoverOAuthConfig', async () => {
@@ -1854,7 +2142,10 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => {
         vi.spyOn(console, 'warn').mockImplementation(() => { });
         vi.spyOn(console, 'error').mockImplementation(() => { });
     });
-    afterEach(() => {
+    afterEach(async () => {
+        vi.useRealTimers();
+        await cleanupTmpDir(testWorkspace);
+        workspaceContext = null;
         vi.clearAllMocks();
     });
     it('should NOT trigger fallback when type="http" is explicit', async () => {
@@ -1942,7 +2233,10 @@ describe('connectToMcpServer - OAuth with transport fallback', () => {
             scopes: ['test-scope'],
         });
     });
-    afterEach(() => {
+    afterEach(async () => {
+        vi.useRealTimers();
+        await cleanupTmpDir(testWorkspace);
+        workspaceContext = null;
         vi.clearAllMocks();
         vi.unstubAllGlobals();
     });