mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 05:45:26 +08:00
248 lines
11 KiB
Go
248 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func TestIsOpenAIOAuthLike(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
account *Account
|
|
want bool
|
|
codex bool
|
|
}{
|
|
{name: "openai_oauth", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}, want: true, codex: true},
|
|
{name: "openai_setup_token", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeSetupToken}, want: true, codex: true},
|
|
{name: "openai_api_key", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, want: false, codex: false},
|
|
{name: "anthropic_setup_token", account: &Account{Platform: PlatformAnthropic, Type: AccountTypeSetupToken}, want: false, codex: false},
|
|
{name: "grok_setup_token", account: &Account{Platform: PlatformGrok, Type: AccountTypeSetupToken}, want: false, codex: false},
|
|
{name: "nil", account: nil, want: false, codex: false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
require.Equal(t, tt.want, tt.account.IsOpenAIOAuthLike())
|
|
require.Equal(t, tt.codex, tt.account.UsesOpenAICodexProtocol())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceGetAccessTokenSetupToken(t *testing.T) {
|
|
svc := &OpenAIGatewayService{openAITokenProvider: &OpenAITokenProvider{}}
|
|
account := &Account{
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeSetupToken,
|
|
Credentials: map[string]any{"access_token": "setup-token-value"},
|
|
}
|
|
|
|
token, tokenType, err := svc.GetAccessToken(context.Background(), account)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "setup-token-value", token)
|
|
require.Equal(t, "oauth", tokenType)
|
|
|
|
delete(account.Credentials, "access_token")
|
|
_, _, err = svc.GetAccessToken(context.Background(), account)
|
|
require.EqualError(t, err, "access_token not found in credentials")
|
|
|
|
for _, platform := range []string{PlatformAnthropic, PlatformGrok} {
|
|
foreign := &Account{
|
|
Platform: platform,
|
|
Type: AccountTypeSetupToken,
|
|
Credentials: map[string]any{"access_token": "foreign-token"},
|
|
}
|
|
_, _, err = svc.GetAccessToken(context.Background(), foreign)
|
|
require.EqualError(t, err, "unsupported account type: setup-token")
|
|
}
|
|
}
|
|
|
|
func TestOpenAISetupTokenImagesUsesOAuthResponsesPath(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
|
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusTooManyRequests,
|
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
|
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
|
|
}}
|
|
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
|
account := &Account{
|
|
ID: 73,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeSetupToken,
|
|
Credentials: map[string]any{"access_token": "setup-token"},
|
|
}
|
|
parsed := &OpenAIImagesRequest{
|
|
Endpoint: openAIImagesGenerationsEndpoint,
|
|
Model: "gpt-image-2",
|
|
Prompt: "draw a square",
|
|
N: 1,
|
|
ResponseFormat: "b64_json",
|
|
}
|
|
|
|
result, err := svc.ForwardImages(context.Background(), c, account, nil, parsed, "")
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
|
require.True(t, failoverErr.RetryableOnSameAccount)
|
|
require.False(t, failoverErr.SameAccountRetryDeadline.IsZero())
|
|
require.Contains(t, upstream.lastReq.URL.String(), "/backend-api/codex/responses")
|
|
}
|
|
|
|
func TestOpenAISetupTokenWSCompatibility(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{}`))
|
|
c.Request.Header.Set("session_id", "session-one")
|
|
c.Set("api_key", &APIKey{ID: 17})
|
|
|
|
account := &Account{
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeSetupToken,
|
|
Credentials: map[string]any{
|
|
"access_token": "setup-token-value",
|
|
"chatgpt_account_id": "chatgpt-setup",
|
|
},
|
|
}
|
|
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
|
|
|
wsURL, err := svc.buildOpenAIResponsesWSURL(account)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "wss://chatgpt.com/backend-api/codex/responses", wsURL)
|
|
foreignURL, err := svc.buildOpenAIResponsesWSURL(&Account{Platform: PlatformGrok, Type: AccountTypeSetupToken})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "wss://api.openai.com/v1/responses", foreignURL)
|
|
|
|
headers, session, err := svc.buildOpenAIWSHeaders(
|
|
context.Background(), c, account, "setup-token-value",
|
|
OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2},
|
|
true, "", "", "", "gpt-5.1-codex", "",
|
|
)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "Bearer setup-token-value", headers.Get("authorization"))
|
|
require.Equal(t, "chatgpt-setup", headers.Get("chatgpt-account-id"))
|
|
require.NotEmpty(t, headers.Get("originator"))
|
|
require.Equal(t, "session-one", session.SessionID)
|
|
require.NotEqual(t, session.SessionID, headers.Get("session_id"))
|
|
|
|
payload := svc.buildOpenAIWSCreatePayload(map[string]any{"store": true}, account)
|
|
require.Equal(t, false, payload["store"])
|
|
}
|
|
|
|
func TestOpenAISetupTokenChatCompletionsUsesCodexTransform(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"setup instructions"},{"role":"user","content":"hello"}],"stream":false}`)
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusBadRequest,
|
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
|
Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop after request capture"}}`)),
|
|
}}
|
|
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
|
account := openAISetupTokenCompatAccount(71)
|
|
|
|
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4")
|
|
|
|
require.Error(t, err)
|
|
require.Nil(t, result)
|
|
require.NotNil(t, upstream.lastReq)
|
|
require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
|
|
require.Equal(t, "Bearer setup-token-value", upstream.lastReq.Header.Get("Authorization"))
|
|
require.Equal(t, "chatgpt-setup", upstream.lastReq.Header.Get("chatgpt-account-id"))
|
|
require.NotEmpty(t, upstream.lastReq.Header.Get("originator"))
|
|
require.Equal(t, "setup instructions", gjson.GetBytes(upstream.lastBody, "instructions").String())
|
|
require.Equal(t, int64(1), gjson.GetBytes(upstream.lastBody, "input.#").Int())
|
|
require.Equal(t, "user", gjson.GetBytes(upstream.lastBody, "input.0.role").String())
|
|
require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
|
|
}
|
|
|
|
func TestOpenAISetupTokenMessagesUsesCodexBridgeAndTurnState(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
firstResp := openAICompatSSECompletedResponse("resp_setup_first", "gpt-5.4")
|
|
firstResp.Header.Set("x-codex-turn-state", "turn_state_setup")
|
|
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
|
firstResp,
|
|
openAICompatSSECompletedResponse("resp_setup_second", "gpt-5.4"),
|
|
}}
|
|
svc := &OpenAIGatewayService{
|
|
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
|
|
httpUpstream: upstream,
|
|
}
|
|
account := openAISetupTokenCompatAccount(72)
|
|
|
|
messages := make([]string, 0, openAICompatAnthropicReplayMaxTailMessages+3)
|
|
for i := 0; i < openAICompatAnthropicReplayMaxTailMessages+3; i++ {
|
|
messages = append(messages, `{"role":"user","content":"message-`+fmt.Sprintf("%02d", i)+`"}`)
|
|
}
|
|
firstBody := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[` + strings.Join(messages, ",") + `],"stream":false}`)
|
|
firstRec := httptest.NewRecorder()
|
|
firstCtx, _ := gin.CreateTestContext(firstRec)
|
|
firstCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(firstBody))
|
|
firstCtx.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
firstResult, err := svc.ForwardAsAnthropic(context.Background(), firstCtx, account, firstBody, "stable-cache-key", "gpt-5.4")
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, firstResult)
|
|
require.True(t, isOpenAICompatMessagesBridgeContext(firstCtx))
|
|
require.Equal(t, int64(openAICompatAnthropicReplayMaxTailMessages+4), gjson.GetBytes(upstream.bodies[0], "input.#").Int())
|
|
require.Equal(t, "developer", gjson.GetBytes(upstream.bodies[0], "input.0.role").String())
|
|
require.Contains(t, gjson.GetBytes(upstream.bodies[0], "input.0.content.0.text").String(), openAICompatClaudeCodeTodoGuardMarker)
|
|
require.Equal(t, "message-00", gjson.GetBytes(upstream.bodies[0], "input.1.content.0.text").String())
|
|
require.False(t, gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").Exists())
|
|
require.Equal(t, chatgptCodexURL, upstream.requests[0].URL.String())
|
|
require.Equal(t, "Bearer setup-token-value", upstream.requests[0].Header.Get("Authorization"))
|
|
require.Equal(t, "chatgpt-setup", upstream.requests[0].Header.Get("chatgpt-account-id"))
|
|
requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], codexCLIUserAgent, "codex-tui")
|
|
require.Empty(t, upstream.requests[0].Header.Get("x-codex-turn-state"))
|
|
|
|
secondBody := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[{"role":"user","content":"next"}],"stream":false}`)
|
|
secondRec := httptest.NewRecorder()
|
|
secondCtx, _ := gin.CreateTestContext(secondRec)
|
|
secondCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(secondBody))
|
|
secondCtx.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
secondResult, err := svc.ForwardAsAnthropic(context.Background(), secondCtx, account, secondBody, "stable-cache-key", "gpt-5.4")
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, secondResult)
|
|
require.True(t, isOpenAICompatMessagesBridgeContext(secondCtx))
|
|
require.Equal(t, "turn_state_setup", upstream.requests[1].Header.Get("x-codex-turn-state"))
|
|
require.Equal(t, generateSessionUUID(isolateOpenAIUpstreamSessionID(0, account, "stable-cache-key")), upstream.requests[1].Header.Get("session_id"))
|
|
require.Empty(t, upstream.requests[1].Header.Get("conversation_id"))
|
|
requireOpenAIMessagesCodexIdentity(t, upstream.requests[1], codexCLIUserAgent, "codex-tui")
|
|
}
|
|
|
|
func openAISetupTokenCompatAccount(id int64) *Account {
|
|
return &Account{
|
|
ID: id,
|
|
Name: "openai-setup-token",
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeSetupToken,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"access_token": "setup-token-value",
|
|
"chatgpt_account_id": "chatgpt-setup",
|
|
},
|
|
}
|
|
}
|