diff --git a/internal/auth/xai/types.go b/internal/auth/xai/types.go index 93a5ebab..ffd9830a 100644 --- a/internal/auth/xai/types.go +++ b/internal/auth/xai/types.go @@ -5,10 +5,11 @@ import "time" const ( // DefaultAPIBaseURL is the default official xAI API base URL. - // Used for OAuth credential defaults, websocket, and media (image/video). - // Non-media HTTP chat is rewritten to CLIChatProxyBaseURL by the executor. + // Used for OAuth credential defaults, websocket, media (image/video), + // and non-media HTTP chat when auth using_api is true or non-OAuth. DefaultAPIBaseURL = "https://api.x.ai/v1" - // CLIChatProxyBaseURL is the Grok CLI chat-proxy base URL for non-image/video HTTP chat. + // CLIChatProxyBaseURL is the Grok CLI chat-proxy base URL for non-image/video + // HTTP chat when auth using_api is false, including the OAuth default. CLIChatProxyBaseURL = "https://cli-chat-proxy.grok.com/v1" // Issuer is xAI's OAuth issuer. Issuer = "https://auth.x.ai" diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 3c857cee..db3e06f6 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -10,6 +10,7 @@ import ( "net/http" "net/url" "sort" + "strconv" "strings" "time" @@ -65,6 +66,8 @@ const ( xaiClientVersionHeader = "x-grok-client-version" // Keep in sync with the current Grok CLI client version that chat-proxy expects. xaiClientVersionValue = "0.2.93" + // xaiUsingAPIAttr enables the official API path for non-media HTTP chat. + xaiUsingAPIAttr = "using_api" ) // XAIExecutor is a stateless executor for xAI Grok's Responses API. @@ -914,13 +917,56 @@ func xaiCreds(auth *cliproxyauth.Auth) (token, baseURL string) { return token, baseURL } +// xaiUsingAPI reports whether this xAI auth should use the official API path +// for non-media HTTP chat. OAuth defaults to false to use Grok Build. +func xaiUsingAPI(auth *cliproxyauth.Auth) bool { + if auth == nil { + return true + } + if len(auth.Attributes) > 0 { + if raw := strings.TrimSpace(auth.Attributes[xaiUsingAPIAttr]); raw != "" { + parsed, errParse := strconv.ParseBool(raw) + if errParse == nil { + return parsed + } + } + } + if len(auth.Metadata) > 0 { + raw, ok := auth.Metadata[xaiUsingAPIAttr] + if ok && raw != nil { + switch v := raw.(type) { + case bool: + return v + case string: + parsed, errParse := strconv.ParseBool(strings.TrimSpace(v)) + if errParse == nil { + return parsed + } + default: + } + } + } + if raw := strings.TrimSpace(auth.Attributes["auth_kind"]); raw != "" { + return !strings.EqualFold(raw, "oauth") + } + return !strings.EqualFold(xaiMetadataString(auth.Metadata, "auth_kind"), "oauth") +} + // xaiChatBaseURL returns the base URL for non-image/video xAI HTTP chat requests. -// Empty or official default base_url is rewritten to the CLI chat-proxy endpoint. -// An explicit non-default base_url (tests / custom gateways) is still honored. +// When auth using_api is true, the official API base URL logic is used. When it +// is false (including its OAuth default), empty or official default base_url is +// rewritten to the CLI chat-proxy endpoint; an explicit non-default base_url is +// still honored. // Websocket transport intentionally does not use this helper: cli-chat-proxy only // accepts HTTP POST and returns 405 for websocket upgrades. func xaiChatBaseURL(auth *cliproxyauth.Auth) string { _, baseURL := xaiCreds(auth) + if xaiUsingAPI(auth) { + if baseURL == "" { + return xaiauth.DefaultAPIBaseURL + } + return baseURL + } if baseURL != "" && !xaiIsDefaultAPIBaseURL(baseURL) { return baseURL } @@ -969,9 +1015,15 @@ func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) { } // applyXAIChatHeaders applies standard xAI headers for non-image/video chat -// requests. CLI chat-proxy identity headers are only attached when the resolved -// chat base URL is the official CLI chat-proxy endpoint. +// requests. When using_api is true, this matches the standard +// applyXAIHeaders behavior. CLI chat-proxy identity headers are only attached +// when using_api is false and 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) { + if xaiUsingAPI(auth) { + applyXAIHeaders(r, auth, token, stream, sessionID) + return + } applyXAIDefaultHeaders(r, token, stream, sessionID) if xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) { r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue) diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 3227a832..c43e8b5d 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -1580,50 +1580,128 @@ func TestXAIChatBaseURL(t *testing.T) { want string }{ { - name: "nil auth defaults to chat proxy", + name: "nil auth defaults to official api", auth: nil, - want: xaiauth.CLIChatProxyBaseURL, + want: xaiauth.DefaultAPIBaseURL, }, { - name: "empty base url defaults to chat proxy", + name: "empty base url defaults to official api without using_api", auth: &cliproxyauth.Auth{Provider: "xai"}, - want: xaiauth.CLIChatProxyBaseURL, + want: xaiauth.DefaultAPIBaseURL, }, { - name: "official default rewrites to chat proxy", + name: "official default stays official without using_api", auth: &cliproxyauth.Auth{ Attributes: map[string]string{"base_url": xaiauth.DefaultAPIBaseURL}, }, - want: xaiauth.CLIChatProxyBaseURL, + want: xaiauth.DefaultAPIBaseURL, }, { - name: "official default with trailing slash rewrites to chat proxy", + name: "OAuth credentials default to chat proxy without using_api", auth: &cliproxyauth.Auth{ - Attributes: map[string]string{"base_url": xaiauth.DefaultAPIBaseURL + "/"}, + Attributes: map[string]string{ + "auth_kind": "oauth", + "base_url": xaiauth.DefaultAPIBaseURL, + }, }, want: xaiauth.CLIChatProxyBaseURL, }, { - name: "metadata official default rewrites to chat proxy", + name: "metadata-only OAuth credentials default to chat proxy without using_api", auth: &cliproxyauth.Auth{ - Metadata: map[string]any{"base_url": xaiauth.DefaultAPIBaseURL}, + Metadata: map[string]any{ + "auth_kind": "oauth", + "base_url": xaiauth.DefaultAPIBaseURL, + }, }, want: xaiauth.CLIChatProxyBaseURL, }, { - name: "custom base url is honored", + name: "using_api false empty base url rewrites to chat proxy", + auth: &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{xaiUsingAPIAttr: "false"}, + }, + want: xaiauth.CLIChatProxyBaseURL, + }, + { + name: "using_api false official default rewrites to chat proxy", + auth: &cliproxyauth.Auth{ + Attributes: map[string]string{ + "base_url": xaiauth.DefaultAPIBaseURL, + xaiUsingAPIAttr: "false", + }, + }, + want: xaiauth.CLIChatProxyBaseURL, + }, + { + name: "using_api false official default with trailing slash rewrites to chat proxy", + auth: &cliproxyauth.Auth{ + Attributes: map[string]string{ + "base_url": xaiauth.DefaultAPIBaseURL + "/", + xaiUsingAPIAttr: "false", + }, + }, + want: xaiauth.CLIChatProxyBaseURL, + }, + { + name: "metadata using_api false official default rewrites to chat proxy", + auth: &cliproxyauth.Auth{ + Metadata: map[string]any{ + "base_url": xaiauth.DefaultAPIBaseURL, + xaiUsingAPIAttr: false, + }, + }, + want: xaiauth.CLIChatProxyBaseURL, + }, + { + name: "using_api false custom base url is honored", + auth: &cliproxyauth.Auth{ + Attributes: map[string]string{ + "base_url": "https://gateway.example.com/v1", + xaiUsingAPIAttr: "false", + }, + }, + want: "https://gateway.example.com/v1", + }, + { + name: "custom base url is honored without using_api", auth: &cliproxyauth.Auth{ Attributes: map[string]string{"base_url": "https://gateway.example.com/v1"}, }, want: "https://gateway.example.com/v1", }, { - name: "explicit chat proxy base url is preserved", + name: "using_api false explicit chat proxy base url is preserved", auth: &cliproxyauth.Auth{ - Attributes: map[string]string{"base_url": xaiauth.CLIChatProxyBaseURL}, + Attributes: map[string]string{ + "base_url": xaiauth.CLIChatProxyBaseURL, + xaiUsingAPIAttr: "false", + }, }, want: xaiauth.CLIChatProxyBaseURL, }, + { + name: "using_api true keeps official api", + auth: &cliproxyauth.Auth{ + Attributes: map[string]string{ + "base_url": xaiauth.DefaultAPIBaseURL, + xaiUsingAPIAttr: "true", + }, + }, + want: xaiauth.DefaultAPIBaseURL, + }, + { + name: "OAuth using_api true keeps official api", + auth: &cliproxyauth.Auth{ + Attributes: map[string]string{ + "auth_kind": "oauth", + "base_url": xaiauth.DefaultAPIBaseURL, + xaiUsingAPIAttr: "true", + }, + }, + want: xaiauth.DefaultAPIBaseURL, + }, } for _, tt := range tests { @@ -1636,13 +1714,37 @@ func TestXAIChatBaseURL(t *testing.T) { } func TestApplyXAIChatHeaders(t *testing.T) { - t.Run("cli chat proxy headers on default base", func(t *testing.T) { + t.Run("non OAuth defaults to official API headers", func(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "https://example.invalid/responses", nil) auth := &cliproxyauth.Auth{ Attributes: map[string]string{"base_url": xaiauth.DefaultAPIBaseURL}, } applyXAIChatHeaders(req, auth, "xai-token", true, "conv-1") + if got := req.Header.Get("Authorization"); got != "Bearer xai-token" { + t.Fatalf("Authorization = %q, want Bearer xai-token", got) + } + if got := req.Header.Get("x-grok-conv-id"); got != "conv-1" { + t.Fatalf("x-grok-conv-id = %q, want conv-1", got) + } + if got := req.Header.Get(xaiTokenAuthHeader); got != "" { + t.Fatalf("%s = %q, want empty for official API", xaiTokenAuthHeader, got) + } + if got := req.Header.Get(xaiClientVersionHeader); got != "" { + t.Fatalf("%s = %q, want empty for official API", xaiClientVersionHeader, got) + } + }) + + t.Run("OAuth defaults to cli chat proxy headers", func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "https://example.invalid/responses", nil) + auth := &cliproxyauth.Auth{ + Attributes: map[string]string{ + "auth_kind": "oauth", + "base_url": xaiauth.DefaultAPIBaseURL, + }, + } + applyXAIChatHeaders(req, auth, "xai-token", true, "conv-1") + if got := req.Header.Get("Authorization"); got != "Bearer xai-token" { t.Fatalf("Authorization = %q, want Bearer xai-token", got) } @@ -1657,10 +1759,13 @@ func TestApplyXAIChatHeaders(t *testing.T) { } }) - t.Run("no cli headers on custom gateway", func(t *testing.T) { + t.Run("no cli headers on custom gateway with using_api false", func(t *testing.T) { req := httptest.NewRequest(http.MethodPost, "https://gateway.example.com/responses", nil) auth := &cliproxyauth.Auth{ - Attributes: map[string]string{"base_url": "https://gateway.example.com/v1"}, + Attributes: map[string]string{ + "base_url": "https://gateway.example.com/v1", + xaiUsingAPIAttr: "false", + }, } applyXAIChatHeaders(req, auth, "xai-token", false, "") @@ -1677,6 +1782,7 @@ func TestApplyXAIChatHeaders(t *testing.T) { auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "base_url": xaiauth.CLIChatProxyBaseURL, + xaiUsingAPIAttr: "false", "header:" + xaiTokenAuthHeader: "custom-token-auth", "header:" + xaiClientVersionHeader: "custom-client-version", }, @@ -1691,10 +1797,13 @@ func TestApplyXAIChatHeaders(t *testing.T) { } }) - t.Run("cli headers on explicit chat proxy base", func(t *testing.T) { + t.Run("cli headers on explicit chat proxy base with using_api false", func(t *testing.T) { req := httptest.NewRequest(http.MethodPost, xaiauth.CLIChatProxyBaseURL+"/responses", nil) auth := &cliproxyauth.Auth{ - Attributes: map[string]string{"base_url": xaiauth.CLIChatProxyBaseURL + "/"}, + Attributes: map[string]string{ + "base_url": xaiauth.CLIChatProxyBaseURL + "/", + xaiUsingAPIAttr: "false", + }, } applyXAIChatHeaders(req, auth, "xai-token", true, "") @@ -1722,7 +1831,8 @@ func TestXAIExecutorExecuteChatUsesProxyHeadersOnlyForChatProxy(t *testing.T) { auth := &cliproxyauth.Auth{ Provider: "xai", Attributes: map[string]string{ - "base_url": server.URL, + "base_url": server.URL, + xaiUsingAPIAttr: "false", }, Metadata: map[string]any{"access_token": "xai-token"}, }