From ca60f93184aeb1ca0d18c4a538c66d900bc9f086 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= Date: Thu, 26 Mar 2026 15:42:51 +0800 Subject: [PATCH] fix(ocm): inject synthetic rate limits inline when intercepting upstream events The initial synthetic event from 6721dff48 arrives before the Codex CLI's response stream reader is active. Additionally, the shouldEmit gate in updateStateFromHeaders suppresses the async replacement when values haven't changed. Send aggregated status inline in proxyWebSocketUpstreamToClient so the client receives it at the exact protocol position it expects. --- service/ocm/service_websocket.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/service/ocm/service_websocket.go b/service/ocm/service_websocket.go index 4f76179ef..a441bdc8d 100644 --- a/service/ocm/service_websocket.go +++ b/service/ocm/service_websocket.go @@ -342,7 +342,7 @@ func (s *Service) handleWebSocket( go func() { defer waitGroup.Done() defer session.Close() - s.proxyWebSocketUpstreamToClient(ctx, upstreamReadWriter, clientConn, &clientWriteAccess, selectedCredential, modelChannel, username, weeklyCycleHint) + s.proxyWebSocketUpstreamToClient(ctx, upstreamReadWriter, clientConn, &clientWriteAccess, selectedCredential, modelChannel, username, weeklyCycleHint, provider, userConfig) }() go func() { defer waitGroup.Done() @@ -407,7 +407,7 @@ func (s *Service) proxyWebSocketClientToUpstream(ctx context.Context, clientConn } } -func (s *Service) proxyWebSocketUpstreamToClient(ctx context.Context, upstreamReadWriter io.ReadWriter, clientConn net.Conn, clientWriteAccess *sync.Mutex, selectedCredential Credential, modelChannel <-chan string, username string, weeklyCycleHint *WeeklyCycleHint) { +func (s *Service) proxyWebSocketUpstreamToClient(ctx context.Context, upstreamReadWriter io.ReadWriter, clientConn net.Conn, clientWriteAccess *sync.Mutex, selectedCredential Credential, modelChannel <-chan string, username string, weeklyCycleHint *WeeklyCycleHint, provider credentialProvider, userConfig *option.OCMUser) { usageTracker := selectedCredential.usageTrackerOrNil() var requestModel string for { @@ -428,6 +428,8 @@ func (s *Service) proxyWebSocketUpstreamToClient(ctx context.Context, upstreamRe switch event.Type { case "codex.rate_limits": s.handleWebSocketRateLimitsEvent(data, selectedCredential) + status := s.computeAggregatedUtilization(provider, userConfig) + writeWebSocketAggregatedStatus(clientConn, clientWriteAccess, status) continue case "error": s.handleWebSocketErrorEvent(data, selectedCredential)