Skip to content

Commit 0b8ced2

Browse files
LucaButBoringLucaButBoring
committed
Correctly check peer task capabilities on receiving end
1 parent 01be32d commit 0b8ced2

File tree

8 files changed

+227
-2
lines changed

8 files changed

+227
-2
lines changed

‎src/client/index.test.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1373,6 +1373,10 @@ describe('Task-based execution', () => {
13731373
tools: {},
13741374
tasks: {
13751375
requests: {
1376+
tools: {
1377+
call: true,
1378+
list: true
1379+
},
13761380
tasks: {
13771381
get: true,
13781382
list: true,
@@ -1461,6 +1465,10 @@ describe('Task-based execution', () => {
14611465
tools: {},
14621466
tasks: {
14631467
requests: {
1468+
tools: {
1469+
call: true,
1470+
list: true
1471+
},
14641472
tasks: {
14651473
get: true,
14661474
list: true,
@@ -1547,6 +1555,10 @@ describe('Task-based execution', () => {
15471555
tools: {},
15481556
tasks: {
15491557
requests: {
1558+
tools: {
1559+
call: true,
1560+
list: true
1561+
},
15501562
tasks: {
15511563
get: true,
15521564
list: true,
@@ -1631,6 +1643,10 @@ describe('Task-based execution', () => {
16311643
tools: {},
16321644
tasks: {
16331645
requests: {
1646+
tools: {
1647+
call: true,
1648+
list: true
1649+
},
16341650
tasks: {
16351651
get: true,
16361652
list: true,
@@ -1738,6 +1754,9 @@ describe('Task-based execution', () => {
17381754
elicitation: {},
17391755
tasks: {
17401756
requests: {
1757+
elicitation: {
1758+
create: true
1759+
},
17411760
tasks: {
17421761
get: true,
17431762
list: true,
@@ -1827,6 +1846,9 @@ describe('Task-based execution', () => {
18271846
elicitation: {},
18281847
tasks: {
18291848
requests: {
1849+
elicitation: {
1850+
create: true
1851+
},
18301852
tasks: {
18311853
get: true,
18321854
list: true,
@@ -1912,6 +1934,9 @@ describe('Task-based execution', () => {
19121934
elicitation: {},
19131935
tasks: {
19141936
requests: {
1937+
elicitation: {
1938+
create: true
1939+
},
19151940
tasks: {
19161941
get: true,
19171942
list: true,
@@ -1996,6 +2021,9 @@ describe('Task-based execution', () => {
19962021
elicitation: {},
19972022
tasks: {
19982023
requests: {
2024+
elicitation: {
2025+
create: true
2026+
},
19992027
tasks: {
20002028
get: true,
20012029
list: true,
@@ -2092,6 +2120,10 @@ describe('Task-based execution', () => {
20922120
tools: {},
20932121
tasks: {
20942122
requests: {
2123+
tools: {
2124+
call: true,
2125+
list: true
2126+
},
20952127
tasks: {
20962128
get: true,
20972129
list: true,
@@ -2202,6 +2234,10 @@ describe('Task-based execution', () => {
22022234
tools: {},
22032235
tasks: {
22042236
requests: {
2237+
tools: {
2238+
call: true,
2239+
list: true
2240+
},
22052241
tasks: {
22062242
get: true,
22072243
list: true,
@@ -2256,6 +2292,10 @@ describe('Task-based execution', () => {
22562292
tools: {},
22572293
tasks: {
22582294
requests: {
2295+
tools: {
2296+
call: true,
2297+
list: true
2298+
},
22592299
tasks: {
22602300
get: true,
22612301
list: true,
@@ -2310,6 +2350,9 @@ describe('Task-based execution', () => {
23102350
elicitation: {},
23112351
tasks: {
23122352
requests: {
2353+
elicitation: {
2354+
create: true
2355+
},
23132356
tasks: {
23142357
get: true,
23152358
list: true,

‎src/client/index.ts‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,68 @@ export class Client<
408408
}
409409
}
410410

411+
protected assertTaskHandlerCapability(method: string): void {
412+
// Task handlers are registered in Protocol constructor before _capabilities is initialized
413+
// Skip capability check for task methods during initialization
414+
if (!this._capabilities) {
415+
return;
416+
}
417+
418+
if (!this._capabilities.tasks?.requests) {
419+
throw new Error(`Client does not support task creation (required for ${method})`);
420+
}
421+
422+
const requests = this._capabilities.tasks.requests;
423+
424+
switch (method) {
425+
case 'sampling/createMessage':
426+
if (!requests.sampling?.createMessage) {
427+
throw new Error(`Client does not support task creation for sampling/createMessage (required for ${method})`);
428+
}
429+
break;
430+
431+
case 'elicitation/create':
432+
if (!requests.elicitation?.create) {
433+
throw new Error(`Client does not support task creation for elicitation/create (required for ${method})`);
434+
}
435+
break;
436+
437+
case 'roots/list':
438+
if (!requests.roots?.list) {
439+
throw new Error(`Client does not support task creation for roots/list (required for ${method})`);
440+
}
441+
break;
442+
443+
case 'tasks/get':
444+
if (!requests.tasks?.get) {
445+
throw new Error(`Client does not support task creation for tasks/get (required for ${method})`);
446+
}
447+
break;
448+
449+
case 'tasks/list':
450+
if (!requests.tasks?.list) {
451+
throw new Error(`Client does not support task creation for tasks/list (required for ${method})`);
452+
}
453+
break;
454+
455+
case 'tasks/result':
456+
if (!requests.tasks?.result) {
457+
throw new Error(`Client does not support task creation for tasks/result (required for ${method})`);
458+
}
459+
break;
460+
461+
case 'tasks/delete':
462+
if (!requests.tasks?.delete) {
463+
throw new Error(`Client does not support task creation for tasks/delete (required for ${method})`);
464+
}
465+
break;
466+
467+
default:
468+
// Method doesn't support tasks, which is fine - no error
469+
break;
470+
}
471+
}
472+
411473
async ping(options?: RequestOptions) {
412474
return this.request({ method: 'ping' }, EmptyResultSchema, options);
413475
}

‎src/examples/server/simpleStreamableHttp.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ const getServer = () => {
3939
websiteUrl: 'https://github.com/modelcontextprotocol/typescript-sdk'
4040
},
4141
{
42-
capabilities: { logging: {} },
42+
capabilities: { logging: {}, tasks: { requests: { tools: { call: true } } } },
4343
taskStore // Enable task support
4444
}
4545
);

‎src/server/index.test.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -970,6 +970,10 @@ describe('Task-based execution', () => {
970970
tools: {},
971971
tasks: {
972972
requests: {
973+
tools: {
974+
call: true,
975+
list: true
976+
},
973977
tasks: {
974978
get: true,
975979
list: true,
@@ -1140,6 +1144,10 @@ describe('Task-based execution', () => {
11401144
tools: {},
11411145
tasks: {
11421146
requests: {
1147+
tools: {
1148+
call: true,
1149+
list: true
1150+
},
11431151
tasks: {
11441152
get: true,
11451153
list: true,
@@ -1671,6 +1679,10 @@ describe('Task-based execution', () => {
16711679
tools: {},
16721680
tasks: {
16731681
requests: {
1682+
tools: {
1683+
call: true,
1684+
list: true
1685+
},
16741686
tasks: {
16751687
get: true,
16761688
list: true,

‎src/server/index.ts‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,6 +349,86 @@ export class Server<
349349
}
350350
}
351351

352+
protected assertTaskHandlerCapability(method: string): void {
353+
// Task handlers are registered in Protocol constructor before _capabilities is initialized
354+
// Skip capability check for task methods during initialization
355+
if (!this._capabilities) {
356+
return;
357+
}
358+
359+
if (!this._capabilities.tasks?.requests) {
360+
throw new Error(`Server does not support task creation (required for ${method})`);
361+
}
362+
363+
const requests = this._capabilities.tasks.requests;
364+
365+
switch (method) {
366+
case 'tools/call':
367+
if (!requests.tools?.call) {
368+
throw new Error(`Server does not support task creation for tools/call (required for ${method})`);
369+
}
370+
break;
371+
372+
case 'tools/list':
373+
if (!requests.tools?.list) {
374+
throw new Error(`Server does not support task creation for tools/list (required for ${method})`);
375+
}
376+
break;
377+
378+
case 'resources/read':
379+
if (!requests.resources?.read) {
380+
throw new Error(`Server does not support task creation for resources/read (required for ${method})`);
381+
}
382+
break;
383+
384+
case 'resources/list':
385+
if (!requests.resources?.list) {
386+
throw new Error(`Server does not support task creation for resources/list (required for ${method})`);
387+
}
388+
break;
389+
390+
case 'prompts/get':
391+
if (!requests.prompts?.get) {
392+
throw new Error(`Server does not support task creation for prompts/get (required for ${method})`);
393+
}
394+
break;
395+
396+
case 'prompts/list':
397+
if (!requests.prompts?.list) {
398+
throw new Error(`Server does not support task creation for prompts/list (required for ${method})`);
399+
}
400+
break;
401+
402+
case 'tasks/get':
403+
if (!requests.tasks?.get) {
404+
throw new Error(`Server does not support task creation for tasks/get (required for ${method})`);
405+
}
406+
break;
407+
408+
case 'tasks/list':
409+
if (!requests.tasks?.list) {
410+
throw new Error(`Server does not support task creation for tasks/list (required for ${method})`);
411+
}
412+
break;
413+
414+
case 'tasks/result':
415+
if (!requests.tasks?.result) {
416+
throw new Error(`Server does not support task creation for tasks/result (required for ${method})`);
417+
}
418+
break;
419+
420+
case 'tasks/delete':
421+
if (!requests.tasks?.delete) {
422+
throw new Error(`Server does not support task creation for tasks/delete (required for ${method})`);
423+
}
424+
break;
425+
426+
default:
427+
// Method doesn't support tasks, which is fine - no error
428+
break;
429+
}
430+
}
431+
352432
private async _oninitialize(request: InitializeRequest): Promise<InitializeResult> {
353433
const requestedVersion = request.params.protocolVersion;
354434

‎src/shared/protocol-transport-handling.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ describe('Protocol transport handling bug', () => {
3838
protected assertNotificationCapability(): void {}
3939
protected assertRequestHandlerCapability(): void {}
4040
protected assertTaskCapability(): void {}
41+
protected assertTaskHandlerCapability(): void {}
4142
})();
4243

4344
transportA = new MockTransport('A');

0 commit comments

Comments
 (0)