Add Kimi model support for opencode (#178)

Add OSS model support for OpenCode

Co-authored-by: Lilly <lilly.gray@tecton.ai>

Author: lilly-luo

src/ucode/agents/__init__.py

@@ -344,7 +344,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"),

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:

src/ucode/cli.py

@@ -90,6 +90,7 @@
     "claude": ("claude", "opencode", "copilot", "pi"),
     "codex": ("codex", "copilot", "pi"),
     "gemini": ("gemini", "opencode", "pi"),
+    "oss": ("opencode",),
 }
 
 
@@ -99,7 +100,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)
@@ -267,13 +273,16 @@ 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 = []
     web_search_model: str | None = None
     if skip_model_discovery:
         # Provider mode: the agent routes through a Model Provider Service and
@@ -291,7 +300,9 @@ def configure_shared_state(
         # 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:
@@ -304,11 +315,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()
@@ -339,6 +354,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)
@@ -352,6 +369,7 @@ def configure_shared_state(
         "claude": claude_reason,
         "gemini": gemini_reason,
         "codex": codex_reason,
+        "oss": oss_reason,
     }
     return state
 

src/ucode/databricks.py

@@ -1193,23 +1193,24 @@ 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
     expose per-model API dialects.
     """
     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) -----------------------------
@@ -2081,6 +2084,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",
     }
 
 

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

tests/test_cli.py

@@ -1045,7 +1045,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))
@@ -1117,7 +1117,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(
@@ -1137,7 +1143,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,
@@ -1215,7 +1221,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))