feat(llm): forward a custom provider's extra_headers as default_headers (#3680)

graphify/llm.py

@@ -364,6 +364,34 @@ def _openai_max_tokens_kwarg(backend: str, max_tokens: int) -> dict[str, int]:
     return {key: max_tokens}
 
 
+def _provider_extra_headers(raw: object, name: str = "") -> dict[str, str] | None:
+    """Validate a custom provider's providers.json ``extra_headers`` map (#3680).
+
+    Some OpenAI-compatible gateways need request headers beyond Authorization
+    (OpenCode Go rejects a request without ``x-opencode-session``). They are
+    passed to the OpenAI client as ``default_headers``, the header counterpart
+    of ``extra_body``. Only string-to-string entries are kept, and an
+    ``Authorization`` entry is dropped: the API key always comes from the
+    provider's ``env_key``, so the config cannot swap in a different credential.
+    Returns None when nothing usable remains.
+    """
+    if raw is None:
+        return None
+    if not isinstance(raw, dict):
+        print(f"[graphify] WARNING: provider {name!r} extra_headers is not an object; ignoring.", file=sys.stderr)
+        return None
+    headers: dict[str, str] = {}
+    for key, value in raw.items():
+        if not (isinstance(key, str) and isinstance(value, str)):
+            print(f"[graphify] WARNING: provider {name!r} extra_headers entry {key!r} is not a string pair; ignoring it.", file=sys.stderr)
+            continue
+        if key.strip().lower() == "authorization":
+            print(f"[graphify] WARNING: provider {name!r} extra_headers cannot set Authorization; ignoring it.", file=sys.stderr)
+            continue
+        headers[key] = value
+    return headers or None
+
+
 # Model-name fragments for OpenAI-compatible "reasoning" models that reject an
 # explicit temperature: the API returns 400 "Unsupported value: 'temperature'
 # does not support 0 with this model. Only the default (1) value is supported."
@@ -1557,6 +1585,7 @@ def _call_openai_compat(
     deep_mode: bool = False,
     images: list[_ImageRef] | None = None,
     extra_body: dict | None = None,
+    extra_headers: dict[str, str] | None = None,
 ) -> dict:
     """Call any OpenAI-compatible API (Kimi, OpenAI, etc.) and return parsed JSON."""
     ollama_seed = _resolve_ollama_seed() if backend == "ollama" else None
@@ -1582,7 +1611,8 @@ def _call_openai_compat(
     if backend == "ollama" and not os.environ.get("GRAPHIFY_MAX_RETRIES", "").strip():
         _retries = 0
     client = OpenAI(api_key=api_key, base_url=base_url, timeout=_resolve_api_timeout(),
-                    max_retries=_retries)
+                    max_retries=_retries,
+                    **({"default_headers": extra_headers} if extra_headers else {}))
     kwargs: dict = {
         "model": model,
         "messages": [
@@ -2214,6 +2244,7 @@ def extract_files_direct(
             deep_mode=deep_mode,
             images=image_refs,
             extra_body=cfg.get("extra_body"),
+            extra_headers=_provider_extra_headers(cfg.get("extra_headers"), backend),
         )
 
     # Verify code-typed nodes against the source the model read and downgrade the
@@ -3260,7 +3291,9 @@ def _rec(inp, out) -> None:
         from openai import OpenAI
     except ImportError as exc:
         raise ImportError(_backend_pkg_hint("openai", "openai")) from exc
-    client = OpenAI(api_key=key, base_url=cfg["base_url"], timeout=_resolve_api_timeout(), max_retries=_resolve_max_retries())
+    _headers = _provider_extra_headers(cfg.get("extra_headers"), backend)
+    client = OpenAI(api_key=key, base_url=cfg["base_url"], timeout=_resolve_api_timeout(), max_retries=_resolve_max_retries(),
+                    **({"default_headers": _headers} if _headers else {}))
     kwargs: dict = {
         "model": mdl,
         "messages": [{"role": "user", "content": prompt}],

tests/test_llm_backends.py

@@ -976,6 +976,124 @@ def test_call_openai_compat_extra_body_wins_over_moonshot_default(monkeypatch):
     assert captured["extra_body"] == {"thinking": {"type": "enabled"}}
 
 
+# ---------------------------------------------------------------------------
+# Custom-provider extra_headers (#3680): gateways such as OpenCode Go reject a
+# request without their own header (x-opencode-session). providers.json passes
+# them to the OpenAI client as default_headers, at extraction and at _call_llm.
+# ---------------------------------------------------------------------------
+
+
+def _install_header_capturing_openai(monkeypatch):
+    """Record the OpenAI client's constructor kwargs (where default_headers go)."""
+    import sys
+    import types
+
+    clients: list[dict] = []
+
+    class _FakeOpenAI:
+        def __init__(self, *_, **kwargs):
+            clients.append(kwargs)
+            self.chat = self
+            self.completions = self
+
+        def create(self, **_kwargs):
+            return _fake_openai_response(
+                '{"nodes":[{"id":"x"}],"edges":[],"hyperedges":[]}',
+                finish_reason="stop",
+                completion_tokens=100,
+            )
+
+    fake_module = types.ModuleType("openai")
+    setattr(fake_module, "OpenAI", _FakeOpenAI)
+    setattr(fake_module, "APIStatusError", _FakeAPIStatusError)
+    monkeypatch.setitem(sys.modules, "openai", fake_module)
+    return clients
+
+
+def _register_header_provider(monkeypatch, extra_headers):
+    monkeypatch.setattr(llm, "_get_backend_api_key", lambda _backend: "go-key")
+    monkeypatch.setitem(
+        llm.BACKENDS,
+        "opencode-go",
+        {
+            "base_url": "https://opencode.ai/zen/go/v1",
+            "default_model": "deepseek-v4-flash",
+            "env_key": "OPENCODE_API_KEY",
+            "pricing": {"input": 0.0, "output": 0.0},
+            "temperature": 0,
+            "extra_headers": extra_headers,
+        },
+    )
+
+
+def test_call_openai_compat_passes_extra_headers_as_default_headers(monkeypatch):
+    clients = _install_header_capturing_openai(monkeypatch)
+
+    llm._call_openai_compat(
+        "https://opencode.ai/zen/go/v1", "go-key", "deepseek-v4-flash",
+        "u", temperature=0, max_completion_tokens=8192, backend="opencode-go",
+        extra_headers={"x-opencode-session": "s-1"},
+    )
+
+    assert clients[-1]["default_headers"] == {"x-opencode-session": "s-1"}
+
+
+def test_call_openai_compat_without_extra_headers_sets_no_default_headers(monkeypatch):
+    clients = _install_header_capturing_openai(monkeypatch)
+
+    llm._call_openai_compat(
+        "https://api.moonshot.ai/v1", "tk", "kimi-k2",
+        "u", temperature=0, max_completion_tokens=8192, backend="kimi",
+    )
+
+    assert "default_headers" not in clients[-1]
+
+
+def test_extraction_forwards_provider_extra_headers(monkeypatch, tmp_path):
+    clients = _install_header_capturing_openai(monkeypatch)
+    _register_header_provider(monkeypatch, {"x-opencode-session": "s-1", "User-Agent": "my-agent/1.0"})
+    doc = tmp_path / "a.md"
+    doc.write_text("# A\n", encoding="utf-8")
+
+    llm.extract_files_direct([doc], backend="opencode-go", root=tmp_path)
+
+    assert clients[-1]["default_headers"] == {"x-opencode-session": "s-1", "User-Agent": "my-agent/1.0"}
+
+
+def test_call_llm_forwards_provider_extra_headers(monkeypatch):
+    clients = _install_header_capturing_openai(monkeypatch)
+    _register_header_provider(monkeypatch, {"x-opencode-session": "s-1"})
+
+    llm._call_llm("label this", backend="opencode-go")
+
+    assert clients[-1]["default_headers"] == {"x-opencode-session": "s-1"}
+
+
+def test_provider_extra_headers_cannot_override_authorization(monkeypatch, capsys):
+    """The key always comes from env_key; a providers.json header must not swap it."""
+    clients = _install_header_capturing_openai(monkeypatch)
+    _register_header_provider(
+        monkeypatch,
+        {"Authorization": "Bearer other", "x-opencode-session": "s-1", "x-bad": 7},
+    )
+
+    llm._call_llm("label this", backend="opencode-go")
+
+    assert clients[-1]["default_headers"] == {"x-opencode-session": "s-1"}
+    err = capsys.readouterr().err
+    assert "cannot set Authorization" in err and "'x-bad'" in err
+
+
+@pytest.mark.parametrize("raw", [None, {}, "x-opencode-session: s-1", {"Authorization": "Bearer x"}])
+def test_provider_extra_headers_unusable_values_send_no_default_headers(monkeypatch, raw):
+    clients = _install_header_capturing_openai(monkeypatch)
+    _register_header_provider(monkeypatch, raw)
+
+    llm._call_llm("label this", backend="opencode-go")
+
+    assert "default_headers" not in clients[-1]
+
+
 # ---------------------------------------------------------------------------
 # GRAPHIFY_DISABLE_THINKING: opt-in disable-thinking for reasoning models like
 # deepseek-v4-flash. Off by default — disabling thinking trades a rare reasoning