From 26bf7b60df937f6bc5f8d549f81fca3912901974 Mon Sep 17 00:00:00 2001 From: GcsSloop Date: Wed, 24 Jun 2026 02:01:17 +0800 Subject: [PATCH 1/4] =?UTF-8?q?feat:=20=E4=BC=98=E5=8C=96=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E7=AB=AF=E8=BD=AC=E5=8F=91=E8=B7=AF=E7=94=B1=E6=95=88?= =?UTF-8?q?=E7=8E=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/internal/api/async_usage.go | 130 ++++++++ backend/internal/api/async_usage_test.go | 46 +++ backend/internal/api/gateway_handler.go | 123 ++++---- backend/internal/api/gateway_handler_test.go | 242 ++++++++++++-- backend/internal/api/responses_handler.go | 163 +++++----- .../internal/api/responses_handler_test.go | 248 +++++++++++++++ backend/internal/api/usage_events.go | 53 +++- backend/internal/bootstrap/bootstrap.go | 32 +- backend/internal/routing/sticky.go | 108 +++++++ backend/internal/routing/sticky_test.go | 52 +++ backend/internal/serverauth/session.go | 47 +++ backend/internal/serverauth/session_test.go | 31 ++ backend/internal/serverusers/repository.go | 295 ++++++------------ .../internal/serverusers/repository_test.go | 141 +++++---- backend/internal/serverusers/types.go | 62 +--- backend/internal/store/sqlite/migrations.go | 2 + backend/internal/store/sqlite/store.go | 2 + .../2026-06-23-server-routing-fast-forward.md | 98 ++++++ 18 files changed, 1383 insertions(+), 492 deletions(-) create mode 100644 backend/internal/api/async_usage.go create mode 100644 backend/internal/api/async_usage_test.go create mode 100644 backend/internal/routing/sticky.go create mode 100644 backend/internal/routing/sticky_test.go create mode 100644 docs/plans/2026-06-23-server-routing-fast-forward.md diff --git a/backend/internal/api/async_usage.go b/backend/internal/api/async_usage.go new file mode 100644 index 0000000..d55bdc1 --- /dev/null +++ b/backend/internal/api/async_usage.go @@ -0,0 +1,130 @@ +package api + +import ( + "sync" + + "github.com/gcssloop/codex-router/backend/internal/usage" +) + +type AsyncUsageStoreOptions struct { + QueueSize int +} + +type AsyncUsageStore struct { + base usageEventStore + tasks chan func() + done chan struct{} + + cacheMu sync.RWMutex + cache map[int64]usage.Snapshot + + wg sync.WaitGroup +} + +func NewAsyncUsageStore(base usageEventStore, opts AsyncUsageStoreOptions) *AsyncUsageStore { + queueSize := opts.QueueSize + if queueSize <= 0 { + queueSize = 4096 + } + store := &AsyncUsageStore{ + base: base, + tasks: make(chan func(), queueSize), + done: make(chan struct{}), + cache: map[int64]usage.Snapshot{}, + } + store.wg.Add(1) + go store.run() + return store +} + +func (s *AsyncUsageStore) GetLatest(accountID int64) (usage.Snapshot, error) { + s.cacheMu.RLock() + snapshot, ok := s.cache[accountID] + s.cacheMu.RUnlock() + if ok { + return snapshot, nil + } + snapshot, err := s.base.GetLatest(accountID) + if err != nil { + return usage.Snapshot{}, err + } + s.rememberSnapshot(snapshot) + return snapshot, nil +} + +func (s *AsyncUsageStore) Save(snapshot usage.Snapshot) error { + s.rememberSnapshot(snapshot) + s.enqueue(func() { + _ = s.base.Save(snapshot) + }) + return nil +} + +func (s *AsyncUsageStore) SaveEvent(event usage.Event) error { + s.enqueue(func() { + _ = s.base.SaveEvent(event) + }) + return nil +} + +func (s *AsyncUsageStore) EnqueueUsageEvent(task usageEventTask) bool { + return s.enqueue(func() { + if saved, ok := persistUsageEventSync(s.base, task); ok { + s.rememberSnapshot(saved) + } + }) +} + +func (s *AsyncUsageStore) Drain() { + done := make(chan struct{}) + if !s.enqueue(func() { + close(done) + }) { + return + } + <-done +} + +func (s *AsyncUsageStore) Close() { + close(s.done) + s.wg.Wait() +} + +func (s *AsyncUsageStore) enqueue(task func()) bool { + select { + case <-s.done: + return false + case s.tasks <- task: + return true + default: + return false + } +} + +func (s *AsyncUsageStore) run() { + defer s.wg.Done() + for { + select { + case <-s.done: + for { + select { + case task := <-s.tasks: + task() + default: + return + } + } + case task := <-s.tasks: + task() + } + } +} + +func (s *AsyncUsageStore) rememberSnapshot(snapshot usage.Snapshot) { + if snapshot.AccountID == 0 { + return + } + s.cacheMu.Lock() + s.cache[snapshot.AccountID] = snapshot + s.cacheMu.Unlock() +} diff --git a/backend/internal/api/async_usage_test.go b/backend/internal/api/async_usage_test.go new file mode 100644 index 0000000..6c64e70 --- /dev/null +++ b/backend/internal/api/async_usage_test.go @@ -0,0 +1,46 @@ +package api + +import ( + "testing" + "time" + + "github.com/gcssloop/codex-router/backend/internal/usage" +) + +type slowUsageStore struct { + saveEventDelay time.Duration + events []usage.Event +} + +func (s *slowUsageStore) GetLatest(accountID int64) (usage.Snapshot, error) { + return usage.Snapshot{AccountID: accountID, HealthScore: 1}, nil +} + +func (s *slowUsageStore) Save(snapshot usage.Snapshot) error { + return nil +} + +func (s *slowUsageStore) SaveEvent(event usage.Event) error { + time.Sleep(s.saveEventDelay) + s.events = append(s.events, event) + return nil +} + +func TestAsyncUsageStoreSaveEventDoesNotBlockCaller(t *testing.T) { + base := &slowUsageStore{saveEventDelay: 100 * time.Millisecond} + store := NewAsyncUsageStore(base, AsyncUsageStoreOptions{QueueSize: 4}) + defer store.Close() + + startedAt := time.Now() + if err := store.SaveEvent(usage.Event{AccountID: 1, RequestKind: "responses", Status: "completed"}); err != nil { + t.Fatalf("SaveEvent returned error: %v", err) + } + if elapsed := time.Since(startedAt); elapsed >= 50*time.Millisecond { + t.Fatalf("SaveEvent blocked for %s, want non-blocking enqueue", elapsed) + } + + store.Drain() + if len(base.events) != 1 { + t.Fatalf("persisted events = %d, want 1", len(base.events)) + } +} diff --git a/backend/internal/api/gateway_handler.go b/backend/internal/api/gateway_handler.go index 4ad2df0..523bd7f 100644 --- a/backend/internal/api/gateway_handler.go +++ b/backend/internal/api/gateway_handler.go @@ -15,14 +15,11 @@ import ( "github.com/gcssloop/codex-router/backend/internal/providers" provideropenai "github.com/gcssloop/codex-router/backend/internal/providers/openai" "github.com/gcssloop/codex-router/backend/internal/routing" - "github.com/gcssloop/codex-router/backend/internal/serverusers" "github.com/gcssloop/codex-router/backend/internal/settings" "github.com/gcssloop/codex-router/backend/internal/usage" "github.com/gcssloop/codex-router/backend/internal/usage/normalize" ) -var errServerUserNoAssignedAccounts = errors.New("no upstream accounts assigned to this user") - type GatewayAccounts interface { List() ([]accounts.Account, error) Update(account accounts.Account) error @@ -45,18 +42,14 @@ type GatewayRoutingSettings interface { ListFailoverQueue() ([]int64, error) } -type GatewayServerUsers interface { - ListAssignedAccounts(userID int64) ([]serverusers.AssignedAccount, error) -} - type GatewayHandler struct { accounts GatewayAccounts usage GatewayUsage conversations GatewayRuns settings GatewayRoutingSettings - serverUsers GatewayServerUsers client *http.Client stateEvents *StateEventBus + sticky *routing.StickySelector } type GatewayHandlerOption func(*GatewayHandler) @@ -81,9 +74,18 @@ func WithGatewayStateEvents(bus *StateEventBus) GatewayHandlerOption { } } -func WithGatewayServerUsers(repo GatewayServerUsers) GatewayHandlerOption { +func WithGatewayStickySelector(sticky *routing.StickySelector) GatewayHandlerOption { + return func(handler *GatewayHandler) { + if sticky != nil { + handler.sticky = sticky + } + } +} + +func WithGatewayServerUsers(_ any) GatewayHandlerOption { return func(handler *GatewayHandler) { - handler.serverUsers = repo + // Deprecated: server user account pools were removed. Server users now + // share the global upstream account pool and are used only for auth/usage. } } @@ -93,6 +95,7 @@ func NewGatewayHandler(accounts GatewayAccounts, usage GatewayUsage, conversatio usage: usage, conversations: conversations, client: http.DefaultClient, + sticky: routing.NewStickySelector(time.Minute, time.Now), } for _, opt := range opts { if opt != nil { @@ -127,10 +130,6 @@ func (h *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } candidates, err := h.candidatesForContext(r.Context(), accountList) if err != nil { - if errors.Is(err, errServerUserNoAssignedAccounts) { - http.Error(w, err.Error(), http.StatusForbidden) - return - } http.Error(w, err.Error(), http.StatusInternalServerError) return } @@ -148,17 +147,20 @@ func (h *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, err.Error(), http.StatusInternalServerError) return } + stickyScope := gatewayStickyScope(r.Context(), "chat_completions") executor := routing.NewExecutor(noopRunRecorder{}, func(ctx context.Context, candidate routing.Candidate) error { account := candidate.Account startedAt := time.Now() logUpstreamSummary("gateway", conversationID, account, "/chat/completions", req.Model) if err := ensureOfficialAccountSession(ctx, h.client, h.accounts, &account); err != nil { logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "ensure_session", startedAt, err) + h.invalidateStickyOnFailover(stickyScope, account.ID, err) return err } credential, err := resolveCredential(account) if err != nil { logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "resolve_credential", startedAt, err) + h.invalidateStickyOnFailover(stickyScope, account.ID, err) return err } @@ -171,12 +173,14 @@ func (h *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { }) if err != nil { logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "build_request", startedAt, err) + h.invalidateStickyOnFailover(stickyScope, account.ID, err) return err } resp, err := doAccountRequest(h.client, upstreamReq, account) if err != nil { logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "upstream_request", startedAt, err) + h.invalidateStickyOnFailover(stickyScope, account.ID, err) return err } defer resp.Body.Close() @@ -189,6 +193,7 @@ func (h *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } upstreamErr := buildUpstreamStatusError(resp.StatusCode, responseBody) logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "upstream_status", startedAt, upstreamErr) + h.invalidateStickyOnFailover(stickyScope, account.ID, upstreamErr) return upstreamErr } @@ -198,12 +203,7 @@ func (h *GatewayHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } else { logResultSummary("gateway", conversationID, account.ID, resp.StatusCode, startedAt, string(upstreamResponse)) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, "completed", parseChatCompletionsUsage(upstreamResponse, account.ID), startedAt) - if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } - if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } + h.recordSuccessfulRoute(ctx, stickyScope, account) } return err }) @@ -233,6 +233,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, http.Error(w, err.Error(), http.StatusInternalServerError) return } + stickyScope := gatewayStickyScope(ctx, "chat_completions") var lastErr error @@ -256,6 +257,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "ensure_session", startedAt, err) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) lastErr = err + h.invalidateStickyOnFailover(stickyScope, account.ID, err) if shouldFailoverOnGatewayStreamError(err) { continue } @@ -267,6 +269,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "resolve_credential", startedAt, err) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) lastErr = err + h.invalidateStickyOnFailover(stickyScope, account.ID, err) if shouldFailoverOnGatewayStreamError(err) { continue } @@ -285,6 +288,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "build_request", startedAt, err) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) lastErr = err + h.invalidateStickyOnFailover(stickyScope, account.ID, err) if shouldFailoverOnGatewayStreamError(err) { continue } @@ -297,6 +301,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "upstream_request", startedAt, err) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) lastErr = err + h.invalidateStickyOnFailover(stickyScope, account.ID, err) if shouldFailoverOnGatewayStreamError(err) { continue } @@ -311,6 +316,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "read_response", startedAt, readErr) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(readErr)), usage.Snapshot{AccountID: account.ID}, startedAt) lastErr = readErr + h.invalidateStickyOnFailover(stickyScope, account.ID, readErr) if shouldFailoverOnGatewayStreamError(readErr) { continue } @@ -321,6 +327,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "upstream_status", startedAt, err) lastErr = err + h.invalidateStickyOnFailover(stickyScope, account.ID, err) if shouldFailoverOnGatewayStreamError(err) { continue } @@ -336,16 +343,12 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, if err != nil { logFailureSummary("gateway", conversationID, account.ID, account.AccountName, "read_stream", startedAt, err) persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, runStatusForErrorClass(classifyRunError(err)), usage.Snapshot{AccountID: account.ID}, startedAt) + h.invalidateStickyOnFailover(stickyScope, account.ID, err) return } logResultSummary("gateway", conversationID, account.ID, resp.StatusCode, startedAt, "") persistUsageEvent(ctx, h.usage, account, "chat_completions", req.Model, "completed", collector.snapshotOrDefault(), startedAt) - if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } - if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } + h.recordSuccessfulRoute(ctx, stickyScope, account) return } @@ -355,43 +358,7 @@ func (h *GatewayHandler) serveStream(ctx context.Context, w http.ResponseWriter, http.Error(w, lastErr.Error(), http.StatusBadGateway) } -func (h *GatewayHandler) candidatesForContext(ctx context.Context, accountList []accounts.Account) ([]routing.Candidate, error) { - if user, ok := ServerUserFromContext(ctx); ok { - if h.serverUsers == nil { - return nil, errors.New("server user account pool is not configured") - } - assigned, err := h.serverUsers.ListAssignedAccounts(user.ID) - if err != nil { - return nil, err - } - if len(assigned) == 0 { - return nil, errServerUserNoAssignedAccounts - } - byID := make(map[int64]accounts.Account, len(accountList)) - for _, account := range accountList { - byID[account.ID] = account - } - candidates := make([]routing.Candidate, 0, len(assigned)) - for _, item := range assigned { - account, ok := byID[item.AccountID] - if !ok { - continue - } - account.Priority = len(assigned) - item.Position - account.IsActive = item.IsActive - account.IsLocked = item.IsLocked - snapshot, err := h.usage.GetLatest(account.ID) - if err != nil { - snapshot = normalize.DefaultFallbackSnapshot(account.ID) - } - candidates = append(candidates, routing.Candidate{Account: account, Snapshot: snapshot}) - } - if len(candidates) == 0 { - return nil, errServerUserNoAssignedAccounts - } - return candidates, nil - } - +func (h *GatewayHandler) candidatesForContext(_ context.Context, accountList []accounts.Account) ([]routing.Candidate, error) { candidates := make([]routing.Candidate, 0, len(accountList)) for _, account := range accountList { snapshot, err := h.usage.GetLatest(account.ID) @@ -415,8 +382,8 @@ func (h *GatewayHandler) publishAccountRoutingStateChanged() { } func (h *GatewayHandler) orderedCandidates(ctx context.Context, candidates []routing.Candidate) ([]routing.Candidate, error) { - if user, ok := ServerUserFromContext(ctx); ok { - return routing.RotateCandidatesForUser(orderCandidatesByPriority(candidates), user.ID), nil + if serverUserFromContextExists(ctx) { + return orderServerUserCandidates(ctx, h.sticky, "chat_completions", candidates), nil } if !autoFailoverEnabled(h.settings) { if candidate, ok := activeCandidate(candidates); ok { @@ -427,6 +394,32 @@ func (h *GatewayHandler) orderedCandidates(ctx context.Context, candidates []rou return expandActiveCandidateRetries(orderCandidatesActiveFirst(candidates), activeAccountFailoverRetryAttempts), nil } +func (h *GatewayHandler) rememberSticky(scope string, accountID int64) { + if scope != "" { + h.sticky.Remember(scope, accountID) + } +} + +func (h *GatewayHandler) recordSuccessfulRoute(ctx context.Context, scope string, account accounts.Account) { + if serverUserFromContextExists(ctx) { + h.rememberSticky(scope, account.ID) + return + } + if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { + h.publishAccountRoutingStateChanged() + } + if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { + h.publishAccountRoutingStateChanged() + } +} + +func (h *GatewayHandler) invalidateStickyOnFailover(scope string, accountID int64, err error) { + if scope == "" || !shouldFailoverOnGatewayStreamError(err) { + return + } + h.sticky.Invalidate(scope, accountID) +} + func resolveCredential(account accounts.Account) (string, error) { if account.AuthMode != accounts.AuthModeLocalImport { return account.CredentialRef, nil diff --git a/backend/internal/api/gateway_handler_test.go b/backend/internal/api/gateway_handler_test.go index 7934c86..81d2c2a 100644 --- a/backend/internal/api/gateway_handler_test.go +++ b/backend/internal/api/gateway_handler_test.go @@ -189,19 +189,12 @@ func TestGatewayHandlerUsesAccountTLSVerificationBypass(t *testing.T) { } } -func TestGatewayHandlerInServerModeUsesOnlyAssignedUserAccounts(t *testing.T) { +func TestGatewayHandlerInServerModeUsesGlobalAccountsWithoutAssignments(t *testing.T) { t.Parallel() var firstHits int firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { firstHits++ - http.Error(w, "first should not be used", http.StatusInternalServerError) - })) - defer firstUpstream.Close() - - var secondHits int - secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - secondHits++ w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{ "object": "chat.completion", @@ -211,6 +204,13 @@ func TestGatewayHandlerInServerModeUsesOnlyAssignedUserAccounts(t *testing.T) { }, }) })) + defer firstUpstream.Close() + + var secondHits int + secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondHits++ + http.Error(w, "lower-priority account should not be used", http.StatusInternalServerError) + })) defer secondUpstream.Close() store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) @@ -246,12 +246,6 @@ func TestGatewayHandlerInServerModeUsesOnlyAssignedUserAccounts(t *testing.T) { if err != nil { t.Fatalf("List accounts returned error: %v", err) } - var assignedID int64 - for _, account := range accountList { - if account.AccountName == "assigned" { - assignedID = account.ID - } - } usageRepo := usage.NewSQLiteRepository(store.DB()) for _, account := range accountList { @@ -280,35 +274,229 @@ func TestGatewayHandlerInServerModeUsesOnlyAssignedUserAccounts(t *testing.T) { api.WithGatewayServerUsers(userRepo), )) - noPoolReq := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{ + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{ "model":"gpt-test", "messages":[{"role":"user","content":"ping"}] }`)) - noPoolReq.Header.Set("Content-Type", "application/json") - noPoolReq.Header.Set("Authorization", "Bearer "+created.Token) - noPoolRec := httptest.NewRecorder() - handler.ServeHTTP(noPoolRec, noPoolReq) - if noPoolRec.Code != http.StatusForbidden { - t.Fatalf("no pool status = %d, want %d; body=%s", noPoolRec.Code, http.StatusForbidden, noPoolRec.Body.String()) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+created.Token) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("server gateway status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) } + if firstHits != 1 || secondHits != 0 { + t.Fatalf("upstream hits first=%d second=%d, want only best global account", firstHits, secondHits) + } +} + +func TestGatewayHandlerServerModeRemembersSuccessWithoutRoutingDBWrites(t *testing.T) { + t.Parallel() + + hits := 0 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "object": "chat.completion", + "model": "gpt-test", + "choices": []map[string]any{ + {"message": map[string]any{"role": "assistant", "content": "ok"}}, + }, + }) + })) + defer upstream.Close() - if err := userRepo.SetAccountAssignments(created.User.ID, []int64{assignedID}); err != nil { - t.Fatalf("SetAccountAssignments returned error: %v", err) + accountStore := &stickyResponsesAccounts{items: []accounts.Account{ + { + ID: 1, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "best", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: upstream.URL + "/v1", + CredentialRef: "sk-best", + Status: accounts.StatusActive, + Priority: 100, + }, + }} + usageStore := stickyResponsesUsage{snapshots: map[int64]usage.Snapshot{ + 1: {AccountID: 1, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.9}, + }} + handler := api.WithServerGatewayAuth( + stickyResponsesUsers{user: serverusers.User{ID: 7, Name: "alice", Status: serverusers.StatusActive}}, + api.NewGatewayHandler(accountStore, usageStore, noopResponsesRuns{}), + ) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{ + "model":"gpt-test", + "messages":[{"role":"user","content":"ping"}] + }`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer valid") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("server gateway status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if hits != 1 { + t.Fatalf("hits = %d, want 1", hits) + } + if accountStore.updateCalls != 0 { + t.Fatalf("Update calls = %d, want 0 routing DB writes in server mode", accountStore.updateCalls) + } + if accountStore.setActiveCalls != 0 { + t.Fatalf("SetActive calls = %d, want 0 routing DB writes in server mode", accountStore.setActiveCalls) } +} + +func TestGatewayHandlerServerModeUsesPreferredAccountFromAuthContext(t *testing.T) { + t.Parallel() + + primaryHits := 0 + primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + primaryHits++ + http.Error(w, "high-priority account should not be used while user route is locked", http.StatusInternalServerError) + })) + defer primary.Close() + + secondaryHits := 0 + secondary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondaryHits++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "object": "chat.completion", + "model": "gpt-test", + "choices": []map[string]any{ + {"message": map[string]any{"role": "assistant", "content": "preferred"}}, + }, + }) + })) + defer secondary.Close() + + accountStore := &stickyResponsesAccounts{items: []accounts.Account{ + { + ID: 1, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "best-global", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: primary.URL + "/v1", + CredentialRef: "sk-primary", + Status: accounts.StatusActive, + Priority: 100, + }, + { + ID: 2, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "user-preferred", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: secondary.URL + "/v1", + CredentialRef: "sk-secondary", + Status: accounts.StatusActive, + Priority: 1, + }, + }} + usageStore := stickyResponsesUsage{snapshots: map[int64]usage.Snapshot{ + 1: {AccountID: 1, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.9}, + 2: {AccountID: 2, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.5}, + }} + preferredID := int64(2) + handler := api.WithServerGatewayAuth( + stickyResponsesUsers{user: serverusers.User{ID: 7, Name: "alice", Status: serverusers.StatusActive, PreferredAccountID: &preferredID, RouteLocked: true}}, + api.NewGatewayHandler(accountStore, usageStore, noopResponsesRuns{}), + ) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{ "model":"gpt-test", "messages":[{"role":"user","content":"ping"}] }`)) req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+created.Token) + req.Header.Set("Authorization", "Bearer valid") rec := httptest.NewRecorder() handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("server gateway status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if primaryHits != 0 || secondaryHits != 1 { + t.Fatalf("upstream hits primary=%d secondary=%d, want only preferred account", primaryHits, secondaryHits) + } + if accountStore.setActiveCalls != 0 || accountStore.updateCalls != 0 { + t.Fatalf("routing DB writes update=%d setActive=%d, want none", accountStore.updateCalls, accountStore.setActiveCalls) + } +} + +func TestGatewayHandlerServerModeAllowsManuallyPreferredLockedAccount(t *testing.T) { + t.Parallel() + + primaryHits := 0 + primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + primaryHits++ + http.Error(w, "automatic account should not be used when locked preferred account is manual", http.StatusInternalServerError) + })) + defer primary.Close() + + lockedHits := 0 + locked := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + lockedHits++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "object": "chat.completion", + "model": "gpt-test", + "choices": []map[string]any{ + {"message": map[string]any{"role": "assistant", "content": "manual locked"}}, + }, + }) + })) + defer locked.Close() + + accountStore := &stickyResponsesAccounts{items: []accounts.Account{ + { + ID: 1, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "best-global", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: primary.URL + "/v1", + CredentialRef: "sk-primary", + Status: accounts.StatusActive, + Priority: 100, + }, + { + ID: 2, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "manual-locked", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: locked.URL + "/v1", + CredentialRef: "sk-locked", + Status: accounts.StatusActive, + Priority: 1, + IsLocked: true, + }, + }} + usageStore := stickyResponsesUsage{snapshots: map[int64]usage.Snapshot{ + 1: {AccountID: 1, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.9}, + 2: {AccountID: 2, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.5}, + }} + preferredID := int64(2) + handler := api.WithServerGatewayAuth( + stickyResponsesUsers{user: serverusers.User{ID: 7, Name: "alice", Status: serverusers.StatusActive, PreferredAccountID: &preferredID}}, + api.NewGatewayHandler(accountStore, usageStore, noopResponsesRuns{}), + ) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{ + "model":"gpt-test", + "messages":[{"role":"user","content":"ping"}] + }`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer valid") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { - t.Fatalf("assigned pool status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf("server gateway status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) } - if firstHits != 0 || secondHits != 1 { - t.Fatalf("upstream hits first=%d second=%d, want only assigned second", firstHits, secondHits) + if primaryHits != 0 || lockedHits != 1 { + t.Fatalf("upstream hits primary=%d locked=%d, want only manual locked account", primaryHits, lockedHits) } } diff --git a/backend/internal/api/responses_handler.go b/backend/internal/api/responses_handler.go index d558daa..13c17f2 100644 --- a/backend/internal/api/responses_handler.go +++ b/backend/internal/api/responses_handler.go @@ -24,7 +24,6 @@ import ( providercodex "github.com/gcssloop/codex-router/backend/internal/providers/codex" provideropenai "github.com/gcssloop/codex-router/backend/internal/providers/openai" "github.com/gcssloop/codex-router/backend/internal/routing" - "github.com/gcssloop/codex-router/backend/internal/serverusers" "github.com/gcssloop/codex-router/backend/internal/settings" "github.com/gcssloop/codex-router/backend/internal/usage" "github.com/gcssloop/codex-router/backend/internal/usage/normalize" @@ -59,18 +58,14 @@ type ResponsesRuns interface { ListMessages(conversationID int64) ([]conversations.Message, error) } -type ResponsesServerUsers interface { - ListAssignedAccounts(userID int64) ([]serverusers.AssignedAccount, error) -} - type ResponsesHandler struct { accounts ResponsesAccounts usage ResponsesUsage conversations ResponsesRuns settings settings.ReadRepository - serverUsers ResponsesServerUsers client *http.Client stateEvents *StateEventBus + sticky *routing.StickySelector } type ResponsesHandlerOption func(*ResponsesHandler) @@ -95,9 +90,18 @@ func WithResponsesStateEvents(bus *StateEventBus) ResponsesHandlerOption { } } -func WithResponsesServerUsers(repo ResponsesServerUsers) ResponsesHandlerOption { +func WithResponsesStickySelector(sticky *routing.StickySelector) ResponsesHandlerOption { return func(handler *ResponsesHandler) { - handler.serverUsers = repo + if sticky != nil { + handler.sticky = sticky + } + } +} + +func WithResponsesServerUsers(_ any) ResponsesHandlerOption { + return func(handler *ResponsesHandler) { + // Deprecated: server user account pools were removed. Server users now + // share the global upstream account pool and are used only for auth/usage. } } @@ -113,6 +117,7 @@ func NewResponsesHandler(accounts ResponsesAccounts, usage ResponsesUsage, conve usage: usage, conversations: conversations, client: http.DefaultClient, + sticky: routing.NewStickySelector(time.Minute, time.Now), } for _, opt := range opts { if opt != nil { @@ -155,10 +160,6 @@ func (h *ResponsesHandler) handleResponses(w http.ResponseWriter, r *http.Reques func (h *ResponsesHandler) handleResponsesTransparentSubpath(w http.ResponseWriter, r *http.Request) { account, err := h.selectThinGatewayAccount(r.Context()) if err != nil { - if errors.Is(err, errServerUserNoAssignedAccounts) { - http.Error(w, err.Error(), http.StatusForbidden) - return - } if errors.Is(err, errThinGatewayRequiresResponsesAccount) || errors.Is(err, errThinGatewayActiveAccountUnsupported) { writeThinGatewayUnsupported(w, err.Error()) return @@ -202,10 +203,6 @@ func (h *ResponsesHandler) handleResponsesTransparentSubpath(w http.ResponseWrit func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Request, req gatewayopenai.ResponsesRequest, rawBody []byte) { candidates, err := h.orderedThinGatewayCandidates(r.Context()) if err != nil { - if errors.Is(err, errServerUserNoAssignedAccounts) { - http.Error(w, err.Error(), http.StatusForbidden) - return - } if errors.Is(err, errThinGatewayRequiresResponsesAccount) || errors.Is(err, errThinGatewayActiveAccountUnsupported) { writeThinGatewayUnsupported(w, err.Error()) return @@ -213,6 +210,8 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re http.Error(w, err.Error(), http.StatusBadGateway) return } + stickyScope := gatewayStickyScope(r.Context(), "responses") + serverMode := serverUserFromContextExists(r.Context()) conversationID := int64(0) var lastErr error var fallbackStatusCode int @@ -232,6 +231,7 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re } if reason, ok := h.skipReasonForThinCandidate(candidate); ok { logThinGatewayCandidate(account, "skip", reason) + h.invalidateSticky(stickyScope, account.ID) cooldownUntil := computeThinCandidateCooldownUntil(candidate.Snapshot, reason) if shouldCooldownThinCandidate(candidate, reason) { changed, err := h.markThinCandidateCooldown(account, candidate.Snapshot, reason) @@ -256,9 +256,11 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logFailureSummary("responses", conversationID, account.ID, account.AccountName, "ensure_session", startedAt, err) lastErr = err if shouldFailoverOnThinError(err) && attemptIndex+1 < maxAttempts { + h.invalidateStickyOnThinError(stickyScope, account.ID, err) continue } if shouldFailoverOnThinError(err) { + h.invalidateStickyOnThinError(stickyScope, account.ID, err) break } writeThinGatewayFailure(w, req.Stream, http.StatusBadGateway, err) @@ -272,9 +274,11 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logFailureSummary("responses", conversationID, account.ID, account.AccountName, "resolve_credential", startedAt, err) lastErr = err if shouldFailoverOnThinError(err) && attemptIndex+1 < maxAttempts { + h.invalidateStickyOnThinError(stickyScope, account.ID, err) continue } if shouldFailoverOnThinError(err) { + h.invalidateStickyOnThinError(stickyScope, account.ID, err) break } writeThinGatewayFailure(w, req.Stream, http.StatusBadGateway, err) @@ -289,6 +293,7 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logFailureSummary("responses", conversationID, account.ID, account.AccountName, "upstream_request", startedAt, err) lastErr = err if shouldFailoverOnThinError(err) { + h.invalidateStickyOnThinError(stickyScope, account.ID, err) continue } writeThinGatewayFailure(w, req.Stream, http.StatusBadGateway, err) @@ -303,9 +308,11 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logFailureSummary("responses", conversationID, account.ID, account.AccountName, "read_response", startedAt, readErr) lastErr = readErr if shouldFailoverOnThinError(readErr) && attemptIndex+1 < maxAttempts { + h.invalidateStickyOnThinError(stickyScope, account.ID, readErr) continue } if shouldFailoverOnThinError(readErr) { + h.invalidateStickyOnThinError(stickyScope, account.ID, readErr) break } writeThinGatewayFailure(w, req.Stream, http.StatusBadGateway, readErr) @@ -318,8 +325,9 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re persistUsageEvent(r.Context(), h.usage, account, "responses", req.Model, runStatus, usage.Snapshot{AccountID: account.ID}, startedAt) if shouldFailoverOnThinStatus(runStatus) { + h.invalidateStickyOnThinStatus(stickyScope, account.ID, runStatus) cooldownUntil := computeThinCandidateCooldownUntil(candidate.Snapshot, runStatus) - if shouldCooldownForRunStatus(runStatus) { + if shouldCooldownForRunStatus(runStatus) && !serverMode { changed, err := h.markThinCandidateCooldown(account, candidate.Snapshot, runStatus) if err != nil { lastErr = err @@ -364,6 +372,7 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logFailureSummary("responses", conversationID, account.ID, account.AccountName, "read_stream", startedAt, err) persistUsageEvent(r.Context(), h.usage, account, "responses", req.Model, runStatus, usage.Snapshot{AccountID: account.ID}, startedAt) _ = resp.Body.Close() + h.invalidateStickyOnThinStatus(stickyScope, account.ID, runStatus) writeThinGatewayFailure(w, true, http.StatusBadGateway, err) return } @@ -371,12 +380,7 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re logResultSummary("responses", conversationID, account.ID, resp.StatusCode, startedAt, "") collector.Save(h.usage) persistUsageEvent(r.Context(), h.usage, account, "responses", req.Model, runStatus, collector.snapshotOrDefault(), startedAt) - if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } - if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } + h.recordSuccessfulThinRoute(r.Context(), stickyScope, account) return } @@ -399,12 +403,7 @@ func (h *ResponsesHandler) handleResponsesThin(w http.ResponseWriter, r *http.Re result := parseResponsesJSONResponse(responseBody, account.ID) logResultSummary("responses", conversationID, account.ID, resp.StatusCode, startedAt, result.Text) persistUsageEvent(r.Context(), h.usage, account, "responses", req.Model, runStatus, result.Snapshot, startedAt) - if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } - if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { - h.publishAccountRoutingStateChanged() - } + h.recordSuccessfulThinRoute(r.Context(), stickyScope, account) w.WriteHeader(resp.StatusCode) _, _ = w.Write(responseBody) return @@ -515,12 +514,22 @@ func (h *ResponsesHandler) selectThinGatewayAccount(ctx context.Context) (accoun if err != nil { return accounts.Account{}, err } - accountList, err = h.filterAccountsForServerUser(ctx, accountList) - if err != nil { - return accounts.Account{}, err - } if len(accountList) == 0 { - return accounts.Account{}, errServerUserNoAssignedAccounts + return accounts.Account{}, errThinGatewayRequiresResponsesAccount + } + if serverUserFromContextExists(ctx) { + for _, candidate := range orderServerUserCandidates(ctx, h.sticky, "responses", h.buildCandidates(accountList)) { + if !routing.CanAttemptCandidate(candidate) { + logThinGatewayCandidate(candidate.Account, "skip", "account_locked") + continue + } + if candidate.Account.NativeResponsesCapable() { + logThinGatewayCandidate(candidate.Account, "select", "server_sticky_or_scored") + return candidate.Account, nil + } + logThinGatewayCandidate(candidate.Account, "skip", "supports_responses=false") + } + return accounts.Account{}, errThinGatewayRequiresResponsesAccount } if h.autoFailoverEnabled() { ordered, err := settings.OrderCandidates(h.settings, h.buildCandidates(accountList)) @@ -574,16 +583,12 @@ func (h *ResponsesHandler) orderedThinGatewayCandidates(ctx context.Context) ([] if err != nil { return nil, err } - accountList, err = h.filterAccountsForServerUser(ctx, accountList) - if err != nil { - return nil, err - } if len(accountList) == 0 { - return nil, errServerUserNoAssignedAccounts + return nil, errThinGatewayRequiresResponsesAccount } candidates := h.buildCandidates(accountList) - if user, ok := ServerUserFromContext(ctx); ok { - return routing.RotateCandidatesForUser(orderCandidatesByPriority(candidates), user.ID), nil + if serverUserFromContextExists(ctx) { + return orderServerUserCandidates(ctx, h.sticky, "responses", candidates), nil } if !h.autoFailoverEnabled() { for _, candidate := range candidates { @@ -625,39 +630,6 @@ func (h *ResponsesHandler) nextThinFailoverTarget(candidates []routing.Candidate return routing.Candidate{}, false } -func (h *ResponsesHandler) filterAccountsForServerUser(ctx context.Context, accountList []accounts.Account) ([]accounts.Account, error) { - user, ok := ServerUserFromContext(ctx) - if !ok { - return accountList, nil - } - if h.serverUsers == nil { - return nil, errors.New("server user account pool is not configured") - } - assigned, err := h.serverUsers.ListAssignedAccounts(user.ID) - if err != nil { - return nil, err - } - if len(assigned) == 0 { - return nil, errServerUserNoAssignedAccounts - } - byID := make(map[int64]accounts.Account, len(accountList)) - for _, account := range accountList { - byID[account.ID] = account - } - filtered := make([]accounts.Account, 0, len(assigned)) - for _, item := range assigned { - account, ok := byID[item.AccountID] - if !ok { - continue - } - account.Priority = len(assigned) - item.Position - account.IsActive = item.IsActive - account.IsLocked = item.IsLocked - filtered = append(filtered, account) - } - return filtered, nil -} - func copyResponseHeaders(dst, src http.Header) { for key, values := range src { dst.Del(key) @@ -852,10 +824,6 @@ func (h *ResponsesHandler) handleModelDetail(w http.ResponseWriter, r *http.Requ func (h *ResponsesHandler) handleModelsPassthrough(w http.ResponseWriter, r *http.Request, endpointPath string) { account, err := h.selectThinGatewayAccount(r.Context()) if err != nil { - if errors.Is(err, errServerUserNoAssignedAccounts) { - http.Error(w, err.Error(), http.StatusForbidden) - return - } if errors.Is(err, errThinGatewayRequiresResponsesAccount) || errors.Is(err, errThinGatewayActiveAccountUnsupported) { writeThinGatewayUnsupported(w, err.Error()) return @@ -1012,6 +980,45 @@ func (h *ResponsesHandler) publishAccountRoutingStateChanged() { } } +func (h *ResponsesHandler) recordSuccessfulThinRoute(ctx context.Context, scope string, account accounts.Account) { + if serverUserFromContextExists(ctx) { + h.rememberSticky(scope, account.ID) + return + } + if changed, err := clearRoutingCooldownIfNeeded(h.accounts, account); err == nil && changed { + h.publishAccountRoutingStateChanged() + } + if changed, err := syncActiveAccount(h.accounts, account); err == nil && changed { + h.publishAccountRoutingStateChanged() + } +} + +func (h *ResponsesHandler) rememberSticky(scope string, accountID int64) { + if scope != "" { + h.sticky.Remember(scope, accountID) + } +} + +func (h *ResponsesHandler) invalidateSticky(scope string, accountID int64) { + if scope != "" { + h.sticky.Invalidate(scope, accountID) + } +} + +func (h *ResponsesHandler) invalidateStickyOnThinError(scope string, accountID int64, err error) { + if scope == "" || !shouldFailoverOnThinError(err) { + return + } + h.sticky.Invalidate(scope, accountID) +} + +func (h *ResponsesHandler) invalidateStickyOnThinStatus(scope string, accountID int64, status string) { + if scope == "" || !shouldFailoverOnThinStatus(status) { + return + } + h.sticky.Invalidate(scope, accountID) +} + func computeThinCandidateCooldownUntil(snapshot usage.Snapshot, reason string) *time.Time { now := time.Now().UTC() switch reason { diff --git a/backend/internal/api/responses_handler_test.go b/backend/internal/api/responses_handler_test.go index 6053772..041312c 100644 --- a/backend/internal/api/responses_handler_test.go +++ b/backend/internal/api/responses_handler_test.go @@ -16,6 +16,7 @@ import ( "github.com/gcssloop/codex-router/backend/internal/accounts" "github.com/gcssloop/codex-router/backend/internal/api" "github.com/gcssloop/codex-router/backend/internal/conversations" + "github.com/gcssloop/codex-router/backend/internal/serverusers" "github.com/gcssloop/codex-router/backend/internal/settings" sqlitestore "github.com/gcssloop/codex-router/backend/internal/store/sqlite" "github.com/gcssloop/codex-router/backend/internal/usage" @@ -1103,6 +1104,253 @@ func newResponsesHandlerTestHandler(t *testing.T, account accounts.Account) http return api.NewResponsesHandler(accountRepo, usageRepo, conversations.NewSQLiteRepository(store.DB())) } +func TestResponsesHandlerInServerModeUsesGlobalAccountsWithoutAssignments(t *testing.T) { + t.Parallel() + + var firstHits int + firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + firstHits++ + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"resp_best","object":"response","status":"completed","output_text":"best"}`) + })) + defer firstUpstream.Close() + + var secondHits int + secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondHits++ + http.Error(w, "lower-priority account should not be used", http.StatusInternalServerError) + })) + defer secondUpstream.Close() + + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + accountRepo := accounts.NewSQLiteRepository(store.DB()) + for _, account := range []accounts.Account{ + { + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "best", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: firstUpstream.URL + "/v1", + CredentialRef: "sk-best", + Status: accounts.StatusActive, + Priority: 100, + SupportsResponses: true, + }, + { + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "fallback", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: secondUpstream.URL + "/v1", + CredentialRef: "sk-fallback", + Status: accounts.StatusActive, + Priority: 1, + SupportsResponses: true, + }, + } { + if err := accountRepo.Create(account); err != nil { + t.Fatalf("Create account returned error: %v", err) + } + } + + accountList, err := accountRepo.List() + if err != nil { + t.Fatalf("List returned error: %v", err) + } + usageRepo := usage.NewSQLiteRepository(store.DB()) + for _, account := range accountList { + if err := usageRepo.Save(usage.Snapshot{ + AccountID: account.ID, + Balance: 100, + QuotaRemaining: 100000, + RPMRemaining: 100, + TPMRemaining: 100000, + HealthScore: 0.9, + }); err != nil { + t.Fatalf("Save snapshot returned error: %v", err) + } + } + + userRepo := serverusers.NewSQLiteRepository(store.DB()) + created, err := userRepo.Create("alice") + if err != nil { + t.Fatalf("Create user returned error: %v", err) + } + handler := api.WithServerGatewayAuth(userRepo, api.NewResponsesHandler( + accountRepo, + usageRepo, + conversations.NewSQLiteRepository(store.DB()), + api.WithResponsesServerUsers(userRepo), + )) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewBufferString(`{"model":"gpt-5.4","input":"ping"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+created.Token) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("server responses status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if firstHits != 1 || secondHits != 0 { + t.Fatalf("upstream hits first=%d second=%d, want only best global account", firstHits, secondHits) + } +} + +func TestResponsesHandlerServerModeSticksToSuccessfulFailoverWithoutRoutingDBWrites(t *testing.T) { + t.Parallel() + + primaryCalls := 0 + primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + primaryCalls++ + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"error":{"message":"rate limit"}}`) + })) + defer primary.Close() + + secondaryCalls := 0 + secondary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondaryCalls++ + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"resp_fallback","object":"response","status":"completed","output_text":"fallback"}`) + })) + defer secondary.Close() + + accountStore := &stickyResponsesAccounts{items: []accounts.Account{ + { + ID: 1, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "primary-rate-limited", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: primary.URL + "/v1", + CredentialRef: "sk-primary", + Status: accounts.StatusActive, + Priority: 100, + SupportsResponses: true, + }, + { + ID: 2, + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "secondary-success", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: secondary.URL + "/v1", + CredentialRef: "sk-secondary", + Status: accounts.StatusActive, + Priority: 1, + SupportsResponses: true, + }, + }} + usageStore := stickyResponsesUsage{snapshots: map[int64]usage.Snapshot{ + 1: {AccountID: 1, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.9}, + 2: {AccountID: 2, Balance: 100, QuotaRemaining: 100000, RPMRemaining: 100, TPMRemaining: 100000, HealthScore: 0.5}, + }} + handler := api.WithServerGatewayAuth( + stickyResponsesUsers{user: serverusers.User{ID: 7, Name: "alice", Status: serverusers.StatusActive, CreatedAt: time.Now().UTC()}}, + api.NewResponsesHandler(accountStore, usageStore, noopResponsesRuns{}), + ) + + for i := 0; i < 2; i++ { + req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewBufferString(`{"model":"gpt-5.4","input":"ping"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer valid") + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("request %d status = %d, want %d; body=%s", i+1, rec.Code, http.StatusOK, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), `"output_text":"fallback"`) { + t.Fatalf("request %d body = %s, want fallback response", i+1, rec.Body.String()) + } + } + + if primaryCalls != 1 { + t.Fatalf("primaryCalls = %d, want 1; successful fallback should be sticky", primaryCalls) + } + if secondaryCalls != 2 { + t.Fatalf("secondaryCalls = %d, want 2", secondaryCalls) + } + if accountStore.updateCalls != 0 { + t.Fatalf("Update calls = %d, want 0 routing DB writes in server mode", accountStore.updateCalls) + } + if accountStore.setActiveCalls != 0 { + t.Fatalf("SetActive calls = %d, want 0 routing DB writes in server mode", accountStore.setActiveCalls) + } +} + +type stickyResponsesAccounts struct { + items []accounts.Account + updateCalls int + setActiveCalls int +} + +func (s *stickyResponsesAccounts) List() ([]accounts.Account, error) { + items := make([]accounts.Account, len(s.items)) + copy(items, s.items) + return items, nil +} + +func (s *stickyResponsesAccounts) Update(accounts.Account) error { + s.updateCalls++ + return nil +} + +func (s *stickyResponsesAccounts) SetActive(int64) error { + s.setActiveCalls++ + return nil +} + +type stickyResponsesUsage struct { + snapshots map[int64]usage.Snapshot +} + +func (s stickyResponsesUsage) GetLatest(accountID int64) (usage.Snapshot, error) { + if snapshot, ok := s.snapshots[accountID]; ok { + return snapshot, nil + } + return usage.Snapshot{AccountID: accountID, HealthScore: 1}, nil +} + +func (s stickyResponsesUsage) Save(usage.Snapshot) error { + return nil +} + +func (s stickyResponsesUsage) SaveEvent(usage.Event) error { + return nil +} + +type stickyResponsesUsers struct { + user serverusers.User +} + +func (s stickyResponsesUsers) Authenticate(token string) (serverusers.User, error) { + if token != "valid" { + return serverusers.User{}, sql.ErrNoRows + } + return s.user, nil +} + +type noopResponsesRuns struct{} + +func (noopResponsesRuns) CreateConversation(conversations.Conversation) (int64, error) { + return 0, nil +} + +func (noopResponsesRuns) CreateRun(conversations.Run) (int64, error) { + return 0, nil +} + +func (noopResponsesRuns) AppendMessage(conversations.Message) error { + return nil +} + +func (noopResponsesRuns) ListMessages(int64) ([]conversations.Message, error) { + return nil, nil +} + func TestResponsesHandlerThinModePassesThroughPreviousResponseID(t *testing.T) { t.Parallel() diff --git a/backend/internal/api/usage_events.go b/backend/internal/api/usage_events.go index 58b4c8e..a58059d 100644 --- a/backend/internal/api/usage_events.go +++ b/backend/internal/api/usage_events.go @@ -23,6 +23,20 @@ func (noopRunRecorder) CreateRun(_ conversations.Run) (int64, error) { return 0, nil } +type usageEventTask struct { + Account accounts.Account + RequestKind string + Model string + Status string + Snapshot usage.Snapshot + ServerUserID *int64 + LatencyMS float64 +} + +type asyncUsageEventStore interface { + EnqueueUsageEvent(task usageEventTask) bool +} + func latestSnapshotOrEmpty(repo interface { GetLatest(accountID int64) (usage.Snapshot, error) }, accountID int64) usage.Snapshot { @@ -37,32 +51,54 @@ func persistUsageEvent(ctx context.Context, repo usageEventStore, account accoun if account.ID == 0 { return } + task := usageEventTask{ + Account: account, + RequestKind: requestKind, + Model: model, + Status: status, + Snapshot: snapshot, + LatencyMS: time.Since(startedAt).Seconds() * 1000, + } + if user, ok := ServerUserFromContext(ctx); ok { + userID := user.ID + task.ServerUserID = &userID + } + if asyncRepo, ok := repo.(asyncUsageEventStore); ok { + asyncRepo.EnqueueUsageEvent(task) + return + } + persistUsageEventSync(repo, task) +} +func persistUsageEventSync(repo usageEventStore, task usageEventTask) (usage.Snapshot, bool) { + account := task.Account before := latestSnapshotOrEmpty(repo, account.ID) + snapshot := task.Snapshot snapshot.AccountID = account.ID if snapshot.CheckedAt.IsZero() { snapshot.CheckedAt = time.Now().UTC() } after := mergeUsageSnapshotWithLatest(repo, snapshot) - if status == "completed" && hasUsageDelta(after) { + if task.Status == "completed" && hasUsageDelta(after) { _ = repo.Save(after) } event := usage.Event{ AccountID: account.ID, ProviderType: string(account.ProviderType), - RequestKind: requestKind, - Model: model, - Status: status, + RequestKind: task.RequestKind, + Model: task.Model, + Status: task.Status, InputTokens: int64(after.LastInputTokens), OutputTokens: int64(after.LastOutputTokens), TotalTokens: int64(after.LastTotalTokens), - EstimatedCost: estimateModelCostUSD(model, after.LastInputTokens, after.LastOutputTokens), - LatencyMS: time.Since(startedAt).Seconds() * 1000, + EstimatedCost: estimateModelCostUSD(task.Model, after.LastInputTokens, after.LastOutputTokens), + LatencyMS: task.LatencyMS, CreatedAt: time.Now().UTC(), } - if user, ok := ServerUserFromContext(ctx); ok { - event.ServerUserID = &user.ID + if task.ServerUserID != nil { + userID := *task.ServerUserID + event.ServerUserID = &userID } if before.Balance != 0 || after.Balance != 0 { event.BalanceBefore = float64Ptr(before.Balance) @@ -73,6 +109,7 @@ func persistUsageEvent(ctx context.Context, repo usageEventStore, account accoun event.QuotaAfter = float64Ptr(after.QuotaRemaining) } _ = repo.SaveEvent(event) + return after, true } func mergeUsageSnapshotWithLatest(repo interface { diff --git a/backend/internal/bootstrap/bootstrap.go b/backend/internal/bootstrap/bootstrap.go index 9dbd9e7..d2a1be4 100644 --- a/backend/internal/bootstrap/bootstrap.go +++ b/backend/internal/bootstrap/bootstrap.go @@ -21,6 +21,7 @@ import ( "github.com/gcssloop/codex-router/backend/internal/conversations" "github.com/gcssloop/codex-router/backend/internal/netproxy" "github.com/gcssloop/codex-router/backend/internal/policy" + "github.com/gcssloop/codex-router/backend/internal/routing" "github.com/gcssloop/codex-router/backend/internal/scheduler" "github.com/gcssloop/codex-router/backend/internal/secrets" "github.com/gcssloop/codex-router/backend/internal/serverauth" @@ -54,6 +55,7 @@ type App struct { store *sqlite.Store cancel context.CancelFunc background sync.WaitGroup + closeHooks []func() } type backgroundTask func(context.Context, time.Time) @@ -102,6 +104,8 @@ func NewApp(_ context.Context, cfg Config) (*App, error) { upstreamHTTPClient := netproxy.NewHTTPClient(settingsRepo) conversationRepo := conversations.NewSQLiteRepository(store.DB()) serverUserRepo := serverusers.NewSQLiteRepository(store.DB()) + gatewayUsageRepo := api.NewAsyncUsageStore(usageRepo, api.AsyncUsageStoreOptions{QueueSize: 8192}) + serverUserSticky := routing.NewStickySelector(time.Minute, time.Now) policyRepo := policy.NewMemoryRepository() authConnector := auth.NewOAuthConnector(auth.Config{ ClientID: "app_EMoamEEZ73f0CkXaXp7hrann", @@ -201,20 +205,22 @@ func NewApp(_ context.Context, cfg Config) (*App, error) { apiMux.Handle("/tooling/mcp/apply", toolingHandler) gatewayHandler := api.NewGatewayHandler( accountRepo, - usageRepo, + gatewayUsageRepo, conversationRepo, api.WithGatewaySettings(settingsRepo), api.WithGatewayHTTPClient(upstreamHTTPClient), api.WithGatewayStateEvents(stateEvents), + api.WithGatewayStickySelector(serverUserSticky), api.WithGatewayServerUsers(serverUserRepo), ) responsesHandler := api.NewResponsesHandler( accountRepo, - usageRepo, + gatewayUsageRepo, conversationRepo, api.WithResponsesSettings(settingsRepo), api.WithResponsesHTTPClient(upstreamHTTPClient), api.WithResponsesStateEvents(stateEvents), + api.WithResponsesStickySelector(serverUserSticky), api.WithResponsesServerUsers(serverUserRepo), ) gatewayMux := http.NewServeMux() @@ -261,10 +267,15 @@ func NewApp(_ context.Context, cfg Config) (*App, error) { mux.Handle(httpPrefix+"/auth/", withCORS(http.StripPrefix(httpPrefix+"/auth", authManager))) apiHandler = authManager.RequireSession(apiHandler) userAPIMux := http.NewServeMux() - meHandler := api.NewServerMeHandler(serverUserRepo) + meHandler := api.NewServerMeHandler( + serverUserRepo, + api.WithServerMeAccounts(accountRepo), + api.WithServerMeUsage(usageRepo), + api.WithServerMeStickySelector(serverUserSticky), + ) userAPIMux.Handle("/me", meHandler) userAPIMux.Handle("/me/", meHandler) - userAPIHandler = authManager.RequireUserSession(http.StripPrefix(httpPrefix+"/api", userAPIMux)) + userAPIHandler = authManager.RequireUserSessionOrToken(http.StripPrefix(httpPrefix+"/api", userAPIMux)) } mux.Handle(httpPrefix+"/webui/", webuiHandler) if cfg.ServerMode { @@ -278,7 +289,13 @@ func NewApp(_ context.Context, cfg Config) (*App, error) { } appCtx, cancel := context.WithCancel(context.Background()) - app := &App{listenAddr: cfg.ListenAddr, handler: mux, store: store, cancel: cancel} + app := &App{ + listenAddr: cfg.ListenAddr, + handler: mux, + store: store, + cancel: cancel, + closeHooks: []func(){gatewayUsageRepo.Close}, + } interval := cfg.SchedulerInterval if interval <= 0 { @@ -541,6 +558,11 @@ func (a *App) Close() error { a.cancel() } a.background.Wait() + for _, closeHook := range a.closeHooks { + if closeHook != nil { + closeHook() + } + } if a.store != nil { return a.store.Close() } diff --git a/backend/internal/routing/sticky.go b/backend/internal/routing/sticky.go new file mode 100644 index 0000000..eb50109 --- /dev/null +++ b/backend/internal/routing/sticky.go @@ -0,0 +1,108 @@ +package routing + +import ( + "sync" + "time" +) + +type StickySelector struct { + mu sync.RWMutex + ttl time.Duration + now func() time.Time + entries map[string]stickyEntry +} + +type stickyEntry struct { + AccountID int64 + ExpiresAt time.Time +} + +func NewStickySelector(ttl time.Duration, now func() time.Time) *StickySelector { + if ttl <= 0 { + ttl = time.Minute + } + if now == nil { + now = time.Now + } + return &StickySelector{ + ttl: ttl, + now: now, + entries: map[string]stickyEntry{}, + } +} + +func (s *StickySelector) Apply(scope string, candidates []Candidate) []Candidate { + ordered := ScoreCandidates(candidates) + if len(ordered) <= 1 || s == nil { + return ordered + } + + entry, ok := s.entry(scope) + if !ok { + return ordered + } + for index, candidate := range ordered { + if candidate.Account.ID != entry.AccountID { + continue + } + if index == 0 { + return ordered + } + remembered := ordered[index] + copy(ordered[1:index+1], ordered[0:index]) + ordered[0] = remembered + return ordered + } + return ordered +} + +func (s *StickySelector) Remember(scope string, accountID int64) { + if s == nil || scope == "" || accountID <= 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.entries[scope] = stickyEntry{AccountID: accountID, ExpiresAt: s.now().UTC().Add(s.ttl)} +} + +func (s *StickySelector) Current(scope string) (int64, bool) { + if s == nil || scope == "" { + return 0, false + } + entry, ok := s.entry(scope) + if !ok { + return 0, false + } + return entry.AccountID, true +} + +func (s *StickySelector) Invalidate(scope string, accountID int64) { + if s == nil || scope == "" || accountID <= 0 { + return + } + s.mu.Lock() + defer s.mu.Unlock() + entry, ok := s.entries[scope] + if ok && entry.AccountID == accountID { + delete(s.entries, scope) + } +} + +func (s *StickySelector) entry(scope string) (stickyEntry, bool) { + s.mu.RLock() + entry, ok := s.entries[scope] + s.mu.RUnlock() + if !ok { + return stickyEntry{}, false + } + if !s.now().UTC().Before(entry.ExpiresAt) { + s.mu.Lock() + current, stillPresent := s.entries[scope] + if stillPresent && current.AccountID == entry.AccountID && current.ExpiresAt.Equal(entry.ExpiresAt) { + delete(s.entries, scope) + } + s.mu.Unlock() + return stickyEntry{}, false + } + return entry, true +} diff --git a/backend/internal/routing/sticky_test.go b/backend/internal/routing/sticky_test.go new file mode 100644 index 0000000..3a06985 --- /dev/null +++ b/backend/internal/routing/sticky_test.go @@ -0,0 +1,52 @@ +package routing_test + +import ( + "testing" + "time" + + "github.com/gcssloop/codex-router/backend/internal/accounts" + "github.com/gcssloop/codex-router/backend/internal/routing" +) + +func TestStickySelectorPrefersRememberedAccountWithinTTL(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) + selector := routing.NewStickySelector(time.Minute, func() time.Time { return now }) + candidates := []routing.Candidate{ + {Account: accounts.Account{ID: 1, AccountName: "best", Priority: 100, Status: accounts.StatusActive}}, + {Account: accounts.Account{ID: 2, AccountName: "remembered", Priority: 1, Status: accounts.StatusActive}}, + } + + selector.Remember("responses", 2) + ordered := selector.Apply("responses", candidates) + + if ordered[0].Account.ID != 2 { + t.Fatalf("first account = %d, want remembered account 2", ordered[0].Account.ID) + } +} + +func TestStickySelectorFallsBackAfterInvalidateOrExpiry(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) + selector := routing.NewStickySelector(time.Minute, func() time.Time { return now }) + candidates := []routing.Candidate{ + {Account: accounts.Account{ID: 1, AccountName: "best", Priority: 100, Status: accounts.StatusActive}}, + {Account: accounts.Account{ID: 2, AccountName: "remembered", Priority: 1, Status: accounts.StatusActive}}, + } + + selector.Remember("responses", 2) + selector.Invalidate("responses", 2) + ordered := selector.Apply("responses", candidates) + if ordered[0].Account.ID != 1 { + t.Fatalf("first account after invalidate = %d, want best account 1", ordered[0].Account.ID) + } + + selector.Remember("responses", 2) + now = now.Add(2 * time.Minute) + ordered = selector.Apply("responses", candidates) + if ordered[0].Account.ID != 1 { + t.Fatalf("first account after expiry = %d, want best account 1", ordered[0].Account.ID) + } +} diff --git a/backend/internal/serverauth/session.go b/backend/internal/serverauth/session.go index 164ea21..42d5e5a 100644 --- a/backend/internal/serverauth/session.go +++ b/backend/internal/serverauth/session.go @@ -20,6 +20,10 @@ type UserStore interface { AuthenticateLogin(username string, token string) (serverusers.User, error) } +type TokenUserStore interface { + Authenticate(token string) (serverusers.User, error) +} + type PasswordStore interface { LoadOrInitialize(initialPassword string) (string, error) Save(password string) error @@ -115,6 +119,41 @@ func (m *Manager) RequireUserSession(next http.Handler) http.Handler { }) } +func (m *Manager) RequireUserSessionOrToken(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + session, ok := m.sessionFromRequest(r) + if ok && session.Role == serverusers.RoleUser && session.UserID > 0 { + next.ServeHTTP(w, r.WithContext(ContextWithSession(r.Context(), session))) + return + } + if ok { + http.Error(w, "user session required", http.StatusForbidden) + return + } + token := bearerToken(r.Header.Get("Authorization")) + if token == "" { + token = strings.TrimSpace(r.Header.Get("X-AI-Gate-Token")) + } + tokenStore, tokenAuthEnabled := m.users.(TokenUserStore) + if token == "" || !tokenAuthEnabled { + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + user, err := tokenStore.Authenticate(token) + if err != nil { + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + session = Session{ + Authenticated: true, + Role: serverusers.RoleUser, + UserID: user.ID, + Username: user.Username, + } + next.ServeHTTP(w, r.WithContext(ContextWithSession(r.Context(), session))) + }) +} + func (m *Manager) login(w http.ResponseWriter, r *http.Request) { var payload struct { Password string `json:"password"` @@ -301,6 +340,14 @@ func randomToken() (string, error) { return hex.EncodeToString(raw[:]), nil } +func bearerToken(header string) string { + fields := strings.Fields(strings.TrimSpace(header)) + if len(fields) != 2 || !strings.EqualFold(fields[0], "bearer") { + return "" + } + return fields[1] +} + func writeJSON(w http.ResponseWriter, status int, value any) { w.Header().Set("Content-Type", "application/json; charset=utf-8") w.WriteHeader(status) diff --git a/backend/internal/serverauth/session_test.go b/backend/internal/serverauth/session_test.go index 0f22f8d..99b68f1 100644 --- a/backend/internal/serverauth/session_test.go +++ b/backend/internal/serverauth/session_test.go @@ -113,6 +113,16 @@ func (s fakeUserStore) AuthenticateLogin(username string, token string) (serveru return s.user, nil } +func (s fakeUserStore) Authenticate(token string) (serverusers.User, error) { + if s.err != nil { + return serverusers.User{}, s.err + } + if token != "agt-test" { + return serverusers.User{}, http.ErrNoCookie + } + return s.user, nil +} + func TestManagerUserLoginSessionAndAdminProtection(t *testing.T) { manager := NewManagerWithUsers("secret-password", time.Minute, fakeUserStore{ user: serverusers.User{ID: 7, Username: "alice", Name: "alice", Role: serverusers.RoleUser, Status: serverusers.StatusActive}, @@ -167,3 +177,24 @@ func TestManagerUserLoginSessionAndAdminProtection(t *testing.T) { t.Fatalf("ordinary user /me status = %d, want %d", meRec.Code, http.StatusNoContent) } } + +func TestManagerRequireUserSessionOrTokenAcceptsBearerToken(t *testing.T) { + manager := NewManagerWithUsers("secret-password", time.Minute, fakeUserStore{ + user: serverusers.User{ID: 7, Username: "alice", Name: "alice", Role: serverusers.RoleUser, Status: serverusers.StatusActive}, + }) + userOnly := manager.RequireUserSessionOrToken(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + session, ok := SessionFromContext(r.Context()) + if !ok || session.UserID != 7 || session.Role != serverusers.RoleUser { + t.Fatalf("session from context = %+v, ok=%v; want user 7", session, ok) + } + w.WriteHeader(http.StatusNoContent) + })) + + req := httptest.NewRequest(http.MethodGet, "/api/me/upstreams", nil) + req.Header.Set("Authorization", "Bearer agt-test") + rec := httptest.NewRecorder() + userOnly.ServeHTTP(rec, req) + if rec.Code != http.StatusNoContent { + t.Fatalf("bearer token /me status = %d, want %d; body=%s", rec.Code, http.StatusNoContent, rec.Body.String()) + } +} diff --git a/backend/internal/serverusers/repository.go b/backend/internal/serverusers/repository.go index de37e7d..4604814 100644 --- a/backend/internal/serverusers/repository.go +++ b/backend/internal/serverusers/repository.go @@ -6,6 +6,7 @@ import ( "database/sql" "encoding/hex" "fmt" + "strconv" "strings" "time" ) @@ -14,6 +15,8 @@ type Repository struct { db *sql.DB } +const lastUsedAtUpdateInterval = time.Minute + func NewSQLiteRepository(db *sql.DB) *Repository { return &Repository{db: db} } @@ -45,20 +48,15 @@ func (r *Repository) Create(name string) (CreatedUser, error) { func (r *Repository) List() ([]User, error) { rows, err := r.db.Query( `SELECT u.id, u.name, u.username, u.token_hash, u.role, u.status, u.created_at, u.last_used_at, + u.preferred_account_id, u.route_locked, COALESCE(usage_stats.request_count, 0) AS request_count, - COALESCE(usage_stats.total_tokens, 0) AS total_tokens, - COALESCE(assignment_stats.assigned_accounts, 0) AS assigned_accounts + COALESCE(usage_stats.total_tokens, 0) AS total_tokens FROM server_users u LEFT JOIN ( SELECT server_user_id, COUNT(id) AS request_count, COALESCE(SUM(total_tokens), 0) AS total_tokens FROM usage_events GROUP BY server_user_id ) usage_stats ON usage_stats.server_user_id = u.id - LEFT JOIN ( - SELECT user_id, COUNT(account_id) AS assigned_accounts - FROM server_user_accounts - GROUP BY user_id - ) assignment_stats ON assignment_stats.user_id = u.id ORDER BY u.id ASC`, ) if err != nil { @@ -68,9 +66,11 @@ func (r *Repository) List() ([]User, error) { users := make([]User, 0) for rows.Next() { var user User - if err := rows.Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt), &user.RequestCount, &user.TotalTokens, &user.AssignedAccounts); err != nil { + var routeLocked int + if err := rows.Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt), nullInt64Dest(&user.PreferredAccountID), &routeLocked, &user.RequestCount, &user.TotalTokens); err != nil { return nil, fmt.Errorf("scan server user: %w", err) } + user.RouteLocked = routeLocked == 1 normalizeUserFields(&user) users = append(users, user) } @@ -96,20 +96,20 @@ func (r *Repository) Get(id int64) (User, error) { func (r *Repository) Authenticate(token string) (User, error) { hash := HashToken(token) var user User + var routeLocked int err := r.db.QueryRow( - `SELECT id, name, username, token_hash, role, status, created_at, last_used_at FROM server_users WHERE token_hash = ?`, + `SELECT id, name, username, token_hash, role, status, created_at, last_used_at, preferred_account_id, route_locked FROM server_users WHERE token_hash = ?`, hash, - ).Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt)) + ).Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt), nullInt64Dest(&user.PreferredAccountID), &routeLocked) if err != nil { return User{}, fmt.Errorf("select server user by token: %w", err) } + user.RouteLocked = routeLocked == 1 normalizeUserFields(&user) if user.Status != StatusActive { return User{}, fmt.Errorf("server user is disabled") } - now := time.Now().UTC() - _, _ = r.db.Exec(`UPDATE server_users SET last_used_at = ? WHERE id = ?`, now, user.ID) - user.LastUsedAt = &now + user.LastUsedAt = r.touchLastUsedAt(user.ID, user.LastUsedAt) return user, nil } @@ -120,26 +120,50 @@ func (r *Repository) AuthenticateLogin(username string, token string) (User, err } hash := HashToken(token) var user User + var routeLocked int err := r.db.QueryRow( - `SELECT id, name, username, token_hash, role, status, created_at, last_used_at + `SELECT id, name, username, token_hash, role, status, created_at, last_used_at, preferred_account_id, route_locked FROM server_users WHERE username = ? AND token_hash = ?`, normalized, hash, - ).Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt)) + ).Scan(&user.ID, &user.Name, &user.Username, &user.TokenHash, &user.Role, &user.Status, &user.CreatedAt, nullTimeDest(&user.LastUsedAt), nullInt64Dest(&user.PreferredAccountID), &routeLocked) if err != nil { return User{}, fmt.Errorf("select server user by username and token: %w", err) } + user.RouteLocked = routeLocked == 1 normalizeUserFields(&user) if user.Status != StatusActive { return User{}, fmt.Errorf("server user is disabled") } - now := time.Now().UTC() - _, _ = r.db.Exec(`UPDATE server_users SET last_used_at = ? WHERE id = ?`, now, user.ID) - user.LastUsedAt = &now + user.LastUsedAt = r.touchLastUsedAt(user.ID, user.LastUsedAt) return user, nil } +func (r *Repository) touchLastUsedAt(userID int64, current *time.Time) *time.Time { + now := time.Now().UTC() + if current != nil && now.Sub(current.UTC()) < lastUsedAtUpdateInterval { + return current + } + _, _ = r.db.Exec(`UPDATE server_users SET last_used_at = ? WHERE id = ?`, now, userID) + return &now +} + +func (r *Repository) UpdateRoute(id int64, accountID *int64, locked bool) (User, error) { + result, err := r.db.Exec(`UPDATE server_users SET preferred_account_id = ?, route_locked = ? WHERE id = ?`, nullInt64(accountID), boolToInt(locked), id) + if err != nil { + return User{}, fmt.Errorf("update server user route: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return User{}, fmt.Errorf("read updated server user route rows: %w", err) + } + if affected == 0 { + return User{}, sql.ErrNoRows + } + return r.Get(id) +} + func (r *Repository) Disable(id int64) error { result, err := r.db.Exec(`UPDATE server_users SET status = ? WHERE id = ?`, StatusDisabled, id) if err != nil { @@ -155,6 +179,21 @@ func (r *Repository) Disable(id int64) error { return nil } +func (r *Repository) Delete(id int64) error { + result, err := r.db.Exec(`DELETE FROM server_users WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("delete server user: %w", err) + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("read deleted server user rows: %w", err) + } + if affected == 0 { + return sql.ErrNoRows + } + return nil +} + func (r *Repository) RotateToken(id int64) (CreatedUser, error) { token, err := newToken() if err != nil { @@ -183,181 +222,6 @@ func (r *Repository) RotateToken(id int64) (CreatedUser, error) { return CreatedUser{}, sql.ErrNoRows } -func (r *Repository) ListAssignedAccounts(userID int64) ([]AssignedAccount, error) { - rows, err := r.db.Query( - `SELECT sua.user_id, a.id, a.account_name, a.provider_type, a.source_icon, a.base_url, - a.status, sua.position, sua.is_active, sua.is_locked, a.supports_responses, a.cooldown_until, a.cooldown_reason - FROM server_user_accounts sua - JOIN accounts a ON a.id = sua.account_id - WHERE sua.user_id = ? - ORDER BY sua.position ASC, a.id ASC`, - userID, - ) - if err != nil { - return nil, fmt.Errorf("query assigned accounts: %w", err) - } - defer rows.Close() - - assigned := make([]AssignedAccount, 0) - for rows.Next() { - var item AssignedAccount - var isActive int - var isLocked int - var supportsResponses int - var cooldown sql.NullTime - if err := rows.Scan( - &item.UserID, - &item.AccountID, - &item.AccountName, - &item.ProviderType, - &item.SourceIcon, - &item.BaseURL, - &item.Status, - &item.Position, - &isActive, - &isLocked, - &supportsResponses, - &cooldown, - &item.CooldownReason, - ); err != nil { - return nil, fmt.Errorf("scan assigned account: %w", err) - } - item.IsActive = isActive == 1 - item.IsLocked = isLocked == 1 - item.SupportsResponses = supportsResponses == 1 - if cooldown.Valid { - value := cooldown.Time.UTC() - item.CooldownUntil = &value - } - assigned = append(assigned, item) - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("iterate assigned accounts: %w", err) - } - return assigned, nil -} - -func (r *Repository) ListAccountAssignments(userID int64) ([]AccountAssignment, error) { - rows, err := r.db.Query( - `SELECT a.id, a.account_name, a.provider_type, a.source_icon, a.base_url, a.status, - CASE WHEN sua.account_id IS NULL THEN 0 ELSE 1 END AS assigned, - COALESCE(sua.position, 0) AS position, - COALESCE(sua.is_active, 0) AS is_active, - COALESCE(sua.is_locked, 0) AS is_locked, - a.supports_responses, a.cooldown_until, a.cooldown_reason - FROM accounts a - LEFT JOIN server_user_accounts sua ON sua.account_id = a.id AND sua.user_id = ? - ORDER BY assigned DESC, position ASC, a.id ASC`, - userID, - ) - if err != nil { - return nil, fmt.Errorf("query account assignments: %w", err) - } - defer rows.Close() - - assignments := make([]AccountAssignment, 0) - for rows.Next() { - var item AccountAssignment - var assigned int - var isActive int - var isLocked int - var supportsResponses int - var cooldown sql.NullTime - if err := rows.Scan( - &item.AccountID, - &item.AccountName, - &item.ProviderType, - &item.SourceIcon, - &item.BaseURL, - &item.Status, - &assigned, - &item.Position, - &isActive, - &isLocked, - &supportsResponses, - &cooldown, - &item.CooldownReason, - ); err != nil { - return nil, fmt.Errorf("scan account assignment: %w", err) - } - item.Assigned = assigned == 1 - item.IsActive = isActive == 1 - item.IsLocked = isLocked == 1 - item.SupportsResponses = supportsResponses == 1 - if cooldown.Valid { - value := cooldown.Time.UTC() - item.CooldownUntil = &value - } - assignments = append(assignments, item) - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("iterate account assignments: %w", err) - } - return assignments, nil -} - -func (r *Repository) SetAccountAssignments(userID int64, accountIDs []int64) error { - tx, err := r.db.Begin() - if err != nil { - return fmt.Errorf("begin account assignment transaction: %w", err) - } - defer func() { - _ = tx.Rollback() - }() - - if _, err := tx.Exec(`DELETE FROM server_user_accounts WHERE user_id = ?`, userID); err != nil { - return fmt.Errorf("clear account assignments: %w", err) - } - now := time.Now().UTC() - seen := map[int64]bool{} - for position, accountID := range accountIDs { - if accountID <= 0 || seen[accountID] { - continue - } - seen[accountID] = true - if _, err := tx.Exec( - `INSERT INTO server_user_accounts (user_id, account_id, position, is_active, is_locked, created_at, updated_at) - VALUES (?, ?, ?, 1, 0, ?, ?)`, - userID, - accountID, - position, - now, - now, - ); err != nil { - return fmt.Errorf("insert account assignment: %w", err) - } - } - if err := tx.Commit(); err != nil { - return fmt.Errorf("commit account assignments: %w", err) - } - return nil -} - -func (r *Repository) UpdateAccountState(userID int64, accountID int64, update AccountStateUpdate) error { - result, err := r.db.Exec( - `UPDATE server_user_accounts - SET position = ?, is_active = ?, is_locked = ?, updated_at = ? - WHERE user_id = ? AND account_id = ?`, - update.Position, - boolToInt(update.IsActive), - boolToInt(update.IsLocked), - time.Now().UTC(), - userID, - accountID, - ) - if err != nil { - return fmt.Errorf("update assigned account state: %w", err) - } - affected, err := result.RowsAffected() - if err != nil { - return fmt.Errorf("read assigned account state rows: %w", err) - } - if affected == 0 { - return sql.ErrNoRows - } - return nil -} - func HashToken(token string) string { sum := sha256.Sum256([]byte(strings.TrimSpace(token))) return hex.EncodeToString(sum[:]) @@ -383,6 +247,21 @@ func normalizeUserFields(user *User) { } } +func nullTimeDest(target **time.Time) any { + return &nullableTimeScanner{target: target} +} + +func nullInt64(value *int64) any { + if value == nil || *value <= 0 { + return nil + } + return *value +} + +func nullInt64Dest(target **int64) any { + return &nullableInt64Scanner{target: target} +} + func boolToInt(value bool) int { if value { return 1 @@ -390,8 +269,34 @@ func boolToInt(value bool) int { return 0 } -func nullTimeDest(target **time.Time) any { - return &nullableTimeScanner{target: target} +type nullableInt64Scanner struct { + target **int64 +} + +func (s *nullableInt64Scanner) Scan(value any) error { + if value == nil { + *s.target = nil + return nil + } + switch typed := value.(type) { + case int64: + copy := typed + *s.target = © + case int: + copy := int64(typed) + *s.target = © + case string: + parsed, err := strconv.ParseInt(typed, 10, 64) + if err != nil { + return err + } + *s.target = &parsed + case []byte: + return s.Scan(string(typed)) + default: + return fmt.Errorf("unsupported nullable int64 type %T", value) + } + return nil } type nullableTimeScanner struct { diff --git a/backend/internal/serverusers/repository_test.go b/backend/internal/serverusers/repository_test.go index 5e2f8df..0fe1200 100644 --- a/backend/internal/serverusers/repository_test.go +++ b/backend/internal/serverusers/repository_test.go @@ -1,11 +1,11 @@ package serverusers_test import ( + "database/sql" "path/filepath" "testing" "time" - "github.com/gcssloop/codex-router/backend/internal/accounts" "github.com/gcssloop/codex-router/backend/internal/serverusers" sqlitestore "github.com/gcssloop/codex-router/backend/internal/store/sqlite" "github.com/gcssloop/codex-router/backend/internal/usage" @@ -97,6 +97,38 @@ func TestRepositoryDisablesAndRotatesToken(t *testing.T) { } } +func TestRepositoryDeletesUser(t *testing.T) { + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + t.Cleanup(func() { + _ = store.Close() + }) + + repo := serverusers.NewSQLiteRepository(store.DB()) + created, err := repo.Create("delete-me") + if err != nil { + t.Fatalf("Create returned error: %v", err) + } + if err := repo.Delete(created.User.ID); err != nil { + t.Fatalf("Delete returned error: %v", err) + } + users, err := repo.List() + if err != nil { + t.Fatalf("List returned error: %v", err) + } + if len(users) != 0 { + t.Fatalf("users = %+v, want deleted user omitted", users) + } + if _, err := repo.Authenticate(created.Token); err == nil { + t.Fatal("Authenticate returned nil error for deleted user") + } + if err := repo.Delete(created.User.ID); err != sql.ErrNoRows { + t.Fatalf("Delete missing user error = %v, want sql.ErrNoRows", err) + } +} + func TestRepositoryAuthenticatesLoginByUsernameAndToken(t *testing.T) { store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) if err != nil { @@ -127,7 +159,7 @@ func TestRepositoryAuthenticatesLoginByUsernameAndToken(t *testing.T) { } } -func TestRepositoryAssignsAccountPoolWithoutGlobalAccountMutation(t *testing.T) { +func TestRepositoryThrottlesLastUsedAtWrites(t *testing.T) { store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) if err != nil { t.Fatalf("Open returned error: %v", err) @@ -136,90 +168,71 @@ func TestRepositoryAssignsAccountPoolWithoutGlobalAccountMutation(t *testing.T) _ = store.Close() }) - accountRepo := accounts.NewSQLiteRepository(store.DB()) - createAccount := func(name string, priority int, active bool, locked bool) int64 { - t.Helper() - if err := accountRepo.Create(accounts.Account{ - ProviderType: accounts.ProviderOpenAICompatible, - AccountName: name, - AuthMode: accounts.AuthModeAPIKey, - CredentialRef: "sk-" + name, - BaseURL: "https://example.invalid/v1", - Status: accounts.StatusActive, - Priority: priority, - IsActive: active, - IsLocked: locked, - }); err != nil { - t.Fatalf("Create(account %s) returned error: %v", name, err) - } - list, err := accountRepo.List() - if err != nil { - t.Fatalf("List accounts returned error: %v", err) - } - for _, account := range list { - if account.AccountName == name { - return account.ID - } - } - t.Fatalf("created account %s not found", name) - return 0 - } - firstID := createAccount("first", 100, true, false) - secondID := createAccount("second", 10, false, false) - repo := serverusers.NewSQLiteRepository(store.DB()) created, err := repo.Create("alice") if err != nil { t.Fatalf("Create returned error: %v", err) } - assigned, err := repo.ListAssignedAccounts(created.User.ID) + first, err := repo.Authenticate(created.Token) if err != nil { - t.Fatalf("ListAssignedAccounts returned error: %v", err) + t.Fatalf("first Authenticate returned error: %v", err) } - if len(assigned) != 0 { - t.Fatalf("new user assigned accounts = %+v, want none", assigned) + if first.LastUsedAt == nil { + t.Fatal("first LastUsedAt = nil, want timestamp") } + second, err := repo.Authenticate(created.Token) + if err != nil { + t.Fatalf("second Authenticate returned error: %v", err) + } + if second.LastUsedAt == nil || !second.LastUsedAt.Equal(*first.LastUsedAt) { + t.Fatalf("second LastUsedAt = %v, want unchanged %v", second.LastUsedAt, first.LastUsedAt) + } +} - if err := repo.SetAccountAssignments(created.User.ID, []int64{secondID, firstID}); err != nil { - t.Fatalf("SetAccountAssignments returned error: %v", err) +func TestRepositoryPersistsRoutePreference(t *testing.T) { + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) } - assigned, err = repo.ListAssignedAccounts(created.User.ID) + t.Cleanup(func() { + _ = store.Close() + }) + + repo := serverusers.NewSQLiteRepository(store.DB()) + created, err := repo.Create("alice") if err != nil { - t.Fatalf("ListAssignedAccounts returned error: %v", err) + t.Fatalf("Create returned error: %v", err) } - if len(assigned) != 2 { - t.Fatalf("assigned accounts = %+v, want two", assigned) + accountID := int64(42) + updated, err := repo.UpdateRoute(created.User.ID, &accountID, true) + if err != nil { + t.Fatalf("UpdateRoute returned error: %v", err) } - if assigned[0].AccountID != secondID || assigned[0].Position != 0 || assigned[1].AccountID != firstID || assigned[1].Position != 1 { - t.Fatalf("assigned accounts = %+v, want saved order second, first", assigned) + if updated.PreferredAccountID == nil || *updated.PreferredAccountID != accountID || !updated.RouteLocked { + t.Fatalf("updated route = account:%v locked:%v, want account 42 locked", updated.PreferredAccountID, updated.RouteLocked) } - if !assigned[0].IsActive || !assigned[1].IsActive { - t.Fatalf("assigned accounts = %+v, want newly assigned accounts active by default", assigned) + + authenticated, err := repo.Authenticate(created.Token) + if err != nil { + t.Fatalf("Authenticate returned error: %v", err) } - if assigned[0].CredentialRef != "" || assigned[1].CredentialRef != "" { - t.Fatalf("assigned accounts leaked credentials: %+v", assigned) + if authenticated.PreferredAccountID == nil || *authenticated.PreferredAccountID != accountID || !authenticated.RouteLocked { + t.Fatalf("authenticated route = account:%v locked:%v, want account 42 locked", authenticated.PreferredAccountID, authenticated.RouteLocked) } - if err := repo.UpdateAccountState(created.User.ID, firstID, serverusers.AccountStateUpdate{ - Position: 0, - IsActive: true, - IsLocked: true, - }); err != nil { - t.Fatalf("UpdateAccountState returned error: %v", err) - } - assigned, err = repo.ListAssignedAccounts(created.User.ID) + users, err := repo.List() if err != nil { - t.Fatalf("ListAssignedAccounts returned error: %v", err) + t.Fatalf("List returned error: %v", err) } - if assigned[0].AccountID != firstID || !assigned[0].IsActive || !assigned[0].IsLocked { - t.Fatalf("assigned accounts after state update = %+v, want first active and locked first", assigned) + if users[0].PreferredAccountID == nil || *users[0].PreferredAccountID != accountID || !users[0].RouteLocked { + t.Fatalf("listed route = account:%v locked:%v, want account 42 locked", users[0].PreferredAccountID, users[0].RouteLocked) } - globalFirst, err := accountRepo.GetByID(firstID) + cleared, err := repo.UpdateRoute(created.User.ID, nil, false) if err != nil { - t.Fatalf("GetByID returned error: %v", err) + t.Fatalf("UpdateRoute clear returned error: %v", err) } - if globalFirst.Priority != 100 || !globalFirst.IsActive || globalFirst.IsLocked { - t.Fatalf("global first account mutated = %+v, want original priority/active/locked", globalFirst) + if cleared.PreferredAccountID != nil || cleared.RouteLocked { + t.Fatalf("cleared route = account:%v locked:%v, want automatic route", cleared.PreferredAccountID, cleared.RouteLocked) } } diff --git a/backend/internal/serverusers/types.go b/backend/internal/serverusers/types.go index 627257f..b4b50de 100644 --- a/backend/internal/serverusers/types.go +++ b/backend/internal/serverusers/types.go @@ -8,59 +8,21 @@ const RoleUser = "user" const RoleAdmin = "admin" type User struct { - ID int64 `json:"id"` - Name string `json:"name"` - Username string `json:"username"` - TokenHash string `json:"-"` - Role string `json:"role"` - Status string `json:"status"` - CreatedAt time.Time `json:"created_at"` - LastUsedAt *time.Time `json:"last_used_at,omitempty"` - RequestCount int64 `json:"request_count,omitempty"` - TotalTokens int64 `json:"total_tokens,omitempty"` - AssignedAccounts int64 `json:"assigned_accounts,omitempty"` + ID int64 `json:"id"` + Name string `json:"name"` + Username string `json:"username"` + TokenHash string `json:"-"` + Role string `json:"role"` + Status string `json:"status"` + CreatedAt time.Time `json:"created_at"` + LastUsedAt *time.Time `json:"last_used_at,omitempty"` + PreferredAccountID *int64 `json:"preferred_account_id,omitempty"` + RouteLocked bool `json:"route_locked"` + RequestCount int64 `json:"request_count,omitempty"` + TotalTokens int64 `json:"total_tokens,omitempty"` } type CreatedUser struct { User User `json:"user"` Token string `json:"token"` } - -type AssignedAccount struct { - UserID int64 `json:"user_id"` - AccountID int64 `json:"account_id"` - AccountName string `json:"account_name"` - ProviderType string `json:"provider_type"` - SourceIcon string `json:"source_icon"` - BaseURL string `json:"base_url,omitempty"` - Status string `json:"status"` - Position int `json:"position"` - IsActive bool `json:"is_active"` - IsLocked bool `json:"is_locked"` - SupportsResponses bool `json:"supports_responses"` - CooldownUntil *time.Time `json:"cooldown_until,omitempty"` - CooldownReason string `json:"cooldown_reason,omitempty"` - CredentialRef string `json:"-"` -} - -type AccountAssignment struct { - AccountID int64 `json:"account_id"` - AccountName string `json:"account_name"` - ProviderType string `json:"provider_type"` - SourceIcon string `json:"source_icon"` - BaseURL string `json:"base_url,omitempty"` - Status string `json:"status"` - Assigned bool `json:"assigned"` - Position int `json:"position"` - IsActive bool `json:"is_active"` - IsLocked bool `json:"is_locked"` - SupportsResponses bool `json:"supports_responses"` - CooldownUntil *time.Time `json:"cooldown_until,omitempty"` - CooldownReason string `json:"cooldown_reason,omitempty"` -} - -type AccountStateUpdate struct { - Position int `json:"position"` - IsActive bool `json:"is_active"` - IsLocked bool `json:"is_locked"` -} diff --git a/backend/internal/store/sqlite/migrations.go b/backend/internal/store/sqlite/migrations.go index eed1615..e5288d8 100644 --- a/backend/internal/store/sqlite/migrations.go +++ b/backend/internal/store/sqlite/migrations.go @@ -78,6 +78,8 @@ var schemaStatements = []string{ token_hash TEXT NOT NULL UNIQUE, role TEXT NOT NULL DEFAULT 'user', status TEXT NOT NULL DEFAULT 'active', + preferred_account_id INTEGER, + route_locked INTEGER NOT NULL DEFAULT 0, last_used_at DATETIME, created_at DATETIME DEFAULT CURRENT_TIMESTAMP );`, diff --git a/backend/internal/store/sqlite/store.go b/backend/internal/store/sqlite/store.go index 2185036..5d62e49 100644 --- a/backend/internal/store/sqlite/store.go +++ b/backend/internal/store/sqlite/store.go @@ -127,6 +127,8 @@ func (s *Store) migrate() error { {table: "accounts", name: "cooldown_reason", definition: "TEXT NOT NULL DEFAULT ''"}, {table: "server_users", name: "username", definition: "TEXT NOT NULL DEFAULT ''"}, {table: "server_users", name: "role", definition: "TEXT NOT NULL DEFAULT 'user'"}, + {table: "server_users", name: "preferred_account_id", definition: "INTEGER"}, + {table: "server_users", name: "route_locked", definition: "INTEGER NOT NULL DEFAULT 0"}, {table: "app_settings", name: "show_home_update_indicator", definition: "INTEGER NOT NULL DEFAULT 1"}, {table: "app_settings", name: "status_refresh_interval_seconds", definition: "INTEGER NOT NULL DEFAULT 60"}, {table: "app_settings", name: "usage_request_timeout_seconds", definition: "INTEGER NOT NULL DEFAULT 15"}, diff --git a/docs/plans/2026-06-23-server-routing-fast-forward.md b/docs/plans/2026-06-23-server-routing-fast-forward.md new file mode 100644 index 0000000..587603c --- /dev/null +++ b/docs/plans/2026-06-23-server-routing-fast-forward.md @@ -0,0 +1,98 @@ +# 服务端路由快速转发实施计划 + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** 移除服务端用户账户池逻辑，让服务端网关请求默认使用全部上游账户，并通过全局最优路径加短期粘性缓存提升多并发转发效率。 + +**Architecture:** 服务端用户仅承担 token 鉴权、登录会话和用量归属，不再拥有独立账户池、排序、激活或锁定状态。网关候选从全局账户池构造，按健康度、优先级和容量评分选出最优路径；选中后写入进程内 sticky selector，在 TTL 内优先复用，只有上游不可达、限流、额度不足或冷却时失效并切换到下一个候选。请求热路径不查询 `server_user_accounts`，服务模式成功/失败切换不依赖 `SetActive` 或 cooldown 同步写库，token 鉴权不做每请求同步写入；gateway/responses 使用内存缓存和异步 usage 写入队列，队列压力过高时优先保障转发而不是回退同步写库。 + +**Tech Stack:** Go `net/http`、SQLite repository、现有 `routing.Candidate` 评分模型、React + Ant Design。 + +--- + +### Task 1: 后端服务用户不再需要账户池 + +**Files:** +- Modify: `backend/internal/api/gateway_handler_test.go` +- Modify: `backend/internal/api/responses_handler_test.go` +- Modify: `backend/internal/api/server_users_handler_test.go` +- Modify: `backend/internal/api/server_me_handler_test.go` +- Modify: `backend/internal/api/gateway_handler.go` +- Modify: `backend/internal/api/responses_handler.go` +- Modify: `backend/internal/api/server_users_handler.go` +- Modify: `backend/internal/api/server_me_handler.go` +- Modify: `backend/internal/serverusers/repository.go` + +**Steps:** +1. 写失败测试：新建服务用户没有账户分配时，合法 token 请求仍能命中全局上游账户。 +2. 写失败测试：服务用户管理 API 不再暴露 `/server-users/{id}/accounts`。 +3. 写失败测试：普通用户 `/me/accounts` 和 `/me/accounts/{id}/state` 不再可用。 +4. 移除 gateway 和 responses handler 中的服务用户账户池过滤。 +5. 移除 handler interface 中的账户池方法依赖。 +6. 保留 legacy repository 方法和 SQLite 表作为兼容残留，但运行时不调用。 + +**Verification:** +- `cd backend && go test ./internal/api ./internal/serverusers -run 'ServerUser|Gateway|Responses|Assigned|ServerMe' -count=1` + +### Task 2: 全局最优路径短期粘性 + +**Files:** +- Create: `backend/internal/routing/sticky.go` +- Create: `backend/internal/routing/sticky_test.go` +- Modify: `backend/internal/api/account_routing_state.go` +- Modify: `backend/internal/api/gateway_handler.go` +- Modify: `backend/internal/api/responses_handler.go` + +**Steps:** +1. 写失败测试：sticky selector 在 TTL 内把上次成功账户移动到候选首位。 +2. 写失败测试：失效 sticky 后会重新选择评分最高候选。 +3. 实现并发安全的进程内 selector，按 scope 记录 `account_id/expires_at`。 +4. 成功请求后记住命中账户，容量/限流/连接错误时失效该 scope。 +5. 服务用户请求只更新进程内 sticky，不在转发热路径同步写入全局 active/cooldown 状态。 +6. 在 chat 和 responses 两条转发路径复用相同排序逻辑。 + +**Verification:** +- `cd backend && go test ./internal/routing ./internal/api -run 'Sticky|Gateway|Responses' -count=1` + +### Task 3: 前端移除账户池界面 + +**Files:** +- Modify: `frontend/src/features/server-users/ServerUsersPage.tsx` +- Modify: `frontend/src/features/server-users/ServerUsersPage.test.tsx` +- Modify: `frontend/src/features/server-users/UserPoolPage.tsx` +- Modify: `frontend/src/features/server-users/UserPoolPage.test.tsx` +- Modify: `frontend/src/lib/api.ts` + +**Steps:** +1. 写失败测试：服务用户页面不再显示“账户池/分配账户”。 +2. 写失败测试：普通用户页面只显示个人用量，不显示账户池状态编辑表。 +3. 删除前端账户池 API 调用和相关类型。 +4. 保留服务用户创建、禁用、轮换 token 和用量展示。 + +**Verification:** +- `npm --prefix frontend test -- ServerUsersPage UserPoolPage` + +### Task 4: 收敛验证 + +**Verification:** +- `cd backend && go test ./internal/api ./internal/routing ./internal/serverusers -count=1` +- `npm --prefix frontend test -- ServerUsersPage UserPoolPage` + +No commit is created unless explicitly requested. + +### Task 5: 转发热路径异步写库 + +**Files:** +- Create: `backend/internal/api/async_usage.go` +- Create: `backend/internal/api/async_usage_test.go` +- Modify: `backend/internal/api/usage_events.go` +- Modify: `backend/internal/bootstrap/bootstrap.go` + +**Steps:** +1. 写失败测试：慢 `SaveEvent` 不阻塞 `AsyncUsageStore.SaveEvent` 调用方。 +2. 实现内存 snapshot cache 和异步写入队列。 +3. `persistUsageEvent` 遇到异步 repo 时只入队，不在队列满时回退同步写库。 +4. bootstrap 只给 gateway/responses handler 使用异步 usage wrapper，其余 dashboard/monitoring/refresh 继续使用原始 repository。 + +**Verification:** +- `cd backend && go test ./internal/api ./internal/bootstrap -run 'AsyncUsage|NewApp|Gateway|Responses' -count=1` From 0564d88aa032787981462565d464e0ff9896c386 Mon Sep 17 00:00:00 2001 From: GcsSloop Date: Wed, 24 Jun 2026 02:04:27 +0800 Subject: [PATCH 2/4] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E8=BF=9C?= =?UTF-8?q?=E7=AB=AF=E4=B8=8A=E6=B8=B8=E8=B4=A6=E6=88=B7=E8=B7=AF=E7=94=B1?= =?UTF-8?q?=E6=8E=A7=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/internal/api/accounts_handler.go | 248 ++++++++++++ backend/internal/api/accounts_handler_test.go | 122 ++++++ backend/internal/api/server_me_handler.go | 366 ++++++++++++++++-- .../internal/api/server_me_handler_test.go | 285 +++++++++++--- backend/internal/api/server_user_route.go | 91 +++++ backend/internal/api/server_users_handler.go | 63 ++- .../internal/api/server_users_handler_test.go | 98 ++--- backend/internal/webui/embed.go | 23 +- backend/internal/webui/embed_test.go | 27 ++ ...6-23-user-upstream-route-control-design.md | 28 ++ .../2026-06-23-user-upstream-route-control.md | 75 ++++ 11 files changed, 1228 insertions(+), 198 deletions(-) create mode 100644 backend/internal/api/server_user_route.go create mode 100644 backend/internal/webui/embed_test.go create mode 100644 docs/plans/2026-06-23-user-upstream-route-control-design.md create mode 100644 docs/plans/2026-06-23-user-upstream-route-control.md diff --git a/backend/internal/api/accounts_handler.go b/backend/internal/api/accounts_handler.go index 6738613..eb4f780 100644 --- a/backend/internal/api/accounts_handler.go +++ b/backend/internal/api/accounts_handler.go @@ -142,6 +142,12 @@ func (h *AccountsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.listAccountsUsage(w, r) case r.Method == http.MethodGet && r.URL.Path == "/accounts": h.listAccounts(w, r) + case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/accounts/") && strings.HasSuffix(r.URL.Path, "/upstreams/route"): + h.updateAIGateAccountRoute(w, r) + case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/accounts/") && strings.Contains(r.URL.Path, "/upstreams/") && strings.HasSuffix(r.URL.Path, "/lock"): + h.updateAIGateUpstreamLock(w, r) + case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/accounts/") && strings.HasSuffix(r.URL.Path, "/upstreams"): + h.getAIGateAccountUpstreams(w, r) case r.Method == http.MethodPost && r.URL.Path == "/accounts/auth/authorize": h.createAuthSession(w, r) case r.Method == http.MethodPost && r.URL.Path == "/accounts/auth/device/complete": @@ -801,6 +807,191 @@ func (h *AccountsHandler) shareAccount(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, map[string]string{"payload": string(payload)}) } +func (h *AccountsHandler) getAIGateAccountUpstreams(w http.ResponseWriter, r *http.Request) { + id, err := accountIDFromPath(strings.TrimSuffix(strings.TrimSuffix(r.URL.Path, "/upstreams"), "/")) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + account, err := h.repo.GetByID(id) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + upstreamsURL, err := aiGateUpstreamsURL(account.BaseURL) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + token := strings.TrimSpace(account.CredentialRef) + if token == "" { + http.Error(w, "account credential is required", http.StatusBadRequest) + return + } + ctx, cancel := context.WithTimeout(r.Context(), 4*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, upstreamsURL, nil) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("X-AI-Gate-Token", token) + resp, err := doAccountRequest(h.client, req, account) + if err != nil { + http.Error(w, "fetch ai-gate upstreams: "+err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + details, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + http.Error(w, strings.TrimSpace(string(details)), resp.StatusCode) + return + } + var payload serverMeUpstreamsResponse + decoder := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)) + if err := decoder.Decode(&payload); err != nil { + http.Error(w, "decode ai-gate upstreams: "+err.Error(), http.StatusBadGateway) + return + } + writeJSON(w, http.StatusOK, payload) +} + +func (h *AccountsHandler) updateAIGateAccountRoute(w http.ResponseWriter, r *http.Request) { + id, err := accountIDFromPath(strings.TrimSuffix(strings.TrimSuffix(r.URL.Path, "/upstreams/route"), "/")) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + account, err := h.repo.GetByID(id) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + routeURL, err := aiGateRouteURL(account.BaseURL) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + token := strings.TrimSpace(account.CredentialRef) + if token == "" { + http.Error(w, "account credential is required", http.StatusBadRequest) + return + } + var payload struct { + AccountID *int64 `json:"account_id"` + Locked bool `json:"locked"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + http.Error(w, "invalid route payload", http.StatusBadRequest) + return + } + if payload.AccountID != nil && *payload.AccountID <= 0 { + http.Error(w, "account_id must be positive", http.StatusBadRequest) + return + } + if payload.AccountID == nil && payload.Locked { + http.Error(w, "locked route requires account_id", http.StatusBadRequest) + return + } + body, err := json.Marshal(payload) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + ctx, cancel := context.WithTimeout(r.Context(), 4*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPut, routeURL, bytes.NewReader(body)) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("X-AI-Gate-Token", token) + resp, err := doAccountRequest(h.client, req, account) + if err != nil { + http.Error(w, "update ai-gate upstream route: "+err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + details, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + http.Error(w, strings.TrimSpace(string(details)), resp.StatusCode) + return + } + var response map[string]any + decoder := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)) + if err := decoder.Decode(&response); err != nil { + http.Error(w, "decode ai-gate route response: "+err.Error(), http.StatusBadGateway) + return + } + writeJSON(w, http.StatusOK, response) +} + +func (h *AccountsHandler) updateAIGateUpstreamLock(w http.ResponseWriter, r *http.Request) { + id, upstreamAccountID, err := aiGateUpstreamLockIDsFromPath(r.URL.Path) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + account, err := h.repo.GetByID(id) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + routeURL, err := aiGateUpstreamLockURL(account.BaseURL, upstreamAccountID) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + token := strings.TrimSpace(account.CredentialRef) + if token == "" { + http.Error(w, "account credential is required", http.StatusBadRequest) + return + } + var payload struct { + Locked bool `json:"locked"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + http.Error(w, "invalid lock payload", http.StatusBadRequest) + return + } + body, err := json.Marshal(payload) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + ctx, cancel := context.WithTimeout(r.Context(), 4*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPut, routeURL, bytes.NewReader(body)) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("X-AI-Gate-Token", token) + resp, err := doAccountRequest(h.client, req, account) + if err != nil { + http.Error(w, "update ai-gate upstream lock: "+err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + details, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + http.Error(w, strings.TrimSpace(string(details)), resp.StatusCode) + return + } + var response map[string]any + decoder := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)) + if err := decoder.Decode(&response); err != nil { + http.Error(w, "decode ai-gate upstream lock response: "+err.Error(), http.StatusBadGateway) + return + } + writeJSON(w, http.StatusOK, response) +} + func (h *AccountsHandler) importSharedAccount(w http.ResponseWriter, r *http.Request) { var req importSharedAccountRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { @@ -1620,6 +1811,26 @@ func accountIDFromPath(path string) (int64, error) { return strconv.ParseInt(trimmed, 10, 64) } +func aiGateUpstreamLockIDsFromPath(raw string) (int64, int64, error) { + trimmed := strings.Trim(strings.TrimPrefix(strings.TrimSpace(raw), "/accounts/"), "/") + parts := strings.Split(trimmed, "/") + if len(parts) != 4 || parts[1] != "upstreams" || parts[3] != "lock" { + return 0, 0, errors.New("invalid upstream lock path") + } + accountID, err := strconv.ParseInt(parts[0], 10, 64) + if err != nil { + return 0, 0, err + } + upstreamAccountID, err := strconv.ParseInt(parts[2], 10, 64) + if err != nil { + return 0, 0, err + } + if accountID <= 0 || upstreamAccountID <= 0 { + return 0, 0, errors.New("account id must be positive") + } + return accountID, upstreamAccountID, nil +} + func parseSharedAccountPayload(raw string) (accountShareEnvelope, error) { var envelope accountShareEnvelope if err := json.Unmarshal([]byte(strings.TrimSpace(raw)), &envelope); err != nil { @@ -1691,6 +1902,43 @@ func validateAbsoluteBaseURL(raw string) error { return nil } +func aiGateUpstreamsURL(raw string) (string, error) { + return aiGateAPIURL(raw, "/api/me/upstreams") +} + +func aiGateRouteURL(raw string) (string, error) { + return aiGateAPIURL(raw, "/api/me/route") +} + +func aiGateUpstreamLockURL(raw string, accountID int64) (string, error) { + return aiGateAPIURL(raw, fmt.Sprintf("/api/me/upstreams/%d/lock", accountID)) +} + +func aiGateAPIURL(raw string, suffix string) (string, error) { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil { + return "", fmt.Errorf("invalid base_url: %w", err) + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return "", fmt.Errorf("invalid base_url scheme: %q", parsed.Scheme) + } + if parsed.Host == "" { + return "", errors.New("invalid base_url: missing host") + } + segments := strings.Split(strings.Trim(parsed.Path, "/"), "/") + for index, segment := range segments { + if segment != "ai-gate" { + continue + } + parsed.Path = "/" + strings.Join(segments[:index+1], "/") + suffix + parsed.RawPath = "" + parsed.RawQuery = "" + parsed.Fragment = "" + return parsed.String(), nil + } + return "", errors.New("account is not an ai-gate upstream") +} + func validateUsageConfigJSON(raw string) error { trimmed := strings.TrimSpace(raw) if trimmed == "" { diff --git a/backend/internal/api/accounts_handler_test.go b/backend/internal/api/accounts_handler_test.go index 7af1033..5e76993 100644 --- a/backend/internal/api/accounts_handler_test.go +++ b/backend/internal/api/accounts_handler_test.go @@ -385,6 +385,128 @@ func TestAccountsHandler(t *testing.T) { } } +func TestAccountsHandlerFetchesAIGateUpstreams(t *testing.T) { + t.Parallel() + + var upstreamRoutePayload map[string]any + var upstreamLockPayload map[string]any + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer agt-client" { + t.Fatalf("authorization = %q, want Bearer agt-client", got) + } + if r.Method == http.MethodPut && r.URL.Path == "/ai-gate/api/me/route" { + if err := json.NewDecoder(r.Body).Decode(&upstreamRoutePayload); err != nil { + t.Fatalf("decode upstream route payload: %v", err) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"preferred_account_id":12,"route_locked":true}`) + return + } + if r.Method == http.MethodPut && r.URL.Path == "/ai-gate/api/me/upstreams/12/lock" { + if err := json.NewDecoder(r.Body).Decode(&upstreamLockPayload); err != nil { + t.Fatalf("decode upstream lock payload: %v", err) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":12,"account_locked":true}`) + return + } + if r.Method != http.MethodGet || r.URL.Path != "/ai-gate/api/me/upstreams" { + t.Fatalf("method/path = %s %q, want GET /ai-gate/api/me/upstreams", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{ + "total_accounts": 3, + "available_accounts": 2, + "current_account_id": 11, + "route_locked": false, + "accounts": [ + {"id": 11, "provider_type": "openai-compatible", "account_name": "upstream-a", "auth_mode": "api_key", "base_url": "https://up-a.example/v1", "status": "active", "available": true, "current": true, "preferred": false, "account_locked": false, "supports_responses": true, "balance": 0, "quota_remaining": 0}, + {"id": 12, "provider_type": "openai-compatible", "account_name": "upstream-b", "auth_mode": "api_key", "base_url": "https://up-b.example/v1", "status": "disabled", "available": false, "current": false, "preferred": false, "account_locked": false, "supports_responses": true, "balance": 0, "quota_remaining": 0} + ] + }`) + })) + t.Cleanup(upstream.Close) + + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + t.Cleanup(func() { _ = store.Close() }) + + repo := accounts.NewSQLiteRepository(store.DB()) + if err := repo.Create(accounts.Account{ + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "team-gate", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: upstream.URL + "/ai-gate/v1", + CredentialRef: "agt-client", + Status: accounts.StatusActive, + SupportsResponses: true, + }); err != nil { + t.Fatalf("Create returned error: %v", err) + } + handler := api.NewAccountsHandler( + repo, + nil, + auth.NewOAuthConnector(auth.Config{}), + auth.NewStateStore(5*time.Minute), + api.WithAccountsHTTPClient(upstream.Client()), + ) + + req := httptest.NewRequest(http.MethodGet, "/accounts/1/upstreams", nil) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("GET /accounts/1/upstreams status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + body := rec.Body.String() + if strings.Contains(body, "agt-client") || strings.Contains(body, "credential_ref") { + t.Fatalf("GET /accounts/1/upstreams leaked credential: %s", body) + } + var payload struct { + TotalAccounts int `json:"total_accounts"` + AvailableAccounts int `json:"available_accounts"` + Accounts []struct { + AccountName string `json:"account_name"` + Current bool `json:"current"` + } `json:"accounts"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil { + t.Fatalf("json.Unmarshal returned error: %v", err) + } + if payload.TotalAccounts != 3 || payload.AvailableAccounts != 2 { + t.Fatalf("summary = %+v, want available 2 total 3", payload) + } + if len(payload.Accounts) != 2 || payload.Accounts[0].AccountName != "upstream-a" || !payload.Accounts[0].Current { + t.Fatalf("accounts = %+v, want upstream-a current first", payload.Accounts) + } + + routeReq := httptest.NewRequest(http.MethodPut, "/accounts/1/upstreams/route", bytes.NewBufferString(`{"account_id":12,"locked":true}`)) + routeReq.Header.Set("Content-Type", "application/json") + routeRec := httptest.NewRecorder() + handler.ServeHTTP(routeRec, routeReq) + if routeRec.Code != http.StatusOK { + t.Fatalf("PUT /accounts/1/upstreams/route status = %d, want %d; body=%s", routeRec.Code, http.StatusOK, routeRec.Body.String()) + } + if upstreamRoutePayload["account_id"].(float64) != 12 || upstreamRoutePayload["locked"].(bool) != true { + t.Fatalf("upstream route payload = %+v, want account 12 locked", upstreamRoutePayload) + } + + lockReq := httptest.NewRequest(http.MethodPut, "/accounts/1/upstreams/12/lock", bytes.NewBufferString(`{"locked":true}`)) + lockRec := httptest.NewRecorder() + handler.ServeHTTP(lockRec, lockReq) + if lockRec.Code != http.StatusOK { + t.Fatalf("PUT /accounts/1/upstreams/12/lock status = %d, want %d; body=%s", lockRec.Code, http.StatusOK, lockRec.Body.String()) + } + if upstreamLockPayload["locked"].(bool) != true { + t.Fatalf("upstream lock payload = %+v, want locked true", upstreamLockPayload) + } + if strings.Contains(lockRec.Body.String(), "agt-client") || strings.Contains(lockRec.Body.String(), "credential_ref") { + t.Fatalf("PUT /accounts/1/upstreams/12/lock leaked credential: %s", lockRec.Body.String()) + } +} + func testJWT(t *testing.T, claims map[string]any) string { t.Helper() headerRaw, err := json.Marshal(map[string]any{"alg": "none", "typ": "JWT"}) diff --git a/backend/internal/api/server_me_handler.go b/backend/internal/api/server_me_handler.go index 7942872..4800a18 100644 --- a/backend/internal/api/server_me_handler.go +++ b/backend/internal/api/server_me_handler.go @@ -1,38 +1,81 @@ package api import ( - "database/sql" "encoding/json" + "errors" "net/http" "strconv" "strings" + "time" + "github.com/gcssloop/codex-router/backend/internal/accounts" + "github.com/gcssloop/codex-router/backend/internal/routing" "github.com/gcssloop/codex-router/backend/internal/serverauth" "github.com/gcssloop/codex-router/backend/internal/serverusers" + "github.com/gcssloop/codex-router/backend/internal/usage" ) type ServerMeStore interface { Get(id int64) (serverusers.User, error) - ListAssignedAccounts(userID int64) ([]serverusers.AssignedAccount, error) - UpdateAccountState(userID int64, accountID int64, update serverusers.AccountStateUpdate) error + UpdateRoute(id int64, accountID *int64, locked bool) (serverusers.User, error) +} + +type ServerMeAccounts interface { + List() ([]accounts.Account, error) + Update(account accounts.Account) error +} + +type ServerMeUsage interface { + ListLatest() ([]usage.Snapshot, error) } type ServerMeHandler struct { - store ServerMeStore + store ServerMeStore + accounts ServerMeAccounts + usage ServerMeUsage + sticky *routing.StickySelector } -func NewServerMeHandler(store ServerMeStore) *ServerMeHandler { - return &ServerMeHandler{store: store} +type ServerMeHandlerOption func(*ServerMeHandler) + +func WithServerMeAccounts(repo ServerMeAccounts) ServerMeHandlerOption { + return func(handler *ServerMeHandler) { + handler.accounts = repo + } +} + +func WithServerMeUsage(repo ServerMeUsage) ServerMeHandlerOption { + return func(handler *ServerMeHandler) { + handler.usage = repo + } +} + +func WithServerMeStickySelector(sticky *routing.StickySelector) ServerMeHandlerOption { + return func(handler *ServerMeHandler) { + handler.sticky = sticky + } +} + +func NewServerMeHandler(store ServerMeStore, opts ...ServerMeHandlerOption) *ServerMeHandler { + handler := &ServerMeHandler{store: store} + for _, opt := range opts { + if opt != nil { + opt(handler) + } + } + return handler } func (h *ServerMeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { switch { case r.Method == http.MethodGet && r.URL.Path == "/me": h.me(w, r) - case r.Method == http.MethodGet && r.URL.Path == "/me/accounts": - h.accounts(w, r) - case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/me/accounts/") && strings.HasSuffix(r.URL.Path, "/state"): - h.updateAccountState(w, r) + case r.Method == http.MethodGet && r.URL.Path == "/me/upstreams": + h.upstreams(w, r) + case r.Method == http.MethodPut && r.URL.Path == "/me/route": + h.updateRoute(w, r) + case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/me/upstreams/") && strings.HasSuffix(r.URL.Path, "/lock"): + h.updateUpstreamLock(w, r) default: http.NotFound(w, r) } @@ -56,51 +99,316 @@ func (h *ServerMeHandler) me(w http.ResponseWriter, r *http.Request) { }) } -func (h *ServerMeHandler) accounts(w http.ResponseWriter, r *http.Request) { +func (h *ServerMeHandler) upstreams(w http.ResponseWriter, r *http.Request) { session, ok := serverauth.SessionFromContext(r.Context()) if !ok || session.UserID <= 0 { http.Error(w, "user session required", http.StatusUnauthorized) return } - accounts, err := h.store.ListAssignedAccounts(session.UserID) + if h.accounts == nil { + http.Error(w, "upstream accounts are not configured", http.StatusInternalServerError) + return + } + user, err := h.store.Get(session.UserID) if err != nil { writeServerUserStoreError(w, err) return } - writeJSON(w, http.StatusOK, accounts) + accountList, err := h.accounts.List() + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + usageByAccount := map[int64]usage.Snapshot{} + if h.usage != nil { + if snapshots, err := h.usage.ListLatest(); err == nil { + for _, snapshot := range snapshots { + usageByAccount[snapshot.AccountID] = snapshot + } + } + } + currentAccountID := h.currentAccountID(user, accountList, usageByAccount) + response := h.buildUpstreamsResponse(user, accountList, usageByAccount, currentAccountID) + writeJSON(w, http.StatusOK, response) } -func (h *ServerMeHandler) updateAccountState(w http.ResponseWriter, r *http.Request) { +func (h *ServerMeHandler) updateRoute(w http.ResponseWriter, r *http.Request) { session, ok := serverauth.SessionFromContext(r.Context()) if !ok || session.UserID <= 0 { http.Error(w, "user session required", http.StatusUnauthorized) return } - accountID, ok := meAccountIDFromStatePath(r.URL.Path) - if !ok { - http.NotFound(w, r) - return + var payload struct { + AccountID *int64 `json:"account_id"` + Locked bool `json:"locked"` } - var payload serverusers.AccountStateUpdate if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { - http.Error(w, "invalid account state payload", http.StatusBadRequest) + http.Error(w, "invalid route payload", http.StatusBadRequest) return } - if err := h.store.UpdateAccountState(session.UserID, accountID, payload); err != nil { - if err == sql.ErrNoRows { - http.Error(w, "assigned account not found", http.StatusNotFound) + if payload.AccountID != nil { + if *payload.AccountID <= 0 { + http.Error(w, "account_id must be positive", http.StatusBadRequest) + return + } + if h.accounts == nil { + http.Error(w, "upstream accounts are not configured", http.StatusInternalServerError) + return + } + account, err := h.findAccount(*payload.AccountID) + if err != nil { + http.Error(w, "upstream account not found", http.StatusNotFound) + return + } + if !serverMeAccountManuallySelectable(account) { + http.Error(w, "upstream account is not available", http.StatusBadRequest) return } + } + if payload.AccountID == nil && payload.Locked { + http.Error(w, "locked route requires account_id", http.StatusBadRequest) + return + } + user, err := h.store.UpdateRoute(session.UserID, payload.AccountID, payload.Locked) + if err != nil { writeServerUserStoreError(w, err) return } - writeJSON(w, http.StatusOK, map[string]any{"updated": true}) + if h.sticky != nil && payload.AccountID != nil { + rememberServerUserSticky(h.sticky, session.UserID, *payload.AccountID) + } + writeJSON(w, http.StatusOK, map[string]any{ + "user": user, + "preferred_account_id": user.PreferredAccountID, + "route_locked": user.RouteLocked, + }) +} + +func (h *ServerMeHandler) updateUpstreamLock(w http.ResponseWriter, r *http.Request) { + session, ok := serverauth.SessionFromContext(r.Context()) + if !ok || session.UserID <= 0 { + http.Error(w, "user session required", http.StatusUnauthorized) + return + } + if h.accounts == nil { + http.Error(w, "upstream accounts are not configured", http.StatusInternalServerError) + return + } + accountID, err := serverMeUpstreamAccountIDFromPath(r.URL.Path) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + var payload struct { + Locked bool `json:"locked"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + http.Error(w, "invalid lock payload", http.StatusBadRequest) + return + } + account, err := h.findAccount(accountID) + if err != nil { + http.Error(w, "upstream account not found", http.StatusNotFound) + return + } + account.IsLocked = payload.Locked + if err := h.accounts.Update(account); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "id": account.ID, + "account_locked": account.IsLocked, + }) } -func meAccountIDFromStatePath(path string) (int64, bool) { - trimmed := strings.TrimPrefix(path, "/me/accounts/") - trimmed = strings.TrimSuffix(trimmed, "/state") +func serverMeUpstreamAccountIDFromPath(raw string) (int64, error) { + trimmed := strings.TrimPrefix(strings.TrimSpace(raw), "/me/upstreams/") + trimmed = strings.TrimSuffix(trimmed, "/lock") trimmed = strings.Trim(trimmed, "/") - id, err := strconv.ParseInt(trimmed, 10, 64) - return id, err == nil && id > 0 + if trimmed == "" { + return 0, errors.New("missing account id") + } + return strconv.ParseInt(trimmed, 10, 64) +} + +type serverMeUpstreamsResponse struct { + TotalAccounts int `json:"total_accounts"` + AvailableAccounts int `json:"available_accounts"` + CurrentAccountID int64 `json:"current_account_id,omitempty"` + PreferredAccountID *int64 `json:"preferred_account_id,omitempty"` + RouteLocked bool `json:"route_locked"` + Accounts []serverMeUpstreamAccount `json:"accounts"` +} + +type serverMeUpstreamAccount struct { + ID int64 `json:"id"` + ProviderType accounts.ProviderType `json:"provider_type"` + AccountName string `json:"account_name"` + AuthMode accounts.AuthMode `json:"auth_mode"` + SourceIcon string `json:"source_icon"` + BaseURL string `json:"base_url"` + Status accounts.Status `json:"status"` + Available bool `json:"available"` + Current bool `json:"current"` + Preferred bool `json:"preferred"` + AccountLocked bool `json:"account_locked"` + SupportsResponses bool `json:"supports_responses"` + CooldownRemainingSeconds *int64 `json:"cooldown_remaining_seconds,omitempty"` + RoutingCooldownRemainingSeconds *int64 `json:"routing_cooldown_remaining_seconds,omitempty"` + RoutingCooldownReason string `json:"routing_cooldown_reason,omitempty"` + Balance float64 `json:"balance"` + QuotaRemaining float64 `json:"quota_remaining"` + RPMRemaining float64 `json:"rpm_remaining"` + TPMRemaining float64 `json:"tpm_remaining"` + HealthScore float64 `json:"health_score"` + RecentErrorRate float64 `json:"recent_error_rate"` + LastTotalTokens float64 `json:"last_total_tokens"` + LastInputTokens float64 `json:"last_input_tokens"` + LastOutputTokens float64 `json:"last_output_tokens"` + ModelContextWindow float64 `json:"model_context_window"` + PrimaryUsedPercent float64 `json:"primary_used_percent"` + SecondaryUsedPercent float64 `json:"secondary_used_percent"` + PrimaryResetsAt *time.Time `json:"primary_resets_at,omitempty"` + SecondaryResetsAt *time.Time `json:"secondary_resets_at,omitempty"` + CheckedAt *time.Time `json:"checked_at,omitempty"` + Stale bool `json:"stale"` + LastError string `json:"last_error,omitempty"` + PPChatTodayUsedQuota float64 `json:"ppchat_today_used_quota,omitempty"` + PPChatTodayAddedQuota float64 `json:"ppchat_today_added_quota,omitempty"` + PPChatTodayRemainingQuota float64 `json:"ppchat_today_remaining_quota,omitempty"` + UsageDisplay map[string]any `json:"usage_display,omitempty"` +} + +func (h *ServerMeHandler) buildUpstreamsResponse(user serverusers.User, accountList []accounts.Account, usageByAccount map[int64]usage.Snapshot, currentAccountID int64) serverMeUpstreamsResponse { + now := time.Now().UTC() + items := make([]serverMeUpstreamAccount, 0, len(accountList)) + availableCount := 0 + for _, account := range accountList { + account = applyBuiltInDriverDefaults(account) + snapshot := usageByAccount[account.ID] + available := serverMeAccountAvailable(account, now) + if available { + availableCount++ + } + preferred := user.PreferredAccountID != nil && *user.PreferredAccountID == account.ID + ppchatSummary := parsePPChatUsageSummary(snapshot.ProviderSnapshotJSON) + item := serverMeUpstreamAccount{ + ID: account.ID, + ProviderType: account.ProviderType, + AccountName: account.AccountName, + AuthMode: account.AuthMode, + SourceIcon: normalizeAccountSourceIcon(account.SourceIcon), + BaseURL: account.BaseURL, + Status: account.Status, + Available: available, + Current: currentAccountID == account.ID, + Preferred: preferred, + AccountLocked: account.IsLocked, + SupportsResponses: account.NativeResponsesCapable(), + Balance: snapshot.Balance, + QuotaRemaining: snapshot.QuotaRemaining, + RPMRemaining: snapshot.RPMRemaining, + TPMRemaining: snapshot.TPMRemaining, + HealthScore: snapshot.HealthScore, + RecentErrorRate: snapshot.RecentErrorRate, + LastTotalTokens: snapshot.LastTotalTokens, + LastInputTokens: snapshot.LastInputTokens, + LastOutputTokens: snapshot.LastOutputTokens, + ModelContextWindow: snapshot.ModelContextWindow, + PrimaryUsedPercent: snapshot.PrimaryUsedPercent, + SecondaryUsedPercent: snapshot.SecondaryUsedPercent, + PrimaryResetsAt: snapshot.PrimaryResetsAt, + SecondaryResetsAt: snapshot.SecondaryResetsAt, + CheckedAt: nilIfZeroTime(snapshot.CheckedAt), + Stale: snapshot.Stale, + LastError: snapshot.LastError, + PPChatTodayUsedQuota: ppchatSummary.TodayUsedQuota, + PPChatTodayAddedQuota: ppchatSummary.TodayAddedQuota, + PPChatTodayRemainingQuota: ppchatSummary.TodayRemainingQuota, + UsageDisplay: parseUsageDisplay(snapshot.ProviderSnapshotJSON), + } + if account.CooldownUntil != nil { + remaining := int64(account.CooldownUntil.Sub(now).Seconds()) + if remaining < 0 { + remaining = 0 + } + item.CooldownRemainingSeconds = &remaining + item.RoutingCooldownRemainingSeconds = &remaining + item.RoutingCooldownReason = account.CooldownReason + } + items = append(items, item) + } + return serverMeUpstreamsResponse{ + TotalAccounts: len(items), + AvailableAccounts: availableCount, + CurrentAccountID: currentAccountID, + PreferredAccountID: user.PreferredAccountID, + RouteLocked: user.RouteLocked, + Accounts: items, + } +} + +func (h *ServerMeHandler) currentAccountID(user serverusers.User, accountList []accounts.Account, usageByAccount map[int64]usage.Snapshot) int64 { + if user.RouteLocked && user.PreferredAccountID != nil { + return *user.PreferredAccountID + } + if h.sticky != nil { + if accountID, ok := h.sticky.Current(serverUserRouteScope(user.ID, "responses")); ok { + return accountID + } + if accountID, ok := h.sticky.Current(serverUserRouteScope(user.ID, "chat_completions")); ok { + return accountID + } + } + if user.PreferredAccountID != nil { + return *user.PreferredAccountID + } + candidates := make([]routing.Candidate, 0, len(accountList)) + for _, account := range accountList { + if !serverMeAccountAvailable(account, time.Now().UTC()) { + continue + } + candidates = append(candidates, routing.Candidate{Account: account, Snapshot: usageByAccount[account.ID]}) + } + scored := routing.ScoreCandidates(candidates) + if len(scored) == 0 { + return 0 + } + return scored[0].Account.ID +} + +func (h *ServerMeHandler) findAccount(accountID int64) (accounts.Account, error) { + accountList, err := h.accounts.List() + if err != nil { + return accounts.Account{}, err + } + for _, account := range accountList { + if account.ID == accountID { + return account, nil + } + } + return accounts.Account{}, errors.New("account not found") +} + +func serverMeAccountAvailable(account accounts.Account, now time.Time) bool { + switch account.Status { + case accounts.StatusDisabled, accounts.StatusInvalid: + return false + } + if account.IsLocked { + return false + } + if account.RoutingCooldownActive(now) { + return false + } + return true +} + +func serverMeAccountManuallySelectable(account accounts.Account) bool { + switch account.Status { + case accounts.StatusDisabled, accounts.StatusInvalid: + return false + } + return true } diff --git a/backend/internal/api/server_me_handler_test.go b/backend/internal/api/server_me_handler_test.go index 6edc4f9..344d2a8 100644 --- a/backend/internal/api/server_me_handler_test.go +++ b/backend/internal/api/server_me_handler_test.go @@ -6,7 +6,7 @@ import ( "net/http" "net/http/httptest" "path/filepath" - "strconv" + "strings" "testing" "time" @@ -17,7 +17,7 @@ import ( "github.com/gcssloop/codex-router/backend/internal/usage" ) -func TestServerMeHandlerReturnsOwnUsageAndAssignedAccounts(t *testing.T) { +func TestServerMeHandlerReturnsOwnUsageWithoutAccountPool(t *testing.T) { store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) if err != nil { t.Fatalf("Open returned error: %v", err) @@ -26,46 +26,15 @@ func TestServerMeHandlerReturnsOwnUsageAndAssignedAccounts(t *testing.T) { _ = store.Close() }) - accountRepo := accounts.NewSQLiteRepository(store.DB()) - if err := accountRepo.Create(accounts.Account{ - ProviderType: accounts.ProviderOpenAICompatible, - AccountName: "assigned", - AuthMode: accounts.AuthModeAPIKey, - CredentialRef: "sk-assigned", - BaseURL: "https://assigned.example.invalid/v1", - Status: accounts.StatusActive, - }); err != nil { - t.Fatalf("Create assigned account returned error: %v", err) - } - if err := accountRepo.Create(accounts.Account{ - ProviderType: accounts.ProviderOpenAICompatible, - AccountName: "unassigned", - AuthMode: accounts.AuthModeAPIKey, - CredentialRef: "sk-unassigned", - BaseURL: "https://unassigned.example.invalid/v1", - Status: accounts.StatusActive, - }); err != nil { - t.Fatalf("Create unassigned account returned error: %v", err) - } - accountsList, err := accountRepo.List() - if err != nil { - t.Fatalf("List accounts returned error: %v", err) - } - assignedID := accountsList[0].ID - unassignedID := accountsList[1].ID - userRepo := serverusers.NewSQLiteRepository(store.DB()) created, err := userRepo.Create("alice") if err != nil { t.Fatalf("Create user returned error: %v", err) } - if err := userRepo.SetAccountAssignments(created.User.ID, []int64{assignedID}); err != nil { - t.Fatalf("SetAccountAssignments returned error: %v", err) - } userID := created.User.ID usageRepo := usage.NewSQLiteRepository(store.DB()) if err := usageRepo.SaveEvent(usage.Event{ - AccountID: assignedID, + AccountID: 1, ServerUserID: &userID, ProviderType: "openai-compatible", RequestKind: "responses", @@ -106,32 +75,240 @@ func TestServerMeHandlerReturnsOwnUsageAndAssignedAccounts(t *testing.T) { accountsReq = accountsReq.WithContext(serverauth.ContextWithSession(accountsReq.Context(), session)) accountsRec := httptest.NewRecorder() handler.ServeHTTP(accountsRec, accountsReq) - if accountsRec.Code != http.StatusOK { - t.Fatalf("GET /me/accounts status = %d, want %d; body=%s", accountsRec.Code, http.StatusOK, accountsRec.Body.String()) - } - var assigned []serverusers.AssignedAccount - if err := json.Unmarshal(accountsRec.Body.Bytes(), &assigned); err != nil { - t.Fatalf("unmarshal assigned accounts: %v", err) - } - if len(assigned) != 1 || assigned[0].AccountID != assignedID || assigned[0].CredentialRef != "" { - t.Fatalf("assigned = %+v, want only assigned account without credential", assigned) + if accountsRec.Code != http.StatusNotFound { + t.Fatalf("GET /me/accounts status = %d, want %d; body=%s", accountsRec.Code, http.StatusNotFound, accountsRec.Body.String()) } - stateReq := httptest.NewRequest(http.MethodPut, "/me/accounts/"+strconv.FormatInt(assignedID, 10)+"/state", bytes.NewBufferString(`{"position":3,"is_active":true,"is_locked":true}`)) - stateReq.Header.Set("Content-Type", "application/json") + stateReq := httptest.NewRequest(http.MethodPut, "/me/accounts/1/state", nil) stateReq = stateReq.WithContext(serverauth.ContextWithSession(stateReq.Context(), session)) stateRec := httptest.NewRecorder() handler.ServeHTTP(stateRec, stateReq) - if stateRec.Code != http.StatusOK { - t.Fatalf("PUT assigned state status = %d, want %d; body=%s", stateRec.Code, http.StatusOK, stateRec.Body.String()) + if stateRec.Code != http.StatusNotFound { + t.Fatalf("PUT /me/accounts/1/state status = %d, want %d", stateRec.Code, http.StatusNotFound) } +} - unassignedReq := httptest.NewRequest(http.MethodPut, "/me/accounts/"+strconv.FormatInt(unassignedID, 10)+"/state", bytes.NewBufferString(`{"position":0,"is_active":true,"is_locked":false}`)) - unassignedReq.Header.Set("Content-Type", "application/json") - unassignedReq = unassignedReq.WithContext(serverauth.ContextWithSession(unassignedReq.Context(), session)) - unassignedRec := httptest.NewRecorder() - handler.ServeHTTP(unassignedRec, unassignedReq) - if unassignedRec.Code != http.StatusNotFound { - t.Fatalf("PUT unassigned state status = %d, want %d", unassignedRec.Code, http.StatusNotFound) +func TestServerMeHandlerReturnsSanitizedUpstreamsAndUpdatesRoute(t *testing.T) { + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + t.Cleanup(func() { + _ = store.Close() + }) + + userRepo := serverusers.NewSQLiteRepository(store.DB()) + created, err := userRepo.Create("alice") + if err != nil { + t.Fatalf("Create user returned error: %v", err) + } + accountRepo := accounts.NewSQLiteRepository(store.DB()) + for _, account := range []accounts.Account{ + { + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "account-a", + SourceIcon: "openai", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: "https://upstream-a.example/v1", + CredentialRef: "sk-secret-a", + UsageConfigJSON: `{"secret":"do-not-leak"}`, + Status: accounts.StatusActive, + Priority: 100, + SupportsResponses: true, + }, + { + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "account-b", + SourceIcon: "ppchat", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: "https://upstream-b.example/v1", + CredentialRef: "sk-secret-b", + Status: accounts.StatusDisabled, + Priority: 10, + SupportsResponses: true, + }, + } { + if err := accountRepo.Create(account); err != nil { + t.Fatalf("Create account returned error: %v", err) + } + } + usageRepo := usage.NewSQLiteRepository(store.DB()) + if err := usageRepo.Save(usage.Snapshot{ + AccountID: 1, + Balance: 10, + QuotaRemaining: 1000, + RPMRemaining: 20, + TPMRemaining: 3000, + HealthScore: 0.9, + LastTotalTokens: 120, + LastInputTokens: 40, + LastOutputTokens: 80, + CheckedAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("Save usage returned error: %v", err) + } + if err := usageRepo.Save(usage.Snapshot{ + AccountID: 2, + QuotaRemaining: -2, + HealthScore: 1, + CheckedAt: time.Now().UTC(), + ProviderSnapshotJSON: `{ + "payload": { + "data": { + "token_info": { + "today_used_quota": 22, + "today_added_quota": 20, + "remain_quota_display": -2 + } + } + } + }`, + }); err != nil { + t.Fatalf("Save ppchat usage returned error: %v", err) + } + + handler := NewServerMeHandler(userRepo, WithServerMeAccounts(accountRepo), WithServerMeUsage(usageRepo)) + session := serverauth.Session{Authenticated: true, Role: serverusers.RoleUser, UserID: created.User.ID, Username: "alice"} + + upstreamsReq := httptest.NewRequest(http.MethodGet, "/me/upstreams", nil) + upstreamsReq = upstreamsReq.WithContext(serverauth.ContextWithSession(upstreamsReq.Context(), session)) + upstreamsRec := httptest.NewRecorder() + handler.ServeHTTP(upstreamsRec, upstreamsReq) + if upstreamsRec.Code != http.StatusOK { + t.Fatalf("GET /me/upstreams status = %d, want %d; body=%s", upstreamsRec.Code, http.StatusOK, upstreamsRec.Body.String()) + } + body := upstreamsRec.Body.String() + for _, leaked := range []string{"sk-secret-a", "sk-secret-b", "credential_ref", "usage_config_json", "do-not-leak"} { + if strings.Contains(body, leaked) { + t.Fatalf("GET /me/upstreams body leaked %q: %s", leaked, body) + } + } + var upstreams struct { + TotalAccounts int `json:"total_accounts"` + AvailableAccounts int `json:"available_accounts"` + CurrentAccountID int64 `json:"current_account_id"` + RouteLocked bool `json:"route_locked"` + Accounts []struct { + ID int64 `json:"id"` + AccountName string `json:"account_name"` + BaseURL string `json:"base_url"` + Status string `json:"status"` + Available bool `json:"available"` + Current bool `json:"current"` + Balance float64 `json:"balance"` + QuotaRemaining float64 `json:"quota_remaining"` + LastTotalTokens float64 `json:"last_total_tokens"` + PPChatTodayUsedQuota float64 `json:"ppchat_today_used_quota"` + PPChatTodayAddedQuota float64 `json:"ppchat_today_added_quota"` + PPChatTodayRemainingQuota float64 `json:"ppchat_today_remaining_quota"` + } `json:"accounts"` + } + if err := json.Unmarshal(upstreamsRec.Body.Bytes(), &upstreams); err != nil { + t.Fatalf("unmarshal upstreams: %v", err) + } + if upstreams.TotalAccounts != 2 || upstreams.AvailableAccounts != 1 || upstreams.CurrentAccountID != 1 { + t.Fatalf("upstreams summary = %+v, want total 2 available 1 current 1", upstreams) + } + if len(upstreams.Accounts) != 2 || !upstreams.Accounts[0].Current || upstreams.Accounts[0].Balance != 10 || upstreams.Accounts[0].LastTotalTokens != 120 { + t.Fatalf("upstream accounts = %+v, want first account current with usage", upstreams.Accounts) + } + if upstreams.Accounts[1].PPChatTodayUsedQuota != 22 || upstreams.Accounts[1].PPChatTodayAddedQuota != 20 || upstreams.Accounts[1].PPChatTodayRemainingQuota != -2 { + t.Fatalf("ppchat upstream usage = %+v, want cached ppchat quota fields", upstreams.Accounts[1]) + } + + routeReq := httptest.NewRequest(http.MethodPut, "/me/route", bytes.NewBufferString(`{"account_id":1,"locked":true}`)) + routeReq.Header.Set("Content-Type", "application/json") + routeReq = routeReq.WithContext(serverauth.ContextWithSession(routeReq.Context(), session)) + routeRec := httptest.NewRecorder() + handler.ServeHTTP(routeRec, routeReq) + if routeRec.Code != http.StatusOK { + t.Fatalf("PUT /me/route status = %d, want %d; body=%s", routeRec.Code, http.StatusOK, routeRec.Body.String()) + } + updated, err := userRepo.Get(created.User.ID) + if err != nil { + t.Fatalf("Get updated user returned error: %v", err) + } + if updated.PreferredAccountID == nil || *updated.PreferredAccountID != 1 || !updated.RouteLocked { + t.Fatalf("updated route = account:%v locked:%v, want account 1 locked", updated.PreferredAccountID, updated.RouteLocked) + } +} + +func TestServerMeHandlerLocksUpstreamAccountAndAllowsManualRoute(t *testing.T) { + store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) + if err != nil { + t.Fatalf("Open returned error: %v", err) + } + t.Cleanup(func() { + _ = store.Close() + }) + + userRepo := serverusers.NewSQLiteRepository(store.DB()) + created, err := userRepo.Create("alice") + if err != nil { + t.Fatalf("Create user returned error: %v", err) + } + accountRepo := accounts.NewSQLiteRepository(store.DB()) + if err := accountRepo.Create(accounts.Account{ + ProviderType: accounts.ProviderOpenAICompatible, + AccountName: "manual-locked", + SourceIcon: "openai", + AuthMode: accounts.AuthModeAPIKey, + BaseURL: "https://locked.example/v1", + CredentialRef: "sk-locked", + Status: accounts.StatusActive, + Priority: 100, + SupportsResponses: true, + }); err != nil { + t.Fatalf("Create account returned error: %v", err) + } + + handler := NewServerMeHandler(userRepo, WithServerMeAccounts(accountRepo)) + session := serverauth.Session{Authenticated: true, Role: serverusers.RoleUser, UserID: created.User.ID, Username: "alice"} + + lockReq := httptest.NewRequest(http.MethodPut, "/me/upstreams/1/lock", bytes.NewBufferString(`{"locked":true}`)) + lockReq.Header.Set("Content-Type", "application/json") + lockReq = lockReq.WithContext(serverauth.ContextWithSession(lockReq.Context(), session)) + lockRec := httptest.NewRecorder() + handler.ServeHTTP(lockRec, lockReq) + if lockRec.Code != http.StatusOK { + t.Fatalf("PUT /me/upstreams/1/lock status = %d, want %d; body=%s", lockRec.Code, http.StatusOK, lockRec.Body.String()) + } + + upstreamsReq := httptest.NewRequest(http.MethodGet, "/me/upstreams", nil) + upstreamsReq = upstreamsReq.WithContext(serverauth.ContextWithSession(upstreamsReq.Context(), session)) + upstreamsRec := httptest.NewRecorder() + handler.ServeHTTP(upstreamsRec, upstreamsReq) + if upstreamsRec.Code != http.StatusOK { + t.Fatalf("GET /me/upstreams status = %d, want %d; body=%s", upstreamsRec.Code, http.StatusOK, upstreamsRec.Body.String()) + } + var upstreams struct { + AvailableAccounts int `json:"available_accounts"` + Accounts []struct { + ID int64 `json:"id"` + Available bool `json:"available"` + AccountLocked bool `json:"account_locked"` + } `json:"accounts"` + } + if err := json.Unmarshal(upstreamsRec.Body.Bytes(), &upstreams); err != nil { + t.Fatalf("unmarshal upstreams: %v", err) + } + if upstreams.AvailableAccounts != 0 || len(upstreams.Accounts) != 1 || upstreams.Accounts[0].Available || !upstreams.Accounts[0].AccountLocked { + t.Fatalf("upstreams = %+v, want locked unavailable account", upstreams) + } + + routeReq := httptest.NewRequest(http.MethodPut, "/me/route", bytes.NewBufferString(`{"account_id":1,"locked":false}`)) + routeReq.Header.Set("Content-Type", "application/json") + routeReq = routeReq.WithContext(serverauth.ContextWithSession(routeReq.Context(), session)) + routeRec := httptest.NewRecorder() + handler.ServeHTTP(routeRec, routeReq) + if routeRec.Code != http.StatusOK { + t.Fatalf("PUT /me/route locked upstream status = %d, want %d; body=%s", routeRec.Code, http.StatusOK, routeRec.Body.String()) + } + updated, err := userRepo.Get(created.User.ID) + if err != nil { + t.Fatalf("Get updated user returned error: %v", err) + } + if updated.PreferredAccountID == nil || *updated.PreferredAccountID != 1 || updated.RouteLocked { + t.Fatalf("updated route = account:%v locked:%v, want unlocked manual account 1", updated.PreferredAccountID, updated.RouteLocked) } } diff --git a/backend/internal/api/server_user_route.go b/backend/internal/api/server_user_route.go new file mode 100644 index 0000000..f433de91 --- /dev/null +++ b/backend/internal/api/server_user_route.go @@ -0,0 +1,91 @@ +package api + +import ( + "context" + "fmt" + + "github.com/gcssloop/codex-router/backend/internal/routing" +) + +func serverUserRouteScope(userID int64, requestKind string) string { + if userID <= 0 || requestKind == "" { + return "" + } + return fmt.Sprintf("%s:user:%d", requestKind, userID) +} + +func gatewayStickyScope(ctx context.Context, requestKind string) string { + if user, ok := ServerUserFromContext(ctx); ok && user.ID > 0 { + return serverUserRouteScope(user.ID, requestKind) + } + return "" +} + +func rememberServerUserSticky(sticky *routing.StickySelector, userID int64, accountID int64) { + if sticky == nil || userID <= 0 || accountID <= 0 { + return + } + sticky.Remember(serverUserRouteScope(userID, "responses"), accountID) + sticky.Remember(serverUserRouteScope(userID, "chat_completions"), accountID) +} + +func orderServerUserCandidates(ctx context.Context, sticky *routing.StickySelector, requestKind string, candidates []routing.Candidate) []routing.Candidate { + scope := gatewayStickyScope(ctx, requestKind) + ordered := sticky.Apply(scope, candidates) + user, ok := ServerUserFromContext(ctx) + if !ok || user.PreferredAccountID == nil || *user.PreferredAccountID <= 0 { + return ordered + } + ordered = markServerUserManualCandidate(ordered, *user.PreferredAccountID) + if user.RouteLocked { + return moveServerUserCandidateToFront(ordered, *user.PreferredAccountID) + } + if sticky == nil { + return moveServerUserCandidateToFront(ordered, *user.PreferredAccountID) + } + if _, ok := sticky.Current(scope); ok { + return ordered + } + return moveServerUserCandidateToFront(ordered, *user.PreferredAccountID) +} + +func markServerUserManualCandidate(candidates []routing.Candidate, accountID int64) []routing.Candidate { + if accountID <= 0 || len(candidates) == 0 { + return candidates + } + marked := make([]routing.Candidate, len(candidates)) + copy(marked, candidates) + for index := range marked { + if marked[index].Account.ID == accountID { + marked[index].Account.IsActive = true + break + } + } + return marked +} + +func moveServerUserCandidateToFront(candidates []routing.Candidate, accountID int64) []routing.Candidate { + if accountID <= 0 || len(candidates) <= 1 { + return candidates + } + ordered := make([]routing.Candidate, len(candidates)) + copy(ordered, candidates) + for index, candidate := range ordered { + if candidate.Account.ID != accountID { + continue + } + if index == 0 { + return ordered + } + selected := ordered[index] + copy(ordered[1:index+1], ordered[0:index]) + ordered[0] = selected + return ordered + } + return ordered +} + +func serverUserFromContextExists(ctx context.Context) bool { + _, ok := ServerUserFromContext(ctx) + return ok +} diff --git a/backend/internal/api/server_users_handler.go b/backend/internal/api/server_users_handler.go index acd24f0..6e8f058 100644 --- a/backend/internal/api/server_users_handler.go +++ b/backend/internal/api/server_users_handler.go @@ -14,9 +14,8 @@ type ServerUsersStore interface { Create(name string) (serverusers.CreatedUser, error) List() ([]serverusers.User, error) Disable(id int64) error + Delete(id int64) error RotateToken(id int64) (serverusers.CreatedUser, error) - ListAccountAssignments(userID int64) ([]serverusers.AccountAssignment, error) - SetAccountAssignments(userID int64, accountIDs []int64) error } type ServerUsersHandler struct { @@ -33,10 +32,8 @@ func (h *ServerUsersHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.list(w) case r.Method == http.MethodPost && r.URL.Path == "/server-users": h.create(w, r) - case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/server-users/") && strings.HasSuffix(r.URL.Path, "/accounts"): - h.listAccountAssignments(w, r) - case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/server-users/") && strings.HasSuffix(r.URL.Path, "/accounts"): - h.setAccountAssignments(w, r) + case r.Method == http.MethodDelete && strings.HasPrefix(r.URL.Path, "/server-users/") && countPathSegments(r.URL.Path) == 2: + h.delete(w, r) case r.Method == http.MethodPost && strings.HasPrefix(r.URL.Path, "/server-users/") && strings.HasSuffix(r.URL.Path, "/disable"): h.disable(w, r) case r.Method == http.MethodPost && strings.HasPrefix(r.URL.Path, "/server-users/") && strings.HasSuffix(r.URL.Path, "/rotate-token"): @@ -46,40 +43,6 @@ func (h *ServerUsersHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } } -func (h *ServerUsersHandler) listAccountAssignments(w http.ResponseWriter, r *http.Request) { - id, ok := serverUserIDFromActionPath(r.URL.Path, "/accounts") - if !ok { - http.NotFound(w, r) - return - } - assignments, err := h.store.ListAccountAssignments(id) - if err != nil { - writeServerUserStoreError(w, err) - return - } - writeJSON(w, http.StatusOK, assignments) -} - -func (h *ServerUsersHandler) setAccountAssignments(w http.ResponseWriter, r *http.Request) { - id, ok := serverUserIDFromActionPath(r.URL.Path, "/accounts") - if !ok { - http.NotFound(w, r) - return - } - var payload struct { - AccountIDs []int64 `json:"account_ids"` - } - if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { - http.Error(w, "invalid server user account assignment payload", http.StatusBadRequest) - return - } - if err := h.store.SetAccountAssignments(id, payload.AccountIDs); err != nil { - writeServerUserStoreError(w, err) - return - } - writeJSON(w, http.StatusOK, map[string]any{"assigned": true}) -} - func (h *ServerUsersHandler) list(w http.ResponseWriter) { users, err := h.store.List() if err != nil { @@ -118,6 +81,19 @@ func (h *ServerUsersHandler) disable(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, map[string]any{"disabled": true}) } +func (h *ServerUsersHandler) delete(w http.ResponseWriter, r *http.Request) { + id, ok := serverUserIDFromPath(r.URL.Path) + if !ok { + http.NotFound(w, r) + return + } + if err := h.store.Delete(id); err != nil { + writeServerUserStoreError(w, err) + return + } + writeJSON(w, http.StatusOK, map[string]any{"deleted": true}) +} + func (h *ServerUsersHandler) rotate(w http.ResponseWriter, r *http.Request) { id, ok := serverUserIDFromActionPath(r.URL.Path, "/rotate-token") if !ok { @@ -132,6 +108,13 @@ func (h *ServerUsersHandler) rotate(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, created) } +func serverUserIDFromPath(path string) (int64, bool) { + trimmed := strings.TrimPrefix(path, "/server-users/") + trimmed = strings.Trim(trimmed, "/") + id, err := strconv.ParseInt(trimmed, 10, 64) + return id, err == nil && id > 0 +} + func serverUserIDFromActionPath(path string, suffix string) (int64, bool) { trimmed := strings.TrimSuffix(strings.TrimPrefix(path, "/server-users/"), suffix) trimmed = strings.Trim(trimmed, "/") diff --git a/backend/internal/api/server_users_handler_test.go b/backend/internal/api/server_users_handler_test.go index 2a66fa2..d24fc54 100644 --- a/backend/internal/api/server_users_handler_test.go +++ b/backend/internal/api/server_users_handler_test.go @@ -6,10 +6,8 @@ import ( "net/http" "net/http/httptest" "path/filepath" - "strconv" "testing" - "github.com/gcssloop/codex-router/backend/internal/accounts" "github.com/gcssloop/codex-router/backend/internal/serverusers" sqlitestore "github.com/gcssloop/codex-router/backend/internal/store/sqlite" ) @@ -67,9 +65,29 @@ func TestServerUsersHandlerCreateListDisableAndRotate(t *testing.T) { if rotated.Token == "" || rotated.Token == created.Token { t.Fatalf("rotated token = %q, original token = %q", rotated.Token, created.Token) } + + deleteReq := httptest.NewRequest(http.MethodDelete, "/server-users/1", nil) + deleteRec := httptest.NewRecorder() + handler.ServeHTTP(deleteRec, deleteReq) + if deleteRec.Code != http.StatusOK { + t.Fatalf("DELETE /server-users/1 status = %d, want %d; body=%s", deleteRec.Code, http.StatusOK, deleteRec.Body.String()) + } + listAfterDeleteReq := httptest.NewRequest(http.MethodGet, "/server-users", nil) + listAfterDeleteRec := httptest.NewRecorder() + handler.ServeHTTP(listAfterDeleteRec, listAfterDeleteReq) + if listAfterDeleteRec.Code != http.StatusOK { + t.Fatalf("GET /server-users after delete status = %d, want %d; body=%s", listAfterDeleteRec.Code, http.StatusOK, listAfterDeleteRec.Body.String()) + } + var usersAfterDelete []serverusers.User + if err := json.Unmarshal(listAfterDeleteRec.Body.Bytes(), &usersAfterDelete); err != nil { + t.Fatalf("unmarshal users after delete: %v", err) + } + if len(usersAfterDelete) != 0 { + t.Fatalf("users after delete = %+v, want empty", usersAfterDelete) + } } -func TestServerUsersHandlerManagesAccountAssignments(t *testing.T) { +func TestServerUsersHandlerDoesNotExposeAccountAssignments(t *testing.T) { store, err := sqlitestore.Open(filepath.Join(t.TempDir(), "router.sqlite")) if err != nil { t.Fatalf("Open returned error: %v", err) @@ -78,34 +96,6 @@ func TestServerUsersHandlerManagesAccountAssignments(t *testing.T) { _ = store.Close() }) - accountRepo := accounts.NewSQLiteRepository(store.DB()) - createAccount := func(name string) int64 { - t.Helper() - if err := accountRepo.Create(accounts.Account{ - ProviderType: accounts.ProviderOpenAICompatible, - AccountName: name, - AuthMode: accounts.AuthModeAPIKey, - CredentialRef: "sk-" + name, - BaseURL: "https://example.invalid/v1", - Status: accounts.StatusActive, - }); err != nil { - t.Fatalf("Create(account %s) returned error: %v", name, err) - } - accountsList, err := accountRepo.List() - if err != nil { - t.Fatalf("List accounts returned error: %v", err) - } - for _, account := range accountsList { - if account.AccountName == name { - return account.ID - } - } - t.Fatalf("created account %s not found", name) - return 0 - } - firstID := createAccount("first") - secondID := createAccount("second") - repo := serverusers.NewSQLiteRepository(store.DB()) created, err := repo.Create("alice") if err != nil { @@ -123,36 +113,23 @@ func TestServerUsersHandlerManagesAccountAssignments(t *testing.T) { if err := json.Unmarshal(listBeforeRec.Body.Bytes(), &usersBefore); err != nil { t.Fatalf("unmarshal users before: %v", err) } - if len(usersBefore) != 1 || usersBefore[0].AssignedAccounts != 0 { - t.Fatalf("users before = %+v, want zero assigned accounts", usersBefore) + if len(usersBefore) != 1 { + t.Fatalf("users before = %+v, want one user", usersBefore) } - assignReq := httptest.NewRequest(http.MethodPut, "/server-users/1/accounts", bytes.NewBufferString(`{"account_ids":[`+strconv.FormatInt(secondID, 10)+`,`+strconv.FormatInt(firstID, 10)+`]}`)) + assignReq := httptest.NewRequest(http.MethodPut, "/server-users/1/accounts", bytes.NewBufferString(`{"account_ids":[1,2]}`)) assignReq.Header.Set("Content-Type", "application/json") assignRec := httptest.NewRecorder() handler.ServeHTTP(assignRec, assignReq) - if assignRec.Code != http.StatusOK { - t.Fatalf("PUT /server-users/1/accounts status = %d, want %d; body=%s", assignRec.Code, http.StatusOK, assignRec.Body.String()) + if assignRec.Code != http.StatusNotFound { + t.Fatalf("PUT /server-users/1/accounts status = %d, want %d; body=%s", assignRec.Code, http.StatusNotFound, assignRec.Body.String()) } accountsReq := httptest.NewRequest(http.MethodGet, "/server-users/1/accounts", nil) accountsRec := httptest.NewRecorder() handler.ServeHTTP(accountsRec, accountsReq) - if accountsRec.Code != http.StatusOK { - t.Fatalf("GET /server-users/1/accounts status = %d, want %d; body=%s", accountsRec.Code, http.StatusOK, accountsRec.Body.String()) - } - var assignmentView []serverusers.AccountAssignment - if err := json.Unmarshal(accountsRec.Body.Bytes(), &assignmentView); err != nil { - t.Fatalf("unmarshal assignment view: %v", err) - } - if len(assignmentView) != 2 { - t.Fatalf("assignment view = %+v, want two accounts", assignmentView) - } - if !assignmentView[0].Assigned || !assignmentView[1].Assigned { - t.Fatalf("assignment view = %+v, want both assigned", assignmentView) - } - if assignmentView[0].AccountID != secondID || assignmentView[0].Position != 0 { - t.Fatalf("assignment view = %+v, want second account first", assignmentView) + if accountsRec.Code != http.StatusNotFound { + t.Fatalf("GET /server-users/1/accounts status = %d, want %d; body=%s", accountsRec.Code, http.StatusNotFound, accountsRec.Body.String()) } listAfterReq := httptest.NewRequest(http.MethodGet, "/server-users", nil) @@ -165,22 +142,7 @@ func TestServerUsersHandlerManagesAccountAssignments(t *testing.T) { if err := json.Unmarshal(listAfterRec.Body.Bytes(), &usersAfter); err != nil { t.Fatalf("unmarshal users after: %v", err) } - if len(usersAfter) != 1 || usersAfter[0].ID != created.User.ID || usersAfter[0].AssignedAccounts != 2 { - t.Fatalf("users after = %+v, want two assigned accounts", usersAfter) - } - - clearReq := httptest.NewRequest(http.MethodPut, "/server-users/1/accounts", bytes.NewBufferString(`{"account_ids":[]}`)) - clearReq.Header.Set("Content-Type", "application/json") - clearRec := httptest.NewRecorder() - handler.ServeHTTP(clearRec, clearReq) - if clearRec.Code != http.StatusOK { - t.Fatalf("clear assignments status = %d, want %d; body=%s", clearRec.Code, http.StatusOK, clearRec.Body.String()) - } - assigned, err := repo.ListAssignedAccounts(created.User.ID) - if err != nil { - t.Fatalf("ListAssignedAccounts returned error: %v", err) - } - if len(assigned) != 0 { - t.Fatalf("assigned after clear = %+v, want none", assigned) + if len(usersAfter) != 1 || usersAfter[0].ID != created.User.ID { + t.Fatalf("users after = %+v, want created user", usersAfter) } } diff --git a/backend/internal/webui/embed.go b/backend/internal/webui/embed.go index 36f0b03..1f4eba5 100644 --- a/backend/internal/webui/embed.go +++ b/backend/internal/webui/embed.go @@ -16,14 +16,14 @@ func Handler(prefix string) http.Handler { if err != nil { return http.NotFoundHandler() } + trimmedPrefix := "/" + strings.Trim(strings.TrimSpace(prefix), "/") + if trimmedPrefix == "/" { + trimmedPrefix = "" + } + webPrefix := trimmedPrefix + "/webui/" fileServer := http.FileServer(http.FS(dist)) - index := serveIndex(dist) + index := serveIndex(dist, webPrefix) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - trimmedPrefix := "/" + strings.Trim(strings.TrimSpace(prefix), "/") - if trimmedPrefix == "/" { - trimmedPrefix = "" - } - webPrefix := trimmedPrefix + "/webui/" if r.URL.Path == trimmedPrefix+"/webui" { http.Redirect(w, r, webPrefix, http.StatusTemporaryRedirect) return @@ -48,15 +48,24 @@ func Handler(prefix string) http.Handler { }) } -func serveIndex(dist fs.FS) http.Handler { +func serveIndex(dist fs.FS, webPrefix string) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { raw, err := fs.ReadFile(dist, "index.html") if err != nil { http.NotFound(w, r) return } + raw = rewriteIndexWebPrefix(raw, webPrefix) w.Header().Set("Content-Type", "text/html; charset=utf-8") w.WriteHeader(http.StatusOK) _, _ = w.Write(raw) }) } + +func rewriteIndexWebPrefix(raw []byte, webPrefix string) []byte { + normalized := "/" + strings.Trim(strings.TrimSpace(webPrefix), "/") + "/" + body := string(raw) + body = strings.ReplaceAll(body, "/ai-gate/webui/", normalized) + body = strings.ReplaceAll(body, "/ai-router/webui/", normalized) + return []byte(body) +} diff --git a/backend/internal/webui/embed_test.go b/backend/internal/webui/embed_test.go new file mode 100644 index 0000000..aab2f34 --- /dev/null +++ b/backend/internal/webui/embed_test.go @@ -0,0 +1,27 @@ +package webui + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestHandlerRewritesEmbeddedIndexAssetPrefix(t *testing.T) { + handler := Handler("/ai-router") + req := httptest.NewRequest(http.MethodGet, "/ai-router/webui/", nil) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("GET /ai-router/webui/ status = %d, want %d", rec.Code, http.StatusOK) + } + body := rec.Body.String() + if !strings.Contains(body, "/ai-router/webui/assets/") { + t.Fatalf("index body missing ai-router asset prefix: %s", body) + } + if strings.Contains(body, "/ai-gate/webui/assets/") { + t.Fatalf("index body still contains ai-gate asset prefix: %s", body) + } +} diff --git a/docs/plans/2026-06-23-user-upstream-route-control-design.md b/docs/plans/2026-06-23-user-upstream-route-control-design.md new file mode 100644 index 0000000..1f76399 --- /dev/null +++ b/docs/plans/2026-06-23-user-upstream-route-control-design.md @@ -0,0 +1,28 @@ +# 用户上游路由控制设计 + +## 目标 + +普通服务用户通过 AI Gate server 版本登录后，可以在“我的网关”页面查看所有上游账号的脱敏基础状态和用量摘要，并能手动切换或锁定自己的上游路由。 + +## 架构 + +服务用户路由状态是用户级状态，不复用全局 `accounts.is_active` 或 `accounts.is_locked`，避免一个用户的切换影响其他用户。`server_users` 新增 `preferred_account_id` 和 `route_locked` 两个字段；转发时这些字段随现有 token 鉴权查询进入 request context，不新增热路径 DB 查询。 + +服务端新增 `/api/me/upstreams` 和 `/api/me/route`。前者返回账号名、供应商、base_url、状态、可用性、用量摘要、当前使用中账号和总数统计，不返回 `credential_ref`、`usage_config_json` 或任何 token。后者只在用户手动切换、锁定、解锁时写库，并同步更新进程内 sticky selector，让后续转发立即生效。 + +## 路由规则 + +- 默认：继续使用全局最优路径加用户级 sticky。 +- 手动切换：将指定账号写入用户偏好，并为该用户的 chat/responses scope 写入 sticky；后续请求优先命中该账号。 +- 锁定：在用户偏好存在时，每次排序都把指定账号放在首位；如果该账号不可尝试，仍允许跳过并由后续候选兜底。 +- 失败：限流、容量不足、上游不可达时失效该用户对应 scope 的 sticky，不影响其他用户。 + +## 前端 + +`UserPoolPage` 在自助用量下方增加可折叠“上游账号”区域。折叠标题右侧显示“可用 / 总数”；展开后逐个展示账号名、base_url、状态、用量摘要和“当前使用中”标记，并提供“切换”“锁定/解锁”按钮。页面不展示任何上游 token 或配置 JSON。 + +## 验证 + +- 后端：`cd backend && go test ./internal/api ./internal/serverusers ./internal/routing -count=1` +- 前端：`npm --prefix frontend test -- UserPoolPage` +- 构建：`npm --prefix frontend run build` diff --git a/docs/plans/2026-06-23-user-upstream-route-control.md b/docs/plans/2026-06-23-user-upstream-route-control.md new file mode 100644 index 0000000..2a5d597 --- /dev/null +++ b/docs/plans/2026-06-23-user-upstream-route-control.md @@ -0,0 +1,75 @@ +# 用户上游路由控制实施计划 + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** 让普通服务用户能查看脱敏上游账号状态，并控制自己的手动切换和锁定路由。 + +**Architecture:** 用户级路由偏好持久化在 `server_users`，随现有 token 鉴权进入 context，转发热路径不新增 DB 查询。页面 API 返回脱敏账号和用量摘要，手动操作才写库并同步进程内 sticky。 + +**Tech Stack:** Go `net/http`、SQLite、现有 `routing.StickySelector`、React + Ant Design。 + +--- + +### Task 1: 用户路由状态模型 + +**Files:** +- Modify: `backend/internal/store/sqlite/migrations.go` +- Modify: `backend/internal/store/sqlite/store.go` +- Modify: `backend/internal/serverusers/types.go` +- Modify: `backend/internal/serverusers/repository.go` +- Modify: `backend/internal/serverusers/repository_test.go` + +**Steps:** +1. 写失败测试：服务用户可保存 `preferred_account_id` 和 `route_locked`。 +2. 增加 SQLite 列和 repository 读写。 +3. 回跑 `cd backend && go test ./internal/serverusers -count=1`。 + +### Task 2: 用户自助上游 API + +**Files:** +- Modify: `backend/internal/api/server_me_handler.go` +- Modify: `backend/internal/api/server_me_handler_test.go` +- Modify: `backend/internal/bootstrap/bootstrap.go` + +**Steps:** +1. 写失败测试：`GET /me/upstreams` 返回脱敏账号状态、用量、当前账号和统计。 +2. 写失败测试：`PUT /me/route` 可切换和锁定指定账号。 +3. 实现 handler 和 bootstrap 注入账号、用量、sticky selector。 + +### Task 3: 转发路由偏好 + +**Files:** +- Modify: `backend/internal/routing/sticky.go` +- Modify: `backend/internal/routing/sticky_test.go` +- Modify: `backend/internal/api/gateway_handler.go` +- Modify: `backend/internal/api/responses_handler.go` +- Modify: `backend/internal/api/gateway_handler_test.go` +- Modify: `backend/internal/api/responses_handler_test.go` + +**Steps:** +1. 写失败测试：用户偏好账号优先于全局评分，且不写全局 active/cooldown。 +2. 写失败测试：不同用户 sticky scope 互不影响。 +3. 实现用户级 scope、偏好排序和 sticky 查询。 + +### Task 4: 前端自助页面 + +**Files:** +- Modify: `frontend/src/lib/api.ts` +- Modify: `frontend/src/features/server-users/UserPoolPage.tsx` +- Modify: `frontend/src/features/server-users/UserPoolPage.test.tsx` +- Modify: `frontend/src/styles.css` + +**Steps:** +1. 写失败测试：页面展示可用/总数、账号行、当前使用中、切换和锁定按钮。 +2. 实现 API wrapper 和 UI。 +3. 确认不渲染 token 或 credential 字段。 + +### Task 5: 收敛验证 + +**Verification:** +- `cd backend && go test ./... -count=1` +- `npm --prefix frontend test -- UserPoolPage` +- `npm --prefix frontend run build` +- `git diff --check` + +No commit is created unless explicitly requested. From e9e228bc3e50b6bb629a46de14235e6cc6abf3c0 Mon Sep 17 00:00:00 2001 From: GcsSloop Date: Wed, 24 Jun 2026 02:07:28 +0800 Subject: [PATCH 3/4] =?UTF-8?q?feat:=20=E5=B1=95=E7=A4=BA=E5=B9=B6?= =?UTF-8?q?=E6=8E=A7=E5=88=B6=E8=BF=9C=E7=AB=AF=E4=B8=8A=E6=B8=B8=E8=B4=A6?= =?UTF-8?q?=E6=88=B7=E7=8A=B6=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../features/accounts/AccountsPage.test.tsx | 307 +++++++++++ .../src/features/accounts/AccountsPage.tsx | 491 ++++++++++++++++-- .../server-users/ServerUsersPage.test.tsx | 98 +++- .../features/server-users/ServerUsersPage.tsx | 153 +++--- .../server-users/UserPoolPage.test.tsx | 158 +++++- .../features/server-users/UserPoolPage.tsx | 210 ++++++-- frontend/src/lib/api.ts | 187 +++++-- frontend/src/lib/i18n.ts | 4 + frontend/src/styles.css | 324 ++++++++++++ 9 files changed, 1657 insertions(+), 275 deletions(-) diff --git a/frontend/src/features/accounts/AccountsPage.test.tsx b/frontend/src/features/accounts/AccountsPage.test.tsx index 21d67be..0d722bf 100644 --- a/frontend/src/features/accounts/AccountsPage.test.tsx +++ b/frontend/src/features/accounts/AccountsPage.test.tsx @@ -13,6 +13,7 @@ import { openExternalUrl, writeDesktopClipboardText, } from "../../lib/desktop-shell"; +import { setAPIBase } from "../../lib/paths"; import { AccountsPage } from "./AccountsPage"; vi.mock("../../lib/desktop-shell", () => ({ @@ -35,9 +36,17 @@ function renderAccountsPage() { describe("AccountsPage", () => { beforeEach(() => { vi.clearAllMocks(); + setAPIBase("/ai-router/api"); + window.history.pushState({}, "", "/ai-router/webui/"); vi.mocked(isDesktopShell).mockReturnValue(false); }); + afterEach(() => { + window.history.pushState({}, "", "/"); + setAPIBase("/ai-router/api"); + vi.unstubAllGlobals(); + }); + it("supports official oauth branch and keeps local import as a separate branch", async () => { let completeCalls = 0; const fetchMock = vi.fn((input: RequestInfo | URL, init?: RequestInit) => { @@ -742,6 +751,216 @@ describe("AccountsPage", () => { expect(within(detailModal).queryByText("1 周剩余")).not.toBeInTheDocument(); }); + it("shows ai-gate upstream counts and toggles inline details", async () => { + const accountList = [ + { + id: 1, + provider_type: "openai-compatible", + account_name: "t1", + source_icon: "openai", + auth_mode: "api_key", + base_url: "http://127.0.0.1:6791/ai-gate/v1", + account_driver: "builtin_api_key", + usage_driver: "", + usage_config_json: "", + status: "active", + is_active: false, + priority: 1, + supports_responses: true, + balance: 1, + quota_remaining: 0, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 0, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + ]; + + const routeRequests: Array<{ account_id?: number | null; locked: boolean }> = []; + const lockRequests: Array<{ locked: boolean }> = []; + const fetchMock = vi.fn((input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + if ( + url === "/ai-router/api/accounts" && + (!init?.method || init.method === "GET") + ) { + return Promise.resolve( + new Response(JSON.stringify(accountList), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + } + if ( + url === "/ai-router/api/accounts/usage" && + (!init?.method || init.method === "GET") + ) { + return Promise.resolve( + new Response(JSON.stringify([]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + } + if ( + url === "/ai-router/api/accounts/1/upstreams" && + (!init?.method || init.method === "GET") + ) { + return Promise.resolve( + new Response( + JSON.stringify({ + total_accounts: 3, + available_accounts: 2, + current_account_id: 11, + route_locked: false, + accounts: [ + { + id: 11, + provider_type: "openai-compatible", + account_name: "upstream-a", + source_icon: "openai", + auth_mode: "api_key", + base_url: "https://up-a.example/v1", + status: "active", + available: true, + current: true, + preferred: false, + account_locked: false, + supports_responses: true, + balance: 0, + quota_remaining: 100, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 0, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + { + id: 12, + provider_type: "openai-compatible", + account_name: "upstream-b", + source_icon: "ppchat", + auth_mode: "api_key", + base_url: "https://code.ppchat.vip/v1", + status: "active", + available: true, + current: false, + preferred: false, + account_locked: true, + supports_responses: true, + balance: 0, + quota_remaining: -2, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 0, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + ppchat_today_used_quota: 22, + ppchat_today_added_quota: 20, + ppchat_today_remaining_quota: -2, + }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ), + ); + } + if ( + url === "/ai-router/api/accounts/1/upstreams/route" && + init?.method === "PUT" + ) { + routeRequests.push(JSON.parse(String(init.body))); + return Promise.resolve( + new Response( + JSON.stringify({ + preferred_account_id: routeRequests.at(-1)?.account_id, + route_locked: routeRequests.at(-1)?.locked, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ), + ); + } + if ( + url === "/ai-router/api/accounts/1/upstreams/12/lock" && + init?.method === "PUT" + ) { + lockRequests.push(JSON.parse(String(init.body))); + return Promise.resolve( + new Response( + JSON.stringify({ + id: 12, + account_locked: lockRequests.at(-1)?.locked, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ), + ); + } + return Promise.resolve(new Response(null, { status: 404 })); + }); + vi.stubGlobal("fetch", fetchMock); + + renderAccountsPage(); + + const card = (await screen.findByText("t1")).closest(".account-card-item"); + expect(card).not.toBeNull(); + await waitFor(() => { + expect(within(card as HTMLElement).getByText("可用/全部")).toBeInTheDocument(); + expect(within(card as HTMLElement).getByText("2/3")).toBeInTheDocument(); + }); + expect( + within(card as HTMLElement).getByRole("img", { name: "AI Gate" }), + ).toBeInTheDocument(); + expect(within(card as HTMLElement).queryByText("余额")).not.toBeInTheDocument(); + + fireEvent.click(card as HTMLElement); + expect(await screen.findByText("upstream-a")).toBeInTheDocument(); + expect(screen.getByText("https://up-a.example/v1")).toBeInTheDocument(); + expect(screen.getByText("当前使用中")).toBeInTheDocument(); + const ppchatRow = screen + .getByText("upstream-b") + .closest(".account-remote-upstream-row") as HTMLElement; + expect(ppchatRow).not.toBeNull(); + expect(within(ppchatRow).getByText("1D")).toBeInTheDocument(); + expect(within(ppchatRow).getByText("0%")).toBeInTheDocument(); + expect(within(ppchatRow).getByText("已锁定")).toBeInTheDocument(); + + fireEvent.mouseEnter(ppchatRow); + fireEvent.click( + within(ppchatRow).getByRole("button", { name: "启用-upstream-b" }), + ); + await waitFor(() => { + expect(routeRequests).toContainEqual({ account_id: 12, locked: false }); + }); + fireEvent.click( + within(ppchatRow).getByRole("button", { name: "解除锁定-upstream-b" }), + ); + await waitFor(() => { + expect(lockRequests).toContainEqual({ locked: false }); + }); + expect(routeRequests).toEqual([{ account_id: 12, locked: false }]); + + fireEvent.click(card as HTMLElement); + await waitFor(() => { + expect(screen.queryByText("upstream-a")).not.toBeInTheDocument(); + }); + }); + it("confirms before sharing and only copies after explicit approval", async () => { const clipboardWriteText = vi.fn().mockResolvedValue(undefined); Object.defineProperty(window.navigator, "clipboard", { @@ -1935,6 +2154,94 @@ describe("AccountsPage", () => { }); }); + it("hides global activation state in server webui", async () => { + window.history.pushState({}, "", "/ai-gate/webui/"); + setAPIBase("/ai-gate/api"); + const accountList = [ + { + id: 1, + provider_type: "openai-compatible", + account_name: "server-active", + source_icon: "openai", + auth_mode: "api_key", + base_url: "https://ai.nodeseek.in", + status: "active", + is_active: true, + priority: 2, + balance: 226.31, + quota_remaining: 0, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 1, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + { + id: 2, + provider_type: "openai-compatible", + account_name: "server-standby", + source_icon: "openai", + auth_mode: "api_key", + base_url: "https://code.ppchat.vip/v1", + status: "active", + is_active: false, + priority: 1, + balance: 0, + quota_remaining: 0, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 0, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + ]; + + const fetchMock = vi.fn((input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + if ( + url === "/ai-gate/api/accounts" && + (!init?.method || init.method === "GET") + ) { + return Promise.resolve( + new Response(JSON.stringify(accountList), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + } + if ( + url === "/ai-gate/api/accounts/usage" && + (!init?.method || init.method === "GET") + ) { + return Promise.resolve( + new Response(JSON.stringify([]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + } + return Promise.resolve(new Response(null, { status: 404 })); + }); + vi.stubGlobal("fetch", fetchMock); + + const { container } = renderAccountsPage(); + + expect(await screen.findByText("server-active")).toBeInTheDocument(); + expect(screen.queryByText("当前使用中")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "设为激活-server-standby" })).not.toBeInTheDocument(); + expect(container.querySelector(".active-account-card")).toBeNull(); + }); + it("toggles account lock state and only shows lock marker for locked accounts", async () => { const accountList = [ { diff --git a/frontend/src/features/accounts/AccountsPage.tsx b/frontend/src/features/accounts/AccountsPage.tsx index f1cd186..c4ca007 100644 --- a/frontend/src/features/accounts/AccountsPage.tsx +++ b/frontend/src/features/accounts/AccountsPage.tsx @@ -65,6 +65,7 @@ import { createAccount, completeOfficialAuthDevice, deleteAccount, + getAccountUpstreams, startOfficialAuth, getLuaUsageScript, importCurrentCodexAuth, @@ -78,13 +79,19 @@ import { testAccount, testLuaUsage, updateAccount, + updateAccountUpstreamLock, + updateAccountUpstreamRoute, type AccountRecord, type AccountTestResult, type OfficialAuthSession, + type ServerUpstreamAccount, + type ServerUpstreams, } from "../../lib/api"; import { writeClipboardText } from "../../lib/clipboard"; import { openExternalUrl, refreshDesktopTrayState } from "../../lib/desktop-shell"; import type { AppLanguage, Translator } from "../../lib/i18n"; +import { isServerWebUI } from "../../lib/paths"; +import sourceAIGateIcon from "../../assets/aigate_1024_1024.png"; import sourceClaudeCodeIcon from "../../assets/providers/claude-code.png"; import sourceOpenAIIcon from "../../assets/providers/openai.png"; import sourcePPChatIcon from "../../assets/providers/ppchat.png"; @@ -116,14 +123,17 @@ const authModeTextMap: Record = { codex_local_import: "本地导入", }; -type SourceIcon = "openai" | "claude_code" | "ppchat"; +type SourceIcon = "openai" | "claude_code" | "ppchat" | "aigate"; const sourceIconMap: Record = { openai: { label: "OpenAI", icon: sourceOpenAIIcon }, claude_code: { label: "Claude Code", icon: sourceClaudeCodeIcon }, ppchat: { label: "PPChat", icon: sourcePPChatIcon }, + aigate: { label: "AI Gate", icon: sourceAIGateIcon }, }; +const selectableSourceIcons: SourceIcon[] = ["openai", "claude_code", "ppchat"]; + const TEST_MODEL_SUGGESTIONS = [ "gpt-5.5", "gpt-5.5-pro", @@ -815,6 +825,87 @@ function isPPChatAccount(record: AccountRecord): boolean { ); } +function isPPChatUpstream(account: ServerUpstreamAccount): boolean { + return ( + normalizeSourceIcon(account.source_icon) === "ppchat" || + /ppchat\.vip/i.test(account.base_url) + ); +} + +function isRemoteUpstreamStatusSelectable(account: ServerUpstreamAccount): boolean { + return account.status !== "disabled" && account.status !== "invalid"; +} + +function canManuallySelectRemoteUpstream(account: ServerUpstreamAccount): boolean { + return isRemoteUpstreamStatusSelectable(account); +} + +function isAIGateAccount(record: AccountRecord): boolean { + const raw = (record.base_url || "").trim(); + if (raw === "") { + return false; + } + try { + const parsed = new URL(raw); + return parsed.pathname.split("/").filter(Boolean).includes("ai-gate"); + } catch { + return /(^|\/)ai-gate(\/|$)/i.test(raw); + } +} + +function getDisplaySourceIcon(record: AccountRecord): SourceIcon { + return isAIGateAccount(record) ? "aigate" : normalizeSourceIcon(record.source_icon); +} + +function isInteractiveTarget( + target: EventTarget | null, + root?: Element | null, +): boolean { + if (!(target instanceof Element)) { + return false; + } + const interactive = target.closest( + "button,a,input,textarea,select,[role='button'],[data-no-row-toggle='true']", + ); + return Boolean(interactive && interactive !== root); +} + +function formatRemoteUpstreamUsage( + account: ServerUpstreamAccount, + language: AppLanguage, +): string { + const display = account.usage_display?.summary; + if (display?.label || display?.value) { + return `${display.label ?? ""} ${display.value ?? ""}`.trim(); + } + if (account.balance > 0) { + return `${language === "en-US" ? "Balance" : "余额"} ${formatUsageAmount(language, account.balance)}`; + } + if (account.quota_remaining > 0) { + return `${language === "en-US" ? "Quota" : "额度"} ${formatUsageAmount(language, account.quota_remaining)}`; + } + return language === "en-US" ? "Usage unknown" : "用量未知"; +} + +function buildRemoteUpstreamUsageWindow( + account: ServerUpstreamAccount, + language: AppLanguage, +) { + const addedQuota = account.ppchat_today_added_quota ?? 0; + if (isPPChatUpstream(account) && addedQuota > 0) { + return { + label: "1D", + remainingPercent: clampPercent( + ((account.ppchat_today_remaining_quota ?? 0) / + Math.max(addedQuota, 1)) * + 100, + ), + resetLabel: formatTomorrowMidnight(language), + }; + } + return null; +} + function buildGenericUsageWindows( record: AccountRecord, language: AppLanguage, @@ -963,6 +1054,12 @@ type AccountCardRenderOptions = { style?: CSSProperties; }; +type RemoteUpstreamsState = { + loading: boolean; + data?: ServerUpstreams; + error?: string; +}; + type SortableAccountCardProps = { id: number; record: AccountRecord; @@ -1051,6 +1148,12 @@ export function AccountsPage({ const [visibleActionAccountID, setVisibleActionAccountID] = useState< number | null >(null); + const [remoteUpstreamsByAccountID, setRemoteUpstreamsByAccountID] = useState< + Record + >({}); + const [expandedRemoteAccountID, setExpandedRemoteAccountID] = useState< + number | null + >(null); const [thirdPartyForm] = Form.useForm(); const [officialForm] = Form.useForm(); @@ -1061,6 +1164,7 @@ export function AccountsPage({ const editBaseURL = Form.useWatch("base_url", editForm); const accountsRef = useRef([]); const dragSnapshotRef = useRef(null); + const serverWebUI = isServerWebUI(); const sensors = useSensors( useSensor(PointerSensor, { activationConstraint: { distance: 4 }, @@ -1200,6 +1304,7 @@ export function AccountsPage({ accountsRef.current = accountItems; setAccounts(accountItems); void refreshUsage(); + void refreshRemoteUpstreams(accountItems); } async function refreshUsage() { @@ -1227,6 +1332,56 @@ export function AccountsPage({ } } + async function refreshRemoteUpstreams(accountItems: AccountRecord[]) { + const aiGateAccounts = accountItems.filter(isAIGateAccount); + const aiGateIDs = new Set(aiGateAccounts.map((item) => item.id)); + setRemoteUpstreamsByAccountID((previous) => { + const next: Record = {}; + for (const account of aiGateAccounts) { + next[account.id] = { + loading: true, + data: previous[account.id]?.data, + error: undefined, + }; + } + return next; + }); + setExpandedRemoteAccountID((current) => + current !== null && aiGateIDs.has(current) ? current : null, + ); + await Promise.all( + aiGateAccounts.map(async (account) => { + try { + const data = await getAccountUpstreams(account.id); + setRemoteUpstreamsByAccountID((previous) => { + if (!accountsRef.current.some((item) => item.id === account.id && isAIGateAccount(item))) { + return previous; + } + return { + ...previous, + [account.id]: { loading: false, data }, + }; + }); + } catch (error) { + const message = + error instanceof Error ? error.message : t("无法读取上游状态"); + setRemoteUpstreamsByAccountID((previous) => { + if (!accountsRef.current.some((item) => item.id === account.id && isAIGateAccount(item))) { + return previous; + } + return { + ...previous, + [account.id]: { loading: false, error: message }, + }; + }); + setExpandedRemoteAccountID((current) => + current === account.id ? null : current, + ); + } + }), + ); + } + async function handleCreateThirdParty(values: { account_name: string; base_url: string; @@ -1696,6 +1851,112 @@ export function AccountsPage({ } } + async function handleUpdateRemoteUpstreamRoute( + account: AccountRecord, + upstream: ServerUpstreamAccount, + locked: boolean, + ) { + if (!canManuallySelectRemoteUpstream(upstream)) { + return; + } + try { + const response = await updateAccountUpstreamRoute(account.id, { + account_id: upstream.id, + locked, + }); + const preferredAccountID = response.preferred_account_id ?? upstream.id; + setRemoteUpstreamsByAccountID((previous) => { + const state = previous[account.id]; + if (!state?.data) { + return previous; + } + return { + ...previous, + [account.id]: { + ...state, + data: { + ...state.data, + current_account_id: preferredAccountID, + preferred_account_id: preferredAccountID, + route_locked: response.route_locked, + accounts: state.data.accounts.map((item) => ({ + ...item, + current: item.id === preferredAccountID, + preferred: item.id === preferredAccountID, + })), + }, + }, + }; + }); + void messageApi.success( + language === "en-US" + ? locked + ? `Locked upstream ${upstream.account_name}` + : `Enabled upstream ${upstream.account_name}` + : locked + ? `已锁定上游账户 ${upstream.account_name}` + : `已启用上游账户 ${upstream.account_name}`, + ); + } catch (error) { + void messageApi.error( + error instanceof Error ? t(error.message) : t("更新上游路由失败"), + ); + } + } + + async function handleUpdateRemoteUpstreamLock( + account: AccountRecord, + upstream: ServerUpstreamAccount, + locked: boolean, + ) { + try { + const response = await updateAccountUpstreamLock(account.id, upstream.id, locked); + setRemoteUpstreamsByAccountID((previous) => { + const state = previous[account.id]; + if (!state?.data) { + return previous; + } + return { + ...previous, + [account.id]: { + ...state, + data: { + ...state.data, + available_accounts: state.data.accounts.reduce((count, item) => { + if (item.id === upstream.id) { + return count + (locked ? 0 : isRemoteUpstreamStatusSelectable(item) ? 1 : 0); + } + return count + (item.available ? 1 : 0); + }, 0), + accounts: state.data.accounts.map((item) => + item.id === upstream.id + ? { + ...item, + account_locked: response.account_locked, + available: !response.account_locked && isRemoteUpstreamStatusSelectable(item), + } + : item, + ), + }, + }, + }; + }); + void messageApi.success( + language === "en-US" + ? locked + ? `Locked upstream account ${upstream.account_name}` + : `Unlocked upstream account ${upstream.account_name}` + : locked + ? `已锁定上游账户 ${upstream.account_name}` + : `已解除锁定上游账户 ${upstream.account_name}`, + ); + } catch (error) { + void messageApi.error( + error instanceof Error ? t(error.message) : t("更新上游锁定状态失败"), + ); + } + } + async function handleShareAccount(record: AccountRecord) { await modal.confirm({ title: t("分享账户"), @@ -1855,8 +2116,12 @@ export function AccountsPage({ options: AccountCardRenderOptions = {}, ) { const actionsVisible = visibleActionAccountID === record.id; - const sourceIcon = sourceIconMap[normalizeSourceIcon(record.source_icon)]; + const sourceIcon = sourceIconMap[getDisplaySourceIcon(record)]; const usageHealthState = getUsageHealthState(record); + const aiGateAccount = isAIGateAccount(record); + const remoteState = remoteUpstreamsByAccountID[record.id]; + const remoteUpstreams = remoteState?.data; + const remoteExpanded = expandedRemoteAccountID === record.id && Boolean(remoteUpstreams); const usageWindows = record.usage_display?.summary ? [] : isOfficialAccount(record) @@ -1887,13 +2152,24 @@ export function AccountsPage({ : buildGenericUsageWindows(record, language); const usageAmount = usageWindows.length === 0 ? buildUsageAmount(record, language) : null; + const canToggleRemote = aiGateAccount && Boolean(remoteUpstreams) && !options.actionsDisabled; + const toggleRemote = () => { + if (!canToggleRemote) { + return; + } + setExpandedRemoteAccountID((current) => (current === record.id ? null : record.id)); + }; return (
handleCardFocus(record.id)} onBlur={ options.actionsDisabled @@ -1908,6 +2184,28 @@ export function AccountsPage({ ? undefined : (event) => handleCardMouseLeave(record.id, event) } + onClick={ + canToggleRemote + ? (event) => { + if (!isInteractiveTarget(event.target, event.currentTarget)) { + toggleRemote(); + } + } + : undefined + } + onKeyDown={ + canToggleRemote + ? (event) => { + if (isInteractiveTarget(event.target, event.currentTarget)) { + return; + } + if (event.key === "Enter" || event.key === " ") { + event.preventDefault(); + toggleRemote(); + } + } + : undefined + } >
@@ -1928,6 +2226,7 @@ export function AccountsPage({
) : null} - {record.is_active ? ( + {record.is_active && !serverWebUI ? ( {t("当前使用中")} ) : null}
@@ -1973,57 +2272,74 @@ export function AccountsPage({
-
- {usageAmount ? ( -
- - {usageAmount.label} - - - {usageAmount.value} + {aiGateAccount ? ( + document.body} + > +
+ {t("可用/全部")} + + {remoteUpstreams ? `${remoteUpstreams.available_accounts}/${remoteUpstreams.total_accounts}` : "--"}
- ) : null} - {usageWindows.map((item) => ( -
-
- - {item.label} + + ) : ( +
+ {usageAmount ? ( +
+ + {usageAmount.label} - - {item.resetLabel || ""} + + {usageAmount.value} - {`${Math.round(item.remainingPercent)}%`}
-
-
+ ) : null} + {usageWindows.map((item) => ( +
+
+ + {item.label} + + + {item.resetLabel || ""} + + {`${Math.round(item.remainingPercent)}%`} +
+
+ className="account-usage-mini-track" + aria-label={`${record.account_name}-${item.label}`} + > +
+
-
- ))} -
+ ))} +
+ )}
- + {!serverWebUI ? ( + + ) : null}
+ {remoteExpanded && remoteUpstreams ? ( +
+
+ {remoteUpstreams.accounts.map((account) => { + const remoteUsageWindow = buildRemoteUpstreamUsageWindow( + account, + language, + ); + return ( +
+
+
+ {account.account_name} + {account.current ? {t("当前使用中")} : null} + {account.account_locked ? {t("已锁定")} : null} + {!account.available && !account.account_locked ? {t("不可用")} : null} +
+ + {account.base_url || t("OpenAI 官方")} + +
+
+ {remoteUsageWindow ? ( +
+
+ + {remoteUsageWindow.label} + + + {remoteUsageWindow.resetLabel} + + {`${Math.round(remoteUsageWindow.remainingPercent)}%`} +
+
+
+
+
+
+
+ ) : ( + {formatRemoteUpstreamUsage(account, language)} + )} + {t(statusTextMap[account.status] ?? account.status)} +
+
+
+
+ ); + })} + {remoteUpstreams.accounts.length === 0 ? ( + {t("暂无上游账号")} + ) : null} +
+
+ ) : null}
); } @@ -2193,7 +2588,7 @@ export function AccountsPage({ { sourceIconMap[ - normalizeSourceIcon(detailAccount.source_icon) + getDisplaySourceIcon(detailAccount) ].label } @@ -2469,7 +2864,7 @@ export function AccountsPage({ rules={[{ required: true, message: t("请选择来源图标") }]} > -
); } + +function buildServerUserImportPayload(issued: CreatedServerUser): string { + return JSON.stringify( + { + kind: "aigate-account-share", + schema_version: 1, + exported_at: new Date().toISOString(), + account: { + provider_type: "openai-compatible", + account_name: issued.user.name || issued.user.username || "AI Gate Server", + source_icon: "", + auth_mode: "api_key", + base_url: serverGatewayBaseURL(), + credential_ref: issued.token, + account_driver: "builtin_api_key", + usage_driver: "", + usage_config_json: "", + supports_responses: true, + skip_tls_verify: false, + }, + }, + null, + 2, + ); +} + +function serverGatewayBaseURL(): string { + if (typeof window === "undefined") { + return "/ai-gate/v1"; + } + const pathname = window.location.pathname || ""; + const webuiIndex = pathname.indexOf("/webui"); + const routePrefix = webuiIndex > 0 ? pathname.slice(0, webuiIndex).replace(/\/+$/, "") : "/ai-gate"; + const prefix = routePrefix || "/ai-gate"; + return `${window.location.origin}${prefix}/v1`; +} diff --git a/frontend/src/features/server-users/UserPoolPage.test.tsx b/frontend/src/features/server-users/UserPoolPage.test.tsx index c4595f7..65ab0ad 100644 --- a/frontend/src/features/server-users/UserPoolPage.test.tsx +++ b/frontend/src/features/server-users/UserPoolPage.test.tsx @@ -1,15 +1,16 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { UserPoolPage } from "./UserPoolPage"; -import { getServerMe, listMyServerAccounts, updateMyServerAccountState } from "../../lib/api"; +import { getServerMe, getServerUpstreams, updateServerRoute, updateServerUpstreamLock } from "../../lib/api"; vi.mock("../../lib/api", async () => { const actual = await vi.importActual("../../lib/api"); return { ...actual, getServerMe: vi.fn(), - listMyServerAccounts: vi.fn(), - updateMyServerAccountState: vi.fn(), + getServerUpstreams: vi.fn(), + updateServerRoute: vi.fn(), + updateServerUpstreamLock: vi.fn(), }; }); @@ -21,30 +22,141 @@ describe("UserPoolPage", () => { request_count: 3, total_tokens: 120, }); - vi.mocked(listMyServerAccounts).mockResolvedValue([ - { - user_id: 1, - account_id: 10, - account_name: "pool-a", - provider_type: "openai-compatible", - status: "active", - position: 0, - is_active: false, - is_locked: false, - supports_responses: true, - }, - ]); - vi.mocked(updateMyServerAccountState).mockResolvedValue(undefined); + vi.mocked(getServerUpstreams).mockResolvedValue({ + total_accounts: 2, + available_accounts: 1, + current_account_id: 1, + route_locked: false, + accounts: [ + { + id: 1, + provider_type: "openai-compatible", + account_name: "账号 A", + source_icon: "openai", + auth_mode: "api_key", + base_url: "https://upstream-a.example/v1", + status: "active", + available: true, + current: true, + preferred: false, + account_locked: false, + supports_responses: true, + balance: 10, + quota_remaining: 1000, + rpm_remaining: 20, + tpm_remaining: 3000, + health_score: 0.9, + recent_error_rate: 0, + last_total_tokens: 120, + last_input_tokens: 40, + last_output_tokens: 80, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + { + id: 2, + provider_type: "openai-compatible", + account_name: "账号 B", + source_icon: "ppchat", + auth_mode: "api_key", + base_url: "https://upstream-b.example/v1", + status: "disabled", + available: false, + current: false, + preferred: false, + account_locked: false, + supports_responses: true, + balance: 0, + quota_remaining: 0, + rpm_remaining: 0, + tpm_remaining: 0, + health_score: 0, + recent_error_rate: 0, + last_total_tokens: 0, + last_input_tokens: 0, + last_output_tokens: 0, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + ], + }); + vi.mocked(updateServerRoute).mockResolvedValue({ + preferred_account_id: 1, + route_locked: true, + }); + vi.mocked(updateServerUpstreamLock).mockResolvedValue({ + id: 1, + account_locked: true, + }); }); - it("shows own usage and updates assigned account state", async () => { - render( value} />); + it("shows own usage and sanitized upstream route controls", async () => { + const { container } = render( value} />); + + expect(await screen.findByText(/请求数 3/)).toBeInTheDocument(); + expect(await screen.findByText("120")).toBeInTheDocument(); + expect(await screen.findByText("可用 1 / 2")).toBeInTheDocument(); + expect(container.querySelector(".user-upstreams-receipt")).toBeInTheDocument(); + expect(await screen.findByText("账号 A")).toBeInTheDocument(); + expect(screen.getByText("当前使用中")).toBeInTheDocument(); + expect(screen.getByText("https://upstream-a.example/v1")).toBeInTheDocument(); + expect(screen.queryByText(/sk-/)).not.toBeInTheDocument(); + expect(screen.queryByText("credential_ref")).not.toBeInTheDocument(); + + fireEvent.click(screen.getByRole("button", { name: "锁定-账号 A" })); + await waitFor(() => expect(updateServerUpstreamLock).toHaveBeenCalledWith(1, true)); + expect(updateServerRoute).not.toHaveBeenCalledWith({ account_id: 1, locked: true }); + await waitFor(() => expect(getServerUpstreams).toHaveBeenCalledTimes(2)); + await waitFor(() => expect(screen.getByRole("button", { name: "切换-账号 A" })).not.toBeDisabled()); - expect(await screen.findByText("pool-a")).toBeInTheDocument(); - expect(screen.getByText(/请求数 3/)).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "切换-账号 A" })); + await waitFor(() => expect(updateServerRoute).toHaveBeenCalledWith({ account_id: 1, locked: false })); + }); + + it("shows account lock state and keeps manual switch available", async () => { + vi.mocked(getServerUpstreams).mockResolvedValue({ + total_accounts: 1, + available_accounts: 0, + current_account_id: 1, + route_locked: false, + accounts: [ + { + id: 1, + provider_type: "openai-compatible", + account_name: "账号 A", + source_icon: "openai", + auth_mode: "api_key", + base_url: "https://upstream-a.example/v1", + status: "active", + available: false, + current: true, + preferred: false, + account_locked: true, + supports_responses: true, + balance: 10, + quota_remaining: 1000, + rpm_remaining: 20, + tpm_remaining: 3000, + health_score: 0.9, + recent_error_rate: 0, + last_total_tokens: 120, + last_input_tokens: 40, + last_output_tokens: 80, + model_context_window: 0, + primary_used_percent: 0, + secondary_used_percent: 0, + }, + ], + }); + + render( value} />); - fireEvent.click(screen.getByRole("switch", { name: "激活 pool-a" })); + expect(await screen.findByText("已锁定")).toBeInTheDocument(); + expect(await screen.findByRole("button", { name: "切换-账号 A" })).not.toBeDisabled(); - await waitFor(() => expect(updateMyServerAccountState).toHaveBeenCalledWith(10, { position: 0, is_active: true, is_locked: false })); + fireEvent.click(screen.getByRole("button", { name: "解锁-账号 A" })); + await waitFor(() => expect(updateServerUpstreamLock).toHaveBeenCalledWith(1, false)); }); }); diff --git a/frontend/src/features/server-users/UserPoolPage.tsx b/frontend/src/features/server-users/UserPoolPage.tsx index bfd494d..98d7e32 100644 --- a/frontend/src/features/server-users/UserPoolPage.tsx +++ b/frontend/src/features/server-users/UserPoolPage.tsx @@ -1,8 +1,8 @@ -import { ReloadOutlined } from "@ant-design/icons"; -import { Button, InputNumber, Space, Switch, Table, Tag, Typography } from "antd"; +import { CheckCircleOutlined, LockOutlined, ReloadOutlined, SwapOutlined, UnlockOutlined } from "@ant-design/icons"; +import { Button, Collapse, Space, Statistic, Tag, Tooltip, Typography } from "antd"; import { useEffect, useState } from "react"; -import { getServerMe, listMyServerAccounts, updateMyServerAccountState, type ServerAssignedAccount, type ServerMe } from "../../lib/api"; +import { getServerMe, getServerUpstreams, updateServerRoute, updateServerUpstreamLock, type ServerMe, type ServerUpstreamAccount, type ServerUpstreams } from "../../lib/api"; type UserPoolPageProps = { t: (value: string) => string; @@ -10,34 +10,50 @@ type UserPoolPageProps = { export function UserPoolPage({ t }: UserPoolPageProps) { const [me, setMe] = useState(null); - const [accounts, setAccounts] = useState([]); + const [upstreams, setUpstreams] = useState(null); const [loading, setLoading] = useState(false); + const [routeUpdatingID, setRouteUpdatingID] = useState(null); async function refresh() { setLoading(true); try { - const [nextMe, nextAccounts] = await Promise.all([getServerMe(), listMyServerAccounts()]); + const [nextMe, nextUpstreams] = await Promise.all([getServerMe(), getServerUpstreams()]); setMe(nextMe); - setAccounts(nextAccounts); + setUpstreams(nextUpstreams); } finally { setLoading(false); } } + async function refreshUpstreams() { + setUpstreams(await getServerUpstreams()); + } + + async function handleSwitch(account: ServerUpstreamAccount) { + setRouteUpdatingID(account.id); + try { + await updateServerRoute({ account_id: account.id, locked: false }); + await refreshUpstreams(); + } finally { + setRouteUpdatingID(null); + } + } + + async function handleToggleLock(account: ServerUpstreamAccount) { + const locked = !account.account_locked; + setRouteUpdatingID(account.id); + try { + await updateServerUpstreamLock(account.id, locked); + await refreshUpstreams(); + } finally { + setRouteUpdatingID(null); + } + } + useEffect(() => { void refresh(); }, []); - async function patchAccount(account: ServerAssignedAccount, patch: Partial>) { - const next = { - position: patch.position ?? account.position, - is_active: patch.is_active ?? account.is_active, - is_locked: patch.is_locked ?? account.is_locked, - }; - setAccounts((items) => items.map((item) => (item.account_id === account.account_id ? { ...item, ...next } : item))); - await updateMyServerAccountState(account.account_id, next); - } - return (
@@ -49,45 +65,129 @@ export function UserPoolPage({ t }: UserPoolPageProps) {
- {status}, - }, - { - title: t("顺序"), - dataIndex: "position", - width: 120, - render: (_: unknown, account: ServerAssignedAccount) => ( - void patchAccount(account, { position: Number(value ?? 0) })} /> - ), - }, - { - title: t("激活"), - dataIndex: "is_active", - render: (_: unknown, account: ServerAssignedAccount) => ( - void patchAccount(account, { is_active: checked })} /> - ), - }, - { - title: t("锁定"), - dataIndex: "is_locked", - render: (_: unknown, account: ServerAssignedAccount) => ( - void patchAccount(account, { is_locked: checked })} /> - ), - }, - ]} - locale={{ emptyText: t("未分配账户池") }} - /> - +
+ + +
+
+ + {t("上游账号")} + + {t("可用")} {upstreams?.available_accounts ?? 0} / {upstreams?.total_accounts ?? 0} + +
+ ), + children: ( +
+ {(upstreams?.accounts ?? []).map((account) => ( + + ))} + {upstreams && upstreams.accounts.length === 0 ? ( + {t("暂无上游账号")} + ) : null} +
+ ), + }, + ]} + /> + ); } + +type UpstreamRowProps = { + account: ServerUpstreamAccount; + updating: boolean; + t: (value: string) => string; + onSwitch: (account: ServerUpstreamAccount) => Promise; + onToggleLock: (account: ServerUpstreamAccount) => Promise; +}; + +function UpstreamRow({ account, updating, t, onSwitch, onToggleLock }: UpstreamRowProps) { + const selectable = canManuallySelectServerUpstream(account); + return ( +
+
+
+ {account.account_name} + {account.current ? ( + }> + {t("当前使用中")} + + ) : null} + {account.account_locked ? {t("已锁定")} : null} + {!account.available && !account.account_locked ? {t("不可用")} : null} +
+ + {account.base_url || t("OpenAI 官方")} + +
+
+ {usageSummary(account)} + {t("Tokens")} {formatCompactNumber(account.last_total_tokens)} +
+ + +
+ ); +} + +function canManuallySelectServerUpstream(account: ServerUpstreamAccount): boolean { + return account.status !== "disabled" && account.status !== "invalid"; +} + +function usageSummary(account: ServerUpstreamAccount): string { + const display = account.usage_display?.summary; + if (display?.label || display?.value) { + return `${display.label ?? ""} ${display.value ?? ""}`.trim(); + } + if (account.balance > 0) { + return `余额 ${formatCompactNumber(account.balance)}`; + } + if (account.quota_remaining > 0) { + return `额度 ${formatCompactNumber(account.quota_remaining)}`; + } + return "用量未知"; +} + +function formatCompactNumber(value: number): string { + if (!Number.isFinite(value)) { + return "0"; + } + return new Intl.NumberFormat(undefined, { maximumFractionDigits: value >= 100 ? 0 : 2 }).format(value); +} diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 4f8bc49..1cf6fc1 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -126,9 +126,10 @@ export type ServerUser = { status: string; created_at: string; last_used_at?: string; + preferred_account_id?: number; + route_locked?: boolean; request_count?: number; total_tokens?: number; - assigned_accounts?: number; }; export type CreatedServerUser = { @@ -142,36 +143,68 @@ export type ServerMe = { total_tokens: number; }; -export type ServerAssignedAccount = { - user_id: number; - account_id: number; - account_name: string; +export type ServerUpstreamAccount = { + id: number; provider_type: string; - source_icon?: string; - base_url?: string; + account_name: string; + source_icon?: "openai" | "claude_code" | "ppchat"; + auth_mode: string; + base_url: string; status: string; - position: number; - is_active: boolean; - is_locked: boolean; + available: boolean; + current: boolean; + preferred: boolean; + account_locked: boolean; supports_responses: boolean; - cooldown_until?: string; - cooldown_reason?: string; + cooldown_remaining_seconds?: number; + routing_cooldown_remaining_seconds?: number; + routing_cooldown_reason?: string; + balance: number; + quota_remaining: number; + rpm_remaining: number; + tpm_remaining: number; + health_score: number; + recent_error_rate: number; + last_total_tokens: number; + last_input_tokens: number; + last_output_tokens: number; + model_context_window: number; + primary_used_percent: number; + secondary_used_percent: number; + primary_resets_at?: string; + secondary_resets_at?: string; + checked_at?: string; + stale?: boolean; + last_error?: string; + ppchat_today_used_quota?: number; + ppchat_today_added_quota?: number; + ppchat_today_remaining_quota?: number; + usage_display?: AccountUsageDisplay; }; -export type ServerUserAccountAssignment = { - account_id: number; - account_name: string; - provider_type: string; - source_icon?: string; - base_url?: string; - status: string; - assigned: boolean; - position: number; - is_active: boolean; - is_locked: boolean; - supports_responses: boolean; - cooldown_until?: string; - cooldown_reason?: string; +export type ServerUpstreams = { + total_accounts: number; + available_accounts: number; + current_account_id?: number; + preferred_account_id?: number; + route_locked: boolean; + accounts: ServerUpstreamAccount[]; +}; + +export type ServerRoutePayload = { + account_id?: number | null; + locked: boolean; +}; + +export type ServerRouteResponse = { + user?: ServerUser; + preferred_account_id?: number | null; + route_locked: boolean; +}; + +export type ServerUpstreamLockResponse = { + id: number; + account_locked: boolean; }; export type UsageTrendPoint = { @@ -527,6 +560,48 @@ export async function listAccountUsage(): Promise { return response.json(); } +export async function getAccountUpstreams(id: number): Promise { + const response = await fetch(apiPath(`/accounts/${id}/upstreams`)); + if (!response.ok) { + const details = await response.text(); + throw new Error(details || "failed to load account upstreams"); + } + return response.json(); +} + +export async function updateAccountUpstreamRoute( + id: number, + payload: ServerRoutePayload, +): Promise { + const response = await fetch(apiPath(`/accounts/${id}/upstreams/route`), { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + if (!response.ok) { + const details = await response.text(); + throw new Error(details || "failed to update account upstream route"); + } + return response.json(); +} + +export async function updateAccountUpstreamLock( + id: number, + upstreamAccountID: number, + locked: boolean, +): Promise { + const response = await fetch(apiPath(`/accounts/${id}/upstreams/${upstreamAccountID}/lock`), { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ locked }), + }); + if (!response.ok) { + const details = await response.text(); + throw new Error(details || "failed to update account upstream lock"); + } + return response.json(); +} + export async function refreshAccountUsage(): Promise { const response = await fetch(apiPath("/accounts/usage/refresh"), { method: "POST", @@ -1183,6 +1258,17 @@ export async function disableServerUser(id: number): Promise { } } +export async function deleteServerUser(id: number): Promise { + const response = await fetch(apiPath(`/server-users/${id}`), { + method: "DELETE", + credentials: "include", + }); + if (!response.ok) { + const details = await response.text(); + throw new Error(details || "failed to delete server user"); + } +} + export async function rotateServerUserToken(id: number): Promise { const response = await fetch(apiPath(`/server-users/${id}/rotate-token`), { method: "POST", @@ -1194,57 +1280,52 @@ export async function rotateServerUserToken(id: number): Promise { - const response = await fetch(apiPath(`/server-users/${id}/accounts`), { credentials: "include" }); +export async function getServerMe(): Promise { + const response = await fetch(apiPath("/me"), { credentials: "include" }); if (!response.ok) { - throw new Error("failed to list server user accounts"); + throw new Error("failed to load current server user"); } return response.json(); } -export async function setServerUserAccounts(id: number, accountIDs: number[]): Promise { - const response = await fetch(apiPath(`/server-users/${id}/accounts`), { - method: "PUT", - headers: { "Content-Type": "application/json" }, - credentials: "include", - body: JSON.stringify({ account_ids: accountIDs }), - }); +export async function getServerUpstreams(): Promise { + const response = await fetch(apiPath("/me/upstreams"), { credentials: "include" }); if (!response.ok) { const details = await response.text(); - throw new Error(details || "failed to set server user accounts"); - } -} - -export async function getServerMe(): Promise { - const response = await fetch(apiPath("/me"), { credentials: "include" }); - if (!response.ok) { - throw new Error("failed to load current server user"); + throw new Error(details || "failed to load upstream accounts"); } return response.json(); } -export async function listMyServerAccounts(): Promise { - const response = await fetch(apiPath("/me/accounts"), { credentials: "include" }); +export async function updateServerRoute(payload: ServerRoutePayload): Promise { + const response = await fetch(apiPath("/me/route"), { + method: "PUT", + headers: { "Content-Type": "application/json" }, + credentials: "include", + body: JSON.stringify(payload), + }); if (!response.ok) { - throw new Error("failed to list assigned accounts"); + const details = await response.text(); + throw new Error(details || "failed to update upstream route"); } return response.json(); } -export async function updateMyServerAccountState( +export async function updateServerUpstreamLock( accountID: number, - payload: { position: number; is_active: boolean; is_locked: boolean }, -): Promise { - const response = await fetch(apiPath(`/me/accounts/${accountID}/state`), { + locked: boolean, +): Promise { + const response = await fetch(apiPath(`/me/upstreams/${accountID}/lock`), { method: "PUT", headers: { "Content-Type": "application/json" }, credentials: "include", - body: JSON.stringify(payload), + body: JSON.stringify({ locked }), }); if (!response.ok) { const details = await response.text(); - throw new Error(details || "failed to update assigned account state"); + throw new Error(details || "failed to update upstream lock"); } + return response.json(); } export async function enableProxy(): Promise { diff --git a/frontend/src/lib/i18n.ts b/frontend/src/lib/i18n.ts index 758d6e5..ba55aa0 100644 --- a/frontend/src/lib/i18n.ts +++ b/frontend/src/lib/i18n.ts @@ -108,6 +108,8 @@ const enUSMessages: Record = { "删除备份": "Delete backup", "确认删除此备份": "Delete backup", "确认删除": "Delete", + "禁用用户": "Disable user", + "删除用户": "Delete user", "暂无数据库备份": "No database backups", "创建于": "Created at", "恢复此备份": "Restore this backup", @@ -285,6 +287,8 @@ const enUSMessages: Record = { "账户已更新": "Account updated", "切换激活账户失败": "Failed to switch the active account", "已复制账户名称和 API 地址": "Copied account name and API address", + "复制 AI Gate 导入配置": "Copy AI Gate import config", + "AI Gate 导入配置已复制": "AI Gate import config copied", "复制失败，请检查系统剪贴板权限": "Copy failed. Check clipboard permissions.", "排序已更新到界面，但保存顺序失败，请稍后重试": "The order updated in the UI, but saving the new order failed. Try again later.", "拖拽排序": "Drag sort", diff --git a/frontend/src/styles.css b/frontend/src/styles.css index c6f0f15..acdeec1 100644 --- a/frontend/src/styles.css +++ b/frontend/src/styles.css @@ -69,6 +69,156 @@ width: min(720px, 100%); } +.user-upstreams-receipt { + position: relative; + overflow: hidden; + border: 1px solid #e1e4e8; + border-radius: 8px; + background: #f4f5f7; + box-shadow: 0 10px 24px rgba(17, 24, 39, 0.05); +} + +.user-upstreams-receipt::before { + position: absolute; + top: 0; + right: 16px; + left: 16px; + height: 1px; + background: repeating-linear-gradient(to right, #d5d9df 0 8px, transparent 8px 14px); + content: ""; +} + +.user-upstreams-collapse.ant-collapse { + border: 0; + background: transparent; +} + +.user-upstreams-collapse .ant-collapse-item { + border: 0; +} + +.user-upstreams-collapse .ant-collapse-header { + min-height: 54px; + align-items: center; + padding: 14px 18px; + color: var(--text-primary); +} + +.user-upstreams-collapse .ant-collapse-expand-icon { + color: var(--text-secondary); + transition: transform 180ms ease; +} + +.user-upstreams-collapse .ant-collapse-item-active .ant-collapse-expand-icon { + transform: translateY(1px); +} + +.user-upstreams-collapse .ant-collapse-content { + border-top: 1px dashed #d8dce2; + background: transparent; + transition: + height 220ms ease, + opacity 180ms ease; +} + +.user-upstreams-collapse .ant-collapse-content-box { + padding: 0 18px 14px; +} + +.user-upstreams-collapse .ant-collapse-content-box > .user-upstream-list { + transform: translateY(-4px); + opacity: 0; + transition: + transform 220ms ease, + opacity 180ms ease; +} + +.user-upstreams-collapse .ant-collapse-content-active .ant-collapse-content-box > .user-upstream-list { + transform: translateY(0); + opacity: 1; +} + +.user-upstreams-header { + width: 100%; + display: flex; + align-items: center; + justify-content: space-between; + gap: 16px; +} + +.user-upstreams-count { + color: var(--text-secondary); + font-size: 13px; +} + +.user-upstream-list { + display: grid; +} + +.user-upstream-row { + min-height: 62px; + display: grid; + grid-template-columns: minmax(220px, 1fr) minmax(160px, 240px) auto; + align-items: center; + gap: 12px; + padding: 12px 0; + border-bottom: 1px dashed #d8dce2; + background: transparent; +} + +.user-upstream-row:last-child { + border-bottom: 0; +} + +.user-upstream-row.is-current { + position: relative; +} + +.user-upstream-row.is-current::before { + position: absolute; + top: 14px; + bottom: 14px; + left: -10px; + width: 3px; + border-radius: 999px; + background: #67a876; + content: ""; +} + +.user-upstream-main { + min-width: 0; + display: grid; + gap: 4px; +} + +.user-upstream-title-row { + min-width: 0; + display: flex; + align-items: center; + gap: 8px; + flex-wrap: wrap; +} + +.user-upstream-base-url { + display: block; + max-width: 100%; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.user-upstream-usage { + display: grid; + gap: 2px; + color: var(--text-secondary); + font-size: 13px; + text-align: right; +} + +.user-upstream-actions { + justify-content: flex-end; +} + .app-theme-shell[data-theme-mode="dark"], body[data-theme-mode="dark"], html[data-theme-mode="dark"] { @@ -2234,6 +2384,15 @@ html[data-theme-mode="dark"] .stats-chart-shell { opacity: 0.42; } +.account-card-item.has-remote-upstreams { + cursor: pointer; +} + +.account-card-item.has-remote-upstreams:focus-visible { + outline: 2px solid rgba(62, 91, 232, 0.32); + outline-offset: 3px; +} + .account-card-shell { display: flex; align-items: center; @@ -2469,6 +2628,153 @@ html[data-theme-mode="dark"] .stats-chart-shell { background: linear-gradient(90deg, #f87171 0%, #dc2626 100%); } +.account-upstream-mini { + display: grid; + justify-items: end; + align-content: center; + min-height: 52px; + gap: 5px; + opacity: 1; + transition: opacity 0.16s ease; +} + +.account-upstream-mini.is-loading { + opacity: 0.72; +} + +.account-upstream-mini-label { + color: var(--text-secondary); + font-size: 11px; + font-weight: 700; + line-height: 1; +} + +.account-upstream-mini-value { + color: var(--text-primary); + font-size: 18px; + font-weight: 700; + line-height: 1; + font-variant-numeric: tabular-nums; +} + +.account-remote-upstreams-panel { + margin: -2px 18px 16px 58px; + overflow: hidden; + border: 1px solid #e1e4e8; + border-radius: 8px; + background: #f4f5f7; + animation: accountRemoteUpstreamsExpand 180ms ease-out; +} + +.account-remote-upstreams-list { + display: grid; + padding: 2px 16px; +} + +.account-remote-upstream-row { + position: relative; + min-height: 56px; + display: grid; + grid-template-columns: minmax(180px, 1fr) minmax(130px, 210px); + align-items: center; + gap: 12px; + padding: 11px 0; + border-bottom: 1px dashed #d8dce2; +} + +.account-remote-upstream-row:last-child { + border-bottom: 0; +} + +.account-remote-upstream-row.is-current { + position: relative; +} + +.account-remote-upstream-row.is-current::before { + position: absolute; + top: 13px; + bottom: 13px; + left: -10px; + width: 3px; + border-radius: 999px; + background: #67a876; + content: ""; +} + +.account-remote-upstream-main { + min-width: 0; + display: grid; + gap: 4px; +} + +.account-remote-upstream-title { + min-width: 0; + display: flex; + align-items: center; + gap: 8px; + flex-wrap: wrap; +} + +.account-remote-upstream-url { + display: block; + max-width: 100%; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.account-remote-upstream-meta { + display: grid; + gap: 3px; + color: var(--text-secondary); + font-size: 13px; + text-align: right; + transition: opacity 0.16s ease; +} + +.account-remote-upstream-usage-mini { + display: grid; + gap: 4px; +} + +.account-remote-upstream-actions { + position: absolute; + top: 50%; + right: 0; + display: inline-flex; + align-items: center; + gap: 4px; + transform: translateY(-50%); + opacity: 0; + pointer-events: none; + transition: opacity 0.16s ease; +} + +.account-remote-upstream-row:hover .account-remote-upstream-meta, +.account-remote-upstream-row:focus-within .account-remote-upstream-meta { + opacity: 0; + pointer-events: none; +} + +.account-remote-upstream-row:hover .account-remote-upstream-actions, +.account-remote-upstream-row:focus-within .account-remote-upstream-actions { + opacity: 1; + pointer-events: auto; +} + +@keyframes accountRemoteUpstreamsExpand { + from { + max-height: 0; + opacity: 0; + transform: translateY(-4px); + } + to { + max-height: 420px; + opacity: 1; + transform: translateY(0); + } +} + .account-actions { position: absolute; top: 50%; @@ -2487,6 +2793,11 @@ html[data-theme-mode="dark"] .stats-chart-shell { pointer-events: none; } +.account-card-item[data-actions-visible="true"] .account-upstream-mini { + opacity: 0; + pointer-events: none; +} + .account-card-item[data-actions-visible="true"] .account-actions { opacity: 1; pointer-events: auto; @@ -3081,6 +3392,19 @@ body[data-theme-mode="dark"] .ant-modal-confirm .ant-modal-confirm-body > .antic flex-direction: column; } + .user-upstream-row { + grid-template-columns: 1fr; + align-items: stretch; + } + + .user-upstream-usage { + text-align: left; + } + + .user-upstream-actions { + justify-content: flex-start; + } + .ppchat-metrics-grid { grid-template-columns: 1fr; } From 5c8b13874d9908e9dbf786edd72302597e15d1a4 Mon Sep 17 00:00:00 2001 From: GcsSloop Date: Wed, 24 Jun 2026 02:08:31 +0800 Subject: [PATCH 4/4] =?UTF-8?q?chore:=20=E5=90=8C=E6=AD=A5=E5=8F=91?= =?UTF-8?q?=E5=B8=83=E7=89=88=E6=9C=AC=201.6.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- desktop/package-lock.json | 4 ++-- desktop/package.json | 2 +- desktop/src-tauri/Cargo.lock | 2 +- desktop/src-tauri/Cargo.toml | 2 +- desktop/src-tauri/tauri.conf.json | 2 +- frontend/package-lock.json | 4 ++-- frontend/package.json | 2 +- 7 files changed, 9 insertions(+), 9 deletions(-) diff --git a/desktop/package-lock.json b/desktop/package-lock.json index f7060e1..603b741 100644 --- a/desktop/package-lock.json +++ b/desktop/package-lock.json @@ -1,12 +1,12 @@ { "name": "aigate-desktop", - "version": "1.5.14", + "version": "1.6.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "aigate-desktop", - "version": "1.5.14", + "version": "1.6.0", "devDependencies": { "@tauri-apps/cli": "^2.8.0" } diff --git a/desktop/package.json b/desktop/package.json index 9de755e..1c61553 100644 --- a/desktop/package.json +++ b/desktop/package.json @@ -1,7 +1,7 @@ { "name": "aigate-desktop", "private": true, - "version": "1.5.14", + "version": "1.6.0", "type": "module", "scripts": { "dev": "tauri dev", diff --git a/desktop/src-tauri/Cargo.lock b/desktop/src-tauri/Cargo.lock index 038044c..a8f3e7a 100644 --- a/desktop/src-tauri/Cargo.lock +++ b/desktop/src-tauri/Cargo.lock @@ -19,7 +19,7 @@ dependencies = [ [[package]] name = "aigate-desktop" -version = "1.5.14" +version = "1.6.0" dependencies = [ "arboard", "base64 0.22.1", diff --git a/desktop/src-tauri/Cargo.toml b/desktop/src-tauri/Cargo.toml index e55e0f1..ae72a3b 100644 --- a/desktop/src-tauri/Cargo.toml +++ b/desktop/src-tauri/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "aigate-desktop" -version = "1.5.14" +version = "1.6.0" edition = "2021" [build-dependencies] diff --git a/desktop/src-tauri/tauri.conf.json b/desktop/src-tauri/tauri.conf.json index d590ffc..a40c5c1 100644 --- a/desktop/src-tauri/tauri.conf.json +++ b/desktop/src-tauri/tauri.conf.json @@ -1,7 +1,7 @@ { "$schema": "https://schema.tauri.app/config/2", "productName": "AI Gate", - "version": "1.5.14", + "version": "1.6.0", "identifier": "com.aigate.desktop", "build": { "frontendDist": "../../frontend/dist", diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 6503cd5..1c45c6e 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -1,12 +1,12 @@ { "name": "aigate-frontend", - "version": "1.5.14", + "version": "1.6.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "aigate-frontend", - "version": "1.5.14", + "version": "1.6.0", "dependencies": { "@ant-design/icons": "^6.1.0", "@dnd-kit/core": "^6.3.1", diff --git a/frontend/package.json b/frontend/package.json index c87a7b3..e92ade0 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -1,7 +1,7 @@ { "name": "aigate-frontend", "private": true, - "version": "1.5.14", + "version": "1.6.0", "type": "module", "scripts": { "dev": "vite",