Merge remote-tracking branch 'upstream/main' into request-auth-info

src/shared/protocol.test.ts

@@ -71,6 +71,44 @@ describe("protocol tests", () => {
       jest.useRealTimers();
     });
 
+    test("should not reset timeout when resetTimeoutOnProgress is false", async () => {
+      await protocol.connect(transport);
+      const request = { method: "example", params: {} };
+      const mockSchema: ZodType<{ result: string }> = z.object({
+        result: z.string(),
+      });
+      const onProgressMock = jest.fn();
+      const requestPromise = protocol.request(request, mockSchema, {
+        timeout: 1000,
+        resetTimeoutOnProgress: false,
+        onprogress: onProgressMock,
+      });
+      
+      jest.advanceTimersByTime(800);
+      
+      if (transport.onmessage) {
+        transport.onmessage({
+          jsonrpc: "2.0",
+          method: "notifications/progress",
+          params: {
+            progressToken: 0,
+            progress: 50,
+            total: 100,
+          },
+        });
+      }
+      await Promise.resolve();
+      
+      expect(onProgressMock).toHaveBeenCalledWith({
+        progress: 50,
+        total: 100,
+      });
+      
+      jest.advanceTimersByTime(201);
+      
+      await expect(requestPromise).rejects.toThrow("Request timed out");
+    });
+
     test("should reset timeout when progress notification is received", async () => {
       await protocol.connect(transport);
       const request = { method: "example", params: {} };

src/shared/protocol.ts

@@ -94,6 +94,11 @@ export type RequestHandlerExtra = {
    * Information about a validated access token, provided to request handlers.
    */
   authInfo?: AuthInfo;
+
+  /**
+   * The session ID from the transport, if available.
+   */
+  sessionId?: string;
 };
 
 /**
@@ -104,6 +109,7 @@ type TimeoutInfo = {
   startTime: number;
   timeout: number;
   maxTotalTimeout?: number;
+  resetTimeoutOnProgress: boolean;
   onTimeout: () => void;
 };
 
@@ -185,13 +191,15 @@ export abstract class Protocol<
     messageId: number,
     timeout: number,
     maxTotalTimeout: number | undefined,
-    onTimeout: () => void
+    onTimeout: () => void,
+    resetTimeoutOnProgress: boolean = false
   ) {
     this._timeoutInfo.set(messageId, {
       timeoutId: setTimeout(onTimeout, timeout),
       startTime: Date.now(),
       timeout,
       maxTotalTimeout,
+      resetTimeoutOnProgress,
       onTimeout
     });
   }
@@ -313,9 +321,16 @@ 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 = {
+      signal: abortController.signal,
+      sessionId: this._transport?.sessionId,
+      authInfo,
+    };
+
     // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
     Promise.resolve()
-      .then(() => handler(request, { signal: abortController.signal, authInfo }))
+      .then(() => handler(request, extra))
       .then(
         (result) => {
           if (abortController.signal.aborted) {
@@ -364,7 +379,9 @@ export abstract class Protocol<
     }
 
     const responseHandler = this._responseHandlers.get(messageId);
-    if (this._timeoutInfo.has(messageId) && responseHandler) {
+    const timeoutInfo = this._timeoutInfo.get(messageId);
+    
+    if (timeoutInfo && responseHandler && timeoutInfo.resetTimeoutOnProgress) {
       try {
         this._resetTimeout(messageId);
       } catch (error) {
@@ -526,7 +543,7 @@ export abstract class Protocol<
         { timeout }
       ));
 
-      this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler);
+      this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false);
 
       this._transport.send(jsonrpcRequest).catch((error) => {
         this._cleanupTimeout(messageId);

src/shared/transport.ts

@@ -45,4 +45,9 @@ export interface Transport {
    * 
    */
   onmessage?: (message: JSONRPCMessage, authInfo?: AuthInfo) => void;
+
+  /**
+   * The session ID generated for this connection.
+   */
+  sessionId?: string;
 }