Merge branch 'main' into main

Author: nanbingxyz

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

@@ -103,6 +103,7 @@ type TimeoutInfo = {
   startTime: number;
   timeout: number;
   maxTotalTimeout?: number;
+  resetTimeoutOnProgress: boolean;
   onTimeout: () => void;
 };
 
@@ -184,13 +185,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
     });
   }
@@ -369,7 +372,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) {
@@ -531,7 +536,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/types.ts

@@ -666,6 +666,23 @@ export const ImageContentSchema = z
   })
   .passthrough();
 
+/**
+ * An Audio provided to or from an LLM.
+ */
+export const AudioContentSchema = z
+  .object({
+    type: z.literal("audio"),
+    /**
+     * The base64-encoded audio data.
+     */
+    data: z.string().base64(),
+    /**
+     * The MIME type of the audio. Different providers may support different audio types.
+     */
+    mimeType: z.string(),
+  })
+  .passthrough();
+
 /**
  * The contents of a resource, embedded into a prompt or tool call result.
  */
@@ -685,6 +702,7 @@ export const PromptMessageSchema = z
     content: z.union([
       TextContentSchema,
       ImageContentSchema,
+      AudioContentSchema,
       EmbeddedResourceSchema,
     ]),
   })
@@ -753,7 +771,7 @@ export const ListToolsResultSchema = PaginatedResultSchema.extend({
  */
 export const CallToolResultSchema = ResultSchema.extend({
   content: z.array(
-    z.union([TextContentSchema, ImageContentSchema, EmbeddedResourceSchema]),
+    z.union([TextContentSchema, ImageContentSchema, AudioContentSchema, EmbeddedResourceSchema]),
   ),
   isError: z.boolean().default(false).optional(),
 });
@@ -877,7 +895,7 @@ export const ModelPreferencesSchema = z
 export const SamplingMessageSchema = z
   .object({
     role: z.enum(["user", "assistant"]),
-    content: z.union([TextContentSchema, ImageContentSchema]),
+    content: z.union([TextContentSchema, ImageContentSchema, AudioContentSchema]),
   })
   .passthrough();
 
@@ -931,6 +949,7 @@ export const CreateMessageResultSchema = ResultSchema.extend({
   content: z.discriminatedUnion("type", [
     TextContentSchema,
     ImageContentSchema,
+    AudioContentSchema
   ]),
 });
 
@@ -1195,6 +1214,7 @@ export type ListPromptsResult = Infer<typeof ListPromptsResultSchema>;
 export type GetPromptRequest = Infer<typeof GetPromptRequestSchema>;
 export type TextContent = Infer<typeof TextContentSchema>;
 export type ImageContent = Infer<typeof ImageContentSchema>;
+export type AudioContent = Infer<typeof AudioContentSchema>;
 export type EmbeddedResource = Infer<typeof EmbeddedResourceSchema>;
 export type PromptMessage = Infer<typeof PromptMessageSchema>;
 export type GetPromptResult = Infer<typeof GetPromptResultSchema>;