From e85295347708784be3b1a79f2e82e4e1e04cbe4c Mon Sep 17 00:00:00 2001 From: Kruttika Nadig Date: Sun, 30 Aug 2026 22:55:18 +0530 Subject: [PATCH 1/2] feat(gemini): route Gemini CLI through a Model Provider Service `ucode gemini` could only run Databricks-hosted models. Provider routing (`--provider`) was wired for claude and codex but not gemini, so a Gemini Enterprise Model Provider Service registered in Unity Catalog was unreachable. Extend the existing `--provider` path to gemini: add `gemini_enterprise` to the tool/provider-type map, emit the `Databricks-Model-Provider-Service` header via `GEMINI_CLI_CUSTOM_HEADERS`, and pin the service's target model (chosen with `--model` when it declares several). The routing header is re-emitted on token refresh, and `ucode configure` offers gemini the provider picker. Co-authored-by: Isaac --- src/ucode/agents/__init__.py | 38 +++++++++++++++++++++++- src/ucode/agents/gemini.py | 57 +++++++++++++++++++++++++++++------- src/ucode/cli.py | 36 +++++++++++++++++++---- src/ucode/databricks.py | 6 ++-- tests/test_agent_gemini.py | 16 ++++++++++ tests/test_agents_init.py | 43 +++++++++++++++++++++++++++ tests/test_databricks.py | 40 +++++++++++++++++++++++++ 7 files changed, 217 insertions(+), 19 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index d2480390..d72b68f7 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -326,6 +326,40 @@ def resolve_provider_models( 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 +) -> 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. Uses ``explicit_model`` (from ``--model``) when it names a target; + the sole target when the service declares exactly one; otherwise asks the + user to choose. Returns ``(model, error)``. + """ + 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 len(targets) == 1: + return targets[0], None + if not targets: + return None, f"Provider service '{provider}' exposes no models to launch." + return None, ( + f"Provider service '{provider}' exposes several models " + f"({', '.join(targets)}); pass --model to choose one." + ) + + def configure_tool( tool: str, state: dict, @@ -359,7 +393,9 @@ def configure_tool( if not model: raise RuntimeError(f"A {tool} model must be selected before configuration.") if tool == "gemini": - result = gemini.write_tool_config(state, model) + # Gemini routes through a provider by header (like codex), but must also + # pin the service's target model in the URL — `model` already holds it. + result = gemini.write_tool_config(state, model, provider=provider) elif tool == "copilot": result = copilot.write_tool_config(state, model) elif tool == "pi": diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index 5e54571a..6aedb490 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" @@ -135,12 +140,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), @@ -151,10 +162,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" @@ -168,16 +181,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 @@ -190,11 +208,26 @@ def default_model(state: dict) -> str | None: return gemini_models[0] if gemini_models else None +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 = parse_dotenv(GEMINI_ENV_PATH).get("GEMINI_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 @@ -207,11 +240,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( @@ -247,8 +281,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 c352e51b..60d2194d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -22,6 +22,7 @@ install_tool_binary, normalize_tool, provider_permission_error, + resolve_gemini_provider_model, resolve_launch_model, resolve_provider_models, validate_all_tools, @@ -776,10 +777,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"] @@ -1863,8 +1864,10 @@ def _launch_tool( try: tool = normalize_tool(tool_name) # A provider service routes by header and pins no model id, so pairing it with an explicit - # model is contradictory — reject rather than silently ignore one. - if model and provider: + # model is contradictory — reject rather than silently ignore one. Gemini is the exception: + # its provider still names a concrete target model in the URL, so `--model` selects which + # of the service's targets to launch on. + if model and provider and tool != "gemini": raise RuntimeError("Use either --model or --provider, not both.") # An explicit --workspace targets that workspace for this launch (and # auto-configures it if unseen), so `ucode claude --provider ... --workspace ...` @@ -2015,6 +2018,12 @@ def _launch_tool( # provider). Skip model resolution, which would otherwise fail when # the workspace has no matching Databricks models. resolved_model = None + 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 @@ -2455,12 +2464,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 3350e99c..1fd57bd8 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2060,11 +2060,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 5cd70fa3..9f60e860 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -59,6 +59,22 @@ 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" + ) + + def test_no_provider_omits_routing_header(self): + env = gemini.render_env_overlay(WS, "gemini-2", "tok") + assert "Databricks-Model-Provider-Service" not in env["GEMINI_CLI_CUSTOM_HEADERS"] + class TestBuildRuntimeEnv: def test_merges_os_environment(self): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index de40affb..da8825e4 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -340,6 +340,49 @@ def test_invalid_provider_returns_error(self, monkeypatch): assert relayed is False +class TestResolveGeminiProviderModel: + _STATE = {"workspace": "https://ws.databricks.com", "profile": None} + + def _patch(self, monkeypatch, service, error=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) + ) + + 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_selects_matching_target(self, monkeypatch): + self._patch( + monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-2.5-pro"]} + ) + model, error = agents_mod.resolve_gemini_provider_model( + self._STATE, "c.s.g", "gemini-2.5-pro" + ) + assert (model, error) == ("gemini-2.5-pro", 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_multiple_targets_without_model_asks_to_choose(self, monkeypatch): + self._patch( + monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-2.5-pro"]} + ) + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert model is None + assert "pass --model" in error + + def test_unresolvable_service_surfaces_error(self, monkeypatch): + self._patch(monkeypatch, None, "boom") + model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) + assert (model, error) == (None, "boom") + + class TestInstallToolBinary: def test_non_strict_returns_false_when_npm_missing(self, monkeypatch): monkeypatch.setattr("ucode.agents.shutil.which", lambda _: None) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 392e59da..5df0b0c7 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -843,6 +843,46 @@ 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"] + + def test_gemini_enterprise_rejected_for_claude(self, monkeypatch): + payload = { + "model_provider_services": [ + { + "name": "model-provider-services/main.schema1.gemini-svc", + "config": {"provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_GEMINI_ENTERPRISE"}, + } + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + service, error = db_mod.resolve_provider_service( + "claude", "main.schema1.gemini-svc", WS, "token" + ) + assert service is None + assert "can't route to" in error + class TestModelProviderFeatureUnavailable: def test_detects_feature_not_available(self): From 3ac5e55edfe4095c44c966ffe5217da1989b4ee2 Mon Sep 17 00:00:00 2001 From: Kruttika Nadig Date: Sun, 30 Aug 2026 23:42:11 +0530 Subject: [PATCH 2/2] test(gemini): trim to a focused minimum Keep one test per new code path: the MPS routing header, target-model resolution (happy + not-a-target), and gemini_enterprise resolving for gemini. Co-authored-by: Isaac --- tests/test_agent_gemini.py | 4 ---- tests/test_agents_init.py | 22 ---------------------- tests/test_databricks.py | 18 ------------------ 3 files changed, 44 deletions(-) diff --git a/tests/test_agent_gemini.py b/tests/test_agent_gemini.py index 9f60e860..e9668768 100644 --- a/tests/test_agent_gemini.py +++ b/tests/test_agent_gemini.py @@ -71,10 +71,6 @@ def test_provider_adds_routing_header_and_pins_target(self, monkeypatch): "Databricks-Model-Provider-Service:cat.sch.gemini-enterprise" ) - def test_no_provider_omits_routing_header(self): - env = gemini.render_env_overlay(WS, "gemini-2", "tok") - assert "Databricks-Model-Provider-Service" not in env["GEMINI_CLI_CUSTOM_HEADERS"] - class TestBuildRuntimeEnv: def test_merges_os_environment(self): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index da8825e4..ea0eb487 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -354,34 +354,12 @@ def test_sole_target_used_by_default(self, monkeypatch): 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_selects_matching_target(self, monkeypatch): - self._patch( - monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-2.5-pro"]} - ) - model, error = agents_mod.resolve_gemini_provider_model( - self._STATE, "c.s.g", "gemini-2.5-pro" - ) - assert (model, error) == ("gemini-2.5-pro", 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_multiple_targets_without_model_asks_to_choose(self, monkeypatch): - self._patch( - monkeypatch, {"name": "c.s.g", "targets": ["gemini-3.5-flash", "gemini-2.5-pro"]} - ) - model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) - assert model is None - assert "pass --model" in error - - def test_unresolvable_service_surfaces_error(self, monkeypatch): - self._patch(monkeypatch, None, "boom") - model, error = agents_mod.resolve_gemini_provider_model(self._STATE, "c.s.g", None) - assert (model, error) == (None, "boom") - class TestInstallToolBinary: def test_non_strict_returns_false_when_npm_missing(self, monkeypatch): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 5df0b0c7..25690528 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -865,24 +865,6 @@ def test_gemini_enterprise_ok_for_gemini(self, monkeypatch): assert service["provider_type"] == "gemini_enterprise" assert service["targets"] == ["gemini-3.5-flash"] - def test_gemini_enterprise_rejected_for_claude(self, monkeypatch): - payload = { - "model_provider_services": [ - { - "name": "model-provider-services/main.schema1.gemini-svc", - "config": {"provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_GEMINI_ENTERPRISE"}, - } - ] - } - monkeypatch.setattr( - db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) - ) - service, error = db_mod.resolve_provider_service( - "claude", "main.schema1.gemini-svc", WS, "token" - ) - assert service is None - assert "can't route to" in error - class TestModelProviderFeatureUnavailable: def test_detects_feature_not_available(self):