src/vs/workbench/contrib/mcp/common/mcpServerRequestHandler.ts
769 LOC · 493 covered · 276 uncovered · 113 ranges · 140 concepts · 26 introducers · 77 tests
File neighbourhood
The centred file is linked to every concept that introduces one of its ranges, every test that runs code from the file, and the gray connector concepts standing between those tests and the file's own introducer concepts. Undirected links join concepts to every file where they introduce source and concepts to the tests they introduce; arrows show specialization between the displayed concepts and bridge only concepts omitted from this view. Concept colors match the source ranges below; connector concepts have no source color and are shown in gray.
Focused file, its introducer and connector concepts, their introduced files, and tests that run code from the file
In the embedded map, ordinary wheel input scrolls the page; use the visible controls to zoom and drag to pan. Open the full-screen map for canvas navigation: wheel pans, Ctrl/Command plus wheel zooms, and arrow keys pan when this region is focused. On touch screens, open the full-screen map to pan or pinch. If JavaScript or WebGL is unavailable, use the related-file, concept, and source links on this page.
Graph controls are ready.
Interactive rendering requires JavaScript and WebGL. Use the related-file, concept, and source links on this page while the interactive map is unavailable.
1
>
/*---------------------------------------------------------------------------------------------
mcpServerRequestHandler.ts ×42
2
>
* Copyright (c) Microsoft Corporation. All rights reserved.
3
>
* Licensed under the MIT License. See License.txt in the project root for license information.
4
>
*--------------------------------------------------------------------------------------------*/
5
>
6
>
import { equals } from '../../../../base/common/arrays.js';
7
>
import { assertNever, softAssertNever } from '../../../../base/common/assert.js';
8
>
import { DeferredPromise, disposableTimeout, IntervalTimer, isThenable } from '../../../../base/common/async.js';
9
>
import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
10
>
import { CancellationError } from '../../../../base/common/errors.js';
11
>
import { Emitter } from '../../../../base/common/event.js';
12
>
import { Iterable } from '../../../../base/common/iterator.js';
13
>
import { JsonRpcError, JsonRpcProtocol } from '../../../../base/common/jsonRpcProtocol.js';
14
>
import { Disposable, DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js';
15
>
import { autorun, ISettableObservable, ObservablePromise, observableValue, transaction } from '../../../../base/common/observable.js';
16
>
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
17
>
import { canLog, ILogger, log, LogLevel } from '../../../../platform/log/common/log.js';
18
>
import { IProductService } from '../../../../platform/product/common/productService.js';
19
>
import { IMcpMessageTransport } from './mcpRegistryTypes.js';
20
>
import { IMcpTaskInternal, McpTaskManager } from './mcpTaskManager.js';
21
>
import { IMcpClientMethods, McpConnectionState, McpError, MpcResponseError } from './mcpTypes.js';
22
>
import { isTaskResult, translateMcpLogMessage } from './mcpTypesUtils.js';
23
>
import { MCP } from './modelContextProtocol.js';
24
>
25
>
export interface McpRoot {
26
>
uri: string;
27
>
name?: string;
28
>
}
29
>
30
>
export interface IMcpServerRequestHandlerOptions extends IMcpClientMethods {
31
>
/** MCP message transport */
32
>
launch: IMcpMessageTransport;
33
>
/** Logger instance. */
34
>
logger: ILogger;
35
>
/** Log level MCP messages is logged at */
36
>
requestLogLevel?: LogLevel;
37
>
/** Task manager for server-side MCP tasks (shared across reconnections) */
38
>
taskManager: McpTaskManager;
39
>
}
40
>
41
>
/**
42
>
* Request handler for communicating with an MCP server.
43
>
*
44
>
* Handles sending requests and receiving responses, with automatic
45
>
* handling of ping requests and typed client request methods.
46
>
*/
47
>
export class McpServerRequestHandler extends Disposable {
48
>
private readonly _rpc: JsonRpcProtocol;
49
>
50
>
private _hasAnnouncedRoots = false;
51
>
private _roots: MCP.Root[] = [];
52
>
53
>
public set roots(roots: MCP.Root[]) {
54
>
if (!equals(this._roots, roots)) {
55
>
this._roots = roots;
56
>
if (this._hasAnnouncedRoots) {
57
>
this.sendNotification({ method: 'notifications/roots/list_changed' });
mcpServerRequestHandler.ts ×1
58
>
this._hasAnnouncedRoots = false;
59
>
}
61
>
}
62
>
63
>
private _serverInit!: MCP.InitializeResult;
64
>
public get capabilities(): MCP.ServerCapabilities {
66
>
}
68
>
public get serverInfo(): MCP.Implementation {
70
>
}
72
>
public get serverInstructions(): string | undefined {
74
>
}
76
>
// Event emitters for server notifications
77
>
private readonly _onDidReceiveCancelledNotification = this._register(new Emitter<MCP.CancelledNotification>());
78
>
readonly onDidReceiveCancelledNotification = this._onDidReceiveCancelledNotification.event;
79
>
80
>
private readonly _onDidReceiveProgressNotification = this._register(new Emitter<MCP.ProgressNotification>());
81
>
readonly onDidReceiveProgressNotification = this._onDidReceiveProgressNotification.event;
82
>
83
>
private readonly _onDidReceiveElicitationCompleteNotification = this._register(new Emitter<MCP.ElicitationCompleteNotification>());
84
>
readonly onDidReceiveElicitationCompleteNotification = this._onDidReceiveElicitationCompleteNotification.event;
85
>
86
>
private readonly _onDidChangeResourceList = this._register(new Emitter<void>());
87
>
readonly onDidChangeResourceList = this._onDidChangeResourceList.event;
88
>
89
>
private readonly _onDidUpdateResource = this._register(new Emitter<MCP.ResourceUpdatedNotification>());
90
>
readonly onDidUpdateResource = this._onDidUpdateResource.event;
91
>
92
>
private readonly _onDidChangeToolList = this._register(new Emitter<void>());
93
>
readonly onDidChangeToolList = this._onDidChangeToolList.event;
94
>
95
>
private readonly _onDidChangePromptList = this._register(new Emitter<void>());
96
>
readonly onDidChangePromptList = this._onDidChangePromptList.event;
97
>
98
>
/**
99
>
* Connects to the MCP server and does the initialization handshake.
100
>
* @throws MpcResponseError if the server fails to initialize.
101
>
*/
102
>
public static async create(instaService: IInstantiationService, opts: IMcpServerRequestHandlerOptions, token?: CancellationToken) {
104
>
const store = new DisposableStore();
105
>
try {
106
>
const timer = store.add(new IntervalTimer());
107
>
timer.cancelAndSet(() => {
108
opts.logger.info('Waiting for server to respond to `initialize` request...');
110
>
111
>
await instaService.invokeFunction(async accessor => {
112
>
const productService = accessor.get(IProductService);
113
>
const initialized = await mcp.sendRequest<MCP.InitializeRequest, MCP.InitializeResult>({
114
>
method: 'initialize',
115
>
params: {
116
>
protocolVersion: MCP.LATEST_PROTOCOL_VERSION,
117
>
capabilities: {
118
>
roots: { listChanged: true },
119
>
sampling: opts.createMessageRequestHandler ? {} : undefined,
120
>
elicitation: opts.elicitationRequestHandler ? { form: {}, url: {} } : undefined,
121
>
tasks: {
122
>
list: {},
123
>
cancel: {},
124
>
requests: {
125
>
sampling: opts.createMessageRequestHandler ? { createMessage: {} } : undefined,
126
>
elicitation: opts.elicitationRequestHandler ? { create: {} } : undefined,
127
>
},
128
>
},
129
>
extensions: {
130
>
'io.modelcontextprotocol/ui': {
131
>
mimeTypes: ['text/html;profile=mcp-app']
132
>
}
133
>
}
134
>
},
135
>
clientInfo: {
136
>
name: productService.nameLong,
137
>
version: productService.version,
138
>
}
139
>
}
140
>
}, token);
142
>
mcp._sendLogLevelToServer(opts.logger.getLevel());
143
>
144
>
mcp.sendNotification<MCP.InitializedNotification>({
145
>
method: 'notifications/initialized'
146
>
});
149
>
return mcp;
152
>
throw e;
154
>
store.dispose();
155
>
}
156
>
}
158
>
public readonly logger: ILogger;
159
>
private readonly _launch: IMcpMessageTransport;
160
>
private readonly _requestLogLevel: LogLevel;
161
>
private readonly _createMessageRequestHandler: IMcpServerRequestHandlerOptions['createMessageRequestHandler'];
162
>
private readonly _elicitationRequestHandler: IMcpServerRequestHandlerOptions['elicitationRequestHandler'];
163
>
private readonly _taskManager: McpTaskManager;
164
>
165
>
protected constructor({
167
>
logger,
168
>
createMessageRequestHandler,
169
>
elicitationRequestHandler,
170
>
requestLogLevel = LogLevel.Debug,
171
>
taskManager,
172
>
}: IMcpServerRequestHandlerOptions) {
173
>
super();
174
>
this._launch = launch;
175
>
this.logger = logger;
176
>
this._requestLogLevel = requestLogLevel;
177
>
this._createMessageRequestHandler = createMessageRequestHandler;
178
>
this._elicitationRequestHandler = elicitationRequestHandler;
179
>
this._taskManager = taskManager;
180
>
181
>
this._rpc = this._register(new JsonRpcProtocol(
182
>
message => this.send(message as MCP.JSONRPCMessage),
183
>
{
184
>
handleRequest: (request, token) => this.handleServerRequest(request as MCP.JSONRPCRequest & MCP.ServerRequest, token),
185
>
handleNotification: notification => this.handleServerNotification(notification as MCP.JSONRPCNotification & MCP.ServerNotification),
186
>
}
187
>
));
188
>
189
>
// Attach this handler to the task manager
190
>
this._taskManager.setHandler(this);
191
>
this._register(this._taskManager.onDidUpdateTask(task => {
192
this.send({
193
jsonrpc: MCP.JSONRPC_VERSION,
194
method: 'notifications/tasks/status',
195
params: task
196
} satisfies MCP.TaskStatusNotification);
198
>
this._register(toDisposable(() => this._taskManager.setHandler(undefined)));
199
>
200
>
this._register(launch.onDidReceiveMessage(message => {
202
log(this.logger, this._requestLogLevel, `[server -> editor] ${JSON.stringify(message)}`);
203
}
206
>
this._register(autorun(reader => {
207
>
const state = launch.state.read(reader).state;
208
>
// the handler will get disposed when the launch stops, but if we're still
209
>
// create()'ing we need to make sure to cancel the initialize request.
210
>
if (state === McpConnectionState.Kind.Error || state === McpConnectionState.Kind.Stopped) {
212
>
}
214
>
215
>
// Listen for log level changes and forward them to the MCP server
216
>
this._register(logger.onDidChangeLogLevel((logLevel) => {
217
this._sendLogLevelToServer(logLevel);
219
>
}
221
>
/**
222
>
* Send a client request to the server and return the response.
223
>
*
224
>
* @param request The request to send
225
>
* @param token Cancellation token
226
>
* @param timeoutMs Optional timeout in milliseconds
227
>
* @returns A promise that resolves with the response
228
>
*/
229
>
private async sendRequest<T extends MCP.ClientRequest, R extends MCP.ServerResult>(
231
>
token: CancellationToken = CancellationToken.None
232
>
): Promise<R> {
233
>
if (this._store.isDisposed) {
234
return Promise.reject(new CancellationError());
235
}
237
>
return this._rpc.sendRequest<R>(
238
>
request,
239
>
token,
240
>
id => this.sendNotification({ method: 'notifications/cancelled', params: { requestId: id } })
241
>
).catch(error => {
243
>
throw new MpcResponseError(error.message, error.code, error.data);
mcpServerRequestHandler.ts ×1
244
>
}
247
>
}
249
>
private send(mcp: MCP.JSONRPCMessage) {
250
>
if (canLog(this.logger.getLevel(), this._requestLogLevel)) { // avoid building the string if we don't need to
mcpServerRequestHandler.ts ×17
251
>
log(this.logger, this._requestLogLevel, `[editor -> server] ${JSON.stringify(mcp)}`);
mcpServerConnection.ts ×1
252
>
}
254
>
this._launch.send(mcp);
257
>
/**
258
>
* Handles paginated requests by making multiple requests until all items are retrieved.
259
>
*
260
>
* @param method The method name to call
261
>
* @param getItems Function to extract the array of items from a result
262
>
* @param initialParams Initial parameters
263
>
* @param token Cancellation token
264
>
* @returns Promise with all items combined
265
>
*/
266
>
private async *sendRequestPaginated<T extends MCP.PaginatedRequest & MCP.ClientRequest, R extends MCP.PaginatedResult, I>(method: T['method'], getItems: (result: R) => I[], initialParams?: Omit<T['params'], 'jsonrpc' | 'id'>, token: CancellationToken = CancellationToken.None): AsyncIterable<I[]> {
268
>
269
>
do {
270
>
const params: T['params'] = {
271
>
...initialParams,
272
>
cursor: nextCursor
273
>
};
274
>
275
>
const result: R = await this.sendRequest<T, R>({ method, params }, token);
277
>
nextCursor = result.nextCursor;
279
>
}
281
>
private sendNotification<N extends MCP.ClientNotification>(notification: Omit<N, 'jsonrpc'>): void {
283
>
}
285
>
/**
286
>
* Handle incoming server requests
287
>
*/
288
>
private handleServerRequest(request: MCP.JSONRPCRequest & MCP.ServerRequest, token: CancellationToken): MCP.Result | Promise<MCP.Result> {
290
if (error instanceof McpError) {
291
return new JsonRpcError(error.code, error.message, error.data);
292
}
293
294
this.logger.error(`Error handling request ${request.method}:`, error);
295
const mcpError = McpError.unknown(error instanceof Error ? error : new Error(String(error)));
296
return new JsonRpcError(mcpError.code, mcpError.message, mcpError.data);
297
};
299
>
try {
300
>
let result: MCP.Result | Promise<MCP.Result>;
301
>
if (request.method === 'ping') {
305
>
} else if (request.method === 'sampling/createMessage' && this._createMessageRequestHandler) {
306
// Check if this is a task-augmented request
307
if (request.params.task) {
308
const taskResult = this._taskManager.createTask(
309
request.params.task.ttl ?? null,
310
(token) => this._createMessageRequestHandler!(request.params, token)
311
);
312
taskResult._meta ??= {};
313
taskResult._meta['io.modelcontextprotocol/related-task'] = { taskId: taskResult.task.taskId };
314
result = taskResult;
315
} else {
316
result = this._createMessageRequestHandler(request.params, token);
317
}
318
} else if (request.method === 'elicitation/create' && this._elicitationRequestHandler) {
319
// Check if this is a task-augmented request
320
if (request.params.task) {
321
const taskResult = this._taskManager.createTask(
322
request.params.task.ttl ?? null,
323
(token) => this._elicitationRequestHandler!(request.params, token)
324
);
325
taskResult._meta ??= {};
326
taskResult._meta['io.modelcontextprotocol/related-task'] = { taskId: taskResult.task.taskId };
327
result = taskResult;
328
} else {
329
result = this._elicitationRequestHandler(request.params, token);
330
}
331
} else if (request.method === 'tasks/get') {
332
result = this._taskManager.getTask(request.params.taskId);
333
} else if (request.method === 'tasks/result') {
334
result = this._taskManager.getTaskResult(request.params.taskId);
335
} else if (request.method === 'tasks/cancel') {
336
result = this._taskManager.cancelTask(request.params.taskId);
337
} else if (request.method === 'tasks/list') {
338
result = this._taskManager.listTasks();
339
} else {
340
throw McpError.methodNotFound(request.method);
341
}
343
>
if (isThenable(result)) {
344
return result.then(undefined, (error: unknown) => {
345
throw mapError(error);
346
});
347
}
349
>
return result;
350
>
} catch (e) {
351
throw mapError(e);
352
}
355
>
* Handle incoming server notifications
356
>
*/
357
>
private handleServerNotification(request: MCP.JSONRPCNotification & MCP.ServerNotification): void {
359
>
switch (request.method) {
360
>
case 'notifications/message':
361
return this.handleLoggingNotification(request);
364
>
return this.handleCancelledNotification(request);
367
>
return;
369
this._onDidChangeResourceList.fire();
370
return;
373
>
return;
375
this._onDidChangeToolList.fire();
376
return;
378
this._onDidChangePromptList.fire();
379
return;
381
this._onDidReceiveElicitationCompleteNotification.fire(request);
382
return;
384
this._taskManager.getClientTask(request.params.taskId)?.onDidUpdateState(request.params);
385
return;
387
softAssertNever(request);
389
>
} catch (error) {
390
this.logger.error(`Error handling notification ${request.method}:`, error);
391
}
394
>
private handleCancelledNotification(request: MCP.CancelledNotification): void {
396
>
this._rpc.cancelPendingRequest(request.params.requestId);
397
>
}
398
>
}
400
>
private handleLoggingNotification(request: MCP.LoggingMessageNotification): void {
401
translateMcpLogMessage(this.logger, request.params);
402
}
404
>
/**
405
>
* Send a response to a ping request
406
>
*/
407
>
private handlePing(_request: MCP.PingRequest): {} {
409
>
}
411
>
/**
412
>
* Send a response to a roots/list request
413
>
*/
414
>
private handleRootsList(_request: MCP.ListRootsRequest): MCP.ListRootsResult {
416
>
return { roots: this._roots };
417
>
}
419
>
private cancelAllRequests() {
421
>
}
423
>
public override dispose(): void {
425
>
super.dispose();
426
>
}
428
>
/**
429
>
* Forwards log level changes to the MCP server if it supports logging
430
>
*/
431
>
private async _sendLogLevelToServer(logLevel: LogLevel): Promise<void> {
433
>
// Only send if the server supports logging capabilities
434
>
if (!this.capabilities.logging) {
435
>
return;
436
>
}
437
438
await this.setLevel({ level: mapLogLevelToMcp(logLevel) });
439
} catch (error) {
440
this.logger.error(`Failed to set MCP server log level: ${error}`);
441
}
444
>
/**
445
>
* Send an initialize request
446
>
*/
447
>
initialize(params: MCP.InitializeRequest['params'], token?: CancellationToken): Promise<MCP.InitializeResult> {
448
return this.sendRequest<MCP.InitializeRequest, MCP.InitializeResult>({ method: 'initialize', params }, token);
449
}
451
>
/**
452
>
* List available resources
453
>
*/
454
>
listResources(params?: MCP.ListResourcesRequest['params'], token?: CancellationToken): Promise<MCP.Resource[]> {
455
>
return Iterable.asyncToArrayFlat(this.listResourcesIterable(params, token));
mcpServerRequestHandler.ts ×2
456
>
}
458
>
/**
459
>
* List available resources (iterable)
460
>
*/
461
>
listResourcesIterable(params?: MCP.ListResourcesRequest['params'], token?: CancellationToken): AsyncIterable<MCP.Resource[]> {
462
>
return this.sendRequestPaginated<MCP.ListResourcesRequest, MCP.ListResourcesResult, MCP.Resource>('resources/list', result => result.resources, params, token);
mcpServerRequestHandler.ts ×2
463
>
}
465
>
/**
466
>
* Read a specific resource
467
>
*/
468
>
readResource(params: MCP.ReadResourceRequest['params'], token?: CancellationToken): Promise<MCP.ReadResourceResult> {
469
>
return this.sendRequest<MCP.ReadResourceRequest, MCP.ReadResourceResult>({ method: 'resources/read', params }, token);
mcpServerRequestHandler.ts ×1
470
>
}
472
>
/**
473
>
* List available resource templates
474
>
*/
475
>
listResourceTemplates(params?: MCP.ListResourceTemplatesRequest['params'], token?: CancellationToken): Promise<MCP.ResourceTemplate[]> {
476
return Iterable.asyncToArrayFlat(this.sendRequestPaginated<MCP.ListResourceTemplatesRequest, MCP.ListResourceTemplatesResult, MCP.ResourceTemplate>('resources/templates/list', result => result.resourceTemplates, params, token));
477
}
479
>
/**
480
>
* Subscribe to resource updates
481
>
*/
482
>
subscribe(params: MCP.SubscribeRequest['params'], token?: CancellationToken): Promise<MCP.EmptyResult> {
483
>
return this.sendRequest<MCP.SubscribeRequest, MCP.EmptyResult>({ method: 'resources/subscribe', params }, token);
mcpResourceFilesystem.ts ×3
484
>
}
486
>
/**
487
>
* Unsubscribe from resource updates
488
>
*/
489
>
unsubscribe(params: MCP.UnsubscribeRequest['params'], token?: CancellationToken): Promise<MCP.EmptyResult> {
490
return this.sendRequest<MCP.UnsubscribeRequest, MCP.EmptyResult>({ method: 'resources/unsubscribe', params }, token);
491
}
493
>
/**
494
>
* List available prompts
495
>
*/
496
>
listPrompts(params?: MCP.ListPromptsRequest['params'], token?: CancellationToken): Promise<MCP.Prompt[]> {
497
return Iterable.asyncToArrayFlat(this.sendRequestPaginated<MCP.ListPromptsRequest, MCP.ListPromptsResult, MCP.Prompt>('prompts/list', result => result.prompts, params, token));
498
}
500
>
/**
501
>
* Get a specific prompt
502
>
*/
503
>
getPrompt(params: MCP.GetPromptRequest['params'], token?: CancellationToken): Promise<MCP.GetPromptResult> {
504
return this.sendRequest<MCP.GetPromptRequest, MCP.GetPromptResult>({ method: 'prompts/get', params }, token);
505
}
507
>
/**
508
>
* List available tools
509
>
*/
510
>
listTools(params?: MCP.ListToolsRequest['params'], token?: CancellationToken): Promise<MCP.Tool[]> {
511
>
return Iterable.asyncToArrayFlat(this.sendRequestPaginated<MCP.ListToolsRequest, MCP.ListToolsResult, MCP.Tool>('tools/list', result => result.tools, params, token));
mcpServerRequestHandler.ts ×1
512
>
}
514
>
/**
515
>
* Call a specific tool. Supports tasks automatically if `task` is set on the request.
516
>
*/
517
>
async callTool(params: MCP.CallToolRequest['params'] & MCP.Request['params'], token?: CancellationToken, onStatusMessage?: (message: string) => void): Promise<MCP.CallToolResult> {
518
>
const response = await this.sendRequest<MCP.CallToolRequest, MCP.CallToolResult | MCP.CreateTaskResult>({ method: 'tools/call', params }, token);
mcpServerRequestHandler.ts ×2
519
>
520
>
if (isTaskResult(response)) {
521
const task = new McpTask<MCP.CallToolResult>(response.task, token, onStatusMessage);
522
this._taskManager.adoptClientTask(task);
523
task.setHandler(this);
524
return task.result.finally(() => {
525
this._taskManager.abandonClientTask(task.id);
526
});
527
}
529
>
return response;
530
>
531
>
}
533
>
/**
534
>
* Set the logging level
535
>
*/
536
>
setLevel(params: MCP.SetLevelRequest['params'], token?: CancellationToken): Promise<MCP.EmptyResult> {
537
return this.sendRequest<MCP.SetLevelRequest, MCP.EmptyResult>({ method: 'logging/setLevel', params }, token);
538
}
540
>
/**
541
>
* Find completions for an argument
542
>
*/
543
>
complete(params: MCP.CompleteRequest['params'], token?: CancellationToken): Promise<MCP.CompleteResult> {
544
return this.sendRequest<MCP.CompleteRequest, MCP.CompleteResult>({ method: 'completion/complete', params }, token);
545
}
547
>
/**
548
>
* Get task status
549
>
*/
550
>
getTask(params: { taskId: string }, token?: CancellationToken): Promise<MCP.GetTaskResult> {
551
return this.sendRequest<MCP.GetTaskRequest, MCP.GetTaskResult>({ method: 'tasks/get', params }, token);
552
}
554
>
/**
555
>
* Get task result
556
>
*/
557
>
getTaskResult(params: { taskId: string }, token?: CancellationToken): Promise<MCP.GetTaskPayloadResult> {
558
return this.sendRequest<MCP.GetTaskPayloadRequest, MCP.GetTaskPayloadResult>({ method: 'tasks/result', params }, token);
559
}
561
>
/**
562
>
* Cancel a task
563
>
*/
564
>
cancelTask(params: { taskId: string }, token?: CancellationToken): Promise<MCP.CancelTaskResult> {
565
return this.sendRequest<MCP.CancelTaskRequest, MCP.CancelTaskResult>({ method: 'tasks/cancel', params }, token);
566
}
568
>
/**
569
>
* List all tasks
570
>
*/
571
>
listTasks(params?: MCP.ListTasksRequest['params'], token?: CancellationToken): Promise<MCP.Task[]> {
572
return Iterable.asyncToArrayFlat(
573
this.sendRequestPaginated<MCP.ListTasksRequest, MCP.ListTasksResult, MCP.Task>(
574
'tasks/list', result => result.tasks, params, token
575
)
576
);
577
}
579
>
580
function isTaskInTerminalState(task: MCP.Task): boolean {
581
return task.status === 'completed' || task.status === 'failed' || task.status === 'cancelled';
582
}
584
>
/**
585
>
* Implementation of a task that handles polling, status notifications, and handler reconnections. It implements the task polling loop internally and can also be
586
>
* updated externally via `onDidUpdateState`, when notifications are received
587
>
* for example.
588
>
* @internal
589
>
*/
590
>
export class McpTask<T extends MCP.Result> extends Disposable implements IMcpTaskInternal {
591
>
private readonly promise = new DeferredPromise<T>();
592
>
593
>
public get result(): Promise<T> {
594
>
return this.promise.p;
595
>
}
596
>
597
>
public get id() {
598
return this._task.taskId;
599
}
601
>
private _lastTaskState: ISettableObservable<MCP.Task>;
602
>
private _handler = observableValue<McpServerRequestHandler | undefined>('mcpTaskHandler', undefined);
603
>
604
>
constructor(
605
private readonly _task: MCP.Task,
606
_token: CancellationToken = CancellationToken.None,
607
private readonly _onStatusMessage?: (message: string) => void,
608
) {
609
super();
610
611
const expiresAt = _task.ttl ? (Date.now() + _task.ttl) : undefined;
612
this._lastTaskState = observableValue('lastTaskState', this._task);
613
614
const store = this._register(new DisposableStore());
615
616
// Handle external cancellation token
617
if (_token.isCancellationRequested) {
618
this._lastTaskState.set({ ...this._task, status: 'cancelled' }, undefined);
619
} else {
620
store.add(_token.onCancellationRequested(() => {
621
const current = this._lastTaskState.get();
622
if (!isTaskInTerminalState(current)) {
623
this._lastTaskState.set({ ...current, status: 'cancelled' }, undefined);
624
}
625
}));
626
}
627
628
// Handle TTL expiration with an explicit timeout
629
if (expiresAt) {
630
const ttlTimeout = expiresAt - Date.now();
631
if (ttlTimeout <= 0) {
632
this._lastTaskState.set({ ...this._task, status: 'cancelled', statusMessage: 'Task timed out.' }, undefined);
633
} else {
634
store.add(disposableTimeout(() => {
635
const current = this._lastTaskState.get();
636
if (!isTaskInTerminalState(current)) {
637
this._lastTaskState.set({ ...current, status: 'cancelled', statusMessage: 'Task timed out.' }, undefined);
638
}
639
}, ttlTimeout));
640
}
641
}
642
643
// A `tasks/result` call triggered by an input_required state.
644
const inputRequiredLookup = observableValue<ObservablePromise<MCP.Task> | undefined>('activeResultLookup', undefined);
645
646
// 1. Poll for task updates when the task isn't in a terminal state
647
store.add(autorun(reader => {
648
const current = this._lastTaskState.read(reader);
649
if (isTaskInTerminalState(current)) {
650
return;
651
}
652
653
// When a task goes into the input_required state, by spec we should call
654
// `tasks/result` which can return an SSE stream of task updates. No need
655
// to poll while such a lookup is going on, but once it resolves we should
656
// clear and update our state.
657
const lookup = inputRequiredLookup.read(reader);
658
if (lookup) {
659
const result = lookup.promiseResult.read(reader);
660
return transaction(tx => {
661
if (!result) {
662
// still ongoing
663
} else if (result.data) {
664
inputRequiredLookup.set(undefined, tx);
665
this._lastTaskState.set(result.data, tx);
666
} else {
667
inputRequiredLookup.set(undefined, tx);
668
if (result.error instanceof McpError && result.error.code === MCP.INVALID_PARAMS) {
669
this._lastTaskState.set({ ...current, status: 'cancelled' }, undefined);
670
} else {
671
// Maybe a connection error -- start polling again
672
this._lastTaskState.set({ ...current, status: 'working' }, undefined);
673
}
674
}
675
});
676
}
677
678
const handler = this._handler.read(reader);
679
if (!handler) {
680
return;
681
}
682
683
const pollInterval = _task.pollInterval ?? 2000;
684
const cts = new CancellationTokenSource(_token);
685
reader.store.add(toDisposable(() => cts.dispose(true)));
686
reader.store.add(disposableTimeout(() => {
687
handler.getTask({ taskId: current.taskId }, cts.token)
688
.catch((e): MCP.Task | undefined => {
689
if (e instanceof McpError && e.code === MCP.INVALID_PARAMS) {
690
return { ...current, status: 'cancelled' };
691
} else {
692
return { ...current }; // errors are already logged, keep in current state
693
}
694
})
695
.then(r => {
696
if (r && !cts.token.isCancellationRequested) {
697
this._lastTaskState.set(r, undefined);
698
}
699
});
700
}, pollInterval));
701
}));
702
703
// 2. Get the result once it's available (or propagate errors). Trigger
704
// input_required handling as needed. Only react when the status itself changes.
705
const lastStatus = this._lastTaskState.map(task => task.status);
706
store.add(autorun(reader => {
707
const status = lastStatus.read(reader);
708
if (status === 'failed') {
709
const current = this._lastTaskState.read(undefined);
710
this.promise.error(new Error(`Task ${current.taskId} failed: ${current.statusMessage ?? 'unknown error'}`));
711
store.dispose();
712
} else if (status === 'cancelled') {
713
this.promise.cancel();
714
store.dispose();
715
} else if (status === 'input_required') {
716
const handler = this._handler.read(reader);
717
if (handler) {
718
const current = this._lastTaskState.read(undefined);
719
const cts = new CancellationTokenSource(_token);
720
reader.store.add(toDisposable(() => cts.dispose(true)));
721
inputRequiredLookup.set(new ObservablePromise<MCP.Task>(handler.getTask({ taskId: current.taskId }, cts.token)), undefined);
722
}
723
} else if (status === 'completed') {
724
const handler = this._handler.read(reader);
725
if (handler) {
726
this.promise.settleWith(handler.getTaskResult({ taskId: _task.taskId }, _token) as Promise<T>);
727
store.dispose();
728
}
729
} else if (status === 'working') {
730
// no-op
731
} else {
732
softAssertNever(status);
733
}
734
}));
735
}
737
>
onDidUpdateState(task: MCP.Task) {
738
this._lastTaskState.set(task, undefined);
739
if (task.statusMessage && this._onStatusMessage) {
740
this._onStatusMessage(task.statusMessage);
741
}
742
}
744
>
setHandler(handler: McpServerRequestHandler | undefined): void {
745
this._handler.set(handler, undefined);
746
}
748
>
749
>
/**
750
>
* Maps VSCode LogLevel to MCP LoggingLevel
751
>
*/
752
function mapLogLevelToMcp(logLevel: LogLevel): MCP.LoggingLevel {
753
switch (logLevel) {
754
case LogLevel.Trace:
755
return 'debug'; // MCP doesn't have trace, use debug
756
case LogLevel.Debug:
757
return 'debug';
758
case LogLevel.Info:
759
return 'info';
760
case LogLevel.Warning:
761
return 'warning';
762
case LogLevel.Error:
763
return 'error';
764
case LogLevel.Off:
765
return 'emergency'; // MCP doesn't have off, use emergency
766
default:
767
return assertNever(logLevel); // Off and other levels are not supported
768
}
769
}