From da40238de582e7084d327a3c97f845e27b900afa Mon Sep 17 00:00:00 2001
From: Kun Song <kun.song@databricks.com>
Date: Thu, 6 Aug 2026 16:33:13 +0200
Subject: [PATCH] Support custom Model Provider Services in Pi
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A Unity Catalog Model Provider Service with provider_type
EXTERNAL_MODEL_PROVIDER_TYPE_CUSTOM — a self-hosted, OpenAI-compatible
model — could not be used by any agent. `_TOOL_PROVIDER_TYPES` had no
`custom` entry and no `pi` key, so `resolve_provider_service` rejected it
with "which pi can't route to (supported: none)", and `ucode pi` had no
`--provider` flag to begin with.

Such a service is addressed differently from a vendor one: the
`Databricks-Model-Provider-Service` header selects it and the request
body's `model` carries the bare target name (the fully-qualified name
returns NOT_FOUND). The dialect each target serves is declared in
`targets[].native_api_types`, which the listing parsed and discarded.

- databricks.py: keep `native_api_types` per target; add
  `build_native_api_base_url` mapping a dialect to its gateway path,
  `custom_openai_chat_targets`, `CUSTOM_PROVIDER_TYPES`, and the
  usability/resolve guards. All agent-agnostic, so a second agent is
  additive.
- agents/pi.py: emit a `databricks-custom` provider (openai-completions
  on /ai-gateway/openai/v1) with the service header and the three compat
  flags that actually change pi's behavior for an unknown backend. Listed
  in PROVIDER_NAMES so it's stripped when a later launch drops
  --provider.
- agents/pi.py: fix `_refresh_token_once` raising under a provider-only
  launch. `_refresh_forever` swallowed the error, so the token silently
  stopped refreshing and the session died at expiry.
- cli.py: `--provider` on `ucode pi` (extracted to a shared
  ProviderOption), pi in the interactive picker, and
  `configure --provider-context-window`.

The Model Provider Service API exposes no context-window metadata, so a
value has to be assumed. The directions are not symmetric: understating
costs earlier compaction, while overstating is unrecoverable — pi
compacts to `contextWindow - reserveTokens`, so a window above the
server's real limit makes the compact-and-retry overflow again and the
turn ends. Default to a conservative 32768, print the assumption at
launch, and make it overridable.

Verified end to end against a live custom service: discovery, resolve,
generated config, a real tool round-trip through pi, token refresh with
zero Databricks models on the workspace, the context-window override, and
cleanup leaving a hand-added provider intact.

Co-authored-by: Isaac
---
 src/ucode/agents/__init__.py |  37 ++++--
 src/ucode/agents/pi.py       | 140 +++++++++++++++++++--
 src/ucode/cli.py             |  95 ++++++++++----
 src/ucode/databricks.py      |  89 ++++++++++++-
 tests/test_agent_pi.py       | 234 +++++++++++++++++++++++++++++++++++
 tests/test_agents_init.py    |  56 +++++++++
 tests/test_cli.py            |  71 +++++++++++
 tests/test_databricks.py     | 135 ++++++++++++++++++++
 tests/test_e2e.py            |  69 +++++++++++
 9 files changed, 876 insertions(+), 50 deletions(-)

diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py
index 3afd4fea..a9163555 100644
--- a/src/ucode/agents/__init__.py
+++ b/src/ucode/agents/__init__.py
@@ -19,6 +19,8 @@
 from ucode.config_io import ToolSpec
 from ucode.databricks import (
     BEDROCK_PROVIDER_TYPES,
+    CUSTOM_PROVIDER_TYPES,
+    custom_openai_chat_targets,
     get_databricks_token,
     install_ai_tools,
     install_databricks_cli,
@@ -282,16 +284,18 @@ def resolve_launch_model(
 
 def resolve_provider_models(
     tool: str, state: dict, provider: str | None
-) -> tuple[dict | None, str | None, bool]:
+) -> tuple[dict[str, str] | list[str] | None, str | None, bool]:
     """Validate ``provider`` for ``tool`` and return the model ids to pin.
 
     Returns ``(provider_models, error, relayed)``. ``provider_models`` is a
     ``{family: model_id}`` dict for a Bedrock-backed claude service (whose
-    provider-side ids must be pinned explicitly), or None for an Anthropic/
-    canonical service or when ``provider`` is None. ``relayed`` is True for a
-    credential-less Anthropic subscription relay, which the launch path wires
-    with the relayed overlay + refresh proxy. A non-None ``error`` means the
-    provider is invalid for the tool and the caller should not launch.
+    provider-side ids must be pinned explicitly), a list of target ids for a
+    custom-backed service (routed by header, with the bare target name as the
+    request body's `model`), or None for an Anthropic/canonical service or when
+    ``provider`` is None. ``relayed`` is True for a credential-less Anthropic
+    subscription relay, which the launch path wires with the relayed overlay +
+    refresh proxy. A non-None ``error`` means the provider is invalid for the
+    tool and the caller should not launch.
     """
     if not provider:
         return None, None, False
@@ -302,6 +306,8 @@ def resolve_provider_models(
     relayed = bool(service.get("relayed"))
     if service["provider_type"] in BEDROCK_PROVIDER_TYPES:
         return map_bedrock_claude_models(service.get("targets") or []), None, relayed
+    if service["provider_type"] in CUSTOM_PROVIDER_TYPES:
+        return custom_openai_chat_targets(service), None, relayed
     return None, None, relayed
 
 
@@ -310,7 +316,7 @@ def configure_tool(
     state: dict,
     model: str | None = None,
     provider: str | None = None,
-    provider_models: dict[str, str] | None = None,
+    provider_models: dict[str, str] | list[str] | None = None,
     relayed: bool = False,
     route_root_model: str | None = None,
 ) -> dict:
@@ -326,20 +332,29 @@ def configure_tool(
             state,
             model,
             provider=provider,
-            provider_models=provider_models,
+            provider_models=provider_models if isinstance(provider_models, dict) else None,
             relayed=relayed,
             route_root_model=route_root_model,
         )
+    elif tool == "pi":
+        # As with claude, a Model Provider Service pins no Databricks model —
+        # the service's targets are the models.
+        if not model and not provider:
+            raise RuntimeError(f"A {tool} model must be selected before configuration.")
+        result = pi.write_tool_config(
+            state,
+            model,
+            provider=provider,
+            provider_models=provider_models if isinstance(provider_models, list) else None,
+        )
     else:
-        # provider routing is claude/codex-only; every other tool needs a model.
+        # provider routing is claude/codex/pi-only; every other tool needs a model.
         if not model:
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         if tool == "gemini":
             result = gemini.write_tool_config(state, model)
         elif tool == "copilot":
             result = copilot.write_tool_config(state, model)
-        elif tool == "pi":
-            result = pi.write_tool_config(state, model)
         else:
             result = opencode.write_tool_config(state, model)
     # gemini/opencode/copilot/pi return (state, token); codex/claude return state
diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py
index e7c17609..c0bdf3c3 100644
--- a/src/ucode/agents/pi.py
+++ b/src/ucode/agents/pi.py
@@ -8,6 +8,12 @@
 - `databricks-openai`  (api: openai-responses)         → /ai-gateway/codex/v1
 - `databricks-gemini`  (api: google-generative-ai)     → /ai-gateway/gemini/v1beta
 
+A fourth provider, `databricks-custom` (api: openai-completions →
+/ai-gateway/openai/v1), is written only when launching through a `custom` Model
+Provider Service — a self-hosted, OpenAI-compatible model registered in Unity
+Catalog. The service is selected by the `Databricks-Model-Provider-Service`
+header and its targets are the models, so no Databricks model is pinned.
+
 Per-provider `compat` flags work around fields the gateway translators reject:
 
 - claude: `supportsEagerToolInputStreaming: false` — the Anthropic translator
@@ -15,6 +21,9 @@
   pi uses for every request. With this flag pi omits the per-tool field and
   sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header
   instead, which the gateway accepts.
+- custom: three flags turn off assumptions pi makes for openai.com. Pi's own
+  auto-detection already gets the rest right for an unrecognized base URL, so
+  restating them would be noise.
 
 OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via
 pi today — they live behind /ai-gateway/mlflow/v1 with per-model
@@ -42,7 +51,9 @@
     write_json_file,
 )
 from ucode.databricks import (
+    OPENAI_CHAT_NATIVE_API_TYPE,
     TOKEN_REFRESH_INTERVAL_SECONDS,
+    build_native_api_base_url,
     build_pi_base_urls,
     get_databricks_token,
 )
@@ -64,14 +75,33 @@
     "backup_path": PI_BACKUP_PATH,
 }
 
+CUSTOM_PROVIDER_NAME = "databricks-custom"
+
 PROVIDER_NAMES = (
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    # Written only under a custom Model Provider Service, but listed here
+    # unconditionally so it's stripped on every write: a later launch without
+    # `--provider` (or against a different service) must not leave a stale
+    # provider pointing at the old service's header.
+    CUSTOM_PROVIDER_NAME,
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
 
+# The Model Provider Service API exposes no context-window metadata, so ucode has
+# to assume one for a custom target. The two directions are not symmetric:
+# understating costs earlier compaction and shorter replies, while overstating is
+# unrecoverable — pi compacts to `contextWindow - reserveTokens`, so a window
+# above the server's real limit makes the compact-and-retry overflow again and
+# the turn ends with "Context overflow recovery failed after one
+# compact-and-retry attempt." Assume the conservative floor common to
+# self-hosted OpenAI-compatible servers and let users raise it with
+# `ucode configure --provider-context-window`.
+PROVIDER_CONTEXT_WINDOW = 32768
+PROVIDER_MAX_OUTPUT_TOKENS = 8192
+
 # Old provider names earlier ucode versions wrote; cleaned up on each write so
 # users don't end up with stale entries pointing at routes that 400.
 LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss")
@@ -82,15 +112,28 @@ def is_update_available() -> tuple[str, str] | None:
 
 
 def _resolve_model_selector(
-    model: str,
+    model: str | None,
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    provider_models: list[str] | None = None,
 ) -> str:
-    """Return a Pi model selector in `<provider>/<model>` form when possible."""
+    """Return a Pi model selector in `<provider>/<model>` form when possible.
+
+    The provider-qualified form matters: a bare target id would match Pi's own
+    built-in provider of the same name (e.g. `deepseek`) and fail with "No API
+    key found for deepseek" rather than routing through the gateway.
+    """
+    provider_models = provider_models or []
+    # Under a custom Model Provider Service no Databricks model is resolved, so
+    # the service's first target is the default.
+    if not model:
+        return f"{CUSTOM_PROVIDER_NAME}/{provider_models[0]}" if provider_models else ""
     for name in PROVIDER_NAMES:
         if model.startswith(f"{name}/"):
             return model
+    if model in provider_models:
+        return f"{CUSTOM_PROVIDER_NAME}/{model}"
     if model in claude_models.values():
         return f"databricks-claude/{model}"
     if model in codex_models:
@@ -101,14 +144,25 @@ def _resolve_model_selector(
 
 
 def render_overlay(
-    model: str,
+    model: str | None,
     token: str,
     pi_base_urls: dict[str, str],
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    *,
+    provider: str | None = None,
+    provider_models: list[str] | None = None,
+    provider_base_url: str | None = None,
+    context_window: int | None = None,
 ) -> tuple[dict, list[list[str]]]:
-    """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
+    """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.
+
+    ``provider`` is a `custom` Model Provider Service name; ``provider_models``
+    its routable targets. When both are set (plus a base URL for the dialect) a
+    `databricks-custom` provider is emitted alongside whichever Databricks
+    providers the workspace exposes.
+    """
     providers: dict = {}
     keys: list[list[str]] = [["model"]]
     # Pi expands header values that match an env var name. Our UA contains
@@ -150,8 +204,47 @@ def render_overlay(
             "models": [{"id": m} for m in gemini_models],
         }
         keys.append(["providers", "databricks-gemini"])
+    if provider and provider_models and provider_base_url:
+        window = context_window or PROVIDER_CONTEXT_WINDOW
+        providers[CUSTOM_PROVIDER_NAME] = {
+            "baseUrl": provider_base_url,
+            "api": "openai-completions",
+            "apiKey": token,
+            "authHeader": True,
+            # The header selects the service; the body's `model` carries the bare
+            # target name (the fully-qualified service name returns NOT_FOUND).
+            "headers": {**ua_headers, "Databricks-Model-Provider-Service": provider},
+            # Only the flags that change pi's behavior against an unknown
+            # OpenAI-compatible backend. Pi already auto-detects
+            # supportsReasoningEffort / supportsUsageInStreaming /
+            # supportsStrictMode correctly for an unrecognized base URL.
+            "compat": {
+                # Older OSS servers accept `max_tokens`, not the newer
+                # `max_completion_tokens` pi would otherwise send.
+                "maxTokensField": "max_tokens",
+                # The `developer` role is an OpenAI-ism; `system` is universal.
+                "supportsDeveloperRole": False,
+                # Many OSS servers reject unknown body fields such as `store`.
+                "supportsStore": False,
+            },
+            # No `reasoning`/`thinkingLevelMap`: the service exposes no capability
+            # metadata, and claiming reasoning support would make pi send
+            # `reasoning_effort` on every request, which a non-reasoning server
+            # can reject outright.
+            "models": [
+                {
+                    "id": target,
+                    "contextWindow": window,
+                    "maxTokens": min(PROVIDER_MAX_OUTPUT_TOKENS, window // 4),
+                }
+                for target in provider_models
+            ],
+        }
+        keys.append(["providers", CUSTOM_PROVIDER_NAME])
     overlay: dict = {
-        "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
+        "model": _resolve_model_selector(
+            model, claude_models, codex_models, gemini_models, provider_models
+        ),
     }
     if providers:
         overlay["providers"] = providers
@@ -160,9 +253,11 @@ def render_overlay(
 
 def write_tool_config(
     state: dict,
-    model: str,
+    model: str | None,
     token: str | None = None,
     *,
+    provider: str | None = None,
+    provider_models: list[str] | None = None,
     force_refresh: bool = False,
 ) -> tuple[dict, str]:
     backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH)
@@ -178,7 +273,22 @@ def write_tool_config(
         state.get("claude_models") or {},
         state.get("codex_models") or [],
         state.get("gemini_models") or [],
+        provider=provider,
+        provider_models=provider_models,
+        provider_base_url=build_native_api_base_url(
+            state["workspace"], OPENAI_CHAT_NATIVE_API_TYPE
+        ),
+        context_window=state.get("provider_context_window"),
     )
+    # Persist the resolved pair so the background refresh thread can rewrite the
+    # same provider config without re-resolving it, and clear it on a launch
+    # without a provider so the next refresh doesn't resurrect a stale one.
+    if provider and provider_models:
+        state["pi_provider"] = provider
+        state["pi_provider_models"] = list(provider_models)
+    else:
+        state.pop("pi_provider", None)
+        state.pop("pi_provider_models", None)
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
     if isinstance(providers, dict):
@@ -219,10 +329,22 @@ def default_model(state: dict) -> str | None:
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
-    model = default_model(state)
-    if not model:
+    # Under a Model Provider Service the workspace may expose no Databricks model
+    # at all — the service's targets are the models. Requiring one here would
+    # raise, and `_refresh_forever` swallows that, so the token would silently
+    # stop refreshing and the session would die when it expired.
+    provider = state.get("pi_provider")
+    provider_models = state.get("pi_provider_models") or []
+    model = None if provider else default_model(state)
+    if not model and not (provider and provider_models):
         raise RuntimeError("No Pi model is available on this workspace.")
-    _, token = write_tool_config(state, model, force_refresh=force_refresh)
+    _, token = write_tool_config(
+        state,
+        model,
+        provider=provider,
+        provider_models=provider_models or None,
+        force_refresh=force_refresh,
+    )
     return token
 
 
diff --git a/src/ucode/cli.py b/src/ucode/cli.py
index 58a4e4a3..71a06ebe 100644
--- a/src/ucode/cli.py
+++ b/src/ucode/cli.py
@@ -32,7 +32,11 @@
     launch as launch_agent,
 )
 from ucode.agents.codex import revert_legacy_shared_config
-from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH
+from ucode.agents.pi import (
+    PI_SETTINGS_BACKUP_PATH,
+    PI_SETTINGS_PATH,
+    PROVIDER_CONTEXT_WINDOW,
+)
 from ucode.config_io import restore_file, set_dry_run
 from ucode.databricks import (
     apply_pat_environment,
@@ -257,6 +261,7 @@ def configure_shared_state(
     skip_preflight: bool = False,
     fable_enabled: bool | None = None,
     databricks_ai_tools_enabled: bool | None = None,
+    provider_context_window: int | None = None,
 ) -> dict:
     """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state.
 
@@ -288,6 +293,8 @@ def configure_shared_state(
         use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace
     if fable_enabled is None:
         fable_enabled = bool(prior_state.get("fable_enabled")) and previous_workspace == workspace
+    if provider_context_window is None and previous_workspace == workspace:
+        provider_context_window = prior_state.get("provider_context_window")
     if databricks_ai_tools_enabled is None:
         # Opt-out: on by default. With no flag, keep this workspace's prior
         # choice but don't inherit another workspace's opt-out.
@@ -325,6 +332,13 @@ def configure_shared_state(
         state["fable_enabled"] = True
     else:
         state.pop("fable_enabled", None)
+    # Persist the custom-provider context window so both launches and the
+    # background token refresh reuse it; the API exposes no such metadata, so
+    # this override is the only way past the conservative default.
+    if provider_context_window:
+        state["provider_context_window"] = provider_context_window
+    else:
+        state.pop("provider_context_window", None)
     state["databricks_ai_tools_enabled"] = databricks_ai_tools_enabled
     state["base_urls"] = build_shared_base_urls(workspace)
 
@@ -485,6 +499,7 @@ def _configure_shared_workspace_states(
     use_pat: bool = False,
     fable_enabled: bool | None = None,
     databricks_ai_tools_enabled: bool | None = None,
+    provider_context_window: int | None = None,
 ) -> list[dict]:
     if not workspaces:
         raise RuntimeError("At least one workspace must be provided.")
@@ -499,6 +514,7 @@ def _configure_shared_workspace_states(
                 use_pat=use_pat,
                 fable_enabled=fable_enabled,
                 databricks_ai_tools_enabled=databricks_ai_tools_enabled,
+                provider_context_window=provider_context_window,
             )
         )
     return states
@@ -514,13 +530,15 @@ def _provider_summary(tool: str, state: dict) -> str:
 
 
 def _maybe_select_provider_service(tool: str, state: dict) -> dict:
-    """Interactively let the user route claude/codex through a Model Provider
+    """Interactively let the user route claude/codex/pi 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
-    matching provider services are found or the listing fails.
+    No-op for tools other than claude/codex/pi. Falls back to Databricks when no
+    matching provider services are found or the listing fails. The listing is
+    already filtered to services each tool can route to, so pi only ever sees
+    custom services with a routable target.
     """
-    if tool not in ("claude", "codex"):
+    if tool not in ("claude", "codex", "pi"):
         return state
     display = TOOL_SPECS[tool]["display"]
 
@@ -580,6 +598,7 @@ def configure_workspace_command(
     skip_validate: bool = False,
     fable_enabled: bool | None = None,
     databricks_ai_tools_enabled: bool | None = None,
+    provider_context_window: int | None = None,
 ) -> int:
     if tool is not None and selected_tools is not None:
         raise RuntimeError("Use either --agent or --agents, not both.")
@@ -599,6 +618,7 @@ def configure_workspace_command(
             use_pat=use_pat,
             fable_enabled=fable_enabled,
             databricks_ai_tools_enabled=databricks_ai_tools_enabled,
+            provider_context_window=provider_context_window,
         )
         state = states[0]
         state = configure_single_tool(tool, state)
@@ -637,6 +657,7 @@ def configure_workspace_command(
         use_pat=use_pat,
         fable_enabled=fable_enabled,
         databricks_ai_tools_enabled=databricks_ai_tools_enabled,
+        provider_context_window=provider_context_window,
     )
     state = states[0]
     save_state(state)
@@ -1317,6 +1338,16 @@ def _launch_tool(
             print_kv("Config", "workspace-managed")
         if provider:
             print_kv("Provider", provider)
+            if tool == "pi":
+                # Say the guess out loud: the Model Provider Service API exposes
+                # no context window, so an unannounced assumption would surface
+                # later as an unexplained mid-session overflow.
+                window = state.get("provider_context_window") or PROVIDER_CONTEXT_WINDOW
+                print_note(
+                    f"Assuming a {window}-token context window (the Model Provider Service "
+                    "API exposes none). Change it with "
+                    "`ucode configure --provider-context-window`."
+                )
         elif route_root_model:
             print_kv("Model", route_root_model)
         elif resolved_model:
@@ -1371,19 +1402,23 @@ def _launch_tool(
     ),
 ]
 
+# Route this launch through an external Model Provider Service rather than
+# Databricks-hosted models. Shared by every tool that can route to one.
+ProviderOption = Annotated[
+    str | None,
+    typer.Option(
+        "--provider",
+        help="Route through a Unity Catalog Model Provider Service "
+        "(<catalog>.<schema>.<name>). Skips Databricks model pinning; pass "
+        "before any `--` separator.",
+    ),
+]
+
 
 @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
 def codex_cmd(
     ctx: typer.Context,
-    provider: Annotated[
-        str | None,
-        typer.Option(
-            "--provider",
-            help="Route through a Unity Catalog Model Provider Service "
-            "(<catalog>.<schema>.<name>). Skips Databricks model pinning; pass "
-            "before any `--` separator.",
-        ),
-    ] = None,
+    provider: ProviderOption = None,
     skip_preflight: SkipPreflightOption = False,
     workspace: WorkspaceOption = None,
     enable_smart_routing_flag: Annotated[
@@ -1422,15 +1457,7 @@ def codex_cmd(
 @app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
 def claude_cmd(
     ctx: typer.Context,
-    provider: Annotated[
-        str | None,
-        typer.Option(
-            "--provider",
-            help="Route through a Unity Catalog Model Provider Service "
-            "(<catalog>.<schema>.<name>). Skips Databricks model pinning; pass "
-            "before any `--` separator.",
-        ),
-    ] = None,
+    provider: ProviderOption = None,
     skip_preflight: SkipPreflightOption = False,
     workspace: WorkspaceOption = None,
     enable_smart_routing_flag: Annotated[
@@ -1487,9 +1514,13 @@ def copilot_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False)
 
 
 @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
-def pi_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None:
+def pi_cmd(
+    ctx: typer.Context,
+    provider: ProviderOption = None,
+    skip_preflight: SkipPreflightOption = False,
+) -> None:
     """Launch Pi coding agent via Databricks."""
-    _launch_tool("pi", ctx, skip_preflight=skip_preflight)
+    _launch_tool("pi", ctx, provider=provider, skip_preflight=skip_preflight)
 
 
 @app.command("cursor", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
@@ -1602,6 +1633,16 @@ def configure(
             "--disable-databricks-ai-tools to opt out.",
         ),
     ] = None,
+    provider_context_window: Annotated[
+        int | None,
+        typer.Option(
+            "--provider-context-window",
+            help="Context window, in tokens, to assume for a custom Model Provider Service "
+            "(the API exposes none). Defaults to a conservative 32768. Setting this above "
+            "the endpoint's real limit makes long sessions fail unrecoverably, so raise it "
+            "only to a value the endpoint actually supports.",
+        ),
+    ] = None,
     mcp: Annotated[
         str | None,
         typer.Option(
@@ -1680,6 +1721,10 @@ def configure(
             agent = "claude"
         if enable_databricks_ai_tools is not None:
             skip_kwargs["databricks_ai_tools_enabled"] = enable_databricks_ai_tools
+        # Same inherit-on-None rule: only forward an explicitly passed window so a
+        # plain re-configure keeps the workspace's existing override.
+        if provider_context_window is not None:
+            skip_kwargs["provider_context_window"] = provider_context_window
         # Set True only in the fully-interactive branch below; gates the optional
         # MCP setup prompt so flag-driven / scripted runs are never interrupted.
         fully_interactive = False
diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index 0a004a15..2ed10fa1 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -1473,6 +1473,11 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = {
     "claude": ("anthropic", "amazon_bedrock"),
     "codex": ("openai",),
+    # pi speaks OpenAI chat completions to a `custom` service (a self-hosted,
+    # OpenAI-compatible model). Which dialect a target actually serves is
+    # declared per-target in `native_api_types`, so usability is narrowed
+    # further by `custom_openai_chat_targets`.
+    "pi": ("custom",),
 }
 
 # Provider types that expose Bedrock-style model ids (e.g.
@@ -1480,6 +1485,12 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str:
 # names, so ucode must pin them explicitly.
 BEDROCK_PROVIDER_TYPES: tuple[str, ...] = ("amazon_bedrock",)
 
+# Provider types backed by a caller-supplied endpoint rather than a known vendor
+# API. The service is selected by the `Databricks-Model-Provider-Service` header
+# and the request body's `model` carries the bare target name, so the target ids
+# must be pinned explicitly (the fully-qualified service name is not routable).
+CUSTOM_PROVIDER_TYPES: tuple[str, ...] = ("custom",)
+
 
 def tool_supports_provider_type(tool: str, provider_type: str) -> bool:
     """True when ``tool``'s API dialect can be backed by ``provider_type``."""
@@ -1559,7 +1570,12 @@ def list_model_provider_services(
 
 
 def _provider_service_entry(raw_service: object) -> dict | None:
-    """Normalize one listing entry, or None when it isn't usable."""
+    """Normalize one listing entry, or None when it isn't usable.
+
+    ``target_api_types`` maps each target model id to the request dialects it
+    declares (``native_api_types``), which is how ucode picks the gateway path
+    for a `custom` service. A target that declares none maps to an empty list.
+    """
     if not isinstance(raw_service, dict):
         return None
     # A bare isinstance narrows to dict[Never, Never], which rejects string keys.
@@ -1572,13 +1588,25 @@ def _provider_service_entry(raw_service: object) -> dict | None:
     raw_config = service.get("config")
     config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {}
     targets: list[str] = []
+    # Per-target `native_api_types` kept alongside `targets` (rather than folded
+    # into it) because callers like `map_bedrock_claude_models` take the plain id
+    # list. It's the only signal for which request dialect a target serves.
+    target_api_types: dict[str, list[str]] = {}
     raw_targets = config.get("targets")
     for target in raw_targets if isinstance(raw_targets, list) else []:
         if not isinstance(target, dict):
             continue
-        model_id = cast("dict[str, object]", target).get("model")
-        if isinstance(model_id, str) and model_id:
-            targets.append(model_id)
+        entry = cast("dict[str, object]", target)
+        model_id = entry.get("model")
+        if not isinstance(model_id, str) or not model_id:
+            continue
+        targets.append(model_id)
+        raw_api_types = entry.get("native_api_types")
+        target_api_types[model_id] = [
+            api_type
+            for api_type in (raw_api_types if isinstance(raw_api_types, list) else [])
+            if isinstance(api_type, str)
+        ]
     # Relayed = credential-less Anthropic (subscription relay). Only whether
     # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed
     # server-side, so both launch identically.
@@ -1589,6 +1617,7 @@ def _provider_service_entry(raw_service: object) -> dict | None:
         "name": full_name,
         "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None),
         "targets": targets,
+        "target_api_types": target_api_types,
         "allow_all_targets": bool(config.get("allow_all_targets")),
         "relayed": relayed,
     }
@@ -1643,13 +1672,16 @@ def service_usable_for_tool(tool: str, service: dict) -> bool:
     Beyond the provider-type match, a Bedrock service is only usable for claude
     if it exposes at least one Claude model in its targets — otherwise there's no
     routable model id to pin. (Anthropic services use canonical names, so any
-    match is usable.)
+    match is usable.) A custom service likewise needs at least one target serving
+    a dialect ucode can route.
     """
     provider_type = service.get("provider_type", "")
     if not tool_supports_provider_type(tool, provider_type):
         return False
     if provider_type in BEDROCK_PROVIDER_TYPES:
         return bool(map_bedrock_claude_models(service.get("targets") or []))
+    if provider_type in CUSTOM_PROVIDER_TYPES:
+        return bool(custom_openai_chat_targets(service))
     return True
 
 
@@ -1696,9 +1728,33 @@ def resolve_provider_service(
             f"Model provider service '{service_name}' exposes no Claude models — "
             f"add Claude targets to it or pick a different service."
         )
+    if provider_type in CUSTOM_PROVIDER_TYPES and not custom_openai_chat_targets(match):
+        return None, (
+            f"Model provider service '{service_name}' exposes no targets serving "
+            f"'{OPENAI_CHAT_NATIVE_API_TYPE}' — add one, or pick a different service."
+        )
     return match, None
 
 
+def custom_openai_chat_targets(service: dict) -> list[str]:
+    """Target model ids on a `custom` service that serve OpenAI chat completions.
+
+    A custom service routes by header and takes the bare target name as the
+    request body's `model`, so these ids are what an agent has to pin — the
+    fully-qualified service name is not routable.
+
+    Targets that declare a different dialect, or declare none at all, are
+    skipped: ucode only routes dialects it has a gateway path for, and guessing
+    turns a clear configure-time error into a 404 mid-session.
+    """
+    api_types = service.get("target_api_types") or {}
+    return [
+        target
+        for target in (service.get("targets") or [])
+        if OPENAI_CHAT_NATIVE_API_TYPE in (api_types.get(target) or [])
+    ]
+
+
 # Bedrock exposes Claude under provider-side ids like
 # `us.anthropic.claude-sonnet-4-6`, `global.anthropic.claude-opus-4-8`, or the
 # region-less `anthropic.claude-opus-4-8`. We map each service target to a
@@ -2445,6 +2501,29 @@ def _is_usage_table_access_error(exc: BaseException) -> bool:
 # URL builders (AI Gateway v2 only — no fallback to /serving-endpoints)
 # ---------------------------------------------------------------------------
 
+# A Model Provider Service target advertises the request dialect it serves in
+# `native_api_types`, and each dialect has one AI Gateway path that speaks it.
+# This mapping is per-dialect rather than per-tool because the service is
+# addressed the same way whichever agent is routing: the
+# `Databricks-Model-Provider-Service` header selects the service and the body's
+# `model` carries the bare target name. Add an entry to route another dialect.
+OPENAI_CHAT_NATIVE_API_TYPE = "openai/v1/chat/completions"
+_NATIVE_API_GATEWAY_PATHS: dict[str, str] = {
+    OPENAI_CHAT_NATIVE_API_TYPE: "/ai-gateway/openai/v1",
+}
+
+
+def build_native_api_base_url(workspace: str, native_api_type: str) -> str | None:
+    """Gateway base URL for a target's ``native_api_types`` entry.
+
+    Returns None when ucode has no path for that dialect, so callers can reject
+    the service with an actionable error instead of emitting a config that 404s
+    on the first request. The URL stops before the suffix the client appends
+    (pi's ``openai-completions`` appends ``/chat/completions``).
+    """
+    path = _NATIVE_API_GATEWAY_PATHS.get(native_api_type)
+    return f"{workspace}{path}" if path else None
+
 
 def build_tool_base_url(tool: str, workspace: str) -> str:
     if tool == "codex":
diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py
index 0afc5fb3..a75e3906 100644
--- a/tests/test_agent_pi.py
+++ b/tests/test_agent_pi.py
@@ -207,6 +207,133 @@ def test_unknown_model_passes_through_unprefixed(self):
         assert overlay["model"] == "custom/whatever"
 
 
+PROVIDER = "main.gateway.custom-svc"
+PROVIDER_BASE_URL = f"{WS}/ai-gateway/openai/v1"
+
+
+def _provider_overlay(model: str | None = None, token: str = "tok", **kwargs):
+    """render_overlay under a custom Model Provider Service.
+
+    Separate from `_overlay` so the provider-specific keyword arguments stay in
+    one place.
+    """
+    bundle = {**_empty(), **kwargs}
+    return pi.render_overlay(
+        model,
+        token,
+        _base_urls(),
+        bundle["claude_models"],
+        bundle["codex_models"],
+        bundle["gemini_models"],
+        provider=kwargs.get("provider", PROVIDER),
+        provider_models=kwargs.get("provider_models", ["deepseek-v4-flash"]),
+        provider_base_url=kwargs.get("provider_base_url", PROVIDER_BASE_URL),
+        context_window=kwargs.get("context_window"),
+    )
+
+
+class TestRenderOverlayCustomProvider:
+    def _provider(self, **kwargs) -> dict:
+        overlay, _ = _provider_overlay(**kwargs)
+        return overlay["providers"][pi.CUSTOM_PROVIDER_NAME]
+
+    def test_uses_openai_completions_on_the_openai_gateway_path(self):
+        provider = self._provider()
+        assert provider["api"] == "openai-completions"
+        assert provider["baseUrl"] == PROVIDER_BASE_URL
+
+    def test_routes_by_provider_service_header(self):
+        # The header selects the service; without it the gateway can't tell which
+        # provider to forward to.
+        headers = self._provider()["headers"]
+        assert headers["Databricks-Model-Provider-Service"] == PROVIDER
+        assert headers["User-Agent"].startswith("ucode/")
+
+    def test_compat_is_exactly_the_behavior_changing_flags(self):
+        # Pinned deliberately: pi already auto-detects supportsReasoningEffort,
+        # supportsUsageInStreaming and supportsStrictMode correctly for an
+        # unrecognized base URL, and `thinkingFormat` has a closed enum that
+        # "reasoning_effort" is not a member of. Restating either would be dead
+        # config at best and invalid at worst.
+        assert self._provider()["compat"] == {
+            "maxTokensField": "max_tokens",
+            "supportsDeveloperRole": False,
+            "supportsStore": False,
+        }
+
+    def test_targets_become_models_with_conservative_context_window(self):
+        assert self._provider()["models"] == [
+            {
+                "id": "deepseek-v4-flash",
+                "contextWindow": pi.PROVIDER_CONTEXT_WINDOW,
+                "maxTokens": pi.PROVIDER_MAX_OUTPUT_TOKENS,
+            }
+        ]
+
+    def test_explicit_context_window_is_honored(self):
+        models = self._provider(context_window=327680)["models"]
+        assert models[0]["contextWindow"] == 327680
+        assert models[0]["maxTokens"] == pi.PROVIDER_MAX_OUTPUT_TOKENS
+
+    def test_max_tokens_clamped_for_a_small_window(self):
+        models = self._provider(context_window=8192)["models"]
+        assert models[0]["maxTokens"] == 2048
+
+    def test_no_reasoning_claimed_without_capability_metadata(self):
+        # The service exposes no capability metadata, and claiming reasoning would
+        # make pi send `reasoning_effort` to a server that may reject it.
+        model = self._provider()["models"][0]
+        assert "reasoning" not in model
+        assert "thinkingLevelMap" not in model
+
+    def test_all_chat_targets_are_exposed(self):
+        provider = self._provider(provider_models=["model-a", "model-b"])
+        assert [m["id"] for m in provider["models"]] == ["model-a", "model-b"]
+
+    def test_databricks_providers_absent_when_workspace_has_no_models(self):
+        overlay, _ = _provider_overlay()
+        assert set(overlay["providers"]) == {pi.CUSTOM_PROVIDER_NAME}
+
+    def test_coexists_with_databricks_providers(self):
+        overlay, _ = _provider_overlay(claude_models={"sonnet": "claude-sonnet"})
+        assert set(overlay["providers"]) == {
+            pi.CUSTOM_PROVIDER_NAME,
+            "databricks-claude",
+        }
+
+    def test_provider_is_a_managed_key(self):
+        _, keys = _provider_overlay()
+        assert ["providers", pi.CUSTOM_PROVIDER_NAME] in keys
+
+    def test_provider_name_is_always_stripped_on_write(self):
+        # Membership in PROVIDER_NAMES is what removes a stale provider on a
+        # later launch without --provider.
+        assert pi.CUSTOM_PROVIDER_NAME in pi.PROVIDER_NAMES
+
+    def test_omitted_without_a_routable_base_url(self):
+        overlay, keys = _provider_overlay(provider_base_url=None)
+        assert pi.CUSTOM_PROVIDER_NAME not in overlay.get("providers", {})
+        assert ["providers", pi.CUSTOM_PROVIDER_NAME] not in keys
+
+
+class TestRenderOverlayCustomProviderSelector:
+    def test_defaults_to_first_target_when_no_model_resolved(self):
+        # Under a provider no Databricks model is resolved, so the selector has to
+        # come from the service's targets.
+        overlay, _ = _provider_overlay(None, provider_models=["model-a", "model-b"])
+        assert overlay["model"] == f"{pi.CUSTOM_PROVIDER_NAME}/model-a"
+
+    def test_prefixes_a_bare_target_id(self):
+        # A bare id would match pi's own built-in provider of the same name and
+        # fail with "No API key found for <provider>".
+        overlay, _ = _provider_overlay("deepseek-v4-flash")
+        assert overlay["model"] == f"{pi.CUSTOM_PROVIDER_NAME}/deepseek-v4-flash"
+
+    def test_preserves_already_prefixed_target(self):
+        overlay, _ = _provider_overlay(f"{pi.CUSTOM_PROVIDER_NAME}/deepseek-v4-flash")
+        assert overlay["model"] == f"{pi.CUSTOM_PROVIDER_NAME}/deepseek-v4-flash"
+
+
 class TestPiDefaultModel:
     def test_prefers_claude_opus(self):
         state = {"claude_models": {"opus": "o4", "sonnet": "s4", "haiku": "h4"}}
@@ -391,6 +518,113 @@ def test_pre_existing_settings_are_backed_up_before_first_write(self, tmp_path,
         assert merged["defaultProvider"] == "databricks-claude"
         assert merged["theme"] == "Default Dark"
 
+    def _write_with_provider(self, pi_mod, state):
+        with (
+            patch("ucode.agents.pi.get_databricks_token", return_value="tok"),
+            patch("ucode.agents.pi.save_state"),
+        ):
+            return pi_mod.write_tool_config(
+                state,
+                None,
+                token="tok",
+                provider=PROVIDER,
+                provider_models=["deepseek-v4-flash"],
+            )
+
+    def test_provider_write_emits_custom_provider_and_pins_settings(self, tmp_path, monkeypatch):
+        pi_mod, config_file, settings_file, _ = self._setup(tmp_path, monkeypatch)
+        state = self._state(claude_models={}, codex_models=[], gemini_models=[])
+
+        self._write_with_provider(pi_mod, state)
+
+        providers = json.loads(config_file.read_text())["providers"]
+        assert pi_mod.CUSTOM_PROVIDER_NAME in providers
+        settings = json.loads(settings_file.read_text())
+        assert settings["defaultProvider"] == pi_mod.CUSTOM_PROVIDER_NAME
+        assert settings["defaultModel"] == "deepseek-v4-flash"
+
+    def test_provider_write_persists_state_for_the_refresh_thread(self, tmp_path, monkeypatch):
+        pi_mod, _, _, _ = self._setup(tmp_path, monkeypatch)
+        state = self._state()
+
+        new_state, _ = self._write_with_provider(pi_mod, state)
+
+        assert new_state["pi_provider"] == PROVIDER
+        assert new_state["pi_provider_models"] == ["deepseek-v4-flash"]
+
+    def test_context_window_override_flows_from_state(self, tmp_path, monkeypatch):
+        pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch)
+        state = self._state(provider_context_window=327680)
+
+        self._write_with_provider(pi_mod, state)
+
+        providers = json.loads(config_file.read_text())["providers"]
+        models = providers[pi_mod.CUSTOM_PROVIDER_NAME]["models"]
+        assert models[0]["contextWindow"] == 327680
+
+    def test_later_launch_without_provider_clears_it(self, tmp_path, monkeypatch):
+        """A stale custom provider would keep routing to the old service's header,
+        so a plain `ucode pi` has to remove it and clear the persisted state."""
+        pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch)
+        state = self._state()
+        state, _ = self._write_with_provider(pi_mod, state)
+        # A hand-added provider must survive; only ucode's own are managed.
+        written = json.loads(config_file.read_text())
+        written["providers"]["user-provider"] = {"keep": True}
+        config_file.write_text(json.dumps(written), encoding="utf-8")
+
+        with (
+            patch("ucode.agents.pi.get_databricks_token", return_value="tok"),
+            patch("ucode.agents.pi.save_state"),
+        ):
+            state, _ = pi_mod.write_tool_config(state, "claude-sonnet", token="tok")
+
+        providers = json.loads(config_file.read_text())["providers"]
+        assert pi_mod.CUSTOM_PROVIDER_NAME not in providers
+        assert providers["user-provider"] == {"keep": True}
+        assert "pi_provider" not in state
+        assert "pi_provider_models" not in state
+
+
+class TestRefreshTokenOnce:
+    def _state(self, **overrides) -> dict:
+        state = {
+            "workspace": WS,
+            "base_urls": {"pi": _base_urls()},
+            "claude_models": {},
+            "codex_models": [],
+            "gemini_models": [],
+            "managed_configs": {},
+        }
+        state.update(overrides)
+        return state
+
+    def test_succeeds_under_a_provider_with_no_databricks_models(self, tmp_path, monkeypatch):
+        """The regression this guards: requiring a Databricks model here raised
+        RuntimeError, `_refresh_forever` swallowed it, and a provider-only session
+        silently stopped refreshing until the token expired mid-session."""
+        import ucode.agents.pi as pi_mod
+
+        monkeypatch.setattr(pi_mod, "PI_CONFIG_PATH", tmp_path / "models.json")
+        monkeypatch.setattr(pi_mod, "PI_SETTINGS_PATH", tmp_path / "settings.json")
+        monkeypatch.setattr(pi_mod, "PI_BACKUP_PATH", tmp_path / "backup.json")
+        monkeypatch.setattr(pi_mod, "PI_SETTINGS_BACKUP_PATH", tmp_path / "s-backup.json")
+        state = self._state(pi_provider=PROVIDER, pi_provider_models=["deepseek-v4-flash"])
+
+        with (
+            patch("ucode.agents.pi.get_databricks_token", return_value="fresh-tok"),
+            patch("ucode.agents.pi.save_state"),
+        ):
+            assert pi_mod._refresh_token_once(state) == "fresh-tok"
+
+    def test_raises_without_a_model_or_a_provider(self):
+        import pytest
+
+        import ucode.agents.pi as pi_mod
+
+        with pytest.raises(RuntimeError, match="No Pi model is available"):
+            pi_mod._refresh_token_once(self._state())
+
 
 class TestValidateAllToolsPiRollback:
     def test_failed_pi_validation_rolls_back_settings(self, tmp_path, monkeypatch):
diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py
index 97acf307..bb4f6973 100644
--- a/tests/test_agents_init.py
+++ b/tests/test_agents_init.py
@@ -325,6 +325,62 @@ def test_invalid_provider_returns_error(self, monkeypatch):
         assert error == "boom"
         assert relayed is False
 
+    def test_custom_returns_chat_completion_targets(self, monkeypatch):
+        # A custom service is routed by header with the bare target name as
+        # `model`, so the routable targets are returned as a list.
+        service = {
+            "provider_type": "custom",
+            "targets": ["chat-model", "responses-model"],
+            "target_api_types": {
+                "chat-model": ["openai/v1/chat/completions"],
+                "responses-model": ["openai/v1/responses"],
+            },
+        }
+        self._patch(monkeypatch, service, None)
+        models, error, relayed = agents_mod.resolve_provider_models("pi", self._STATE, "main.c.svc")
+        assert error is None
+        assert relayed is False
+        assert models == ["chat-model"]
+
+
+class TestConfigureToolPiProvider:
+    _STATE = {"workspace": "https://ws.databricks.com", "profile": None}
+
+    def test_forwards_provider_and_targets_without_a_model(self, monkeypatch):
+        captured: dict = {}
+
+        def fake_write(state, model, **kwargs):
+            captured["model"] = model
+            captured.update(kwargs)
+            return state, "tok"
+
+        monkeypatch.setattr(agents_mod.pi, "write_tool_config", fake_write)
+        agents_mod.configure_tool(
+            "pi", dict(self._STATE), None, provider="main.c.svc", provider_models=["chat-model"]
+        )
+        assert captured["model"] is None
+        assert captured["provider"] == "main.c.svc"
+        assert captured["provider_models"] == ["chat-model"]
+
+    def test_ignores_a_bedrock_style_dict(self, monkeypatch):
+        # claude's Bedrock path passes a {family: id} dict; pi must not treat it
+        # as a target list.
+        captured: dict = {}
+
+        def fake_write(state, model, **kwargs):
+            captured.update(kwargs)
+            return state, "tok"
+
+        monkeypatch.setattr(agents_mod.pi, "write_tool_config", fake_write)
+        agents_mod.configure_tool(
+            "pi", dict(self._STATE), None, provider="main.c.svc", provider_models={"opus": "x"}
+        )
+        assert captured["provider_models"] is None
+
+    def test_still_requires_a_model_without_a_provider(self):
+        with pytest.raises(RuntimeError, match="model must be selected"):
+            agents_mod.configure_tool("pi", dict(self._STATE), 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 eb0cad45..472253bc 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -198,6 +198,21 @@ def test_no_workspace_flag_leaves_current_workspace(self):
         assert result.exit_code == 0, result.output
         mock_set.assert_not_called()
 
+    def test_pi_accepts_provider_and_threads_it_through(self):
+        with patch("ucode.cli._launch_tool") as mock_launch:
+            result = runner.invoke(app, ["pi", "--provider", "main.gateway.custom-svc"])
+
+        assert result.exit_code == 0, result.output
+        assert mock_launch.call_args.kwargs["provider"] == "main.gateway.custom-svc"
+
+    def test_pi_declares_provider_option(self):
+        # Asserted on the registered option rather than `--help` text, which Rich
+        # truncates at the terminal width.
+        import typer
+
+        pi_command = typer.main.get_command(app).commands["pi"]
+        assert "--provider" in [opt for p in pi_command.params for opt in (p.opts or [])]
+
     def test_codex_enable_smart_routing_is_consumed_by_ucode(self):
         with patch("ucode.cli._launch_tool") as mock_launch:
             result = runner.invoke(app, ["codex", "--enable-smart-routing"])
@@ -1348,6 +1363,31 @@ def test_provider_picker_gated_by_interactive_path(self, monkeypatch):
         cli_mod.configure_workspace_command()
         assert picked_for == ["claude"]
 
+    def test_provider_picker_probes_for_pi_but_not_gemini(self, monkeypatch):
+        # pi can route to a custom Model Provider Service, so the interactive
+        # picker must actually probe for one instead of skipping it.
+        import ucode.cli as cli_mod
+
+        state = {**MINIMAL_STATE, "available_tools": []}
+        monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "tok")
+        monkeypatch.setattr(cli_mod, "save_state", lambda s: None)
+        monkeypatch.setattr(cli_mod, "set_provider_service", lambda s, tool, name: s)
+        probed: list[str] = []
+
+        def fake_list(tool, workspace, token):
+            probed.append(tool)
+            # No services on the workspace: falls back to Databricks.
+            return [], None
+
+        monkeypatch.setattr(cli_mod, "list_tool_provider_services", fake_list)
+
+        cli_mod._maybe_select_provider_service("pi", state)
+        assert probed == ["pi"]
+
+        # gemini has no dialect for any provider type, so it stays a pass-through.
+        cli_mod._maybe_select_provider_service("gemini", state)
+        assert probed == ["pi"]
+
     def test_unavailable_selected_tool_errors_before_configure(self, monkeypatch):
         import ucode.cli as cli_mod
 
@@ -1386,6 +1426,7 @@ def fake_configure_shared_state(
             use_pat=False,
             fable_enabled=None,
             databricks_ai_tools_enabled=None,
+            provider_context_window=None,
         ):
             configured_shared.append(
                 (workspace, profile, tuple(tools) if tools is not None else None, force_login)
@@ -1780,6 +1821,36 @@ def test_reconfigure_with_disable_fable_clears_opt_in(self, monkeypatch):
         assert "fable_enabled" not in state
         assert "fable" not in state["claude_models"]
 
+    def test_provider_context_window_persists(self, monkeypatch):
+        cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
+        state = cli_mod.configure_shared_state(
+            self.WS, profile="DEFAULT", provider_context_window=327680
+        )
+        assert state["provider_context_window"] == 327680
+
+    def test_provider_context_window_absent_by_default(self, monkeypatch):
+        # Absent means "use the conservative default" in agents/pi.py.
+        cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
+        state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT")
+        assert "provider_context_window" not in state
+
+    def test_launch_inherits_persisted_provider_context_window(self, monkeypatch):
+        # A launch passes None; the same workspace's stored override still applies,
+        # so the background token refresh keeps writing the right window.
+        cli_mod, *_ = self._stub_deps(
+            monkeypatch,
+            pat_token="dapi-pat",
+            existing_state={
+                "workspace": self.WS,
+                "profile": "DEFAULT",
+                "provider_context_window": 327680,
+            },
+        )
+
+        state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT")
+
+        assert state["provider_context_window"] == 327680
+
     def test_ai_tools_disable_persists(self, monkeypatch):
         cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
         state = cli_mod.configure_shared_state(
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index 5a3a4cac..29d7ab9d 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -390,6 +390,29 @@ class TestListModelProviderServices:
                     "targets": [{"model": "amazon.titan-text-express-v1"}],
                 },
             },
+            {
+                "name": "model-provider-services/main.schema3.custom-svc",
+                "config": {
+                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_CUSTOM",
+                    "targets": [
+                        {
+                            "model": "deepseek-v4-flash",
+                            "native_api_types": ["openai/v1/chat/completions"],
+                        },
+                        # A dialect ucode has no gateway path for, plus a target
+                        # that declares none at all — neither is routable.
+                        {"model": "responses-only", "native_api_types": ["openai/v1/responses"]},
+                        {"model": "undeclared"},
+                    ],
+                },
+            },
+            {
+                "name": "model-provider-services/main.schema3.custom-unroutable-svc",
+                "config": {
+                    "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_CUSTOM",
+                    "targets": [{"model": "mystery", "native_api_types": "not-a-list"}],
+                },
+            },
         ]
     }
 
@@ -403,6 +426,7 @@ def test_strips_prefix_and_tags_provider_type(self, monkeypatch):
             "name": "main.schema1.anthropic-svc",
             "provider_type": "anthropic",
             "targets": [],
+            "target_api_types": {},
             "allow_all_targets": False,
             "relayed": False,
         }
@@ -410,7 +434,31 @@ def test_strips_prefix_and_tags_provider_type(self, monkeypatch):
             "anthropic",
             "openai",
             "amazon_bedrock",
+            "custom",
+        }
+
+    def test_keeps_per_target_native_api_types(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
+        )
+        services, _ = db_mod.list_model_provider_services(WS, "token")
+        custom = next(s for s in services if s["name"] == "main.schema3.custom-svc")
+        assert custom["target_api_types"] == {
+            "deepseek-v4-flash": ["openai/v1/chat/completions"],
+            "responses-only": ["openai/v1/responses"],
+            # Declared nothing: an empty list, not a missing key.
+            "undeclared": [],
         }
+        # The plain id list keeps its shape for callers like map_bedrock_claude_models.
+        assert custom["targets"] == ["deepseek-v4-flash", "responses-only", "undeclared"]
+
+    def test_tolerates_non_list_native_api_types(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
+        )
+        services, _ = db_mod.list_model_provider_services(WS, "token")
+        unroutable = next(s for s in services if s["name"] == "main.schema3.custom-unroutable-svc")
+        assert unroutable["target_api_types"] == {"mystery": []}
 
     def test_flags_relayed_anthropic(self, monkeypatch):
         monkeypatch.setattr(
@@ -462,6 +510,67 @@ def test_codex_filters_to_openai(self, monkeypatch):
         names, _ = db_mod.list_tool_provider_services("codex", WS, "token")
         assert names == ["main.schema1.openai-svc"]
 
+    def test_pi_filters_to_routable_custom(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None)
+        )
+        names, _ = db_mod.list_tool_provider_services("pi", WS, "token")
+        # The custom service with no chat-completions target is hidden, so the
+        # interactive picker never offers something that would fail at launch.
+        assert names == ["main.schema3.custom-svc"]
+
+
+class TestBuildNativeApiBaseUrl:
+    def test_openai_chat_completions(self):
+        # Stops before the `/chat/completions` pi's openai-completions appends.
+        assert (
+            db_mod.build_native_api_base_url(WS, "openai/v1/chat/completions")
+            == f"{WS}/ai-gateway/openai/v1"
+        )
+
+    def test_unrouted_dialect_returns_none(self):
+        assert db_mod.build_native_api_base_url(WS, "anthropic/v1/messages") is None
+
+
+class TestCustomProviderSupport:
+    _CUSTOM = {
+        "name": "main.schema3.custom-svc",
+        "provider_type": "custom",
+        "targets": ["chat-model", "responses-model"],
+        "target_api_types": {
+            "chat-model": ["openai/v1/chat/completions"],
+            "responses-model": ["openai/v1/responses"],
+        },
+    }
+
+    def test_pi_supports_custom(self):
+        assert db_mod.tool_supports_provider_type("pi", "custom")
+
+    def test_claude_does_not_support_custom(self):
+        assert not db_mod.tool_supports_provider_type("claude", "custom")
+
+    def test_pi_does_not_support_openai(self):
+        # pi speaks chat completions to a custom service, not a vendor OpenAI one.
+        assert not db_mod.tool_supports_provider_type("pi", "openai")
+
+    def test_targets_filtered_to_chat_completions(self):
+        assert db_mod.custom_openai_chat_targets(self._CUSTOM) == ["chat-model"]
+
+    def test_no_chat_targets_yields_empty(self):
+        service = {**self._CUSTOM, "target_api_types": {"responses-model": []}}
+        assert db_mod.custom_openai_chat_targets(service) == []
+
+    def test_usable_for_pi_when_a_chat_target_exists(self):
+        assert db_mod.service_usable_for_tool("pi", self._CUSTOM)
+
+    def test_not_usable_for_pi_without_a_chat_target(self):
+        service = {
+            **self._CUSTOM,
+            "targets": ["responses-model"],
+            "target_api_types": {"responses-model": ["openai/v1/responses"]},
+        }
+        assert not db_mod.service_usable_for_tool("pi", service)
+
 
 class TestMapBedrockClaudeModels:
     def test_maps_families(self):
@@ -669,6 +778,32 @@ def test_bedrock_without_claude_rejected(self, monkeypatch):
         assert service is None
         assert "no Claude models" in error
 
+    def test_custom_for_pi_ok(self, monkeypatch):
+        self._patch(monkeypatch)
+        service, error = db_mod.resolve_provider_service(
+            "pi", "main.schema3.custom-svc", WS, "token"
+        )
+        assert error is None
+        assert service["provider_type"] == "custom"
+        assert db_mod.custom_openai_chat_targets(service) == ["deepseek-v4-flash"]
+
+    def test_custom_without_chat_target_rejected(self, monkeypatch):
+        self._patch(monkeypatch)
+        service, error = db_mod.resolve_provider_service(
+            "pi", "main.schema3.custom-unroutable-svc", WS, "token"
+        )
+        assert service is None
+        assert "openai/v1/chat/completions" in error
+
+    def test_pi_rejects_non_custom_type(self, monkeypatch):
+        self._patch(monkeypatch)
+        service, error = db_mod.resolve_provider_service(
+            "pi", "main.schema1.openai-svc", WS, "token"
+        )
+        assert service is None
+        assert "can't route to" in error
+        assert "supported: custom" in error
+
     def test_not_found_lists_usable(self, monkeypatch):
         self._patch(monkeypatch)
         service, error = db_mod.resolve_provider_service("claude", "main.x.missing", WS, "token")
diff --git a/tests/test_e2e.py b/tests/test_e2e.py
index 3f47d954..3f2562da 100644
--- a/tests/test_e2e.py
+++ b/tests/test_e2e.py
@@ -576,6 +576,14 @@ def _skip_if_no_permission(combined: str, provider: str) -> None:
         if "USE CONNECTION" in combined or "EXECUTE" in combined:
             pytest.skip(f"no permission on provider {provider}: {combined[:200]}")
 
+    @staticmethod
+    def _skip_if_upstream_unavailable(combined: str, provider: str) -> None:
+        """A custom service fronts a caller-operated endpoint, which can be down
+        independently of ucode. The gateway surfaces that as a 502/503, so treat it
+        as an environment skip rather than a ucode failure."""
+        if "502" in combined or "503" in combined or "temporarily unavailable" in combined:
+            pytest.skip(f"provider {provider} upstream is unavailable: {combined[:200]}")
+
     def test_launch_claude_through_provider(
         self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token
     ):
@@ -655,6 +663,67 @@ def test_launch_codex_through_provider(
             f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"
         )
 
+    def test_launch_pi_through_custom_provider(
+        self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token
+    ):
+        """A `custom` service exercises a path no vendor service covers: the
+        request is OpenAI chat completions on /ai-gateway/openai/v1, selected by
+        header, with the bare target name as the body's `model`."""
+        import ucode.config_io as config_io_mod
+        from ucode.agents import pi, resolve_provider_models
+
+        _require_binary("pi")
+        provider = self._first_service("pi", e2e_workspace, e2e_token)
+        state = {**e2e_state, "workspace": e2e_workspace}
+        # A custom service pins no Databricks model — its targets are the models,
+        # and only those declaring the chat-completions dialect are routable.
+        provider_models, error, _relayed = resolve_provider_models("pi", state, provider)
+        assert error is None, f"provider={provider} could not resolve models: {error}"
+        assert provider_models, f"provider={provider} exposed no routable targets"
+
+        monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path)
+        # Pi reads models.json below HOME/.pi/agent; point its runtime HOME and
+        # our writer at the same isolated tmp home.
+        pi_home = tmp_path / "pi-home"
+        pi_dir = pi_home / ".pi" / "agent"
+        monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home)
+        monkeypatch.setattr(pi, "PI_CONFIG_PATH", pi_dir / "models.json")
+        monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json")
+        monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "pi-models.backup.json")
+        monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json")
+
+        with pytest.MonkeyPatch().context() as mp:
+            mp.setattr("ucode.state.save_state", lambda s: None)
+            mp.setattr(
+                "ucode.agents.pi.get_databricks_token",
+                lambda ws, profile=None, **kwargs: e2e_token,
+            )
+            pi.write_tool_config(
+                state,
+                None,
+                token=e2e_token,
+                provider=provider,
+                provider_models=provider_models,
+            )
+
+        written = json.loads((pi_dir / "models.json").read_text())
+        custom = written["providers"][pi.CUSTOM_PROVIDER_NAME]
+        assert custom["headers"]["Databricks-Model-Provider-Service"] == provider
+        assert custom["baseUrl"] == f"{e2e_workspace}/ai-gateway/openai/v1"
+        # `--print` with no --model resolves through the pinned defaults, so the
+        # settings pin is what makes the launch below exercise the provider.
+        settings = json.loads((pi_dir / "settings.json").read_text())
+        assert settings["defaultProvider"] == pi.CUSTOM_PROVIDER_NAME
+
+        result = _run_agent(pi.validate_cmd("pi"), env=pi.build_runtime_env(e2e_token), timeout=120)
+        combined = (result.stdout + result.stderr).strip()
+        self._skip_if_no_permission(combined, provider)
+        self._skip_if_upstream_unavailable(combined, provider)
+        assert result.returncode == 0 and combined, (
+            f"provider={provider} rc={result.returncode} "
+            f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"
+        )
+
 
 class TestGeminiLaunch:
     """Run gemini against every available gemini model."""