diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..5ad691cc 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -338,9 +338,63 @@ def resolve_provider_models( # model selection server-side for that tier, so there's nothing to reconcile. if relayed: return None, None, relayed + # Only Claude pins per-family model ids. Codex ignores this map, and gemini resolves + # its target through resolve_gemini_provider_model instead — so mapping their targets + # through Claude-family logic would be meaningless (see docstring). + if tool != "claude": + return None, None, relayed return map_claude_family_models(service.get("targets") or []) or None, None, relayed +def resolve_gemini_provider_model( + state: dict, + provider: str, + explicit_model: str | None, + *, + service: dict | None = None, +) -> tuple[str | None, str | None]: + """Pick the Gemini model to pin for a provider-service launch. + + A Gemini Enterprise service routes by header but the request still names a + concrete model in the URL, so one of the service's declared targets must be + pinned. In precedence order: ``explicit_model`` (from ``--model``) when it + names a target; the model already pinned in the env file when it is still a + target (so a bare relaunch or reconfigure needn't re-pass ``--model``); the + sole target when the service declares exactly one; otherwise ask the user to + choose. Returns ``(model, error)``. + + Pass ``service`` to reuse an already-fetched service dict and skip the + control-plane lookup (the launch/configure paths hold one). + """ + if service is None: + token = get_databricks_token(state["workspace"], state.get("profile")) + service, error = resolve_provider_service("gemini", provider, state["workspace"], token) + if error or service is None: + return None, error or f"Model provider service '{provider}' was not found." + targets = [t for t in (service.get("targets") or []) if isinstance(t, str) and t] + if explicit_model: + if explicit_model in targets: + return explicit_model, None + available = ", ".join(targets) or "none" + return None, ( + f"Model '{explicit_model}' is not a target of provider service '{provider}'. " + f"Available: {available}." + ) + if not targets: + return None, f"Provider service '{provider}' exposes no models to launch." + # Reuse a previously pinned target so a bare relaunch/reconfigure keeps working without + # --model — but only when it is still one of the service's declared targets. + persisted = gemini.persisted_provider_model() + if persisted in targets: + return persisted, None + if len(targets) == 1: + return targets[0], None + return None, ( + f"Provider service '{provider}' exposes several models " + f"({', '.join(targets)}); pass --model to choose one." + ) + + def configure_tool( tool: str, state: dict, @@ -370,11 +424,12 @@ def configure_tool( custom_model=custom_model, ) else: - # provider routing is claude/codex-only; every other tool needs a model. + # Every tool in this branch needs a model — including gemini under a provider, + # which still pins the service's target model in the URL. if not model: raise RuntimeError(f"A {tool} model must be selected before configuration.") if tool == "gemini": - result = gemini.write_tool_config(state, model) + result = gemini.write_tool_config(state, model, provider=provider) elif tool == "copilot": result = copilot.write_tool_config(state, model) elif tool == "pi": @@ -457,6 +512,14 @@ def configure_single_tool(tool: str, state: dict) -> dict: def _configure_one(tool: str, state: dict, provider: str | None) -> dict: """Write one tool's config, routing through ``provider`` when set.""" if provider: + if tool == "gemini": + # Gemini pins a concrete target in the URL, so configure must resolve one now — + # unlike claude/codex, its config writer requires a model. This also validates the + # service in a single lookup (no resolve_provider_models family map for gemini). + model, error = resolve_gemini_provider_model(state, provider, None) + if error: + raise RuntimeError(error) + return configure_tool(tool, state, model, provider=provider) provider_models, error, relayed = resolve_provider_models(tool, state, provider) if error: raise RuntimeError(error) diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index a45850c9..5f14f1a3 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -25,7 +25,12 @@ build_tool_base_url, get_databricks_token, ) -from ucode.state import mark_tool_managed, save_state +from ucode.state import ( + get_provider_service, + mark_tool_managed, + save_state, + set_provider_service, +) from ucode.telemetry import agent_version, ucode_version GEMINI_CONFIG_DIR = Path.home() / ".gemini" @@ -115,12 +120,18 @@ def _ensure_local_settings_selected_type() -> None: write_json_file(GEMINI_SETTINGS_PATH, settings) -def render_env_overlay(workspace: str, model: str, token: str) -> dict[str, str]: +def render_env_overlay( + workspace: str, model: str, token: str, *, provider: str | None = None +) -> 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, # so a key named `User-Agent` overrides the default. Resolved via # upstream issue google-gemini/gemini-cli#10088. custom_headers = f"User-Agent:ucode/{ucode_version()} gemini/{agent_version('gemini')}" + if provider: + # A Model Provider Service routes by this header; the request still names + # the service's target model in `GEMINI_MODEL` (pinned by the launch path). + custom_headers += f",Databricks-Model-Provider-Service:{provider}" return { "GEMINI_MODEL": model, "GOOGLE_GEMINI_BASE_URL": build_tool_base_url("gemini", workspace), @@ -131,10 +142,12 @@ 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]: +def build_runtime_env( + workspace: str, model: str, token: str, *, provider: str | None = None +) -> dict[str, str]: _ensure_local_settings_selected_type() env = os.environ.copy() - env.update(render_env_overlay(workspace, model, token)) + env.update(render_env_overlay(workspace, model, token, provider=provider)) # 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" @@ -148,16 +161,21 @@ def write_tool_config( token: str | None = None, *, force_refresh: bool = False, + provider: str | None = None, ) -> tuple[dict, str]: backup_existing_file(GEMINI_ENV_PATH, GEMINI_BACKUP_PATH) if token is None: token = get_databricks_token( state["workspace"], state.get("profile"), force_refresh=force_refresh ) - overlay = render_env_overlay(state["workspace"], model, token) + overlay = render_env_overlay(state["workspace"], model, token, provider=provider) existing = parse_dotenv(GEMINI_ENV_PATH) existing.update(overlay) write_dotenv(GEMINI_ENV_PATH, existing) + if provider: + # Persist so the token-refresh thread and later bare `ucode gemini` re-emit + # the routing header; `--provider` on a launch overrides this saved choice. + state = set_provider_service(state, "gemini", provider) state = mark_tool_managed(state, "gemini", MANAGED_KEYS) save_state(state) return state, token @@ -170,11 +188,35 @@ def default_model(state: dict) -> str | None: return gemini_models[0] if gemini_models else None +def persisted_provider_model() -> str | None: + """The Gemini model currently pinned in the env file, if any. + + A provider launch writes the service's target here; reusing it lets a bare + ``ucode gemini`` relaunch (or reconfigure) run without re-passing ``--model``. + """ + return parse_dotenv(GEMINI_ENV_PATH).get("GEMINI_MODEL") + + +def _launch_model(state: dict, provider: str | None) -> str | None: + """The model this session runs on. + + Under a provider it is the service's target model, pinned into the env file + by ``write_tool_config`` (``default_model`` would return a Databricks id the + gateway can't resolve behind the provider header). Otherwise the usual default. + """ + if provider: + written = persisted_provider_model() + if written: + return written + return default_model(state) + + def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: - model = default_model(state) + provider = get_provider_service(state, "gemini") + model = _launch_model(state, provider) if not model: raise RuntimeError("No Gemini model is configured.") - _, token = write_tool_config(state, model, force_refresh=force_refresh) + _, token = write_tool_config(state, model, force_refresh=force_refresh, provider=provider) return token @@ -187,11 +229,12 @@ def _refresh_forever(state: dict, stop_event: threading.Event) -> None: def launch(state: dict, tool_args: list[str]) -> None: + provider = get_provider_service(state, "gemini") token = _refresh_token_once(state) - model = default_model(state) + model = _launch_model(state, provider) if not model: raise RuntimeError("No Gemini model is configured.") - env = build_runtime_env(state["workspace"], model, token) + env = build_runtime_env(state["workspace"], model, token, provider=provider) stop_event = threading.Event() refresher = threading.Thread( @@ -227,8 +270,9 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - model = default_model(state) + provider = get_provider_service(state, "gemini") + model = _launch_model(state, provider) if not model: raise RuntimeError("No Gemini model is configured.") token = get_databricks_token(workspace, state.get("profile")) - return build_runtime_env(workspace, model, token) + return build_runtime_env(workspace, model, token, provider=provider) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..7ae0b2e2 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -25,6 +25,7 @@ install_tool_binary, normalize_tool, provider_permission_error, + resolve_gemini_provider_model, resolve_launch_model, resolve_provider_models, validate_all_tools, @@ -777,10 +778,10 @@ def _maybe_select_provider_service(tool: str, state: dict) -> dict: """Interactively let the user route claude/codex through a Model Provider Service instead of Databricks models, and persist (or clear) the choice. - No-op for tools other than claude/codex. Falls back to Databricks when no + No-op for tools other than claude/codex/gemini. Falls back to Databricks when no matching provider services are found or the listing fails. """ - if tool not in ("claude", "codex"): + if tool not in ("claude", "codex", "gemini"): return state display = TOOL_SPECS[tool]["display"] @@ -1910,9 +1911,9 @@ def _launch_tool( forwarded_model = ( explicit_model_arg_value(ctx.args) if tool in {"claude", "codex"} else None ) - # `--model` is claude-only (no other launch command exposes it). Under a provider it selects - # which tier the service offers to launch on, rather than being rejected — see the provider - # branch below. + # `--model` is exposed by the claude and gemini launch commands. Under a provider it selects + # which of the service's targets/tiers to launch on, rather than being rejected — see the + # provider branch below. # An explicit --workspace targets that workspace for this launch (and # auto-configures it if unseen), so `ucode claude --provider ... --workspace ...` # works without a prior `ucode configure`. @@ -2014,9 +2015,11 @@ def _launch_tool( # Validate the provider service before launching — it must exist, be a # provider type this tool can route to (e.g. claude can't use an OpenAI # or Foundry service), and, for Bedrock, expose Claude models to pin. + # Gemini is exempt: it validates the service and resolves its target in a single + # lookup via resolve_gemini_provider_model (below), and uses no family model map. provider_models = None relayed = False - if provider: + if provider and tool != "gemini": provider_models, error, relayed = resolve_provider_models(tool, state, provider) if error: if managed is not None and provider == managed_provider_service(managed, tool): @@ -2051,6 +2054,12 @@ def _launch_tool( # Relayed services forward --model to Claude Code's own flag at launch (below), not env. if tool == "claude" and not relayed and (model or provider_models): route_root_model = resolve_provider_launch_model(model, provider_models or {}) + if tool == "gemini": + # Gemini is the exception: the request still names a concrete model + # in the URL, so pin one of the service's targets (--model or default). + resolved_model, gemini_error = resolve_gemini_provider_model(state, provider, model) + if gemini_error: + raise RuntimeError(gemini_error) else: # A managed default_model is the model the admin wants sessions to start on, so it goes # in as the explicit model rather than being applied afterwards: for codex the proto has @@ -2102,6 +2111,9 @@ def _launch_tool( print_kv("Model", forwarded_model) elif route_root_model: print_kv("Model", route_root_model) + # Gemini pins a concrete target under a provider (held in resolved_model). + elif resolved_model: + print_kv("Model", resolved_model) elif forwarded_model: print_kv("Model", forwarded_model) elif model and tool == "claude": @@ -2465,12 +2477,29 @@ def claude_cmd( @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def gemini_cmd( ctx: typer.Context, + provider: Annotated[ + str | None, + typer.Option( + "--provider", + help="Route through a Unity Catalog Model Provider Service " + "(..) that serves a Gemini model. Pass before any " + "`--` separator.", + ), + ] = None, + model: Annotated[ + str | None, + typer.Option( + "--model", + help="Model to launch on. Under --provider, selects which of the service's " + "target models to use. Pass before any `--` separator.", + ), + ] = None, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, ) -> None: """Launch Gemini CLI via Databricks.""" _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("gemini", ctx, skip_preflight=skip_preflight) + _launch_tool("gemini", ctx, provider=provider, model=model, skip_preflight=skip_preflight) @app.command( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f875a5c7..dcfe0acf 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2092,11 +2092,13 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: # Maps the gateway routing dialect a coding tool speaks to the Model Provider # Service `provider_type`s it can be backed by. claude speaks Anthropic's API, # which both the `anthropic` and `amazon_bedrock` provider types serve (Bedrock -# just exposes different model ids); codex speaks OpenAI's. Tags are the short -# form produced by `_provider_type_tag` (e.g. `amazon_bedrock`). +# just exposes different model ids); codex speaks OpenAI's; gemini speaks +# Google's, served by a Gemini Enterprise provider. Tags are the short form +# produced by `_provider_type_tag` (e.g. `amazon_bedrock`). _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = { "claude": ("anthropic", "amazon_bedrock"), "codex": ("openai",), + "gemini": ("gemini_enterprise",), } # Provider types that expose Bedrock-style model ids (e.g. diff --git a/tests/test_agent_gemini.py b/tests/test_agent_gemini.py index 1f3f67c5..8c6235a1 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -59,6 +59,18 @@ def test_sets_user_agent_via_custom_headers(self, monkeypatch): env = gemini.render_env_overlay(WS, "gemini-2", "tok") assert env["GEMINI_CLI_CUSTOM_HEADERS"] == "User-Agent:ucode/0.1.0 gemini/0.40.0" + def test_provider_adds_routing_header_and_pins_target(self, monkeypatch): + monkeypatch.setattr(gemini, "ucode_version", lambda: "0.1.0") + monkeypatch.setattr(gemini, "agent_version", lambda binary: "0.40.0") + env = gemini.render_env_overlay( + WS, "gemini-3.5-flash", "tok", provider="cat.sch.gemini-enterprise" + ) + assert env["GEMINI_MODEL"] == "gemini-3.5-flash" + assert env["GEMINI_CLI_CUSTOM_HEADERS"] == ( + "User-Agent:ucode/0.1.0 gemini/0.40.0," + "Databricks-Model-Provider-Service:cat.sch.gemini-enterprise" + ) + class TestBuildRuntimeEnv: def test_merges_os_environment(self): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 5977cc0c..8a0128ea 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -367,6 +367,118 @@ def test_invalid_provider_returns_error(self, monkeypatch): assert error == "boom" assert relayed is False + @pytest.mark.parametrize("tool", ["gemini", "codex"]) + def test_non_claude_pins_no_family_map(self, monkeypatch, tool): + # Only claude pins a per-family map; codex ignores it and gemini resolves its own + # target, so a non-claude service must not be run through Claude-family logic. + self._patch( + monkeypatch, + {"provider_type": "gemini_enterprise", "targets": ["gemini-3.5-flash"]}, + None, + ) + models, error, relayed = agents_mod.resolve_provider_models(tool, self._STATE, "c.s.svc") + assert (models, error, relayed) == (None, None, False) + + +class TestConfigureOneGeminiProvider: + _STATE = {"workspace": "https://ws.databricks.com", "profile": None} + + def test_gemini_provider_resolves_target_before_configure(self, monkeypatch): + # Regression: configuring gemini through a provider must resolve a target model rather + # than passing model=None into configure_tool (whose gemini branch requires one). + monkeypatch.setattr( + agents_mod, + "resolve_gemini_provider_model", + lambda state, provider, model, **kw: ("gemini-3.5-flash", None), + ) + captured = {} + + def _fake_configure_tool(tool, state, model=None, **kwargs): + captured["tool"] = tool + captured["model"] = model + captured["provider"] = kwargs.get("provider") + return state + + monkeypatch.setattr(agents_mod, "configure_tool", _fake_configure_tool) + agents_mod._configure_one("gemini", self._STATE, "c.s.g") + assert captured == {"tool": "gemini", "model": "gemini-3.5-flash", "provider": "c.s.g"} + + def test_gemini_provider_resolution_error_raises(self, monkeypatch): + monkeypatch.setattr( + agents_mod, + "resolve_gemini_provider_model", + lambda state, provider, model, **kw: (None, "pick a model"), + ) + with pytest.raises(RuntimeError, match="pick a model"): + agents_mod._configure_one("gemini", self._STATE, "c.s.g") + + +class TestResolveGeminiProviderModel: + _STATE = {"workspace": "https://ws.databricks.com", "profile": None} + + def _patch(self, monkeypatch, service, error=None, persisted=None): + monkeypatch.setattr(agents_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr( + agents_mod, "resolve_provider_service", lambda t, n, w, tok: (service, error) + ) + # Hermetic: never read the developer's real ~/.gemini/ucode.env. + monkeypatch.setattr(agents_mod.gemini, "persisted_provider_model", lambda: persisted) + + def test_sole_target_used_by_default(self, monkeypatch): + self._patch(monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash"]}) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert (model, error) == ("gemini-3.5-flash", None) + + def test_explicit_model_not_a_target_errors(self, monkeypatch): + self._patch(monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash"]}) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", "gpt-5") + assert model is None + assert "is not a target" in error + + def test_persisted_target_reused_for_multi_target(self, monkeypatch): + # A bare relaunch (no --model) of a multi-target service reuses the pinned model. + self._patch( + monkeypatch, + {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-3.5-pro"]}, + persisted="gemini-3.5-pro", + ) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert (model, error) == ("gemini-3.5-pro", None) + + def test_multi_target_without_choice_errors(self, monkeypatch): + # Multiple targets, nothing pinned, no --model → ask the user to choose. + self._patch( + monkeypatch, + {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-3.5-pro"]}, + persisted=None, + ) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert model is None + assert "exposes several models" in error + + def test_stale_persisted_ignored_falls_back_to_sole(self, monkeypatch): + # A pinned model that is no longer a target (e.g. after switching services) is ignored. + self._patch( + monkeypatch, + {"name": "c.s.g", "targets": ["gemini-3.5-flash"]}, + persisted="gemini-3.5-pro", + ) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert (model, error) == ("gemini-3.5-flash", None) + + def test_prefetched_service_skips_lookup(self, monkeypatch): + # Passing a service dict must not trigger a token fetch or control-plane lookup. + def _boom(*a, **k): # pragma: no cover - must never run + raise AssertionError("should not fetch when service is provided") + + monkeypatch.setattr(agents_mod, "get_databricks_token", _boom) + monkeypatch.setattr(agents_mod, "resolve_provider_service", _boom) + monkeypatch.setattr(agents_mod.gemini, "persisted_provider_model", lambda: None) + model, error = agents_mod.resolve_gemini_provider_model( + self._STATE, "c.s.g", None, service={"name": "c.s.g", "targets": ["gemini-3.5-flash"]} + ) + assert (model, error) == ("gemini-3.5-flash", None) + class TestInstallToolBinary: def test_non_strict_returns_false_when_npm_missing(self, monkeypatch): diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..16fb9a50 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -679,6 +679,39 @@ def test_provider_sets_transient_claude_launch_marker(self): assert mock_launch.call_args.args[1]["_claude_launch_provider"] == "main.default.anthropic" +class TestGeminiProviderLaunch: + @staticmethod + def _launch(monkeypatch, resolve_provider_models): + state = dict(MINIMAL_STATE) + monkeypatch.setattr("ucode.cli.ensure_bootstrap_dependencies", lambda *a, **k: None) + monkeypatch.setattr("ucode.cli.load_state", lambda: state) + monkeypatch.setattr("ucode.cli.ensure_provider_state", lambda t: state) + monkeypatch.setattr("ucode.cli.configure_shared_state", lambda *a, **k: state) + monkeypatch.setattr("ucode.cli._fetch_managed_config", lambda s: (None, False)) + monkeypatch.setattr("ucode.cli.resolve_provider_models", resolve_provider_models) + monkeypatch.setattr("ucode.cli.configure_tool", lambda *a, **k: state) + monkeypatch.setattr( + "ucode.cli.resolve_gemini_provider_model", + lambda s, p, m: ("gemini-3.5-flash", None), + ) + mock_launch = MagicMock() + monkeypatch.setattr("ucode.cli.launch_agent", mock_launch) + return runner.invoke(app, ["gemini", "--provider", "cat.schema.svc"]) + + def test_prints_resolved_target_model(self, monkeypatch): + # The concrete target gemini pins must show in the launch summary (not just the provider). + result = self._launch(monkeypatch, lambda t, s, p: (None, None, False)) + assert result.exit_code == 0, result.output + assert "Model: gemini-3.5-flash" in _strip_ansi(result.output) + + def test_skips_resolve_provider_models(self, monkeypatch): + # Gemini resolves its own target, so the claude-family lookup must not run (no double fetch). + mock_rpm = MagicMock(return_value=(None, None, False)) + result = self._launch(monkeypatch, mock_rpm) + assert result.exit_code == 0, result.output + mock_rpm.assert_not_called() + + class TestMcpSubcommands: def test_web_search_subcommand_help(self): result = runner.invoke(app, ["mcp", "web-search", "--help"]) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..9f0e40a4 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -886,6 +886,28 @@ def test_feature_unavailable(self, monkeypatch): assert service is None assert "not available" in error + def test_gemini_enterprise_ok_for_gemini(self, monkeypatch): + payload = { + "model_provider_services": [ + { + "name": "model-provider-services/main.schema1.gemini-svc", + "config": { + "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_GEMINI_ENTERPRISE", + "targets": [{"model": "gemini-3.5-flash"}], + }, + } + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + service, error = db_mod.resolve_provider_service( + "gemini", "main.schema1.gemini-svc", WS, "token" + ) + assert error is None + assert service["provider_type"] == "gemini_enterprise" + assert service["targets"] == ["gemini-3.5-flash"] + class TestModelProviderFeatureUnavailable: def test_detects_feature_not_available(self):