zelos.execution_engine

  1"""
  2Execution Engine — Dispatches Tasks to Agents, monitors lifecycle, enforces timeouts.
  3"""
  4
  5import threading
  6import time
  7import uuid
  8from collections.abc import Callable
  9from dataclasses import dataclass, field
 10
 11from .event_bus import Event, EventBus
 12from .task_graph import Task, TaskGraphEngine, TaskStatus
 13
 14
 15@dataclass
 16class InFlightTask:
 17    task_id: str
 18    agent_id: str
 19    agent_name: str
 20    started_at: float
 21    timeout_at: float
 22    heartbeat_at: float = 0.0  # v0.8.0: last heartbeat timestamp
 23    heartbeat_timeout_ms: int = 30000  # v0.8.0: heartbeat timeout in ms
 24
 25
 26@dataclass
 27class AgentState:
 28    agent_id: str
 29    agent_name: str
 30    status: str = "registered"  # registered → connected → heartbeating → disconnected → shutdown
 31    operational_state: str = "idle"
 32    last_heartbeat_at: float = 0.0
 33    heartbeat_interval_ms: int = 30000
 34    endpoint: str | None = None
 35    max_concurrent_tasks: int = 5
 36    current_tasks: list[str] = field(default_factory=list)
 37    capabilities: list[dict] = field(default_factory=list)
 38    required_credentials: list[str] = field(default_factory=list)  # v1.1.0
 39    historical_success_rate: float = 0.0
 40    total_completed: int = 0
 41    total_failed: int = 0
 42
 43
 44class ExecutionEngine:
 45    """Kernel component — Task dispatch, lifecycle, timeouts, heartbeat tracking.
 46
 47    v1.3.0: MPC replan_check hook, incremental verification, diagnosis trigger.
 48    """
 49
 50    def __init__(self, task_graph: TaskGraphEngine, event_bus: EventBus):
 51        self._task_graph = task_graph
 52        self._event_bus = event_bus
 53        task_graph._event_bus = event_bus  # v0.9.0: wire lifecycle events
 54        self._in_flight: dict[str, InFlightTask] = {}  # task_id → InFlightTask
 55        self._agents: dict[str, AgentState] = {}
 56        self._agent_dispatch: Callable | None = None  # Callback: (agent_id, task) → bool
 57        self._agent_cancel: Callable | None = None
 58        self._lock = threading.RLock()
 59        self._monitor_thread: threading.Thread | None = None
 60        self._running = False
 61        self._task_inputs: dict[str, dict] = {}  # v0.9.0: task input context
 62        self._task_start_times: dict[str, float] = {}  # v0.9.0: task start timestamps
 63        self._credential_injector = None  # v1.1.0: set by runtime
 64
 65        # v1.3.0: MPC Adaptive Loop
 66        self._replan_callback: Callable | None = None  # Runtime._on_replan
 67        self._replan_rules: list = []  # ReplanRule instances
 68        self._current_plan = None  # PlannerPlan reference
 69        self._replan_count: dict[str, int] = {}  # goal_id → count
 70        self._max_replans: int = 5
 71        self._incremental_verifier = None  # SchemaVerifier for per-task check
 72        self._diagnosis_engine = None  # DiagnosisEngine instance (from feature flag)
 73
 74    # ── Agent Management ──
 75
 76    def register_agent(self, agent_id: str, agent_name: str, **kwargs) -> AgentState:
 77        state = AgentState(agent_id=agent_id, agent_name=agent_name, **kwargs)
 78        with self._lock:
 79            self._agents[agent_id] = state
 80        state.status = "connected"
 81        return state
 82
 83    def heartbeat(self, agent_id: str) -> bool:
 84        with self._lock:
 85            agent = self._agents.get(agent_id)
 86            if not agent:
 87                return False
 88            agent.last_heartbeat_at = time.time()
 89            if agent.status == "connected":
 90                agent.status = "heartbeating"
 91            elif agent.status == "disconnected":
 92                agent.status = "heartbeating"
 93            # Update operational state
 94            agent.operational_state = "busy" if agent.current_tasks else "idle"
 95            return True
 96
 97    def remove_agent(self, agent_id: str) -> None:
 98        with self._lock:
 99            self._agents.pop(agent_id, None)
100            # Cancel all in-flight tasks for this agent
101            for tid in list(self._in_flight.keys()):
102                if self._in_flight[tid].agent_id == agent_id:
103                    self.cancel_task(tid)
104
105    def get_agent(self, agent_id: str) -> AgentState | None:
106        return self._agents.get(agent_id)
107
108    def list_agents(self) -> list[AgentState]:
109        return list(self._agents.values())
110
111    # ── Dispatch ──
112
113    def set_task_input_context(self, task_id: str, input_context: dict) -> None:
114        """v0.9.0: Set the input context for a task before dispatch."""
115        with self._lock:
116            self._task_inputs[task_id] = input_context
117
118    def dispatch(self, task_id: str, agent_id: str) -> bool:
119        """Dispatch a task to an agent. Returns True if agent accepted."""
120        task = self._task_graph.get_task(task_id)
121        agent = self._agents.get(agent_id)
122        if not task or not agent:
123            return False
124
125        with self._lock:
126            self._task_graph.transition(task_id, TaskStatus.STARTED, agent_id=agent_id)
127            hb_timeout = getattr(task, 'heartbeat_timeout_ms', 0) or agent.heartbeat_interval_ms * 3
128            now = time.time()
129            in_flight = InFlightTask(
130                task_id=task_id,
131                agent_id=agent_id,
132                agent_name=agent.agent_name,
133                started_at=now,
134                timeout_at=now + (task.timeout_ms / 1000),
135                heartbeat_at=now,
136                heartbeat_timeout_ms=hb_timeout,
137            )
138            self._in_flight[task_id] = in_flight
139            self._task_start_times[task_id] = now  # v0.9.0
140            agent.current_tasks.append(task_id)
141            agent.operational_state = "busy"
142
143        # v1.1.0: Inject credentials before dispatch
144        if self._credential_injector and agent.required_credentials:
145            try:
146                self._credential_injector.inject(task, agent_id, agent.required_credentials)
147            except Exception:
148                # Credential failure → reject task for retry
149                self._in_flight.pop(task_id, None)
150                agent.current_tasks.remove(task_id)
151                try:
152                    self._task_graph.transition(task_id, TaskStatus.FAILED)
153                except ValueError:
154                    pass
155                return False
156
157        # v0.9.0: Publish task.started event with input context
158        input_ctx = self._task_inputs.pop(task_id, None)
159        self._event_bus.publish(Event(
160            event_id=str(uuid.uuid4()),
161            event_type="task.started",
162            source="execution_engine",
163            timestamp=now,
164            correlation_id=task.plan_id,
165            payload={
166                "task_id": task_id,
167                "agent_id": agent_id,
168                "agent_name": agent.agent_name,
169                "required_capability": task.required_capability,
170                "input_context": input_ctx or {},
171            },
172        ))
173
174        # Call the agent dispatch callback (in-process: direct function call)
175        if self._agent_dispatch:
176            self._agent_dispatch(agent_id, task)
177
178        return True
179
180    def reject_task(self, task_id: str) -> None:
181        """Agent rejected the task → re-schedule."""
182        with self._lock:
183            self._in_flight.pop(task_id, None)
184            try:
185                self._task_graph.transition(task_id, TaskStatus.READY)
186            except ValueError:
187                pass
188
189    # ── Result Handling ──
190
191    def submit_result(self, task_id: str, agent_id: str, result: dict) -> bool:
192        """
193        Agent returns task result. result = {status: "completed"|"failed", artifact?: ..., error?: ...}
194
195        v0.8.0: If error code matches task.non_retryable_errors, transition to FATAL_FAILED.
196        """
197        with self._lock:
198            inflight = self._in_flight.pop(task_id, None)
199            if not inflight:
200                return False
201
202            agent = self._agents.get(agent_id)
203            if agent and task_id in agent.current_tasks:
204                agent.current_tasks.remove(task_id)
205                if not agent.current_tasks:
206                    agent.operational_state = "idle"
207
208            status = result.get("status")
209            duration_ms = (time.time() - self._task_start_times.pop(task_id, time.time())) * 1000
210            task = self._task_graph.get_task(task_id)
211
212            if status == "completed":
213                try:
214                    self._task_graph.transition(task_id, TaskStatus.COMPLETED)
215                except ValueError:
216                    return False
217                if agent:
218                    agent.total_completed += 1
219                # v0.9.0: Publish task.completed with artifact
220                artifact = result.get("artifact", {})
221                payload_size = len(str(artifact))
222                event_payload = {
223                    "task_id": task_id, "agent_id": agent_id,
224                    "duration_ms": int(duration_ms),
225                }
226                if payload_size > 100000:
227                    event_payload["content_ref"] = f"storage:task:{task_id}:artifact"
228                    event_payload["output_artifact_summary"] = f"Large artifact ({payload_size} bytes)"
229                else:
230                    event_payload["output_artifact"] = artifact
231                self._event_bus.publish(Event(
232                    event_id=str(uuid.uuid4()),
233                    event_type="task.completed",
234                    source="execution_engine",
235                    timestamp=time.time(),
236                    correlation_id=task.plan_id if task else "",
237                    payload=event_payload,
238                ))
239            else:
240                # v0.8.0: Check non_retryable_errors
241                error_code = (result.get("error") or {}).get("code", "")
242                if task and error_code and task.non_retryable_errors:
243                    if error_code in task.non_retryable_errors:
244                        try:
245                            self._task_graph.transition(task_id, TaskStatus.FATAL_FAILED)
246                        except ValueError:
247                            self._task_graph.transition(task_id, TaskStatus.FAILED)
248                        if agent:
249                            agent.total_failed += 1
250                        return True
251                try:
252                    self._task_graph.transition(task_id, TaskStatus.FAILED)
253                except ValueError:
254                    return False
255                if agent:
256                    agent.total_failed += 1
257                # v0.9.0: Publish task.failed with error
258                self._event_bus.publish(Event(
259                    event_id=str(uuid.uuid4()),
260                    event_type="task.failed",
261                    source="execution_engine",
262                    timestamp=time.time(),
263                    correlation_id=task.plan_id if task else "",
264                    payload={
265                        "task_id": task_id, "agent_id": agent_id,
266                        "error": result.get("error", {}),
267                        "duration_ms": int(duration_ms),
268                    },
269                ))
270
271        # ── v1.3.0: MPC Replan Check (outside lock to avoid deadlock) ──
272        self._mpc_replan_check(task_id, result)
273        return True
274
275    def _mpc_replan_check(self, task_id: str, result: dict) -> None:
276        """v1.3.0: After task result is handled, check if replan is needed."""
277        if not self._replan_callback or not self._replan_rules:
278            return
279        if not self._current_plan:
280            return
281
282        # Only trigger replan on failed tasks
283        status = result.get("status")
284        if status == "completed":
285            return
286
287        goal_id = getattr(self._current_plan, 'goal_id', 'unknown')
288        count = self._replan_count.get(goal_id, 0)
289        if count >= self._max_replans:
290            self._event_bus.publish(Event(
291                event_id=str(uuid.uuid4()),
292                event_type="goal.replan_limit_exceeded",
293                source="execution_engine",
294                timestamp=time.time(),
295                correlation_id=goal_id,
296                payload={
297                    "goal_id": goal_id,
298                    "replan_count": count,
299                    "max_replans": self._max_replans,
300                    "reason": f"Exceeded max replans ({self._max_replans})",
301                },
302            ))
303            return
304
305        artifact = result.get("artifact", {})
306        verdict = self._run_incremental_verify(task_id, artifact)
307
308        from .replan_rules import ReplanContext
309        ctx = ReplanContext(
310            task_id=task_id,
311            artifact=artifact,
312            verdict=verdict,
313            current_plan=self._current_plan,
314        )
315
316        diagnosis = self._run_diagnosis_if_applicable(task_id, result)
317
318        for rule in self._replan_rules:
319            if rule.should_replan(ctx):
320                new_plan = self._replan_callback(
321                    plan_id=self._current_plan.plan_id,
322                    trigger=rule.trigger_reason,
323                    context={
324                        "failed_task_id": task_id,
325                        "verdict": verdict,
326                        "diagnosis": diagnosis,
327                        "completed_tasks": self._get_completed_task_ids(),
328                        "failed_task": self._task_graph.get_task(task_id),
329                    },
330                )
331                if new_plan:
332                    self._current_plan = new_plan
333                    self._replan_count[goal_id] = count + 1
334                    self._event_bus.publish(Event(
335                        event_id=str(uuid.uuid4()),
336                        event_type="execution_plan.modified",
337                        source="execution_engine",
338                        timestamp=time.time(),
339                        correlation_id=new_plan.plan_id,
340                        payload={
341                            "plan_id": new_plan.plan_id,
342                            "trigger_reason": rule.trigger_reason,
343                            "failed_task_id": task_id,
344                            "replan_count": count + 1,
345                        },
346                    ))
347                break
348
349    # ── v1.3.0: MPC Support Methods ──
350
351    def set_replan_callback(self, callback: Callable) -> None:
352        """v1.3.0: Inject Runtime._on_replan as the replan callback."""
353        self._replan_callback = callback
354
355    def set_current_plan(self, plan) -> None:
356        """v1.3.0: Set the current ExecutionPlan reference."""
357        self._current_plan = plan
358
359    def set_diagnosis_engine(self, engine) -> None:
360        """v1.3.0: Inject DiagnosisEngine for failure analysis."""
361        self._diagnosis_engine = engine
362
363    def _run_incremental_verify(self, task_id: str, artifact: dict | None):
364        """v1.3.0: Lightweight per-task verification after completion."""
365        verifier = self._incremental_verifier
366        if not verifier:
367            from .verifier import Verdict
368            return Verdict(verdict="passed", score=1.0, verifier_id="incremental")
369
370        if artifact is None or (isinstance(artifact, dict) and not artifact):
371            from .verifier import Verdict
372            return Verdict(verdict="failed", score=0.0, verifier_id="incremental",
373                          summary="Empty artifact")
374
375        task = self._task_graph.get_task(task_id)
376        rules = ["non_empty"]
377        expected_schema = getattr(task, 'expected_output_schema', None) if task else None
378        if expected_schema:
379            rules.append("schema")
380
381        from .verifier import VerificationCriteria
382        criteria = VerificationCriteria(
383            expected_output_schema=expected_schema or {},
384            rules=rules,
385        )
386        return verifier.verify(artifact, criteria)
387
388    def _run_diagnosis_if_applicable(self, task_id: str, result: dict):
389        """v1.3.0: Run DiagnosisEngine if test_output is present in result."""
390        if not self._diagnosis_engine:
391            return None
392        test_output = (
393            result.get("test_output")
394            or (result.get("error") or {}).get("test_output", "")
395        )
396        if not test_output:
397            return None
398        return self._diagnosis_engine.diagnose(test_output)
399
400    def _get_completed_task_ids(self) -> list[str]:
401        """v1.3.0: Return task_ids of all COMPLETED tasks in current plan."""
402        if not self._current_plan:
403            return []
404        completed = []
405        for task in self._task_graph.list_tasks():
406            if task.status == TaskStatus.COMPLETED:
407                completed.append(task.task_id)
408        return completed
409
410    # ── Cancellation ──
411
412    def cancel_task(self, task_id: str) -> bool:
413        with self._lock:
414            inflight = self._in_flight.pop(task_id, None)
415            if inflight:
416                agent = self._agents.get(inflight.agent_id)
417                if agent and task_id in agent.current_tasks:
418                    agent.current_tasks.remove(task_id)
419                try:
420                    self._task_graph.transition(task_id, TaskStatus.CANCELLED)
421                except ValueError:
422                    pass
423                if self._agent_cancel:
424                    self._agent_cancel(inflight.agent_id, task_id)
425                return True
426        return False
427
428    # ── v0.8.0: Heartbeat ──
429
430    def submit_heartbeat(self, task_id: str, agent_id: str = "") -> bool:
431        """v0.8.0: Update heartbeat timestamp for an in-flight task."""
432        with self._lock:
433            ft = self._in_flight.get(task_id)
434            if not ft:
435                return False
436            if agent_id and ft.agent_id != agent_id:
437                return False
438            ft.heartbeat_at = time.time()
439            return True
440
441    def _check_heartbeat_timeouts(self) -> list[str]:
442        """v0.8.0: Check for heartbeat timeouts, transition tasks to FAILED.
443        Returns list of task_ids that timed out due to heartbeat.
444        """
445        now = time.time()
446        timed_out = []
447        with self._lock:
448            for tid, ft in list(self._in_flight.items()):
449                if ft.heartbeat_timeout_ms > 0:
450                    timeout_at = ft.heartbeat_at + (ft.heartbeat_timeout_ms / 1000)
451                    if now >= timeout_at:
452                        try:
453                            self._task_graph.transition(tid, TaskStatus.FAILED)
454                        except ValueError:
455                            pass
456                        self._in_flight.pop(tid, None)
457                        agent = self._agents.get(ft.agent_id)
458                        if agent and tid in agent.current_tasks:
459                            agent.current_tasks.remove(tid)
460                        timed_out.append(tid)
461        return timed_out
462
463    # ── Timeout Monitor ──
464
465    def start_monitor(self) -> None:
466        self._running = True
467        self._monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
468        self._monitor_thread.start()
469
470    def stop_monitor(self) -> None:
471        self._running = False
472
473    def _monitor_loop(self) -> None:
474        while self._running:
475            now = time.time()
476            with self._lock:
477                timed_out = [tid for tid, ft in self._in_flight.items() if now >= ft.timeout_at]
478            for tid in timed_out:
479                try:
480                    self._task_graph.transition(tid, TaskStatus.TIMED_OUT)
481                except ValueError:
482                    pass
483                self._in_flight.pop(tid, None)
484
485            # v0.8.0: Heartbeat timeout check
486            self._check_heartbeat_timeouts()
487
488            # Agent heartbeat check
489            with self._lock:
490                for agent in list(self._agents.values()):
491                    if agent.status == "heartbeating":
492                        elapsed = now - agent.last_heartbeat_at
493                        if elapsed > agent.heartbeat_interval_ms * 3 / 1000:
494                            agent.status = "disconnected"
495            time.sleep(1.0)
496
497    @property
498    def in_flight_count(self) -> int:
499        return len(self._in_flight)
500
501    @property
502    def in_flight_task_ids(self) -> list[str]:
503        return list(self._in_flight.keys())