From dc162b938e434d28061b1660a39b0f9983b7e03c Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Sat, 11 Jul 2026 02:12:57 +0800
Subject: [PATCH] feat(executor): refactor XAI header application logic and
 extend test coverage

- Extracted `applyXAIDefaultHeaders` and `applyXAICustomHeaders` for cleaner separation of concerns in API header application.
- Updated `applyXAIChatHeaders` to ensure proper header overrides for CLI chat-proxy-specific configurations.
- Added test case to validate custom headers' ability to override CLI chat-proxy defaults.
---
 internal/runtime/executor/xai_executor.go     | 18 +++++++++++++-----
 .../runtime/executor/xai_executor_test.go     | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go
index 432b9a83..3c857cee 100644
--- a/internal/runtime/executor/xai_executor.go
+++ b/internal/runtime/executor/xai_executor.go
@@ -940,6 +940,11 @@ func xaiIsCLIChatProxyBaseURL(baseURL string) bool {
 }
 
 func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) {
+	applyXAIDefaultHeaders(r, token, stream, sessionID)
+	applyXAICustomHeaders(r, auth)
+}
+
+func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionID string) {
 	r.Header.Set("Content-Type", "application/json")
 	if strings.TrimSpace(token) != "" {
 		r.Header.Set("Authorization", "Bearer "+token)
@@ -953,6 +958,9 @@ func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, str
 	if sessionID != "" {
 		r.Header.Set("x-grok-conv-id", sessionID)
 	}
+}
+
+func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) {
 	var attrs map[string]string
 	if auth != nil {
 		attrs = auth.Attributes
@@ -964,12 +972,12 @@ func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, str
 // requests. CLI chat-proxy identity headers are only attached when the resolved
 // chat base URL is the official CLI chat-proxy endpoint.
 func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) {
-	applyXAIHeaders(r, auth, token, stream, sessionID)
-	if !xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) {
-		return
+	applyXAIDefaultHeaders(r, token, stream, sessionID)
+	if xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) {
+		r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue)
+		r.Header.Set(xaiClientVersionHeader, xaiClientVersionValue)
 	}
-	r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue)
-	r.Header.Set(xaiClientVersionHeader, xaiClientVersionValue)
+	applyXAICustomHeaders(r, auth)
 }
 
 func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) {
diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go
index 4becf49e..3227a832 100644
--- a/internal/runtime/executor/xai_executor_test.go
+++ b/internal/runtime/executor/xai_executor_test.go
@@ -1672,6 +1672,25 @@ func TestApplyXAIChatHeaders(t *testing.T) {
 		}
 	})
 
+	t.Run("custom headers override cli chat proxy defaults", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodPost, xaiauth.CLIChatProxyBaseURL+"/responses", nil)
+		auth := &cliproxyauth.Auth{
+			Attributes: map[string]string{
+				"base_url":                         xaiauth.CLIChatProxyBaseURL,
+				"header:" + xaiTokenAuthHeader:     "custom-token-auth",
+				"header:" + xaiClientVersionHeader: "custom-client-version",
+			},
+		}
+		applyXAIChatHeaders(req, auth, "xai-token", true, "")
+
+		if got := req.Header.Get(xaiTokenAuthHeader); got != "custom-token-auth" {
+			t.Fatalf("%s = %q, want custom-token-auth", xaiTokenAuthHeader, got)
+		}
+		if got := req.Header.Get(xaiClientVersionHeader); got != "custom-client-version" {
+			t.Fatalf("%s = %q, want custom-client-version", xaiClientVersionHeader, got)
+		}
+	})
+
 	t.Run("cli headers on explicit chat proxy base", func(t *testing.T) {
 		req := httptest.NewRequest(http.MethodPost, xaiauth.CLIChatProxyBaseURL+"/responses", nil)
 		auth := &cliproxyauth.Auth{
-- 
2.51.2