feat(cloud): attribute planner token usage

apps/device-host-agent/host_agent/assignment.py

@@ -7,6 +7,7 @@ from typing import Any
 from cloud.internal_api.models import AssignmentModel
 from core.models import Task
 from host_agent.execution import ExecutionFactories
+from host_agent.planner_context import bind_planner_execution_context
 
 
 @dataclass(frozen=True)
@@ -26,19 +27,20 @@ class AssignmentExecutor:
         *,
         should_stop: Callable[[], bool] | None = None,
     ) -> AssignmentExecutionResult:
-        if should_stop is not None and should_stop():
+        with bind_planner_execution_context(assignment):
+            if should_stop is not None and should_stop():
+                return AssignmentExecutionResult(
+                    status="failed",
+                    failure_reason="execution interrupted",
+                )
+            if assignment.workflow_definition_id is not None:
+                return self._execute_workflow(assignment, should_stop=should_stop)
+            if assignment.goal is not None:
+                return self._execute_goal(assignment, should_stop=should_stop)
             return AssignmentExecutionResult(
                 status="failed",
-                failure_reason="execution interrupted",
+                failure_reason="assignment has neither goal nor workflow definition",
             )
-        if assignment.workflow_definition_id is not None:
-            return self._execute_workflow(assignment, should_stop=should_stop)
-        if assignment.goal is not None:
-            return self._execute_goal(assignment, should_stop=should_stop)
-        return AssignmentExecutionResult(
-            status="failed",
-            failure_reason="assignment has neither goal nor workflow definition",
-        )
 
     def _execute_goal(
         self,

apps/device-host-agent/host_agent/cloud_planner_client.py

@@ -27,6 +27,7 @@ import httpx
 
 from cloud.internal_api.models import PlannerDecisionError, PlannerDecisionResponse
 from host_agent.config import HostAgentConfig
+from host_agent.planner_context import current_planner_execution_context
 from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable, ToolCallUsage
 from runtime.tool_specs import ToolSpec
 
@@ -70,6 +71,15 @@ class CloudProxyToolCallingClient:
             ],
             "timeout_seconds": timeout,
         }
+        context = current_planner_execution_context()
+        if context is not None:
+            payload.update(
+                {
+                    "task_id": context.task_id,
+                    "attempt": context.attempt,
+                    "lease_id": context.lease_id,
+                }
+            )
         try:
             response = self._client.post(
                 f"/internal/v1/hosts/{self.config.host_id}/planner/decide",

apps/device-host-agent/host_agent/planner_context.py

@@ -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)