From cd4c13a2421262a7868e9bbdfa124db3947dd868 Mon Sep 17 00:00:00 2001
From: Ahmed Allam <ahmed39652003@gmail.com>
Date: Wed, 28 Oct 2026 01:17:39 +0000
Subject: [PATCH] feat(llm): send app attribution headers to Vercel AI Gateway

---
 strix/config/models.py | 33 +++++++++++++++++++++------------
 strix/core/inputs.py   |  7 ++-----
 tests/test_inputs.py   |  5 +++++
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/strix/config/models.py b/strix/config/models.py
index d9127ee8f5..102d132cfd 100644
--- a/strix/config/models.py
+++ b/strix/config/models.py
@@ -713,7 +713,7 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
     if codex.subscription_model(llm.model):
         return
     _configure_litellm_compatibility()
-    _configure_openrouter_attribution(llm.model)
+    _configure_attribution(llm.model)
     if llm.api_key:
         set_default_openai_key(llm.api_key, use_for_tracing=False)
         _configure_litellm_default("api_key", llm.api_key)
@@ -884,9 +884,13 @@ def _openrouter_error_type(error: object) -> object:
     return None
 
 
-OPENROUTER_ATTRIBUTION_HEADERS = {
+APP_ATTRIBUTION_HEADERS = {
     "HTTP-Referer": "https://strix.ai",
     "X-Title": "Strix",
+}
+
+OPENROUTER_ATTRIBUTION_HEADERS = {
+    **APP_ATTRIBUTION_HEADERS,
     "X-OpenRouter-Categories": "cli-agent",
 }
 
@@ -895,20 +899,25 @@ def is_openrouter_model(model_name: str | None) -> bool:
     return bool(model_name) and "openrouter/" in (model_name or "").strip().lower()
 
 
-def _configure_openrouter_attribution(model_name: str | None) -> None:
+def is_vercel_ai_gateway_model(model_name: str | None) -> bool:
+    return (model_name or "").strip().lower().startswith("vercel_ai_gateway/")
+
+
+def attribution_headers(model_name: str | None) -> dict[str, str]:
+    if is_openrouter_model(model_name):
+        return OPENROUTER_ATTRIBUTION_HEADERS
+    if is_vercel_ai_gateway_model(model_name):
+        return APP_ATTRIBUTION_HEADERS
+    return {}
+
+
+def _configure_attribution(model_name: str | None) -> None:
     import litellm
 
     current: object = litellm.headers
     existing: dict[str, str] = current if isinstance(current, dict) else {}
-    if not is_openrouter_model(model_name):
-        if any(key in existing for key in OPENROUTER_ATTRIBUTION_HEADERS):
-            remaining = {
-                k: v for k, v in existing.items() if k not in OPENROUTER_ATTRIBUTION_HEADERS
-            }
-            litellm.headers = remaining or None  # type: ignore[assignment]
-        return
-
-    litellm.headers = {**existing, **OPENROUTER_ATTRIBUTION_HEADERS}  # type: ignore[assignment]
+    remaining = {k: v for k, v in existing.items() if k not in OPENROUTER_ATTRIBUTION_HEADERS}
+    litellm.headers = {**remaining, **attribution_headers(model_name)} or None  # type: ignore[assignment]
 
 
 def _configure_extra_headers(llm: LlmSettings) -> None:
diff --git a/strix/core/inputs.py b/strix/core/inputs.py
index 6dbe468a3f..b5dad68bce 100644
--- a/strix/core/inputs.py
+++ b/strix/core/inputs.py
@@ -10,12 +10,11 @@
 
 from strix.config.models import (
     DEFAULT_MODEL_RETRY,
-    OPENROUTER_ATTRIBUTION_HEADERS,
+    attribution_headers,
     bedrock_route_supports_prompt_caching,
     is_bedrock_route,
     is_claude_model,
     is_known_openai_bare_model,
-    is_openrouter_model,
     model_supports_reasoning,
     request_timeout_extra_args,
     routes_through_litellm,
@@ -287,9 +286,7 @@ def make_model_settings(
 def _request_headers(
     model_name: str, extra_headers: dict[str, str] | None
 ) -> dict[str, str] | None:
-    headers: dict[str, str] = {}
-    if is_openrouter_model(model_name):
-        headers.update(OPENROUTER_ATTRIBUTION_HEADERS)
+    headers: dict[str, str] = dict(attribution_headers(model_name))
     if extra_headers:
         headers.update(extra_headers)
     return headers or None
diff --git a/tests/test_inputs.py b/tests/test_inputs.py
index 34bd65c784..870dc4d1de 100644
--- a/tests/test_inputs.py
+++ b/tests/test_inputs.py
@@ -474,6 +474,11 @@ def test_openrouter_attribution_rides_on_the_request_headers() -> None:
     }
 
 
+def test_vercel_ai_gateway_attribution_rides_on_the_request_headers() -> None:
+    headers = make_model_settings(None, model_name="vercel_ai_gateway/zai/glm-5.3").extra_headers
+    assert headers == {"HTTP-Referer": "https://strix.ai", "X-Title": "Strix"}
+
+
 def test_openrouter_attribution_absent_for_other_providers() -> None:
     assert make_model_settings(None, model_name="anthropic/claude-sonnet-4-5").extra_headers is None