From f353979e0a6f4d5a9fcc4a1a8d4fb616a710852a Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 2 Jun 2026 02:52:27 +0800 Subject: [PATCH] feat(watcher, redisqueue): add usage refresh notification support - Introduced `NotifyUsageRefresh` in `redisqueue` to notify subscribers of usage refresh events. - Enhanced `Watcher` logic to trigger usage refresh notifications on client changes (add/update/remove). - Updated tests to validate proper broadcast of usage refresh messages to subscribers. - Added support for initial `support_refresh` payload upon subscription initialization. --- .../redis_queue_protocol_integration_test.go | 124 ++++++++++++++++++ internal/redisqueue/queue.go | 8 ++ internal/redisqueue/queue_test.go | 23 ++++ internal/watcher/clients.go | 4 + internal/watcher/watcher_test.go | 64 +++++++++ 5 files changed, 223 insertions(+) diff --git a/internal/api/redis_queue_protocol_integration_test.go b/internal/api/redis_queue_protocol_integration_test.go index 834e4a86..7d443f67 100644 --- a/internal/api/redis_queue_protocol_integration_test.go +++ b/internal/api/redis_queue_protocol_integration_test.go @@ -159,6 +159,68 @@ func readRESPArrayOfBulkStrings(r *bufio.Reader) ([][]byte, error) { return out, nil } +func readTestRESPPubSubSubscribe(r *bufio.Reader) (string, int, error) { + prefix, errRead := r.ReadByte() + if errRead != nil { + return "", 0, errRead + } + if prefix != '*' { + return "", 0, fmt.Errorf("expected array prefix '*', got %q", prefix) + } + line, errLine := readTestRESPLine(r) + if errLine != nil { + return "", 0, errLine + } + count, errParse := strconv.Atoi(line) + if errParse != nil { + return "", 0, fmt.Errorf("invalid array length %q: %v", line, errParse) + } + if count != 3 { + return "", 0, fmt.Errorf("subscribe ack length = %d, want 3", count) + } + kind, errKind := readTestRESPBulkString(r) + if errKind != nil { + return "", 0, errKind + } + if string(kind) != "subscribe" { + return "", 0, fmt.Errorf("subscribe ack kind = %q", string(kind)) + } + channel, errChannel := readTestRESPBulkString(r) + if errChannel != nil { + return "", 0, errChannel + } + prefix, errRead = r.ReadByte() + if errRead != nil { + return "", 0, errRead + } + if prefix != ':' { + return "", 0, fmt.Errorf("expected integer prefix ':', got %q", prefix) + } + line, errLine = readTestRESPLine(r) + if errLine != nil { + return "", 0, errLine + } + subscriptions, errParse := strconv.Atoi(line) + if errParse != nil { + return "", 0, fmt.Errorf("invalid subscription count %q: %v", line, errParse) + } + return string(channel), subscriptions, nil +} + +func readTestRESPPubSubMessage(r *bufio.Reader) (string, []byte, error) { + items, errItems := readRESPArrayOfBulkStrings(r) + if errItems != nil { + return "", nil, errItems + } + if len(items) != 3 { + return "", nil, fmt.Errorf("pubsub message length = %d, want 3", len(items)) + } + if string(items[0]) != "message" { + return "", nil, fmt.Errorf("pubsub message kind = %q", string(items[0])) + } + return string(items[1]), items[2], nil +} + func TestRedisProtocol_ManagementDisabled_RejectsConnection(t *testing.T) { t.Setenv("MANAGEMENT_PASSWORD", "") redisqueue.SetEnabled(false) @@ -235,6 +297,68 @@ func TestRedisProtocol_HomeEnabled_DisablesConnection(t *testing.T) { } } +func TestRedisProtocol_SUBSCRIBE_UsageSendsSupportRefresh(t *testing.T) { + const managementPassword = "test-management-password" + + t.Setenv("MANAGEMENT_PASSWORD", managementPassword) + redisqueue.SetEnabled(false) + t.Cleanup(func() { redisqueue.SetEnabled(false) }) + + server := newTestServer(t) + if !server.managementRoutesEnabled.Load() { + t.Fatalf("expected managementRoutesEnabled to be true") + } + + addr, stop := startRedisMuxListener(t, server) + t.Cleanup(stop) + + conn, errDial := net.DialTimeout("tcp", addr, time.Second) + if errDial != nil { + t.Fatalf("failed to dial redis listener: %v", errDial) + } + t.Cleanup(func() { _ = conn.Close() }) + + reader := bufio.NewReader(conn) + _ = conn.SetDeadline(time.Now().Add(5 * time.Second)) + + if errWrite := writeTestRESPCommand(conn, "AUTH", managementPassword); errWrite != nil { + t.Fatalf("failed to write AUTH command: %v", errWrite) + } + if msg, errRead := readTestRESPSimpleString(reader); errRead != nil { + t.Fatalf("failed to read AUTH response: %v", errRead) + } else if msg != "OK" { + t.Fatalf("unexpected AUTH response: %q", msg) + } + + if errWrite := writeTestRESPCommand(conn, "SUBSCRIBE", "usage"); errWrite != nil { + t.Fatalf("failed to write SUBSCRIBE command: %v", errWrite) + } + channel, subscriptions, errSubscribe := readTestRESPPubSubSubscribe(reader) + if errSubscribe != nil { + t.Fatalf("failed to read subscribe response: %v", errSubscribe) + } + if channel != "usage" || subscriptions != 1 { + t.Fatalf("unexpected subscribe response channel=%q subscriptions=%d", channel, subscriptions) + } + + channel, payload, errMessage := readTestRESPPubSubMessage(reader) + if errMessage != nil { + t.Fatalf("failed to read support refresh message: %v", errMessage) + } + if channel != "usage" || string(payload) != `{"support_refresh":true}` { + t.Fatalf("unexpected support refresh message channel=%q payload=%q", channel, string(payload)) + } + + redisqueue.Enqueue([]byte(`{"id":1}`)) + channel, payload, errMessage = readTestRESPPubSubMessage(reader) + if errMessage != nil { + t.Fatalf("failed to read usage message: %v", errMessage) + } + if channel != "usage" || string(payload) != `{"id":1}` { + t.Fatalf("unexpected usage message channel=%q payload=%q", channel, string(payload)) + } +} + func TestRedisProtocol_AUTH_And_PopContracts(t *testing.T) { const managementPassword = "test-management-password" diff --git a/internal/redisqueue/queue.go b/internal/redisqueue/queue.go index 6a2a594e..60aecdff 100644 --- a/internal/redisqueue/queue.go +++ b/internal/redisqueue/queue.go @@ -10,6 +10,9 @@ const ( defaultRetentionSeconds int64 = 60 maxRetentionSeconds int64 = 3600 usageSubscriberBuffer = 256 + + usageSupportRefreshPayload = `{"support_refresh":true}` + usageRefreshPayload = `{"refresh":true}` ) type queueItem struct { @@ -83,6 +86,10 @@ func SubscribeUsage() (<-chan []byte, func()) { return global.subscribeUsage() } +func NotifyUsageRefresh() { + global.publishToSubscribers([]byte(usageRefreshPayload)) +} + func (q *queue) clear() { q.mu.Lock() @@ -137,6 +144,7 @@ func (q *queue) publishToSubscribers(payload []byte) bool { func (q *queue) subscribeUsage() (<-chan []byte, func()) { subscriber := make(chan []byte, usageSubscriberBuffer) + subscriber <- []byte(usageSupportRefreshPayload) q.mu.Lock() if q.subscribers == nil { diff --git a/internal/redisqueue/queue_test.go b/internal/redisqueue/queue_test.go index f40c8826..1bc0fc30 100644 --- a/internal/redisqueue/queue_test.go +++ b/internal/redisqueue/queue_test.go @@ -12,6 +12,9 @@ func TestEnqueueBroadcastsToUsageSubscribersAndSkipsQueue(t *testing.T) { second, unsubscribeSecond := SubscribeUsage() defer unsubscribeSecond() + requireUsageSubscriberPayload(t, first, usageSupportRefreshPayload) + requireUsageSubscriberPayload(t, second, usageSupportRefreshPayload) + Enqueue([]byte("usage-record")) requireUsageSubscriberPayload(t, first, "usage-record") @@ -37,6 +40,8 @@ func TestSetEnabledFalseClosesUsageSubscribers(t *testing.T) { subscriber, unsubscribe := SubscribeUsage() defer unsubscribe() + requireUsageSubscriberPayload(t, subscriber, usageSupportRefreshPayload) + SetEnabled(false) select { @@ -50,6 +55,24 @@ func TestSetEnabledFalseClosesUsageSubscribers(t *testing.T) { }) } +func TestNotifyUsageRefreshBroadcastsOnlyToUsageSubscribers(t *testing.T) { + withEnabledQueue(t, func() { + subscriber, unsubscribe := SubscribeUsage() + defer unsubscribe() + + requireUsageSubscriberPayload(t, subscriber, usageSupportRefreshPayload) + + NotifyUsageRefresh() + requireUsageSubscriberPayload(t, subscriber, usageRefreshPayload) + + unsubscribe() + NotifyUsageRefresh() + if items := PopOldest(1); len(items) != 0 { + t.Fatalf("PopOldest() items = %q, want empty after refresh notification without subscribers", items) + } + }) +} + func requireUsageSubscriberPayload(t *testing.T, subscriber <-chan []byte, want string) { t.Helper() diff --git a/internal/watcher/clients.go b/internal/watcher/clients.go index 0a46660e..be6738ce 100644 --- a/internal/watcher/clients.go +++ b/internal/watcher/clients.go @@ -14,6 +14,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/diff" "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" @@ -134,6 +135,7 @@ func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string } w.refreshAuthState(forceAuthRefresh) + redisqueue.NotifyUsageRefresh() log.Infof("full client load complete - %d clients (%d auth files + %d Gemini API keys + %d Vertex API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)", totalNewClients, @@ -233,6 +235,7 @@ func (w *Watcher) addOrUpdateClient(path string) { w.persistAuthAsync(fmt.Sprintf("Sync auth %s", filepath.Base(path)), path) w.dispatchAuthUpdates(updates) + redisqueue.NotifyUsageRefresh() } func (w *Watcher) removeClient(path string) { @@ -251,6 +254,7 @@ func (w *Watcher) removeClient(path string) { w.persistAuthAsync(fmt.Sprintf("Remove auth %s", filepath.Base(path)), path) w.dispatchAuthUpdates(updates) + redisqueue.NotifyUsageRefresh() } func (w *Watcher) computePerPathUpdatesLocked(oldByID, newByID map[string]*coreauth.Auth) []AuthUpdate { diff --git a/internal/watcher/watcher_test.go b/internal/watcher/watcher_test.go index bb3b5577..d93c2233 100644 --- a/internal/watcher/watcher_test.go +++ b/internal/watcher/watcher_test.go @@ -15,6 +15,7 @@ import ( "github.com/fsnotify/fsnotify" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/diff" "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" @@ -441,6 +442,34 @@ func TestRemoveClientRemovesHash(t *testing.T) { } } +func TestAuthFileClientChangesNotifyUsageSubscribersToRefresh(t *testing.T) { + tmpDir := t.TempDir() + authFile := filepath.Join(tmpDir, "sample.json") + if err := os.WriteFile(authFile, []byte(`{"type":"demo","api_key":"k"}`), 0o644); err != nil { + t.Fatalf("failed to create auth file: %v", err) + } + + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(true) + t.Cleanup(func() { redisqueue.SetEnabled(false) }) + + subscriber, unsubscribe := redisqueue.SubscribeUsage() + defer unsubscribe() + requireWatcherUsagePayload(t, subscriber, `{"support_refresh":true}`) + + w := &Watcher{ + authDir: tmpDir, + lastAuthHashes: make(map[string]string), + } + w.SetConfig(&config.Config{AuthDir: tmpDir}) + + w.addOrUpdateClient(authFile) + requireWatcherUsagePayload(t, subscriber, `{"refresh":true}`) + + w.removeClient(authFile) + requireWatcherUsagePayload(t, subscriber, `{"refresh":true}`) +} + func TestAuthFileEventsDoNotInvokeSnapshotCoreAuths(t *testing.T) { tmpDir := t.TempDir() authFile := filepath.Join(tmpDir, "sample.json") @@ -699,6 +728,25 @@ func TestReloadClientsHandlesNilConfig(t *testing.T) { w.reloadClients(true, nil, false) } +func TestReloadClientsNotifiesUsageSubscribersToRefresh(t *testing.T) { + tmp := t.TempDir() + redisqueue.SetEnabled(false) + redisqueue.SetEnabled(true) + t.Cleanup(func() { redisqueue.SetEnabled(false) }) + + subscriber, unsubscribe := redisqueue.SubscribeUsage() + defer unsubscribe() + requireWatcherUsagePayload(t, subscriber, `{"support_refresh":true}`) + + w := &Watcher{ + authDir: tmp, + config: &config.Config{AuthDir: tmp}, + } + w.reloadClients(false, nil, false) + + requireWatcherUsagePayload(t, subscriber, `{"refresh":true}`) +} + func TestReloadClientsFiltersProvidersWithNilCurrentAuths(t *testing.T) { tmp := t.TempDir() w := &Watcher{ @@ -711,6 +759,22 @@ func TestReloadClientsFiltersProvidersWithNilCurrentAuths(t *testing.T) { } } +func requireWatcherUsagePayload(t *testing.T, subscriber <-chan []byte, want string) { + t.Helper() + + select { + case got, ok := <-subscriber: + if !ok { + t.Fatalf("subscriber closed before receiving %q", want) + } + if string(got) != want { + t.Fatalf("subscriber payload = %q, want %q", string(got), want) + } + case <-time.After(time.Second): + t.Fatalf("timeout waiting for subscriber payload %q", want) + } +} + func TestSetAuthUpdateQueueNilResetsDispatch(t *testing.T) { w := &Watcher{} queue := make(chan AuthUpdate, 1) -- 2.51.2