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)

apps/device-host-agent/tests/test_cloud_planner_client.py

@@ -7,6 +7,7 @@ import pytest
 
 from host_agent.cloud_planner_client import CloudProxyToolCallingClient
 from host_agent.config import HostAgentConfig
+from host_agent.planner_context import PlannerExecutionContext, _context
 from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable
 from runtime.tool_specs import ToolSpec
 
@@ -78,6 +79,34 @@ def test_decide_base64_encodes_screenshot() -> None:
     assert body["screenshot_base64"] == "aGVsbG8="
 
 
+def test_decide_includes_bound_assignment_context() -> None:
+    seen_requests: list[httpx.Request] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen_requests.append(request)
+        return httpx.Response(200, json={"tool_name": "tap", "arguments": {}})
+
+    client = _client(handler)
+    token = _context.set(
+        PlannerExecutionContext(task_id="task-a", attempt=2, lease_id="lease-a")
+    )
+    try:
+        client.decide(
+            system_prompt="sp",
+            user_prompt="up",
+            screenshot=None,
+            tools=_TOOLS,
+            timeout=10.0,
+        )
+    finally:
+        _context.reset(token)
+
+    body = json.loads(seen_requests[0].content)
+    assert body["task_id"] == "task-a"
+    assert body["attempt"] == 2
+    assert body["lease_id"] == "lease-a"
+
+
 def test_decide_raises_tool_call_unavailable_on_network_error() -> None:
     def handler(request: httpx.Request) -> httpx.Response:
         raise httpx.ConnectError("connection refused", request=request)