feat(gemini): route Gemini CLI through a Model Provider Service

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":

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)

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 "
+            "(<catalog>.<schema>.<name>) 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(

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.

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):

tests/test_agents_init.py

@@ -340,6 +340,27 @@ 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_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
+
+
 class TestInstallToolBinary:
     def test_non_strict_returns_false_when_npm_missing(self, monkeypatch):
         monkeypatch.setattr("ucode.agents.shutil.which", lambda _: None)

tests/test_databricks.py

@@ -843,6 +843,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):