raw request propagation in tools - implementation, unit tests, types

src/shared/protocol.test.ts

@@ -27,9 +27,11 @@ class MockTransport implements Transport {
 describe("protocol tests", () => {
   let protocol: Protocol<Request, Notification, Result>;
   let transport: MockTransport;
+  let sendSpy: jest.SpyInstance;
 
   beforeEach(() => {
     transport = new MockTransport();
+    sendSpy = jest.spyOn(transport, 'send');
     protocol = new (class extends Protocol<Request, Notification, Result> {
       protected assertCapabilityForMethod(): void {}
       protected assertNotificationCapability(): void {}
@@ -63,6 +65,130 @@ describe("protocol tests", () => {
     expect(oncloseMock).toHaveBeenCalled();
   });
 
+  describe("_meta preservation with onprogress", () => {
+    test("should preserve existing _meta when adding progressToken", async () => {
+      await protocol.connect(transport);
+      const request = { 
+        method: "example", 
+        params: {
+          data: "test",
+          _meta: {
+            customField: "customValue",
+            anotherField: 123
+          }
+        }
+      };
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      const onProgressMock = jest.fn();
+      
+      protocol.request(request, mockSchema, {
+        onprogress: onProgressMock,
+      });
+      
+      expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({
+        method: "example",
+        params: {
+          data: "test",
+          _meta: {
+            customField: "customValue",
+            anotherField: 123,
+            progressToken: expect.any(Number)
+          }
+        },
+        jsonrpc: "2.0",
+        id: expect.any(Number)
+      }), expect.any(Object));
+    });
+
+    test("should create _meta with progressToken when no _meta exists", async () => {
+      await protocol.connect(transport);
+      const request = { 
+        method: "example", 
+        params: {
+          data: "test"
+        }
+      };
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      const onProgressMock = jest.fn();
+      
+      protocol.request(request, mockSchema, {
+        onprogress: onProgressMock,
+      });
+      
+      expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({
+        method: "example",
+        params: {
+          data: "test",
+          _meta: {
+            progressToken: expect.any(Number)
+          }
+        },
+        jsonrpc: "2.0",
+        id: expect.any(Number)
+      }), expect.any(Object));
+    });
+
+    test("should not modify _meta when onprogress is not provided", async () => {
+      await protocol.connect(transport);
+      const request = { 
+        method: "example", 
+        params: {
+          data: "test",
+          _meta: {
+            customField: "customValue"
+          }
+        }
+      };
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      
+      protocol.request(request, mockSchema);
+      
+      expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({
+        method: "example",
+        params: {
+          data: "test",
+          _meta: {
+            customField: "customValue"
+          }
+        },
+        jsonrpc: "2.0",
+        id: expect.any(Number)
+      }), expect.any(Object));
+    });
+
+    test("should handle params being undefined with onprogress", async () => {
+      await protocol.connect(transport);
+      const request = { 
+        method: "example"
+      };
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      const onProgressMock = jest.fn();
+      
+      protocol.request(request, mockSchema, {
+        onprogress: onProgressMock,
+      });
+      
+      expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({
+        method: "example",
+        params: {
+          _meta: {
+            progressToken: expect.any(Number)
+          }
+        },
+        jsonrpc: "2.0",
+        id: expect.any(Number)
+      }), expect.any(Object));
+    });
+  });
+
   describe("progress notification timeout behavior", () => {
     beforeEach(() => {
       jest.useFakeTimers();

src/shared/protocol.ts

@@ -25,6 +25,7 @@ import {
 } from "../types.js";
 import { Transport, TransportSendOptions } from "./transport.js";
 import { AuthInfo } from "../server/auth/types.js";
+import { MessageExtraInfo, RequestInfo } from "../server/types/types.js";
 
 /**
  * Callback for progress notifications.
@@ -127,6 +128,11 @@ export type RequestHandlerExtra<SendRequestT extends Request,
      */
     requestId: RequestId;
 
+    /**
+     * The original HTTP request.
+     */
+    requestInfo: RequestInfo;
+
     /**
      * Sends a notification that relates to the current request being handled.
      * 
@@ -339,7 +345,7 @@ export abstract class Protocol<
       );
   }
 
-  private _onrequest(request: JSONRPCRequest, extra?: { authInfo?: AuthInfo }): void {
+  private _onrequest(request: JSONRPCRequest, extra: MessageExtraInfo): void {
     const handler =
       this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;
 
@@ -375,6 +381,7 @@ export abstract class Protocol<
         this.request(r, resultSchema, { ...options, relatedRequestId: request.id }),
       authInfo: extra?.authInfo,
       requestId: request.id,
+      requestInfo: extra.requestInfo
     };
 
     // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
@@ -541,7 +548,10 @@ export abstract class Protocol<
         this._progressHandlers.set(messageId, options.onprogress);
         jsonrpcRequest.params = {
           ...request.params,
-          _meta: { progressToken: messageId },
+          _meta: {
+            ...(request.params?._meta || {}),
+            progressToken: messageId
+          },
         };
       }
 

src/shared/transport.ts

@@ -1,4 +1,4 @@
-import { AuthInfo } from "../server/auth/types.js";
+import { MessageExtraInfo } from "../server/types/types.js";
 import { JSONRPCMessage, RequestId } from "../types.js";
 
 /**
@@ -66,10 +66,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.
+   * Includes the requestInfo and authInfo if the transport is authenticated.
    * 
+   * The requestInfo can be used to get the original request information (headers, etc.)
    */
-  onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
+  onmessage?: (message: JSONRPCMessage, extra: MessageExtraInfo) => void;
 
   /**
    * The session ID generated for this connection.