88 lines
3.7 KiB
Python
88 lines
3.7 KiB
Python
"""Delegation guardrails."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .errors import ConflictError, ForbiddenError, NotFoundError
|
|
from .models import Budget, TaskPlan, TaskRecord
|
|
from .store import TaskPlaneStore
|
|
|
|
|
|
class DelegationEngine:
|
|
def __init__(self, store: TaskPlaneStore):
|
|
self.store = store
|
|
|
|
def can_delegate(self, parent_id: str, child_plan: TaskPlan) -> bool:
|
|
parent = self.store.get_task(parent_id)
|
|
if parent.status == "COMPLETED":
|
|
raise ConflictError("cannot delegate completed task")
|
|
if parent.task_plan.max_delegation_depth <= 0:
|
|
raise ConflictError("max delegation depth not configured")
|
|
|
|
depth = self._depth(parent_id)
|
|
if depth >= parent.task_plan.max_delegation_depth:
|
|
raise ConflictError("max delegation depth reached")
|
|
|
|
if child_plan.parent_task_id not in {None, parent_id}:
|
|
raise ConflictError("delegation parent mismatch")
|
|
current = parent_id
|
|
while current:
|
|
if child_plan.task_id and current == child_plan.task_id:
|
|
raise ConflictError("delegation cycle")
|
|
current = self.store.parent_of(current)
|
|
return True
|
|
|
|
def delegate(self, parent_id: str, child_plan: TaskPlan, *, caller: str) -> TaskRecord:
|
|
self.can_delegate(parent_id, child_plan)
|
|
parent = self.store.get_task(parent_id)
|
|
if child_plan.required_capabilities and not set(child_plan.required_capabilities).issubset(
|
|
set(parent.task_plan.required_capabilities or [])
|
|
):
|
|
raise ForbiddenError("child capabilities exceed parent constraints")
|
|
risk_rank = {"low": 1, "medium": 2, "high": 3}
|
|
if risk_rank.get(child_plan.risk, 99) > risk_rank.get(parent.task_plan.risk, -1):
|
|
raise ForbiddenError("child risk exceeds parent constraints")
|
|
child_plan.parent_task_id = parent_id
|
|
child_plan.correlation_id = parent.task_plan.correlation_id or parent.task_id
|
|
child_plan.preferred_workers = parent.task_plan.preferred_workers
|
|
child_plan.max_delegation_depth = min(
|
|
child_plan.max_delegation_depth,
|
|
parent.task_plan.max_delegation_depth,
|
|
)
|
|
child_plan.approval_policy = parent.task_plan.approval_policy
|
|
|
|
parent_budget = self.store.get_task_budget(parent_id)
|
|
required = Budget()
|
|
for name in ("time_seconds", "attempts", "tokens", "cost"):
|
|
available = getattr(parent_budget, name)
|
|
requested = getattr(child_plan.budgets, name)
|
|
if name == "attempts" and requested is None:
|
|
requested = 1
|
|
if available is not None:
|
|
requested = 0 if requested is None else requested
|
|
if requested > available:
|
|
raise ConflictError(f"budget exceeded for {name}")
|
|
setattr(child_plan.budgets, name, requested)
|
|
setattr(required, name, requested)
|
|
|
|
# A retried identical delegation is a read, not a second budget charge.
|
|
if child_plan.task_id:
|
|
try:
|
|
existing = self.store.get_task(child_plan.task_id)
|
|
except NotFoundError:
|
|
existing = None
|
|
if existing is not None:
|
|
if (
|
|
existing.task_plan.parent_task_id == parent_id
|
|
and existing.task_plan.operation_digest()
|
|
== child_plan.operation_digest()
|
|
):
|
|
return existing
|
|
raise ConflictError("delegation task_id conflict")
|
|
|
|
return self.store.delegate_task_with_budget_and_events(
|
|
parent_id, child_plan, caller=caller
|
|
)
|
|
|
|
def _depth(self, task_id: str) -> int:
|
|
return self.store._compute_depth(task_id)
|