feat(mcp): stream workflow task status (ENG-161)

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);
       });
     });
 
@@ -366,4 +412,91 @@ describe('StudioMcpService Unit Tests', () => {
       expect(workflowsService.list).toHaveBeenNthCalledWith(2, authContext2);
     });
   });
+
+  describe('monitorWorkflowRun', () => {
+    it('polls status and saves result on completion', async () => {
+      const mockTaskStore = {
+        updateTaskStatus: jest.fn().mockResolvedValue(true),
+        storeTaskResult: jest.fn().mockResolvedValue(true),
+      };
+
+      const mockServer = {} as McpServer;
+      const taskId = 'test-task-id';
+      const runId = 'test-run-id';
+
+      // Mock getRunStatus to return RUNNING first, then COMPLETED
+      let callCount = 0;
+      (workflowsService.getRunStatus as jest.Mock).mockImplementation(() => {
+        callCount++;
+        return Promise.resolve({
+          status: callCount === 1 ? 'RUNNING' : 'COMPLETED',
+        });
+      });
+
+      (workflowsService.getRunResult as jest.Mock).mockResolvedValue({
+        output: 'test-output',
+      });
+
+      // We overwrite the 2000ms timeout temporarily for the test to avoid slow running loop
+      const originalSetTimeout = global.setTimeout;
+      (global as any).setTimeout = (fn: any) => originalSetTimeout(fn, 1);
+
+      try {
+        await (service as any).monitorWorkflowRun(
+          runId,
+          undefined,
+          taskId,
+          mockTaskStore,
+          mockServer,
+          mockAuthContext,
+        );
+      } finally {
+        global.setTimeout = originalSetTimeout as any;
+      }
+
+      // updateTaskStatus is only called for non-terminal states (RUNNING → working).
+      // For COMPLETED, storeTaskResult handles the terminal transition directly.
+      expect(mockTaskStore.updateTaskStatus).toHaveBeenCalledTimes(1);
+      expect(mockTaskStore.updateTaskStatus).toHaveBeenCalledWith(taskId, 'working', 'RUNNING');
+      expect(mockTaskStore.updateTaskStatus).not.toHaveBeenCalledWith(
+        taskId,
+        'completed',
+        'COMPLETED',
+      );
+      expect(workflowsService.getRunResult).toHaveBeenCalledWith(runId, undefined, mockAuthContext);
+      expect(mockTaskStore.storeTaskResult).toHaveBeenCalledWith(taskId, 'completed', {
+        content: [{ type: 'text', text: JSON.stringify({ output: 'test-output' }, null, 2) }],
+      });
+    });
+
+    it('handles failures by storing the failure reason', async () => {
+      const mockTaskStore = {
+        updateTaskStatus: jest.fn().mockResolvedValue(true),
+        storeTaskResult: jest.fn().mockResolvedValue(true),
+      };
+
+      const taskId = 'test-task-id';
+      const runId = 'test-run-id';
+
+      (workflowsService.getRunStatus as jest.Mock).mockResolvedValue({
+        status: 'FAILED',
+        failure: { message: 'boom' },
+      });
+
+      await (service as any).monitorWorkflowRun(
+        runId,
+        undefined,
+        taskId,
+        mockTaskStore,
+        {} as McpServer,
+        mockAuthContext,
+      );
+
+      // updateTaskStatus is NOT called for terminal states — storeTaskResult handles it.
+      expect(mockTaskStore.updateTaskStatus).not.toHaveBeenCalled();
+      expect(mockTaskStore.storeTaskResult).toHaveBeenCalledWith(taskId, 'failed', {
+        content: [{ type: 'text', text: JSON.stringify({ message: 'boom' }, null, 2) }],
+      });
+    });
+  });
 });