From fa125988c47ce292b1488638a5e2615f1aca3402 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Sat, 25 Jul 2026 08:51:30 +1000 Subject: [PATCH 1/8] fix(pi): point pi at config dir via PI_CODING_AGENT_DIR instead of redirecting HOME Pi honors the PI_CODING_AGENT_DIR env var to resolve its config directory (~/.pi/agent), so redirecting /Users/david.okeeffe to APP_DIR/pi-home was unnecessary. The HOME redirect broke macOS keychain default resolution under ucode: the Security framework looks for the login keychain under the redirected HOME, finds none, and security default-keychain returns 'A default keychain could not be found'. As a result gh auth, the git credential helper, and any keychain-backed tool failed inside pi. Setting PI_CODING_AGENT_DIR to the existing PI_CONFIG_DIR preserves config isolation (models.json/settings.json/sessions still land under APP_DIR/pi-home/.pi/agent) while leaving /Users/david.okeeffe as the user's real home, so the login keychain stays discoverable. Tests: the two pi e2e sites monkeypatched PI_UCODE_HOME/PI_CONFIG_PATH to redirect pi at a tmp home. They now also patch PI_CONFIG_DIR (read by build_runtime_env) and PI_SETTINGS_PATH/PI_SETTINGS_BACKUP_PATH (previously masked because the HOME redirect made pi read settings from the un-patched real APP_DIR path). --- src/ucode/agents/pi.py | 6 +++--- tests/test_agent_pi.py | 8 ++++++-- tests/test_e2e.py | 7 +++++-- tests/test_e2e_user_agent.py | 3 +++ 4 files changed, 17 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c17609..b27e427c 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -1,4 +1,4 @@ -"""Pi coding agent: writes ~/.pi/agent/models.json with Databricks-backed providers. +"""Pi coding agent: writes a ucode-private models.json with Databricks-backed providers. Pi (https://pi.dev) is a multi-provider coding agent. We register three providers in its `models.json`, each speaking the API dialect best suited to @@ -108,7 +108,7 @@ def render_overlay( codex_models: list[str], gemini_models: list[str], ) -> tuple[dict, list[list[str]]]: - """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.""" + """Return (overlay, managed_key_paths) for Pi's private agent config.""" providers: dict = {} keys: list[list[str]] = [["model"]] # Pi expands header values that match an env var name. Our UA contains @@ -237,7 +237,7 @@ def _refresh_forever(state: dict, stop_event: threading.Event) -> None: def build_runtime_env(token: str) -> dict[str, str]: env = os.environ.copy() env["OAUTH_TOKEN"] = token - env["HOME"] = str(PI_UCODE_HOME) + env["PI_CODING_AGENT_DIR"] = str(PI_CONFIG_DIR) return env diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb3..0d3be801 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -240,9 +240,13 @@ def test_sets_oauth_token(self): env = pi.build_runtime_env("tok") assert env["OAUTH_TOKEN"] == "tok" - def test_sets_ucode_home(self): + def test_sets_private_agent_dir_without_replacing_home(self, monkeypatch): + monkeypatch.setenv("HOME", "/real-user-home") + env = pi.build_runtime_env("tok") - assert env["HOME"] == str(pi.PI_UCODE_HOME) + + assert env["PI_CODING_AGENT_DIR"] == str(pi.PI_CONFIG_DIR) + assert env["HOME"] == "/real-user-home" class TestPiValidateCmd: diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 30ce2876..80eb04a0 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -903,14 +903,17 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa pytest.skip("No Pi-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - # Pi reads models.json below HOME/.pi/agent. Point both pi's runtime - # HOME and our writer at the same isolated tmp home. + # Point PI_CODING_AGENT_DIR and ucode's config writer at the same + # isolated directory without changing the process HOME. pi_home = tmp_path / "pi-home" pi_dir = pi_home / ".pi" / "agent" config_path = pi_dir / "models.json" backup_path = tmp_path / "pi-models.backup.json" monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) + monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) failures = [] diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index 884e663d..e6cec214 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -325,7 +325,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) + monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "pi.backup.json") state = { From 65b4609f28ce4515a86e38caa7806bb1d634d090 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:16:37 +1000 Subject: [PATCH 2/8] 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. --- src/ucode/agents/__init__.py | 3 +- src/ucode/agents/pi.py | 69 ++++++++++++++++++++++++++---- src/ucode/cli.py | 4 +- src/ucode/databricks.py | 59 ++++++++++++++++++++++---- tests/test_agent_opencode.py | 15 +++++-- tests/test_agent_pi.py | 60 +++++++++++++++++++++++++- tests/test_agents_init.py | 17 ++++++++ tests/test_cli.py | 11 +++++ tests/test_databricks.py | 82 ++++++++++++++++++++++++++++++------ 9 files changed, 282 insertions(+), 38 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 88cb9e00..71c0665b 100644 --- a/src/ucode/agents/__init__.py +++ b/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"), } diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c17609..7f782932 100644 --- a/src/ucode/agents/pi.py +++ b/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 `/` form when possible.""" for name in PROVIDER_NAMES: @@ -97,9 +108,29 @@ 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, @@ -107,6 +138,7 @@ def render_overlay( 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: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 975eb0b7..57e75bf1 100644 --- a/src/ucode/cli.py +++ b/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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f9e3f1b6..9e6d331a 100644 --- a/src/ucode/databricks.py +++ b/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__MODEL). Add an entry to # support a new family in both discovery paths (`claude--*` 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 "); 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", } diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0960c640..3cc7348b 100644 --- a/tests/test_agent_opencode.py +++ b/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"]} diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb3..4c5444ee 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -17,6 +17,7 @@ def _base_urls() -> dict[str, str]: "claude": f"{WS}/ai-gateway/anthropic", "openai": f"{WS}/ai-gateway/codex/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -26,6 +27,7 @@ def _empty() -> dict: "claude_models": {}, "codex_models": [], "gemini_models": [], + "oss_models": [], } @@ -39,6 +41,7 @@ def _overlay(model: str, token: str = "tok", **kwargs): bundle["claude_models"], bundle["codex_models"], bundle["gemini_models"], + bundle["oss_models"], ) @@ -81,17 +84,57 @@ def test_gemini_provider_uses_google_generative_ai(self): assert provider["api"] == "google-generative-ai" assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta" - def test_all_three_providers_when_all_present(self): + def test_mlflow_provider_uses_openai_completions(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + provider = overlay["providers"]["databricks-mlflow"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{WS}/ai-gateway/mlflow/v1" + assert provider["compat"] == {"supportsStore": False, "supportsStrictMode": False} + + def test_no_mlflow_provider_when_no_oss_models(self): + overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"]) + assert "databricks-mlflow" not in overlay.get("providers", {}) + + def test_all_four_providers_when_all_present(self): overlay, _ = _overlay( "claude-sonnet", claude_models={"sonnet": "claude-sonnet"}, codex_models=["gpt-5"], gemini_models=["gemini-2"], + oss_models=["system.ai.glm-5-2"], ) assert set(overlay["providers"].keys()) == { "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-mlflow", + } + + +class TestRenderOverlayOssEnrichment: + """OSS mlflow model entries carry reasoning + contextWindow + maxTokens + from the shared databricks.model_token_limits / model_is_reasoning tables.""" + + def test_reasoning_model_enriched(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry["id"] == "system.ai.glm-5-2" + assert entry["reasoning"] is True + assert entry["contextWindow"] == 1_000_000 + assert entry["maxTokens"] == 65_536 + + def test_unvalidated_model_has_no_inferred_metadata(self): + # Discovery does not offer this model; even if supplied directly, Pi + # must not infer capabilities for an unvalidated coding model. + overlay, _ = _overlay("system.ai.inkling", oss_models=["system.ai.inkling"]) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry == {"id": "system.ai.inkling"} + + def test_unknown_oss_model_bare(self): + # No limits/reasoning table entry -> only id, client keeps defaults. + overlay, _ = _overlay("system.ai.mystery-7b", oss_models=["system.ai.mystery-7b"]) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.mystery-7b" } @@ -193,6 +236,10 @@ def test_prefixes_gemini_model(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) assert overlay["model"] == "databricks-gemini/gemini-2" + def test_prefixes_oss_model(self): + overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"]) + assert overlay["model"] == "databricks-mlflow/system.ai.glm-5-2" + def test_preserves_already_prefixed_model(self): overlay, _ = _overlay( "databricks-claude/claude-sonnet", @@ -228,6 +275,15 @@ def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert pi.default_model(state) == "gemini-2" + def test_falls_back_to_oss_last(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.glm-5-2"], + } + assert pi.default_model(state) == "system.ai.glm-5-2" + def test_returns_none_when_empty(self): assert pi.default_model({}) is None assert ( @@ -296,6 +352,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc "databricks-claude": {"old": True}, "databricks-openai": {"old": True}, "databricks-gemini": {"old": True}, + "databricks-mlflow": {"old": True}, "user-provider": {"keep": True}, } } @@ -311,6 +368,7 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc providers = written.get("providers", {}) assert providers.get("databricks-claude") != {"old": True} assert "old" not in providers.get("databricks-claude", {}) + assert "databricks-mlflow" not in providers assert providers.get("user-provider") == {"keep": True} def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 3443aad8..59f68579 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -189,6 +189,14 @@ def test_pi_available_with_codex(self): def test_pi_available_with_gemini(self): assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "pi") is True + def test_pi_available_with_oss(self): + assert check_gateway_endpoint({"oss_models": ["system.ai.glm-5-2"]}, "pi") is True + + def test_pi_oss_discovery_reason_is_reported(self): + state = {"_discovery_reasons": {"oss": "no validated OSS models"}} + detail = agents_mod._availability_failure_detail("pi", state) + assert detail == " (oss discovery: no validated OSS models)" + def test_pi_unavailable_when_no_models(self): assert check_gateway_endpoint({}, "pi") is False @@ -243,6 +251,15 @@ def test_pi_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} assert default_model_for_tool("pi", state) == "gemini-2" + def test_pi_falls_back_to_oss(self): + state = { + "claude_models": {}, + "codex_models": [], + "gemini_models": [], + "oss_models": ["system.ai.glm-5-2"], + } + assert default_model_for_tool("pi", state) == "system.ai.glm-5-2" + def test_pi_returns_none_when_no_models(self): assert default_model_for_tool("pi", {}) is None diff --git a/tests/test_cli.py b/tests/test_cli.py index 8e7473db..9639d172 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -27,6 +27,17 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] +def test_oss_discovery_diagnostic_names_all_consumers(monkeypatch): + import ucode.cli as cli_mod + + notes = [] + monkeypatch.setattr(cli_mod, "print_note", notes.append) + + cli_mod._print_discovery_diagnostics({"_discovery_reasons": {"oss": "not found"}}) + + assert notes[0] == "OSS models (needed for: opencode, pi): not found" + + @pytest.fixture(autouse=True) def no_state_writes(): """Prevent any test from writing to the real state file on disk.""" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 212ca626..0adf379d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -173,19 +173,52 @@ def _model_service(model_id: str) -> dict: class TestModelTokenLimits: def test_glm_is_capped(self): + # Probed 2026-07-16: glm-5-2 accepts 1M context / 65536 output. assert db_mod.model_token_limits("system.ai.glm-5-2") == { - "context": 200_000, - "output": 25_000, + "context": 1_000_000, + "output": 65_536, } - def test_glm_matches_any_version(self): - assert db_mod.model_token_limits("system.ai.glm-4-6-flash") == { + @pytest.mark.parametrize( + "model_id", + ["system.ai.glm-4-6-flash", "system.ai.glm-future"], + ) + def test_other_glm_versions_keep_conservative_limits(self, model_id): + assert db_mod.model_token_limits(model_id) == { "context": 200_000, "output": 25_000, } - def test_uncapped_model_returns_none(self): - assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") is None + def test_kimi_is_capped(self): + assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") == { + "context": 128_000, + "output": 65_536, + } + + def test_unvalidated_families_return_none(self): + for model_id in ( + "system.ai.inkling", + "system.ai.gpt-oss-120b", + "system.ai.llama-4-maverick", + "system.ai.qwen35-122b-a10b", + "system.ai.gemma-3-12b", + "system.ai.deepseek-v3", + ): + assert db_mod.model_token_limits(model_id) is None + + def test_embedding_model_returns_none_not_fallback(self): + assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None + + +class TestModelIsReasoning: + def test_reasoning_families(self): + assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True + assert db_mod.model_is_reasoning("system.ai.kimi-k2-7-code") is True + + def test_unvalidated_families_are_not_marked_reasoning(self): + assert db_mod.model_is_reasoning("system.ai.inkling") is False + assert db_mod.model_is_reasoning("system.ai.qwen35-122b-a10b") is False + assert db_mod.model_is_reasoning("system.ai.gpt-oss-120b") is False class TestDiscoverModelServices: @@ -220,19 +253,20 @@ def test_buckets_families_by_name(self, monkeypatch): assert codex == ["system.ai.gpt-5"] # Gemini ordered newest-first via the shared sort key. assert gemini[0] == "system.ai.gemini-3-5-flash" - # kimi and glm are the allowlisted OSS families; llama is not. + # Only coding-harness-validated OSS families are offered. assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] - def test_oss_allowlist_drops_unsupported_families(self, monkeypatch): - # Only kimi/glm are allowlisted; other families are dropped. + def test_oss_allowlist_drops_unvalidated_families(self, monkeypatch): payload = { "model_services": [ _model_service("system.ai.glm-5-2"), _model_service("system.ai.kimi-k2-7-code"), - _model_service("system.ai.qwen-3-coder"), + _model_service("system.ai.qwen35-122b-a10b"), + _model_service("system.ai.inkling"), + _model_service("system.ai.llama-4-maverick"), + _model_service("system.ai.gemma-3-12b"), _model_service("system.ai.deepseek-v3"), - _model_service("system.ai.gte-large-embed"), - _model_service("system.ai.bge-reranker-v2"), + _model_service("system.ai.qwen3-embedding-0-6b"), ] } monkeypatch.setattr( @@ -245,6 +279,25 @@ def test_oss_allowlist_drops_unsupported_families(self, monkeypatch): assert (claude, codex, gemini) == ({}, [], []) assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"] + def test_gpt_oss_is_neither_selectable_oss_nor_codex(self, monkeypatch): + # Keep the independent codex exclusion: gpt-oss contains "gpt-" but + # cannot use the Responses API, even though it is not an offered model. + payload = { + "model_services": [ + _model_service("system.ai.gpt-5"), + _model_service("system.ai.gpt-oss-120b"), + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) + ) + + _, codex, _, oss, _ = db_mod.discover_model_services(WS, "token") + + assert codex == ["system.ai.gpt-5"] + assert "system.ai.gpt-oss-120b" not in oss + assert "system.ai.gpt-oss-120b" not in codex + def test_paginates_via_next_page_token(self, monkeypatch): pages = { None: { @@ -281,7 +334,8 @@ def test_http_failure_returns_reason(self, monkeypatch): assert reason == "HTTP 500 Server Error" def test_no_matching_families_reports_sample(self, monkeypatch): - payload = {"model_services": [_model_service("system.ai.llama-4-maverick")]} + # deepseek is outside every claude/gpt/gemini/oss family bucket. + payload = {"model_services": [_model_service("system.ai.deepseek-v3")]} monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) @@ -289,7 +343,7 @@ def test_no_matching_families_reports_sample(self, monkeypatch): claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") assert (claude, codex, gemini, oss) == ({}, [], [], []) - assert reason is not None and "llama-4-maverick" in reason + assert reason is not None and "deepseek-v3" in reason def test_ignores_non_system_ai_schemas(self, monkeypatch): # The metastore listing returns services from every schema; only From 636c3b625ec5b32f7827b9c6ec96119c01cfa2c5 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:20:17 +1000 Subject: [PATCH 3/8] fix: share Claude context capability policy Centralize Claude family/version parsing so Pi metadata, adaptive-thinking compatibility, and Claude Code's [1m] selector cannot drift. Cover Sonnet 4.5, Opus 4.6, future major versions, Fable fallback, and prefixed model IDs. --- src/ucode/agents/claude.py | 20 +----- src/ucode/agents/pi.py | 49 ++++++++++++-- src/ucode/databricks.py | 131 +++++++++++++++++++++++++++++++++++++ tests/test_agent_claude.py | 37 +++++++++++ tests/test_agent_pi.py | 75 ++++++++++++++++++++- tests/test_databricks.py | 102 +++++++++++++++++++++++++++++ 6 files changed, 389 insertions(+), 25 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 2e22b240..40efd48e 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -23,6 +23,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 @@ -65,11 +66,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. @@ -242,19 +238,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: diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 7f782932..806c26f8 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -51,9 +51,12 @@ from ucode.databricks import ( TOKEN_REFRESH_INTERVAL_SECONDS, build_pi_base_urls, + claude_model_capabilities, get_databricks_token, + gpt_model_token_limits, model_is_reasoning, model_token_limits, + preferred_gpt_model, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -113,6 +116,25 @@ def _resolve_model_selector( return model +def _pi_claude_model_entry(model_id: str) -> dict: + """Build a Claude entry with explicit limits. + + Databricks model ids do not match Pi's built-in Anthropic ids, so a bare + custom entry silently gets Pi's 128k context / 4k output defaults. + """ + capabilities = claude_model_capabilities(model_id) + entry: dict = { + "id": model_id, + "reasoning": True, + "input": ["text", "image"], + "contextWindow": capabilities.context, + "maxTokens": capabilities.output, + } + if capabilities.force_adaptive_thinking: + entry["compat"] = {"forceAdaptiveThinking": True} + return entry + + 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 @@ -131,6 +153,23 @@ def _pi_oss_model_entry(model_id: str) -> dict: return entry +def _pi_gpt_model_entry(model_id: str) -> dict: + """Build a Pi openai (codex) model entry with `contextWindow`/`maxTokens` + from `databricks.gpt_model_token_limits`. GPT ids aren't in Pi's built-in + catalog, so without an explicit window Pi falls back to a small default and + truncates long sessions.""" + limits = gpt_model_token_limits(model_id) + entry: dict = { + "id": model_id, + "contextWindow": limits["context"], + "maxTokens": limits["output"], + } + if "gpt-5" in model_id.lower().replace(".", "-"): + entry["reasoning"] = True + entry["input"] = ["text", "image"] + return entry + + def render_overlay( model: str, token: str, @@ -159,7 +198,7 @@ def render_overlay( # the legacy beta header instead when this is false. "compat": {"supportsEagerToolInputStreaming": False}, "headers": ua_headers, - "models": [{"id": m} for m in claude_ids], + "models": [_pi_claude_model_entry(m) for m in claude_ids], } keys.append(["providers", "databricks-claude"]) if codex_models: @@ -169,7 +208,7 @@ def render_overlay( "apiKey": token, "authHeader": True, "headers": ua_headers, - "models": [{"id": m} for m in codex_models], + "models": [_pi_gpt_model_entry(m) for m in codex_models], } keys.append(["providers", "databricks-openai"]) if gemini_models: @@ -259,9 +298,9 @@ def default_model(state: dict) -> str | None: for family in ("opus", "sonnet", "haiku"): if claude_models.get(family): return claude_models[family] - codex_models = state.get("codex_models") or [] - if codex_models: - return codex_models[0] + codex_model = preferred_gpt_model(state.get("codex_models") or []) + if codex_model: + return codex_model gemini_models = state.get("gemini_models") or [] if gemini_models: return gemini_models[0] diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9e6d331a..089fe05f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -22,6 +22,7 @@ from concurrent.futures import ( TimeoutError as FutureTimeoutError, ) +from dataclasses import dataclass from pathlib import Path from typing import Literal, cast, overload from urllib import error as urllib_error @@ -1235,6 +1236,136 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None +# Pi treats every custom model without explicit metadata as 128k context / 4k +# output. Gateway ids are custom ids (not Pi's built-ins), so preserve the +# upstream windows explicitly. Entries are ordered most-specific first after +# normalizing dotted OpenAI ids and hyphenated Databricks ids to one form. +# GPT-5.6 Sol/Terra/Luna support the opt-in 1.05M window; 272k is merely Pi's +# built-in short-context pricing default, not the model's hard context limit. +_GPT_TOKEN_LIMITS: tuple[tuple[str, dict[str, int]], ...] = ( + ("gpt-5-6-sol", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-terra", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-6-luna", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-4-pro", {"context": 1_050_000, "output": 128_000}), + ("gpt-5-5", {"context": 272_000, "output": 128_000}), + ("gpt-5-4-mini", {"context": 400_000, "output": 128_000}), + ("gpt-5-4-nano", {"context": 400_000, "output": 128_000}), + ("gpt-5-4", {"context": 272_000, "output": 128_000}), + ("gpt-5", {"context": 400_000, "output": 128_000}), + ("gpt-4-1", {"context": 1_047_576, "output": 32_768}), + ("gpt-4o", {"context": 128_000, "output": 16_384}), + ("gpt-4-turbo", {"context": 128_000, "output": 4_096}), + ("gpt-4", {"context": 8_192, "output": 8_192}), +) +_GPT_FALLBACK_LIMITS = {"context": 128_000, "output": 16_384} + + +def _normalized_foundation_model_id(model_id: str) -> str: + """Strip route prefixes case-insensitively and normalize dotted versions.""" + tail = model_id.split("/")[-1].lower() + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + if tail.startswith("databricks-"): + tail = tail[len("databricks-") :] + return tail.replace(".", "-") + + +def gpt_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits for a GPT (codex/openai) gateway model.""" + tail = _normalized_foundation_model_id(model_id) + for family, limits in _GPT_TOKEN_LIMITS: + if tail == family or tail.startswith(f"{family}-"): + return dict(limits) + return dict(_GPT_FALLBACK_LIMITS) + + +def preferred_gpt_model(model_ids: list[str]) -> str | None: + """Prefer the newest numeric GPT id, then any other Responses endpoint. + + ``gpt-oss`` is chat-completions-only and must never be selected for the + Responses route, even if stale state supplies it here. + """ + eligible = [ + model_id + for model_id in model_ids + if not _normalized_foundation_model_id(model_id).startswith("gpt-oss") + ] + numeric_gpt = [ + model_id + for model_id in eligible + if re.match(r"^gpt-\d(?:-|$)", _normalized_foundation_model_id(model_id)) + ] + if numeric_gpt: + return min( + numeric_gpt, + key=lambda model_id: model_version_sort_key(_normalized_foundation_model_id(model_id)), + ) + return eligible[0] if eligible else None + + +@dataclass(frozen=True) +class ClaudeModelCapabilities: + context: int + output: int + supports_1m: bool = False + force_adaptive_thinking: bool = False + + +_CLAUDE_FALLBACK_CAPABILITIES = ClaudeModelCapabilities(context=200_000, output=64_000) +_CLAUDE_MODEL_RE = re.compile(r"^claude-(fable|opus|sonnet|haiku)-(\d+)(?:-(\d+))?") + + +def claude_model_capabilities(model_id: str) -> ClaudeModelCapabilities: + """Return the shared Claude capability policy for every agent. + + Opus gained the opt-in 1M window in 4.6; Sonnet gained it in 4.5. + Sonnet's verified 1M tiers retain a conservative 64k output cap. Fable 5 + is 1M by default (so it needs no ``[1m]`` suffix) with a 128k output cap. + Opus 4.5, Haiku, and unrecognized ids use the conservative 200k fallback. + """ + tail = _normalized_foundation_model_id(model_id) + match = _CLAUDE_MODEL_RE.match(tail) + if not match: + return _CLAUDE_FALLBACK_CAPABILITIES + family, major_raw, minor_raw = match.groups() + version = (int(major_raw), int(minor_raw or 0)) + if family == "opus" and version >= (4, 6): + return ClaudeModelCapabilities( + context=1_000_000, + output=128_000, + supports_1m=True, + force_adaptive_thinking=True, + ) + if family == "sonnet" and version >= (4, 6): + return ClaudeModelCapabilities( + context=1_000_000, + output=64_000, + supports_1m=True, + force_adaptive_thinking=True, + ) + if family == "sonnet" and version >= (4, 5): + return ClaudeModelCapabilities(context=1_000_000, output=64_000, supports_1m=True) + if family == "fable" and version >= (5, 0): + return ClaudeModelCapabilities( + context=1_000_000, + output=128_000, + force_adaptive_thinking=True, + ) + return _CLAUDE_FALLBACK_CAPABILITIES + + +def claude_model_supports_1m(model_id: str) -> bool: + """Whether Claude Code should request the model's opt-in ``[1m]`` tier.""" + return claude_model_capabilities(model_id).supports_1m + + +def claude_model_token_limits(model_id: str) -> dict[str, int]: + """Return Pi metadata limits from the shared Claude capability policy.""" + capabilities = claude_model_capabilities(model_id) + return {"context": capabilities.context, "output": capabilities.output} + + def _model_service_id(service: dict) -> str | None: """Extract the `system.ai.` id from one model-service entry. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b6952ee2..cc062608 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -48,6 +48,30 @@ def test_adds_1m_suffix_for_sonnet_4_6_and_later(self): overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-7[1m]" ) + def test_adds_1m_suffix_for_sonnet_4_5(self): + # Sonnet 4.5 supports the 1M context window (its 1M beta shipped before + # Opus's), so it must get the [1m] suffix even though it predates the + # Opus 4.6 floor. + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-5"} + ) + assert ( + overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-5[1m]" + ) + + def test_does_not_add_1m_suffix_for_sonnet_4_4(self): + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-4"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-4" + + def test_does_not_add_1m_suffix_for_opus_4_5(self): + # Opus's 1M window starts at 4.6, so 4.5 stays on the default context. + overlay, _ = claude.render_overlay( + WS, "s4", claude_models={"opus": "databricks-claude-opus-4-5"} + ) + assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-5" + def test_does_not_add_1m_suffix_for_haiku(self): overlay, _ = claude.render_overlay( WS, "s4", claude_models={"haiku": "databricks-claude-haiku-4-6"} @@ -72,6 +96,19 @@ def test_no_1m_suffix_for_model_services_haiku(self): ) assert overlay["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-4-6" + @pytest.mark.parametrize( + ("model_id", "expected"), + [ + ("system.ai.claude-opus-5", "system.ai.claude-opus-5[1m]"), + ("databricks-claude-sonnet-5", "databricks-claude-sonnet-5[1m]"), + ("system.ai.claude-opus-4-5", "system.ai.claude-opus-4-5"), + ("system.ai.claude-fable-5", "system.ai.claude-fable-5"), + ("not-a-claude-model", "not-a-claude-model"), + ], + ) + def test_suffix_uses_shared_capability_policy(self, model_id, expected): + assert claude._maybe_add_1m_suffix(model_id) == expected + def test_sets_anthropic_base_url(self): overlay, _ = claude.render_overlay(WS, "s4") assert overlay["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 4c5444ee..b0711eee 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -78,6 +78,60 @@ def test_openai_provider_uses_openai_responses(self): assert provider["api"] == "openai-responses" assert provider["baseUrl"] == f"{WS}/ai-gateway/codex/v1" + def test_gpt56_sol_model_entry_pins_1m_context(self): + # Gateway ids are custom to Pi, so explicit metadata is required to + # avoid its 128k custom-model default. + overlay, _ = _overlay("gpt-5-6-sol", codex_models=["gpt-5-6-sol"]) + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert entry["id"] == "gpt-5-6-sol" + assert entry["contextWindow"] == 1_050_000 + assert entry["maxTokens"] == 128_000 + assert entry["reasoning"] is True + assert entry["input"] == ["text", "image"] + + def test_gpt_model_entries_use_model_specific_windows(self): + overlay, _ = _overlay( + "system.ai.gpt-5-2", + codex_models=[ + "system.ai.gpt-5-2", + "databricks-gpt-5-4-nano", + "databricks-gpt-5-6-sol", + ], + ) + windows = { + m["id"]: m["contextWindow"] for m in overlay["providers"]["databricks-openai"]["models"] + } + assert windows == { + "system.ai.gpt-5-2": 400_000, + "databricks-gpt-5-4-nano": 400_000, + "databricks-gpt-5-6-sol": 1_050_000, + } + + def test_claude_entries_pin_limits_and_capabilities(self): + overlay, _ = _overlay( + "databricks-claude-opus-4-8", + claude_models={ + "opus": "databricks-claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-5", + "haiku": "databricks-claude-haiku-4-5", + "fable": "system.ai.claude-fable-5", + }, + ) + entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]} + opus = entries["databricks-claude-opus-4-8"] + assert opus["contextWindow"] == 1_000_000 + assert opus["maxTokens"] == 128_000 + assert opus["reasoning"] is True + assert opus["input"] == ["text", "image"] + assert opus["compat"] == {"forceAdaptiveThinking": True} + assert entries["system.ai.claude-sonnet-4-5"]["contextWindow"] == 1_000_000 + assert entries["system.ai.claude-sonnet-4-5"]["maxTokens"] == 64_000 + assert entries["databricks-claude-haiku-4-5"]["contextWindow"] == 200_000 + fable = entries["system.ai.claude-fable-5"] + assert fable["contextWindow"] == 1_000_000 + assert fable["maxTokens"] == 128_000 + assert fable["compat"] == {"forceAdaptiveThinking": True} + def test_gemini_provider_uses_google_generative_ai(self): overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"]) provider = overlay["providers"]["databricks-gemini"] @@ -267,9 +321,24 @@ def test_falls_back_to_haiku(self): state = {"claude_models": {"haiku": "h4"}} assert pi.default_model(state) == "h4" - def test_falls_back_to_codex(self): - state = {"claude_models": {}, "codex_models": ["gpt-5"]} - assert pi.default_model(state) == "gpt-5" + def test_falls_back_to_newest_codex_model(self): + state = { + "claude_models": {}, + "codex_models": ["databricks-gpt-5", "system.ai.gpt-5-6-sol", "gpt-5-5"], + } + assert pi.default_model(state) == "system.ai.gpt-5-6-sol" + + def test_falls_back_to_generic_responses_endpoint(self): + state = {"claude_models": {}, "codex_models": ["c1"]} + assert pi.default_model(state) == "c1" + + def test_does_not_route_gpt_oss_to_responses(self): + state = { + "claude_models": {}, + "codex_models": ["gpt-oss-120b"], + "gemini_models": ["gemini-2"], + } + assert pi.default_model(state) == "gemini-2" def test_falls_back_to_gemini(self): state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]} diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0adf379d..a2af5f9f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -210,6 +210,108 @@ def test_embedding_model_returns_none_not_fallback(self): assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None +class TestGptModelTokenLimits: + def test_gpt56_sol_serves_1m_context_across_id_forms(self): + for model_id in ( + "gpt-5.6-sol", + "system.ai.gpt-5-6-sol", + "databricks-gpt-5-6-sol", + "databricks-openai/gpt-5.6-sol", + ): + assert db_mod.gpt_model_token_limits(model_id) == { + "context": 1_050_000, + "output": 128_000, + } + + def test_gpt5_windows_are_model_specific(self): + assert db_mod.gpt_model_token_limits("system.ai.gpt-5-2")["context"] == 400_000 + assert db_mod.gpt_model_token_limits("databricks-gpt-5-4")["context"] == 272_000 + assert db_mod.gpt_model_token_limits("databricks-gpt-5-4-nano")["context"] == 400_000 + assert db_mod.gpt_model_token_limits("gpt-5.5-pro")["context"] == 1_050_000 + + def test_gpt41_window(self): + assert db_mod.gpt_model_token_limits("databricks-gpt-4-1") == { + "context": 1_047_576, + "output": 32_768, + } + + @pytest.mark.parametrize( + "model_id", + ["gpt-5.40", "gpt-4.10", "gpt-6-turbo"], + ) + def test_unknown_specific_gpt_uses_family_or_conservative_fallback(self, model_id): + expected = ( + {"context": 400_000, "output": 128_000} + if model_id == "gpt-5.40" + else {"context": 8_192, "output": 8_192} + if model_id == "gpt-4.10" + else {"context": 128_000, "output": 16_384} + ) + assert db_mod.gpt_model_token_limits(model_id) == expected + + def test_route_prefixes_are_stripped_case_insensitively(self): + assert db_mod.gpt_model_token_limits("SYSTEM.AI.GPT-5-6-SOL") == { + "context": 1_050_000, + "output": 128_000, + } + assert db_mod.gpt_model_token_limits("DATABRICKS-GPT-4-1") == { + "context": 1_047_576, + "output": 32_768, + } + + def test_preferred_gpt_model_uses_semantic_version_across_prefixes(self): + assert ( + db_mod.preferred_gpt_model( + ["databricks-gpt-5", "system.ai.gpt-5-6-sol", "databricks-gpt-5-5"] + ) + == "system.ai.gpt-5-6-sol" + ) + assert db_mod.preferred_gpt_model(["not-gpt", "claude-opus-4-8"]) == "not-gpt" + + def test_preferred_gpt_model_falls_back_to_generic_responses_endpoint(self): + assert db_mod.preferred_gpt_model(["c1", "custom-responses"]) == "c1" + + def test_preferred_gpt_model_excludes_gpt_oss(self): + assert db_mod.preferred_gpt_model(["gpt-oss-120b", "c1"]) == "c1" + assert db_mod.preferred_gpt_model(["system.ai.gpt-oss-120b"]) is None + + +class TestClaudeModelCapabilities: + @pytest.mark.parametrize( + ("model_id", "context", "output", "supports_1m", "adaptive"), + [ + ("databricks-claude-opus-4-5", 200_000, 64_000, False, False), + ("databricks-claude-opus-4-6", 1_000_000, 128_000, True, True), + ("system.ai.claude-opus-5", 1_000_000, 128_000, True, True), + ("claude-sonnet-4-4", 200_000, 64_000, False, False), + ("system.ai.claude-sonnet-4-5", 1_000_000, 64_000, True, False), + ("claude-sonnet-4-6[1m]", 1_000_000, 64_000, True, True), + ("claude-sonnet-5", 1_000_000, 64_000, True, True), + ("claude-haiku-4-5", 200_000, 64_000, False, False), + ("system.ai.claude-fable-5", 1_000_000, 128_000, False, True), + ("claude-future", 200_000, 64_000, False, False), + ], + ) + def test_shared_capability_policy( + self, + model_id, + context, + output, + supports_1m, + adaptive, + ): + capabilities = db_mod.claude_model_capabilities(model_id) + assert capabilities.context == context + assert capabilities.output == output + assert capabilities.supports_1m is supports_1m + assert capabilities.force_adaptive_thinking is adaptive + assert db_mod.claude_model_supports_1m(model_id) is supports_1m + assert db_mod.claude_model_token_limits(model_id) == { + "context": context, + "output": output, + } + + class TestModelIsReasoning: def test_reasoning_families(self): assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True From 3f1c227aa418a1ee4a2bb4339261e72a9ad1fa47 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:24:47 +1000 Subject: [PATCH 4/8] fix(opencode): route GPT models and discover OSS fallbacks Configure the Databricks OpenAI Responses provider alongside the validated GLM/Kimi provider, and fall back to foundation-model serving endpoints when UC model services are unavailable. --- src/ucode/agents/opencode.py | 45 +++++++++++- src/ucode/cli.py | 11 ++- src/ucode/databricks.py | 71 ++++++++++++++++-- tests/conftest.py | 32 +++++++-- tests/test_agent_opencode.py | 117 +++++++++++++++++++++++++++++- tests/test_cli.py | 24 +++++++ tests/test_databricks.py | 136 +++++++++++++++++++++++++++++++++++ tests/test_e2e.py | 22 +++++- 8 files changed, 438 insertions(+), 20 deletions(-) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 4b614f32..9c905d1e 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -20,7 +20,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 @@ -42,6 +44,7 @@ PROVIDER_KEYS: list[list[str]] = [ ["provider", "databricks-anthropic"], ["provider", "databricks-google"], + ["provider", "databricks-openai"], ["provider", "databricks-oss"], ] @@ -52,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 [] @@ -63,6 +73,10 @@ 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}" @@ -84,6 +98,15 @@ def _oss_model_overlay(model: str, ua_header: dict[str, str]) -> dict: 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, @@ -102,6 +125,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 = {} @@ -137,6 +161,22 @@ 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..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: providers["databricks-oss"] = { "npm": "@ai-sdk/openai", @@ -233,6 +273,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] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 57e75bf1..43dd947e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -37,6 +37,7 @@ discover_codex_models, discover_gemini_models, discover_model_services, + discover_oss_models, ensure_ai_gateway_v2, ensure_databricks_auth, ensure_pat_bearer, @@ -92,7 +93,7 @@ _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = { "claude": ("claude", "opencode", "copilot", "pi"), - "codex": ("codex", "copilot", "pi"), + "codex": ("codex", "copilot", "opencode", "pi"), "gemini": ("gemini", "opencode", "pi"), "oss": ("opencode", "pi"), } @@ -371,7 +372,9 @@ 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_codex = ( + fetch_all or "codex" in tools or "copilot" in tools or "opencode" in tools or "pi" in tools + ) want_oss = fetch_all or "opencode" in tools or "pi" in tools claude_reason: str | None = None @@ -424,10 +427,14 @@ def configure_shared_state( codex_models, codex_reason = discover_codex_models(workspace, token) if want_oss: oss_models, oss_reason = ms_oss, ms_reason + if not oss_models: + oss_models, oss_reason = discover_oss_models(workspace, token) if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if codex_models: + opencode_models["openai"] = codex_models if oss_models: opencode_models["oss"] = oss_models diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 089fe05f..6bd9cac3 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2128,27 +2128,53 @@ def discover_endpoints_with_api_type( return [], reason data = cast(dict, payload) if isinstance(payload, dict) else {} - endpoints = data.get("endpoints", []) + raw_endpoints = data.get("endpoints", []) + endpoints = raw_endpoints if isinstance(raw_endpoints, list) else [] out: list[str] = [] saw_endpoint_without_v2 = False + saw_malformed = not isinstance(raw_endpoints, list) for ep in endpoints: - name = ep.get("name", "") - entities = ep.get("config", {}).get("served_entities", []) + if not isinstance(ep, dict): + saw_malformed = True + continue + name = ep.get("name") + config = ep.get("config") + if not isinstance(name, str) or not name or not isinstance(config, dict): + saw_malformed = True + continue + raw_entities = config.get("served_entities", []) + if not isinstance(raw_entities, list): + saw_malformed = True + continue api_types: set[str] = set() any_v2 = False - for se in entities: - fm = se.get("foundation_model", {}) + for se in raw_entities: + if not isinstance(se, dict): + saw_malformed = True + continue + fm = se.get("foundation_model") + if not isinstance(fm, dict): + saw_malformed = True + continue if fm.get("ai_gateway_v2_supported") is True: any_v2 = True - api_types.update(fm.get("api_types", [])) - if not any_v2 and entities: + raw_api_types = fm.get("api_types", []) + if isinstance(raw_api_types, list): + api_types.update(value for value in raw_api_types if isinstance(value, str)) + else: + saw_malformed = True + if not any_v2 and raw_entities: saw_endpoint_without_v2 = True if api_type in api_types: out.append(name) if out: return sorted(out, key=sort_key), None if not endpoints: + if saw_malformed: + return [], "foundation-models listing returned malformed `endpoints`" return [], "foundation-models listing returned no endpoints" + if saw_malformed: + return [], "foundation-models listing contained no valid matching endpoints" if saw_endpoint_without_v2: return [], ( f"no endpoint exposes api_type `{api_type}` with " @@ -2175,6 +2201,34 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | return discover_endpoints_with_api_type(workspace, token, "openai/v1/responses") +def discover_oss_models(workspace: str, token: str) -> tuple[list[str], str | None]: + """Discover OSS chat models served as AI Gateway foundation-model endpoints. + + Fallback for workspaces that don't register OSS foundation models as + `system.ai.*` UC model-services (see `discover_model_services`): those + workspaces expose the same models as regular `databricks-*` serving + endpoints instead. Lists every endpoint advertising the + `mlflow/v1/chat/completions` dialect, then keeps only the OSS chat families + (`_is_oss_chat_model`) — on some workspaces the Claude/Gemini endpoints also + advertise that dialect, so the family filter is what separates the OSS + cohort from them. Mirrors the AI-Gateway fallback the other families use + when the UC model-services listing is empty. + """ + endpoints, reason = discover_endpoints_with_api_type( + workspace, token, "mlflow/v1/chat/completions" + ) + if not endpoints: + return [], reason + oss = [e for e in endpoints if _is_oss_chat_model(e)] + if oss: + return oss, None + sample = ", ".join(endpoints[:5]) + return [], ( + "foundation-models exposing `mlflow/v1/chat/completions` matched no OSS " + f"chat family (got: {sample})" + ) + + def fetch_gemini_models(workspace: str, token: str) -> list[str]: models, _ = discover_gemini_models(workspace, token) return models @@ -2372,6 +2426,9 @@ def build_opencode_base_urls(workspace: str) -> dict[str, str]: return { "anthropic": build_tool_base_url("claude", workspace) + "/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + # @ai-sdk/openai appends "/responses" (or "/chat/completions") to baseURL, + # so stop just before that — matches the Pi adapter's build_pi_base_urls. + "openai": build_tool_base_url("codex", workspace), "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/tests/conftest.py b/tests/conftest.py index 0cc7932b..e2051af7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,9 +8,11 @@ from ucode.databricks import ( build_shared_base_urls, - fetch_ai_gateway_claude_models, - fetch_codex_models, - fetch_gemini_models, + discover_claude_models, + discover_codex_models, + discover_gemini_models, + discover_model_services, + discover_oss_models, get_databricks_token, ) from ucode.ui import normalize_workspace_url @@ -52,22 +54,38 @@ def e2e_token(e2e_workspace): @pytest.fixture(scope="session") def e2e_state(e2e_workspace, e2e_token): - """Full state dict mirroring what configure_shared_state produces.""" - claude_models = fetch_ai_gateway_claude_models(e2e_workspace, e2e_token) - gemini_models = fetch_gemini_models(e2e_workspace, e2e_token) - codex_models = fetch_codex_models(e2e_workspace, e2e_token) + """Full state dict mirroring configure's UC-first family discovery.""" + claude_models, codex_models, gemini_models, oss_models, _ = discover_model_services( + e2e_workspace, e2e_token + ) + if not claude_models: + claude_models, _ = discover_claude_models(e2e_workspace, e2e_token) + if not gemini_models: + gemini_models, _ = discover_gemini_models(e2e_workspace, e2e_token) + if not codex_models: + codex_models, _ = discover_codex_models(e2e_workspace, e2e_token) + if not oss_models: + oss_models, _ = discover_oss_models(e2e_workspace, e2e_token) + + # E2E mirrors configure's default (Fable is premium and opt-in). + claude_models.pop("fable", None) opencode_models: dict = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if codex_models: + opencode_models["openai"] = codex_models + if oss_models: + opencode_models["oss"] = oss_models return { "workspace": e2e_workspace, "claude_models": claude_models, "gemini_models": gemini_models, "codex_models": codex_models, + "oss_models": oss_models, "opencode_models": opencode_models, "base_urls": build_shared_base_urls(e2e_workspace), "managed_configs": {}, diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 3cc7348b..b070f652 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -14,6 +14,7 @@ def _base_urls() -> dict[str, str]: return { "anthropic": f"{WS}/ai-gateway/anthropic/v1", "gemini": f"{WS}/ai-gateway/gemini/v1beta", + "openai": f"{WS}/ai-gateway/codex/v1", "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -211,6 +212,89 @@ def test_prefixes_oss_model_with_provider_id(self): assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" +class TestOpenAIProvider: + """OpenCode reaches Databricks GPT-5 / GPT-5.6 / Codex models through the + databricks-openai provider (@ai-sdk/openai against /ai-gateway/codex/v1). + Without this wiring the codex/openai family is unreachable from OpenCode.""" + + def test_openai_provider_added_when_codex_models_present(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + assert "databricks-openai" in overlay["provider"] + + def test_openai_provider_uses_ai_sdk_openai_npm(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + assert overlay["provider"]["databricks-openai"]["npm"] == "@ai-sdk/openai" + + def test_openai_base_url_points_at_codex_gateway(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + options = overlay["provider"]["databricks-openai"]["options"] + assert options["baseURL"] == f"{WS}/ai-gateway/codex/v1" + + def test_use_responses_api_set_on_every_codex_model(self): + models = {"openai": ["databricks-gpt-5-6-sol", "databricks-gpt-codex"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + provider_models = overlay["provider"]["databricks-openai"]["models"] + for m in ("databricks-gpt-5-6-sol", "databricks-gpt-codex"): + assert provider_models[m]["options"]["useResponsesApi"] is True + + def test_openai_models_include_explicit_token_limits(self): + models = {"openai": ["databricks-gpt-5-6-sol", "databricks-gpt-4-1"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + provider_models = overlay["provider"]["databricks-openai"]["models"] + assert provider_models["databricks-gpt-5-6-sol"]["limit"] == { + "context": 1_050_000, + "output": 128_000, + } + assert provider_models["databricks-gpt-4-1"]["limit"] == { + "context": 1_047_576, + "output": 32_768, + } + + def test_openai_authorization_header(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + headers = overlay["provider"]["databricks-openai"]["options"]["headers"] + assert headers["Authorization"] == "Bearer tok" + + def test_managed_keys_include_openai_provider(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + _, keys = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + assert ["provider", "databricks-openai"] in keys + + def test_prefixes_openai_model_with_provider_id(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay("databricks-gpt-5-6-sol", "tok", _base_urls(), models) + assert overlay["model"] == "databricks-openai/databricks-gpt-5-6-sol" + + def test_already_prefixed_openai_model_is_preserved(self): + models = {"openai": ["databricks-gpt-5-6-sol"]} + overlay, _ = opencode.render_overlay( + "databricks-openai/databricks-gpt-5-6-sol", "tok", _base_urls(), models + ) + assert overlay["model"] == "databricks-openai/databricks-gpt-5-6-sol" + + def test_all_four_providers_when_all_present(self): + models = { + "anthropic": ["claude-sonnet"], + "gemini": ["gemini-2"], + "openai": ["databricks-gpt-5-6-sol"], + "oss": ["system.ai.kimi-k2-7-code"], + } + overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + assert set(overlay["provider"].keys()) == { + "databricks-anthropic", + "databricks-google", + "databricks-openai", + "databricks-oss", + } + + def test_provider_keys_listed_in_module(self): + assert ["provider", "databricks-openai"] in opencode.PROVIDER_KEYS + + class TestMcpServerConfig: def test_builds_remote_server_entry_with_oauth_token_env_header(self): entry = opencode.build_mcp_server_entry(f"{WS}/api/2.0/mcp/external/github") @@ -323,6 +407,32 @@ def test_prefers_anthropic(self): state = {"opencode_models": {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "claude-sonnet" + def test_falls_back_to_openai_before_gemini(self): + state = { + "opencode_models": { + "anthropic": [], + "openai": ["databricks-gpt-5-6-sol"], + "gemini": ["gemini-2"], + } + } + assert opencode.default_model(state) == "databricks-gpt-5-6-sol" + + def test_openai_fallback_chooses_newest_gpt(self): + state = { + "opencode_models": { + "openai": ["databricks-gpt-4-1", "databricks-gpt-5-5", "databricks-gpt-5-4"] + } + } + assert opencode.default_model(state) == "databricks-gpt-5-5" + + def test_openai_falls_back_to_generic_responses_endpoint(self): + state = {"opencode_models": {"openai": ["c1"]}} + assert opencode.default_model(state) == "c1" + + def test_gpt_oss_is_not_selected_for_responses(self): + state = {"opencode_models": {"openai": ["gpt-oss-120b"], "gemini": ["gemini-2"]}} + assert opencode.default_model(state) == "gemini-2" + def test_falls_back_to_gemini(self): state = {"opencode_models": {"anthropic": [], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "gemini-2" @@ -371,6 +481,8 @@ def test_stale_providers_removed_before_merge(self, tmp_path, monkeypatch): "provider": { "databricks-anthropic": {"old": True}, "databricks-google": {"old": True}, + "databricks-openai": {"old": True}, + "databricks-oss": {"old": True}, "other-provider": {"keep": True}, } } @@ -393,7 +505,10 @@ def test_stale_providers_removed_before_merge(self, tmp_path, monkeypatch): providers = written.get("provider", {}) # stale entry is replaced with new data, not kept as-is assert providers.get("databricks-anthropic") != {"old": True} - # unmanaged provider entry survives + # Providers no longer discovered are removed rather than left stale. + assert "databricks-openai" not in providers + assert "databricks-oss" not in providers + # Unmanaged provider entries survive. assert providers.get("other-provider") == {"keep": True} def test_config_written_with_correct_model(self, tmp_path, monkeypatch): diff --git a/tests/test_cli.py b/tests/test_cli.py index 9639d172..33cd70d3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1396,6 +1396,7 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_oss_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) return cli_mod, logins, ensures, saved @@ -1602,6 +1603,7 @@ def test_ai_tools_disable_does_not_leak_across_workspaces(self, monkeypatch): def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): # No UC model-services: each family falls back to the legacy listing. cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + calls: list[str] = [] monkeypatch.setattr( cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], "no model services") ) @@ -1613,13 +1615,34 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): None, ), ) + monkeypatch.setattr( + cli_mod, + "discover_codex_models", + lambda w, t: (calls.append("codex") or ["databricks-gpt-5-6-sol"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_models", + lambda w, t: (calls.append("oss") or ["databricks-glm-5-2"], None), + ) state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + assert calls == ["codex", "oss"] assert state["claude_models"] == { "opus": "databricks-claude-opus-4-8", "sonnet": "databricks-claude-sonnet-4-6", } + assert state["codex_models"] == ["databricks-gpt-5-6-sol"] + assert state["oss_models"] == ["databricks-glm-5-2"] + assert state["opencode_models"] == { + "anthropic": [ + "databricks-claude-opus-4-8", + "databricks-claude-sonnet-4-6", + ], + "openai": ["databricks-gpt-5-6-sol"], + "oss": ["databricks-glm-5-2"], + } class TestConfigureSkipValidate: @@ -1685,6 +1708,7 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_oss_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) def test_purges_residue_when_workspace_changes(self, monkeypatch): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index a2af5f9f..3e222482 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -102,6 +102,12 @@ def test_returns_anthropic_gemini_and_oss(self): assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" + def test_returns_openai_codex_gateway(self): + # @ai-sdk/openai appends /responses (Responses API) or /chat/completions + # to baseURL, so stop just before that suffix. Mirrors build_pi_base_urls. + urls = build_opencode_base_urls(WS) + assert urls["openai"] == f"{WS}/ai-gateway/codex/v1" + class TestBuildSharedBaseUrls: def test_contains_all_tools(self): @@ -900,6 +906,65 @@ def test_unversioned_names_sort_last_alphabetically(self): assert ordered[1:] == ["another-endpoint", "custom-endpoint"] +class TestDiscoverEndpointsWithApiType: + @pytest.mark.parametrize( + "payload", + [ + {"endpoints": None}, + {"endpoints": "not-a-list"}, + {"endpoints": [None, "not-an-endpoint"]}, + {"endpoints": [{"name": 123, "config": {}}]}, + {"endpoints": [{"name": "model", "config": None}]}, + {"endpoints": [{"name": "model", "config": {"served_entities": None}}]}, + {"endpoints": [{"name": "model", "config": {"served_entities": [None, "bad"]}}]}, + { + "endpoints": [ + { + "name": "model", + "config": {"served_entities": [{"foundation_model": None}]}, + } + ] + }, + { + "endpoints": [ + { + "name": "model", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": "openai/v1/responses", + } + } + ] + }, + } + ] + }, + ], + ) + def test_malformed_payload_records_are_skipped(self, monkeypatch, payload): + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_endpoints_with_api_type(WS, "token", "openai/v1/responses") + + assert models == [] + assert reason and ("malformed" in reason or "no valid" in reason) + + def test_malformed_records_do_not_hide_valid_endpoint(self, monkeypatch): + payload = _foundation_models_payload(["databricks-gemini-3-5-flash"]) + payload["endpoints"].insert(0, None) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_endpoints_with_api_type( + WS, "token", "gemini/v1/generateContent" + ) + + assert models == ["databricks-gemini-3-5-flash"] + assert reason is None + + class TestDiscoverGeminiModels: def test_returns_newest_flash_first(self, monkeypatch): payload = _foundation_models_payload( @@ -945,6 +1010,77 @@ def test_codex_discovery_keeps_alphabetical_order(self, monkeypatch): assert models == ["databricks-gpt-4-1", "databricks-gpt-5-2-codex"] +def _mlflow_chat_payload(names, *, api_type="mlflow/v1/chat/completions", v2=True): + return { + "endpoints": [ + { + "name": name, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": v2, + "api_types": [api_type], + } + } + ] + }, + } + for name in names + ] + } + + +class TestDiscoverOssModels: + def test_finds_oss_endpoints_via_foundation_models(self, monkeypatch): + # Mirrors a workspace with no system.ai UC model-services: OSS models are + # plain databricks-* serving endpoints under the mlflow chat dialect. + payload = _mlflow_chat_payload( + [ + "databricks-glm-5-2", + "databricks-kimi-k2-7-code", + "databricks-inkling", + "databricks-qwen35-122b-a10b", + "databricks-gemma-3-12b", + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert reason is None + assert models == ["databricks-glm-5-2", "databricks-kimi-k2-7-code"] + + def test_excludes_claude_and_gemini_sharing_the_mlflow_dialect(self, monkeypatch): + # On some workspaces every foundation model advertises the mlflow chat + # dialect, so the api_type filter alone is too broad — the OSS family + # filter must drop Claude/Gemini and keep only the OSS cohort. + payload = _mlflow_chat_payload( + [ + "databricks-claude-opus-4-8", + "databricks-gemini-2-5-pro", + "databricks-glm-5-2", + "databricks-qwen3-embedding-0-6b", + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert reason is None + assert models == ["databricks-glm-5-2"] + + def test_reports_reason_when_no_oss_family_matches(self, monkeypatch): + payload = _mlflow_chat_payload(["databricks-claude-opus-4-8"]) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_oss_models(WS, "token") + + assert models == [] + assert reason is not None + assert "no OSS" in reason + + class TestResolvePatToken: def test_reads_pat_profile_token_from_cfg(self, monkeypatch, tmp_path): cfg = tmp_path / "databrickscfg" diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 30ce2876..efadecbd 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -891,8 +891,23 @@ def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: out.append(("codex", model)) for model in e2e_state.get("gemini_models") or []: out.append(("gemini", model)) + for model in e2e_state.get("oss_models") or []: + out.append(("oss", model)) return out + def test_all_models_includes_oss_provider(self): + models = self._all_models( + { + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": ["gpt-5"], + "gemini_models": ["gemini-3"], + "oss_models": ["system.ai.glm-5-2"], + } + ) + + assert ("oss", "system.ai.glm-5-2") in models + assert len(models) == 4 + def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token): import ucode.config_io as config_io_mod from ucode.agents import pi @@ -903,14 +918,17 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa pytest.skip("No Pi-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - # Pi reads models.json below HOME/.pi/agent. Point both pi's runtime - # HOME and our writer at the same isolated tmp home. + # Point Pi's runtime config and ucode's writers at the same isolated + # directory so models/settings never touch the developer's real config. pi_home = tmp_path / "pi-home" pi_dir = pi_home / ".pi" / "agent" config_path = pi_dir / "models.json" backup_path = tmp_path / "pi-models.backup.json" monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) + monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) failures = [] From da2f96903ff7fc34300fe54752ea0a17467d6030 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:26:28 +1000 Subject: [PATCH 5/8] fix(pi): pin gpt-5 off-thinking level so the codex gateway stops 400ing `_pi_gpt_model_entry` declared `reasoning: True` without an off-state, so for the thinking-off case Pi's Responses builder fell back to `reasoning: {effort: "none"}` (pi-ai openai-responses.js, the `thinkingLevelMap?.off !== null` branch). `"none"` is only valid on gpt-5.1+, so every request to gpt-5, gpt-5-mini, gpt-5-nano and gpt-5-5-pro was rejected: BAD_REQUEST: Unsupported value: 'none' is not supported with the 'gpt-5' model. Supported values are: 'minimal', 'low', 'medium', and 'high'. Setting `thinkingLevelMap: {"off": None}` makes Pi omit `reasoning` entirely, which the gateway accepts for all 14 codex ids. Verified against /ai-gateway/codex/v1/responses: effort="none" 400s on gpt-5/-mini/-nano/-5-5-pro and 200s on gpt-5-1..-5-6; omitting `reasoning` is 200 everywhere. `{"off": "minimal"}` was rejected as an alternative because gpt-5-5-pro 400s on it too. Same pattern already used for the Gemini 3.x entries. The rest of Pi's Responses payload was bisected against the gateway and is fine: store:false, prompt_cache_key, prompt_cache_retention:"24h", prompt_cache_options, include:["reasoning.encrypted_content"], developer role, flat tool schemas, and the session_id / x-client-request-id affinity headers. Regression was hard to spot because the gateway returns {"error_code","message"} rather than OpenAI's {"error":...}, so Pi's error-body.js recovery no-ops and every 400 renders as "OpenAI API error (400): 400 status code (no body)". Reported upstream as earendil-works/pi#7748. Refs #286 --- src/ucode/agents/pi.py | 18 +++++++++++++++++- tests/test_agent_pi.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e9efd05c..66bcf41a 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -19,6 +19,10 @@ - mlflow: `supportsStore: false` and `supportsStrictMode: false` — the MLflow chat-completions gateway rejects OpenAI's `store` field and `tools[].function.strict`. +- openai: no `compat` flags needed, but the per-model `thinkingLevelMap` + matters — see `_pi_gpt_model_entry`. Declaring `reasoning: true` without an + off-state makes Pi send `reasoning: {effort: "none"}`, which `gpt-5`, + `gpt-5-mini`, `gpt-5-nano` and `gpt-5-5-pro` reject with a 400. The `databricks-mlflow` provider carries the validated OSS coding models (GLM and Kimi) discovered upstream. Per model it sets @@ -157,7 +161,18 @@ def _pi_gpt_model_entry(model_id: str) -> dict: """Build a Pi openai (codex) model entry with `contextWindow`/`maxTokens` from `databricks.gpt_model_token_limits`. GPT ids aren't in Pi's built-in catalog, so without an explicit window Pi falls back to a small default and - truncates long sessions.""" + truncates long sessions. + + `thinkingLevelMap: {"off": None}` is required alongside `reasoning: True`. + When a model declares `reasoning` but no off-state, Pi's Responses builder + falls back to `reasoning: {effort: "none"}` for the thinking-off case + (`pi-ai/dist/api/openai-responses.js`, the `thinkingLevelMap?.off !== null` + branch). `"none"` is only valid on gpt-5.1+, so `gpt-5`, `gpt-5-mini`, + `gpt-5-nano` and `gpt-5-5-pro` reject every request with + `BAD_REQUEST: Unsupported value: 'none' is not supported with the 'gpt-5' + model`. An explicit `None` makes Pi omit `reasoning` entirely, which the + gateway accepts for all ids (verified against /ai-gateway/codex/v1). + """ limits = gpt_model_token_limits(model_id) entry: dict = { "id": model_id, @@ -167,6 +182,7 @@ def _pi_gpt_model_entry(model_id: str) -> dict: if "gpt-5" in model_id.lower().replace(".", "-"): entry["reasoning"] = True entry["input"] = ["text", "image"] + entry["thinkingLevelMap"] = {"off": None} return entry diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 896cc31c..1daed53b 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -89,6 +89,34 @@ def test_gpt56_sol_model_entry_pins_1m_context(self): assert entry["reasoning"] is True assert entry["input"] == ["text", "image"] + def test_gpt_entries_pin_off_thinking_level_to_none(self): + # `reasoning: True` without an off-state makes Pi send + # `reasoning: {effort: "none"}`, which gpt-5 / -mini / -nano / -5-5-pro + # reject with a 400. `{"off": None}` makes Pi omit `reasoning`. + overlay, _ = _overlay( + "system.ai.gpt-5", + codex_models=[ + "system.ai.gpt-5", + "system.ai.gpt-5-mini", + "system.ai.gpt-5-nano", + "system.ai.gpt-5-5-pro", + "system.ai.gpt-5-6-luna", + ], + ) + entries = overlay["providers"]["databricks-openai"]["models"] + assert entries, "expected gpt entries" + for entry in entries: + assert entry["reasoning"] is True + assert entry["thinkingLevelMap"] == {"off": None}, entry["id"] + + def test_non_gpt_codex_entry_has_no_thinking_level_map(self): + # Only the gpt-5 family declares `reasoning`, so only it needs the + # off-state override. + overlay, _ = _overlay("gpt-oss-120b", codex_models=["gpt-oss-120b"]) + entry = overlay["providers"]["databricks-openai"]["models"][0] + assert "reasoning" not in entry + assert "thinkingLevelMap" not in entry + def test_gpt_model_entries_use_model_specific_windows(self): overlay, _ = _overlay( "system.ai.gpt-5-2", From 347f81d2a45077631dc922aebad39e3e139d5499 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:22:27 +1000 Subject: [PATCH 6/8] feat: discover OSS MLflow capabilities and model metadata --- src/ucode/agents/opencode.py | 74 ++++++++-- src/ucode/agents/pi.py | 92 ++++++++++--- src/ucode/cli.py | 22 ++- src/ucode/databricks.py | 252 ++++++++++++++++++++++++++++++----- tests/conftest.py | 10 +- tests/test_agent_opencode.py | 117 ++++++++++++++++ tests/test_agent_pi.py | 99 ++++++++++++++ tests/test_cli.py | 102 +++++++++++++- tests/test_databricks.py | 202 +++++++++++++++++++++++++++- 9 files changed, 896 insertions(+), 74 deletions(-) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 10eda8a9..0056ceeb 100644 --- a/src/ucode/agents/opencode.py +++ b/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 ( @@ -83,17 +84,67 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) - 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} - 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 _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 + + +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 @@ -111,6 +162,7 @@ def render_overlay( 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}"} @@ -177,6 +229,7 @@ def render_overlay( } 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": { @@ -184,7 +237,7 @@ def render_overlay( "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"]) @@ -214,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") diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 7950b23a..a6d6a1aa 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -24,13 +24,9 @@ off-state makes Pi send `reasoning: {effort: "none"}`, which `gpt-5`, `gpt-5-mini`, `gpt-5-nano` and `gpt-5-5-pro` reject with a 400. -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 `databricks-mlflow` provider carries validated MLflow chat-completions +models discovered upstream. Per-model reasoning and token metadata comes from +the persisted gateway capability specs, with conservative static fallback. The bearer token is baked into the file and refreshed by a background thread while the session runs (same pattern as OpenCode/Copilot). @@ -42,6 +38,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 ( @@ -141,21 +138,71 @@ def _pi_claude_model_entry(model_id: str) -> dict: return entry -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.""" +_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 + + +def _pi_oss_model_entry(model_id: str, spec: dict[str, object] | None = None) -> dict: + """Build a Pi MLflow model entry from discovered or static capabilities. + + A valid discovered boolean overrides static reasoning. Any discovered spec + receives a complete conservative limit pair, so missing capability fields + cannot leave a validated model effectively uncapped. Missing specs retain + the existing static GLM/Kimi behavior, while unknown models remain bare. + """ entry: dict = {"id": model_id} - if model_is_reasoning(model_id): + static_limits = model_token_limits(model_id) + static_reasoning = model_is_reasoning(model_id) + + reasoning = spec.get("reasoning") if isinstance(spec, dict) else None + if not isinstance(reasoning, bool): + reasoning = static_reasoning + if reasoning: 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"] + + 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): + entry["contextWindow"] = context or ( + static_limits.get("context") if static_limits else _OSS_SAFE_LIMITS["context"] + ) + entry["maxTokens"] = output or ( + static_limits.get("output") if static_limits else _OSS_SAFE_LIMITS["output"] + ) + elif static_limits: + entry["contextWindow"] = static_limits["context"] + entry["maxTokens"] = static_limits["output"] return entry @@ -196,6 +243,7 @@ def render_overlay( codex_models: list[str], gemini_models: list[str], oss_models: list[str], + oss_specs: list[dict] | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Pi's private agent config.""" providers: dict = {} @@ -240,6 +288,7 @@ def render_overlay( } keys.append(["providers", "databricks-gemini"]) if oss_models: + specs_by_id = _oss_specs_by_id(oss_specs) providers["databricks-mlflow"] = { "baseUrl": pi_base_urls["oss"], "api": "openai-completions", @@ -249,7 +298,7 @@ def render_overlay( # 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], + "models": [_pi_oss_model_entry(m, specs_by_id.get(m)) for m in oss_models], } keys.append(["providers", "databricks-mlflow"]) overlay: dict = { @@ -289,6 +338,7 @@ def write_tool_config( codex_models, gemini_models, state.get("oss_models") or [], + state.get("oss_model_specs") or [], ) existing = read_json_safe(PI_CONFIG_PATH) providers = existing.get("providers") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0e8b6a6f..75e9c7c4 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -41,7 +41,7 @@ discover_codex_models, discover_gemini_models, discover_model_services, - discover_oss_models, + discover_oss_model_specs, ensure_ai_gateway_v2, ensure_databricks_auth, ensure_pat_bearer, @@ -492,6 +492,7 @@ def configure_shared_state( gemini_models = [] codex_models = [] oss_models = [] + oss_specs: list[dict] = [] opencode_models: dict[str, list[str]] = {} web_search_model: str | None = None if skip_model_discovery: @@ -534,8 +535,22 @@ def configure_shared_state( codex_models, codex_reason = discover_codex_models(workspace, token) if want_oss: oss_models, oss_reason = ms_oss, ms_reason - if not oss_models: - oss_models, oss_reason = discover_oss_models(workspace, token) + if oss_models: + oss_specs, specs_reason = discover_oss_model_specs(workspace, token, oss_models) + # Keep IDs and specs aligned. Broad OSS families are admitted + # only by live capability validation; if that refresh fails, + # offering the stale IDs without safe metadata would regress + # them to uncapped client defaults. Static GLM/Kimi fallback + # specs are still returned by discover_oss_model_specs. + oss_models = [spec["id"] for spec in oss_specs] + if not oss_specs and specs_reason: + oss_reason = specs_reason + else: + # The endpoint fallback returns ids and capabilities from + # the same validated listing, avoiding a second request + # whose transient failure could leave broad models uncapped. + oss_specs, oss_reason = discover_oss_model_specs(workspace, token) + oss_models = [spec["id"] for spec in oss_specs] if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: @@ -560,6 +575,7 @@ def configure_shared_state( state["codex_models"] = codex_models if want_oss: state["oss_models"] = oss_models + state["oss_model_specs"] = oss_specs if fetch_all or "opencode" in tools: state["opencode_models"] = opencode_models save_state(state) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 96c17fc0..e5b883c7 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1316,14 +1316,20 @@ def build_auth_shell_command( # Databricks-managed foundation models under `system.ai`. _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai." -# 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 families with statically validated coding-agent behavior. They remain the +# safe fallback when the workspace's foundation-model capability listing is +# unavailable. Other model families are offered only after the listing confirms +# that they are served exclusively through MLflow chat completions. _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. +# Native routes take precedence over the generic MLflow chat-completions route. +# A foundation model advertising any of these must not be duplicated as OSS. +_NATIVE_PROVIDER_API_TYPES = frozenset( + {"anthropic/v1/messages", "openai/v1/responses", "gemini/v1/generateContent"} +) + +# Non-chat services must never be offered to a chat agent, even if malformed +# metadata happens to advertise a chat API type. _OSS_NON_CHAT_SUBSTRINGS = ("embedding", "embed", "rerank") @@ -1411,6 +1417,202 @@ def model_token_limits(model_id: str) -> dict[str, int] | None: return None +# The foundation-model API exposes context windows only in free-text +# descriptions (for example "context length of 1M tokens"). Keep parsing +# deliberately narrow: unrecognized or invalid text simply yields no override. +_CONTEXT_LENGTH_RE = re.compile(r"context (?:length|window) of ([\d.,]+)\s*([MK])", re.IGNORECASE) + + +def _parse_context_window(description: str) -> int | None: + if not isinstance(description, str): + return None + match = _CONTEXT_LENGTH_RE.search(description) + if not match: + return None + try: + value = float(match.group(1).replace(",", "")) + multiplier = 1_000_000 if match.group(2).upper() == "M" else 1_000 + tokens = int(value * multiplier) + except (OverflowError, ValueError): + return None + return tokens if tokens > 0 else None + + +# Per-model output ceilings enforced by the MLflow gateway. There is no +# structured metadata field for these values; they were established by probing +# oversized requests. Keys omit route prefixes so the same entry applies to +# both `databricks-*` endpoint ids and `system.ai.*` model-service ids. +_OSS_MAX_OUTPUT_TOKENS: dict[str, int] = { + "glm-5-2": 65_536, + "inkling": 65_536, + "kimi-k2-7-code": 65_536, + "gpt-oss-120b": 25_000, + "gpt-oss-20b": 25_000, + "qwen35-122b-a10b": 25_000, + "qwen3-next-80b-a3b-instruct": 10_000, + "llama-4-maverick": 8_192, + "meta-llama-3-1-8b-instruct": 8_192, + "meta-llama-3-3-70b-instruct": 8_192, + "gemma-3-12b": 8_192, +} + + +def _canonical_oss_model_id(model_id: str) -> str: + """Normalize endpoint/model-service ids for capability matching.""" + tail = model_id.rsplit("/", 1)[-1].strip().lower() + if tail.startswith("system.ai."): + tail = tail[len("system.ai.") :] + if tail.startswith("databricks-"): + tail = tail[len("databricks-") :] + return tail + + +def _static_oss_spec(model_id: str) -> dict | None: + """Capability fallback for the statically validated GLM/Kimi families.""" + if not _is_oss_chat_model(model_id.lower()): + return None + limits = model_token_limits(model_id) or {} + return { + "id": model_id, + "reasoning": model_is_reasoning(model_id.lower()), + "context_window": limits.get("context"), + "max_tokens": limits.get("output"), + } + + +def _oss_specs_from_foundation_models(payload: object) -> list[dict]: + """Parse validated chat-completions-only model specs from a listing.""" + if not isinstance(payload, dict): + return [] + payload_dict = cast(dict[str, object], payload) + raw_endpoints = payload_dict.get("endpoints") + if not isinstance(raw_endpoints, list): + return [] + + specs: list[dict] = [] + for endpoint in raw_endpoints: + if not isinstance(endpoint, dict): + continue + endpoint_dict = cast(dict[str, object], endpoint) + name = endpoint_dict.get("name") + config = endpoint_dict.get("config") + if not isinstance(name, str) or not name.strip() or not isinstance(config, dict): + continue + name = name.strip() + lowered_name = _canonical_oss_model_id(name) + is_native_family = ( + lowered_name.startswith("claude-") + or lowered_name.startswith("gemini-") + or re.match(r"^gpt-\d(?:-|$)", lowered_name) is not None + ) + if is_native_family or any(bad in lowered_name for bad in _OSS_NON_CHAT_SUBSTRINGS): + continue + config_dict = cast(dict[str, object], config) + entities = config_dict.get("served_entities") + if not isinstance(entities, list): + continue + + api_types: set[str] = set() + description = "" + has_v2_entity = False + for entity in entities: + if not isinstance(entity, dict): + continue + entity_dict = cast(dict[str, object], entity) + foundation_model = entity_dict.get("foundation_model") + if not isinstance(foundation_model, dict): + continue + foundation_model_dict = cast(dict[str, object], foundation_model) + if foundation_model_dict.get("ai_gateway_v2_supported") is not True: + continue + has_v2_entity = True + raw_api_types = foundation_model_dict.get("api_types") + if isinstance(raw_api_types, list): + api_types.update(value for value in raw_api_types if isinstance(value, str)) + raw_description = foundation_model_dict.get("description") + if not description and isinstance(raw_description, str): + description = raw_description + + if not has_v2_entity or "mlflow/v1/chat/completions" not in api_types: + continue + if api_types & _NATIVE_PROVIDER_API_TYPES: + continue + capabilities = endpoint_dict.get("capabilities") + capabilities_dict = ( + cast(dict[str, object], capabilities) if isinstance(capabilities, dict) else {} + ) + reasoning = capabilities_dict.get("openai_reasoning") is True + canonical_id = _canonical_oss_model_id(name) + context_window = _parse_context_window(description) + max_tokens = _OSS_MAX_OUTPUT_TOKENS.get(canonical_id) + static_fallback = _static_oss_spec(name) + if static_fallback is not None: + # Missing/partial metadata must not regress the statically verified + # GLM/Kimi capabilities already used by existing installations. + reasoning = reasoning or static_fallback["reasoning"] + context_window = context_window or static_fallback["context_window"] + max_tokens = max_tokens or static_fallback["max_tokens"] + specs.append( + { + "id": name, + "reasoning": reasoning, + "context_window": context_window, + "max_tokens": max_tokens, + } + ) + # Foundation listings can repeat a served endpoint. Emit one stable spec + # per canonical model so downstream model lists/configs stay deduplicated. + deduped: dict[str, dict] = {} + for spec in sorted(specs, key=lambda item: item["id"]): + deduped.setdefault(_canonical_oss_model_id(spec["id"]), spec) + return list(deduped.values()) + + +def discover_oss_model_specs( + workspace: str, + token: str, + model_ids: list[str] | None = None, +) -> tuple[list[dict], str | None]: + """Discover validated MLflow chat-completions models and capabilities. + + With ``model_ids`` (the UC-first path), endpoint capabilities are projected + back onto those exact ids using their normalized model name. Statically + validated GLM/Kimi ids remain available when capability discovery fails or + omits them. Without ``model_ids`` (the serving-endpoint fallback), only + models validated by the live API metadata are returned. + """ + hostname = workspace_hostname(workspace) + payload, reason = _http_get_json( + f"https://{hostname}/api/2.0/serving-endpoints:foundation-models", token + ) + discovered = _oss_specs_from_foundation_models(payload) + + if model_ids is None: + if discovered: + return discovered, None + if payload is None: + return [], reason + return [], "no validated chat-completions-only OSS endpoints" + + discovered_by_id = {_canonical_oss_model_id(spec["id"]): spec for spec in discovered} + specs: list[dict] = [] + for model_id in model_ids: + if not isinstance(model_id, str) or not model_id.strip(): + continue + dynamic = discovered_by_id.get(_canonical_oss_model_id(model_id)) + if dynamic is not None: + specs.append({**dynamic, "id": model_id}) + continue + fallback = _static_oss_spec(model_id) + if fallback is not None: + specs.append(fallback) + if specs: + return specs, None + if payload is None: + return [], reason + return [], "requested model ids matched no validated OSS endpoint" + + # Pi treats every custom model without explicit metadata as 128k context / 4k # output. Gateway ids are custom ids (not Pi's built-ins), so preserve the # upstream windows explicitly. Entries are ordered most-specific first after @@ -1728,7 +1930,11 @@ def discover_model_services( 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 _is_oss_chat_model(m)] + # Project the live endpoint capability listing onto the UC ids. This + # broadens discovery beyond the static GLM/Kimi fallback only when the + # corresponding endpoint is validated as MLflow chat-completions-only. + oss_specs, _ = discover_oss_model_specs(workspace, token, ids) + oss_models = [spec["id"] for spec in oss_specs] if not (claude_models or codex_models or gemini_models or oss_models): sample = ", ".join(ids[:5]) @@ -2804,31 +3010,15 @@ def discover_codex_models(workspace: str, token: str) -> tuple[list[str], str | def discover_oss_models(workspace: str, token: str) -> tuple[list[str], str | None]: - """Discover OSS chat models served as AI Gateway foundation-model endpoints. - - Fallback for workspaces that don't register OSS foundation models as - `system.ai.*` UC model-services (see `discover_model_services`): those - workspaces expose the same models as regular `databricks-*` serving - endpoints instead. Lists every endpoint advertising the - `mlflow/v1/chat/completions` dialect, then keeps only the OSS chat families - (`_is_oss_chat_model`) — on some workspaces the Claude/Gemini endpoints also - advertise that dialect, so the family filter is what separates the OSS - cohort from them. Mirrors the AI-Gateway fallback the other families use - when the UC model-services listing is empty. + """Discover validated chat-completions-only serving endpoints. + + This is the fallback for workspaces without UC model-services. Unlike the + static family fallback used for UC ids, every endpoint returned here has + live metadata confirming AI Gateway v2 MLflow chat completions and no + competing native Anthropic, Responses, or Gemini route. """ - endpoints, reason = discover_endpoints_with_api_type( - workspace, token, "mlflow/v1/chat/completions" - ) - if not endpoints: - return [], reason - oss = [e for e in endpoints if _is_oss_chat_model(e)] - if oss: - return oss, None - sample = ", ".join(endpoints[:5]) - return [], ( - "foundation-models exposing `mlflow/v1/chat/completions` matched no OSS " - f"chat family (got: {sample})" - ) + specs, reason = discover_oss_model_specs(workspace, token) + return [spec["id"] for spec in specs], reason def fetch_gemini_models(workspace: str, token: str) -> list[str]: diff --git a/tests/conftest.py b/tests/conftest.py index 65ae3683..08b5ec82 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,7 +12,7 @@ discover_codex_models, discover_gemini_models, discover_model_services, - discover_oss_models, + discover_oss_model_specs, get_databricks_token, ) from ucode.ui import normalize_workspace_url @@ -68,8 +68,11 @@ def e2e_state(e2e_workspace, e2e_token): gemini_models, _ = discover_gemini_models(e2e_workspace, e2e_token) if not codex_models: codex_models, _ = discover_codex_models(e2e_workspace, e2e_token) - if not oss_models: - oss_models, _ = discover_oss_models(e2e_workspace, e2e_token) + if oss_models: + oss_model_specs, _ = discover_oss_model_specs(e2e_workspace, e2e_token, oss_models) + else: + oss_model_specs, _ = discover_oss_model_specs(e2e_workspace, e2e_token) + oss_models = [spec["id"] for spec in oss_model_specs] # E2E mirrors configure's default (Fable is premium and opt-in). claude_models.pop("fable", None) @@ -90,6 +93,7 @@ def e2e_state(e2e_workspace, e2e_token): "gemini_models": gemini_models, "codex_models": codex_models, "oss_models": oss_models, + "oss_model_specs": oss_model_specs, "opencode_models": opencode_models, "base_urls": build_shared_base_urls(e2e_workspace), "managed_configs": {}, diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 6d424382..29780fc3 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -117,6 +117,93 @@ def test_uncapped_oss_model_has_no_limit(self): 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 + assert "reasoning" not in entry + + def test_dynamic_full_spec_sets_reasoning_and_limits(self): + models = {"oss": ["system.ai.qwen35-122b-a10b"]} + specs = [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 262_144, + "max_tokens": 25_000, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.qwen35-122b-a10b", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.qwen35-122b-a10b"] + assert entry["reasoning"] is True + assert entry["limit"] == {"context": 262_144, "output": 25_000} + + def test_dynamic_reasoning_false_is_respected(self): + models = {"oss": ["system.ai.glm-5-2"]} + specs = [ + { + "id": "system.ai.glm-5-2", + "reasoning": False, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.glm-5-2", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"] + assert entry["reasoning"] is False + assert entry["limit"] == {"context": 1_000_000, "output": 65_536} + + def test_unknown_dynamic_spec_gets_safe_complete_limit_pair(self): + models = {"oss": ["system.ai.deepseek-v3"]} + specs = [ + { + "id": "system.ai.deepseek-v3", + "reasoning": False, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.deepseek-v3", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.deepseek-v3"] + assert entry["reasoning"] is False + assert entry["limit"] == {"context": 128_000, "output": 8_192} + + def test_partial_dynamic_limit_is_completed_as_valid_pair(self): + models = {"oss": ["system.ai.inkling"]} + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": 65_536, + } + ] + overlay, _ = opencode.render_overlay( + "system.ai.inkling", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.inkling"] + assert entry["limit"] == {"context": 128_000, "output": 65_536} + + def test_malformed_dynamic_spec_is_ignored_safely(self): + models = {"oss": ["system.ai.mystery-7b"]} + specs = [ + None, + {"id": 12, "reasoning": True}, + { + "id": "system.ai.mystery-7b", + "reasoning": "true", + "context_window": 0, + "max_tokens": True, + }, + ] + overlay, _ = opencode.render_overlay( + "system.ai.mystery-7b", "tok", _base_urls(), models, specs + ) + entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.mystery-7b"] + assert "reasoning" not in entry + assert "limit" not in entry def test_token_in_api_key(self): models = {"anthropic": ["claude-sonnet"]} @@ -539,3 +626,33 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): written = json.loads(config_file.read_text()) assert written["model"] == "databricks-anthropic/claude-sonnet" + + def test_state_oss_specs_reach_written_model_entry(self, tmp_path, monkeypatch): + import ucode.agents.opencode as oc_mod + + config_file = tmp_path / "opencode.json" + monkeypatch.setattr(oc_mod, "OPENCODE_CONFIG_PATH", config_file) + monkeypatch.setattr(oc_mod, "OPENCODE_BACKUP_PATH", tmp_path / "opencode-backup.json") + state = { + "workspace": WS, + "base_urls": {"opencode": _base_urls()}, + "opencode_models": {"oss": ["system.ai.inkling"]}, + "oss_model_specs": [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": 256_000, + "max_tokens": 65_536, + } + ], + "managed_configs": {}, + } + + with patch("ucode.agents.opencode.save_state"): + oc_mod.write_tool_config(state, "system.ai.inkling", token="tok") + + entry = json.loads(config_file.read_text())["provider"]["databricks-oss"]["models"][ + "system.ai.inkling" + ] + assert entry["reasoning"] is True + assert entry["limit"] == {"context": 256_000, "output": 65_536} diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 42b39b2d..87e9c6c8 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -28,6 +28,7 @@ def _empty() -> dict: "codex_models": [], "gemini_models": [], "oss_models": [], + "oss_specs": [], } @@ -42,6 +43,7 @@ def _overlay(model: str, token: str = "tok", **kwargs): bundle["codex_models"], bundle["gemini_models"], bundle["oss_models"], + bundle["oss_specs"], ) @@ -219,6 +221,80 @@ def test_unknown_oss_model_bare(self): "id": "system.ai.mystery-7b" } + def test_dynamic_full_spec_overrides_static_metadata(self): + specs = [ + { + "id": "system.ai.glm-5-2", + "reasoning": False, + "context_window": 256_000, + "max_tokens": 12_345, + } + ] + overlay, _ = _overlay( + "system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"], oss_specs=specs + ) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry == { + "id": "system.ai.glm-5-2", + "contextWindow": 256_000, + "maxTokens": 12_345, + } + + def test_dynamic_reasoning_true_is_applied_with_safe_unknown_limits(self): + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": None, + } + ] + overlay, _ = _overlay( + "system.ai.inkling", oss_models=["system.ai.inkling"], oss_specs=specs + ) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.inkling", + "reasoning": True, + "contextWindow": 128_000, + "maxTokens": 8_192, + } + + def test_partial_dynamic_limits_are_completed_conservatively(self): + specs = [ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": None, + "max_tokens": 65_536, + } + ] + overlay, _ = _overlay( + "system.ai.inkling", oss_models=["system.ai.inkling"], oss_specs=specs + ) + entry = overlay["providers"]["databricks-mlflow"]["models"][0] + assert entry["contextWindow"] == 128_000 + assert entry["maxTokens"] == 65_536 + + def test_malformed_spec_is_ignored_safely(self): + specs = [ + None, + {"id": 12, "reasoning": True}, + { + "id": "system.ai.mystery-7b", + "reasoning": "yes", + "context_window": -1, + "max_tokens": True, + }, + ] + overlay, _ = _overlay( + "system.ai.mystery-7b", + oss_models=["system.ai.mystery-7b"], + oss_specs=specs, + ) + assert overlay["providers"]["databricks-mlflow"]["models"][0] == { + "id": "system.ai.mystery-7b" + } + class TestRenderOverlayUserAgent: def test_user_agent_set_on_all_three_providers(self, monkeypatch): @@ -515,6 +591,29 @@ def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch assert written["model"] == "databricks-claude/claude-sonnet" assert written["providers"]["databricks-claude"]["apiKey"] == "tok" + def test_state_oss_specs_reach_written_model_entry(self, tmp_path, monkeypatch): + pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch) + state = self._state( + claude_models={}, + oss_models=["system.ai.inkling"], + oss_model_specs=[ + { + "id": "system.ai.inkling", + "reasoning": True, + "context_window": 256_000, + "max_tokens": 65_536, + } + ], + ) + + with patch("ucode.agents.pi.save_state"): + pi_mod.write_tool_config(state, "system.ai.inkling", token="tok") + + entry = json.loads(config_file.read_text())["providers"]["databricks-mlflow"]["models"][0] + assert entry["reasoning"] is True + assert entry["contextWindow"] == 256_000 + assert entry["maxTokens"] == 65_536 + def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): # Without this, Pi's `findInitialModel` can fall through to a built-in # provider when an unrelated env var (e.g. HF_TOKEN) makes one look diff --git a/tests/test_cli.py b/tests/test_cli.py index c4307630..9015af57 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1634,7 +1634,22 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) - monkeypatch.setattr(cli_mod, "discover_oss_models", lambda w, t: ([], None)) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + [ + { + "id": model_id, + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + for model_id in (model_ids or []) + ], + None, + ), + ) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) return cli_mod, logins, ensures, saved @@ -1725,6 +1740,62 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): assert legacy_called == [] assert "uc_enabled" not in state + def test_uc_oss_ids_persist_matching_capability_specs(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({}, [], [], ["system.ai.qwen35-122b-a10b"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ], + None, + ), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["oss_models"] == ["system.ai.qwen35-122b-a10b"] + assert state["oss_model_specs"] == [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ] + assert state["opencode_models"]["oss"] == ["system.ai.qwen35-122b-a10b"] + + def test_uc_dynamic_oss_ids_are_dropped_when_spec_refresh_fails(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({}, [], [], ["system.ai.inkling"], None), + ) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ([], "HTTP 503 unavailable"), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["oss_models"] == [] + assert state["oss_model_specs"] == [] + assert "oss" not in state["opencode_models"] + assert state["_discovery_reasons"]["oss"] == "HTTP 503 unavailable" + def _stub_with_fable(self, monkeypatch): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") monkeypatch.setattr( @@ -1860,8 +1931,19 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): ) monkeypatch.setattr( cli_mod, - "discover_oss_models", - lambda w, t: (calls.append("oss") or ["databricks-glm-5-2"], None), + "discover_oss_model_specs", + lambda w, t, model_ids=None: ( + calls.append("oss") + or [ + { + "id": "databricks-glm-5-2", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + ], + None, + ), ) state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") @@ -1873,6 +1955,14 @@ def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): } assert state["codex_models"] == ["databricks-gpt-5-6-sol"] assert state["oss_models"] == ["databricks-glm-5-2"] + assert state["oss_model_specs"] == [ + { + "id": "databricks-glm-5-2", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 8_192, + } + ] assert state["opencode_models"] == { "anthropic": [ "databricks-claude-opus-4-8", @@ -1946,7 +2036,11 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) - monkeypatch.setattr(cli_mod, "discover_oss_models", lambda w, t: ([], None)) + monkeypatch.setattr( + cli_mod, + "discover_oss_model_specs", + lambda w, t, model_ids=None: ([], None), + ) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) def test_purges_residue_when_workspace_changes(self, monkeypatch): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 22489d38..98cbae4b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1326,7 +1326,13 @@ def test_finds_oss_endpoints_via_foundation_models(self, monkeypatch): models, reason = db_mod.discover_oss_models(WS, "token") assert reason is None - assert models == ["databricks-glm-5-2", "databricks-kimi-k2-7-code"] + assert models == [ + "databricks-gemma-3-12b", + "databricks-glm-5-2", + "databricks-inkling", + "databricks-kimi-k2-7-code", + "databricks-qwen35-122b-a10b", + ] def test_excludes_claude_and_gemini_sharing_the_mlflow_dialect(self, monkeypatch): # On some workspaces every foundation model advertises the mlflow chat @@ -1355,7 +1361,199 @@ def test_reports_reason_when_no_oss_family_matches(self, monkeypatch): assert models == [] assert reason is not None - assert "no OSS" in reason + assert "OSS" in reason + + +class TestDiscoverOssModelSpecs: + def test_parses_reasoning_context_and_known_output_cap(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-inkling", + "capabilities": {"openai_reasoning": True}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "Supports a context window of 1.5M tokens.", + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert specs == [ + { + "id": "databricks-inkling", + "reasoning": True, + "context_window": 1_500_000, + "max_tokens": 65_536, + } + ] + + def test_excludes_endpoint_with_native_api_and_malformed_entries(self, monkeypatch): + payload = { + "endpoints": [ + None, + {"name": "broken", "config": {"served_entities": "bad"}}, + { + "name": "databricks-qwen35-122b-a10b", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [ + "mlflow/v1/chat/completions", + "openai/v1/responses", + ], + } + } + ] + }, + }, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert specs == [] + assert reason is not None + + def test_v2_and_mlflow_type_must_belong_to_same_entity(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-inkling", + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": [], + } + }, + { + "foundation_model": { + "ai_gateway_v2_supported": False, + "api_types": ["mlflow/v1/chat/completions"], + } + }, + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert specs == [] + assert reason is not None + + def test_duplicate_endpoint_ids_are_deduplicated(self, monkeypatch): + payload = _mlflow_chat_payload(["databricks-inkling", "databricks-inkling"]) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token") + + assert reason is None + assert [spec["id"] for spec in specs] == ["databricks-inkling"] + + def test_uc_ids_receive_matching_endpoint_capabilities(self, monkeypatch): + payload = { + "endpoints": [ + { + "name": "databricks-qwen35-122b-a10b", + "capabilities": {"openai_reasoning": True}, + "config": { + "served_entities": [ + { + "foundation_model": { + "ai_gateway_v2_supported": True, + "api_types": ["mlflow/v1/chat/completions"], + "description": "context length of 128K tokens", + } + } + ] + }, + } + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, reason = db_mod.discover_oss_model_specs(WS, "token", ["system.ai.qwen35-122b-a10b"]) + + assert reason is None + assert specs == [ + { + "id": "system.ai.qwen35-122b-a10b", + "reasoning": True, + "context_window": 128_000, + "max_tokens": 25_000, + } + ] + + def test_unavailable_metadata_keeps_static_glm_kimi_fallback(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token: (None, "HTTP 503 unavailable") + ) + + specs, reason = db_mod.discover_oss_model_specs( + WS, + "token", + ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code", "system.ai.inkling"], + ) + + assert reason is None + assert [spec["id"] for spec in specs] == [ + "system.ai.glm-5-2", + "system.ai.kimi-k2-7-code", + ] + + @pytest.mark.parametrize("description", ["", "context length of nope", "context window of 0K"]) + def test_malformed_context_description_degrades_to_none(self, monkeypatch, description): + payload = _mlflow_chat_payload(["databricks-inkling"]) + payload["endpoints"][0]["config"]["served_entities"][0]["foundation_model"][ + "description" + ] = description + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + specs, _ = db_mod.discover_oss_model_specs(WS, "token") + + assert specs[0]["context_window"] is None + + +class TestDiscoverModelServicesDynamicOss: + def test_uc_first_broad_model_requires_matching_endpoint_validation(self, monkeypatch): + model_services = { + "model_services": [ + _model_service("system.ai.qwen35-122b-a10b"), + _model_service("system.ai.inkling"), + ] + } + foundation_models = _mlflow_chat_payload(["databricks-qwen35-122b-a10b"]) + + def fake_get(url, token, timeout=10): + if "model-services" in url: + return model_services, None + return foundation_models, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + _, _, _, oss, reason = db_mod.discover_model_services(WS, "token") + + assert reason is None + assert oss == ["system.ai.qwen35-122b-a10b"] class TestResolvePatToken: From 2a731444789f40014efe98e6b787ec833a71107b Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:24:08 +1000 Subject: [PATCH 7/8] fix(pi): repair MLflow SSE stream termination --- src/ucode/agents/_mlflow_proxy.py | 255 ++++++++++++++++++++++++ src/ucode/agents/pi.py | 131 ++++++++++-- tests/test_agent_pi.py | 208 ++++++++++++++++++- tests/test_mlflow_proxy.py | 319 ++++++++++++++++++++++++++++++ 4 files changed, 897 insertions(+), 16 deletions(-) create mode 100644 src/ucode/agents/_mlflow_proxy.py create mode 100644 tests/test_mlflow_proxy.py diff --git a/src/ucode/agents/_mlflow_proxy.py b/src/ucode/agents/_mlflow_proxy.py new file mode 100644 index 00000000..04973c98 --- /dev/null +++ b/src/ucode/agents/_mlflow_proxy.py @@ -0,0 +1,255 @@ +"""Loopback SSE-repair proxy for Pi's MLflow chat-completions provider. + +Some MLflow-served models omit the terminal OpenAI ``finish_reason``. Pi's +strict ``openai-completions`` parser rejects those streams. This loopback-only +proxy forwards requests without logging credentials or bodies and repairs only +successful SSE responses that have already produced data. Healthy SSE and all +non-streaming/error responses pass through unchanged. +""" + +from __future__ import annotations + +import json +from email.message import Message +from http.client import IncompleteRead +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import IO +from urllib import error as urllib_error +from urllib import request as urllib_request +from urllib.parse import urlsplit + +from ucode.gateway_proxy import _HOP_BY_HOP +from ucode.ui import print_warning + +_STREAM_CHUNK = 8192 +_CHAT_COMPLETIONS_PATH = "/ai-gateway/mlflow/v1/chat/completions" +_SKIP_REQUEST_HEADERS = _HOP_BY_HOP | {"accept-encoding"} +_ERROR_BODY = b'{"error":"MLflow proxy upstream unavailable"}\n' + + +class _NoRedirect(urllib_request.HTTPRedirectHandler): + """Keep authenticated requests pinned to the configured workspace origin.""" + + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + +def _finish_chunk(chunk_id: str | None) -> bytes: + payload: dict = { + "object": "chat.completion.chunk", + "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}], + } + if chunk_id is not None: + payload["id"] = chunk_id + return json.dumps(payload, separators=(",", ":")).encode("utf-8") + + +def _data_payload(raw_line: bytes) -> bytes | None: + """Return an SSE data field's payload, accepting the optional one space.""" + stripped = raw_line.rstrip(b"\r\n") + if not stripped.startswith(b"data:"): + return None + payload = stripped[5:] + return payload[1:] if payload.startswith(b" ") else payload + + +def _forwarded_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]: + return { + key: value + for key, value in handler.headers.items() + if key.lower() not in _SKIP_REQUEST_HEADERS + } + + +def _safe_response_headers(headers: Message, *, streaming: bool) -> list[tuple[str, str]]: + safe: list[tuple[str, str]] = [] + for key, value in headers.items(): + lowered = key.lower() + if lowered in _HOP_BY_HOP: + # A non-streaming body is unchanged, so preserving Content-Length + # avoids relying on EOF framing. Repaired streams can change size. + if lowered == "content-length" and not streaming: + safe.append((key, value)) + continue + safe.append((key, value)) + return safe + + +class _ProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + upstream_origin: str + + def log_message(self, format: str, *args: object) -> None: + return + + def do_POST(self) -> None: # noqa: N802 (stdlib handler API) + raw_length = self.headers.get("Content-Length", "0") + try: + length = int(raw_length) + if length < 0: + raise ValueError + except (TypeError, ValueError): + self._fixed_response(400, b'{"error":"invalid Content-Length"}\n') + return + + try: + body = self.rfile.read(length) if length else b"" + except OSError: + return + parsed_target = urlsplit(self.path) + if ( + parsed_target.scheme + or parsed_target.netloc + or parsed_target.fragment + or parsed_target.path != _CHAT_COMPLETIONS_PATH + ): + self._fixed_response(400, b'{"error":"invalid MLflow proxy request target"}\n') + return + target = self.upstream_origin.rstrip("/") + parsed_target.path + if parsed_target.query: + target += f"?{parsed_target.query}" + request = urllib_request.Request( + target, + data=body, + method="POST", + headers=_forwarded_request_headers(self), + ) + try: + opener = urllib_request.build_opener(_NoRedirect) + with opener.open(request, timeout=600) as response: # noqa: S310 + content_type = response.headers.get_content_type().lower() + if content_type == "text/event-stream": + self._relay_sse(response.status, response.headers, response) + else: + self._relay_verbatim(response.status, response.headers, response) + except urllib_error.HTTPError as exc: + # Relay upstream status, headers, and bytes verbatim. Never turn an + # upstream rejection into a successful repaired stream. + self._relay_verbatim(exc.code, exc.headers, exc) + except (urllib_error.URLError, OSError): + self._fixed_response(502, _ERROR_BODY) + + def _send_headers(self, status: int, headers: Message, *, streaming: bool) -> bool: + try: + self.send_response(status) + for key, value in _safe_response_headers(headers, streaming=streaming): + self.send_header(key, value) + # The proxy never reuses downstream connections. EOF framing is + # therefore safe for responses without Content-Length (including + # 204s and repaired SSE), and shutdown cannot leave a keep-alive + # client waiting on an otherwise complete response. + self.send_header("Connection", "close") + self.close_connection = True + self.end_headers() + return True + except (BrokenPipeError, ConnectionResetError, OSError): + return False + + def _relay_verbatim(self, status: int, headers: Message, stream: IO[bytes]) -> None: + if not self._send_headers(status, headers, streaming=False): + return + try: + while chunk := stream.read(_STREAM_CHUNK): + self.wfile.write(chunk) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError, OSError, IncompleteRead): + return + + def _relay_sse(self, status: int, headers: Message, stream: IO[bytes]) -> None: + if not self._send_headers(status, headers, streaming=True): + return + saw_data = False + saw_finish = False + saw_done = False + saw_error = False + last_id: str | None = None + try: + for raw_line in stream: + payload = _data_payload(raw_line) + event_line = raw_line.rstrip(b"\r\n") + if event_line.lower().startswith(b"event:"): + event_name = event_line[6:] + if event_name.startswith(b" "): + event_name = event_name[1:] + if event_name.lower() == b"error": + saw_error = True + if payload == b"[DONE]": + if saw_data and not saw_finish and not saw_error: + self._write(b"data: " + _finish_chunk(last_id) + b"\n\n") + saw_finish = True + self._write(raw_line) + saw_done = True + continue + if payload is not None and payload: + saw_data = True + try: + event = json.loads(payload) + if isinstance(event, dict): + if "error" in event: + saw_error = True + event_id = event.get("id") + if isinstance(event_id, str): + last_id = event_id + choices = event.get("choices") + if isinstance(choices, list) and any( + isinstance(choice, dict) and choice.get("finish_reason") is not None + for choice in choices + ): + saw_finish = True + except (UnicodeDecodeError, json.JSONDecodeError): + pass + self._write(raw_line) + except (BrokenPipeError, ConnectionResetError): + return + except (OSError, IncompleteRead): + # A dropped upstream after data is treated as a truncated stream; + # the repair below gives Pi a structurally valid terminator. + pass + + if saw_data and not saw_error: + try: + if not saw_finish: + self._write(b"data: " + _finish_chunk(last_id) + b"\n\n") + if not saw_done: + self._write(b"data: [DONE]\n\n") + except (BrokenPipeError, ConnectionResetError, OSError): + return + + def _write(self, data: bytes) -> None: + self.wfile.write(data) + self.wfile.flush() + + def _fixed_response(self, status: int, body: bytes) -> None: + try: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError, OSError): + return + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + +def start(upstream_origin: str) -> tuple[ThreadingHTTPServer, str] | None: + """Bind a repair proxy to a fresh loopback port; the caller owns its lifecycle.""" + if not isinstance(upstream_origin, str) or not upstream_origin: + print_warning("MLflow stream repair proxy was not started: invalid upstream URL.") + return None + parsed_origin = urlsplit(upstream_origin) + if parsed_origin.scheme not in {"http", "https"} or not parsed_origin.netloc: + print_warning("MLflow stream repair proxy was not started: invalid upstream URL.") + return None + handler = type("_BoundProxyHandler", (_ProxyHandler,), {"upstream_origin": upstream_origin}) + try: + server = _Server(("127.0.0.1", 0), handler) + except OSError as exc: + print_warning(f"MLflow stream repair proxy was not started ({exc}).") + return None + port = int(server.server_address[1]) + return server, f"http://127.0.0.1:{port}" diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index a6d6a1aa..c43ec468 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -27,6 +27,8 @@ The `databricks-mlflow` provider carries validated MLflow chat-completions models discovered upstream. Per-model reasoning and token metadata comes from the persisted gateway capability specs, with conservative static fallback. +At launch this provider alone is routed through a loopback repair proxy because +some models (notably Inkling) omit the terminal `finish_reason` Pi requires. The bearer token is baked into the file and refreshed by a background thread while the session runs (same pattern as OpenCode/Copilot). @@ -34,13 +36,16 @@ from __future__ import annotations +import ipaddress import os import signal import subprocess import threading from typing import cast +from urllib.parse import urlparse from ucode.agent_updates import available_npm_package_update +from ucode.agents import _mlflow_proxy from ucode.config_io import ( APP_DIR, ToolSpec, @@ -63,6 +68,7 @@ ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version +from ucode.ui import print_warning PI_UCODE_HOME = APP_DIR / "pi-home" PI_CONFIG_DIR = PI_UCODE_HOME / ".pi" / "agent" @@ -443,27 +449,122 @@ def build_runtime_env(token: str) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: - token = _refresh_token_once(state) - env = build_runtime_env(token) +def _is_loopback_origin(origin: str) -> bool: + hostname = urlparse(origin).hostname + if not hostname: + return True + if hostname.lower() == "localhost": + return True + try: + return ipaddress.ip_address(hostname).is_loopback + except ValueError: + return False - stop_event = threading.Event() - refresher = threading.Thread( - target=_refresh_forever, - args=(state, stop_event), - daemon=True, + +def _start_oss_proxy( + state: dict, +) -> tuple[threading.Thread, _mlflow_proxy.ThreadingHTTPServer] | None: + """Start Pi's MLflow repair proxy and rewrite only the in-memory base URL. + + The upstream is always derived from ``workspace`` rather than an existing + base URL, so a stale loopback port from an old config can never become the + next proxy's upstream. + """ + if not (state.get("oss_models") or []): + return None + direct_oss_url = build_pi_base_urls(state["workspace"])["oss"] + origin = direct_oss_url.split("/ai-gateway/", 1)[0] + if _is_loopback_origin(origin): + print_warning("MLflow stream repair proxy skipped for a loopback workspace URL.") + return None + started = _mlflow_proxy.start(origin) + if started is None: + return None + server, proxy_origin = started + thread = threading.Thread(target=server.serve_forever, daemon=True) + try: + thread.start() + except RuntimeError: + server.server_close() + print_warning("MLflow stream repair proxy could not start; using the direct gateway URL.") + return None + pi_urls = state.setdefault("base_urls", {}).setdefault( + "pi", build_pi_base_urls(state["workspace"]) + ) + pi_urls["oss"] = f"{proxy_origin}/ai-gateway/mlflow/v1" + return thread, server + + +def _restore_direct_oss_config(state: dict, token: str | None) -> None: + """Replace the session-only proxy URL before its listener is released.""" + pi_urls = state.setdefault("base_urls", {}).setdefault( + "pi", build_pi_base_urls(state["workspace"]) ) - refresher.start() + pi_urls["oss"] = build_pi_base_urls(state["workspace"])["oss"] + model = default_model(state) + if model and token is not None: + write_tool_config(state, model, token=token) + return + + # Token acquisition can fail before the normal config rewrite returns a + # token. Repair an existing generated provider in place without changing + # its credential, so a stale loopback port is never left behind. + existing = read_json_safe(PI_CONFIG_PATH) + providers = existing.get("providers") + mlflow = providers.get("databricks-mlflow") if isinstance(providers, dict) else None + if isinstance(mlflow, dict): + mlflow["baseUrl"] = pi_urls["oss"] + write_json_file(PI_CONFIG_PATH, existing) - proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env) + +def launch(state: dict, tool_args: list[str]) -> None: + proxy: tuple[threading.Thread, _mlflow_proxy.ThreadingHTTPServer] | None = None + stop_event = threading.Event() + refresher: threading.Thread | None = None + proc: subprocess.Popen | None = None + token: str | None = None + primary_error: BaseException | None = None try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + # The proxy must be live and its URL in state before the first config + # write; refreshes then keep writing the same live loopback endpoint. + proxy = _start_oss_proxy(state) + token = _refresh_token_once(state) + env = build_runtime_env(token) + + refresher = threading.Thread( + target=_refresh_forever, + args=(state, stop_event), + daemon=True, + ) + refresher.start() + + proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + except BaseException as exc: + primary_error = exc + raise finally: stop_event.set() - refresher.join(timeout=1) + if refresher is not None: + refresher.join(timeout=1) + if proxy is not None: + proxy_thread, server = proxy + restore_error: Exception | None = None + try: + _restore_direct_oss_config(state, token) + except Exception as exc: + restore_error = exc + print_warning(f"Pi MLflow direct configuration could not be restored ({exc}).") + finally: + server.shutdown() + server.server_close() + proxy_thread.join(timeout=1) + if restore_error is not None and primary_error is None: + raise restore_error raise SystemExit(returncode) diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 87e9c6c8..06b99d07 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -4,7 +4,9 @@ import json from contextlib import nullcontext -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +import pytest from ucode.agents import pi @@ -716,6 +718,210 @@ def test_partially_servable_list_still_splits(self): assert families == ({"opus": "system.ai.claude-opus-4-8"}, [], []) +class TestMlflowProxyLifecycle: + def test_not_started_without_oss_models(self): + state = {"workspace": WS, "oss_models": [], "base_urls": {"pi": _base_urls()}} + with patch.object(pi._mlflow_proxy, "start") as start: + assert pi._start_oss_proxy(state) is None + start.assert_not_called() + + def test_stale_loopback_url_is_replaced_and_real_workspace_is_upstream(self): + server = MagicMock() + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object( + pi._mlflow_proxy, + "start", + return_value=(server, "http://127.0.0.1:60000"), + ) as start: + running = pi._start_oss_proxy(state) + assert running is not None + start.assert_called_once_with(WS) + assert state["base_urls"]["pi"]["oss"] == ("http://127.0.0.1:60000/ai-gateway/mlflow/v1") + + def test_loopback_workspace_is_not_recursively_proxied(self): + state = { + "workspace": "http://127.0.0.1:9999", + "oss_models": ["system.ai.inkling"], + } + with patch.object(pi._mlflow_proxy, "start") as start: + assert pi._start_oss_proxy(state) is None + start.assert_not_called() + + @staticmethod + def _proxy_pair(): + proxy_thread = MagicMock() + server = MagicMock() + return (proxy_thread, server), proxy_thread, server + + def test_restore_rewrites_persistent_config_to_direct_gateway(self): + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object(pi, "write_tool_config") as write: + pi._restore_direct_oss_config(state, "tok") + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + write.assert_called_once_with(state, "system.ai.inkling", token="tok") + + def test_restore_without_token_clears_state_and_existing_config_url( + self, tmp_path, monkeypatch + ): + config_path = tmp_path / "models.json" + config_path.write_text( + json.dumps( + { + "providers": { + "databricks-mlflow": { + "baseUrl": "http://127.0.0.1:54321/ai-gateway/mlflow/v1", + "apiKey": "existing-token", + } + } + } + ) + ) + monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with patch.object(pi, "write_tool_config") as write: + pi._restore_direct_oss_config(state, None) + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + write.assert_not_called() + restored = json.loads(config_path.read_text()) + assert restored["providers"]["databricks-mlflow"]["baseUrl"] == ( + f"{WS}/ai-gateway/mlflow/v1" + ) + assert restored["providers"]["databricks-mlflow"]["apiKey"] == "existing-token" + + def test_proxy_precedes_first_config_write_and_is_cleaned_on_normal_exit(self): + proxy, proxy_thread, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + order = [] + + def start(current): + current.setdefault("base_urls", {}).setdefault("pi", {})["oss"] = "http://live" + order.append("proxy") + return proxy + + def refresh(current, *, force_refresh=False): + assert current["base_urls"]["pi"]["oss"] == "http://live" + order.append("config") + return "tok" + + proc = MagicMock() + proc.wait.return_value = 0 + with ( + patch.object(pi, "_start_oss_proxy", side_effect=start), + patch.object(pi, "_refresh_token_once", side_effect=refresh), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + assert exit_info.value.code == 0 + assert order[:2] == ["proxy", "config"] + restore.assert_called_once_with(state, "tok") + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + proxy_thread.join.assert_called_once_with(timeout=1) + + def test_interrupt_forwards_sigint_and_cleans_proxy(self): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.wait.side_effect = [KeyboardInterrupt, 130] + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(SystemExit) as exit_info, + ): + pi.launch(state, []) + assert exit_info.value.code == 130 + proc.send_signal.assert_called_once_with(pi.signal.SIGINT) + restore.assert_called_once_with(state, "tok") + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + @pytest.mark.parametrize("failure_stage", ["config", "popen"]) + def test_setup_failure_still_cleans_proxy(self, failure_stage): + proxy, _, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + refresh = MagicMock(return_value="tok") + popen = MagicMock(return_value=MagicMock()) + if failure_stage == "config": + refresh.side_effect = RuntimeError("token failed") + else: + popen.side_effect = OSError("binary missing") + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", refresh), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config") as restore, + patch.object(pi.subprocess, "Popen", popen), + pytest.raises((RuntimeError, OSError)) as exc_info, + ): + pi.launch(state, []) + expected = "token failed" if failure_stage == "config" else "binary missing" + assert str(exc_info.value) == expected + expected_token = None if failure_stage == "config" else "tok" + restore.assert_called_once_with(state, expected_token) + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + def test_setup_failure_remains_primary_when_restore_also_fails(self): + proxy, _, server = self._proxy_pair() + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object( + pi, "_restore_direct_oss_config", side_effect=RuntimeError("restore failed") + ), + patch.object(pi, "print_warning") as warning, + patch.object(pi.subprocess, "Popen", side_effect=OSError("binary missing")), + pytest.raises(OSError, match="binary missing"), + ): + pi.launch(state, []) + warning.assert_called_once() + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + def test_restore_failure_does_not_skip_proxy_shutdown(self): + proxy, _, server = self._proxy_pair() + proc = MagicMock() + proc.wait.return_value = 0 + state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} + with ( + patch.object(pi, "_start_oss_proxy", return_value=proxy), + patch.object(pi, "_refresh_token_once", return_value="tok"), + patch.object(pi, "_refresh_forever", return_value=None), + patch.object(pi, "_restore_direct_oss_config", side_effect=OSError("restore failed")), + patch.object(pi.subprocess, "Popen", return_value=proc), + pytest.raises(OSError, match="restore failed"), + ): + pi.launch(state, []) + server.shutdown.assert_called_once() + server.server_close.assert_called_once() + + class TestManagedDefaultModel: """A managed config's `pi_default_model` takes priority over the allowlist.""" diff --git a/tests/test_mlflow_proxy.py b/tests/test_mlflow_proxy.py new file mode 100644 index 00000000..c4d1c859 --- /dev/null +++ b/tests/test_mlflow_proxy.py @@ -0,0 +1,319 @@ +"""Behavioral tests for Pi's MLflow SSE-repair proxy.""" + +from __future__ import annotations + +import json +import socket +import threading +import urllib.error +import urllib.request +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +from ucode.agents import _mlflow_proxy + +_STREAM_NO_FINISH = ( + b'data: {"id":"c1","choices":[{"delta":{"content":"ok"},"index":0}]}\n\ndata: [DONE]\n\n' +) +_STREAM_WITH_FINISH = ( + b'data: {"id":"c2","choices":[{"delta":{"content":"ok"},"index":0}]}\n\n' + b'data: {"id":"c2","choices":[{"delta":{},"finish_reason":"stop","index":0}]}\n\n' + b"data: [DONE]\n\n" +) + + +class _Gateway(HTTPServer): + response_status = 200 + response_type = "text/event-stream" + response_body = b"" + truncate = False + received_headers: dict[str, str] + + +def _gateway( + body: bytes, + *, + status: int = 200, + content_type: str = "text/event-stream", + truncate: bool = False, +) -> tuple[str, _Gateway, threading.Thread]: + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self): # noqa: N802 + self.server.received_headers = dict(self.headers.items()) # type: ignore[attr-defined] + length = int(self.headers.get("Content-Length", 0)) + self.rfile.read(length) + self.send_response(self.server.response_status) # type: ignore[attr-defined] + self.send_header("Content-Type", self.server.response_type) # type: ignore[attr-defined] + advertised = len(self.server.response_body) + (20 if self.server.truncate else 0) # type: ignore[attr-defined] + self.send_header("Content-Length", str(advertised)) + self.send_header("X-Upstream", "yes") + self.end_headers() + self.wfile.write(self.server.response_body) # type: ignore[attr-defined] + self.wfile.flush() + if self.server.truncate: # type: ignore[attr-defined] + self.close_connection = True + + def log_message(self, format: str, *args: object) -> None: + return + + server = _Gateway(("127.0.0.1", 0), Handler) + server.response_status = status + server.response_type = content_type + server.response_body = body + server.truncate = truncate + server.received_headers = {} + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return f"http://127.0.0.1:{server.server_address[1]}", server, thread + + +def _proxy(upstream: str): + started = _mlflow_proxy.start(upstream) + assert started is not None + server, base = started + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return base, server, thread + + +def _stop(server: HTTPServer, thread: threading.Thread) -> None: + server.shutdown() + server.server_close() + thread.join(timeout=2) + assert not thread.is_alive() + + +def _post(base: str, *, authorization: str | None = None) -> tuple[int, dict[str, str], bytes]: + headers = {"Content-Type": "application/json"} + if authorization: + headers["Authorization"] = authorization + request = urllib.request.Request( + f"{base}/ai-gateway/mlflow/v1/chat/completions", + data=b'{"stream":true}', + headers=headers, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=5) as response: + return response.status, dict(response.headers.items()), response.read() + except urllib.error.HTTPError as exc: + return exc.code, dict(exc.headers.items()), exc.read() + + +class TestSseRepair: + def test_healthy_stream_body_is_byte_identical(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + status, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 200 + assert body == _STREAM_WITH_FINISH + assert body.count(b"finish_reason") == 1 + + def test_missing_finish_is_injected_before_done(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_NO_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body.count(b"finish_reason") == 1 + assert body.index(b"finish_reason") < body.index(b"[DONE]") + assert "Content-Length" not in headers + + def test_data_field_without_space_and_absent_id(self): + stream = b'data:{"choices":[{"delta":{"content":"ok"}}]}\n\ndata:[DONE]\n\n' + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert b"finish_reason" in body + assert b'"id":null' not in body + + def test_truncated_content_stream_gets_finish_and_done(self): + stream = b'data: {"id":"c3","choices":[{"delta":{"content":"ok"}}]}\n\n' + upstream, gateway, gateway_thread = _gateway(stream, truncate=True) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert b"finish_reason" in body + assert body.rstrip().endswith(b"[DONE]") + + def test_any_choice_finish_reason_suppresses_injection(self): + stream = ( + b'data: {"choices":[{"delta":{}},{"delta":{},"finish_reason":"length"}]}\n\n' + b"data: [DONE]\n\n" + ) + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert body.count(b"finish_reason") == 1 + + @pytest.mark.parametrize( + "stream", + [ + b'data: {"error":{"message":"rate limited"}}\n\n', + b'event: error\ndata: {"message":"rate limited"}\n\n', + b'event:error\ndata: {"message":"rate limited"}\n\n', + ], + ) + def test_explicit_sse_error_is_not_turned_into_success(self, stream): + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert b"finish_reason" not in body + assert b"[DONE]" not in body + + +class TestPassthroughAndFailures: + def test_non_streaming_json_status_headers_and_body_preserved(self): + payload = b'{"choices":[{"message":{"content":"ok"}}]}' + upstream, gateway, gateway_thread = _gateway(payload, content_type="application/json") + base, proxy, proxy_thread = _proxy(upstream) + try: + status, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 200 + assert headers["Content-Type"] == "application/json" + assert headers["Content-Length"] == str(len(payload)) + assert headers["X-Upstream"] == "yes" + assert body == payload + + def test_http_error_is_relayed_without_repair(self): + payload = b'{"error":"rate limited"}' + upstream, gateway, gateway_thread = _gateway( + payload, status=429, content_type="application/json" + ) + base, proxy, proxy_thread = _proxy(upstream) + try: + status, headers, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert status == 429 + assert headers["Content-Type"] == "application/json" + assert body == payload + assert b"finish_reason" not in body + + def test_connection_refused_returns_controlled_502(self): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + base, proxy, proxy_thread = _proxy(f"http://127.0.0.1:{port}") + try: + status, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + assert status == 502 + assert json.loads(body)["error"] + assert b"finish_reason" not in body + + def test_absolute_request_target_is_rejected_without_forwarding_auth(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + proxy_port = int(base.rsplit(":", 1)[1]) + attacker, attacker_server, attacker_thread = _gateway( + b"captured", content_type="text/plain" + ) + request = ( + f"POST {attacker}/capture HTTP/1.1\r\n" + "Host: ignored\r\n" + "Authorization: Bearer secret-value\r\n" + "Content-Length: 2\r\n" + "Connection: close\r\n\r\n{}" + ).encode() + sock = socket.create_connection(("127.0.0.1", proxy_port), timeout=5) + try: + sock.sendall(request) + response = b"" + while chunk := sock.recv(4096): + response += chunk + finally: + sock.close() + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + _stop(attacker_server, attacker_thread) + assert b" 400 " in response.split(b"\r\n", 1)[0] + assert gateway.received_headers == {} + assert attacker_server.received_headers == {} + + def test_authorization_forwarded_and_hop_by_hop_headers_removed(self): + upstream, gateway, gateway_thread = _gateway(_STREAM_WITH_FINISH) + base, proxy, proxy_thread = _proxy(upstream) + try: + _post(base, authorization="Bearer secret-value") + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + lowered = {key.lower(): value for key, value in gateway.received_headers.items()} + assert lowered["authorization"] == "Bearer secret-value" + # urllib regenerates identity after the client value is stripped, so + # the parseable SSE cannot arrive gzip-compressed. + assert lowered["accept-encoding"] == "identity" + assert lowered["host"].startswith("127.0.0.1:") # regenerated for upstream + + +class TestLifecycle: + def test_shutdown_and_server_close_release_port(self): + started = _mlflow_proxy.start("https://example.com") + assert started is not None + server, _ = started + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + _stop(server, thread) + rebound = socket.socket() + try: + rebound.bind(("127.0.0.1", port)) + finally: + rebound.close() + + def test_repeated_start_stop_uses_fresh_live_servers(self): + ports = [] + for _ in range(3): + started = _mlflow_proxy.start("https://example.com") + assert started is not None + server, _ = started + ports.append(server.server_address[1]) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + _stop(server, thread) + assert all(isinstance(port, int) and port > 0 for port in ports) + + def test_bind_failure_warns_and_degrades_to_direct_gateway(self, monkeypatch): + warnings = [] + monkeypatch.setattr( + _mlflow_proxy, + "_Server", + lambda *args, **kwargs: (_ for _ in ()).throw(OSError("no sockets")), + ) + monkeypatch.setattr(_mlflow_proxy, "print_warning", warnings.append) + + assert _mlflow_proxy.start("https://example.com") is None + assert warnings and "not started" in warnings[0] From 8e78937c5546d4ccb9d00b7db7a704010ebaa756 Mon Sep 17 00:00:00 2001 From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:46:05 +1000 Subject: [PATCH 8/8] fix(pi): preserve proxy safety across truncation and shutdown --- src/ucode/agents/_mlflow_proxy.py | 66 ++++++++++++++++++++---------- src/ucode/agents/pi.py | 19 ++++++++- tests/test_agent_pi.py | 67 +++++++++++++++++++++++++++++++ tests/test_mlflow_proxy.py | 25 ++++++++++-- 4 files changed, 150 insertions(+), 27 deletions(-) diff --git a/src/ucode/agents/_mlflow_proxy.py b/src/ucode/agents/_mlflow_proxy.py index 04973c98..7cb6dad9 100644 --- a/src/ucode/agents/_mlflow_proxy.py +++ b/src/ucode/agents/_mlflow_proxy.py @@ -163,17 +163,46 @@ def _relay_sse(self, status: int, headers: Message, stream: IO[bytes]) -> None: saw_done = False saw_error = False last_id: str | None = None + event_data: list[bytes] = [] + + def inspect_event() -> None: + nonlocal saw_error, saw_finish, last_id + if not event_data: + return + payload = b"\n".join(event_data) + event_data.clear() + try: + event = json.loads(payload) + except (UnicodeDecodeError, json.JSONDecodeError): + return + if not isinstance(event, dict): + return + if "error" in event: + saw_error = True + event_id = event.get("id") + if isinstance(event_id, str): + last_id = event_id + choices = event.get("choices") + if isinstance(choices, list) and any( + isinstance(choice, dict) and choice.get("finish_reason") is not None + for choice in choices + ): + saw_finish = True + try: for raw_line in stream: payload = _data_payload(raw_line) event_line = raw_line.rstrip(b"\r\n") - if event_line.lower().startswith(b"event:"): + if not event_line: + inspect_event() + elif event_line.lower().startswith(b"event:"): event_name = event_line[6:] if event_name.startswith(b" "): event_name = event_name[1:] if event_name.lower() == b"error": saw_error = True if payload == b"[DONE]": + inspect_event() if saw_data and not saw_finish and not saw_error: self._write(b"data: " + _finish_chunk(last_id) + b"\n\n") saw_finish = True @@ -182,30 +211,23 @@ def _relay_sse(self, status: int, headers: Message, stream: IO[bytes]) -> None: continue if payload is not None and payload: saw_data = True - try: - event = json.loads(payload) - if isinstance(event, dict): - if "error" in event: - saw_error = True - event_id = event.get("id") - if isinstance(event_id, str): - last_id = event_id - choices = event.get("choices") - if isinstance(choices, list) and any( - isinstance(choice, dict) and choice.get("finish_reason") is not None - for choice in choices - ): - saw_finish = True - except (UnicodeDecodeError, json.JSONDecodeError): - pass + event_data.append(payload) self._write(raw_line) - except (BrokenPipeError, ConnectionResetError): + except (BrokenPipeError, ConnectionResetError, OSError, IncompleteRead): + # Never turn a transport-failed partial stream into a successful + # synthetic completion. EOF without a transport exception remains + # repairable below because affected gateways can end cleanly after + # their final data event. + return + + # ``HTTPResponse`` line iteration can end without raising even when a + # declared Content-Length was not satisfied. A positive remainder is + # still a transport truncation, not a clean finish-reason omission. + remaining = getattr(stream, "length", None) + if isinstance(remaining, int) and remaining > 0: return - except (OSError, IncompleteRead): - # A dropped upstream after data is treated as a truncated stream; - # the repair below gives Pi a structurally valid terminator. - pass + inspect_event() if saw_data and not saw_error: try: if not saw_finish: diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index c43ec468..7096a930 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -76,6 +76,7 @@ PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" PI_BACKUP_PATH = APP_DIR / "pi-models.backup.json" PI_SETTINGS_BACKUP_PATH = APP_DIR / "pi-settings.backup.json" +_CONFIG_WRITE_LOCK = threading.RLock() SPEC: ToolSpec = { "binary": "pi", @@ -323,6 +324,17 @@ def write_tool_config( token: str | None = None, *, force_refresh: bool = False, +) -> tuple[dict, str]: + with _CONFIG_WRITE_LOCK: + return _write_tool_config_unlocked(state, model, token, force_refresh=force_refresh) + + +def _write_tool_config_unlocked( + state: dict, + model: str, + token: str | None = None, + *, + force_refresh: bool = False, ) -> tuple[dict, str]: backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH) if token is None: @@ -496,7 +508,12 @@ def _start_oss_proxy( def _restore_direct_oss_config(state: dict, token: str | None) -> None: - """Replace the session-only proxy URL before its listener is released.""" + """Replace the session-only proxy URL after any in-flight config write.""" + with _CONFIG_WRITE_LOCK: + _restore_direct_oss_config_unlocked(state, token) + + +def _restore_direct_oss_config_unlocked(state: dict, token: str | None) -> None: pi_urls = state.setdefault("base_urls", {}).setdefault( "pi", build_pi_base_urls(state["workspace"]) ) diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 06b99d07..2e881316 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import threading from contextlib import nullcontext from unittest.mock import MagicMock, patch @@ -806,6 +807,72 @@ def test_restore_without_token_clears_state_and_existing_config_url( ) assert restored["providers"]["databricks-mlflow"]["apiKey"] == "existing-token" + def test_refresh_keeps_live_proxy_url(self, tmp_path, monkeypatch): + config_path = tmp_path / "models.json" + settings_path = tmp_path / "settings.json" + monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", settings_path) + monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "models.backup.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "settings.backup.json") + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + with ( + patch.object(pi, "get_databricks_token", return_value="refreshed-token"), + patch.object(pi, "save_state"), + ): + pi._refresh_token_once(state, force_refresh=True) + provider = json.loads(config_path.read_text())["providers"]["databricks-mlflow"] + assert provider["baseUrl"] == "http://127.0.0.1:54321/ai-gateway/mlflow/v1" + assert provider["apiKey"] == "refreshed-token" + + def test_restore_waits_for_inflight_refresh_and_writes_direct_last(self, monkeypatch): + state = { + "workspace": WS, + "oss_models": ["system.ai.inkling"], + "base_urls": { + "pi": {**_base_urls(), "oss": "http://127.0.0.1:54321/ai-gateway/mlflow/v1"} + }, + } + refresh_started = threading.Event() + release_refresh = threading.Event() + restore_done = threading.Event() + write_order: list[str] = [] + + def fake_write(current, model, token=None, *, force_refresh=False): + write_order.append(str(token)) + if token == "refresh-token": + refresh_started.set() + assert release_refresh.wait(timeout=2) + return current, str(token) + + monkeypatch.setattr(pi, "_write_tool_config_unlocked", fake_write) + refresher = threading.Thread( + target=pi.write_tool_config, + args=(state, "system.ai.inkling", "refresh-token"), + ) + refresher.start() + assert refresh_started.wait(timeout=2) + restorer = threading.Thread( + target=lambda: ( + pi._restore_direct_oss_config(state, "restore-token"), + restore_done.set(), + ) + ) + restorer.start() + assert not restore_done.wait(timeout=0.05) + release_refresh.set() + refresher.join(timeout=2) + restorer.join(timeout=2) + assert not refresher.is_alive() + assert not restorer.is_alive() + assert write_order == ["refresh-token", "restore-token"] + assert state["base_urls"]["pi"]["oss"] == f"{WS}/ai-gateway/mlflow/v1" + def test_proxy_precedes_first_config_write_and_is_cleaned_on_normal_exit(self): proxy, proxy_thread, server = self._proxy_pair() state = {"workspace": WS, "oss_models": ["system.ai.inkling"]} diff --git a/tests/test_mlflow_proxy.py b/tests/test_mlflow_proxy.py index c4d1c859..6338606b 100644 --- a/tests/test_mlflow_proxy.py +++ b/tests/test_mlflow_proxy.py @@ -140,8 +140,8 @@ def test_data_field_without_space_and_absent_id(self): assert b"finish_reason" in body assert b'"id":null' not in body - def test_truncated_content_stream_gets_finish_and_done(self): - stream = b'data: {"id":"c3","choices":[{"delta":{"content":"ok"}}]}\n\n' + def test_transport_truncated_stream_is_not_turned_into_success(self): + stream = b'data: {"id":"c3","choices":[{"delta":{"content":"partial"}}]}\n\n' upstream, gateway, gateway_thread = _gateway(stream, truncate=True) base, proxy, proxy_thread = _proxy(upstream) try: @@ -149,8 +149,25 @@ def test_truncated_content_stream_gets_finish_and_done(self): finally: _stop(proxy, proxy_thread) _stop(gateway, gateway_thread) - assert b"finish_reason" in body - assert body.rstrip().endswith(b"[DONE]") + assert body == stream + assert b"finish_reason" not in body + assert b"[DONE]" not in body + + def test_multiline_finish_event_is_not_repaired_twice(self): + stream = ( + b'data: {"id":"c4","choices":[{"delta":{},\n' + b'data: "index":0,"finish_reason":"stop"}]}\n\n' + b"data: [DONE]\n\n" + ) + upstream, gateway, gateway_thread = _gateway(stream) + base, proxy, proxy_thread = _proxy(upstream) + try: + _, _, body = _post(base) + finally: + _stop(proxy, proxy_thread) + _stop(gateway, gateway_thread) + assert body == stream + assert body.count(b"finish_reason") == 1 def test_any_choice_finish_reason_suppresses_injection(self): stream = (