wangjia/sing-box
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ocm: rewrite codex.rate_limits WebSocket events for external users

Browse Source 
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.
This commit is contained in:
世界
2026-03-13 21:54:47 +08:00
parent 5516d7b045
commit 63d4cdffef
1 changed files with 40 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/ocm/service_websocket.go
+40 -2
View File
@@ -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 {