| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424 |
- package service
- import (
- "context"
- "io"
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/gin-gonic/gin"
- "github.com/stretchr/testify/require"
- )
- type openAIStream403AccountRepo struct {
- AccountRepository
- setErrorCalls int
- }
- func (r *openAIStream403AccountRepo) SetError(context.Context, int64, string) error {
- r.setErrorCalls++
- return nil
- }
- type openAIAuthPolicyAccountRepo struct {
- AccountRepository
- tempCalls int
- setErrorCalls int
- }
- func (r *openAIAuthPolicyAccountRepo) SetTempUnschedulable(context.Context, int64, time.Time, string) error {
- r.tempCalls++
- return nil
- }
- func (r *openAIAuthPolicyAccountRepo) SetError(context.Context, int64, string) error {
- r.setErrorCalls++
- return nil
- }
- type openAIAuthPolicy403Counter struct {
- counts []int64
- }
- func (s *openAIAuthPolicy403Counter) IncrementOpenAI403Count(context.Context, int64, int) (int64, error) {
- if len(s.counts) == 0 {
- return 1, nil
- }
- count := s.counts[0]
- s.counts = s.counts[1:]
- return count, nil
- }
- func (*openAIAuthPolicy403Counter) ResetOpenAI403Count(context.Context, int64) error {
- return nil
- }
- func TestOpenAIUpstreamAccessStateClassification(t *testing.T) {
- tests := []struct {
- name string
- body string
- want bool
- }{
- {"workspace_code", `{"detail":{"code":"deactivated_workspace"}}`, true},
- {"disabled_account_message", `{"error":{"message":"Your account is disabled"}}`, false},
- {"suspended_workspace_message", `{"response":{"error":{"message":"This workspace has been suspended"}}}`, false},
- {"deactivated_organization_message", `{"detail":{"message":"The organization is deactivated"}}`, false},
- {"scalar_detail", `{"detail":"This workspace has been disabled"}`, false},
- {"suspended_org_code", `{"error":{"code":"org_suspended"}}`, true},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- body := []byte(tt.body)
- require.Equal(t, tt.want, isOpenAIUpstreamAccessStateError("", body))
- if !tt.want {
- return
- }
- require.True(t, (&OpenAIGatewayService{}).shouldFailoverOpenAIUpstreamResponse(newOpenAIUpstreamErrorTestAccount(), http.StatusForbidden, "", body))
- require.True(t, shouldFailoverOpenAIPassthroughResponse(&Account{Type: AccountTypeOAuth}, http.StatusForbidden, body))
- err := newOpenAIUpstreamFailoverError(http.StatusForbidden, nil, body, "", true)
- require.True(t, err.IsCredentialFailure())
- require.Equal(t, GatewayFailureScopeAccount, err.Scope)
- require.Equal(t, OpenAIUpstreamAccessStateReason, err.Reason)
- require.Equal(t, NextAccountRetry, err.NextAccountAction)
- require.False(t, err.RetryableOnSameAccount)
- require.False(t, err.RequestScopedTransient)
- require.Equal(t, http.StatusBadGateway, err.ClientStatusCode)
- require.Equal(t, openAIUpstreamAccessUnavailableClientMessage, err.ClientMessage)
- })
- }
- }
- func TestOpenAIUpstreamAccessStateDoesNotScanEchoedJSON(t *testing.T) {
- body := []byte(`{"error":{"code":"invalid_request_error","message":"Invalid input"},"echo":{"prompt":"my account is disabled"}}`)
- require.False(t, isOpenAIUpstreamAccessStateError("", body))
- require.False(t, shouldFailoverOpenAIPassthroughResponse(&Account{Type: AccountTypeOAuth}, http.StatusBadRequest, body))
- }
- func TestOpenAIHTTPAccessStateDoesNotTrustBadRequestMessage(t *testing.T) {
- body := []byte(`{"error":{"type":"invalid_request_error","code":"unknown_parameter","message":"Unknown parameter: account disabled"}}`)
- svc := &OpenAIGatewayService{}
- require.False(t, isOpenAIUpstreamAccessStateError("", body), "free-form stream messages are not durable account evidence")
- require.False(t, isOpenAIHTTPUpstreamAccessStateError(http.StatusBadRequest, "", body))
- require.False(t, svc.shouldFailoverOpenAIUpstreamResponse(newOpenAIUpstreamErrorTestAccount(), http.StatusBadRequest, "", body))
- require.False(t, shouldFailoverOpenAIPassthroughResponse(&Account{Type: AccountTypeOAuth}, http.StatusBadRequest, body))
- err := newOpenAIUpstreamFailoverError(http.StatusBadRequest, nil, body, "", false)
- require.False(t, err.IsCredentialFailure())
- }
- func TestOpenAIHTTPAccessStateBadRequestDoesNotDisableAccount(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
- account := &Account{ID: 925, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
- body := []byte(`{"error":{"code":"unknown_parameter","message":"Unknown parameter: account disabled"}}`)
- disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, nil, body)
- require.False(t, disabled)
- require.Zero(t, repo.setErrorCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIStreamEchoedAccessStateMessageDoesNotDisableOrFailover(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
- account := &Account{ID: 926, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
- payload := []byte(`{"type":"response.failed","response":{"error":{"type":"invalid_request_error","code":"unknown_parameter","message":"Unknown parameter: account disabled"}}}`)
- message := extractOpenAISSEErrorMessage(payload)
- require.False(t, isOpenAIUpstreamAccessStateError(message, payload))
- require.False(t, openAIStreamFailedEventShouldFailover(payload, message))
- status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, message, nil)
- require.Equal(t, http.StatusBadGateway, status)
- require.False(t, disabled)
- require.Zero(t, repo.setErrorCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIHTTPAccessStateTrustsStructuredCode(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
- account := &Account{ID: 930, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
- body := []byte(`{"error":{"code":"organization_deactivated","message":"request rejected"}}`)
- require.True(t, isOpenAIHTTPUpstreamAccessStateError(http.StatusBadRequest, "", body))
- require.True(t, svc.shouldFailoverOpenAIUpstreamResponse(newOpenAIUpstreamErrorTestAccount(), http.StatusBadRequest, "", body))
- require.True(t, svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, nil, body))
- require.Equal(t, 1, repo.setErrorCalls)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIHTTPAuthMessagesUseExistingStatusPolicies(t *testing.T) {
- t.Run("oauth 401 remains recoverable", func(t *testing.T) {
- repo := &openAIAuthPolicyAccountRepo{}
- rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
- svc := &OpenAIGatewayService{rateLimitService: rateLimits}
- account := &Account{ID: 931, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true,
- Credentials: map[string]any{"refresh_token": "refreshable"}}
- body := []byte(`{"error":{"message":"account is disabled"}}`)
- require.False(t, isOpenAIHTTPUpstreamAccessStateError(http.StatusUnauthorized, "", body))
- require.True(t, svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusUnauthorized, nil, body))
- require.Zero(t, repo.setErrorCalls)
- require.Equal(t, 1, repo.tempCalls)
- })
- t.Run("403 uses counter cooldown", func(t *testing.T) {
- repo := &openAIAuthPolicyAccountRepo{}
- counter := &openAIAuthPolicy403Counter{counts: []int64{1}}
- rateLimits := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
- rateLimits.openAI403CounterCache = counter
- svc := &OpenAIGatewayService{rateLimitService: rateLimits}
- rateLimits.SetAccountRuntimeBlocker(svc)
- account := &Account{ID: 932, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true}
- body := []byte(`{"error":{"message":"workspace has been suspended"}}`)
- require.False(t, isOpenAIHTTPUpstreamAccessStateError(http.StatusForbidden, "", body))
- require.True(t, svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body))
- require.Zero(t, repo.setErrorCalls)
- require.Equal(t, 1, repo.tempCalls)
- })
- }
- func TestOpenAICyberPolicyWrapped5xxNeverFailsOver(t *testing.T) {
- body := []byte(`{"error":{"code":"cyber_policy","message":"blocked"}}`)
- svc := &OpenAIGatewayService{}
- require.False(t, svc.shouldFailoverOpenAIUpstreamResponse(newOpenAIUpstreamErrorTestAccount(), http.StatusBadGateway, "wrapped upstream failure", body))
- require.False(t, shouldFailoverOpenAIPassthroughResponse(&Account{Type: AccountTypeOAuth}, http.StatusBadGateway, body))
- }
- func TestOpenAICapacityFailoverCarriesSafeTerminalResponse(t *testing.T) {
- message := "Our servers are currently overloaded. Please try again later."
- body := []byte(`{"error":{"code":"server_is_overloaded","message":"` + message + `"}}`)
- err := newOpenAIUpstreamFailoverError(http.StatusBadRequest, nil, body, message, false)
- require.True(t, err.IsOpenAICapacityShed())
- require.Equal(t, http.StatusServiceUnavailable, err.ClientStatusCode)
- require.Equal(t, message, err.ClientMessage)
- require.NotContains(t, err.ClientMessage, "server_is_overloaded")
- }
- func TestOpenAIStreamSemanticStatusesPreservedAcrossTerminalShapes(t *testing.T) {
- tests := []struct {
- name string
- body string
- status int
- wantFailover bool
- }{
- {"unauthorized", `{"type":"error","error":{"type":"authentication_error","code":"invalid_api_key","message":"unauthorized"}}`, http.StatusUnauthorized, true},
- {"forbidden", `{"type":"response.failed","response":{"error":{"type":"permission_error","message":"forbidden"}}}`, http.StatusForbidden, false},
- {"rate_limit", `{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down"}}`, http.StatusTooManyRequests, true},
- {"overload_529", `{"type":"error","error":{"status_code":529,"code":"overloaded","message":"overloaded"}}`, 529, true},
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- payload := []byte(tt.body)
- message := extractOpenAISSEErrorMessage(payload)
- require.Equal(t, tt.status, openAIStreamFailureStatus(payload, message))
- require.Equal(t, tt.wantFailover, openAIStreamErrorEventShouldFailover(payload, message))
- })
- }
- }
- func TestOpenAIStreamBareErrorUsesSemanticFailover(t *testing.T) {
- payload := []byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down"}}`)
- require.True(t, openAIStreamErrorEventShouldFailover(payload, "slow down"))
- }
- func TestOpenAIStream403FailoverRequiresStructuredAccountCredentialSignal(t *testing.T) {
- tests := []struct {
- name string
- payload string
- want bool
- }{
- {
- name: "ordinary permission error",
- payload: `{"type":"response.failed","response":{"error":{"type":"permission_error","code":"forbidden","message":"access denied for this request"}}}`,
- },
- {
- name: "explicit 403 request status",
- payload: `{"type":"error","error":{"type":"permission_error","code":"forbidden","status_code":403,"message":"forbidden content"}}`,
- },
- {
- name: "structured access state",
- payload: `{"type":"response.failed","response":{"error":{"code":"workspace_suspended","message":"workspace is suspended"}}}`,
- want: true,
- },
- {
- name: "explicit credential auth code",
- payload: `{"type":"error","error":{"type":"permission_error","code":"invalid_api_key","status_code":403,"message":"credential rejected"}}`,
- want: true,
- },
- {
- name: "explicit authentication type",
- payload: `{"type":"error","error":{"type":"authentication_error","status_code":403,"message":"credential rejected"}}`,
- want: true,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- payload := []byte(tt.payload)
- message := extractOpenAISSEErrorMessage(payload)
- require.Equal(t, http.StatusForbidden, openAIStreamFailureStatus(payload, message))
- require.Equal(t, tt.want, openAIStreamFailedEventShouldFailover(payload, message))
- require.Equal(t, tt.want, openAIStreamErrorEventShouldFailover(payload, message))
- })
- }
- }
- func TestOpenAIStream403PostOutputAccountSideEffectsIgnoreRequestPermissionErrors(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- rateLimits := &RateLimitService{accountRepo: repo}
- svc := &OpenAIGatewayService{rateLimitService: rateLimits}
- account := &Account{ID: 918, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- payload := []byte(`{"type":"error","error":{"type":"permission_error","code":"forbidden","status_code":403,"message":"access denied for this request"}}`)
- status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, "access denied for this request", nil)
- require.Equal(t, http.StatusForbidden, status)
- require.False(t, disabled)
- require.Zero(t, repo.setErrorCalls)
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIStream403ExplicitCredentialAuthAppliesAccountSideEffects(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- rateLimits := &RateLimitService{accountRepo: repo}
- svc := &OpenAIGatewayService{rateLimitService: rateLimits}
- account := &Account{ID: 917, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- payload := []byte(`{"type":"error","error":{"type":"permission_error","code":"invalid_api_key","status_code":403,"message":"credential rejected"}}`)
- status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, "credential rejected", nil)
- require.Equal(t, http.StatusForbidden, status)
- require.True(t, disabled)
- require.Equal(t, 1, repo.setErrorCalls)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIWSStandaloneFailedStructured403AppliesAccountSideEffectsOnce(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
- account := &Account{ID: 923, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- failed := []byte(`{"type":"response.failed","response":{"error":{"type":"permission_error","code":"invalid_api_key","status_code":403,"message":"credential rejected"}}}`)
- require.True(t, svc.handleOpenAIWSFailureAccountSideEffects(context.Background(), account, "gpt-5", nil, failed))
- require.Equal(t, 1, repo.setErrorCalls)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIWSPairedStructured403SideEffectsCanBeDeduplicated(t *testing.T) {
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
- account := &Account{ID: 924, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- errorEvent := []byte(`{"type":"error","error":{"code":"workspace_suspended","status_code":403,"message":"workspace is suspended"}}`)
- failedEvent := []byte(`{"type":"response.failed","response":{"error":{"code":"workspace_suspended","status_code":403,"message":"workspace is suspended"}}}`)
- applied := svc.handleOpenAIWSFailureAccountSideEffects(context.Background(), account, "gpt-5", nil, errorEvent)
- if !applied {
- applied = svc.handleOpenAIWSFailureAccountSideEffects(context.Background(), account, "gpt-5", nil, failedEvent)
- }
- require.True(t, applied)
- require.Equal(t, 1, repo.setErrorCalls)
- }
- func TestOpenAIStreamAccessStateAppliesAccountHealthBeforeFailover(t *testing.T) {
- svc := &OpenAIGatewayService{}
- account := &Account{ID: 919, Platform: PlatformOpenAI, Type: AccountTypeSetupToken}
- payload := []byte(`{"type":"response.failed","response":{"error":{"code":"workspace_suspended","message":"workspace is suspended"}}}`)
- status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, "workspace is suspended", nil)
- require.Equal(t, http.StatusForbidden, status)
- require.True(t, disabled)
- require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
- }
- func TestOpenAIStreamPairedFailureAppliesAccountSideEffectsOnce(t *testing.T) {
- const upstream = "data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n" +
- "data: {\"type\":\"error\",\"error\":{\"status_code\":403,\"code\":\"workspace_suspended\",\"message\":\"workspace is suspended\"}}\n\n" +
- "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_failed\",\"error\":{\"status_code\":403,\"code\":\"workspace_suspended\",\"message\":\"workspace is suspended\"}}}\n\n"
- t.Run("native", func(t *testing.T) {
- gin.SetMode(gin.TestMode)
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{
- cfg: &config.Config{},
- toolCorrector: NewCodexToolCorrector(),
- rateLimitService: &RateLimitService{accountRepo: repo},
- }
- account := &Account{ID: 921, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- recorder := newOpenAIResponseFlushRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- Body: io.NopCloser(strings.NewReader(upstream)),
- }
- result, err := svc.handleStreamingResponse(context.Background(), resp, c, account, time.Now(), "gpt-5", "gpt-5")
- require.Error(t, err)
- require.NotNil(t, result)
- require.Equal(t, 1, repo.setErrorCalls)
- })
- t.Run("passthrough", func(t *testing.T) {
- gin.SetMode(gin.TestMode)
- repo := &openAIStream403AccountRepo{}
- svc := &OpenAIGatewayService{
- cfg: &config.Config{},
- rateLimitService: &RateLimitService{accountRepo: repo},
- }
- account := &Account{ID: 922, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- recorder := httptest.NewRecorder()
- c, _ := gin.CreateTestContext(recorder)
- c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
- writer := &passthroughFlushTestWriter{
- ResponseWriter: c.Writer,
- recorder: recorder,
- failAfterWrites: -1,
- }
- c.Writer = writer
- resp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{"Content-Type": []string{"text/event-stream"}},
- Body: io.NopCloser(strings.NewReader(upstream)),
- }
- result, err := svc.handleStreamingResponsePassthrough(
- context.Background(), resp, c, account, time.Now(), "gpt-5", "gpt-5",
- )
- require.Error(t, err)
- require.NotNil(t, result)
- require.Equal(t, 1, repo.setErrorCalls)
- })
- }
- func TestOpenAIStreamOAuthLike429GetsDeadlineWithoutImmediateRuntimeBlock(t *testing.T) {
- for _, accountType := range []string{AccountTypeOAuth, AccountTypeSetupToken} {
- t.Run(accountType, func(t *testing.T) {
- svc := &OpenAIGatewayService{}
- account := &Account{ID: 920, Platform: PlatformOpenAI, Type: accountType}
- payload := []byte(`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"slow down"}}`)
- status, disabled := svc.handleOpenAIStreamTerminalAccountSideEffects(nil, account, payload, "slow down", nil)
- err := svc.newOpenAIAccountFailoverError(account, status, nil, payload, "slow down", disabled, false)
- require.Equal(t, http.StatusTooManyRequests, status)
- require.False(t, disabled)
- require.True(t, err.RetryableOnSameAccount)
- require.False(t, err.SameAccountRetryDeadline.IsZero())
- require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
- })
- }
- }
|