mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
feat(grok-cli): read captured usage at finalize; keep interactive serve working
The provider mounts the per-agent data dir and points the entrypoint's usage file at it; the orchestrator's grok finalize reads that usage.json first (the grok-CLI total, priced at the output rate) and falls back to opencode.db for the still-opencode interactive intake/secretary path. Re-add _reasoning_effort_for to grok.py as a clearly-temporary shim for that interactive path (it needs opencode's "minimal" variant, distinct from the CLI's --effort) until it is converted too.
This commit is contained in:
@@ -31,6 +31,7 @@ import os
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Protocol
|
from typing import TYPE_CHECKING, Protocol
|
||||||
|
|
||||||
|
from roboco.agents_config import get_agent_role
|
||||||
from roboco.llm.providers._docker import container_running, stop_container
|
from roboco.llm.providers._docker import container_running, stop_container
|
||||||
from roboco.llm.providers.base import AgentProvider, ProviderError, SpawnResult
|
from roboco.llm.providers.base import AgentProvider, ProviderError, SpawnResult
|
||||||
|
|
||||||
@@ -56,12 +57,55 @@ GROK_AUTH_HOST_PATH = os.environ.get(
|
|||||||
# In-container paths.
|
# In-container paths.
|
||||||
_MCP_CONFIG_IN_CONTAINER = "/app/mcp-config.json"
|
_MCP_CONFIG_IN_CONTAINER = "/app/mcp-config.json"
|
||||||
_GROK_AUTH_IN_CONTAINER = "/home/agent/.grok/auth.json"
|
_GROK_AUTH_IN_CONTAINER = "/home/agent/.grok/auth.json"
|
||||||
|
# Per-agent data dir (the host side is reused from the shared assembly): the
|
||||||
|
# entrypoint writes the captured token usage here so the orchestrator reads it
|
||||||
|
# back at finalize, the grok analogue of the mounted Claude transcript.
|
||||||
|
_GROK_USAGE_DIR_IN_CONTAINER = "/home/agent/.grok-usage"
|
||||||
|
_GROK_USAGE_FILE_IN_CONTAINER = f"{_GROK_USAGE_DIR_IN_CONTAINER}/usage.json"
|
||||||
|
|
||||||
|
|
||||||
def _container_name(agent_id: str) -> str:
|
def _container_name(agent_id: str) -> str:
|
||||||
return f"roboco-agent-{agent_id}"
|
return f"roboco-agent-{agent_id}"
|
||||||
|
|
||||||
|
|
||||||
|
# --- interactive (opencode-serve) shim — PENDING conversion to the grok CLI ---
|
||||||
|
# The intake / secretary roles still run on the opencode-serve interactive path,
|
||||||
|
# which needs the opencode ``--variant`` reasoning effort ("minimal") — distinct
|
||||||
|
# from the one-shot CLI's ``--effort`` ("low") computed in grok_cli_config. Kept
|
||||||
|
# here until the interactive path is moved onto the grok CLI too; the orchestrator
|
||||||
|
# interactive-spawn methods import it.
|
||||||
|
_MINIMAL_REASONING_ROLES = frozenset(
|
||||||
|
{
|
||||||
|
"cell_pm",
|
||||||
|
"main_pm",
|
||||||
|
"documenter",
|
||||||
|
"product_owner",
|
||||||
|
"head_marketing",
|
||||||
|
"auditor",
|
||||||
|
"prompter",
|
||||||
|
"secretary",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
_FULL_REASONING_OVERRIDES = frozenset({"default", "full", "none", ""})
|
||||||
|
|
||||||
|
|
||||||
|
def _reasoning_effort_for(agent_id: str) -> str | None:
|
||||||
|
"""opencode ``--variant`` for the interactive serve path (pending conversion).
|
||||||
|
|
||||||
|
Returns ``None`` (opencode default / full reasoning) for code-quality roles;
|
||||||
|
``"minimal"`` for coordination / docs / board roles. A global
|
||||||
|
``ROBOCO_GROK_REASONING_EFFORT`` override wins.
|
||||||
|
"""
|
||||||
|
override = os.environ.get("ROBOCO_GROK_REASONING_EFFORT", "").strip()
|
||||||
|
if override:
|
||||||
|
return None if override.lower() in _FULL_REASONING_OVERRIDES else override
|
||||||
|
return (
|
||||||
|
"minimal"
|
||||||
|
if (get_agent_role(agent_id) or "") in _MINIMAL_REASONING_ROLES
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class _GrokHost(Protocol):
|
class _GrokHost(Protocol):
|
||||||
"""The orchestrator surface GrokCliProvider reuses for container assembly.
|
"""The orchestrator surface GrokCliProvider reuses for container assembly.
|
||||||
|
|
||||||
@@ -71,6 +115,8 @@ class _GrokHost(Protocol):
|
|||||||
|
|
||||||
async def _remove_container(self, container_name: str) -> None: ...
|
async def _remove_container(self, container_name: str) -> None: ...
|
||||||
|
|
||||||
|
def _ensure_opencode_data_dir(self, agent_id: str) -> None: ...
|
||||||
|
|
||||||
def _resolve_host_paths(
|
def _resolve_host_paths(
|
||||||
self, config: AgentConfig, agent_settings_path: Path | None
|
self, config: AgentConfig, agent_settings_path: Path | None
|
||||||
) -> dict[str, str | None]: ...
|
) -> dict[str, str | None]: ...
|
||||||
@@ -108,6 +154,9 @@ class GrokCliProvider(AgentProvider):
|
|||||||
|
|
||||||
container_name = _container_name(config.agent_id)
|
container_name = _container_name(config.agent_id)
|
||||||
await self._host._remove_container(container_name)
|
await self._host._remove_container(container_name)
|
||||||
|
# Pre-create the per-agent data dir (world-writable) before the bind
|
||||||
|
# mount so the non-root agent can write the usage file (else EACCES).
|
||||||
|
self._host._ensure_opencode_data_dir(config.agent_id)
|
||||||
|
|
||||||
# Reuse the orchestrator's mount/auth/git assembly so the agent gets the
|
# Reuse the orchestrator's mount/auth/git assembly so the agent gets the
|
||||||
# full MCP gateway + identity wiring. Blank the provider routing fields
|
# full MCP gateway + identity wiring. Blank the provider routing fields
|
||||||
@@ -122,6 +171,7 @@ class GrokCliProvider(AgentProvider):
|
|||||||
self._host._append_agent_auth_env(cmd, config)
|
self._host._append_agent_auth_env(cmd, config)
|
||||||
self._host._append_git_context_env(cmd, config)
|
self._host._append_git_context_env(cmd, config)
|
||||||
self._append_grok_auth_mount(cmd)
|
self._append_grok_auth_mount(cmd)
|
||||||
|
self._append_usage_mount(cmd, hosts)
|
||||||
self._append_grok_env(cmd, config, initial_prompt)
|
self._append_grok_env(cmd, config, initial_prompt)
|
||||||
cmd.append(self._image)
|
cmd.append(self._image)
|
||||||
|
|
||||||
@@ -154,6 +204,18 @@ class GrokCliProvider(AgentProvider):
|
|||||||
if auth_json.exists():
|
if auth_json.exists():
|
||||||
cmd.extend(["-v", f"{auth_json}:{_GROK_AUTH_IN_CONTAINER}:ro"])
|
cmd.extend(["-v", f"{auth_json}:{_GROK_AUTH_IN_CONTAINER}:ro"])
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _append_usage_mount(cmd: list[str], hosts: dict[str, str | None]) -> None:
|
||||||
|
"""Mount the per-agent data dir so the orchestrator reads usage back.
|
||||||
|
|
||||||
|
Reuses the shared per-agent host dir (``hosts["opencode"]``); the
|
||||||
|
entrypoint writes ``usage.json`` here after the run and the orchestrator
|
||||||
|
reads it at finalize. Without it a Grok agent finalizes at 0 tokens / $0.
|
||||||
|
"""
|
||||||
|
data_host = hosts.get("opencode")
|
||||||
|
if data_host:
|
||||||
|
cmd.extend(["-v", f"{data_host}:{_GROK_USAGE_DIR_IN_CONTAINER}"])
|
||||||
|
|
||||||
def _append_grok_env(
|
def _append_grok_env(
|
||||||
self, cmd: list[str], config: AgentConfig, initial_prompt: str | None
|
self, cmd: list[str], config: AgentConfig, initial_prompt: str | None
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -173,6 +235,8 @@ class GrokCliProvider(AgentProvider):
|
|||||||
f"ROBOCO_MCP_CONFIG={_MCP_CONFIG_IN_CONTAINER}",
|
f"ROBOCO_MCP_CONFIG={_MCP_CONFIG_IN_CONTAINER}",
|
||||||
"-e",
|
"-e",
|
||||||
f"ROBOCO_INITIAL_PROMPT={initial_prompt or ''}",
|
f"ROBOCO_INITIAL_PROMPT={initial_prompt or ''}",
|
||||||
|
"-e",
|
||||||
|
f"ROBOCO_GROK_USAGE_FILE={_GROK_USAGE_FILE_IN_CONTAINER}",
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
if config.claude_session_id:
|
if config.claude_session_id:
|
||||||
|
|||||||
@@ -3933,23 +3933,28 @@ class AgentOrchestrator:
|
|||||||
return str(Path(OPENCODE_DATA_DIR) / agent_id / "opencode.db")
|
return str(Path(OPENCODE_DATA_DIR) / agent_id / "opencode.db")
|
||||||
|
|
||||||
def _grok_usage_from_opencode(self, agent_id: str) -> tuple[int, int, int, int]:
|
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.
|
"""Sum a GROK agent's token usage from its per-agent data dir.
|
||||||
|
|
||||||
A GROK agent runs opencode — no SDK ``/usage/status`` server and no
|
The grok-CLI one-shot path writes a ``usage.json`` (``total_tokens``)
|
||||||
Claude transcript — so its usage lands in opencode.db. Reasoning is
|
there post-run — read that first. The interactive opencode-serve path
|
||||||
folded into output (it bills at the output rate, matching
|
still lands its usage in ``opencode.db``, so fall back to it. The grok
|
||||||
|
total folds into output (it bills at the output rate, matching
|
||||||
``calculate_cost``). A WARNING is logged on a 0-token read because a
|
``calculate_cost``). A WARNING is logged on a 0-token read because a
|
||||||
silent mount/uid failure is otherwise indistinguishable from a genuine
|
silent mount/uid failure is otherwise indistinguishable from a genuine
|
||||||
zero-cost run. Returns ``(input, output, cache_read, cache_write)``.
|
zero-cost run. Returns ``(input, output, cache_read, cache_write)``.
|
||||||
"""
|
"""
|
||||||
|
cli_tokens = self._grok_cli_total_tokens(agent_id)
|
||||||
|
if cli_tokens is not None:
|
||||||
|
return (0, cli_tokens, 0, 0)
|
||||||
|
|
||||||
from roboco.llm.providers.opencode_usage import read_session_usage
|
from roboco.llm.providers.opencode_usage import read_session_usage
|
||||||
|
|
||||||
db_path = self._opencode_db_path(agent_id)
|
db_path = self._opencode_db_path(agent_id)
|
||||||
usage = read_session_usage(db_path)
|
usage = read_session_usage(db_path)
|
||||||
if usage is None:
|
if usage is None:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"GROK agent finalized with no readable opencode usage "
|
"GROK agent finalized with no readable usage "
|
||||||
"(0 tokens / $0) — check the opencode db mount",
|
"(0 tokens / $0) — check the data dir mount",
|
||||||
agent_id=agent_id,
|
agent_id=agent_id,
|
||||||
db_path=db_path,
|
db_path=db_path,
|
||||||
)
|
)
|
||||||
@@ -3961,6 +3966,20 @@ class AgentOrchestrator:
|
|||||||
usage.tokens_cache_write,
|
usage.tokens_cache_write,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _grok_cli_total_tokens(agent_id: str) -> int | None:
|
||||||
|
"""Total tokens from a grok-CLI ``usage.json``, or None if absent.
|
||||||
|
|
||||||
|
The grok-cli entrypoint writes ``{model, total_tokens, cost_usd}`` to the
|
||||||
|
per-agent data dir; the orchestrator sees it at ``OPENCODE_DATA_DIR``.
|
||||||
|
"""
|
||||||
|
usage_json = Path(OPENCODE_DATA_DIR) / agent_id / "usage.json"
|
||||||
|
try:
|
||||||
|
data = json.loads(usage_json.read_text(encoding="utf-8"))
|
||||||
|
return int(data.get("total_tokens", 0))
|
||||||
|
except (OSError, json.JSONDecodeError, ValueError, TypeError):
|
||||||
|
return None
|
||||||
|
|
||||||
async def _enforce_grok_cost_budget(self) -> None:
|
async def _enforce_grok_cost_budget(self) -> None:
|
||||||
"""Kill a live GROK container whose cumulative opencode cost exceeds the cap.
|
"""Kill a live GROK container whose cumulative opencode cost exceeds the cap.
|
||||||
|
|
||||||
|
|||||||
@@ -67,6 +67,7 @@ class _FakeHost:
|
|||||||
self.removed: list[str] = []
|
self.removed: list[str] = []
|
||||||
self.spawn_args: tuple[object, ...] | None = None
|
self.spawn_args: tuple[object, ...] | None = None
|
||||||
self.mount_config: OrchestratorAgentConfig | None = None
|
self.mount_config: OrchestratorAgentConfig | None = None
|
||||||
|
self.data_dirs_ensured: list[str] = []
|
||||||
|
|
||||||
async def _spawn_container(
|
async def _spawn_container(
|
||||||
self,
|
self,
|
||||||
@@ -80,6 +81,9 @@ class _FakeHost:
|
|||||||
async def _remove_container(self, container_name: str) -> None:
|
async def _remove_container(self, container_name: str) -> None:
|
||||||
self.removed.append(container_name)
|
self.removed.append(container_name)
|
||||||
|
|
||||||
|
def _ensure_opencode_data_dir(self, agent_id: str) -> None:
|
||||||
|
self.data_dirs_ensured.append(agent_id)
|
||||||
|
|
||||||
def _resolve_host_paths(
|
def _resolve_host_paths(
|
||||||
self, config: OrchestratorAgentConfig, agent_settings_path: Path | None
|
self, config: OrchestratorAgentConfig, agent_settings_path: Path | None
|
||||||
) -> dict[str, str | None]:
|
) -> dict[str, str | None]:
|
||||||
@@ -88,6 +92,7 @@ class _FakeHost:
|
|||||||
if config.mcp_config_path
|
if config.mcp_config_path
|
||||||
else None,
|
else None,
|
||||||
"settings": str(agent_settings_path) if agent_settings_path else None,
|
"settings": str(agent_settings_path) if agent_settings_path else None,
|
||||||
|
"opencode": f"/host/data/opencode/{config.agent_id}",
|
||||||
}
|
}
|
||||||
|
|
||||||
def _build_mount_args(
|
def _build_mount_args(
|
||||||
@@ -215,6 +220,10 @@ async def test_grok_spawn_wires_gateway_env_and_image_last() -> None:
|
|||||||
assert "ROBOCO_AGENT_MODEL=grok-build" in cmd
|
assert "ROBOCO_AGENT_MODEL=grok-build" in cmd
|
||||||
# Fixed session id so usage capture can locate the run's session store.
|
# Fixed session id so usage capture can locate the run's session store.
|
||||||
assert "ROBOCO_AGENT_SESSION_ID=sess-1" in cmd
|
assert "ROBOCO_AGENT_SESSION_ID=sess-1" in cmd
|
||||||
|
# Usage capture: per-agent data dir mounted + the entrypoint's usage file.
|
||||||
|
assert host.data_dirs_ensured == ["be-dev-1"]
|
||||||
|
assert "/host/data/opencode/be-dev-1:/home/agent/.grok-usage" in cmd
|
||||||
|
assert "ROBOCO_GROK_USAGE_FILE=/home/agent/.grok-usage/usage.json" in cmd
|
||||||
# Identity wiring from the shared host helpers is present.
|
# Identity wiring from the shared host helpers is present.
|
||||||
assert "ROBOCO_AGENT_TOKEN=hmac-be-dev-1" in cmd
|
assert "ROBOCO_AGENT_TOKEN=hmac-be-dev-1" in cmd
|
||||||
# The image is the final docker-run argument.
|
# The image is the final docker-run argument.
|
||||||
|
|||||||
Reference in New Issue
Block a user