diff --git a/src/inMemory.test.ts b/src/inMemory.test.ts index b6906cc8..baf43446 100644 --- a/src/inMemory.test.ts +++ b/src/inMemory.test.ts @@ -52,12 +52,12 @@ describe("InMemoryTransport", () => { let receivedMessage: JSONRPCMessage | undefined; let receivedAuthInfo: AuthInfo | undefined; - serverTransport.onmessage = (msg, auth) => { + serverTransport.onmessage = (msg, extra) => { receivedMessage = msg; - receivedAuthInfo = auth; + receivedAuthInfo = extra?.authInfo; }; - await clientTransport.sendWithAuth(message, authInfo); + await clientTransport.send(message, { authInfo }); expect(receivedMessage).toEqual(message); expect(receivedAuthInfo).toEqual(authInfo); }); diff --git a/src/inMemory.ts b/src/inMemory.ts index 8889fdd8..558ca7ad 100644 --- a/src/inMemory.ts +++ b/src/inMemory.ts @@ -4,7 +4,7 @@ import { AuthInfo } from "./server/auth/types.js"; interface QueuedMessage { message: JSONRPCMessage; - authInfo?: AuthInfo; + extra?: { authInfo?: AuthInfo }; } /** @@ -16,7 +16,7 @@ export class InMemoryTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; - onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void; + onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void; sessionId?: string; /** @@ -34,7 +34,7 @@ export class InMemoryTransport implements Transport { // Process any messages that were queued before start was called while (this._messageQueue.length > 0) { const queuedMessage = this._messageQueue.shift()!; - this.onmessage?.(queuedMessage.message, queuedMessage.authInfo); + this.onmessage?.(queuedMessage.message, queuedMessage.extra); } } @@ -45,23 +45,19 @@ export class InMemoryTransport implements Transport { this.onclose?.(); } - async send(message: JSONRPCMessage): Promise { - return this.sendWithAuth(message); - } - /** * Sends a message with optional auth info. * This is useful for testing authentication scenarios. */ - async send(message: JSONRPCMessage, options?: { authInfo?: AuthInfo }): Promise + async send(message: JSONRPCMessage, options?: { authInfo?: AuthInfo }): Promise { if (!this._otherTransport) { throw new Error("Not connected"); } if (this._otherTransport.onmessage) { - this._otherTransport.onmessage(message, authInfo); + this._otherTransport.onmessage(message, { authInfo: options?.authInfo }); } else { - this._otherTransport._messageQueue.push({ message, authInfo }); + this._otherTransport._messageQueue.push({ message, extra: { authInfo: options?.authInfo } }); } } } diff --git a/src/server/sse.ts b/src/server/sse.ts index 848043ba..ae993dd7 100644 --- a/src/server/sse.ts +++ b/src/server/sse.ts @@ -19,7 +19,7 @@ export class SSEServerTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; - onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void; + onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void; /** * Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL identified by `_endpoint`. @@ -96,7 +96,7 @@ export class SSEServerTransport implements Transport { } try { - await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, authInfo); + await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, { authInfo }); } catch { res.writeHead(400).end(`Invalid message: ${body}`); return; @@ -108,7 +108,7 @@ export class SSEServerTransport implements Transport { /** * Handle a client message, regardless of how it arrived. This can be used to inform the server of messages that arrive via a means different than HTTP POST. */ - async handleMessage(message: unknown, authInfo?: AuthInfo): Promise { + async handleMessage(message: unknown, extra?: { authInfo?: AuthInfo }): Promise { let parsedMessage: JSONRPCMessage; try { parsedMessage = JSONRPCMessageSchema.parse(message); @@ -117,7 +117,7 @@ export class SSEServerTransport implements Transport { throw error; } - this.onmessage?.(parsedMessage, authInfo); + this.onmessage?.(parsedMessage, extra); } async close(): Promise { diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 789f5a7e..a39eba68 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -246,11 +246,11 @@ export abstract class Protocol< this._onerror(error); }; - this._transport.onmessage = (message, authInfo) => { + this._transport.onmessage = (message, extra) => { if (!("method" in message)) { this._onresponse(message); } else if ("id" in message) { - this._onrequest(message, authInfo); + this._onrequest(message, extra); } else { this._onnotification(message); } @@ -296,7 +296,7 @@ export abstract class Protocol< ); } - private _onrequest(request: JSONRPCRequest, authInfo?: AuthInfo): void { + private _onrequest(request: JSONRPCRequest, extra?: { authInfo?: AuthInfo }): void { const handler = this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler; @@ -321,16 +321,17 @@ export abstract class Protocol< const abortController = new AbortController(); this._requestHandlerAbortControllers.set(request.id, abortController); - // Create extra object with both abort signal and sessionId from transport - const extra: RequestHandlerExtra = { + // Create fullExtra object with both abort signal and sessionId from transport + // in addition to any authInfo provided + const fullExtra: RequestHandlerExtra = { signal: abortController.signal, sessionId: this._transport?.sessionId, - authInfo, + authInfo: extra?.authInfo, }; // Starting with Promise.resolve() puts any synchronous errors into the monad as well. Promise.resolve() - .then(() => handler(request, extra)) + .then(() => handler(request, fullExtra)) .then( (result) => { if (abortController.signal.aborted) {