From 63d4cdffefa71469ed556f3d7edfb4a18a192b14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Fri, 13 Mar 2026 21:54:47 +0800
Subject: [PATCH] ocm: rewrite codex.rate_limits WebSocket events for external
 users

The HTTP path rewrites utilization headers for external users via
rewriteResponseHeadersForExternalUser to show aggregated values.
The WebSocket upgrade headers were also rewritten, but in-band
codex.rate_limits events were forwarded unmodified, leaking
per-credential utilization to external users.
---
 service/ocm/service_websocket.go | 42 ++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/service/ocm/service_websocket.go b/service/ocm/service_websocket.go
index d3f2535c0..b35733b69 100644
--- a/service/ocm/service_websocket.go
+++ b/service/ocm/service_websocket.go
@@ -249,7 +249,7 @@ func (s *Service) handleWebSocket(
 	go func() {
 		defer waitGroup.Done()
 		defer session.Close()
-		s.proxyWebSocketUpstreamToClient(upstreamReadWriter, clientConn, selectedCredential, modelChannel, username, weeklyCycleHint)
+		s.proxyWebSocketUpstreamToClient(upstreamReadWriter, clientConn, selectedCredential, userConfig, provider, modelChannel, username, weeklyCycleHint)
 	}()
 	waitGroup.Wait()
 }
@@ -287,7 +287,7 @@ func (s *Service) proxyWebSocketClientToUpstream(clientConn net.Conn, upstreamCo
 	}
 }
 
-func (s *Service) proxyWebSocketUpstreamToClient(upstreamReadWriter io.ReadWriter, clientConn net.Conn, selectedCredential credential, modelChannel <-chan string, username string, weeklyCycleHint *WeeklyCycleHint) {
+func (s *Service) proxyWebSocketUpstreamToClient(upstreamReadWriter io.ReadWriter, clientConn net.Conn, selectedCredential credential, userConfig *option.OCMUser, provider credentialProvider, modelChannel <-chan string, username string, weeklyCycleHint *WeeklyCycleHint) {
 	usageTracker := selectedCredential.usageTrackerOrNil()
 	var requestModel string
 	for {
@@ -308,6 +308,12 @@ func (s *Service) proxyWebSocketUpstreamToClient(upstreamReadWriter io.ReadWrite
 				switch event.Type {
 				case "codex.rate_limits":
 					s.handleWebSocketRateLimitsEvent(data, selectedCredential)
+					if userConfig != nil && userConfig.ExternalCredential != "" {
+						rewritten, rewriteErr := s.rewriteWebSocketRateLimitsForExternalUser(data, provider, userConfig)
+						if rewriteErr == nil {
+							data = rewritten
+						}
+					}
 				case "error":
 					if event.StatusCode == http.StatusTooManyRequests {
 						s.handleWebSocketErrorRateLimited(data, selectedCredential)
@@ -397,6 +403,38 @@ func (s *Service) handleWebSocketErrorRateLimited(data []byte, selectedCredentia
 	selectedCredential.markRateLimited(resetAt)
 }
 
+func (s *Service) rewriteWebSocketRateLimitsForExternalUser(data []byte, provider credentialProvider, userConfig *option.OCMUser) ([]byte, error) {
+	var event struct {
+		Type       string `json:"type"`
+		RateLimits struct {
+			Primary *struct {
+				UsedPercent   float64 `json:"used_percent"`
+				WindowMinutes int64   `json:"window_minutes,omitempty"`
+				ResetAt       int64   `json:"reset_at,omitempty"`
+			} `json:"primary,omitempty"`
+			Secondary *struct {
+				UsedPercent   float64 `json:"used_percent"`
+				WindowMinutes int64   `json:"window_minutes,omitempty"`
+				ResetAt       int64   `json:"reset_at,omitempty"`
+			} `json:"secondary,omitempty"`
+		} `json:"rate_limits"`
+		LimitName        string `json:"limit_name,omitempty"`
+		MeteredLimitName string `json:"metered_limit_name,omitempty"`
+	}
+	err := json.Unmarshal(data, &event)
+	if err != nil {
+		return nil, err
+	}
+	averageFiveHour, averageWeekly := s.computeAggregatedUtilization(provider, userConfig)
+	if event.RateLimits.Primary != nil {
+		event.RateLimits.Primary.UsedPercent = averageFiveHour
+	}
+	if event.RateLimits.Secondary != nil {
+		event.RateLimits.Secondary.UsedPercent = averageWeekly
+	}
+	return json.Marshal(event)
+}
+
 func (s *Service) handleWebSocketResponseCompleted(data []byte, usageTracker *AggregatedUsage, requestModel string, username string, weeklyCycleHint *WeeklyCycleHint) {
 	var streamEvent responses.ResponseStreamEventUnion
 	if json.Unmarshal(data, &streamEvent) != nil {