Pass extra object instead of authInfo directly

Author: allenzhou101

Diff for: 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 for: 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<void> {
-    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<void>
+  async send(message: JSONRPCMessage, options?: { authInfo?: AuthInfo }): Promise<void> {
     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 for: 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<void> {
+  async handleMessage(message: unknown, extra?: { authInfo?: AuthInfo }): Promise<void> {
     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<void> {

Diff for: 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) {