001-migration-ui-redesign (#3)

Reviewed-on: #3

backend/src/core/task_manager.py

@@ -23,6 +23,7 @@ class TaskStatus(str, Enum):
     RUNNING = "RUNNING"
     SUCCESS = "SUCCESS"
     FAILED = "FAILED"
+    AWAITING_MAPPING = "AWAITING_MAPPING"
 
 # [/DEF]
 
@@ -64,6 +65,7 @@ class TaskManager:
         self.subscribers: Dict[str, List[asyncio.Queue]] = {}
         self.executor = ThreadPoolExecutor(max_workers=5) # For CPU-bound plugin execution
         self.loop = asyncio.get_event_loop()
+        self.task_futures: Dict[str, asyncio.Future] = {}
 # [/DEF]
 
     async def create_task(self, plugin_id: str, params: Dict[str, Any], user_id: Optional[str] = None) -> Task:
@@ -99,9 +101,11 @@ class TaskManager:
             # Execute plugin in a separate thread to avoid blocking the event loop
             # if the plugin's execute method is synchronous and potentially CPU-bound.
             # If the plugin's execute method is already async, this can be simplified.
+            # Pass task_id to plugin so it can signal pause
+            params = {**task.params, "_task_id": task_id}
             await self.loop.run_in_executor(
                 self.executor,
-                lambda: asyncio.run(plugin.execute(task.params)) if asyncio.iscoroutinefunction(plugin.execute) else plugin.execute(task.params)
+                lambda: asyncio.run(plugin.execute(params)) if asyncio.iscoroutinefunction(plugin.execute) else plugin.execute(params)
             )
             task.status = TaskStatus.SUCCESS
             self._add_log(task_id, "INFO", f"Task completed successfully for plugin '{plugin.name}'")
@@ -112,6 +116,38 @@ class TaskManager:
             task.finished_at = datetime.utcnow()
             # In a real system, you might notify clients via WebSocket here
 
+    async def resolve_task(self, task_id: str, resolution_params: Dict[str, Any]):
+        """
+        Resumes a task that is awaiting mapping.
+        """
+        task = self.tasks.get(task_id)
+        if not task or task.status != TaskStatus.AWAITING_MAPPING:
+            raise ValueError("Task is not awaiting mapping.")
+        
+        # Update task params with resolution
+        task.params.update(resolution_params)
+        task.status = TaskStatus.RUNNING
+        self._add_log(task_id, "INFO", "Task resumed after mapping resolution.")
+        
+        # Signal the future to continue
+        if task_id in self.task_futures:
+            self.task_futures[task_id].set_result(True)
+
+    async def wait_for_resolution(self, task_id: str):
+        """
+        Pauses execution and waits for a resolution signal.
+        """
+        task = self.tasks.get(task_id)
+        if not task: return
+        
+        task.status = TaskStatus.AWAITING_MAPPING
+        self.task_futures[task_id] = self.loop.create_future()
+        
+        try:
+            await self.task_futures[task_id]
+        finally:
+            del self.task_futures[task_id]
+
     def get_task(self, task_id: str) -> Optional[Task]:
         """
         Retrieves a task by its ID.