Include Authorization Info in Tool Calls (and request handlers generally)

src/server/sse.ts

@@ -4,6 +4,7 @@ import { Transport } from "../shared/transport.js";
 import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 import getRawBody from "raw-body";
 import contentType from "content-type";
+import { AuthInfo } from "./auth/types.js";
 
 const MAXIMUM_MESSAGE_SIZE = "4mb";
 
@@ -18,7 +19,7 @@ export class SSEServerTransport implements Transport {
 
   onclose?: () => void;
   onerror?: (error: Error) => void;
-  onmessage?: (message: JSONRPCMessage) => void;
+  onmessage?: (message: JSONRPCMessage, 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`.
@@ -66,7 +67,7 @@ export class SSEServerTransport implements Transport {
    * This should be called when a POST request is made to send a message to the server.
    */
   async handlePostMessage(
-    req: IncomingMessage,
+    req: IncomingMessage & { auth?: AuthInfo },
     res: ServerResponse,
     parsedBody?: unknown,
   ): Promise<void> {
@@ -75,6 +76,7 @@ export class SSEServerTransport implements Transport {
       res.writeHead(500).end(message);
       throw new Error(message);
     }
+    const authInfo: AuthInfo | undefined = req.auth;
 
     let body: string | unknown;
     try {
@@ -94,7 +96,7 @@ export class SSEServerTransport implements Transport {
     }
 
     try {
-      await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body);
+      await this.handleMessage(typeof body === 'string' ? JSON.parse(body) : body, authInfo);
     } catch {
       res.writeHead(400).end(`Invalid message: ${body}`);
       return;
@@ -106,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): Promise<void> {
+  async handleMessage(message: unknown, authInfo?: AuthInfo): Promise<void> {
     let parsedMessage: JSONRPCMessage;
     try {
       parsedMessage = JSONRPCMessageSchema.parse(message);
@@ -115,7 +117,7 @@ export class SSEServerTransport implements Transport {
       throw error;
     }
 
-    this.onmessage?.(parsedMessage);
+    this.onmessage?.(parsedMessage, authInfo);
   }
 
   async close(): Promise<void> {

src/shared/protocol.ts

@@ -19,6 +19,7 @@ import {
   ServerCapabilities,
 } from "../types.js";
 import { Transport } from "./transport.js";
+import { AuthInfo } from "../server/auth/types.js";
 
 /**
  * Callback for progress notifications.
@@ -89,6 +90,11 @@ export type RequestHandlerExtra = {
    */
   signal: AbortSignal;
 
+  /**
+   * Information about a validated access token, provided to request handlers.
+   */
+  authInfo?: AuthInfo;
+
   /**
    * The session ID from the transport, if available.
    */
@@ -240,11 +246,11 @@ export abstract class Protocol<
       this._onerror(error);
     };
 
-    this._transport.onmessage = (message) => {
+    this._transport.onmessage = (message, authInfo) => {
       if (!("method" in message)) {
         this._onresponse(message);
       } else if ("id" in message) {
-        this._onrequest(message);
+        this._onrequest(message, authInfo);
       } else {
         this._onnotification(message);
       }
@@ -290,7 +296,7 @@ export abstract class Protocol<
       );
   }
 
-  private _onrequest(request: JSONRPCRequest): void {
+  private _onrequest(request: JSONRPCRequest, authInfo?: AuthInfo): void {
     const handler =
       this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;
 
@@ -319,6 +325,7 @@ export abstract class Protocol<
     const extra: RequestHandlerExtra = {
       signal: abortController.signal,
       sessionId: this._transport?.sessionId,
+      authInfo,
     };
 
     // Starting with Promise.resolve() puts any synchronous errors into the monad as well.

src/shared/transport.ts

@@ -1,3 +1,4 @@
+import { AuthInfo } from "../server/auth/types.js";
 import { JSONRPCMessage } from "../types.js";
 
 /**
@@ -39,8 +40,11 @@ export interface Transport {
 
   /**
    * Callback for when a message (request or response) is received over the connection.
+   * 
+   * Includes the authInfo if the transport is authenticated.
+   * 
    */
-  onmessage?: (message: JSONRPCMessage) => void;
+  onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
 
   /**
    * The session ID generated for this connection.