feat: discover OSS MLflow capabilities and model metadata

src/ucode/agents/__init__.py

@@ -387,6 +387,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
 
@@ -397,7 +398,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/claude.py

@@ -27,6 +27,7 @@
 from ucode.databricks import (
     build_auth_shell_command,
     build_tool_base_url,
+    claude_model_supports_1m,
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
@@ -83,11 +84,6 @@ def _resolve_web_search_model(state: dict) -> str | None:
 
 
 WEB_SEARCH_MCP_NAME = "web_search"
-# Matches both the AI Gateway form (`databricks-claude-opus-4-8`) and the UC
-# model-services form (`system.ai.claude-opus-4-8`).
-_CLAUDE_MODEL_RE = re.compile(
-    r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)-(\d+)(.*)$"
-)
 
 # Env keys the MLflow Stop hook reads to route traces. Written into the
 # settings `env` block alongside the hook itself.
@@ -390,19 +386,9 @@ def render_overlay(
 
 
 def _maybe_add_1m_suffix(model: str) -> str:
-    if model.endswith("[1m]"):
-        return model
-    match = _CLAUDE_MODEL_RE.match(model)
-    if not match:
+    if model.endswith("[1m]") or not claude_model_supports_1m(model):
         return model
-
-    family, major_raw, minor_raw, _ = match.groups()
-    major = int(major_raw)
-    minor = int(minor_raw)
-    should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or (
-        family == "sonnet" and (major, minor) >= (4, 6)
-    )
-    return f"{model}[1m]" if should_suffix else model
+    return f"{model}[1m]"
 
 
 def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool:

src/ucode/agents/opencode.py

@@ -6,6 +6,7 @@
 import signal
 import subprocess
 import threading
+from typing import cast
 
 from ucode.agent_updates import available_npm_package_update
 from ucode.config_io import (
@@ -20,7 +21,9 @@
     TOKEN_REFRESH_INTERVAL_SECONDS,
     build_opencode_base_urls,
     get_databricks_token,
+    gpt_model_token_limits,
     model_token_limits,
+    preferred_gpt_model,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -41,6 +44,7 @@
 PROVIDER_KEYS: list[list[str]] = [
     ["provider", "databricks-anthropic"],
     ["provider", "databricks-google"],
+    ["provider", "databricks-openai"],
     ["provider", "databricks-oss"],
 ]
 
@@ -51,7 +55,14 @@ 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-google/",
+            "databricks-openai/",
+            "databricks-oss/",
+        )
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
@@ -62,32 +73,96 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -
     if model in gemini_models:
         return f"databricks-google/{model}"
 
+    openai_models = opencode_models.get("openai") or []
+    if model in openai_models:
+        return f"databricks-openai/{model}"
+
     oss_models = opencode_models.get("oss") or []
     if model in oss_models:
         return f"databricks-oss/{model}"
 
     return model
 
 
-def _oss_model_overlay(model: str, ua_header: dict[str, str]) -> dict:
-    """Per-model overlay for an OSS model entry.
+_OSS_SAFE_LIMITS = {"context": 128_000, "output": 8_192}
+
+
+def _positive_int(value: object) -> int | None:
+    return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None
+
+
+def _oss_specs_by_id(raw_specs: object) -> dict[str, dict[str, object]]:
+    if not isinstance(raw_specs, list):
+        return {}
+    specs: dict[str, dict[str, object]] = {}
+    for raw_spec in raw_specs:
+        if not isinstance(raw_spec, dict):
+            continue
+        typed_spec = cast(dict[str, object], raw_spec)
+        model_id = typed_spec.get("id")
+        reasoning = typed_spec.get("reasoning")
+        context = typed_spec.get("context_window")
+        output = typed_spec.get("max_tokens")
+        valid_limits = all(
+            value is None or _positive_int(value) is not None for value in (context, output)
+        )
+        if (
+            isinstance(model_id, str)
+            and model_id
+            and isinstance(reasoning, bool)
+            and "context_window" in typed_spec
+            and "max_tokens" in typed_spec
+            and valid_limits
+            and model_id not in specs
+        ):
+            specs[model_id] = typed_spec
+    return specs
 
-    All OSS models carry the User-Agent header; models with known token limits
-    also pin `limit` (context + output) so OpenCode clamps `max_tokens` to a
-    value the gateway accepts. OpenCode's schema requires both fields together,
-    so the limits table always supplies both."""
+
+def _oss_model_overlay(
+    model: str, ua_header: dict[str, str], spec: dict[str, object] | None = None
+) -> dict:
+    """Per-model OSS overlay from discovered or static capabilities.
+
+    OpenCode requires context and output limits together. Every discovered spec
+    therefore receives a complete conservative pair. Missing specs retain
+    static GLM/Kimi metadata, and unknown no-spec models remain uncapped.
+    """
     overlay: dict = {"headers": ua_header}
-    limits = model_token_limits(model)
-    if limits is not None:
-        overlay["limit"] = limits
+    static_limits = model_token_limits(model)
+    context = _positive_int(spec.get("context_window")) if isinstance(spec, dict) else None
+    output = _positive_int(spec.get("max_tokens")) if isinstance(spec, dict) else None
+    if isinstance(spec, dict):
+        overlay["limit"] = {
+            "context": context
+            or (static_limits.get("context") if static_limits else _OSS_SAFE_LIMITS["context"]),
+            "output": output
+            or (static_limits.get("output") if static_limits else _OSS_SAFE_LIMITS["output"]),
+        }
+    elif static_limits is not None:
+        overlay["limit"] = static_limits
+
+    reasoning = spec.get("reasoning") if isinstance(spec, dict) else None
+    if isinstance(reasoning, bool):
+        overlay["reasoning"] = reasoning
     return overlay
 
 
+def _openai_model_overlay(model: str, ua_header: dict[str, str]) -> dict:
+    """Per-model Responses API options and explicit GPT token limits."""
+    return {
+        "headers": ua_header,
+        "limit": gpt_model_token_limits(model),
+        "options": {"useResponsesApi": True},
+    }
+
+
 def render_overlay(
     model: str,
     token: str,
     opencode_base_urls: dict[str, str],
     opencode_models: dict[str, list[str]],
+    oss_specs: list[dict] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for opencode.json."""
     auth_headers = {"Authorization": f"Bearer {token}"}
@@ -101,6 +176,7 @@ def render_overlay(
 
     anthropic_models = opencode_models.get("anthropic") or []
     gemini_models = opencode_models.get("gemini") or []
+    openai_models = opencode_models.get("openai") or []
     oss_models = opencode_models.get("oss") or []
 
     providers: dict = {}
@@ -136,15 +212,32 @@ def render_overlay(
             "models": {m: {"headers": ua_header} for m in gemini_models},
         }
         keys.append(["provider", "databricks-google"])
+    if openai_models:
+        # @ai-sdk/openai supports both the Responses API and the legacy
+        # chat-completions API. Databricks GPT-5 / GPT-5.6 / Codex models are
+        # Responses-only on /ai-gateway/codex/v1, so the per-model flag
+        # `useResponsesApi: true` lives in models.<m>.options where opencode
+        # reads it (provider-level options is read by the SDK only).
+        providers["databricks-openai"] = {
+            "npm": "@ai-sdk/openai",
+            "options": {
+                "baseURL": opencode_base_urls["openai"],
+                "apiKey": token,
+                "headers": auth_headers,
+            },
+            "models": {m: _openai_model_overlay(m, ua_header) for m in openai_models},
+        }
+        keys.append(["provider", "databricks-openai"])
     if oss_models:
+        specs_by_id = _oss_specs_by_id(oss_specs)
         providers["databricks-oss"] = {
             "npm": "@ai-sdk/openai",
             "options": {
                 "baseURL": opencode_base_urls["oss"],
                 "apiKey": token,
                 "headers": auth_headers,
             },
-            "models": {m: _oss_model_overlay(m, ua_header) for m in oss_models},
+            "models": {m: _oss_model_overlay(m, ua_header, specs_by_id.get(m)) for m in oss_models},
         }
         keys.append(["provider", "databricks-oss"])
 
@@ -174,6 +267,7 @@ def write_tool_config(
         token,
         opencode_base_urls,
         state.get("opencode_models") or {},
+        state.get("oss_model_specs") or [],
     )
     existing = read_json_safe(OPENCODE_CONFIG_PATH)
     providers = existing.get("provider")
@@ -234,6 +328,9 @@ def default_model(state: dict) -> str | None:
     anthropic = opencode_models.get("anthropic") or []
     if anthropic:
         return anthropic[0]
+    openai = preferred_gpt_model(opencode_models.get("openai") or [])
+    if openai:
+        return openai
     gemini = opencode_models.get("gemini") or []
     if gemini:
         return gemini[0]