Skip to content

Commit fd14b33

Browse files
LucaButBoringfelixweinberger
LucaButBoring
authored and
felixweinberger
committed
chore: replace timed events with getTask logic in task tests
This removes the setTimeout logic we had in tests, which was masking an issue where the getTask handlers weren't being called. The appropriate logic has been moved into the getTask handlers themselves.
1 parent 8ae7217 commit fd14b33

File tree

1 file changed

+27
-103
lines changed

1 file changed

+27
-103
lines changed

‎test/server/mcp.test.ts‎

Lines changed: 27 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -26,22 +26,6 @@ import { McpServer, ResourceTemplate } from '../../src/server/mcp.js';
2626
import { InMemoryTaskStore } from '../../src/experimental/tasks/stores/in-memory.js';
2727
import { zodTestMatrix, type ZodMatrixEntry } from '../../src/__fixtures__/zodTestMatrix.js';
2828

29-
function createLatch() {
30-
let latch = false;
31-
const waitForLatch = async () => {
32-
while (!latch) {
33-
await new Promise(resolve => setTimeout(resolve, 0));
34-
}
35-
};
36-
37-
return {
38-
releaseLatch: () => {
39-
latch = true;
40-
},
41-
waitForLatch
42-
};
43-
}
44-
4529
describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
4630
const { z } = entry;
4731

@@ -6334,31 +6318,15 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
63346318
}
63356319
},
63366320
{
6337-
createTask: async ({ input }, extra) => {
6321+
createTask: async (_args, extra) => {
63386322
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6339-
6340-
// Capture taskStore for use in setTimeout
6341-
const store = extra.taskStore;
6342-
6343-
// Simulate async work
6344-
setTimeout(async () => {
6345-
await store.storeTaskResult(task.taskId, 'completed', {
6346-
content: [{ type: 'text' as const, text: `Processed: ${input}` }]
6347-
});
6348-
}, 200);
6349-
63506323
return { task };
63516324
},
63526325
getTask: async extra => {
6353-
const task = await extra.taskStore.getTask(extra.taskId);
6354-
if (!task) {
6355-
throw new Error('Task not found');
6356-
}
6357-
return task;
6326+
return await extra.taskStore.getTask(extra.taskId);
63586327
},
63596328
getTaskResult: async extra => {
6360-
const result = await extra.taskStore.getTaskResult(extra.taskId);
6361-
return result as CallToolResult;
6329+
return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
63626330
}
63636331
}
63646332
);
@@ -6386,7 +6354,6 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
63866354

63876355
test('should automatically poll and return CallToolResult for tool with taskSupport "optional" called without task augmentation', async () => {
63886356
const taskStore = new InMemoryTaskStore();
6389-
const { releaseLatch, waitForLatch } = createLatch();
63906357

63916358
// Spies to verify handler invocations
63926359
const createTaskSpy = vi.fn();
@@ -6447,26 +6414,21 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
64476414
createTask: async ({ value }, extra) => {
64486415
createTaskSpy({ value });
64496416
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6450-
6451-
// Capture taskStore for use in setTimeout
6452-
const store = extra.taskStore;
6453-
6454-
// Simulate async work
6455-
setTimeout(async () => {
6456-
await store.storeTaskResult(task.taskId, 'completed', {
6457-
content: [{ type: 'text' as const, text: `Result: ${value * 2}` }]
6458-
});
6459-
releaseLatch();
6460-
}, 150);
6461-
6462-
return { task };
6417+
return { task, _value: value };
64636418
},
64646419
getTask: async extra => {
64656420
getTaskSpy(extra.taskId);
64666421
const task = await extra.taskStore.getTask(extra.taskId);
64676422
if (!task) {
64686423
throw new Error('Task not found');
64696424
}
6425+
// Complete the task when polled - proves getTask is actually called
6426+
if (task.status === 'working') {
6427+
await extra.taskStore.storeTaskResult(extra.taskId, 'completed', {
6428+
content: [{ type: 'text' as const, text: 'Result: 42' }]
6429+
});
6430+
return await extra.taskStore.getTask(extra.taskId);
6431+
}
64706432
return task;
64716433
},
64726434
getTaskResult: async extra => {
@@ -6501,14 +6463,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
65016463
expect(getTaskSpy).toHaveBeenCalled(); // Called at least once during polling
65026464
expect(getTaskResultSpy).toHaveBeenCalledOnce();
65036465

6504-
// Wait for async operations to complete
6505-
await waitForLatch();
65066466
taskStore.cleanup();
65076467
});
65086468

65096469
test('should return CreateTaskResult when tool with taskSupport "required" is called WITH task augmentation', async () => {
65106470
const taskStore = new InMemoryTaskStore();
6511-
const { releaseLatch, waitForLatch } = createLatch();
65126471

65136472
const mcpServer = new McpServer(
65146473
{
@@ -6561,32 +6520,15 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
65616520
}
65626521
},
65636522
{
6564-
createTask: async ({ data }, extra) => {
6523+
createTask: async (_args, extra) => {
65656524
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6566-
6567-
// Capture taskStore for use in setTimeout
6568-
const store = extra.taskStore;
6569-
6570-
// Simulate async work
6571-
setTimeout(async () => {
6572-
await store.storeTaskResult(task.taskId, 'completed', {
6573-
content: [{ type: 'text' as const, text: `Completed: ${data}` }]
6574-
});
6575-
releaseLatch();
6576-
}, 200);
6577-
65786525
return { task };
65796526
},
65806527
getTask: async extra => {
6581-
const task = await extra.taskStore.getTask(extra.taskId);
6582-
if (!task) {
6583-
throw new Error('Task not found');
6584-
}
6585-
return task;
6528+
return await extra.taskStore.getTask(extra.taskId);
65866529
},
65876530
getTaskResult: async extra => {
6588-
const result = await extra.taskStore.getTaskResult(extra.taskId);
6589-
return result as CallToolResult;
6531+
return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
65906532
}
65916533
}
65926534
);
@@ -6621,14 +6563,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
66216563
expect(result.task).toHaveProperty('taskId');
66226564
expect(result.task.status).toBe('working');
66236565

6624-
// Wait for async operations to complete
6625-
await waitForLatch();
66266566
taskStore.cleanup();
66276567
});
66286568

66296569
test('should handle task failures during automatic polling', async () => {
66306570
const taskStore = new InMemoryTaskStore();
6631-
const { releaseLatch, waitForLatch } = createLatch();
66326571

66336572
const mcpServer = new McpServer(
66346573
{
@@ -6680,26 +6619,21 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
66806619
{
66816620
createTask: async extra => {
66826621
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6683-
6684-
// Capture taskStore for use in setTimeout
6685-
const store = extra.taskStore;
6686-
6687-
// Simulate async failure
6688-
setTimeout(async () => {
6689-
await store.storeTaskResult(task.taskId, 'failed', {
6690-
content: [{ type: 'text' as const, text: 'Error occurred' }],
6691-
isError: true
6692-
});
6693-
releaseLatch();
6694-
}, 150);
6695-
66966622
return { task };
66976623
},
66986624
getTask: async extra => {
66996625
const task = await extra.taskStore.getTask(extra.taskId);
67006626
if (!task) {
67016627
throw new Error('Task not found');
67026628
}
6629+
// Fail the task when polled - proves getTask is actually called
6630+
if (task.status === 'working') {
6631+
await extra.taskStore.storeTaskResult(extra.taskId, 'failed', {
6632+
content: [{ type: 'text' as const, text: 'Error occurred' }],
6633+
isError: true
6634+
});
6635+
return await extra.taskStore.getTask(extra.taskId);
6636+
}
67036637
return task;
67046638
},
67056639
getTaskResult: async extra => {
@@ -6727,14 +6661,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
67276661
expect(result.content).toEqual([{ type: 'text' as const, text: 'Error occurred' }]);
67286662
expect(result.isError).toBe(true);
67296663

6730-
// Wait for async operations to complete
6731-
await waitForLatch();
67326664
taskStore.cleanup();
67336665
});
67346666

67356667
test('should handle task cancellation during automatic polling', async () => {
67366668
const taskStore = new InMemoryTaskStore();
6737-
const { releaseLatch, waitForLatch } = createLatch();
67386669

67396670
const mcpServer = new McpServer(
67406671
{
@@ -6786,23 +6717,18 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
67866717
{
67876718
createTask: async extra => {
67886719
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6789-
6790-
// Capture taskStore for use in setTimeout
6791-
const store = extra.taskStore;
6792-
6793-
// Simulate async cancellation
6794-
setTimeout(async () => {
6795-
await store.updateTaskStatus(task.taskId, 'cancelled', 'Task was cancelled');
6796-
releaseLatch();
6797-
}, 150);
6798-
67996720
return { task };
68006721
},
68016722
getTask: async extra => {
68026723
const task = await extra.taskStore.getTask(extra.taskId);
68036724
if (!task) {
68046725
throw new Error('Task not found');
68056726
}
6727+
// Cancel the task when polled - proves getTask is actually called
6728+
if (task.status === 'working') {
6729+
await extra.taskStore.updateTaskStatus(extra.taskId, 'cancelled', 'Task was cancelled');
6730+
return await extra.taskStore.getTask(extra.taskId);
6731+
}
68066732
return task;
68076733
},
68086734
getTaskResult: async extra => {
@@ -6829,8 +6755,6 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
68296755
expect(result).toHaveProperty('content');
68306756
expect(result.content).toEqual([{ type: 'text' as const, text: expect.stringContaining('has no result stored') }]);
68316757

6832-
// Wait for async operations to complete
6833-
await waitForLatch();
68346758
taskStore.cleanup();
68356759
});
68366760

0 commit comments

Comments
 (0)