11import { InMemoryTaskStore } from '@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js' ;
22import { describe , expect , it } from 'vitest' ;
33
4+ import { storeTaskResultWithMessage } from '../../src/mcp/utils.js' ;
5+
46/**
57 * Regression tests for #638: tasks/list shows stale statusMessage for completed tasks.
68 *
79 * storeTaskResult() does not accept a statusMessage parameter, so we must call
810 * updateTaskStatus() with the final message before calling storeTaskResult().
9- * These tests verify that statusMessage survives the storeTaskResult() transition.
11+ *
12+ * These tests verify:
13+ * 1. statusMessage survives the storeTaskResult() transition (SDK invariant)
14+ * 2. The correct final message is set for each outcome (success, payment required, failed)
15+ * 3. Non-success paths do NOT get a misleading ": completed" message
1016 */
1117describe ( 'Task statusMessage after terminal transition' , ( ) => {
1218 const REQUEST_ID = 'req-1' ;
@@ -17,39 +23,80 @@ describe('Task statusMessage after terminal transition', () => {
1723 return task . taskId ;
1824 }
1925
20- it ( 'should retain final statusMessage after storeTaskResult completes the task ' , async ( ) => {
26+ it ( 'should set statusMessage and status atomically via storeTaskResultWithMessage ' , async ( ) => {
2127 const store = new InMemoryTaskStore ( ) ;
2228 const taskId = await createWorkingTask ( store ) ;
2329
24- // Intermediate progress message (as set by ProgressTracker during execution)
25- await store . updateTaskStatus ( taskId , 'working' , 'apify/rag-web-browser: Starting the crawler.' ) ;
26-
27- // Final message set just before storeTaskResult (the fix)
28- await store . updateTaskStatus ( taskId , 'working' , 'apify/rag-web-browser: completed' ) ;
29-
30- await store . storeTaskResult ( taskId , 'completed' , {
30+ await storeTaskResultWithMessage ( store , taskId , 'completed' , {
3131 content : [ { type : 'text' , text : 'result data' } ] ,
32- } ) ;
32+ } , 'apify/rag-web-browser: completed' ) ;
3333
3434 const task = await store . getTask ( taskId ) ;
3535 expect ( task ) . not . toBeNull ( ) ;
3636 expect ( task ! . status ) . toBe ( 'completed' ) ;
3737 expect ( task ! . statusMessage ) . toBe ( 'apify/rag-web-browser: completed' ) ;
3838 } ) ;
3939
40- it ( 'should retain final statusMessage after storeTaskResult marks task as failed ' , async ( ) => {
40+ it ( 'should set statusMessage for failed tasks via storeTaskResultWithMessage ' , async ( ) => {
4141 const store = new InMemoryTaskStore ( ) ;
4242 const taskId = await createWorkingTask ( store ) ;
4343
44- await store . updateTaskStatus ( taskId , 'working' , 'my-tool: failed' ) ;
45-
46- await store . storeTaskResult ( taskId , 'failed' , {
44+ await storeTaskResultWithMessage ( store , taskId , 'failed' , {
4745 content : [ { type : 'text' , text : 'error' } ] ,
4846 isError : true ,
49- } ) ;
47+ } , 'my-tool: failed' ) ;
5048
5149 const task = await store . getTask ( taskId ) ;
5250 expect ( task ! . status ) . toBe ( 'failed' ) ;
5351 expect ( task ! . statusMessage ) . toBe ( 'my-tool: failed' ) ;
5452 } ) ;
53+
54+ it ( 'should set "payment required" statusMessage for pre-flight 402 path, not "completed"' , async ( ) => {
55+ const store = new InMemoryTaskStore ( ) ;
56+ const taskId = await createWorkingTask ( store ) ;
57+
58+ // Simulates the pre-flight payment failure path in executeToolAndUpdateTask:
59+ // paymentRequiredResult is set, toolStatus = SOFT_FAIL, no actor execution happens.
60+ // The fix guards the updateTaskStatus call so it stamps ": payment required" instead of ": completed".
61+ await store . updateTaskStatus ( taskId , 'working' , 'apify/rag-web-browser: payment required' ) ;
62+
63+ await store . storeTaskResult ( taskId , 'completed' , {
64+ content : [ { type : 'text' , text : 'Payment required' } ] ,
65+ } ) ;
66+
67+ const task = await store . getTask ( taskId ) ;
68+ expect ( task ! . status ) . toBe ( 'completed' ) ;
69+ expect ( task ! . statusMessage ) . toBe ( 'apify/rag-web-browser: payment required' ) ;
70+ expect ( task ! . statusMessage ) . not . toContain ( 'completed' ) ;
71+ } ) ;
72+
73+ it ( 'should set status "cancelled" when task was aborted via signal, not "completed"' , async ( ) => {
74+ const store = new InMemoryTaskStore ( ) ;
75+ const taskId = await createWorkingTask ( store ) ;
76+
77+ // Simulates the signal-based abort path: executeActorTool returns null,
78+ // toolStatus = ABORTED. The fix transitions to 'cancelled' instead of
79+ // storing as 'completed' with empty result.
80+ await store . updateTaskStatus ( taskId , 'working' , 'apify/rag-web-browser: Running the crawler.' ) ;
81+ await store . updateTaskStatus ( taskId , 'cancelled' , 'apify/rag-web-browser: aborted by client' ) ;
82+
83+ const task = await store . getTask ( taskId ) ;
84+ expect ( task ! . status ) . toBe ( 'cancelled' ) ;
85+ expect ( task ! . statusMessage ) . toBe ( 'apify/rag-web-browser: aborted by client' ) ;
86+ } ) ;
87+
88+ it ( 'should not call updateTaskStatus on cancelled task (402 catch path)' , async ( ) => {
89+ const store = new InMemoryTaskStore ( ) ;
90+ const taskId = await createWorkingTask ( store ) ;
91+
92+ // Simulate: task is cancelled while actor is running
93+ await store . updateTaskStatus ( taskId , 'cancelled' , 'Cancelled by client' ) ;
94+
95+ const task = await store . getTask ( taskId ) ;
96+ expect ( task ! . status ) . toBe ( 'cancelled' ) ;
97+
98+ // Trying to transition from cancelled → working should throw (SDK enforces terminal state).
99+ // The fix adds an isTaskCancelled guard before updateTaskStatus in the 402 catch path.
100+ await expect ( store . updateTaskStatus ( taskId , 'working' , 'some-tool: payment required' ) ) . rejects . toThrow ( ) ;
101+ } ) ;
55102} ) ;
0 commit comments