diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index 3415104..0df9b41 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -13,8 +13,11 @@ APP_DIR, ToolSpec, backup_existing_file, + deep_merge_dict, parse_dotenv, + read_json_safe, write_dotenv, + write_json_file, ) from ucode.databricks import ( TOKEN_REFRESH_INTERVAL_SECONDS, @@ -27,6 +30,8 @@ GEMINI_CONFIG_DIR = Path.home() / ".gemini" GEMINI_ENV_PATH = GEMINI_CONFIG_DIR / "ucode.env" GEMINI_BACKUP_PATH = APP_DIR / "gemini-ucode-env.backup" +GEMINI_HOME_DIR = APP_DIR / ".gemini-home" +GEMINI_SETTINGS_PATH = GEMINI_HOME_DIR / ".gemini" / "settings.json" SPEC: ToolSpec = { "binary": "gemini", @@ -50,6 +55,15 @@ def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) +def _ensure_local_settings_selected_type() -> None: + settings = read_json_safe(GEMINI_SETTINGS_PATH) + deep_merge_dict( + settings, + {"security": {"auth": {"selectedType": "gemini-api-key"}}}, + ) + write_json_file(GEMINI_SETTINGS_PATH, settings) + + def render_env_overlay(workspace: str, model: str, token: str) -> dict[str, str]: # Gemini CLI parses GEMINI_CLI_CUSTOM_HEADERS as comma-separated # `Key:Value` pairs and spreads them after the SDK's default User-Agent, @@ -67,11 +81,13 @@ def render_env_overlay(workspace: str, model: str, token: str) -> dict[str, str] def build_runtime_env(workspace: str, model: str, token: str) -> dict[str, str]: + _ensure_local_settings_selected_type() env = os.environ.copy() env.update(render_env_overlay(workspace, model, token)) # Newer Gemini CLI releases refuse to run in untrusted directories; # opt every launch into trust so `ucode gemini` works in any folder. env["GEMINI_CLI_TRUST_WORKSPACE"] = "true" + env["GEMINI_CLI_HOME"] = str(GEMINI_HOME_DIR) return env diff --git a/tests/test_agent_gemini.py b/tests/test_agent_gemini.py index fef70f5..b44172c 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -4,11 +4,20 @@ import json +import pytest + from ucode.agents import gemini WS = "https://example.databricks.com" +@pytest.fixture(autouse=True) +def _redirect_gemini_home(tmp_path, monkeypatch): + gemini_home = tmp_path / ".gemini-home" + monkeypatch.setattr(gemini, "GEMINI_HOME_DIR", gemini_home) + monkeypatch.setattr(gemini, "GEMINI_SETTINGS_PATH", gemini_home / ".gemini" / "settings.json") + + class TestGeminiSpec: def test_binary(self): assert gemini.SPEC["binary"] == "gemini" @@ -70,6 +79,35 @@ def test_sets_oauth_token_for_mcp(self): env = gemini.build_runtime_env(WS, "gemini-2", "tok") assert env["OAUTH_TOKEN"] == "tok" + def test_sets_private_gemini_home(self, tmp_path, monkeypatch): + gemini_home = tmp_path / "private-gemini-home" + monkeypatch.setattr(gemini, "GEMINI_HOME_DIR", gemini_home) + monkeypatch.setattr( + gemini, "GEMINI_SETTINGS_PATH", gemini_home / ".gemini" / "settings.json" + ) + + env = gemini.build_runtime_env(WS, "gemini-2", "tok") + + assert env["GEMINI_CLI_HOME"] == str(gemini_home) + + def test_writes_private_auth_settings(self): + gemini.build_runtime_env(WS, "gemini-2", "tok") + + settings = json.loads(gemini.GEMINI_SETTINGS_PATH.read_text()) + assert settings == {"security": {"auth": {"selectedType": "gemini-api-key"}}} + + def test_preserves_existing_private_settings(self): + gemini.GEMINI_SETTINGS_PATH.parent.mkdir(parents=True) + gemini.GEMINI_SETTINGS_PATH.write_text( + json.dumps({"theme": "dark", "security": {"auth": {"selectedType": "oauth"}}}) + ) + + gemini.build_runtime_env(WS, "gemini-2", "tok") + + settings = json.loads(gemini.GEMINI_SETTINGS_PATH.read_text()) + assert settings["theme"] == "dark" + assert settings["security"]["auth"]["selectedType"] == "gemini-api-key" + class TestGeminiDefaultModel: def test_returns_first_model(self): diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 22a8d3b..8bb3b85 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -10,9 +10,12 @@ from __future__ import annotations +import json import os import shutil import subprocess +from urllib import error as urllib_error +from urllib import request as urllib_request import pytest @@ -57,6 +60,39 @@ def _run_agent( ) +def _run_gemini_gateway_smoke(workspace: str, model: str, token: str) -> str: + """Call the Gemini gateway directly with a text-only prompt. + + This keeps auth recovery coverage focused on the recovered Databricks token + instead of Gemini CLI's separate tool-calling request shape. + """ + url = f"{build_tool_base_url('gemini', workspace)}/v1beta/models/{model}:generateContent" + payload = { + "contents": [ + {"role": "user", "parts": [{"text": "say hi in 5 words or less"}]}, + ], + } + req = urllib_request.Request( + url, + data=json.dumps(payload).encode("utf-8"), + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "application/json", + }, + method="POST", + ) + try: + with urllib_request.urlopen(req, timeout=30) as response: + body = response.read().decode("utf-8") + except urllib_error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") if exc.fp else "" + raise AssertionError(f"Gemini gateway smoke failed: HTTP {exc.code}: {body[:500]}") from exc + + data = json.loads(body) + return data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "") + + # --------------------------------------------------------------------------- # Databricks auth / token # --------------------------------------------------------------------------- @@ -442,6 +478,10 @@ def test_launch_gemini_per_model( monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(gemini, "GEMINI_ENV_PATH", tmp_path / "ucode.env") monkeypatch.setattr(gemini, "GEMINI_BACKUP_PATH", tmp_path / "gemini-ucode-env.backup") + monkeypatch.setattr(gemini, "GEMINI_HOME_DIR", tmp_path / ".gemini-home") + monkeypatch.setattr( + gemini, "GEMINI_SETTINGS_PATH", tmp_path / ".gemini-home" / ".gemini" / "settings.json" + ) # Run from tmp_path so Gemini sees an untrusted folder — that mirrors # what users hit on a fresh checkout and exercises the trust + .env # discovery code paths that previously broke validation. @@ -853,6 +893,10 @@ def test_recovers_when_initial_token_empty( monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(gemini, "GEMINI_ENV_PATH", tmp_path / "ucode.env") monkeypatch.setattr(gemini, "GEMINI_BACKUP_PATH", tmp_path / "gemini-ucode-env.backup") + monkeypatch.setattr(gemini, "GEMINI_HOME_DIR", tmp_path / ".gemini-home") + monkeypatch.setattr( + gemini, "GEMINI_SETTINGS_PATH", tmp_path / ".gemini-home" / ".gemini" / "settings.json" + ) model = gemini_models[0] fake_db_dir = _make_reauth_fake_databricks(tmp_path / "fake_db", e2e_token) @@ -870,11 +914,4 @@ def test_recovers_when_initial_token_empty( "get_databricks_token may not be retrying after auth login." ) - env = gemini.build_runtime_env(e2e_workspace, model, recovered_token) - cmd = gemini.validate_cmd("gemini") - result = _run_agent(cmd, env=env, timeout=90) - combined = (result.stdout + result.stderr).strip() - assert result.returncode == 0 and combined, ( - f"Gemini failed after auth recovery: rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + assert _run_gemini_gateway_smoke(e2e_workspace, model, recovered_token).strip() diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index cea8830..884e663 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -297,6 +297,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(gemini, "GEMINI_ENV_PATH", tmp_path / "ucode.env") monkeypatch.setattr(gemini, "GEMINI_BACKUP_PATH", tmp_path / "gemini-ucode-env.backup") + monkeypatch.setattr(gemini, "GEMINI_HOME_DIR", tmp_path / ".gemini-home") + monkeypatch.setattr( + gemini, "GEMINI_SETTINGS_PATH", tmp_path / ".gemini-home" / ".gemini" / "settings.json" + ) # Run from tmp_path so Gemini sees an untrusted folder (the trust env # var built into build_runtime_env handles it). monkeypatch.chdir(tmp_path)