mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
fix(openai): align alpha search PAT forwarding
This commit is contained in:
@@ -118,7 +118,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
requestedModel,
|
||||
failedAccountIDs,
|
||||
service.OpenAIUpstreamTransportHTTPSSE,
|
||||
service.OpenAIEndpointCapabilityChatCompletions,
|
||||
service.OpenAIEndpointCapabilityAlphaSearch,
|
||||
false,
|
||||
false,
|
||||
service.PlatformOpenAI,
|
||||
|
||||
@@ -87,6 +87,7 @@ const openAILongContextBillingEnabledKey = "openai_long_context_billing_enabled"
|
||||
const (
|
||||
OpenAIEndpointCapabilityChatCompletions OpenAIEndpointCapability = "chat_completions"
|
||||
OpenAIEndpointCapabilityEmbeddings OpenAIEndpointCapability = "embeddings"
|
||||
OpenAIEndpointCapabilityAlphaSearch OpenAIEndpointCapability = "alpha_search"
|
||||
)
|
||||
|
||||
const openAIEndpointCapabilitiesCredentialKey = "openai_capabilities"
|
||||
@@ -1397,6 +1398,13 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa
|
||||
}
|
||||
switch capability {
|
||||
case OpenAIEndpointCapabilityChatCompletions:
|
||||
case OpenAIEndpointCapabilityAlphaSearch:
|
||||
// Codex alpha/search 是 ChatGPT/Codex 后端工具端点,必须使用
|
||||
// OAuth/PAT/AgentIdentity 这类 ChatGPT 账号凭据;API key 被发往
|
||||
// chatgpt.com/backend-api/codex/alpha/search 会稳定 401。
|
||||
if a.Type != AccountTypeOAuth {
|
||||
return false
|
||||
}
|
||||
case OpenAIEndpointCapabilityEmbeddings:
|
||||
if a.Type != AccountTypeAPIKey {
|
||||
return false
|
||||
@@ -1409,6 +1417,9 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa
|
||||
if !found {
|
||||
return true
|
||||
}
|
||||
if capability == OpenAIEndpointCapabilityAlphaSearch && configured[string(OpenAIEndpointCapabilityChatCompletions)] {
|
||||
return true
|
||||
}
|
||||
return configured[string(capability)]
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -39,21 +40,30 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co
|
||||
if upstreamModel != "" && upstreamModel != requestedModel {
|
||||
body = ReplaceModelInBody(body, upstreamModel)
|
||||
}
|
||||
sanitizedBody, err := sanitizeOpenAIAlphaSearchBody(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("sanitize alpha search request body: %w", err)
|
||||
}
|
||||
body = sanitizedBody
|
||||
|
||||
token, _, err := s.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := s.buildOpenAIAlphaSearchRequest(ctx, c, account, body, token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
if err := s.ensureOpenAIAlphaSearchAuthMetadata(ctx, account, token, proxyURL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := s.buildOpenAIAlphaSearchRequest(ctx, c, account, body, token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
upstreamStart := time.Now()
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
|
||||
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
|
||||
@@ -71,7 +81,14 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co
|
||||
upstreamMessage := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
||||
if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) {
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel)
|
||||
// alpha/search 是独立的工具端点,单次 401 不能证明账号的模型调用
|
||||
// 凭据全局失效。若沿用通用 401 逻辑,PAT 会因没有 refresh_token
|
||||
// 被永久标记为 error;历史导入且缺少 auth_mode 标记的 at- token 也会
|
||||
// 漏过 PAT 类型判断。这里仍允许本次请求换号,但不修改任何账号状态;
|
||||
// 真正的凭据失效由普通 Responses 请求或 whoami 校验判定。
|
||||
if shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(resp.StatusCode) {
|
||||
s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel)
|
||||
}
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
@@ -103,15 +120,6 @@ func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Co
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) {
|
||||
clientBeta := ""
|
||||
if c != nil {
|
||||
clientBeta = c.GetHeader("OpenAI-Beta")
|
||||
}
|
||||
req, err := s.buildUpstreamRequestOpenAIPassthrough(ctx, c, account, body, token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
targetURL, err := s.openAIAlphaSearchURL(account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -129,19 +137,165 @@ func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context
|
||||
}
|
||||
parsedURL.RawQuery = query.Encode()
|
||||
}
|
||||
req.URL = parsedURL
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, parsedURL.String(), bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
|
||||
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build openai authentication headers: %w", err)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
if clientBeta == "" {
|
||||
req.Header.Del("OpenAI-Beta")
|
||||
}
|
||||
if version := strings.TrimSpace(c.GetHeader("Version")); version != "" {
|
||||
req.Header.Set("Version", version)
|
||||
} else if account.Type == AccountTypeOAuth {
|
||||
req.Header.Set("Version", codexCLIVersion)
|
||||
|
||||
if account.Type == AccountTypeOAuth {
|
||||
req.Host = "chatgpt.com"
|
||||
if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, req.Header, account); err != nil {
|
||||
return nil, fmt.Errorf("resolve chatgpt account headers: %w", err)
|
||||
}
|
||||
|
||||
if turnMetadata := openAIAlphaSearchInboundHeader(c, "X-Codex-Turn-Metadata"); turnMetadata != "" {
|
||||
req.Header.Set("X-Codex-Turn-Metadata", turnMetadata)
|
||||
}
|
||||
if version := openAIAlphaSearchInboundHeader(c, "Version"); version != "" {
|
||||
req.Header.Set("Version", version)
|
||||
} else {
|
||||
req.Header.Set("Version", codexCLIVersion)
|
||||
}
|
||||
if originator := openAIAlphaSearchInboundHeader(c, "Originator"); originator != "" {
|
||||
req.Header.Set("Originator", originator)
|
||||
} else {
|
||||
req.Header.Set("Originator", "codex_cli_rs")
|
||||
}
|
||||
if customUA := account.GetOpenAIUserAgent(); customUA != "" {
|
||||
req.Header.Set("User-Agent", customUA)
|
||||
} else if userAgent := openAIAlphaSearchInboundHeader(c, "User-Agent"); userAgent != "" {
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
} else {
|
||||
req.Header.Set("User-Agent", codexCLIUserAgent)
|
||||
}
|
||||
if s.cfg != nil && s.cfg.Gateway.ForceCodexCLI {
|
||||
req.Header.Set("User-Agent", codexCLIUserAgent)
|
||||
}
|
||||
s.overrideBrowserUserAgent(ctx, account, req)
|
||||
enforceCodexIdentityHeaders(req.Header)
|
||||
}
|
||||
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
stripOpenAIAlphaSearchResponsesHeaders(req.Header)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// stripOpenAIAlphaSearchResponsesHeaders 让独立搜索请求与官方 Codex
|
||||
// SearchClient 的线协议保持一致。alpha/search 不是 /responses 的子请求:官方
|
||||
// 客户端仅在 Provider/Auth 基础头之外附加 x-codex-turn-metadata,不发送
|
||||
// OpenAI-Beta、会话隔离或 Responses Lite 状态头。originator 与 User-Agent
|
||||
// 属于官方默认客户端头,必须保留。
|
||||
//
|
||||
// alpha/search 使用专用构造器生成官方 SearchClient 的最小线协议形态;
|
||||
// 该函数作为最后一道防线,避免账号 header 覆写或后续改动重新带入
|
||||
// Responses 专用头,使 PAT 的 alpha/search 被上游按错误认证路径处理。
|
||||
func stripOpenAIAlphaSearchResponsesHeaders(headers http.Header) {
|
||||
if headers == nil {
|
||||
return
|
||||
}
|
||||
for _, key := range []string{
|
||||
"OpenAI-Beta",
|
||||
"Session_ID",
|
||||
"Conversation_ID",
|
||||
"X-Codex-Beta-Features",
|
||||
"X-Codex-Turn-State",
|
||||
responsesLiteHeaderKey,
|
||||
} {
|
||||
headers.Del(key)
|
||||
}
|
||||
}
|
||||
|
||||
func openAIAlphaSearchInboundHeader(c *gin.Context, key string) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(c.GetHeader(key))
|
||||
}
|
||||
|
||||
var openAIAlphaSearchUnsupportedBodyFields = [...]string{
|
||||
// Codex alpha/search 是 SearchRequest 独立协议,不是 /responses 子请求。
|
||||
// 新版 Codex/第三方代理可能把 Responses 公共字段误带到搜索请求里;ChatGPT
|
||||
// alpha/search 会对这些字段返回 Unknown parameter(例如 prompt_cache_key)。
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_retention",
|
||||
}
|
||||
|
||||
func sanitizeOpenAIAlphaSearchBody(body []byte) ([]byte, error) {
|
||||
if len(body) == 0 {
|
||||
return body, nil
|
||||
}
|
||||
var obj map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &obj); err != nil || obj == nil {
|
||||
return body, nil
|
||||
}
|
||||
changed := false
|
||||
for _, field := range openAIAlphaSearchUnsupportedBodyFields {
|
||||
if _, ok := obj[field]; ok {
|
||||
delete(obj, field)
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return body, nil
|
||||
}
|
||||
out, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) ensureOpenAIAlphaSearchAuthMetadata(ctx context.Context, account *Account, token string, proxyURL string) error {
|
||||
if s == nil || account == nil || !account.IsOpenAIPersonalAccessToken() {
|
||||
return nil
|
||||
}
|
||||
if strings.TrimSpace(account.GetChatGPTAccountID()) != "" {
|
||||
return nil
|
||||
}
|
||||
var oauthService *OpenAIOAuthService
|
||||
if s.openAITokenProvider != nil {
|
||||
oauthService = s.openAITokenProvider.openAIOAuthService
|
||||
}
|
||||
if oauthService == nil {
|
||||
return nil
|
||||
}
|
||||
tokenInfo, err := oauthService.ValidateCodexPersonalAccessToken(ctx, token, proxyURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("validate Codex PAT metadata for alpha/search: %w", err)
|
||||
}
|
||||
credentials := shallowCopyMap(account.Credentials)
|
||||
for key, value := range oauthService.BuildAccountCredentials(tokenInfo) {
|
||||
credentials[key] = value
|
||||
}
|
||||
credentials = NormalizeOpenAIPersonalAccessTokenCredentials(account, tokenInfo, credentials)
|
||||
account.Credentials = shallowCopyMap(credentials)
|
||||
if s.accountRepo != nil {
|
||||
if err := persistAccountCredentials(ctx, s.accountRepo, account, credentials); err != nil {
|
||||
return fmt.Errorf("persist Codex PAT metadata for alpha/search: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(statusCode int) bool {
|
||||
return statusCode != http.StatusUnauthorized
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIAlphaSearchURL(account *Account) (string, error) {
|
||||
if account == nil {
|
||||
return "", fmt.Errorf("account is required")
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -15,6 +16,18 @@ import (
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type alphaSearchAccountStateRepo struct {
|
||||
AccountRepository
|
||||
setErrorCalls int
|
||||
lastError string
|
||||
}
|
||||
|
||||
func (r *alphaSearchAccountStateRepo) SetError(_ context.Context, _ int64, errorMsg string) error {
|
||||
r.setErrorCalls++
|
||||
r.lastError = errorMsg
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{
|
||||
@@ -70,6 +83,136 @@ func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) {
|
||||
require.JSONEq(t, string(body), string(upstream.lastBody))
|
||||
}
|
||||
|
||||
func TestForwardAlphaSearchPATUsesStandaloneHeaderShape(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{
|
||||
"id":"search-session",
|
||||
"model":"gpt-5.6-sol",
|
||||
"commands":{"search_query":[{"q":"OpenAI news"}]},
|
||||
"prompt_cache_key":"responses-cache-key",
|
||||
"prompt_cache_retention":"24h"
|
||||
}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Request.Header.Set("User-Agent", codexCLIUserAgent)
|
||||
c.Request.Header.Set("Originator", "codex_cli_rs")
|
||||
c.Request.Header.Set("Version", "0.144.1")
|
||||
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
c.Request.Header.Set("Accept-Language", "zh-CN")
|
||||
c.Request.Header.Set("Authorization", "Bearer client-token")
|
||||
c.Request.Header.Set("Session_ID", "session-client")
|
||||
c.Request.Header.Set("Conversation_ID", "conversation-client")
|
||||
c.Request.Header.Set("X-Codex-Beta-Features", "feature-a")
|
||||
c.Request.Header.Set("X-Codex-Turn-State", "turn-state")
|
||||
c.Request.Header.Set(responsesLiteHeaderKey, "true")
|
||||
c.Request.Header.Set("X-Codex-Turn-Metadata", `{"turn_id":"turn-1"}`)
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"output":"search result"}`)),
|
||||
}}
|
||||
service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{
|
||||
ID: 43,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "at-test-token",
|
||||
"auth_mode": OpenAIAuthModePersonalAccessToken,
|
||||
"chatgpt_account_id": "chatgpt-account",
|
||||
"chatgpt_account_is_fedramp": true,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "Bearer at-test-token", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Equal(t, "chatgpt-account", upstream.lastReq.Header.Get("ChatGPT-Account-ID"))
|
||||
require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp"))
|
||||
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
|
||||
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Accept"))
|
||||
require.Equal(t, "0.144.1", upstream.lastReq.Header.Get("Version"))
|
||||
require.Equal(t, `{"turn_id":"turn-1"}`, upstream.lastReq.Header.Get("X-Codex-Turn-Metadata"))
|
||||
require.Equal(t, "codex_cli_rs", upstream.lastReq.Header.Get("Originator"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("OpenAI-Beta"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("Session_ID"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("Conversation_ID"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Beta-Features"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Turn-State"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get(responsesLiteHeaderKey))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("Accept-Language"))
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_retention").Exists())
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
}
|
||||
|
||||
func TestForwardAlphaSearchPATBackfillsMissingChatGPTAccountMetadata(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"OpenAI news"}]}}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
var whoamiCalls int32
|
||||
whoamiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt32(&whoamiCalls, 1)
|
||||
require.Equal(t, "Bearer at-test-token", r.Header.Get("Authorization"))
|
||||
require.Equal(t, "application/json", r.Header.Get("Accept"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"email":"pat@example.com",
|
||||
"chatgpt_user_id":"user-123",
|
||||
"chatgpt_account_id":"acct-123",
|
||||
"chatgpt_plan_type":"plus",
|
||||
"chatgpt_account_is_fedramp":true
|
||||
}`))
|
||||
}))
|
||||
defer whoamiServer.Close()
|
||||
oldWhoamiURL := openAICodexPATWhoamiURL
|
||||
openAICodexPATWhoamiURL = whoamiServer.URL
|
||||
defer func() { openAICodexPATWhoamiURL = oldWhoamiURL }()
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"output":"search result"}`)),
|
||||
}}
|
||||
oauthService := NewOpenAIOAuthService(nil, nil)
|
||||
service := &OpenAIGatewayService{
|
||||
cfg: &config.Config{},
|
||||
httpUpstream: upstream,
|
||||
openAITokenProvider: NewOpenAITokenProvider(nil, nil, oauthService),
|
||||
}
|
||||
account := &Account{
|
||||
ID: 45,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "at-test-token",
|
||||
"auth_mode": OpenAIAuthModePersonalAccessToken,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, int32(1), atomic.LoadInt32(&whoamiCalls))
|
||||
require.Equal(t, "acct-123", upstream.lastReq.Header.Get("ChatGPT-Account-ID"))
|
||||
require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp"))
|
||||
require.Equal(t, "acct-123", account.Credentials["chatgpt_account_id"])
|
||||
require.Equal(t, "user-123", account.Credentials["chatgpt_user_id"])
|
||||
require.Equal(t, OpenAIAuthModePersonalAccessToken, account.Credentials["auth_mode"])
|
||||
}
|
||||
|
||||
func TestForwardAlphaSearchAPIKeyMapsModelAndPassesThroughError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
|
||||
@@ -143,3 +286,52 @@ func TestForwardAlphaSearchReturnsFailoverBeforeWriting(t *testing.T) {
|
||||
require.False(t, c.Writer.Written())
|
||||
require.Empty(t, recorder.Body.String())
|
||||
}
|
||||
|
||||
func TestForwardAlphaSearchUnauthorizedDoesNotMarkAccountError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"detail":"Unauthorized"}`)),
|
||||
}}
|
||||
repo := &alphaSearchAccountStateRepo{}
|
||||
cfg := &config.Config{}
|
||||
service := &OpenAIGatewayService{
|
||||
cfg: cfg,
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil),
|
||||
}
|
||||
account := &Account{
|
||||
ID: 44,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
// 刻意不设置 auth_mode:覆盖历史上把 at- token 当普通 OAuth 导入的账号。
|
||||
"access_token": "at-test-token",
|
||||
"chatgpt_account_id": "chatgpt-account",
|
||||
},
|
||||
}
|
||||
|
||||
result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
|
||||
require.Zero(t, repo.setErrorCalls)
|
||||
require.Empty(t, repo.lastError)
|
||||
require.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestShouldApplyOpenAIAlphaSearchAccountErrorSideEffects(t *testing.T) {
|
||||
require.False(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusUnauthorized))
|
||||
require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusForbidden))
|
||||
require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusTooManyRequests))
|
||||
}
|
||||
|
||||
@@ -492,9 +492,24 @@ func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) {
|
||||
}
|
||||
|
||||
require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
|
||||
require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
|
||||
require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
|
||||
})
|
||||
|
||||
t.Run("alpha search 仅允许 OpenAI OAuth/PAT 类账号", func(t *testing.T) {
|
||||
apiKey := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
}
|
||||
oauth := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
}
|
||||
|
||||
require.False(t, apiKey.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
|
||||
require.True(t, oauth.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
|
||||
})
|
||||
|
||||
t.Run("显式列表支持同时声明 chat 和 embeddings", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
@@ -518,9 +533,22 @@ func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) {
|
||||
}
|
||||
|
||||
require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityChatCompletions))
|
||||
require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
|
||||
require.False(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityEmbeddings))
|
||||
})
|
||||
|
||||
t.Run("OAuth 显式列表沿用 chat 能力放行 alpha search", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"openai_capabilities": []any{"chat_completions"},
|
||||
},
|
||||
}
|
||||
|
||||
require.True(t, account.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityAlphaSearch))
|
||||
})
|
||||
|
||||
t.Run("显式 map 支持单独关闭 chat 并开启 embeddings", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
|
||||
Reference in New Issue
Block a user