feat(cloud): attribute planner token usage
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user