Merge pull request #152 from Ozamatash/feature/progress-timeout-reset

Add timeout reset on progress notifications

Author: jspahrsummers

src/shared/protocol.test.ts

@@ -62,6 +62,162 @@ describe("protocol tests", () => {
  await transport.close();
  expect(oncloseMock).toHaveBeenCalled();
  });
+
+ describe("progress notification timeout behavior", () => {
+ beforeEach(() => {
+ jest.useFakeTimers();
+ });
+ afterEach(() => {
+ jest.useRealTimers();
+ });
+
+ test("should reset timeout when progress notification is received", 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: true,
+ 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(800);
+ if (transport.onmessage) {
+ transport.onmessage({
+ jsonrpc: "2.0",
+ id: 0,
+ result: { result: "success" },
+ });
+ }
+ await Promise.resolve();
+ await expect(requestPromise).resolves.toEqual({ result: "success" });
+ });
+
+ test("should respect maxTotalTimeout", 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,
+ maxTotalTimeout: 150,
+ resetTimeoutOnProgress: true,
+ onprogress: onProgressMock,
+ });
+
+ // First progress notification should work
+ jest.advanceTimersByTime(80);
+ 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(80);
+ if (transport.onmessage) {
+ transport.onmessage({
+ jsonrpc: "2.0",
+ method: "notifications/progress",
+ params: {
+ progressToken: 0,
+ progress: 75,
+ total: 100,
+ },
+ });
+ }
+ await expect(requestPromise).rejects.toThrow("Maximum total timeout exceeded");
+ expect(onProgressMock).toHaveBeenCalledTimes(1);
+ });
+
+ test("should timeout if no progress received within timeout period", async () => {
+ await protocol.connect(transport);
+ const request = { method: "example", params: {} };
+ const mockSchema: ZodType<{ result: string }> = z.object({
+ result: z.string(),
+ });
+ const requestPromise = protocol.request(request, mockSchema, {
+ timeout: 100,
+ resetTimeoutOnProgress: true,
+ });
+ jest.advanceTimersByTime(101);
+ await expect(requestPromise).rejects.toThrow("Request timed out");
+ });
+
+ test("should handle multiple progress notifications correctly", 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: true,
+ onprogress: onProgressMock,
+ });
+
+ // Simulate multiple progress updates
+ for (let i = 1; i <= 3; i++) {
+ jest.advanceTimersByTime(800);
+ if (transport.onmessage) {
+ transport.onmessage({
+ jsonrpc: "2.0",
+ method: "notifications/progress",
+ params: {
+ progressToken: 0,
+ progress: i * 25,
+ total: 100,
+ },
+ });
+ }
+ await Promise.resolve();
+ expect(onProgressMock).toHaveBeenNthCalledWith(i, {
+ progress: i * 25,
+ total: 100,
+ });
+ }
+ if (transport.onmessage) {
+ transport.onmessage({
+ jsonrpc: "2.0",
+ id: 0,
+ result: { result: "success" },
+ });
+ }
+ await Promise.resolve();
+ await expect(requestPromise).resolves.toEqual({ result: "success" });
+ });
+ });
 });
 
 describe("mergeCapabilities", () => {