Feat: Align Pi and OpenCode model family options

src/ucode/agents/__init__.py

@@ -369,17 +369,18 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool:
             bool(state.get("claude_models"))
             or bool(state.get("codex_models"))
             or bool(state.get("gemini_models"))
+            or bool(state.get("oss_models"))
         )
     return False
 
 
 _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = {
     "claude": ("claude",),
-    "opencode": ("claude", "gemini", "oss"),
+    "opencode": ("claude", "codex", "gemini", "oss"),
     "codex": ("codex",),
     "gemini": ("gemini",),
     "copilot": ("claude", "codex"),
-    "pi": ("claude", "codex", "gemini"),
+    "pi": ("claude", "codex", "gemini", "oss"),
 }
 
 

src/ucode/agents/opencode.py

@@ -40,6 +40,7 @@
 
 PROVIDER_KEYS: list[list[str]] = [
     ["provider", "databricks-anthropic"],
+    ["provider", "databricks-openai"],
     ["provider", "databricks-google"],
     ["provider", "databricks-oss"],
 ]
@@ -51,13 +52,24 @@ def is_update_available() -> tuple[str, str] | None:
 
 def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str:
     """Return an OpenCode model selector in provider/model form when possible."""
-    if model.startswith(("databricks-anthropic/", "databricks-google/", "databricks-oss/")):
+    if model.startswith(
+        (
+            "databricks-anthropic/",
+            "databricks-openai/",
+            "databricks-google/",
+            "databricks-oss/",
+        )
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
     if model in anthropic_models:
         return f"databricks-anthropic/{model}"
 
+    openai_models = opencode_models.get("openai") or []
+    if model in openai_models:
+        return f"databricks-openai/{model}"
+
     gemini_models = opencode_models.get("gemini") or []
     if model in gemini_models:
         return f"databricks-google/{model}"
@@ -100,6 +112,7 @@ def render_overlay(
     }
 
     anthropic_models = opencode_models.get("anthropic") or []
+    openai_models = opencode_models.get("openai") or []
     gemini_models = opencode_models.get("gemini") or []
     oss_models = opencode_models.get("oss") or []
 
@@ -125,6 +138,17 @@ def render_overlay(
             "models": dict.fromkeys(anthropic_models, anthropic_model_overlay),
         }
         keys.append(["provider", "databricks-anthropic"])
+    if openai_models:
+        providers["databricks-openai"] = {
+            "npm": "@ai-sdk/openai",
+            "options": {
+                "baseURL": opencode_base_urls["openai"],
+                "apiKey": token,
+                "headers": auth_headers,
+            },
+            "models": {m: {"headers": ua_header} for m in openai_models},
+        }
+        keys.append(["provider", "databricks-openai"])
     if gemini_models:
         providers["databricks-google"] = {
             "npm": "@ai-sdk/google",
@@ -234,6 +258,9 @@ def default_model(state: dict) -> str | None:
     anthropic = opencode_models.get("anthropic") or []
     if anthropic:
         return anthropic[0]
+    openai = opencode_models.get("openai") or []
+    if openai:
+        return openai[0]
     gemini = opencode_models.get("gemini") or []
     if gemini:
         return gemini[0]

src/ucode/agents/pi.py

@@ -1,25 +1,26 @@
 """Pi coding agent: writes ~/.pi/agent/models.json with Databricks-backed providers.
 
-Pi (https://pi.dev) is a multi-provider coding agent. We register three
+Pi (https://pi.dev) is a multi-provider coding agent. We register four
 providers in its `models.json`, each speaking the API dialect best suited to
 that family's gateway path:
 
 - `databricks-claude`  (api: anthropic-messages)       → /ai-gateway/anthropic
 - `databricks-openai`  (api: openai-responses)         → /ai-gateway/codex/v1
 - `databricks-gemini`  (api: google-generative-ai)     → /ai-gateway/gemini/v1beta
+- `databricks-oss`     (api: openai-responses)         → /ai-gateway/mlflow/v1
 
-Per-provider `compat` flags work around fields the gateway translators reject:
+One provider-specific `compat` flag works around a field the gateway translator rejects:
 
 - claude: `supportsEagerToolInputStreaming: false` — the Anthropic translator
   rejects `tools[].eager_input_streaming` on the streaming + tools path that
   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.
 
-OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via
-pi today — they live behind /ai-gateway/mlflow/v1 with per-model
-`max_tokens` caps that pi has no global way to honor without per-model
-config we don't currently maintain.
+- OSS models use the MLflow Responses route because its chat-completions stream
+  can end without a `finish_reason`, which Pi treats as an error. They also
+  carry per-model `contextWindow` and `maxTokens` from the shared token-limits
+  table.
 
 The bearer token is baked into the file and refreshed by a background thread
 while the session runs (same pattern as OpenCode/Copilot).
@@ -47,6 +48,7 @@
     build_pi_base_urls,
     classify_model_family,
     get_databricks_token,
+    model_token_limits,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -70,13 +72,14 @@
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    "databricks-oss",
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
 
 # 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")
+LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex")
 
 
 def is_update_available() -> tuple[str, str] | None:
@@ -88,6 +91,7 @@ def _resolve_model_selector(
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    oss_models: list[str],
 ) -> str:
     """Return a Pi model selector in `<provider>/<model>` form when possible."""
     for name in PROVIDER_NAMES:
@@ -99,16 +103,28 @@ def _resolve_model_selector(
         return f"databricks-openai/{model}"
     if model in gemini_models:
         return f"databricks-gemini/{model}"
+    if model in oss_models:
+        return f"databricks-oss/{model}"
     return model
 
 
+def _oss_model_entry(model: str) -> dict:
+    entry: dict = {"id": model}
+    limits = model_token_limits(model)
+    if limits is not None:
+        entry["contextWindow"] = limits["context"]
+        entry["maxTokens"] = limits["output"]
+    return entry
+
+
 def render_overlay(
     model: str,
     token: str,
     pi_base_urls: dict[str, str],
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    oss_models: list[str],
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
     providers: dict = {}
@@ -152,8 +168,20 @@ def render_overlay(
             "models": [{"id": m} for m in gemini_models],
         }
         keys.append(["providers", "databricks-gemini"])
+    if oss_models:
+        providers["databricks-oss"] = {
+            "baseUrl": pi_base_urls["oss"],
+            "api": "openai-responses",
+            "apiKey": token,
+            "authHeader": True,
+            "headers": ua_headers,
+            "models": [_oss_model_entry(m) for m in oss_models],
+        }
+        keys.append(["providers", "databricks-oss"])
     overlay: dict = {
-        "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
+        "model": _resolve_model_selector(
+            model, claude_models, codex_models, gemini_models, oss_models
+        ),
     }
     if providers:
         overlay["providers"] = providers
@@ -174,10 +202,11 @@ def write_tool_config(
         )
     pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"])
     managed_families = _managed_model_families(state)
-    claude_models, codex_models, gemini_models = managed_families or (
+    claude_models, codex_models, gemini_models, oss_models = managed_families or (
         state.get("claude_models") or {},
         state.get("codex_models") or [],
         state.get("gemini_models") or [],
+        state.get("oss_models") or [],
     )
     overlay, managed_keys = render_overlay(
         model,
@@ -186,6 +215,7 @@ def write_tool_config(
         claude_models,
         codex_models,
         gemini_models,
+        oss_models,
     )
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
@@ -213,7 +243,9 @@ def _write_settings(model_selector: str) -> None:
     write_json_file(PI_SETTINGS_PATH, merged)
 
 
-def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], list[str]] | None:
+def _managed_model_families(
+    state: dict,
+) -> tuple[dict[str, str], list[str], list[str], list[str]] | None:
     """Split a managed config's ``pi_models`` into the per-family inputs Pi's providers need.
 
     Pi builds one provider block per family, so a flat list has to be classified back out. Returns
@@ -226,6 +258,7 @@ def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], lis
     claude: dict[str, str] = {}
     codex: list[str] = []
     gemini: list[str] = []
+    oss: list[str] = []
     for model in managed:
         if not isinstance(model, str) or not model.strip():
             continue
@@ -236,13 +269,15 @@ def _managed_model_families(state: dict) -> tuple[dict[str, str], list[str], lis
             codex.append(model)
         elif family == "gemini":
             gemini.append(model)
-    if not (claude or codex or gemini):
+        elif family == "oss":
+            oss.append(model)
+    if not (claude or codex or gemini or oss):
         return None
-    return claude, codex, gemini
+    return claude, codex, gemini, oss
 
 
 def default_model(state: dict) -> str | None:
-    """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini.
+    """Prefer Claude opus → sonnet → haiku; then codex, gemini, OSS.
 
     A managed config's ``pi_default_model`` and ``pi_models`` both win outright: the former is
     the admin's chosen session start, the latter their allowlist. Workspace-wide discovery falls back.
@@ -260,7 +295,10 @@ def default_model(state: dict) -> str | None:
     if codex_models:
         return codex_models[0]
     gemini_models = state.get("gemini_models") or []
-    return gemini_models[0] if gemini_models else None
+    if gemini_models:
+        return gemini_models[0]
+    oss_models = state.get("oss_models") or []
+    return oss_models[0] if oss_models else None
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:

src/ucode/cli.py

@@ -122,9 +122,9 @@
 
 _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = {
     "claude": ("claude", "opencode", "copilot", "pi"),
-    "codex": ("codex", "copilot", "pi"),
+    "codex": ("codex", "opencode", "copilot", "pi"),
     "gemini": ("gemini", "opencode", "pi"),
-    "oss": ("opencode",),
+    "oss": ("opencode", "pi"),
 }
 
 
@@ -478,8 +478,10 @@ def configure_shared_state(
         fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools
     )
     want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools
-    want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools
-    want_oss = fetch_all or "opencode" in tools
+    want_codex = (
+        fetch_all or "codex" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools
+    )
+    want_oss = fetch_all or "opencode" in tools or "pi" in tools
 
     claude_reason: str | None = None
     gemini_reason: str | None = None
@@ -533,6 +535,8 @@ def configure_shared_state(
                 oss_models, oss_reason = ms_oss, ms_reason
         if claude_models:
             opencode_models["anthropic"] = list(claude_models.values())
+        if codex_models:
+            opencode_models["openai"] = codex_models
         if gemini_models:
             opencode_models["gemini"] = gemini_models
         if oss_models:

src/ucode/databricks.py

@@ -1493,7 +1493,9 @@ def discover_model_services(
     - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the
       newest matching ``system.ai.claude-*`` id (mirrors
       ``discover_claude_models``).
-    - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids.
+    - ``codex_models`` is the list of Responses-compatible ``system.ai.*gpt-*``
+      ids; ``gpt-oss`` is excluded because its gateway rejects the session and
+      prompt-caching fields sent by Pi and OpenCode.
     - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first.
     - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids.
 
@@ -1514,7 +1516,7 @@ def discover_model_services(
         if candidates:
             claude_models[family] = candidates[0]
 
-    codex_models = [m for m in ids if "gpt-" in m]
+    codex_models = [m for m in ids if "gpt-" in m and "gpt-oss-" not in m]
     gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key)
 
     oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)]
@@ -2466,7 +2468,8 @@ def discover_gemini_models(workspace: str, token: str) -> tuple[list[str], str |
 
 
 def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | None]:
-    return discover_endpoints_with_api_type(workspace, token, "openai/v1/responses")
+    models, reason = discover_endpoints_with_api_type(workspace, token, "openai/v1/responses")
+    return [model for model in models if "gpt-oss-" not in model], reason
 
 
 def fetch_gemini_models(workspace: str, token: str) -> list[str]:
@@ -2665,6 +2668,7 @@ def build_tool_base_url(tool: str, workspace: str) -> str:
 def build_opencode_base_urls(workspace: str) -> dict[str, str]:
     return {
         "anthropic": build_tool_base_url("claude", workspace) + "/v1",
+        "openai": build_tool_base_url("codex", workspace),
         "gemini": build_tool_base_url("gemini", workspace) + "/v1beta",
         "oss": f"{workspace}/ai-gateway/mlflow/v1",
     }
@@ -2675,17 +2679,15 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]:
     # path (verified end-to-end). Each `api` type appends its own path suffix:
     #
     # - anthropic-messages       appends `/v1/messages`
-    # - openai-responses         appends `/responses`
+    # - openai-responses         appends `/responses` (codex and OSS providers)
     # - google-generative-ai     appends `/v1beta/models/{id}:streamGenerateContent`
-    # - openai-completions       appends `/chat/completions`
     #
     # So the baseUrls below stop just before the suffix Pi will tack on.
-    # Compat flags applied per-provider in agents/pi.py; required for `oss`
-    # only (MLflow rejects `store` and `tools[].function.strict`).
     return {
         "claude": build_tool_base_url("claude", workspace),
         "openai": build_tool_base_url("codex", workspace),
         "gemini": build_tool_base_url("gemini", workspace) + "/v1beta",
+        "oss": f"{workspace}/ai-gateway/mlflow/v1",
     }