feat(mcp): implement experimental tasks API for streaming workflow status to chat

Signed-off-by: betterclever <paliwal.pranjal83@gmail.com>

Author: betterclever

backend/src/studio-mcp/__tests__/studio-mcp.service.spec.ts

@@ -4,9 +4,9 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import type { AuthContext } from '../../auth/types';
 import type { WorkflowsService } from '../../workflows/workflows.service';
 
-// Helper to access private _registeredTools on McpServer (plain object at runtime)
-type ToolHandler = (...args: unknown[]) => unknown;
-type RegisteredToolsMap = Record<string, { handler: ToolHandler }>;
+// Helper to access private _registeredTools and experimental tasks on McpServer (plain object at runtime)
+type RegisteredToolsMap = Record<string, any>;
+
 function getRegisteredTools(server: McpServer): RegisteredToolsMap {
   return (server as unknown as { _registeredTools: RegisteredToolsMap })._registeredTools;
 }
@@ -60,12 +60,11 @@ describe('StudioMcpService Unit Tests', () => {
       expect(server).toBeInstanceOf(McpServer);
     });
 
-    it('registers all 9 expected tools', () => {
+    it('registers all expected tools and tasks', () => {
       const server = service.createServer(mockAuthContext);
       const registeredTools = getRegisteredTools(server);
 
       expect(registeredTools).toBeDefined();
-      expect(Object.keys(registeredTools).length).toBe(9);
 
       const toolNames = Object.keys(registeredTools).sort();
       expect(toolNames).toEqual([
@@ -110,16 +109,27 @@ describe('StudioMcpService Unit Tests', () => {
       expect(workflowsService.findById).toHaveBeenCalledWith(workflowId, mockAuthContext);
     });
 
-    it('run_workflow tool uses auth context passed at creation time', async () => {
+    it('run_workflow task uses auth context passed at creation time', async () => {
       const workflowId = '11111111-1111-4111-8111-111111111111';
       const inputs = { key: 'value' };
 
       const server = service.createServer(mockAuthContext);
       const registeredTools = getRegisteredTools(server);
-      const runWorkflowTool = registeredTools['run_workflow'];
+      const runWorkflowTask = registeredTools['run_workflow'];
+
+      expect(runWorkflowTask).toBeDefined();
+
+      // Need to mock the extra params for the experimental tasks
+      const mockExtra = {
+        taskStore: {
+          createTask: jest.fn().mockResolvedValue({ taskId: 'mockTaskId', status: 'working' }),
+          getTask: jest.fn().mockResolvedValue({ taskId: 'mockTaskId', status: 'working' }),
+          updateTaskStatus: jest.fn().mockResolvedValue(true),
+          storeTaskResult: jest.fn().mockResolvedValue(true),
+        },
+      };
 
-      expect(runWorkflowTool).toBeDefined();
-      await runWorkflowTool.handler({ workflowId, inputs });
+      await runWorkflowTask.handler.createTask({ workflowId, inputs }, mockExtra);
 
       expect(workflowsService.run).toHaveBeenCalledWith(
         workflowId,
@@ -129,7 +139,7 @@ describe('StudioMcpService Unit Tests', () => {
           trigger: {
             type: 'api',
             sourceId: mockAuthContext.userId,
-            label: 'Studio MCP',
+            label: 'Studio MCP Task',
           },
         },
       );
@@ -230,12 +240,21 @@ describe('StudioMcpService Unit Tests', () => {
 
       it('denies run_workflow when workflows.run is false', async () => {
         const server = service.createServer(restrictedAuth);
-        const tools = getRegisteredTools(server);
-        const result = (await tools['run_workflow'].handler({
-          workflowId: '11111111-1111-4111-8111-111111111111',
-        })) as { isError?: boolean; content: { text: string }[] };
-        expect(result.isError).toBe(true);
-        expect(result.content[0].text).toContain('workflows.run');
+        const tasks = getRegisteredTools(server);
+
+        let errorThrown = false;
+        try {
+          await tasks['run_workflow'].handler.createTask(
+            {
+              workflowId: '11111111-1111-4111-8111-111111111111',
+            },
+            {} as any,
+          );
+        } catch (_e: any) {
+          errorThrown = true;
+          expect(_e.message).toContain('workflows.run');
+        }
+        expect(errorThrown).toBe(true);
       });
 
       it('denies cancel_run when runs.cancel is false', async () => {
@@ -260,15 +279,28 @@ describe('StudioMcpService Unit Tests', () => {
       it('allows all tools when no apiKeyPermissions (non-API-key auth)', async () => {
         const server = service.createServer(mockAuthContext); // no apiKeyPermissions
         const tools = getRegisteredTools(server);
+        const tasks = getRegisteredTools(server);
 
         // All workflow/run tools should work without permission errors
         const listResult = (await tools['list_workflows'].handler({})) as { isError?: boolean };
         expect(listResult.isError).toBeUndefined();
 
-        const runResult = (await tools['run_workflow'].handler({
-          workflowId: '11111111-1111-4111-8111-111111111111',
-        })) as { isError?: boolean };
-        expect(runResult.isError).toBeUndefined();
+        const mockExtra = {
+          taskStore: {
+            createTask: jest.fn().mockResolvedValue({ taskId: 'mock', status: 'working' }),
+            getTask: jest.fn().mockResolvedValue({ taskId: 'mock', status: 'working' }),
+            updateTaskStatus: jest.fn().mockResolvedValue(true),
+            storeTaskResult: jest.fn().mockResolvedValue(true),
+          },
+        };
+
+        const runResult = await tasks['run_workflow'].handler.createTask(
+          {
+            workflowId: '11111111-1111-4111-8111-111111111111',
+          },
+          mockExtra,
+        );
+        expect(runResult.task.taskId).toEqual('mock');
 
         const cancelResult = (await tools['cancel_run'].handler({
           runId: 'test-run-id',
@@ -308,11 +340,11 @@ describe('StudioMcpService Unit Tests', () => {
         };
         const server = service.createServer(noPermsAuth);
         const tools = getRegisteredTools(server);
+        const tasks = getRegisteredTools(server);
 
         const gatedTools = [
           'list_workflows',
           'get_workflow',
-          'run_workflow',
           'list_runs',
           'get_run_status',
           'get_run_result',
@@ -326,6 +358,20 @@ describe('StudioMcpService Unit Tests', () => {
           })) as { isError?: boolean };
           expect(result.isError).toBe(true);
         }
+
+        // Test run_workflow separately since it's a task now
+        let errorThrown = false;
+        try {
+          await tasks['run_workflow'].handler.createTask(
+            {
+              workflowId: '11111111-1111-4111-8111-111111111111',
+            },
+            {} as any,
+          );
+        } catch (_e: any) {
+          errorThrown = true;
+        }
+        expect(errorThrown).toBe(true);
       });
     });
 

backend/src/studio-mcp/studio-mcp.service.ts

@@ -138,57 +138,71 @@ export class StudioMcpService {
       },
     );
 
-    server.registerTool(
+    const runWorkflowSchema = {
+      workflowId: z.string().uuid(),
+      inputs: z.record(z.string(), z.unknown()).optional(),
+      versionId: z.string().uuid().optional(),
+    };
+
+    server.experimental.tasks.registerToolTask(
       'run_workflow',
       {
         description:
-          'Start a workflow execution. Returns the run ID and initial status. Use get_run_status to poll for completion.',
-        inputSchema: {
-          workflowId: z.string().uuid(),
-          inputs: z.record(z.string(), z.unknown()).optional(),
-          versionId: z.string().uuid().optional(),
-        },
+          'Start a workflow execution as a background task. The task handle can be monitored for status updates, and finally retrieved for the workflow result. Also supports legacy polling via get_run_status.',
+        inputSchema: runWorkflowSchema,
+        execution: { taskSupport: 'optional' },
       },
-      async (args: {
-        workflowId: string;
-        inputs?: Record<string, unknown>;
-        versionId?: string;
-      }) => {
-        const gate = this.checkPermission(auth, 'workflows.run');
-        if (!gate.allowed) return gate.error;
-        try {
+      {
+        createTask: async (args, extra) => {
+          const gate = this.checkPermission(auth, 'workflows.run');
+          if (!gate.allowed) throw new Error(gate.error.content[0].text);
+
+          const task = await extra.taskStore.createTask({ ttl: 12 * 60 * 60 * 1000 });
+
           const handle = await this.workflowsService.run(
             args.workflowId,
-            { inputs: args.inputs ?? {}, versionId: args.versionId },
+            {
+              inputs: args.inputs ?? {},
+              versionId: args.versionId,
+            },
             auth,
             {
               trigger: {
                 type: 'api',
                 sourceId: auth.userId ?? 'api-key',
-                label: 'Studio MCP',
+                label: 'Studio MCP Task',
               },
             },
           );
-          return {
-            content: [
-              {
-                type: 'text' as const,
-                text: JSON.stringify(
-                  {
-                    runId: handle.runId,
-                    workflowId: handle.workflowId,
-                    status: handle.status,
-                    workflowVersion: handle.workflowVersion,
-                  },
-                  null,
-                  2,
-                ),
-              },
-            ],
-          };
-        } catch (error) {
-          return this.errorResult(error);
-        }
+
+          this.monitorWorkflowRun(
+            handle.runId,
+            handle.temporalRunId,
+            task.taskId,
+            extra.taskStore,
+            server,
+            auth,
+          ).catch((err) => {
+            this.logger.error(`Error monitoring workflow run task for run ${handle.runId}: ${err}`);
+          });
+
+          return { task };
+        },
+        getTask: async (args, extra) => {
+          const gate = this.checkPermission(auth, 'runs.read');
+          if (!gate.allowed) throw new Error(gate.error.content[0].text);
+          const task = await extra.taskStore.getTask(extra.taskId);
+          if (!task) {
+            throw new Error(`Task ${extra.taskId} not found`);
+          }
+          return task;
+        },
+        getTaskResult: async (args, extra) => {
+          const gate = this.checkPermission(auth, 'runs.read');
+          if (!gate.allowed) throw new Error(gate.error.content[0].text);
+          const result = await extra.taskStore.getTaskResult(extra.taskId);
+          return result as any;
+        },
       },
     );
   }
@@ -397,6 +411,85 @@ export class StudioMcpService {
     );
   }
 
+  private async monitorWorkflowRun(
+    runId: string,
+    temporalRunId: string | undefined,
+    taskId: string,
+    taskStore: any,
+    server: McpServer,
+    auth: AuthContext,
+  ): Promise<void> {
+    const isTerminal = (status: string) =>
+      ['COMPLETED', 'FAILED', 'CANCELLED', 'TERMINATED', 'TIMED_OUT'].includes(status);
+
+    const mapStatus = (status: string): 'working' | 'completed' | 'cancelled' | 'failed' => {
+      switch (status) {
+        case 'RUNNING':
+        case 'QUEUED':
+        case 'AWAITING_INPUT':
+          return 'working';
+        case 'COMPLETED':
+          return 'completed';
+        case 'CANCELLED':
+        case 'TERMINATED':
+        case 'TIMED_OUT':
+          return 'cancelled';
+        case 'FAILED':
+          return 'failed';
+        default:
+          return 'working';
+      }
+    };
+
+    while (true) {
+      try {
+        const runStatusPayload = await this.workflowsService.getRunStatus(
+          runId,
+          temporalRunId,
+          auth,
+        );
+        const taskState = mapStatus(runStatusPayload.status);
+
+        await taskStore.updateTaskStatus(taskId, taskState, runStatusPayload.status);
+
+        if (isTerminal(runStatusPayload.status)) {
+          let resultData: any;
+          if (taskState === 'completed') {
+            try {
+              resultData = await this.workflowsService.getRunResult(runId, temporalRunId, auth);
+            } catch (err) {
+              resultData = { error: String(err) };
+            }
+          } else {
+            resultData = runStatusPayload.failure || { reason: runStatusPayload.status };
+          }
+
+          const resultPayload = {
+            content: [{ type: 'text', text: JSON.stringify(resultData, null, 2) }],
+          };
+
+          const storeStatus = taskState === 'completed' ? 'completed' : 'failed';
+          await taskStore.storeTaskResult(taskId, storeStatus, resultPayload);
+          break;
+        }
+
+        await new Promise((res) => setTimeout(res, 2000));
+      } catch (err) {
+        this.logger.error(`Error monitoring task ${taskId} (run: ${runId}): ${err}`);
+        try {
+          await taskStore.updateTaskStatus(taskId, 'failed', String(err));
+          await taskStore.storeTaskResult(taskId, 'failed', {
+            content: [{ type: 'text', text: `Failed to monitor workflow run: ${String(err)}` }],
+            isError: true,
+          });
+        } catch (_updateErr) {
+          // Ignore
+        }
+        break;
+      }
+    }
+  }
+
   // ---------------------------------------------------------------------------
   // Helpers
   // ---------------------------------------------------------------------------