From f45180bfc11eee72f9891066fdeeab816dfac4e1 Mon Sep 17 00:00:00 2001 From: Lilly Date: Wed, 1 Jul 2026 17:11:55 +0000 Subject: [PATCH] Add OSS model support for OpenCode --- src/ucode/agents/__init__.py | 2 +- src/ucode/agents/opencode.py | 33 ++++++++++++++++++++++---- src/ucode/cli.py | 22 ++++++++++++++++-- src/ucode/databricks.py | 16 ++++++++----- tests/test_agent_opencode.py | 45 ++++++++++++++++++++++++++++++++++++ tests/test_cli.py | 14 +++++++---- tests/test_databricks.py | 25 ++++++++++++-------- tests/test_e2e_uc.py | 16 +++++++++---- 8 files changed, 141 insertions(+), 32 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b94e855..d218dd7 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -304,7 +304,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = { "claude": ("claude",), - "opencode": ("claude", "gemini"), + "opencode": ("claude", "gemini", "oss"), "codex": ("codex",), "gemini": ("gemini",), "copilot": ("claude", "codex"), diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 8792625..f6e999b 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -1,4 +1,4 @@ -"""OpenCode agent: writes opencode.json with two Databricks-backed providers.""" +"""OpenCode agent: writes opencode.json with Databricks-backed providers.""" from __future__ import annotations @@ -41,6 +41,7 @@ PROVIDER_KEYS: list[list[str]] = [ ["provider", "databricks-anthropic"], ["provider", "databricks-google"], + ["provider", "databricks-oss"], ] @@ -50,7 +51,7 @@ 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/") or model.startswith("databricks-google/"): + if model.startswith(("databricks-anthropic/", "databricks-google/", "databricks-oss/")): return model anthropic_models = opencode_models.get("anthropic") or [] @@ -61,6 +62,10 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) - if model in gemini_models: return f"databricks-google/{model}" + oss_models = opencode_models.get("oss") or [] + if model in oss_models: + return f"databricks-oss/{model}" + return model @@ -82,6 +87,7 @@ def render_overlay( anthropic_models = opencode_models.get("anthropic") or [] gemini_models = opencode_models.get("gemini") or [] + oss_models = opencode_models.get("oss") or [] providers: dict = {} keys: list[list[str]] = [["model"]] @@ -116,6 +122,17 @@ def render_overlay( "models": {m: {"headers": ua_header} for m in gemini_models}, } keys.append(["provider", "databricks-google"]) + if oss_models: + providers["databricks-oss"] = { + "npm": "@ai-sdk/openai", + "options": { + "baseURL": opencode_base_urls["oss"], + "apiKey": token, + "headers": auth_headers, + }, + "models": {m: {"headers": ua_header} for m in oss_models}, + } + keys.append(["provider", "databricks-oss"]) overlay: dict = {"model": _resolve_model_selector(model, opencode_models)} if providers: @@ -147,7 +164,12 @@ def write_tool_config( existing = read_json_safe(OPENCODE_CONFIG_PATH) providers = existing.get("provider") if isinstance(providers, dict): - for stale in ("databricks-anthropic", "databricks-google", "databricks-openai"): + for stale in ( + "databricks-anthropic", + "databricks-google", + "databricks-openai", + "databricks-oss", + ): providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(OPENCODE_CONFIG_PATH, merged) @@ -197,7 +219,10 @@ def default_model(state: dict) -> str | None: if anthropic: return anthropic[0] gemini = opencode_models.get("gemini") or [] - return gemini[0] if gemini else None + if gemini: + return gemini[0] + oss = opencode_models.get("oss") or [] + return oss[0] if oss 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 ba2b790..b397088 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -76,6 +76,7 @@ "claude": ("claude", "opencode", "copilot", "pi"), "codex": ("codex", "copilot", "pi"), "gemini": ("gemini", "opencode", "pi"), + "oss": ("opencode",), } @@ -85,7 +86,12 @@ def _print_discovery_diagnostics(state: dict) -> None: reasons = state.get("_discovery_reasons") or {} if not reasons: return - labels = {"claude": "Claude models", "codex": "Codex models", "gemini": "Gemini models"} + labels = { + "claude": "Claude models", + "codex": "Codex models", + "gemini": "Gemini models", + "oss": "OSS models", + } for source, reason in reasons.items(): consumers = ", ".join(_DISCOVERY_CONSUMERS.get(source, ())) label = labels.get(source, source) @@ -252,19 +258,24 @@ 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 claude_reason: str | None = None gemini_reason: str | None = None codex_reason: str | None = None + oss_reason: str | None = None claude_models = {} gemini_models = [] codex_models = [] + oss_models = [] # UC-first, best-effort: one UC model-services call yields all families as # `system.ai.` ids, bucketed by name. If a family comes back # empty (workspace without UC model-services, or the listing failed), fall # back to the per-family AI Gateway listing for that family only. with spinner("Fetching available models..."): - ms_claude, ms_codex, ms_gemini, ms_reason = discover_model_services(workspace, token) + ms_claude, ms_codex, ms_gemini, ms_oss, ms_reason = discover_model_services( + workspace, token + ) if want_claude: claude_models, claude_reason = ms_claude, ms_reason if not claude_models: @@ -277,11 +288,15 @@ def configure_shared_state( codex_models, codex_reason = ms_codex, ms_reason if not codex_models: codex_models, codex_reason = discover_codex_models(workspace, token) + if want_oss: + oss_models, oss_reason = ms_oss, ms_reason opencode_models: dict[str, list[str]] = {} if claude_models: opencode_models["anthropic"] = list(claude_models.values()) if gemini_models: opencode_models["gemini"] = gemini_models + if oss_models: + opencode_models["oss"] = oss_models # Merge into existing workspace state so prior tool configs are preserved. state = load_state() @@ -305,6 +320,8 @@ def configure_shared_state( state["gemini_models"] = gemini_models if want_codex: state["codex_models"] = codex_models + if want_oss: + state["oss_models"] = oss_models if fetch_all or "opencode" in tools: state["opencode_models"] = opencode_models save_state(state) @@ -318,6 +335,7 @@ def configure_shared_state( "claude": claude_reason, "gemini": gemini_reason, "codex": codex_reason, + "oss": oss_reason, } return state diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 250e9da..e6078ca 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1193,15 +1193,16 @@ def list_model_services( def discover_model_services( workspace: str, token: str -) -> tuple[dict[str, str], list[str], list[str], str | None]: +) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: """Discover models via UC model-services and bucket them by family name. - Returns (claude_models, codex_models, gemini_models, reason): + Returns (claude_models, codex_models, gemini_models, oss_models, reason): - ``claude_models`` maps ``opus``/``sonnet``/``haiku`` to the newest matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. + - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. ``reason`` is None on success, else explains why nothing was found. Family bucketing is by name substring because the model-services API does not @@ -1209,7 +1210,7 @@ def discover_model_services( """ ids, reason = list_model_services(workspace, token) if not ids: - return {}, [], [], reason + return {}, [], [], [], reason claude_models: dict[str, str] = {} for family in ("opus", "sonnet", "haiku"): @@ -1222,19 +1223,21 @@ def discover_model_services( codex_models = [m for m in ids if "gpt-" 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 "kimi-" in m] - if not (claude_models or codex_models or gemini_models): + if not (claude_models or codex_models or gemini_models or oss_models): sample = ", ".join(ids[:5]) return ( {}, [], [], + [], ( "model-services returned model ids but none matched " - f"claude/gpt/gemini families (got: {sample})" + f"claude/gpt/gemini/oss families (got: {sample})" ), ) - return claude_models, codex_models, gemini_models, None + return claude_models, codex_models, gemini_models, oss_models, None # --- MCP services (parallel to model services) ----------------------------- @@ -1879,6 +1882,7 @@ 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", + "oss": f"{workspace}/ai-gateway/mlflow/v1", } diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0f32f4d..9d355e7 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", + "oss": f"{WS}/ai-gateway/mlflow/v1", } @@ -48,6 +49,20 @@ def test_gemini_provider_added_when_models_present(self): overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert "databricks-google" in overlay["provider"] + def test_oss_provider_added_when_models_present(self): + models = {"oss": ["system.ai.kimi-k2-7-code"]} + overlay, _ = opencode.render_overlay( + "system.ai.kimi-k2-7-code", "tok", _base_urls(), models + ) + assert "databricks-oss" in overlay["provider"] + + def test_oss_provider_uses_ai_sdk_openai_package(self): + models = {"oss": ["system.ai.kimi-k2-7-code"]} + overlay, _ = opencode.render_overlay( + "system.ai.kimi-k2-7-code", "tok", _base_urls(), models + ) + assert overlay["provider"]["databricks-oss"]["npm"] == "@ai-sdk/openai" + def test_both_providers_when_both_present(self): models = {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]} overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) @@ -70,6 +85,14 @@ def test_gemini_base_url(self): options = overlay["provider"]["databricks-google"]["options"] assert options["baseURL"] == f"{WS}/ai-gateway/gemini/v1beta" + def test_oss_base_url(self): + models = {"oss": ["system.ai.kimi-k2-7-code"]} + overlay, _ = opencode.render_overlay( + "system.ai.kimi-k2-7-code", "tok", _base_urls(), models + ) + options = overlay["provider"]["databricks-oss"]["options"] + assert options["baseURL"] == f"{WS}/ai-gateway/mlflow/v1" + def test_token_in_api_key(self): models = {"anthropic": ["claude-sonnet"]} overlay, _ = opencode.render_overlay("claude-sonnet", "mytoken", _base_urls(), models) @@ -134,6 +157,11 @@ def test_managed_keys_include_gemini_provider(self): _, keys = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert ["provider", "databricks-google"] in keys + def test_managed_keys_include_oss_provider(self): + models = {"oss": ["system.ai.kimi-k2-7-code"]} + _, keys = opencode.render_overlay("system.ai.kimi-k2-7-code", "tok", _base_urls(), models) + assert ["provider", "databricks-oss"] in keys + def test_anthropic_models_listed(self): models = {"anthropic": ["claude-sonnet", "claude-haiku"]} overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) @@ -151,6 +179,13 @@ def test_prefixes_gemini_model_with_provider_id(self): overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) assert overlay["model"] == "databricks-google/gemini-2" + def test_prefixes_oss_model_with_provider_id(self): + models = {"oss": ["system.ai.kimi-k2-7-code"]} + overlay, _ = opencode.render_overlay( + "system.ai.kimi-k2-7-code", "tok", _base_urls(), models + ) + assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" + class TestMcpServerConfig: def test_builds_remote_server_entry_with_oauth_token_env_header(self): @@ -268,6 +303,16 @@ def test_falls_back_to_gemini(self): state = {"opencode_models": {"anthropic": [], "gemini": ["gemini-2"]}} assert opencode.default_model(state) == "gemini-2" + def test_falls_back_to_oss(self): + state = { + "opencode_models": { + "anthropic": [], + "gemini": [], + "oss": ["system.ai.kimi-k2-7-code"], + } + } + assert opencode.default_model(state) == "system.ai.kimi-k2-7-code" + def test_returns_none_when_empty(self): assert opencode.default_model({}) is None assert opencode.default_model({"opencode_models": {}}) is None diff --git a/tests/test_cli.py b/tests/test_cli.py index cda0721..23de09c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1012,7 +1012,7 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) - monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], None)) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], 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)) @@ -1084,7 +1084,13 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): monkeypatch.setattr( cli_mod, "discover_model_services", - lambda w, t: ({"opus": "system.ai.claude-opus-4-8"}, ["system.ai.gpt-5"], [], None), + lambda w, t: ( + {"opus": "system.ai.claude-opus-4-8"}, + ["system.ai.gpt-5"], + [], + [], + None, + ), ) legacy_called: list[str] = [] monkeypatch.setattr( @@ -1104,7 +1110,7 @@ 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") monkeypatch.setattr( - cli_mod, "discover_model_services", lambda w, t: ({}, [], [], "no model services") + cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], "no model services") ) monkeypatch.setattr( cli_mod, @@ -1182,7 +1188,7 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", lambda w, t: None) - monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], None)) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], 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)) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index da630b0..a52fa4e 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -94,10 +94,11 @@ def test_unsupported_tool_raises(self): class TestBuildOpencodeBaseUrls: - def test_returns_anthropic_and_gemini(self): + def test_returns_anthropic_gemini_and_oss(self): urls = build_opencode_base_urls(WS) assert urls["anthropic"] == f"{WS}/ai-gateway/anthropic/v1" assert urls["gemini"] == f"{WS}/ai-gateway/gemini/v1beta" + assert urls["oss"] == f"{WS}/ai-gateway/mlflow/v1" class TestBuildSharedBaseUrls: @@ -149,6 +150,7 @@ def test_buckets_families_by_name(self, monkeypatch): _model_service("system.ai.gpt-5"), _model_service("system.ai.gemini-2-5-flash"), _model_service("system.ai.gemini-3-5-flash"), + _model_service("system.ai.kimi-k2-7-code"), _model_service("system.ai.llama-4-maverick"), ] } @@ -156,7 +158,7 @@ def test_buckets_families_by_name(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") assert reason is None # Newest opus wins; sonnet bucketed; haiku absent. @@ -167,8 +169,9 @@ 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" - # llama is not bucketed into any of the three families. - assert "system.ai.llama-4-maverick" not in codex + gemini + assert oss == ["system.ai.kimi-k2-7-code"] + # llama is not bucketed into any of the four families. + assert "system.ai.llama-4-maverick" not in codex + gemini + oss def test_paginates_via_next_page_token(self, monkeypatch): pages = { @@ -189,7 +192,7 @@ def fake_get(url, token, timeout=10): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - claude, codex, _, reason = db_mod.discover_model_services(WS, "token") + claude, codex, _, _, reason = db_mod.discover_model_services(WS, "token") assert reason is None assert codex == ["system.ai.gpt-5"] @@ -200,9 +203,9 @@ def test_http_failure_returns_reason(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (None, "HTTP 500 Server Error") ) - claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") - assert (claude, codex, gemini) == ({}, [], []) + assert (claude, codex, gemini, oss) == ({}, [], [], []) assert reason == "HTTP 500 Server Error" def test_no_matching_families_reports_sample(self, monkeypatch): @@ -211,9 +214,9 @@ def test_no_matching_families_reports_sample(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") - assert (claude, codex, gemini) == ({}, [], []) + assert (claude, codex, gemini, oss) == ({}, [], [], []) assert reason is not None and "llama-4-maverick" in reason def test_ignores_non_system_ai_schemas(self, monkeypatch): @@ -223,6 +226,7 @@ def test_ignores_non_system_ai_schemas(self, monkeypatch): "model_services": [ _model_service("system.ai.gpt-5"), _model_service("main.svenwb.gpt-5-5"), + _model_service("temp.erni.kimi-k2-7-code"), _model_service("temp.erni.claude-opus-4-8"), _model_service("dnasi_agent_cuj.default.dnasi-gpt55-test"), ] @@ -231,12 +235,13 @@ def test_ignores_non_system_ai_schemas(self, monkeypatch): db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None) ) - claude, codex, gemini, reason = db_mod.discover_model_services(WS, "token") + claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") assert reason is None assert codex == ["system.ai.gpt-5"] assert claude == {} # temp.erni.claude-* must not be bucketed assert gemini == [] + assert oss == [] def test_requests_bounded_page_size(self, monkeypatch): # The endpoint 499s without a bounded page_size, so every request must diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index 32dbf63..c716dcb 100644 --- a/tests/test_e2e_uc.py +++ b/tests/test_e2e_uc.py @@ -23,14 +23,15 @@ def _has_uc_models(workspace: str, token: str) -> bool: - claude, codex, gemini, _reason = discover_model_services(workspace, token) - return bool(claude or codex or gemini) + claude, codex, gemini, oss, _reason = discover_model_services(workspace, token) + return bool(claude or codex or gemini or oss) def _all_resolved_model_ids(state: dict) -> list[str]: ids: list[str] = list((state.get("claude_models") or {}).values()) ids += state.get("codex_models") or [] ids += state.get("gemini_models") or [] + ids += state.get("oss_models") or [] return ids @@ -42,14 +43,19 @@ def _all_resolved_model_ids(state: dict) -> list[str]: class TestDiscoverModelServicesE2E: def test_returns_only_system_ai_models(self, e2e_workspace, e2e_token): - claude, codex, gemini, reason = discover_model_services(e2e_workspace, e2e_token) - if not (claude or codex or gemini): + claude, codex, gemini, oss, reason = discover_model_services(e2e_workspace, e2e_token) + if not (claude or codex or gemini or oss): pytest.skip(f"No system.ai.* model services on workspace: {reason}") non_system = sorted( { m for m in _all_resolved_model_ids( - {"claude_models": claude, "codex_models": codex, "gemini_models": gemini} + { + "claude_models": claude, + "codex_models": codex, + "gemini_models": gemini, + "oss_models": oss, + } ) if not m.startswith("system.ai.") }