mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
feat(providers): Gemini CLI provider — ModelProvider.GEMINI (#660)
* feat(providers): Gemini CLI provider — ModelProvider.GEMINI Mirrors the grok blueprint with source-verified divergences (all facts pinned against google-gemini/gemini-cli @ 9681621c): no refresher daemon — Google's refresh tokens are reusable, so the RO host mount is COPIED to a writable container-local ~/.gemini and each container refreshes in-process independently (the write-back crash risk on RO never triggers); settings.json renders security.auth.selectedType 'oauth-personal', experimental.enableAgents=false (subagent ban), autoConfigureMemory=false with a bounded heap; tool scoping rides the tiered TOML Policy Engine (deny-only rules that yolo mode structurally cannot beat); gemini -p with --output-format stream-json; usage parsed from the run's own stdout stats — the adversarial pass caught the parser reading the json-mode nested shape while the entrypoint runs stream-json's FLAT shape (every real run would have priced $0 forever, hidden by fixtures sharing the assumption) — now flat-primary with the nested shape as cited fallback; rate-limit classified from structured error.type only (model-echo immune), native exit 41 auth passthrough; per-model pricing for the three GA models; migrations 084 (enum) + 085 (seed) complete the 082-085 finale chain. V1 excludes interactive intake/secretary. Stack-merge required two behavior-preserving complexity refactors in the shared park/usage plumbing (a park-pair loop; a usage-reader dispatch dict). * fix(providers): route gemini usage read through the containment barrier Mirrors the codex/grok fix — _gemini_usage_json now delegates to _read_usage_json_contained, so CodeQL's path-injection alert on the gemini read is resolved by the same resolve-and-contain guard. --------- Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
@@ -0,0 +1,158 @@
|
||||
"""GEMINI agents capture token usage/cost from their captured ``usage.json``.
|
||||
|
||||
A Gemini agent runs the gemini CLI — no SDK /usage/status server and no
|
||||
Claude transcript — so finalize reads the ``usage.json`` the entrypoint wrote
|
||||
to the per-agent data dir (mounted into the orchestrator). Mirrors
|
||||
test_grok_usage_finalize.py; gemini's usage.json is priced per-model
|
||||
server-side (gemini_cli_usage.usage_and_cost) but flattens to the SAME
|
||||
``{model, total_tokens, cost_usd}`` shape, so the read side is identical to
|
||||
grok's: the whole total folds into output.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from roboco.models.runtime import AgentInstance
|
||||
from roboco.runtime import orchestrator as orch_mod
|
||||
from roboco.runtime.orchestrator import AgentOrchestrator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _write_usage(path: Path, total_tokens: int, cost_usd: float) -> None:
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"model": "gemini-2.5-pro",
|
||||
"total_tokens": total_tokens,
|
||||
"cost_usd": cost_usd,
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def test_gemini_usage_folds_total_into_output(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
usage = tmp_path / "usage.json"
|
||||
_write_usage(usage, total_tokens=180, cost_usd=0.02)
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
monkeypatch.setattr(
|
||||
orch, "_gemini_usage_json", lambda _aid: json.loads(usage.read_text())
|
||||
)
|
||||
|
||||
assert orch._gemini_usage_tokens("be-dev-1") == (0, 180, 0, 0)
|
||||
|
||||
|
||||
def test_gemini_usage_zero_when_store_missing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
monkeypatch.setattr(orch, "_gemini_usage_json", lambda _aid: None)
|
||||
assert orch._gemini_usage_tokens("be-dev-1") == (0, 0, 0, 0)
|
||||
|
||||
|
||||
def test_gemini_cost_read_from_usage_json(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured_cost = 3.25
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
monkeypatch.setattr(
|
||||
orch,
|
||||
"_gemini_usage_json",
|
||||
lambda _aid: {"cost_usd": captured_cost, "total_tokens": 9},
|
||||
)
|
||||
assert orch._gemini_cost_usd("be-dev-1") == captured_cost
|
||||
monkeypatch.setattr(orch, "_gemini_usage_json", lambda _aid: None)
|
||||
assert orch._gemini_cost_usd("be-dev-1") == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_final_usage_routes_gemini_to_usage_json(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
monkeypatch.setattr(
|
||||
orch, "_gemini_usage_json", lambda _aid: {"total_tokens": 12, "cost_usd": 0.01}
|
||||
)
|
||||
cfg = type("C", (), {"provider_type": "gemini"})()
|
||||
orch._instances = {"be-dev-1": AgentInstance(agent_id="be-dev-1", config=cfg)}
|
||||
|
||||
assert await orch._resolve_final_token_usage("be-dev-1") == (0, 12, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_final_turns_tools_gemini_has_neither() -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
cfg = type("C", (), {"provider_type": "gemini"})()
|
||||
orch._instances = {"be-dev-1": AgentInstance(agent_id="be-dev-1", config=cfg)}
|
||||
assert await orch._resolve_final_turns_tools("be-dev-1") == (0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_active_tokens_routes_gemini_to_usage_json(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
monkeypatch.setattr(
|
||||
orch, "_gemini_usage_json", lambda _aid: {"total_tokens": 12, "cost_usd": 0.01}
|
||||
)
|
||||
cfg = type("C", (), {"provider_type": "gemini"})()
|
||||
orch._instances = {"be-dev-1": AgentInstance(agent_id="be-dev-1", config=cfg)}
|
||||
async with httpx.AsyncClient() as client:
|
||||
assert await orch._resolve_active_tokens(client, "be-dev-1") == (0, 12, 0, 0)
|
||||
|
||||
|
||||
def test_gemini_usage_dir_branches_compose_vs_local(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(orch_mod, "PROJECT_HOST_PATH", "")
|
||||
local = AgentOrchestrator._gemini_usage_dir("be-dev-1")
|
||||
assert "roboco-gemini-usage" in str(local)
|
||||
assert local.name == "be-dev-1"
|
||||
|
||||
monkeypatch.setattr(orch_mod, "PROJECT_HOST_PATH", "/volume1/roboco")
|
||||
monkeypatch.setattr(orch_mod, "GEMINI_USAGE_DATA_DIR", "/data/gemini-usage")
|
||||
assert str(AgentOrchestrator._gemini_usage_dir("be-dev-1")) == (
|
||||
"/data/gemini-usage/be-dev-1"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad",
|
||||
["..", ".", "../etc", "a/b", "a\\b", "", "be-dev-1/../x", "x\x00y"],
|
||||
)
|
||||
def test_gemini_usage_dir_rejects_path_traversal(bad: str) -> None:
|
||||
with pytest.raises(ValueError, match="unsafe agent id"):
|
||||
AgentOrchestrator._gemini_usage_dir(bad)
|
||||
|
||||
|
||||
def test_gemini_usage_json_reads_the_real_local_dir(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
# The un-mocked read path must find usage.json in the SAME branched dir the
|
||||
# writer mounts (mirrors _ensure_gemini_usage_dir's create path).
|
||||
monkeypatch.setattr(orch_mod, "PROJECT_HOST_PATH", "")
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
udir = tmp_path / "roboco-gemini-usage" / "be-dev-1"
|
||||
udir.mkdir(parents=True)
|
||||
(udir / "usage.json").write_text(
|
||||
json.dumps({"total_tokens": 55, "cost_usd": 0.1}), encoding="utf-8"
|
||||
)
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
assert orch._gemini_usage_tokens("be-dev-1") == (0, 55, 0, 0)
|
||||
assert orch._gemini_cost_usd("be-dev-1") == 0.1 # noqa: PLR2004
|
||||
|
||||
|
||||
def test_ensure_gemini_usage_dir_creates_world_writable(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(orch_mod, "PROJECT_HOST_PATH", "")
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
orch._ensure_gemini_usage_dir("be-dev-1")
|
||||
target = tmp_path / "roboco-gemini-usage" / "be-dev-1"
|
||||
assert target.is_dir()
|
||||
Reference in New Issue
Block a user