From 37ca02d36df56487430f13a6b4b6c0889647933b Mon Sep 17 00:00:00 2001
From: Nikhil Taneja <80670724+its-nikhil@users.noreply.github.com>
Date: Thu, 24 Apr 2025 19:22:55 +0530
Subject: [PATCH] [FEAT] added support for AuthInfo in extra for
 StreamableHTTPServerTransport

---
 src/server/streamableHttp.ts | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts
index c9051073..f88842fa 100644
--- a/src/server/streamableHttp.ts
+++ b/src/server/streamableHttp.ts
@@ -4,6 +4,7 @@ import { isInitializeRequest, isJSONRPCError, isJSONRPCRequest, isJSONRPCRespons
 import getRawBody from "raw-body";
 import contentType from "content-type";
 import { randomUUID } from "node:crypto";
+import { AuthInfo } from "./auth/types.js";
 
 const MAXIMUM_MESSAGE_SIZE = "4mb";
 
@@ -112,7 +113,7 @@ export class StreamableHTTPServerTransport implements Transport {
   sessionId?: string | undefined;
   onclose?: () => void;
   onerror?: (error: Error) => void;
-  onmessage?: (message: JSONRPCMessage) => void;
+  onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
 
   constructor(options: StreamableHTTPServerTransportOptions) {
     this.sessionIdGenerator = options.sessionIdGenerator;
@@ -286,7 +287,7 @@ export class StreamableHTTPServerTransport implements Transport {
   /**
    * Handles POST requests containing JSON-RPC messages
    */
-  private async handlePostRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise<void> {
+  private async handlePostRequest(req: IncomingMessage & { auth?: AuthInfo }, res: ServerResponse, parsedBody?: unknown): Promise<void> {
     try {
       // Validate the Accept header
       const acceptHeader = req.headers.accept;
@@ -316,6 +317,8 @@ export class StreamableHTTPServerTransport implements Transport {
         return;
       }
 
+      const authInfo: AuthInfo | undefined = req.auth;
+
       let rawMessage;
       if (parsedBody !== undefined) {
         rawMessage = parsedBody;
@@ -392,7 +395,7 @@ export class StreamableHTTPServerTransport implements Transport {
 
         // handle each message
         for (const message of messages) {
-          this.onmessage?.(message);
+          this.onmessage?.(message, { authInfo });
         }
       } else if (hasRequests) {
         // The default behavior is to use SSE streaming
@@ -427,7 +430,7 @@ export class StreamableHTTPServerTransport implements Transport {
 
         // handle each message
         for (const message of messages) {
-          this.onmessage?.(message);
+          this.onmessage?.(message, { authInfo });
         }
         // The server SHOULD NOT close the SSE stream before sending all JSON-RPC responses
         // This will be handled by the send() method when responses are ready