protocol.js 51 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107
  1. import { safeParse } from '../server/zod-compat.js';
  2. import { CancelledNotificationSchema, CreateTaskResultSchema, ErrorCode, GetTaskRequestSchema, GetTaskResultSchema, GetTaskPayloadRequestSchema, ListTasksRequestSchema, ListTasksResultSchema, CancelTaskRequestSchema, CancelTaskResultSchema, isJSONRPCErrorResponse, isJSONRPCRequest, isJSONRPCResultResponse, isJSONRPCNotification, McpError, PingRequestSchema, ProgressNotificationSchema, RELATED_TASK_META_KEY, TaskStatusNotificationSchema, isTaskAugmentedRequestParams } from '../types.js';
  3. import { isTerminal } from '../experimental/tasks/interfaces.js';
  4. import { getMethodLiteral, parseWithCompat } from '../server/zod-json-schema-compat.js';
  5. /**
  6. * The default request timeout, in miliseconds.
  7. */
  8. export const DEFAULT_REQUEST_TIMEOUT_MSEC = 60000;
  9. /**
  10. * Implements MCP protocol framing on top of a pluggable transport, including
  11. * features like request/response linking, notifications, and progress.
  12. */
  13. export class Protocol {
  14. constructor(_options) {
  15. this._options = _options;
  16. this._requestMessageId = 0;
  17. this._requestHandlers = new Map();
  18. this._requestHandlerAbortControllers = new Map();
  19. this._notificationHandlers = new Map();
  20. this._responseHandlers = new Map();
  21. this._progressHandlers = new Map();
  22. this._timeoutInfo = new Map();
  23. this._pendingDebouncedNotifications = new Set();
  24. // Maps task IDs to progress tokens to keep handlers alive after CreateTaskResult
  25. this._taskProgressTokens = new Map();
  26. this._requestResolvers = new Map();
  27. this.setNotificationHandler(CancelledNotificationSchema, notification => {
  28. this._oncancel(notification);
  29. });
  30. this.setNotificationHandler(ProgressNotificationSchema, notification => {
  31. this._onprogress(notification);
  32. });
  33. this.setRequestHandler(PingRequestSchema,
  34. // Automatic pong by default.
  35. _request => ({}));
  36. // Install task handlers if TaskStore is provided
  37. this._taskStore = _options?.taskStore;
  38. this._taskMessageQueue = _options?.taskMessageQueue;
  39. if (this._taskStore) {
  40. this.setRequestHandler(GetTaskRequestSchema, async (request, extra) => {
  41. const task = await this._taskStore.getTask(request.params.taskId, extra.sessionId);
  42. if (!task) {
  43. throw new McpError(ErrorCode.InvalidParams, 'Failed to retrieve task: Task not found');
  44. }
  45. // Per spec: tasks/get responses SHALL NOT include related-task metadata
  46. // as the taskId parameter is the source of truth
  47. // @ts-expect-error SendResultT cannot contain GetTaskResult, but we include it in our derived types everywhere else
  48. return {
  49. ...task
  50. };
  51. });
  52. this.setRequestHandler(GetTaskPayloadRequestSchema, async (request, extra) => {
  53. const handleTaskResult = async () => {
  54. const taskId = request.params.taskId;
  55. // Deliver queued messages
  56. if (this._taskMessageQueue) {
  57. let queuedMessage;
  58. while ((queuedMessage = await this._taskMessageQueue.dequeue(taskId, extra.sessionId))) {
  59. // Handle response and error messages by routing them to the appropriate resolver
  60. if (queuedMessage.type === 'response' || queuedMessage.type === 'error') {
  61. const message = queuedMessage.message;
  62. const requestId = message.id;
  63. // Lookup resolver in _requestResolvers map
  64. const resolver = this._requestResolvers.get(requestId);
  65. if (resolver) {
  66. // Remove resolver from map after invocation
  67. this._requestResolvers.delete(requestId);
  68. // Invoke resolver with response or error
  69. if (queuedMessage.type === 'response') {
  70. resolver(message);
  71. }
  72. else {
  73. // Convert JSONRPCError to McpError
  74. const errorMessage = message;
  75. const error = new McpError(errorMessage.error.code, errorMessage.error.message, errorMessage.error.data);
  76. resolver(error);
  77. }
  78. }
  79. else {
  80. // Handle missing resolver gracefully with error logging
  81. const messageType = queuedMessage.type === 'response' ? 'Response' : 'Error';
  82. this._onerror(new Error(`${messageType} handler missing for request ${requestId}`));
  83. }
  84. // Continue to next message
  85. continue;
  86. }
  87. // Send the message on the response stream by passing the relatedRequestId
  88. // This tells the transport to write the message to the tasks/result response stream
  89. await this._transport?.send(queuedMessage.message, { relatedRequestId: extra.requestId });
  90. }
  91. }
  92. // Now check task status
  93. const task = await this._taskStore.getTask(taskId, extra.sessionId);
  94. if (!task) {
  95. throw new McpError(ErrorCode.InvalidParams, `Task not found: ${taskId}`);
  96. }
  97. // Block if task is not terminal (we've already delivered all queued messages above)
  98. if (!isTerminal(task.status)) {
  99. // Wait for status change or new messages
  100. await this._waitForTaskUpdate(taskId, extra.signal);
  101. // After waking up, recursively call to deliver any new messages or result
  102. return await handleTaskResult();
  103. }
  104. // If task is terminal, return the result
  105. if (isTerminal(task.status)) {
  106. const result = await this._taskStore.getTaskResult(taskId, extra.sessionId);
  107. this._clearTaskQueue(taskId);
  108. return {
  109. ...result,
  110. _meta: {
  111. ...result._meta,
  112. [RELATED_TASK_META_KEY]: {
  113. taskId: taskId
  114. }
  115. }
  116. };
  117. }
  118. return await handleTaskResult();
  119. };
  120. return await handleTaskResult();
  121. });
  122. this.setRequestHandler(ListTasksRequestSchema, async (request, extra) => {
  123. try {
  124. const { tasks, nextCursor } = await this._taskStore.listTasks(request.params?.cursor, extra.sessionId);
  125. // @ts-expect-error SendResultT cannot contain ListTasksResult, but we include it in our derived types everywhere else
  126. return {
  127. tasks,
  128. nextCursor,
  129. _meta: {}
  130. };
  131. }
  132. catch (error) {
  133. throw new McpError(ErrorCode.InvalidParams, `Failed to list tasks: ${error instanceof Error ? error.message : String(error)}`);
  134. }
  135. });
  136. this.setRequestHandler(CancelTaskRequestSchema, async (request, extra) => {
  137. try {
  138. // Get the current task to check if it's in a terminal state, in case the implementation is not atomic
  139. const task = await this._taskStore.getTask(request.params.taskId, extra.sessionId);
  140. if (!task) {
  141. throw new McpError(ErrorCode.InvalidParams, `Task not found: ${request.params.taskId}`);
  142. }
  143. // Reject cancellation of terminal tasks
  144. if (isTerminal(task.status)) {
  145. throw new McpError(ErrorCode.InvalidParams, `Cannot cancel task in terminal status: ${task.status}`);
  146. }
  147. await this._taskStore.updateTaskStatus(request.params.taskId, 'cancelled', 'Client cancelled task execution.', extra.sessionId);
  148. this._clearTaskQueue(request.params.taskId);
  149. const cancelledTask = await this._taskStore.getTask(request.params.taskId, extra.sessionId);
  150. if (!cancelledTask) {
  151. // Task was deleted during cancellation (e.g., cleanup happened)
  152. throw new McpError(ErrorCode.InvalidParams, `Task not found after cancellation: ${request.params.taskId}`);
  153. }
  154. return {
  155. _meta: {},
  156. ...cancelledTask
  157. };
  158. }
  159. catch (error) {
  160. // Re-throw McpError as-is
  161. if (error instanceof McpError) {
  162. throw error;
  163. }
  164. throw new McpError(ErrorCode.InvalidRequest, `Failed to cancel task: ${error instanceof Error ? error.message : String(error)}`);
  165. }
  166. });
  167. }
  168. }
  169. async _oncancel(notification) {
  170. if (!notification.params.requestId) {
  171. return;
  172. }
  173. // Handle request cancellation
  174. const controller = this._requestHandlerAbortControllers.get(notification.params.requestId);
  175. controller?.abort(notification.params.reason);
  176. }
  177. _setupTimeout(messageId, timeout, maxTotalTimeout, onTimeout, resetTimeoutOnProgress = false) {
  178. this._timeoutInfo.set(messageId, {
  179. timeoutId: setTimeout(onTimeout, timeout),
  180. startTime: Date.now(),
  181. timeout,
  182. maxTotalTimeout,
  183. resetTimeoutOnProgress,
  184. onTimeout
  185. });
  186. }
  187. _resetTimeout(messageId) {
  188. const info = this._timeoutInfo.get(messageId);
  189. if (!info)
  190. return false;
  191. const totalElapsed = Date.now() - info.startTime;
  192. if (info.maxTotalTimeout && totalElapsed >= info.maxTotalTimeout) {
  193. this._timeoutInfo.delete(messageId);
  194. throw McpError.fromError(ErrorCode.RequestTimeout, 'Maximum total timeout exceeded', {
  195. maxTotalTimeout: info.maxTotalTimeout,
  196. totalElapsed
  197. });
  198. }
  199. clearTimeout(info.timeoutId);
  200. info.timeoutId = setTimeout(info.onTimeout, info.timeout);
  201. return true;
  202. }
  203. _cleanupTimeout(messageId) {
  204. const info = this._timeoutInfo.get(messageId);
  205. if (info) {
  206. clearTimeout(info.timeoutId);
  207. this._timeoutInfo.delete(messageId);
  208. }
  209. }
  210. /**
  211. * Attaches to the given transport, starts it, and starts listening for messages.
  212. *
  213. * The Protocol object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward.
  214. */
  215. async connect(transport) {
  216. if (this._transport) {
  217. throw new Error('Already connected to a transport. Call close() before connecting to a new transport, or use a separate Protocol instance per connection.');
  218. }
  219. this._transport = transport;
  220. const _onclose = this.transport?.onclose;
  221. this._transport.onclose = () => {
  222. _onclose?.();
  223. this._onclose();
  224. };
  225. const _onerror = this.transport?.onerror;
  226. this._transport.onerror = (error) => {
  227. _onerror?.(error);
  228. this._onerror(error);
  229. };
  230. const _onmessage = this._transport?.onmessage;
  231. this._transport.onmessage = (message, extra) => {
  232. _onmessage?.(message, extra);
  233. if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) {
  234. this._onresponse(message);
  235. }
  236. else if (isJSONRPCRequest(message)) {
  237. this._onrequest(message, extra);
  238. }
  239. else if (isJSONRPCNotification(message)) {
  240. this._onnotification(message);
  241. }
  242. else {
  243. this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`));
  244. }
  245. };
  246. await this._transport.start();
  247. }
  248. _onclose() {
  249. const responseHandlers = this._responseHandlers;
  250. this._responseHandlers = new Map();
  251. this._progressHandlers.clear();
  252. this._taskProgressTokens.clear();
  253. this._pendingDebouncedNotifications.clear();
  254. for (const info of this._timeoutInfo.values()) {
  255. clearTimeout(info.timeoutId);
  256. }
  257. this._timeoutInfo.clear();
  258. // Abort all in-flight request handlers so they stop sending messages
  259. for (const controller of this._requestHandlerAbortControllers.values()) {
  260. controller.abort();
  261. }
  262. this._requestHandlerAbortControllers.clear();
  263. const error = McpError.fromError(ErrorCode.ConnectionClosed, 'Connection closed');
  264. this._transport = undefined;
  265. this.onclose?.();
  266. for (const handler of responseHandlers.values()) {
  267. handler(error);
  268. }
  269. }
  270. _onerror(error) {
  271. this.onerror?.(error);
  272. }
  273. _onnotification(notification) {
  274. const handler = this._notificationHandlers.get(notification.method) ?? this.fallbackNotificationHandler;
  275. // Ignore notifications not being subscribed to.
  276. if (handler === undefined) {
  277. return;
  278. }
  279. // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
  280. Promise.resolve()
  281. .then(() => handler(notification))
  282. .catch(error => this._onerror(new Error(`Uncaught error in notification handler: ${error}`)));
  283. }
  284. _onrequest(request, extra) {
  285. const handler = this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler;
  286. // Capture the current transport at request time to ensure responses go to the correct client
  287. const capturedTransport = this._transport;
  288. // Extract taskId from request metadata if present (needed early for method not found case)
  289. const relatedTaskId = request.params?._meta?.[RELATED_TASK_META_KEY]?.taskId;
  290. if (handler === undefined) {
  291. const errorResponse = {
  292. jsonrpc: '2.0',
  293. id: request.id,
  294. error: {
  295. code: ErrorCode.MethodNotFound,
  296. message: 'Method not found'
  297. }
  298. };
  299. // Queue or send the error response based on whether this is a task-related request
  300. if (relatedTaskId && this._taskMessageQueue) {
  301. this._enqueueTaskMessage(relatedTaskId, {
  302. type: 'error',
  303. message: errorResponse,
  304. timestamp: Date.now()
  305. }, capturedTransport?.sessionId).catch(error => this._onerror(new Error(`Failed to enqueue error response: ${error}`)));
  306. }
  307. else {
  308. capturedTransport
  309. ?.send(errorResponse)
  310. .catch(error => this._onerror(new Error(`Failed to send an error response: ${error}`)));
  311. }
  312. return;
  313. }
  314. const abortController = new AbortController();
  315. this._requestHandlerAbortControllers.set(request.id, abortController);
  316. const taskCreationParams = isTaskAugmentedRequestParams(request.params) ? request.params.task : undefined;
  317. const taskStore = this._taskStore ? this.requestTaskStore(request, capturedTransport?.sessionId) : undefined;
  318. const fullExtra = {
  319. signal: abortController.signal,
  320. sessionId: capturedTransport?.sessionId,
  321. _meta: request.params?._meta,
  322. sendNotification: async (notification) => {
  323. if (abortController.signal.aborted)
  324. return;
  325. // Include related-task metadata if this request is part of a task
  326. const notificationOptions = { relatedRequestId: request.id };
  327. if (relatedTaskId) {
  328. notificationOptions.relatedTask = { taskId: relatedTaskId };
  329. }
  330. await this.notification(notification, notificationOptions);
  331. },
  332. sendRequest: async (r, resultSchema, options) => {
  333. if (abortController.signal.aborted) {
  334. throw new McpError(ErrorCode.ConnectionClosed, 'Request was cancelled');
  335. }
  336. // Include related-task metadata if this request is part of a task
  337. const requestOptions = { ...options, relatedRequestId: request.id };
  338. if (relatedTaskId && !requestOptions.relatedTask) {
  339. requestOptions.relatedTask = { taskId: relatedTaskId };
  340. }
  341. // Set task status to input_required when sending a request within a task context
  342. // Use the taskId from options (explicit) or fall back to relatedTaskId (inherited)
  343. const effectiveTaskId = requestOptions.relatedTask?.taskId ?? relatedTaskId;
  344. if (effectiveTaskId && taskStore) {
  345. await taskStore.updateTaskStatus(effectiveTaskId, 'input_required');
  346. }
  347. return await this.request(r, resultSchema, requestOptions);
  348. },
  349. authInfo: extra?.authInfo,
  350. requestId: request.id,
  351. requestInfo: extra?.requestInfo,
  352. taskId: relatedTaskId,
  353. taskStore: taskStore,
  354. taskRequestedTtl: taskCreationParams?.ttl,
  355. closeSSEStream: extra?.closeSSEStream,
  356. closeStandaloneSSEStream: extra?.closeStandaloneSSEStream
  357. };
  358. // Starting with Promise.resolve() puts any synchronous errors into the monad as well.
  359. Promise.resolve()
  360. .then(() => {
  361. // If this request asked for task creation, check capability first
  362. if (taskCreationParams) {
  363. // Check if the request method supports task creation
  364. this.assertTaskHandlerCapability(request.method);
  365. }
  366. })
  367. .then(() => handler(request, fullExtra))
  368. .then(async (result) => {
  369. if (abortController.signal.aborted) {
  370. // Request was cancelled
  371. return;
  372. }
  373. const response = {
  374. result,
  375. jsonrpc: '2.0',
  376. id: request.id
  377. };
  378. // Queue or send the response based on whether this is a task-related request
  379. if (relatedTaskId && this._taskMessageQueue) {
  380. await this._enqueueTaskMessage(relatedTaskId, {
  381. type: 'response',
  382. message: response,
  383. timestamp: Date.now()
  384. }, capturedTransport?.sessionId);
  385. }
  386. else {
  387. await capturedTransport?.send(response);
  388. }
  389. }, async (error) => {
  390. if (abortController.signal.aborted) {
  391. // Request was cancelled
  392. return;
  393. }
  394. const errorResponse = {
  395. jsonrpc: '2.0',
  396. id: request.id,
  397. error: {
  398. code: Number.isSafeInteger(error['code']) ? error['code'] : ErrorCode.InternalError,
  399. message: error.message ?? 'Internal error',
  400. ...(error['data'] !== undefined && { data: error['data'] })
  401. }
  402. };
  403. // Queue or send the error response based on whether this is a task-related request
  404. if (relatedTaskId && this._taskMessageQueue) {
  405. await this._enqueueTaskMessage(relatedTaskId, {
  406. type: 'error',
  407. message: errorResponse,
  408. timestamp: Date.now()
  409. }, capturedTransport?.sessionId);
  410. }
  411. else {
  412. await capturedTransport?.send(errorResponse);
  413. }
  414. })
  415. .catch(error => this._onerror(new Error(`Failed to send response: ${error}`)))
  416. .finally(() => {
  417. // Only delete if the stored controller is still ours; after close()+connect(),
  418. // a new connection may have reused the same request ID with a different controller.
  419. if (this._requestHandlerAbortControllers.get(request.id) === abortController) {
  420. this._requestHandlerAbortControllers.delete(request.id);
  421. }
  422. });
  423. }
  424. _onprogress(notification) {
  425. const { progressToken, ...params } = notification.params;
  426. const messageId = Number(progressToken);
  427. const handler = this._progressHandlers.get(messageId);
  428. if (!handler) {
  429. this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));
  430. return;
  431. }
  432. const responseHandler = this._responseHandlers.get(messageId);
  433. const timeoutInfo = this._timeoutInfo.get(messageId);
  434. if (timeoutInfo && responseHandler && timeoutInfo.resetTimeoutOnProgress) {
  435. try {
  436. this._resetTimeout(messageId);
  437. }
  438. catch (error) {
  439. // Clean up if maxTotalTimeout was exceeded
  440. this._responseHandlers.delete(messageId);
  441. this._progressHandlers.delete(messageId);
  442. this._cleanupTimeout(messageId);
  443. responseHandler(error);
  444. return;
  445. }
  446. }
  447. handler(params);
  448. }
  449. _onresponse(response) {
  450. const messageId = Number(response.id);
  451. // Check if this is a response to a queued request
  452. const resolver = this._requestResolvers.get(messageId);
  453. if (resolver) {
  454. this._requestResolvers.delete(messageId);
  455. if (isJSONRPCResultResponse(response)) {
  456. resolver(response);
  457. }
  458. else {
  459. const error = new McpError(response.error.code, response.error.message, response.error.data);
  460. resolver(error);
  461. }
  462. return;
  463. }
  464. const handler = this._responseHandlers.get(messageId);
  465. if (handler === undefined) {
  466. this._onerror(new Error(`Received a response for an unknown message ID: ${JSON.stringify(response)}`));
  467. return;
  468. }
  469. this._responseHandlers.delete(messageId);
  470. this._cleanupTimeout(messageId);
  471. // Keep progress handler alive for CreateTaskResult responses
  472. let isTaskResponse = false;
  473. if (isJSONRPCResultResponse(response) && response.result && typeof response.result === 'object') {
  474. const result = response.result;
  475. if (result.task && typeof result.task === 'object') {
  476. const task = result.task;
  477. if (typeof task.taskId === 'string') {
  478. isTaskResponse = true;
  479. this._taskProgressTokens.set(task.taskId, messageId);
  480. }
  481. }
  482. }
  483. if (!isTaskResponse) {
  484. this._progressHandlers.delete(messageId);
  485. }
  486. if (isJSONRPCResultResponse(response)) {
  487. handler(response);
  488. }
  489. else {
  490. const error = McpError.fromError(response.error.code, response.error.message, response.error.data);
  491. handler(error);
  492. }
  493. }
  494. get transport() {
  495. return this._transport;
  496. }
  497. /**
  498. * Closes the connection.
  499. */
  500. async close() {
  501. await this._transport?.close();
  502. }
  503. /**
  504. * Sends a request and returns an AsyncGenerator that yields response messages.
  505. * The generator is guaranteed to end with either a 'result' or 'error' message.
  506. *
  507. * @example
  508. * ```typescript
  509. * const stream = protocol.requestStream(request, resultSchema, options);
  510. * for await (const message of stream) {
  511. * switch (message.type) {
  512. * case 'taskCreated':
  513. * console.log('Task created:', message.task.taskId);
  514. * break;
  515. * case 'taskStatus':
  516. * console.log('Task status:', message.task.status);
  517. * break;
  518. * case 'result':
  519. * console.log('Final result:', message.result);
  520. * break;
  521. * case 'error':
  522. * console.error('Error:', message.error);
  523. * break;
  524. * }
  525. * }
  526. * ```
  527. *
  528. * @experimental Use `client.experimental.tasks.requestStream()` to access this method.
  529. */
  530. async *requestStream(request, resultSchema, options) {
  531. const { task } = options ?? {};
  532. // For non-task requests, just yield the result
  533. if (!task) {
  534. try {
  535. const result = await this.request(request, resultSchema, options);
  536. yield { type: 'result', result };
  537. }
  538. catch (error) {
  539. yield {
  540. type: 'error',
  541. error: error instanceof McpError ? error : new McpError(ErrorCode.InternalError, String(error))
  542. };
  543. }
  544. return;
  545. }
  546. // For task-augmented requests, we need to poll for status
  547. // First, make the request to create the task
  548. let taskId;
  549. try {
  550. // Send the request and get the CreateTaskResult
  551. const createResult = await this.request(request, CreateTaskResultSchema, options);
  552. // Extract taskId from the result
  553. if (createResult.task) {
  554. taskId = createResult.task.taskId;
  555. yield { type: 'taskCreated', task: createResult.task };
  556. }
  557. else {
  558. throw new McpError(ErrorCode.InternalError, 'Task creation did not return a task');
  559. }
  560. // Poll for task completion
  561. while (true) {
  562. // Get current task status
  563. const task = await this.getTask({ taskId }, options);
  564. yield { type: 'taskStatus', task };
  565. // Check if task is terminal
  566. if (isTerminal(task.status)) {
  567. if (task.status === 'completed') {
  568. // Get the final result
  569. const result = await this.getTaskResult({ taskId }, resultSchema, options);
  570. yield { type: 'result', result };
  571. }
  572. else if (task.status === 'failed') {
  573. yield {
  574. type: 'error',
  575. error: new McpError(ErrorCode.InternalError, `Task ${taskId} failed`)
  576. };
  577. }
  578. else if (task.status === 'cancelled') {
  579. yield {
  580. type: 'error',
  581. error: new McpError(ErrorCode.InternalError, `Task ${taskId} was cancelled`)
  582. };
  583. }
  584. return;
  585. }
  586. // When input_required, call tasks/result to deliver queued messages
  587. // (elicitation, sampling) via SSE and block until terminal
  588. if (task.status === 'input_required') {
  589. const result = await this.getTaskResult({ taskId }, resultSchema, options);
  590. yield { type: 'result', result };
  591. return;
  592. }
  593. // Wait before polling again
  594. const pollInterval = task.pollInterval ?? this._options?.defaultTaskPollInterval ?? 1000;
  595. await new Promise(resolve => setTimeout(resolve, pollInterval));
  596. // Check if cancelled
  597. options?.signal?.throwIfAborted();
  598. }
  599. }
  600. catch (error) {
  601. yield {
  602. type: 'error',
  603. error: error instanceof McpError ? error : new McpError(ErrorCode.InternalError, String(error))
  604. };
  605. }
  606. }
  607. /**
  608. * Sends a request and waits for a response.
  609. *
  610. * Do not use this method to emit notifications! Use notification() instead.
  611. */
  612. request(request, resultSchema, options) {
  613. const { relatedRequestId, resumptionToken, onresumptiontoken, task, relatedTask } = options ?? {};
  614. // Send the request
  615. return new Promise((resolve, reject) => {
  616. const earlyReject = (error) => {
  617. reject(error);
  618. };
  619. if (!this._transport) {
  620. earlyReject(new Error('Not connected'));
  621. return;
  622. }
  623. if (this._options?.enforceStrictCapabilities === true) {
  624. try {
  625. this.assertCapabilityForMethod(request.method);
  626. // If task creation is requested, also check task capabilities
  627. if (task) {
  628. this.assertTaskCapability(request.method);
  629. }
  630. }
  631. catch (e) {
  632. earlyReject(e);
  633. return;
  634. }
  635. }
  636. options?.signal?.throwIfAborted();
  637. const messageId = this._requestMessageId++;
  638. const jsonrpcRequest = {
  639. ...request,
  640. jsonrpc: '2.0',
  641. id: messageId
  642. };
  643. if (options?.onprogress) {
  644. this._progressHandlers.set(messageId, options.onprogress);
  645. jsonrpcRequest.params = {
  646. ...request.params,
  647. _meta: {
  648. ...(request.params?._meta || {}),
  649. progressToken: messageId
  650. }
  651. };
  652. }
  653. // Augment with task creation parameters if provided
  654. if (task) {
  655. jsonrpcRequest.params = {
  656. ...jsonrpcRequest.params,
  657. task: task
  658. };
  659. }
  660. // Augment with related task metadata if relatedTask is provided
  661. if (relatedTask) {
  662. jsonrpcRequest.params = {
  663. ...jsonrpcRequest.params,
  664. _meta: {
  665. ...(jsonrpcRequest.params?._meta || {}),
  666. [RELATED_TASK_META_KEY]: relatedTask
  667. }
  668. };
  669. }
  670. const cancel = (reason) => {
  671. this._responseHandlers.delete(messageId);
  672. this._progressHandlers.delete(messageId);
  673. this._cleanupTimeout(messageId);
  674. this._transport
  675. ?.send({
  676. jsonrpc: '2.0',
  677. method: 'notifications/cancelled',
  678. params: {
  679. requestId: messageId,
  680. reason: String(reason)
  681. }
  682. }, { relatedRequestId, resumptionToken, onresumptiontoken })
  683. .catch(error => this._onerror(new Error(`Failed to send cancellation: ${error}`)));
  684. // Wrap the reason in an McpError if it isn't already
  685. const error = reason instanceof McpError ? reason : new McpError(ErrorCode.RequestTimeout, String(reason));
  686. reject(error);
  687. };
  688. this._responseHandlers.set(messageId, response => {
  689. if (options?.signal?.aborted) {
  690. return;
  691. }
  692. if (response instanceof Error) {
  693. return reject(response);
  694. }
  695. try {
  696. const parseResult = safeParse(resultSchema, response.result);
  697. if (!parseResult.success) {
  698. // Type guard: if success is false, error is guaranteed to exist
  699. reject(parseResult.error);
  700. }
  701. else {
  702. resolve(parseResult.data);
  703. }
  704. }
  705. catch (error) {
  706. reject(error);
  707. }
  708. });
  709. options?.signal?.addEventListener('abort', () => {
  710. cancel(options?.signal?.reason);
  711. });
  712. const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC;
  713. const timeoutHandler = () => cancel(McpError.fromError(ErrorCode.RequestTimeout, 'Request timed out', { timeout }));
  714. this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false);
  715. // Queue request if related to a task
  716. const relatedTaskId = relatedTask?.taskId;
  717. if (relatedTaskId) {
  718. // Store the response resolver for this request so responses can be routed back
  719. const responseResolver = (response) => {
  720. const handler = this._responseHandlers.get(messageId);
  721. if (handler) {
  722. handler(response);
  723. }
  724. else {
  725. // Log error when resolver is missing, but don't fail
  726. this._onerror(new Error(`Response handler missing for side-channeled request ${messageId}`));
  727. }
  728. };
  729. this._requestResolvers.set(messageId, responseResolver);
  730. this._enqueueTaskMessage(relatedTaskId, {
  731. type: 'request',
  732. message: jsonrpcRequest,
  733. timestamp: Date.now()
  734. }).catch(error => {
  735. this._cleanupTimeout(messageId);
  736. reject(error);
  737. });
  738. // Don't send through transport - queued messages are delivered via tasks/result only
  739. // This prevents duplicate delivery for bidirectional transports
  740. }
  741. else {
  742. // No related task - send through transport normally
  743. this._transport.send(jsonrpcRequest, { relatedRequestId, resumptionToken, onresumptiontoken }).catch(error => {
  744. this._cleanupTimeout(messageId);
  745. reject(error);
  746. });
  747. }
  748. });
  749. }
  750. /**
  751. * Gets the current status of a task.
  752. *
  753. * @experimental Use `client.experimental.tasks.getTask()` to access this method.
  754. */
  755. async getTask(params, options) {
  756. // @ts-expect-error SendRequestT cannot directly contain GetTaskRequest, but we ensure all type instantiations contain it anyways
  757. return this.request({ method: 'tasks/get', params }, GetTaskResultSchema, options);
  758. }
  759. /**
  760. * Retrieves the result of a completed task.
  761. *
  762. * @experimental Use `client.experimental.tasks.getTaskResult()` to access this method.
  763. */
  764. async getTaskResult(params, resultSchema, options) {
  765. // @ts-expect-error SendRequestT cannot directly contain GetTaskPayloadRequest, but we ensure all type instantiations contain it anyways
  766. return this.request({ method: 'tasks/result', params }, resultSchema, options);
  767. }
  768. /**
  769. * Lists tasks, optionally starting from a pagination cursor.
  770. *
  771. * @experimental Use `client.experimental.tasks.listTasks()` to access this method.
  772. */
  773. async listTasks(params, options) {
  774. // @ts-expect-error SendRequestT cannot directly contain ListTasksRequest, but we ensure all type instantiations contain it anyways
  775. return this.request({ method: 'tasks/list', params }, ListTasksResultSchema, options);
  776. }
  777. /**
  778. * Cancels a specific task.
  779. *
  780. * @experimental Use `client.experimental.tasks.cancelTask()` to access this method.
  781. */
  782. async cancelTask(params, options) {
  783. // @ts-expect-error SendRequestT cannot directly contain CancelTaskRequest, but we ensure all type instantiations contain it anyways
  784. return this.request({ method: 'tasks/cancel', params }, CancelTaskResultSchema, options);
  785. }
  786. /**
  787. * Emits a notification, which is a one-way message that does not expect a response.
  788. */
  789. async notification(notification, options) {
  790. if (!this._transport) {
  791. throw new Error('Not connected');
  792. }
  793. this.assertNotificationCapability(notification.method);
  794. // Queue notification if related to a task
  795. const relatedTaskId = options?.relatedTask?.taskId;
  796. if (relatedTaskId) {
  797. // Build the JSONRPC notification with metadata
  798. const jsonrpcNotification = {
  799. ...notification,
  800. jsonrpc: '2.0',
  801. params: {
  802. ...notification.params,
  803. _meta: {
  804. ...(notification.params?._meta || {}),
  805. [RELATED_TASK_META_KEY]: options.relatedTask
  806. }
  807. }
  808. };
  809. await this._enqueueTaskMessage(relatedTaskId, {
  810. type: 'notification',
  811. message: jsonrpcNotification,
  812. timestamp: Date.now()
  813. });
  814. // Don't send through transport - queued messages are delivered via tasks/result only
  815. // This prevents duplicate delivery for bidirectional transports
  816. return;
  817. }
  818. const debouncedMethods = this._options?.debouncedNotificationMethods ?? [];
  819. // A notification can only be debounced if it's in the list AND it's "simple"
  820. // (i.e., has no parameters and no related request ID or related task that could be lost).
  821. const canDebounce = debouncedMethods.includes(notification.method) && !notification.params && !options?.relatedRequestId && !options?.relatedTask;
  822. if (canDebounce) {
  823. // If a notification of this type is already scheduled, do nothing.
  824. if (this._pendingDebouncedNotifications.has(notification.method)) {
  825. return;
  826. }
  827. // Mark this notification type as pending.
  828. this._pendingDebouncedNotifications.add(notification.method);
  829. // Schedule the actual send to happen in the next microtask.
  830. // This allows all synchronous calls in the current event loop tick to be coalesced.
  831. Promise.resolve().then(() => {
  832. // Un-mark the notification so the next one can be scheduled.
  833. this._pendingDebouncedNotifications.delete(notification.method);
  834. // SAFETY CHECK: If the connection was closed while this was pending, abort.
  835. if (!this._transport) {
  836. return;
  837. }
  838. let jsonrpcNotification = {
  839. ...notification,
  840. jsonrpc: '2.0'
  841. };
  842. // Augment with related task metadata if relatedTask is provided
  843. if (options?.relatedTask) {
  844. jsonrpcNotification = {
  845. ...jsonrpcNotification,
  846. params: {
  847. ...jsonrpcNotification.params,
  848. _meta: {
  849. ...(jsonrpcNotification.params?._meta || {}),
  850. [RELATED_TASK_META_KEY]: options.relatedTask
  851. }
  852. }
  853. };
  854. }
  855. // Send the notification, but don't await it here to avoid blocking.
  856. // Handle potential errors with a .catch().
  857. this._transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error));
  858. });
  859. // Return immediately.
  860. return;
  861. }
  862. let jsonrpcNotification = {
  863. ...notification,
  864. jsonrpc: '2.0'
  865. };
  866. // Augment with related task metadata if relatedTask is provided
  867. if (options?.relatedTask) {
  868. jsonrpcNotification = {
  869. ...jsonrpcNotification,
  870. params: {
  871. ...jsonrpcNotification.params,
  872. _meta: {
  873. ...(jsonrpcNotification.params?._meta || {}),
  874. [RELATED_TASK_META_KEY]: options.relatedTask
  875. }
  876. }
  877. };
  878. }
  879. await this._transport.send(jsonrpcNotification, options);
  880. }
  881. /**
  882. * Registers a handler to invoke when this protocol object receives a request with the given method.
  883. *
  884. * Note that this will replace any previous request handler for the same method.
  885. */
  886. setRequestHandler(requestSchema, handler) {
  887. const method = getMethodLiteral(requestSchema);
  888. this.assertRequestHandlerCapability(method);
  889. this._requestHandlers.set(method, (request, extra) => {
  890. const parsed = parseWithCompat(requestSchema, request);
  891. return Promise.resolve(handler(parsed, extra));
  892. });
  893. }
  894. /**
  895. * Removes the request handler for the given method.
  896. */
  897. removeRequestHandler(method) {
  898. this._requestHandlers.delete(method);
  899. }
  900. /**
  901. * Asserts that a request handler has not already been set for the given method, in preparation for a new one being automatically installed.
  902. */
  903. assertCanSetRequestHandler(method) {
  904. if (this._requestHandlers.has(method)) {
  905. throw new Error(`A request handler for ${method} already exists, which would be overridden`);
  906. }
  907. }
  908. /**
  909. * Registers a handler to invoke when this protocol object receives a notification with the given method.
  910. *
  911. * Note that this will replace any previous notification handler for the same method.
  912. */
  913. setNotificationHandler(notificationSchema, handler) {
  914. const method = getMethodLiteral(notificationSchema);
  915. this._notificationHandlers.set(method, notification => {
  916. const parsed = parseWithCompat(notificationSchema, notification);
  917. return Promise.resolve(handler(parsed));
  918. });
  919. }
  920. /**
  921. * Removes the notification handler for the given method.
  922. */
  923. removeNotificationHandler(method) {
  924. this._notificationHandlers.delete(method);
  925. }
  926. /**
  927. * Cleans up the progress handler associated with a task.
  928. * This should be called when a task reaches a terminal status.
  929. */
  930. _cleanupTaskProgressHandler(taskId) {
  931. const progressToken = this._taskProgressTokens.get(taskId);
  932. if (progressToken !== undefined) {
  933. this._progressHandlers.delete(progressToken);
  934. this._taskProgressTokens.delete(taskId);
  935. }
  936. }
  937. /**
  938. * Enqueues a task-related message for side-channel delivery via tasks/result.
  939. * @param taskId The task ID to associate the message with
  940. * @param message The message to enqueue
  941. * @param sessionId Optional session ID for binding the operation to a specific session
  942. * @throws Error if taskStore is not configured or if enqueue fails (e.g., queue overflow)
  943. *
  944. * Note: If enqueue fails, it's the TaskMessageQueue implementation's responsibility to handle
  945. * the error appropriately (e.g., by failing the task, logging, etc.). The Protocol layer
  946. * simply propagates the error.
  947. */
  948. async _enqueueTaskMessage(taskId, message, sessionId) {
  949. // Task message queues are only used when taskStore is configured
  950. if (!this._taskStore || !this._taskMessageQueue) {
  951. throw new Error('Cannot enqueue task message: taskStore and taskMessageQueue are not configured');
  952. }
  953. const maxQueueSize = this._options?.maxTaskQueueSize;
  954. await this._taskMessageQueue.enqueue(taskId, message, sessionId, maxQueueSize);
  955. }
  956. /**
  957. * Clears the message queue for a task and rejects any pending request resolvers.
  958. * @param taskId The task ID whose queue should be cleared
  959. * @param sessionId Optional session ID for binding the operation to a specific session
  960. */
  961. async _clearTaskQueue(taskId, sessionId) {
  962. if (this._taskMessageQueue) {
  963. // Reject any pending request resolvers
  964. const messages = await this._taskMessageQueue.dequeueAll(taskId, sessionId);
  965. for (const message of messages) {
  966. if (message.type === 'request' && isJSONRPCRequest(message.message)) {
  967. // Extract request ID from the message
  968. const requestId = message.message.id;
  969. const resolver = this._requestResolvers.get(requestId);
  970. if (resolver) {
  971. resolver(new McpError(ErrorCode.InternalError, 'Task cancelled or completed'));
  972. this._requestResolvers.delete(requestId);
  973. }
  974. else {
  975. // Log error when resolver is missing during cleanup for better observability
  976. this._onerror(new Error(`Resolver missing for request ${requestId} during task ${taskId} cleanup`));
  977. }
  978. }
  979. }
  980. }
  981. }
  982. /**
  983. * Waits for a task update (new messages or status change) with abort signal support.
  984. * Uses polling to check for updates at the task's configured poll interval.
  985. * @param taskId The task ID to wait for
  986. * @param signal Abort signal to cancel the wait
  987. * @returns Promise that resolves when an update occurs or rejects if aborted
  988. */
  989. async _waitForTaskUpdate(taskId, signal) {
  990. // Get the task's poll interval, falling back to default
  991. let interval = this._options?.defaultTaskPollInterval ?? 1000;
  992. try {
  993. const task = await this._taskStore?.getTask(taskId);
  994. if (task?.pollInterval) {
  995. interval = task.pollInterval;
  996. }
  997. }
  998. catch {
  999. // Use default interval if task lookup fails
  1000. }
  1001. return new Promise((resolve, reject) => {
  1002. if (signal.aborted) {
  1003. reject(new McpError(ErrorCode.InvalidRequest, 'Request cancelled'));
  1004. return;
  1005. }
  1006. // Wait for the poll interval, then resolve so caller can check for updates
  1007. const timeoutId = setTimeout(resolve, interval);
  1008. // Clean up timeout and reject if aborted
  1009. signal.addEventListener('abort', () => {
  1010. clearTimeout(timeoutId);
  1011. reject(new McpError(ErrorCode.InvalidRequest, 'Request cancelled'));
  1012. }, { once: true });
  1013. });
  1014. }
  1015. requestTaskStore(request, sessionId) {
  1016. const taskStore = this._taskStore;
  1017. if (!taskStore) {
  1018. throw new Error('No task store configured');
  1019. }
  1020. return {
  1021. createTask: async (taskParams) => {
  1022. if (!request) {
  1023. throw new Error('No request provided');
  1024. }
  1025. return await taskStore.createTask(taskParams, request.id, {
  1026. method: request.method,
  1027. params: request.params
  1028. }, sessionId);
  1029. },
  1030. getTask: async (taskId) => {
  1031. const task = await taskStore.getTask(taskId, sessionId);
  1032. if (!task) {
  1033. throw new McpError(ErrorCode.InvalidParams, 'Failed to retrieve task: Task not found');
  1034. }
  1035. return task;
  1036. },
  1037. storeTaskResult: async (taskId, status, result) => {
  1038. await taskStore.storeTaskResult(taskId, status, result, sessionId);
  1039. // Get updated task state and send notification
  1040. const task = await taskStore.getTask(taskId, sessionId);
  1041. if (task) {
  1042. const notification = TaskStatusNotificationSchema.parse({
  1043. method: 'notifications/tasks/status',
  1044. params: task
  1045. });
  1046. await this.notification(notification);
  1047. if (isTerminal(task.status)) {
  1048. this._cleanupTaskProgressHandler(taskId);
  1049. // Don't clear queue here - it will be cleared after delivery via tasks/result
  1050. }
  1051. }
  1052. },
  1053. getTaskResult: taskId => {
  1054. return taskStore.getTaskResult(taskId, sessionId);
  1055. },
  1056. updateTaskStatus: async (taskId, status, statusMessage) => {
  1057. // Check if task exists
  1058. const task = await taskStore.getTask(taskId, sessionId);
  1059. if (!task) {
  1060. throw new McpError(ErrorCode.InvalidParams, `Task "${taskId}" not found - it may have been cleaned up`);
  1061. }
  1062. // Don't allow transitions from terminal states
  1063. if (isTerminal(task.status)) {
  1064. throw new McpError(ErrorCode.InvalidParams, `Cannot update task "${taskId}" from terminal status "${task.status}" to "${status}". Terminal states (completed, failed, cancelled) cannot transition to other states.`);
  1065. }
  1066. await taskStore.updateTaskStatus(taskId, status, statusMessage, sessionId);
  1067. // Get updated task state and send notification
  1068. const updatedTask = await taskStore.getTask(taskId, sessionId);
  1069. if (updatedTask) {
  1070. const notification = TaskStatusNotificationSchema.parse({
  1071. method: 'notifications/tasks/status',
  1072. params: updatedTask
  1073. });
  1074. await this.notification(notification);
  1075. if (isTerminal(updatedTask.status)) {
  1076. this._cleanupTaskProgressHandler(taskId);
  1077. // Don't clear queue here - it will be cleared after delivery via tasks/result
  1078. }
  1079. }
  1080. },
  1081. listTasks: cursor => {
  1082. return taskStore.listTasks(cursor, sessionId);
  1083. }
  1084. };
  1085. }
  1086. }
  1087. function isPlainObject(value) {
  1088. return value !== null && typeof value === 'object' && !Array.isArray(value);
  1089. }
  1090. export function mergeCapabilities(base, additional) {
  1091. const result = { ...base };
  1092. for (const key in additional) {
  1093. const k = key;
  1094. const addValue = additional[k];
  1095. if (addValue === undefined)
  1096. continue;
  1097. const baseValue = result[k];
  1098. if (isPlainObject(baseValue) && isPlainObject(addValue)) {
  1099. result[k] = { ...baseValue, ...addValue };
  1100. }
  1101. else {
  1102. result[k] = addValue;
  1103. }
  1104. }
  1105. return result;
  1106. }
  1107. //# sourceMappingURL=protocol.js.map