feat(cloud): attribute planner token usage

2026-07-13 23:10:21 +08:00
parent b4803f90e6
commit 40efa53411
16 changed files with 259 additions and 17 deletions
@@ -0,0 +1,47 @@
+"""Host-local execution metadata for Cloud planner accounting."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from cloud.internal_api.models import AssignmentModel
+
+
+@dataclass(frozen=True)
+class PlannerExecutionContext:
+    task_id: str
+    attempt: int
+    lease_id: str
+
+
+_context: ContextVar[PlannerExecutionContext | None] = ContextVar(
+    "host_agent_planner_execution_context",
+    default=None,
+)
+
+
+def current_planner_execution_context() -> PlannerExecutionContext | None:
+    return _context.get()
+
+
+@contextmanager
+def bind_planner_execution_context(
+    assignment: "AssignmentModel",
+) -> "Iterator[None]":
+    token = _context.set(
+        PlannerExecutionContext(
+            task_id=assignment.task_id,
+            attempt=assignment.attempt,
+            lease_id=assignment.lease_id,
+        )
+    )
+    try:
+        yield
+    finally:
+        _context.reset(token)