diff --git a/.downstream/native-domains.md b/.downstream/native-domains.md index 50ec349c4738..9e3bf1934061 100644 --- a/.downstream/native-domains.md +++ b/.downstream/native-domains.md @@ -19,17 +19,18 @@ nothing. While no Codex runtime row is stored, ## Codex gateway borrowing -The Extensions sidebar has one Codex Borrowing entry. Its in-page navigation -switches between three independently addressable administrator pages: settings at +The Extensions sidebar has separate Codex Borrowing and Pelican Tests entries. +Borrowing's in-page navigation switches between settings at `/admin/codex-gateway-borrow`, route status at -`/admin/codex-gateway-borrow/status`, and manual comparison/history at -`/admin/codex-pelican-comparison`. Opening or navigating these pages only reads +`/admin/codex-gateway-borrow/status`. Pelican generation/history lives at +`/admin/pelican-tests`; `/admin/codex-pelican-comparison` redirects there. +Opening or navigating these pages only reads local data. Settings drafts contain configuration IDs/options only, persist in the browser session, and never affect the saved configuration used by status or generation. Only explicit actions save, prepare, verify or generate. The status page polls read-only status every five seconds while preparation/verification is active; local countdowns expire displayed routes without renewing them. Leaving -the comparison page cancels its active generation. +the Pelican page cancels its queued and active generation. `GET/PUT /admin/codex-gateway-borrow/config` owns one JSON setting, `codex_gateway_borrow_config`: `enabled`, `source_account_ids`, @@ -54,17 +55,44 @@ one finite preparation; business cache misses prepare synchronously. There is no renewal timer. `/prepare` and `/verify` explicitly prepare or revalidate. All probes only record observations, including failures before business dispatch. -`POST /admin/codex-gateway-borrow/tests` accepts a timestamped UUIDv7 and explicit -account/model pairs. It sends the fixed pelican prompt once per pair, only with -matching unexpired qualification, using at most three shared generation slots -and serializing each account. It does not prepare missing routes. An expired -or future UUID is rejected; replaying a saved UUID cannot dispatch again. -Migration 275 adds only this feature's task/result tables. Records expire 24 -hours after task creation, disappear from reads immediately, and are removed -by scoped minute cleanup. Output/errors retain their original bytes. +The historical `/admin/codex-gateway-borrow/tests` endpoints retain their +cache-qualified, 90-second request contract. They share execution admission with +the independent tests described below. + +## Independent Pelican tests + +`/admin/pelican-tests` provides local account/model options, explicit generation +with SSE progress, and stored history/detail. Candidates include all platforms, +account types and account states, using local mappings, saved catalogs and +platform defaults. Option reads never refresh credentials or request an upstream +model directory. Each selected account/model receives the fixed Pelican prompt +once. The server chooses its actual existing account sending path, model mapping, +credentials, identity headers, proxy, TLS, protocol and borrowing configuration; +it creates no API Key, group policy or client handshake, and never switches to a +different account on failure. Unsupported text capabilities retain a skip reason. + +All batches share ten model execution slots, with one request per account. +Nested borrowing probes participate in the same coordination and normal business +account capacity. Generation starts after preparation and admission; its default +budget is 600 seconds, configurable from 60 to 1800 seconds. Queue, preparation +and generation durations are recorded separately. Observed business account slots +renew every 30 seconds; lost capacity cancels the operation. Existing credential +refresh remains available, while account health, quota, scheduling, affinity, +usage and billing writes are suppressed, including background callbacks. + +Task IDs are timestamped UUIDv7 values. Expired/future IDs are rejected and replay +cannot dispatch again. Migration 275 owns the task/result tables; migration 276 +adds budgets, execution phases and actual invocation metadata in place, preserving +old records and raw BYTEA output. Existing records retain their 90-second budget. +Records expire 24 hours after task creation, immediately disappear from reads, +and are removed by scoped minute cleanup. Output/errors retain original bytes; +binary upstream envelopes use lossless Base64. Truncated or incomplete output is +recorded without another generation or automatic repair. Result previews use short random capability URLs under -`/codex-gateway-borrow/preview/`, an opaque script-enabled sandbox and their own +`/pelican-tests/preview/` (the historical borrowing path remains an alias), +expire within 30 minutes and the task's remaining lifetime, and use an opaque +script-enabled sandbox and their own response CSP. Inline animation and external scripts/styles/fonts/images are allowed. Preview documents receive no administrator token or UI bridge. Adaptation sources and retained notices are in diff --git a/.github/workflows/downstream-verify.yml b/.github/workflows/downstream-verify.yml index e3255bae73a5..5e9664d2d378 100644 --- a/.github/workflows/downstream-verify.yml +++ b/.github/workflows/downstream-verify.yml @@ -44,7 +44,7 @@ jobs: - name: Native gateway borrowing contracts working-directory: backend run: | - go test -p=1 -tags unit ./internal/service ./internal/repository ./internal/handler ./internal/handler/admin -run '^TestCodexGatewayBorrow' -count=1 + go test -p=1 -tags unit ./internal/service ./internal/repository ./internal/handler ./internal/handler/admin -run '^Test(CodexGatewayBorrow|Pelican|AccountObservation|ObservedAccountLease)' -count=1 go test -p=1 -tags integration ./internal/repository -run '^TestCodexGatewayBorrowTestRepositoryIntegrationLifecycle$' -count=1 - name: Upstream automation classifier tests run: | @@ -145,6 +145,7 @@ jobs: src/views/admin/__tests__/CodexGatewayBorrowView.spec.ts \ src/views/admin/__tests__/CodexGatewayBorrowStatusView.spec.ts \ src/views/admin/__tests__/CodexPelicanComparisonView.spec.ts \ + src/api/admin/__tests__/pelicanTests.spec.ts \ src/components/admin/__tests__/CodexBorrowNav.spec.ts \ src/router/__tests__/codex-borrow-routes.spec.ts \ src/components/account/__tests__/EditAccountModal.spec.ts \ diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index bfec3a57bc70..abfb409a7db2 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -310,14 +310,16 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { return nil, err } codexGatewayBorrowTestRepository := repository.NewCodexGatewayBorrowTestRepository(db) - codexGatewayBorrowTestRunner := service.NewCodexGatewayBorrowTestRunner(codexGatewayBorrowTestRepository, codexGatewayBorrowService, accountRepository) + codexGatewayBorrowTestRunner := service.ProvidePelicanTestRunner(codexGatewayBorrowTestRepository, codexGatewayBorrowService, accountRepository, accountTestService, gatewayService, geminiMessagesCompatService, concurrencyService) codexGatewayBorrowPreviewService := service.NewCodexGatewayBorrowPreviewService(codexGatewayBorrowTestRepository) codexGatewayBorrowHandler := admin.NewCodexGatewayBorrowHandler(codexGatewayBorrowService, codexGatewayBorrowTestRunner, codexGatewayBorrowPreviewService) + pelicanTestCatalog := service.NewPelicanTestCatalog(accountRepository) + pelicanTestHandler := admin.NewPelicanTestHandler(pelicanTestCatalog, codexGatewayBorrowTestRunner, codexGatewayBorrowPreviewService) upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db) openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db) idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig) claudeResetCreditService := service.ProvideClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService, idempotencyCoordinator, leaderLockCache) - adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, accountJobHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, systemPromptHandler, reasoningRecoveryHandler, codexGatewayBorrowHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService) + adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, accountJobHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, systemPromptHandler, reasoningRecoveryHandler, codexGatewayBorrowHandler, pelicanTestHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService) usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig) userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient) userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig) diff --git a/backend/internal/handler/admin/codex_gateway_borrow_handler_test.go b/backend/internal/handler/admin/codex_gateway_borrow_handler_test.go index 7d1380a29527..d48b7e6181fd 100644 --- a/backend/internal/handler/admin/codex_gateway_borrow_handler_test.go +++ b/backend/internal/handler/admin/codex_gateway_borrow_handler_test.go @@ -84,6 +84,7 @@ func TestCodexGatewayBorrowHandlerPreviewUsesOnlyCapabilityAndOverridesCSP(t *te r := gin.New() r.Use(middleware.SecurityHeaders(config.CSPConfig{Enabled: true}, nil)) r.GET("/api/v1/codex-gateway-borrow/preview/:cap/index.html", h.ServePreview) + r.GET("/api/v1/pelican-tests/preview/:cap/index.html", h.ServePreview) r.GET("/", func(c *gin.Context) { c.String(http.StatusOK, "main app") }) w := httptest.NewRecorder() request := httptest.NewRequest(http.MethodGet, decorated.PreviewURL, nil) diff --git a/backend/internal/handler/admin/pelican_test_handler.go b/backend/internal/handler/admin/pelican_test_handler.go new file mode 100644 index 000000000000..71b5611d3daf --- /dev/null +++ b/backend/internal/handler/admin/pelican_test_handler.go @@ -0,0 +1,218 @@ +package admin + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/response" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +// PelicanTestHandler owns the independent administrator entry point. Borrowing +// configuration and preparation remain in CodexGatewayBorrowHandler; this entry +// selects only an account, a model, an effort, and a generation budget. +type PelicanTestHandler struct { + catalog *service.PelicanTestCatalog + runner *service.CodexGatewayBorrowTestRunner + preview *service.CodexGatewayBorrowPreviewService +} + +func NewPelicanTestHandler(catalog *service.PelicanTestCatalog, runner *service.CodexGatewayBorrowTestRunner, preview *service.CodexGatewayBorrowPreviewService) *PelicanTestHandler { + return &PelicanTestHandler{catalog: catalog, runner: runner, preview: preview} +} + +func (h *PelicanTestHandler) GetOptions(c *gin.Context) { + h.options(c, nil) +} + +func (h *PelicanTestHandler) GetAccountOptions(c *gin.Context) { + codexGatewayBorrowPrivateHeaders(c) + var request struct { + AccountIDs []int64 `json:"account_ids"` + } + if err := c.ShouldBindJSON(&request); err != nil { + response.BadRequest(c, err.Error()) + return + } + h.options(c, request.AccountIDs) +} + +func (h *PelicanTestHandler) options(c *gin.Context, accountIDs []int64) { + codexGatewayBorrowPrivateHeaders(c) + options, err := h.catalog.Options(c.Request.Context(), accountIDs) + if err != nil { + if errors.Is(err, service.ErrPelicanTestOptionsInvalidRequest) { + response.BadRequest(c, err.Error()) + } else { + response.Error(c, http.StatusInternalServerError, err.Error()) + } + return + } + response.Success(c, options) +} + +func (h *PelicanTestHandler) StartTests(c *gin.Context) { + codexGatewayBorrowPrivateHeaders(c) + var request service.CodexGatewayBorrowTestRequest + if err := c.ShouldBindJSON(&request); err != nil { + response.BadRequest(c, err.Error()) + return + } + if _, err := uuid.Parse(strings.TrimSpace(request.ClientTaskID)); err != nil { + response.BadRequest(c, "client_task_id must be a UUID") + return + } + if len(request.Targets) == 0 { + response.BadRequest(c, "targets is required") + return + } + // Only this server-side entry point can select account execution. A JSON + // field cannot change the legacy borrowing handler's existing contract. + request.Standalone = true + var createdBy int64 + if subject, ok := middleware.GetAuthSubjectFromContext(c); ok { + createdBy = subject.UserID + } + task, replayed, err := h.runner.Create(c.Request.Context(), createdBy, request) + if err != nil { + codexGatewayBorrowTestError(c, err) + return + } + middleware.SetAuditExtra(c, map[string]any{"result": "pelican_manual_test", "task_id": task.ID, "replayed": replayed, "targets": task.Total}) + ctx, cancel := context.WithCancel(c.Request.Context()) + defer cancel() + c.Header("Content-Type", "text/event-stream") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Status(http.StatusOK) + var writeMu sync.Mutex + emit := func(event service.CodexGatewayBorrowTestEvent) { + if ctx.Err() != nil { + return + } + if event.Task != nil { + event.Task = h.preview.DecorateTask(event.Task) + } + if event.Result != nil { + event.Result = h.preview.DecorateResult(event.Result) + } + payload, marshalErr := json.Marshal(event) + if marshalErr != nil { + cancel() + return + } + writeMu.Lock() + defer writeMu.Unlock() + if ctx.Err() != nil { + return + } + if _, writeErr := fmt.Fprintf(c.Writer, "data: %s\n\n", payload); writeErr != nil { + cancel() + return + } + c.Writer.Flush() + } + emit(service.CodexGatewayBorrowTestEvent{Type: "task_start", Task: task}) + if replayed { + emit(service.CodexGatewayBorrowTestEvent{Type: "task_complete", Task: task}) + return + } + keepAliveDone, keepAliveStopped := make(chan struct{}), make(chan struct{}) + go func() { + defer close(keepAliveStopped) + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + for { + select { + case <-keepAliveDone: + return + case <-ctx.Done(): + return + case <-ticker.C: + writeMu.Lock() + if ctx.Err() == nil { + if _, writeErr := fmt.Fprint(c.Writer, ": keep-alive\n\n"); writeErr != nil { + cancel() + } else { + c.Writer.Flush() + } + } + writeMu.Unlock() + } + } + }() + runErr := h.runner.Run(ctx, task, emit) + close(keepAliveDone) + <-keepAliveStopped + if runErr != nil && ctx.Err() == nil { + task.Status, task.Error = "incomplete", runErr.Error() + emit(service.CodexGatewayBorrowTestEvent{Type: "task_complete", Task: task}) + } +} + +func (h *PelicanTestHandler) ListTests(c *gin.Context) { + codexGatewayBorrowPrivateHeaders(c) + page, size := response.ParsePagination(c) + if rawSize := c.Query("size"); rawSize != "" { + if parsed, err := strconv.Atoi(rawSize); err == nil && parsed > 0 { + size = parsed + } + } + list, err := h.runner.List(c.Request.Context(), page, size) + if err != nil { + codexGatewayBorrowTestError(c, err) + return + } + response.Success(c, list) +} + +func (h *PelicanTestHandler) GetTest(c *gin.Context) { + codexGatewayBorrowPrivateHeaders(c) + id := strings.TrimSpace(c.Param("id")) + if id == "last" { + list, err := h.runner.List(c.Request.Context(), 1, 1) + if err != nil { + codexGatewayBorrowTestError(c, err) + return + } + if len(list.Items) == 0 { + response.NotFound(c, "no unexpired pelican test") + return + } + id = list.Items[0].ID + } + parsed, err := uuid.Parse(id) + if err != nil { + response.BadRequest(c, "test ID must be a UUID or last") + return + } + task, err := h.runner.Get(c.Request.Context(), parsed.String()) + if err != nil { + codexGatewayBorrowTestError(c, err) + return + } + response.Success(c, h.preview.DecorateTask(task)) +} + +func (h *PelicanTestHandler) ServePreview(c *gin.Context) { + codexGatewayBorrowPrivateHeaders(c) + c.Header("Content-Security-Policy", service.CodexGatewayBorrowPreviewCSP) + c.Header("X-Frame-Options", "SAMEORIGIN") + c.Header("Cross-Origin-Resource-Policy", "cross-origin") + html, err := h.preview.Resolve(c.Request.Context(), strings.TrimSpace(c.Param("cap"))) + if err != nil { + c.Status(http.StatusGone) + return + } + c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(html)) +} diff --git a/backend/internal/handler/admin/pelican_test_handler_test.go b/backend/internal/handler/admin/pelican_test_handler_test.go new file mode 100644 index 000000000000..b3c5f7abfc5f --- /dev/null +++ b/backend/internal/handler/admin/pelican_test_handler_test.go @@ -0,0 +1,160 @@ +//go:build unit + +package admin + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type pelicanTestHandlerLocalAccountReader struct { + service.AccountRepository + reads int +} + +type pelicanTestHandlerReplayCaptureRepo struct { + *codexGatewayBorrowHandlerTestRepo + incoming *service.CodexGatewayBorrowTestTask +} + +func (r *pelicanTestHandlerReplayCaptureRepo) Create(ctx context.Context, task *service.CodexGatewayBorrowTestTask) (*service.CodexGatewayBorrowTestTask, bool, error) { + r.incoming = task + return r.codexGatewayBorrowHandlerTestRepo.Create(ctx, task) +} + +func (r *pelicanTestHandlerLocalAccountReader) GetByIDs(_ context.Context, _ []int64) ([]*service.Account, error) { + r.reads++ + return []*service.Account{{ID: 42, Name: "paused account", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, Status: service.StatusDisabled}}, nil +} + +func TestPelicanTestHandlerOptionsDefaultsAndSelectedLocalModels(t *testing.T) { + gin.SetMode(gin.TestMode) + accounts := &pelicanTestHandlerLocalAccountReader{} + h := NewPelicanTestHandler(service.NewPelicanTestCatalog(accounts), nil, nil) + router := gin.New() + router.GET("/options", h.GetOptions) + router.POST("/options", h.GetAccountOptions) + request := func(method, body string) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + r := httptest.NewRequest(method, "/options", strings.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, r) + return w + } + defaults := request(http.MethodGet, "") + require.Equal(t, http.StatusOK, defaults.Code) + require.Zero(t, accounts.reads) + require.Contains(t, defaults.Body.String(), `"default_generation_timeout_seconds":600`) + require.Equal(t, "private, no-store", defaults.Header().Get("Cache-Control")) + selected := request(http.MethodPost, `{"account_ids":[42,999]}`) + require.Equal(t, http.StatusOK, selected.Code) + var envelope struct { + Data service.PelicanTestOptions `json:"data"` + } + require.NoError(t, json.Unmarshal(selected.Body.Bytes(), &envelope)) + require.Equal(t, 1, accounts.reads) + require.Len(t, envelope.Data.Accounts, 2) + require.Equal(t, service.StatusDisabled, envelope.Data.Accounts[0].Status) + require.Equal(t, "missing", envelope.Data.Accounts[1].Status) + ids := make([]string, 0, len(envelope.Data.Accounts[0].Models)) + for _, model := range envelope.Data.Accounts[0].Models { + ids = append(ids, model.ID) + } + require.Contains(t, ids, "gpt-6-astra") + require.Contains(t, ids, "gpt-6.1-sol") + invalid := request(http.MethodPost, `{"account_ids":[0]}`) + require.Equal(t, http.StatusBadRequest, invalid.Code) + require.Equal(t, 1, accounts.reads) +} + +func TestPelicanTestHandlerSSEReplayAndHistoryNeverGenerateAndPreserveRawContent(t *testing.T) { + gin.SetMode(gin.TestMode) + legacy, repo := newCodexGatewayBorrowHandlerFixture(t) + h := NewPelicanTestHandler(service.NewPelicanTestCatalog(nil), legacy.runner, legacy.preview) + router := gin.New() + router.GET("/tests/:id", h.GetTest) + router.GET("/tests", h.ListTests) + router.POST("/tests", h.StartTests) + router.GET("/api/v1/pelican-tests/preview/:cap/index.html", h.ServePreview) + history := httptest.NewRecorder() + router.ServeHTTP(history, httptest.NewRequest(http.MethodGet, "/tests/last", nil)) + require.Equal(t, http.StatusOK, history.Code) + var envelope struct { + Data service.CodexGatewayBorrowTestTask `json:"data"` + } + require.NoError(t, json.Unmarshal(history.Body.Bytes(), &envelope)) + require.Equal(t, repo.result.Error, envelope.Data.Results[0].Error) + require.Equal(t, repo.result.RawAnswer, envelope.Data.Results[0].RawAnswer) + require.True(t, strings.HasPrefix(envelope.Data.Results[0].PreviewURL, "/api/v1/pelican-tests/preview/")) + preview := httptest.NewRecorder() + router.ServeHTTP(preview, httptest.NewRequest(http.MethodGet, envelope.Data.Results[0].PreviewURL, nil)) + require.Equal(t, http.StatusOK, preview.Code) + require.Equal(t, repo.result.HTML, preview.Body.String()) + require.Equal(t, service.CodexGatewayBorrowPreviewCSP, preview.Header().Get("Content-Security-Policy")) + stream := httptest.NewRecorder() + body := fmt.Sprintf(`{"client_task_id":%q,"generation_timeout_seconds":600,"targets":[{"account_id":42,"model_id":"gpt-6-astra"}]}`, repo.task.ClientTaskID) + request := httptest.NewRequest(http.MethodPost, "/tests", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(stream, request) + require.Equal(t, http.StatusOK, stream.Code) + require.Equal(t, "text/event-stream", stream.Header().Get("Content-Type")) + require.Contains(t, stream.Body.String(), `"type":"task_start"`) + require.Contains(t, stream.Body.String(), `"type":"task_complete"`) + require.NotContains(t, stream.Body.String(), `"type":"result_started"`) + require.Contains(t, stream.Body.String(), "original upstream error token=admin-visible") + require.Equal(t, 1, repo.creates) + require.Zero(t, repo.writes) +} + +func TestPelicanTestHandlerForcesAccountModeAndFixedPromptBeforeReplay(t *testing.T) { + gin.SetMode(gin.TestMode) + _, existing := newCodexGatewayBorrowHandlerFixture(t) + repo := &pelicanTestHandlerReplayCaptureRepo{codexGatewayBorrowHandlerTestRepo: existing} + runner := service.NewCodexGatewayBorrowTestRunner(repo, nil, nil) + preview := service.NewCodexGatewayBorrowPreviewService(repo) + t.Cleanup(runner.Stop) + t.Cleanup(preview.Stop) + h := NewPelicanTestHandler(service.NewPelicanTestCatalog(nil), runner, preview) + router := gin.New() + router.POST("/tests", h.StartTests) + body := fmt.Sprintf(`{"client_task_id":%q,"standalone":false,"execution_mode":"legacy_cache","prompt":"client modified prompt","targets":[{"account_id":42,"model_id":"gpt-6-astra"}]}`, existing.task.ClientTaskID) + request := httptest.NewRequest(http.MethodPost, "/tests", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + router.ServeHTTP(w, request) + require.Equal(t, http.StatusOK, w.Code) + require.NotNil(t, repo.incoming) + require.Equal(t, service.PelicanExecutionModeAccount, repo.incoming.ExecutionMode) + require.Equal(t, service.CodexGatewayBorrowPelicanPrompt, repo.incoming.Prompt) + require.Equal(t, 600, repo.incoming.GenerationTimeoutSeconds) + require.Equal(t, int64(42), repo.incoming.Results[0].AccountID) + require.Empty(t, repo.incoming.Results[0].Effort, "the generator applies only the selected model's supported default") + require.Zero(t, existing.writes, "replay must not dispatch the forced-account execution") +} + +func TestPelicanTestHandlerRejectsGenerationBudgetBeforeCreate(t *testing.T) { + gin.SetMode(gin.TestMode) + legacy, repo := newCodexGatewayBorrowHandlerFixture(t) + h := NewPelicanTestHandler(service.NewPelicanTestCatalog(nil), legacy.runner, legacy.preview) + router := gin.New() + router.POST("/tests", h.StartTests) + for _, seconds := range []int{59, 1801} { + w := httptest.NewRecorder() + body := fmt.Sprintf(`{"client_task_id":%q,"generation_timeout_seconds":%d,"targets":[{"account_id":42,"model_id":"gpt-6-astra"}]}`, repo.task.ClientTaskID, seconds) + r := httptest.NewRequest(http.MethodPost, "/tests", strings.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, r) + require.Equal(t, http.StatusBadRequest, w.Code) + } + require.Zero(t, repo.creates, "invalid budget must not leave a pending task") + require.Zero(t, repo.writes) +} diff --git a/backend/internal/handler/handler.go b/backend/internal/handler/handler.go index 3b493f1c7450..6d01386b2010 100644 --- a/backend/internal/handler/handler.go +++ b/backend/internal/handler/handler.go @@ -47,6 +47,7 @@ type AdminHandlers struct { SystemPrompt *admin.SystemPromptHandler ReasoningRecovery *admin.ReasoningRecoveryHandler CodexGatewayBorrow *admin.CodexGatewayBorrowHandler + PelicanTest *admin.PelicanTestHandler } // Handlers contains all HTTP handlers diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 90a4e357ad52..2f027c6003d9 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -51,6 +51,7 @@ func ProvideAdminHandlers( systemPromptHandler *admin.SystemPromptHandler, reasoningRecoveryHandler *admin.ReasoningRecoveryHandler, codexGatewayBorrowHandler *admin.CodexGatewayBorrowHandler, + pelicanTestHandler *admin.PelicanTestHandler, upstreamBillingProbe *service.UpstreamBillingProbeService, ollamaCloudUsage *service.OllamaCloudUsageService, opencodeGoUsage *service.OpenCodeGoUsageService, @@ -101,6 +102,7 @@ func ProvideAdminHandlers( SystemPrompt: systemPromptHandler, ReasoningRecovery: reasoningRecoveryHandler, CodexGatewayBorrow: codexGatewayBorrowHandler, + PelicanTest: pelicanTestHandler, } } @@ -302,6 +304,7 @@ var ProviderSet = wire.NewSet( admin.NewSystemPromptHandler, admin.NewReasoningRecoveryHandler, admin.NewCodexGatewayBorrowHandler, + admin.NewPelicanTestHandler, ProvideAccountJobRuntime, // AdminHandlers and Handlers constructors diff --git a/backend/internal/repository/codex_gateway_borrow_test_repo.go b/backend/internal/repository/codex_gateway_borrow_test_repo.go index ddf8f78a6fe1..b154fcbcd6a3 100644 --- a/backend/internal/repository/codex_gateway_borrow_test_repo.go +++ b/backend/internal/repository/codex_gateway_borrow_test_repo.go @@ -16,11 +16,12 @@ func NewCodexGatewayBorrowTestRepository(db *sql.DB) service.CodexGatewayBorrowT } const codexGatewayBorrowTestTaskColumns = `id, client_task_id, created_by, request_hash, status, prompt, - created_at, expires_at, started_at, finished_at, total, completed, error` + created_at, expires_at, started_at, finished_at, total, completed, error, generation_timeout_seconds, execution_mode` const codexGatewayBorrowTestResultColumns = `r.id, r.task_id, r.ordinal, r.account_id, r.account_name, r.model_id, r.upstream_model, r.effort, r.status, r.raw_answer, r.raw_response, r.raw_html, r.html, r.error, - r.started_at, r.finished_at, r.duration_ms, t.expires_at` + r.started_at, r.finished_at, r.duration_ms, t.expires_at, r.platform, r.actual_endpoint, r.actual_protocol, + r.actual_transport, r.borrow_applied, r.queue_duration_ms, r.preparation_duration_ms, r.generation_duration_ms, r.generation_started_at` type codexGatewayBorrowTestScanner interface{ Scan(...any) error } @@ -29,7 +30,7 @@ func scanCodexGatewayBorrowTestTask(row codexGatewayBorrowTestScanner) (*service var started, finished sql.NullTime var rawError []byte if err := row.Scan(&task.ID, &task.ClientTaskID, &task.CreatedBy, &task.RequestHash, &task.Status, &task.Prompt, - &task.CreatedAt, &task.ExpiresAt, &started, &finished, &task.Total, &task.Completed, &rawError); err != nil { + &task.CreatedAt, &task.ExpiresAt, &started, &finished, &task.Total, &task.Completed, &rawError, &task.GenerationTimeoutSeconds, &task.ExecutionMode); err != nil { return nil, err } if started.Valid { @@ -45,10 +46,11 @@ func scanCodexGatewayBorrowTestTask(row codexGatewayBorrowTestScanner) (*service func scanCodexGatewayBorrowTestResult(row codexGatewayBorrowTestScanner) (*service.CodexGatewayBorrowTestResult, error) { result := &service.CodexGatewayBorrowTestResult{} var answer, rawResponse, rawHTML, html, rawError []byte - var started, finished sql.NullTime + var started, finished, generationStarted sql.NullTime if err := row.Scan(&result.ID, &result.TaskID, &result.Ordinal, &result.AccountID, &result.AccountName, &result.ModelID, &result.UpstreamModel, &result.Effort, &result.Status, &answer, &rawResponse, &rawHTML, &html, &rawError, - &started, &finished, &result.DurationMS, &result.ExpiresAt); err != nil { + &started, &finished, &result.DurationMS, &result.ExpiresAt, &result.Platform, &result.ActualEndpoint, &result.ActualProtocol, + &result.ActualTransport, &result.BorrowApplied, &result.QueueDurationMS, &result.PreparationDurationMS, &result.GenerationDurationMS, &generationStarted); err != nil { return nil, err } result.RawAnswer, result.RawResponse, result.RawHTML, result.HTML, result.Error = string(answer), string(rawResponse), string(rawHTML), string(html), string(rawError) @@ -58,6 +60,9 @@ func scanCodexGatewayBorrowTestResult(row codexGatewayBorrowTestScanner) (*servi if finished.Valid { result.FinishedAt = &finished.Time } + if generationStarted.Valid { + result.GenerationStartedAt = &generationStarted.Time + } if result.Status != "complete" { result.PreviewUnavailable = "test " + result.Status } else if result.HTML == "" { @@ -75,10 +80,10 @@ func (r *codexGatewayBorrowTestRepository) Create(ctx context.Context, task *ser } defer func() { _ = tx.Rollback() }() stored, err := scanCodexGatewayBorrowTestTask(tx.QueryRowContext(ctx, `INSERT INTO codex_gateway_borrow_test_tasks - (id,client_task_id,created_by,request_hash,status,prompt,total,created_at,expires_at) - VALUES ($1,$2,$3,$4,'pending',$5,$6,$7,$8) + (id,client_task_id,created_by,request_hash,status,prompt,total,created_at,expires_at,generation_timeout_seconds,execution_mode) + VALUES ($1,$2,$3,$4,'pending',$5,$6,$7,$8,$9,$10) ON CONFLICT (client_task_id) DO NOTHING RETURNING `+codexGatewayBorrowTestTaskColumns, - task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Prompt, task.Total, task.CreatedAt, task.ExpiresAt)) + task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Prompt, task.Total, task.CreatedAt, task.ExpiresAt, task.GenerationTimeoutSeconds, task.ExecutionMode)) if errors.Is(err, sql.ErrNoRows) { stored, err = scanCodexGatewayBorrowTestTask(tx.QueryRowContext(ctx, `SELECT `+codexGatewayBorrowTestTaskColumns+` FROM codex_gateway_borrow_test_tasks WHERE scope='manual_pelican' AND client_task_id=$1 FOR UPDATE`, task.ClientTaskID)) @@ -102,8 +107,8 @@ func (r *codexGatewayBorrowTestRepository) Create(ctx context.Context, task *ser } for _, result := range task.Results { if _, err = tx.ExecContext(ctx, `INSERT INTO codex_gateway_borrow_test_results - (id,task_id,ordinal,account_id,account_name,model_id,effort,status) - VALUES ($1,$2,$3,$4,$5,$6,$7,'pending')`, result.ID, stored.ID, result.Ordinal, result.AccountID, result.AccountName, result.ModelID, result.Effort); err != nil { + (id,task_id,ordinal,account_id,account_name,model_id,effort,status,platform) + VALUES ($1,$2,$3,$4,$5,$6,$7,'pending',$8)`, result.ID, stored.ID, result.Ordinal, result.AccountID, result.AccountName, result.ModelID, result.Effort, result.Platform); err != nil { return nil, false, err } } @@ -221,12 +226,15 @@ func (r *codexGatewayBorrowTestRepository) SaveResult(ctx context.Context, resul defer func() { _ = tx.Rollback() }() if err = codexGatewayBorrowRequireAffected(tx.ExecContext(ctx, `UPDATE codex_gateway_borrow_test_results r SET account_name=$3,upstream_model=$4,effort=$5,status=$6,raw_answer=$7,raw_response=$8,raw_html=$9,html=$10,error=$11, - started_at=$12,finished_at=$13,duration_ms=$14,updated_at=CURRENT_TIMESTAMP + started_at=$12,finished_at=$13,duration_ms=$14,platform=$15,actual_endpoint=$16,actual_protocol=$17, + actual_transport=$18,borrow_applied=$19,queue_duration_ms=$20,preparation_duration_ms=$21, + generation_duration_ms=$22,generation_started_at=$23,updated_at=CURRENT_TIMESTAMP FROM codex_gateway_borrow_test_tasks t WHERE r.id=$1 AND r.task_id=$2 AND t.id=r.task_id AND t.scope='manual_pelican' AND t.expires_at>CURRENT_TIMESTAMP`, result.ID, result.TaskID, result.AccountName, result.UpstreamModel, result.Effort, result.Status, []byte(result.RawAnswer), []byte(result.RawResponse), []byte(result.RawHTML), []byte(result.HTML), []byte(result.Error), - result.StartedAt, result.FinishedAt, result.DurationMS)); err != nil { + result.StartedAt, result.FinishedAt, result.DurationMS, result.Platform, result.ActualEndpoint, result.ActualProtocol, + result.ActualTransport, result.BorrowApplied, result.QueueDurationMS, result.PreparationDurationMS, result.GenerationDurationMS, result.GenerationStartedAt)); err != nil { return err } if _, err = tx.ExecContext(ctx, `UPDATE codex_gateway_borrow_test_tasks t diff --git a/backend/internal/repository/codex_gateway_borrow_test_repo_integration_test.go b/backend/internal/repository/codex_gateway_borrow_test_repo_integration_test.go index ff93f6f8fbbd..e277e70c1353 100644 --- a/backend/internal/repository/codex_gateway_borrow_test_repo_integration_test.go +++ b/backend/internal/repository/codex_gateway_borrow_test_repo_integration_test.go @@ -63,7 +63,42 @@ func codexGatewayBorrowPrivateIntegrationDB(t *testing.T, ctx context.Context) * ddl, err := migrations.FS.ReadFile("275_codex_gateway_borrow_tests.sql") require.NoError(t, err) _, err = privateDB.ExecContext(ctx, string(ddl)) - require.NoError(t, err, "apply the actual new migration inside the private fixture") + require.NoError(t, err, "apply the original observation schema inside the private fixture") + // Put one legacy record in the original schema before applying the new + // migration. Its existing output and retention must survive unchanged. + legacyTask, legacyResult := uuid.NewString(), uuid.NewString() + legacyClient, err := uuid.NewV7() + require.NoError(t, err) + created := time.Now().UTC().Truncate(time.Microsecond) + legacyAnswer := []byte("legacy original answer\x00tail") + _, err = privateDB.ExecContext(ctx, `INSERT INTO codex_gateway_borrow_test_tasks + (id,client_task_id,created_by,request_hash,status,prompt,total,created_at,expires_at) + VALUES ($1,$2,7,'legacy-hash','complete',$3,1,$4,$5)`, legacyTask, legacyClient.String(), service.CodexGatewayBorrowPelicanPrompt, created, created.Add(service.CodexGatewayBorrowTestTTL)) + require.NoError(t, err) + _, err = privateDB.ExecContext(ctx, `INSERT INTO codex_gateway_borrow_test_results + (id,task_id,ordinal,account_id,model_id,effort,status,raw_answer,error) + VALUES ($1,$2,1,42,'gpt-6-astra','high','complete',$3,$3)`, legacyResult, legacyTask, legacyAnswer) + require.NoError(t, err) + ddl, err = migrations.FS.ReadFile("276_pelican_test_execution.sql") + require.NoError(t, err) + _, err = privateDB.ExecContext(ctx, string(ddl)) + require.NoError(t, err, "apply the actual additive execution migration") + var answer, originalError []byte + var executionMode string + var timeout int + var retention int64 + err = privateDB.QueryRowContext(ctx, `SELECT r.raw_answer,r.error,t.execution_mode,t.generation_timeout_seconds, + EXTRACT(EPOCH FROM (t.expires_at-t.created_at))::BIGINT + FROM codex_gateway_borrow_test_results r JOIN codex_gateway_borrow_test_tasks t ON t.id=r.task_id + WHERE r.id=$1`, legacyResult).Scan(&answer, &originalError, &executionMode, &timeout, &retention) + require.NoError(t, err) + require.Equal(t, legacyAnswer, answer) + require.Equal(t, legacyAnswer, originalError) + require.Equal(t, service.PelicanExecutionModeLegacyCache, executionMode) + require.Equal(t, 90, timeout, "historical observations retain their actual legacy budget") + require.Equal(t, int64(service.CodexGatewayBorrowTestTTL/time.Second), retention) + _, err = privateDB.ExecContext(ctx, `DELETE FROM codex_gateway_borrow_test_tasks WHERE id=$1`, legacyTask) + require.NoError(t, err) return privateDB } @@ -76,10 +111,11 @@ func codexGatewayBorrowIntegrationTask(t *testing.T, accountID int64) *service.C ID: uuid.NewString(), ClientTaskID: clientID.String(), CreatedBy: 7, RequestHash: strings.Repeat("a", 64), Status: "pending", Prompt: service.CodexGatewayBorrowPelicanPrompt, CreatedAt: created, ExpiresAt: created.Add(service.CodexGatewayBorrowTestTTL), Total: 1, + GenerationTimeoutSeconds: service.PelicanDefaultGenerationTimeoutSeconds, ExecutionMode: service.PelicanExecutionModeAccount, } task.Results = []*service.CodexGatewayBorrowTestResult{{ ID: uuid.NewString(), TaskID: task.ID, Ordinal: 1, AccountID: accountID, - AccountName: "private integration fixture", ModelID: "gpt-6-astra", Effort: "high", + AccountName: "private integration fixture", Platform: "openai", ModelID: "gpt-6-astra", Effort: "high", Status: "pending", ExpiresAt: task.ExpiresAt, }} return task @@ -100,6 +136,8 @@ func TestCodexGatewayBorrowTestRepositoryIntegrationLifecycle(t *testing.T) { require.Equal(t, task.ID, stored.ID) require.Equal(t, "pending", stored.Status) require.Len(t, stored.Results, 1) + require.Equal(t, service.PelicanDefaultGenerationTimeoutSeconds, stored.GenerationTimeoutSeconds) + require.Equal(t, service.PelicanExecutionModeAccount, stored.ExecutionMode) require.Equal(t, service.CodexGatewayBorrowTestTTL, stored.ExpiresAt.Sub(stored.CreatedAt)) started := time.Now().UTC().Truncate(time.Microsecond) @@ -108,6 +146,10 @@ func TestCodexGatewayBorrowTestRepositoryIntegrationLifecycle(t *testing.T) { finished := started.Add(250 * time.Millisecond) result.Status, result.UpstreamModel = "complete", "gpt-6-astra-reported" result.StartedAt, result.FinishedAt, result.DurationMS = &started, &finished, 250 + generationStarted := started.Add(50 * time.Millisecond) + result.GenerationStartedAt = &generationStarted + result.QueueDurationMS, result.PreparationDurationMS, result.GenerationDurationMS = 20, 30, 200 + result.ActualEndpoint, result.ActualProtocol, result.ActualTransport, result.BorrowApplied = "https://local.test/responses", "responses", "http", true result.RawAnswer = "```svg\n鹈鹕\x00tail\n```" result.RawResponse = "data: {\"original\":\"full stream\"}\n\n\x00end" result.RawHTML = "鹈鹕\x00tail" @@ -127,6 +169,15 @@ func TestCodexGatewayBorrowTestRepositoryIntegrationLifecycle(t *testing.T) { require.Equal(t, "gpt-6-astra", readResult.ModelID) require.Equal(t, result.UpstreamModel, readResult.UpstreamModel) require.Equal(t, result.DurationMS, readResult.DurationMS) + require.Equal(t, result.Platform, readResult.Platform) + require.Equal(t, result.ActualEndpoint, readResult.ActualEndpoint) + require.Equal(t, result.ActualProtocol, readResult.ActualProtocol) + require.Equal(t, result.ActualTransport, readResult.ActualTransport) + require.True(t, readResult.BorrowApplied) + require.Equal(t, result.QueueDurationMS, readResult.QueueDurationMS) + require.Equal(t, result.PreparationDurationMS, readResult.PreparationDurationMS) + require.Equal(t, result.GenerationDurationMS, readResult.GenerationDurationMS) + require.True(t, readResult.GenerationStartedAt.Equal(generationStarted)) require.True(t, readResult.StartedAt.Equal(started)) require.True(t, readResult.FinishedAt.Equal(finished)) require.True(t, readResult.ExpiresAt.Equal(task.ExpiresAt)) diff --git a/backend/internal/repository/codex_gateway_borrow_test_repo_test.go b/backend/internal/repository/codex_gateway_borrow_test_repo_test.go index 32d3a5b1fc9e..f06f0d55a0cd 100644 --- a/backend/internal/repository/codex_gateway_borrow_test_repo_test.go +++ b/backend/internal/repository/codex_gateway_borrow_test_repo_test.go @@ -24,13 +24,14 @@ func newCodexGatewayBorrowRepoTest(t *testing.T) (*codexGatewayBorrowTestReposit } func codexGatewayBorrowTaskTestRows(task *service.CodexGatewayBorrowTestTask) *sqlmock.Rows { - return sqlmock.NewRows([]string{"id", "client_task_id", "created_by", "request_hash", "status", "prompt", "created_at", "expires_at", "started_at", "finished_at", "total", "completed", "error"}).AddRow( - task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Status, task.Prompt, task.CreatedAt, task.ExpiresAt, nil, nil, task.Total, task.Completed, []byte(task.Error)) + return sqlmock.NewRows([]string{"id", "client_task_id", "created_by", "request_hash", "status", "prompt", "created_at", "expires_at", "started_at", "finished_at", "total", "completed", "error", "generation_timeout_seconds", "execution_mode"}).AddRow( + task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Status, task.Prompt, task.CreatedAt, task.ExpiresAt, nil, nil, task.Total, task.Completed, []byte(task.Error), task.GenerationTimeoutSeconds, task.ExecutionMode) } func codexGatewayBorrowResultTestRows(result *service.CodexGatewayBorrowTestResult) *sqlmock.Rows { - return sqlmock.NewRows([]string{"id", "task_id", "ordinal", "account_id", "account_name", "model_id", "upstream_model", "effort", "status", "raw_answer", "raw_response", "raw_html", "html", "error", "started_at", "finished_at", "duration_ms", "expires_at"}).AddRow( - result.ID, result.TaskID, result.Ordinal, result.AccountID, result.AccountName, result.ModelID, result.UpstreamModel, result.Effort, result.Status, []byte(result.RawAnswer), []byte(result.RawResponse), []byte(result.RawHTML), []byte(result.HTML), []byte(result.Error), nil, nil, result.DurationMS, result.ExpiresAt) + return sqlmock.NewRows([]string{"id", "task_id", "ordinal", "account_id", "account_name", "model_id", "upstream_model", "effort", "status", "raw_answer", "raw_response", "raw_html", "html", "error", "started_at", "finished_at", "duration_ms", "expires_at", "platform", "actual_endpoint", "actual_protocol", "actual_transport", "borrow_applied", "queue_duration_ms", "preparation_duration_ms", "generation_duration_ms", "generation_started_at"}).AddRow( + result.ID, result.TaskID, result.Ordinal, result.AccountID, result.AccountName, result.ModelID, result.UpstreamModel, result.Effort, result.Status, []byte(result.RawAnswer), []byte(result.RawResponse), []byte(result.RawHTML), []byte(result.HTML), []byte(result.Error), result.StartedAt, result.FinishedAt, result.DurationMS, result.ExpiresAt, + result.Platform, result.ActualEndpoint, result.ActualProtocol, result.ActualTransport, result.BorrowApplied, result.QueueDurationMS, result.PreparationDurationMS, result.GenerationDurationMS, result.GenerationStartedAt) } func codexGatewayBorrowRepoFixture(t *testing.T) *service.CodexGatewayBorrowTestTask { @@ -38,8 +39,9 @@ func codexGatewayBorrowRepoFixture(t *testing.T) *service.CodexGatewayBorrowTest now := time.Now().UTC() clientID, err := uuid.NewV7() require.NoError(t, err) - task := &service.CodexGatewayBorrowTestTask{ID: uuid.NewString(), ClientTaskID: clientID.String(), CreatedBy: 7, RequestHash: "hash", Status: "pending", Prompt: service.CodexGatewayBorrowPelicanPrompt, CreatedAt: now, ExpiresAt: now.Add(service.CodexGatewayBorrowTestTTL), Total: 1} - task.Results = []*service.CodexGatewayBorrowTestResult{{ID: uuid.NewString(), TaskID: task.ID, Ordinal: 1, AccountID: 42, AccountName: "paused target", ModelID: "gpt-6-astra", Effort: "high", Status: "pending", ExpiresAt: task.ExpiresAt}} + task := &service.CodexGatewayBorrowTestTask{ID: uuid.NewString(), ClientTaskID: clientID.String(), CreatedBy: 7, RequestHash: "hash", Status: "pending", Prompt: service.CodexGatewayBorrowPelicanPrompt, CreatedAt: now, ExpiresAt: now.Add(service.CodexGatewayBorrowTestTTL), Total: 1, + GenerationTimeoutSeconds: service.PelicanDefaultGenerationTimeoutSeconds, ExecutionMode: service.PelicanExecutionModeAccount} + task.Results = []*service.CodexGatewayBorrowTestResult{{ID: uuid.NewString(), TaskID: task.ID, Ordinal: 1, AccountID: 42, AccountName: "paused target", Platform: "openai", ModelID: "gpt-6-astra", Effort: "high", Status: "pending", ExpiresAt: task.ExpiresAt}} return task } @@ -47,15 +49,17 @@ func TestCodexGatewayBorrowTestRepositoryCreateAtomicAndReplayReadsWithoutWrites repo, mock := newCodexGatewayBorrowRepoTest(t) task := codexGatewayBorrowRepoFixture(t) mock.ExpectBegin() - mock.ExpectQuery("INSERT INTO codex_gateway_borrow_test_tasks").WithArgs(task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Prompt, task.Total, task.CreatedAt, task.ExpiresAt).WillReturnRows(codexGatewayBorrowTaskTestRows(task)) + mock.ExpectQuery("INSERT INTO codex_gateway_borrow_test_tasks").WithArgs(task.ID, task.ClientTaskID, task.CreatedBy, task.RequestHash, task.Prompt, task.Total, task.CreatedAt, task.ExpiresAt, task.GenerationTimeoutSeconds, task.ExecutionMode).WillReturnRows(codexGatewayBorrowTaskTestRows(task)) result := task.Results[0] - mock.ExpectExec("INSERT INTO codex_gateway_borrow_test_results").WithArgs(result.ID, task.ID, 1, result.AccountID, result.AccountName, result.ModelID, result.Effort).WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec("INSERT INTO codex_gateway_borrow_test_results").WithArgs(result.ID, task.ID, 1, result.AccountID, result.AccountName, result.ModelID, result.Effort, result.Platform).WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() stored, replayed, err := repo.Create(context.Background(), task) require.NoError(t, err) require.False(t, replayed) require.Equal(t, task.ID, stored.ID) require.Len(t, stored.Results, 1) + require.Equal(t, task.GenerationTimeoutSeconds, stored.GenerationTimeoutSeconds) + require.Equal(t, task.ExecutionMode, stored.ExecutionMode) // A repeated client UUID gets the original full observation, not another child. task.Status = "complete" result.Status, result.RawAnswer, result.RawResponse, result.Error = "complete", "answer\x00tail", "full raw stream\x00end", "untruncated upstream diagnostic" @@ -71,6 +75,7 @@ func TestCodexGatewayBorrowTestRepositoryCreateAtomicAndReplayReadsWithoutWrites require.Equal(t, result.RawAnswer, stored.Results[0].RawAnswer) require.Equal(t, result.RawResponse, stored.Results[0].RawResponse) require.Equal(t, result.Error, stored.Results[0].Error) + require.Equal(t, result.Platform, stored.Results[0].Platform) require.NoError(t, mock.ExpectationsWereMet()) } @@ -121,6 +126,11 @@ func TestCodexGatewayBorrowTestRepositoryExpiredReadsAndScopedCleanup(t *testing require.Contains(t, string(data), "client_task_id UUID NOT NULL UNIQUE") require.Contains(t, string(data), "expires_at - created_at = INTERVAL '24 hours'") require.Contains(t, string(data), "raw_answer BYTEA NOT NULL") + data, err = migrations.FS.ReadFile("276_pelican_test_execution.sql") + require.NoError(t, err) + require.Contains(t, string(data), "generation_timeout_seconds BETWEEN 60 AND 1800") + require.Contains(t, string(data), "DEFAULT 'legacy_cache'") + require.NotContains(t, string(data), "DROP") } func TestCodexGatewayBorrowTestRepositorySavesFullRawValuesAndMarksRestartIncomplete(t *testing.T) { @@ -128,11 +138,22 @@ func TestCodexGatewayBorrowTestRepositorySavesFullRawValuesAndMarksRestartIncomp task := codexGatewayBorrowRepoFixture(t) result := task.Results[0] result.Status, result.UpstreamModel, result.RawAnswer, result.RawResponse, result.RawHTML, result.HTML, result.Error = "incomplete", "reported", "answer\x00tail", "SSE\x00raw", "source\x00raw", "html\x00raw", "API returned 403: entire original body\x00tail" + started := time.Now().UTC() + result.ActualEndpoint, result.ActualProtocol, result.ActualTransport, result.BorrowApplied = "https://upstream.test/responses", "responses", "http", true + result.QueueDurationMS, result.PreparationDurationMS, result.GenerationDurationMS, result.GenerationStartedAt = 1200, 2300, 3400, &started mock.ExpectBegin() - mock.ExpectExec("UPDATE codex_gateway_borrow_test_results r.*t.expires_at>CURRENT_TIMESTAMP").WithArgs(result.ID, result.TaskID, result.AccountName, result.UpstreamModel, result.Effort, result.Status, []byte(result.RawAnswer), []byte(result.RawResponse), []byte(result.RawHTML), []byte(result.HTML), []byte(result.Error), nil, nil, int64(0)).WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec("UPDATE codex_gateway_borrow_test_results r.*t.expires_at>CURRENT_TIMESTAMP").WithArgs(result.ID, result.TaskID, result.AccountName, result.UpstreamModel, result.Effort, result.Status, []byte(result.RawAnswer), []byte(result.RawResponse), []byte(result.RawHTML), []byte(result.HTML), []byte(result.Error), nil, nil, int64(0), + result.Platform, result.ActualEndpoint, result.ActualProtocol, result.ActualTransport, result.BorrowApplied, result.QueueDurationMS, result.PreparationDurationMS, result.GenerationDurationMS, result.GenerationStartedAt).WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("UPDATE codex_gateway_borrow_test_tasks t.*completed=.*status NOT IN").WithArgs(task.ID).WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectCommit() require.NoError(t, repo.SaveResult(context.Background(), result)) + mock.ExpectQuery("WHERE t.scope='manual_pelican' AND r.id=\\$1 AND t.expires_at>CURRENT_TIMESTAMP").WithArgs(result.ID).WillReturnRows(codexGatewayBorrowResultTestRows(result)) + stored, err := repo.GetResult(context.Background(), result.ID) + require.NoError(t, err) + result.PreviewUnavailable = "test incomplete" + require.Equal(t, result, stored, "actual invocation and phase durations round-trip alongside untruncated raw observations") + // The reader derives the preview reason from the terminal status. + require.Equal(t, "test incomplete", stored.PreviewUnavailable) mock.ExpectBegin() mock.ExpectExec("UPDATE codex_gateway_borrow_test_results r.*status='incomplete'.*t.scope='manual_pelican'.*r.status IN").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("UPDATE codex_gateway_borrow_test_tasks t.*status='incomplete'.*t.scope='manual_pelican'").WillReturnResult(sqlmock.NewResult(0, 1)) diff --git a/backend/internal/repository/concurrency_account_lease_test.go b/backend/internal/repository/concurrency_account_lease_test.go new file mode 100644 index 000000000000..156d236a75cb --- /dev/null +++ b/backend/internal/repository/concurrency_account_lease_test.go @@ -0,0 +1,96 @@ +//go:build unit + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestObservedAccountLeaseCacheRefreshRetainsThirtyMinuteRequestAndActiveIndex(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { require.NoError(t, client.Close()) }) + cache := NewConcurrencyCache(client, 0, 0) + refresher, ok := cache.(service.AccountSlotRefreshCache) + require.True(t, ok) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + server.SetTime(now) + acquired, err := cache.AcquireAccountSlot(ctx, 42, 1, "observed-request") + require.NoError(t, err) + require.True(t, acquired) + for i := 0; i < 60; i++ { + now = now.Add(30 * time.Second) + server.FastForward(30 * time.Second) + server.SetTime(now) + owned, refreshErr := refresher.RefreshAccountSlot(ctx, 42, "observed-request") + require.NoError(t, refreshErr) + require.True(t, owned) + count, countErr := cache.GetAccountConcurrency(ctx, 42) + require.NoError(t, countErr) + require.Equal(t, 1, count) + acquired, err = cache.AcquireAccountSlot(ctx, 42, 1, "business-must-wait") + require.NoError(t, err) + require.False(t, acquired, "the same regular account capacity remains reserved for the entire 30-minute observation") + } + score, err := client.ZScore(ctx, accountSlotKey(42), "observed-request").Result() + require.NoError(t, err) + require.Equal(t, float64(now.Unix()), score) + indexExpiry, err := client.ZScore(ctx, accountActiveIndexKey, "42").Result() + require.NoError(t, err) + require.Equal(t, float64(now.Unix()+15*60), indexExpiry) + require.Equal(t, 15*time.Minute, server.TTL(accountSlotKey(42))) + require.NoError(t, cache.ReleaseAccountSlot(ctx, 42, "observed-request")) + _, err = client.ZScore(ctx, accountSlotKey(42), "observed-request").Result() + require.ErrorIs(t, err, redis.Nil) + _, err = client.ZScore(ctx, accountActiveIndexKey, "42").Result() + require.ErrorIs(t, err, redis.Nil) + acquired, err = cache.AcquireAccountSlot(ctx, 42, 1, "business-after-release") + require.NoError(t, err) + require.True(t, acquired) +} + +func TestObservedAccountLeaseCacheCannotRecreateExpiredOrDeletedOwner(t *testing.T) { + for _, expired := range []bool{false, true} { + t.Run(map[bool]string{false: "deleted", true: "expired"}[expired], func(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { require.NoError(t, client.Close()) }) + cache := NewConcurrencyCache(client, 1, 60) + refresher := cache.(service.AccountSlotRefreshCache) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + server.SetTime(now) + acquired, err := cache.AcquireAccountSlot(ctx, 42, 1, "lost-owner") + require.NoError(t, err) + require.True(t, acquired) + if expired { + now = now.Add(time.Minute) + server.FastForward(time.Minute) + server.SetTime(now) + } else { + require.NoError(t, client.ZRem(ctx, accountSlotKey(42), "lost-owner").Err()) + } + acquired, err = cache.AcquireAccountSlot(ctx, 42, 1, "new-business-owner") + require.NoError(t, err) + require.True(t, acquired) + owned, err := refresher.RefreshAccountSlot(ctx, 42, "lost-owner") + require.NoError(t, err) + require.False(t, owned, "refresh must not silently reacquire beyond the account's business limit") + _, err = client.ZScore(ctx, accountSlotKey(42), "lost-owner").Result() + require.ErrorIs(t, err, redis.Nil) + count, err := cache.GetAccountConcurrency(ctx, 42) + require.NoError(t, err) + require.Equal(t, 1, count) + _, err = client.ZScore(ctx, accountSlotKey(42), "new-business-owner").Result() + require.NoError(t, err) + }) + } +} diff --git a/backend/internal/repository/concurrency_cache.go b/backend/internal/repository/concurrency_cache.go index c21268f52611..abaddf8c5c05 100644 --- a/backend/internal/repository/concurrency_cache.go +++ b/backend/internal/repository/concurrency_cache.go @@ -106,6 +106,23 @@ var ( return {0, now} `) + // Renew one existing regular account-slot member. A removed or expired + // request ID is never recreated, even when the account has spare capacity. + refreshAccountSlotScript = redis.NewScript(` + redis.replicate_commands() + local key = KEYS[1] + local ttl = tonumber(ARGV[1]) + local requestID = ARGV[2] + local now = tonumber(redis.call('TIME')[1]) + redis.call('ZREMRANGEBYSCORE', key, '-inf', now - ttl) + if redis.call('ZSCORE', key, requestID) == false then + return {0, now} + end + redis.call('ZADD', key, now, requestID) + redis.call('EXPIRE', key, ttl) + return {1, now} + `) + // getCountScript 统计有序集合中的槽位数量并清理过期条目 // 使用 Redis TIME 命令获取服务器时间 // KEYS[1] = 普通槽位键，KEYS[2] = 对应 Live 槽位键 @@ -652,6 +669,19 @@ func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int return nil } +func (c *concurrencyCache) RefreshAccountSlot(ctx context.Context, accountID int64, requestID string) (bool, error) { + result, now, err := runScriptInt64Pair(ctx, c.rdb, refreshAccountSlotScript, []string{accountSlotKey(accountID)}, c.slotTTLSeconds, requestID) + if err != nil { + return false, err + } + if result == 1 { + c.touchActiveIndexAt(ctx, accountActiveIndexKey, accountID, now+int64(c.slotTTLSeconds)) + } else { + c.refreshAccountActiveIndex(ctx, accountID) + } + return result == 1, nil +} + func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) { key := accountSlotKey(accountID) // 时间戳在 Lua 脚本内使用 Redis TIME 命令获取 diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index 757607d97fc5..771c0c72adc0 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -203,11 +203,18 @@ func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream { // 注意: // - 调用方必须关闭 resp.Body，否则会导致 inFlight 计数泄漏 // - inFlight > 0 的客户端不会被淘汰，确保活跃请求不被中断 -func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { +func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (response *http.Response, requestErr error) { applyGrokCLIProxyHeaders(req) if err := s.validateRequestHost(req); err != nil { return nil, err } + var finishPelican func(*http.Response, error) (*http.Response, error) + var err error + req, finishPelican, err = service.PreparePelicanHTTPRequest(req, accountID, accountConcurrency, "http") + if err != nil { + return nil, err + } + defer func() { response, requestErr = finishPelican(response, requestErr) }() profile := service.HTTPUpstreamProfileDefault if req != nil { profile = service.HTTPUpstreamProfileFromContext(req.Context()) @@ -246,7 +253,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i // // profile 为 nil 时不启用 TLS 指纹，行为与 Do 方法相同。 // profile 非 nil 时使用指定的 Profile 进行 TLS 指纹伪装。 -func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) { +func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (response *http.Response, requestErr error) { if profile == nil { return s.Do(req, proxyURL, accountID, accountConcurrency) } @@ -274,6 +281,13 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco if err := s.validateRequestHost(req); err != nil { return nil, err } + var finishPelican func(*http.Response, error) (*http.Response, error) + var err error + req, finishPelican, err = service.PreparePelicanHTTPRequest(req, accountID, accountConcurrency, "http") + if err != nil { + return nil, err + } + defer func() { response, requestErr = finishPelican(response, requestErr) }() entry, err := s.acquireClientWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile) if err != nil { diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index 1dd14476a020..4e2007bd58e5 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -23,6 +23,7 @@ func RegisterAdminRoutes( // 插件 UI 使用短时能力 URL，仅提供经过安装校验的静态资源。 v1.GET("/plugin-ui/:token/*path", h.Admin.Plugin.ServeUIAsset) v1.GET("/codex-gateway-borrow/preview/:cap/index.html", h.Admin.CodexGatewayBorrow.ServePreview) + v1.GET("/pelican-tests/preview/:cap/index.html", h.Admin.PelicanTest.ServePreview) admin := v1.Group("/admin") admin.Use(gin.HandlerFunc(adminAuth)) @@ -133,6 +134,7 @@ func RegisterAdminRoutes( // 推理恢复全局开关 registerReasoningRecoveryRoutes(admin, h) registerCodexGatewayBorrowRoutes(admin, h) + registerPelicanTestRoutes(admin, h) // 邀请返利（专属用户管理） registerAffiliateRoutes(admin, h) diff --git a/backend/internal/server/routes/pelican_tests.go b/backend/internal/server/routes/pelican_tests.go new file mode 100644 index 000000000000..492d580f167f --- /dev/null +++ b/backend/internal/server/routes/pelican_tests.go @@ -0,0 +1,15 @@ +package routes + +import ( + "github.com/Wei-Shaw/sub2api/internal/handler" + "github.com/gin-gonic/gin" +) + +func registerPelicanTestRoutes(admin *gin.RouterGroup, h *handler.Handlers) { + pelican := admin.Group("/pelican-tests") + pelican.GET("/options", h.Admin.PelicanTest.GetOptions) + pelican.POST("/options", h.Admin.PelicanTest.GetAccountOptions) + pelican.POST("/tests", h.Admin.PelicanTest.StartTests) + pelican.GET("/tests", h.Admin.PelicanTest.ListTests) + pelican.GET("/tests/:id", h.Admin.PelicanTest.GetTest) +} diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index b9691c94a697..980025e7fe67 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -164,6 +164,10 @@ type AccountTestService struct { settingService *SettingService tlsFPProfileService *TLSFingerprintProfileService openaiGatewayService *OpenAIGatewayService + pelicanGatewayService *GatewayService + pelicanGeminiService *GeminiMessagesCompatService + pelicanConcurrencyService *ConcurrencyService + pelicanBorrowService *CodexGatewayBorrowService modelMetadataRegistryMu sync.Mutex modelMetadataRegistry map[string]modelsDevProvider modelMetadataRegistryAt time.Time @@ -194,7 +198,36 @@ func (s *AccountTestService) SetPluginManager(pluginManager *PluginManager) { func (s *AccountTestService) SetOpenAIGatewayService(gateway *OpenAIGatewayService) { if s != nil { s.openaiGatewayService = gateway + if gateway != nil && s.pelicanBorrowService != nil { + gateway.SetCodexGatewayBorrowService(s.pelicanBorrowService) + } + } +} + +func (s *AccountTestService) SetCodexGatewayBorrowService(borrow *CodexGatewayBorrowService) { + if s != nil { + s.pelicanBorrowService = borrow + if s.openaiGatewayService != nil { + s.openaiGatewayService.SetCodexGatewayBorrowService(borrow) + } + } +} + +// SetPelicanBusinessServices binds the same account senders used by ordinary +// gateway requests. Dependencies are installed once during application wiring. +func (s *AccountTestService) SetPelicanBusinessServices(gateway *GatewayService, gemini *GeminiMessagesCompatService, concurrency *ConcurrencyService) { + if s != nil { + s.pelicanGatewayService = gateway + s.pelicanGeminiService = gemini + s.pelicanConcurrencyService = concurrency + } +} + +func (s *AccountTestService) setAccountTestError(ctx context.Context, accountID int64, message string) error { + if IsAccountObservation(ctx) { + return nil } + return s.accountRepo.SetError(ctx, accountID, message) } // FetchOpenAIAccountModels projects shared discovery data into the account UI @@ -430,10 +463,13 @@ func (s *AccountTestService) TestAccountConnection(c *gin.Context, accountID int } // Route to platform-specific test method - ctx, finishObservation := s.quotaActivity.Attach(ctx) - c.Request = c.Request.WithContext(ctx) - ObserveQuotaAccount(ctx, account.ID) - defer finishObservation() + if !IsAccountObservation(ctx) { + var finishObservation func() + ctx, finishObservation = s.quotaActivity.Attach(ctx) + c.Request = c.Request.WithContext(ctx) + ObserveQuotaAccount(ctx, account.ID) + defer finishObservation() + } if account.IsOpenCodeGo() { return s.testOpenCodeGoConnection(c, account, modelID, prompt) } @@ -633,7 +669,7 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account // 403 表示账号被上游封禁，标记为 error 状态 if resp.StatusCode == http.StatusForbidden { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) @@ -701,7 +737,7 @@ func (s *AccountTestService) testClaudeVertexServiceAccountConnection(c *gin.Con body, _ := io.ReadAll(resp.Body) errMsg := fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)) if resp.StatusCode == http.StatusForbidden { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) } @@ -993,7 +1029,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account } defer func() { _ = resp.Body.Close() }() - if isOAuth && !account.IsShadow() && s.accountRepo != nil { + if isOAuth && !account.IsShadow() && s.accountRepo != nil && !IsAccountObservation(ctx) { if updates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(updates) > 0 { _ = s.accountRepo.UpdateExtra(ctx, account.ID, updates) mergeAccountExtra(account, updates) @@ -1018,7 +1054,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account // 401 Unauthorized: 标记账号为永久错误 if resp.StatusCode == http.StatusUnauthorized && !budgetExceeded && s.accountRepo != nil { errMsg := fmt.Sprintf("Authentication failed (401): %s", string(body)) - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body))) } @@ -1191,6 +1227,9 @@ func (s *AccountTestService) applyGrokTestRequestHeaders(req *http.Request, acco } func (s *AccountTestService) observeGrokTestResponse(ctx context.Context, account *Account, resp *http.Response) { + if IsAccountObservation(ctx) { + return + } if resp == nil { return } @@ -2211,7 +2250,7 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection( } if resp.StatusCode == http.StatusUnauthorized && !budgetExceeded && s.accountRepo != nil { errMsg := fmt.Sprintf("Chat Completions authentication failed (401): %s", string(body)) - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, fmt.Sprintf("Chat Completions API (/v1/chat/completions) returned %d: %s", resp.StatusCode, string(body))) } @@ -2335,7 +2374,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account resp, err := s.doOpenAIAccountTestUpstream(req, proxyURL, account, true) if err != nil { - if s.accountRepo != nil { + if s.accountRepo != nil && !IsAccountObservation(ctx) { updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, false, time.Now()) _ = s.accountRepo.UpdateExtra(ctx, account.ID, updates) mergeAccountExtra(account, updates) @@ -2357,7 +2396,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account } compactionFound := openAICompactProbeFoundCompactionItem(body) - if s.accountRepo != nil { + if s.accountRepo != nil && !IsAccountObservation(ctx) { updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, compactionFound, time.Now()) if !account.IsShadow() { if codexUpdates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(codexUpdates) > 0 { @@ -2377,7 +2416,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account if resp.StatusCode != http.StatusOK { if resp.StatusCode == http.StatusUnauthorized && !budgetExceeded && s.accountRepo != nil { errMsg := fmt.Sprintf("Authentication failed (401): %s", string(body)) - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body))) } @@ -2392,7 +2431,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account } func (s *AccountTestService) reconcileOpenAI429State(ctx context.Context, account *Account, headers http.Header, body []byte) { - if s == nil || s.accountRepo == nil || account == nil { + if IsAccountObservation(ctx) || s == nil || s.accountRepo == nil || account == nil { return } // Spark quota is observed through its own /wham scope, not global probe 429s. @@ -2463,7 +2502,7 @@ func (s *AccountTestService) markOpenAIBudgetExceededFromTest(ctx context.Contex (statusCode != http.StatusOK || !isOpenAIBudgetExceededTerminalEvent(body)) { return false } - if s != nil { + if s != nil && !IsAccountObservation(ctx) { s.openaiGatewayService.handleOpenAIBudgetExceeded(ctx, account, body) } return true diff --git a/backend/internal/service/account_test_service_cn_adaptive.go b/backend/internal/service/account_test_service_cn_adaptive.go index d73665bdc0d6..079ac6857ccd 100644 --- a/backend/internal/service/account_test_service_cn_adaptive.go +++ b/backend/internal/service/account_test_service_cn_adaptive.go @@ -96,7 +96,7 @@ func (s *AccountTestService) testCNProviderAdaptiveAnthropicConnection(c *gin.Co body, _ := io.ReadAll(resp.Body) errMsg := fmt.Sprintf("Adaptive Anthropic endpoint returned %d: %s", resp.StatusCode, string(body)) if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) } @@ -150,7 +150,7 @@ func (s *AccountTestService) testCNProviderAdaptiveResponsesConnection(c *gin.Co body, _ := io.ReadAll(resp.Body) errMsg := fmt.Sprintf("Adaptive Responses endpoint returned %d: %s", resp.StatusCode, string(body)) if resp.StatusCode == http.StatusUnauthorized && s.accountRepo != nil { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) } @@ -240,7 +240,7 @@ func (s *AccountTestService) testCNProviderAnthropicConnection(c *gin.Context, a body, _ := io.ReadAll(resp.Body) errMsg := fmt.Sprintf("Anthropic endpoint returned %d: %s", resp.StatusCode, string(body)) if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && s.accountRepo != nil { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) } diff --git a/backend/internal/service/account_test_service_typesafe.go b/backend/internal/service/account_test_service_typesafe.go index 27d1093cec1e..a55a16ee575b 100644 --- a/backend/internal/service/account_test_service_typesafe.go +++ b/backend/internal/service/account_test_service_typesafe.go @@ -83,7 +83,7 @@ func (s *AccountTestService) testTypeSafeAccountConnection(c *gin.Context, accou errMsg := fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)) // 401/403 表示 API Key 无效或被上游拒绝，标记为 error 状态。 if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && s.accountRepo != nil { - _ = s.accountRepo.SetError(ctx, account.ID, errMsg) + _ = s.setAccountTestError(ctx, account.ID, errMsg) } return s.sendErrorAndEnd(c, errMsg) } diff --git a/backend/internal/service/antigravity_credits_overages.go b/backend/internal/service/antigravity_credits_overages.go index ec3650859ce6..10b45f6174ef 100644 --- a/backend/internal/service/antigravity_credits_overages.go +++ b/backend/internal/service/antigravity_credits_overages.go @@ -59,6 +59,9 @@ func (a *Account) isCreditsExhausted() bool { // setCreditsExhausted 标记账号积分耗尽：写入 model_rate_limits["AICredits"] + 更新缓存。 func (s *AntigravityGatewayService) setCreditsExhausted(ctx context.Context, account *Account) { + if IsAccountObservation(ctx) { + return + } if account == nil || account.ID == 0 { return } @@ -74,6 +77,9 @@ func (s *AntigravityGatewayService) setCreditsExhausted(ctx context.Context, acc // clearCreditsExhausted 清除账号的 AICredits 限流 key。 func (s *AntigravityGatewayService) clearCreditsExhausted(ctx context.Context, account *Account) { + if IsAccountObservation(ctx) { + return + } if account == nil || account.ID == 0 || account.Extra == nil { return } diff --git a/backend/internal/service/antigravity_gateway_retry.go b/backend/internal/service/antigravity_gateway_retry.go index ff3862569097..6ca62db41465 100644 --- a/backend/internal/service/antigravity_gateway_retry.go +++ b/backend/internal/service/antigravity_gateway_retry.go @@ -710,7 +710,7 @@ urlFallbackLoop: } } - if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { + if !IsAccountObservation(p.ctx) && resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL) } @@ -807,6 +807,9 @@ const googleConfigErrorCooldown = 1 * time.Minute // tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁， // 避免短时间内反复调度到同一个有问题的账号。 func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + if IsAccountObservation(ctx) { + return + } until := time.Now().Add(googleConfigErrorCooldown) reason := "400: invalid project resource name (auto temp-unschedule 1m)" if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { @@ -822,6 +825,9 @@ const emptyResponseCooldown = 1 * time.Minute // tempUnscheduleEmptyResponse 对空流式响应触发临时封禁， // 避免短时间内反复调度到同一个返回空响应的账号。 func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + if IsAccountObservation(ctx) { + return + } until := time.Now().Add(emptyResponseCooldown) reason := "empty stream response (auto temp-unschedule 1m)" if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { @@ -867,6 +873,9 @@ func isSingleAccountRetry(ctx context.Context) bool { // 直接使用上游返回的模型 ID（如 claude-sonnet-4-5）作为限流 key // 返回是否已成功设置（若模型名为空或 repo 为 nil 将返回 false） func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if IsAccountObservation(ctx) { + return false + } if repo == nil || modelName == "" { return false } @@ -903,6 +912,9 @@ func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Co } func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { + if IsAccountObservation(ctx) { + return + } if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { return } @@ -1159,6 +1171,9 @@ func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimit // setModelRateLimitAndClearSession 设置模型限流并清除粘性会话 func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) { + if IsAccountObservation(p.ctx) { + return + } resetAt := time.Now().Add(info.RetryDelay) logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) @@ -1173,6 +1188,9 @@ func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleMo // updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态 func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) { + if IsAccountObservation(ctx) { + return + } if s.schedulerSnapshot == nil || account == nil || modelKey == "" { return } @@ -1205,6 +1223,9 @@ func (s *AntigravityGatewayService) handleUpstreamError( requestedModel string, groupID int64, sessionHash string, isStickySession bool, ) *handleModelRateLimitResult { + if IsAccountObservation(ctx) { + return nil + } // 遵守自定义错误码策略：未命中则跳过所有限流处理 if !account.ShouldHandleErrorCode(statusCode) { return nil diff --git a/backend/internal/service/antigravity_internal500_penalty.go b/backend/internal/service/antigravity_internal500_penalty.go index 747a4d4eb329..1938800fa892 100644 --- a/backend/internal/service/antigravity_internal500_penalty.go +++ b/backend/internal/service/antigravity_internal500_penalty.go @@ -35,6 +35,9 @@ func isAntigravityInternalServerError(statusCode int, body []byte) bool { func (s *AntigravityGatewayService) applyInternal500Penalty( ctx context.Context, prefix string, account *Account, count int64, ) { + if IsAccountObservation(ctx) { + return + } switch { case count >= int64(internal500PenaltyTier3Threshold): reason := fmt.Sprintf("INTERNAL 500 consecutive failures: %d rounds", count) @@ -71,6 +74,9 @@ func (s *AntigravityGatewayService) applyInternal500Penalty( func (s *AntigravityGatewayService) handleInternal500RetryExhausted( ctx context.Context, prefix string, account *Account, ) { + if IsAccountObservation(ctx) { + return + } if s.internal500Cache == nil { return } @@ -87,6 +93,9 @@ func (s *AntigravityGatewayService) handleInternal500RetryExhausted( func (s *AntigravityGatewayService) resetInternal500Counter( ctx context.Context, prefix string, accountID int64, ) { + if IsAccountObservation(ctx) { + return + } if s.internal500Cache == nil { return } diff --git a/backend/internal/service/antigravity_token_provider.go b/backend/internal/service/antigravity_token_provider.go index a0d0e1c26e8a..f7fcdd568d67 100644 --- a/backend/internal/service/antigravity_token_provider.go +++ b/backend/internal/service/antigravity_token_provider.go @@ -104,7 +104,9 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account * result, err := p.refreshAPI.RefreshIfNeeded(refreshCtx, account, p.executor, antigravityTokenRefreshSkew) if err != nil { // 标记账号临时不可调度，避免后续请求继续命中 - p.markTempUnschedulable(account, err) + if !IsAccountObservation(ctx) { + p.markTempUnschedulable(account, err) + } if p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn { return "", err } diff --git a/backend/internal/service/codex_gateway_borrow.go b/backend/internal/service/codex_gateway_borrow.go index 7ac331caadb0..2aa5233bae3d 100644 --- a/backend/internal/service/codex_gateway_borrow.go +++ b/backend/internal/service/codex_gateway_borrow.go @@ -396,7 +396,7 @@ type codexGatewayBorrowObservationContextKey struct{} // The token provider may refresh credentials normally, but must not attach an // observation failure to account health, scheduling or quota. func WithCodexGatewayBorrowObservation(ctx context.Context) context.Context { - return context.WithValue(ctx, codexGatewayBorrowObservationContextKey{}, true) + return context.WithValue(WithAccountObservation(ctx), codexGatewayBorrowObservationContextKey{}, true) } func IsCodexGatewayBorrowObservation(ctx context.Context) bool { return ctx != nil && ctx.Value(codexGatewayBorrowObservationContextKey{}) == true diff --git a/backend/internal/service/codex_gateway_borrow_http.go b/backend/internal/service/codex_gateway_borrow_http.go index dc01e0731054..f073bdad44ba 100644 --- a/backend/internal/service/codex_gateway_borrow_http.go +++ b/backend/internal/service/codex_gateway_borrow_http.go @@ -188,6 +188,9 @@ func withCodexGatewayBorrowHTTPFirstOutputStart(ctx context.Context, start time. } func codexGatewayBorrowHTTPFirstOutputStart(ctx context.Context, fallback time.Time) time.Time { + if IsPelicanGeneration(ctx) { + return PelicanFirstOutputStart(ctx, fallback) + } if ctx != nil { if start, ok := ctx.Value(codexGatewayBorrowHTTPFirstOutputContextKey{}).(time.Time); ok { return start diff --git a/backend/internal/service/codex_gateway_borrow_preview.go b/backend/internal/service/codex_gateway_borrow_preview.go index a9a5ea17b93c..1f36381024b9 100644 --- a/backend/internal/service/codex_gateway_borrow_preview.go +++ b/backend/internal/service/codex_gateway_borrow_preview.go @@ -196,7 +196,7 @@ func (p *CodexGatewayBorrowPreviewService) DecorateResult(result *CodexGatewayBo } p.capabilities[token] = codexGatewayBorrowPreviewCapability{ResultID: decorated.ID, ExpiresAt: expires} p.mu.Unlock() - decorated.PreviewURL = "/api/v1/codex-gateway-borrow/preview/" + token + "/index.html" + decorated.PreviewURL = "/api/v1/pelican-tests/preview/" + token + "/index.html" decorated.PreviewExpiresAt, decorated.PreviewUnavailable = &expires, "" return &decorated } diff --git a/backend/internal/service/codex_gateway_borrow_preview_test.go b/backend/internal/service/codex_gateway_borrow_preview_test.go index 71ba1ce8d109..1f2eed0f0980 100644 --- a/backend/internal/service/codex_gateway_borrow_preview_test.go +++ b/backend/internal/service/codex_gateway_borrow_preview_test.go @@ -56,7 +56,8 @@ func TestCodexGatewayBorrowPreviewCapabilityExpiryAndRecordRecheck(t *testing.T) require.Empty(t, decorated.PreviewUnavailable) require.Equal(t, result.ExpiresAt, *decorated.PreviewExpiresAt) require.Empty(t, result.PreviewURL, "minting must not mutate the stored result") - token := strings.TrimSuffix(strings.TrimPrefix(decorated.PreviewURL, "/api/v1/codex-gateway-borrow/preview/"), "/index.html") + require.True(t, strings.HasPrefix(decorated.PreviewURL, "/api/v1/pelican-tests/preview/")) + token := strings.TrimSuffix(strings.TrimPrefix(decorated.PreviewURL, "/api/v1/pelican-tests/preview/"), "/index.html") require.Len(t, token, 43) html, err := preview.Resolve(context.Background(), token) require.NoError(t, err) @@ -68,7 +69,7 @@ func TestCodexGatewayBorrowPreviewCapabilityExpiryAndRecordRecheck(t *testing.T) result.ExpiresAt = now.Add(time.Hour) decorated = preview.DecorateResult(result) require.Equal(t, now.Add(CodexGatewayBorrowPreviewTTL), *decorated.PreviewExpiresAt) - token = strings.TrimSuffix(strings.TrimPrefix(decorated.PreviewURL, "/api/v1/codex-gateway-borrow/preview/"), "/index.html") + token = strings.TrimSuffix(strings.TrimPrefix(decorated.PreviewURL, "/api/v1/pelican-tests/preview/"), "/index.html") now = now.Add(CodexGatewayBorrowPreviewTTL) reads := repo.resultReads _, err = preview.Resolve(context.Background(), token) diff --git a/backend/internal/service/codex_gateway_borrow_probe.go b/backend/internal/service/codex_gateway_borrow_probe.go index fcbaf59cd00f..de7f25a291e3 100644 --- a/backend/internal/service/codex_gateway_borrow_probe.go +++ b/backend/internal/service/codex_gateway_borrow_probe.go @@ -301,6 +301,27 @@ func (s *CodexGatewayBorrowService) fireObservation(ctx context.Context, headers func (s *CodexGatewayBorrowService) sendObservation(req *http.Request, account *Account, proxy string, profile *tlsfingerprint.Profile, expectedModel string, maxBody int64) (codexGatewayBorrowObservation, error) { var shot codexGatewayBorrowObservation + leasedCtx, release, acquireErr := AcquirePelicanExecution(req.Context(), account.ID) + if acquireErr != nil { + shot.errorText = acquireErr.Error() + return shot, acquireErr + } + defer release() + req = req.WithContext(leasedCtx) + if pelicanExecutionFromContext(leasedCtx) != nil { + var capacity *ConcurrencyService + if capture := pelicanCaptureFromContext(leasedCtx); capture != nil { + capacity = capture.concurrency + } else if s.gateway != nil { + capacity = s.gateway.concurrencyService + } + releaseCapacity, capacityErr := acquirePelicanBusinessAccountSlot(leasedCtx, capacity, account.ID, account.Concurrency) + if capacityErr != nil { + shot.errorText = capacityErr.Error() + return shot, capacityErr + } + defer releaseCapacity() + } wire, err := prepareOpenAICodexWireRequest(req, account) if err != nil { return shot, err diff --git a/backend/internal/service/codex_gateway_borrow_test_runner.go b/backend/internal/service/codex_gateway_borrow_test_runner.go index f42b04b317d2..b6adbfbc5b26 100644 --- a/backend/internal/service/codex_gateway_borrow_test_runner.go +++ b/backend/internal/service/codex_gateway_borrow_test_runner.go @@ -17,13 +17,12 @@ import ( const ( CodexGatewayBorrowTestTTL = 24 * time.Hour CodexGatewayBorrowPelicanPrompt = "创建一个HTML，内容是SVG绘制一个鹈鹕骑自行车的2D动画，你不需要任何测试" - codexGatewayBorrowTestTimeout = 90 * time.Second codexGatewayBorrowTestMaxTargets = 5000 ) var ( ErrCodexGatewayBorrowTestNotFound = errors.New("codex gateway borrow test not found or expired") - ErrCodexGatewayBorrowTestReplayConflict = errors.New("client_task_id was reused with different targets") + ErrCodexGatewayBorrowTestReplayConflict = errors.New("client_task_id was reused with different test parameters") ErrCodexGatewayBorrowTestExpired = errors.New("codex gateway borrow test has expired") ErrCodexGatewayBorrowTestInvalidRequest = errors.New("invalid manual pelican test request") ) @@ -35,50 +34,63 @@ type CodexGatewayBorrowTestTarget struct { } type CodexGatewayBorrowTestRequest struct { - ClientTaskID string `json:"client_task_id"` - Targets []CodexGatewayBorrowTestTarget `json:"targets"` + ClientTaskID string `json:"client_task_id"` + GenerationTimeoutSeconds int `json:"generation_timeout_seconds"` + Targets []CodexGatewayBorrowTestTarget `json:"targets"` + Standalone bool `json:"-"` } type CodexGatewayBorrowTestResult struct { - ID string `json:"id"` - TaskID string `json:"task_id"` - Ordinal int `json:"ordinal"` - AccountID int64 `json:"account_id"` - AccountName string `json:"account_name"` - ModelID string `json:"model_id"` - UpstreamModel string `json:"upstream_model"` - Effort string `json:"effort"` - Status string `json:"status"` - RawAnswer string `json:"raw_answer"` - RawResponse string `json:"raw_response"` - RawHTML string `json:"raw_html"` - HTML string `json:"html"` - Error string `json:"error"` - StartedAt *time.Time `json:"started_at,omitempty"` - FinishedAt *time.Time `json:"finished_at,omitempty"` - DurationMS int64 `json:"duration_ms"` - ExpiresAt time.Time `json:"expires_at"` - PreviewURL string `json:"preview_url,omitempty"` - PreviewExpiresAt *time.Time `json:"preview_expires_at,omitempty"` - PreviewUnavailable string `json:"preview_unavailable"` + ID string `json:"id"` + TaskID string `json:"task_id"` + Ordinal int `json:"ordinal"` + AccountID int64 `json:"account_id"` + AccountName string `json:"account_name"` + Platform string `json:"platform"` + ModelID string `json:"model_id"` + UpstreamModel string `json:"upstream_model"` + Effort string `json:"effort"` + Status string `json:"status"` + RawAnswer string `json:"raw_answer"` + RawResponse string `json:"raw_response"` + RawHTML string `json:"raw_html"` + HTML string `json:"html"` + Error string `json:"error"` + StartedAt *time.Time `json:"started_at,omitempty"` + FinishedAt *time.Time `json:"finished_at,omitempty"` + DurationMS int64 `json:"duration_ms"` + GenerationStartedAt *time.Time `json:"generation_started_at,omitempty"` + QueueDurationMS int64 `json:"queue_duration_ms"` + PreparationDurationMS int64 `json:"preparation_duration_ms"` + GenerationDurationMS int64 `json:"generation_duration_ms"` + ActualEndpoint string `json:"actual_endpoint"` + ActualProtocol string `json:"actual_protocol"` + ActualTransport string `json:"actual_transport"` + BorrowApplied bool `json:"borrow_applied"` + ExpiresAt time.Time `json:"expires_at"` + PreviewURL string `json:"preview_url,omitempty"` + PreviewExpiresAt *time.Time `json:"preview_expires_at,omitempty"` + PreviewUnavailable string `json:"preview_unavailable"` } type CodexGatewayBorrowTestTask struct { - ID string `json:"id"` - ClientTaskID string `json:"client_task_id"` - CreatedBy int64 `json:"created_by"` - RequestHash string `json:"-"` - Status string `json:"status"` - Prompt string `json:"prompt"` - CreatedAt time.Time `json:"created_at"` - ExpiresAt time.Time `json:"expires_at"` - StartedAt *time.Time `json:"started_at,omitempty"` - FinishedAt *time.Time `json:"finished_at,omitempty"` - Total int `json:"total"` - Completed int `json:"completed"` - Error string `json:"error"` - Results []*CodexGatewayBorrowTestResult `json:"results,omitempty"` - Replayed bool `json:"replayed"` + ID string `json:"id"` + ClientTaskID string `json:"client_task_id"` + CreatedBy int64 `json:"created_by"` + RequestHash string `json:"-"` + GenerationTimeoutSeconds int `json:"generation_timeout_seconds"` + ExecutionMode string `json:"execution_mode"` + Status string `json:"status"` + Prompt string `json:"prompt"` + CreatedAt time.Time `json:"created_at"` + ExpiresAt time.Time `json:"expires_at"` + StartedAt *time.Time `json:"started_at,omitempty"` + FinishedAt *time.Time `json:"finished_at,omitempty"` + Total int `json:"total"` + Completed int `json:"completed"` + Error string `json:"error"` + Results []*CodexGatewayBorrowTestResult `json:"results,omitempty"` + Replayed bool `json:"replayed"` } type CodexGatewayBorrowTestList struct { @@ -111,30 +123,26 @@ type codexGatewayBorrowAccountReader interface { GetByID(context.Context, int64) (*Account, error) } -type codexGatewayBorrowAccountSlot struct { - slot chan struct{} - refs int -} - -// All batches share these three slots and account locks. No native test service, -// recovery service, rate-limit writer or account-state writer is used here. +// All batches, including legacy cache-only tests, share ten outbound execution +// slots. Account leases surround actual requests rather than route preparation. type CodexGatewayBorrowTestRunner struct { - repo CodexGatewayBorrowTestRepository - generator codexGatewayBorrowPelicanGenerator - accounts codexGatewayBorrowAccountReader - slots chan struct{} - accountMu sync.Mutex - accountSlots map[int64]*codexGatewayBorrowAccountSlot - initMu sync.Mutex - initialized bool - stopCtx context.Context - stop context.CancelFunc - stopOnce sync.Once - cleanupDone chan struct{} - runMu sync.Mutex - running sync.WaitGroup - stopped bool - now func() time.Time + repo CodexGatewayBorrowTestRepository + generator codexGatewayBorrowPelicanGenerator + standaloneGenerator codexGatewayBorrowPelicanGenerator + accounts codexGatewayBorrowAccountReader + coordinator *pelicanExecutionCoordinator + initMu sync.Mutex + initialized bool + stopCtx context.Context + stop context.CancelFunc + stopOnce sync.Once + cleanupDone chan struct{} + runMu sync.Mutex + running sync.WaitGroup + stopped bool + activeTasks map[string]bool + now func() time.Time + generationContext func(context.Context, time.Duration) (context.Context, context.CancelFunc) } func NewCodexGatewayBorrowTestRunner(repo CodexGatewayBorrowTestRepository, core *CodexGatewayBorrowService, accounts AccountRepository) *CodexGatewayBorrowTestRunner { @@ -144,12 +152,18 @@ func NewCodexGatewayBorrowTestRunner(repo CodexGatewayBorrowTestRepository, core func newCodexGatewayBorrowTestRunner(repo CodexGatewayBorrowTestRepository, generator codexGatewayBorrowPelicanGenerator, accounts codexGatewayBorrowAccountReader) *CodexGatewayBorrowTestRunner { ctx, cancel := context.WithCancel(context.Background()) r := &CodexGatewayBorrowTestRunner{repo: repo, generator: generator, accounts: accounts, - slots: make(chan struct{}, 3), accountSlots: make(map[int64]*codexGatewayBorrowAccountSlot), - stopCtx: ctx, stop: cancel, cleanupDone: make(chan struct{}), now: time.Now} + coordinator: sharedPelicanExecution, activeTasks: make(map[string]bool), + stopCtx: ctx, stop: cancel, cleanupDone: make(chan struct{}), now: time.Now, generationContext: context.WithTimeout} go r.cleanupLoop() return r } +// SetStandaloneGenerator is configured once before serving. Both entry points +// use the same runner so restart interruption and outbound capacity stay shared. +func (r *CodexGatewayBorrowTestRunner) SetStandaloneGenerator(generator codexGatewayBorrowPelicanGenerator) { + r.standaloneGenerator = generator +} + func (r *CodexGatewayBorrowTestRunner) Stop() { r.stopOnce.Do(func() { r.runMu.Lock() @@ -204,19 +218,26 @@ func normalizeCodexGatewayBorrowTestRequest(request CodexGatewayBorrowTestReques return request, fmt.Errorf("%w: client_task_id timestamp is in the future", ErrCodexGatewayBorrowTestInvalidRequest) } request.ClientTaskID = id.String() + if request.GenerationTimeoutSeconds == 0 { + request.GenerationTimeoutSeconds = PelicanDefaultGenerationTimeoutSeconds + } + if request.GenerationTimeoutSeconds < PelicanMinGenerationTimeoutSeconds || request.GenerationTimeoutSeconds > PelicanMaxGenerationTimeoutSeconds { + return request, fmt.Errorf("%w: generation_timeout_seconds must be between %d and %d", ErrCodexGatewayBorrowTestInvalidRequest, + PelicanMinGenerationTimeoutSeconds, PelicanMaxGenerationTimeoutSeconds) + } targets := make([]CodexGatewayBorrowTestTarget, 0, len(request.Targets)) seen := make(map[string]bool, len(request.Targets)) for _, target := range request.Targets { target.ModelID = strings.TrimSpace(target.ModelID) target.Effort = strings.ToLower(strings.TrimSpace(target.Effort)) - if target.Effort == "" { + if target.Effort == "" && !request.Standalone { target.Effort = "high" } if target.AccountID <= 0 || target.ModelID == "" || len(target.ModelID) > 512 || strings.ContainsRune(target.ModelID, '\x00') { return request, fmt.Errorf("%w: each target needs a positive account_id and a valid model_id", ErrCodexGatewayBorrowTestInvalidRequest) } switch target.Effort { - case "none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra": + case "", "none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra": default: return request, fmt.Errorf("%w: invalid reasoning effort %q", ErrCodexGatewayBorrowTestInvalidRequest, target.Effort) } @@ -244,24 +265,43 @@ func (r *CodexGatewayBorrowTestRunner) Create(ctx context.Context, createdBy int if err = r.initialize(ctx); err != nil { return nil, false, err } + mode := PelicanExecutionModeLegacyCache + if request.Standalone { + mode = PelicanExecutionModeAccount + } + // Default legacy requests retain their existing hash so bookmarks/replays + // continue to read records created before this migration. New execution + // parameters participate in every standalone request's idempotency check. raw, _ := json.Marshal(request.Targets) + if request.Standalone || request.GenerationTimeoutSeconds != PelicanDefaultGenerationTimeoutSeconds { + raw, _ = json.Marshal(struct { + Targets []CodexGatewayBorrowTestTarget `json:"targets"` + GenerationTimeoutSeconds int `json:"generation_timeout_seconds"` + ExecutionMode string `json:"execution_mode"` + }{request.Targets, request.GenerationTimeoutSeconds, mode}) + } hash := sha256.Sum256(raw) now := r.now().UTC() task := &CodexGatewayBorrowTestTask{ID: uuid.NewString(), ClientTaskID: request.ClientTaskID, CreatedBy: createdBy, RequestHash: hex.EncodeToString(hash[:]), Status: "pending", + GenerationTimeoutSeconds: request.GenerationTimeoutSeconds, ExecutionMode: mode, Prompt: CodexGatewayBorrowPelicanPrompt, CreatedAt: now, ExpiresAt: now.Add(CodexGatewayBorrowTestTTL), Total: len(request.Targets), Results: make([]*CodexGatewayBorrowTestResult, 0, len(request.Targets))} - names := make(map[int64]string) + accountInfo := make(map[int64]*Account) for i, target := range request.Targets { - name, found := names[target.AccountID] + account, found := accountInfo[target.AccountID] if !found && r.accounts != nil { - if account, readErr := r.accounts.GetByID(ctx, target.AccountID); readErr == nil && account != nil { - name = account.Name + if read, readErr := r.accounts.GetByID(ctx, target.AccountID); readErr == nil { + account = read } - names[target.AccountID] = name + accountInfo[target.AccountID] = account + } + name, platform := "", "" + if account != nil { + name, platform = account.Name, account.Platform } task.Results = append(task.Results, &CodexGatewayBorrowTestResult{ID: uuid.NewString(), TaskID: task.ID, - Ordinal: i + 1, AccountID: target.AccountID, AccountName: name, ModelID: target.ModelID, + Ordinal: i + 1, AccountID: target.AccountID, AccountName: name, Platform: platform, ModelID: target.ModelID, Effort: target.Effort, Status: "pending", ExpiresAt: task.ExpiresAt, PreviewUnavailable: "test has not completed"}) } @@ -288,62 +328,34 @@ func (r *CodexGatewayBorrowTestRunner) Get(ctx context.Context, id string) (*Cod type CodexGatewayBorrowTestEvent struct { Type string `json:"type"` + Phase string `json:"phase,omitempty"` TaskID string `json:"task_id,omitempty"` Task *CodexGatewayBorrowTestTask `json:"task,omitempty"` Result *CodexGatewayBorrowTestResult `json:"result,omitempty"` } -func (r *CodexGatewayBorrowTestRunner) acquire(ctx context.Context, accountID int64) (func(), bool) { - r.accountMu.Lock() - entry := r.accountSlots[accountID] - if entry == nil { - entry = &codexGatewayBorrowAccountSlot{slot: make(chan struct{}, 1)} - r.accountSlots[accountID] = entry - } - entry.refs++ - r.accountMu.Unlock() - dropReference := func() { - r.accountMu.Lock() - entry.refs-- - if entry.refs == 0 { - delete(r.accountSlots, accountID) - } - r.accountMu.Unlock() - } - select { - case entry.slot <- struct{}{}: - case <-ctx.Done(): - dropReference() - return nil, false - } - select { - case r.slots <- struct{}{}: - case <-ctx.Done(): - <-entry.slot - dropReference() - return nil, false - } - if ctx.Err() != nil { - <-r.slots - <-entry.slot - dropReference() - return nil, false - } - return func() { <-r.slots; <-entry.slot; dropReference() }, true -} - // Run is synchronous with the SSE request. Cancellation stops both queued and // in-flight targets; detached, short writes retain their final observations. func (r *CodexGatewayBorrowTestRunner) Run(parent context.Context, task *CodexGatewayBorrowTestTask, emit func(CodexGatewayBorrowTestEvent)) error { r.runMu.Lock() + if task.Replayed || task.Status != "pending" || r.activeTasks[task.ID] { + r.runMu.Unlock() + return nil + } if r.stopped { r.runMu.Unlock() return r.finishUnstarted(parent, task, "cancelled", "test runner stopped", emit) } r.running.Add(1) + r.activeTasks[task.ID] = true r.runMu.Unlock() - defer r.running.Done() - ctx, cancel := context.WithCancel(parent) + defer func() { + r.runMu.Lock() + delete(r.activeTasks, task.ID) + r.runMu.Unlock() + r.running.Done() + }() + ctx, cancel := context.WithDeadline(parent, task.ExpiresAt) defer cancel() stopLink := context.AfterFunc(r.stopCtx, cancel) defer stopLink() @@ -352,6 +364,9 @@ func (r *CodexGatewayBorrowTestRunner) Run(parent context.Context, task *CodexGa return ErrCodexGatewayBorrowTestExpired } if err := r.repo.StartTask(ctx, task.ID, started); err != nil { + if errors.Is(err, ErrCodexGatewayBorrowTestNotFound) { + return err + } status := "incomplete" if ctx.Err() != nil { status = "cancelled" @@ -367,13 +382,48 @@ func (r *CodexGatewayBorrowTestRunner) Run(parent context.Context, task *CodexGa wg.Add(1) go func() { defer wg.Done() - if release, ok := r.acquire(ctx, result.AccountID); ok { - defer release() - r.runOne(ctx, result, emit) + budget := task.GenerationTimeoutSeconds + if budget == 0 { + budget = PelicanDefaultGenerationTimeoutSeconds + } + state := newPelicanExecutionState(ctx, r.coordinator, time.Duration(budget)*time.Second, r.now) + state.withTimeout = r.generationContext + state.onPhase = func(phase string, snapshot pelicanExecutionSnapshot) { + if emit != nil { + copy := *result + applyPelicanExecutionSnapshot(©, snapshot) + emit(CodexGatewayBorrowTestEvent{Type: "result_phase", Phase: phase, TaskID: task.ID, Result: ©}) + } + } + testCtx := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + if emit != nil { + copy := *result + emit(CodexGatewayBorrowTestEvent{Type: "result_phase", Phase: "queued", TaskID: task.ID, Result: ©}) + } + if task.ExecutionMode == PelicanExecutionModeAccount { + r.runOne(testCtx, result, r.standaloneGenerator, false, emit) } else { - finished := r.now().UTC() - result.Status, result.Error, result.FinishedAt = "cancelled", "test cancelled before dispatch", &finished - result.PreviewUnavailable = "test cancelled" + // The existing cache-only generator never prepares nested routes. + // Its outer lease also covers legacy callers without sender hooks. + leasedCtx, release, acquireErr := AcquirePelicanExecution(testCtx, result.AccountID) + if acquireErr == nil { + r.runOne(leasedCtx, result, r.generator, true, emit) + release() + } else { + finished := r.now().UTC() + result.Status, result.Error, result.FinishedAt = "cancelled", "test cancelled before dispatch", &finished + result.PreviewUnavailable = "test cancelled" + } + } + snapshot := state.finish() + applyPelicanExecutionSnapshot(result, snapshot) + if snapshot.generationTimedOut && ctx.Err() == nil && result.Status != "skipped" { + result.Status = "incomplete" + result.PreviewUnavailable = "test incomplete" + if result.Error != "" { + result.Error += "\n" + } + result.Error += "generation budget exhausted" } writeCtx, writeCancel := context.WithTimeout(context.WithoutCancel(parent), 10*time.Second) writeErr := r.repo.SaveResult(writeCtx, result) @@ -408,7 +458,8 @@ func (r *CodexGatewayBorrowTestRunner) Run(parent context.Context, task *CodexGa return nil } -func (r *CodexGatewayBorrowTestRunner) runOne(ctx context.Context, result *CodexGatewayBorrowTestResult, emit func(CodexGatewayBorrowTestEvent)) { +func (r *CodexGatewayBorrowTestRunner) runOne(ctx context.Context, result *CodexGatewayBorrowTestResult, generator codexGatewayBorrowPelicanGenerator, legacy bool, emit func(CodexGatewayBorrowTestEvent)) { + callerCtx := ctx now := r.now().UTC() if !result.ExpiresAt.After(now) { result.Status, result.Error, result.FinishedAt = "skipped", "test record expired before dispatch", &now @@ -432,11 +483,18 @@ func (r *CodexGatewayBorrowTestRunner) runOne(ctx context.Context, result *Codex result.PreviewUnavailable = "test cancelled" return } - testCtx, cancel := context.WithTimeout(ctx, codexGatewayBorrowTestTimeout) - generated, err := r.generator.GeneratePelican(testCtx, result.AccountID, result.ModelID, result.Effort) - cancel() + BeginPelicanPreparation(ctx) + if generator == nil { + result.Status, result.Error, result.FinishedAt = "skipped", "account Pelican generator is unavailable", &now + result.PreviewUnavailable = "test skipped" + return + } + if legacy { + ctx = BeginPelicanGeneration(ctx) + } + generated, err := generator.GeneratePelican(ctx, result.AccountID, result.ModelID, result.Effort) finished := r.now().UTC() - result.FinishedAt, result.DurationMS = &finished, finished.Sub(now).Milliseconds() + result.FinishedAt, result.DurationMS = &finished, nonnegativePelicanDuration(finished.Sub(now)).Milliseconds() if generated != nil { result.RawAnswer, result.RawResponse, result.Error = generated.RawAnswer, generated.RawResponse, generated.Error result.UpstreamModel = generated.ModelID @@ -454,11 +512,24 @@ func (r *CodexGatewayBorrowTestRunner) runOne(ctx context.Context, result *Codex if result.Status != "skipped" && result.Status != "incomplete" && result.Status != "failed" { result.Status = "failed" } + if errors.Is(err, context.DeadlineExceeded) && result.Status != "skipped" { + result.Status = "incomplete" + } } - if ctx.Err() != nil { + if callerCtx.Err() != nil { result.Status = "cancelled" + if errors.Is(context.Cause(callerCtx), context.DeadlineExceeded) { + result.Status = "incomplete" + } + if cause := context.Cause(callerCtx); errors.Is(cause, ErrObservedAccountLeaseLost) { + result.Status = "incomplete" + if result.Error != "" { + result.Error += "\n" + } + result.Error += cause.Error() + } if result.Error == "" { - result.Error = ctx.Err().Error() + result.Error = callerCtx.Err().Error() } } switch result.Status { diff --git a/backend/internal/service/codex_gateway_borrow_test_runner_test.go b/backend/internal/service/codex_gateway_borrow_test_runner_test.go index 7a027b5c8579..cc66411f4843 100644 --- a/backend/internal/service/codex_gateway_borrow_test_runner_test.go +++ b/backend/internal/service/codex_gateway_borrow_test_runner_test.go @@ -4,6 +4,7 @@ package service import ( "context" + "encoding/json" "errors" "sync" "testing" @@ -134,6 +135,7 @@ func TestCodexGatewayBorrowRunnerPersistsBeforeDispatchDedupsAndReplays(t *testi require.NoError(t, err) require.True(t, replayed) require.Equal(t, task.ID, replayedTask.ID) + require.NoError(t, runner.Run(context.Background(), replayedTask, nil)) require.Equal(t, 1, calls, "reading an existing task must never regenerate") request.Targets[0].Effort = "medium" _, _, err = runner.Create(context.Background(), 7, request) @@ -174,7 +176,7 @@ func TestCodexGatewayBorrowRunnerRejectsExpiredV7AfterRowsAreGone(t *testing.T) require.ErrorIs(t, err, ErrCodexGatewayBorrowTestInvalidRequest) } -func TestCodexGatewayBorrowRunnerSharesThreeSlotsAndSerializesAccountsAcrossBatches(t *testing.T) { +func TestCodexGatewayBorrowRunnerSharesTenSlotsAndSerializesAccountsAcrossBatches(t *testing.T) { repo := newCodexGatewayBorrowMemoryTestRepo() var mu sync.Mutex active, maxActive, calls := 0, 0, 0 @@ -207,24 +209,30 @@ func TestCodexGatewayBorrowRunnerSharesThreeSlotsAndSerializesAccountsAcrossBatc }} runner := newCodexGatewayBorrowTestRunner(repo, gen, nil) t.Cleanup(runner.Stop) - first, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "gpt-6-astra"}, {AccountID: 2, ModelID: "gpt-6-astra"}, {AccountID: 3, ModelID: "gpt-6-astra"}}}) + firstTargets, secondTargets := make([]CodexGatewayBorrowTestTarget, 7), make([]CodexGatewayBorrowTestTarget, 7) + for i := range firstTargets { + firstTargets[i] = CodexGatewayBorrowTestTarget{AccountID: int64(i + 1), ModelID: "gpt-6-astra"} + secondTargets[i] = CodexGatewayBorrowTestTarget{AccountID: int64(i + 7), ModelID: "gpt-6.1-sol"} + } + secondTargets[0].AccountID = 1 + first, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Targets: firstTargets}) require.NoError(t, err) - second, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "gpt-6.1-sol"}, {AccountID: 4, ModelID: "gpt-6-astra"}, {AccountID: 5, ModelID: "gpt-6-astra"}}}) + second, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Targets: secondTargets}) require.NoError(t, err) done := make(chan error, 2) go func() { done <- runner.Run(context.Background(), first, nil) }() go func() { done <- runner.Run(context.Background(), second, nil) }() - for i := 0; i < 3; i++ { + for i := 0; i < PelicanExecutionConcurrency; i++ { select { case <-started: case <-time.After(2 * time.Second): - t.Fatal("three dispatches did not start") + t.Fatal("ten dispatches did not start") } } mu.Lock() activeBeforeRelease := active mu.Unlock() - require.Equal(t, 3, activeBeforeRelease) + require.Equal(t, PelicanExecutionConcurrency, activeBeforeRelease) close(gate) for i := 0; i < 2; i++ { select { @@ -236,8 +244,8 @@ func TestCodexGatewayBorrowRunnerSharesThreeSlotsAndSerializesAccountsAcrossBatc } mu.Lock() defer mu.Unlock() - require.Equal(t, 6, calls) - require.Equal(t, 3, maxActive) + require.Equal(t, 14, calls) + require.Equal(t, PelicanExecutionConcurrency, maxActive) for _, max := range maxAccountActive { require.Equal(t, 1, max) } @@ -247,7 +255,7 @@ func TestCodexGatewayBorrowRunnerSharesThreeSlotsAndSerializesAccountsAcrossBatc func TestCodexGatewayBorrowRunnerCancellationStopsQueuedAndInflightAndRetainsRawOutput(t *testing.T) { repo := newCodexGatewayBorrowMemoryTestRepo() - started := make(chan struct{}, 4) + started := make(chan struct{}, PelicanExecutionConcurrency+1) gen := codexGatewayBorrowFakeGenerator{generate: func(ctx context.Context, _ int64, _, _ string) (*CodexGatewayBorrowPelicanResult, error) { started <- struct{}{} <-ctx.Done() @@ -255,7 +263,7 @@ func TestCodexGatewayBorrowRunnerCancellationStopsQueuedAndInflightAndRetainsRaw }} runner := newCodexGatewayBorrowTestRunner(repo, gen, nil) t.Cleanup(runner.Stop) - targets := make([]CodexGatewayBorrowTestTarget, 4) + targets := make([]CodexGatewayBorrowTestTarget, PelicanExecutionConcurrency+1) for i := range targets { targets[i] = CodexGatewayBorrowTestTarget{AccountID: int64(i + 1), ModelID: "gpt-6-astra"} } @@ -264,7 +272,7 @@ func TestCodexGatewayBorrowRunnerCancellationStopsQueuedAndInflightAndRetainsRaw ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- runner.Run(ctx, task, nil) }() - for i := 0; i < 3; i++ { + for i := 0; i < PelicanExecutionConcurrency; i++ { select { case <-started: case <-time.After(2 * time.Second): @@ -290,7 +298,287 @@ func TestCodexGatewayBorrowRunnerCancellationStopsQueuedAndInflightAndRetainsRaw require.Contains(t, result.Error, "full upstream failure detail") } } - require.Equal(t, 3, withPartial) + require.Equal(t, PelicanExecutionConcurrency, withPartial) +} + +func TestPelicanRequestBudgetBoundsDefaultsAndServerMode(t *testing.T) { + base := CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Standalone: true, + Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "local-model"}}} + for _, seconds := range []int{0, 60, 600, 1800} { + request := base + request.GenerationTimeoutSeconds = seconds + normalized, err := normalizeCodexGatewayBorrowTestRequest(request, time.Now().UTC()) + require.NoError(t, err) + if seconds == 0 { + require.Equal(t, PelicanDefaultGenerationTimeoutSeconds, normalized.GenerationTimeoutSeconds) + } else { + require.Equal(t, seconds, normalized.GenerationTimeoutSeconds) + } + require.Empty(t, normalized.Targets[0].Effort, "models without high must retain their normal default") + } + for _, seconds := range []int{-1, 1, 59, 1801} { + request := base + request.GenerationTimeoutSeconds = seconds + _, err := normalizeCodexGatewayBorrowTestRequest(request, time.Now().UTC()) + require.ErrorIs(t, err, ErrCodexGatewayBorrowTestInvalidRequest) + } + var decoded CodexGatewayBorrowTestRequest + require.NoError(t, json.Unmarshal([]byte(`{"standalone":true,"client_task_id":"ignored","generation_timeout_seconds":600}`), &decoded)) + require.False(t, decoded.Standalone, "the legacy endpoint cannot be promoted to account mode by request JSON") +} + +func TestPelicanRunnerReplayIncludesBudgetAndExecutionMode(t *testing.T) { + repo := newCodexGatewayBorrowMemoryTestRepo() + runner := newCodexGatewayBorrowTestRunner(repo, nil, nil) + t.Cleanup(runner.Stop) + request := CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Standalone: true, + Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "gpt-6-astra", Effort: "high"}}} + task, replayed, err := runner.Create(context.Background(), 1, request) + require.NoError(t, err) + require.False(t, replayed) + require.Equal(t, PelicanExecutionModeAccount, task.ExecutionMode) + require.Equal(t, PelicanDefaultGenerationTimeoutSeconds, task.GenerationTimeoutSeconds) + request.GenerationTimeoutSeconds = 60 + _, replayed, err = runner.Create(context.Background(), 1, request) + require.True(t, replayed) + require.ErrorIs(t, err, ErrCodexGatewayBorrowTestReplayConflict) + request.GenerationTimeoutSeconds, request.Standalone = 0, false + _, replayed, err = runner.Create(context.Background(), 1, request) + require.True(t, replayed) + require.ErrorIs(t, err, ErrCodexGatewayBorrowTestReplayConflict) +} + +func TestPelicanRunnerGenerationClockStartsAfterPreparationAndRecordsActualInvocation(t *testing.T) { + repo := newCodexGatewayBorrowMemoryTestRepo() + runner := newCodexGatewayBorrowTestRunner(repo, nil, nil) + t.Cleanup(runner.Stop) + now := time.Now().UTC().Add(time.Second) + runner.now = func() time.Time { return now } + budgetStarts := 0 + runner.generationContext = func(ctx context.Context, budget time.Duration) (context.Context, context.CancelFunc) { + budgetStarts++ + require.Equal(t, 10*time.Minute, budget) + return context.WithCancel(ctx) + } + runner.SetStandaloneGenerator(codexGatewayBorrowFakeGenerator{generate: func(ctx context.Context, id int64, _, _ string) (*CodexGatewayBorrowPelicanResult, error) { + require.Zero(t, budgetStarts, "neither queue nor preparation may start generation's timer") + now = now.Add(2 * time.Minute) + requestCtx, release, err := AcquirePelicanExecution(ctx, id) + require.NoError(t, err) + defer release() + finishBusinessWait := BeginPelicanQueueWait(requestCtx) + now = now.Add(10 * time.Second) + finishBusinessWait() + finishBusinessWait() + requestCtx = BeginPelicanGeneration(requestCtx) + require.Equal(t, 1, budgetStarts) + RecordPelicanInvocation(requestCtx, PelicanInvocation{Platform: "openai", Model: "gpt-6.1-sol", Effort: "high", + Endpoint: "https://upstream.test/backend-api/codex/responses", Protocol: "responses", Transport: "http", BorrowApplied: true}) + now = now.Add(30 * time.Second) + _ = BeginPelicanGeneration(requestCtx) + require.Equal(t, 1, budgetStarts, "same-account compatibility retries use the original generation clock") + now = now.Add(5 * time.Second) + return &CodexGatewayBorrowPelicanResult{Status: "complete", ModelID: "response-reported-model", RawAnswer: "original", RawResponse: "original response\x00tail"}, nil + }}) + task, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Standalone: true, + Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "mapped-alias"}}}) + require.NoError(t, err) + var phases []string + require.NoError(t, runner.Run(context.Background(), task, func(event CodexGatewayBorrowTestEvent) { + if event.Type == "result_phase" { + phases = append(phases, event.Phase) + } + })) + result := task.Results[0] + require.Equal(t, []string{"queued", "preparing", "generating"}, phases) + require.Equal(t, "complete", result.Status) + require.Equal(t, "mapped-alias", result.ModelID) + require.Equal(t, "gpt-6.1-sol", result.UpstreamModel) + require.Equal(t, "openai", result.Platform) + require.Equal(t, "high", result.Effort) + require.Equal(t, "responses", result.ActualProtocol) + require.Equal(t, "http", result.ActualTransport) + require.Equal(t, "https://upstream.test/backend-api/codex/responses", result.ActualEndpoint) + require.True(t, result.BorrowApplied) + require.Equal(t, int64(10000), result.QueueDurationMS) + require.Equal(t, int64(120000), result.PreparationDurationMS) + require.Equal(t, int64(35000), result.GenerationDurationMS) + require.Equal(t, int64(165000), result.DurationMS) + require.NotNil(t, result.GenerationStartedAt) + require.Equal(t, "original response\x00tail", result.RawResponse) +} + +func TestPelicanRunnerNestedPreparationDoesNotReserveAllExecutionSlots(t *testing.T) { + repo := newCodexGatewayBorrowMemoryTestRepo() + runner := newCodexGatewayBorrowTestRunner(repo, nil, nil) + t.Cleanup(runner.Stop) + preparing := make(chan struct{}, PelicanExecutionConcurrency+1) + allowPreparation := make(chan struct{}) + var mu sync.Mutex + active, maxActive, probes, generations := 0, 0, 0, 0 + accountActive := map[int64]int{} + maxAccountActive := map[int64]int{} + observe := func(ctx context.Context, id int64, generation bool) error { + requestCtx, release, err := AcquirePelicanExecution(ctx, id) + if err != nil { + return err + } + defer release() + // Sender wrappers for this exact account/request are reentrant and do + // not take another global permit or wait on their own account lease. + _, nestedRelease, err := AcquirePelicanExecution(requestCtx, id) + if err != nil { + return err + } + defer nestedRelease() + mu.Lock() + active++ + accountActive[id]++ + if active > maxActive { + maxActive = active + } + if accountActive[id] > maxAccountActive[id] { + maxAccountActive[id] = accountActive[id] + } + if generation { + generations++ + } else { + probes++ + } + accountActive[id]-- + active-- + mu.Unlock() + return nil + } + runner.SetStandaloneGenerator(codexGatewayBorrowFakeGenerator{generate: func(ctx context.Context, id int64, model, effort string) (*CodexGatewayBorrowPelicanResult, error) { + preparing <- struct{}{} + select { + case <-allowPreparation: + case <-ctx.Done(): + return nil, ctx.Err() + } + // Every target needs the same source; that source is also selected as + // its own Pelican target. None of these outer preparations may retain + // the final target's account/global lease while requesting the source. + if err := observe(ctx, 1, false); err != nil { + return nil, err + } + if err := observe(ctx, id, true); err != nil { + return nil, err + } + return &CodexGatewayBorrowPelicanResult{Status: "complete", ModelID: model, Effort: effort, RawAnswer: "synthetic"}, nil + }}) + targets := make([]CodexGatewayBorrowTestTarget, PelicanExecutionConcurrency+1) + for i := range targets { + targets[i] = CodexGatewayBorrowTestTarget{AccountID: int64(i + 1), ModelID: "local"} + } + task, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Standalone: true, Targets: targets}) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- runner.Run(ctx, task, nil) }() + for range targets { + select { + case <-preparing: + case <-time.After(2 * time.Second): + cancel() + t.Fatal("outer preparations reserved execution permits before their nested source requests") + } + } + close(allowPreparation) + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("source/target preparation deadlocked with the selected source's own Pelican generation") + } + mu.Lock() + defer mu.Unlock() + require.Equal(t, len(targets), probes) + require.Equal(t, len(targets), generations) + require.LessOrEqual(t, maxActive, PelicanExecutionConcurrency) + for _, maximum := range maxAccountActive { + require.Equal(t, 1, maximum) + } + require.Equal(t, "complete", task.Status) +} + +func TestPelicanExecutionIsScopedAndRejectsIncorrectCrossAccountNesting(t *testing.T) { + ordinary := context.WithValue(context.Background(), struct{}{}, "ordinary context") + ctx, release, err := AcquirePelicanExecution(ordinary, 1) + require.NoError(t, err) + require.Equal(t, ordinary, ctx) + release() + require.Equal(t, ordinary, BeginPelicanGeneration(ordinary)) + BeginPelicanPreparation(ordinary) + RecordPelicanInvocation(ordinary, PelicanInvocation{Model: "ignored"}) + BeginPelicanQueueWait(ordinary)() + coordinator := newPelicanExecutionCoordinator(1) + state := newPelicanExecutionState(context.Background(), coordinator, 10*time.Minute, time.Now) + ctx = context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + leasing, release, err := AcquirePelicanExecution(ctx, 1) + require.NoError(t, err) + _, nestedRelease, err := AcquirePelicanExecution(leasing, 1) + require.NoError(t, err) + nestedRelease() + coordinator.mu.Lock() + active := coordinator.active + coordinator.mu.Unlock() + require.Equal(t, 1, active) + _, _, err = AcquirePelicanExecution(leasing, 2) + require.ErrorContains(t, err, "preparation must finish") + release() + release() + coordinator.mu.Lock() + require.Zero(t, coordinator.active) + require.Empty(t, coordinator.accounts) + require.Empty(t, coordinator.waiting) + coordinator.mu.Unlock() + state.finish() +} + +func TestPelicanRunnerGenerationTimeoutRetainsIncompleteOutput(t *testing.T) { + repo := newCodexGatewayBorrowMemoryTestRepo() + runner := newCodexGatewayBorrowTestRunner(repo, nil, nil) + t.Cleanup(runner.Stop) + now := time.Now().UTC().Add(time.Second) + runner.now = func() time.Time { return now } + var timeout context.CancelCauseFunc + runner.generationContext = func(ctx context.Context, budget time.Duration) (context.Context, context.CancelFunc) { + require.Equal(t, time.Minute, budget) + budgetCtx, cancel := context.WithCancelCause(ctx) + timeout = cancel + return budgetCtx, func() { cancel(context.Canceled) } + } + runner.SetStandaloneGenerator(codexGatewayBorrowFakeGenerator{generate: func(ctx context.Context, id int64, _, _ string) (*CodexGatewayBorrowPelicanResult, error) { + now = now.Add(90 * time.Second) + requestCtx, release, err := AcquirePelicanExecution(ctx, id) + require.NoError(t, err) + defer release() + requestCtx = BeginPelicanGeneration(requestCtx) + now = now.Add(time.Minute) + timeout(context.DeadlineExceeded) + <-requestCtx.Done() + <-ctx.Done() + require.ErrorIs(t, context.Cause(ctx), context.DeadlineExceeded, "the entire generator, including retries and later preparation, inherits the budget") + _, _, err = AcquirePelicanExecution(ctx, id+1) + require.ErrorIs(t, err, context.Canceled, "budget expiry cannot start another probe or compatibility retry") + return &CodexGatewayBorrowPelicanResult{Status: "incomplete", RawAnswer: "partial", RawResponse: "original partial\x00tail"}, requestCtx.Err() + }}) + task, _, err := runner.Create(context.Background(), 1, CodexGatewayBorrowTestRequest{ClientTaskID: codexGatewayBorrowV7(t), Standalone: true, GenerationTimeoutSeconds: 60, + Targets: []CodexGatewayBorrowTestTarget{{AccountID: 1, ModelID: "local"}}}) + require.NoError(t, err) + require.NoError(t, runner.Run(context.Background(), task, nil)) + require.Equal(t, "incomplete", task.Status) + result := task.Results[0] + require.Equal(t, "incomplete", result.Status) + require.Contains(t, result.Error, "generation budget exhausted") + require.Equal(t, "partial", result.RawAnswer) + require.Equal(t, "original partial\x00tail", result.RawResponse) + require.Equal(t, int64(90000), result.PreparationDurationMS) + require.Equal(t, int64(60000), result.GenerationDurationMS) + require.NotEmpty(t, result.PreviewUnavailable) } func TestCodexGatewayBorrowRunnerDoesNotDispatchWhenRunningRecordCannotPersist(t *testing.T) { diff --git a/backend/internal/service/codex_gateway_borrow_ws_anchor.go b/backend/internal/service/codex_gateway_borrow_ws_anchor.go index 62984ae8db30..c1bdc5c07283 100644 --- a/backend/internal/service/codex_gateway_borrow_ws_anchor.go +++ b/backend/internal/service/codex_gateway_borrow_ws_anchor.go @@ -46,6 +46,7 @@ type codexGatewayBorrowWSTurn struct { headers http.Header qualified bool finished sync.Once + pelicanOneShot bool } // Passthrough owns a direct socket for the entire client session. Its anchor is @@ -227,6 +228,26 @@ func (s *OpenAIGatewayService) prepareCodexGatewayBorrowWSTurn(ctx context.Conte return nil, err } selected := s.codexGatewayBorrowWSSelected(account, model, req) + if IsPelicanGeneration(ctx) { + if !selected { + return nil, nil + } + if previousID != "" { + return nil, errors.New("pelican generation starts a new conversation") + } + // A trusted administrator invocation has no customer's API key or + // continuation identity. It borrows through the normal preparation path + // but never creates a persistent client continuation anchor. + borrowed, application, applyErr := s.gatewayBorrow.Apply(req, account, model, proxyURL, nil, false) + if applyErr != nil { + return nil, applyErr + } + if borrowed == nil || application == nil || !application.Applied { + return nil, &CodexGatewayBorrowFailure{Cause: ErrCodexGatewayBorrowChanged} + } + return &codexGatewayBorrowWSTurn{key: codexGatewayBorrowWSAnchorKey{account: account.ID, model: model}, + model: model, headers: cloneHeader(borrowed.Header), expires: application.ExpiresAt, pelicanOneShot: true}, nil + } key, err := codexGatewayBorrowWSKey(c, account, headers, rawClientBody, model, proxyURL) if err != nil { if !selected { @@ -312,6 +333,9 @@ func (turn *codexGatewayBorrowWSTurn) bindLease(lease *openAIWSConnLease) error return NewOpenAIContinuationStateUnavailableError(http.StatusBadRequest, nil, nil) } turn.connID = lease.ConnID() + if turn.pelicanOneShot { + return nil + } lease.conn.anchorUntilNano.Store(turn.expires.UnixNano()) return nil } @@ -326,7 +350,18 @@ func (turn *codexGatewayBorrowWSTurn) observe(message []byte) { } func (turn *codexGatewayBorrowWSTurn) finish(pool *openAIWSConnPool, result *OpenAIForwardResult, resultErr error) { - if turn == nil || turn.store == nil { + if turn == nil { + return + } + if turn.pelicanOneShot { + turn.finished.Do(func() { + if pool != nil && turn.connID != "" { + pool.evictConn(turn.key.account, turn.connID) + } + }) + return + } + if turn.store == nil { return } turn.finished.Do(func() { diff --git a/backend/internal/service/concurrency_service.go b/backend/internal/service/concurrency_service.go index df18379704bf..1b92aea297e6 100644 --- a/backend/internal/service/concurrency_service.go +++ b/backend/internal/service/concurrency_service.go @@ -7,6 +7,7 @@ import ( "encoding/binary" "encoding/hex" "errors" + "fmt" "os" "strconv" "sync" @@ -61,6 +62,160 @@ type APIKeyConcurrencyCache interface { GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) } +// AccountSlotRefreshCache is an optional extension for explicitly retained +// administrator observations. Refresh must only renew an existing request ID; +// recreating a lost member could exceed ordinary business account capacity. +type AccountSlotRefreshCache interface { + RefreshAccountSlot(ctx context.Context, accountID int64, requestID string) (bool, error) +} + +const ( + // The configured regular-slot TTL is an integer number of minutes, at + // least one minute. Thirty seconds also covers that minimum configuration. + observedAccountLeaseRefreshInterval = 30 * time.Second + observedAccountLeaseOperationTO = 2 * time.Second +) + +var ErrObservedAccountLeaseLost = errors.New("observed account concurrency lease lost") + +// ObservedAccountLease occupies the same regular account slot as a business +// request, while renewing it for an explicitly long observation. Normal +// AcquireAccountSlot callers do not create a timer or change their behavior. +type ObservedAccountLease struct { + ctx context.Context + cancel context.CancelCauseFunc + cache ConcurrencyCache + refresher AccountSlotRefreshCache + accountID int64 + requestID string + onLost func(error) + ticks <-chan time.Time + stopTimer func() + stopCh chan struct{} + done chan struct{} + + mu sync.Mutex + released bool + stopOnce sync.Once + releaseOnce sync.Once +} + +func (l *ObservedAccountLease) Context() context.Context { + if l == nil || l.ctx == nil { + return context.Background() + } + return l.ctx +} + +func (l *ObservedAccountLease) Release() { + if l == nil { + return + } + l.stopOnce.Do(func() { + l.mu.Lock() + l.released = true + l.mu.Unlock() + if l.stopCh != nil { + close(l.stopCh) + } + if l.cancel != nil { + l.cancel(context.Canceled) + } + if l.done != nil { + <-l.done + } + l.releaseSlot() + }) +} + +func (l *ObservedAccountLease) releaseSlot() { + l.releaseOnce.Do(func() { + if l.cache == nil || l.accountID <= 0 || l.requestID == "" { + return + } + ctx, cancel := context.WithTimeout(context.Background(), observedAccountLeaseOperationTO) + defer cancel() + if err := l.cache.ReleaseAccountSlot(ctx, l.accountID, l.requestID); err != nil { + logger.L().Warn("observed_account_lease_release_failed", zap.Int64("account_id", l.accountID), zap.Error(err)) + } + }) +} + +func (l *ObservedAccountLease) refreshLoop() { + defer close(l.done) + defer l.stopTimer() + defer l.releaseSlot() + for { + select { + case <-l.ctx.Done(): + return + case <-l.stopCh: + return + case <-l.ticks: + if l.ctx.Err() != nil { + return + } + ctx, cancel := context.WithTimeout(l.ctx, observedAccountLeaseOperationTO) + owned, err := l.refresher.RefreshAccountSlot(ctx, l.accountID, l.requestID) + cancel() + if err == nil && owned { + continue + } + // A failed renewal is fail-closed. Cancel before releasing capacity + // rather than continuing an unconfirmed 30-minute model request. + cause := ErrObservedAccountLeaseLost + if err != nil { + cause = fmt.Errorf("%w: %v", cause, err) + } + l.mu.Lock() + if !l.released && l.ctx.Err() == nil { + l.cancel(cause) + if l.onLost != nil { + l.onLost(cause) + } + } + l.mu.Unlock() + return + } + } +} + +func newObservedAccountLease(ctx context.Context, cache ConcurrencyCache, refresher AccountSlotRefreshCache, accountID int64, requestID string, + onLost func(error), ticks <-chan time.Time, stopTimer func()) *ObservedAccountLease { + leaseCtx, cancel := context.WithCancelCause(ctx) + lease := &ObservedAccountLease{ctx: leaseCtx, cancel: cancel, cache: cache, refresher: refresher, + accountID: accountID, requestID: requestID, onLost: onLost, ticks: ticks, stopTimer: stopTimer, + stopCh: make(chan struct{}), done: make(chan struct{})} + go lease.refreshLoop() + return lease +} + +// AcquireObservedAccountLease is opt-in for long administrator observations. +// Acquisition uses the unchanged account limit and slot namespace; the cache +// extension only keeps that exact owner alive until release/cancellation/loss. +func (s *ConcurrencyService) AcquireObservedAccountLease(ctx context.Context, accountID int64, maxConcurrency int, onLost func(error)) (*ObservedAccountLease, bool, error) { + if maxConcurrency <= 0 { + return nil, true, nil + } + if s == nil || s.cache == nil || accountID <= 0 { + return nil, false, errors.New("observed account concurrency cache is unavailable") + } + refresher, ok := s.cache.(AccountSlotRefreshCache) + if !ok { + return nil, false, errors.New("observed account concurrency lease renewal is unsupported") + } + if ctx == nil { + ctx = context.Background() + } + requestID := generateRequestID() + acquired, err := s.cache.AcquireAccountSlot(ctx, accountID, maxConcurrency, requestID) + if err != nil || !acquired { + return nil, acquired, err + } + ticker := time.NewTicker(observedAccountLeaseRefreshInterval) + return newObservedAccountLease(ctx, s.cache, refresher, accountID, requestID, onLost, ticker.C, ticker.Stop), true, nil +} + // OpenAIWSIngressLeaseCache owns the short-lived distributed lease used to // bound live client WebSocket sessions. It is deliberately independent of the // request-slot namespace: idle ingress connections do not occupy turn slots. diff --git a/backend/internal/service/gateway_observation.go b/backend/internal/service/gateway_observation.go new file mode 100644 index 000000000000..2b4e073b25a3 --- /dev/null +++ b/backend/internal/service/gateway_observation.go @@ -0,0 +1,34 @@ +package service + +import "context" + +type accountObservationContextKey struct{} + +// WithAccountObservation marks an administrator's account invocation as +// observation-only. Outbound requests, compatibility handling and credential +// refresh still run normally; account health, quota and billing observations +// must not be published by this invocation. +func WithAccountObservation(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, accountObservationContextKey{}, true) +} + +func IsAccountObservation(ctx context.Context) bool { + if ctx == nil { + return false + } + marked, _ := ctx.Value(accountObservationContextKey{}).(bool) + return marked +} + +// CopyAccountObservationContext preserves the mode when a callback needs a +// separate lifecycle. Copy before starting a goroutine or queueing a task: +// replacing a request context with Background alone loses the write boundary. +func CopyAccountObservationContext(parent, base context.Context) context.Context { + if IsAccountObservation(parent) { + return WithAccountObservation(base) + } + return base +} diff --git a/backend/internal/service/gateway_observation_test.go b/backend/internal/service/gateway_observation_test.go new file mode 100644 index 000000000000..77557cee2175 --- /dev/null +++ b/backend/internal/service/gateway_observation_test.go @@ -0,0 +1,418 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/stretchr/testify/require" +) + +type accountObservationRepo struct { + AccountRepository + errorCalls atomic.Int64 + tempCalls atomic.Int64 + rateCalls atomic.Int64 + modelCalls atomic.Int64 + extraCalls atomic.Int64 + clearedCalls atomic.Int64 + extraWritten chan struct{} +} + +func (r *accountObservationRepo) SetError(context.Context, int64, string) error { + r.errorCalls.Add(1) + return nil +} + +func (r *accountObservationRepo) SetTempUnschedulable(context.Context, int64, time.Time, string) error { + r.tempCalls.Add(1) + return nil +} + +func (r *accountObservationRepo) SetRateLimited(context.Context, int64, time.Time) error { + r.rateCalls.Add(1) + return nil +} + +func (r *accountObservationRepo) SetModelRateLimit(context.Context, int64, string, time.Time, ...string) error { + r.modelCalls.Add(1) + return nil +} + +func (r *accountObservationRepo) UpdateExtra(context.Context, int64, map[string]any) error { + r.extraCalls.Add(1) + if r.extraWritten != nil { + r.extraWritten <- struct{}{} + } + return nil +} + +func (r *accountObservationRepo) ClearRateLimitIfObserved(context.Context, int64, time.Time, time.Time) (bool, error) { + r.clearedCalls.Add(1) + return true, nil +} + +func requireNoAccountObservationWrites(t *testing.T, r *accountObservationRepo) { + t.Helper() + require.Zero(t, r.errorCalls.Load()) + require.Zero(t, r.tempCalls.Load()) + require.Zero(t, r.rateCalls.Load()) + require.Zero(t, r.modelCalls.Load()) + require.Zero(t, r.extraCalls.Load()) + require.Zero(t, r.clearedCalls.Load()) +} + +func TestAccountObservationBusinessForwardPreservesErrorsWithoutPublishingState(t *testing.T) { + for _, tc := range []struct { + name string + status int + body string + }{ + {"http_auth", http.StatusUnauthorized, `{"error":{"message":"synthetic token expired"}}`}, + {"http_budget", http.StatusTooManyRequests, `{"error":{"type":"budget_exceeded","message":"synthetic budget exhausted"}}`}, + {"stream_budget", http.StatusOK, "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_observation_error\",\"status\":\"failed\",\"error\":{\"type\":\"budget_exceeded\",\"message\":\"synthetic budget exhausted\"}}}\n\n"}, + } { + t.Run(tc.name, func(t *testing.T) { + for _, observation := range []bool{true, false} { + mode := "normal" + if observation { + mode = "observation" + } + t.Run(mode, func(t *testing.T) { + setup := newAstraOAuthSetup(t, false) + repo := &accountObservationRepo{} + setup.svc.accountRepo = repo + setup.svc.rateLimitService = NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + setup.svc.rateLimitService.SetAccountRuntimeBlocker(setup.svc) + setup.account.Credentials["refresh_token"] = "synthetic-refresh-token" + setup.upstream.resp = &http.Response{ + StatusCode: tc.status, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(tc.body)), + } + ctx := context.Background() + if observation { + ctx = WithAccountObservation(ctx) + } + setup.c.Request = setup.c.Request.WithContext(ctx) + _, err := setup.svc.Forward(ctx, setup.c, setup.account, astraRequestBody("gpt-6-astra", true, "", "high")) + require.Error(t, err) + require.Len(t, setup.upstream.requests, 1, "observation must send to this account exactly once") + var failure *UpstreamFailoverError + require.ErrorAs(t, err, &failure) + require.Contains(t, string(failure.ResponseBody), "synthetic") + if observation { + requireNoAccountObservationWrites(t, repo) + require.False(t, setup.svc.isOpenAIAccountRuntimeBlocked(setup.account)) + } else { + require.Positive(t, repo.errorCalls.Load()+repo.tempCalls.Load(), "ordinary forwarding retains its account-state contract") + } + }) + } + }) + } +} + +func TestAccountObservationSuccessSuppressesBackgroundQuotaAndProxyHealth(t *testing.T) { + setup := newAstraOAuthSetup(t, false) + repo := &accountObservationRepo{extraWritten: make(chan struct{}, 1)} + setup.svc.accountRepo = repo + setup.upstream.resp = &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": {"text/event-stream"}, + "X-Codex-Primary-Used-Percent": {"50"}, + "X-Codex-Primary-Window-Minutes": {"300"}, + }, + Body: io.NopCloser(strings.NewReader(codexCompletedSSE(`{"id":"resp_observation_success","status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"synthetic html"}]}],"usage":{"input_tokens":1,"output_tokens":1}}`))), + } + proxyID := int64(71) + setup.account.ProxyID = &proxyID + setup.svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{ + failureThreshold: 1, failureWindow: time.Minute, quarantineTTL: time.Hour, maxEntries: 2, + }) + setup.svc.openaiProxyStreamCircuit.recordFailure(proxyID, time.Now()) + ctx := WithAccountObservation(context.Background()) + setup.c.Request = setup.c.Request.WithContext(ctx) + result, err := setup.svc.Forward(ctx, setup.c, setup.account, astraRequestBody("gpt-6-astra", true, "", "high")) + require.NoError(t, err) + require.NotNil(t, result) + require.Contains(t, setup.rec.Body.String(), "synthetic html") + requireNoAccountObservationWrites(t, repo) + require.Nil(t, setup.svc.codexSnapshotThrottle, "suppression happens before throttle state and goroutine creation") + require.True(t, setup.svc.openaiProxyStreamCircuit.isBlocked(proxyID, time.Now()), "a diagnostic success cannot recover the business proxy") + + // Normal success remains an explicit positive control for the detached write. + snapshot := ParseCodexRateLimitHeaders(setup.upstream.resp.Header) + require.NotNil(t, snapshot) + setup.svc.updateCodexUsageSnapshot(context.Background(), setup.account.ID, snapshot) + select { + case <-repo.extraWritten: + case <-time.After(time.Second): + t.Fatal("ordinary success did not publish its quota snapshot") + } + require.EqualValues(t, 1, repo.extraCalls.Load()) + setup.svc.clearOpenAIProxyStreamDisconnectInContext(context.Background(), setup.account) + require.False(t, setup.svc.openaiProxyStreamCircuit.isBlocked(proxyID, time.Now())) +} + +func TestAccountObservationTokenRefreshPersistsCredentialsAcrossProviders(t *testing.T) { + for _, platform := range []string{PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformGrok, PlatformAntigravity} { + t.Run(platform, func(t *testing.T) { + account := &Account{ID: 91, Platform: platform, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, + Credentials: map[string]any{"access_token": "old-token", "refresh_token": "old-refresh", "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339)}, + } + repo := &refreshAPIAccountRepo{account: account} + executor := &refreshAPIExecutorStub{needsRefresh: true, credentials: map[string]any{ + "access_token": "new-token", "refresh_token": "new-refresh", "expires_at": time.Now().Add(time.Hour).Format(time.RFC3339), + }} + api := NewOAuthRefreshAPI(repo, nil) + var token string + var err error + ctx := WithAccountObservation(context.Background()) + switch platform { + case PlatformOpenAI: + provider := NewOpenAITokenProvider(repo, nil, nil) + provider.SetRefreshAPI(api, executor) + token, err = provider.GetAccessToken(ctx, account) + case PlatformAnthropic: + provider := NewClaudeTokenProvider(repo, nil, nil) + provider.SetRefreshAPI(api, executor) + token, err = provider.GetAccessToken(ctx, account) + case PlatformGemini: + provider := NewGeminiTokenProvider(repo, nil, nil) + provider.SetRefreshAPI(api, executor) + token, err = provider.GetAccessToken(ctx, account) + case PlatformGrok: + provider := NewGrokTokenProvider(repo, nil) + provider.SetRefreshAPI(api, executor) + token, err = provider.GetAccessToken(ctx, account) + case PlatformAntigravity: + provider := NewAntigravityTokenProvider(repo, nil, nil) + provider.SetRefreshAPI(api, executor) + token, err = provider.GetAccessToken(ctx, account) + } + require.NoError(t, err) + require.Equal(t, "new-token", token) + require.Equal(t, 1, executor.refreshCalls) + require.Equal(t, 1, repo.updateCredentialsCalls) + require.Equal(t, "new-refresh", repo.account.GetCredential("refresh_token")) + require.Equal(t, StatusActive, repo.account.Status) + }) + } +} + +type accountObservationTempCache struct { + TempUnschedCache + setCalls []int64 +} + +func (c *accountObservationTempCache) SetTempUnsched(_ context.Context, id int64, _ *TempUnschedState) error { + c.setCalls = append(c.setCalls, id) + return nil +} + +func TestAccountObservationProviderFailuresSkipDetachedStateWrites(t *testing.T) { + ctx := WithAccountObservation(context.Background()) + repo := &rateLimitAccountRepoStub{} + openAI := NewOpenAITokenProvider(repo, nil, nil) + blocker := &runtimeBlockRecorder{} + openAI.SetAccountRuntimeBlocker(blocker) + _, err := openAI.GetAccessToken(ctx, &Account{ID: 92, Platform: PlatformOpenAI, Type: AccountTypeOAuth, + Credentials: map[string]any{"access_token": "expired", "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339)}, + }) + require.ErrorContains(t, err, "refresh_token is missing") + require.Zero(t, repo.setErrorCalls) + require.Empty(t, blocker.accounts) + + account := &Account{ID: 93, Platform: PlatformAntigravity, Type: AccountTypeOAuth, Status: StatusActive, + Credentials: map[string]any{"access_token": "expired", "refresh_token": "synthetic-refresh", "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339)}, + } + refreshRepo := &refreshAPIAccountRepo{account: account} + executor := &refreshAPIExecutorStub{needsRefresh: true, err: errors.New("synthetic refresh failure")} + provider := NewAntigravityTokenProvider(repo, nil, nil) + provider.SetRefreshAPI(NewOAuthRefreshAPI(refreshRepo, nil), executor) + cache := &accountObservationTempCache{} + provider.SetTempUnschedCache(cache) + _, err = provider.GetAccessToken(ctx, account) + require.Error(t, err) + require.Zero(t, repo.tempCalls) + require.Empty(t, cache.setCalls) + _, err = provider.GetAccessToken(context.Background(), account) + require.Error(t, err) + require.Equal(t, 1, repo.tempCalls) + require.Len(t, cache.setCalls, 1, "ordinary refresh failure must retain its detached DB and cache writes") +} + +type accountObservationBillingRepo struct { + UsageBillingRepository + calls int +} + +func (r *accountObservationBillingRepo) Apply(context.Context, *UsageBillingCommand) (*UsageBillingApplyResult, error) { + r.calls++ + return &UsageBillingApplyResult{Applied: false}, nil +} + +type accountObservationLogRepo struct { + UsageLogRepository + calls int +} + +func (r *accountObservationLogRepo) Create(context.Context, *UsageLog) (bool, error) { + r.calls++ + return true, nil +} + +func TestAccountObservationQueuedUsagePreservesMarkerAndDoesNotWrite(t *testing.T) { + store := "aActivityMemoryStore{} + activity := NewQuotaActivityService(store) + ctx, finish := activity.Attach(WithAccountObservation(context.WithValue(context.Background(), ctxkey.AccountID, int64(95)))) + ObserveQuotaAccount(ctx, 95) + var ran bool + task, discard := TrackQuotaUsageTask(ctx, func(workerCtx context.Context) { + ran = true + require.True(t, IsAccountObservation(workerCtx)) + require.NoError(t, (&GatewayService{}).RecordUsage(workerCtx, nil)) + require.NoError(t, (&OpenAIGatewayService{}).RecordUsage(workerCtx, nil)) + MarkQuotaLogPersisted(workerCtx, 95) + }) + finish() + task(context.Background()) + discard() + require.True(t, ran) + stamp, known := activity.Read(context.Background(), 95) + require.True(t, known) + require.Equal(t, QuotaActivityStamp{}, stamp) + + billing := &accountObservationBillingRepo{} + deferred := &DeferredService{} + params := &postUsageBillingParams{User: &User{ID: 1}, APIKey: &APIKey{ID: 2}, Account: &Account{ID: 95, Type: AccountTypeAPIKey}, Cost: &CostBreakdown{ActualCost: 1}} + applied, err := applyUsageBilling(CopyAccountObservationContext(ctx, context.Background()), "synthetic-bill", nil, params, &billingDeps{deferredService: deferred}, billing) + require.NoError(t, err) + require.False(t, applied) + require.Zero(t, billing.calls) + _, scheduled := deferred.lastUsedUpdates.Load(int64(95)) + require.False(t, scheduled) + _, err = applyUsageBilling(context.Background(), "synthetic-bill", nil, params, &billingDeps{deferredService: deferred}, billing) + require.NoError(t, err) + require.Equal(t, 1, billing.calls) + _, scheduled = deferred.lastUsedUpdates.Load(int64(95)) + require.True(t, scheduled) + + logs := &accountObservationLogRepo{} + require.False(t, writeUsageLogBestEffort(ctx, logs, &UsageLog{AccountID: 95}, "observation-test")) + require.Zero(t, logs.calls) + require.True(t, writeUsageLogBestEffort(context.Background(), logs, &UsageLog{AccountID: 95}, "observation-test")) + require.Equal(t, 1, logs.calls) +} + +func TestAccountObservationTimeoutRecoveryAndAntigravityRetryBoundaries(t *testing.T) { + ctx := WithAccountObservation(context.Background()) + repo := &accountObservationRepo{} + rateLimit := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ID: 96, Platform: PlatformAntigravity, Type: AccountTypeOAuth, + Credentials: map[string]any{ + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{map[string]any{ + "error_code": float64(http.StatusInternalServerError), + "keywords": []any{"synthetic failure"}, + "duration_minutes": float64(1), + }}, + }, + } + require.Equal(t, ErrorPolicyTempUnscheduled, rateLimit.CheckErrorPolicy(ctx, account, http.StatusInternalServerError, []byte(`{"error":{"message":"synthetic failure"}}`), "synthetic-model"), "request-local policy matching survives observation isolation") + require.False(t, rateLimit.HandleStreamTimeout(ctx, account, "synthetic-model")) + result, err := rateLimit.RecoverAccountState(ctx, account.ID, AccountRecoveryOptions{InvalidateToken: true}) + require.NoError(t, err) + require.Equal(t, &SuccessfulTestRecoveryResult{}, result) + require.NoError(t, rateLimit.ClearRateLimit(ctx, account.ID)) + require.NoError(t, rateLimit.ClearTempUnschedulable(ctx, account.ID)) + tempUnscheduleAccountForTransportError(ctx, repo, account, "synthetic transport failure") + tempUnscheduleGoogleConfigError(ctx, repo, account.ID, "observation-test") + tempUnscheduleEmptyResponse(ctx, repo, account.ID, "observation-test") + require.False(t, setModelRateLimitByModelName(ctx, repo, account.ID, "synthetic-model", "observation-test", 429, time.Now().Add(time.Hour), true)) + internalCounter := &mockInternal500Cache{incrementCount: 3} + antigravity := &AntigravityGatewayService{accountRepo: repo, internal500Cache: internalCounter} + antigravity.handleInternal500RetryExhausted(ctx, "observation-test", account) + antigravity.resetInternal500Counter(ctx, "observation-test", account.ID) + antigravity.setCreditsExhausted(ctx, account) + require.Empty(t, internalCounter.incrementCalls) + require.Empty(t, internalCounter.resetCalls) + requireNoAccountObservationWrites(t, repo) + antigravity.handleInternal500RetryExhausted(context.Background(), "observation-test", account) + require.EqualValues(t, 1, repo.errorCalls.Load()) +} + +func TestAccountObservationWebSocketFailureCallbacksDoNotPublishCooldown(t *testing.T) { + ctx := WithAccountObservation(context.Background()) + repo := &accountObservationRepo{} + gateway := &OpenAIGatewayService{accountRepo: repo, cfg: &config.Config{}} + gateway.rateLimitService = NewRateLimitService(repo, nil, gateway.cfg, nil, nil) + gateway.rateLimitService.SetAccountRuntimeBlocker(gateway) + account := &Account{ID: 97, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + payload := []byte(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"rate_limit_error","message":"synthetic rate limit exceeded"}}`) + require.True(t, gateway.handleOpenAIWSFailureAccountSideEffects(ctx, account, "gpt-6.1-sol", nil, payload)) + require.False(t, gateway.isOpenAIAccountRuntimeBlocked(account)) + gateway.persistOpenAIWSRateLimitSignal(ctx, account, nil, payload, "rate_limit_exceeded", "rate_limit_error", "synthetic rate limit exceeded", "gpt-6.1-sol") + requireNoAccountObservationWrites(t, repo) + require.True(t, gateway.handleOpenAIWSFailureAccountSideEffects(context.Background(), account, "gpt-6.1-sol", nil, payload)) + require.True(t, gateway.isOpenAIAccountRuntimeBlocked(account), "ordinary WS failure retains its retry-exhausted cooldown") +} + +type accountObservationBlockingUpstream struct { + entered chan context.Context +} + +func (u *accountObservationBlockingUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + u.entered <- req.Context() + <-req.Context().Done() + return nil, req.Context().Err() +} + +func (u *accountObservationBlockingUpstream) DoWithTLS(req *http.Request, proxy string, id int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.Do(req, proxy, id, concurrency) +} + +func TestAccountObservationCancellationReachesBusinessUpstream(t *testing.T) { + setup := newAstraOAuthSetup(t, false) + upstream := &accountObservationBlockingUpstream{entered: make(chan context.Context, 1)} + setup.svc.httpUpstream = upstream + ctx, cancel := context.WithCancel(WithAccountObservation(context.Background())) + defer cancel() + setup.c.Request = setup.c.Request.WithContext(ctx) + done := make(chan error, 1) + go func() { + _, err := setup.svc.Forward(ctx, setup.c, setup.account, astraRequestBody("gpt-6-astra", true, "", "high")) + done <- err + }() + select { + case actualCtx := <-upstream.entered: + require.True(t, IsAccountObservation(actualCtx)) + case <-time.After(time.Second): + t.Fatal("business sender did not reach local fake upstream") + } + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("observation request detached its cancellation") + } + + ordinary, ordinaryCancel := context.WithCancel(context.Background()) + detached, release := detachStreamUpstreamContext(ordinary, true) + defer release() + ordinaryCancel() + require.NoError(t, detached.Err(), "ordinary stream keeps the existing detached usage-collection contract") +} diff --git a/backend/internal/service/gateway_upstream_response.go b/backend/internal/service/gateway_upstream_response.go index 6aaaaf4b16d5..a502f7c86321 100644 --- a/backend/internal/service/gateway_upstream_response.go +++ b/backend/internal/service/gateway_upstream_response.go @@ -363,8 +363,10 @@ func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, err func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) { // Upstream returned a non-success HTTP status; count Ollama Cloud / OpenCode Go activity. - scheduleOllamaCloudUsageActivity(s.deferredService, account) - scheduleOpenCodeGoUsageActivity(s.deferredService, account) + if !IsAccountObservation(ctx) { + scheduleOllamaCloudUsageActivity(s.deferredService, account) + scheduleOpenCodeGoUsageActivity(s.deferredService, account) + } body, readErr := s.readUpstreamErrorBody(resp) if readErr != nil { // 读取失败时 body 可能被截断，错误分类会基于不完整数据；记录日志以便排查， diff --git a/backend/internal/service/gateway_upstream_transport_error.go b/backend/internal/service/gateway_upstream_transport_error.go index a9f459fed94e..32ca7a760ace 100644 --- a/backend/internal/service/gateway_upstream_transport_error.go +++ b/backend/internal/service/gateway_upstream_transport_error.go @@ -61,7 +61,9 @@ func (s *GatewayService) handleUpstreamTransportError(ctx context.Context, c *gi } // Transport attempt left local validation; count Ollama Cloud activity. - scheduleOllamaCloudUsageActivity(s.deferredService, account) + if !IsAccountObservation(ctx) { + scheduleOllamaCloudUsageActivity(s.deferredService, account) + } if classifyUpstreamTransportError(err).Persistent { s.tempUnscheduleTransportError(ctx, account, safeErr) @@ -95,6 +97,9 @@ func (s *GatewayService) tempUnscheduleTransportError(ctx context.Context, accou // shared by every forward path whose scheduler reads the persisted // temp-unschedulable state (Anthropic/Bedrock and Gemini). func tempUnscheduleAccountForTransportError(ctx context.Context, repo AccountRepository, account *Account, safeErr string) { + if IsAccountObservation(ctx) { + return + } if account == nil || repo == nil { return } diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index e4945431baba..1884f520d4c2 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -144,6 +144,9 @@ func (p *postUsageBillingParams) shouldUpdateAccountQuota() bool { // billing repo is unavailable (nil). Production uses applyUsageBilling → repo.Apply // for atomic billing. This path only runs in tests or degraded mode. func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps) { + if IsAccountObservation(ctx) { + return + } billingCtx, cancel := detachedBillingContext(ctx) defer cancel() @@ -351,6 +354,9 @@ func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsage } func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog, p *postUsageBillingParams, deps *billingDeps, repo UsageBillingRepository) (bool, error) { + if IsAccountObservation(ctx) { + return false, nil + } if p == nil || deps == nil { return false, nil } @@ -390,6 +396,9 @@ func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog } func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + if IsAccountObservation(ctx) { + return + } if p == nil || p.Cost == nil || deps == nil { return } @@ -574,6 +583,9 @@ func detachStreamUpstreamContext(ctx context.Context, stream bool) (context.Cont if ctx == nil { return context.Background(), func() {} } + if IsAccountObservation(ctx) { + return ctx, func() {} + } if !stream { return ctx, func() {} } @@ -584,6 +596,9 @@ func detachUpstreamContext(ctx context.Context) (context.Context, context.Cancel if ctx == nil { return context.Background(), func() {} } + if IsAccountObservation(ctx) { + return ctx, func() {} + } return context.WithoutCancel(ctx), func() {} } @@ -616,6 +631,9 @@ func (s *GatewayService) billingDeps() *billingDeps { // successful synchronous fallback. An enqueue, timeout, or failed write alone // must not publish a log-backed statistics change. func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usageLog *UsageLog, logKey string) bool { + if IsAccountObservation(ctx) { + return false + } if repo == nil || usageLog == nil { return false } @@ -660,6 +678,9 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage // RecordUsage 记录使用量并扣费（或更新订阅用量） func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error { + if IsAccountObservation(ctx) { + return nil + } return s.recordUsageCore(ctx, &recordUsageCoreInput{ Result: input.Result, APIKey: input.APIKey, @@ -776,6 +797,9 @@ func logResponseModelBillingApplied(component string, account *Account, requestI // recordUsageCore 是 RecordUsage 的核心实现。 func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput) error { + if IsAccountObservation(ctx) { + return nil + } result := input.Result apiKey := input.APIKey user := input.User diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index cd3f97a0a3dd..daf7bb08cc91 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -3086,6 +3086,9 @@ func asInt(v any) (int, bool) { } func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, body []byte) { + if IsAccountObservation(ctx) { + return + } // 遵守自定义错误码策略：未命中则跳过所有限流处理 if !account.ShouldHandleErrorCode(statusCode) { return diff --git a/backend/internal/service/grok_credential_failure.go b/backend/internal/service/grok_credential_failure.go index c988cc698439..3a4a227148ba 100644 --- a/backend/internal/service/grok_credential_failure.go +++ b/backend/internal/service/grok_credential_failure.go @@ -291,6 +291,9 @@ func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFa } func (s *OpenAIGatewayService) applyGrokCredentialAccountFailure(ctx context.Context, account *Account, class grokCredentialFailureClass) (string, error) { + if IsAccountObservation(ctx) { + return "", nil + } if s == nil || account == nil || ctx == nil || ctx.Err() != nil { if ctx != nil { return "", ctx.Err() diff --git a/backend/internal/service/grok_token_provider.go b/backend/internal/service/grok_token_provider.go index 984f56cc2cb4..6ee5f67697d9 100644 --- a/backend/internal/service/grok_token_provider.go +++ b/backend/internal/service/grok_token_provider.go @@ -60,6 +60,9 @@ func (p *GrokTokenProvider) SetTempUnschedCache(cache TempUnschedCache) { } func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account) (string, error) { + if IsAccountObservation(ctx) { + return p.GetAccessTokenForManualTest(ctx, account) + } if account == nil { return "", errors.New("account is nil") } diff --git a/backend/internal/service/grok_upstream_errors.go b/backend/internal/service/grok_upstream_errors.go index 816ac2fb974b..f211d840e1e9 100644 --- a/backend/internal/service/grok_upstream_errors.go +++ b/backend/internal/service/grok_upstream_errors.go @@ -272,6 +272,9 @@ func grokStructuredErrorMessageCandidates(body []byte) []string { // unschedulable rules to a non-content 403. It reports true only when a rule // matched; unmatched responses retain the legacy entitlement cooldown. func (s *OpenAIGatewayService) applyGrokForbiddenPolicy(ctx context.Context, account *Account, responseBody []byte) bool { + if IsAccountObservation(ctx) { + return false + } if account == nil || !account.IsTempUnschedulableEnabled() { return false } diff --git a/backend/internal/service/grok_upstream_failure.go b/backend/internal/service/grok_upstream_failure.go index 53d9d5500ac0..d0754e51b1fe 100644 --- a/backend/internal/service/grok_upstream_failure.go +++ b/backend/internal/service/grok_upstream_failure.go @@ -556,6 +556,9 @@ func (s *OpenAIGatewayService) applyGrokUpstreamFailureDecision( account *Account, decision GrokUpstreamFailureDecision, ) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || account == nil || !decision.ShouldCooldown || decision.Cooldown <= 0 { return false } diff --git a/backend/internal/service/observed_account_lease_test.go b/backend/internal/service/observed_account_lease_test.go new file mode 100644 index 000000000000..3097bc95bdae --- /dev/null +++ b/backend/internal/service/observed_account_lease_test.go @@ -0,0 +1,203 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type observedAccountLeaseCacheForTest struct { + stubConcurrencyCacheForTest + mu sync.Mutex + now time.Time + slots map[int64]map[string]time.Time + releases map[int64]int + refreshErr error + refreshes atomic.Int64 + refreshed chan struct{} +} + +func newObservedAccountLeaseCacheForTest() *observedAccountLeaseCacheForTest { + return &observedAccountLeaseCacheForTest{now: time.Now().UTC(), slots: make(map[int64]map[string]time.Time), + releases: make(map[int64]int), refreshed: make(chan struct{}, 64)} +} + +func (c *observedAccountLeaseCacheForTest) AcquireAccountSlot(_ context.Context, id int64, limit int, requestID string) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.slots[id] == nil { + c.slots[id] = make(map[string]time.Time) + } + for member, expires := range c.slots[id] { + if !expires.After(c.now) { + delete(c.slots[id], member) + } + } + if len(c.slots[id]) >= limit { + return false, nil + } + c.slots[id][requestID] = c.now.Add(15 * time.Minute) + return true, nil +} + +func (c *observedAccountLeaseCacheForTest) RefreshAccountSlot(ctx context.Context, id int64, requestID string) (bool, error) { + defer func() { c.refreshed <- struct{}{} }() + c.refreshes.Add(1) + c.mu.Lock() + defer c.mu.Unlock() + if ctx.Err() != nil { + return false, ctx.Err() + } + if c.refreshErr != nil { + return false, c.refreshErr + } + if expires, found := c.slots[id][requestID]; !found || !expires.After(c.now) { + return false, nil + } + c.slots[id][requestID] = c.now.Add(15 * time.Minute) + return true, nil +} + +func (c *observedAccountLeaseCacheForTest) ReleaseAccountSlot(_ context.Context, id int64, requestID string) error { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.slots[id], requestID) + c.releases[id]++ + return nil +} + +func TestObservedAccountLeaseKeepsThirtyMinuteCapacityAndStopsOnRelease(t *testing.T) { + cache := newObservedAccountLeaseCacheForTest() + svc := NewConcurrencyService(cache) + // Ordinary acquisition retains its current non-renewing behavior even + // though this cache also supports the optional observation extension. + normal, err := svc.AcquireAccountSlot(context.Background(), 2, 1) + require.NoError(t, err) + require.True(t, normal.Acquired) + defer normal.ReleaseFunc() + owned, err := cache.AcquireAccountSlot(context.Background(), 1, 1, "observed") + require.NoError(t, err) + require.True(t, owned) + ticks := make(chan time.Time, 1) + var stopped atomic.Int64 + lease := newObservedAccountLease(context.Background(), cache, cache, 1, "observed", nil, ticks, func() { stopped.Add(1) }) + t.Cleanup(lease.Release) + for i := 0; i < 60; i++ { + cache.mu.Lock() + cache.now = cache.now.Add(observedAccountLeaseRefreshInterval) + now := cache.now + cache.mu.Unlock() + ticks <- now + <-cache.refreshed + allowed, acquireErr := cache.AcquireAccountSlot(context.Background(), 1, 1, "other-business") + require.NoError(t, acquireErr) + require.False(t, allowed, "a competing ordinary request remains blocked beyond the original fifteen-minute TTL") + } + allowed, err := cache.AcquireAccountSlot(context.Background(), 2, 1, "ordinary-after-original-ttl") + require.NoError(t, err) + require.True(t, allowed, "ordinary slots must keep the existing expiry contract") + lease.Release() + lease.Release() + require.Equal(t, int64(1), stopped.Load()) + require.Equal(t, int64(60), cache.refreshes.Load()) + ticks <- cache.now + cache.mu.Lock() + releases, remaining := cache.releases[1], len(cache.slots[1]) + cache.mu.Unlock() + require.Equal(t, 1, releases) + require.Zero(t, remaining) + require.ErrorIs(t, lease.Context().Err(), context.Canceled) +} + +func TestObservedAccountLeaseCancellationAndLossStopWorkerAndOperation(t *testing.T) { + for _, mode := range []string{"cancel", "missing", "redis_error"} { + t.Run(mode, func(t *testing.T) { + cache := newObservedAccountLeaseCacheForTest() + owned, err := cache.AcquireAccountSlot(context.Background(), 1, 1, "observed") + require.NoError(t, err) + require.True(t, owned) + parent, cancel := context.WithCancel(context.Background()) + defer cancel() + state := newPelicanExecutionState(parent, newPelicanExecutionCoordinator(1), 30*time.Minute, time.Now) + operation := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + defer state.finish() + ticks := make(chan time.Time, 1) + var stopped, losses atomic.Int64 + lease := newObservedAccountLease(operation, cache, cache, 1, "observed", func(cause error) { + losses.Add(1) + CancelPelicanExecution(operation, cause) + }, ticks, func() { stopped.Add(1) }) + t.Cleanup(lease.Release) + if mode == "cancel" { + cancel() + } else { + cache.mu.Lock() + if mode == "missing" { + delete(cache.slots[1], "observed") + } else { + cache.refreshErr = errors.New("synthetic redis unavailable") + } + cache.mu.Unlock() + ticks <- cache.now + } + select { + case <-lease.done: + case <-time.After(2 * time.Second): + t.Fatal("cancelled or lost account lease left its renewal worker running") + } + require.Equal(t, int64(1), stopped.Load()) + cache.mu.Lock() + releases, remaining := cache.releases[1], len(cache.slots[1]) + cache.mu.Unlock() + require.Equal(t, 1, releases) + require.Zero(t, remaining) + if mode == "cancel" { + require.Zero(t, losses.Load()) + require.Zero(t, cache.refreshes.Load()) + } else { + require.Equal(t, int64(1), losses.Load()) + require.ErrorIs(t, context.Cause(operation), ErrObservedAccountLeaseLost) + require.ErrorIs(t, context.Cause(lease.Context()), ErrObservedAccountLeaseLost) + _, _, acquireErr := AcquirePelicanExecution(operation, 2) + require.ErrorIs(t, acquireErr, context.Canceled, "lost account ownership cannot continue with a later model request") + } + lease.Release() + require.Equal(t, int64(1), stopped.Load()) + }) + } +} + +func TestObservedAccountLeaseAcquisitionUsesNormalCapacityAndRequiresRenewal(t *testing.T) { + cache := newObservedAccountLeaseCacheForTest() + svc := NewConcurrencyService(cache) + lease, acquired, err := svc.AcquireObservedAccountLease(nil, 1, 1, nil) + require.NoError(t, err) + require.True(t, acquired) + require.NotNil(t, lease) + second, acquired, err := svc.AcquireObservedAccountLease(context.Background(), 1, 1, nil) + require.NoError(t, err) + require.False(t, acquired) + require.Nil(t, second) + lease.Release() + cache.mu.Lock() + releases := cache.releases[1] + cache.mu.Unlock() + require.Equal(t, 1, releases) + require.Zero(t, cache.refreshes.Load(), "acquiring/releasing a short observation need not wait for a timer") + unsupported, acquired, err := NewConcurrencyService(&stubConcurrencyCacheForTest{}).AcquireObservedAccountLease(context.Background(), 1, 1, nil) + require.ErrorContains(t, err, "renewal is unsupported") + require.False(t, acquired) + require.Nil(t, unsupported) + var empty *ConcurrencyService + unlimited, acquired, err := empty.AcquireObservedAccountLease(nil, 1, 0, nil) + require.NoError(t, err) + require.True(t, acquired) + require.Nil(t, unlimited) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 91791621419c..795542162877 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -448,6 +448,9 @@ func isOpenAIAccount(account *Account) bool { // handleOpenAIAccountUpstreamError expects canonicalModel to be the model used // for scheduling after applying account mapping exactly once. func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, canonicalModel ...string) bool { + if IsAccountObservation(ctx) { + return false + } // Classify request-local failures before any policy can mutate account health. // The same predicates are used by the HTTP failover and error constructors. if isOpenAIRequestScopedSafetyRejection(responseBody) { @@ -612,6 +615,9 @@ func shouldCooldownOpenAITransientUpstreamError(statusCode int, responseBody []b } func (s *OpenAIGatewayService) markOpenAIOAuth429RateLimited(ctx context.Context, account *Account, headers http.Header, responseBody []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || !isOpenAIOAuthAccount(account) { return } @@ -1221,6 +1227,9 @@ func (s *OpenAIGatewayService) CooldownOpenAIRetryExhausted( canonicalModel string, failoverErr *UpstreamFailoverError, ) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil || failoverErr == nil || !isOpenAIAccount(account) { return } diff --git a/backend/internal/service/openai_apikey_health_breaker.go b/backend/internal/service/openai_apikey_health_breaker.go index 65581ae3d283..cc56a93f5ecb 100644 --- a/backend/internal/service/openai_apikey_health_breaker.go +++ b/backend/internal/service/openai_apikey_health_breaker.go @@ -51,6 +51,9 @@ func classifyOpenAIAPIKeyHealthFailure(err error) (int, []byte, bool) { } func (s *RateLimitService) ObserveOpenAIAPIKeyHealthFailure(ctx context.Context, account *Account, upstreamErr error) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || s.openAIAPIKeyHealth == nil || s.settingService == nil || s.accountRepo == nil || !isOpenAIAPIKeyHealthBreakerAccount(account) { return false } diff --git a/backend/internal/service/openai_budget_exceeded.go b/backend/internal/service/openai_budget_exceeded.go index c06e7131f0c1..78875e5821d2 100644 --- a/backend/internal/service/openai_budget_exceeded.go +++ b/backend/internal/service/openai_budget_exceeded.go @@ -102,6 +102,9 @@ func openAIBudgetExceededAccountError(payload []byte) string { // handleOpenAIBudgetExceeded moves the account to the error state. func (s *RateLimitService) handleOpenAIBudgetExceeded(ctx context.Context, account *Account, payload []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil { return } @@ -118,6 +121,9 @@ func (s *RateLimitService) handleOpenAIBudgetExceeded(ctx context.Context, accou } func (s *OpenAIGatewayService) handleOpenAIBudgetExceeded(ctx context.Context, account *Account, payload []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil { return } diff --git a/backend/internal/service/openai_codex_turn_state.go b/backend/internal/service/openai_codex_turn_state.go index 774041f18a9e..4614df3c47ca 100644 --- a/backend/internal/service/openai_codex_turn_state.go +++ b/backend/internal/service/openai_codex_turn_state.go @@ -139,6 +139,9 @@ func (s *OpenAIGatewayService) commitOpenAIWSSessionTurnState( sessionHash string, turnState string, ) { + if c != nil && c.Request != nil && IsAccountObservation(c.Request.Context()) { + return + } turnState = strings.TrimSpace(turnState) if stateStore != nil && sessionHash != "" { if turnState == "" { diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 2be6f3e74ca3..626dc4eff541 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -993,7 +993,8 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco // Build upstream request upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) var headerGuard *openAIFirstOutputHeaderGuard - if firstOutputTimeout > 0 && !s.codexGatewayBorrowHTTPConfigured(account, upstreamModel) { + var pelicanHeaderPolicy *pelicanFirstOutputPolicy + if firstOutputTimeout > 0 && !IsPelicanGeneration(ctx) && !s.codexGatewayBorrowHTTPConfigured(account, upstreamModel) { upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard( upstreamCtx, releaseUpstreamCtx, startTime.Add(borrowPreparationElapsed).Add(firstOutputTimeout), ) @@ -1056,8 +1057,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } firstOutputStartedAt := startTime.Add(borrowPreparationElapsed) if firstOutputTimeout > 0 && headerGuard == nil { - upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard(upstreamReq.Context(), releaseUpstreamCtx, firstOutputStartedAt.Add(firstOutputTimeout)) - upstreamReq = upstreamReq.WithContext(upstreamCtx) + if IsPelicanGeneration(ctx) { + pelicanHeaderPolicy = &pelicanFirstOutputPolicy{timeout: firstOutputTimeout} + upstreamReq = upstreamReq.WithContext(context.WithValue(upstreamReq.Context(), pelicanFirstOutputPolicyContextKey{}, pelicanHeaderPolicy)) + } else { + upstreamCtx, headerGuard = newOpenAIFirstOutputHeaderGuard(upstreamReq.Context(), releaseUpstreamCtx, firstOutputStartedAt.Add(firstOutputTimeout)) + upstreamReq = upstreamReq.WithContext(upstreamCtx) + } } if _, borrowed := upstreamReq.Context().Value(codexGatewayBorrowHTTPPreparationContextKey{}).(codexGatewayBorrowHTTPPreparation); borrowed { finalized, finalizationErr := s.applyCodexGatewayBorrowHTTP(upstreamReq, account, upstreamModel, proxyURL) @@ -1078,6 +1084,12 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco upstreamStart := time.Now() reasoningRecovery.MarkAttemptDispatched() resp, err := s.doOpenAICodexUpstream(upstreamReq, account, proxyURL, upstreamModel) + if pelicanHeaderPolicy != nil { + headerGuard = pelicanHeaderPolicy.guard + if !pelicanHeaderPolicy.startedAt.IsZero() { + firstOutputStartedAt = pelicanHeaderPolicy.startedAt + } + } reasoningRecovery.ObserveResponse(resp) SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) if IsCodexGatewayBorrowRequestFailure(err) { @@ -1089,7 +1101,13 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco } return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) } - if headerGuard != nil && headerGuard.stopHeaderWait() { + headerTimedOut := false + if pelicanHeaderPolicy != nil { + headerTimedOut = pelicanHeaderPolicy.stopHeaderWait() + } else if headerGuard != nil { + headerTimedOut = headerGuard.stopHeaderWait() + } + if headerTimedOut { if resp != nil && resp.Body != nil { _ = resp.Body.Close() } @@ -1280,7 +1298,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco if reqStream { setOpenAIRefusalEarlyStreamEligibility(c, account, body) streamCtx := ctx - if borrowPreparationElapsed > 0 { + if borrowPreparationElapsed > 0 || IsPelicanGeneration(ctx) { streamCtx = withCodexGatewayBorrowHTTPFirstOutputStart(ctx, firstOutputStartedAt) } streamResult, err := s.handleStreamingResponseWithReasoning(streamCtx, resp, c, account, startTime, originalModel, upstreamModel, reasoningEffortValue) diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 9d57a201db7e..0a06f068aa41 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -184,7 +184,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( s.handleGrokAccountUpstreamError(errCtx, account, resp.StatusCode, resp.Header, respBody) // Quota/rate-limit responses stamp the team+model overlay. Capacity is // request pressure and must not hide sibling accounts. - if shouldMarkGrokTeamModelRateLimit(resp.StatusCode, respBody) { + if !IsAccountObservation(ctx) && shouldMarkGrokTeamModelRateLimit(resp.StatusCode, respBody) { markGrokTeamModelRateLimit(account, upstreamModel, resolveGrokTeamRateLimitUntil(time.Now().Add(grokTeamRateLimitDefaultTTL), time.Now())) } if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) { @@ -1632,6 +1632,9 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco } func (s *OpenAIGatewayService) updateGrokUsageSnapshotWithRateLimit(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot, installRateLimit bool) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil || account.ID <= 0 || snapshot == nil { return } @@ -1687,6 +1690,9 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshotWithRateLimit(ctx context. } func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) { + if IsAccountObservation(ctx) { + return + } snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now()) if snapshot != nil { stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx)) @@ -1848,6 +1854,9 @@ func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnap } func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) { + if IsAccountObservation(ctx) { + return + } if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil { return } @@ -1862,6 +1871,9 @@ func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository } func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) { + if IsAccountObservation(ctx) { + return + } if repo == nil || account == nil || account.ID <= 0 { return } @@ -1880,6 +1892,9 @@ func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account * } func (s *OpenAIGatewayService) rateLimitGrok(ctx context.Context, account *Account, resetAt time.Time) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil { return } @@ -2019,6 +2034,9 @@ func persistGrokTransientModelCooldown(account *Account, decision GrokUpstreamFa } func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil { return } @@ -2113,6 +2131,9 @@ func isGrokSpendingLimitError(responseBody []byte) bool { } func (s *OpenAIGatewayService) tempUnscheduleGrok(ctx context.Context, account *Account, cooldown time.Duration, reason string) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil { return } diff --git a/backend/internal/service/openai_gateway_ops_visibility.go b/backend/internal/service/openai_gateway_ops_visibility.go index 45b7465473fd..a766ad67a3b9 100644 --- a/backend/internal/service/openai_gateway_ops_visibility.go +++ b/backend/internal/service/openai_gateway_ops_visibility.go @@ -297,6 +297,9 @@ var ( // written, was written inside the repeat window, or too many writes are in // flight. It does not change status, schedulability or any cooldown. func (s *OpenAIGatewayService) recordOpenAILastUpstreamError(ctx context.Context, account *Account, statusCode int, body []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || s.accountRepo == nil || account == nil || account.ID <= 0 { return } diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index adaa57e2d46e..2a9a91c373af 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -2546,7 +2546,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms } - s.clearOpenAIProxyStreamDisconnect(account) + s.clearOpenAIProxyStreamDisconnectInContext(ctx, account) recordOpenAIRefusalRecovery(ctx, c, account, "sse", false, refusalEarlyEmitted, openAIRefusalActionRewritten, refusalStream.evidence()) return resultWithUsage(), nil case openAIRefusalStreamHold: @@ -2623,7 +2623,7 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( if clientDisconnected { return resultWithUsage(), fmt.Errorf("stream usage incomplete after disconnect: %w", err) } - s.recordOpenAIProxyStreamDisconnect(account, err, upstreamRequestID) + s.recordOpenAIProxyStreamDisconnectInContext(ctx, account, err, upstreamRequestID) logger.LegacyPrintf("service.openai_gateway", "[OpenAI passthrough] 流读取异常中断: account=%d request_id=%s err=%v", account.ID, @@ -2652,14 +2652,14 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough( s.newOpenAIStreamFailoverError(c, account, true, upstreamRequestID, nil, "OpenAI stream ended before a terminal event", resp.StatusCode) } s.recordOpenAIStreamUpstreamError(c, account, true, upstreamRequestID, "stream_error", nil, "OpenAI stream ended before an authoritative terminal event") - s.recordOpenAIProxyStreamDisconnect(account, errors.New("stream ended before terminal event"), upstreamRequestID) + s.recordOpenAIProxyStreamDisconnectInContext(ctx, account, errors.New("stream ended before terminal event"), upstreamRequestID) return resultWithUsage(), errors.New("stream usage incomplete: missing terminal event") } if terminalEventType != "response.completed" && terminalEventType != "response.done" { return resultWithUsage(), fmt.Errorf("upstream response terminated with %s", terminalEventType) } if sawTerminalEvent && !sawFailedEvent { - s.clearOpenAIProxyStreamDisconnect(account) + s.clearOpenAIProxyStreamDisconnectInContext(ctx, account) } logOpenAISuccessMissingUsage(ctx, c, account, resp, usage, terminalEventType, clientDisconnected) diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 8c0c242b8b67..65b083665719 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -520,7 +520,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. terminalEventType = "response.failed" } if sawTerminalEvent && !sawFailedEvent && terminalResponseErr == nil { - s.clearOpenAIProxyStreamDisconnect(account) + s.clearOpenAIProxyStreamDisconnectInContext(ctx, account) } if !sawTerminalEvent && !openAIStreamClientOutputStarted(c, clientOutputStarted) && !eventShouldFlush { return resultWithUsage(), s.newOpenAIStreamFailoverError( @@ -537,7 +537,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. if !sawTerminalEvent { s.recordOpenAIStreamUpstreamError(c, account, false, upstreamRequestID, "stream_error", nil, "OpenAI stream ended before an authoritative terminal event") if openAIStreamClientOutputStarted(c, clientOutputStarted) && !clientDisconnected { - s.recordOpenAIProxyStreamDisconnect(account, errors.New("stream ended before terminal event"), upstreamRequestID) + s.recordOpenAIProxyStreamDisconnectInContext(ctx, account, errors.New("stream ended before terminal event"), upstreamRequestID) } return resultWithUsage(), fmt.Errorf("stream usage incomplete: missing terminal event") } @@ -606,7 +606,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. if clientDisconnected { return resultWithUsage(), fmt.Errorf("stream usage incomplete after disconnect: %w", scanErr), true } - s.recordOpenAIProxyStreamDisconnect(account, scanErr, upstreamRequestID) + s.recordOpenAIProxyStreamDisconnectInContext(ctx, account, scanErr, upstreamRequestID) code, message := classifyOpenAIUpstreamStreamReadError(scanErr) sendErrorEvent(code, message) return resultWithUsage(), fmt.Errorf("stream read error: %w", scanErr), true @@ -1639,6 +1639,9 @@ func extractOpenAIResponseIDFromJSONBytes(body []byte) string { } func (s *OpenAIGatewayService) bindHTTPResponseAccount(ctx context.Context, c *gin.Context, account *Account, responseID string) { + if IsAccountObservation(ctx) { + return + } if s == nil || account == nil || account.ID <= 0 { return } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 0bfce8037740..2157012c0d9b 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -1093,6 +1093,12 @@ func SnapshotOpenAICompatibilityFallbackMetrics() OpenAICompatibilityFallbackMet } func (s *OpenAIGatewayService) detectCodexClientRestriction(c *gin.Context, account *Account, body []byte) CodexClientRestrictionDetectionResult { + // A trusted administrator operation has no downstream client to admit. + // Its purpose is installed internally; upstream identity/authentication and + // every account/global sending policy still use the ordinary gateway path. + if c != nil && c.Request != nil && IsPelicanGeneration(c.Request.Context()) && !IsCodexGatewayBorrowObservation(c.Request.Context()) { + return CodexClientRestrictionDetectionResult{Reason: CodexClientRestrictionReasonDisabled} + } // 安全默认：即便缺 settingService（仅测试/误配可达）也保持指纹门为默认种子， // 避免零值 policy（nil 信号）让指纹门失败开放。有 settingService 时整体覆盖为全局策略。 policy := CodexRestrictionPolicy{EngineFingerprintSignals: openai.DefaultEngineFingerprintSignals} diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index f9df82abe82f..a9ddac975b33 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -152,6 +152,9 @@ func groupBillsOpenAIFastAtStandard(apiKey *APIKey, account *Account, serviceTie // RecordUsage records usage and deducts balance func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRecordUsageInput) error { + if IsAccountObservation(ctx) { + return nil + } if input == nil { return errors.New("openai usage input is nil") } @@ -1143,6 +1146,9 @@ func buildCodexUsageExtraUpdates(snapshot *OpenAICodexUsageSnapshot, fallbackNow // (/wham/usage bengalfox 道)更新,不能被全局头口径污染(外审第7轮 P1)。本函数仅持 accountID, // 无法在此自检影子,故守卫前置到各调用点。 func (s *OpenAIGatewayService) updateCodexUsageSnapshot(ctx context.Context, accountID int64, snapshot *OpenAICodexUsageSnapshot) { + if IsAccountObservation(ctx) { + return + } if snapshot == nil { return } diff --git a/backend/internal/service/openai_official_http_failover.go b/backend/internal/service/openai_official_http_failover.go index 5bd2daa836fc..eb653ea9b085 100644 --- a/backend/internal/service/openai_official_http_failover.go +++ b/backend/internal/service/openai_official_http_failover.go @@ -42,6 +42,9 @@ func (s *OpenAIGatewayService) ReportOpenAIOfficialHTTPAccountScheduleResult( firstTokenMs *int, observedErr error, ) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || !IsOpenAIOfficialHTTPFailover(ctx, account) { return false } @@ -70,6 +73,9 @@ func (s *OpenAIGatewayService) ReportOpenAIOfficialHTTPAccountScheduleResult( // ObserveOpenAIOfficialHTTPAccountHealthFailure is also used when semantic // output has already committed, where observing health must not replay output. func (s *OpenAIGatewayService) ObserveOpenAIOfficialHTTPAccountHealthFailure(ctx context.Context, account *Account, observedErr error) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || s.rateLimitService == nil || observedErr == nil || !IsOpenAIOfficialHTTPFailover(ctx, account) { return false } diff --git a/backend/internal/service/openai_plugin_transport.go b/backend/internal/service/openai_plugin_transport.go index bc24feac7646..5dfaad6e84d9 100644 --- a/backend/internal/service/openai_plugin_transport.go +++ b/backend/internal/service/openai_plugin_transport.go @@ -9,11 +9,17 @@ func (s *OpenAIGatewayService) SetPluginManager(manager *PluginManager) { // doOpenAIUpstream 只在 OpenAI OAuth 能力绑定已启用时把真实请求交给插件。 // 插件返回标准 http.Response，响应解析、错误映射、SSE 和计费仍由现有核心链处理。 func (s *OpenAIGatewayService) doOpenAIUpstream(request *http.Request, proxyURL string, account *Account) (*http.Response, error) { - if s.pluginManager != nil { + if s.pluginManager != nil && s.pluginManager.ShouldRouteOpenAIOAuth(account) { + prepared, finish, prepareErr := PreparePelicanHTTPRequest(request, account.ID, account.Concurrency, "plugin") + if prepareErr != nil { + return nil, prepareErr + } + request = prepared response, handled, err := s.pluginManager.RoundTripOpenAIOAuth(request.Context(), request, proxyURL, account) if handled { - return response, err + return finish(response, err) } + _, _ = finish(nil, nil) } return s.httpUpstream.Do(request, proxyURL, account.ID, account.Concurrency) } @@ -28,11 +34,17 @@ func (s *AccountTestService) doOpenAIAccountTestUpstream( account *Account, useTLSFallback bool, ) (*http.Response, error) { - if s.pluginManager != nil { + if s.pluginManager != nil && s.pluginManager.ShouldRouteOpenAIOAuth(account) { + prepared, finish, prepareErr := PreparePelicanHTTPRequest(request, account.ID, account.Concurrency, "plugin") + if prepareErr != nil { + return nil, prepareErr + } + request = prepared response, handled, err := s.pluginManager.RoundTripOpenAIOAuth(request.Context(), request, proxyURL, account) if handled { - return response, err + return finish(response, err) } + _, _ = finish(nil, nil) } // 插件 round-trip 有意拿到明文请求；zstd 压缩是网关传输层的事，只在真正经 // HTTPUpstream 发出的副本上做。 diff --git a/backend/internal/service/openai_proxy_stream_circuit.go b/backend/internal/service/openai_proxy_stream_circuit.go index 575fc134bb76..0e0064dddba9 100644 --- a/backend/internal/service/openai_proxy_stream_circuit.go +++ b/backend/internal/service/openai_proxy_stream_circuit.go @@ -235,6 +235,18 @@ func openAIProxyStreamCircuitProxyID(account *Account) (int64, bool) { return *account.ProxyID, true } +func (s *OpenAIGatewayService) recordOpenAIProxyStreamDisconnectInContext(ctx context.Context, account *Account, streamErr error, upstreamRequestID string) { + if !IsAccountObservation(ctx) { + s.recordOpenAIProxyStreamDisconnect(account, streamErr, upstreamRequestID) + } +} + +func (s *OpenAIGatewayService) clearOpenAIProxyStreamDisconnectInContext(ctx context.Context, account *Account) { + if !IsAccountObservation(ctx) { + s.clearOpenAIProxyStreamDisconnect(account) + } +} + func (s *OpenAIGatewayService) recordOpenAIProxyStreamDisconnect(account *Account, streamErr error, upstreamRequestID string) { proxyID, ok := openAIProxyStreamCircuitProxyID(account) if !ok || streamErr == nil || errors.Is(streamErr, context.Canceled) || errors.Is(streamErr, context.DeadlineExceeded) { diff --git a/backend/internal/service/openai_team_linked_error.go b/backend/internal/service/openai_team_linked_error.go index 2225e59a5b5b..8f690ab26dc8 100644 --- a/backend/internal/service/openai_team_linked_error.go +++ b/backend/internal/service/openai_team_linked_error.go @@ -22,6 +22,9 @@ const ( // 的其余 active 账户一并置为 error 并立即熔断。触发账户自身不在 fan-out 范围内， // 仍由常规 402 处理标记。 func (s *RateLimitService) maybeHandleOpenAITeamLinkedError(ctx context.Context, account *Account, statusCode int, responseBody []byte) { + if IsAccountObservation(ctx) { + return + } if s == nil || s.accountRepo == nil || statusCode != http.StatusPaymentRequired || !isOpenAIOAuthAccount(account) { return } diff --git a/backend/internal/service/openai_token_provider.go b/backend/internal/service/openai_token_provider.go index f2b8aafa20a3..9d24e54eb7e0 100644 --- a/backend/internal/service/openai_token_provider.go +++ b/backend/internal/service/openai_token_provider.go @@ -162,7 +162,7 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou const reason = "openai access_token expired and refresh_token is missing" // 永久故障：缺失 refresh_token 时账号无法自愈，必须立即从调度池剔除， // 否则会被反复选中、每次都在 token 阶段直接返回错误，对用户呈现持续 502。 - if !IsCodexGatewayBorrowObservation(ctx) { + if !IsAccountObservation(ctx) { p.disableAccountMissingRefreshToken(account, reason) } return "", errors.New(reason) diff --git a/backend/internal/service/openai_upstream_transport_error.go b/backend/internal/service/openai_upstream_transport_error.go index 0585136313c6..b9e4761ba9ef 100644 --- a/backend/internal/service/openai_upstream_transport_error.go +++ b/backend/internal/service/openai_upstream_transport_error.go @@ -143,7 +143,7 @@ func (s *OpenAIGatewayService) handleOpenAIUpstreamTransportError(ctx context.Co } // Transport attempt reached the network path; count as Ollama Cloud / OpenCode Go activity. - if s != nil { + if s != nil && !IsAccountObservation(ctx) { scheduleOllamaCloudUsageActivity(s.deferredService, account) scheduleOpenCodeGoUsageActivity(s.deferredService, account) } @@ -180,6 +180,9 @@ func (s *OpenAIGatewayService) handleOpenAIUpstreamTransportError(ctx context.Co // tempUnscheduleOpenAITransportError restores the upstream ten-minute durable // transport policy only for the marked ordinary HTTP account path. func (s *OpenAIGatewayService) tempUnscheduleOpenAITransportError(ctx context.Context, account *Account, safeErr string) { + if IsAccountObservation(ctx) { + return + } if s == nil || !IsOpenAIOfficialHTTPFailover(ctx, account) { return } diff --git a/backend/internal/service/openai_ws_forwarder_support.go b/backend/internal/service/openai_ws_forwarder_support.go index e2a176661e4f..8790a0550aa6 100644 --- a/backend/internal/service/openai_ws_forwarder_support.go +++ b/backend/internal/service/openai_ws_forwarder_support.go @@ -186,7 +186,7 @@ func (s *OpenAIGatewayService) performOpenAIWSGeneratePrewarm( } lease.MarkPrewarmed() - if prewarmResponseID != "" && stateStore != nil { + if !IsAccountObservation(ctx) && prewarmResponseID != "" && stateStore != nil { ttl := s.openAIWSResponseStickyTTL() logOpenAIWSBindResponseAccountWarn(groupID, account.ID, prewarmResponseID, stateStore.BindResponseAccount(ctx, groupID, prewarmResponseID, account.ID, ttl)) stateStore.BindResponseConn(prewarmResponseID, lease.ConnID(), ttl) @@ -878,6 +878,9 @@ func isOpenAIWSRateLimitError(codeRaw, errTypeRaw, msgRaw string) bool { } func (s *OpenAIGatewayService) persistOpenAIWSRateLimitSignal(ctx context.Context, account *Account, headers http.Header, responseBody []byte, codeRaw, errTypeRaw, msgRaw string, canonicalModel ...string) { + if IsAccountObservation(ctx) { + return + } if s == nil || s.rateLimitService == nil || account == nil || account.Platform != PlatformOpenAI { return } diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 579d66ad0b14..69fc1f33c241 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -8,6 +8,7 @@ import ( "fmt" "net/http" "net/url" + "slices" "strings" "time" @@ -300,12 +301,11 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } return nil, wrapOpenAIWSFallback(classifyOpenAIWSAcquireError(err), err) } - // cleanExit 标记正常终端事件退出，此时上游不会再发送帧，连接可安全归还复用。 - // 所有异常路径（读写错误、error 事件等）已在各自分支中提前调用 MarkBroken， - // 因此 defer 中只需处理正常退出时不 MarkBroken 即可。 + // 正常业务终态可归还连接。管理员借用的一次性连接必须先移出池再 + // 释放租约，避免等待中的业务请求取得它后又被借用清理关闭。 cleanExit := false defer func() { - if !cleanExit { + if !cleanExit || (borrowTurn != nil && borrowTurn.pelicanOneShot) { lease.MarkBroken() } lease.Release() @@ -386,6 +386,37 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( s.commitOpenAIWSSessionTurnState(c, account, stateStore, groupID, sessionHash, handshakeTurnState) } + var observePelicanFrame func([]byte) error + var releasePelican func() + pelicanPreparationCtx := ctx + ctx, observePelicanFrame, releasePelican, err = preparePelicanWSSend(ctx, account, wsURL, payload, borrowTurn != nil) + if err != nil { + return nil, err + } + defer releasePelican() + if IsPelicanGeneration(ctx) && borrowTurn != nil && !time.Now().Before(borrowTurn.expires) { + // No prompt has been written. Release the expired route and all execution + // resources, then perform the ordinary preparation on this same account. + releasePelican() + lease.MarkBroken() + lease.Release() + borrowTurn.finish(s.getOpenAIWSConnPool(), nil, ErrCodexGatewayBorrowChanged) + return s.forwardOpenAIWSV2(pelicanPreparationCtx, c, account, reqBody, clientPromptCacheKey, executionScope, token, decision, + isCodexCLI, reqStream, originalModel, mappedModel, startTime, attempt, lastFailureReason, agentTaskRecoveryTried) + } + if IsPelicanGeneration(ctx) && borrowTurn != nil { + prepared, prepareErr := codexGatewayBorrowWSRequest(ctx, wsURL, wsHeaders) + if prepareErr != nil { + return nil, codexGatewayBorrowWSPreparationError(c, account, prepareErr) + } + borrowed, application, applyErr := s.gatewayBorrow.Apply(prepared, account, openAIWSPayloadString(payload, "model"), proxyURL, nil, true) + if applyErr != nil { + return nil, codexGatewayBorrowWSPreparationError(c, account, applyErr) + } + if application == nil || !application.Applied || !application.ExpiresAt.Equal(borrowTurn.expires) || !slices.Equal(prepared.Header.Values("Cookie"), borrowed.Header.Values("Cookie")) { + return nil, codexGatewayBorrowWSPreparationError(c, account, ErrCodexGatewayBorrowChanged) + } + } if err := s.performOpenAIWSGeneratePrewarm( ctx, lease, @@ -399,6 +430,12 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( ); err != nil { return nil, err } + if IsPelicanGeneration(ctx) { + ctx = BeginPelicanGeneration(ctx) + if err := ctx.Err(); err != nil { + return nil, err + } + } if err := lease.WriteJSONWithContextTimeout(ctx, payload, s.openAIWSWriteTimeout()); err != nil { lease.MarkBroken() @@ -456,7 +493,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( } clientDisconnected = true clientDisconnectDrainStartedAt = time.Now() - if !upstreamReadDetached { + if !upstreamReadDetached && !IsAccountObservation(ctx) { upstreamReadCtx = context.WithoutCancel(ctx) upstreamReadDetached = true } @@ -624,6 +661,12 @@ readLoop: } } message, readErr = lease.ReadMessageWithContextTimeout(upstreamReadCtx, currentReadTimeout) + if readErr == nil { + if observeErr := observePelicanFrame(message); observeErr != nil { + lease.MarkBroken() + return nil, observeErr + } + } if readErr == nil { if documents, repaired := splitOpenAIConcatenatedJSONDocuments(message); repaired { logOpenAIWSModeInfo( @@ -660,6 +703,9 @@ readLoop: } if readErr != nil { lease.MarkBroken() + if IsAccountObservation(ctx) && ctx.Err() != nil { + return nil, ctx.Err() + } closeStatus, closeReason := summarizeOpenAIWSReadCloseError(readErr) logOpenAIWSModeInfo( "read_fail account_id=%d conn_id=%s wrote_downstream=%v close_status=%s close_reason=%s cause=%s events=%d token_events=%d terminal_events=%d buffered_pending=%d buffered_flushed=%d first_event=%s last_event=%s", @@ -1032,12 +1078,12 @@ readLoop: flushStreamWriter(true) } - if responseID != "" && stateStore != nil { + if responseID != "" && stateStore != nil && !IsAccountObservation(ctx) { ttl := s.openAIWSResponseStickyTTL() logOpenAIWSBindResponseAccountWarn(groupID, account.ID, responseID, stateStore.BindResponseAccount(ctx, groupID, responseID, account.ID, ttl)) stateStore.BindResponseConn(responseID, lease.ConnID(), ttl) } - if stateStore != nil && storeDisabled && sessionHash != "" { + if stateStore != nil && storeDisabled && sessionHash != "" && !IsAccountObservation(ctx) { stateStore.BindSessionConn(groupID, sessionHash, lease.ConnID(), s.openAIWSSessionStickyTTL()) } firstTokenMsValue := -1 diff --git a/backend/internal/service/pelican_account_generator.go b/backend/internal/service/pelican_account_generator.go new file mode 100644 index 000000000000..da3145672205 --- /dev/null +++ b/backend/internal/service/pelican_account_generator.go @@ -0,0 +1,669 @@ +package service + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "slices" + "strings" + "sync" + "time" + "unicode/utf8" + + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + "github.com/gin-gonic/gin" + "github.com/klauspost/compress/zstd" + "github.com/tidwall/gjson" +) + +type pelicanGenerationContextKey struct{} +type pelicanFirstOutputPolicyContextKey struct{} + +// Retain the existing first-output policy, but arm its header phase only once +// preparation and account/global admission have completed. +type pelicanFirstOutputPolicy struct { + timeout time.Duration + guard *openAIFirstOutputHeaderGuard + startedAt time.Time + headersOnce sync.Once + headersTimedOut bool + newGuard func(context.Context, context.CancelFunc, time.Time) (context.Context, *openAIFirstOutputHeaderGuard) +} + +func (p *pelicanFirstOutputPolicy) arm(ctx context.Context) context.Context { + if p == nil || p.timeout <= 0 { + return ctx + } + p.startedAt = PelicanFirstOutputStart(ctx, time.Now()) + create := p.newGuard + if create == nil { + create = newOpenAIFirstOutputHeaderGuard + } + var guarded context.Context + guarded, p.guard = create(ctx, func() {}, p.startedAt.Add(p.timeout)) + return guarded +} + +func (p *pelicanFirstOutputPolicy) stopHeaderWait() bool { + if p == nil || p.guard == nil { + return false + } + p.headersOnce.Do(func() { p.headersTimedOut = p.guard.stopHeaderWait() }) + return p.headersTimedOut +} + +// PelicanFirstOutputStart is shared by header and semantic-output guards. The +// first actual send remains the origin for ordinary same-account retries. +func PelicanFirstOutputStart(ctx context.Context, fallback time.Time) time.Time { + if ctx != nil && IsPelicanGeneration(ctx) { + if state := pelicanExecutionFromContext(ctx); state != nil { + state.mu.Lock() + started := state.generationAt + state.mu.Unlock() + if !started.IsZero() { + return started + } + } + } + return fallback +} + +// The capture belongs to one administrator's fixed-account invocation. It does +// not replace gateway dependencies, create an API key, or enter a selector. +type pelicanGenerationCapture struct { + account *Account + concurrency *ConcurrencyService + openaiGateway *OpenAIGatewayService + ginContext *gin.Context + fallbackModel string + fallbackEffort string + mu sync.Mutex + attempts []*pelicanHTTPAttempt +} + +type pelicanHTTPAttempt struct { + protocol string + contentType string + status int + model string + effort string + raw bytes.Buffer + err error + wsFrames bool +} + +// IsPelicanGeneration is an internal purpose marker. It is installed only by +// the dedicated generator, never from an HTTP header or public request body. +func IsPelicanGeneration(ctx context.Context) bool { + return pelicanCaptureFromContext(ctx) != nil +} + +func pelicanCaptureFromContext(ctx context.Context) *pelicanGenerationCapture { + if ctx == nil { + return nil + } + capture, _ := ctx.Value(pelicanGenerationContextKey{}).(*pelicanGenerationCapture) + return capture +} + +// PreparePelicanHTTPRequest is the shared outbound boundary. Ordinary requests +// retain their exact pointer and context. Borrow probes coordinate separately; +// their old observation marker must never start the work-generation budget. +// The returned completion function owns the execution/account leases until the +// response body closes, including error responses and same-account retries. +func PreparePelicanHTTPRequest(req *http.Request, accountID int64, accountConcurrency int, transport string) (*http.Request, func(*http.Response, error) (*http.Response, error), error) { + finish := func(resp *http.Response, err error) (*http.Response, error) { return resp, err } + if req == nil { + return req, finish, nil + } + capture := pelicanCaptureFromContext(req.Context()) + if capture == nil || capture.account.ID != accountID || IsCodexGatewayBorrowObservation(req.Context()) { + return req, finish, nil + } + var ctx context.Context + var release func() + for { + requestCtx := req.Context() + var acquireErr error + var releaseExecution func() + ctx, releaseExecution, acquireErr = AcquirePelicanExecution(requestCtx, accountID) + if acquireErr != nil { + return nil, finish, acquireErr + } + releaseAccount, acquireErr := acquirePelicanBusinessAccountSlot(ctx, capture.concurrency, accountID, accountConcurrency) + if acquireErr != nil { + releaseExecution() + return nil, finish, acquireErr + } + var releaseOnce sync.Once + release = func() { releaseOnce.Do(func() { releaseAccount(); releaseExecution() }) } + // Admission can outlast a borrowed cookie. Recheck the prepared cache + // binding after all waits, without a nested probe under this lease. + prepared, hasPreparation := requestCtx.Value(codexGatewayBorrowHTTPPreparationContextKey{}).(codexGatewayBorrowHTTPPreparation) + if hasPreparation && capture.openaiGateway != nil { + checked, prepareErr := capture.openaiGateway.applyCodexGatewayBorrowHTTP(req.WithContext(ctx), capture.account, prepared.model, prepared.proxy) + if prepareErr != nil { + release() + if prepared.application != nil && !time.Now().Before(prepared.application.ExpiresAt) && + capture.openaiGateway.codexGatewayBorrowHTTPConfigured(capture.account, prepared.model) && errors.Is(prepareErr, ErrCodexGatewayBorrowChanged) { + // Refresh an expired route through the normal preparation path + // after releasing every account/global execution resource. + refreshed, _, refreshErr := capture.openaiGateway.prepareCodexGatewayBorrowHTTP(requestCtx, req.WithContext(requestCtx), capture.account, prepared.model, prepared.proxy) + if refreshErr != nil { + return nil, finish, refreshErr + } + req = refreshed + continue + } + return nil, finish, prepareErr + } + req = checked + } + break + } + req = req.WithContext(ctx) + fallbackModel := capture.fallbackModel + if capture.ginContext != nil { + if observed := capture.ginContext.GetString(OpsUpstreamModelKey); observed != "" { + fallbackModel = observed + } + } + model, effort, protocol, metadataErr := pelicanFinalHTTPRequest(req, fallbackModel, capture.fallbackEffort) + if metadataErr != nil { + release() + return nil, finish, metadataErr + } + ctx = BeginPelicanGeneration(ctx) + if err := ctx.Err(); err != nil { + release() + return nil, finish, err + } + firstOutputPolicy, _ := ctx.Value(pelicanFirstOutputPolicyContextKey{}).(*pelicanFirstOutputPolicy) + ctx = firstOutputPolicy.arm(ctx) + req = req.WithContext(ctx) + borrowApplied := false + if prepared, ok := ctx.Value(codexGatewayBorrowHTTPPreparationContextKey{}).(codexGatewayBorrowHTTPPreparation); ok { + borrowApplied = prepared.application != nil && prepared.application.Applied + } + invocation := PelicanInvocation{Platform: capture.account.Platform, Model: model, Effort: effort, + Endpoint: req.URL.Scheme + "://" + req.URL.Host + req.URL.EscapedPath(), Protocol: protocol, Transport: transport, BorrowApplied: borrowApplied} + RecordPelicanInvocation(ctx, invocation) + attempt := &pelicanHTTPAttempt{protocol: protocol, model: model, effort: effort} + capture.mu.Lock() + capture.attempts = append(capture.attempts, attempt) + capture.mu.Unlock() + finish = func(resp *http.Response, sendErr error) (*http.Response, error) { + firstOutputPolicy.stopHeaderWait() + attempt.err = sendErr + if resp == nil || resp.Body == nil { + release() + return resp, sendErr + } + attempt.status, attempt.contentType = resp.StatusCode, resp.Header.Get("Content-Type") + if transport == "http" && resp.Proto != "" { + invocation.Transport = resp.Proto + RecordPelicanInvocation(ctx, invocation) + } + // Gateway error readers can have a smaller diagnostic cap. Capture the + // actual error body first so the administrator keeps its complete text. + if resp.StatusCode >= http.StatusBadRequest { + raw, readErr := io.ReadAll(io.LimitReader(resp.Body, codexGatewayBorrowPelicanMaxBody+1)) + _ = resp.Body.Close() + _, _ = attempt.raw.Write(raw) + attempt.err = readErr + release() + resp.Body = io.NopCloser(bytes.NewReader(raw)) + if int64(len(raw)) > codexGatewayBorrowPelicanMaxBody { + attempt.err = errors.New("upstream response exceeds the pelican record size limit") + } + return resp, sendErr + } + resp.Body = &pelicanCapturingBody{ReadCloser: resp.Body, attempt: attempt, release: release} + return resp, sendErr + } + return req, finish, nil +} + +func acquirePelicanBusinessAccountSlot(ctx context.Context, concurrency *ConcurrencyService, accountID int64, limit int) (func(), error) { + if concurrency == nil || limit <= 0 { + return func() {}, nil + } + finishQueueWait := BeginPelicanQueueWait(ctx) + defer finishQueueWait() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + lease, acquired, err := concurrency.AcquireObservedAccountLease(ctx, accountID, limit, func(cause error) { CancelPelicanExecution(ctx, cause) }) + if err != nil { + return nil, fmt.Errorf("acquire account concurrency capacity: %w", err) + } + if acquired { + return lease.Release, nil + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-ticker.C: + } + } +} + +// The WS sender calls this after normal route/socket preparation. Generation +// starts separately, immediately before response.create, so a pool wait or +// generate:false prewarm does not spend the work-generation budget. +func preparePelicanWSSend(ctx context.Context, account *Account, endpoint string, payload map[string]any, borrowed bool) (context.Context, func([]byte) error, func(), error) { + observe := func([]byte) error { return nil } + capture := pelicanCaptureFromContext(ctx) + if capture == nil || account == nil || capture.account.ID != account.ID { + return ctx, observe, func() {}, nil + } + ctx, releaseExecution, err := AcquirePelicanExecution(ctx, account.ID) + if err != nil { + return ctx, observe, func() {}, err + } + releaseAccount, err := acquirePelicanBusinessAccountSlot(ctx, capture.concurrency, account.ID, account.Concurrency) + if err != nil { + releaseExecution() + return ctx, observe, func() {}, err + } + var once sync.Once + release := func() { once.Do(func() { releaseAccount(); releaseExecution() }) } + model := openAIWSPayloadString(payload, "model") + effort := gjson.GetBytes(payloadAsJSONBytes(payload), "reasoning.effort").String() + RecordPelicanInvocation(ctx, PelicanInvocation{Platform: account.Platform, Model: model, Effort: effort, + Endpoint: endpoint, Protocol: "responses", Transport: "websocket", BorrowApplied: borrowed}) + attempt := &pelicanHTTPAttempt{protocol: "responses", status: http.StatusOK, model: model, effort: effort, wsFrames: true} + capture.mu.Lock() + capture.attempts = append(capture.attempts, attempt) + capture.mu.Unlock() + observe = func(message []byte) error { + if attempt.raw.Len()+len(message)+1 > int(codexGatewayBorrowPelicanMaxBody) { + attempt.err = errors.New("upstream response exceeds the pelican record size limit") + return attempt.err + } + _, _ = attempt.raw.Write(message) + _ = attempt.raw.WriteByte('\n') + return nil + } + return ctx, observe, release, nil +} + +type pelicanCapturingBody struct { + io.ReadCloser + attempt *pelicanHTTPAttempt + release func() + once sync.Once +} + +func (b *pelicanCapturingBody) Read(data []byte) (int, error) { + n, err := b.ReadCloser.Read(data) + if n > 0 { + if b.attempt.raw.Len()+n > int(codexGatewayBorrowPelicanMaxBody) { + b.attempt.err = errors.New("upstream response exceeds the pelican record size limit") + return 0, b.attempt.err + } + _, _ = b.attempt.raw.Write(data[:n]) + } + if err != nil && err != io.EOF { + b.attempt.err = err + } + return n, err +} + +func (b *pelicanCapturingBody) Close() error { + err := b.ReadCloser.Close() + b.once.Do(b.release) + return err +} + +func pelicanFinalHTTPRequest(req *http.Request, fallbackModel, fallbackEffort string) (model, effort, protocol string, err error) { + model, effort = fallbackModel, fallbackEffort + path := strings.ToLower(req.URL.Path) + switch { + case strings.Contains(path, "chat/completions"): + protocol = "chat" + case strings.Contains(path, "responses"): + protocol = "responses" + case strings.Contains(path, "generatecontent"): + protocol = "gemini" + default: + protocol = "claude" + } + if req.GetBody != nil || req.Body != nil { + body := req.Body + if req.GetBody != nil { + body, err = req.GetBody() + if err != nil { + return model, effort, protocol, err + } + } + { + raw, readErr := io.ReadAll(io.LimitReader(body, codexGatewayBorrowPelicanMaxBody+1)) + _ = body.Close() + if readErr != nil { + return model, effort, protocol, readErr + } + if int64(len(raw)) > codexGatewayBorrowPelicanMaxBody { + return model, effort, protocol, errors.New("pelican request body exceeds its record size limit") + } + if req.GetBody == nil { + // Reasoning recovery deliberately disables transparent replay. + // Preserve that nil GetBody contract; only restore the exact bytes + // consumed for this fixed request's final model/effort observation. + req.Body = io.NopCloser(bytes.NewReader(raw)) + } + if readErr == nil && strings.EqualFold(req.Header.Get("Content-Encoding"), "zstd") { + if decoder, decodeErr := zstd.NewReader(nil, zstd.WithDecoderMaxMemory(codexGatewayBorrowPelicanMaxBody)); decodeErr == nil { + raw, readErr = decoder.DecodeAll(raw, nil) + decoder.Close() + } + } + if readErr == nil && json.Valid(raw) { + for _, field := range []string{"model", "request.model"} { + if value := gjson.GetBytes(raw, field).String(); value != "" { + model = value + break + } + } + effort = "" + for _, field := range []string{"reasoning.effort", "reasoning_effort", "output_config.effort"} { + if value := gjson.GetBytes(raw, field).String(); value != "" { + effort = value + break + } + } + } + } + } + if protocol == "gemini" { + if _, after, found := strings.Cut(req.URL.Path, "/models/"); found { + if id, _, found := strings.Cut(after, ":"); found { + if unescaped, err := url.PathUnescape(id); err == nil { + model = unescaped + } + } + } + } + return model, effort, protocol, nil +} + +// GeneratePelican sends the fixed题面 using the selected account's ordinary +// business sender. Only that sender owns credential refresh and compatibility +// retries. An UpstreamFailoverError is recorded here and never enters selection. +func (s *AccountTestService) GeneratePelican(ctx context.Context, accountID int64, model, effort string) (*CodexGatewayBorrowPelicanResult, error) { + result := &CodexGatewayBorrowPelicanResult{ModelID: model, Effort: effort, Status: "failed"} + ctx = WithAccountObservation(ctx) + if pelicanExecutionFromContext(ctx) == nil { + state := newPelicanExecutionState(ctx, sharedPelicanExecution, time.Duration(PelicanDefaultGenerationTimeoutSeconds)*time.Second, time.Now) + ctx = context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + defer state.finish() + } + BeginPelicanPreparation(ctx) + if s == nil || s.accountRepo == nil { + result.Error = "account sender is unavailable" + return result, errors.New(result.Error) + } + account, err := s.accountRepo.GetByID(ctx, accountID) + if err != nil { + result.Error = "load selected account: " + err.Error() + return result, err + } + if model == "" { + model = PickConnectionTestModel(account) + result.ModelID = model + } + option := PelicanTestModelOptions(account, model) + if reason := PelicanTestTextCapabilityReason(account, model); reason != "" { + result.Status, result.Error = "skipped", reason + return result, nil + } + if effort == "" { + effort = option.DefaultEffort + } + if effort != "" && !slices.Contains(option.ReasoningEfforts, effort) { + result.Error = "reasoning effort is not supported by the selected account model" + return result, errors.New(result.Error) + } + result.Effort = effort + capture := &pelicanGenerationCapture{account: account, concurrency: s.pelicanConcurrencyService, + openaiGateway: s.openaiGatewayService, fallbackModel: option.UpstreamModel, fallbackEffort: effort} + ctx = context.WithValue(ctx, pelicanGenerationContextKey{}, capture) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + capture.ginContext = c + c.Request = httptest.NewRequest(http.MethodPost, "/internal/admin/pelican-tests", nil).WithContext(ctx) + // This represents the server's internal generation operation. No client + // UA, protocol handshake, API-key subject or group policy is synthesized. + sendErr := s.sendPelicanWithAccount(ctx, c, account, model, option.UpstreamModel, effort) + capture.mu.Lock() + var last *pelicanHTTPAttempt + if len(capture.attempts) > 0 { + last = capture.attempts[len(capture.attempts)-1] + } + capture.mu.Unlock() + if last == nil { + result.RawResponse = w.Body.String() + if sendErr == nil { + sendErr = errors.New("account sender completed without an upstream generation request") + } + result.Error = sendErr.Error() + return result, sendErr + } + result.RawResponse = last.raw.String() + if !utf8.ValidString(result.RawResponse) || strings.Contains(strings.ToLower(last.contentType), "eventstream") { + // AWS event-stream envelopes are binary. Keep the complete original + // bytes in a lossless representation and parse the shared adapter output. + result.RawResponse = "base64:" + base64.StdEncoding.EncodeToString(last.raw.Bytes()) + } + result.ModelID, result.Effort = last.model, last.effort + if last.wsFrames && last.raw.Len() == 0 && sendErr != nil { + // A normal generate:false prewarm can reject a route before the题面. + // Its complete native error is useful even though no work was generated. + var eventErr *openAIWSUpstreamEventError + if errors.As(sendErr, &eventErr) && len(eventErr.payload) > 0 { + result.RawResponse = string(eventErr.payload) + result.Error = string(eventErr.payload) + } else { + result.Error = sendErr.Error() + } + return result, sendErr + } + if last.status >= http.StatusBadRequest { + result.Error = fmt.Sprintf("API returned %d: %s", last.status, last.raw.String()) + return result, sendErr + } + parseBody := last.raw.String() + if strings.Contains(strings.ToLower(last.contentType), "eventstream") { + parseBody = w.Body.String() + } + answer, limited, responseModel, parseErr := parsePelicanTextResponse(last.protocol, parseBody) + if last.wsFrames { + answer, limited, responseModel, parseErr = parsePelicanWSFrames(parseBody) + } + result.RawAnswer = answer + if responseModel != "" { + result.ModelID = responseModel + } + if last.err != nil { + parseErr = fmt.Errorf("%w: %v", ErrAccountTestIncomplete, last.err) + } + if limited { + parseErr = fmt.Errorf("%w: output limit reached", ErrAccountTestIncomplete) + } + if parseErr != nil { + result.Error = parseErr.Error() + if errors.Is(parseErr, ErrAccountTestIncomplete) || errors.Is(parseErr, ErrAccountTestEmpty) { + result.Status = "incomplete" + } + return result, sendErr + } + if sendErr != nil { + result.Error = sendErr.Error() + return result, sendErr + } + result.Status = "complete" + return result, nil +} + +func pelicanOutputBudget(account *Account, upstreamModel string) int64 { + capacity := ResolveAccountModelContextCapacity(account, upstreamModel) + if capacity.MaxOutputTokens > 0 { + return capacity.MaxOutputTokens + } + return 64000 +} + +func pelicanMessagesBody(model, upstreamModel, effort string, budget int64) []byte { + payload := map[string]any{"model": model, "stream": true, "max_tokens": budget, + "messages": []map[string]any{{"role": "user", "content": CodexGatewayBorrowPelicanPrompt}}} + if effort != "" { + payload["output_config"] = map[string]any{"effort": effort} + if claude.EffortUsesAdaptiveThinking(upstreamModel) { + payload["thinking"] = map[string]any{"type": "adaptive"} + } + } + body, _ := json.Marshal(payload) + return body +} + +func (s *AccountTestService) sendPelicanWithAccount(ctx context.Context, c *gin.Context, account *Account, model, upstreamModel, effort string) error { + budget := pelicanOutputBudget(account, upstreamModel) + switch { + case account.IsGemini(), account.Platform == PlatformAntigravity && account.Type != AccountTypeUpstream && isAntigravityGeminiModel(upstreamModel): + payload := map[string]any{"contents": []map[string]any{{"role": "user", "parts": []map[string]any{{"text": CodexGatewayBorrowPelicanPrompt}}}}, + "generationConfig": map[string]any{"maxOutputTokens": budget}} + body, _ := json.Marshal(payload) + if account.IsGemini() { + if s.pelicanGeminiService == nil { + return errors.New("gemini account sender is unavailable") + } + _, err := s.pelicanGeminiService.ForwardNative(ctx, c, account, model, "streamGenerateContent", true, body) + return err + } + if s.antigravityGatewayService == nil { + return errors.New("antigravity account sender is unavailable") + } + _, err := s.antigravityGatewayService.ForwardGemini(ctx, c, account, model, "streamGenerateContent", true, body, false) + return err + case account.Platform == PlatformAnthropic: + if s.pelicanGatewayService == nil { + return errors.New("anthropic account sender is unavailable") + } + body := pelicanMessagesBody(model, upstreamModel, effort, budget) + parsed, err := ParseGatewayRequest(NewRequestBodyRef(body), PlatformAnthropic) + if err != nil { + return err + } + _, err = s.pelicanGatewayService.Forward(ctx, c, account, parsed) + return err + case account.Platform == PlatformAntigravity: + if s.antigravityGatewayService == nil { + return errors.New("antigravity account sender is unavailable") + } + _, err := s.antigravityGatewayService.Forward(ctx, c, account, pelicanMessagesBody(model, upstreamModel, effort, budget), false) + return err + case account.IsOpenAICompatible(): + if s.openaiGatewayService == nil { + return errors.New("OpenAI-compatible account sender is unavailable") + } + payload := map[string]any{"model": model, "stream": true, "max_output_tokens": budget, + "input": []map[string]any{{"role": "user", "content": []map[string]any{{"type": "input_text", "text": CodexGatewayBorrowPelicanPrompt}}}}} + if effort != "" { + payload["reasoning"] = map[string]any{"effort": effort} + } + body, _ := json.Marshal(payload) + _, err := s.openaiGatewayService.Forward(ctx, c, account, body) + return err + default: + return fmt.Errorf("platform %s has no supported text generation sender", account.Platform) + } +} + +func parsePelicanTextResponse(protocol, raw string) (answer string, limited bool, model string, err error) { + var text strings.Builder + emit := func(event TestEvent) { + if event.Type == "content" { + _, _ = text.WriteString(event.Text) + } + } + if !json.Valid([]byte(raw)) && pelicanHasSSEDataLine(raw) { + limited, model, err = parseAccountConnectionStream(protocol, strings.NewReader(raw), false, emit) + return text.String(), limited, model, err + } + var payload map[string]any + if json.Unmarshal([]byte(raw), &payload) != nil { + return "", false, "", fmt.Errorf("%w: %s", ErrAccountTestProtocol, raw) + } + p := connectionStreamState{protocol: protocol, emit: emit} + if protocol == "claude" { + if payload["error"] != nil { + return "", false, "", p.upstreamError(payload["error"]) + } + if value, _ := payload["model"].(string); value != "" { + p.upstreamModel = value + } + content, _ := payload["content"].([]any) + for _, rawBlock := range content { + block, _ := rawBlock.(map[string]any) + if block["type"] == "text" { + p.text(block["text"]) + } + } + p.stopReason, _ = payload["stop_reason"].(string) + if p.stopReason == "" { + return text.String(), false, p.upstreamModel, p.incomplete() + } + payload = map[string]any{"type": "message_stop"} + } else if protocol == "responses" && payload["type"] == nil { + event := "response.completed" + if payload["status"] == "incomplete" { + event = "response.incomplete" + } + payload = map[string]any{"type": event, "response": payload} + } + done, err := p.consume(payload) + if err == nil && !done { + err = p.incomplete() + } + return text.String(), p.limited, p.upstreamModel, err +} + +func pelicanHasSSEDataLine(raw string) bool { + for _, line := range strings.Split(raw, "\n") { + if strings.HasPrefix(strings.TrimSpace(line), "data:") { + return true + } + } + return false +} + +func parsePelicanWSFrames(raw string) (answer string, limited bool, model string, err error) { + var text strings.Builder + p := connectionStreamState{protocol: "responses", emit: func(event TestEvent) { + if event.Type == "content" { + _, _ = text.WriteString(event.Text) + } + }} + decoder := json.NewDecoder(strings.NewReader(raw)) + for { + var payload map[string]any + if err := decoder.Decode(&payload); err != nil { + if err == io.EOF { + break + } + return text.String(), false, p.upstreamModel, fmt.Errorf("%w: %v", ErrAccountTestProtocol, err) + } + done, parseErr := p.consume(payload) + if done || parseErr != nil { + return text.String(), p.limited, p.upstreamModel, parseErr + } + } + return text.String(), p.limited, p.upstreamModel, p.incomplete() +} diff --git a/backend/internal/service/pelican_account_generator_test.go b/backend/internal/service/pelican_account_generator_test.go new file mode 100644 index 000000000000..bda282215832 --- /dev/null +++ b/backend/internal/service/pelican_account_generator_test.go @@ -0,0 +1,295 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type pelicanGeneratorAccountRepo struct { + AccountRepository + accounts map[int64]*Account + reads []int64 +} + +func (r *pelicanGeneratorAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) { + r.reads = append(r.reads, id) + if account := r.accounts[id]; account != nil { + return account, nil + } + return nil, errors.New("missing local test account") +} + +type pelicanGeneratorUpstream struct { + status int + body string + contentType string + requests []*http.Request + bodies [][]byte + accountIDs []int64 + proxies []string + profiles []*tlsfingerprint.Profile + beforeSend func(*http.Request) +} + +func (s *pelicanGeneratorUpstream) Do(req *http.Request, proxy string, id int64, concurrency int) (*http.Response, error) { + return s.DoWithTLS(req, proxy, id, concurrency, nil) +} + +func (s *pelicanGeneratorUpstream) DoWithTLS(req *http.Request, proxy string, id int64, concurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) { + if s.beforeSend != nil { + s.beforeSend(req) + } + prepared, finish, err := PreparePelicanHTTPRequest(req, id, concurrency, "http") + if err != nil { + return nil, err + } + body, err := io.ReadAll(prepared.Body) + if err != nil { + return finish(nil, err) + } + prepared.Body = io.NopCloser(bytes.NewReader(body)) + s.requests = append(s.requests, prepared) + s.bodies = append(s.bodies, body) + s.accountIDs = append(s.accountIDs, id) + s.proxies = append(s.proxies, proxy) + s.profiles = append(s.profiles, profile) + status := s.status + if status == 0 { + status = http.StatusOK + } + contentType := s.contentType + if contentType == "" { + contentType = "text/event-stream" + } + return finish(&http.Response{StatusCode: status, Proto: "HTTP/1.1", Header: http.Header{"Content-Type": {contentType}}, Body: io.NopCloser(strings.NewReader(s.body))}, nil) +} + +func TestPelicanAccountGeneratorQueuePreservesSharedFirstOutputPolicy(t *testing.T) { + clock := time.Now().UTC() + state := newPelicanExecutionState(context.Background(), newPelicanExecutionCoordinator(1), 10*time.Minute, func() time.Time { return clock }) + defer state.finish() + ctx := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + account := &Account{ID: 70, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "local-key", "base_url": "http://fake.example"}} + upstream := &pelicanGeneratorUpstream{body: pelicanResponsesSSE("gpt-6.1-sol", "completed", "")} + svc, _ := newPelicanGeneratorForTest(account, upstream) + svc.openaiGatewayService.cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 20 + var policy *pelicanFirstOutputPolicy + upstream.beforeSend = func(req *http.Request) { + policy, _ = req.Context().Value(pelicanFirstOutputPolicyContextKey{}).(*pelicanFirstOutputPolicy) + require.NotNil(t, policy) + require.Equal(t, 20*time.Second, policy.timeout) + require.Nil(t, policy.guard, "account/global admission precedes the normal header guard") + require.True(t, state.generationAt.IsZero()) + clock = clock.Add(8 * time.Minute) + require.NoError(t, req.Context().Err()) + } + result, err := svc.GeneratePelican(ctx, account.ID, "gpt-6.1-sol", "high") + require.NoError(t, err) + require.Equal(t, "complete", result.Status, result.Error) + require.Equal(t, clock, policy.startedAt) + require.Equal(t, clock, state.generationAt) +} + +func newPelicanGeneratorForTest(account *Account, upstream *pelicanGeneratorUpstream) (*AccountTestService, *pelicanGeneratorAccountRepo) { + repo := &pelicanGeneratorAccountRepo{accounts: map[int64]*Account{account.ID: account}} + cfg := &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{AllowInsecureHTTP: true}}, Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}} + svc := &AccountTestService{accountRepo: repo, cfg: cfg, httpUpstream: upstream} + svc.openaiGatewayService = &OpenAIGatewayService{accountRepo: repo, cfg: cfg, httpUpstream: upstream} + svc.pelicanGatewayService = &GatewayService{accountRepo: repo, cfg: cfg, httpUpstream: upstream, tlsFPProfileService: &TLSFingerprintProfileService{}} + svc.pelicanGeminiService = &GeminiMessagesCompatService{accountRepo: repo, cfg: cfg, httpUpstream: upstream} + svc.antigravityGatewayService = &AntigravityGatewayService{accountRepo: repo, httpUpstream: upstream, settingService: &SettingService{cfg: cfg}} + return svc, repo +} + +func pelicanResponsesSSE(model, status, suffix string) string { + completed := map[string]any{"type": "response.completed", "response": map[string]any{"id": "resp_local", "model": model, "status": status, "output": []any{}, "usage": map[string]any{"input_tokens": 4, "output_tokens": 8}}} + if status == "incomplete" { + completed["type"] = "response.incomplete" + completed["response"].(map[string]any)["incomplete_details"] = map[string]any{"reason": "max_output_tokens"} + } + raw, _ := json.Marshal(completed) + return "data: {\"type\":\"response.output_text.delta\",\"delta\":\"" + suffix + "\"}\n\ndata: " + string(raw) + "\n\n" +} + +func TestPelicanAccountGeneratorSharedOpenAISettings(t *testing.T) { + for _, target := range []string{"gpt-6-astra", "gpt-6.1-sol"} { + t.Run(target, func(t *testing.T) { + proxyID := int64(37) + account := &Account{ID: 12, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1, Status: StatusError, Schedulable: false, + ProxyID: &proxyID, Proxy: &Proxy{Protocol: "socks5", Host: "127.0.0.1", Port: 19090}, + Credentials: map[string]any{"api_key": "local-key", "base_url": "http://fake.example", "model_mapping": map[string]any{"public-model": target}, + credKeyHeaderOverrideEnabled: true, credKeyHeaderOverrides: map[string]any{"X-Account-Identity": "saved-identity"}}} + upstream := &pelicanGeneratorUpstream{body: pelicanResponsesSSE(target, "completed", "")} + svc, repo := newPelicanGeneratorForTest(account, upstream) + prompts := NewSystemPromptService(nil, nil, nil) + prompts.publish(SystemPromptConfig{Enabled: true, DefaultPromptID: "site", Prompts: []SystemPrompt{{ID: "site", Body: "saved site instructions", Position: SystemPromptPositionPrepend, Role: SystemPromptRoleDeveloper}}}) + svc.openaiGatewayService.systemPrompts = prompts + result, err := svc.GeneratePelican(context.Background(), account.ID, "public-model", "") + require.NoError(t, err) + require.Equal(t, "complete", result.Status) + require.Equal(t, "", result.RawAnswer) + require.Equal(t, upstream.body, result.RawResponse) + require.Len(t, upstream.requests, 1) + require.Equal(t, []int64{account.ID}, repo.reads) + require.Equal(t, []int64{account.ID}, upstream.accountIDs) + require.Equal(t, "socks5://127.0.0.1:19090", upstream.proxies[0]) + require.Equal(t, []string{"saved-identity"}, upstream.requests[0].Header[resolveWireCasing("x-account-identity")]) + require.Equal(t, target, gjson.GetBytes(upstream.bodies[0], "model").String()) + require.Equal(t, "high", gjson.GetBytes(upstream.bodies[0], "reasoning.effort").String()) + require.Contains(t, string(upstream.bodies[0]), CodexGatewayBorrowPelicanPrompt) + require.Contains(t, string(upstream.bodies[0]), "saved site instructions") + require.Greater(t, gjson.GetBytes(upstream.bodies[0], "max_output_tokens").Int(), int64(1024)) + require.True(t, IsAccountObservation(upstream.requests[0].Context())) + }) + } +} + +func TestPelicanAccountGeneratorHTTPFailureKeepsSelectedAccountAndRawBody(t *testing.T) { + raw := `{"error":{"type":"rate_limit_error","message":"` + strings.Repeat("complete original detail ", 2000) + `"}}` + account := &Account{ID: 28, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "local-key", "base_url": "http://fake.example"}} + upstream := &pelicanGeneratorUpstream{status: 429, contentType: "application/json", body: raw} + svc, repo := newPelicanGeneratorForTest(account, upstream) + repo.accounts[29] = &Account{ID: 29, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + result, _ := svc.GeneratePelican(context.Background(), account.ID, "gpt-6.1-sol", "high") + require.Equal(t, "failed", result.Status) + require.Equal(t, raw, result.RawResponse) + require.Equal(t, "API returned 429: "+raw, result.Error) + require.Equal(t, []int64{account.ID}, repo.reads) + require.Equal(t, []int64{account.ID}, upstream.accountIDs) +} + +func TestPelicanAccountGeneratorTruncatedWorkIsIncomplete(t *testing.T) { + account := &Account{ID: 31, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "local-key", "base_url": "http://fake.example"}} + upstream := &pelicanGeneratorUpstream{body: pelicanResponsesSSE("gpt-6.1-sol", "incomplete", "")} + svc, _ := newPelicanGeneratorForTest(account, upstream) + result, _ := svc.GeneratePelican(context.Background(), account.ID, "gpt-6.1-sol", "high") + require.Equal(t, "incomplete", result.Status) + require.Equal(t, "", result.RawAnswer) + require.Equal(t, upstream.body, result.RawResponse) + require.Contains(t, result.Error, "output limit reached") + require.Len(t, upstream.requests, 1) +} + +func TestPelicanAccountGeneratorNativeTextDispatch(t *testing.T) { + claudeSSE := "data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"work\"}}\n\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\ndata: {\"type\":\"message_stop\"}\n\n" + chatSSE := "data: {\"model\":\"text-model\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"work\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n" + geminiSSE := "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"work\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":1,\"candidatesTokenCount\":1}}\n\n" + for _, scenario := range []struct{ name, platform, kind, model, protocol, path, response string }{ + {"cn-chat", PlatformDeepseek, AccountTypeAPIKey, "text-model", APIProtocolChatCompletions, "/chat/completions", chatSSE}, + {"cn-anthropic", PlatformMiniMax, AccountTypeAPIKey, "text-model", APIProtocolAnthropic, "/messages", claudeSSE}, + {"opencode", PlatformOpenCodeGo, AccountTypeAPIKey, "text-model", APIProtocolChatCompletions, "/chat/completions", chatSSE}, + {"anthropic", PlatformAnthropic, AccountTypeAPIKey, "text-model", "", "/messages", claudeSSE}, + {"gemini", PlatformGemini, AccountTypeAPIKey, "gemini-3.1-pro", "", ":streamGenerateContent", geminiSSE}, + {"antigravity-upstream", PlatformAntigravity, AccountTypeUpstream, "gemini-text-model", "", "/messages", claudeSSE}, + } { + t.Run(scenario.name, func(t *testing.T) { + account := &Account{ID: 44, Platform: scenario.platform, Type: scenario.kind, Credentials: map[string]any{"api_key": "local-key", "base_url": "http://fake.example", "api_protocol": scenario.protocol}} + upstream := &pelicanGeneratorUpstream{body: scenario.response} + svc, _ := newPelicanGeneratorForTest(account, upstream) + result, err := svc.GeneratePelican(context.Background(), account.ID, scenario.model, "") + require.NoError(t, err) + require.Equal(t, "complete", result.Status, result.Error) + require.Equal(t, "work", result.RawAnswer) + require.Len(t, upstream.requests, 1) + require.Contains(t, upstream.requests[0].URL.Path, scenario.path) + require.Contains(t, string(upstream.bodies[0]), CodexGatewayBorrowPelicanPrompt) + }) + } +} + +func TestPelicanParsersKeepJSONDataURLsAndPrettyWSFrames(t *testing.T) { + jsonReply := `{"model":"text-model","choices":[{"index":0,"message":{"content":""},"finish_reason":"stop"}]}` + text, limited, _, err := parsePelicanTextResponse("chat", jsonReply) + require.NoError(t, err) + require.False(t, limited) + require.Equal(t, "", text) + frames := "{\n\"type\":\"response.output_text.delta\",\n\"delta\":\"whole\"\n}\n" + `{"type":"response.completed","response":{"status":"completed","model":"gpt-6-astra","output":[]}}` + "\n" + text, limited, model, err := parsePelicanWSFrames(frames) + require.NoError(t, err) + require.False(t, limited) + require.Equal(t, "gpt-6-astra", model) + require.Equal(t, "whole", text) +} + +func TestPelicanHTTPRequestOrdinaryAndPreparationPurposeAreNoops(t *testing.T) { + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "http://fake.example/responses", strings.NewReader(`{}`)) + require.NoError(t, err) + prepared, finish, err := PreparePelicanHTTPRequest(req, 17, 1, "http") + require.NoError(t, err) + require.Same(t, req, prepared) + resp := &http.Response{} + actual, err := finish(resp, nil) + require.NoError(t, err) + require.Same(t, resp, actual) + capture := &pelicanGenerationCapture{account: &Account{ID: 17}} + state := newPelicanExecutionState(context.Background(), newPelicanExecutionCoordinator(1), time.Minute, time.Now) + defer state.finish() + ctx := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + ctx = context.WithValue(ctx, pelicanGenerationContextKey{}, capture) + ctx = WithCodexGatewayBorrowObservation(ctx) + req = req.WithContext(ctx) + prepared, _, err = PreparePelicanHTTPRequest(req, 17, 1, "http") + require.NoError(t, err) + require.Same(t, req, prepared) + require.True(t, state.generationAt.IsZero(), "borrow preparation must keep its own clock") + require.Empty(t, capture.attempts) +} + +func TestPelicanFirstOutputPolicyStartsAfterQueueAndPreservesNilReplay(t *testing.T) { + now := time.Now().UTC() + clock := now + state := newPelicanExecutionState(context.Background(), newPelicanExecutionCoordinator(1), 10*time.Minute, func() time.Time { return clock }) + defer state.finish() + capture := &pelicanGenerationCapture{account: &Account{ID: 67, Platform: PlatformOpenAI}, fallbackModel: "gpt-6.1-sol-high", fallbackEffort: "medium"} + ctx := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + ctx = context.WithValue(ctx, pelicanGenerationContextKey{}, capture) + var armedDeadline time.Time + policy := &pelicanFirstOutputPolicy{timeout: 30 * time.Second, newGuard: func(ctx context.Context, release context.CancelFunc, deadline time.Time) (context.Context, *openAIFirstOutputHeaderGuard) { + armedDeadline = deadline + return newOpenAIFirstOutputHeaderGuard(ctx, release, time.Now().Add(time.Hour)) + }} + ctx = context.WithValue(ctx, pelicanFirstOutputPolicyContextKey{}, policy) + BeginPelicanPreparation(ctx) + clock = clock.Add(8 * time.Minute) // A modeled queue, never a wall-clock sleep. + require.True(t, state.generationAt.IsZero()) + require.Nil(t, policy.guard, "queued work must not arm the upstream header guard") + body := `{"model":"gpt-6.1-sol","reasoning":{"effort":"high"},"input":"fixed","stream":true}` + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://fake.example/responses", strings.NewReader(body)) + require.NoError(t, err) + req.GetBody = nil // The normal reasoning-recovery dispatch contract. + prepared, finish, err := PreparePelicanHTTPRequest(req, 67, 1, "http") + require.NoError(t, err) + defer policy.guard.close() + require.Equal(t, clock.Add(policy.timeout), armedDeadline) + require.Equal(t, clock, PelicanFirstOutputStart(prepared.Context(), now)) + require.Nil(t, prepared.GetBody, "observation must not enable transport replay") + actualBody, err := io.ReadAll(prepared.Body) + require.NoError(t, err) + require.Equal(t, body, string(actualBody)) + require.Equal(t, "gpt-6.1-sol", capture.attempts[0].model) + require.Equal(t, "high", capture.attempts[0].effort) + _, err = finish(nil, nil) + require.NoError(t, err) + require.False(t, policy.stopHeaderWait(), "header completion is safe to query again by Forward") + ordinaryPolicy := &pelicanFirstOutputPolicy{timeout: time.Minute} + ordinaryCtx := context.WithValue(context.Background(), pelicanFirstOutputPolicyContextKey{}, ordinaryPolicy) + req = req.WithContext(ordinaryCtx) + prepared, _, err = PreparePelicanHTTPRequest(req, 67, 1, "http") + require.NoError(t, err) + require.Same(t, req, prepared) + require.Nil(t, ordinaryPolicy.guard, "ordinary requests keep their existing guard owner") +} diff --git a/backend/internal/service/pelican_borrow_ws_test.go b/backend/internal/service/pelican_borrow_ws_test.go new file mode 100644 index 000000000000..f061685e531f --- /dev/null +++ b/backend/internal/service/pelican_borrow_ws_test.go @@ -0,0 +1,138 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// Capture the shared pool's actual lease state when the socket closes. A +// one-shot test socket must be evicted while still leased, before it can be +// offered to a waiting business request. +type pelicanBorrowClosingConn struct { + *openAIWSCaptureConn + pool *openAIWSConnPool + accountID int64 + leasedConn *openAIWSConn + closedAfterRelease bool +} + +func (c *pelicanBorrowClosingConn) WriteJSON(ctx context.Context, value any) error { + ap := c.pool.getOrCreateAccountPool(c.accountID) + ap.mu.Lock() + for _, conn := range ap.conns { + if conn.ws == c { + c.leasedConn = conn + break + } + } + ap.mu.Unlock() + return c.openAIWSCaptureConn.WriteJSON(ctx, value) +} + +func (c *pelicanBorrowClosingConn) Close() error { + if c.leasedConn != nil { + c.closedAfterRelease = len(c.leasedConn.leaseCh) > 0 + } + return c.openAIWSCaptureConn.Close() +} + +func TestPelicanBorrowWSUsesAccountWithoutCustomerAnchor(t *testing.T) { + for _, model := range []string{"gpt-6-astra", "gpt-6.1-sol"} { + t.Run(model, func(t *testing.T) { + svc, account, borrow, upstream := borrowWSFixture(t, 1) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + capture := &pelicanGenerationCapture{account: account, openaiGateway: svc} + ctx := context.WithValue(WithAccountObservation(context.Background()), pelicanGenerationContextKey{}, capture) + c.Request = httptest.NewRequest("POST", "/internal/admin/pelican-tests", nil).WithContext(ctx) + seedBorrowWSTarget(t, svc, borrow, account, c, model) + headers, _, err := svc.buildOpenAIWSHeaders(ctx, c, account, "target-token", + svc.getOpenAIWSProtocolResolver().Resolve(account), true, "", "", "", model, "") + require.NoError(t, err) + wsURL, err := svc.buildOpenAIResponsesWSURL(account) + require.NoError(t, err) + turn, err := svc.prepareCodexGatewayBorrowWSTurn(ctx, c, account, wsURL, headers, + []byte(`{"model":"`+model+`","input":"one question"}`), model, "", account.Proxy.URL()) + require.NoError(t, err) + require.NotNil(t, turn) + require.True(t, turn.pelicanOneShot) + require.Nil(t, turn.store) + require.Equal(t, account.ID, turn.key.account) + require.Contains(t, turn.headers.Get("Cookie"), "__oailb=shared-native-route") + require.Zero(t, getAPIKeyIDFromContext(c)) + require.Empty(t, upstream.requests, "hot qualification must not prepare another route") + _, err = svc.prepareCodexGatewayBorrowWSTurn(ctx, c, account, wsURL, headers, + []byte(`{"model":"`+model+`"}`), model, "previous-response", account.Proxy.URL()) + require.ErrorContains(t, err, "new conversation") + }) + } +} + +func TestPelicanBorrowWSSharedForwardGeneratesOnceAndReleasesItsSocket(t *testing.T) { + for _, model := range []string{"gpt-6-astra", "gpt-6.1-sol"} { + t.Run(model, func(t *testing.T) { + gateway, account, borrow, probes := borrowWSFixture(t, 1) + gateway.cfg.Gateway.MaxLineSize = defaultMaxLineSize + account.Status, account.Schedulable = StatusError, false + repo := &pelicanGeneratorAccountRepo{accounts: map[int64]*Account{account.ID: account}} + gateway.accountRepo = repo + sender := &AccountTestService{accountRepo: repo, cfg: gateway.cfg, openaiGatewayService: gateway} + seed, _ := gin.CreateTestContext(httptest.NewRecorder()) + seed.Request = httptest.NewRequest("POST", "/internal/admin/pelican-tests", nil) + seedBorrowWSTarget(t, gateway, borrow, account, seed, model) + answer := "" + delta, err := json.Marshal(map[string]any{"type": "response.output_text.delta", "delta": answer}) + require.NoError(t, err) + completed, err := json.Marshal(map[string]any{"type": "response.completed", "response": map[string]any{ + "id": "resp_pelican", "model": model, "status": "completed", + "output": []any{map[string]any{"type": "message", "role": "assistant", "content": []any{map[string]any{"type": "output_text", "text": answer}}}}, + "usage": map[string]any{"input_tokens": 4, "output_tokens": 8}, + }}) + require.NoError(t, err) + conn := &pelicanBorrowClosingConn{openAIWSCaptureConn: &openAIWSCaptureConn{events: [][]byte{delta, completed}}, + pool: gateway.openaiWSPool, accountID: account.ID} + dialer := &borrowWSDialer{conns: []openAIWSClientConn{conn}} + gateway.openaiWSPool.setClientDialerForTest(dialer) + state := newPelicanExecutionState(context.Background(), newPelicanExecutionCoordinator(10), 10*time.Minute, time.Now) + defer state.finish() + ctx := context.WithValue(state.parent, pelicanExecutionContextKey{}, state) + result, err := sender.GeneratePelican(ctx, account.ID, model, "high") + require.NoError(t, err) + require.Equal(t, "complete", result.Status, result.Error) + require.Equal(t, answer, result.RawAnswer) + require.Equal(t, string(delta)+"\n"+string(completed)+"\n", result.RawResponse) + require.Equal(t, []int64{account.ID}, repo.reads) + require.Len(t, conn.writes, 1, "the real shared Forward sends the fixed prompt only once") + wire, err := json.Marshal(conn.writes[0]) + require.NoError(t, err) + require.Contains(t, string(wire), CodexGatewayBorrowPelicanPrompt) + require.Equal(t, model, conn.writes[0]["model"]) + require.Len(t, dialer.headers, 1) + require.Equal(t, "Bearer target-token", dialer.headers[0].Get("Authorization")) + require.Equal(t, "target-chatgpt", dialer.headers[0].Get("Chatgpt-Account-Id")) + require.Contains(t, dialer.headers[0].Get("Cookie"), "__oailb=shared-native-route") + require.Equal(t, []string{account.Proxy.URL()}, dialer.proxies) + require.Empty(t, probes.requests, "qualified routes must not issue extra probes") + snapshot := state.finish() + require.NotNil(t, snapshot.generationStartedAt) + require.Equal(t, "websocket", snapshot.invocation.Transport) + require.Equal(t, "responses", snapshot.invocation.Protocol) + require.Equal(t, model, snapshot.invocation.Model) + require.Equal(t, "high", snapshot.invocation.Effort) + require.True(t, snapshot.invocation.BorrowApplied) + require.True(t, strings.HasPrefix(snapshot.invocation.Endpoint, "wss://")) + require.True(t, conn.closed, "a one-shot admin test must release its borrowed socket") + require.NotNil(t, conn.leasedConn) + require.False(t, conn.closedAfterRelease, "evict before exposing the socket to a waiting business request") + require.Empty(t, borrow.wsAnchors.entries, "admin tests must not create customer continuation anchors") + }) + } +} diff --git a/backend/internal/service/pelican_execution.go b/backend/internal/service/pelican_execution.go new file mode 100644 index 000000000000..167df5903cd9 --- /dev/null +++ b/backend/internal/service/pelican_execution.go @@ -0,0 +1,373 @@ +package service + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" +) + +const ( + PelicanExecutionConcurrency = 10 + PelicanDefaultGenerationTimeoutSeconds = 600 + PelicanMinGenerationTimeoutSeconds = 60 + PelicanMaxGenerationTimeoutSeconds = 1800 + PelicanExecutionModeAccount = "account" + PelicanExecutionModeLegacyCache = "legacy_cache" +) + +// Every manual Pelican runner in this process shares the same outbound capacity. +// Preparation and generation acquire it only while an actual model request is +// in flight. A waiting preparation therefore cannot retain a target's permit +// while it waits for a source account or the remaining global permits. +var sharedPelicanExecution = newPelicanExecutionCoordinator(PelicanExecutionConcurrency) + +type pelicanExecutionContextKey struct{} +type pelicanExecutionLeaseContextKey struct{} + +type pelicanExecutionWaiter struct { + accountID int64 + ready chan struct{} + granted bool +} + +type pelicanExecutionCoordinator struct { + mu sync.Mutex + limit int + active int + accounts map[int64]bool + waiting []*pelicanExecutionWaiter +} + +func newPelicanExecutionCoordinator(limit int) *pelicanExecutionCoordinator { + return &pelicanExecutionCoordinator{limit: limit, accounts: make(map[int64]bool)} +} + +// scheduleLocked grants account and global capacity atomically. No waiter holds +// one resource while waiting for the other, and each account keeps FIFO order. +func (c *pelicanExecutionCoordinator) scheduleLocked() { + for i := 0; i < len(c.waiting) && c.active < c.limit; { + waiter := c.waiting[i] + if c.accounts[waiter.accountID] { + i++ + continue + } + c.waiting = append(c.waiting[:i], c.waiting[i+1:]...) + c.accounts[waiter.accountID] = true + c.active++ + waiter.granted = true + close(waiter.ready) + } +} + +func (c *pelicanExecutionCoordinator) acquire(ctx context.Context, accountID int64) (func(), error) { + if err := ctx.Err(); err != nil { + return nil, err + } + waiter := &pelicanExecutionWaiter{accountID: accountID, ready: make(chan struct{})} + c.mu.Lock() + c.waiting = append(c.waiting, waiter) + c.scheduleLocked() + c.mu.Unlock() + select { + case <-waiter.ready: + case <-ctx.Done(): + } + c.mu.Lock() + if err := ctx.Err(); err != nil { + if waiter.granted { + delete(c.accounts, accountID) + c.active-- + } else { + for i, pending := range c.waiting { + if pending == waiter { + c.waiting = append(c.waiting[:i], c.waiting[i+1:]...) + break + } + } + } + c.scheduleLocked() + c.mu.Unlock() + return nil, err + } + c.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + c.mu.Lock() + delete(c.accounts, accountID) + c.active-- + c.scheduleLocked() + c.mu.Unlock() + }) + }, nil +} + +type pelicanExecutionLease struct { + accountID int64 + active atomic.Bool +} + +// PelicanInvocation records the resolved request at the actual sending boundary. +// It deliberately contains no group/client simulation settings or credentials. +type PelicanInvocation struct { + Platform string + Model string + Effort string + Endpoint string + Protocol string + Transport string + BorrowApplied bool +} + +type pelicanExecutionSnapshot struct { + invocation PelicanInvocation + generationStartedAt *time.Time + queueDurationMS int64 + preparationDurationMS int64 + generationDurationMS int64 + generationTimedOut bool +} + +type pelicanExecutionState struct { + mu sync.Mutex + coordinator *pelicanExecutionCoordinator + parent context.Context + operationCancel context.CancelCauseFunc + now func() time.Time + withTimeout func(context.Context, time.Duration) (context.Context, context.CancelFunc) + budget time.Duration + preparationAt time.Time + preparationWait time.Duration + queueDuration time.Duration + generationAt time.Time + finishedAt time.Time + invocation PelicanInvocation + budgetCtx context.Context + budgetCancel context.CancelFunc + budgetStopLink func() bool + requestCancels []context.CancelFunc + onPhase func(string, pelicanExecutionSnapshot) +} + +func newPelicanExecutionState(parent context.Context, coordinator *pelicanExecutionCoordinator, budget time.Duration, now func() time.Time) *pelicanExecutionState { + operationCtx, operationCancel := context.WithCancelCause(parent) + return &pelicanExecutionState{parent: operationCtx, operationCancel: operationCancel, coordinator: coordinator, budget: budget, now: now, + withTimeout: context.WithTimeout} +} + +func pelicanExecutionFromContext(ctx context.Context) *pelicanExecutionState { + if ctx == nil { + return nil + } + state, _ := ctx.Value(pelicanExecutionContextKey{}).(*pelicanExecutionState) + return state +} + +// CancelPelicanExecution stops the whole fixed-account operation when a shared +// sender loses its distributed account slot. The batch remains able to persist +// the observation. Outside a Pelican operation this helper has no side effects. +func CancelPelicanExecution(ctx context.Context, cause error) { + if state := pelicanExecutionFromContext(ctx); state != nil { + state.operationCancel(cause) + } +} + +// AcquirePelicanExecution coordinates an actual model request, including a +// nested borrow source/target probe, with all other manual Pelican requests. +// Call it after route preparation and retain its lease until the response body +// has been consumed/closed. Nested wrappers for the same request/account reuse +// the lease; a leased context must not be shared by parallel model requests. +// Ordinary business requests have no Pelican state and retain their usual path. +func AcquirePelicanExecution(ctx context.Context, accountID int64) (context.Context, func(), error) { + state := pelicanExecutionFromContext(ctx) + if state == nil { + return ctx, func() {}, nil + } + if err := ctx.Err(); err != nil { + return ctx, nil, err + } + if lease, _ := ctx.Value(pelicanExecutionLeaseContextKey{}).(*pelicanExecutionLease); lease != nil && lease.active.Load() { + if lease.accountID == accountID { + return ctx, func() {}, nil + } + // Route preparation belongs before the final request lease. Refusing an + // incorrectly nested account prevents every global permit being held by + // parents waiting for child requests, rather than silently deadlocking. + return ctx, nil, errors.New("pelican route preparation must finish before acquiring another account's request lease") + } + finishWait := BeginPelicanQueueWait(ctx) + release, err := state.coordinator.acquire(ctx, accountID) + finishWait() + if err != nil { + return ctx, nil, err + } + lease := &pelicanExecutionLease{accountID: accountID} + lease.active.Store(true) + return context.WithValue(ctx, pelicanExecutionLeaseContextKey{}, lease), func() { + if lease.active.Swap(false) { + release() + } + }, nil +} + +// BeginPelicanQueueWait accounts for ordinary business-account capacity waits +// using the same clock as the global/account coordinator. Before the first +// question is sent, waiting time is excluded from route preparation. Once the +// generation clock starts, retry waits remain inside its original budget. +func BeginPelicanQueueWait(ctx context.Context) func() { + state := pelicanExecutionFromContext(ctx) + if state == nil { + return func() {} + } + started := state.now() + var once sync.Once + return func() { + once.Do(func() { + finished := state.now() + state.mu.Lock() + if state.generationAt.IsZero() { + waited := nonnegativePelicanDuration(finished.Sub(started)) + state.queueDuration += waited + if !state.preparationAt.IsZero() { + state.preparationWait += waited + } + } + state.mu.Unlock() + }) + } +} + +// BeginPelicanPreparation marks account/credential/route preparation. Borrow +// preparation keeps its own existing budgets; none of this starts generation. +func BeginPelicanPreparation(ctx context.Context) { + state := pelicanExecutionFromContext(ctx) + if state == nil { + return + } + state.mu.Lock() + if !state.preparationAt.IsZero() { + state.mu.Unlock() + return + } + state.preparationAt = state.now().UTC() + callback, snapshot := state.onPhase, state.snapshotLocked(state.preparationAt) + state.mu.Unlock() + if callback != nil { + callback("preparing", snapshot) + } +} + +// BeginPelicanGeneration starts the one generation clock immediately before the +// first actual question request, after route and outbound capacity preparation. +// Repeated same-account compatibility sends inherit the first clock. It returns +// a context retaining the current sender's values and cancellation boundaries. +func BeginPelicanGeneration(ctx context.Context) context.Context { + state := pelicanExecutionFromContext(ctx) + if state == nil { + return ctx + } + state.mu.Lock() + first := state.generationAt.IsZero() + if first { + state.generationAt = state.now().UTC() + state.budgetCtx, state.budgetCancel = state.withTimeout(state.parent, state.budget) + state.budgetStopLink = context.AfterFunc(state.budgetCtx, func() { + if errors.Is(context.Cause(state.budgetCtx), context.DeadlineExceeded) { + state.operationCancel(context.DeadlineExceeded) + } + }) + } + budgetCtx := state.budgetCtx + deadline, hasDeadline := budgetCtx.Deadline() + var generationCtx context.Context + var cancel context.CancelFunc + if hasDeadline { + generationCtx, cancel = context.WithDeadline(ctx, deadline) + } else { + generationCtx, cancel = context.WithCancel(ctx) + } + stopLink := context.AfterFunc(budgetCtx, cancel) + state.requestCancels = append(state.requestCancels, func() { stopLink(); cancel() }) + callback, snapshot := state.onPhase, state.snapshotLocked(state.generationAt) + state.mu.Unlock() + if budgetCtx.Err() != nil { + cancel() + } + if first && callback != nil { + callback("generating", snapshot) + } + return generationCtx +} + +func RecordPelicanInvocation(ctx context.Context, invocation PelicanInvocation) { + if state := pelicanExecutionFromContext(ctx); state != nil { + state.mu.Lock() + state.invocation = invocation + state.mu.Unlock() + } +} + +func nonnegativePelicanDuration(value time.Duration) time.Duration { + if value < 0 { + return 0 + } + return value +} + +func (s *pelicanExecutionState) snapshotLocked(now time.Time) pelicanExecutionSnapshot { + snapshot := pelicanExecutionSnapshot{invocation: s.invocation, queueDurationMS: s.queueDuration.Milliseconds()} + endPreparation := now + if !s.generationAt.IsZero() { + started := s.generationAt + snapshot.generationStartedAt = &started + endPreparation = started + snapshot.generationDurationMS = nonnegativePelicanDuration(now.Sub(started)).Milliseconds() + snapshot.generationTimedOut = errors.Is(context.Cause(s.budgetCtx), context.DeadlineExceeded) + } + if !s.preparationAt.IsZero() { + snapshot.preparationDurationMS = nonnegativePelicanDuration(endPreparation.Sub(s.preparationAt) - s.preparationWait).Milliseconds() + } + return snapshot +} + +func (s *pelicanExecutionState) finish() pelicanExecutionSnapshot { + s.mu.Lock() + s.finishedAt = s.now().UTC() + snapshot := s.snapshotLocked(s.finishedAt) + cancels := s.requestCancels + budgetCancel := s.budgetCancel + budgetStopLink := s.budgetStopLink + s.requestCancels = nil + s.mu.Unlock() + for _, cancel := range cancels { + cancel() + } + if budgetCancel != nil { + budgetCancel() + } + if budgetStopLink != nil { + budgetStopLink() + } + s.operationCancel(context.Canceled) + return snapshot +} + +func applyPelicanExecutionSnapshot(result *CodexGatewayBorrowTestResult, snapshot pelicanExecutionSnapshot) { + if snapshot.invocation.Platform != "" { + result.Platform = snapshot.invocation.Platform + } + if snapshot.invocation.Model != "" { + result.UpstreamModel = snapshot.invocation.Model + result.Effort = snapshot.invocation.Effort + } + result.ActualEndpoint = snapshot.invocation.Endpoint + result.ActualProtocol = snapshot.invocation.Protocol + result.ActualTransport = snapshot.invocation.Transport + result.BorrowApplied = snapshot.invocation.BorrowApplied + result.GenerationStartedAt = snapshot.generationStartedAt + result.QueueDurationMS = snapshot.queueDurationMS + result.PreparationDurationMS = snapshot.preparationDurationMS + result.GenerationDurationMS = snapshot.generationDurationMS +} diff --git a/backend/internal/service/pelican_test_catalog.go b/backend/internal/service/pelican_test_catalog.go new file mode 100644 index 000000000000..704c3f2fbab0 --- /dev/null +++ b/backend/internal/service/pelican_test_catalog.go @@ -0,0 +1,413 @@ +package service + +import ( + "context" + "errors" + "fmt" + "slices" + "sort" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" +) + +const ( + PelicanTestMaxConcurrency = PelicanExecutionConcurrency + PelicanTestDefaultGenerationTimeoutSeconds = PelicanDefaultGenerationTimeoutSeconds + PelicanTestMinGenerationTimeoutSeconds = PelicanMinGenerationTimeoutSeconds + PelicanTestMaxGenerationTimeoutSeconds = PelicanMaxGenerationTimeoutSeconds + pelicanTestCatalogMaxAccounts = 5000 +) + +var ErrPelicanTestOptionsInvalidRequest = errors.New("invalid pelican test options request") + +type PelicanTestModelOption struct { + ID string `json:"id"` + DisplayName string `json:"display_name"` + UpstreamModel string `json:"upstream_model"` + ReasoningEfforts []string `json:"reasoning_efforts"` + DefaultEffort string `json:"default_effort"` + TextSupported bool `json:"text_supported"` + CapabilityReason string `json:"capability_reason"` +} + +type PelicanTestAccountOption struct { + ID int64 `json:"id"` + Name string `json:"name"` + Platform string `json:"platform"` + Type string `json:"type"` + Status string `json:"status"` + Schedulable bool `json:"schedulable"` + ParentAccountID *int64 `json:"parent_account_id"` + ProxyID *int64 `json:"proxy_id"` + ProxyName string `json:"proxy_name"` + RateLimitedUntil *time.Time `json:"rate_limited_until"` + OverloadUntil *time.Time `json:"overload_until"` + TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until"` + TempUnschedulableReason string `json:"temp_unschedulable_reason"` + ErrorMessage string `json:"error_message"` + CapabilityReason string `json:"capability_reason"` + ManualModelAllowed bool `json:"manual_model_allowed"` + DefaultModelID string `json:"default_model_id"` + Models []PelicanTestModelOption `json:"models"` +} + +type PelicanTestOptions struct { + GeneratedAt time.Time `json:"generated_at"` + MaxConcurrency int `json:"max_concurrency"` + DefaultGenerationTimeoutSeconds int `json:"default_generation_timeout_seconds"` + MinGenerationTimeoutSeconds int `json:"min_generation_timeout_seconds"` + MaxGenerationTimeoutSeconds int `json:"max_generation_timeout_seconds"` + Accounts []PelicanTestAccountOption `json:"accounts"` +} + +type pelicanTestCatalogAccountReader interface { + GetByIDs(context.Context, []int64) ([]*Account, error) +} + +// PelicanTestCatalog depends only on a local account reader. It cannot refresh a +// token, discover a remote catalog, prepare a borrowed route, or mutate accounts. +type PelicanTestCatalog struct { + accounts pelicanTestCatalogAccountReader + now func() time.Time +} + +func NewPelicanTestCatalog(accounts AccountRepository) *PelicanTestCatalog { + return &PelicanTestCatalog{accounts: accounts, now: time.Now} +} + +func (s *PelicanTestCatalog) Options(ctx context.Context, accountIDs []int64) (*PelicanTestOptions, error) { + options := &PelicanTestOptions{ + GeneratedAt: s.now().UTC(), MaxConcurrency: PelicanTestMaxConcurrency, + DefaultGenerationTimeoutSeconds: PelicanTestDefaultGenerationTimeoutSeconds, + MinGenerationTimeoutSeconds: PelicanTestMinGenerationTimeoutSeconds, + MaxGenerationTimeoutSeconds: PelicanTestMaxGenerationTimeoutSeconds, + Accounts: make([]PelicanTestAccountOption, 0, len(accountIDs)), + } + if len(accountIDs) > pelicanTestCatalogMaxAccounts { + return nil, fmt.Errorf("%w: account_ids must contain at most %d accounts", ErrPelicanTestOptionsInvalidRequest, pelicanTestCatalogMaxAccounts) + } + ids := make([]int64, 0, len(accountIDs)) + seen := make(map[int64]bool, len(accountIDs)) + for _, id := range accountIDs { + if id <= 0 { + return nil, fmt.Errorf("%w: account_ids must be positive", ErrPelicanTestOptionsInvalidRequest) + } + if !seen[id] { + ids = append(ids, id) + seen[id] = true + } + } + if len(ids) == 0 { + return options, nil + } + if s.accounts == nil { + return nil, errors.New("pelican test local account reader is unavailable") + } + accounts, err := s.accounts.GetByIDs(ctx, ids) + if err != nil { + return nil, fmt.Errorf("read pelican test local accounts: %w", err) + } + byID := make(map[int64]*Account, len(accounts)) + for _, account := range accounts { + if account != nil { + byID[account.ID] = account + } + } + for _, id := range ids { + account := byID[id] + if account == nil { + options.Accounts = append(options.Accounts, PelicanTestAccountOption{ + ID: id, Status: "missing", CapabilityReason: "selected account was not found", Models: []PelicanTestModelOption{}, + }) + continue + } + options.Accounts = append(options.Accounts, pelicanTestAccountOptions(account)) + } + return options, nil +} + +func pelicanTestAccountOptions(account *Account) PelicanTestAccountOption { + option := PelicanTestAccountOption{ + ID: account.ID, Name: account.Name, Platform: account.Platform, Type: account.Type, + Status: account.Status, Schedulable: account.Schedulable, ParentAccountID: account.ParentAccountID, + ProxyID: account.ProxyID, RateLimitedUntil: account.RateLimitResetAt, OverloadUntil: account.OverloadUntil, + TempUnschedulableUntil: account.TempUnschedulableUntil, TempUnschedulableReason: account.TempUnschedulableReason, + ErrorMessage: account.ErrorMessage, CapabilityReason: pelicanTestAccountCapabilityReason(account), + ManualModelAllowed: true, Models: []PelicanTestModelOption{}, + } + if account.Proxy != nil { + option.ProxyName = account.Proxy.Name + } + for _, id := range pelicanTestLocalModelIDs(account) { + model := PelicanTestModelOptions(account, id) + option.Models = append(option.Models, model) + if option.DefaultModelID == "" && model.TextSupported { + option.DefaultModelID = id + } + } + if option.CapabilityReason == "" && len(option.Models) > 0 && option.DefaultModelID == "" { + option.CapabilityReason = "the account's locally configured models do not support text generation" + } + return option +} + +// These IDs are candidates, not evidence of a successful live model call. Exact +// saved mapping keys are retained even when they have no text capability. A +// wildcard is expanded only against local platform data and saved source-bound +// directory metadata; administrators can also enter a concrete ID themselves. +func pelicanTestLocalModelIDs(account *Account) []string { + defaults := pelicanTestPlatformModelIDs(account) + mapping := account.GetModelMapping() + useMapping := len(mapping) > 0 && !account.IsOpenAIPassthroughEnabled() && (account.Platform != PlatformAntigravity || account.Type != AccountTypeUpstream) + ids := make(map[string]bool) + add := func(id string) { + id = strings.TrimSpace(id) + if id == "" || strings.ContainsAny(id, "*\x00") || len(id) > 512 { + return + } + if !useMapping || account.IsModelSupported(id) { + ids[id] = true + } + } + for id := range mapping { + add(id) + } + for _, id := range defaults { + add(id) + } + if snapshot := account.GetUpstreamModelMetadataSnapshot(); upstreamModelMetadataSourceMatches(account, snapshot) { + for id := range snapshot.Models { + add(id) + } + } + ordered := make([]string, 0, len(ids)) + for id := range ids { + ordered = append(ordered, id) + } + sort.Strings(ordered) + return ordered +} + +func pelicanTestPlatformModelIDs(account *Account) []string { + switch account.Platform { + case PlatformKimi, PlatformZhipu: + return nil // These platforms publish no local default catalog. + case PlatformDeepseek: + return []string{"deepseek-v4-pro", "deepseek-v4-flash", "deepseek-flash"} + case PlatformMiniMax: + return []string{"MiniMax-M3", "MiniMax-M2.7", "MiniMax-M2.5"} + case PlatformGemini: + if account.IsGeminiGoogleOne() { + ids := make([]string, 0, len(geminicli.GoogleOneModels)) + for _, model := range geminicli.GoogleOneModels { + ids = append(ids, model.ID) + } + return ids + } + case PlatformOpenAI, PlatformAnthropic, PlatformAntigravity, PlatformGrok, PlatformOpenCodeGo, PlatformTypeSafe: + default: + return nil + } + return defaultModelsListCandidateIDs(account.Platform) +} + +func pelicanTestAccountCapabilityReason(account *Account) string { + if account == nil { + return "selected account was not found" + } + if account.Platform == PlatformTypeSafe { + return "TypeSafe System One provides a structured native response and does not support arbitrary HTML text generation" + } + switch account.Platform { + case PlatformOpenAI, PlatformAnthropic, PlatformGemini, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo: + default: + return fmt.Sprintf("platform %q has no text generation sender", account.Platform) + } + supported := false + switch account.Platform { + case PlatformAnthropic: + supported = slices.Contains([]string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken, AccountTypeBedrock, AccountTypeServiceAccount}, account.Type) + case PlatformOpenAI: + supported = slices.Contains([]string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken}, account.Type) + case PlatformGemini: + supported = slices.Contains([]string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeServiceAccount}, account.Type) + case PlatformGrok: + supported = account.Type == AccountTypeAPIKey || account.Type == AccountTypeOAuth + case PlatformAntigravity: + supported = account.Type == AccountTypeOAuth || account.Type == AccountTypeUpstream + case PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax, PlatformOpenCodeGo: + supported = account.Type == AccountTypeAPIKey + } + if !supported { + return fmt.Sprintf("account type %q has no text generation sender for platform %q", account.Type, account.Platform) + } + return "" +} + +func pelicanTestUpstreamModel(account *Account, model string) (string, bool) { + if account == nil { + return model, false + } + switch { + case account.IsOpenAI(): + return resolveOpenAIAccountUpstreamModelForRequest(account, model, false), true + case account.Platform == PlatformAntigravity && account.Type == AccountTypeUpstream: + return model, true + case account.Platform == PlatformAnthropic && account.IsBedrock(): + return ResolveBedrockModelID(account, model) + case account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount: + if mapped, matched := account.ResolveMappedModel(model); matched { + return mapped, true + } + return normalizeVertexAnthropicModelID(claude.NormalizeModelID(model)), true + case account.Platform == PlatformAnthropic && account.IsOAuth(): + return claude.NormalizeModelID(model), true + case account.Platform == PlatformAntigravity: + if mapped, ok := resolveGeminiThinkingVariantForLevel(account, model, ""); ok { + return mapped, true + } + mapped := mapAntigravityModel(account, model) + return mapped, mapped != "" + case account.Platform == PlatformGemini && account.IsOAuth(): + return model, true + default: + return account.GetMappedModel(model), true + } +} + +func pelicanTestLocalMetadata(account *Account, model string) (UpstreamModelMetadata, bool) { + snapshot := account.GetUpstreamModelMetadataSnapshot() + if !upstreamModelMetadataSourceMatches(account, snapshot) { + return UpstreamModelMetadata{}, false + } + metadata, ok := snapshot.Models[model] + return metadata, ok +} + +// PelicanTestTextCapabilityReason records unsupported choices without hiding an +// account, changing its configuration, or applying a group's routing policy. +func PelicanTestTextCapabilityReason(account *Account, model string) string { + if reason := pelicanTestAccountCapabilityReason(account); reason != "" { + return reason + } + model = strings.TrimSpace(model) + if model == "" || strings.ContainsAny(model, "*\x00") || len(model) > 512 { + return "a concrete model ID is required" + } + mapped, ok := pelicanTestUpstreamModel(account, model) + if !ok || strings.TrimSpace(mapped) == "" { + return fmt.Sprintf("model %q is unsupported by the account's native text sender", model) + } + // Test both IDs independently, using no second mapping of the resolved target. + for _, id := range []string{model, mapped} { + if !AccountTestSupportsTextConversation(&Account{Platform: PlatformOpenAI}, id) { + return fmt.Sprintf("model %q does not support text generation", id) + } + } + if metadata, known := pelicanTestLocalMetadata(account, mapped); known && len(metadata.InputModalities) > 0 && !slices.Contains(metadata.InputModalities, "text") { + return fmt.Sprintf("model %q declares no text input capability in the account's local directory", mapped) + } + return "" +} + +func pelicanTestUsesClaudeEffort(account *Account, model string) bool { + if account.Platform == PlatformAnthropic { + return true + } + if !claude.IsClaudeModel(model) { + return false + } + if account.Platform == PlatformAntigravity && account.Type == AccountTypeUpstream { + return true + } + if account.IsOpenCodeGo() { + return openCodeGoNativeProtocol(account, model) == APIProtocolAnthropic + } + return account.IsCNProvider() && account.GetAPIProtocol() == APIProtocolAnthropic +} + +func pelicanTestSupportsReasoningWire(account *Account, model string) bool { + if pelicanTestUsesClaudeEffort(account, model) { + return !account.IsBedrock() || bedrockKeepsOutputConfigEffort(model) + } + switch { + case account.IsOpenAI(): + return true + case account.IsOpenCodeGo(): + protocol := openCodeGoNativeProtocol(account, model) + return protocol == APIProtocolChatCompletions || protocol == APIProtocolResponses + case account.IsCNProvider(): + // A Pelican task sends one native request. Adaptive accounts do not have + // the ordinary connection test's every-endpoint effort restriction. + return account.GetAPIProtocol() != APIProtocolAnthropic + case account.Platform == PlatformGrok: + return grokSupportsReasoningEffort(accountTestGrokUpstreamModel(model)) + default: + return false + } +} + +// PelicanTestModelOptions uses the account's actual native mapping and locally +// saved capability metadata. An omitted effort keeps the platform's own default. +func PelicanTestModelOptions(account *Account, model string) PelicanTestModelOption { + model = strings.TrimSpace(model) + mapped, _ := pelicanTestUpstreamModel(account, model) + option := PelicanTestModelOption{ID: model, DisplayName: model, UpstreamModel: mapped, ReasoningEfforts: []string{}} + option.CapabilityReason = PelicanTestTextCapabilityReason(account, model) + option.TextSupported = option.CapabilityReason == "" + if account == nil { + return option + } + metadata, known := pelicanTestLocalMetadata(account, mapped) + if known && mapped == model && strings.TrimSpace(metadata.DisplayName) != "" { + option.DisplayName = metadata.DisplayName + } else if mapped == model { + for _, local := range openai.DefaultModels { + if local.ID == model { + option.DisplayName = local.DisplayName + break + } + } + } + if !option.TextSupported || !pelicanTestSupportsReasoningWire(account, mapped) { + return option + } + if pelicanTestUsesClaudeEffort(account, mapped) { + if !known || metadata.Reasoning == nil || *metadata.Reasoning { + option.ReasoningEfforts = claude.EffortLevelsForModel(mapped) + } + } else if known && (metadata.Reasoning != nil || len(metadata.SupportedReasoningLevels) > 0) { + if metadata.Reasoning == nil || *metadata.Reasoning { + option.ReasoningEfforts = normalizeReasoningLevels(metadata.SupportedReasoningLevels) + if isOpenAIGPT6SolOrLunaModel(mapped) { + option.ReasoningEfforts = intersectOrderedStrings(option.ReasoningEfforts, openai.GPT6APIReasoningEfforts()) + } + } + } else if account.IsOpenAIApiKey() && isOpenAIGPT6SolOrLunaModel(mapped) { + option.ReasoningEfforts = openai.GPT6APIReasoningEfforts() + } else { + var levels []configuredCodexReasoningLevel + switch { + case account.Platform == PlatformGrok: + levels = configuredCodexGrokReasoningLevels(accountTestGrokUpstreamModel(mapped)) + case account.IsOpenAI() && isOpenAICodexReasoningGPTModel(mapped): + levels = configuredCodexGPTReasoningLevels(mapped) + } + for _, level := range levels { + if level.Effort != "ultra" { + option.ReasoningEfforts = append(option.ReasoningEfforts, level.Effort) + } + } + } + if option.ReasoningEfforts == nil { + option.ReasoningEfforts = []string{} + } + if slices.Contains(option.ReasoningEfforts, "high") { + option.DefaultEffort = "high" + } + return option +} diff --git a/backend/internal/service/pelican_test_catalog_test.go b/backend/internal/service/pelican_test_catalog_test.go new file mode 100644 index 000000000000..993f35a0a914 --- /dev/null +++ b/backend/internal/service/pelican_test_catalog_test.go @@ -0,0 +1,214 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type pelicanTestCatalogReaderStub struct { + accounts []*Account + reads int + ids []int64 +} + +func (r *pelicanTestCatalogReaderStub) GetByIDs(_ context.Context, ids []int64) ([]*Account, error) { + r.reads++ + r.ids = append([]int64(nil), ids...) + return r.accounts, nil +} + +func TestPelicanTestCatalogReadsOnlyRequestedLocalAccountsAndRetainsStates(t *testing.T) { + now := time.Date(2026, 10, 8, 3, 0, 0, 0, time.UTC) + limitedUntil := now.Add(time.Hour) + reader := &pelicanTestCatalogReaderStub{accounts: []*Account{ + {ID: 2, Name: "paused", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusDisabled, Schedulable: false}, + {ID: 1, Name: "limited", Platform: PlatformGemini, Type: AccountTypeServiceAccount, Status: StatusError, Schedulable: true, ErrorMessage: "full error body token=administrator-visible", RateLimitResetAt: &limitedUntil}, + }} + catalog := &PelicanTestCatalog{accounts: reader, now: func() time.Time { return now }} + defaults, err := catalog.Options(context.Background(), nil) + require.NoError(t, err) + require.Equal(t, now, defaults.GeneratedAt) + require.Equal(t, 10, defaults.MaxConcurrency) + require.Equal(t, 600, defaults.DefaultGenerationTimeoutSeconds) + require.Equal(t, 60, defaults.MinGenerationTimeoutSeconds) + require.Equal(t, 1800, defaults.MaxGenerationTimeoutSeconds) + require.Empty(t, defaults.Accounts) + require.Zero(t, reader.reads, "opening the page cannot even enumerate accounts implicitly") + + options, err := catalog.Options(context.Background(), []int64{2, 1, 999, 2}) + require.NoError(t, err) + require.Equal(t, 1, reader.reads) + require.Equal(t, []int64{2, 1, 999}, reader.ids, "one local batch read, without scheduling filters") + require.Len(t, options.Accounts, 3) + require.Equal(t, StatusDisabled, options.Accounts[0].Status) + require.False(t, options.Accounts[0].Schedulable) + require.Equal(t, AccountTypeServiceAccount, options.Accounts[1].Type) + require.Equal(t, reader.accounts[1].ErrorMessage, options.Accounts[1].ErrorMessage) + require.Equal(t, &limitedUntil, options.Accounts[1].RateLimitedUntil) + require.Equal(t, int64(999), options.Accounts[2].ID) + require.Equal(t, "missing", options.Accounts[2].Status) + require.Contains(t, options.Accounts[2].CapabilityReason, "not found") + _, err = catalog.Options(context.Background(), []int64{0}) + require.ErrorIs(t, err, ErrPelicanTestOptionsInvalidRequest) + require.Equal(t, 1, reader.reads) +} + +func TestPelicanTestCatalogUsesMappingAndSourceBoundLocalMetadataWithoutMutation(t *testing.T) { + reasoning := true + account := &Account{ID: 5, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "base_url": "https://catalog.example", "model_mapping": map[string]any{ + "sol-alias": "gpt-6.1-sol", "custom": "custom-upstream", "media": "gpt-image-2", "gpt-6-*": "gpt-6-astra", + }, + }} + account.SetUpstreamModelMetadataSnapshot(UpstreamModelMetadataSnapshot{ + SourceIdentity: UpstreamModelMetadataSourceIdentity(account), Models: map[string]UpstreamModelMetadata{ + "custom-upstream": {ID: "custom-upstream", Reasoning: &reasoning, SupportedReasoningLevels: []string{"low", "high"}, InputModalities: []string{"text"}}, + "not-mapped": {ID: "not-mapped", InputModalities: []string{"text"}}, + }, + }) + beforeCredentials, err := json.Marshal(account.Credentials) + require.NoError(t, err) + beforeExtra, err := json.Marshal(account.Extra) + require.NoError(t, err) + options := pelicanTestAccountOptions(account) + byID := make(map[string]PelicanTestModelOption) + for _, model := range options.Models { + byID[model.ID] = model + } + require.Contains(t, byID, "gpt-6-astra") + require.Contains(t, byID, "sol-alias") + require.NotContains(t, byID, "gpt-6-*") + require.NotContains(t, byID, "not-mapped") + require.Equal(t, "gpt-6.1-sol", byID["sol-alias"].UpstreamModel) + require.Equal(t, "high", byID["sol-alias"].DefaultEffort) + require.Equal(t, []string{"low", "high"}, byID["custom"].ReasoningEfforts) + require.Equal(t, "high", byID["custom"].DefaultEffort) + require.False(t, byID["media"].TextSupported) + require.Contains(t, byID["media"].CapabilityReason, "gpt-image-2") + afterCredentials, err := json.Marshal(account.Credentials) + require.NoError(t, err) + afterExtra, err := json.Marshal(account.Extra) + require.NoError(t, err) + require.JSONEq(t, string(beforeCredentials), string(afterCredentials)) + require.JSONEq(t, string(beforeExtra), string(afterExtra)) + + account.Credentials["base_url"] = "https://different-source.example" + changed := PelicanTestModelOptions(account, "custom") + require.Empty(t, changed.ReasoningEfforts, "old endpoint metadata must not claim current reasoning capabilities") + require.Empty(t, changed.DefaultEffort) +} + +func TestPelicanTestCatalogDefaultsNativeCapabilitiesAndSpecificUnsupportedReasons(t *testing.T) { + for _, test := range []struct { + name string + platform string + kind string + model string + high bool + }{ + {"astra api", PlatformOpenAI, AccountTypeAPIKey, "gpt-6-astra", true}, + {"sol oauth", PlatformOpenAI, AccountTypeOAuth, "gpt-6.1-sol", true}, + {"claude setup token", PlatformAnthropic, AccountTypeSetupToken, "claude-opus-5-5", true}, + {"vertex", PlatformAnthropic, AccountTypeServiceAccount, "claude-opus-5-5", true}, + {"gemini", PlatformGemini, AccountTypeAPIKey, "gemini-2.5-flash", false}, + {"antigravity", PlatformAntigravity, AccountTypeOAuth, "gemini-3-flash", false}, + {"kimi", PlatformKimi, AccountTypeAPIKey, "kimi-k2.5", false}, + {"zhipu", PlatformZhipu, AccountTypeAPIKey, "glm-5", false}, + {"deepseek", PlatformDeepseek, AccountTypeAPIKey, "deepseek-v4-pro", false}, + {"minimax", PlatformMiniMax, AccountTypeAPIKey, "MiniMax-M3", false}, + {"opencode", PlatformOpenCodeGo, AccountTypeAPIKey, DefaultOpenCodeGoTestModel, false}, + } { + t.Run(test.name, func(t *testing.T) { + account := &Account{Platform: test.platform, Type: test.kind} + option := PelicanTestModelOptions(account, test.model) + require.True(t, option.TextSupported, option.CapabilityReason) + if test.high { + require.Contains(t, option.ReasoningEfforts, "high") + require.Equal(t, "high", option.DefaultEffort) + } else { + require.Empty(t, option.DefaultEffort, "models without high keep their normal default") + } + }) + } + unsupported := PelicanTestModelOptions(&Account{Platform: PlatformTypeSafe, Type: AccountTypeAPIKey}, "jev-latest") + require.False(t, unsupported.TextSupported) + require.Contains(t, unsupported.CapabilityReason, "System One") + unknown := PelicanTestModelOptions(&Account{Platform: "legacy-unknown", Type: AccountTypeAPIKey}, "model") + require.False(t, unknown.TextSupported) + require.Contains(t, unknown.CapabilityReason, "legacy-unknown") + manual := PelicanTestModelOptions(&Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, "vendor/concrete-model") + require.True(t, manual.TextSupported) + require.Empty(t, manual.DefaultEffort) +} + +func TestPelicanTestCatalogPassthroughAndDeclaredNoTextUseEffectiveModel(t *testing.T) { + account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "model_mapping": map[string]any{"gpt-6.1-sol": "gpt-image-2"}, + }, Extra: map[string]any{"openai_passthrough": true}} + require.True(t, account.IsOpenAIPassthroughEnabled()) + option := PelicanTestModelOptions(account, "gpt-6.1-sol") + require.Equal(t, "gpt-6.1-sol", option.UpstreamModel) + require.True(t, option.TextSupported) + require.Equal(t, "high", option.DefaultEffort) + account.SetUpstreamModelMetadataSnapshot(UpstreamModelMetadataSnapshot{ + SourceIdentity: UpstreamModelMetadataSourceIdentity(account), Models: map[string]UpstreamModelMetadata{ + "gpt-6.1-sol": {ID: "gpt-6.1-sol", InputModalities: []string{"image"}}, + }, + }) + option = PelicanTestModelOptions(account, "gpt-6.1-sol") + require.False(t, option.TextSupported) + require.Contains(t, option.CapabilityReason, "no text input capability") +} + +func TestPelicanTestCatalogNativeUpstreamAndAdaptiveUseBusinessEffortCapabilities(t *testing.T) { + upstream := &Account{Platform: PlatformAntigravity, Type: AccountTypeUpstream, Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-opus-5-5": "gemini-3-flash"}, + }} + option := PelicanTestModelOptions(upstream, "claude-opus-5-5") + require.True(t, option.TextSupported, option.CapabilityReason) + require.Equal(t, "claude-opus-5-5", option.UpstreamModel, "native upstream forwarding preserves the request model") + require.Contains(t, option.ReasoningEfforts, "high") + require.Equal(t, "high", option.DefaultEffort) + + adaptive := &Account{Platform: PlatformKimi, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_protocol": APIProtocolAdaptive}} + reasoning := true + adaptive.SetUpstreamModelMetadataSnapshot(UpstreamModelMetadataSnapshot{ + SourceIdentity: UpstreamModelMetadataSourceIdentity(adaptive), Models: map[string]UpstreamModelMetadata{ + "kimi-k2.5": {ID: "kimi-k2.5", Reasoning: &reasoning, SupportedReasoningLevels: []string{"medium", "high"}}, + }, + }) + option = PelicanTestModelOptions(adaptive, "kimi-k2.5") + require.Equal(t, []string{"medium", "high"}, option.ReasoningEfforts, "single native request can use the declared effort on adaptive accounts") + require.Equal(t, "high", option.DefaultEffort) + adaptive.Credentials["api_protocol"] = APIProtocolAnthropic + option = PelicanTestModelOptions(adaptive, "claude-opus-5-5") + require.True(t, option.TextSupported) + require.Contains(t, option.ReasoningEfforts, "high", "native Anthropic adapter retains Claude output_config.effort") + + invalidType := PelicanTestModelOptions(&Account{Platform: PlatformOpenAI, Type: AccountTypeBedrock}, "gpt-6-astra") + require.False(t, invalidType.TextSupported) + require.Contains(t, invalidType.CapabilityReason, AccountTypeBedrock) +} + +func TestPelicanTestCatalogOpenAIModelCapabilityUsesNativeResolvedAlias(t *testing.T) { + account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ + "base_url": "https://local-catalog.example", "model_mapping": map[string]any{"public-model": "gpt-6.1-sol-high"}, + }} + reasoning := false + account.SetUpstreamModelMetadataSnapshot(UpstreamModelMetadataSnapshot{ + SourceIdentity: UpstreamModelMetadataSourceIdentity(account), Models: map[string]UpstreamModelMetadata{ + "gpt-6.1-sol": {ID: "gpt-6.1-sol", Reasoning: &reasoning, InputModalities: []string{"text"}}, + }, + }) + option := PelicanTestModelOptions(account, "public-model") + require.True(t, option.TextSupported, option.CapabilityReason) + require.Equal(t, "gpt-6.1-sol", option.UpstreamModel, "native OpenAI forwarding removes the finite effort suffix") + require.Empty(t, option.ReasoningEfforts, "capabilities belong to the actual wire model, not its mapped spelling") + require.Empty(t, option.DefaultEffort, "the saved non-reasoning declaration must not be bypassed by local defaults") +} diff --git a/backend/internal/service/pelican_wire.go b/backend/internal/service/pelican_wire.go new file mode 100644 index 000000000000..f61b56f6174b --- /dev/null +++ b/backend/internal/service/pelican_wire.go @@ -0,0 +1,19 @@ +package service + +// One runner serves both the historical borrowed-cache endpoints and the +// independent account test endpoints, sharing admission and account leases. +func ProvidePelicanTestRunner( + repo CodexGatewayBorrowTestRepository, + core *CodexGatewayBorrowService, + accounts AccountRepository, + tests *AccountTestService, + gateway *GatewayService, + gemini *GeminiMessagesCompatService, + concurrency *ConcurrencyService, +) *CodexGatewayBorrowTestRunner { + tests.SetCodexGatewayBorrowService(core) + tests.SetPelicanBusinessServices(gateway, gemini, concurrency) + runner := NewCodexGatewayBorrowTestRunner(repo, core, accounts) + runner.SetStandaloneGenerator(tests) + return runner +} diff --git a/backend/internal/service/quota_activity.go b/backend/internal/service/quota_activity.go index 22c9ff8bc6af..b4c632b01303 100644 --- a/backend/internal/service/quota_activity.go +++ b/backend/internal/service/quota_activity.go @@ -73,6 +73,9 @@ type quotaActivityTrace struct { } func (s *QuotaActivityService) Attach(ctx context.Context) (context.Context, func()) { + if IsAccountObservation(ctx) { + return ctx, func() {} + } if quotaActivity(ctx) != nil { return ctx, func() {} } @@ -97,6 +100,7 @@ func quotaActivity(ctx context.Context) *quotaActivityTrace { } func CopyQuotaActivityContext(parent, base context.Context) context.Context { + base = CopyAccountObservationContext(parent, base) if trace := quotaActivity(parent); trace != nil { base = context.WithValue(base, quotaActivityContextKey{}, trace) } @@ -104,6 +108,9 @@ func CopyQuotaActivityContext(parent, base context.Context) context.Context { } func ObserveQuotaAccount(ctx context.Context, accountID int64) { + if IsAccountObservation(ctx) { + return + } trace := quotaActivity(ctx) if trace == nil || accountID <= 0 { return @@ -124,6 +131,9 @@ func ObserveQuotaAccount(ctx context.Context, accountID int64) { } func MarkQuotaLogPersisted(ctx context.Context, accountID int64) { + if IsAccountObservation(ctx) { + return + } if ctx != nil { if receipt, ok := ctx.Value(quotaUsageTaskContextKey{}).(*quotaUsageTaskReceipt); ok && receipt.accountID == accountID { receipt.logged.Store(true) @@ -141,6 +151,9 @@ func MarkQuotaLogPersisted(ctx context.Context, accountID int64) { } func MarkQuotaLogFailed(ctx context.Context, accountID int64) { + if IsAccountObservation(ctx) { + return + } trace := quotaActivity(ctx) if trace == nil { return @@ -155,6 +168,11 @@ func MarkQuotaLogFailed(ctx context.Context, accountID int64) { // Pending is registered before queue submission, closing the gap between // releasing a network slot and a worker eventually committing the usage log. func TrackQuotaUsageTask(parent context.Context, task UsageRecordTask) (UsageRecordTask, func()) { + if IsAccountObservation(parent) { + return func(ctx context.Context) { + task(CopyAccountObservationContext(parent, ctx)) + }, func() {} + } if parent == nil { return task, func() {} } diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index c80311e60a85..f6c89e7a3dde 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -184,6 +184,9 @@ func (s *RateLimitService) notifyAccountSchedulingBlockCleared(accountID int64) // unschedulable until the winning window resets. Returns true when the account // is blocked (either newly or already paused for the same threshold reason). func (s *RateLimitService) ApplyAccountSchedulingThreshold(ctx context.Context, account *Account) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || s.settingService == nil || s.accountRepo == nil || account == nil || account.ID <= 0 { return false } @@ -349,6 +352,9 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun // HandleUpstreamError 处理上游错误响应，标记账号状态 // 返回是否应该停止该账号的调度 func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte, requestedModel ...string) (shouldDisable bool) { + if IsAccountObservation(ctx) { + return false + } ctx = withTempUnschedulableModel(ctx, requestedModel) // Team 联动熔断必须先于池模式/自定义错误码/临时不可调度的各类早退； // 同请求内与 fastpath 调用点的重复触发由方法内去重吸收。 @@ -980,6 +986,9 @@ func (s *RateLimitService) GeminiCooldown(ctx context.Context, account *Account) // handleAuthError 处理认证类错误(401/403)，停止账号调度 func (s *RateLimitService) handleAuthError(ctx context.Context, account *Account, errorMsg string) { + if IsAccountObservation(ctx) { + return + } s.notifyAccountSchedulingBlocked(account, time.Time{}, "auth_error") if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil { slog.Warn("account_set_error_failed", "account_id", account.ID, "error", err) @@ -2006,6 +2015,9 @@ func (s *RateLimitService) handle529(ctx context.Context, account *Account) { // UpdateSessionWindow 从成功响应更新5h窗口状态 func (s *RateLimitService) UpdateSessionWindow(ctx context.Context, account *Account, headers http.Header) { + if IsAccountObservation(ctx) { + return + } status := headers.Get("anthropic-ratelimit-unified-5h-status") if status == "" { return @@ -2129,6 +2141,9 @@ func (s *RateLimitService) samplePassiveUsageFromHeaders(ctx context.Context, ac // ClearRateLimit 清除账号的限流状态 func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) error { + if IsAccountObservation(ctx) { + return nil + } if err := s.accountRepo.ClearRateLimit(ctx, accountID); err != nil { return err } @@ -2153,6 +2168,9 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) } func (s *RateLimitService) ResetOpenAI403Counter(ctx context.Context, accountID int64) { + if IsAccountObservation(ctx) { + return + } if s == nil || s.openAI403CounterCache == nil || accountID <= 0 { return } @@ -2163,6 +2181,9 @@ func (s *RateLimitService) ResetOpenAI403Counter(ctx context.Context, accountID // RecoverAccountState 按需恢复账号的可恢复运行时状态。 func (s *RateLimitService) RecoverAccountState(ctx context.Context, accountID int64, options AccountRecoveryOptions) (*SuccessfulTestRecoveryResult, error) { + if IsAccountObservation(ctx) { + return &SuccessfulTestRecoveryResult{}, nil + } account, err := s.accountRepo.GetByID(ctx, accountID) if err != nil { return nil, err @@ -2204,6 +2225,9 @@ func (s *RateLimitService) RecoverAccountAfterSuccessfulTest(ctx context.Context } func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error { + if IsAccountObservation(ctx) { + return nil + } if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil { return err } @@ -2314,6 +2338,9 @@ func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account } func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || account == nil || s.accountRepo == nil { return false } @@ -2341,6 +2368,9 @@ func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, accou // Spark 的 x-codex-* 使用率和 reset 时间只代表 Spark 模型维度，不能写入账号级 // RateLimitResetAt，否则同一 OAuth 账号上的其他模型也会被错误停调。 func (s *RateLimitService) HandleOpenAICodexSparkRateLimit(ctx context.Context, account *Account, requestedModel string, statusCode int, headers http.Header, responseBody []byte) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || account == nil || s.accountRepo == nil || statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) { return false } @@ -2375,6 +2405,9 @@ func (s *RateLimitService) HandleOpenAICodexSparkRateLimit(ctx context.Context, } func (s *RateLimitService) HandleOpenAIImageCapabilityLoss(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || account == nil || s.accountRepo == nil { return false } @@ -2523,6 +2556,9 @@ func modelRateLimitUpstreamMessage(body []byte) string { // cooldown expires, instead of re-selecting an account that can never serve // the model. func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string, statusCode int, responseBody []byte) bool { + if IsAccountObservation(ctx) { + return false + } if s == nil || account == nil || s.accountRepo == nil { return false } @@ -2671,6 +2707,10 @@ func matchTempUnschedulableRules(account *Account, statusCode int, responseBody } func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte, requestedModel ...string) bool { + if IsAccountObservation(ctx) { + // Preserve request-local policy matching without publishing its cooldown. + return len(matchTempUnschedulableRules(account, statusCode, responseBody)) > 0 + } if account == nil { return false } @@ -2808,6 +2848,9 @@ func truncateTempUnschedMessage(body []byte, maxBytes int) string { // 根据系统设置决定是否标记账户为临时不可调度或错误状态 // 返回是否应该停止该账号的调度 func (s *RateLimitService) HandleStreamTimeout(ctx context.Context, account *Account, model string) bool { + if IsAccountObservation(ctx) { + return false + } if account == nil { return false } diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index ab0c9cc36226..506084d9a655 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -1021,8 +1021,9 @@ var ProviderSet = wire.NewSet( ProvideSystemPromptService, ProvideReasoningRecoveryService, ProvideCodexGatewayBorrowService, - NewCodexGatewayBorrowTestRunner, + ProvidePelicanTestRunner, NewCodexGatewayBorrowPreviewService, + NewPelicanTestCatalog, ProvideOpenAIGatewayService, ProvideImageStorageSettingService, ProvideImageTaskService, diff --git a/backend/migrations/276_pelican_test_execution.sql b/backend/migrations/276_pelican_test_execution.sql new file mode 100644 index 000000000000..bfad449720e3 --- /dev/null +++ b/backend/migrations/276_pelican_test_execution.sql @@ -0,0 +1,23 @@ +-- Extend the existing observation records in place. Original BYTEA answers and +-- errors, client UUID uniqueness, and the 24-hour retention boundary are kept. +ALTER TABLE codex_gateway_borrow_test_tasks + ADD COLUMN generation_timeout_seconds INTEGER NOT NULL DEFAULT 90 + CHECK (generation_timeout_seconds BETWEEN 60 AND 1800), + ADD COLUMN execution_mode TEXT NOT NULL DEFAULT 'legacy_cache' + CHECK (execution_mode IN ('legacy_cache','account')); + +-- Existing legacy observations used a 90-second generation clock. Future tasks +-- default to the new 10-minute budget; the runner always writes its saved budget. +ALTER TABLE codex_gateway_borrow_test_tasks + ALTER COLUMN generation_timeout_seconds SET DEFAULT 600; + +ALTER TABLE codex_gateway_borrow_test_results + ADD COLUMN platform TEXT NOT NULL DEFAULT '', + ADD COLUMN actual_endpoint TEXT NOT NULL DEFAULT '', + ADD COLUMN actual_protocol TEXT NOT NULL DEFAULT '', + ADD COLUMN actual_transport TEXT NOT NULL DEFAULT '', + ADD COLUMN borrow_applied BOOLEAN NOT NULL DEFAULT FALSE, + ADD COLUMN queue_duration_ms BIGINT NOT NULL DEFAULT 0 CHECK (queue_duration_ms >= 0), + ADD COLUMN preparation_duration_ms BIGINT NOT NULL DEFAULT 0 CHECK (preparation_duration_ms >= 0), + ADD COLUMN generation_duration_ms BIGINT NOT NULL DEFAULT 0 CHECK (generation_duration_ms >= 0), + ADD COLUMN generation_started_at TIMESTAMPTZ; diff --git a/frontend/src/api/admin/__tests__/pelicanTests.spec.ts b/frontend/src/api/admin/__tests__/pelicanTests.spec.ts new file mode 100644 index 000000000000..55e3b35e83c6 --- /dev/null +++ b/frontend/src/api/admin/__tests__/pelicanTests.spec.ts @@ -0,0 +1,83 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { pelicanTestsAPI, streamPelicanTests, type PelicanTestRequest } from '../pelicanTests' + +const mocks = vi.hoisted(() => ({ get: vi.fn(), post: vi.fn() })) +vi.mock('@/api/client', () => ({ apiClient: mocks, buildApiUrl: (path: string) => `/api/v1${path}` })) + +const request: PelicanTestRequest = { + client_task_id: '019ae333-ffff-7000-8000-000000000001', generation_timeout_seconds: 600, + targets: [{ account_id: 2, model_id: 'gpt-6.1-sol', effort: 'high' }] +} +function response(chunks: string[]) { + const encoder = new TextEncoder() + return new Response(new ReadableStream({ + start(controller) { for (const chunk of chunks) controller.enqueue(encoder.encode(chunk)); controller.close() } + }), { headers: { 'Content-Type': 'text/event-stream' } }) +} + +describe('Standalone Pelican transport', () => { + beforeEach(() => { vi.clearAllMocks(); localStorage.setItem('auth_token', 'admin-test-token') }) + afterEach(() => { vi.unstubAllGlobals(); localStorage.clear() }) + + it('uses only independent local metadata and stored test endpoints for page reads', async () => { + mocks.get.mockResolvedValue({ data: { accounts: [] } }) + mocks.post.mockResolvedValue({ data: { accounts: [] } }) + const signal = new AbortController().signal + await pelicanTestsAPI.getOptions(signal) + await pelicanTestsAPI.getAccountOptions([1, 2], signal) + await pelicanTestsAPI.listTests(2, signal) + await pelicanTestsAPI.getTest('saved-id', signal) + expect(mocks.get.mock.calls).toEqual([ + ['/admin/pelican-tests/options', { signal }], + ['/admin/pelican-tests/tests', { params: { page: 2, size: 12 }, signal }], + ['/admin/pelican-tests/tests/saved-id', { signal }] + ]) + expect(mocks.post).toHaveBeenCalledOnce() + expect(mocks.post).toHaveBeenCalledWith('/admin/pelican-tests/options', { account_ids: [1, 2] }, { signal }) + }) + + it('posts one budgeted submission, parses phase events across CRLF chunks, and ends at the stored task terminal event', async () => { + const fetch = vi.fn().mockResolvedValue(response([ + ': keep-alive\r\n\r\ndata: {"type":"task_start","task":{"id":"server-id"}}\r\n\r\n', + 'data: {"type":"result_phase","task_id":"server-id","phase":"pre', + 'paring","result":{"id":"one","status":"running"}}\r\n\r\ndata: {"type":"task_complete",\r\ndata: "task":{"id":"server-id","status":"complete"}}\r\n\r\n' + ])) + vi.stubGlobal('fetch', fetch) + const event = vi.fn() + const signal = new AbortController().signal + await streamPelicanTests(request, event, signal) + expect(fetch).toHaveBeenCalledOnce() + expect(fetch).toHaveBeenCalledWith('/api/v1/admin/pelican-tests/tests', expect.objectContaining({ + method: 'POST', body: JSON.stringify(request), signal, + headers: expect.objectContaining({ Authorization: 'Bearer admin-test-token', 'X-Admin-UI-Request': '1' }) + })) + expect(event.mock.calls.map(([value]) => value.type)).toEqual(['task_start', 'result_phase', 'task_complete']) + expect(event.mock.calls[1][0].phase).toBe('preparing') + }) + + it('preserves full HTTP and in-stream errors and rejects unfinished streams without retrying', async () => { + const body = 'upstream full error\n' + 'x'.repeat(14000) + const fetch = vi.fn().mockResolvedValueOnce(new Response(body, { status: 429 })) + .mockResolvedValueOnce(response([`data: ${JSON.stringify({ type: 'error', error: body })}\n\n`])) + .mockResolvedValueOnce(response(['data: {"type":"task_start","task":{"id":"one"}}\n\n'])) + vi.stubGlobal('fetch', fetch) + await expect(streamPelicanTests(request, vi.fn(), new AbortController().signal)).rejects.toThrow(body) + await expect(streamPelicanTests(request, vi.fn(), new AbortController().signal)).rejects.toThrow(body) + await expect(streamPelicanTests(request, vi.fn(), new AbortController().signal)).rejects.toThrow('ended before task_complete') + expect(fetch).toHaveBeenCalledTimes(3) + }) + + it('cancels the pending stream when stop or page leave aborts the request', async () => { + const cancel = vi.fn() + const fetch = vi.fn().mockResolvedValue(new Response(new ReadableStream({ cancel }))) + vi.stubGlobal('fetch', fetch) + const controller = new AbortController() + const pending = streamPelicanTests(request, vi.fn(), controller.signal) + await Promise.resolve() + await Promise.resolve() + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(cancel).toHaveBeenCalledTimes(1) + expect(fetch).toHaveBeenCalledTimes(1) + }) +}) diff --git a/frontend/src/api/admin/pelicanTests.ts b/frontend/src/api/admin/pelicanTests.ts new file mode 100644 index 000000000000..bb0372ab43ef --- /dev/null +++ b/frontend/src/api/admin/pelicanTests.ts @@ -0,0 +1,163 @@ +import { apiClient, buildApiUrl } from '../client' +import { ADMIN_UI_REQUEST_HEADER } from '../adminUIRequest' +import type { BorrowTestResult, BorrowTestTask } from './codexGatewayBorrow' + +export { createBorrowClientTaskId as createPelicanClientTaskId, PELICAN_BORROW_PROMPT as PELICAN_PROMPT } from './codexGatewayBorrow' + +export interface PelicanModelOption { + id: string + display_name: string + upstream_model: string + reasoning_efforts: string[] + default_effort: string + text_supported: boolean + capability_reason: string +} + +export interface PelicanAccountOption { + id: number + name: string + platform: string + type: string + status: string + schedulable: boolean + parent_account_id: number | null + proxy_id: number | null + proxy_name: string + rate_limited_until: string | null + overload_until: string | null + temp_unschedulable_until: string | null + error_message?: string + capability_reason: string + manual_model_allowed: boolean + default_model_id: string + models: PelicanModelOption[] +} + +export interface PelicanOptions { + generated_at: string + max_concurrency: number + default_generation_timeout_seconds: number + min_generation_timeout_seconds: number + max_generation_timeout_seconds: number + accounts: PelicanAccountOption[] +} + +export type PelicanPhase = 'queued' | 'preparing' | 'generating' + +export interface PelicanTestResult extends BorrowTestResult { + platform?: string + actual_endpoint?: string + actual_protocol?: string + actual_transport?: string + borrow_applied?: boolean + queue_duration_ms?: number + preparation_duration_ms?: number + generation_duration_ms?: number + generation_started_at?: string + phase?: PelicanPhase +} + +export interface PelicanTestTask extends Omit { + generation_timeout_seconds?: number + execution_mode?: string + results?: PelicanTestResult[] +} + +export interface PelicanTestRequest { + client_task_id: string + generation_timeout_seconds: number + targets: Array<{ account_id: number; model_id: string; effort: string }> +} + +export type PelicanTestEvent = + | { type: 'task_start' | 'task_complete'; task: PelicanTestTask } + | { type: 'result_started' | 'result_complete'; task_id: string; result: PelicanTestResult } + | { type: 'result_phase'; task_id: string; result: PelicanTestResult; phase: PelicanPhase } + | { type: 'error'; error?: string; message?: string } + +export interface PelicanTestHistory { + items: PelicanTestTask[] + total: number + page: number + size: number +} + +const basePath = '/admin/pelican-tests' + +/** One user submission has one identity and one request; an interrupted stream is not retried. */ +export async function streamPelicanTests( + request: PelicanTestRequest, + onEvent: (event: PelicanTestEvent) => void, + signal: AbortSignal +): Promise { + const token = localStorage.getItem('auth_token') + const response = await fetch(buildApiUrl(`${basePath}/tests`), { + method: 'POST', + headers: { + ...(token ? { Authorization: `Bearer ${token}` } : {}), + 'Content-Type': 'application/json', + Accept: 'text/event-stream', + [ADMIN_UI_REQUEST_HEADER]: '1' + }, + credentials: 'include', + body: JSON.stringify(request), + signal + }) + if (!response.ok) { + const body = await response.text() + throw new Error(`HTTP ${response.status}${body ? `: ${body}` : ''}`) + } + if (!response.body) throw new Error('Pelican test stream has no response body') + const reader = response.body.getReader() + const decoder = new TextDecoder() + let buffer = '' + let complete = false + const abort = () => { void reader.cancel().catch(() => {}) } + signal.addEventListener('abort', abort, { once: true }) + + function consume(block: string) { + const data = block.split(/\r?\n/).filter(line => line.startsWith('data:')) + .map(line => line.slice(5).replace(/^ /, '')).join('\n') + if (!data) return + const event = JSON.parse(data) as PelicanTestEvent + if (event.type === 'error') throw new Error(event.error || event.message || data) + onEvent(event) + if (event.type === 'task_complete') complete = true + } + + try { + if (signal.aborted) throw new DOMException('Aborted', 'AbortError') + while (!complete) { + const { value, done } = await reader.read() + if (signal.aborted) throw new DOMException('Aborted', 'AbortError') + buffer += decoder.decode(value ?? new Uint8Array(), { stream: !done }) + let boundary: RegExpExecArray | null + while ((boundary = /\r?\n\r?\n/.exec(buffer))) { + const block = buffer.slice(0, boundary.index) + buffer = buffer.slice(boundary.index + boundary[0].length) + consume(block) + if (complete) break + } + if (done) { + if (buffer.trim() && !complete) consume(buffer) + break + } + } + if (!complete) throw new Error('Pelican test stream ended before task_complete') + } finally { + signal.removeEventListener('abort', abort) + await reader.cancel().catch(() => {}) + reader.releaseLock() + } +} + +export const pelicanTestsAPI = { + getOptions: async (signal?: AbortSignal) => (await apiClient.get(`${basePath}/options`, { signal })).data, + getAccountOptions: async (accountIds: number[], signal?: AbortSignal) => (await apiClient.post(`${basePath}/options`, { account_ids: accountIds }, { signal })).data, + listTests: async (page = 1, signal?: AbortSignal) => (await apiClient.get(`${basePath}/tests`, { params: { page, size: 12 }, signal })).data, + getTest: async (id: string, signal?: AbortSignal) => (await apiClient.get(`${basePath}/tests/${encodeURIComponent(id)}`, { signal })).data, + streamTests: streamPelicanTests +} + +export default pelicanTestsAPI diff --git a/frontend/src/components/admin/CodexBorrowNav.vue b/frontend/src/components/admin/CodexBorrowNav.vue index 48d959a5995d..f1ca58fe3cdc 100644 --- a/frontend/src/components/admin/CodexBorrowNav.vue +++ b/frontend/src/components/admin/CodexBorrowNav.vue @@ -1,6 +1,6 @@