File size: 6,738 Bytes
81e5fe7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0721bb4
81e5fe7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0721bb4
81e5fe7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0721bb4
81e5fe7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
"""TaskRunner — deterministic execution of a static `TaskList`. Zero LLM.

Executes tasks in dependency order, parallelizing each ready "wave" with
`asyncio.gather`. For each task it resolves `${t<id>}` placeholders from upstream
results, does an internal `validate_args`, invokes each tool via the `ToolInvoker`
seam, and records a `TaskResult`. On failure it **degrades and continues**: the
task is marked failed, its dependents are skipped, independent branches keep
running. There is no replanning and no mid-run LLM (INV-6).

`success_criteria` is *not* machine-evaluated here (it is free text); task status
is derived from tool execution outcomes and carried to the Assembler to report.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.4.
"""

from __future__ import annotations

import asyncio
from typing import Any

from src.middlewares.logging import get_logger

from ..planner.contracts import ToolOutput, ToolRegistry
from ..planner.schemas import PLACEHOLDER_RE, Task
from ..planner.schemas import TaskList as PlanTaskList
from .invoker import ToolInvoker
from .schemas import RunState, TaskResult, TaskStatus

logger = get_logger("task_runner")


class TaskRunner:
    """Runs a `TaskList` against a `ToolInvoker`, producing a `RunState`."""

    def __init__(self, invoker: ToolInvoker, registry: ToolRegistry) -> None:
        self._invoker = invoker
        self._registry = registry

    async def run(self, task_list: PlanTaskList, business_context_id: str) -> RunState:
        tasks_by_id: dict[str, Task] = {t.id: t for t in task_list.tasks}
        results: dict[str, TaskResult] = {}
        remaining: set[str] = set(tasks_by_id)

        while remaining:
            ready = [
                tid
                for tid in remaining
                if all(dep in results for dep in tasks_by_id[tid].depends_on)
            ]
            if not ready:
                # A dependency points outside the plan (or a cycle slipped past the
                # planner validator): nothing more can run. Fail the rest honestly.
                for tid in list(remaining):
                    results[tid] = TaskResult(
                        task_id=tid,
                        stage=tasks_by_id[tid].stage,
                        status="failure",
                        objective=tasks_by_id[tid].objective,
                        error="unresolved dependency; task could not run",
                    )
                    remaining.discard(tid)
                break

            # Skip any ready task whose dependency failed (degrade-and-continue).
            to_run: list[Task] = []
            for tid in ready:
                task = tasks_by_id[tid]
                failed = [d for d in task.depends_on if results[d].status == "failure"]
                if failed:
                    results[tid] = TaskResult(
                        task_id=tid,
                        stage=task.stage,
                        status="failure",
                        objective=task.objective,
                        error=f"skipped: upstream {failed} did not succeed",
                    )
                    remaining.discard(tid)
                else:
                    to_run.append(task)

            if not to_run:
                continue  # remaining dependents will be re-evaluated (and skipped)

            wave = await asyncio.gather(
                *(self._run_task(task, results) for task in to_run)
            )
            for tr in wave:
                results[tr.task_id] = tr
                remaining.discard(tr.task_id)

        return RunState(
            plan_id=task_list.plan_id,
            business_context_id=business_context_id,
            results=results,
            open_questions=list(task_list.open_questions),
        )

    async def _run_task(self, task: Task, results: dict[str, TaskResult]) -> TaskResult:
        outputs: list[ToolOutput] = []
        for call in task.tool_calls:
            resolved = self._resolve_args(call.args, results)
            arg_error = self._validate_args(call.tool, resolved)
            if arg_error is not None:
                outputs.append(ToolOutput(tool=call.tool, kind="error", error=arg_error))
                continue
            outputs.append(await self._safe_invoke(call.tool, resolved))

        status = _label(outputs)
        error: str | None = None
        if status == "failure":
            errs = [o.error for o in outputs if o.kind == "error" and o.error]
            error = errs[0] if errs else "all tool calls failed"
        return TaskResult(
            task_id=task.id,
            stage=task.stage,
            status=status,
            objective=task.objective,
            outputs=outputs,
            error=error,
        )

    def _resolve_args(
        self, args: dict[str, Any], results: dict[str, TaskResult]
    ) -> dict[str, Any]:
        return {k: self._resolve_value(v, results) for k, v in args.items()}

    @staticmethod
    def _resolve_value(value: Any, results: dict[str, TaskResult]) -> Any:
        # A data arg is exactly a "${t<id>}" placeholder (Pattern A); resolve it to
        # the referenced task's representative output (its last ToolOutput).
        # Materializing that envelope into a DataFrame is the invoker's job.
        if isinstance(value, str):
            match = PLACEHOLDER_RE.fullmatch(value.strip())
            if match:
                upstream = results.get(match.group(1))
                if upstream is None or not upstream.outputs:
                    return None
                return upstream.outputs[-1]
        return value

    def _validate_args(self, tool: str, resolved: dict[str, Any]) -> str | None:
        spec = self._registry.get(tool)
        if spec is None:
            return f"tool {tool!r} not in registry"
        required = spec.input_schema.get("required", [])
        missing = [r for r in required if resolved.get(r) is None]
        if missing:
            return f"missing required arg(s): {sorted(missing)}"
        return None

    async def _safe_invoke(self, tool: str, args: dict[str, Any]) -> ToolOutput:
        try:
            return await self._invoker.invoke(tool, args)
        except Exception as exc:  # noqa: BLE001 — backstop; the invoker is never-throw (§8.4)
            logger.warning("tool invoker raised", tool=tool, error=str(exc))
            return ToolOutput(tool=tool, kind="error", error=f"invoker raised: {exc}")


def _label(outputs: list[ToolOutput]) -> TaskStatus:
    if not outputs:
        return "failure"
    errors = sum(1 for o in outputs if o.kind == "error")
    if errors == 0:
        return "success"
    if errors == len(outputs):
        return "failure"
    return "partial"