Extend MCP tests (#4562)

packages/gitbook/tests/mcp-utils.ts

@@ -0,0 +1,191 @@
+import { Client } from '@modelcontextprotocol/sdk/client';
+import type { OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js';
+import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import type {
+    OAuthClientInformation,
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthTokens,
+} from '@modelcontextprotocol/sdk/shared/auth.js';
+import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js';
+
+import { getContentTestURL } from './utils';
+
+/**
+ * Client ID returned by the stubbed registration endpoint.
+ */
+export const STUBBED_OAUTH_CLIENT_ID = 'gitbook-open-tests-client-id';
+
+/**
+ * Redirect URI the test client registers. It is a dummy address that nothing serves: the flow is
+ * asserted at the point where a real client would open a browser, so the authorization code never
+ * comes back and this is never called.
+ */
+const TEST_REDIRECT_URI = 'https://auth.acme.org/mcp/callback';
+
+/**
+ * Everything an MCP client observes while discovering how to authenticate against a site, recorded
+ * so tests can assert on each document of the chain.
+ */
+export type RecordedOAuthDiscovery = {
+    /** `WWW-Authenticate` header of the 401 challenge that started the flow. */
+    challenge?: string;
+    /** `resource_metadata` URL advertised by the challenge. */
+    resourceMetadataURL?: string;
+    /** Protected Resource Metadata document, RFC 9728. */
+    protectedResourceMetadata?: Record<string, unknown>;
+    /** Authorization Server Metadata document, RFC 8414. */
+    authorizationServerMetadata?: Record<string, unknown>;
+    /** Endpoint the client posted its registration to. */
+    registrationURL?: string;
+    /** Client ID the client ended up using. */
+    clientId?: string;
+    /** Authorization URL the client would have opened in the visitor's browser. */
+    authorizationURL?: URL;
+};
+
+/**
+ * An MCP client playing the OAuth client role without a browser.
+ *
+ * It walks the real discovery chain — 401 challenge, protected resource metadata, authorization
+ * server metadata — recording each document, and stops where a real client would hand the visitor
+ * over to their browser to log in and consent.
+ *
+ * Dynamic client registration is stubbed: what GitBook Open owns is advertising the endpoint, while
+ * the registration itself and the rest of the flow belong to the OAuth server and are covered by
+ * its own tests. Registering for real would also create a client record on every run.
+ */
+export class McpOAuthTestClient implements OAuthClientProvider {
+    readonly redirectUrl = TEST_REDIRECT_URI;
+    readonly discovery: RecordedOAuthDiscovery = {};
+
+    private clientInformationFull: OAuthClientInformationFull | undefined;
+    private savedCodeVerifier: string | undefined;
+
+    /**
+     * Connect to an MCP endpoint. On a protected endpoint this rejects with `UnauthorizedError`
+     * once the discovery chain has been walked and recorded.
+     */
+    async connect(url: string) {
+        const client = new Client({ name: 'gitbook-open-tests', version: '1.0.0' });
+        await client.connect(
+            new StreamableHTTPClientTransport(new URL(url), {
+                authProvider: this,
+                fetch: this.fetchAndRecord,
+            })
+        );
+        return client;
+    }
+
+    get clientMetadata(): OAuthClientMetadata {
+        return {
+            client_name: 'gitbook-open-tests',
+            redirect_uris: [TEST_REDIRECT_URI],
+            grant_types: ['authorization_code', 'refresh_token'],
+            response_types: ['code'],
+            token_endpoint_auth_method: 'none',
+        };
+    }
+
+    /** Returning nothing on the first call is what sends the SDK down the registration path. */
+    clientInformation(): OAuthClientInformation | undefined {
+        return this.clientInformationFull;
+    }
+
+    saveClientInformation(clientInformation: OAuthClientInformationFull) {
+        this.clientInformationFull = clientInformation;
+        this.discovery.clientId = clientInformation.client_id;
+    }
+
+    /** No token and no refresh token, so the SDK always starts a fresh authorization. */
+    tokens(): OAuthTokens | undefined {
+        return undefined;
+    }
+
+    saveTokens() {}
+
+    saveCodeVerifier(codeVerifier: string) {
+        this.savedCodeVerifier = codeVerifier;
+    }
+
+    codeVerifier() {
+        if (!this.savedCodeVerifier) {
+            throw new Error('No code verifier saved');
+        }
+        return this.savedCodeVerifier;
+    }
+
+    /** Where a real client opens a browser. */
+    redirectToAuthorization(authorizationUrl: URL) {
+        this.discovery.authorizationURL = authorizationUrl;
+    }
+
+    /**
+     * Fetch used for every request the client makes, recording the discovery documents on the way
+     * through and standing in for the authorization server's registration endpoint.
+     */
+    private readonly fetchAndRecord: FetchLike = async (input, init) => {
+        const url = new URL(input.toString());
+
+        const registrationEndpoint =
+            this.discovery.authorizationServerMetadata?.registration_endpoint;
+        if (registrationEndpoint && url.href === registrationEndpoint) {
+            this.discovery.registrationURL = url.href;
+            return Response.json({
+                ...JSON.parse(String(init?.body)),
+                client_id: STUBBED_OAUTH_CLIENT_ID,
+            });
+        }
+
+        const response = await fetch(input, init);
+
+        const challenge = response.headers.get('WWW-Authenticate');
+        if (response.status === 401 && challenge) {
+            this.discovery.challenge = challenge;
+            this.discovery.resourceMetadataURL = challenge.match(
+                /resource_metadata="([^"]*)"/
+            )?.[1];
+        }
+
+        if (response.ok && url.pathname.includes('/.well-known/')) {
+            const document = (await response.clone().json()) as Record<string, unknown>;
+            if (url.pathname.includes('/.well-known/oauth-protected-resource')) {
+                this.discovery.protectedResourceMetadata = document;
+            } else if (!this.discovery.authorizationServerMetadata) {
+                // The SDK tries several well-known locations and uses the first that answers.
+                this.discovery.authorizationServerMetadata = document;
+            }
+        }
+
+        return response;
+    };
+}
+
+/**
+ * Connect an MCP client to an endpoint, optionally presenting a visitor token as the OAuth bearer
+ * token — which is what an MCP client does once it holds an access token.
+ */
+export async function connectMCPClient(url: string, options: { token?: string } = {}) {
+    const client = new Client({ name: 'test', version: '1.0.0' });
+    await client.connect(
+        new StreamableHTTPClientTransport(new URL(url), {
+            ...(options.token
+                ? { requestInit: { headers: { Authorization: `Bearer ${options.token}` } } }
+                : {}),
+        })
+    );
+    return client;
+}
+
+/**
+ * Fetch the OAuth 2.0 Protected Resource Metadata (RFC 9728) document for an MCP endpoint.
+ */
+export async function fetchProtectedResourceMetadata(siteURL: string, endpoint: string) {
+    const response = await fetch(
+        getContentTestURL(`${siteURL}/.well-known/oauth-protected-resource/${endpoint}`)
+    );
+    return {
+        status: response.status,
+        body: response.ok ? ((await response.json()) as Record<string, unknown>) : undefined,
+    };
+}

packages/gitbook/tests/mcp.test.ts

@@ -1,21 +1,20 @@
-import { Client } from '@modelcontextprotocol/sdk/client';
-import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { expect, it } from 'bun:test';
+import { UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js';
+import { describe, expect, it } from 'bun:test';
+import jwt from 'jsonwebtoken';
 
+import {
+    McpOAuthTestClient,
+    STUBBED_OAUTH_CLIENT_ID,
+    connectMCPClient,
+    fetchProtectedResourceMetadata,
+} from './mcp-utils';
 import { getContentTestURL } from './utils';
 
 it(
     'should expose a MCP server',
     async () => {
-        const client = new Client({
-            name: 'test',
-            version: '1.0.0',
-        });
-
-        await client.connect(
-            new StreamableHTTPClientTransport(
-                new URL(getContentTestURL('https://gitbook.com/docs/~gitbook/mcp'))
-            )
+        const client = await connectMCPClient(
+            getContentTestURL('https://gitbook.com/docs/~gitbook/mcp')
         );
 
         const tools = await client.listTools();
@@ -34,43 +33,12 @@ it(
     { timeout: 10_000 }
 );
 
-it(
-    'should expose a MCP server on the authenticated path',
-    async () => {
-        const client = new Client({
-            name: 'test',
-            version: '1.0.0',
-        });
-
-        await client.connect(
-            new StreamableHTTPClientTransport(
-                new URL(getContentTestURL('https://gitbook.com/docs/~gitbook/mcp/auth'))
-            )
-        );
-
-        const tools = await client.listTools();
-        expect(tools.tools[0]?.name).toBe('searchDocumentation');
-        expect(tools.tools[1]?.name).toBe('getPage');
-        expect(tools.tools.some((tool) => tool.name === 'sendFeedback')).toBe(true);
-    },
-    { timeout: 10_000 }
-);
-
 it(
     'should get a page from another site space through MCP',
     async () => {
-        const client = new Client({
-            name: 'test',
-            version: '1.0.0',
-        });
-
-        await client.connect(
-            new StreamableHTTPClientTransport(
-                new URL(
-                    getContentTestURL(
-                        'https://gitbook-open-e2e-sites.gitbook.io/api-multi-versions-share-links/8tNo6MeXg7CkFMzSSz81/~gitbook/mcp/auth'
-                    )
-                )
+        const client = await connectMCPClient(
+            getContentTestURL(
+                'https://gitbook-open-e2e-sites.gitbook.io/api-multi-versions-share-links/8tNo6MeXg7CkFMzSSz81/~gitbook/mcp/auth'
             )
         );
 
@@ -86,3 +54,248 @@ it(
     },
     { timeout: 15_000 }
 );
+
+describe('MCP on a site behind visitor authentication', () => {
+    const VA_SITE_URL = 'https://gitbook-open-e2e-sites.gitbook.io/va-site-redirects-fallback';
+
+    // Both endpoints are protected on a VA site: there is no public content to serve.
+    const protectedEndpoints = ['~gitbook/mcp', '~gitbook/mcp/auth'];
+
+    it.each(protectedEndpoints)(
+        'should challenge unauthenticated requests to %s with a pointer to the PRM document',
+        async (endpoint) => {
+            const response = await fetch(getContentTestURL(`${VA_SITE_URL}/${endpoint}`));
+
+            expect(response.status).toBe(401);
+            expect(response.headers.get('WWW-Authenticate')).toMatch(
+                new RegExp(
+                    `^Bearer realm="mcp", resource_metadata="https?://.+/\\.well-known/oauth-protected-resource/${endpoint}"$`
+                )
+            );
+        },
+        15_000
+    );
+
+    it.each(protectedEndpoints)(
+        'should serve protected resource metadata for %s',
+        async (endpoint) => {
+            const { status, body } = await fetchProtectedResourceMetadata(VA_SITE_URL, endpoint);
+
+            expect(status).toBe(200);
+            expect(body?.resource).toMatch(new RegExp(`/${endpoint}$`));
+            // The authorization server is the site's own OAuth server, which is what the client
+            // resolves the registration and authorization endpoints from.
+            expect(body?.authorization_servers).toEqual([
+                expect.stringMatching(/\/oauth2\/v1\/site_[\w-]+$/),
+            ]);
+        },
+        15_000
+    );
+
+    it(
+        'should reject an MCP client that connects without a token',
+        async () => {
+            // A client without OAuth support just sees the 401; the transport surfaces it either
+            // from the POST or from the SSE stream, so we only assert on the message.
+            await expect(
+                connectMCPClient(getContentTestURL(`${VA_SITE_URL}/~gitbook/mcp`))
+            ).rejects.toThrow(/Unauthorized/);
+        },
+        { timeout: 15_000 }
+    );
+
+    describe('OAuth discovery by an MCP client', () => {
+        const client = new McpOAuthTestClient();
+        let connecting: Promise<void> | undefined;
+
+        /**
+         * Walk the discovery chain once and share the recording across the assertions below. The
+         * client gives up with an `UnauthorizedError` because completing the flow requires the
+         * visitor to log in and consent in a browser.
+         */
+        async function discover() {
+            connecting ??= (async () => {
+                await expect(
+                    client.connect(getContentTestURL(`${VA_SITE_URL}/~gitbook/mcp`))
+                ).rejects.toThrow(UnauthorizedError);
+            })();
+            await connecting;
+            return client.discovery;
+        }
+
+        it(
+            'should point the client at the protected resource metadata',
+            async () => {
+                const { challenge, resourceMetadataURL } = await discover();
+
+                expect(challenge).toStartWith('Bearer realm="mcp", ');
+                expect(resourceMetadataURL).toMatch(
+                    /\/\.well-known\/oauth-protected-resource\/~gitbook\/mcp$/
+                );
+            },
+            { timeout: 30_000 }
+        );
+
+        it(
+            'should describe the resource and its authorization server',
+            async () => {
+                const { protectedResourceMetadata } = await discover();
+
+                expect(protectedResourceMetadata?.resource).toMatch(/\/~gitbook\/mcp$/);
+                expect(protectedResourceMetadata?.authorization_servers).toEqual([
+                    expect.stringMatching(/\/oauth2\/v1\/site_[\w-]+$/),
+                ]);
+            },
+            { timeout: 30_000 }
+        );
+
+        it(
+            'should lead to an authorization server supporting registration and PKCE',
+            async () => {
+                const { protectedResourceMetadata, authorizationServerMetadata } = await discover();
+
+                expect(authorizationServerMetadata?.issuer).toBe(
+                    (protectedResourceMetadata?.authorization_servers as string[])[0]!
+                );
+                expect(authorizationServerMetadata?.registration_endpoint).toMatch(
+                    /\/oauth2\/v1\/site_[\w-]+\/register$/
+                );
+                expect(authorizationServerMetadata?.code_challenge_methods_supported).toContain(
+                    'S256'
+                );
+            },
+            { timeout: 30_000 }
+        );
+
+        it(
+            'should register itself against the advertised registration endpoint',
+            async () => {
+                const { registrationURL, authorizationServerMetadata, clientId } = await discover();
+
+                expect(registrationURL).toBe(
+                    authorizationServerMetadata?.registration_endpoint as string
+                );
+                expect(clientId).toBe(STUBBED_OAUTH_CLIENT_ID);
+            },
+            { timeout: 30_000 }
+        );
+
+        it(
+            'should build an authorization URL for the registered client',
+            async () => {
+                const { authorizationURL, authorizationServerMetadata, protectedResourceMetadata } =
+                    await discover();
+
+                expect(authorizationURL?.href).toStartWith(
+                    authorizationServerMetadata?.authorization_endpoint as string
+                );
+                expect(authorizationURL?.searchParams.get('response_type')).toBe('code');
+                expect(authorizationURL?.searchParams.get('client_id')).toBe(
+                    STUBBED_OAUTH_CLIENT_ID
+                );
+                expect(authorizationURL?.searchParams.get('code_challenge_method')).toBe('S256');
+                // RFC 8707 resource indicator, naming the MCP endpoint the token is meant for.
+                expect(authorizationURL?.searchParams.get('resource')).toBe(
+                    protectedResourceMetadata?.resource as string
+                );
+            },
+            { timeout: 30_000 }
+        );
+    });
+});
+
+/**
+ * A public site serving adaptive content through a custom backend: `~gitbook/mcp` is open to
+ * everyone, while `~gitbook/mcp/auth` opts into the visitor-specific content.
+ */
+describe('MCP on a public site with adaptive content', () => {
+    it(
+        'should let an MCP client connect to ~gitbook/mcp without authenticating',
+        async () => {
+            const client = await connectMCPClient(
+                getContentTestURL(
+                    'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/~gitbook/mcp'
+                )
+            );
+
+            const tools = await client.listTools();
+            expect(tools.tools.map((tool) => tool.name)).toContain('searchDocumentation');
+        },
+        { timeout: 15_000 }
+    );
+
+    it(
+        'should not advertise protected resource metadata for the public ~gitbook/mcp endpoint',
+        async () => {
+            // Advertising a PRM document for an endpoint that never challenges would send clients
+            // doing proactive discovery into an OAuth flow they don't need.
+            const { status } = await fetchProtectedResourceMetadata(
+                'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public',
+                '~gitbook/mcp'
+            );
+
+            expect(status).toBe(404);
+        },
+        { timeout: 15_000 }
+    );
+
+    it(
+        'should advertise protected resource metadata for the ~gitbook/mcp/auth endpoint',
+        async () => {
+            const { status, body } = await fetchProtectedResourceMetadata(
+                'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public',
+                '~gitbook/mcp/auth'
+            );
+
+            expect(status).toBe(200);
+            expect(body?.resource).toMatch(/\/~gitbook\/mcp\/auth$/);
+            expect(body?.authorization_servers).toEqual([
+                expect.stringMatching(/\/oauth2\/v1\/site_[\w-]+$/),
+            ]);
+        },
+        { timeout: 15_000 }
+    );
+
+    it(
+        'should serve the visitor-specific content once the client presents a token',
+        async () => {
+            // The anonymous half goes through the public endpoint: `~gitbook/mcp/auth` now
+            // challenges unauthenticated clients, so there is no anonymous session to compare with.
+            const anonymousClient = await connectMCPClient(
+                getContentTestURL(
+                    'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/~gitbook/mcp'
+                )
+            );
+            const anonymousResponse = await anonymousClient.callTool({
+                name: 'getPage',
+                arguments: {
+                    url: 'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/alpha-user',
+                },
+            });
+            expect(anonymousResponse.isError).toBe(true);
+
+            const alphaClient = await connectMCPClient(
+                getContentTestURL(
+                    'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/~gitbook/mcp/auth'
+                ),
+                {
+                    // The token an adaptive content backend would issue for this visitor.
+                    token: jwt.sign(
+                        { name: 'gitbook-open-tests', isAlphaUser: true },
+                        '4ddd3c2f-e4b7-4e73-840b-526c3be19746',
+                        { expiresIn: '1h' }
+                    ),
+                }
+            );
+            const alphaResponse = await alphaClient.callTool({
+                name: 'getPage',
+                arguments: {
+                    url: 'https://gitbook-open-e2e-sites.gitbook.io/adaptive-content-public/alpha-user',
+                },
+            });
+            // @ts-expect-error - response.content is of type unknown
+            expect(alphaResponse.content[0]?.text).toContain('# Alpha User');
+        },
+        { timeout: 30_000 }
+    );
+});