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,164 @@
|
||||
"""gemini_cli_config — mcp-config -> settings.json + per-role Policy Engine TOML."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tomllib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from roboco.llm.providers import gemini_cli_config as gc
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_SAMPLE_MCP = {
|
||||
"mcpServers": {
|
||||
"roboco-flow": {
|
||||
"command": "uv",
|
||||
"args": ["run", "--no-sync", "python", "-m", "roboco.mcp.flow_server"],
|
||||
"env": {"ROBOCO_AGENT_ID": "be-dev-1", "ROBOCO_AGENT_TOKEN": "tok-123"},
|
||||
},
|
||||
"roboco-do": {"command": "uv", "args": ["run", "x"]},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _rules_by_tool(rules: list[dict], tool: str) -> list[dict]:
|
||||
return [r for r in rules if r.get("toolName") == tool]
|
||||
|
||||
|
||||
def test_render_settings_json_injects_mcp_servers_and_env() -> None:
|
||||
rendered = gc.render_settings_json(_SAMPLE_MCP)
|
||||
flow = rendered["mcpServers"]["roboco-flow"]
|
||||
assert flow["command"] == "uv"
|
||||
assert flow["args"][:2] == ["run", "--no-sync"]
|
||||
assert flow["env"]["ROBOCO_AGENT_TOKEN"] == "tok-123"
|
||||
assert "env" not in rendered["mcpServers"]["roboco-do"]
|
||||
|
||||
|
||||
def test_render_settings_json_fixed_flags() -> None:
|
||||
rendered = gc.render_settings_json({})
|
||||
assert rendered["security"]["auth"]["selectedType"] == "oauth-personal"
|
||||
assert rendered["experimental"]["enableAgents"] is False
|
||||
assert rendered["advanced"]["autoConfigureMemory"] is False
|
||||
assert rendered["mcpServers"] == {}
|
||||
|
||||
|
||||
def test_write_gemini_memory_installs_the_blueprint(tmp_path: Path) -> None:
|
||||
src = tmp_path / "system-prompt.md"
|
||||
src.write_text("You are a RoboCo backend developer.", encoding="utf-8")
|
||||
dest = tmp_path / ".gemini" / "GEMINI.md"
|
||||
assert gc.write_gemini_memory(source=src, dest=dest) is True
|
||||
assert dest.read_text(encoding="utf-8") == "You are a RoboCo backend developer."
|
||||
|
||||
|
||||
def test_write_gemini_memory_noops_when_source_absent(tmp_path: Path) -> None:
|
||||
dest = tmp_path / ".gemini" / "GEMINI.md"
|
||||
assert gc.write_gemini_memory(source=tmp_path / "absent.md", dest=dest) is False
|
||||
assert not dest.exists()
|
||||
|
||||
|
||||
def test_developer_policy_only_denies_bash_capable_hazards() -> None:
|
||||
rules = gc.policy_rules_for_role("developer")
|
||||
# Developer writes code + runs a shell -> no edit-tool / shell-blanket deny.
|
||||
assert _rules_by_tool(rules, "write_file") == []
|
||||
assert _rules_by_tool(rules, "replace") == []
|
||||
shell_rules = _rules_by_tool(rules, "run_shell_command")
|
||||
assert shell_rules # bash-capable: command-scoped denies exist
|
||||
assert all("commandPrefix" in r for r in shell_rules)
|
||||
prefixes = {r["commandPrefix"] for r in shell_rules}
|
||||
assert "git push" in prefixes
|
||||
assert "rm -rf" in prefixes
|
||||
|
||||
|
||||
def test_pr_reviewer_policy_blanket_denies_shell_and_edit() -> None:
|
||||
rules = gc.policy_rules_for_role("pr_reviewer")
|
||||
assert _rules_by_tool(rules, "write_file")
|
||||
assert _rules_by_tool(rules, "replace")
|
||||
shell_rules = _rules_by_tool(rules, "run_shell_command")
|
||||
# A read-only reviewer gets ONE blanket shell deny, no command scoping.
|
||||
assert len(shell_rules) == 1
|
||||
assert "commandPrefix" not in shell_rules[0]
|
||||
|
||||
|
||||
def test_main_pm_keeps_shell_but_denies_git_and_edit() -> None:
|
||||
rules = gc.policy_rules_for_role("main_pm")
|
||||
assert _rules_by_tool(rules, "write_file") # PM doesn't write code
|
||||
shell_rules = _rules_by_tool(rules, "run_shell_command")
|
||||
prefixes = {r.get("commandPrefix") for r in shell_rules}
|
||||
assert "git push" in prefixes
|
||||
assert None not in prefixes # no blanket deny — PM keeps its shell
|
||||
|
||||
|
||||
def test_render_policy_toml_is_valid_toml() -> None:
|
||||
parsed = tomllib.loads(gc.render_policy_toml("developer"))
|
||||
assert isinstance(parsed["rule"], list)
|
||||
assert all(r["decision"] == "deny" for r in parsed["rule"])
|
||||
|
||||
|
||||
def test_render_policy_toml_empty_when_no_rules(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Every REAL role currently produces at least one rule (write or shell
|
||||
# denies), but render_policy_toml must still degrade to "" rather than
|
||||
# emit an empty [[rule]] table for the hypothetical case it doesn't.
|
||||
monkeypatch.setattr(gc, "policy_rules_for_role", lambda _role: [])
|
||||
assert gc.render_policy_toml("anything") == ""
|
||||
|
||||
|
||||
def test_unknown_role_gets_every_deny_category() -> None:
|
||||
# An unrecognised role name fails _allows_write's role_config lookup (->
|
||||
# False) and isn't in _BASH_ROLES, so it gets edit denies PLUS the blanket
|
||||
# shell deny — the most restrictive combination.
|
||||
rules = gc.policy_rules_for_role("unknown-role-xyz")
|
||||
assert _rules_by_tool(rules, "write_file")
|
||||
assert _rules_by_tool(rules, "replace")
|
||||
assert len(_rules_by_tool(rules, "run_shell_command")) == 1
|
||||
|
||||
|
||||
def test_write_policy_toml_writes_file(tmp_path: Path) -> None:
|
||||
policies_dir = tmp_path / "policies"
|
||||
assert gc.write_policy_toml("developer", policies_dir=policies_dir) is True
|
||||
written = (policies_dir / "roboco.toml").read_text(encoding="utf-8")
|
||||
assert "run_shell_command" in written
|
||||
|
||||
|
||||
def test_gemini_cli_args_is_yolo_only() -> None:
|
||||
assert gc.gemini_cli_args() == ["--approval-mode", "yolo"]
|
||||
|
||||
|
||||
def test_main_writes_settings_and_args(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
mcp_path = tmp_path / "mcp-config.json"
|
||||
mcp_path.write_text(json.dumps(_SAMPLE_MCP), encoding="utf-8")
|
||||
settings_path = tmp_path / ".gemini" / "settings.json"
|
||||
memory_path = tmp_path / ".gemini" / "GEMINI.md"
|
||||
policies_dir = tmp_path / ".gemini" / "policies"
|
||||
args_path = tmp_path / "gemini-args"
|
||||
system_prompt = tmp_path / "system-prompt.md"
|
||||
system_prompt.write_text("blueprint", encoding="utf-8")
|
||||
|
||||
monkeypatch.setattr(gc, "GEMINI_SETTINGS_PATH", settings_path)
|
||||
monkeypatch.setattr(gc, "GEMINI_MEMORY_PATH", memory_path)
|
||||
monkeypatch.setattr(gc, "GEMINI_POLICIES_DIR", policies_dir)
|
||||
monkeypatch.setattr(gc, "GEMINI_ARGS_PATH", args_path)
|
||||
monkeypatch.setattr(gc, "SYSTEM_PROMPT_PATH", system_prompt)
|
||||
monkeypatch.setenv("ROBOCO_AGENT_ID", "be-dev-1")
|
||||
monkeypatch.setenv("ROBOCO_MCP_CONFIG", str(mcp_path))
|
||||
|
||||
assert gc.main() == 0
|
||||
|
||||
rendered = json.loads(settings_path.read_text(encoding="utf-8"))
|
||||
assert rendered["mcpServers"]["roboco-flow"]["env"]["ROBOCO_AGENT_TOKEN"] == (
|
||||
"tok-123"
|
||||
)
|
||||
assert memory_path.read_text(encoding="utf-8") == "blueprint"
|
||||
assert (policies_dir / "roboco.toml").exists()
|
||||
# One flag token per line — the entrypoint reads it via bash `mapfile -t`.
|
||||
assert args_path.read_text(encoding="utf-8").splitlines() == [
|
||||
"--approval-mode",
|
||||
"yolo",
|
||||
]
|
||||
@@ -0,0 +1,239 @@
|
||||
"""gemini_cli_usage — stats-from-stdout usage capture + exit classification."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
from roboco.llm.providers import gemini_cli_usage as gu
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _single_json(stats: dict) -> str:
|
||||
return json.dumps({"response": "ok", "stats": stats, "error": None})
|
||||
|
||||
|
||||
def _stream_json(events: list[dict]) -> str:
|
||||
return "\n".join(json.dumps(e) for e in events)
|
||||
|
||||
|
||||
# Flat ModelStreamStats — the REAL shape our entrypoint actually parses,
|
||||
# transcribed verbatim from the terminal `result` event's `stats.models.<name>`
|
||||
# entry (--output-format stream-json), per
|
||||
# packages/core/src/output/types.ts's ModelStreamStats interface:
|
||||
# {total_tokens, input_tokens, output_tokens, cached, input} — NO nested
|
||||
# "tokens" key. `input_tokens` is already the full billable prompt count.
|
||||
_FLAT_MODEL_STATS = {
|
||||
"models": {
|
||||
"gemini-2.5-pro": {
|
||||
"total_tokens": 1500,
|
||||
"input_tokens": 1000,
|
||||
"output_tokens": 500,
|
||||
"cached": 0,
|
||||
"input": 1000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Nested SessionMetrics.ModelMetrics — the --output-format json FALLBACK
|
||||
# shape (never actually emitted by our stream-json entrypoint, but tolerated
|
||||
# defensively), transcribed verbatim from
|
||||
# packages/core/src/telemetry/uiTelemetry.ts's ModelMetrics interface:
|
||||
# tokens: {input, prompt, candidates, total, cached, thoughts, tool}.
|
||||
_NESTED_MODEL_STATS = {
|
||||
"models": {
|
||||
"gemini-2.5-pro": {
|
||||
"tokens": {
|
||||
"input": 1000,
|
||||
"prompt": 1000,
|
||||
"candidates": 500,
|
||||
"total": 1700,
|
||||
"cached": 0,
|
||||
"thoughts": 200,
|
||||
"tool": 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_extract_model_stats_reads_flat_stream_json_shape() -> None:
|
||||
# The PRIMARY path: this is the real shape produced by our entrypoint's
|
||||
# --output-format stream-json — no "tokens" nesting, no thoughts/tool
|
||||
# fields to fold (they aren't broken out in this flat shape at all).
|
||||
result = gu.extract_model_stats(_FLAT_MODEL_STATS)
|
||||
assert result == {"gemini-2.5-pro": (1000, 500)}
|
||||
|
||||
|
||||
def test_extract_model_stats_empty_for_missing_models() -> None:
|
||||
assert gu.extract_model_stats({}) == {}
|
||||
assert gu.extract_model_stats({"models": "not-a-dict"}) == {}
|
||||
|
||||
|
||||
def test_extract_model_stats_reads_nested_json_mode_fallback() -> None:
|
||||
# The regression test for the shape bug: a fixture in the OTHER mode's
|
||||
# (--output-format json) shape must still produce sane non-zero usage via
|
||||
# the nested-"tokens" fallback branch, even though our entrypoint never
|
||||
# actually emits this shape. thoughts folds into output: 500 + 200 = 700.
|
||||
result = gu.extract_model_stats(_NESTED_MODEL_STATS)
|
||||
assert result == {"gemini-2.5-pro": (1000, 700)}
|
||||
|
||||
|
||||
def test_usage_and_cost_prices_each_model_at_its_own_rate() -> None:
|
||||
stats = {
|
||||
"models": {
|
||||
# pro: $1.25/$10.00 per 1M
|
||||
"gemini-2.5-pro": {"input_tokens": 1_000_000, "output_tokens": 0},
|
||||
# flash-lite: $0.10/$0.40 per 1M
|
||||
"gemini-2.5-flash-lite": {"input_tokens": 0, "output_tokens": 1_000_000},
|
||||
}
|
||||
}
|
||||
tokens, cost = gu.usage_and_cost(stats)
|
||||
assert tokens == 2_000_000 # noqa: PLR2004
|
||||
assert cost == pytest.approx(1.25 + 0.40)
|
||||
|
||||
|
||||
def test_usage_and_cost_zero_for_empty_stats() -> None:
|
||||
assert gu.usage_and_cost({}) == (0, 0.0)
|
||||
|
||||
|
||||
def test_stats_from_run_log_single_json(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.json"
|
||||
log.write_text(_single_json(_FLAT_MODEL_STATS), encoding="utf-8")
|
||||
assert gu.stats_from_run_log(log) == _FLAT_MODEL_STATS
|
||||
|
||||
|
||||
def test_stats_from_run_log_stream_json_terminal_result_wins(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json(
|
||||
[
|
||||
{"type": "init"},
|
||||
{"type": "message", "data": "hi"},
|
||||
{"type": "result", "stats": _FLAT_MODEL_STATS},
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert gu.stats_from_run_log(log) == _FLAT_MODEL_STATS
|
||||
|
||||
|
||||
def test_stats_from_run_log_missing_or_empty(tmp_path: Path) -> None:
|
||||
assert gu.stats_from_run_log(tmp_path / "absent.json") == {}
|
||||
empty = tmp_path / "empty.json"
|
||||
empty.write_text("", encoding="utf-8")
|
||||
assert gu.stats_from_run_log(empty) == {}
|
||||
|
||||
|
||||
def test_is_quota_error_detects_terminal_and_retryable(tmp_path: Path) -> None:
|
||||
terminal = tmp_path / "terminal.json"
|
||||
terminal.write_text(
|
||||
_single_json({}).replace(
|
||||
'"error": null', '"error": {"type": "TerminalQuotaError"}'
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert gu.is_quota_error(terminal) is True
|
||||
|
||||
retryable = tmp_path / "retryable.ndjson"
|
||||
retryable.write_text(
|
||||
_stream_json([{"type": "error", "error": {"type": "RetryableQuotaError"}}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert gu.is_quota_error(retryable) is True
|
||||
|
||||
|
||||
def test_is_quota_error_false_for_unrelated_error(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json([{"type": "error", "error": {"type": "SomeOtherError"}}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert gu.is_quota_error(log) is False
|
||||
assert gu.is_quota_error(tmp_path / "absent.ndjson") is False
|
||||
|
||||
|
||||
def test_classify_exit_code_auth_passes_through(tmp_path: Path) -> None:
|
||||
# 41 is returned unchanged regardless of what the log carries.
|
||||
log = tmp_path / "run.json"
|
||||
log.write_text(_single_json({}), encoding="utf-8")
|
||||
assert gu.classify_exit_code(41, log) == 41 # noqa: PLR2004
|
||||
|
||||
|
||||
def test_classify_exit_code_remaps_quota_to_75(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json([{"type": "error", "error": {"type": "TerminalQuotaError"}}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert gu.classify_exit_code(1, log) == 75 # noqa: PLR2004
|
||||
|
||||
|
||||
def test_classify_exit_code_passes_through_other_codes(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.json"
|
||||
log.write_text(_single_json({}), encoding="utf-8")
|
||||
for code in (0, 42, 52, 53, 54, 130):
|
||||
assert gu.classify_exit_code(code, log) == code
|
||||
|
||||
|
||||
def test_capture_run_usage_writes_usage_json(tmp_path: Path) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json([{"type": "result", "stats": _FLAT_MODEL_STATS}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
out = tmp_path / "usage.json"
|
||||
tokens = gu.capture_run_usage(
|
||||
run_log=log, fallback_model="gemini-2.5-pro", out_path=out
|
||||
)
|
||||
assert tokens == 1500 # noqa: PLR2004 — 1000 input + 500 output
|
||||
data = json.loads(out.read_text())
|
||||
assert data["model"] == "gemini-2.5-pro"
|
||||
assert data["total_tokens"] == 1500 # noqa: PLR2004
|
||||
assert data["cost_usd"] > 0.0
|
||||
|
||||
|
||||
def test_capture_run_usage_zero_when_log_absent(tmp_path: Path) -> None:
|
||||
out = tmp_path / "usage.json"
|
||||
tokens = gu.capture_run_usage(
|
||||
run_log=tmp_path / "absent.ndjson",
|
||||
fallback_model="gemini-2.5-pro",
|
||||
out_path=out,
|
||||
)
|
||||
assert tokens == 0
|
||||
assert json.loads(out.read_text())["total_tokens"] == 0
|
||||
|
||||
|
||||
def test_main_writes_usage_file(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json([{"type": "result", "stats": _FLAT_MODEL_STATS}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
out = tmp_path / "usage.json"
|
||||
monkeypatch.setattr(gu, "USAGE_OUT_PATH", out)
|
||||
monkeypatch.setenv("ROBOCO_GEMINI_RUN_LOG", str(log))
|
||||
monkeypatch.setenv("ROBOCO_AGENT_MODEL", "gemini-2.5-pro")
|
||||
assert gu.main([]) == 0
|
||||
assert json.loads(out.read_text())["total_tokens"] == 1500 # noqa: PLR2004
|
||||
|
||||
|
||||
def test_main_classify_exit_prints_remapped_code(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
log = tmp_path / "run.ndjson"
|
||||
log.write_text(
|
||||
_stream_json([{"type": "error", "error": {"type": "RetryableQuotaError"}}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setenv("ROBOCO_GEMINI_RUN_LOG", str(log))
|
||||
monkeypatch.setenv("ROBOCO_GEMINI_CLI_EXIT_CODE", "1")
|
||||
assert gu.main(["--classify-exit"]) == 0
|
||||
assert capsys.readouterr().out.strip() == "75"
|
||||
@@ -0,0 +1,258 @@
|
||||
"""Tests for GeminiCliProvider (Google Gemini via the official ``gemini`` CLI).
|
||||
|
||||
Mirrors ``tests/unit/llm/test_providers.py``'s Grok coverage — the same safety
|
||||
properties matter here:
|
||||
|
||||
* the agent gets the MCP gateway wiring (reuses the orchestrator mount path);
|
||||
* the OAuth credential (~/.gemini) is mounted, and the provider routing
|
||||
fields are blanked so the gemini endpoint is never mislabelled ANTHROPIC_*;
|
||||
* the prompt travels via env, so a leading ``--`` cannot become a CLI flag.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from roboco.llm.providers import GeminiCliProvider, ProviderError, SpawnResult
|
||||
from roboco.models.runtime import OrchestratorAgentConfig
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_gemini_auth(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
"""Point GEMINI_AUTH_HOST_PATH at a fresh tmp dir so tests never mount the
|
||||
real ~/.gemini. Tests that exercise the auth mount create oauth_creds.json
|
||||
themselves."""
|
||||
monkeypatch.setattr(
|
||||
"roboco.llm.providers.gemini.GEMINI_AUTH_HOST_PATH", str(tmp_path)
|
||||
)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _config(
|
||||
*,
|
||||
agent_id: str = "be-dev-1",
|
||||
provider_type: str = "gemini",
|
||||
provider_base_url: str | None = None,
|
||||
provider_auth_token: str | None = None,
|
||||
mcp_config_path: Path | None = Path("/host/mcp-configs/be-dev-1.json"),
|
||||
) -> OrchestratorAgentConfig:
|
||||
return OrchestratorAgentConfig(
|
||||
agent_id=agent_id,
|
||||
blueprint_path=Path("/app/system-prompt.md"),
|
||||
model="gemini-2.5-pro",
|
||||
mcp_config_path=mcp_config_path,
|
||||
claude_session_id="sess-1",
|
||||
provider_type=provider_type,
|
||||
provider_base_url=provider_base_url,
|
||||
provider_auth_token=provider_auth_token,
|
||||
)
|
||||
|
||||
|
||||
class _FakeHost:
|
||||
"""Implements the orchestrator surface the provider delegates to."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.removed: list[str] = []
|
||||
self.remove_stop_reasons: list[str | None] = []
|
||||
self.mount_config: OrchestratorAgentConfig | None = None
|
||||
self.data_dirs_ensured: list[str] = []
|
||||
|
||||
async def _remove_container(
|
||||
self, container_name: str, *, stop_reason: str | None = None
|
||||
) -> None:
|
||||
self.removed.append(container_name)
|
||||
self.remove_stop_reasons.append(stop_reason)
|
||||
|
||||
def _ensure_gemini_usage_dir(self, agent_id: str) -> None:
|
||||
self.data_dirs_ensured.append(agent_id)
|
||||
|
||||
def _resolve_host_paths(
|
||||
self, config: OrchestratorAgentConfig, agent_settings_path: Path | None
|
||||
) -> dict[str, str | None]:
|
||||
return {
|
||||
"mcp_config": str(config.mcp_config_path)
|
||||
if config.mcp_config_path
|
||||
else None,
|
||||
"settings": str(agent_settings_path) if agent_settings_path else None,
|
||||
"gemini_usage": f"/host/data/gemini-usage/{config.agent_id}",
|
||||
}
|
||||
|
||||
def _build_mount_args(
|
||||
self,
|
||||
container_name: str,
|
||||
config: OrchestratorAgentConfig,
|
||||
hosts: dict[str, str | None],
|
||||
) -> list[str]:
|
||||
# Record the config the mount step saw, and MIMIC the real
|
||||
# _append_provider_env so a missed blanking would leak ANTHROPIC_*.
|
||||
self.mount_config = config
|
||||
cmd = ["docker", "run", "-d", "--name", container_name]
|
||||
mcp = hosts.get("mcp_config")
|
||||
if mcp:
|
||||
cmd += ["-v", f"{mcp}:/app/mcp-config.json:ro"]
|
||||
if config.provider_base_url:
|
||||
cmd += ["-e", f"ANTHROPIC_BASE_URL={config.provider_base_url}"]
|
||||
if config.provider_auth_token:
|
||||
cmd += ["-e", f"ANTHROPIC_AUTH_TOKEN={config.provider_auth_token}"]
|
||||
return cmd
|
||||
|
||||
def _append_agent_auth_env(
|
||||
self, cmd: list[str], config: OrchestratorAgentConfig
|
||||
) -> None:
|
||||
cmd += ["-e", f"ROBOCO_AGENT_TOKEN=hmac-{config.agent_id}"]
|
||||
|
||||
def _append_git_context_env(
|
||||
self, cmd: list[str], config: OrchestratorAgentConfig
|
||||
) -> None:
|
||||
cmd += ["-e", f"ROBOCO_GIT_AGENT={config.agent_id}"]
|
||||
|
||||
|
||||
def _proc(
|
||||
returncode: int = 0, stdout: bytes = b"cid\n", stderr: bytes = b""
|
||||
) -> MagicMock:
|
||||
proc = MagicMock()
|
||||
proc.returncode = returncode
|
||||
proc.communicate = AsyncMock(return_value=(stdout, stderr))
|
||||
return proc
|
||||
|
||||
|
||||
async def test_gemini_spawn_requires_mcp_config() -> None:
|
||||
provider = GeminiCliProvider(_FakeHost())
|
||||
with pytest.raises(ProviderError, match="MCP config"):
|
||||
await provider.spawn(_config(mcp_config_path=None))
|
||||
|
||||
|
||||
async def test_gemini_spawn_does_not_require_api_key() -> None:
|
||||
# OAuth login (mounted ~/.gemini) — a missing provider key/url is fine.
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())):
|
||||
result = await provider.spawn(_config(provider_auth_token=None))
|
||||
assert result.instance_id == "roboco-agent-be-dev-1"
|
||||
|
||||
|
||||
async def test_gemini_spawn_no_anthropic_leak() -> None:
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host, image="roboco-agent-gemini:test")
|
||||
with patch(
|
||||
"asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())
|
||||
) as exec_mock:
|
||||
await provider.spawn(
|
||||
_config(provider_base_url="https://ignored", provider_auth_token="ignored"),
|
||||
initial_prompt="do the work",
|
||||
)
|
||||
cmd = list(exec_mock.call_args.args)
|
||||
# The provider endpoint must NOT be injected as an Anthropic var.
|
||||
assert not any(c.startswith("ANTHROPIC_BASE_URL=") for c in cmd)
|
||||
assert not any(c.startswith("ANTHROPIC_AUTH_TOKEN=") for c in cmd)
|
||||
# Provider fields were blanked before the shared mount step.
|
||||
assert host.mount_config is not None
|
||||
assert host.mount_config.provider_base_url is None
|
||||
assert host.mount_config.provider_auth_token is None
|
||||
|
||||
|
||||
async def test_gemini_spawn_wires_gateway_env_and_image_last() -> None:
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host, image="roboco-agent-gemini:test")
|
||||
with patch(
|
||||
"asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())
|
||||
) as exec_mock:
|
||||
result = await provider.spawn(_config())
|
||||
cmd = list(exec_mock.call_args.args)
|
||||
assert "ROBOCO_MCP_CONFIG=/app/mcp-config.json" in cmd
|
||||
assert "ROBOCO_AGENT_ID=be-dev-1" in cmd
|
||||
assert "ROBOCO_AGENT_MODEL=gemini-2.5-pro" 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/gemini-usage/be-dev-1:/home/agent/.gemini-usage" in cmd
|
||||
assert "ROBOCO_GEMINI_USAGE_FILE=/home/agent/.gemini-usage/usage.json" 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.
|
||||
assert cmd[-1] == "roboco-agent-gemini:test"
|
||||
assert host.removed == ["roboco-agent-be-dev-1"]
|
||||
assert host.remove_stop_reasons == ["pre_spawn_stale_clear"]
|
||||
assert result == SpawnResult(
|
||||
instance_id="roboco-agent-be-dev-1",
|
||||
extra={"container_id": "cid", "model": "gemini-2.5-pro"},
|
||||
)
|
||||
|
||||
|
||||
async def test_gemini_spawn_mounts_auth_when_present(
|
||||
_isolate_gemini_auth: Path,
|
||||
) -> None:
|
||||
(_isolate_gemini_auth / "oauth_creds.json").write_text("{}", encoding="utf-8")
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
with patch(
|
||||
"asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())
|
||||
) as exec_mock:
|
||||
await provider.spawn(_config())
|
||||
cmd = list(exec_mock.call_args.args)
|
||||
expected = f"{_isolate_gemini_auth}:/home/agent/.gemini-auth-ro:ro"
|
||||
assert expected in cmd
|
||||
|
||||
|
||||
async def test_gemini_spawn_omits_auth_mount_when_absent() -> None:
|
||||
# No oauth_creds.json in the (tmp) GEMINI_AUTH_HOST_PATH -> no mount, no crash.
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
with patch(
|
||||
"asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())
|
||||
) as exec_mock:
|
||||
await provider.spawn(_config())
|
||||
cmd = list(exec_mock.call_args.args)
|
||||
assert not any("/home/agent/.gemini-auth-ro" in c for c in cmd)
|
||||
|
||||
|
||||
async def test_gemini_spawn_warns_when_auth_absent(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""A missing host oauth_creds.json must not be silent — the spawn is doomed
|
||||
to exit 41, so the operator gets a spawn-time WARNING naming the missing
|
||||
file and the remediation."""
|
||||
caplog.set_level("WARNING", logger="roboco.llm.providers.gemini")
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())):
|
||||
await provider.spawn(_config())
|
||||
warnings = [r for r in caplog.records if r.levelname == "WARNING"]
|
||||
assert warnings, "expected a spawn-time WARNING for the missing oauth_creds.json"
|
||||
msg = warnings[0].getMessage()
|
||||
assert "oauth_creds.json" in msg
|
||||
assert "gemini" in msg # names the remediation
|
||||
|
||||
|
||||
async def test_gemini_spawn_prompt_is_injection_safe() -> None:
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
nasty = "--model evil --approval-mode yolo-override"
|
||||
with patch(
|
||||
"asyncio.create_subprocess_exec", AsyncMock(return_value=_proc())
|
||||
) as exec_mock:
|
||||
await provider.spawn(_config(), initial_prompt=nasty)
|
||||
cmd = list(exec_mock.call_args.args)
|
||||
# Passed only as an env value, never as a bare argv token.
|
||||
assert f"ROBOCO_INITIAL_PROMPT={nasty}" in cmd
|
||||
assert nasty not in cmd
|
||||
|
||||
|
||||
async def test_gemini_spawn_raises_on_docker_failure() -> None:
|
||||
provider = GeminiCliProvider(_FakeHost())
|
||||
with (
|
||||
patch(
|
||||
"asyncio.create_subprocess_exec",
|
||||
AsyncMock(return_value=_proc(returncode=1, stderr=b"boom")),
|
||||
),
|
||||
pytest.raises(ProviderError, match="boom"),
|
||||
):
|
||||
await provider.spawn(_config())
|
||||
|
||||
|
||||
async def test_gemini_remove_delegates_to_host() -> None:
|
||||
host = _FakeHost()
|
||||
provider = GeminiCliProvider(host)
|
||||
await provider.remove("roboco-agent-be-dev-1")
|
||||
assert host.removed == ["roboco-agent-be-dev-1"]
|
||||
@@ -0,0 +1,253 @@
|
||||
"""GEMINI quota/auth parking: break the exit -> respawn cost loop.
|
||||
|
||||
A one-shot gemini run that hits a quota error is remapped to exit 75 by the
|
||||
entrypoint wrapper (see gemini_cli_usage.classify_exit_code); a missing/empty
|
||||
OAuth credential exits 41 (the CLI's own dedicated auth-failure code). Both
|
||||
park the GEMINI provider instead of crash-retrying, mirroring grok's exit-75 /
|
||||
exit-78 parks (see test_grok_rate_limit.py) but tracked independently.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from roboco.models.runtime import AgentInstance
|
||||
from roboco.runtime.orchestrator import (
|
||||
_GEMINI_AUTH_EXIT_CODE,
|
||||
_GEMINI_RATE_LIMIT_EXIT_CODE,
|
||||
_GEMINI_REPARK_BACKOFF_CAP,
|
||||
AgentOrchestrator,
|
||||
AgentState,
|
||||
)
|
||||
|
||||
|
||||
def _gemini_instance(provider_type: str = "gemini") -> AgentInstance:
|
||||
cfg = type("C", (), {"provider_type": provider_type, "model": "gemini-2.5-pro"})()
|
||||
inst = AgentInstance(agent_id="be-dev-1", state=AgentState.ACTIVE, config=cfg)
|
||||
inst.current_task_id = "task-1"
|
||||
inst.container_id = "cid"
|
||||
return inst
|
||||
|
||||
|
||||
class _FakeTracker:
|
||||
def __init__(self) -> None:
|
||||
self.activated_with: dict[str, object] | None = None
|
||||
|
||||
async def activate(
|
||||
self,
|
||||
*,
|
||||
retry_after: float,
|
||||
affected_agents: list[str],
|
||||
kind: str = "rate_limited",
|
||||
) -> None:
|
||||
self.activated_with = {
|
||||
"retry_after": retry_after,
|
||||
"affected_agents": affected_agents,
|
||||
"kind": kind,
|
||||
}
|
||||
|
||||
|
||||
class _RecordingTracker:
|
||||
"""Records every activate() retry_after across multiple re-parks."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.retry_afters: list[float] = []
|
||||
self.kinds: list[str] = []
|
||||
|
||||
async def activate(
|
||||
self, *, retry_after: float, affected_agents: list[str], kind: str
|
||||
) -> None:
|
||||
del affected_agents
|
||||
self.retry_afters.append(retry_after)
|
||||
self.kinds.append(kind)
|
||||
|
||||
|
||||
def _orch() -> AgentOrchestrator:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
orch._waiting_records = {}
|
||||
orch._rate_limit_ceo_notified = set()
|
||||
orch._gemini_last_park_at = None
|
||||
orch._gemini_repark_count = 0
|
||||
orch._gemini_rate_limit_retry_after_s = 60.0
|
||||
orch._gemini_auth_retry_after_s = 60.0
|
||||
return orch
|
||||
|
||||
|
||||
def test_is_gemini_rate_limit_exit() -> None:
|
||||
inst = _gemini_instance()
|
||||
assert AgentOrchestrator._is_gemini_rate_limit_exit(
|
||||
inst, _GEMINI_RATE_LIMIT_EXIT_CODE
|
||||
)
|
||||
assert not AgentOrchestrator._is_gemini_rate_limit_exit(inst, 0)
|
||||
assert not AgentOrchestrator._is_gemini_rate_limit_exit(inst, 1)
|
||||
assert not AgentOrchestrator._is_gemini_rate_limit_exit(
|
||||
_gemini_instance(provider_type="anthropic"), _GEMINI_RATE_LIMIT_EXIT_CODE
|
||||
)
|
||||
# Same numeric exit code as grok's own detector, but provider-scoped: a
|
||||
# grok instance exiting 75 is NOT a gemini rate-limit exit.
|
||||
assert not AgentOrchestrator._is_gemini_rate_limit_exit(
|
||||
_gemini_instance(provider_type="grok"), _GEMINI_RATE_LIMIT_EXIT_CODE
|
||||
)
|
||||
|
||||
|
||||
def test_is_gemini_auth_exit() -> None:
|
||||
inst = _gemini_instance()
|
||||
assert AgentOrchestrator._is_gemini_auth_exit(inst, _GEMINI_AUTH_EXIT_CODE)
|
||||
assert not AgentOrchestrator._is_gemini_auth_exit(inst, 0)
|
||||
assert not AgentOrchestrator._is_gemini_auth_exit(inst, 1)
|
||||
assert not AgentOrchestrator._is_gemini_auth_exit(
|
||||
_gemini_instance(provider_type="anthropic"), _GEMINI_AUTH_EXIT_CODE
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_park_gemini_rate_limited_activates_and_offlines(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = _orch()
|
||||
inst = _gemini_instance()
|
||||
inst.error_count = 2 # pretend prior crashes — parking must NOT count one
|
||||
tracker = _FakeTracker()
|
||||
monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker)
|
||||
finalize = AsyncMock()
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", finalize)
|
||||
monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock())
|
||||
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
|
||||
finalize.assert_awaited_once()
|
||||
assert inst.state == AgentState.OFFLINE
|
||||
assert inst.container_id is None
|
||||
assert inst.error_count == 0 # a quota park is not a crash
|
||||
assert tracker.activated_with == {
|
||||
"retry_after": pytest.approx(60.0),
|
||||
"affected_agents": ["be-dev-1"],
|
||||
"kind": "rate_limited",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_park_gemini_auth_unavailable_activates_with_auth_missing_kind(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = _orch()
|
||||
inst = _gemini_instance()
|
||||
inst.error_count = 2
|
||||
tracker = _FakeTracker()
|
||||
monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", AsyncMock())
|
||||
monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock())
|
||||
|
||||
await orch._park_gemini_auth_unavailable("be-dev-1", inst)
|
||||
|
||||
assert inst.state == AgentState.OFFLINE
|
||||
assert inst.container_id is None
|
||||
assert inst.error_count == 0
|
||||
assert tracker.activated_with == {
|
||||
"retry_after": pytest.approx(60.0),
|
||||
"affected_agents": ["be-dev-1"],
|
||||
"kind": "auth_missing",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_stopped_container_parks_on_gemini_quota_exit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
inst = _gemini_instance()
|
||||
park = AsyncMock()
|
||||
finalize = AsyncMock()
|
||||
monkeypatch.setattr(orch, "_park_gemini_rate_limited", park)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", finalize)
|
||||
|
||||
await orch._handle_stopped_container("be-dev-1", inst, _GEMINI_RATE_LIMIT_EXIT_CODE)
|
||||
|
||||
park.assert_awaited_once_with("be-dev-1", inst)
|
||||
finalize.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_stopped_container_parks_on_gemini_auth_exit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
inst = _gemini_instance()
|
||||
park = AsyncMock()
|
||||
finalize = AsyncMock()
|
||||
monkeypatch.setattr(orch, "_park_gemini_auth_unavailable", park)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", finalize)
|
||||
|
||||
await orch._handle_stopped_container("be-dev-1", inst, _GEMINI_AUTH_EXIT_CODE)
|
||||
|
||||
park.assert_awaited_once_with("be-dev-1", inst)
|
||||
finalize.assert_not_awaited()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Gemini has no real recovery probe either (an OAuth-login daily quota cap has
|
||||
# no cheap balance-check API) — mirrors grok's repark-backoff tests exactly.
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _backoff_orchestrator() -> AgentOrchestrator:
|
||||
return _orch()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_repark_backs_off_within_episode(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = _backoff_orchestrator()
|
||||
tracker = _RecordingTracker()
|
||||
monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", AsyncMock())
|
||||
monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock())
|
||||
inst = _gemini_instance()
|
||||
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
|
||||
assert tracker.retry_afters == [60.0, 120.0, 240.0]
|
||||
assert tracker.kinds == ["rate_limited", "rate_limited", "rate_limited"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_repark_resets_after_episode_gap(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
orch = _backoff_orchestrator()
|
||||
orch._gemini_repark_count = 3
|
||||
orch._gemini_last_park_at = datetime.now(UTC) - timedelta(hours=2)
|
||||
tracker = _RecordingTracker()
|
||||
monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", AsyncMock())
|
||||
monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock())
|
||||
inst = _gemini_instance()
|
||||
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
|
||||
assert tracker.retry_afters == [60.0]
|
||||
assert orch._gemini_repark_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_repark_backoff_caps(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
orch = _backoff_orchestrator()
|
||||
tracker = _RecordingTracker()
|
||||
monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker)
|
||||
monkeypatch.setattr(orch, "_finalize_spawn_session", AsyncMock())
|
||||
monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock())
|
||||
inst = _gemini_instance()
|
||||
|
||||
for _ in range(_GEMINI_REPARK_BACKOFF_CAP + 3):
|
||||
await orch._park_gemini_rate_limited("be-dev-1", inst)
|
||||
|
||||
max_expected = 60.0 * (2**_GEMINI_REPARK_BACKOFF_CAP)
|
||||
assert all(
|
||||
r == max_expected for r in tracker.retry_afters[_GEMINI_REPARK_BACKOFF_CAP:]
|
||||
)
|
||||
assert max(tracker.retry_afters) == max_expected
|
||||
@@ -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()
|
||||
@@ -1,15 +1,16 @@
|
||||
"""The orchestrator routes only dedicated-backend providers through the registry.
|
||||
|
||||
GROK gets the GrokCliProvider; Anthropic / Ollama Cloud / self-hosted (and any
|
||||
unknown value) return None so ``_spawn_container`` runs its built-in Claude Code
|
||||
path unchanged. This keeps the GROK addition purely additive.
|
||||
GROK gets the GrokCliProvider, GEMINI gets the GeminiCliProvider; Anthropic /
|
||||
Ollama Cloud / self-hosted (and any unknown value) return None so
|
||||
``_spawn_container`` runs its built-in Claude Code path unchanged. This keeps
|
||||
the GROK / GEMINI additions purely additive.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from roboco.llm.providers import GrokCliProvider
|
||||
from roboco.llm.providers import GeminiCliProvider, GrokCliProvider
|
||||
from roboco.runtime.orchestrator import AgentOrchestrator
|
||||
|
||||
|
||||
@@ -24,6 +25,10 @@ def test_provider_for_grok_returns_grok_provider() -> None:
|
||||
assert isinstance(_make_orch()._provider_for("grok"), GrokCliProvider)
|
||||
|
||||
|
||||
def test_provider_for_gemini_returns_gemini_provider() -> None:
|
||||
assert isinstance(_make_orch()._provider_for("gemini"), GeminiCliProvider)
|
||||
|
||||
|
||||
def test_provider_for_anthropic_returns_none() -> None:
|
||||
assert _make_orch()._provider_for("anthropic") is None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user