fix(llm): scope prompt cache keys to Moonshot APIs

src/kimi_cli/llm.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Protocol, Self, cast, get_args
+from urllib.parse import urlparse
 
 from kosong.chat_provider import ChatProvider, StreamedMessage, ThinkingEffort
 from kosong.message import (
@@ -47,6 +48,7 @@
 DEFAULT_UNKNOWN_CONTEXT_COMPLETION_TOKENS = 32_000
 DEFAULT_COMPLETION_TOKEN_SAFETY_MARGIN = 1_024
 MEDIA_TOKEN_ESTIMATE = 2_000
+MOONSHOT_API_HOSTS = {"api.kimi.com", "api.moonshot.ai", "api.moonshot.cn"}
 
 
 @dataclass(slots=True)
@@ -323,6 +325,16 @@ def _kimi_default_headers(provider: LLMProvider, oauth: OAuthManager | None) ->
     return headers
 
 
+def _supports_prompt_cache_key(base_url: str) -> bool:
+    url = base_url.strip()
+    if "//" not in url:
+        # urlparse() only fills in `hostname` when a netloc is present, so a
+        # scheme-less base_url such as "api.moonshot.ai/v1" parses entirely as
+        # a path. Give it a netloc before parsing.
+        url = f"//{url}"
+    return urlparse(url).hostname in MOONSHOT_API_HOSTS
+
+
 def create_llm(
     provider: LLMProvider,
     model: LLMModel,
@@ -358,7 +370,7 @@ def create_llm(
             )
 
             gen_kwargs: Kimi.GenerationKwargs = {}
-            if session_id:
+            if session_id and _supports_prompt_cache_key(provider.base_url):
                 gen_kwargs["prompt_cache_key"] = session_id
             if temperature := os.getenv("KIMI_MODEL_TEMPERATURE"):
                 gen_kwargs["temperature"] = float(temperature)

tests/core/test_create_llm.py

@@ -84,6 +84,49 @@ def test_create_llm_kimi_model_parameters(monkeypatch):
     )
 
 
+def test_create_llm_omits_prompt_cache_key_for_third_party_kimi_endpoint():
+    provider = LLMProvider(
+        type="kimi",
+        base_url="https://integrate.api.nvidia.com/v1",
+        api_key=SecretStr("test-key"),
+    )
+    model = LLMModel(
+        provider="nvidia",
+        model="meta/llama",
+        max_context_size=4096,
+    )
+
+    llm = create_llm(provider, model, session_id="session-123")
+
+    assert llm is not None
+    assert isinstance(llm.chat_provider, Kimi)
+    assert "prompt_cache_key" not in llm.chat_provider.model_parameters
+
+
+def test_create_llm_keeps_prompt_cache_key_for_moonshot_endpoints():
+    model = LLMModel(provider="kimi", model="kimi-base", max_context_size=4096)
+
+    for base_url in (
+        "https://api.kimi.com/coding/v1",
+        "https://api.moonshot.ai/v1",
+        "https://api.moonshot.cn/v1",
+        # scheme-less forms are accepted by the OpenAI-compatible clients too
+        "api.moonshot.ai/v1",
+        "api.kimi.com/coding/v1",
+    ):
+        provider = LLMProvider(
+            type="kimi",
+            base_url=base_url,
+            api_key=SecretStr("test-key"),
+        )
+
+        llm = create_llm(provider, model, session_id="session-123")
+
+        assert llm is not None
+        assert isinstance(llm.chat_provider, Kimi)
+        assert llm.chat_provider.model_parameters["prompt_cache_key"] == "session-123"
+
+
 def test_create_llm_kimi_prefers_max_completion_tokens_env(monkeypatch):
     provider = LLMProvider(
         type="kimi",