diff --git a/src/client/sse.ts b/src/client/sse.ts index 14921f57..173da8ab 100644 --- a/src/client/sse.ts +++ b/src/client/sse.ts @@ -1,6 +1,7 @@ import { EventSource, type ErrorEvent, type EventSourceInit } from "eventsource"; import { Transport } from "../shared/transport.js"; import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; +import { auth, AuthResult, OAuthClientProvider } from "./auth.js"; export class SseError extends Error { constructor( @@ -12,6 +13,34 @@ export class SseError extends Error { } } +/** + * Configuration options for the `SSEClientTransport`. + */ +export type SSEClientTransportOptions = { + /** + * An OAuth client provider to use for authentication. + * + * If given, the transport will automatically attach an `Authorization` header + * if an access token is available, or begin the authorization flow if not. + */ + authProvider?: OAuthClientProvider; + + /** + * Customizes the initial SSE request to the server (the request that begins the stream). + * + * NOTE: Setting this property will prevent an `Authorization` header from + * being automatically attached to the SSE request, if an `authProvider` is + * also given. This can be worked around by setting the `Authorization` header + * manually. + */ + eventSourceInit?: EventSourceInit; + + /** + * Customizes recurring POST requests to the server. + */ + requestInit?: RequestInit; +}; + /** * Client transport for SSE: this will connect to a server using Server-Sent Events for receiving * messages and make separate POST requests for sending messages. @@ -23,6 +52,7 @@ export class SSEClientTransport implements Transport { private _url: URL; private _eventSourceInit?: EventSourceInit; private _requestInit?: RequestInit; + private _authProvider?: OAuthClientProvider; onclose?: () => void; onerror?: (error: Error) => void; @@ -30,28 +60,62 @@ export class SSEClientTransport implements Transport { constructor( url: URL, - opts?: { eventSourceInit?: EventSourceInit; requestInit?: RequestInit }, + opts?: SSEClientTransportOptions, ) { this._url = url; this._eventSourceInit = opts?.eventSourceInit; this._requestInit = opts?.requestInit; + this._authProvider = opts?.authProvider; } - start(): Promise { - if (this._eventSource) { - throw new Error( - "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically.", - ); + private async _authThenStart(): Promise { + if (!this._authProvider) { + throw new Error("No auth provider"); } + let result: AuthResult; + try { + result = await auth(this._authProvider, { serverUrl: this._url }); + } catch (error) { + this.onerror?.(error as Error); + throw error; + } + + if (result !== "AUTHORIZED") { + throw new Error("Unauthorized"); + } + + return await this._startOrAuth(); + } + + private async _commonHeaders(): Promise { + const headers: HeadersInit = {}; + if (this._authProvider) { + const tokens = await this._authProvider.tokens(); + if (tokens) { + headers["Authorization"] = `Bearer ${tokens.access_token}`; + } + } + + return headers; + } + + private _startOrAuth(): Promise { return new Promise((resolve, reject) => { this._eventSource = new EventSource( this._url.href, - this._eventSourceInit, + this._eventSourceInit ?? { + fetch: (url, init) => this._commonHeaders().then((headers) => fetch(url, { ...init, headers })), + }, ); this._abortController = new AbortController(); this._eventSource.onerror = (event) => { + if (event.code === 401 && this._authProvider) { + this._authThenStart().then(resolve, reject); + return; + } + const error = new SseError(event.code, event.message, event); reject(error); this.onerror?.(error); @@ -97,6 +161,16 @@ export class SSEClientTransport implements Transport { }); } + async start() { + if (this._eventSource) { + throw new Error( + "SSEClientTransport already started! If using Client class, note that connect() calls start() automatically.", + ); + } + + return await this._startOrAuth(); + } + async close(): Promise { this._abortController?.abort(); this._eventSource?.close(); @@ -109,7 +183,8 @@ export class SSEClientTransport implements Transport { } try { - const headers = new Headers(this._requestInit?.headers); + const commonHeaders = await this._commonHeaders(); + const headers = new Headers({ ...commonHeaders, ...this._requestInit?.headers }); headers.set("content-type", "application/json"); const init = { ...this._requestInit,