From e36549f01fe8c002f8d42f87b2c9ccae62a7c928 Mon Sep 17 00:00:00 2001 From: Renn F Date: Thu, 18 Jun 2026 11:54:58 +0200 Subject: [PATCH] feat(grok): capture one-shot Grok usage/cost from the opencode store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A GROK agent runs opencode, not Claude Code: it has no SDK /usage/status server and writes no Claude transcript, so _resolve_final_token_usage found nothing and every Grok agent finalized at 0 tokens / $0 — the opencode_usage reader existed but had no caller. - Mount a per-agent opencode data dir ($DATA/opencode/ → /home/agent/.local/share/opencode) so opencode.db is captured, and mount the same host dir into the orchestrator (/data/opencode) in all three compose files so the finalizer can read it back — the opencode analogue of the mounted Claude transcript. - _resolve_final_token_usage branches on provider_type: GROK reads opencode.db via opencode_usage (reasoning folded into output, billed at the output rate) and skips the SDK/transcript path. A 0-token read logs a WARNING so a silent mount failure isn't mistaken for a real zero-cost run. - ROBOCO_OPENCODE_DATA_DIR overrides the in-orchestrator path for local runs. --- docker-compose.registry.yml | 2 + docker-compose.yaml | 3 + docker-compose.yml | 3 + roboco/llm/providers/grok.py | 20 ++++ roboco/runtime/orchestrator.py | 63 +++++++++++- tests/unit/llm/test_providers.py | 4 + .../unit/runtime/test_grok_usage_finalize.py | 99 +++++++++++++++++++ 7 files changed, 190 insertions(+), 4 deletions(-) create mode 100644 tests/unit/runtime/test_grok_usage_finalize.py diff --git a/docker-compose.registry.yml b/docker-compose.registry.yml index f2088150..16f938d6 100644 --- a/docker-compose.registry.yml +++ b/docker-compose.registry.yml @@ -231,6 +231,8 @@ services: - ${ROBOCO_DATA_DIR:-./data}/prompts-generated:/app/prompts-generated - ${ROBOCO_DATA_DIR:-./data}/agent-settings:/app/agent-settings - ${ROBOCO_DATA_DIR:-./data}/workspaces:/data/workspaces + # Per-agent opencode stores (GROK usage/cost capture). + - ${ROBOCO_DATA_DIR:-./data}/opencode:/data/opencode - ${ROBOCO_DATA_DIR:-./data}/logs:/data/logs - ${ROBOCO_DATA_DIR:-./data}/briefings:/app/briefings - ${ROBOCO_DATA_DIR:-./data}/manifests:/app/manifests diff --git a/docker-compose.yaml b/docker-compose.yaml index 33bda725..18d3fbfa 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -330,6 +330,9 @@ services: - ${ROBOCO_DATA_DIR:-./data}/agent-settings:/app/agent-settings # Agent workspaces (git clones) - persisted across restarts - ${ROBOCO_DATA_DIR:-./data}/workspaces:/data/workspaces + # Per-agent opencode stores (GROK usage/cost capture): each Grok agent + # writes opencode.db under /; the finalizer reads it back here. + - ${ROBOCO_DATA_DIR:-./data}/opencode:/data/opencode # Persistent logs — survive `docker compose down/up`. Orchestrator and # each spawned agent write structured logs here so we can audit past # runs instead of relying on ephemeral `docker logs`. diff --git a/docker-compose.yml b/docker-compose.yml index 33bda725..18d3fbfa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -330,6 +330,9 @@ services: - ${ROBOCO_DATA_DIR:-./data}/agent-settings:/app/agent-settings # Agent workspaces (git clones) - persisted across restarts - ${ROBOCO_DATA_DIR:-./data}/workspaces:/data/workspaces + # Per-agent opencode stores (GROK usage/cost capture): each Grok agent + # writes opencode.db under /; the finalizer reads it back here. + - ${ROBOCO_DATA_DIR:-./data}/opencode:/data/opencode # Persistent logs — survive `docker compose down/up`. Orchestrator and # each spawned agent write structured logs here so we can audit past # runs instead of relying on ephemeral `docker logs`. diff --git a/roboco/llm/providers/grok.py b/roboco/llm/providers/grok.py index dad38faa..fd0725a4 100644 --- a/roboco/llm/providers/grok.py +++ b/roboco/llm/providers/grok.py @@ -63,6 +63,10 @@ _DEFAULT_XAI_BASE_URL = "https://api.x.ai/v1" # In-container paths mounted by the orchestrator's `_build_mount_args`. _MCP_CONFIG_IN_CONTAINER = "/app/mcp-config.json" _SYSTEM_PROMPT_IN_CONTAINER = "/app/system-prompt.md" +# opencode's data dir inside the agent (HOME=/home/agent); opencode.db lands +# here. Mounted to a per-agent host dir so the orchestrator can read usage back +# (mirror of roboco.llm.providers.opencode_usage.DEFAULT_DB_PATH's parent). +_OPENCODE_DATA_DIR_IN_CONTAINER = "/home/agent/.local/share/opencode" # Reasoning effort by role. grok-build-0.1 reasons heavily by default, and # reasoning bills at the output rate — it dominates cost (a live "say ok" call @@ -168,6 +172,7 @@ class GrokProvider(AgentProvider): cmd = self._host._build_mount_args(container_name, mount_config, hosts) self._host._append_agent_auth_env(cmd, config) self._host._append_git_context_env(cmd, config) + self._append_opencode_data_mount(cmd, hosts) self._append_grok_env(cmd, config, initial_prompt) cmd.append(self._image) @@ -187,6 +192,21 @@ class GrokProvider(AgentProvider): extra={"container_id": stdout.decode().strip(), "model": config.model}, ) + def _append_opencode_data_mount( + self, cmd: list[str], hosts: dict[str, str | None] + ) -> None: + """Mount the per-agent opencode data dir so the orchestrator can read it. + + opencode persists token usage to ``opencode.db`` under its data dir + (``$HOME/.local/share/opencode``). Binding a per-agent host dir there + lets the finalizer read the store back over the shared data volume — + the opencode analogue of the mounted Claude transcript. Without this a + Grok agent finalizes at 0 tokens / $0. + """ + opencode_host = hosts.get("opencode") + if opencode_host: + cmd.extend(["-v", f"{opencode_host}:{_OPENCODE_DATA_DIR_IN_CONTAINER}"]) + def _append_grok_env( self, cmd: list[str], config: AgentConfig, initial_prompt: str | None ) -> None: diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 1b76f985..4e924645 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -163,6 +163,11 @@ CLAUDE_AUTH_HOST_PATH = os.environ.get( ) PROJECT_HOST_PATH = os.environ.get("ROBOCO_HOST_PROJECT_DIR", "") DATA_HOST_PATH = os.environ.get("ROBOCO_HOST_DATA_DIR", "") +# In-orchestrator path where each GROK agent's opencode store is visible. The +# agent writes /opencode/; the compose file mounts the +# same host dir here so the finalizer can read opencode.db back (mirrors how the +# Claude transcript is read from the mounted ~/.claude). Override for local runs. +OPENCODE_DATA_DIR = os.environ.get("ROBOCO_OPENCODE_DATA_DIR", "/data/opencode") # ============================================================================= @@ -1692,6 +1697,9 @@ class AgentOrchestrator: "workspaces": f"{DATA_HOST_PATH}/workspaces", "claude": CLAUDE_AUTH_HOST_PATH, "mcp_config": f"{DATA_HOST_PATH}/mcp-configs/{mcp_name}", + # Per-agent opencode store (GROK only); the orchestrator reads it + # back at finalize via the shared data volume (see OPENCODE_DATA_DIR). + "opencode": f"{DATA_HOST_PATH}/opencode/{config.agent_id}", "prompt": ( f"{DATA_HOST_PATH}/prompts-generated/{config.agent_id}-prompt.md" ), @@ -1711,6 +1719,9 @@ class AgentOrchestrator: "workspaces": str(Path(settings.workspaces_root)), "claude": CLAUDE_AUTH_HOST_PATH, "mcp_config": str(config.mcp_config_path), + "opencode": str( + Path(tempfile.gettempdir()) / "roboco-opencode" / config.agent_id + ), "prompt": str( Path(tempfile.gettempdir()) / "roboco-prompts" @@ -3671,17 +3682,61 @@ class AgentOrchestrator: except OSError: return (0, 0, 0, 0) + def _opencode_db_path(self, agent_id: str) -> str: + """In-orchestrator path to a GROK agent's opencode SQLite store. + + The agent writes opencode.db under the shared data volume; the compose + file mounts that host dir at ``OPENCODE_DATA_DIR`` here, so finalize can + read it back — the opencode analogue of the mounted Claude transcript. + """ + return str(Path(OPENCODE_DATA_DIR) / agent_id / "opencode.db") + + def _grok_usage_from_opencode(self, agent_id: str) -> tuple[int, int, int, int]: + """Sum a GROK agent's token usage from its opencode SQLite store. + + A GROK agent runs opencode — no SDK ``/usage/status`` server and no + Claude transcript — so its usage lands in opencode.db. Reasoning is + folded into output (it bills at the output rate, matching + ``calculate_cost``). A WARNING is logged on a 0-token read because a + silent mount/uid failure is otherwise indistinguishable from a genuine + zero-cost run. Returns ``(input, output, cache_read, cache_write)``. + """ + from roboco.llm.providers.opencode_usage import read_session_usage + + db_path = self._opencode_db_path(agent_id) + usage = read_session_usage(db_path) + if usage is None: + logger.warning( + "GROK agent finalized with no readable opencode usage " + "(0 tokens / $0) — check the opencode db mount", + agent_id=agent_id, + db_path=db_path, + ) + return (0, 0, 0, 0) + return ( + usage.tokens_input, + usage.tokens_output + usage.tokens_reasoning, + usage.tokens_cache_read, + usage.tokens_cache_write, + ) + async def _resolve_final_token_usage( self, agent_id: str ) -> tuple[int, int, int, int]: """Resolve final token counts for a stopping agent. - Tries the live SDK ``/usage/status`` first; if that misses — the SDK's - in-memory counts race container teardown for short-lived agents — it - falls back to the agent's Claude Code transcript, which is durable and - mounted into this container. Returns + For a GROK agent, reads the opencode SQLite store (no SDK server / Claude + transcript exists). Otherwise tries the live SDK ``/usage/status`` first; + if that misses — the SDK's in-memory counts race container teardown for + short-lived agents — it falls back to the agent's Claude Code transcript, + which is durable and mounted into this container. Returns ``(input, output, cache_read, cache_write)``. """ + from roboco.models.base import ModelProvider + + if self.get_provider_for_agent(agent_id) == ModelProvider.GROK.value: + return self._grok_usage_from_opencode(agent_id) + tokens = (0, 0, 0, 0) sdk_url = f"http://roboco-agent-{agent_id}:{SDK_PORT}/usage/status" try: diff --git a/tests/unit/llm/test_providers.py b/tests/unit/llm/test_providers.py index 22f02414..a9e34595 100644 --- a/tests/unit/llm/test_providers.py +++ b/tests/unit/llm/test_providers.py @@ -76,6 +76,7 @@ class _FakeHost: if config.mcp_config_path else None, "settings": str(agent_settings_path) if agent_settings_path else None, + "opencode": f"/host/opencode/{config.agent_id}", } def _build_mount_args( @@ -200,6 +201,9 @@ async def test_grok_spawn_wires_gateway_and_image_last() -> None: # Tool restriction lives in the rendered opencode.json (opencode `tools`), # not a spawn env var — no ROBOCO_AGENT_TOOLS is injected. assert not any(c.startswith("ROBOCO_AGENT_TOOLS=") for c in cmd) + # The opencode store is mounted so the orchestrator can read usage/cost + # back at finalize. + assert "/host/opencode/be-dev-1:/home/agent/.local/share/opencode" in cmd # Identity wiring from the shared host helpers is present. assert "ROBOCO_AGENT_TOKEN=hmac-be-dev-1" in cmd # The image is the final docker-run argument. diff --git a/tests/unit/runtime/test_grok_usage_finalize.py b/tests/unit/runtime/test_grok_usage_finalize.py new file mode 100644 index 00000000..368ee001 --- /dev/null +++ b/tests/unit/runtime/test_grok_usage_finalize.py @@ -0,0 +1,99 @@ +"""GROK agents capture token usage/cost from their opencode SQLite store. + +A Grok agent runs opencode — no SDK /usage/status server and no Claude +transcript — so finalize must read opencode.db (mounted into the orchestrator) +instead. Reasoning folds into output (it bills at the output rate). +""" + +from __future__ import annotations + +import sqlite3 +from typing import TYPE_CHECKING + +import pytest +from roboco.models.runtime import AgentInstance +from roboco.runtime.orchestrator import AgentOrchestrator + +if TYPE_CHECKING: + from pathlib import Path + + +def _make_db(path: Path, cols: dict[str, float]) -> None: + con = sqlite3.connect(path) + con.execute( + "CREATE TABLE session (id TEXT, tokens_input INT, tokens_output INT, " + "tokens_cache_read INT, tokens_cache_write INT, tokens_reasoning INT, " + "cost REAL)" + ) + con.execute( + "INSERT INTO session (id, tokens_input, tokens_output, tokens_cache_read, " + "tokens_cache_write, tokens_reasoning, cost) VALUES (?,?,?,?,?,?,?)", + ( + "s1", + cols["tokens_input"], + cols["tokens_output"], + cols["tokens_cache_read"], + cols["tokens_cache_write"], + cols["tokens_reasoning"], + cols["cost"], + ), + ) + con.commit() + con.close() + + +def test_grok_usage_folds_reasoning_into_output( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + db = tmp_path / "opencode.db" + _make_db( + db, + { + "tokens_input": 100, + "tokens_output": 50, + "tokens_reasoning": 30, + "tokens_cache_read": 10, + "tokens_cache_write": 5, + "cost": 0.02, + }, + ) + orch = AgentOrchestrator.__new__(AgentOrchestrator) + monkeypatch.setattr(orch, "_opencode_db_path", lambda _aid: str(db)) + + # reasoning (30) folded into output (50) → 80; bills at the output rate. + assert orch._grok_usage_from_opencode("be-dev-1") == (100, 80, 10, 5) + + +def test_grok_usage_zero_when_store_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + orch = AgentOrchestrator.__new__(AgentOrchestrator) + monkeypatch.setattr( + orch, "_opencode_db_path", lambda _aid: str(tmp_path / "absent.db") + ) + assert orch._grok_usage_from_opencode("be-dev-1") == (0, 0, 0, 0) + + +@pytest.mark.asyncio +async def test_resolve_final_usage_routes_grok_to_opencode( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + db = tmp_path / "opencode.db" + _make_db( + db, + { + "tokens_input": 7, + "tokens_output": 3, + "tokens_reasoning": 2, + "tokens_cache_read": 0, + "tokens_cache_write": 0, + "cost": 0.01, + }, + ) + orch = AgentOrchestrator.__new__(AgentOrchestrator) + monkeypatch.setattr(orch, "_opencode_db_path", lambda _aid: str(db)) + cfg = type("C", (), {"provider_type": "grok"})() + orch._instances = {"be-dev-1": AgentInstance(agent_id="be-dev-1", config=cfg)} + + # No SDK fetch / transcript read for GROK — usage comes from opencode.db. + assert await orch._resolve_final_token_usage("be-dev-1") == (7, 5, 0, 0)