Skip to content

Commit b8b9da9

Browse files
LucaButBoringLucaButBoring
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 6d37b16 commit b8b9da9

1 file changed

Lines changed: 27 additions & 103 deletions

File tree

‎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

@@ -6303,31 +6287,15 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
63036287
}
63046288
},
63056289
{
6306-
createTask: async ({ input }, extra) => {
6290+
createTask: async (_args, extra) => {
63076291
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6308-
6309-
// Capture taskStore for use in setTimeout
6310-
const store = extra.taskStore;
6311-
6312-
// Simulate async work
6313-
setTimeout(async () => {
6314-
await store.storeTaskResult(task.taskId, 'completed', {
6315-
content: [{ type: 'text' as const, text: `Processed: ${input}` }]
6316-
});
6317-
}, 200);
6318-
63196292
return { task };
63206293
},
63216294
getTask: async extra => {
6322-
const task = await extra.taskStore.getTask(extra.taskId);
6323-
if (!task) {
6324-
throw new Error('Task not found');
6325-
}
6326-
return task;
6295+
return await extra.taskStore.getTask(extra.taskId);
63276296
},
63286297
getTaskResult: async extra => {
6329-
const result = await extra.taskStore.getTaskResult(extra.taskId);
6330-
return result as CallToolResult;
6298+
return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
63316299
}
63326300
}
63336301
);
@@ -6355,7 +6323,6 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
63556323

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

63606327
// Spies to verify handler invocations
63616328
const createTaskSpy = vi.fn();
@@ -6416,26 +6383,21 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
64166383
createTask: async ({ value }, extra) => {
64176384
createTaskSpy({ value });
64186385
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6419-
6420-
// Capture taskStore for use in setTimeout
6421-
const store = extra.taskStore;
6422-
6423-
// Simulate async work
6424-
setTimeout(async () => {
6425-
await store.storeTaskResult(task.taskId, 'completed', {
6426-
content: [{ type: 'text' as const, text: `Result: ${value * 2}` }]
6427-
});
6428-
releaseLatch();
6429-
}, 150);
6430-
6431-
return { task };
6386+
return { task, _value: value };
64326387
},
64336388
getTask: async extra => {
64346389
getTaskSpy(extra.taskId);
64356390
const task = await extra.taskStore.getTask(extra.taskId);
64366391
if (!task) {
64376392
throw new Error('Task not found');
64386393
}
6394+
// Complete the task when polled - proves getTask is actually called
6395+
if (task.status === 'working') {
6396+
await extra.taskStore.storeTaskResult(extra.taskId, 'completed', {
6397+
content: [{ type: 'text' as const, text: 'Result: 42' }]
6398+
});
6399+
return await extra.taskStore.getTask(extra.taskId);
6400+
}
64396401
return task;
64406402
},
64416403
getTaskResult: async extra => {
@@ -6470,14 +6432,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
64706432
expect(getTaskSpy).toHaveBeenCalled(); // Called at least once during polling
64716433
expect(getTaskResultSpy).toHaveBeenCalledOnce();
64726434

6473-
// Wait for async operations to complete
6474-
await waitForLatch();
64756435
taskStore.cleanup();
64766436
});
64776437

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

64826441
const mcpServer = new McpServer(
64836442
{
@@ -6530,32 +6489,15 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
65306489
}
65316490
},
65326491
{
6533-
createTask: async ({ data }, extra) => {
6492+
createTask: async (_args, extra) => {
65346493
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6535-
6536-
// Capture taskStore for use in setTimeout
6537-
const store = extra.taskStore;
6538-
6539-
// Simulate async work
6540-
setTimeout(async () => {
6541-
await store.storeTaskResult(task.taskId, 'completed', {
6542-
content: [{ type: 'text' as const, text: `Completed: ${data}` }]
6543-
});
6544-
releaseLatch();
6545-
}, 200);
6546-
65476494
return { task };
65486495
},
65496496
getTask: async extra => {
6550-
const task = await extra.taskStore.getTask(extra.taskId);
6551-
if (!task) {
6552-
throw new Error('Task not found');
6553-
}
6554-
return task;
6497+
return await extra.taskStore.getTask(extra.taskId);
65556498
},
65566499
getTaskResult: async extra => {
6557-
const result = await extra.taskStore.getTaskResult(extra.taskId);
6558-
return result as CallToolResult;
6500+
return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
65596501
}
65606502
}
65616503
);
@@ -6590,14 +6532,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
65906532
expect(result.task).toHaveProperty('taskId');
65916533
expect(result.task.status).toBe('working');
65926534

6593-
// Wait for async operations to complete
6594-
await waitForLatch();
65956535
taskStore.cleanup();
65966536
});
65976537

65986538
test('should handle task failures during automatic polling', async () => {
65996539
const taskStore = new InMemoryTaskStore();
6600-
const { releaseLatch, waitForLatch } = createLatch();
66016540

66026541
const mcpServer = new McpServer(
66036542
{
@@ -6649,26 +6588,21 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
66496588
{
66506589
createTask: async extra => {
66516590
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6652-
6653-
// Capture taskStore for use in setTimeout
6654-
const store = extra.taskStore;
6655-
6656-
// Simulate async failure
6657-
setTimeout(async () => {
6658-
await store.storeTaskResult(task.taskId, 'failed', {
6659-
content: [{ type: 'text' as const, text: 'Error occurred' }],
6660-
isError: true
6661-
});
6662-
releaseLatch();
6663-
}, 150);
6664-
66656591
return { task };
66666592
},
66676593
getTask: async extra => {
66686594
const task = await extra.taskStore.getTask(extra.taskId);
66696595
if (!task) {
66706596
throw new Error('Task not found');
66716597
}
6598+
// Fail the task when polled - proves getTask is actually called
6599+
if (task.status === 'working') {
6600+
await extra.taskStore.storeTaskResult(extra.taskId, 'failed', {
6601+
content: [{ type: 'text' as const, text: 'Error occurred' }],
6602+
isError: true
6603+
});
6604+
return await extra.taskStore.getTask(extra.taskId);
6605+
}
66726606
return task;
66736607
},
66746608
getTaskResult: async extra => {
@@ -6696,14 +6630,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
66966630
expect(result.content).toEqual([{ type: 'text' as const, text: 'Error occurred' }]);
66976631
expect(result.isError).toBe(true);
66986632

6699-
// Wait for async operations to complete
6700-
await waitForLatch();
67016633
taskStore.cleanup();
67026634
});
67036635

67046636
test('should handle task cancellation during automatic polling', async () => {
67056637
const taskStore = new InMemoryTaskStore();
6706-
const { releaseLatch, waitForLatch } = createLatch();
67076638

67086639
const mcpServer = new McpServer(
67096640
{
@@ -6755,23 +6686,18 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
67556686
{
67566687
createTask: async extra => {
67576688
const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
6758-
6759-
// Capture taskStore for use in setTimeout
6760-
const store = extra.taskStore;
6761-
6762-
// Simulate async cancellation
6763-
setTimeout(async () => {
6764-
await store.updateTaskStatus(task.taskId, 'cancelled', 'Task was cancelled');
6765-
releaseLatch();
6766-
}, 150);
6767-
67686689
return { task };
67696690
},
67706691
getTask: async extra => {
67716692
const task = await extra.taskStore.getTask(extra.taskId);
67726693
if (!task) {
67736694
throw new Error('Task not found');
67746695
}
6696+
// Cancel the task when polled - proves getTask is actually called
6697+
if (task.status === 'working') {
6698+
await extra.taskStore.updateTaskStatus(extra.taskId, 'cancelled', 'Task was cancelled');
6699+
return await extra.taskStore.getTask(extra.taskId);
6700+
}
67756701
return task;
67766702
},
67776703
getTaskResult: async extra => {
@@ -6798,8 +6724,6 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
67986724
expect(result).toHaveProperty('content');
67996725
expect(result.content).toEqual([{ type: 'text' as const, text: expect.stringContaining('has no result stored') }]);
68006726

6801-
// Wait for async operations to complete
6802-
await waitForLatch();
68036727
taskStore.cleanup();
68046728
});
68056729

0 commit comments

Comments
 (0)