diff --git a/package.json b/package.json index bf21340f..894081d7 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@modelcontextprotocol/sdk", - "version": "1.15.0", + "version": "1.15.1", "description": "Model Context Protocol implementation for TypeScript", "license": "MIT", "author": "Anthropic, PBC (https://anthropic.com)", diff --git a/src/client/auth.test.ts b/src/client/auth.test.ts index 8155e134..93dd8e94 100644 --- a/src/client/auth.test.ts +++ b/src/client/auth.test.ts @@ -10,6 +10,7 @@ import { auth, type OAuthClientProvider, } from "./auth.js"; +import { OAuthMetadata } from '../shared/auth.js'; // Mock fetch globally const mockFetch = jest.fn(); @@ -177,6 +178,174 @@ describe("OAuth Authorization", () => { await expect(discoverOAuthProtectedResourceMetadata("https://resource.example.com")) .rejects.toThrow(); }); + + it("returns metadata when discovery succeeds with path", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validMetadata, + }); + + const metadata = await discoverOAuthProtectedResourceMetadata("https://resource.example.com/path/name"); + expect(metadata).toEqual(validMetadata); + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(1); + const [url] = calls[0]; + expect(url.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource/path/name"); + }); + + it("preserves query parameters in path-aware discovery", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validMetadata, + }); + + const metadata = await discoverOAuthProtectedResourceMetadata("https://resource.example.com/path?param=value"); + expect(metadata).toEqual(validMetadata); + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(1); + const [url] = calls[0]; + expect(url.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource/path?param=value"); + }); + + it("falls back to root discovery when path-aware discovery returns 404", async () => { + // First call (path-aware) returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + // Second call (root fallback) succeeds + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validMetadata, + }); + + const metadata = await discoverOAuthProtectedResourceMetadata("https://resource.example.com/path/name"); + expect(metadata).toEqual(validMetadata); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(2); + + // First call should be path-aware + const [firstUrl, firstOptions] = calls[0]; + expect(firstUrl.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource/path/name"); + expect(firstOptions.headers).toEqual({ + "MCP-Protocol-Version": LATEST_PROTOCOL_VERSION + }); + + // Second call should be root fallback + const [secondUrl, secondOptions] = calls[1]; + expect(secondUrl.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource"); + expect(secondOptions.headers).toEqual({ + "MCP-Protocol-Version": LATEST_PROTOCOL_VERSION + }); + }); + + it("throws error when both path-aware and root discovery return 404", async () => { + // First call (path-aware) returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + // Second call (root fallback) also returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + await expect(discoverOAuthProtectedResourceMetadata("https://resource.example.com/path/name")) + .rejects.toThrow("Resource server does not implement OAuth 2.0 Protected Resource Metadata."); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(2); + }); + + it("does not fallback when the original URL is already at root path", async () => { + // First call (path-aware for root) returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + await expect(discoverOAuthProtectedResourceMetadata("https://resource.example.com/")) + .rejects.toThrow("Resource server does not implement OAuth 2.0 Protected Resource Metadata."); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(1); // Should not attempt fallback + + const [url] = calls[0]; + expect(url.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource"); + }); + + it("does not fallback when the original URL has no path", async () => { + // First call (path-aware for no path) returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + await expect(discoverOAuthProtectedResourceMetadata("https://resource.example.com")) + .rejects.toThrow("Resource server does not implement OAuth 2.0 Protected Resource Metadata."); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(1); // Should not attempt fallback + + const [url] = calls[0]; + expect(url.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource"); + }); + + it("falls back when path-aware discovery encounters CORS error", async () => { + // First call (path-aware) fails with TypeError (CORS) + mockFetch.mockImplementationOnce(() => Promise.reject(new TypeError("CORS error"))); + + // Retry path-aware without headers (simulating CORS retry) + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + // Second call (root fallback) succeeds + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validMetadata, + }); + + const metadata = await discoverOAuthProtectedResourceMetadata("https://resource.example.com/deep/path"); + expect(metadata).toEqual(validMetadata); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(3); + + // Final call should be root fallback + const [lastUrl, lastOptions] = calls[2]; + expect(lastUrl.toString()).toBe("https://resource.example.com/.well-known/oauth-protected-resource"); + expect(lastOptions.headers).toEqual({ + "MCP-Protocol-Version": LATEST_PROTOCOL_VERSION + }); + }); + + it("does not fallback when resourceMetadataUrl is provided", async () => { + // Call with explicit URL returns 404 + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + }); + + await expect(discoverOAuthProtectedResourceMetadata("https://resource.example.com/path", { + resourceMetadataUrl: "https://custom.example.com/metadata" + })).rejects.toThrow("Resource server does not implement OAuth 2.0 Protected Resource Metadata."); + + const calls = mockFetch.mock.calls; + expect(calls.length).toBe(1); // Should not attempt fallback when explicit URL is provided + + const [url] = calls[0]; + expect(url.toString()).toBe("https://custom.example.com/metadata"); + }); }); describe("discoverOAuthMetadata", () => { @@ -231,7 +400,7 @@ describe("OAuth Authorization", () => { ok: false, status: 404, }); - + // Second call (root fallback) succeeds mockFetch.mockResolvedValueOnce({ ok: true, @@ -241,17 +410,17 @@ describe("OAuth Authorization", () => { const metadata = await discoverOAuthMetadata("https://auth.example.com/path/name"); expect(metadata).toEqual(validMetadata); - + const calls = mockFetch.mock.calls; expect(calls.length).toBe(2); - + // First call should be path-aware const [firstUrl, firstOptions] = calls[0]; expect(firstUrl.toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server/path/name"); expect(firstOptions.headers).toEqual({ "MCP-Protocol-Version": LATEST_PROTOCOL_VERSION }); - + // Second call should be root fallback const [secondUrl, secondOptions] = calls[1]; expect(secondUrl.toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server"); @@ -266,7 +435,7 @@ describe("OAuth Authorization", () => { ok: false, status: 404, }); - + // Second call (root fallback) also returns 404 mockFetch.mockResolvedValueOnce({ ok: false, @@ -275,7 +444,7 @@ describe("OAuth Authorization", () => { const metadata = await discoverOAuthMetadata("https://auth.example.com/path/name"); expect(metadata).toBeUndefined(); - + const calls = mockFetch.mock.calls; expect(calls.length).toBe(2); }); @@ -289,10 +458,10 @@ describe("OAuth Authorization", () => { const metadata = await discoverOAuthMetadata("https://auth.example.com/"); expect(metadata).toBeUndefined(); - + const calls = mockFetch.mock.calls; expect(calls.length).toBe(1); // Should not attempt fallback - + const [url] = calls[0]; expect(url.toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server"); }); @@ -306,10 +475,10 @@ describe("OAuth Authorization", () => { const metadata = await discoverOAuthMetadata("https://auth.example.com"); expect(metadata).toBeUndefined(); - + const calls = mockFetch.mock.calls; expect(calls.length).toBe(1); // Should not attempt fallback - + const [url] = calls[0]; expect(url.toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server"); }); @@ -317,13 +486,13 @@ describe("OAuth Authorization", () => { it("falls back when path-aware discovery encounters CORS error", async () => { // First call (path-aware) fails with TypeError (CORS) mockFetch.mockImplementationOnce(() => Promise.reject(new TypeError("CORS error"))); - + // Retry path-aware without headers (simulating CORS retry) mockFetch.mockResolvedValueOnce({ ok: false, status: 404, }); - + // Second call (root fallback) succeeds mockFetch.mockResolvedValueOnce({ ok: true, @@ -333,10 +502,10 @@ describe("OAuth Authorization", () => { const metadata = await discoverOAuthMetadata("https://auth.example.com/deep/path"); expect(metadata).toEqual(validMetadata); - + const calls = mockFetch.mock.calls; expect(calls.length).toBe(3); - + // Final call should be root fallback const [lastUrl, lastOptions] = calls[2]; expect(lastUrl.toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server"); @@ -600,6 +769,13 @@ describe("OAuth Authorization", () => { refresh_token: "refresh123", }; + const validMetadata = { + issuer: "https://auth.example.com", + authorization_endpoint: "https://auth.example.com/authorize", + token_endpoint: "https://auth.example.com/token", + response_types_supported: ["code"] + }; + const validClientInfo = { client_id: "client123", client_secret: "secret123", @@ -629,9 +805,9 @@ describe("OAuth Authorization", () => { }), expect.objectContaining({ method: "POST", - headers: { + headers: new Headers({ "Content-Type": "application/x-www-form-urlencoded", - }, + }), }) ); @@ -645,6 +821,52 @@ describe("OAuth Authorization", () => { expect(body.get("resource")).toBe("https://api.example.com/mcp-server"); }); + it("exchanges code for tokens with auth", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await exchangeAuthorization("https://auth.example.com", { + metadata: validMetadata, + clientInformation: validClientInfo, + authorizationCode: "code123", + codeVerifier: "verifier123", + redirectUri: "http://localhost:3000/callback", + addClientAuthentication: (headers: Headers, params: URLSearchParams, url: string | URL, metadata: OAuthMetadata) => { + headers.set("Authorization", "Basic " + btoa(validClientInfo.client_id + ":" + validClientInfo.client_secret)); + params.set("example_url", typeof url === 'string' ? url : url.toString()); + params.set("example_metadata", metadata.authorization_endpoint); + params.set("example_param", "example_value"); + }, + }); + + expect(tokens).toEqual(validTokens); + expect(mockFetch).toHaveBeenCalledWith( + expect.objectContaining({ + href: "https://auth.example.com/token", + }), + expect.objectContaining({ + method: "POST", + }) + ); + + const headers = mockFetch.mock.calls[0][1].headers as Headers; + expect(headers.get("Content-Type")).toBe("application/x-www-form-urlencoded"); + expect(headers.get("Authorization")).toBe("Basic Y2xpZW50MTIzOnNlY3JldDEyMw=="); + const body = mockFetch.mock.calls[0][1].body as URLSearchParams; + expect(body.get("grant_type")).toBe("authorization_code"); + expect(body.get("code")).toBe("code123"); + expect(body.get("code_verifier")).toBe("verifier123"); + expect(body.get("client_id")).toBeNull(); + expect(body.get("redirect_uri")).toBe("http://localhost:3000/callback"); + expect(body.get("example_url")).toBe("https://auth.example.com"); + expect(body.get("example_metadata")).toBe("https://auth.example.com/authorize"); + expect(body.get("example_param")).toBe("example_value"); + expect(body.get("client_secret")).toBeNull(); + }); + it("validates token response schema", async () => { mockFetch.mockResolvedValueOnce({ ok: true, @@ -693,6 +915,13 @@ describe("OAuth Authorization", () => { refresh_token: "newrefresh123", }; + const validMetadata = { + issuer: "https://auth.example.com", + authorization_endpoint: "https://auth.example.com/authorize", + token_endpoint: "https://auth.example.com/token", + response_types_supported: ["code"] + }; + const validClientInfo = { client_id: "client123", client_secret: "secret123", @@ -720,9 +949,9 @@ describe("OAuth Authorization", () => { }), expect.objectContaining({ method: "POST", - headers: { + headers: new Headers({ "Content-Type": "application/x-www-form-urlencoded", - }, + }), }) ); @@ -734,6 +963,48 @@ describe("OAuth Authorization", () => { expect(body.get("resource")).toBe("https://api.example.com/mcp-server"); }); + it("exchanges refresh token for new tokens with auth", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokensWithNewRefreshToken, + }); + + const tokens = await refreshAuthorization("https://auth.example.com", { + metadata: validMetadata, + clientInformation: validClientInfo, + refreshToken: "refresh123", + addClientAuthentication: (headers: Headers, params: URLSearchParams, url: string | URL, metadata?: OAuthMetadata) => { + headers.set("Authorization", "Basic " + btoa(validClientInfo.client_id + ":" + validClientInfo.client_secret)); + params.set("example_url", typeof url === 'string' ? url : url.toString()); + params.set("example_metadata", metadata?.authorization_endpoint ?? '?'); + params.set("example_param", "example_value"); + }, + }); + + expect(tokens).toEqual(validTokensWithNewRefreshToken); + expect(mockFetch).toHaveBeenCalledWith( + expect.objectContaining({ + href: "https://auth.example.com/token", + }), + expect.objectContaining({ + method: "POST", + }) + ); + + const headers = mockFetch.mock.calls[0][1].headers as Headers; + expect(headers.get("Content-Type")).toBe("application/x-www-form-urlencoded"); + expect(headers.get("Authorization")).toBe("Basic Y2xpZW50MTIzOnNlY3JldDEyMw=="); + const body = mockFetch.mock.calls[0][1].body as URLSearchParams; + expect(body.get("grant_type")).toBe("refresh_token"); + expect(body.get("refresh_token")).toBe("refresh123"); + expect(body.get("client_id")).toBeNull(); + expect(body.get("example_url")).toBe("https://auth.example.com"); + expect(body.get("example_metadata")).toBe("https://auth.example.com/authorize"); + expect(body.get("example_param")).toBe("example_value"); + expect(body.get("client_secret")).toBeNull(); + }); + it("exchanges refresh token for new tokens and keep existing refresh token if none is returned", async () => { mockFetch.mockResolvedValueOnce({ ok: true, @@ -1476,5 +1747,326 @@ describe("OAuth Authorization", () => { expect(body.get("grant_type")).toBe("refresh_token"); expect(body.get("refresh_token")).toBe("refresh123"); }); + + it("fetches AS metadata with path from serverUrl when PRM returns external AS", async () => { + // Mock PRM discovery that returns an external AS + mockFetch.mockImplementation((url) => { + const urlString = url.toString(); + + if (urlString === "https://my.resource.com/.well-known/oauth-protected-resource/path/name") { + return Promise.resolve({ + ok: true, + status: 200, + json: async () => ({ + resource: "https://my.resource.com/", + authorization_servers: ["https://auth.example.com/"], + }), + }); + } else if (urlString === "https://auth.example.com/.well-known/oauth-authorization-server/path/name") { + // Path-aware discovery on AS with path from serverUrl + return Promise.resolve({ + ok: true, + status: 200, + json: async () => ({ + issuer: "https://auth.example.com", + authorization_endpoint: "https://auth.example.com/authorize", + token_endpoint: "https://auth.example.com/token", + response_types_supported: ["code"], + code_challenge_methods_supported: ["S256"], + }), + }); + } + + return Promise.resolve({ ok: false, status: 404 }); + }); + + // Mock provider methods + (mockProvider.clientInformation as jest.Mock).mockResolvedValue({ + client_id: "test-client", + client_secret: "test-secret", + }); + (mockProvider.tokens as jest.Mock).mockResolvedValue(undefined); + (mockProvider.saveCodeVerifier as jest.Mock).mockResolvedValue(undefined); + (mockProvider.redirectToAuthorization as jest.Mock).mockResolvedValue(undefined); + + // Call auth with serverUrl that has a path + const result = await auth(mockProvider, { + serverUrl: "https://my.resource.com/path/name", + }); + + expect(result).toBe("REDIRECT"); + + // Verify the correct URLs were fetched + const calls = mockFetch.mock.calls; + + // First call should be to PRM + expect(calls[0][0].toString()).toBe("https://my.resource.com/.well-known/oauth-protected-resource/path/name"); + + // Second call should be to AS metadata with the path from serverUrl + expect(calls[1][0].toString()).toBe("https://auth.example.com/.well-known/oauth-authorization-server/path/name"); + }); }); + + describe("exchangeAuthorization with multiple client authentication methods", () => { + const validTokens = { + access_token: "access123", + token_type: "Bearer", + expires_in: 3600, + refresh_token: "refresh123", + }; + + const validClientInfo = { + client_id: "client123", + client_secret: "secret123", + redirect_uris: ["http://localhost:3000/callback"], + client_name: "Test Client", + }; + + const metadataWithBasicOnly = { + issuer: "https://auth.example.com", + authorization_endpoint: "https://auth.example.com/auth", + token_endpoint: "https://auth.example.com/token", + response_types_supported: ["code"], + code_challenge_methods_supported: ["S256"], + token_endpoint_auth_methods_supported: ["client_secret_basic"], + }; + + const metadataWithPostOnly = { + ...metadataWithBasicOnly, + token_endpoint_auth_methods_supported: ["client_secret_post"], + }; + + const metadataWithNoneOnly = { + ...metadataWithBasicOnly, + token_endpoint_auth_methods_supported: ["none"], + }; + + const metadataWithAllBuiltinMethods = { + ...metadataWithBasicOnly, + token_endpoint_auth_methods_supported: ["client_secret_basic", "client_secret_post", "none"], + }; + + it("uses HTTP Basic authentication when client_secret_basic is supported", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await exchangeAuthorization("https://auth.example.com", { + metadata: metadataWithBasicOnly, + clientInformation: validClientInfo, + authorizationCode: "code123", + redirectUri: "http://localhost:3000/callback", + codeVerifier: "verifier123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check Authorization header + const authHeader = request.headers.get("Authorization"); + const expected = "Basic " + btoa("client123:secret123"); + expect(authHeader).toBe(expected); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBeNull(); + expect(body.get("client_secret")).toBeNull(); + }); + + it("includes credentials in request body when client_secret_post is supported", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await exchangeAuthorization("https://auth.example.com", { + metadata: metadataWithPostOnly, + clientInformation: validClientInfo, + authorizationCode: "code123", + redirectUri: "http://localhost:3000/callback", + codeVerifier: "verifier123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check no Authorization header + expect(request.headers.get("Authorization")).toBeNull(); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBe("client123"); + expect(body.get("client_secret")).toBe("secret123"); + }); + + it("it picks client_secret_basic when all builtin methods are supported", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await exchangeAuthorization("https://auth.example.com", { + metadata: metadataWithAllBuiltinMethods, + clientInformation: validClientInfo, + authorizationCode: "code123", + redirectUri: "http://localhost:3000/callback", + codeVerifier: "verifier123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check Authorization header - should use Basic auth as it's the most secure + const authHeader = request.headers.get("Authorization"); + const expected = "Basic " + btoa("client123:secret123"); + expect(authHeader).toBe(expected); + + // Credentials should not be in body when using Basic auth + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBeNull(); + expect(body.get("client_secret")).toBeNull(); + }); + + it("uses public client authentication when none method is specified", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const clientInfoWithoutSecret = { + client_id: "client123", + redirect_uris: ["http://localhost:3000/callback"], + client_name: "Test Client", + }; + + const tokens = await exchangeAuthorization("https://auth.example.com", { + metadata: metadataWithNoneOnly, + clientInformation: clientInfoWithoutSecret, + authorizationCode: "code123", + redirectUri: "http://localhost:3000/callback", + codeVerifier: "verifier123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check no Authorization header + expect(request.headers.get("Authorization")).toBeNull(); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBe("client123"); + expect(body.get("client_secret")).toBeNull(); + }); + + it("defaults to client_secret_post when no auth methods specified", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await exchangeAuthorization("https://auth.example.com", { + clientInformation: validClientInfo, + authorizationCode: "code123", + redirectUri: "http://localhost:3000/callback", + codeVerifier: "verifier123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check headers + expect(request.headers.get("Content-Type")).toBe("application/x-www-form-urlencoded"); + expect(request.headers.get("Authorization")).toBeNull(); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBe("client123"); + expect(body.get("client_secret")).toBe("secret123"); + }); + }); + + describe("refreshAuthorization with multiple client authentication methods", () => { + const validTokens = { + access_token: "newaccess123", + token_type: "Bearer", + expires_in: 3600, + refresh_token: "newrefresh123", + }; + + const validClientInfo = { + client_id: "client123", + client_secret: "secret123", + redirect_uris: ["http://localhost:3000/callback"], + client_name: "Test Client", + }; + + const metadataWithBasicOnly = { + issuer: "https://auth.example.com", + authorization_endpoint: "https://auth.example.com/auth", + token_endpoint: "https://auth.example.com/token", + response_types_supported: ["code"], + token_endpoint_auth_methods_supported: ["client_secret_basic"], + }; + + const metadataWithPostOnly = { + ...metadataWithBasicOnly, + token_endpoint_auth_methods_supported: ["client_secret_post"], + }; + + it("uses client_secret_basic for refresh token", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await refreshAuthorization("https://auth.example.com", { + metadata: metadataWithBasicOnly, + clientInformation: validClientInfo, + refreshToken: "refresh123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check Authorization header + const authHeader = request.headers.get("Authorization"); + const expected = "Basic " + btoa("client123:secret123"); + expect(authHeader).toBe(expected); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBeNull(); // should not be in body + expect(body.get("client_secret")).toBeNull(); // should not be in body + expect(body.get("refresh_token")).toBe("refresh123"); + }); + + it("uses client_secret_post for refresh token", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => validTokens, + }); + + const tokens = await refreshAuthorization("https://auth.example.com", { + metadata: metadataWithPostOnly, + clientInformation: validClientInfo, + refreshToken: "refresh123", + }); + + expect(tokens).toEqual(validTokens); + const request = mockFetch.mock.calls[0][1]; + + // Check no Authorization header + expect(request.headers.get("Authorization")).toBeNull(); + + const body = request.body as URLSearchParams; + expect(body.get("client_id")).toBe("client123"); + expect(body.get("client_secret")).toBe("secret123"); + expect(body.get("refresh_token")).toBe("refresh123"); + }); + }); + }); diff --git a/src/client/auth.ts b/src/client/auth.ts index 71101a42..2b69a5d8 100644 --- a/src/client/auth.ts +++ b/src/client/auth.ts @@ -73,6 +73,26 @@ export interface OAuthClientProvider { */ codeVerifier(): string | Promise; + /** + * Adds custom client authentication to OAuth token requests. + * + * This optional method allows implementations to customize how client credentials + * are included in token exchange and refresh requests. When provided, this method + * is called instead of the default authentication logic, giving full control over + * the authentication mechanism. + * + * Common use cases include: + * - Supporting authentication methods beyond the standard OAuth 2.0 methods + * - Adding custom headers for proprietary authentication schemes + * - Implementing client assertion-based authentication (e.g., JWT bearer tokens) + * + * @param headers - The request headers (can be modified to add authentication) + * @param params - The request body parameters (can be modified to add credentials) + * @param url - The token endpoint URL being called + * @param metadata - Optional OAuth metadata for the server, which may include supported authentication methods + */ + addClientAuthentication?(headers: Headers, params: URLSearchParams, url: string | URL, metadata?: OAuthMetadata): void | Promise; + /** * If defined, overrides the selection and validation of the * RFC 8707 Resource Indicator. If left undefined, default @@ -91,6 +111,114 @@ export class UnauthorizedError extends Error { } } +type ClientAuthMethod = 'client_secret_basic' | 'client_secret_post' | 'none'; + +/** + * Determines the best client authentication method to use based on server support and client configuration. + * + * Priority order (highest to lowest): + * 1. client_secret_basic (if client secret is available) + * 2. client_secret_post (if client secret is available) + * 3. none (for public clients) + * + * @param clientInformation - OAuth client information containing credentials + * @param supportedMethods - Authentication methods supported by the authorization server + * @returns The selected authentication method + */ +function selectClientAuthMethod( + clientInformation: OAuthClientInformation, + supportedMethods: string[] +): ClientAuthMethod { + const hasClientSecret = clientInformation.client_secret !== undefined; + + // If server doesn't specify supported methods, use RFC 6749 defaults + if (supportedMethods.length === 0) { + return hasClientSecret ? "client_secret_post" : "none"; + } + + // Try methods in priority order (most secure first) + if (hasClientSecret && supportedMethods.includes("client_secret_basic")) { + return "client_secret_basic"; + } + + if (hasClientSecret && supportedMethods.includes("client_secret_post")) { + return "client_secret_post"; + } + + if (supportedMethods.includes("none")) { + return "none"; + } + + // Fallback: use what we have + return hasClientSecret ? "client_secret_post" : "none"; +} + +/** + * Applies client authentication to the request based on the specified method. + * + * Implements OAuth 2.1 client authentication methods: + * - client_secret_basic: HTTP Basic authentication (RFC 6749 Section 2.3.1) + * - client_secret_post: Credentials in request body (RFC 6749 Section 2.3.1) + * - none: Public client authentication (RFC 6749 Section 2.1) + * + * @param method - The authentication method to use + * @param clientInformation - OAuth client information containing credentials + * @param headers - HTTP headers object to modify + * @param params - URL search parameters to modify + * @throws {Error} When required credentials are missing + */ +function applyClientAuthentication( + method: ClientAuthMethod, + clientInformation: OAuthClientInformation, + headers: Headers, + params: URLSearchParams +): void { + const { client_id, client_secret } = clientInformation; + + switch (method) { + case "client_secret_basic": + applyBasicAuth(client_id, client_secret, headers); + return; + case "client_secret_post": + applyPostAuth(client_id, client_secret, params); + return; + case "none": + applyPublicAuth(client_id, params); + return; + default: + throw new Error(`Unsupported client authentication method: ${method}`); + } +} + +/** + * Applies HTTP Basic authentication (RFC 6749 Section 2.3.1) + */ +function applyBasicAuth(clientId: string, clientSecret: string | undefined, headers: Headers): void { + if (!clientSecret) { + throw new Error("client_secret_basic authentication requires a client_secret"); + } + + const credentials = btoa(`${clientId}:${clientSecret}`); + headers.set("Authorization", `Basic ${credentials}`); +} + +/** + * Applies POST body authentication (RFC 6749 Section 2.3.1) + */ +function applyPostAuth(clientId: string, clientSecret: string | undefined, params: URLSearchParams): void { + params.set("client_id", clientId); + if (clientSecret) { + params.set("client_secret", clientSecret); + } +} + +/** + * Applies public client authentication (RFC 6749 Section 2.1) + */ +function applyPublicAuth(clientId: string, params: URLSearchParams): void { + params.set("client_id", clientId); +} + /** * Orchestrates the full auth flow with a server. * @@ -107,12 +235,13 @@ export async function auth( serverUrl: string | URL; authorizationCode?: string; scope?: string; - resourceMetadataUrl?: URL }): Promise { + resourceMetadataUrl?: URL + }): Promise { let resourceMetadata: OAuthProtectedResourceMetadata | undefined; let authorizationServerUrl = serverUrl; try { - resourceMetadata = await discoverOAuthProtectedResourceMetadata(serverUrl, {resourceMetadataUrl}); + resourceMetadata = await discoverOAuthProtectedResourceMetadata(serverUrl, { resourceMetadataUrl }); if (resourceMetadata.authorization_servers && resourceMetadata.authorization_servers.length > 0) { authorizationServerUrl = resourceMetadata.authorization_servers[0]; } @@ -122,7 +251,9 @@ export async function auth( const resource: URL | undefined = await selectResourceURL(serverUrl, provider, resourceMetadata); - const metadata = await discoverOAuthMetadata(authorizationServerUrl); + const metadata = await discoverOAuthMetadata(serverUrl, { + authorizationServerUrl + }); // Handle client registration if needed let clientInformation = await Promise.resolve(provider.clientInformation()); @@ -154,6 +285,7 @@ export async function auth( codeVerifier, redirectUri: provider.redirectUrl, resource, + addClientAuthentication: provider.addClientAuthentication, }); await provider.saveTokens(tokens); @@ -171,6 +303,7 @@ export async function auth( clientInformation, refreshToken: tokens.refresh_token, resource, + addClientAuthentication: provider.addClientAuthentication, }); await provider.saveTokens(newTokens); @@ -197,7 +330,7 @@ export async function auth( return "REDIRECT"; } -export async function selectResourceURL(serverUrl: string| URL, provider: OAuthClientProvider, resourceMetadata?: OAuthProtectedResourceMetadata): Promise { +export async function selectResourceURL(serverUrl: string | URL, provider: OAuthClientProvider, resourceMetadata?: OAuthProtectedResourceMetadata): Promise { const defaultResource = resourceUrlFromServerUrl(serverUrl); // If provider has custom validation, delegate to it @@ -256,31 +389,16 @@ export async function discoverOAuthProtectedResourceMetadata( serverUrl: string | URL, opts?: { protocolVersion?: string, resourceMetadataUrl?: string | URL }, ): Promise { + const response = await discoverMetadataWithFallback( + serverUrl, + 'oauth-protected-resource', + { + protocolVersion: opts?.protocolVersion, + metadataUrl: opts?.resourceMetadataUrl, + }, + ); - let url: URL - if (opts?.resourceMetadataUrl) { - url = new URL(opts?.resourceMetadataUrl); - } else { - url = new URL("/.well-known/oauth-protected-resource", serverUrl); - } - - let response: Response; - try { - response = await fetch(url, { - headers: { - "MCP-Protocol-Version": opts?.protocolVersion ?? LATEST_PROTOCOL_VERSION - } - }); - } catch (error) { - // CORS errors come back as TypeError - if (error instanceof TypeError) { - response = await fetch(url); - } else { - throw error; - } - } - - if (response.status === 404) { + if (!response || response.status === 404) { throw new Error(`Resource server does not implement OAuth 2.0 Protected Resource Metadata.`); } @@ -318,8 +436,8 @@ async function fetchWithCorsRetry( /** * Constructs the well-known path for OAuth metadata discovery */ -function buildWellKnownPath(pathname: string): string { - let wellKnownPath = `/.well-known/oauth-authorization-server${pathname}`; +function buildWellKnownPath(wellKnownPrefix: string, pathname: string): string { + let wellKnownPath = `/.well-known/${wellKnownPrefix}${pathname}`; if (pathname.endsWith('/')) { // Strip trailing slash from pathname to avoid double slashes wellKnownPath = wellKnownPath.slice(0, -1); @@ -347,6 +465,38 @@ function shouldAttemptFallback(response: Response | undefined, pathname: string) return !response || response.status === 404 && pathname !== '/'; } +/** + * Generic function for discovering OAuth metadata with fallback support + */ +async function discoverMetadataWithFallback( + serverUrl: string | URL, + wellKnownType: 'oauth-authorization-server' | 'oauth-protected-resource', + opts?: { protocolVersion?: string; metadataUrl?: string | URL, metadataServerUrl?: string | URL }, +): Promise { + const issuer = new URL(serverUrl); + const protocolVersion = opts?.protocolVersion ?? LATEST_PROTOCOL_VERSION; + + let url: URL; + if (opts?.metadataUrl) { + url = new URL(opts.metadataUrl); + } else { + // Try path-aware discovery first + const wellKnownPath = buildWellKnownPath(wellKnownType, issuer.pathname); + url = new URL(wellKnownPath, opts?.metadataServerUrl ?? issuer); + url.search = issuer.search; + } + + let response = await tryMetadataDiscovery(url, protocolVersion); + + // If path-aware discovery fails with 404 and we're not already at root, try fallback to root discovery + if (!opts?.metadataUrl && shouldAttemptFallback(response, issuer.pathname)) { + const rootUrl = new URL(`/.well-known/${wellKnownType}`, issuer); + response = await tryMetadataDiscovery(rootUrl, protocolVersion); + } + + return response; +} + /** * Looks up RFC 8414 OAuth 2.0 Authorization Server Metadata. * @@ -354,22 +504,35 @@ function shouldAttemptFallback(response: Response | undefined, pathname: string) * return `undefined`. Any other errors will be thrown as exceptions. */ export async function discoverOAuthMetadata( - authorizationServerUrl: string | URL, - opts?: { protocolVersion?: string }, + issuer: string | URL, + { + authorizationServerUrl, + protocolVersion, + }: { + authorizationServerUrl?: string | URL, + protocolVersion?: string, + } = {}, ): Promise { - const issuer = new URL(authorizationServerUrl); - const protocolVersion = opts?.protocolVersion ?? LATEST_PROTOCOL_VERSION; - - // Try path-aware discovery first (RFC 8414 compliant) - const wellKnownPath = buildWellKnownPath(issuer.pathname); - const pathAwareUrl = new URL(wellKnownPath, issuer); - let response = await tryMetadataDiscovery(pathAwareUrl, protocolVersion); - - // If path-aware discovery fails with 404, try fallback to root discovery - if (shouldAttemptFallback(response, issuer.pathname)) { - const rootUrl = new URL("/.well-known/oauth-authorization-server", issuer); - response = await tryMetadataDiscovery(rootUrl, protocolVersion); + if (typeof issuer === 'string') { + issuer = new URL(issuer); } + if (!authorizationServerUrl) { + authorizationServerUrl = issuer; + } + if (typeof authorizationServerUrl === 'string') { + authorizationServerUrl = new URL(authorizationServerUrl); + } + protocolVersion ??= LATEST_PROTOCOL_VERSION; + + const response = await discoverMetadataWithFallback( + issuer, + 'oauth-authorization-server', + { + protocolVersion, + metadataServerUrl: authorizationServerUrl, + }, + ); + if (!response || response.status === 404) { return undefined; } @@ -460,6 +623,15 @@ export async function startAuthorization( /** * Exchanges an authorization code for an access token with the given server. + * + * Supports multiple client authentication methods as specified in OAuth 2.1: + * - Automatically selects the best authentication method based on server support + * - Falls back to appropriate defaults when server metadata is unavailable + * + * @param authorizationServerUrl - The authorization server's base URL + * @param options - Configuration object containing client info, auth code, etc. + * @returns Promise resolving to OAuth tokens + * @throws {Error} When token exchange fails or authentication is invalid */ export async function exchangeAuthorization( authorizationServerUrl: string | URL, @@ -470,6 +642,7 @@ export async function exchangeAuthorization( codeVerifier, redirectUri, resource, + addClientAuthentication }: { metadata?: OAuthMetadata; clientInformation: OAuthClientInformation; @@ -477,37 +650,43 @@ export async function exchangeAuthorization( codeVerifier: string; redirectUri: string | URL; resource?: URL; + addClientAuthentication?: OAuthClientProvider["addClientAuthentication"]; }, ): Promise { const grantType = "authorization_code"; - let tokenUrl: URL; - if (metadata) { - tokenUrl = new URL(metadata.token_endpoint); + const tokenUrl = metadata?.token_endpoint + ? new URL(metadata.token_endpoint) + : new URL("/token", authorizationServerUrl); - if ( - metadata.grant_types_supported && + if ( + metadata?.grant_types_supported && !metadata.grant_types_supported.includes(grantType) - ) { - throw new Error( + ) { + throw new Error( `Incompatible auth server: does not support grant type ${grantType}`, - ); - } - } else { - tokenUrl = new URL("/token", authorizationServerUrl); + ); } // Exchange code for tokens + const headers = new Headers({ + "Content-Type": "application/x-www-form-urlencoded", + }); const params = new URLSearchParams({ grant_type: grantType, - client_id: clientInformation.client_id, code: authorizationCode, code_verifier: codeVerifier, redirect_uri: String(redirectUri), }); - if (clientInformation.client_secret) { - params.set("client_secret", clientInformation.client_secret); + if (addClientAuthentication) { + addClientAuthentication(headers, params, authorizationServerUrl, metadata); + } else { + // Determine and apply client authentication method + const supportedMethods = metadata?.token_endpoint_auth_methods_supported ?? []; + const authMethod = selectClientAuthMethod(clientInformation, supportedMethods); + + applyClientAuthentication(authMethod, clientInformation, headers, params); } if (resource) { @@ -516,9 +695,7 @@ export async function exchangeAuthorization( const response = await fetch(tokenUrl, { method: "POST", - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, + headers, body: params, }); @@ -531,6 +708,15 @@ export async function exchangeAuthorization( /** * Exchange a refresh token for an updated access token. + * + * Supports multiple client authentication methods as specified in OAuth 2.1: + * - Automatically selects the best authentication method based on server support + * - Preserves the original refresh token if a new one is not returned + * + * @param authorizationServerUrl - The authorization server's base URL + * @param options - Configuration object containing client info, refresh token, etc. + * @returns Promise resolving to OAuth tokens (preserves original refresh_token if not replaced) + * @throws {Error} When token refresh fails or authentication is invalid */ export async function refreshAuthorization( authorizationServerUrl: string | URL, @@ -539,12 +725,14 @@ export async function refreshAuthorization( clientInformation, refreshToken, resource, + addClientAuthentication, }: { metadata?: OAuthMetadata; clientInformation: OAuthClientInformation; refreshToken: string; resource?: URL; - }, + addClientAuthentication?: OAuthClientProvider["addClientAuthentication"]; + } ): Promise { const grantType = "refresh_token"; @@ -565,14 +753,22 @@ export async function refreshAuthorization( } // Exchange refresh token + const headers = new Headers({ + "Content-Type": "application/x-www-form-urlencoded", + }); const params = new URLSearchParams({ grant_type: grantType, - client_id: clientInformation.client_id, refresh_token: refreshToken, }); - if (clientInformation.client_secret) { - params.set("client_secret", clientInformation.client_secret); + if (addClientAuthentication) { + addClientAuthentication(headers, params, authorizationServerUrl, metadata); + } else { + // Determine and apply client authentication method + const supportedMethods = metadata?.token_endpoint_auth_methods_supported ?? []; + const authMethod = selectClientAuthMethod(clientInformation, supportedMethods); + + applyClientAuthentication(authMethod, clientInformation, headers, params); } if (resource) { @@ -581,9 +777,7 @@ export async function refreshAuthorization( const response = await fetch(tokenUrl, { method: "POST", - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, + headers, body: params, }); if (!response.ok) { diff --git a/src/server/streamableHttp.test.ts b/src/server/streamableHttp.test.ts index e54bea01..3a0a5c06 100644 --- a/src/server/streamableHttp.test.ts +++ b/src/server/streamableHttp.test.ts @@ -29,7 +29,8 @@ interface TestServerConfig { enableJsonResponse?: boolean; customRequestHandler?: (req: IncomingMessage, res: ServerResponse, parsedBody?: unknown) => Promise; eventStore?: EventStore; - onsessionclosed?: (sessionId: string) => void; + onsessioninitialized?: (sessionId: string) => void | Promise; + onsessionclosed?: (sessionId: string) => void | Promise; } /** @@ -59,6 +60,7 @@ async function createTestServer(config: TestServerConfig = { sessionIdGenerator: sessionIdGenerator: config.sessionIdGenerator, enableJsonResponse: config.enableJsonResponse ?? false, eventStore: config.eventStore, + onsessioninitialized: config.onsessioninitialized, onsessionclosed: config.onsessionclosed }); @@ -114,6 +116,7 @@ async function createTestAuthServer(config: TestServerConfig = { sessionIdGenera sessionIdGenerator: config.sessionIdGenerator, enableJsonResponse: config.enableJsonResponse ?? false, eventStore: config.eventStore, + onsessioninitialized: config.onsessioninitialized, onsessionclosed: config.onsessionclosed }); @@ -1666,6 +1669,213 @@ describe("StreamableHTTPServerTransport onsessionclosed callback", () => { }); }); +// Test async callbacks for onsessioninitialized and onsessionclosed +describe("StreamableHTTPServerTransport async callbacks", () => { + it("should support async onsessioninitialized callback", async () => { + const initializationOrder: string[] = []; + + // Create server with async onsessioninitialized callback + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: async (sessionId: string) => { + initializationOrder.push('async-start'); + // Simulate async operation + await new Promise(resolve => setTimeout(resolve, 10)); + initializationOrder.push('async-end'); + initializationOrder.push(sessionId); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize to trigger the callback + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + const tempSessionId = initResponse.headers.get("mcp-session-id"); + + // Give time for async callback to complete + await new Promise(resolve => setTimeout(resolve, 50)); + + expect(initializationOrder).toEqual(['async-start', 'async-end', tempSessionId]); + + // Clean up + tempServer.close(); + }); + + it("should support sync onsessioninitialized callback (backwards compatibility)", async () => { + const capturedSessionId: string[] = []; + + // Create server with sync onsessioninitialized callback + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: (sessionId: string) => { + capturedSessionId.push(sessionId); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize to trigger the callback + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + const tempSessionId = initResponse.headers.get("mcp-session-id"); + + expect(capturedSessionId).toEqual([tempSessionId]); + + // Clean up + tempServer.close(); + }); + + it("should support async onsessionclosed callback", async () => { + const closureOrder: string[] = []; + + // Create server with async onsessionclosed callback + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessionclosed: async (sessionId: string) => { + closureOrder.push('async-close-start'); + // Simulate async operation + await new Promise(resolve => setTimeout(resolve, 10)); + closureOrder.push('async-close-end'); + closureOrder.push(sessionId); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize to get a session ID + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + const tempSessionId = initResponse.headers.get("mcp-session-id"); + expect(tempSessionId).toBeDefined(); + + // DELETE the session + const deleteResponse = await fetch(tempUrl, { + method: "DELETE", + headers: { + "mcp-session-id": tempSessionId || "", + "mcp-protocol-version": "2025-03-26", + }, + }); + + expect(deleteResponse.status).toBe(200); + + // Give time for async callback to complete + await new Promise(resolve => setTimeout(resolve, 50)); + + expect(closureOrder).toEqual(['async-close-start', 'async-close-end', tempSessionId]); + + // Clean up + tempServer.close(); + }); + + it("should propagate errors from async onsessioninitialized callback", async () => { + const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(); + + // Create server with async onsessioninitialized callback that throws + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: async (_sessionId: string) => { + throw new Error('Async initialization error'); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize should fail when callback throws + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + expect(initResponse.status).toBe(400); + + // Clean up + consoleErrorSpy.mockRestore(); + tempServer.close(); + }); + + it("should propagate errors from async onsessionclosed callback", async () => { + const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(); + + // Create server with async onsessionclosed callback that throws + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessionclosed: async (_sessionId: string) => { + throw new Error('Async closure error'); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize to get a session ID + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + const tempSessionId = initResponse.headers.get("mcp-session-id"); + + // DELETE should fail when callback throws + const deleteResponse = await fetch(tempUrl, { + method: "DELETE", + headers: { + "mcp-session-id": tempSessionId || "", + "mcp-protocol-version": "2025-03-26", + }, + }); + + expect(deleteResponse.status).toBe(500); + + // Clean up + consoleErrorSpy.mockRestore(); + tempServer.close(); + }); + + it("should handle both async callbacks together", async () => { + const events: string[] = []; + + // Create server with both async callbacks + const result = await createTestServer({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: async (sessionId: string) => { + await new Promise(resolve => setTimeout(resolve, 5)); + events.push(`initialized:${sessionId}`); + }, + onsessionclosed: async (sessionId: string) => { + await new Promise(resolve => setTimeout(resolve, 5)); + events.push(`closed:${sessionId}`); + }, + }); + + const tempServer = result.server; + const tempUrl = result.baseUrl; + + // Initialize to trigger first callback + const initResponse = await sendPostRequest(tempUrl, TEST_MESSAGES.initialize); + const tempSessionId = initResponse.headers.get("mcp-session-id"); + + // Wait for async callback + await new Promise(resolve => setTimeout(resolve, 20)); + + expect(events).toContain(`initialized:${tempSessionId}`); + + // DELETE to trigger second callback + const deleteResponse = await fetch(tempUrl, { + method: "DELETE", + headers: { + "mcp-session-id": tempSessionId || "", + "mcp-protocol-version": "2025-03-26", + }, + }); + + expect(deleteResponse.status).toBe(200); + + // Wait for async callback + await new Promise(resolve => setTimeout(resolve, 20)); + + expect(events).toContain(`closed:${tempSessionId}`); + expect(events).toHaveLength(2); + + // Clean up + tempServer.close(); + }); +}); + // Test DNS rebinding protection describe("StreamableHTTPServerTransport DNS rebinding protection", () => { let server: Server; diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts index 37164c86..3bf84e43 100644 --- a/src/server/streamableHttp.ts +++ b/src/server/streamableHttp.ts @@ -47,7 +47,7 @@ export interface StreamableHTTPServerTransportOptions { * and need to keep track of them. * @param sessionId The generated session ID */ - onsessioninitialized?: (sessionId: string) => void; + onsessioninitialized?: (sessionId: string) => void | Promise; /** * A callback for session close events @@ -59,7 +59,7 @@ export interface StreamableHTTPServerTransportOptions { * session open/running. * @param sessionId The session ID that was closed */ - onsessionclosed?: (sessionId: string) => void; + onsessionclosed?: (sessionId: string) => void | Promise; /** * If true, the server will return JSON responses instead of starting an SSE stream. @@ -138,8 +138,8 @@ export class StreamableHTTPServerTransport implements Transport { private _enableJsonResponse: boolean = false; private _standaloneSseStreamId: string = '_GET_stream'; private _eventStore?: EventStore; - private _onsessioninitialized?: (sessionId: string) => void; - private _onsessionclosed?: (sessionId: string) => void; + private _onsessioninitialized?: (sessionId: string) => void | Promise; + private _onsessionclosed?: (sessionId: string) => void | Promise; private _allowedHosts?: string[]; private _allowedOrigins?: string[]; private _enableDnsRebindingProtection: boolean; @@ -460,7 +460,7 @@ export class StreamableHTTPServerTransport implements Transport { // If we have a session ID and an onsessioninitialized handler, call it immediately // This is needed in cases where the server needs to keep track of multiple sessions if (this.sessionId && this._onsessioninitialized) { - this._onsessioninitialized(this.sessionId); + await Promise.resolve(this._onsessioninitialized(this.sessionId)); } } @@ -552,7 +552,7 @@ export class StreamableHTTPServerTransport implements Transport { if (!this.validateProtocolVersion(req, res)) { return; } - this._onsessionclosed?.(this.sessionId!); + await Promise.resolve(this._onsessionclosed?.(this.sessionId!)); await this.close(); res.writeHead(200).end(); }