feat(pi): add validated MLflow OSS provider

Expose the GLM and Kimi coding-model cohort through Pi and OpenCode with shared token limits and reasoning metadata. Keep unsupported chat models out of discovery, including Inkling until gateway issue #215 is fixed, and retain the GPT-OSS Responses API routing guard.

Author: dgokeeffe

src/ucode/agents/__init__.py

@@ -356,6 +356,7 @@ 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
 
@@ -366,7 +367,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool:
     "codex": ("codex",),
     "gemini": ("gemini",),
     "copilot": ("claude", "codex"),
-    "pi": ("claude", "codex", "gemini"),
+    "pi": ("claude", "codex", "gemini", "oss"),
 }
 
 

src/ucode/agents/pi.py

@@ -1,12 +1,13 @@
 """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-mlflow`  (api: openai-completions)       → /ai-gateway/mlflow/v1
 
 Per-provider `compat` flags work around fields the gateway translators reject:
 
@@ -15,11 +16,17 @@
   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.
+- mlflow: `supportsStore: false` and `supportsStrictMode: false` — the MLflow
+  chat-completions gateway rejects OpenAI's `store` field and
+  `tools[].function.strict`.
+
+The `databricks-mlflow` provider carries the validated OSS coding models
+(GLM and Kimi) discovered upstream. Per model it sets
+`contextWindow`/`maxTokens` from `databricks.model_token_limits` and
+`reasoning` from `databricks.model_is_reasoning` (so Pi renders the gateway's
+streamed reasoning_content as thinking). Inkling is intentionally not offered
+until the gateway emits a terminal `finish_reason` on natural completion
+(issue #215).
 
 The bearer token is baked into the file and refreshed by a background thread
 while the session runs (same pattern as OpenCode/Copilot).
@@ -45,6 +52,8 @@
     TOKEN_REFRESH_INTERVAL_SECONDS,
     build_pi_base_urls,
     get_databricks_token,
+    model_is_reasoning,
+    model_token_limits,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -68,6 +77,7 @@
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    "databricks-mlflow",
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
@@ -86,6 +96,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:
@@ -97,16 +108,37 @@ 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-mlflow/{model}"
     return model
 
 
+def _pi_oss_model_entry(model_id: str) -> dict:
+    """Build a Pi mlflow model entry enriched from the shared limits/reasoning
+    tables: `reasoning:true` for reasoning models (Pi renders their streamed
+    reasoning_content as thinking), and `contextWindow`/`maxTokens` from
+    `model_token_limits`. Fields are omitted when unknown so Pi keeps its
+    default."""
+    entry: dict = {"id": model_id}
+    if model_is_reasoning(model_id):
+        entry["reasoning"] = True
+    limits = model_token_limits(model_id)
+    if limits:
+        if limits.get("context"):
+            entry["contextWindow"] = limits["context"]
+        if limits.get("output"):
+            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 = {}
@@ -150,8 +182,23 @@ def render_overlay(
             "models": [{"id": m} for m in gemini_models],
         }
         keys.append(["providers", "databricks-gemini"])
+    if oss_models:
+        providers["databricks-mlflow"] = {
+            "baseUrl": pi_base_urls["oss"],
+            "api": "openai-completions",
+            "apiKey": token,
+            "authHeader": True,
+            # MLflow chat-completions gateway rejects OpenAI's `store` field
+            # and per-tool `strict`. Pi omits both when these are false.
+            "compat": {"supportsStore": False, "supportsStrictMode": False},
+            "headers": ua_headers,
+            "models": [_pi_oss_model_entry(m) for m in oss_models],
+        }
+        keys.append(["providers", "databricks-mlflow"])
     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
@@ -178,6 +225,7 @@ def write_tool_config(
         state.get("claude_models") or {},
         state.get("codex_models") or [],
         state.get("gemini_models") or [],
+        state.get("oss_models") or [],
     )
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
@@ -206,7 +254,7 @@ def _write_settings(model_selector: str) -> None:
 
 
 def default_model(state: dict) -> str | None:
-    """Prefer Claude opus → sonnet → haiku; fall back to codex, gemini."""
+    """Prefer Claude opus → sonnet → haiku; fall back to codex, Gemini, then OSS."""
     claude_models = state.get("claude_models") or {}
     for family in ("opus", "sonnet", "haiku"):
         if claude_models.get(family):
@@ -215,7 +263,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

@@ -94,7 +94,7 @@
     "claude": ("claude", "opencode", "copilot", "pi"),
     "codex": ("codex", "copilot", "pi"),
     "gemini": ("gemini", "opencode", "pi"),
-    "oss": ("opencode",),
+    "oss": ("opencode", "pi"),
 }
 
 
@@ -372,7 +372,7 @@ def configure_shared_state(
     )
     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_oss = fetch_all or "opencode" in tools or "pi" in tools
 
     claude_reason: str | None = None
     gemini_reason: str | None = None

src/ucode/databricks.py

@@ -1161,10 +1161,24 @@ def build_auth_shell_command(
 # Databricks-managed foundation models under `system.ai`.
 _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai."
 
-# Supported OSS chat families, matched by name substring. Add an entry to
-# support a new family.
+# OSS families validated as coding models in ucode, matched by name substring.
+# Keep this as an explicit product allowlist rather than exposing every model on
+# the chat-completions route. Inkling remains excluded until gateway issue #215
+# is fixed; other families require coding-harness validation before inclusion.
 _OSS_MODEL_FAMILIES = ("kimi-", "glm-")
 
+# Non-chat services must never be offered to a chat agent if a future supported
+# family also uses one of these substrings.
+_OSS_NON_CHAT_SUBSTRINGS = ("embedding", "embed", "rerank")
+
+
+def _is_oss_chat_model(model_id: str) -> bool:
+    """True if the id matches an OSS chat family and isn't a non-chat service."""
+    if any(bad in model_id for bad in _OSS_NON_CHAT_SUBSTRINGS):
+        return False
+    return any(family in model_id for family in _OSS_MODEL_FAMILIES)
+
+
 # Claude model families ucode buckets, newest tier first. Each maps to a
 # Claude Code family alias (ANTHROPIC_DEFAULT_<FAMILY>_MODEL). Add an entry to
 # support a new family in both discovery paths (`claude-<family>-*` via the
@@ -1178,21 +1192,46 @@ def build_auth_shell_command(
 # config dialect. Both fields are provided because agents like OpenCode require
 # context and output together. Keyed by family substring; add an entry to bound
 # a new model.
+#
+# Output caps probed from the gateway 2026-07-16 (it 400s with "max_tokens (N)
+# cannot exceed <cap>"); context windows from each model's docs/description
+# (conservative when unstated). If the gateway raises a cap or ships a new
+# model, update this table.
 _MODEL_TOKEN_LIMITS: dict[str, dict[str, int]] = {
-    # GLM-4.6: 200k context, but the gateway caps output well below the model's
-    # native 128k — pin 25k so requests aren't rejected.
+    # Keep the version-specific entry before the family fallback: GLM 5.2 has
+    # materially higher probed gateway limits than earlier/unknown variants.
+    "glm-5-2": {"context": 1_000_000, "output": 65_536},
     "glm": {"context": 200_000, "output": 25_000},
+    "kimi": {"context": 128_000, "output": 65_536},
 }
 
+# Conservative fallback for a future variant that matches a validated family
+# but has no specific entry. Pinning a low output ceiling risks truncation, not
+# a gateway 400, so it is the safe failure direction.
+_OSS_FALLBACK_LIMITS = {"context": 128_000, "output": 8_192}
+
+# Validated families that emit reasoning. Pi renders their streamed
+# reasoning_content as thinking when the model entry sets reasoning:true.
+_OSS_REASONING_FAMILIES = ("glm", "kimi")
+
+
+def model_is_reasoning(model_id: str) -> bool:
+    """True if the OSS model reports reasoning output (family-matched)."""
+    return any(family in model_id for family in _OSS_REASONING_FAMILIES)
+
 
 def model_token_limits(model_id: str) -> dict[str, int] | None:
     """Return ``{"context": ..., "output": ...}`` limits for ``model_id``, or None.
 
-    Matches by family substring (e.g. any ``*glm*`` id). None means the model
-    has no known limits and the agent should not pin any."""
+    Prefers a specific `_MODEL_TOKEN_LIMITS` family entry (e.g. any ``*glm*``
+    id). Any other OSS chat model falls back to a conservative floor so it is
+    never offered uncapped (which would 400). None only for non-OSS ids, where
+    the agent should not pin any limit."""
     for family, limits in _MODEL_TOKEN_LIMITS.items():
         if family in model_id:
             return dict(limits)
+    if _is_oss_chat_model(model_id):
+        return dict(_OSS_FALLBACK_LIMITS)
     return None
 
 
@@ -1321,10 +1360,13 @@ def discover_model_services(
         if candidates:
             claude_models[family] = candidates[0]
 
-    codex_models = [m for m in ids if "gpt-" in m]
+    # `gpt-oss-*` also contains "gpt-" but is a chat-completions-only OSS model
+    # (served via /ai-gateway/mlflow/v1), NOT an openai-responses codex model —
+    # exclude it here so it isn't offered under the codex provider (which 400s).
+    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)]
+    oss_models = [m for m in ids if _is_oss_chat_model(m)]
 
     if not (claude_models or codex_models or gemini_models or oss_models):
         sample = ", ".join(ids[:5])
@@ -2219,6 +2261,7 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]:
         "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",
     }
 
 

tests/test_agent_opencode.py

@@ -98,15 +98,24 @@ def test_glm_gets_token_limits(self):
         overlay, _ = opencode.render_overlay("system.ai.glm-5-2", "tok", _base_urls(), models)
         glm = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"]
         # OpenCode's schema requires both context and output on `limit`.
-        assert glm["limit"] == {"context": 200000, "output": 25000}
+        # Probed 2026-07-16: glm-5-2 is 1M context / 65536 output.
+        assert glm["limit"] == {"context": 1_000_000, "output": 65_536}
 
-    def test_non_glm_oss_model_has_no_output_cap(self):
+    def test_kimi_gets_token_limits(self):
+        # kimi is now a capped OSS family (128k context / 65536 output).
         models = {"oss": ["system.ai.kimi-k2-7-code"]}
         overlay, _ = opencode.render_overlay(
             "system.ai.kimi-k2-7-code", "tok", _base_urls(), models
         )
         kimi = overlay["provider"]["databricks-oss"]["models"]["system.ai.kimi-k2-7-code"]
-        assert "limit" not in kimi
+        assert kimi["limit"] == {"context": 128_000, "output": 65_536}
+
+    def test_uncapped_oss_model_has_no_limit(self):
+        # A model outside the limits table gets no `limit` (client default).
+        models = {"oss": ["system.ai.mystery-7b"]}
+        overlay, _ = opencode.render_overlay("system.ai.mystery-7b", "tok", _base_urls(), models)
+        entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.mystery-7b"]
+        assert "limit" not in entry
 
     def test_token_in_api_key(self):
         models = {"anthropic": ["claude-sonnet"]}