From f329b9d10c21bd461ed0bf58a06c6a9004509adb Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 27 Jul 2026 03:46:55 +0800 Subject: [PATCH] feat(store): add PostgreSQL-backed cooldown state persistence - Introduced `PostgresCooldownStateStore` for saving and loading cooldown states in a PostgreSQL database. - Supported runtime cooldown state management with efficient upsert and deletion. - Added configurable integration via `PostgresStoreConfig` and updated schema management to include cooldown-related tables. - Included tests to validate persistence, normalization, and concurrent handling of cooldown state records. - Enabled overriding cooldown state stores via service builder and SDK configuration. Closes: #4254 --- internal/store/postgres_cooldown_store.go | 193 ++++++++++++ .../store/postgres_cooldown_store_test.go | 297 ++++++++++++++++++ internal/store/postgresstore.go | 49 ++- sdk/cliproxy/auth/cooldown_state.go | 5 + sdk/cliproxy/builder.go | 16 + sdk/cliproxy/service.go | 3 + sdk/cliproxy/service_auth.go | 3 + sdk/cliproxy/service_cooldown_store_test.go | 68 ++++ 8 files changed, 620 insertions(+), 14 deletions(-) create mode 100644 internal/store/postgres_cooldown_store.go create mode 100644 internal/store/postgres_cooldown_store_test.go create mode 100644 sdk/cliproxy/service_cooldown_store_test.go diff --git a/internal/store/postgres_cooldown_store.go b/internal/store/postgres_cooldown_store.go new file mode 100644 index 00000000..11cf5ef4 --- /dev/null +++ b/internal/store/postgres_cooldown_store.go @@ -0,0 +1,193 @@ +package store + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +var _ cliproxyauth.CooldownStateStoreProvider = (*PostgresStore)(nil) +var _ cliproxyauth.CooldownStateStore = (*postgresCooldownStateStore)(nil) + +type postgresCooldownStateKey struct { + authID string + model string +} + +type postgresCooldownStateRecord struct { + key postgresCooldownStateKey + content []byte + updatedAt time.Time +} + +type postgresCooldownStateVersion struct { + updatedAt time.Time +} + +type postgresCooldownStateStore struct { + store *PostgresStore + mu sync.Mutex + previous map[postgresCooldownStateKey]postgresCooldownStateVersion +} + +// CooldownStateStore returns the PostgreSQL-backed runtime cooldown store. +func (s *PostgresStore) CooldownStateStore() cliproxyauth.CooldownStateStore { + if s == nil { + return nil + } + return s.cooldownStore +} + +func (s *postgresCooldownStateStore) Load(ctx context.Context) (records []cliproxyauth.CooldownStateRecord, err error) { + if s == nil || s.store == nil || s.store.db == nil { + return nil, fmt.Errorf("postgres cooldown store: not initialized") + } + if ctx == nil { + ctx = context.Background() + } + + s.mu.Lock() + defer s.mu.Unlock() + + table := s.store.fullTableName(s.store.cfg.CooldownTable) + query := fmt.Sprintf("SELECT content, updated_at FROM %s WHERE deleted = FALSE", table) + rows, errQuery := s.store.db.QueryContext(ctx, query) + if errQuery != nil { + return nil, fmt.Errorf("postgres cooldown store: load state: %w", errQuery) + } + defer func() { + if errClose := rows.Close(); errClose != nil { + err = errors.Join(err, fmt.Errorf("postgres cooldown store: close state rows: %w", errClose)) + } + }() + + records = make([]cliproxyauth.CooldownStateRecord, 0) + previous := make(map[postgresCooldownStateKey]postgresCooldownStateVersion) + for rows.Next() { + var content []byte + var updatedAt time.Time + if errScan := rows.Scan(&content, &updatedAt); errScan != nil { + return nil, fmt.Errorf("postgres cooldown store: scan state: %w", errScan) + } + var record cliproxyauth.CooldownStateRecord + if errUnmarshal := json.Unmarshal(content, &record); errUnmarshal != nil { + return nil, fmt.Errorf("postgres cooldown store: decode state: %w", errUnmarshal) + } + key := cooldownStateKey(record) + if key.authID == "" { + return nil, fmt.Errorf("postgres cooldown store: decoded state has empty auth ID") + } + records = append(records, record) + previous[key] = postgresCooldownStateVersion{updatedAt: updatedAt} + } + if errRows := rows.Err(); errRows != nil { + return nil, fmt.Errorf("postgres cooldown store: iterate state: %w", errRows) + } + s.previous = previous + return records, nil +} + +func (s *postgresCooldownStateStore) Save(ctx context.Context, records []cliproxyauth.CooldownStateRecord) error { + if s == nil || s.store == nil || s.store.db == nil { + return fmt.Errorf("postgres cooldown store: not initialized") + } + if ctx == nil { + ctx = context.Background() + } + + now := normalizePostgresCooldownTime(time.Now(), time.Time{}) + current := make(map[postgresCooldownStateKey]postgresCooldownStateVersion, len(records)) + encoded := make([]postgresCooldownStateRecord, 0, len(records)) + for i := range records { + record := records[i] + key := cooldownStateKey(record) + if key.authID == "" { + return fmt.Errorf("postgres cooldown store: state has empty auth ID") + } + record.UpdatedAt = normalizePostgresCooldownTime(record.UpdatedAt, now) + content, errMarshal := json.Marshal(record) + if errMarshal != nil { + return fmt.Errorf("postgres cooldown store: encode state for %q: %w", key.authID, errMarshal) + } + current[key] = postgresCooldownStateVersion{updatedAt: record.UpdatedAt} + encoded = append(encoded, postgresCooldownStateRecord{key: key, content: content, updatedAt: record.UpdatedAt}) + } + + s.mu.Lock() + defer s.mu.Unlock() + + tx, errBegin := s.store.db.BeginTx(ctx, nil) + if errBegin != nil { + return fmt.Errorf("postgres cooldown store: begin save: %w", errBegin) + } + table := s.store.fullTableName(s.store.cfg.CooldownTable) + upsertQuery := fmt.Sprintf(` + INSERT INTO %s AS target (auth_id, model, content, deleted, created_at, updated_at) + VALUES ($1, $2, $3, FALSE, NOW(), $4) + ON CONFLICT (auth_id, model) DO UPDATE SET + content = EXCLUDED.content, + deleted = FALSE, + updated_at = EXCLUDED.updated_at + WHERE target.updated_at <= EXCLUDED.updated_at + `, table) + for i := range encoded { + record := encoded[i] + if _, errExec := tx.ExecContext(ctx, upsertQuery, record.key.authID, record.key.model, record.content, record.updatedAt); errExec != nil { + return rollbackPostgresCooldownTransaction(tx, fmt.Errorf("postgres cooldown store: save state for %q: %w", record.key.authID, errExec)) + } + } + deleteQuery := fmt.Sprintf(` + INSERT INTO %s AS target (auth_id, model, content, deleted, created_at, updated_at) + VALUES ($1, $2, $3, TRUE, NOW(), $4) + ON CONFLICT (auth_id, model) DO UPDATE SET + content = EXCLUDED.content, + deleted = TRUE, + updated_at = EXCLUDED.updated_at + WHERE NOT target.deleted AND target.updated_at <= $5 + `, table) + for key, previous := range s.previous { + if _, ok := current[key]; ok { + continue + } + deletedAt := now + if !deletedAt.After(previous.updatedAt) { + deletedAt = previous.updatedAt.Add(time.Microsecond) + } + if _, errExec := tx.ExecContext(ctx, deleteQuery, key.authID, key.model, []byte(`{}`), deletedAt, previous.updatedAt); errExec != nil { + return rollbackPostgresCooldownTransaction(tx, fmt.Errorf("postgres cooldown store: clear state for %q: %w", key.authID, errExec)) + } + } + if errCommit := tx.Commit(); errCommit != nil { + return fmt.Errorf("postgres cooldown store: commit save: %w", errCommit) + } + s.previous = current + return nil +} + +func cooldownStateKey(record cliproxyauth.CooldownStateRecord) postgresCooldownStateKey { + return postgresCooldownStateKey{ + authID: strings.TrimSpace(record.AuthID), + model: strings.TrimSpace(record.Model), + } +} + +func normalizePostgresCooldownTime(value, fallback time.Time) time.Time { + if value.IsZero() { + value = fallback + } + return value.UTC().Truncate(time.Microsecond) +} + +func rollbackPostgresCooldownTransaction(tx *sql.Tx, operationErr error) error { + if errRollback := tx.Rollback(); errRollback != nil && !errors.Is(errRollback, sql.ErrTxDone) { + return errors.Join(operationErr, fmt.Errorf("postgres cooldown store: rollback save: %w", errRollback)) + } + return operationErr +} diff --git a/internal/store/postgres_cooldown_store_test.go b/internal/store/postgres_cooldown_store_test.go new file mode 100644 index 00000000..16c067aa --- /dev/null +++ b/internal/store/postgres_cooldown_store_test.go @@ -0,0 +1,297 @@ +package store + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "io" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +var cooldownTestDriverID atomic.Uint64 + +type cooldownTestDriver struct { + state *cooldownTestState +} + +type cooldownTestState struct { + mu sync.Mutex + rows map[string]cooldownTestRow + queries []string +} + +type cooldownTestRow struct { + content []byte + deleted bool + updatedAt time.Time +} + +type cooldownTestConn struct { + state *cooldownTestState +} + +type cooldownTestTx struct{} + +type cooldownTestRows struct { + rows []cooldownTestRow + index int +} + +func (d *cooldownTestDriver) Open(string) (driver.Conn, error) { + return &cooldownTestConn{state: d.state}, nil +} + +func (c *cooldownTestConn) Prepare(string) (driver.Stmt, error) { + return nil, errors.New("prepare is not supported") +} + +func (c *cooldownTestConn) Close() error { + return nil +} + +func (c *cooldownTestConn) Begin() (driver.Tx, error) { + return &cooldownTestTx{}, nil +} + +func (c *cooldownTestConn) ExecContext(_ context.Context, query string, args []driver.NamedValue) (driver.Result, error) { + c.state.mu.Lock() + defer c.state.mu.Unlock() + c.state.queries = append(c.state.queries, query) + if !strings.Contains(query, "INSERT INTO") || (len(args) != 4 && len(args) != 5) { + return driver.RowsAffected(1), nil + } + authID, okAuthID := args[0].Value.(string) + model, okModel := args[1].Value.(string) + content, okContent := args[2].Value.([]byte) + updatedAt, okUpdatedAt := args[3].Value.(time.Time) + if !okAuthID || !okModel || !okContent || !okUpdatedAt { + return nil, errors.New("invalid cooldown query arguments") + } + key := authID + "\x00" + model + current, exists := c.state.rows[key] + if len(args) == 4 { + if !exists || !current.updatedAt.After(updatedAt) { + c.state.rows[key] = cooldownTestRow{content: append([]byte(nil), content...), updatedAt: updatedAt} + } + return driver.RowsAffected(1), nil + } + observedAt, okObservedAt := args[4].Value.(time.Time) + if !okObservedAt { + return nil, errors.New("invalid cooldown delete version") + } + if !exists || (!current.deleted && !current.updatedAt.After(observedAt)) { + c.state.rows[key] = cooldownTestRow{content: append([]byte(nil), content...), deleted: true, updatedAt: updatedAt} + } + return driver.RowsAffected(1), nil +} + +func (c *cooldownTestConn) QueryContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + c.state.mu.Lock() + defer c.state.mu.Unlock() + c.state.queries = append(c.state.queries, query) + rows := make([]cooldownTestRow, 0, len(c.state.rows)) + for _, row := range c.state.rows { + if !row.deleted { + row.content = append([]byte(nil), row.content...) + rows = append(rows, row) + } + } + return &cooldownTestRows{rows: rows}, nil +} + +func (*cooldownTestTx) Commit() error { + return nil +} + +func (*cooldownTestTx) Rollback() error { + return nil +} + +func (r *cooldownTestRows) Columns() []string { + return []string{"content", "updated_at"} +} + +func (r *cooldownTestRows) Close() error { + return nil +} + +func (r *cooldownTestRows) Next(dest []driver.Value) error { + if r.index >= len(r.rows) { + return io.EOF + } + dest[0] = r.rows[r.index].content + dest[1] = r.rows[r.index].updatedAt + r.index++ + return nil +} + +func TestPostgresCooldownStateStore_SaveLoad(t *testing.T) { + state := &cooldownTestState{rows: make(map[string]cooldownTestRow)} + driverName := fmt.Sprintf("cliproxy_postgres_cooldown_test_%d", cooldownTestDriverID.Add(1)) + sql.Register(driverName, &cooldownTestDriver{state: state}) + db, errOpen := sql.Open(driverName, "") + if errOpen != nil { + t.Fatalf("sql.Open() error = %v", errOpen) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Errorf("db.Close() error = %v", errClose) + } + }) + + postgresStore := &PostgresStore{ + db: db, + cfg: PostgresStoreConfig{ + ConfigTable: defaultConfigTable, + AuthTable: defaultAuthTable, + CooldownTable: defaultCooldownTable, + }, + } + cooldownStore := &postgresCooldownStateStore{store: postgresStore} + postgresStore.cooldownStore = cooldownStore + + if errSchema := postgresStore.EnsureSchema(context.Background()); errSchema != nil { + t.Fatalf("EnsureSchema() error = %v", errSchema) + } + if got := postgresStore.CooldownStateStore(); got != cooldownStore { + t.Fatalf("CooldownStateStore() = %T, want configured PostgreSQL store", got) + } + + nextRetry := time.Date(2026, time.March, 15, 12, 0, 0, 0, time.UTC) + records := []cliproxyauth.CooldownStateRecord{ + { + Provider: "codex", + AuthID: "account-1", + Model: "gpt-test", + Status: string(cliproxyauth.StatusError), + NextRetryAfter: nextRetry, + Reason: "rate limited", + UpdatedAt: nextRetry.Add(-time.Minute), + }, + } + if errSave := cooldownStore.Save(context.Background(), records); errSave != nil { + t.Fatalf("Save() error = %v", errSave) + } + loaded, errLoad := cooldownStore.Load(context.Background()) + if errLoad != nil { + t.Fatalf("Load() error = %v", errLoad) + } + if !reflect.DeepEqual(loaded, records) { + t.Fatalf("Load() = %#v, want %#v", loaded, records) + } + + zeroTimeRecord := cliproxyauth.CooldownStateRecord{AuthID: "account-2", Model: "gpt-test"} + if errSave := cooldownStore.Save(context.Background(), []cliproxyauth.CooldownStateRecord{zeroTimeRecord}); errSave != nil { + t.Fatalf("Save() with zero UpdatedAt error = %v", errSave) + } + loaded, errLoad = cooldownStore.Load(context.Background()) + if errLoad != nil { + t.Fatalf("Load() after zero UpdatedAt error = %v", errLoad) + } + if len(loaded) != 1 || loaded[0].UpdatedAt.IsZero() { + t.Fatalf("Load() did not persist a normalized UpdatedAt: %#v", loaded) + } + + if errSave := cooldownStore.Save(context.Background(), nil); errSave != nil { + t.Fatalf("Save(nil) error = %v", errSave) + } + loaded, errLoad = cooldownStore.Load(context.Background()) + if errLoad != nil { + t.Fatalf("Load() after Save(nil) error = %v", errLoad) + } + if len(loaded) != 0 { + t.Fatalf("Load() after Save(nil) returned %d records, want 0", len(loaded)) + } + + state.mu.Lock() + queries := strings.Join(state.queries, "\n") + state.mu.Unlock() + if !strings.Contains(queries, `CREATE TABLE IF NOT EXISTS "cooldown_store"`) { + t.Fatalf("EnsureSchema() did not create cooldown table; queries:\n%s", queries) + } +} + +func TestPostgresCooldownStateStore_MergesConcurrentInstances(t *testing.T) { + state := &cooldownTestState{rows: make(map[string]cooldownTestRow)} + driverName := fmt.Sprintf("cliproxy_postgres_cooldown_merge_test_%d", cooldownTestDriverID.Add(1)) + sql.Register(driverName, &cooldownTestDriver{state: state}) + db, errOpen := sql.Open(driverName, "") + if errOpen != nil { + t.Fatalf("sql.Open() error = %v", errOpen) + } + t.Cleanup(func() { + if errClose := db.Close(); errClose != nil { + t.Errorf("db.Close() error = %v", errClose) + } + }) + postgresStore := &PostgresStore{ + db: db, + cfg: PostgresStoreConfig{CooldownTable: defaultCooldownTable}, + } + storeA := &postgresCooldownStateStore{store: postgresStore} + storeB := &postgresCooldownStateStore{store: postgresStore} + staleStore := &postgresCooldownStateStore{store: postgresStore} + + for _, cooldownStore := range []*postgresCooldownStateStore{storeA, storeB} { + if _, errLoad := cooldownStore.Load(context.Background()); errLoad != nil { + t.Fatalf("initial Load() error = %v", errLoad) + } + } + updatedAt := time.Now().UTC().Add(-time.Minute) + recordA := cliproxyauth.CooldownStateRecord{AuthID: "account-a", Model: "model-a", UpdatedAt: updatedAt} + recordB := cliproxyauth.CooldownStateRecord{AuthID: "account-b", Model: "model-b", UpdatedAt: updatedAt} + if errSave := storeA.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordA}); errSave != nil { + t.Fatalf("storeA.Save() error = %v", errSave) + } + if errSave := storeB.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordB}); errSave != nil { + t.Fatalf("storeB.Save() error = %v", errSave) + } + staleRecords, errLoad := staleStore.Load(context.Background()) + if errLoad != nil { + t.Fatalf("staleStore.Load() error = %v", errLoad) + } + if len(staleRecords) != 2 { + t.Fatalf("merged Load() returned %d records, want 2", len(staleRecords)) + } + + newerRecordA := recordA + newerRecordA.UpdatedAt = updatedAt.Add(time.Hour) + if errSave := storeA.Save(context.Background(), []cliproxyauth.CooldownStateRecord{newerRecordA}); errSave != nil { + t.Fatalf("storeA.Save(newer) error = %v", errSave) + } + if errSave := staleStore.Save(context.Background(), []cliproxyauth.CooldownStateRecord{recordB}); errSave != nil { + t.Fatalf("staleStore.Save(without newer record) error = %v", errSave) + } + resurrectStore := &postgresCooldownStateStore{store: postgresStore} + activeRecords, errLoad := resurrectStore.Load(context.Background()) + if errLoad != nil { + t.Fatalf("resurrectStore.Load() error = %v", errLoad) + } + if len(activeRecords) != 2 { + t.Fatalf("Load() after stale delete returned %d records, want 2", len(activeRecords)) + } + + if errSave := storeA.Save(context.Background(), nil); errSave != nil { + t.Fatalf("storeA.Save(nil) error = %v", errSave) + } + if errSave := resurrectStore.Save(context.Background(), activeRecords); errSave != nil { + t.Fatalf("resurrectStore.Save() error = %v", errSave) + } + reader := &postgresCooldownStateStore{store: postgresStore} + loaded, errLoad := reader.Load(context.Background()) + if errLoad != nil { + t.Fatalf("reader.Load() error = %v", errLoad) + } + if len(loaded) != 1 || loaded[0].AuthID != recordB.AuthID { + t.Fatalf("Load() after stale save = %#v, want only account-b", loaded) + } +} diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index 4b979486..46e7515d 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -20,29 +20,32 @@ import ( ) const ( - defaultConfigTable = "config_store" - defaultAuthTable = "auth_store" - defaultConfigKey = "config" + defaultConfigTable = "config_store" + defaultAuthTable = "auth_store" + defaultCooldownTable = "cooldown_store" + defaultConfigKey = "config" ) // PostgresStoreConfig captures configuration required to initialize a Postgres-backed store. type PostgresStoreConfig struct { - DSN string - Schema string - ConfigTable string - AuthTable string - SpoolDir string + DSN string + Schema string + ConfigTable string + AuthTable string + CooldownTable string + SpoolDir string } // PostgresStore persists configuration and authentication metadata using PostgreSQL as backend // while mirroring data to a local workspace so existing file-based workflows continue to operate. type PostgresStore struct { - db *sql.DB - cfg PostgresStoreConfig - spoolRoot string - configPath string - authDir string - mu sync.Mutex + db *sql.DB + cfg PostgresStoreConfig + spoolRoot string + configPath string + authDir string + cooldownStore *postgresCooldownStateStore + mu sync.Mutex } // NewPostgresStore establishes a connection to PostgreSQL and prepares the local workspace. @@ -58,6 +61,9 @@ func NewPostgresStore(ctx context.Context, cfg PostgresStoreConfig) (*PostgresSt if cfg.AuthTable == "" { cfg.AuthTable = defaultAuthTable } + if cfg.CooldownTable == "" { + cfg.CooldownTable = defaultCooldownTable + } spoolRoot := strings.TrimSpace(cfg.SpoolDir) if spoolRoot == "" { @@ -96,6 +102,7 @@ func NewPostgresStore(ctx context.Context, cfg PostgresStoreConfig) (*PostgresSt configPath: filepath.Join(configDir, "config.yaml"), authDir: authDir, } + store.cooldownStore = &postgresCooldownStateStore{store: store} return store, nil } @@ -140,6 +147,20 @@ func (s *PostgresStore) EnsureSchema(ctx context.Context) error { `, authTable)); err != nil { return fmt.Errorf("postgres store: create auth table: %w", err) } + cooldownTable := s.fullTableName(s.cfg.CooldownTable) + if _, err := s.db.ExecContext(ctx, fmt.Sprintf(` + CREATE TABLE IF NOT EXISTS %s ( + auth_id TEXT NOT NULL, + model TEXT NOT NULL DEFAULT '', + content JSONB NOT NULL, + deleted BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (auth_id, model) + ) + `, cooldownTable)); err != nil { + return fmt.Errorf("postgres store: create cooldown table: %w", err) + } return nil } diff --git a/sdk/cliproxy/auth/cooldown_state.go b/sdk/cliproxy/auth/cooldown_state.go index ab43ab0e..830a2e6e 100644 --- a/sdk/cliproxy/auth/cooldown_state.go +++ b/sdk/cliproxy/auth/cooldown_state.go @@ -35,6 +35,11 @@ type CooldownStateStore interface { Save(context.Context, []CooldownStateRecord) error } +// CooldownStateStoreProvider exposes a backend-specific cooldown state store. +type CooldownStateStoreProvider interface { + CooldownStateStore() CooldownStateStore +} + type cooldownStateFile struct { Version int `json:"version"` AuthID string `json:"auth_id,omitempty"` diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go index ada632c4..1081b6a5 100644 --- a/sdk/cliproxy/builder.go +++ b/sdk/cliproxy/builder.go @@ -48,6 +48,9 @@ type Builder struct { // coreManager handles core authentication and execution. coreManager *coreauth.Manager + // cooldownStateStore overrides runtime cooldown persistence. + cooldownStateStore coreauth.CooldownStateStore + // pluginHost owns dynamic plugin lifecycle and adapters. pluginHost *pluginhost.Host @@ -146,6 +149,12 @@ func (b *Builder) WithCoreAuthManager(mgr *coreauth.Manager) *Builder { return b } +// WithCooldownStateStore overrides the store used for runtime cooldown persistence. +func (b *Builder) WithCooldownStateStore(store coreauth.CooldownStateStore) *Builder { + b.cooldownStateStore = store + return b +} + // WithPluginHost overrides the dynamic plugin host used by the service. func (b *Builder) WithPluginHost(host *pluginhost.Host) *Builder { b.pluginHost = host @@ -227,12 +236,18 @@ func (b *Builder) Build() (*Service, error) { accessManager.SetProviders(sdkaccess.RegisteredProviders()) coreManager := b.coreManager + cooldownStateStore := b.cooldownStateStore var appliedRoutingState *routingRuntimeState if coreManager == nil { tokenStore := sdkAuth.GetTokenStore() if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil { dirSetter.SetBaseDir(b.cfg.AuthDir) } + if cooldownStateStore == nil { + if provider, ok := tokenStore.(coreauth.CooldownStateStoreProvider); ok { + cooldownStateStore = provider.CooldownStateStore() + } + } routingState := normalizedRoutingRuntimeState(b.cfg) coreManager = coreauth.NewManager(tokenStore, newRoutingSelector(routingState), nil) @@ -256,6 +271,7 @@ func (b *Builder) Build() (*Service, error) { authManager: authManager, accessManager: accessManager, coreManager: coreManager, + cooldownStateStore: cooldownStateStore, pluginHost: pluginHost, appliedRoutingState: appliedRoutingState, serverOptions: append([]api.ServerOption(nil), b.serverOptions...), diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 0077344f..bc08dbf8 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -89,6 +89,9 @@ type Service struct { // coreManager handles core authentication and execution. coreManager *coreauth.Manager + // cooldownStateStore persists runtime cooldown state when enabled. + cooldownStateStore coreauth.CooldownStateStore + // pluginHost owns dynamic plugin lifecycle and runtime capability adapters. pluginHost *pluginhost.Host diff --git a/sdk/cliproxy/service_auth.go b/sdk/cliproxy/service_auth.go index 0b1990c2..11b1e1d1 100644 --- a/sdk/cliproxy/service_auth.go +++ b/sdk/cliproxy/service_auth.go @@ -384,6 +384,9 @@ func (s *Service) resolveCooldownStateStore(cfg *config.Config) coreauth.Cooldow if cfg == nil || !cfg.SaveCooldownStatus || cfg.Home.Enabled { return nil } + if s != nil && s.cooldownStateStore != nil { + return s.cooldownStateStore + } authDir, errResolve := resolveCooldownStateAuthDir(cfg) if errResolve != nil { log.Warnf("failed to resolve cooldown state directory: %v", errResolve) diff --git a/sdk/cliproxy/service_cooldown_store_test.go b/sdk/cliproxy/service_cooldown_store_test.go new file mode 100644 index 00000000..0c7305ed --- /dev/null +++ b/sdk/cliproxy/service_cooldown_store_test.go @@ -0,0 +1,68 @@ +package cliproxy + +import ( + "context" + "path/filepath" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +type cooldownProviderTokenStore struct { + cooldownStore coreauth.CooldownStateStore +} + +func (s *cooldownProviderTokenStore) List(context.Context) ([]*coreauth.Auth, error) { + return nil, nil +} + +func (s *cooldownProviderTokenStore) Save(context.Context, *coreauth.Auth) (string, error) { + return "", nil +} + +func (s *cooldownProviderTokenStore) Delete(context.Context, string) error { + return nil +} + +func (s *cooldownProviderTokenStore) CooldownStateStore() coreauth.CooldownStateStore { + return s.cooldownStore +} + +type serviceCooldownStateStore struct{} + +func (*serviceCooldownStateStore) Load(context.Context) ([]coreauth.CooldownStateRecord, error) { + return nil, nil +} + +func (*serviceCooldownStateStore) Save(context.Context, []coreauth.CooldownStateRecord) error { + return nil +} + +func TestResolveCooldownStateStoreUsesCapturedBackendProvider(t *testing.T) { + originalStore := sdkAuth.GetTokenStore() + t.Cleanup(func() { + sdkAuth.RegisterTokenStore(originalStore) + }) + + providedStore := &serviceCooldownStateStore{} + sdkAuth.RegisterTokenStore(&cooldownProviderTokenStore{cooldownStore: providedStore}) + cfg := &config.Config{ + AuthDir: t.TempDir(), + SaveCooldownStatus: true, + } + service, errBuild := NewBuilder(). + WithConfig(cfg). + WithConfigPath(filepath.Join(t.TempDir(), "config.yaml")). + Build() + if errBuild != nil { + t.Fatalf("Build() error = %v", errBuild) + } + + sdkAuth.RegisterTokenStore(&cooldownProviderTokenStore{cooldownStore: &serviceCooldownStateStore{}}) + got := service.resolveCooldownStateStore(cfg) + if got != providedStore { + t.Fatalf("resolveCooldownStateStore() = %T, want captured backend-provided store", got) + } +} -- 2.51.2