From 0fc2e73cf856e573e25925f4206d52d120cefef4 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Wed, 21 Oct 2026 15:14:40 -0400 Subject: [PATCH 01/12] feat: Carry Codex's Responses WebSocket and parse every turn on it Codex prefers a WebSocket to /backend-api/codex/responses: one socket per session, with later turns sending only new input. Cortex could not carry it, so it declined it (426) and Codex used HTTPS. Now the forward proxy carries it. It completes the handshake with the origin, takes over the client connection and relays frames unchanged in both directions. Each response.create and the event messages answering it run the outbound pipeline as the equivalent HTTPS request and SSE stream: inference-parser reads the turn, its cost is settled, it records an ordinary POST request/response pair, and Finishers run. The socket itself records a GET open and close with tunnelReason "websocket". A plugin rejecting a turn keeps it from the origin and answers Codex with the Responses error event it reads as a server refusal. A request-body writer's rewrite is what the origin receives. permessage-deflate is stripped from the handshake so frames stay readable. A socket that cannot be carried (not bridged, HTTP/2, or a WritesResponseBody plugin in the chain) is declined with the 426. listener.responses_websockets: parse | decline | forward replaces the decline_responses_websockets boolean; parse is the default. Measured on Codex 0.163.0: three codex exec turns each opened one socket with no handshake error, and every turn was recorded with its model, tokens and cost, in about 3s per turn. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- CLAUDE.md | 14 +- cmd/agentop/tui/detail_pane.go | 2 +- cmd/cortex/main.go | 8 +- core/config/config.go | 56 +- core/config/config_test.go | 46 +- core/config/validate.go | 7 + core/listener/forwardproxy/server.go | 107 +-- .../forwardproxy/upgradedeclined_test.go | 31 +- .../forwardproxy/upgradestripped_test.go | 17 +- core/listener/forwardproxy/websocket.go | 699 +++++++++++++++++ core/listener/forwardproxy/websocket_test.go | 708 ++++++++++++++++++ core/listener/internal/wsframe/wsframe.go | 260 +++++++ .../listener/internal/wsframe/wsframe_test.go | 216 ++++++ core/listener/parity/dispatch_audit_test.go | 8 + core/pipeline/context.go | 7 +- core/pipeline/session.go | 9 +- docs/README.md | 2 +- docs/agents/codex.md | 82 +- docs/framework-architecture.md | 39 +- 19 files changed, 2130 insertions(+), 188 deletions(-) create mode 100644 core/listener/forwardproxy/websocket.go create mode 100644 core/listener/forwardproxy/websocket_test.go create mode 100644 core/listener/internal/wsframe/wsframe.go create mode 100644 core/listener/internal/wsframe/wsframe_test.go diff --git a/CLAUDE.md b/CLAUDE.md index 2a4d02b09..894c3ea69 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -600,11 +600,15 @@ Every event on `/v1/sessions/{id}` and `/v1/events` carries: `api.openai.com`; other hosts keep the strip), because the origin would only answer the stripped GET with a 405, which Codex retries for ~6s, while a 426 sends it straight to HTTPS. Nothing is forwarded, the Finisher sees an allow with - status 426, and `/v1/usage` does not count it as an error. On by default; - `listener.decline_responses_websockets: false` forwards the upgrade instead. - **A request whose body could not be buffered is the other exception:** it never - reached the pipeline, so its request row carries no - parse and no Finisher runs. Over the forward proxy's `maxRequestBodySize` (32 MiB) + status 426, and `/v1/usage` does not count it as an error. It is what + `listener.responses_websockets: decline` does to every such socket, and what the + default, `parse`, does to one it cannot carry (not bridged, HTTP/2, or a + `WritesResponseBody` plugin in the chain); a carried socket records instead a `GET` + open and close with `tunnelReason: "websocket"` and one ordinary `POST` + request/response pair per turn on it (`core/listener/forwardproxy/websocket.go`). + `forward` restores the stripped GET and the 405. **A request whose body could not be + buffered is the other exception:** it never reached the pipeline, so its request row + carries no parse and no Finisher runs. Over the forward proxy's `maxRequestBodySize` (32 MiB) it is a 413 `proxy_error` with an empty `error.code`; a client that hung up mid-upload gets a 499 `client_canceled`, which here does log a WARN (`request body unreadable`); anything else unreadable is a 400 `request_unreadable`. Before these diff --git a/cmd/agentop/tui/detail_pane.go b/cmd/agentop/tui/detail_pane.go index 036073bdc..b309438ec 100644 --- a/cmd/agentop/tui/detail_pane.go +++ b/cmd/agentop/tui/detail_pane.go @@ -120,7 +120,7 @@ func redirectHeader(e *pipeline.SessionEvent) string { // user blames Cortex. Cortex must not forward Upgrade (RFC 9110 makes it hop-by-hop), // so the request either reaches the origin as a bare GET — which an endpoint wanting // a POST or an upgrade answers 405 — or, on a Responses API path, is answered 426 by -// the proxy itself (listener.decline_responses_websockets). Codex falls back to +// the proxy itself (listener.responses_websockets). Codex falls back to // HTTPS after either. // // The wording holds for both because this line is drawn from the REQUEST row, which diff --git a/cmd/cortex/main.go b/cmd/cortex/main.go index caa7774a5..b92990d1d 100644 --- a/cmd/cortex/main.go +++ b/cmd/cortex/main.go @@ -972,10 +972,10 @@ func main() { fpSrv.SkipHosts = skipHosts fpSrv.TLSBridge = bridge fpSrv.Shared = sharedStore - // A Responses API WebSocket is answered 426 so Codex falls back to HTTPS at once, - // rather than after ~6s of retrying the origin's 405. On unless - // listener.decline_responses_websockets: false. - fpSrv.DeclineResponsesWebSockets = cfg.Listener.DeclinesResponsesWebSockets() + // A Responses API WebSocket — Codex's transport — is carried and parsed turn by + // turn by default; listener.responses_websockets: decline answers it 426 so Codex + // uses HTTPS at once, and forward restores the stripped GET. + fpSrv.ResponsesWebSockets = cfg.Listener.ResponsesWebSocketsMode() // Per-session bucketing. Without it every coding-agent session on the // machine records into one shared bucket: two Claude Code windows // interleave and per-session cost cannot be computed at all. Defaults to diff --git a/core/config/config.go b/core/config/config.go index f16cc8bfa..5802e259b 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -935,29 +935,45 @@ type ListenerConfig struct { // in #496. SkipHosts []string `yaml:"skip_hosts" json:"skip_hosts"` - // DeclineResponsesWebSockets: whether the forward proxy answers a WebSocket - // upgrade on an OpenAI Responses API path (one ending /responses, on chatgpt.com - // or api.openai.com) with 426 itself, instead of forwarding the request with its Upgrade stripped. nil means - // unset, and the default is on; explicit false restores the forwarding. See - // DeclinesResponsesWebSockets for why on is the default. - // - // A pointer for the reason SessionConfig.Enabled is one: with a plain bool, - // "unset" and "false" are the same value, and the default could not be on. - DeclineResponsesWebSockets *bool `yaml:"decline_responses_websockets,omitempty" json:"decline_responses_websockets,omitempty"` + // ResponsesWebSockets is what the forward proxy does with a WebSocket upgrade on + // an OpenAI Responses API path (a GET with Upgrade: websocket on a path ending + // /responses — Codex's transport): one of the ResponsesWebSockets* values below. + // Empty means the default, ResponsesWebSocketsParse. See + // ResponsesWebSocketsMode. + ResponsesWebSockets string `yaml:"responses_websockets,omitempty" json:"responses_websockets,omitempty"` } -// DeclinesResponsesWebSockets reports whether the forward proxy answers a WebSocket -// upgrade on a Responses API path with 426 rather than forwarding it. +// Valid ListenerConfig.ResponsesWebSockets values. +const ( + // ResponsesWebSocketsParse carries the WebSocket and parses every turn on it the + // way an HTTPS Responses request is parsed. Where the proxy cannot do that — the + // connection is HTTP/2, the request was not decrypted, or a plugin rewrites + // response bodies — it falls back to ResponsesWebSocketsDecline for that socket. + ResponsesWebSocketsParse = "parse" + // ResponsesWebSocketsDecline answers the upgrade with 426 Upgrade Required and + // contacts no origin; Codex reads that as "use HTTPS" and switches at once. + ResponsesWebSocketsDecline = "decline" + // ResponsesWebSocketsForward forwards the upgrade with Upgrade stripped, as + // every WebSocket was before: the origin sees a bare GET and answers 405, which + // Codex retries for about six seconds before falling back to HTTPS. + ResponsesWebSocketsForward = "forward" +) + +// ResponsesWebSocketsMode is ListenerConfig.ResponsesWebSockets with the default +// applied: parse, in every deployment. // -// On unless set false, in every deployment. The proxy must strip Upgrade (RFC 9110 -// makes it hop-by-hop), so a forwarded upgrade reaches the origin as a bare GET, and -// OpenAI's Responses endpoint answers that 405. Codex reads a 405 as transient and -// retries for about six seconds before falling back to HTTPS; it reads a 426 as "use -// HTTP" and falls back at once. The forwarded request can never become a WebSocket, so -// nothing is lost by answering it here. False is the escape hatch for a client that -// reads 426 some other way. -func (l ListenerConfig) DeclinesResponsesWebSockets() bool { - return l.DeclineResponsesWebSockets == nil || *l.DeclineResponsesWebSockets +// Parse is the default because it is the transport Codex asks for and the one it +// performs best on: one socket for a whole session, with later turns sending only new +// input instead of the whole conversation. Every turn on it is still parsed — model, +// tokens, cost, tool calls — so carrying it costs Cortex no visibility, and where it +// cannot be carried the socket is declined rather than broken. Decline and forward are +// the escape hatches: decline keeps Codex on HTTPS without delay, and forward is the +// behaviour before either existed. +func (l ListenerConfig) ResponsesWebSocketsMode() string { + if l.ResponsesWebSockets == "" { + return ResponsesWebSocketsParse + } + return l.ResponsesWebSockets } // Proxy roles selectable via ListenerConfig.Roles in proxy-sidecar mode. diff --git a/core/config/config_test.go b/core/config/config_test.go index cdfe70113..66e55e854 100644 --- a/core/config/config_test.go +++ b/core/config/config_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "os" "path/filepath" + "strings" "testing" "time" @@ -1043,19 +1044,21 @@ func TestSessionConfig_ClientAffinity(t *testing.T) { } } -// TestListenerConfig_DeclineResponsesWebSockets pins the default (on) and the key that -// turns it off. Unset must mean on with or without a listener block: every laptop config -// already on disk predates the key, and the built-in one is never regenerated. -func TestListenerConfig_DeclineResponsesWebSockets(t *testing.T) { +// TestListenerConfig_ResponsesWebSockets pins the default (parse) and every value an +// operator can write. Unset must mean parse with or without a listener block: every +// laptop config already on disk predates the key, and the built-in one is never +// regenerated. +func TestListenerConfig_ResponsesWebSockets(t *testing.T) { tests := []struct { name string yaml string - want bool + want string }{ - {"no listener block", "mode: proxy-sidecar\n", true}, - {"listener block without the key", "listener:\n roles: [forward]\n", true}, - {"explicit true", "listener:\n decline_responses_websockets: true\n", true}, - {"explicit false", "listener:\n decline_responses_websockets: false\n", false}, + {"no listener block", "mode: proxy-sidecar\n", ResponsesWebSocketsParse}, + {"listener block without the key", "listener:\n roles: [forward]\n", ResponsesWebSocketsParse}, + {"parse", "listener:\n responses_websockets: parse\n", ResponsesWebSocketsParse}, + {"decline", "listener:\n responses_websockets: decline\n", ResponsesWebSocketsDecline}, + {"forward", "listener:\n responses_websockets: forward\n", ResponsesWebSocketsForward}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -1063,13 +1066,34 @@ func TestListenerConfig_DeclineResponsesWebSockets(t *testing.T) { if err := yaml.Unmarshal([]byte(tt.yaml), &cfg); err != nil { t.Fatal(err) } - if got := cfg.Listener.DeclinesResponsesWebSockets(); got != tt.want { - t.Errorf("DeclinesResponsesWebSockets() = %v, want %v", got, tt.want) + if got := cfg.Listener.ResponsesWebSocketsMode(); got != tt.want { + t.Errorf("ResponsesWebSocketsMode() = %q, want %q", got, tt.want) } }) } } +// A misspelt mode is refused at startup rather than read as the default: an operator +// who wrote "declined" asked for something, and silently parsing would ignore them. +func TestValidate_ResponsesWebSockets(t *testing.T) { + for _, v := range []string{"", ResponsesWebSocketsParse, ResponsesWebSocketsDecline, ResponsesWebSocketsForward} { + cfg := &Config{Mode: ModeProxySidecar, Listener: ListenerConfig{ + Roles: []string{RoleForward}, + ResponsesWebSockets: v, + }} + if err := validateListeners(cfg); err != nil { + t.Errorf("%q: unexpected error %v", v, err) + } + } + cfg := &Config{Mode: ModeProxySidecar, Listener: ListenerConfig{ + Roles: []string{RoleForward}, + ResponsesWebSockets: "declined", + }} + if err := validateListeners(cfg); err == nil || !strings.Contains(err.Error(), "responses_websockets") { + t.Errorf("a misspelt mode: err = %v, want one naming listener.responses_websockets", err) + } +} + // TestLoad_PreAffinityLaptopConfigGetsClientAffinity loads a laptop config written before // session.client_affinity existed — no session block at all, which is what every install // made before it still has on disk. With affinity off, Bob Shell's header-less startup diff --git a/core/config/validate.go b/core/config/validate.go index bc24959c3..52c27c3d2 100644 --- a/core/config/validate.go +++ b/core/config/validate.go @@ -112,6 +112,13 @@ func validateListeners(cfg *Config) error { return fmt.Errorf("listener.inbound_interception: %q is not valid (use %q or %q)", cfg.Listener.InboundInterception, InboundInterceptionReverseProxy, InboundInterceptionTransparent) } + switch cfg.Listener.ResponsesWebSockets { + case "", ResponsesWebSocketsParse, ResponsesWebSocketsDecline, ResponsesWebSocketsForward: + // valid + default: + return fmt.Errorf("listener.responses_websockets: %q is not valid (use %q, %q or %q)", + cfg.Listener.ResponsesWebSockets, ResponsesWebSocketsParse, ResponsesWebSocketsDecline, ResponsesWebSocketsForward) + } roles := cfg.Listener.ActiveRoles() if cfg.Listener.InboundTransparent() && !roles[RoleReverse] { return fmt.Errorf("listener.inbound_interception: transparent requires the %q role (it selects how inbound reaches the inbound pipeline)", RoleReverse) diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 212cf163e..e1b9ae642 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -23,6 +23,7 @@ import ( "time" "errors" + "github.com/rossoctl/cortex/core/config" "github.com/rossoctl/cortex/core/listener/httpx" "github.com/rossoctl/cortex/core/listener/internal/bodycount" "github.com/rossoctl/cortex/core/listener/internal/bodyread" @@ -132,12 +133,13 @@ type Server struct { // ClientAffinity; false keeps today's resolution byte for byte. ClientAffinity bool - // DeclineResponsesWebSockets answers a WebSocket upgrade on a Responses API path - // with 426 instead of forwarding it with Upgrade stripped; see - // declineResponsesWebSocket. Set from config.ListenerConfig's - // DeclinesResponsesWebSockets, which defaults on; false keeps the forwarding byte - // for byte. - DeclineResponsesWebSockets bool + // ResponsesWebSockets is what a WebSocket upgrade on a Responses API path gets: + // config.ResponsesWebSocketsParse carries and parses it (serveResponsesWebSocket), + // Decline answers 426 (declineResponsesWebSocket), and Forward — which is also + // what "" means here, so a Server built without it behaves as one always did — + // strips the Upgrade and forwards a bare GET. cmd/cortex sets it from + // config.ListenerConfig.ResponsesWebSocketsMode, which defaults to parse. + ResponsesWebSockets string // Processes, when non-nil, names the process behind each client connection, and a // request with no session header is filed under that process's session — or its @@ -148,6 +150,11 @@ type Server struct { // resolution byte for byte. Processes peerproc.Resolver + // dialTunnel opens a CONNECT tunnel's upstream TCP connection; nil is + // net.DialTimeout. A seam for tests that address an origin by a real hostname — + // chatgpt.com, so the Responses WebSocket path applies — while it listens locally. + dialTunnel func(addr string, timeout time.Duration) (net.Conn, error) + // bufferedFallbackOnce keeps the SSE-buffered-path notice to one line per // process; the condition is a supported chain shape, not an error. bufferedFallbackOnce sync.Once @@ -488,11 +495,24 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne // Upgrade present on the inbound request is one that will be removed. pctx.RecordStrippedUpgrade(r.Header) - // After the pipeline, so a plugin's reject still wins; before any row, since this - // records its own pair. See declineResponsesWebSocket. - if !skipped && s.DeclineResponsesWebSockets && wantsResponsesWebSocket(r) { - s.declineResponsesWebSocket(w, r, pctx, tl, s.recordingSessionID(sessionID, r.Header, chain), chain) - return + // After the pipeline, so a plugin's reject still wins; before any row, since both + // arms record their own. A socket that cannot be carried is declined rather than + // forwarded, so asking for parse never makes Codex sit through the 405 retries. + if !skipped && wantsResponsesWebSocket(r) { + switch s.ResponsesWebSockets { + case config.ResponsesWebSocketsParse: + sid := s.recordingSessionID(sessionID, r.Header, chain) + if why := s.cannotCarryWebSocket(r, tl); why != "" { + slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", r.Host, "reason", why) + s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) + return + } + s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain) + return + case config.ResponsesWebSocketsDecline: + s.declineResponsesWebSocket(w, r, pctx, tl, s.recordingSessionID(sessionID, r.Header, chain), chain) + return + } } if !skipped && s.Sessions != nil { @@ -1010,59 +1030,6 @@ func (s *Server) refuseMisdirected(w http.ResponseWriter, r *http.Request, pctx // client sent is echoed back. const declinedUpgradeBody = `{"detail":"WebSocket transport is not available through Cortex; use HTTPS"}` + "\n" -// wantsResponsesWebSocket reports whether r asks to open a WebSocket on an OpenAI -// Responses API path — the one shape declineResponsesWebSocket answers. -// -// Scoped twice: to OpenAI's own hosts (isOpenAIResponsesHost), and to the Responses -// suffix inference-parser reads (/v1/responses, and Codex's -// /backend-api/codex/responses). A 426 is an answer only Codex is known to read as -// "use HTTPS", and only against OpenAI is that known: every other WebSocket keeps the -// strip and whatever its origin makes of the bare GET. -func wantsResponsesWebSocket(r *http.Request) bool { - return r.Method == http.MethodGet && - upgradesToWebSocket(r.Header) && - isOpenAIResponsesHost(r.URL.Host) && - strings.HasSuffix(r.URL.Path, "/responses") -} - -// isOpenAIResponsesHost reports whether authority is one of OpenAI's own Responses API -// hosts: chatgpt.com, which serves Codex on a ChatGPT login, and api.openai.com, which -// serves it on an API key and is the public /v1/responses. -// -// OpenAI's only, because what a 426 means to a client is only known for Codex — it is -// read from Codex's own source. Another harness speaking the same API through another -// host (a LiteLLM gateway, OpenCode Zen, a local server) may treat a 426 as a hard -// failure where today it falls back, so those keep the strip and the origin's answer. -// Exact names, compared case-insensitively without a port or a trailing dot: a -// subdomain is not the same service. -// -// authority is r.URL.Host, the address the request is dialed to — on a bridged request -// the tunnel's own, which tunnelAuthorityMatches has already held the Host header to, -// and after a redirect the redirect's — so the judgement follows where the bytes go. -func isOpenAIResponsesHost(authority string) bool { - switch strings.TrimSuffix(strings.ToLower(hostOnly(authority)), ".") { - case "chatgpt.com", "api.openai.com": - return true - } - return false -} - -// upgradesToWebSocket reports whether h's Upgrade names websocket. RFC 9110 §7.8 makes -// Upgrade a comma-separated list, matched case-insensitively and in any repeat of the -// header, so "websocket, h2c" asks for a WebSocket as surely as "websocket" does. A -// token may carry a version ("websocket/13"); only the name before it is compared. -func upgradesToWebSocket(h http.Header) bool { - for _, line := range h.Values("Upgrade") { - for _, tok := range strings.Split(line, ",") { - name, _, _ := strings.Cut(strings.TrimSpace(tok), "/") - if strings.EqualFold(name, "websocket") { - return true - } - } - } - return false -} - // declineResponsesWebSocket answers a WebSocket upgrade on a Responses API path with // 426 Upgrade Required, without contacting the origin. // @@ -2043,6 +2010,14 @@ func (s *Server) recordOutboundRejectIn(tl *tunnelLog, pctx *pipeline.Context, a s.appendOutbound(tl, sid, sessionevent.Deny(pctx, action, pipeline.Outbound)) } +// dialUpstream opens a CONNECT tunnel's upstream TCP connection to addr. +func (s *Server) dialUpstream(addr string) (net.Conn, error) { + if s.dialTunnel != nil { + return s.dialTunnel(addr, connectDialTimeout) + } + return net.DialTimeout("tcp", addr, connectDialTimeout) +} + // connectDialTimeout bounds the upstream TCP dial for a CONNECT tunnel. // Once the tunnel is open the timeout no longer applies — the agent's TLS // handshake and subsequent traffic flow at their own pace. @@ -2174,7 +2149,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // Plain TCP dial. See package-level comment on why mTLS doesn't // apply here. r.Host on a CONNECT carries "host:port" already. - upstream, err := net.DialTimeout("tcp", r.Host, connectDialTimeout) + upstream, err := s.dialUpstream(r.Host) if err != nil { slog.Warn("forward-proxy: CONNECT upstream dial failed", "host", r.Host, "error", err) http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) @@ -2229,7 +2204,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { } // fell open → re-dial for the tunnel. That tunnel is opaque, so it closes // like any other opaque one. - up2, derr := net.DialTimeout("tcp", r.Host, connectDialTimeout) + up2, derr := s.dialUpstream(r.Host) if derr != nil { // The client already has its 200, so that stays the status, and the // error says why the tunnel died. Closing the client is what makes diff --git a/core/listener/forwardproxy/upgradedeclined_test.go b/core/listener/forwardproxy/upgradedeclined_test.go index fd60edeef..ff189e436 100644 --- a/core/listener/forwardproxy/upgradedeclined_test.go +++ b/core/listener/forwardproxy/upgradedeclined_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/rossoctl/cortex/core/config" "github.com/rossoctl/cortex/core/pipeline" "github.com/rossoctl/cortex/core/plugins/plugintesting" "github.com/rossoctl/cortex/core/session" @@ -19,8 +20,8 @@ import ( // A WebSocket upgrade on a Responses API path is answered 426 by the proxy, because // the origin's own answer to the stripped request is a 405 that Codex retries for // about six seconds before falling back to HTTPS. See declineResponsesWebSocket. With -// listener.decline_responses_websockets: false the request is forwarded as before; -// upgradestripped_test.go covers that side. +// listener.responses_websockets: forward the request is forwarded as before; +// upgradestripped_test.go covers that side, and websocket_test.go covers parse. // declineHarness is a forward proxy in front of an origin that counts what reaches it. type declineHarness struct { @@ -32,20 +33,20 @@ type declineHarness struct { host string } -// newDeclineHarness builds the proxy with the decline on, as config defaults it, in -// front of an origin addressed as chatgpt.com. +// newDeclineHarness builds the proxy in decline mode, in front of an origin +// addressed as chatgpt.com. func newDeclineHarness(t *testing.T, plugins ...pipeline.Plugin) *declineHarness { - return newDeclineHarnessWith(t, true, plugins...) + return newDeclineHarnessWith(t, config.ResponsesWebSocketsDecline, plugins...) } -func newDeclineHarnessWith(t *testing.T, decline bool, plugins ...pipeline.Plugin) *declineHarness { - return newDeclineHarnessAt(t, "chatgpt.com", decline, plugins...) +func newDeclineHarnessWith(t *testing.T, mode string, plugins ...pipeline.Plugin) *declineHarness { + return newDeclineHarnessAt(t, "chatgpt.com", mode, plugins...) } // newDeclineHarnessAt addresses the origin as host: requests name host and its port, // and the proxy's own client dials the local test server whatever host it was asked // for, so the host the decline judges is the one the request really names. -func newDeclineHarnessAt(t *testing.T, host string, decline bool, plugins ...pipeline.Plugin) *declineHarness { +func newDeclineHarnessAt(t *testing.T, host string, mode string, plugins ...pipeline.Plugin) *declineHarness { t.Helper() h := &declineHarness{store: session.New(5*time.Minute, 100, 0), host: host} t.Cleanup(h.store.Close) @@ -70,7 +71,7 @@ func newDeclineHarnessAt(t *testing.T, host string, decline bool, plugins ...pip return (&net.Dialer{}).DialContext(ctx, network, originAddr) }, }}, - DeclineResponsesWebSockets: decline, + ResponsesWebSockets: mode, } proxy := httptest.NewServer(srv.Handler()) t.Cleanup(proxy.Close) @@ -146,10 +147,10 @@ func TestForwardProxy_ResponsesWebSocketIsDeclinedWith426(t *testing.T) { } } -// The off switch restores the forwarding exactly: the origin is reached and its own -// 405 is what the client gets — the behaviour before the decline existed. -func TestForwardProxy_DeclineOff_ForwardsTheUpgradeAsBefore(t *testing.T) { - h := newDeclineHarnessWith(t, false) +// Forward mode restores the forwarding exactly: the origin is reached and its own 405 +// is what the client gets — the behaviour before the decline existed. +func TestForwardProxy_ForwardMode_ForwardsTheUpgradeAsBefore(t *testing.T) { + h := newDeclineHarnessWith(t, config.ResponsesWebSocketsForward) resp, body := h.do(t, http.MethodGet, "/backend-api/codex/responses", websocketHeaders()) if resp.StatusCode != http.StatusMethodNotAllowed { @@ -167,7 +168,7 @@ func TestForwardProxy_DeclineOff_ForwardsTheUpgradeAsBefore(t *testing.T) { } for _, ev := range v.Events { if ev.Error != nil && ev.Error.Kind == "upgrade_declined" { - t.Errorf("a row says upgrade_declined with the decline off: %+v", ev) + t.Errorf("a row says upgrade_declined in forward mode: %+v", ev) } } } @@ -234,7 +235,7 @@ func TestForwardProxy_OnlyOpenAIHostsAreDeclined(t *testing.T) { } for _, tc := range cases { t.Run(tc.host, func(t *testing.T) { - h := newDeclineHarnessAt(t, tc.host, true) + h := newDeclineHarnessAt(t, tc.host, config.ResponsesWebSocketsDecline) resp, _ := h.do(t, http.MethodGet, "/v1/responses", websocketHeaders()) if tc.decline { if resp.StatusCode != http.StatusUpgradeRequired || h.reached.Load() != 0 { diff --git a/core/listener/forwardproxy/upgradestripped_test.go b/core/listener/forwardproxy/upgradestripped_test.go index bd15ec43a..332c3a60b 100644 --- a/core/listener/forwardproxy/upgradestripped_test.go +++ b/core/listener/forwardproxy/upgradestripped_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/rossoctl/cortex/core/config" "github.com/rossoctl/cortex/core/pipeline" "github.com/rossoctl/cortex/core/session" ) @@ -108,10 +109,10 @@ func TestStripForUpstream_StillRemovesTheUpgrade(t *testing.T) { // Upgrade — so a change that recorded the marker by keeping the header fails here // rather than looking like an improvement. // -// WITH THE DECLINE OFF. On a Responses path a WebSocket is otherwise answered 426 -// before anything is forwarded (upgradedeclined_test.go), so the origin would see -// nothing and the "received no Upgrade" check below would pass without testing the -// strip at all. That is why it also asserts the origin was reached. +// IN FORWARD MODE. On a Responses path a WebSocket is otherwise declined or carried +// (upgradedeclined_test.go, websocket_test.go), so the origin would never see the +// stripped GET and the "received no Upgrade" check below would pass without testing +// the strip at all. That is why it also asserts the origin was reached. func TestForwardProxy_EndToEnd_RecordsAndStripsTheUpgrade(t *testing.T) { store := session.New(5*time.Minute, 100, 0) defer store.Close() @@ -136,10 +137,10 @@ func TestForwardProxy_EndToEnd_RecordsAndStripsTheUpgrade(t *testing.T) { t.Fatal(err) } srv := &Server{ - OutboundPipeline: pipeline.NewHolder(p), - Sessions: store, - Client: http.DefaultClient, - DeclineResponsesWebSockets: false, + OutboundPipeline: pipeline.NewHolder(p), + Sessions: store, + Client: http.DefaultClient, + ResponsesWebSockets: config.ResponsesWebSocketsForward, } proxy := httptest.NewServer(srv.Handler()) defer proxy.Close() diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go new file mode 100644 index 000000000..76d8286ff --- /dev/null +++ b/core/listener/forwardproxy/websocket.go @@ -0,0 +1,699 @@ +package forwardproxy + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/rossoctl/cortex/core/listener/httpx" + "github.com/rossoctl/cortex/core/listener/internal/wsframe" + "github.com/rossoctl/cortex/core/pipeline" + "github.com/rossoctl/cortex/core/session" +) + +// A Responses API WebSocket, carried and parsed. +// +// Codex's preferred transport is a WebSocket to its /responses endpoint. One socket +// lives for a whole session and carries many turns: the client sends a text message +// {"type":"response.create", ...the POST body's fields...} and the server answers with +// one text message per Responses stream event — the same JSON the HTTPS path streams as +// SSE `data:` — ending in response.completed. Later turns send only new input, with +// previous_response_id naming the turn before. +// +// SO EACH TURN IS RUN AS IF IT WERE THE HTTPS REQUEST. The handshake GET runs the +// pipeline once, as every bridged request does, and records a tunnel-style open. Every +// response.create then gets a pipeline.Context of its own, cloned from the handshake's +// identity, with Method POST and Body the message minus "type": inference-parser +// parses it, the turn's server messages are fed to it frame by frame exactly as +// handleStreamingResponse feeds SSE events, and the turn ends — settled, recorded, +// finished — on response.completed, response.failed, an error event, or the socket +// closing. Per-turn rows are an ordinary request/response pair, so cost, usage, the +// session timeline and agentop read them like any other inference. +// +// RELAYED, NEVER RE-ENCODED, with two exceptions. Every frame is forwarded as the bytes +// it arrived as (wsframe.Frame.Raw). A plugin that rewrites a turn's request body +// (tool-prune) is honoured by sending the rewritten message as one new client frame in +// place of the original; and a turn a plugin rejects is answered to the client with an +// error event in the shape Codex already reads for a server refusal, and never sent. +// +// permessage-deflate is stripped from the handshake, so neither end compresses and the +// relay reads plain frames. Codex offers it and does not require it. + +const ( + // wsMessageLimit is the largest message read for parsing. A larger one is relayed + // whole and not read: the turn it belongs to is still recorded, just without a + // parse. Matches the forward proxy's own response-body cap. + wsMessageLimit = maxBodySize + // wsFrameLimit is the largest single frame relayed. Raw is buffered whole before + // it is forwarded, so this bounds what one frame can cost; a frame over it ends the + // socket. Matches the request-body cap, the largest thing the HTTPS path accepts. + wsFrameLimit = maxRequestBodySize + // wsHandshakeTimeout bounds the upstream handshake. Codex's own connect timeout is + // 15s; a stalled origin should not hold the client's handshake longer than that. + wsHandshakeTimeout = 15 * time.Second +) + +// wantsResponsesWebSocket reports whether r asks to open a WebSocket on OpenAI's +// Responses API — the one shape listener.responses_websockets governs. +// +// Scoped twice: to OpenAI's own hosts (isOpenAIResponsesHost), and to the Responses +// suffix inference-parser reads (/v1/responses, and Codex's +// /backend-api/codex/responses). That is the protocol this file can parse, and the one +// whose client — Codex — is known to read a 426 as "use HTTPS": every other WebSocket +// keeps the strip and whatever its origin makes of the bare GET. +func wantsResponsesWebSocket(r *http.Request) bool { + return r.Method == http.MethodGet && + upgradesToWebSocket(r.Header) && + isOpenAIResponsesHost(r.URL.Host) && + strings.HasSuffix(r.URL.Path, "/responses") +} + +// isOpenAIResponsesHost reports whether authority is one of OpenAI's own Responses API +// hosts: chatgpt.com, which serves Codex on a ChatGPT login, and api.openai.com, which +// serves it on an API key and is the public /v1/responses. +// +// OpenAI's only, because what a 426 means to a client is only known for Codex — it is +// read from Codex's own source — and the socket this file carries is Codex's. Another +// harness speaking the same API through another host (a LiteLLM gateway, OpenCode Zen, +// a local server) may treat a 426 as a hard failure where today it falls back, so those +// keep the strip and the origin's answer. Exact names, compared case-insensitively +// without a port or a trailing dot: a subdomain is not the same service. +// +// authority is r.URL.Host, the address the request is dialed to — on a bridged request +// the tunnel's own, which tunnelAuthorityMatches has already held the Host header to, +// and after a redirect the redirect's — so the judgement follows where the bytes go. +func isOpenAIResponsesHost(authority string) bool { + switch strings.TrimSuffix(strings.ToLower(hostOnly(authority)), ".") { + case "chatgpt.com", "api.openai.com": + return true + } + return false +} + +// upgradesToWebSocket reports whether h's Upgrade names websocket. RFC 9110 §7.8 makes +// Upgrade a comma-separated list, matched case-insensitively and in any repeat of the +// header, so "websocket, h2c" asks for a WebSocket as surely as "websocket" does. A +// token may carry a version ("websocket/13"); only the name before it is compared. +func upgradesToWebSocket(h http.Header) bool { + for _, line := range h.Values("Upgrade") { + for _, tok := range strings.Split(line, ",") { + name, _, _ := strings.Cut(strings.TrimSpace(tok), "/") + if strings.EqualFold(name, "websocket") { + return true + } + } + } + return false +} + +// cannotCarryWebSocket says why this socket cannot be carried and parsed, or "" when it +// can. Each reason is a shape the relay below does not handle; the caller declines +// those with a 426 rather than forwarding them, so a parse-mode client is never left +// sitting through the origin's 405. +func (s *Server) cannotCarryWebSocket(r *http.Request, tl *tunnelLog) string { + switch { + case tl == nil || s.TLSBridge == nil: + // A plaintext proxied request: no TLS bridge to re-originate over, and no + // Responses endpoint is served over plain HTTP. + return "not a bridged request" + case r.ProtoMajor != 1: + // An h2 extended CONNECT (RFC 8441) carries the socket as a stream, not a + // connection to take over. tungstenite, Codex's client, does not send one. + return "not HTTP/1.1" + case s.OutboundPipeline.WritesResponseBody(): + // A response-body writer rewrites a whole response; a turn's response is a + // live stream of frames already on their way to the client. The SSE path + // buffers instead, and a WebSocket has no buffered form. + return "a plugin rewrites response bodies" + } + // A connection that cannot be taken over is the one case left, and only the + // ResponseWriter can say; serveResponsesWebSocket declines it there. + return "" +} + +// serveResponsesWebSocket completes the upgrade with the origin, takes over the client's +// connection, and relays the socket until either end closes it, parsing every turn. +// +// pctx is the handshake's, after the pipeline ran on it. Its rows are the socket's open +// (now) and close (when it ends), sharing its RequestID the way an opaque tunnel's do; +// the turns record their own pairs between them. +func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc) { + hj, ok := w.(http.Hijacker) + if !ok { + s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) + return + } + + // The handshake as the origin should see it: the client's headers after the + // pipeline, minus what is the proxy's own — and minus the extension offer, so + // neither end compresses frames the relay has to read. + up := r.Clone(r.Context()) + syncHeaders(up, pctx.Headers) + for _, h := range []string{"Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization", "Proxy-Connection", "TE", "Trailer", "Transfer-Encoding", "Sec-Websocket-Extensions"} { + up.Header.Del(h) + } + up.Header.Set("Connection", "Upgrade") + up.Header.Set("Upgrade", "websocket") + up.RequestURI = "" + up.Body = nil + up.ContentLength = 0 + ctx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), wsHandshakeTimeout) + up = up.WithContext(ctx) + + client := s.TLSBridge.Upstream + resp, err := client.Do(up) + cancel() + if err != nil { + status, fail := pipeline.ExchangeFailure(r.Context(), err) + slog.Warn("forward-proxy: WebSocket handshake with the origin failed", "host", r.Host, "kind", fail.Kind, "error", fail.Message) + http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) + s.recordWebSocketRefusal(pctx, tl, sid, chain, status, fail) + return + } + if resp.StatusCode != http.StatusSwitchingProtocols { + // The origin refused the upgrade (401, 403, 426…): relay its answer as an + // ordinary response, so the client sees what it would have without Cortex. + defer resp.Body.Close() + for k, vs := range resp.Header { + for _, v := range vs { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + n, _ := io.Copy(w, io.LimitReader(resp.Body, maxBodySize)) + pctx.ResponseBytes = n + pctx.StatusCode = resp.StatusCode + s.recordWebSocketRefusal(pctx, tl, sid, chain, resp.StatusCode, nil) + return + } + upConn, ok := resp.Body.(io.ReadWriteCloser) + if !ok { + _ = resp.Body.Close() + http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) + s.recordWebSocketRefusal(pctx, tl, sid, chain, http.StatusBadGateway, &pipeline.EventError{Kind: "upstream_error", Message: "the origin's 101 carried no writable connection"}) + return + } + if ext := resp.Header.Get("Sec-Websocket-Extensions"); ext != "" { + // It accepted an extension nobody offered. Frames would arrive compressed and + // unreadable, so end the socket now rather than relay what cannot be parsed. + _ = upConn.Close() + http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) + s.recordWebSocketRefusal(pctx, tl, sid, chain, http.StatusBadGateway, &pipeline.EventError{Kind: "upstream_error", Message: "the origin negotiated a WebSocket extension that was not offered"}) + return + } + + clientConn, clientBuf, err := hj.Hijack() + if err != nil { + _ = upConn.Close() + slog.Warn("forward-proxy: could not take over the client connection for a WebSocket", "host", r.Host, "error", err) + return + } + // The bridge's handshake deadline was cleared before serving; nothing sets one + // again, so the socket lives as long as both ends keep it. + _ = clientConn.SetDeadline(time.Time{}) + + // The 101 as the origin sent it, minus the hop-by-hop headers the proxy owns. + var hdr bytes.Buffer + fmt.Fprintf(&hdr, "HTTP/1.1 101 Switching Protocols\r\n") + for k, vs := range resp.Header { + switch http.CanonicalHeaderKey(k) { + case "Keep-Alive", "Proxy-Connection", "Transfer-Encoding", "Content-Length": + continue + } + for _, v := range vs { + fmt.Fprintf(&hdr, "%s: %s\r\n", k, v) + } + } + hdr.WriteString("\r\n") + if _, err := clientConn.Write(hdr.Bytes()); err != nil { + _ = clientConn.Close() + _ = upConn.Close() + return + } + + pctx.StatusCode = http.StatusSwitchingProtocols + sock := &wsSocket{ + s: s, r: r, hs: pctx, tl: tl, sid: sid, chain: chain, + client: clientConn, clientR: clientBuf.Reader, up: upConn, + start: time.Now(), + } + s.recordWebSocketOpen(sock) + sock.run() +} + +// recordWebSocketRefusal records the handshake's request row and the answer it got, for +// a socket that never opened. The same pair an ordinary request records. +func (s *Server) recordWebSocketRefusal(pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc, status int, fail *pipeline.EventError) { + if s.Sessions == nil { + return + } + pctx.OutboundSessionID = sid + if chain != nil { + defer s.Sessions.TouchProcess(sid, chain) + } + s.recordOutboundRequestEvent(tl, pctx, sid) + s.recordOutboundResponseEvent(pctx, status, fail) +} + +// recordWebSocketOpen records the socket's open row: the handshake GET, marked as a +// tunnel whose reason is the WebSocket, so agentop pairs it with the close the way it +// pairs any tunnel's. +func (s *Server) recordWebSocketOpen(ws *wsSocket) { + if s.Sessions == nil { + return + } + ws.hs.OutboundSessionID = ws.sid + ev := s.tunnelOpenEvent(ws.hs, pipeline.TunnelWebSocket) + ev.HTTPMethod, ev.HTTPPath = ws.hs.Method, ws.hs.Path + s.appendOutbound(ws.tl, ws.sid, ev) +} + +// wsSocket is one carried Responses WebSocket. +type wsSocket struct { + s *Server + r *http.Request + hs *pipeline.Context // the handshake's, for the socket's own rows + tl *tunnelLog + sid string + chain []session.Proc + + client net.Conn + clientR *bufio.Reader // the hijack's buffer: bytes the server read past the request + up io.ReadWriteCloser + start time.Time + + bytesUp, bytesDown atomic.Int64 + + // clientMu serialises writes to the client: the downstream relay and a rejected + // turn's error event both write there. + clientMu sync.Mutex + // upMu does the same for the origin: the upstream relay, and a rewritten turn + // written in place of the client's frame. + upMu sync.Mutex + + // turnMu guards turn, which both relays touch: the client side opens one, the + // server side feeds and ends it. + turnMu sync.Mutex + turn *wsTurn +} + +// wsTurn is one response.create and the server messages that answer it. +type wsTurn struct { + pctx *pipeline.Context +} + +func (ws *wsSocket) run() { + done := make(chan struct{}, 2) + go func() { ws.relayUp(); done <- struct{}{} }() + go func() { ws.relayDown(); done <- struct{}{} }() + <-done + // One side is gone: close both, so the other relay's read returns. + _ = ws.client.Close() + _ = ws.up.Close() + <-done + + ws.turnMu.Lock() + t := ws.turn + ws.turn = nil + ws.turnMu.Unlock() + if t != nil { + ws.endTurn(t, http.StatusOK, &pipeline.EventError{Kind: "upstream_error", Message: "the WebSocket closed before the turn completed"}) + } + ws.recordClose() +} + +// relayUp carries client frames to the origin, reading each data message on the way. +func (ws *wsSocket) relayUp() { + fr := wsframe.NewReader(ws.clientR, wsMessageLimit, wsFrameLimit) + msg := wsframe.NewMessage(wsMessageLimit) + // held collects a data message's frames until it is whole: a turn a plugin + // rejects or rewrites must not have reached the origin first. + var held [][]byte + for { + f, err := fr.ReadFrame() + if err != nil { + ws.logEnd("client", err) + return + } + ws.bytesUp.Add(int64(len(f.Raw))) + if f.IsControl() { + if !ws.writeUp(f.Raw) { + return + } + continue + } + op, payload, readable, complete, err := msg.Add(f) + if err != nil { + ws.logEnd("client", err) + return + } + held = append(held, f.Raw) + if !complete { + continue + } + out := held + held = nil + if op == wsframe.OpText && readable { + replace, refuse := ws.clientMessage(payload) + if refuse { + continue + } + if replace != nil { + out = [][]byte{wsframe.Encode(wsframe.OpText, replace, true)} + } + } + for _, raw := range out { + if !ws.writeUp(raw) { + return + } + } + } +} + +// relayDown carries origin frames to the client, feeding each data message to the turn +// it belongs to. +func (ws *wsSocket) relayDown() { + fr := wsframe.NewReader(ws.up, wsMessageLimit, wsFrameLimit) + msg := wsframe.NewMessage(wsMessageLimit) + for { + f, err := fr.ReadFrame() + if err != nil { + ws.logEnd("origin", err) + return + } + ws.bytesDown.Add(int64(len(f.Raw))) + // Forwarded first, as handleStreamingResponse writes a block after its + // plugins have seen it only so a Reject can keep it off the wire. A + // WebSocket turn cannot be retracted once it started, so the frame goes + // now and the parse follows. + if !ws.writeClient(f.Raw) { + return + } + if f.IsControl() { + continue + } + op, payload, readable, complete, err := msg.Add(f) + if err != nil { + ws.logEnd("origin", err) + return + } + if complete && op == wsframe.OpText && readable { + ws.serverMessage(payload) + } + } +} + +func (ws *wsSocket) writeUp(b []byte) bool { + ws.upMu.Lock() + defer ws.upMu.Unlock() + _, err := ws.up.Write(b) + return err == nil +} + +func (ws *wsSocket) writeClient(b []byte) bool { + ws.clientMu.Lock() + defer ws.clientMu.Unlock() + _, err := ws.client.Write(b) + return err == nil +} + +func (ws *wsSocket) logEnd(side string, err error) { + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { + return + } + if errors.Is(err, wsframe.ErrProtocol) { + slog.Warn("forward-proxy: ending a WebSocket on a frame the relay cannot read", "host", ws.r.Host, "side", side, "error", err) + return + } + slog.Debug("forward-proxy: WebSocket ended", "host", ws.r.Host, "side", side, "error", err) +} + +// wsEnvelope is the one field every message on the socket carries. +type wsEnvelope struct { + Type string `json:"type"` +} + +// clientMessage handles one complete text message from the client. It returns a +// replacement to send instead of the original when a plugin rewrote the turn, and +// refuse when a plugin rejected it — in which case the client has already been told. +func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool) { + var env wsEnvelope + if json.Unmarshal(payload, &env) != nil || env.Type != "response.create" { + // response.interrupt, and anything else, is protocol between the two ends. + return nil, false + } + body, generate := responseCreateBody(payload) + if !generate { + // A prewarm: generate:false asks the server to set up state and answer at + // once, with no model output. Not inference, so not a turn. + return nil, false + } + + ws.turnMu.Lock() + prev := ws.turn + ws.turn = nil + ws.turnMu.Unlock() + if prev != nil { + // Codex holds the socket for one response at a time, so a second + // response.create means the first one ended without a terminal event. + ws.endTurn(prev, http.StatusOK, &pipeline.EventError{Kind: "upstream_error", Message: "a new turn began before this one completed"}) + } + + pctx := ws.newTurnContext(body) + t := &wsTurn{pctx: pctx} + action := ws.s.OutboundPipeline.Run(ws.ctx(), pctx) + if action.Type == pipeline.Reject { + ws.s.recordOutboundRejectIn(nil, pctx, action, ws.sid) + ws.refuseTurn(action) + ws.s.OutboundPipeline.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) + return nil, true + } + if ws.s.Sessions != nil { + pctx.OutboundSessionID = ws.sid + ws.s.Sessions.Append(ws.sid, ws.turnRequestEvent(pctx)) + } + ws.turnMu.Lock() + ws.turn = t + ws.turnMu.Unlock() + + if pctx.BodyMutated() { + return withType(pctx.Body), false + } + return nil, false +} + +// serverMessage feeds one complete text message from the origin to the open turn, and +// ends the turn on a terminal event. +func (ws *wsSocket) serverMessage(payload []byte) { + ws.turnMu.Lock() + t := ws.turn + ws.turnMu.Unlock() + if t == nil { + // The answer to a prewarm, or to a turn already ended. Nothing to parse into. + return + } + var env struct { + Type string `json:"type"` + Status int `json:"status"` + } + _ = json.Unmarshal(payload, &env) + if env.Type == "error" { + // Codex's wrapped error event — a refused turn (usage limit, bad input), with + // the HTTP status the HTTPS path would have answered. + status := env.Status + if status == 0 { + status = http.StatusInternalServerError + } + ws.finishTurn(t, status, &pipeline.EventError{Kind: "backend_error", Code: fmt.Sprintf("%d", status), Message: "the origin answered the turn with an error event"}) + return + } + ws.s.OutboundPipeline.RunResponseFrame(ws.ctx(), t.pctx, payload, false) + t.pctx.ResponseBytes += int64(len(payload)) + // The three events inference-parser reads usage and the finish reason from; each + // ends a response, as response.completed ends Codex's own read loop. + switch env.Type { + case "response.completed", "response.incomplete", "response.failed": + ws.finishTurn(t, http.StatusOK, nil) + } +} + +// finishTurn ends t if it is still the open turn. +func (ws *wsSocket) finishTurn(t *wsTurn, status int, fail *pipeline.EventError) { + ws.turnMu.Lock() + if ws.turn != t { + ws.turnMu.Unlock() + return + } + ws.turn = nil + ws.turnMu.Unlock() + ws.endTurn(t, status, fail) +} + +// endTurn finalizes a turn: the terminal frame that settles its cost, its response +// row, and the Finisher dispatch — the three things handleStreamingResponse's deferred +// block does at the end of an SSE stream, in the same order. +func (ws *wsSocket) endTurn(t *wsTurn, status int, fail *pipeline.EventError) { + finalCtx, cancelFinal := httpx.TeardownContext(ws.ctx()) + defer cancelFinal() + t.pctx.MarkResponseDelivered() + if fail == nil { + t.pctx.StatusCode = status + } + ws.s.OutboundPipeline.RunResponseFrame(finalCtx, t.pctx, nil, true) + ws.s.recordOutboundResponseEvent(t.pctx, status, fail) + ws.s.OutboundPipeline.RunFinish(finalCtx, t.pctx, pipeline.OutcomeFromContext(t.pctx)) +} + +// ctx is the context every turn's pipeline runs under: the handshake request's values +// (the client's process, for attribution) without its cancellation, since the handshake +// request ended when the upgrade was answered and the socket outlives it. +func (ws *wsSocket) ctx() context.Context { + return context.WithoutCancel(ws.r.Context()) +} + +// newTurnContext builds a turn's pipeline.Context: the handshake's identity — host, +// path, headers, client, session — with the request the turn would have been over +// HTTPS, a POST carrying the message's fields as its body. +func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { + hs := ws.hs + headers := ws.r.Header.Clone() + for _, h := range []string{"Upgrade", "Connection", "Sec-Websocket-Key", "Sec-Websocket-Version", "Sec-Websocket-Extensions", "Sec-Websocket-Protocol"} { + headers.Del(h) + } + headers.Set("Content-Type", "application/json") + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Method: http.MethodPost, + Scheme: hs.Scheme, + Host: hs.Host, + Path: hs.Path, + Headers: headers, + Body: body, + Shared: ws.s.Shared, + StartedAt: time.Now(), + Session: hs.Session, + SessionSource: hs.SessionSource, + SessionHeader: hs.SessionHeader, + ChainSession: hs.ChainSession, + } + pctx.ResolveClient() + return pctx +} + +// turnRequestEvent is a turn's request row: the one an HTTPS POST of the same body +// would have recorded. +func (ws *wsSocket) turnRequestEvent(pctx *pipeline.Context) pipeline.SessionEvent { + return pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Outbound, + Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), + BytesUp: int64(len(pctx.Body)), + MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), + Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), + Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), + Plugins: pipeline.SnapshotPlugins(pctx.Extensions.Custom), + Identity: pipeline.SnapshotIdentity(pctx), + Host: pctx.Host, + HTTPMethod: pctx.Method, + HTTPPath: pctx.Path, + Client: pctx.ClientInfo(), + } +} + +// refuseTurn answers a rejected turn to the client in the shape Codex reads as a server +// refusal — {"type":"error","status":…,"error":{…}} — so it surfaces as a failed turn +// with the policy's reason rather than as a broken socket. +func (ws *wsSocket) refuseTurn(action pipeline.Action) { + status, message, code := http.StatusForbidden, "request rejected by policy", "policy.forbidden" + if v := action.Violation; v != nil { + code = v.Code + status = v.Status + if status == 0 { + status = pipeline.StatusFromCode(v.Code) + } + if v.Reason != "" { + message = v.Reason + } + } + ev, _ := json.Marshal(map[string]any{ + "type": "error", + "status": status, + "error": map[string]string{"type": "cortex_policy", "code": code, "message": message}, + }) + ws.writeClient(wsframe.Encode(wsframe.OpText, ev, false)) +} + +// recordClose records the socket's close row, paired with its open by RequestID. +func (ws *wsSocket) recordClose() { + if ws.s.Sessions == nil { + return + } + ev := pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Outbound, + Phase: pipeline.SessionResponse, + RequestID: ws.hs.RequestID(), + Identity: pipeline.SnapshotIdentity(ws.hs), + Host: ws.hs.Host, + HTTPMethod: ws.hs.Method, + HTTPPath: ws.hs.Path, + StatusCode: http.StatusSwitchingProtocols, + Duration: time.Since(ws.start), + BytesUp: ws.bytesUp.Load(), + BytesDown: ws.bytesDown.Load(), + Tunnel: true, + TunnelReason: pipeline.TunnelWebSocket, + Client: ws.hs.ClientInfo(), + } + ws.s.Sessions.Append(ws.sid, ev) + if ws.chain != nil { + ws.s.Sessions.TouchProcess(ws.sid, ws.chain) + } +} + +// responseCreateBody is a response.create message as the HTTPS request body it stands +// for — the same object without "type" — and whether it asks for generation. Codex's +// prewarm sends generate:false; everything else generates. +func responseCreateBody(payload []byte) (body []byte, generate bool) { + var fields map[string]json.RawMessage + if json.Unmarshal(payload, &fields) != nil { + return payload, true + } + generate = true + if g, ok := fields["generate"]; ok && bytes.Equal(bytes.TrimSpace(g), []byte("false")) { + generate = false + } + delete(fields, "type") + delete(fields, "generate") + out, err := json.Marshal(fields) + if err != nil { + return payload, generate + } + return out, generate +} + +// withType turns a rewritten request body back into a response.create message. +func withType(body []byte) []byte { + var fields map[string]json.RawMessage + if json.Unmarshal(body, &fields) != nil { + return body + } + fields["type"] = json.RawMessage(`"response.create"`) + out, err := json.Marshal(fields) + if err != nil { + return body + } + return out +} diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go new file mode 100644 index 000000000..0b4928d60 --- /dev/null +++ b/core/listener/forwardproxy/websocket_test.go @@ -0,0 +1,708 @@ +package forwardproxy + +import ( + "bufio" + "context" + "crypto/sha1" + "crypto/tls" + "crypto/x509" + "encoding/base64" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rossoctl/cortex/core/config" + "github.com/rossoctl/cortex/core/listener/internal/wsframe" + "github.com/rossoctl/cortex/core/pipeline" + "github.com/rossoctl/cortex/core/plugins/inferenceparser" + "github.com/rossoctl/cortex/core/plugins/plugintesting" + "github.com/rossoctl/cortex/core/session" + "github.com/rossoctl/cortex/core/tlsbridge" +) + +// End to end: a Codex-shaped client opens a WebSocket through a real CONNECT bridge to +// a TLS origin that speaks the Responses WebSocket protocol, and every turn on it comes +// out of the session store parsed — model, tokens, completion — the way the same turn +// over HTTPS would. + +const wsPath = "/backend-api/codex/responses" + +// wsOrigin is a fake chatgpt.com: it accepts a WebSocket on wsPath and answers each +// response.create with a stream of Responses events, one per text message. +type wsOrigin struct { + *httptest.Server + mu sync.Mutex + // creates are the response.create payloads the origin received, in order. + creates []map[string]any + // sawExtensions is the Sec-WebSocket-Extensions offer the handshake reached it with. + sawExtensions string +} + +// wsHost is the name the origin is addressed by. It must be one of OpenAI's: the +// WebSocket path applies only there (isOpenAIResponsesHost). +const wsHost = "chatgpt.com" + +// newWSOrigin starts the origin with a certificate for wsHost, from a CA of its own — +// so the bridge's upstream check verifies it as it would OpenAI's — and returns that +// CA's PEM for the bridge to trust. +func newWSOrigin(t *testing.T) (*wsOrigin, []byte) { + t.Helper() + o := &wsOrigin{} + o.Server = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return // the bridge's upstream-verify probe + } + if r.URL.Path != wsPath || !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + o.mu.Lock() + o.sawExtensions = r.Header.Get("Sec-Websocket-Extensions") + o.mu.Unlock() + conn, brw, err := w.(http.Hijacker).Hijack() + if err != nil { + return + } + defer conn.Close() + accept := wsAccept(r.Header.Get("Sec-Websocket-Key")) + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n") + _ = brw.Flush() + o.serve(conn, brw.Reader) + })) + ca, err := tlsbridge.NewEphemeralSource() + if err != nil { + t.Fatal(err) + } + leaf, err := tlsbridge.NewMinter(ca, tlsbridge.MinterOpts{}).GetCertificateForHost(wsHost) + if err != nil { + t.Fatal(err) + } + o.TLS = &tls.Config{Certificates: []tls.Certificate{*leaf}} + o.StartTLS() + t.Cleanup(o.Close) + return o, ca.CACertPEM() +} + +// originPort is the port the origin listens on, which requests name with wsHost. +func (o *wsOrigin) originPort() string { + _, port, _ := net.SplitHostPort(o.Listener.Addr().String()) + return port +} + +// authority is how a request addresses the origin: wsHost at its port. +func (o *wsOrigin) authority() string { return net.JoinHostPort(wsHost, o.originPort()) } + +func wsAccept(key string) string { + h := sha1.Sum([]byte(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")) + return base64.StdEncoding.EncodeToString(h[:]) +} + +func (o *wsOrigin) serve(conn net.Conn, br *bufio.Reader) { + fr := wsframe.NewReader(br, 1<<20, 1<<20) + msg := wsframe.NewMessage(1 << 20) + turn := 0 + for { + f, err := fr.ReadFrame() + if err != nil { + return + } + if f.Opcode == wsframe.OpPing { + _, _ = conn.Write(wsframe.Encode(wsframe.OpPong, f.Payload, false)) + continue + } + if f.Opcode == wsframe.OpClose { + _, _ = conn.Write(wsframe.Encode(wsframe.OpClose, f.Payload, false)) + return + } + _, payload, readable, ok, _ := msg.Add(f) + if !ok || !readable { + continue + } + var m map[string]any + if json.Unmarshal(payload, &m) != nil || m["type"] != "response.create" { + continue + } + o.mu.Lock() + o.creates = append(o.creates, m) + o.mu.Unlock() + if g, ok := m["generate"].(bool); ok && !g { + // A prewarm: completed at once, no output, no usage. + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.completed","response":{"id":"warm","status":"completed"}}`), false)) + continue + } + turn++ + for _, ev := range turnEvents(turn) { + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) + } + } +} + +// turnEvents is the event sequence live Codex traffic carries for a one-word reply, +// with per-turn usage so two turns are told apart. +func turnEvents(n int) []string { + id := map[int]string{1: "resp_1", 2: "resp_2"}[n] + input := map[int]int{1: 12862, 2: 13002}[n] + return []string{ + `{"type":"response.created","response":{"id":"` + id + `","status":"in_progress"}}`, + `{"type":"codex.rate_limits","primary":{"used_percent":1.0}}`, + `{"type":"response.output_item.added","output_index":0,"item":{"id":"msg","type":"message"}}`, + `{"type":"response.output_text.delta","item_id":"msg","output_index":0,"content_index":0,"delta":"po"}`, + `{"type":"response.output_text.delta","item_id":"msg","output_index":0,"content_index":0,"delta":"ng"}`, + `{"type":"response.output_item.done","output_index":0,"item":{"id":"msg","type":"message"}}`, + `{"type":"response.completed","response":{"id":"` + id + `","status":"completed","usage":{` + + `"input_tokens":` + itoa(input) + `,"input_tokens_details":{"cached_tokens":11008},` + + `"output_tokens":5,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":` + itoa(input+5) + `}}}`, + } +} + +func itoa(n int) string { b, _ := json.Marshal(n); return string(b) } + +// wsHarness is a proxy with the TLS bridge on, in front of a wsOrigin. +type wsHarness struct { + store *session.Store + origin *wsOrigin + proxy *httptest.Server + caPEM []byte +} + +func newWSHarness(t *testing.T, mode string, plugins ...pipeline.Plugin) *wsHarness { + t.Helper() + origin, originCAPEM := newWSOrigin(t) + h := &wsHarness{store: session.New(5*time.Minute, 1000, 0), origin: origin} + t.Cleanup(h.store.Close) + + src, err := tlsbridge.NewEphemeralSource() + if err != nil { + t.Fatal(err) + } + up, err := tlsbridge.NewUpstreamClient(originCAPEM, false) + if err != nil { + t.Fatal(err) + } + // The bridge's upstream client resolves wsHost to the local origin, as DNS would + // resolve chatgpt.com to OpenAI's. + originAddr := h.origin.Listener.Addr().String() + toOrigin := func(ctx context.Context, network, _ string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, network, originAddr) + } + up.Transport.(*http.Transport).DialContext = toOrigin + engine := &tlsbridge.Engine{ + Decision: mustBridgeDecision(t, portOf(h.origin.authority())), + Term: tlsbridge.NewTerminator(tlsbridge.NewMinter(src, tlsbridge.MinterOpts{})), + Skip: tlsbridge.NewSkipSet(), + Upstream: up, + CAPEM: src.CACertPEM(), + } + h.caPEM = engine.CAPEM + + all := append([]pipeline.Plugin{inferenceparser.NewInferenceParser()}, plugins...) + p, err := plugintesting.BuildPipeline(all) + if err != nil { + t.Fatal(err) + } + srv := &Server{ + OutboundPipeline: pipeline.NewHolder(p), + Sessions: h.store, + Client: http.DefaultClient, + TLSBridge: engine, + // Codex's session header, as a laptop install reads it. + SessionIDHeaders: []string{session.CodexSessionHeader}, + ResponsesWebSockets: mode, + dialTunnel: func(string, time.Duration) (net.Conn, error) { + return net.DialTimeout("tcp", originAddr, 5*time.Second) + }, + } + h.proxy = httptest.NewServer(srv.Handler()) + t.Cleanup(h.proxy.Close) + return h +} + +// wsClient is the Codex end: a WebSocket over a CONNECT bridge through the proxy. +type wsClient struct { + t *testing.T + conn *tls.Conn + fr *wsframe.Reader + msg *wsframe.Message + status int + header http.Header + body string +} + +// dial opens the CONNECT, the TLS session against the bridge's forged leaf, and sends +// the WebSocket handshake with the headers Codex sends. A non-101 answer is kept on the +// client for the test to read. +func (h *wsHarness) dial(t *testing.T, extra http.Header) *wsClient { + t.Helper() + pu, _ := url.Parse(h.proxy.URL) + authority := h.origin.authority() + raw, err := net.Dial("tcp", pu.Host) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = raw.Close() }) + _ = raw.SetDeadline(time.Now().Add(15 * time.Second)) + connect, _ := http.NewRequest(http.MethodConnect, "//"+authority, nil) + connect.Host, connect.URL.Host = authority, authority + if err := connect.Write(raw); err != nil { + t.Fatal(err) + } + br := bufio.NewReader(raw) + resp, err := http.ReadResponse(br, connect) + if err != nil || resp.StatusCode != http.StatusOK { + t.Fatalf("CONNECT: %v %v", resp, err) + } + pool := x509.NewCertPool() + pool.AppendCertsFromPEM(h.caPEM) + conn := tls.Client(&bufferedConn{Conn: raw, r: br}, &tls.Config{ + ServerName: wsHost, + RootCAs: pool, + NextProtos: []string{"http/1.1"}, + }) + if err := conn.Handshake(); err != nil { + t.Fatalf("TLS through the bridge: %v", err) + } + + req, _ := http.NewRequest(http.MethodGet, "https://"+authority+wsPath, nil) + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Sec-WebSocket-Version", "13") + req.Header.Set("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") + req.Header.Set("Sec-WebSocket-Extensions", "permessage-deflate; client_max_window_bits") + req.Header.Set("OpenAI-Beta", "responses_websockets=2026-02-06") + req.Header.Set("Session-Id", "11111111-2222-3333-4444-555555555555") + req.Header.Set("User-Agent", "codex_exec/0.163.0 (Mac OS 26.5.1; arm64) unknown") + for k, v := range extra { + req.Header[k] = v + } + if err := req.Write(conn); err != nil { + t.Fatal(err) + } + cbr := bufio.NewReader(conn) + hs, err := http.ReadResponse(cbr, req) + if err != nil { + t.Fatalf("reading the handshake answer: %v", err) + } + c := &wsClient{t: t, conn: conn, status: hs.StatusCode, header: hs.Header} + if hs.StatusCode != http.StatusSwitchingProtocols { + b, _ := io.ReadAll(hs.Body) + c.body = string(b) + return c + } + c.fr = wsframe.NewReader(cbr, 1<<20, 1<<20) + c.msg = wsframe.NewMessage(1 << 20) + return c +} + +func (c *wsClient) send(v string) { + c.t.Helper() + if _, err := c.conn.Write(wsframe.Encode(wsframe.OpText, []byte(v), true)); err != nil { + c.t.Fatal(err) + } +} + +// readUntil reads server messages until one of the given type arrives, returning all +// of them. +func (c *wsClient) readUntil(typ string) []map[string]any { + c.t.Helper() + var out []map[string]any + for { + f, err := c.fr.ReadFrame() + if err != nil { + c.t.Fatalf("reading from the socket: %v (after %d messages)", err, len(out)) + } + _, payload, readable, ok, _ := c.msg.Add(f) + if !ok || !readable { + continue + } + var m map[string]any + _ = json.Unmarshal(payload, &m) + out = append(out, m) + if m["type"] == typ { + return out + } + } +} + +func (c *wsClient) close() { + _, _ = c.conn.Write(wsframe.Encode(wsframe.OpClose, wsframe.ClosePayload(1000, ""), true)) + _ = c.conn.Close() +} + +func responseCreate(text string) string { + return `{"type":"response.create","model":"gpt-6-luna","stream":true,"input":[` + + `{"type":"message","role":"user","content":[{"type":"input_text","text":"` + text + `"}]}]}` +} + +// events reads the session's rows once the socket has closed and recorded its close. +func (h *wsHarness) events(t *testing.T, sid string, wantClose bool) []pipeline.SessionEvent { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for { + v := h.store.View(sid) + if v != nil { + if !wantClose { + return v.Events + } + for _, e := range v.Events { + if e.Tunnel && e.Phase == pipeline.SessionResponse { + return v.Events + } + } + } + if time.Now().After(deadline) { + if v == nil { + t.Fatalf("no session %q", sid) + } + return v.Events + } + time.Sleep(20 * time.Millisecond) + } +} + +const codexSID = "11111111-2222-3333-4444-555555555555" + +// The headline: two turns and a prewarm on one socket become two parsed inference +// pairs, between the socket's own open and close, all in Codex's session. +func TestResponsesWebSocket_TurnsAreParsedLikeHTTPS(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d (%s), want 101", c.status, c.body) + } + + c.send(`{"type":"response.create","model":"gpt-6-luna","stream":true,"generate":false,"input":[]}`) + c.readUntil("response.completed") + c.send(responseCreate("hello")) + first := c.readUntil("response.completed") + c.send(responseCreate("again")) + c.readUntil("response.completed") + c.close() + + // The client sees every event the origin sent, unchanged, codex's own included. + var sawRateLimits bool + for _, m := range first { + if m["type"] == "codex.rate_limits" { + sawRateLimits = true + } + } + if !sawRateLimits { + t.Error("the relay dropped a codex.rate_limits event") + } + + evs := h.events(t, codexSID, true) + // The bridged CONNECT records its own open, folded by agentop into the first + // request it carried — here the handshake. The socket's rows are the ones with the + // websocket reason. + var turns, opens, closes []pipeline.SessionEvent + for _, e := range evs { + switch { + case e.Tunnel && e.TunnelReason == pipeline.TunnelWebSocket && e.Phase == pipeline.SessionRequest: + opens = append(opens, e) + case e.Tunnel && e.TunnelReason == pipeline.TunnelWebSocket && e.Phase == pipeline.SessionResponse: + closes = append(closes, e) + case e.HTTPMethod == http.MethodPost: + turns = append(turns, e) + } + } + if len(opens) != 1 || opens[0].TunnelReason != pipeline.TunnelWebSocket || opens[0].HTTPMethod != http.MethodGet { + t.Fatalf("socket open rows = %+v, want one GET tunnel row with reason websocket", opens) + } + if len(closes) != 1 || closes[0].RequestID != opens[0].RequestID || closes[0].BytesUp == 0 || closes[0].BytesDown == 0 { + t.Fatalf("socket close rows = %+v, want one paired with the open, carrying the bytes", closes) + } + // agentop renders upgradeStripped as "websocket not forwarded"; on a socket that was + // carried that line would be false. + for _, e := range evs { + if e.UpgradeStripped != "" { + t.Errorf("row %d (%v %s) carries upgradeStripped=%q on a carried socket", e.Seq, e.Phase, e.HTTPMethod, e.UpgradeStripped) + } + } + // Two generating turns; the prewarm is not inference and records nothing. + if len(turns) != 4 { + for _, e := range evs { + t.Logf(" %v %s %s status=%d tunnel=%v", e.Phase, e.HTTPMethod, e.HTTPPath, e.StatusCode, e.Tunnel) + } + t.Fatalf("got %d turn rows, want 4 (two request/response pairs)", len(turns)) + } + wantInput := []int{12862, 13002} + for i := 0; i < 2; i++ { + req, res := turns[2*i], turns[2*i+1] + if req.Phase != pipeline.SessionRequest || res.Phase != pipeline.SessionResponse || req.RequestID != res.RequestID { + t.Errorf("turn %d rows do not pair: %+v / %+v", i+1, req, res) + continue + } + if req.Inference == nil || req.Inference.Model != "gpt-6-luna" || len(req.Inference.Messages) != 1 { + t.Errorf("turn %d request parse = %+v, want model gpt-6-luna and its one message", i+1, req.Inference) + } + if res.StatusCode != http.StatusOK || res.Inference == nil { + t.Errorf("turn %d response = status %d inference %+v", i+1, res.StatusCode, res.Inference) + continue + } + inf := res.Inference + // Cached tokens are subtracted from the prompt, as on the HTTPS path. + if inf.Completion != "pong" || inf.CompletionTokens != 5 || inf.TotalTokens != wantInput[i]+5 { + t.Errorf("turn %d response parse: completion=%q out=%d total=%d", i+1, inf.Completion, inf.CompletionTokens, inf.TotalTokens) + } + } + if turns[0].RequestID == turns[2].RequestID { + t.Error("the two turns share a RequestID; each needs its own") + } + // Turns sit between the socket's open and close in the timeline. + if !(opens[0].Seq < turns[0].Seq && turns[3].Seq < closes[0].Seq) { + t.Errorf("order: open %d, turns %d..%d, close %d", opens[0].Seq, turns[0].Seq, turns[3].Seq, closes[0].Seq) + } +} + +// Compression would make every frame unreadable, so the offer never reaches the origin. +func TestResponsesWebSocket_DeflateOfferIsStripped(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d", c.status) + } + c.close() + h.origin.mu.Lock() + defer h.origin.mu.Unlock() + if h.origin.sawExtensions != "" { + t.Errorf("origin saw Sec-WebSocket-Extensions %q; the offer must be stripped", h.origin.sawExtensions) + } +} + +// A turn a plugin rejects never reaches the origin, and the client is told in the +// error-event shape Codex reads as a server refusal. +func TestResponsesWebSocket_RejectedTurnIsAnsweredAndNotForwarded(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &turnDenier{}) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d", c.status) + } + c.send(responseCreate("forbidden")) + msgs := c.readUntil("error") + errEv := msgs[len(msgs)-1] + if errEv["status"] != float64(http.StatusForbidden) { + t.Errorf("error event status = %v, want 403", errEv["status"]) + } + c.close() + + h.origin.mu.Lock() + n := len(h.origin.creates) + h.origin.mu.Unlock() + if n != 0 { + t.Errorf("origin received %d response.create; a rejected turn must not be forwarded", n) + } + var denied bool + for _, e := range h.events(t, codexSID, true) { + if e.Phase == pipeline.SessionDenied { + denied = true + } + } + if !denied { + t.Error("no denied row recorded for the rejected turn") + } +} + +// A request-body writer's rewrite is what the origin receives. +func TestResponsesWebSocket_RewrittenTurnIsWhatTheOriginGets(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &modelRewriter{}) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d", c.status) + } + c.send(responseCreate("hello")) + c.readUntil("response.completed") + c.close() + + h.origin.mu.Lock() + defer h.origin.mu.Unlock() + if len(h.origin.creates) != 1 { + t.Fatalf("origin got %d messages, want 1", len(h.origin.creates)) + } + got := h.origin.creates[0] + if got["model"] != "rewritten-model" || got["type"] != "response.create" { + t.Errorf("origin got %v, want the rewritten model on a response.create", got) + } +} + +// A socket the proxy cannot carry is declined, never forwarded: with a response-body +// writer in the chain the turn's live frames could not be rewritten. +func TestResponsesWebSocket_UncarriableSocketIsDeclined(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &responseRewriter{}) + c := h.dial(t, nil) + if c.status != http.StatusUpgradeRequired { + t.Fatalf("handshake answered %d, want 426", c.status) + } +} + +// The same handshake to a host that is not OpenAI's is neither carried nor declined: +// it is forwarded with Upgrade stripped, as every WebSocket was, and gets its origin's +// own answer. What a 426 means is only known for Codex against OpenAI. +func TestResponsesWebSocket_OtherHostsAreForwardedAsBefore(t *testing.T) { + store := session.New(5*time.Minute, 100, 0) + t.Cleanup(store.Close) + var reached atomic.Int32 + var sawUpgrade atomic.Value + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + reached.Add(1) + sawUpgrade.Store(r.Header.Get("Upgrade")) + w.WriteHeader(http.StatusMethodNotAllowed) + })) + t.Cleanup(origin.Close) + p, err := plugintesting.BuildPipeline([]pipeline.Plugin{inferenceparser.NewInferenceParser()}) + if err != nil { + t.Fatal(err) + } + srv := &Server{ + OutboundPipeline: pipeline.NewHolder(p), + Sessions: store, + Client: http.DefaultClient, + ResponsesWebSockets: config.ResponsesWebSocketsParse, + } + proxy := httptest.NewServer(srv.Handler()) + t.Cleanup(proxy.Close) + pu, _ := url.Parse(proxy.URL) + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(pu)}} + + req, _ := http.NewRequest(http.MethodGet, origin.URL+"/v1/responses", nil) + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusMethodNotAllowed || reached.Load() != 1 { + t.Errorf("status %d, origin reached %d times; want the origin's own 405", resp.StatusCode, reached.Load()) + } + if got, _ := sawUpgrade.Load().(string); got != "" { + t.Errorf("origin received Upgrade %q; it must still be stripped", got) + } +} + +// Decline mode still answers 426 through the bridge, as #1462 shipped it. +func TestResponsesWebSocket_DeclineModeAnswers426(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsDecline) + c := h.dial(t, nil) + if c.status != http.StatusUpgradeRequired { + t.Fatalf("handshake answered %d, want 426", c.status) + } +} + +// The origin refusing the upgrade is relayed as its own answer. +func TestResponsesWebSocket_OriginRefusalIsRelayed(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + w.WriteHeader(http.StatusUnauthorized) + _, _ = io.WriteString(w, `{"detail":"unauthorized"}`) + }) + c := h.dial(t, nil) + if c.status != http.StatusUnauthorized || !strings.Contains(c.body, "unauthorized") { + t.Fatalf("handshake answered %d %q, want the origin's 401", c.status, c.body) + } +} + +// A client that hangs up mid-turn still gets that turn recorded, with what had +// arrived, and the socket's close. +func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + // An origin that answers a create with a partial stream and then stalls. + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { + return + } + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.created","response":{"id":"r","status":"in_progress"}}`), false)) + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.output_text.delta","item_id":"m","output_index":0,"content_index":0,"delta":"par"}`), false)) + _, _ = fr.ReadFrame() // until the proxy closes us + }) + c := h.dial(t, nil) + c.send(responseCreate("hello")) + c.readUntil("response.output_text.delta") + _ = c.conn.Close() + + var res *pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionResponse { + e := e + res = &e + } + } + if res == nil { + t.Fatal("the interrupted turn recorded no response row") + } + if res.Error == nil || res.Inference == nil || res.Inference.Completion != "par" { + t.Errorf("interrupted turn: error=%+v completion=%v, want an error and the partial text", res.Error, res.Inference) + } +} + +// --- test plugins --- + +// turnDenier rejects every inference request, and nothing else, so the handshake GET +// passes and the turn is refused. +type turnDenier struct{} + +func (turnDenier) Name() string { return "turn-denier" } +func (turnDenier) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{} } +func (turnDenier) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + if pctx.Extensions.Inference == nil { + return pipeline.Action{Type: pipeline.Continue} + } + return pctx.DenyAndRecord("test_deny", "policy.forbidden", "denied by the test") +} +func (turnDenier) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// modelRewriter rewrites an inference request's model, as a request-body writer can. +type modelRewriter struct{} + +func (modelRewriter) Name() string { return "model-rewriter" } +func (modelRewriter) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ReadsBody: true, WritesRequestBody: true} +} +func (modelRewriter) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + if pctx.Extensions.Inference == nil { + return pipeline.Action{Type: pipeline.Continue} + } + var m map[string]any + if json.Unmarshal(pctx.Body, &m) == nil { + m["model"] = "rewritten-model" + b, _ := json.Marshal(m) + pctx.SetBody(b) + } + return pipeline.Action{Type: pipeline.Continue} +} +func (modelRewriter) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// responseRewriter declares a response-body write, which a live socket cannot honour. +type responseRewriter struct{} + +func (responseRewriter) Name() string { return "response-rewriter" } +func (responseRewriter) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ReadsBody: true, WritesResponseBody: true} +} +func (responseRewriter) OnRequest(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (responseRewriter) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} diff --git a/core/listener/internal/wsframe/wsframe.go b/core/listener/internal/wsframe/wsframe.go new file mode 100644 index 000000000..261e7e0d3 --- /dev/null +++ b/core/listener/internal/wsframe/wsframe.go @@ -0,0 +1,260 @@ +// Package wsframe reads and writes RFC 6455 WebSocket frames for a proxy that relays a +// connection it also needs to read. +// +// RELAY FIRST, READ SECOND. A Frame keeps the exact bytes it arrived as (Raw), so the +// caller forwards those and never re-encodes what it relays: a re-encode could only +// change what the peer receives. Payload is a separate, unmasked copy for reading. +// +// What it does not do is deliberate. It negotiates no extensions — the caller strips +// Sec-WebSocket-Extensions from the handshake, so a frame with RSV1 set is a protocol +// error here rather than something to inflate — and it does not enforce RFC 6455's +// close handshake, which is the two endpoints' business; it only reads close frames so +// the caller can see one go by. +package wsframe + +import ( + "crypto/rand" + "encoding/binary" + "errors" + "fmt" + "io" +) + +// Opcodes (RFC 6455 §5.2). +const ( + OpContinuation byte = 0x0 + OpText byte = 0x1 + OpBinary byte = 0x2 + OpClose byte = 0x8 + OpPing byte = 0x9 + OpPong byte = 0xA +) + +// Close codes this package writes (RFC 6455 §7.4.1). +const ( + ClosePolicyViolation = 1008 + CloseInternalError = 1011 +) + +// maxControlPayload is RFC 6455 §5.5's cap on a control frame's payload. +const maxControlPayload = 125 + +// ErrProtocol is a frame no conforming endpoint sends. The caller should stop relaying: +// the two ends no longer agree on where frames begin. +var ErrProtocol = errors.New("wsframe: protocol error") + +// Frame is one frame as read off the wire. +type Frame struct { + Fin bool + Opcode byte + Masked bool + // Raw is the frame exactly as it arrived, header and (still masked) payload. + // Forward this, not a re-encoding. + Raw []byte + // Payload is the unmasked payload, or nil when the frame is over the reader's + // limit: such a frame is still relayed whole, but not read. + Payload []byte + // Truncated reports that Payload is nil because the frame exceeded the limit. + Truncated bool +} + +// IsControl reports whether f is a close, ping or pong frame. +func (f *Frame) IsControl() bool { return f.Opcode&0x8 != 0 } + +// Reader reads frames from one direction of a connection. +type Reader struct { + r io.Reader + // limit is the largest payload kept for reading. A larger frame is read and + // relayed in full but its Payload is dropped, so one huge message costs the + // proxy no more than limit bytes of reading memory... beyond the Raw copy the + // relay needs anyway. + limit int64 + // maxFrame bounds Raw: a frame over it is refused, since relaying it would mean + // buffering it whole first. + maxFrame int64 +} + +// NewReader reads frames from r, keeping payloads up to limit bytes for reading and +// refusing any single frame over maxFrame bytes. +func NewReader(r io.Reader, limit, maxFrame int64) *Reader { + return &Reader{r: r, limit: limit, maxFrame: maxFrame} +} + +// ReadFrame reads the next frame. It returns io.EOF only when the connection closed +// cleanly between frames; a close mid-frame is io.ErrUnexpectedEOF. +func (r *Reader) ReadFrame() (*Frame, error) { + var hdr [14]byte + if _, err := io.ReadFull(r.r, hdr[:2]); err != nil { + if errors.Is(err, io.ErrUnexpectedEOF) { + return nil, io.ErrUnexpectedEOF + } + return nil, err + } + b0, b1 := hdr[0], hdr[1] + f := &Frame{Fin: b0&0x80 != 0, Opcode: b0 & 0x0F, Masked: b1&0x80 != 0} + if b0&0x70 != 0 { + // RSV1-3. RSV1 is permessage-deflate's, which the handshake strip should have + // kept off; the other two have no meaning without an extension either. + return nil, fmt.Errorf("%w: reserved bits set (0x%02x) — an extension the proxy did not negotiate", ErrProtocol, b0&0x70) + } + switch f.Opcode { + case OpContinuation, OpText, OpBinary, OpClose, OpPing, OpPong: + default: + return nil, fmt.Errorf("%w: unknown opcode 0x%x", ErrProtocol, f.Opcode) + } + + n := 2 + length := int64(b1 & 0x7F) + switch length { + case 126: + if _, err := io.ReadFull(r.r, hdr[n:n+2]); err != nil { + return nil, unexpected(err) + } + length = int64(binary.BigEndian.Uint16(hdr[n : n+2])) + n += 2 + case 127: + if _, err := io.ReadFull(r.r, hdr[n:n+8]); err != nil { + return nil, unexpected(err) + } + u := binary.BigEndian.Uint64(hdr[n : n+8]) + if u>>63 != 0 { + return nil, fmt.Errorf("%w: 64-bit length with the high bit set", ErrProtocol) + } + length = int64(u) + n += 8 + } + if f.IsControl() && (length > maxControlPayload || !f.Fin) { + return nil, fmt.Errorf("%w: control frame fragmented or over %d bytes", ErrProtocol, maxControlPayload) + } + var key [4]byte + if f.Masked { + if _, err := io.ReadFull(r.r, hdr[n:n+4]); err != nil { + return nil, unexpected(err) + } + copy(key[:], hdr[n:n+4]) + n += 4 + } + if length > r.maxFrame-int64(n) { + return nil, fmt.Errorf("%w: a %d-byte frame is over the relay's %d-byte limit", ErrProtocol, length, r.maxFrame) + } + + f.Raw = make([]byte, n+int(length)) + copy(f.Raw, hdr[:n]) + if _, err := io.ReadFull(r.r, f.Raw[n:]); err != nil { + return nil, unexpected(err) + } + if length > r.limit { + f.Truncated = true + return f, nil + } + f.Payload = append([]byte(nil), f.Raw[n:]...) + if f.Masked { + maskBytes(f.Payload, key) + } + return f, nil +} + +func unexpected(err error) error { + if errors.Is(err, io.EOF) { + return io.ErrUnexpectedEOF + } + return err +} + +// maskBytes XORs b with key in place — masking and unmasking are the same operation. +func maskBytes(b []byte, key [4]byte) { + for i := range b { + b[i] ^= key[i&3] + } +} + +// Encode builds one unfragmented frame. A client must mask what it sends and a server +// must not (RFC 6455 §5.1), so the caller says which side it is writing as. +func Encode(opcode byte, payload []byte, masked bool) []byte { + length := len(payload) + hdr := make([]byte, 0, 14+length) + hdr = append(hdr, 0x80|opcode) + var maskBit byte + if masked { + maskBit = 0x80 + } + switch { + case length <= 125: + hdr = append(hdr, maskBit|byte(length)) + case length <= 0xFFFF: + hdr = append(hdr, maskBit|126) + hdr = binary.BigEndian.AppendUint16(hdr, uint16(length)) + default: + hdr = append(hdr, maskBit|127) + hdr = binary.BigEndian.AppendUint64(hdr, uint64(length)) + } + if !masked { + return append(hdr, payload...) + } + var key [4]byte + _, _ = rand.Read(key[:]) + hdr = append(hdr, key[:]...) + start := len(hdr) + hdr = append(hdr, payload...) + maskBytes(hdr[start:], key) + return hdr +} + +// ClosePayload is a close frame's payload: the two-byte code, then the reason. The +// reason is cut to fit a control frame. +func ClosePayload(code uint16, reason string) []byte { + if len(reason) > maxControlPayload-2 { + reason = reason[:maxControlPayload-2] + } + b := binary.BigEndian.AppendUint16(nil, code) + return append(b, reason...) +} + +// Message reassembles a data message from its frames. +// +// It sees control frames go by without letting them interrupt a message, as RFC 6455 +// §5.4 allows a peer to interleave them between fragments. It reports a message whole +// once its final fragment arrives, and only if every fragment was readable and the total +// stayed within limit; otherwise the message is still relayed, just not read. +type Message struct { + limit int64 + opcode byte + buf []byte + inFlight bool + // lost is set when a fragment of the current message was over the reader's + // limit or the total grew past this one's. + lost bool +} + +// NewMessage reassembles messages of up to limit bytes. +func NewMessage(limit int64) *Message { return &Message{limit: limit} } + +// Add takes the next data frame and, when it completes a message, returns that +// message's opcode and payload with ok true. A message that could not be read returns +// ok true and a nil payload with readable false, so the caller can still count it. +func (m *Message) Add(f *Frame) (opcode byte, payload []byte, readable, ok bool, err error) { + if f.IsControl() { + return 0, nil, false, false, nil + } + switch { + case f.Opcode == OpContinuation && !m.inFlight: + return 0, nil, false, false, fmt.Errorf("%w: continuation frame with no message in progress", ErrProtocol) + case f.Opcode != OpContinuation && m.inFlight: + return 0, nil, false, false, fmt.Errorf("%w: a new message began before the last one finished", ErrProtocol) + case f.Opcode != OpContinuation: + m.opcode, m.buf, m.lost, m.inFlight = f.Opcode, m.buf[:0], false, true + } + if f.Truncated || int64(len(m.buf))+int64(len(f.Payload)) > m.limit { + m.lost = true + } else if !m.lost { + m.buf = append(m.buf, f.Payload...) + } + if !f.Fin { + return 0, nil, false, false, nil + } + m.inFlight = false + if m.lost { + return m.opcode, nil, false, true, nil + } + return m.opcode, append([]byte(nil), m.buf...), true, true, nil +} diff --git a/core/listener/internal/wsframe/wsframe_test.go b/core/listener/internal/wsframe/wsframe_test.go new file mode 100644 index 000000000..9f3497ba0 --- /dev/null +++ b/core/listener/internal/wsframe/wsframe_test.go @@ -0,0 +1,216 @@ +package wsframe + +import ( + "bytes" + "errors" + "io" + "strings" + "testing" +) + +func readOne(t *testing.T, raw []byte, limit int64) *Frame { + t.Helper() + f, err := NewReader(bytes.NewReader(raw), limit, 1<<20).ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + return f +} + +// RFC 6455 §5.7's own examples, so the reader agrees with the spec rather than only +// with this package's encoder. +func TestReadFrame_RFC6455Examples(t *testing.T) { + cases := []struct { + name string + raw []byte + opcode byte + fin bool + want string + }{ + {"unmasked text", []byte{0x81, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f}, OpText, true, "Hello"}, + {"masked text", []byte{0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58}, OpText, true, "Hello"}, + {"first fragment", []byte{0x01, 0x03, 0x48, 0x65, 0x6c}, OpText, false, "Hel"}, + {"last fragment", []byte{0x80, 0x02, 0x6c, 0x6f}, OpContinuation, true, "lo"}, + {"unmasked ping", []byte{0x89, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f}, OpPing, true, "Hello"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + f := readOne(t, tc.raw, 1024) + if f.Opcode != tc.opcode || f.Fin != tc.fin || string(f.Payload) != tc.want { + t.Errorf("got opcode=%x fin=%v payload=%q", f.Opcode, f.Fin, f.Payload) + } + // The relay forwards Raw, so it must be the bytes exactly as sent. + if !bytes.Equal(f.Raw, tc.raw) { + t.Errorf("Raw = % x, want the input % x", f.Raw, tc.raw) + } + }) + } +} + +// The 16- and 64-bit length forms, both directions of masking. +func TestEncodeReadRoundTrip_AllLengthForms(t *testing.T) { + for _, size := range []int{0, 125, 126, 65535, 65536, 200_000} { + for _, masked := range []bool{false, true} { + payload := bytes.Repeat([]byte("x"), size) + raw := Encode(OpText, payload, masked) + f := readOne(t, raw, 1<<20) + if !bytes.Equal(f.Payload, payload) || f.Masked != masked || !bytes.Equal(f.Raw, raw) { + t.Errorf("size %d masked %v: round trip failed (len %d)", size, masked, len(f.Payload)) + } + } + } +} + +// A frame over the reading limit is still read whole — the relay needs its bytes — but +// its payload is not kept. +func TestReadFrame_OverLimitIsRelayedNotRead(t *testing.T) { + raw := Encode(OpText, bytes.Repeat([]byte("y"), 300), true) + f := readOne(t, raw, 100) + if !f.Truncated || f.Payload != nil { + t.Errorf("Truncated=%v Payload len=%d, want truncated with no payload", f.Truncated, len(f.Payload)) + } + if !bytes.Equal(f.Raw, raw) { + t.Error("Raw must still carry the whole frame") + } +} + +func TestReadFrame_RefusesWhatNoEndpointSends(t *testing.T) { + cases := map[string][]byte{ + "RSV1 (permessage-deflate)": {0xC1, 0x00}, + "unknown opcode": {0x83, 0x00}, + "fragmented ping": {0x09, 0x00}, + "ping over 125 bytes": append([]byte{0x89, 126, 0x00, 126}, bytes.Repeat([]byte("p"), 126)...), + "64-bit length high bit": {0x81, 127, 0x80, 0, 0, 0, 0, 0, 0, 0}, + } + for name, raw := range cases { + t.Run(name, func(t *testing.T) { + _, err := NewReader(bytes.NewReader(raw), 1024, 1<<20).ReadFrame() + if !errors.Is(err, ErrProtocol) { + t.Errorf("err = %v, want ErrProtocol", err) + } + }) + } +} + +func TestReadFrame_RefusesAFrameOverTheRelayLimit(t *testing.T) { + raw := Encode(OpBinary, make([]byte, 5000), false) + _, err := NewReader(bytes.NewReader(raw), 1024, 4096).ReadFrame() + if !errors.Is(err, ErrProtocol) { + t.Errorf("err = %v, want ErrProtocol", err) + } +} + +// A clean close between frames is EOF; one partway through a frame is not, so the +// caller can tell a finished connection from a torn one. +func TestReadFrame_EOFShapes(t *testing.T) { + if _, err := NewReader(bytes.NewReader(nil), 10, 1<<20).ReadFrame(); err != io.EOF { + t.Errorf("empty stream: err = %v, want io.EOF", err) + } + torn := Encode(OpText, []byte("hello"), true)[:5] + if _, err := NewReader(bytes.NewReader(torn), 10, 1<<20).ReadFrame(); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Errorf("torn frame: err = %v, want io.ErrUnexpectedEOF", err) + } +} + +func messageFrames(t *testing.T, raws ...[]byte) []*Frame { + t.Helper() + r := NewReader(bytes.NewReader(bytes.Join(raws, nil)), 1<<20, 1<<20) + var out []*Frame + for { + f, err := r.ReadFrame() + if err == io.EOF { + return out + } + if err != nil { + t.Fatal(err) + } + out = append(out, f) + } +} + +func frag(op byte, fin bool, payload string) []byte { + b := Encode(op, []byte(payload), true) + if !fin { + b[0] &^= 0x80 + } + return b +} + +// A fragmented message reassembles across its continuations, and a ping between two +// fragments does not interrupt it (RFC 6455 §5.4). +func TestMessage_ReassemblesAroundInterleavedControlFrames(t *testing.T) { + m := NewMessage(1024) + frames := messageFrames(t, + frag(OpText, false, `{"type":"response.`), + Encode(OpPing, []byte("hi"), true), + frag(OpContinuation, false, `create","model":`), + frag(OpContinuation, true, `"gpt"}`), + ) + var got []string + for _, f := range frames { + op, payload, readable, ok, err := m.Add(f) + if err != nil { + t.Fatal(err) + } + if ok { + if op != OpText || !readable { + t.Fatalf("op=%x readable=%v", op, readable) + } + got = append(got, string(payload)) + } + } + if len(got) != 1 || got[0] != `{"type":"response.create","model":"gpt"}` { + t.Errorf("messages = %q", got) + } +} + +// A message over the limit completes but is not readable, and the next one starts +// fresh rather than inheriting the lost state. +func TestMessage_OverLimitIsCountedNotRead(t *testing.T) { + m := NewMessage(8) + frames := messageFrames(t, + frag(OpText, false, "12345"), + frag(OpContinuation, true, "67890"), + frag(OpText, true, "ok"), + ) + var results []string + for _, f := range frames { + _, payload, readable, ok, err := m.Add(f) + if err != nil { + t.Fatal(err) + } + if ok { + if readable { + results = append(results, string(payload)) + } else { + results = append(results, "") + } + } + } + if strings.Join(results, ",") != ",ok" { + t.Errorf("results = %v", results) + } +} + +func TestMessage_RefusesOutOfOrderFragments(t *testing.T) { + m := NewMessage(1024) + if _, _, _, _, err := m.Add(messageFrames(t, frag(OpContinuation, true, "x"))[0]); !errors.Is(err, ErrProtocol) { + t.Errorf("stray continuation: err = %v", err) + } + m = NewMessage(1024) + fs := messageFrames(t, frag(OpText, false, "a"), frag(OpText, true, "b")) + _, _, _, _, _ = m.Add(fs[0]) + if _, _, _, _, err := m.Add(fs[1]); !errors.Is(err, ErrProtocol) { + t.Errorf("new message mid-message: err = %v", err) + } +} + +func TestClosePayload_FitsAControlFrame(t *testing.T) { + p := ClosePayload(ClosePolicyViolation, strings.Repeat("r", 500)) + if len(p) > maxControlPayload { + t.Errorf("len = %d, over %d", len(p), maxControlPayload) + } + if p[0] != 0x03 || p[1] != 0xF0 { + t.Errorf("code bytes = % x, want 03 f0 (1008)", p[:2]) + } +} diff --git a/core/listener/parity/dispatch_audit_test.go b/core/listener/parity/dispatch_audit_test.go index 27f504ff5..4dca7a966 100644 --- a/core/listener/parity/dispatch_audit_test.go +++ b/core/listener/parity/dispatch_audit_test.go @@ -93,6 +93,14 @@ var dispatchSites = map[string]dispatchClass{ "forwardproxy/server.go:streamFallbackBuffered:phaseCtx": finalization, "forwardproxy/server.go:streamFallbackBuffered:foldCtx": finalization, "forwardproxy/server.go:streamFallbackBuffered:finalCtx": finalization, + // A carried Responses WebSocket: each turn's terminal frame, once the origin's + // terminal event (or the socket's end) has already gone to the client — the SSE + // finish defer's case, one turn at a time. + "forwardproxy/websocket.go:endTurn:finalCtx": finalizationDelivered, + // Each of a turn's server messages, mid-turn: the WebSocket counterpart of the SSE + // path's per-block frames. Detached from the handshake request, which ended when the + // upgrade was answered; the socket's own lifetime bounds it. + "forwardproxy/websocket.go:serverMessage:ws.ctx()": inFlight, // ext_proc's teardown flush: the stream is gone by definition here, and its two dispatches take // separate budgets. The second goes through dispatchTerminalFrame, and the CALL is audited here diff --git a/core/pipeline/context.go b/core/pipeline/context.go index 869633228..e0af5ee50 100644 --- a/core/pipeline/context.go +++ b/core/pipeline/context.go @@ -138,9 +138,10 @@ type Context struct { // looking at `405 Method Not Allowed` in their agent's output with every // reason to think Cortex broke, and an event saying only `backend_error/405` // with nothing connecting the two. (Codex no longer reaches that 405 by default: - // the forward proxy answers a Responses API WebSocket with 426 itself, which it - // reads as "use HTTPS" at once — see listener.decline_responses_websockets. The - // marker is still set there, since the Upgrade is still not forwarded.) + // the forward proxy carries a Responses API WebSocket, or answers it 426 itself, + // which Codex reads as "use HTTPS" at once — see listener.responses_websockets. + // The marker is still set on the handshake row, since the client's Upgrade header + // is never forwarded as it was sent.) // // Recorded on the REQUEST row, where the fact is true, and paired with the // response row's status by RequestID — the same split Tunnel/TunnelReason use. diff --git a/core/pipeline/session.go b/core/pipeline/session.go index 5a9e85656..ff178b71e 100644 --- a/core/pipeline/session.go +++ b/core/pipeline/session.go @@ -250,7 +250,7 @@ type SessionEvent struct { // the sequence worth being able to read off the timeline rather than out of an // agent's stderr. A 426 with error.kind "upgrade_declined" is the proxy answering // a Responses API WebSocket itself, before any origin is contacted - // (listener.decline_responses_websockets). + // (listener.responses_websockets). UpgradeStripped string // HTTPMethod and HTTPPath are the request's HTTP verb and its @@ -404,6 +404,11 @@ const ( // case can never produce the EMPTY string, which a consumer reads as "bridged" — // the exact opposite of what happened, and invisible in a timeline. TunnelPassthroughUnknown TunnelReason = "passthrough-unknown" + // TunnelWebSocket — not opaque at all: a Responses API WebSocket the forward + // proxy carries and reads. The open row is its handshake and the close row its + // end, with the bytes and duration of the whole socket; every turn on it records + // its own request and response rows between the two, parsed like an HTTPS turn. + TunnelWebSocket TunnelReason = "websocket" ) // EventTLS describes the TLS state of a connection that produced a @@ -604,7 +609,7 @@ type EventError struct { } // ErrorKindUpgradeDeclined is the forward proxy answering a Responses API WebSocket -// upgrade with 426 itself (listener.decline_responses_websockets). It rides on an +// upgrade with 426 itself (listener.responses_websockets). It rides on an // EventError because the status is a 4xx and the row should say whose answer it is, // but it is not a failure: the client is expected to switch to HTTPS, and does. // usage does not count it as an error, and the Finisher sees an allow. diff --git a/docs/README.md b/docs/README.md index ba6ab97c2..02f7e15fe 100644 --- a/docs/README.md +++ b/docs/README.md @@ -71,7 +71,7 @@ subsystem each one drives rather than in one file, so this table is the index. | `session.archive:` | The laptop's on-disk session history, its bounds, and clearing it | [`laptop-service.md`](laptop-service.md#session-history-is-kept-in-cortexsessions) | | `stats:` | The diagnostic listener (`/stats`, `/config`, `/reload/status`, `/pricing/table`, `/tls-bridge/unread`), default `:9093` | [`framework-architecture.md`](framework-architecture.md) | | `spiffe:` | SVID sourcing over the Workload API and the `/opt` file mirror | [`architecture.md`](architecture.md#spiffe-configuration) | -| `listener:` | Listener addresses, `roles`, interception mode, `skip_hosts`, `decline_responses_websockets` | [`cmd/README.md`](../cmd/README.md) (ports, roles), [`framework-architecture.md`](framework-architecture.md#skipping-infrastructure-hosts-listenerskip_hosts) (reload rules, `skip_hosts`), [`framework-architecture.md`](framework-architecture.md#declining-responses-api-websockets-listenerdecline_responses_websockets) (`decline_responses_websockets`) | +| `listener:` | Listener addresses, `roles`, interception mode, `skip_hosts`, `responses_websockets` | [`cmd/README.md`](../cmd/README.md) (ports, roles), [`framework-architecture.md`](framework-architecture.md#skipping-infrastructure-hosts-listenerskip_hosts) (reload rules, `skip_hosts`), [`framework-architecture.md`](framework-architecture.md#responses-api-websockets-listenerresponses_websockets) (`responses_websockets`) | | `mtls:` | Transport mTLS on the listeners, both deployment shapes | [`framework-architecture.md`](framework-architecture.md#8a-mtls-layer) | | `tls_bridge:` | The laptop TLS bridge and its CA | [`laptop-service.md`](laptop-service.md#the-tls_bridge-configuration) | diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 1c61ae289..8e534bf74 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -202,42 +202,52 @@ mislabelled `Content-Type: application/json`. Codex's tool calls arrive as is one — so a consumer that parses `ToolCalls[].Arguments` as JSON must tolerate failure for this dialect. -**HTTPS, not WebSocket.** Codex prefers a WebSocket to the same `/responses` path, and -Cortex cannot carry one: `Upgrade` is a hop-by-hop header a proxy must not forward, and -the parser reads HTTP requests, not frames. So the forward proxy answers that handshake -itself with **426 Upgrade Required**, which Codex reads as "this endpoint wants HTTP" and -switches to HTTPS straight away, for the rest of its session. That covers both of -Codex's endpoints — `chatgpt.com` on a ChatGPT login and `api.openai.com` on an API key -— and no other host, since what a 426 means is only known for Codex against OpenAI. It -is not counted as an error in `/v1/usage`. The attempt shows on the -timeline as a `GET` on the Responses path answered 426 with `error.kind: -"upgrade_declined"`, and Codex prints one line for it: - -``` -ERROR codex_api::endpoint::responses_websocket: failed to connect to websocket: - HTTP error: 426 Upgrade Required, url: wss://chatgpt.com/backend-api/codex/responses -``` - -That line is expected and harmless; the turn continues over HTTPS. Before the 426, the -proxy forwarded the stripped `GET` and chatgpt.com answered **405**, which Codex treats -as a transient failure: it retried five times on a doubling backoff, printing -`Reconnecting… n/5`, before falling back. Measured on Codex 0.163.0, a one-word -`codex exec` turn took a median **11.0s** that way and **3.0s** with the 426. There is no -supported Codex setting to skip the WebSocket attempt: the built-in `openai` provider -hardcodes `supports_websockets = true` and Codex refuses a config that overrides a +**WebSocket, carried and parsed.** Codex prefers a WebSocket to the same `/responses` +path — one socket for its whole session, with later turns sending only new input rather +than the whole conversation. Cortex carries it on both of Codex's endpoints — `chatgpt.com` +on a ChatGPT login and `api.openai.com` on an API key — and reads every turn on it: the +forward proxy completes the handshake, relays the frames unchanged, and runs each +`response.create` and the events answering it through the pipeline as if it were the HTTPS +request and its SSE stream. So a turn over the socket is parsed, priced and recorded +exactly as it is over HTTPS — model, messages, tools, completion, tokens, cost. + +On the timeline a socket is a `GET` on the Responses path with `tunnelReason: +"websocket"`, open and close sharing a request id, the close carrying the socket's +duration and bytes. Its turns sit between the two as ordinary `POST` request and response +rows. Plugins act on each turn: one that rejects it keeps it from reaching chatgpt.com +and Codex is answered with the error event it reads for a server refusal; one that +rewrites the request body (`tool-prune`) changes what is sent. + +Where Cortex cannot carry the socket it answers the handshake with **426 Upgrade +Required** instead, which Codex reads as "this endpoint wants HTTP" and switches to HTTPS +at once, for the rest of its session. That happens for a connection the bridge did not +decrypt, an HTTP/2 connection, or a pipeline with a plugin that rewrites response bodies, +since a live socket's frames are on their way to Codex before a whole response exists. The +timeline then shows a 426 with `error.kind: "upgrade_declined"`, and Codex prints one +harmless `failed to connect to websocket: HTTP error: 426 Upgrade Required` line before +carrying on over HTTPS. + +Before either, the proxy forwarded the handshake with `Upgrade` stripped — it is a +hop-by-hop header — and chatgpt.com answered the bare `GET` with **405**. Codex treats a +405 as transient: it retried five times on a doubling backoff, printing `Reconnecting… +n/5`, before falling back. Measured on Codex 0.163.0, a one-word `codex exec` turn took a +median **11.0s** that way, against **~3s** with the socket carried or declined. There is +no supported Codex setting to skip the WebSocket attempt: the built-in `openai` provider +hardcodes `supports_websockets = true`, and Codex refuses a config that overrides a built-in provider. -The 426 is Cortex's setting, not Codex's, and it is on by default, so there is nothing to -configure. To turn it off — which brings back the 405 and its retries — add the key to the -`listener:` block that is already in `~/.cortex/config.yaml`: +The behaviour is Cortex's setting, not Codex's, and the default is `parse`, so there is +nothing to configure. To change it, add the key to the `listener:` block that is already +in `~/.cortex/config.yaml`: ```yaml listener: # ...the lines already here... - decline_responses_websockets: false + responses_websockets: decline # parse (default) | decline | forward ``` -and restart the proxy with `agentop service restart`. Unlike most of that file, a +`decline` keeps Codex on HTTPS with no delay; `forward` brings back the 405 and its +retries. Then restart the proxy with `agentop service restart`. Unlike most of that file, a `listener:` edit is not picked up live, and the restart briefly disconnects every agent routed through Cortex. @@ -327,11 +337,12 @@ with Cortex built from the change that added this page. This was exercised: Codex's source — and including the telemetry exporter staying unrecognised; - the `codex` series in `GET /v1/usage?group=agent`, and the model showing as priceable but unpriced; -- **on Codex 0.163.0**, the WebSocket decline: eight `codex exec` turns each made one - handshake attempt, answered 426, then completed over HTTPS with the model, tokens and a - cost from the downloaded price list — against seven 405'd attempts and a fallback - warning per turn without it. A TUI turn, driven over a pty, did the same: one 426, then - its inference parsed over HTTPS under `codex-tui`. +- **on Codex 0.163.0**, the WebSocket transport: three `codex exec` turns each opened one + socket through the bridge, with no handshake error and no fallback, and each turn on it + was recorded with its model, a 13,188-token total and a cost — in about 3s, against a + median 11.0s through the old 405 retries. The `decline` mode was measured the same way + over eight turns (one 426, then HTTPS), and a TUI turn driven over a pty declined and + parsed over HTTPS under `codex-tui`. Not tested: @@ -345,6 +356,11 @@ Not tested: dialect is modelled from the published schema rather than from captured traffic. - **A `function_call` tool call.** Codex's own `exec` tool is a `custom_tool_call`; the `function_call` shape is modelled from the documented schema. +- **The TUI over a carried WebSocket.** It shares `codex exec`'s transport code, and its + decline path was watched, but a TUI turn over a carried socket was not. +- **A multi-turn session over one socket against the real backend.** The relay's own + tests drive several turns, a prewarm, a rejected turn and a rewritten one over one + socket; the live runs were one turn each. - **Tool pruning**, which is unsupported above rather than merely untested. ## Known issues diff --git a/docs/framework-architecture.md b/docs/framework-architecture.md index e718cc9e0..3454d99b1 100644 --- a/docs/framework-architecture.md +++ b/docs/framework-architecture.md @@ -905,31 +905,32 @@ or token-exchange policy applied — bypass means bypass. It exists for traffic whose recording floods the session store and crowds out the user intent a guardrail needs (see the intent-pin note in [`laptop-service.md`](laptop-service.md)). -#### Declining Responses API WebSockets (`listener.decline_responses_websockets`) +#### Responses API WebSockets (`listener.responses_websockets`) -The forward proxy answers a WebSocket upgrade on OpenAI's Responses API — a `GET` with -`Upgrade: websocket` on a path ending `/responses`, to `chatgpt.com` or `api.openai.com` -— with **426 Upgrade Required** itself, and forwards nothing. On by default in every -deployment. +What the forward proxy does with a WebSocket upgrade on OpenAI's Responses API — a +`GET` with `Upgrade: websocket` on a path ending `/responses`, to `chatgpt.com` or +`api.openai.com` — which is Codex's preferred transport: ```yaml listener: - decline_responses_websockets: false # forward the upgrade instead, as before + responses_websockets: parse # parse (default) | decline | forward ``` -The proxy cannot carry the WebSocket either way: `Upgrade` is hop-by-hop, so a -forwarded handshake reaches the origin as a bare `GET`, which OpenAI's endpoint answers -405. Codex — the one client known to use this transport — retries a 405 five times over -about six seconds before falling back to HTTPS, and reads a 426 as "use HTTP" and falls -back at once. `false` restores the forwarding byte for byte, for a client that reads a -426 some other way. Every other WebSocket is forwarded whatever this says, including -one to another host serving the same API (a LiteLLM gateway, OpenCode Zen, a local -server): what a 426 means is only known for Codex against OpenAI, and another client -might stop where today it falls back. Restart to -apply, like every `listener.*` field. The key goes in the runtime config: on a laptop, -`~/.cortex/config.yaml`, then `agentop service restart`; in a cluster, the -`authbridge-runtime-config` ConfigMap, then a pod restart. See -[`agents/codex.md`](agents/codex.md#what-cortex-shows-for-it) for the measured difference. +| Mode | What happens | +|---|---| +| `parse` | The socket is carried: the handshake is completed with the origin, frames are relayed unchanged, and every turn on it — a `response.create` message and the event messages answering it — runs the outbound pipeline as the equivalent HTTPS request and SSE stream would. Each turn records a request/response pair, is priced, and finishes; the socket records an open and a close with `tunnelReason: "websocket"`. A plugin rejecting a turn keeps it from the origin and answers the client with a Responses `error` event; a request-body writer's rewrite is what the origin receives. `permessage-deflate` is stripped from the handshake so frames stay readable. A socket that cannot be carried — not decrypted by the TLS bridge, HTTP/2, or a pipeline with a `WritesResponseBody` plugin — is declined as below. | +| `decline` | Answered **426 Upgrade Required** by the proxy, contacting no origin. Codex reads that as "use HTTP" and switches to HTTPS at once. Recorded as `error.kind: "upgrade_declined"`, which `/v1/usage` does not count as an error. | +| `forward` | The behaviour before either existed: forwarded with `Upgrade` stripped (it is hop-by-hop), so the origin sees a bare `GET`, which OpenAI's endpoint answers 405. Codex retries a 405 five times over about six seconds before falling back to HTTPS. | + +Every other WebSocket is forwarded with `Upgrade` stripped whatever this says, +including one to another host serving the same API (a LiteLLM gateway, OpenCode Zen, a +local server): what a 426 means, and what the socket carries, is only known for Codex +against OpenAI, and another client might stop where today it falls back. A misspelt +mode is refused at startup. Restart to apply, like every `listener.*` field. The +key goes in the runtime config: on a laptop, `~/.cortex/config.yaml`, then `agentop service +restart`; in a cluster, the `authbridge-runtime-config` ConfigMap, then a pod restart. See +[`agents/codex.md`](agents/codex.md#what-cortex-shows-for-it) for what Codex shows and the +measured difference. --- From 6d4a59b52785d87598cbc8211f18e5bb16bc16b5 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Mon, 26 Oct 2026 19:52:57 -0400 Subject: [PATCH 02/12] fix: Pin a WebSocket turn's pipeline and parse requests up to 32 MiB Three fixes to the Responses WebSocket relay. - Each turn runs every step on the pipeline it began on. A hot reload keeps the old pipeline running for its drain window so in-flight requests finish where they started; an HTTPS turn always does, but a socket outlives the window, and a turn looking the pipeline up per step ended on one whose plugins never saw it begin. RunFinish dispatches by index into the pipeline it is called on, so the turn's own Finishers were skipped or the wrong ones ran. - What the client sends is parsed up to maxRequestBodySize (32 MiB), as an HTTPS request body is. It was capped at maxBodySize (10 MB), the per-event response cap, so a socket's first turn, which carries the whole conversation, went unparsed between 10 MB and 32 MiB where HTTPS parsed it. - A turn no longer inherits the handshake's Content-Encoding or Content-Length, which describe no body a turn carries. session-api.md and CLAUDE.md now say a websocket tunnel row is not opaque. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- CLAUDE.md | 2 +- core/listener/forwardproxy/websocket.go | 45 ++++++---- core/listener/forwardproxy/websocket_test.go | 91 +++++++++++++++++++- docs/session-api.md | 7 +- 4 files changed, 125 insertions(+), 20 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 894c3ea69..83273f238 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -619,7 +619,7 @@ Every event on `/v1/sessions/{id}` and `/v1/events` carries: - `requestedHost` — present only when a plugin redirected the request (`pctx.Redirect`): the host the client asked for. `host` is where the request actually went, and usage, the cost ledger and pricing all follow `host`. agentop's detail pane shows both on a `redirected:` line. - `inference.requestedModel` — present only when a plugin changed the model (`pctx.SetRequestModel`): the model the client asked for. `inference.model` is the model the request was sent for, which settlement prices; agentop's detail pane shows both on a `model:` line. - `bytesUp`, `bytesDown` — **body** bytes, one direction per row: a request row carries `bytesUp`, the body **as forwarded** (so a `tool-prune`-style rewrite is reflected, since the count is taken after the pipeline ran), and a response row carries `bytesDown`, the body **as it came back from the destination** — not as delivered, because extproc records the row before it emits a response mutator's replacement. Headers are not counted. A tunnel close is the one row carrying both. **Zero means the listener counted nothing, not that the body was empty**, and there is no third value for "unknown": every listener buffers bodies only when a plugin asks, so an unbuffered request reports zero, and so does a response relayed straight through — both proxies copy that one after the row is appended. agentop renders a zero side blank rather than as `0B`, and that is why its BYTES column can be trusted as a size when it shows one. All three listeners count the same thing for the same response, which took some doing: extproc sums the chunks Envoy hands it, while the proxies count off the upstream body at one choke point each, so every arm — buffered, block-by-block SSE, or chunk copy — reports the same figure. `core/listener/internal/bodycount` and `core/listener/parity`'s bytes fixtures hold that line. -- `tunnel`, `tunnelReason` — an opaque CONNECT (or transparent-redirect) tunnel records two rows sharing a `requestId`: the open (`phase: "request"`, `tunnelReason` saying why the bytes stayed opaque) and, when the tunnel ends, the close (`phase: "response"`). The close carries the CONNECT's own `statusCode` — 200, or 502 with the dial error in `error` when the destination could not be reached (`tunnelReason: "dial-failed"`) — plus `durationMs` for how long the tunnel stayed open and, in `bytesUp`/`bytesDown`, the bytes it carried each way (up = client to destination; opaque bytes have no body/header split, so these are everything that crossed). It is not the destination's status: that travels inside the client's end-to-end TLS. A bridged tunnel's open is recorded with its first decrypted request — in that request's session, directly before it, stamped with its time — and records no close, because the request carries its own response; a bridged tunnel that recorded no request gets its open and a close when it ends. Tunnel rows are kept out of `/v1/usage`: a tunnel's lifetime is not a request latency. +- `tunnel`, `tunnelReason` — an opaque CONNECT (or transparent-redirect) tunnel records two rows sharing a `requestId`: the open (`phase: "request"`, `tunnelReason` saying why the bytes stayed opaque) and, when the tunnel ends, the close (`phase: "response"`). The close carries the CONNECT's own `statusCode` — 200, or 502 with the dial error in `error` when the destination could not be reached (`tunnelReason: "dial-failed"`) — plus `durationMs` for how long the tunnel stayed open and, in `bytesUp`/`bytesDown`, the bytes it carried each way (up = client to destination; opaque bytes have no body/header split, so these are everything that crossed). It is not the destination's status: that travels inside the client's end-to-end TLS. A bridged tunnel's open is recorded with its first decrypted request — in that request's session, directly before it, stamped with its time — and records no close, because the request carries its own response; a bridged tunnel that recorded no request gets its open and a close when it ends. Tunnel rows are kept out of `/v1/usage`: a tunnel's lifetime is not a request latency. **`tunnelReason: "websocket"` is the one that is not opaque:** a Responses API WebSocket the forward proxy carries (`listener.responses_websockets`, `core/listener/forwardproxy/websocket.go`). Its open is the handshake `GET`, its close is status 101 with the whole socket's duration and bytes, and each turn on it records its own ordinary `POST` request/response pair between them, parsed and priced like an HTTPS turn. A turn runs on the pipeline it began on, so a hot reload never splits one. - `httpMethod`, `httpPath` — the HTTP verb and path, so a request no parser recognized is still identifiable rather than showing only a host. Distinct from the `method` inside `a2a` / `mcp`, which is a protocol method name. On an opaque tunnel `httpMethod` is `CONNECT` and `httpPath` is absent — opaque bytes carry no request line. The path is query-stripped and percent-decoded, so query-borne credentials never reach the timeline, but a secret in a path *segment* (a bot token, a webhook path) does survive on this unauthenticated surface — worth knowing before exporting events off-box. ### Invocation action vocabulary diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index 76d8286ff..a329c737a 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -51,10 +51,14 @@ import ( // relay reads plain frames. Codex offers it and does not require it. const ( - // wsMessageLimit is the largest message read for parsing. A larger one is relayed - // whole and not read: the turn it belongs to is still recorded, just without a - // parse. Matches the forward proxy's own response-body cap. - wsMessageLimit = maxBodySize + // wsClientMessageLimit and wsServerMessageLimit are the largest messages read for + // parsing, each direction matching what the HTTPS path parses: maxRequestBodySize + // for what the client sends — a socket's first turn carries the whole + // conversation, the same body an HTTPS POST would — and maxBodySize for each event + // the origin sends back, the SSE path's per-frame cap. A larger message is relayed + // whole and not read: its turn is still recorded, just without a parse. + wsClientMessageLimit = maxRequestBodySize + wsServerMessageLimit = maxBodySize // wsFrameLimit is the largest single frame relayed. Raw is buffered whole before // it is forwarded, so this bounds what one frame can cost; a frame over it ends the // socket. Matches the request-body cap, the largest thing the HTTPS path accepts. @@ -311,6 +315,13 @@ type wsSocket struct { // wsTurn is one response.create and the server messages that answer it. type wsTurn struct { pctx *pipeline.Context + // p is the pipeline the turn began on, and every later step of the turn runs on + // it. A hot reload swaps OutboundPipeline for new requests and keeps the old one + // running for its drain window so requests already in flight finish where they + // started; an HTTPS turn does that by lasting seconds. A socket outlives the + // window, so a turn looking the pipeline up again at each step would end on one + // whose parser never saw it begin — no tokens, no cost. + p *pipeline.Pipeline } func (ws *wsSocket) run() { @@ -335,8 +346,8 @@ func (ws *wsSocket) run() { // relayUp carries client frames to the origin, reading each data message on the way. func (ws *wsSocket) relayUp() { - fr := wsframe.NewReader(ws.clientR, wsMessageLimit, wsFrameLimit) - msg := wsframe.NewMessage(wsMessageLimit) + fr := wsframe.NewReader(ws.clientR, wsClientMessageLimit, wsFrameLimit) + msg := wsframe.NewMessage(wsClientMessageLimit) // held collects a data message's frames until it is whole: a turn a plugin // rejects or rewrites must not have reached the origin first. var held [][]byte @@ -384,8 +395,8 @@ func (ws *wsSocket) relayUp() { // relayDown carries origin frames to the client, feeding each data message to the turn // it belongs to. func (ws *wsSocket) relayDown() { - fr := wsframe.NewReader(ws.up, wsMessageLimit, wsFrameLimit) - msg := wsframe.NewMessage(wsMessageLimit) + fr := wsframe.NewReader(ws.up, wsServerMessageLimit, wsFrameLimit) + msg := wsframe.NewMessage(wsServerMessageLimit) for { f, err := fr.ReadFrame() if err != nil { @@ -471,12 +482,12 @@ func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool) } pctx := ws.newTurnContext(body) - t := &wsTurn{pctx: pctx} - action := ws.s.OutboundPipeline.Run(ws.ctx(), pctx) + t := &wsTurn{pctx: pctx, p: ws.s.OutboundPipeline.Load()} + action := t.p.Run(ws.ctx(), pctx) if action.Type == pipeline.Reject { ws.s.recordOutboundRejectIn(nil, pctx, action, ws.sid) ws.refuseTurn(action) - ws.s.OutboundPipeline.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) + t.p.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) return nil, true } if ws.s.Sessions != nil { @@ -518,7 +529,7 @@ func (ws *wsSocket) serverMessage(payload []byte) { ws.finishTurn(t, status, &pipeline.EventError{Kind: "backend_error", Code: fmt.Sprintf("%d", status), Message: "the origin answered the turn with an error event"}) return } - ws.s.OutboundPipeline.RunResponseFrame(ws.ctx(), t.pctx, payload, false) + t.p.RunResponseFrame(ws.ctx(), t.pctx, payload, false) t.pctx.ResponseBytes += int64(len(payload)) // The three events inference-parser reads usage and the finish reason from; each // ends a response, as response.completed ends Codex's own read loop. @@ -550,9 +561,9 @@ func (ws *wsSocket) endTurn(t *wsTurn, status int, fail *pipeline.EventError) { if fail == nil { t.pctx.StatusCode = status } - ws.s.OutboundPipeline.RunResponseFrame(finalCtx, t.pctx, nil, true) + t.p.RunResponseFrame(finalCtx, t.pctx, nil, true) ws.s.recordOutboundResponseEvent(t.pctx, status, fail) - ws.s.OutboundPipeline.RunFinish(finalCtx, t.pctx, pipeline.OutcomeFromContext(t.pctx)) + t.p.RunFinish(finalCtx, t.pctx, pipeline.OutcomeFromContext(t.pctx)) } // ctx is the context every turn's pipeline runs under: the handshake request's values @@ -568,7 +579,11 @@ func (ws *wsSocket) ctx() context.Context { func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { hs := ws.hs headers := ws.r.Header.Clone() - for _, h := range []string{"Upgrade", "Connection", "Sec-Websocket-Key", "Sec-Websocket-Version", "Sec-Websocket-Extensions", "Sec-Websocket-Protocol"} { + // The handshake's own framing headers, and what it said about a body it does not + // have: a turn's body is a WebSocket text message, never encoded, and the + // handshake's Content-Encoding and Content-Length describe nothing the turn carries. + // A plugin reading them would be told something false about the body it is shown. + for _, h := range []string{"Upgrade", "Connection", "Sec-Websocket-Key", "Sec-Websocket-Version", "Sec-Websocket-Extensions", "Sec-Websocket-Protocol", "Content-Encoding", "Content-Length"} { headers.Del(h) } headers.Set("Content-Type", "application/json") diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 0b4928d60..7ad3615f9 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -106,8 +106,8 @@ func wsAccept(key string) string { } func (o *wsOrigin) serve(conn net.Conn, br *bufio.Reader) { - fr := wsframe.NewReader(br, 1<<20, 1<<20) - msg := wsframe.NewMessage(1 << 20) + fr := wsframe.NewReader(br, 64<<20, 64<<20) + msg := wsframe.NewMessage(64 << 20) turn := 0 for { f, err := fr.ReadFrame() @@ -171,6 +171,9 @@ type wsHarness struct { origin *wsOrigin proxy *httptest.Server caPEM []byte + // holder is the outbound pipeline the proxy runs, so a test can swap it as a hot + // reload does. + holder *pipeline.Holder } func newWSHarness(t *testing.T, mode string, plugins ...pipeline.Plugin) *wsHarness { @@ -208,8 +211,9 @@ func newWSHarness(t *testing.T, mode string, plugins ...pipeline.Plugin) *wsHarn if err != nil { t.Fatal(err) } + h.holder = pipeline.NewHolder(p) srv := &Server{ - OutboundPipeline: pipeline.NewHolder(p), + OutboundPipeline: h.holder, Sessions: h.store, Client: http.DefaultClient, TLSBridge: engine, @@ -652,6 +656,87 @@ func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { } } +// A hot reload mid-turn must not move the rest of that turn onto the new pipeline. The +// turn's Finishers are the plugins whose OnRequest ran — lineage's span, a budget's +// reservation — and RunFinish dispatches them by their index in the pipeline it is +// called on, so ending the turn on a newer pipeline would finish the wrong plugins, or +// none, and leave the started ones open. The turn ends on the pipeline it began on; the +// next turn takes the new one. +func TestResponsesWebSocket_ReloadMidTurnFinishesOnTheTurnsPipeline(t *testing.T) { + started := &outcomeRecorder{} + h := newWSHarness(t, config.ResponsesWebSocketsParse, started) + midTurn := make(chan struct{}) + resume := make(chan struct{}) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { + return + } + evs := turnEvents(1) + for _, ev := range evs[:3] { + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) + } + close(midTurn) + <-resume + for _, ev := range evs[3:] { + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) + } + _, _ = fr.ReadFrame() + }) + c := h.dial(t, nil) + c.send(responseCreate("hello")) + <-midTurn + + // The reload: a pipeline of different plugins, with a Finisher of its own that + // never saw this turn begin. + other := &outcomeRecorder{} + fresh, err := plugintesting.BuildPipeline([]pipeline.Plugin{inferenceparser.NewInferenceParser(), &turnDenier{}, other}) + if err != nil { + t.Fatal(err) + } + h.holder.Store(fresh) + close(resume) + c.readUntil("response.completed") + c.close() + + if started.seen() == nil { + t.Error("the turn's own Finisher never ran: the turn ended on a pipeline other than the one it began on") + } + if other.got != nil { + t.Error("the reloaded pipeline's Finisher ran for a turn it never saw begin") + } +} + +// A socket's first turn carries the whole conversation, so it can be far over the +// 10 MB per-event cap the origin's messages get. It is parsed up to the same 32 MiB an +// HTTPS request body is. +func TestResponsesWebSocket_LargeFirstTurnIsParsed(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + c := h.dial(t, nil) + big := strings.Repeat("x", 15<<20) + c.send(responseCreate(big)) + c.readUntil("response.completed") + c.close() + + var req *pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionRequest { + e := e + req = &e + } + } + if req == nil || req.Inference == nil || req.Inference.Model != "gpt-6-luna" { + t.Fatalf("a 15 MB first turn was not parsed: %+v", req) + } +} + // --- test plugins --- // turnDenier rejects every inference request, and nothing else, so the handshake GET diff --git a/docs/session-api.md b/docs/session-api.md index f326bb1bb..d41327e9f 100644 --- a/docs/session-api.md +++ b/docs/session-api.md @@ -203,7 +203,12 @@ Every event on `/v1/sessions/{id}` and `/v1/events` carries: "response"`), carrying the CONNECT's own status, the duration, and the bytes each way. It is not the destination's status: that travels inside the client's end-to-end TLS. Tunnel rows are kept out of `/v1/usage`: a tunnel's - lifetime is not a request latency. + lifetime is not a request latency. **`tunnelReason: "websocket"` is the one + that is not opaque:** a Responses API WebSocket the forward proxy carries + (`listener.responses_websockets`). Its open is the handshake `GET`, its close + is status 101 with the whole socket's duration and bytes, and every turn on it + records its own ordinary `POST` request and response rows between the two, + parsed and priced like an HTTPS turn. - `httpMethod`, `httpPath` — the HTTP verb and query-stripped, percent-decoded path, so a request no parser recognized is still identifiable. On an opaque tunnel `httpMethod` is `CONNECT` and `httpPath` From 35d7181ddc0fd7549d680875a0b3f44438952863 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Tue, 27 Oct 2026 00:58:13 -0400 Subject: [PATCH 03/12] test: Pin Codex's WebSocket interrupt, and record the live TUI and multi-turn runs Esc in Codex's TUI sends response.interrupt over the socket and waits for response.interrupt.accepted carrying the same response id, which Codex rejects if it does not arrive intact. The new test drives that exchange in both orders the server may send the acknowledgement and the turn's terminal event, and checks the relay forwards both unchanged and records the interrupted turn once, with its partial text and usage. docs/agents/codex.md moves two items from Not tested to the verified list, both measured on Codex 0.163.0: several turns on one socket, and a TUI turn over a carried socket. It adds the API-key login and a live interrupt to Not tested. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/websocket_test.go | 103 +++++++++++++++++++ docs/agents/codex.md | 21 ++-- 2 files changed, 118 insertions(+), 6 deletions(-) diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 7ad3615f9..27c5f37a5 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -656,6 +656,109 @@ func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { } } +// Esc in Codex's TUI interrupts a turn over the socket: it sends response.interrupt +// naming the response, and waits on the same socket for the server's +// response.interrupt.accepted carrying that id, which it rejects as "unexpected" or +// "invalid" if it does not arrive intact. The relay must carry both unchanged, record +// the interrupted turn once with what had arrived, and not mistake the acknowledgement +// for anything that opens or ends a turn — in either order the server may send the +// acknowledgement and the turn's terminal event. +func TestResponsesWebSocket_InterruptIsRelayedAndTheTurnRecordedOnce(t *testing.T) { + const ( + interrupt = `{"type":"response.interrupt","response_id":"resp_i","mode":"discard_partial_items"}` + ack = `{"type":"response.interrupt.accepted","response_id":"resp_i"}` + terminal = `{"type":"response.incomplete","response":{"id":"resp_i","status":"incomplete","usage":{` + + `"input_tokens":12862,"input_tokens_details":{"cached_tokens":11008},"output_tokens":2,"total_tokens":12864}}}` + ) + for _, order := range []struct { + name string + after []string + }{ + {"ack then terminal", []string{ack, terminal}}, + {"terminal then ack", []string{terminal, ack}}, + } { + t.Run(order.name, func(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + var gotInterrupt atomic.Value + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { // the response.create + return + } + for _, ev := range []string{ + `{"type":"response.created","response":{"id":"resp_i","status":"in_progress"}}`, + `{"type":"response.output_text.delta","item_id":"m","output_index":0,"content_index":0,"delta":"po"}`, + } { + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) + } + f, err := fr.ReadFrame() // the interrupt + if err != nil { + return + } + gotInterrupt.Store(string(f.Payload)) + for _, ev := range order.after { + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) + } + _, _ = fr.ReadFrame() + }) + + c := h.dial(t, nil) + c.send(responseCreate("hello")) + c.readUntil("response.output_text.delta") + c.send(interrupt) + var sawAck, sawTerminal string + for sawAck == "" || sawTerminal == "" { + for _, m := range c.readUntil(map[bool]string{true: "response.incomplete", false: "response.interrupt.accepted"}[sawTerminal == ""]) { + b, _ := json.Marshal(m) + switch m["type"] { + case "response.interrupt.accepted": + sawAck = string(b) + case "response.incomplete": + sawTerminal = string(b) + } + } + } + c.close() + + if got, _ := gotInterrupt.Load().(string); got != interrupt { + t.Errorf("origin received %q, want the interrupt unchanged", got) + } + if !strings.Contains(sawAck, `"response_id":"resp_i"`) { + t.Errorf("client's acknowledgement %s lost the response id Codex matches on", sawAck) + } + + var reqs, resps []pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod != http.MethodPost { + continue + } + if e.Phase == pipeline.SessionRequest { + reqs = append(reqs, e) + } else { + resps = append(resps, e) + } + } + if len(reqs) != 1 || len(resps) != 1 { + t.Fatalf("got %d turn requests and %d responses, want one of each: the interrupt and its acknowledgement are not turns", len(reqs), len(resps)) + } + res := resps[0] + if res.Error != nil { + t.Errorf("an interrupted turn was recorded as an error: %+v", res.Error) + } + if res.Inference == nil || res.Inference.Completion != "po" || res.Inference.FinishReason != "incomplete" || res.Inference.CompletionTokens != 2 { + t.Errorf("interrupted turn = %+v, want the partial text, finish reason incomplete and its usage", res.Inference) + } + }) + } +} + // A hot reload mid-turn must not move the rest of that turn onto the new pipeline. The // turn's Finishers are the plugins whose OnRequest ran — lineage's span, a budget's // reservation — and RunFinish dispatches them by their index in the pipeline it is diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 8e534bf74..33916cf64 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -342,11 +342,25 @@ with Cortex built from the change that added this page. This was exercised: was recorded with its model, a 13,188-token total and a cost — in about 3s, against a median 11.0s through the old 405 retries. The `decline` mode was measured the same way over eight turns (one 426, then HTTPS), and a TUI turn driven over a pty declined and - parsed over HTTPS under `codex-tui`. + parsed over HTTPS under `codex-tui`; +- **several turns on one socket**: a `codex exec` prompt that ran two shell commands + carried two turns on one WebSocket, each recorded with its model, tokens and a cost. + The second turn's request held no messages, only the tool results, which is the + incremental input described in #1515; +- **the TUI over a carried socket**: a prompt typed into the real interactive TUI over + a pty opened a WebSocket under `codex-tui`, and its turn was parsed with model, + tokens and a cost, as was the TUI's own model check on `gpt-5.6-luna`. Not tested: - **Linux.** Everything here was exercised on macOS only. +- **An API-key login over the WebSocket.** Codex then opens it to `api.openai.com` + rather than `chatgpt.com`; Cortex treats the two hosts alike, but only the ChatGPT + login was watched. +- **Interrupting a turn over the WebSocket** (Esc in the TUI). Codex sends a + `response.interrupt` message and waits for the server's acknowledgement on the same + socket; the relay forwards both unchanged and its tests cover the shape, but a live + interrupt was not watched. - **Codex versions other than 0.160.1 and 0.162.1.** Both appeared in one capture (an auto-update landed mid-session) and both parse identically, which is some evidence the User-Agent shape is stable across releases rather than pinned to one build — but the @@ -356,11 +370,6 @@ Not tested: dialect is modelled from the published schema rather than from captured traffic. - **A `function_call` tool call.** Codex's own `exec` tool is a `custom_tool_call`; the `function_call` shape is modelled from the documented schema. -- **The TUI over a carried WebSocket.** It shares `codex exec`'s transport code, and its - decline path was watched, but a TUI turn over a carried socket was not. -- **A multi-turn session over one socket against the real backend.** The relay's own - tests drive several turns, a prewarm, a rejected turn and a rewritten one over one - socket; the live runs were one turn each. - **Tool pruning**, which is unsupported above rather than merely untested. ## Known issues From 05b21b4728f7d551941e18f9757b966dff5b239d Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Tue, 27 Oct 2026 01:19:17 -0400 Subject: [PATCH 04/12] fix: Record a turn the client abandons as client_canceled, not upstream_error A turn still open when its WebSocket ended was always recorded as an upstream_error. Codex closes the socket when the user presses Esc, so every interrupted turn read as OpenAI failing. Each relay now reports which side stopped it, and the first to stop decides: the client closing records the turn as the HTTPS path records a client that hangs up, 499 client_canceled, which /v1/usage keeps out of the 5xx series; the origin closing records 502 upstream_error. Measured live on Codex 0.163.0: Esc mid-answer closed the socket within milliseconds, the turn was recorded with the text that had arrived, and the next prompt opened a new socket and was parsed normally. The Codex doc said Esc sends response.interrupt; it does not. That in-band interrupt is for Codex's own steering and mailbox input, and the doc and test now say so. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/websocket.go | 64 ++++++++++++++++---- core/listener/forwardproxy/websocket_test.go | 61 ++++++++++++++++--- docs/agents/codex.md | 22 +++++-- 3 files changed, 122 insertions(+), 25 deletions(-) diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index a329c737a..19a0aaa41 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -310,6 +310,16 @@ type wsSocket struct { // server side feeds and ends it. turnMu sync.Mutex turn *wsTurn + + // endedBy is the side whose relay stopped first, "client" or "origin", which + // decides how a turn still open at that moment is recorded. Set once, by the + // first relay to return; the other only ever sees the connection closed under it. + endedBy atomic.Value +} + +// wsEnd records who ended the socket, if no one has yet. +func (ws *wsSocket) wsEnd(side string) { + ws.endedBy.CompareAndSwap(nil, side) } // wsTurn is one response.create and the server messages that answer it. @@ -326,8 +336,8 @@ type wsTurn struct { func (ws *wsSocket) run() { done := make(chan struct{}, 2) - go func() { ws.relayUp(); done <- struct{}{} }() - go func() { ws.relayDown(); done <- struct{}{} }() + go func() { ws.wsEnd(ws.relayUp()); done <- struct{}{} }() + go func() { ws.wsEnd(ws.relayDown()); done <- struct{}{} }() <-done // One side is gone: close both, so the other relay's read returns. _ = ws.client.Close() @@ -339,13 +349,40 @@ func (ws *wsSocket) run() { ws.turn = nil ws.turnMu.Unlock() if t != nil { - ws.endTurn(t, http.StatusOK, &pipeline.EventError{Kind: "upstream_error", Message: "the WebSocket closed before the turn completed"}) + status, fail := ws.unfinishedTurn() + ws.endTurn(t, status, fail) } ws.recordClose() } +// unfinishedTurn is how a turn still open when the socket ended is recorded, by who +// ended it — the WebSocket counterpart of what the HTTPS path records for a stream that +// stops before its last event. +// +// THE CLIENT CLOSING IS NOT AN UPSTREAM FAILURE. It is what Codex does on Esc: it drops +// the turn's stream, which closes this socket, and opens another for the next prompt. +// So it is recorded as the HTTPS path records a client that hangs up before its answer +// — 499, client_canceled — which /v1/usage counts as an error but keeps out of the 5xx +// series, rather than as a 502-shaped upstream_error that would read as OpenAI failing +// every time someone interrupts a turn. The origin closing first, or the socket breaking +// on a frame the relay cannot read, is what upstream_error is for. +func (ws *wsSocket) unfinishedTurn() (int, *pipeline.EventError) { + if side, _ := ws.endedBy.Load().(string); side == "client" { + return pipeline.StatusClientClosedRequest, &pipeline.EventError{ + Kind: "client_canceled", + Message: "the client closed the WebSocket before the turn completed", + } + } + return http.StatusBadGateway, &pipeline.EventError{ + Kind: "upstream_error", + Message: "the origin closed the WebSocket before the turn completed", + } +} + // relayUp carries client frames to the origin, reading each data message on the way. -func (ws *wsSocket) relayUp() { +// It returns the side that stopped it: the client when reading from it failed, the +// origin when writing to it did. +func (ws *wsSocket) relayUp() (endedBy string) { fr := wsframe.NewReader(ws.clientR, wsClientMessageLimit, wsFrameLimit) msg := wsframe.NewMessage(wsClientMessageLimit) // held collects a data message's frames until it is whole: a turn a plugin @@ -355,19 +392,19 @@ func (ws *wsSocket) relayUp() { f, err := fr.ReadFrame() if err != nil { ws.logEnd("client", err) - return + return "client" } ws.bytesUp.Add(int64(len(f.Raw))) if f.IsControl() { if !ws.writeUp(f.Raw) { - return + return "origin" } continue } op, payload, readable, complete, err := msg.Add(f) if err != nil { ws.logEnd("client", err) - return + return "client" } held = append(held, f.Raw) if !complete { @@ -386,22 +423,23 @@ func (ws *wsSocket) relayUp() { } for _, raw := range out { if !ws.writeUp(raw) { - return + return "origin" } } } } // relayDown carries origin frames to the client, feeding each data message to the turn -// it belongs to. -func (ws *wsSocket) relayDown() { +// it belongs to. It returns the side that stopped it: the origin when reading from it +// failed, the client when writing to it did. +func (ws *wsSocket) relayDown() (endedBy string) { fr := wsframe.NewReader(ws.up, wsServerMessageLimit, wsFrameLimit) msg := wsframe.NewMessage(wsServerMessageLimit) for { f, err := fr.ReadFrame() if err != nil { ws.logEnd("origin", err) - return + return "origin" } ws.bytesDown.Add(int64(len(f.Raw))) // Forwarded first, as handleStreamingResponse writes a block after its @@ -409,7 +447,7 @@ func (ws *wsSocket) relayDown() { // WebSocket turn cannot be retracted once it started, so the frame goes // now and the parse follows. if !ws.writeClient(f.Raw) { - return + return "client" } if f.IsControl() { continue @@ -417,7 +455,7 @@ func (ws *wsSocket) relayDown() { op, payload, readable, complete, err := msg.Add(f) if err != nil { ws.logEnd("origin", err) - return + return "origin" } if complete && op == wsframe.OpText && readable { ws.serverMessage(payload) diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 27c5f37a5..839a06910 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -615,8 +615,10 @@ func TestResponsesWebSocket_OriginRefusalIsRelayed(t *testing.T) { } } -// A client that hangs up mid-turn still gets that turn recorded, with what had -// arrived, and the socket's close. +// A client that hangs up mid-turn — Codex does this on Esc, closing the socket rather +// than interrupting over it — still gets that turn recorded with what had arrived, as a +// 499 client_canceled the way an HTTPS client hanging up is, not as an upstream +// failure: the origin did nothing wrong. func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { h := newWSHarness(t, config.ResponsesWebSocketsParse) // An origin that answers a create with a partial stream and then stalls. @@ -651,15 +653,60 @@ func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { if res == nil { t.Fatal("the interrupted turn recorded no response row") } - if res.Error == nil || res.Inference == nil || res.Inference.Completion != "par" { - t.Errorf("interrupted turn: error=%+v completion=%v, want an error and the partial text", res.Error, res.Inference) + if res.Inference == nil || res.Inference.Completion != "par" { + t.Errorf("interrupted turn completion = %+v, want the partial text", res.Inference) } + if res.StatusCode != pipeline.StatusClientClosedRequest || res.Error == nil || res.Error.Kind != "client_canceled" { + t.Errorf("interrupted turn = status %d error %+v, want 499 client_canceled", res.StatusCode, res.Error) + } +} + +// The origin dropping the socket mid-turn is the upstream failure: a 502 upstream_error, +// as the HTTPS path records a stream that breaks off from the server's side. +func TestResponsesWebSocket_OriginDropFinalizesTheOpenTurnAsUpstreamError(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { + _ = conn.Close() + return + } + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.created","response":{"id":"r","status":"in_progress"}}`), false)) + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.output_text.delta","item_id":"m","output_index":0,"content_index":0,"delta":"par"}`), false)) + _ = conn.Close() // the origin goes away mid-turn + }) + c := h.dial(t, nil) + c.send(responseCreate("hello")) + c.readUntil("response.output_text.delta") + + var res *pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionResponse { + e := e + res = &e + } + } + if res == nil { + t.Fatal("the dropped turn recorded no response row") + } + if res.StatusCode != http.StatusBadGateway || res.Error == nil || res.Error.Kind != "upstream_error" { + t.Errorf("dropped turn = status %d error %+v, want 502 upstream_error", res.StatusCode, res.Error) + } + _ = c.conn.Close() } -// Esc in Codex's TUI interrupts a turn over the socket: it sends response.interrupt -// naming the response, and waits on the same socket for the server's +// Codex's in-band interrupt: for steering a turn already under way, or for queued +// mailbox input, on models with use_responses_lite set, it sends response.interrupt +// naming the response and waits on the same socket for the server's // response.interrupt.accepted carrying that id, which it rejects as "unexpected" or -// "invalid" if it does not arrive intact. The relay must carry both unchanged, record +// "invalid" if it does not arrive intact. (Esc does not use this: it closes the socket, +// which TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn covers.) The relay must carry both unchanged, record // the interrupted turn once with what had arrived, and not mistake the acknowledgement // for anything that opens or ends a turn — in either order the server may send the // acknowledgement and the turn's terminal event. diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 33916cf64..8d888fbd6 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -218,6 +218,13 @@ rows. Plugins act on each turn: one that rejects it keeps it from reaching chatg and Codex is answered with the error event it reads for a server refusal; one that rewrites the request body (`tool-prune`) changes what is sent. +**Pressing Esc mid-answer closes the socket.** Codex drops the turn's stream, which ends +the WebSocket, and opens a new one for the next prompt. The interrupted turn is recorded +with the text that had arrived, as `statusCode` 499 with `error.kind: +"client_canceled"` — the same as an HTTPS client hanging up — so it is not counted as an +upstream failure. OpenAI sends no usage for a turn cut off this way, so it has no token +count or cost: it shows as priceable but unpriced, the same as a cancelled HTTPS stream. + Where Cortex cannot carry the socket it answers the handshake with **426 Upgrade Required** instead, which Codex reads as "this endpoint wants HTTP" and switches to HTTPS at once, for the rest of its session. That happens for a connection the bridge did not @@ -349,7 +356,10 @@ with Cortex built from the change that added this page. This was exercised: incremental input described in #1515; - **the TUI over a carried socket**: a prompt typed into the real interactive TUI over a pty opened a WebSocket under `codex-tui`, and its turn was parsed with model, - tokens and a cost, as was the TUI's own model check on `gpt-5.6-luna`. + tokens and a cost, as was the TUI's own model check on `gpt-5.6-luna`; +- **Esc mid-answer in the TUI**: the WebSocket closed within milliseconds, the + interrupted turn was recorded with the 1,491 characters that had arrived, and the + next prompt opened a new socket and was answered and parsed normally. Not tested: @@ -357,10 +367,12 @@ Not tested: - **An API-key login over the WebSocket.** Codex then opens it to `api.openai.com` rather than `chatgpt.com`; Cortex treats the two hosts alike, but only the ChatGPT login was watched. -- **Interrupting a turn over the WebSocket** (Esc in the TUI). Codex sends a - `response.interrupt` message and waits for the server's acknowledgement on the same - socket; the relay forwards both unchanged and its tests cover the shape, but a live - interrupt was not watched. +- **Codex's in-band interrupt.** For some internal cases — steering a turn already + under way, or queued "mailbox" input, on models with `use_responses_lite` set — Codex + sends `response.interrupt` on the open socket and waits for the server's + `response.interrupt.accepted` instead of closing it. The relay forwards both unchanged + and its tests cover the exchange in both orders, but it was not seen live: Esc does + not use it. - **Codex versions other than 0.160.1 and 0.162.1.** Both appeared in one capture (an auto-update landed mid-session) and both parse identically, which is some evidence the User-Agent shape is stable across releases rather than pinned to one build — but the From d8297e4a22676ebb7e1d426f46d43e44dcd91e07 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Tue, 27 Oct 2026 01:54:32 -0400 Subject: [PATCH 05/12] fix: Read close frames and error codes when recording a WebSocket turn Two refinements to how a turn's row is recorded. - Who ended a socket mid-turn is read from the first close frame either side sent, falling back to which relay stopped first only when neither sent one. When one side closes cleanly the other answers with a close of its own, and that answer could make its relay stop first: a client closing cleanly was recorded as the origin's fault. The close code and reason go into the row's message. - A refusal the origin sends as an error event takes the event's own status when it has one, and otherwise the status its error code is answered with over HTTPS (insufficient_quota and rate_limit_exceeded 429, context_length_exceeded 400, server_is_overloaded 503), before falling back to 500. A live quota refusal arrived with no status and was recorded as a 500. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/websocket.go | 114 +++++++++++++-- core/listener/forwardproxy/websocket_test.go | 132 ++++++++++++++++++ core/listener/internal/wsframe/wsframe.go | 43 ++++++ .../listener/internal/wsframe/wsframe_test.go | 18 +++ docs/agents/codex.md | 10 ++ 5 files changed, 308 insertions(+), 9 deletions(-) diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index 19a0aaa41..89959daa6 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -11,6 +11,7 @@ import ( "log/slog" "net" "net/http" + "strconv" "strings" "sync" "sync/atomic" @@ -315,6 +316,22 @@ type wsSocket struct { // decides how a turn still open at that moment is recorded. Set once, by the // first relay to return; the other only ever sees the connection closed under it. endedBy atomic.Value + + // closedBy is the side that sent the first close frame, and closeNote its code and + // reason. A close frame is the sender saying it is ending the socket, which is + // better evidence than which relay happened to stop first: when one side closes + // cleanly the other answers with a close of its own, and that answer can make its + // relay stop first. Set once, by the first close frame either relay sees. + closedBy atomic.Value + closeNote atomic.Value +} + +// sawClose records a close frame from side, if none has been seen yet. +func (ws *wsSocket) sawClose(side string, f *wsframe.Frame) { + if !ws.closedBy.CompareAndSwap(nil, side) { + return + } + ws.closeNote.Store(wsframe.DescribeClose(f.Payload)) } // wsEnd records who ended the socket, if no one has yet. @@ -366,17 +383,86 @@ func (ws *wsSocket) run() { // series, rather than as a 502-shaped upstream_error that would read as OpenAI failing // every time someone interrupts a turn. The origin closing first, or the socket breaking // on a frame the relay cannot read, is what upstream_error is for. +// +// A close frame decides first, when either side sent one: whoever sent the first close +// ended the socket, and its code and reason go into the message. Without one — a +// connection dropped, which is what Codex's Esc does — the relay that stopped first +// decides. Either way the statuses are the same two; a close code is evidence about +// who, recorded beside the status, not a status of its own. func (ws *wsSocket) unfinishedTurn() (int, *pipeline.EventError) { - if side, _ := ws.endedBy.Load().(string); side == "client" { + side, _ := ws.closedBy.Load().(string) + note, _ := ws.closeNote.Load().(string) + if side == "" { + side, _ = ws.endedBy.Load().(string) + } + suffix := "" + if note != "" { + suffix = " (" + note + ")" + } + if side == "client" { return pipeline.StatusClientClosedRequest, &pipeline.EventError{ Kind: "client_canceled", - Message: "the client closed the WebSocket before the turn completed", + Message: "the client closed the WebSocket before the turn completed" + suffix, } } return http.StatusBadGateway, &pipeline.EventError{ Kind: "upstream_error", - Message: "the origin closed the WebSocket before the turn completed", + Message: "the origin closed the WebSocket before the turn completed" + suffix, + } +} + +// errorEventStatus is the HTTP status recorded for a turn the origin refused with an +// error event. +// +// AN OPEN WEBSOCKET HAS NO HTTP STATUS: the handshake's 101 was the last one. So the +// origin reports a refused turn inside a message, and the status field is optional — +// Codex's own fixtures carry "status": 429 on a usage limit, while a live quota refusal +// arrived with none. The row records the status the same refusal answers over HTTPS, +// so a WebSocket turn and an HTTPS one count alike in /v1/usage: the event's own status +// when it has one; else the one its error code is answered with over HTTPS; else 500. +// Codes are the Responses API's, as Codex's source maps them (codex-api's api_bridge). +func errorEventStatus(status int, code, errType string) int { + if status >= 400 && status <= 599 { + return status + } + for _, name := range []string{code, errType} { + switch name { + case "rate_limit_exceeded", "insufficient_quota", "usage_limit_reached", + "usage_not_included", "websocket_connection_limit_reached", "slow_down": + return http.StatusTooManyRequests + case "context_length_exceeded", "invalid_prompt", "invalid_request_error", + "previous_response_not_found": + return http.StatusBadRequest + case "server_is_overloaded": + return http.StatusServiceUnavailable + } } + return http.StatusInternalServerError +} + +// firstNonEmpty is the first of vs that is not empty, or "". +func firstNonEmpty(vs ...string) string { + for _, v := range vs { + if v != "" { + return v + } + } + return "" +} + +// capErrorName bounds an origin-supplied error code before it reaches a row an +// operator reads in a terminal: printable ASCII only, 64 bytes at most. +func capErrorName(s string) string { + var b strings.Builder + for _, r := range s { + if b.Len() >= 64 { + break + } + if r >= 0x20 && r < 0x7f { + b.WriteRune(r) + } + } + return b.String() } // relayUp carries client frames to the origin, reading each data message on the way. @@ -396,6 +482,9 @@ func (ws *wsSocket) relayUp() (endedBy string) { } ws.bytesUp.Add(int64(len(f.Raw))) if f.IsControl() { + if f.Opcode == wsframe.OpClose { + ws.sawClose("client", f) + } if !ws.writeUp(f.Raw) { return "origin" } @@ -450,6 +539,9 @@ func (ws *wsSocket) relayDown() (endedBy string) { return "client" } if f.IsControl() { + if f.Opcode == wsframe.OpClose { + ws.sawClose("origin", f) + } continue } op, payload, readable, complete, err := msg.Add(f) @@ -555,16 +647,20 @@ func (ws *wsSocket) serverMessage(payload []byte) { var env struct { Type string `json:"type"` Status int `json:"status"` + Error struct { + Type string `json:"type"` + Code string `json:"code"` + } `json:"error"` } _ = json.Unmarshal(payload, &env) if env.Type == "error" { - // Codex's wrapped error event — a refused turn (usage limit, bad input), with - // the HTTP status the HTTPS path would have answered. - status := env.Status - if status == 0 { - status = http.StatusInternalServerError + // The origin refused the turn (a usage limit, bad input). See errorEventStatus. + status := errorEventStatus(env.Status, env.Error.Code, env.Error.Type) + message := "the origin answered the turn with an error event" + if name := firstNonEmpty(env.Error.Code, env.Error.Type); name != "" { + message += ": " + capErrorName(name) } - ws.finishTurn(t, status, &pipeline.EventError{Kind: "backend_error", Code: fmt.Sprintf("%d", status), Message: "the origin answered the turn with an error event"}) + ws.finishTurn(t, status, &pipeline.EventError{Kind: "backend_error", Code: strconv.Itoa(status), Message: message}) return } t.p.RunResponseFrame(ws.ctx(), t.pctx, payload, false) diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 839a06910..1f3851757 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -701,6 +701,138 @@ func TestResponsesWebSocket_OriginDropFinalizesTheOpenTurnAsUpstreamError(t *tes _ = c.conn.Close() } +// A clean close mid-turn names who ended the socket even when the other side's relay +// stops first. Here the client sends a close frame, the origin answers with one of its +// own and drops the connection at once, so the relay reading from the origin is the +// first to see its connection gone — timing alone would blame the origin. The client +// sent the first close, so the turn is the client's: 499, with the code it closed with. +func TestResponsesWebSocket_FirstCloseFrameDecidesWhoEndedTheTurn(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { + _ = conn.Close() + return + } + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.created","response":{"id":"r","status":"in_progress"}}`), false)) + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(`{"type":"response.output_text.delta","item_id":"m","output_index":0,"content_index":0,"delta":"par"}`), false)) + // Answer the client's close with our own, then drop the connection at once. + for { + f, err := fr.ReadFrame() + if err != nil { + _ = conn.Close() + return + } + if f.Opcode == wsframe.OpClose { + _, _ = conn.Write(wsframe.Encode(wsframe.OpClose, f.Payload, false)) + _ = conn.Close() + return + } + } + }) + c := h.dial(t, nil) + c.send(responseCreate("hello")) + c.readUntil("response.output_text.delta") + _, _ = c.conn.Write(wsframe.Encode(wsframe.OpClose, wsframe.ClosePayload(1001, "user left"), true)) + + var res *pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionResponse { + e := e + res = &e + } + } + _ = c.conn.Close() + if res == nil { + t.Fatal("the closed turn recorded no response row") + } + if res.StatusCode != pipeline.StatusClientClosedRequest || res.Error == nil || res.Error.Kind != "client_canceled" { + t.Fatalf("turn = status %d error %+v, want 499 client_canceled: the client sent the first close", res.StatusCode, res.Error) + } + if !strings.Contains(res.Error.Message, "close code 1001 going away: user left") { + t.Errorf("message %q does not name the client's close code and reason", res.Error.Message) + } +} + +// A refused turn's status follows the error event's own status when it has one, and +// otherwise the status its error code is answered with over HTTPS — so a quota refusal, +// which arrived live with no status, counts as the 429 it is and not as a 500. +func TestErrorEventStatus(t *testing.T) { + cases := []struct { + name string + status int + code, errType string + want int + }{ + {"the event's own status wins", 429, "server_is_overloaded", "", 429}, + {"quota, no status", 0, "insufficient_quota", "", 429}, + {"quota as the type", 0, "", "insufficient_quota", 429}, + {"rate limit", 0, "rate_limit_exceeded", "", 429}, + {"usage limit", 0, "", "usage_limit_reached", 429}, + {"socket limit", 0, "websocket_connection_limit_reached", "", 429}, + {"context too long", 0, "context_length_exceeded", "", 400}, + {"stale previous response", 0, "previous_response_not_found", "", 400}, + {"overloaded", 0, "server_is_overloaded", "", 503}, + {"unknown code", 0, "something_new", "", 500}, + {"nothing at all", 0, "", "", 500}, + {"a non-error status is not taken", 200, "insufficient_quota", "", 429}, + } + for _, tc := range cases { + if got := errorEventStatus(tc.status, tc.code, tc.errType); got != tc.want { + t.Errorf("%s: errorEventStatus(%d, %q, %q) = %d, want %d", tc.name, tc.status, tc.code, tc.errType, got, tc.want) + } + } +} + +// End to end: a quota refusal with no status, as one arrived live, is recorded as 429 +// with its code named, and the client receives the event unchanged. +func TestResponsesWebSocket_QuotaRefusalIsRecordedAs429(t *testing.T) { + const refusal = `{"type":"error","error":{"type":"insufficient_quota","code":"insufficient_quota","message":"You exceeded your current quota."}}` + h := newWSHarness(t, config.ResponsesWebSocketsParse) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 1<<20, 1<<20) + if _, err := fr.ReadFrame(); err != nil { + return + } + _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(refusal), false)) + _, _ = fr.ReadFrame() + }) + c := h.dial(t, nil) + c.send(responseCreate("hello")) + got := c.readUntil("error") + c.close() + if b, _ := json.Marshal(got[len(got)-1]); !strings.Contains(string(b), `"insufficient_quota"`) { + t.Errorf("client received %s, want the refusal unchanged", b) + } + + var res *pipeline.SessionEvent + for _, e := range h.events(t, codexSID, true) { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionResponse { + e := e + res = &e + } + } + if res == nil || res.StatusCode != http.StatusTooManyRequests { + t.Fatalf("refused turn = %+v, want status 429", res) + } + if res.Error == nil || !strings.Contains(res.Error.Message, "insufficient_quota") { + t.Errorf("refused turn error %+v does not name the code", res.Error) + } +} + // Codex's in-band interrupt: for steering a turn already under way, or for queued // mailbox input, on models with use_responses_lite set, it sends response.interrupt // naming the response and waits on the same socket for the server's diff --git a/core/listener/internal/wsframe/wsframe.go b/core/listener/internal/wsframe/wsframe.go index 261e7e0d3..2d8388ac8 100644 --- a/core/listener/internal/wsframe/wsframe.go +++ b/core/listener/internal/wsframe/wsframe.go @@ -210,6 +210,49 @@ func ClosePayload(code uint16, reason string) []byte { return append(b, reason...) } +// closeCodeNames are the RFC 6455 §7.4.1 codes a close frame commonly carries. +var closeCodeNames = map[uint16]string{ + 1000: "normal closure", + 1001: "going away", + 1002: "protocol error", + 1003: "unsupported data", + 1007: "invalid payload", + 1008: "policy violation", + 1009: "message too big", + 1010: "missing extension", + 1011: "internal error", + 1012: "service restart", + 1013: "try again later", +} + +// DescribeClose renders a close frame's payload for a log or a recorded row: its code, +// the code's RFC 6455 name when it has one, and the sender's reason. "close without a +// code" for an empty payload, which RFC 6455 allows. The reason is the sender's own +// text, so it is kept to printable ASCII and cut to fit a short line. +func DescribeClose(payload []byte) string { + if len(payload) < 2 { + return "close without a code" + } + code := binary.BigEndian.Uint16(payload[:2]) + out := fmt.Sprintf("close code %d", code) + if name, ok := closeCodeNames[code]; ok { + out += " " + name + } + var reason []byte + for _, b := range payload[2:] { + if len(reason) >= 80 { + break + } + if b >= 0x20 && b < 0x7f { + reason = append(reason, b) + } + } + if len(reason) > 0 { + out += ": " + string(reason) + } + return out +} + // Message reassembles a data message from its frames. // // It sees control frames go by without letting them interrupt a message, as RFC 6455 diff --git a/core/listener/internal/wsframe/wsframe_test.go b/core/listener/internal/wsframe/wsframe_test.go index 9f3497ba0..af91fd00f 100644 --- a/core/listener/internal/wsframe/wsframe_test.go +++ b/core/listener/internal/wsframe/wsframe_test.go @@ -214,3 +214,21 @@ func TestClosePayload_FitsAControlFrame(t *testing.T) { t.Errorf("code bytes = % x, want 03 f0 (1008)", p[:2]) } } + +func TestDescribeClose(t *testing.T) { + cases := map[string]struct { + payload []byte + want string + }{ + "no code": {nil, "close without a code"}, + "normal": {ClosePayload(1000, ""), "close code 1000 normal closure"}, + "internal, reason": {ClosePayload(1011, "boom"), "close code 1011 internal error: boom"}, + "unnamed code": {ClosePayload(4000, ""), "close code 4000"}, + "escape in reason": {ClosePayload(1008, "a\x1b[2Jb"), "close code 1008 policy violation: a[2Jb"}, + } + for name, tc := range cases { + if got := DescribeClose(tc.payload); got != tc.want { + t.Errorf("%s: DescribeClose = %q, want %q", name, got, tc.want) + } + } +} diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 8d888fbd6..26854bbb1 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -225,6 +225,16 @@ with the text that had arrived, as `statusCode` 499 with `error.kind: upstream failure. OpenAI sends no usage for a turn cut off this way, so it has no token count or cost: it shows as priceable but unpriced, the same as a cancelled HTTPS stream. +**A turn's status follows what the same turn would get over HTTPS.** An open WebSocket +carries no HTTP status after its 101, so each turn row is given the one its HTTPS +equivalent would have: 200 for a completed turn; for a refusal the origin sends as an +`error` event, the event's own status, else the one its error code is answered with +(`insufficient_quota` and `rate_limit_exceeded` are 429, `context_length_exceeded` 400, +`server_is_overloaded` 503), else 500; 499 `client_canceled` when Codex ended the socket +mid-turn and 502 `upstream_error` when OpenAI did. Who ended it is read from the first +close frame either side sent, whose code and reason go into the row's message, and from +which side's connection dropped first when neither sent one. + Where Cortex cannot carry the socket it answers the handshake with **426 Upgrade Required** instead, which Codex reads as "this endpoint wants HTTP" and switches to HTTPS at once, for the rest of its session. That happens for a connection the bridge did not From e72133a043e2e976c82489c9b1ea2372a145472a Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Tue, 27 Oct 2026 09:49:46 -0400 Subject: [PATCH 06/12] docs: Record the API-key and Linux WebSocket runs in the Codex doc Verified depth gains two entries: an API-key login on macOS, whose socket goes to api.openai.com, and Linux (Ubuntu 26.10 arm64, Codex 0.164.0), where a single turn, three turns on one socket and a TUI turn were each carried and parsed with tokens and a cost, and Codex's traffic was attributed to its process. Not tested is narrowed to match: Linux now means what was not run there (configure, dotenv routing, the HTTPS fallback), the version range runs to 0.164.0, and the public /v1/responses API entry covers HTTPS only. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- docs/agents/codex.md | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 26854bbb1..224e2f49f 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -369,27 +369,35 @@ with Cortex built from the change that added this page. This was exercised: tokens and a cost, as was the TUI's own model check on `gpt-5.6-luna`; - **Esc mid-answer in the TUI**: the WebSocket closed within milliseconds, the interrupted turn was recorded with the 1,491 characters that had arrived, and the - next prompt opened a new socket and was answered and parsed normally. + next prompt opened a new socket and was answered and parsed normally; +- **an API-key login**, on macOS: Codex opened the WebSocket to `api.openai.com` + instead of `chatgpt.com`, and a single-turn and a three-turn run were each parsed + with model (`gpt-6.1-sol`), tokens and a cost; +- **Linux**, on Codex 0.164.0 in an Ubuntu 26.10 arm64 VM with an API-key login: a + single turn, three turns on one socket, and a TUI turn, each carried and parsed with + tokens and a cost, and Codex's traffic attributed to its process and filed under its + own session. Not tested: -- **Linux.** Everything here was exercised on macOS only. -- **An API-key login over the WebSocket.** Codex then opens it to `api.openai.com` - rather than `chatgpt.com`; Cortex treats the two hosts alike, but only the ChatGPT - login was watched. +- **Linux beyond the WebSocket.** On Linux only the WebSocket transport above was + exercised; `agentop configure codex`, the dotenv routing and the HTTPS fallback were + tested on macOS only. - **Codex's in-band interrupt.** For some internal cases — steering a turn already under way, or queued "mailbox" input, on models with `use_responses_lite` set — Codex sends `response.interrupt` on the open socket and waits for the server's `response.interrupt.accepted` instead of closing it. The relay forwards both unchanged and its tests cover the exchange in both orders, but it was not seen live: Esc does not use it. -- **Codex versions other than 0.160.1 and 0.162.1.** Both appeared in one capture (an - auto-update landed mid-session) and both parse identically, which is some evidence the - User-Agent shape is stable across releases rather than pinned to one build — but the - `.env` loading behaviour and the `CODEX_` filter were only read from 0.160.1's source. -- **The public `/v1/responses` API** (an OpenAI API key rather than a ChatGPT - subscription). The parser recognises the path, and the non-streaming shape of that - dialect is modelled from the published schema rather than from captured traffic. +- **Codex versions other than 0.160.1 through 0.164.0.** 0.160.1 and 0.162.1 appeared in + one capture (an auto-update landed mid-session) and parse identically, and the + WebSocket transport was run on 0.163.0 and 0.164.0, which is some evidence the wire + shape is stable across releases rather than pinned to one build — but the `.env` + loading behaviour and the `CODEX_` filter were only read from 0.160.1's source. +- **The public `/v1/responses` API over HTTPS.** An API-key login was exercised over the + WebSocket above; over HTTPS — with `responses_websockets: decline`, or a client that + does not ask for the socket — it was not, and the non-streaming shape of that dialect + is modelled from the published schema rather than from captured traffic. - **A `function_call` tool call.** Codex's own `exec` tool is a `custom_tool_call`; the `function_call` shape is modelled from the documented schema. - **Tool pruning**, which is unsupported above rather than merely untested. From a221f0f793834a7ac6282a67dd7b96dad7588b65 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Wed, 28 Oct 2026 12:36:32 -0400 Subject: [PATCH 07/12] test: Give the large-WebSocket-turn test room on a slow runner The test sent a 15 MB turn and gave the client 15s to see the reply. Under -race that turn takes ~3.4s on a Mac and ~6.2s on Linux, and on a shared CI runner it went past 15s twice. The size is what it tests, so send just over the 10 MB per-event cap instead (11 MiB, ~5s on Linux under -race at 2 CPUs) and allow 60s. Measured on Linux with Go 1.26.10 and -race, the old version never hung: it was only slow. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/websocket_test.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 1f3851757..716492db3 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -999,10 +999,15 @@ func TestResponsesWebSocket_ReloadMidTurnFinishesOnTheTurnsPipeline(t *testing.T // A socket's first turn carries the whole conversation, so it can be far over the // 10 MB per-event cap the origin's messages get. It is parsed up to the same 32 MiB an // HTTPS request body is. +// +// The turn is just over that 10 MB, not far over it: the size is what the test is about, +// and parsing it is the costliest thing in this file — under -race it is ~3s on a laptop +// and was over 15s on a shared CI runner. The read deadline is raised for the same reason. func TestResponsesWebSocket_LargeFirstTurnIsParsed(t *testing.T) { h := newWSHarness(t, config.ResponsesWebSocketsParse) c := h.dial(t, nil) - big := strings.Repeat("x", 15<<20) + _ = c.conn.SetDeadline(time.Now().Add(60 * time.Second)) + big := strings.Repeat("x", maxBodySize+1<<20) c.send(responseCreate(big)) c.readUntil("response.completed") c.close() @@ -1015,7 +1020,7 @@ func TestResponsesWebSocket_LargeFirstTurnIsParsed(t *testing.T) { } } if req == nil || req.Inference == nil || req.Inference.Model != "gpt-6-luna" { - t.Fatalf("a 15 MB first turn was not parsed: %+v", req) + t.Fatalf("an 11 MB first turn was not parsed: %+v", req) } } From a2d37f558869ce204be7cd92e896106cf8d6e12c Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Wed, 28 Oct 2026 15:38:37 -0400 Subject: [PATCH 08/12] fix: Fix the WebSocket relay's review findings and CodeQL alerts From the three-way review of #1467: Must-fix - Turns end where their answer is parsed. relayDown alone ends a turn, from a queue in send order, so a response.create arriving while the previous turn's response.completed is still being parsed no longer ends that turn with no usage, and a late terminal event can no longer land on the next turn. This also removes the data race between the two relays on the turn's context. - A message over its limit ends the socket with close code 1009, as HTTPS answers 413, instead of being relayed unread past every plugin and recorded nowhere. Fragments are capped at 1024, so a flood of empty continuations cannot be held without bound either. - The handshake's Finishers run on the pipeline the handshake ran on. The deferred RunFinish loaded the pipeline current when the socket closed, so a reload while Codex was connected finished plugins that never saw the request, or indexed past a shorter pipeline and panicked. handleConnect had the same shape; both now pin the pipeline. - A pipeline with a WritesDestination plugin (inference-router) is declined with a 426, so Codex routed to another server keeps being routed over HTTPS instead of going to OpenAI on the socket. Suggestions - A reload that adds a response-body writer or a destination writer ends an open socket at its next turn with close code 1013, so Codex reconnects and is declined, rather than ignoring the plugin. - A prewarm (generate:false) runs the request phase: it carries the prompt the next turn builds on. It still records no turn. - Each turn gets a fresh session view, as each HTTPS request does. - An Esc'd turn is recorded as HTTPS records a stream the client leaves part way: its 200, with error.kind turn_interrupted, which /v1/usage does not count as an error. - An origin's error event gives Finishers an allow with its status, as the same refusal over HTTPS does, not an error with none. - Mid-turn frames are marked delivered before dispatch, and turns run under a context cancelled when the socket ends. - An origin's refusal body is no longer cut off by the handshake deadline. - The codec grows a frame's buffer as bytes arrive instead of sizing it from the claimed length, copies a one-frame message once, and cuts a close reason at a UTF-8 boundary. - The removed listener.decline_responses_websockets is refused at startup with the matching responses_websockets mode named. CodeQL: log values a client chose are passed through logValue, which replaces line breaks; the socket's session id is cleaned where the socket takes it; Encode bounds its capacity arithmetic. Tests: one per reproduced defect, each confirmed to fail on the code before the fix. The weak assertions the review found are fixed: the rejected-turn check waits for the socket to end, the reload check reads under its lock, and events() fails on timeout. Docs and comments that were wrong (Esc as 499, the unreachable HTTP/2 decline, tool-prune as the example rewriter, the upgrade marker on the handshake row) are corrected. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- CLAUDE.md | 5 +- core/config/config.go | 10 +- core/config/config_test.go | 25 ++ core/config/validate.go | 9 + core/cost/usage/usage.go | 2 +- core/listener/forwardproxy/server.go | 33 +- core/listener/forwardproxy/websocket.go | 264 +++++++++++----- .../forwardproxy/websocket_review_test.go | 294 ++++++++++++++++++ core/listener/forwardproxy/websocket_test.go | 26 +- core/listener/internal/wsframe/wsframe.go | 122 ++++++-- .../listener/internal/wsframe/wsframe_test.go | 144 +++++++-- core/listener/parity/dispatch_audit_test.go | 6 +- core/pipeline/context.go | 5 +- core/pipeline/pipeline.go | 13 + core/pipeline/session.go | 6 + docs/agents/codex.md | 32 +- docs/framework-architecture.md | 2 +- docs/session-api.md | 11 +- 18 files changed, 841 insertions(+), 168 deletions(-) create mode 100644 core/listener/forwardproxy/websocket_review_test.go diff --git a/CLAUDE.md b/CLAUDE.md index 3a835132f..4761820d7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -606,8 +606,9 @@ Every event on `/v1/sessions/{id}` and `/v1/events` carries: 426 sends it straight to HTTPS. Nothing is forwarded, the Finisher sees an allow with status 426, and `/v1/usage` does not count it as an error. It is what `listener.responses_websockets: decline` does to every such socket, and what the - default, `parse`, does to one it cannot carry (not bridged, HTTP/2, or a - `WritesResponseBody` plugin in the chain); a carried socket records instead a `GET` + default, `parse`, does to one it cannot carry (not bridged, or a `WritesResponseBody` + or `WritesDestination` plugin in the chain; an HTTP/2 client is refused 400 by Go's + HTTP/2 server before the handler runs); a carried socket records instead a `GET` open and close with `tunnelReason: "websocket"` and one ordinary `POST` request/response pair per turn on it (`core/listener/forwardproxy/websocket.go`). `forward` restores the stripped GET and the 405. **A request whose body could not be diff --git a/core/config/config.go b/core/config/config.go index 5802e259b..01077b025 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -941,14 +941,20 @@ type ListenerConfig struct { // Empty means the default, ResponsesWebSocketsParse. See // ResponsesWebSocketsMode. ResponsesWebSockets string `yaml:"responses_websockets,omitempty" json:"responses_websockets,omitempty"` + + // LegacyDeclineResponsesWebSockets is the boolean responses_websockets replaced, + // read only so a config that still sets it is refused with the new key named. It + // merged in #1462 and never shipped in a release, but someone who set it to false + // to forward the upgrade would otherwise silently get parse. Never written. + LegacyDeclineResponsesWebSockets *bool `yaml:"decline_responses_websockets,omitempty" json:"-"` } // Valid ListenerConfig.ResponsesWebSockets values. const ( // ResponsesWebSocketsParse carries the WebSocket and parses every turn on it the // way an HTTPS Responses request is parsed. Where the proxy cannot do that — the - // connection is HTTP/2, the request was not decrypted, or a plugin rewrites - // response bodies — it falls back to ResponsesWebSocketsDecline for that socket. + // request was not decrypted, or a plugin rewrites response bodies or redirects + // requests — it falls back to ResponsesWebSocketsDecline for that socket. ResponsesWebSocketsParse = "parse" // ResponsesWebSocketsDecline answers the upgrade with 426 Upgrade Required and // contacts no origin; Codex reads that as "use HTTPS" and switches at once. diff --git a/core/config/config_test.go b/core/config/config_test.go index 66e55e854..3cf786e42 100644 --- a/core/config/config_test.go +++ b/core/config/config_test.go @@ -1094,6 +1094,31 @@ func TestValidate_ResponsesWebSockets(t *testing.T) { } } +// The boolean responses_websockets replaced is refused with the mode that matches it +// named, not ignored: someone who set it false to forward the upgrade would otherwise get +// parse with no word said. +func TestLoad_RemovedDeclineKeyIsRefusedWithItsReplacement(t *testing.T) { + for _, tc := range []struct { + value, want string + }{ + {"false", "responses_websockets: " + ResponsesWebSocketsForward}, + {"true", "responses_websockets: " + ResponsesWebSocketsDecline}, + } { + path := filepath.Join(t.TempDir(), "c.yaml") + body := "mode: proxy-sidecar\nlistener:\n roles: [forward]\n decline_responses_websockets: " + tc.value + "\n" + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := Load(path) + if err == nil { + err = validateListeners(cfg) + } + if err == nil || !strings.Contains(err.Error(), tc.want) { + t.Errorf("decline_responses_websockets: %s: err = %v, want one naming %q", tc.value, err, tc.want) + } + } +} + // TestLoad_PreAffinityLaptopConfigGetsClientAffinity loads a laptop config written before // session.client_affinity existed — no session block at all, which is what every install // made before it still has on disk. With affinity off, Bob Shell's header-less startup diff --git a/core/config/validate.go b/core/config/validate.go index 52c27c3d2..0b8f5a6e7 100644 --- a/core/config/validate.go +++ b/core/config/validate.go @@ -112,6 +112,15 @@ func validateListeners(cfg *Config) error { return fmt.Errorf("listener.inbound_interception: %q is not valid (use %q or %q)", cfg.Listener.InboundInterception, InboundInterceptionReverseProxy, InboundInterceptionTransparent) } + if v := cfg.Listener.LegacyDeclineResponsesWebSockets; v != nil { + want := ResponsesWebSocketsDecline + if !*v { + want = ResponsesWebSocketsForward + } + return fmt.Errorf("listener.decline_responses_websockets was replaced by listener.responses_websockets; "+ + "for decline_responses_websockets: %t use responses_websockets: %s (the default is %s)", + *v, want, ResponsesWebSocketsParse) + } switch cfg.Listener.ResponsesWebSockets { case "", ResponsesWebSocketsParse, ResponsesWebSocketsDecline, ResponsesWebSocketsForward: // valid diff --git a/core/cost/usage/usage.go b/core/cost/usage/usage.go index 9f7b20867..b7ce628df 100644 --- a/core/cost/usage/usage.go +++ b/core/cost/usage/usage.go @@ -1271,7 +1271,7 @@ func (a *Aggregator) evictColdestLocked() { // design and on the default path, so counting it would give every healthy Codex // fleet a nonzero error rate. The Finisher sees it as an allow for the same reason. func countsAsError(e *pipeline.SessionEvent) bool { - if e.Error != nil && e.Error.Kind == pipeline.ErrorKindUpgradeDeclined { + if e.Error != nil && (e.Error.Kind == pipeline.ErrorKindUpgradeDeclined || e.Error.Kind == pipeline.ErrorKindTurnInterrupted) { return false } return e.StatusCode >= 400 || e.Phase == pipeline.SessionDenied diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index f0ab7fca8..b4bc183bb 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -405,6 +405,13 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne skipped = true } + // The pipeline this request runs on, from Run to its Finishers. RunFinish replays the + // plugins Run dispatched, by position, so it must be handed the same pipeline: a + // reload between the two would finish plugins that never saw the request, or index + // past the end of a shorter chain. An HTTPS request spans a reload only by bad luck; + // a carried WebSocket's handshake spans the whole Codex session, so it always can. + p := s.OutboundPipeline.Load() + // Finisher dispatch runs after every exit path. RunFinish is a // no-op when pctx.dispatched is empty (pre-pipeline rejects), so // this defer is safe on every path including the body-too-large @@ -412,7 +419,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne // this request and there is nothing to finalize. if !skipped { defer func() { - s.OutboundPipeline.RunFinish(r.Context(), pctx, pipeline.OutcomeFromContext(pctx)) + p.RunFinish(r.Context(), pctx, pipeline.OutcomeFromContext(pctx)) }() } @@ -450,7 +457,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne } if !skipped { - action := s.OutboundPipeline.Run(r.Context(), pctx) + action := p.Run(r.Context(), pctx) // Send the request where the pipeline pointed it, before any row is recorded, // so the row names the host the bytes went to. The target comes from @@ -501,16 +508,19 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne if !skipped && wantsResponsesWebSocket(r) { switch s.ResponsesWebSockets { case config.ResponsesWebSocketsParse: - sid := s.recordingSessionID(sessionID, r.Header, chain) - if why := s.cannotCarryWebSocket(r, tl); why != "" { - slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", r.Host, "reason", why) + // logValue changes nothing a header id can hold — IDHeaderFrom already refused + // control characters — but the socket carries this id into rows and log lines + // for a whole session, so it is cleaned where the socket takes it. + sid := logValue(s.recordingSessionID(sessionID, r.Header, chain)) + if why := s.cannotCarryWebSocket(tl, p); why != "" { + slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", logValue(r.Host), "reason", why) s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) return } - s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain) + s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain, p) return case config.ResponsesWebSocketsDecline: - s.declineResponsesWebSocket(w, r, pctx, tl, s.recordingSessionID(sessionID, r.Header, chain), chain) + s.declineResponsesWebSocket(w, r, pctx, tl, logValue(s.recordingSessionID(sessionID, r.Header, chain)), chain) return } } @@ -1053,7 +1063,7 @@ const declinedUpgradeBody = `{"detail":"WebSocket transport is not available thr // does for the 421. func (s *Server) declineResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc) { slog.Debug("forward-proxy: declined a WebSocket upgrade on a Responses path; the client falls back to HTTPS", - "host", r.Host, "path", r.URL.Path) + "host", logValue(r.Host), "path", logValue(r.URL.Path)) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusUpgradeRequired) _, _ = io.WriteString(w, declinedUpgradeBody) @@ -2079,9 +2089,12 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { var sessionID string var chain []session.Proc + // One pipeline from Run to RunFinish, for serveOutbound's reason: a tunnel can stay + // open across any number of reloads. + p := s.OutboundPipeline.Load() if !skipped { defer func() { - s.OutboundPipeline.RunFinish(r.Context(), pctx, pipeline.OutcomeFromContext(pctx)) + p.RunFinish(r.Context(), pctx, pipeline.OutcomeFromContext(pctx)) }() // Same one-identity rule as the request path, and the same local-variable @@ -2099,7 +2112,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // Run the outbound pipeline. Plugins that policy on host/identity // (ibac, content gates) still get to allow/deny; plugins that need // HTTP body (parsers) see no body, which they handle gracefully. - action := s.OutboundPipeline.Run(r.Context(), pctx) + action := p.Run(r.Context(), pctx) if action.Type == pipeline.Reject { s.recordOutboundReject(pctx, action, s.tunnelSessionID(sessionID, r.Header, chain)) // Render as a JSON-RPC error frame when the rejected diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index 89959daa6..d8298da43 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -56,8 +56,9 @@ const ( // parsing, each direction matching what the HTTPS path parses: maxRequestBodySize // for what the client sends — a socket's first turn carries the whole // conversation, the same body an HTTPS POST would — and maxBodySize for each event - // the origin sends back, the SSE path's per-frame cap. A larger message is relayed - // whole and not read: its turn is still recorded, just without a parse. + // the origin sends back, the SSE path's per-frame cap. A larger message ends the + // socket with close code 1009, as the HTTPS path refuses an oversized body with a + // 413: relaying it unread would carry a turn no plugin saw and no row records. wsClientMessageLimit = maxRequestBodySize wsServerMessageLimit = maxBodySize // wsFrameLimit is the largest single frame relayed. Raw is buffered whole before @@ -126,24 +127,34 @@ func upgradesToWebSocket(h http.Header) bool { // can. Each reason is a shape the relay below does not handle; the caller declines // those with a 426 rather than forwarding them, so a parse-mode client is never left // sitting through the origin's 405. -func (s *Server) cannotCarryWebSocket(r *http.Request, tl *tunnelLog) string { - switch { - case tl == nil || s.TLSBridge == nil: +// +// An HTTP/2 client never gets here: x/net/http2 refuses any request carrying Upgrade +// with a 400 before the handler runs, so that is what an h2 client sees. +func (s *Server) cannotCarryWebSocket(tl *tunnelLog, p *pipeline.Pipeline) string { + if tl == nil || s.TLSBridge == nil { // A plaintext proxied request: no TLS bridge to re-originate over, and no // Responses endpoint is served over plain HTTP. return "not a bridged request" - case r.ProtoMajor != 1: - // An h2 extended CONNECT (RFC 8441) carries the socket as a stream, not a - // connection to take over. tungstenite, Codex's client, does not send one. - return "not HTTP/1.1" - case s.OutboundPipeline.WritesResponseBody(): + } + return pipelineCannotCarry(p) +} + +// pipelineCannotCarry says why p cannot run a carried socket's turns, or "" when it +// can. Checked at the handshake and again at every turn, since a hot reload can change +// the answer under a socket that is already open. +func pipelineCannotCarry(p *pipeline.Pipeline) string { + switch { + case p.WritesResponseBody(): // A response-body writer rewrites a whole response; a turn's response is a // live stream of frames already on their way to the client. The SSE path // buffers instead, and a WebSocket has no buffered form. return "a plugin rewrites response bodies" + case p.WritesDestination(): + // A destination writer (inference-router) sends a request to another server. + // A turn rides a socket already connected to OpenAI, so it cannot be sent + // elsewhere; over HTTPS each turn is its own request, which can. + return "a plugin redirects requests" } - // A connection that cannot be taken over is the one case left, and only the - // ResponseWriter can say; serveResponsesWebSocket declines it there. return "" } @@ -153,7 +164,7 @@ func (s *Server) cannotCarryWebSocket(r *http.Request, tl *tunnelLog) string { // pctx is the handshake's, after the pipeline ran on it. Its rows are the socket's open // (now) and close (when it ends), sharing its RequestID the way an opaque tunnel's do; // the turns record their own pairs between them. -func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc) { +func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc, p *pipeline.Pipeline) { hj, ok := w.(http.Hijacker) if !ok { s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) @@ -178,10 +189,10 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, client := s.TLSBridge.Upstream resp, err := client.Do(up) - cancel() if err != nil { + cancel() status, fail := pipeline.ExchangeFailure(r.Context(), err) - slog.Warn("forward-proxy: WebSocket handshake with the origin failed", "host", r.Host, "kind", fail.Kind, "error", fail.Message) + slog.Warn("forward-proxy: WebSocket handshake with the origin failed", "host", logValue(r.Host), "kind", fail.Kind, "error", logValue(fail.Message)) http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) s.recordWebSocketRefusal(pctx, tl, sid, chain, status, fail) return @@ -189,6 +200,9 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, if resp.StatusCode != http.StatusSwitchingProtocols { // The origin refused the upgrade (401, 403, 426…): relay its answer as an // ordinary response, so the client sees what it would have without Cortex. + // The handshake's deadline still covers the body: cancelled before the copy, it + // would cut the answer off part way. + defer cancel() defer resp.Body.Close() for k, vs := range resp.Header { for _, v := range vs { @@ -202,6 +216,8 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, s.recordWebSocketRefusal(pctx, tl, sid, chain, resp.StatusCode, nil) return } + // A 101: the deadline was for the handshake, and the socket outlives it. + cancel() upConn, ok := resp.Body.(io.ReadWriteCloser) if !ok { _ = resp.Body.Close() @@ -221,7 +237,7 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, clientConn, clientBuf, err := hj.Hijack() if err != nil { _ = upConn.Close() - slog.Warn("forward-proxy: could not take over the client connection for a WebSocket", "host", r.Host, "error", err) + slog.Warn("forward-proxy: could not take over the client connection for a WebSocket", "host", logValue(r.Host), "error", logValue(err.Error())) return } // The bridge's handshake deadline was cleared before serving; nothing sets one @@ -248,10 +264,12 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, } pctx.StatusCode = http.StatusSwitchingProtocols + ctx, stop := context.WithCancel(context.WithoutCancel(r.Context())) + defer stop() sock := &wsSocket{ s: s, r: r, hs: pctx, tl: tl, sid: sid, chain: chain, client: clientConn, clientR: clientBuf.Reader, up: upConn, - start: time.Now(), + start: time.Now(), ctxv: ctx, stop: stop, hsp: p, } s.recordWebSocketOpen(sock) sock.run() @@ -307,10 +325,23 @@ type wsSocket struct { // written in place of the client's frame. upMu sync.Mutex - // turnMu guards turn, which both relays touch: the client side opens one, the - // server side feeds and ends it. + // turnMu guards turns, which both relays touch. The client side appends a turn + // when it sends a response.create; the server side alone ends one, on that turn's + // terminal event. A turn is ended where its last event is parsed, never from the + // client side, so a response.create that arrives while relayDown is still parsing + // the previous turn's response.completed cannot end that turn early, and an answer + // cannot be fed to a turn it does not belong to: the server answers turns in the + // order they were sent, which is the order of this queue. turnMu sync.Mutex - turn *wsTurn + turns []*wsTurn + + // ctxv is the context every turn's plugins run under: the handshake's values + // without its cancellation, cancelled by stop when the socket ends, so a plugin + // waiting on something (a budget pause, a slow script) is released with it. + ctxv context.Context + stop context.CancelFunc + // hsp is the pipeline the handshake ran on. + hsp *pipeline.Pipeline // endedBy is the side whose relay stopped first, "client" or "origin", which // decides how a turn still open at that moment is recorded. Set once, by the @@ -349,6 +380,9 @@ type wsTurn struct { // window, so a turn looking the pipeline up again at each step would end on one // whose parser never saw it begin — no tokens, no cost. p *pipeline.Pipeline + // prewarm is a generate:false response.create: run through the request phase, not + // recorded, and its answer not parsed. + prewarm bool } func (ws *wsSocket) run() { @@ -361,28 +395,33 @@ func (ws *wsSocket) run() { _ = ws.up.Close() <-done + // Both relays have returned, so nothing else touches the queue: every turn still + // in it was sent and never answered in full. ws.turnMu.Lock() - t := ws.turn - ws.turn = nil + open := ws.turns + ws.turns = nil ws.turnMu.Unlock() - if t != nil { - status, fail := ws.unfinishedTurn() + status, fail := ws.unfinishedTurn() + for _, t := range open { ws.endTurn(t, status, fail) } ws.recordClose() + // Released last: endTurn's teardown contexts derive from it. + ws.stop() } // unfinishedTurn is how a turn still open when the socket ended is recorded, by who // ended it — the WebSocket counterpart of what the HTTPS path records for a stream that // stops before its last event. // -// THE CLIENT CLOSING IS NOT AN UPSTREAM FAILURE. It is what Codex does on Esc: it drops -// the turn's stream, which closes this socket, and opens another for the next prompt. -// So it is recorded as the HTTPS path records a client that hangs up before its answer -// — 499, client_canceled — which /v1/usage counts as an error but keeps out of the 5xx -// series, rather than as a 502-shaped upstream_error that would read as OpenAI failing -// every time someone interrupts a turn. The origin closing first, or the socket breaking -// on a frame the relay cannot read, is what upstream_error is for. +// THE CLIENT CLOSING IS NOT A FAILURE. It is what Codex does on Esc: it drops the turn's +// stream, which closes this socket, and opens another for the next prompt. Over HTTPS +// the same Esc ends a stream the origin had already answered 200, and that is how it is +// recorded — the upstream's 200 and no error — so an interrupted turn counts in +// /v1/usage as priceable-but-unpriced rather than as an error. A turn ended here is +// recorded the same way: 200, with the reason in Message only, under a kind that +// countsAsError leaves alone. The origin closing first, or the socket breaking on a +// frame the relay cannot read, is what upstream_error is for. // // A close frame decides first, when either side sent one: whoever sent the first close // ended the socket, and its code and reason go into the message. Without one — a @@ -400,8 +439,8 @@ func (ws *wsSocket) unfinishedTurn() (int, *pipeline.EventError) { suffix = " (" + note + ")" } if side == "client" { - return pipeline.StatusClientClosedRequest, &pipeline.EventError{ - Kind: "client_canceled", + return http.StatusOK, &pipeline.EventError{ + Kind: pipeline.ErrorKindTurnInterrupted, Message: "the client closed the WebSocket before the turn completed" + suffix, } } @@ -492,9 +531,14 @@ func (ws *wsSocket) relayUp() (endedBy string) { } op, payload, readable, complete, err := msg.Add(f) if err != nil { + if errors.Is(err, wsframe.ErrMessageTooBig) { + // Neither relayed unread nor held any longer: see wsClientMessageLimit. + ws.closeBoth(wsframe.CloseMessageTooBig, "message too big for the proxy to read") + } ws.logEnd("client", err) return "client" } + // Bounded by Message: it ends the message before its frames pass the limit. held = append(held, f.Raw) if !complete { continue @@ -546,6 +590,9 @@ func (ws *wsSocket) relayDown() (endedBy string) { } op, payload, readable, complete, err := msg.Add(f) if err != nil { + if errors.Is(err, wsframe.ErrMessageTooBig) { + ws.closeBoth(wsframe.CloseMessageTooBig, "message too big for the proxy to read") + } ws.logEnd("origin", err) return "origin" } @@ -555,6 +602,15 @@ func (ws *wsSocket) relayDown() (endedBy string) { } } +// closeBoth sends each end a close frame with code and reason before the socket is torn +// down, so both see why rather than a dropped connection. Best effort: a write that fails +// is the connection already gone. +func (ws *wsSocket) closeBoth(code uint16, reason string) { + payload := wsframe.ClosePayload(code, reason) + ws.writeClient(wsframe.Encode(wsframe.OpClose, payload, false)) + ws.writeUp(wsframe.Encode(wsframe.OpClose, payload, true)) +} + func (ws *wsSocket) writeUp(b []byte) bool { ws.upMu.Lock() defer ws.upMu.Unlock() @@ -574,10 +630,10 @@ func (ws *wsSocket) logEnd(side string, err error) { return } if errors.Is(err, wsframe.ErrProtocol) { - slog.Warn("forward-proxy: ending a WebSocket on a frame the relay cannot read", "host", ws.r.Host, "side", side, "error", err) + slog.Warn("forward-proxy: ending a WebSocket on a frame the relay cannot read", "host", logValue(ws.r.Host), "side", side, "error", logValue(err.Error())) return } - slog.Debug("forward-proxy: WebSocket ended", "host", ws.r.Host, "side", side, "error", err) + slog.Debug("forward-proxy: WebSocket ended", "host", logValue(ws.r.Host), "side", side, "error", logValue(err.Error())) } // wsEnvelope is the one field every message on the socket carries. @@ -595,41 +651,43 @@ func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool) return nil, false } body, generate := responseCreateBody(payload) - if !generate { - // A prewarm: generate:false asks the server to set up state and answer at - // once, with no model output. Not inference, so not a turn. - return nil, false - } - ws.turnMu.Lock() - prev := ws.turn - ws.turn = nil - ws.turnMu.Unlock() - if prev != nil { - // Codex holds the socket for one response at a time, so a second - // response.create means the first one ended without a terminal event. - ws.endTurn(prev, http.StatusOK, &pipeline.EventError{Kind: "upstream_error", Message: "a new turn began before this one completed"}) + p := ws.s.OutboundPipeline.Load() + if why := pipelineCannotCarry(p); why != "" { + // A reload since the handshake made this socket one the proxy would have + // declined. End it with "try again later": Codex reconnects, and the new + // handshake is answered 426 and falls back to HTTPS, where the plugin applies. + slog.Info("forward-proxy: ending a carried WebSocket after a reload", "host", logValue(ws.r.Host), "reason", why) + ws.closeBoth(wsframe.CloseTryAgainLater, "the proxy's pipeline changed; reconnect") + return nil, true } pctx := ws.newTurnContext(body) - t := &wsTurn{pctx: pctx, p: ws.s.OutboundPipeline.Load()} - action := t.p.Run(ws.ctx(), pctx) + t := &wsTurn{pctx: pctx, p: p, prewarm: !generate} + action := p.Run(ws.ctx(), pctx) if action.Type == pipeline.Reject { ws.s.recordOutboundRejectIn(nil, pctx, action, ws.sid) ws.refuseTurn(action) - t.p.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) + p.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) return nil, true } - if ws.s.Sessions != nil { + if t.prewarm { + // generate:false asks the server to set up state and answer at once, with no + // model output. It carries the same prompt a real turn does — the turn after it + // sends only what is new — so its request phase runs, and a plugin can refuse + // or rewrite it. But it is not inference: no row, and its answer is not a + // turn's response. + p.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) + } else if ws.s.Sessions != nil { pctx.OutboundSessionID = ws.sid ws.s.Sessions.Append(ws.sid, ws.turnRequestEvent(pctx)) } ws.turnMu.Lock() - ws.turn = t + ws.turns = append(ws.turns, t) ws.turnMu.Unlock() if pctx.BodyMutated() { - return withType(pctx.Body), false + return withType(pctx.Body, !generate), false } return nil, false } @@ -638,10 +696,13 @@ func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool) // ends the turn on a terminal event. func (ws *wsSocket) serverMessage(payload []byte) { ws.turnMu.Lock() - t := ws.turn + var t *wsTurn + if len(ws.turns) > 0 { + t = ws.turns[0] + } ws.turnMu.Unlock() if t == nil { - // The answer to a prewarm, or to a turn already ended. Nothing to parse into. + // Nothing is waiting for an answer. Nothing to parse into. return } var env struct { @@ -653,6 +714,15 @@ func (ws *wsSocket) serverMessage(payload []byte) { } `json:"error"` } _ = json.Unmarshal(payload, &env) + if t.prewarm { + // A prewarm's answer is set-up, not output. It ends on the same events a turn + // does, and is not recorded. + switch env.Type { + case "error", "response.completed", "response.incomplete", "response.failed": + ws.popTurn(t) + } + return + } if env.Type == "error" { // The origin refused the turn (a usage limit, bad input). See errorEventStatus. status := errorEventStatus(env.Status, env.Error.Code, env.Error.Type) @@ -660,9 +730,18 @@ func (ws *wsSocket) serverMessage(payload []byte) { if name := firstNonEmpty(env.Error.Code, env.Error.Type); name != "" { message += ": " + capErrorName(name) } + // The origin answered, so the status is its answer, not a failure of the proxy: + // the row carries the error, and Finishers see what HTTPS would show them for + // the same refusal — an allow with this status, not an error with none. + t.pctx.StatusCode = status ws.finishTurn(t, status, &pipeline.EventError{Kind: "backend_error", Code: strconv.Itoa(status), Message: message}) return } + // Delivered before it is read: relayDown forwarded this frame to the client already, + // and a WebSocket answer cannot be retracted. Marked here, not only at the turn's + // end, so a plugin rejecting a mid-turn frame records a late refusal rather than + // turning a turn the client received into a denial. + t.pctx.MarkResponseDelivered() t.p.RunResponseFrame(ws.ctx(), t.pctx, payload, false) t.pctx.ResponseBytes += int64(len(payload)) // The three events inference-parser reads usage and the finish reason from; each @@ -673,16 +752,24 @@ func (ws *wsSocket) serverMessage(payload []byte) { } } -// finishTurn ends t if it is still the open turn. +// finishTurn ends t, the turn at the head of the queue. Only relayDown calls it, so no +// other path can end t while its last event is still being parsed. func (ws *wsSocket) finishTurn(t *wsTurn, status int, fail *pipeline.EventError) { + if ws.popTurn(t) { + ws.endTurn(t, status, fail) + } +} + +// popTurn removes t from the head of the queue, reporting whether it was there. +func (ws *wsSocket) popTurn(t *wsTurn) bool { ws.turnMu.Lock() - if ws.turn != t { - ws.turnMu.Unlock() - return + defer ws.turnMu.Unlock() + if len(ws.turns) == 0 || ws.turns[0] != t { + return false } - ws.turn = nil - ws.turnMu.Unlock() - ws.endTurn(t, status, fail) + ws.turns[0] = nil + ws.turns = ws.turns[1:] + return true } // endTurn finalizes a turn: the terminal frame that settles its cost, its response @@ -692,7 +779,7 @@ func (ws *wsSocket) endTurn(t *wsTurn, status int, fail *pipeline.EventError) { finalCtx, cancelFinal := httpx.TeardownContext(ws.ctx()) defer cancelFinal() t.pctx.MarkResponseDelivered() - if fail == nil { + if fail == nil || fail.Kind == pipeline.ErrorKindTurnInterrupted { t.pctx.StatusCode = status } t.p.RunResponseFrame(finalCtx, t.pctx, nil, true) @@ -701,11 +788,10 @@ func (ws *wsSocket) endTurn(t *wsTurn, status int, fail *pipeline.EventError) { } // ctx is the context every turn's pipeline runs under: the handshake request's values -// (the client's process, for attribution) without its cancellation, since the handshake -// request ended when the upgrade was answered and the socket outlives it. -func (ws *wsSocket) ctx() context.Context { - return context.WithoutCancel(ws.r.Context()) -} +// (the client's process, for attribution) without its cancellation — the handshake +// request ended when the upgrade was answered, and the socket outlives it — cancelled +// instead when the socket ends. +func (ws *wsSocket) ctx() context.Context { return ws.ctxv } // newTurnContext builds a turn's pipeline.Context: the handshake's identity — host, // path, headers, client, session — with the request the turn would have been over @@ -722,16 +808,19 @@ func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { } headers.Set("Content-Type", "application/json") pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Method: http.MethodPost, - Scheme: hs.Scheme, - Host: hs.Host, - Path: hs.Path, - Headers: headers, - Body: body, - Shared: ws.s.Shared, - StartedAt: time.Now(), - Session: hs.Session, + Direction: pipeline.Outbound, + Method: http.MethodPost, + Scheme: hs.Scheme, + Host: hs.Host, + Path: hs.Path, + Headers: headers, + Body: body, + Shared: ws.s.Shared, + StartedAt: time.Now(), + // Built per turn, as HTTPS builds one per request: the handshake's was taken + // before any turn was recorded, and a plugin reading the session (an intent, a + // request count) would see none of the turns before this one. + Session: ws.s.sessionViewFor(ws.sid), SessionSource: hs.SessionSource, SessionHeader: hs.SessionHeader, ChainSession: hs.ChainSession, @@ -833,16 +922,29 @@ func responseCreateBody(payload []byte) (body []byte, generate bool) { return out, generate } -// withType turns a rewritten request body back into a response.create message. -func withType(body []byte) []byte { +// withType turns a rewritten request body back into a response.create message, with +// generate:false again when it was a prewarm. +func withType(body []byte, prewarm bool) []byte { var fields map[string]json.RawMessage if json.Unmarshal(body, &fields) != nil { return body } fields["type"] = json.RawMessage(`"response.create"`) + if prewarm { + fields["generate"] = json.RawMessage(`false`) + } out, err := json.Marshal(fields) if err != nil { return body } return out } + +// logValue is s as it may go in a log line: line breaks replaced, so a value a client +// chose (a Host, an error that quotes one) cannot start a log line of its own. slog's +// text handler already quotes such a value; this keeps the guarantee independent of +// which handler a binary installs. +func logValue(s string) string { + s = strings.ReplaceAll(s, "\n", `\n`) + return strings.ReplaceAll(s, "\r", `\r`) +} diff --git a/core/listener/forwardproxy/websocket_review_test.go b/core/listener/forwardproxy/websocket_review_test.go new file mode 100644 index 000000000..e87127a15 --- /dev/null +++ b/core/listener/forwardproxy/websocket_review_test.go @@ -0,0 +1,294 @@ +package forwardproxy + +import ( + "bufio" + "context" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/rossoctl/cortex/core/config" + "github.com/rossoctl/cortex/core/listener/internal/wsframe" + "github.com/rossoctl/cortex/core/pipeline" + "github.com/rossoctl/cortex/core/plugins/inferenceparser" + "github.com/rossoctl/cortex/core/plugins/plugintesting" +) + +// Tests for the review of #1467: each pins one defect a reviewer reproduced. + +// scriptedOrigin replaces the harness origin's handler with one that, after the +// handshake, runs script with a reader for client frames and a writer for its own. +func scriptedOrigin(h *wsHarness, script func(read func() *wsframe.Frame, write func(string))) { + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + conn, brw, _ := w.(http.Hijacker).Hijack() + defer conn.Close() + _, _ = brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + wsAccept(r.Header.Get("Sec-Websocket-Key")) + "\r\n\r\n") + _ = brw.Flush() + fr := wsframe.NewReader(brw.Reader, 64<<20, 64<<20) + read := func() *wsframe.Frame { + f, err := fr.ReadFrame() + if err != nil { + return nil + } + return f + } + write := func(ev string) { _, _ = conn.Write(wsframe.Encode(wsframe.OpText, []byte(ev), false)) } + script(read, write) + for read() != nil { + } + }) +} + +// turnRows is each turn's response row, in order. +func turnRows(events []pipeline.SessionEvent) []pipeline.SessionEvent { + var out []pipeline.SessionEvent + for _, e := range events { + if e.HTTPMethod == http.MethodPost && e.Phase == pipeline.SessionResponse { + out = append(out, e) + } + } + return out +} + +// Turn 1's terminal event arrives after the client has already sent turn 2 — what Codex +// does after a fast tool. Each answer belongs to the turn it answers, in order: turn 1 +// keeps its own usage, and turn 2 is not ended by turn 1's terminal event. +func TestResponsesWebSocket_LateTerminalEventEndsItsOwnTurn(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + scriptedOrigin(h, func(read func() *wsframe.Frame, write func(string)) { + read() // create 1 + one, two := turnEvents(1), turnEvents(2) + for _, ev := range one[:len(one)-1] { + write(ev) + } + read() // create 2, before turn 1's terminal event + write(one[len(one)-1]) + for _, ev := range two { + write(ev) + } + }) + c := h.dial(t, nil) + c.send(responseCreate("one")) + c.readUntil("response.output_item.done") + c.send(responseCreate("two")) + c.readUntil("response.completed") + c.readUntil("response.completed") + c.close() + + rows := turnRows(h.events(t, codexSID, true)) + if len(rows) != 2 { + t.Fatalf("%d turn rows, want 2", len(rows)) + } + for i, want := range []int{12867, 13007} { + r := rows[i] + if r.Error != nil { + t.Errorf("turn %d recorded error %+v; it completed", i+1, r.Error) + } + if r.Inference == nil || r.Inference.TotalTokens != want { + t.Errorf("turn %d tokens = %+v, want %d: an answer went to the wrong turn", i+1, r.Inference, want) + } + } +} + +// The handshake's Finishers run on the pipeline the handshake ran on, whatever a reload +// installs while the socket is open. A reload to a shorter pipeline used to index past +// its end. +func TestResponsesWebSocket_HandshakeFinishesOnItsOwnPipeline(t *testing.T) { + first := &outcomeRecorder{} + extra := &outcomeRecorder{} + h := newWSHarness(t, config.ResponsesWebSocketsParse, first, extra) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d", c.status) + } + shorter, err := plugintesting.BuildPipeline([]pipeline.Plugin{inferenceparser.NewInferenceParser()}) + if err != nil { + t.Fatal(err) + } + h.holder.Store(shorter) + c.close() + h.events(t, codexSID, true) + if first.seen() == nil || extra.seen() == nil { + t.Error("the handshake's own Finishers did not run once its socket closed") + } +} + +// A destination writer (inference-router) cannot steer a turn on a socket already +// connected to OpenAI, so the handshake is declined with a 426 and Codex takes HTTPS, +// where it can. +func TestResponsesWebSocket_DestinationWriterIsDeclined(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &redirector{}) + c := h.dial(t, nil) + if c.status != http.StatusUpgradeRequired { + t.Fatalf("handshake answered %d, want 426 with a redirecting plugin in the chain", c.status) + } +} + +// A reload that adds a plugin the socket cannot honour ends the socket at the next +// turn, with "try again later", rather than ignoring the plugin for the socket's life. +func TestResponsesWebSocket_ReloadAddingAnUncarriablePluginEndsTheSocket(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + c := h.dial(t, nil) + if c.status != http.StatusSwitchingProtocols { + t.Fatalf("handshake answered %d", c.status) + } + fresh, err := plugintesting.BuildPipeline([]pipeline.Plugin{inferenceparser.NewInferenceParser(), &responseRewriter{}}) + if err != nil { + t.Fatal(err) + } + h.holder.Store(fresh) + c.send(responseCreate("hello")) + f, err := c.fr.ReadFrame() + if err != nil || f.Opcode != wsframe.OpClose || len(f.Payload) < 2 || int(f.Payload[0])<<8|int(f.Payload[1]) != wsframe.CloseTryAgainLater { + t.Fatalf("next frame = %+v, %v; want a close with code %d", f, err, wsframe.CloseTryAgainLater) + } + c.close() + h.origin.mu.Lock() + n := len(h.origin.creates) + h.origin.mu.Unlock() + if n != 0 { + t.Errorf("origin received %d response.create; the turn must not be sent without its plugin", n) + } +} + +// A turn over the 32 MiB a request body may be ends the socket with 1009, as HTTPS +// answers 413: it is never forwarded unread, past every plugin. +func TestResponsesWebSocket_OverLimitTurnEndsTheSocket(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &turnDenier{}) + c := h.dial(t, nil) + _ = c.conn.SetDeadline(time.Now().Add(60 * time.Second)) + half := strings.Repeat("x", (wsClientMessageLimit/2)+1) + _, _ = c.conn.Write(fragmentFrame(wsframe.OpText, false, []byte(half))) + _, _ = c.conn.Write(fragmentFrame(wsframe.OpContinuation, true, []byte(half))) + f, err := c.fr.ReadFrame() + if err != nil || f.Opcode != wsframe.OpClose || len(f.Payload) < 2 || int(f.Payload[0])<<8|int(f.Payload[1]) != wsframe.CloseMessageTooBig { + t.Fatalf("answer = %+v, %v; want a close with code %d", f, err, wsframe.CloseMessageTooBig) + } + h.events(t, codexSID, true) + h.origin.mu.Lock() + n := len(h.origin.creates) + h.origin.mu.Unlock() + if n != 0 { + t.Errorf("origin received %d response.create; an over-limit turn must not be forwarded", n) + } +} + +// A prewarm carries the prompt the next turn builds on, so it runs through the request +// phase: a plugin that refuses the content refuses the prewarm too. It records no turn. +func TestResponsesWebSocket_PrewarmRunsTheRequestPhase(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse, &turnDenier{}) + c := h.dial(t, nil) + c.send(`{"type":"response.create","generate":false,"model":"gpt-6-luna","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"secret"}]}]}`) + msgs := c.readUntil("error") + if msgs[len(msgs)-1]["status"] != float64(http.StatusForbidden) { + t.Errorf("prewarm answer = %v, want the plugin's refusal", msgs[len(msgs)-1]) + } + c.close() + events := h.events(t, codexSID, true) + h.origin.mu.Lock() + n := len(h.origin.creates) + h.origin.mu.Unlock() + if n != 0 { + t.Errorf("origin received %d response.create; the refused prewarm must not be forwarded", n) + } + if rows := turnRows(events); len(rows) != 0 { + t.Errorf("%d turn rows for a prewarm; a prewarm is not a turn", len(rows)) + } +} + +// Each turn sees the session as it is when the turn starts, as an HTTPS request does, +// not the view taken at the handshake before any turn was recorded. +func TestResponsesWebSocket_EachTurnSeesTheSessionSoFar(t *testing.T) { + seen := &sessionLenRecorder{} + h := newWSHarness(t, config.ResponsesWebSocketsParse, seen) + c := h.dial(t, nil) + c.send(responseCreate("one")) + c.readUntil("response.completed") + c.send(responseCreate("two")) + c.readUntil("response.completed") + c.close() + h.events(t, codexSID, true) + got := seen.lens() + if len(got) < 3 || got[len(got)-1] <= got[len(got)-2] { + t.Errorf("session lengths seen = %v; the second turn should see more of the session than the first", got) + } +} + +// An origin's refusal is relayed whole: the handshake's deadline must not cut its body. +func TestResponsesWebSocket_OriginRefusalBodyIsRelayedWhole(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + body := strings.Repeat("r", 200<<10) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + w.WriteHeader(http.StatusForbidden) + bw := bufio.NewWriterSize(w, 4096) + for i := 0; i < len(body); i += 4096 { + _, _ = bw.WriteString(body[i:min(i+4096, len(body))]) + _ = bw.Flush() + w.(http.Flusher).Flush() + } + }) + c := h.dial(t, nil) + if c.status != http.StatusForbidden || len(c.body) != len(body) { + t.Errorf("refusal = %d with %d body bytes, want 403 with %d", c.status, len(c.body), len(body)) + } +} + +// fragmentFrame is one masked client frame, FIN as given. +func fragmentFrame(op byte, fin bool, payload []byte) []byte { + f := wsframe.Encode(op, payload, true) + if !fin { + f[0] &^= 0x80 + } + return f +} + +// redirector declares WritesDestination, as inference-router does. +type redirector struct{} + +func (redirector) Name() string { return "destination-writer" } +func (redirector) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{WritesDestination: true} +} +func (redirector) OnRequest(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (redirector) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// sessionLenRecorder records how many events the session view held at each request. +type sessionLenRecorder struct { + mu sync.Mutex + got []int +} + +func (*sessionLenRecorder) Name() string { return "session-len-recorder" } +func (*sessionLenRecorder) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{} +} +func (r *sessionLenRecorder) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + r.mu.Lock() + defer r.mu.Unlock() + n := -1 + if pctx.Session != nil { + n = pctx.Session.Len() + } + r.got = append(r.got, n) + return pipeline.Action{Type: pipeline.Continue} +} +func (*sessionLenRecorder) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (r *sessionLenRecorder) lens() []int { + r.mu.Lock() + defer r.mu.Unlock() + return append([]int(nil), r.got...) +} diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 716492db3..7115b8f86 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -348,7 +348,7 @@ func responseCreate(text string) string { // events reads the session's rows once the socket has closed and recorded its close. func (h *wsHarness) events(t *testing.T, sid string, wantClose bool) []pipeline.SessionEvent { t.Helper() - deadline := time.Now().Add(3 * time.Second) + deadline := time.Now().Add(10 * time.Second) for { v := h.store.View(sid) if v != nil { @@ -365,7 +365,9 @@ func (h *wsHarness) events(t *testing.T, sid string, wantClose bool) []pipeline. if v == nil { t.Fatalf("no session %q", sid) } - return v.Events + // Failing here, not returning what there is: a slow run would otherwise + // fail later, on an assertion about rows that had not landed yet. + t.Fatalf("session %q has no WebSocket close row after 10s (%d events)", sid, len(v.Events)) } time.Sleep(20 * time.Millisecond) } @@ -496,6 +498,8 @@ func TestResponsesWebSocket_RejectedTurnIsAnsweredAndNotForwarded(t *testing.T) } c.close() + // After the socket has ended, so a forwarded frame would have been read by now. + events := h.events(t, codexSID, true) h.origin.mu.Lock() n := len(h.origin.creates) h.origin.mu.Unlock() @@ -503,7 +507,7 @@ func TestResponsesWebSocket_RejectedTurnIsAnsweredAndNotForwarded(t *testing.T) t.Errorf("origin received %d response.create; a rejected turn must not be forwarded", n) } var denied bool - for _, e := range h.events(t, codexSID, true) { + for _, e := range events { if e.Phase == pipeline.SessionDenied { denied = true } @@ -656,8 +660,10 @@ func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { if res.Inference == nil || res.Inference.Completion != "par" { t.Errorf("interrupted turn completion = %+v, want the partial text", res.Inference) } - if res.StatusCode != pipeline.StatusClientClosedRequest || res.Error == nil || res.Error.Kind != "client_canceled" { - t.Errorf("interrupted turn = status %d error %+v, want 499 client_canceled", res.StatusCode, res.Error) + // As HTTPS records a stream the client leaves part way: the 200 the origin was + // answering with, and a kind /v1/usage does not count as an error. + if res.StatusCode != http.StatusOK || res.Error == nil || res.Error.Kind != pipeline.ErrorKindTurnInterrupted { + t.Errorf("interrupted turn = status %d error %+v, want 200 %s", res.StatusCode, res.Error, pipeline.ErrorKindTurnInterrupted) } } @@ -752,8 +758,8 @@ func TestResponsesWebSocket_FirstCloseFrameDecidesWhoEndedTheTurn(t *testing.T) if res == nil { t.Fatal("the closed turn recorded no response row") } - if res.StatusCode != pipeline.StatusClientClosedRequest || res.Error == nil || res.Error.Kind != "client_canceled" { - t.Fatalf("turn = status %d error %+v, want 499 client_canceled: the client sent the first close", res.StatusCode, res.Error) + if res.StatusCode != http.StatusOK || res.Error == nil || res.Error.Kind != pipeline.ErrorKindTurnInterrupted { + t.Fatalf("turn = status %d error %+v, want 200 %s: the client sent the first close", res.StatusCode, res.Error, pipeline.ErrorKindTurnInterrupted) } if !strings.Contains(res.Error.Message, "close code 1001 going away: user left") { t.Errorf("message %q does not name the client's close code and reason", res.Error.Message) @@ -991,7 +997,11 @@ func TestResponsesWebSocket_ReloadMidTurnFinishesOnTheTurnsPipeline(t *testing.T if started.seen() == nil { t.Error("the turn's own Finisher never ran: the turn ended on a pipeline other than the one it began on") } - if other.got != nil { + h.events(t, codexSID, true) // the socket has ended, and with it every Finisher + other.mu.Lock() + got := other.got + other.mu.Unlock() + if got != nil { t.Error("the reloaded pipeline's Finisher ran for a turn it never saw begin") } } diff --git a/core/listener/internal/wsframe/wsframe.go b/core/listener/internal/wsframe/wsframe.go index 2d8388ac8..71432cd92 100644 --- a/core/listener/internal/wsframe/wsframe.go +++ b/core/listener/internal/wsframe/wsframe.go @@ -13,11 +13,14 @@ package wsframe import ( + "bytes" "crypto/rand" "encoding/binary" "errors" "fmt" "io" + "math" + "unicode/utf8" ) // Opcodes (RFC 6455 §5.2). @@ -138,22 +141,44 @@ func (r *Reader) ReadFrame() (*Frame, error) { return nil, fmt.Errorf("%w: a %d-byte frame is over the relay's %d-byte limit", ErrProtocol, length, r.maxFrame) } - f.Raw = make([]byte, n+int(length)) - copy(f.Raw, hdr[:n]) - if _, err := io.ReadFull(r.r, f.Raw[n:]); err != nil { - return nil, unexpected(err) + // Raw grows as the payload arrives rather than being sized from the length the + // header claims: a peer that sends a 14-byte header claiming 32 MiB and then + // stalls would otherwise hold that much memory for as long as it stalls. + var raw bytes.Buffer + raw.Grow(n + int(min(length, rawPrealloc))) + raw.Write(hdr[:n]) + if got, err := io.CopyN(&raw, r.r, length); err != nil || got != length { + return nil, unexpected(orEOF(err)) } + f.Raw = raw.Bytes() if length > r.limit { f.Truncated = true return f, nil } - f.Payload = append([]byte(nil), f.Raw[n:]...) - if f.Masked { - maskBytes(f.Payload, key) + // Payload aliases Raw only when there is no mask to undo; otherwise it is the one + // copy reading needs. + if !f.Masked { + f.Payload = f.Raw[n:] + return f, nil } + f.Payload = append([]byte(nil), f.Raw[n:]...) + maskBytes(f.Payload, key) return f, nil } +// rawPrealloc is how much of a frame's claimed length is reserved up front: enough +// for the common frame in one allocation, and no more than that before the bytes +// arrive. +const rawPrealloc = 64 << 10 + +// orEOF is err, or io.EOF when a copy stopped short without one. +func orEOF(err error) error { + if err == nil { + return io.EOF + } + return err +} + func unexpected(err error) error { if errors.Is(err, io.EOF) { return io.ErrUnexpectedEOF @@ -172,7 +197,13 @@ func maskBytes(b []byte, key [4]byte) { // must not (RFC 6455 §5.1), so the caller says which side it is writing as. func Encode(opcode byte, payload []byte, masked bool) []byte { length := len(payload) - hdr := make([]byte, 0, 14+length) + // 14 is the largest header; the check keeps the capacity from overflowing an int, + // which only a payload near the int's own limit could do. + capacity := length + if capacity <= math.MaxInt-14 { + capacity += 14 + } + hdr := make([]byte, 0, capacity) hdr = append(hdr, 0x80|opcode) var maskBit byte if masked { @@ -204,7 +235,13 @@ func Encode(opcode byte, payload []byte, masked bool) []byte { // reason is cut to fit a control frame. func ClosePayload(code uint16, reason string) []byte { if len(reason) > maxControlPayload-2 { - reason = reason[:maxControlPayload-2] + // RFC 6455 §5.5.1 has the receiver fail the connection on a reason that is + // not valid UTF-8, so the cut backs up to a character boundary. + cut := maxControlPayload - 2 + for cut > 0 && !utf8.RuneStart(reason[cut]) { + cut-- + } + reason = reason[:cut] } b := binary.BigEndian.AppendUint16(nil, code) return append(b, reason...) @@ -253,28 +290,49 @@ func DescribeClose(payload []byte) string { return out } +// ErrMessageTooBig is a data message over its limit. The caller should end the socket +// (close code 1009, message too big), as the HTTPS path answers a body over its cap +// with a 413: a message too big to read is also one whose fragments the relay would +// otherwise have to hold without bound, waiting for an end that may never come. +var ErrMessageTooBig = errors.New("wsframe: message over the size limit") + +// CloseMessageTooBig is RFC 6455 §7.4.1's code for ErrMessageTooBig. +const CloseMessageTooBig = 1009 + +// CloseTryAgainLater asks the peer to reconnect (RFC 6455 §7.4.1, registered by IANA). +const CloseTryAgainLater = 1013 + +// maxFragments bounds the frames one message may arrive in. Each fragment is a +// separate allocation the relay holds until the message ends, so a flood of empty +// continuations — six bytes each on the wire, and never over the byte limit — would +// cost memory without bound. Codex sends each message as one frame. +const maxFragments = 1024 + // Message reassembles a data message from its frames. // // It sees control frames go by without letting them interrupt a message, as RFC 6455 // §5.4 allows a peer to interleave them between fragments. It reports a message whole -// once its final fragment arrives, and only if every fragment was readable and the total -// stayed within limit; otherwise the message is still relayed, just not read. +// once its final fragment arrives. A message over limit — in payload, in wire bytes, or +// in fragments — is ErrMessageTooBig, which ends the socket: it is never relayed +// unread, because a turn too big to read is one no plugin would see. type Message struct { - limit int64 - opcode byte - buf []byte - inFlight bool - // lost is set when a fragment of the current message was over the reader's - // limit or the total grew past this one's. - lost bool + limit int64 + opcode byte + buf []byte + wire int64 // wire bytes of the message so far, header and all + fragments int + inFlight bool } // NewMessage reassembles messages of up to limit bytes. func NewMessage(limit int64) *Message { return &Message{limit: limit} } // Add takes the next data frame and, when it completes a message, returns that -// message's opcode and payload with ok true. A message that could not be read returns -// ok true and a nil payload with readable false, so the caller can still count it. +// message's opcode and payload with ok true. It returns ErrMessageTooBig once the +// message passes its limit, and ErrProtocol for fragments out of order. +// +// readable is false only for a binary message, which is relayed and counted but not +// parsed. The payload is the caller's to keep: the next message does not reuse it. func (m *Message) Add(f *Frame) (opcode byte, payload []byte, readable, ok bool, err error) { if f.IsControl() { return 0, nil, false, false, nil @@ -285,19 +343,25 @@ func (m *Message) Add(f *Frame) (opcode byte, payload []byte, readable, ok bool, case f.Opcode != OpContinuation && m.inFlight: return 0, nil, false, false, fmt.Errorf("%w: a new message began before the last one finished", ErrProtocol) case f.Opcode != OpContinuation: - m.opcode, m.buf, m.lost, m.inFlight = f.Opcode, m.buf[:0], false, true + m.opcode, m.buf, m.wire, m.fragments, m.inFlight = f.Opcode, nil, 0, 0, true } - if f.Truncated || int64(len(m.buf))+int64(len(f.Payload)) > m.limit { - m.lost = true - } else if !m.lost { - m.buf = append(m.buf, f.Payload...) + m.wire += int64(len(f.Raw)) + m.fragments++ + if f.Truncated || int64(len(m.buf))+int64(len(f.Payload)) > m.limit || m.wire > m.limit+maxFragments*14 || m.fragments > maxFragments { + m.buf, m.inFlight = nil, false + return 0, nil, false, false, fmt.Errorf("%w: over %d bytes or %d fragments", ErrMessageTooBig, m.limit, maxFragments) } + if f.Fin && m.fragments == 1 { + // The common case, one frame: hand the frame's own payload over rather than + // copying it into buf and out again. + m.inFlight = false + return m.opcode, f.Payload, m.opcode == OpText, true, nil + } + m.buf = append(m.buf, f.Payload...) if !f.Fin { return 0, nil, false, false, nil } m.inFlight = false - if m.lost { - return m.opcode, nil, false, true, nil - } - return m.opcode, append([]byte(nil), m.buf...), true, true, nil + payload, m.buf = m.buf, nil + return m.opcode, payload, m.opcode == OpText, true, nil } diff --git a/core/listener/internal/wsframe/wsframe_test.go b/core/listener/internal/wsframe/wsframe_test.go index af91fd00f..c8957eeb8 100644 --- a/core/listener/internal/wsframe/wsframe_test.go +++ b/core/listener/internal/wsframe/wsframe_test.go @@ -4,8 +4,11 @@ import ( "bytes" "errors" "io" + "runtime" "strings" "testing" + "time" + "unicode/utf8" ) func readOne(t *testing.T, raw []byte, limit int64) *Frame { @@ -164,31 +167,59 @@ func TestMessage_ReassemblesAroundInterleavedControlFrames(t *testing.T) { } } -// A message over the limit completes but is not readable, and the next one starts -// fresh rather than inheriting the lost state. -func TestMessage_OverLimitIsCountedNotRead(t *testing.T) { +// A message over the limit is ErrMessageTooBig as soon as it passes it, not relayed +// unread: a turn too big to read is one no plugin sees, and its fragments would +// otherwise be held until a final one that may never come. +func TestMessage_OverLimitIsRefused(t *testing.T) { m := NewMessage(8) frames := messageFrames(t, frag(OpText, false, "12345"), frag(OpContinuation, true, "67890"), - frag(OpText, true, "ok"), ) - var results []string - for _, f := range frames { - _, payload, readable, ok, err := m.Add(f) - if err != nil { - t.Fatal(err) - } - if ok { - if readable { - results = append(results, string(payload)) - } else { - results = append(results, "") - } - } + if _, _, _, ok, err := m.Add(frames[0]); err != nil || ok { + t.Fatalf("first fragment: ok=%v err=%v, want neither", ok, err) + } + if _, _, _, _, err := m.Add(frames[1]); !errors.Is(err, ErrMessageTooBig) { + t.Errorf("over the limit: err = %v, want ErrMessageTooBig", err) + } +} + +// A frame the reader did not keep, being over its limit, ends the message the same way. +func TestMessage_TruncatedFragmentIsRefused(t *testing.T) { + m := NewMessage(1 << 20) + first := &Frame{Opcode: OpText, Fin: false, Payload: []byte("A"), Raw: []byte{0x01, 0x01, 'A'}} + big := &Frame{Opcode: OpContinuation, Fin: false, Truncated: true, Raw: make([]byte, 16)} + if _, _, _, _, err := m.Add(first); err != nil { + t.Fatal(err) + } + if _, _, _, _, err := m.Add(big); !errors.Is(err, ErrMessageTooBig) { + t.Errorf("a truncated middle fragment: err = %v, want ErrMessageTooBig — it must not be spliced out", err) + } +} + +// Empty continuations cost a frame each and add no payload, so the payload limit alone +// would never stop them. The fragment cap does. +func TestMessage_FragmentFloodIsRefused(t *testing.T) { + m := NewMessage(1 << 20) + if _, _, _, _, err := m.Add(&Frame{Opcode: OpText, Raw: []byte{0x01, 0x00}}); err != nil { + t.Fatal(err) } - if strings.Join(results, ",") != ",ok" { - t.Errorf("results = %v", results) + var err error + for i := 0; i < maxFragments+1 && err == nil; i++ { + _, _, _, _, err = m.Add(&Frame{Opcode: OpContinuation, Raw: []byte{0x00, 0x00}}) + } + if !errors.Is(err, ErrMessageTooBig) { + t.Errorf("after %d empty continuations: err = %v, want ErrMessageTooBig", maxFragments+1, err) + } +} + +// After a refused message the next one starts fresh. +func TestMessage_NextMessageAfterARefusalIsRead(t *testing.T) { + m := NewMessage(4) + _, _, _, _, _ = m.Add(messageFrames(t, frag(OpText, true, "too long"))[0]) + _, payload, readable, ok, err := m.Add(messageFrames(t, frag(OpText, true, "ok"))[0]) + if err != nil || !ok || !readable || string(payload) != "ok" { + t.Errorf("next message: payload=%q readable=%v ok=%v err=%v", payload, readable, ok, err) } } @@ -232,3 +263,78 @@ func TestDescribeClose(t *testing.T) { } } } + +// RFC 6455 §5.7's longer examples: a 256-byte and a 65536-byte unmasked binary frame, +// in the 16-bit and 64-bit length forms. Both read back with the right length, and +// both directions agree on byte order (a little-endian length would misframe them). +func TestReadFrame_RFC6455LongLengthExamples(t *testing.T) { + for _, tc := range []struct { + name string + hdr []byte + n int + }{ + {"256 bytes, 16-bit length", []byte{0x82, 0x7E, 0x01, 0x00}, 256}, + {"64 KiB, 64-bit length", []byte{0x82, 0x7F, 0, 0, 0, 0, 0, 0x01, 0x00, 0x00}, 65536}, + } { + t.Run(tc.name, func(t *testing.T) { + raw := append(append([]byte(nil), tc.hdr...), bytes.Repeat([]byte{'x'}, tc.n)...) + f, err := NewReader(bytes.NewReader(raw), 1<<20, 1<<20).ReadFrame() + if err != nil || len(f.Payload) != tc.n || !bytes.Equal(f.Raw, raw) { + t.Fatalf("payload=%d err=%v rawEqual=%v", len(f.Payload), err, f != nil && bytes.Equal(f.Raw, raw)) + } + if enc := Encode(OpBinary, f.Payload, false); !bytes.Equal(enc[:len(tc.hdr)], tc.hdr) { + t.Errorf("Encode header = % x, want % x", enc[:len(tc.hdr)], tc.hdr) + } + }) + } +} + +// The relay limit is checked against the header's claim before any payload is read, so +// a header claiming more than the limit is refused without waiting for, or reserving +// room for, a body that never comes. +func TestReadFrame_RefusesAnOversizedClaimBeforeReading(t *testing.T) { + claim := []byte{0x82, 0x7F, 0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF} // no body follows + if _, err := NewReader(bytes.NewReader(claim), 1024, 4096).ReadFrame(); !errors.Is(err, ErrProtocol) { + t.Errorf("err = %v, want ErrProtocol", err) + } + // At the limit exactly: header plus payload is maxFrame, which is allowed. + exact := append([]byte{0x82, 6}, "123456"...) + if _, err := NewReader(bytes.NewReader(exact), 1024, 8).ReadFrame(); err != nil { + t.Errorf("a frame exactly at the limit: err = %v", err) + } + if _, err := NewReader(bytes.NewReader(exact), 1024, 7).ReadFrame(); !errors.Is(err, ErrProtocol) { + t.Errorf("a frame one byte over: err = %v, want ErrProtocol", err) + } +} + +// A header claiming a large payload and then stalling holds memory for what arrived, not +// for what it claimed. +func TestReadFrame_ReservesForWhatArrivesNotWhatIsClaimed(t *testing.T) { + hdr := []byte{0x82, 0x7F, 0, 0, 0, 0, 0x01, 0, 0, 0} // 16 MiB claimed + _, err := NewReader(io.MultiReader(bytes.NewReader(hdr), bytes.NewReader([]byte("abc"))), 1<<30, 1<<30).ReadFrame() + if !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("err = %v, want ErrUnexpectedEOF for a body that stopped short", err) + } + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + pr, pw := io.Pipe() + go func() { _, _ = pw.Write(hdr); _, _ = pw.Write([]byte("abc")) }() + done := make(chan struct{}) + go func() { _, _ = NewReader(pr, 1<<30, 1<<30).ReadFrame(); close(done) }() + time.Sleep(50 * time.Millisecond) // the reader is now stalled mid-payload + runtime.ReadMemStats(&after) + _ = pw.Close() + <-done + if grew := int64(after.HeapAlloc) - int64(before.HeapAlloc); grew > 4<<20 { + t.Errorf("a stalled 16 MiB claim grew the heap by %d bytes; want it bounded by what arrived", grew) + } +} + +// A close reason is cut at a character boundary, so the payload stays valid UTF-8. +func TestClosePayload_CutsAtACharacter(t *testing.T) { + p := ClosePayload(1000, strings.Repeat("é", 100)) + if len(p) > maxControlPayload || !utf8.Valid(p[2:]) { + t.Errorf("payload %d bytes, valid UTF-8 %v", len(p), utf8.Valid(p[2:])) + } +} diff --git a/core/listener/parity/dispatch_audit_test.go b/core/listener/parity/dispatch_audit_test.go index 4dca7a966..e4cd5f9e1 100644 --- a/core/listener/parity/dispatch_audit_test.go +++ b/core/listener/parity/dispatch_audit_test.go @@ -98,8 +98,10 @@ var dispatchSites = map[string]dispatchClass{ // finish defer's case, one turn at a time. "forwardproxy/websocket.go:endTurn:finalCtx": finalizationDelivered, // Each of a turn's server messages, mid-turn: the WebSocket counterpart of the SSE - // path's per-block frames. Detached from the handshake request, which ended when the - // upgrade was answered; the socket's own lifetime bounds it. + // path's per-block frames. inFlight for its context — the socket's, cancelled when + // the socket ends, not a teardown budget that would cut a long answer short — but + // unlike an SSE block the frame is already on the wire, so serverMessage marks the + // turn delivered before dispatching: a reject here is late, never a denial. "forwardproxy/websocket.go:serverMessage:ws.ctx()": inFlight, // ext_proc's teardown flush: the stream is gone by definition here, and its two dispatches take diff --git a/core/pipeline/context.go b/core/pipeline/context.go index 1c1636790..ee13dbbb2 100644 --- a/core/pipeline/context.go +++ b/core/pipeline/context.go @@ -140,8 +140,9 @@ type Context struct { // with nothing connecting the two. (Codex no longer reaches that 405 by default: // the forward proxy carries a Responses API WebSocket, or answers it 426 itself, // which Codex reads as "use HTTPS" at once — see listener.responses_websockets. - // The marker is still set on the handshake row, since the client's Upgrade header - // is never forwarded as it was sent.) + // Neither sets the marker: a carried socket forwards the upgrade, and its open + // row is a tunnel row; a declined one forwards nothing. It is set only under + // responses_websockets: forward, and for every other WebSocket.) // // Recorded on the REQUEST row, where the fact is true, and paired with the // response row's status by RequestID — the same split Tunnel/TunnelReason use. diff --git a/core/pipeline/pipeline.go b/core/pipeline/pipeline.go index 6ddc39e16..7e80eb2ef 100644 --- a/core/pipeline/pipeline.go +++ b/core/pipeline/pipeline.go @@ -457,6 +457,19 @@ func (p *Pipeline) WritesRequestBody() bool { return false } +// WritesDestination returns true if any plugin in the pipeline declares +// WritesDestination: it may send a request somewhere other than where the client +// addressed it. A carried WebSocket is already connected to its origin, so the +// forward proxy declines one when this is true and the turns go over HTTPS instead. +func (p *Pipeline) WritesDestination() bool { + for _, plugin := range p.plugins { + if plugin.Capabilities().Normalize().WritesDestination { + return true + } + } + return false +} + // WritesResponseBody returns true if any plugin in the pipeline declares // WritesResponseBody. This is the SSE streaming predicate: a response // mutator needs the whole response to rewrite it, so listeners fall back diff --git a/core/pipeline/session.go b/core/pipeline/session.go index cc81dd95c..3a9715a82 100644 --- a/core/pipeline/session.go +++ b/core/pipeline/session.go @@ -650,6 +650,12 @@ type EventError struct { // usage does not count it as an error, and the Finisher sees an allow. const ErrorKindUpgradeDeclined = "upgrade_declined" +// ErrorKindTurnInterrupted is a carried WebSocket turn the client ended before its last +// event: Codex's Esc, which drops the socket. The row keeps the 200 the origin was +// answering with, as an HTTPS stream the client leaves part way does, and /v1/usage +// does not count it as an error. The kind says what happened; it is not a failure. +const ErrorKindTurnInterrupted = "turn_interrupted" + // SessionSource is how a listener resolved Context.Session. type SessionSource string diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 6ecd44b75..5f3ff0f12 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -249,21 +249,30 @@ On the timeline a socket is a `GET` on the Responses path with `tunnelReason: duration and bytes. Its turns sit between the two as ordinary `POST` request and response rows. Plugins act on each turn: one that rejects it keeps it from reaching chatgpt.com and Codex is answered with the error event it reads for a server refusal; one that -rewrites the request body (`tool-prune`) changes what is sent. +rewrites the request body (a redaction script, `context-guru`) changes what is sent. +Codex's prewarm — a `response.create` with `generate: false`, sent before the first turn +with the prompt the turns build on — goes through the same request plugins, so a plugin +that refuses or rewrites content does so there too. It records no turn: it produces no +output. + +A turn over the **32 MiB** an HTTPS request body may be ends the socket with close code +1009, "message too big" — the WebSocket counterpart of the 413 HTTPS answers. It is never +relayed unread, since no plugin would see it and no row would record it. **Pressing Esc mid-answer closes the socket.** Codex drops the turn's stream, which ends the WebSocket, and opens a new one for the next prompt. The interrupted turn is recorded -with the text that had arrived, as `statusCode` 499 with `error.kind: -"client_canceled"` — the same as an HTTPS client hanging up — so it is not counted as an -upstream failure. OpenAI sends no usage for a turn cut off this way, so it has no token -count or cost: it shows as priceable but unpriced, the same as a cancelled HTTPS stream. +with the text that had arrived, as `statusCode` 200 with `error.kind: +"turn_interrupted"` — the way HTTPS records a stream the client leaves part way, with the +200 the origin was already answering — so `/v1/usage` does not count it as an error. +OpenAI sends no usage for a turn cut off this way, so it has no token count or cost: it +shows as priceable but unpriced, the same as a cancelled HTTPS stream. **A turn's status follows what the same turn would get over HTTPS.** An open WebSocket carries no HTTP status after its 101, so each turn row is given the one its HTTPS equivalent would have: 200 for a completed turn; for a refusal the origin sends as an `error` event, the event's own status, else the one its error code is answered with (`insufficient_quota` and `rate_limit_exceeded` are 429, `context_length_exceeded` 400, -`server_is_overloaded` 503), else 500; 499 `client_canceled` when Codex ended the socket +`server_is_overloaded` 503), else 500; 200 `turn_interrupted` when Codex ended the socket mid-turn and 502 `upstream_error` when OpenAI did. Who ended it is read from the first close frame either side sent, whose code and reason go into the row's message, and from which side's connection dropped first when neither sent one. @@ -271,8 +280,15 @@ which side's connection dropped first when neither sent one. Where Cortex cannot carry the socket it answers the handshake with **426 Upgrade Required** instead, which Codex reads as "this endpoint wants HTTP" and switches to HTTPS at once, for the rest of its session. That happens for a connection the bridge did not -decrypt, an HTTP/2 connection, or a pipeline with a plugin that rewrites response bodies, -since a live socket's frames are on their way to Codex before a whole response exists. The +decrypt; a pipeline with a plugin that rewrites response bodies, since a live socket's +frames are on their way to Codex before a whole response exists; and a pipeline with a +plugin that sends requests elsewhere — `inference-router`, as `agentop server use` sets +up — since a turn rides a socket already connected to OpenAI, where over HTTPS each turn +is its own request and can be routed. A reload that adds either kind of plugin while a +socket is open ends that socket at its next turn with close code 1013, "try again later": +Codex reconnects, is answered 426, and goes on over HTTPS with the plugin applied. (An +HTTP/2 connection never reaches this far: Go's HTTP/2 server answers any request carrying +`Upgrade` with a 400 before Cortex sees it. Codex's client speaks HTTP/1.1.) The timeline then shows a 426 with `error.kind: "upgrade_declined"`, and Codex prints one harmless `failed to connect to websocket: HTTP error: 426 Upgrade Required` line before carrying on over HTTPS. diff --git a/docs/framework-architecture.md b/docs/framework-architecture.md index 255125423..39c09515a 100644 --- a/docs/framework-architecture.md +++ b/docs/framework-architecture.md @@ -918,7 +918,7 @@ listener: | Mode | What happens | |---|---| -| `parse` | The socket is carried: the handshake is completed with the origin, frames are relayed unchanged, and every turn on it — a `response.create` message and the event messages answering it — runs the outbound pipeline as the equivalent HTTPS request and SSE stream would. Each turn records a request/response pair, is priced, and finishes; the socket records an open and a close with `tunnelReason: "websocket"`. A plugin rejecting a turn keeps it from the origin and answers the client with a Responses `error` event; a request-body writer's rewrite is what the origin receives. `permessage-deflate` is stripped from the handshake so frames stay readable. A socket that cannot be carried — not decrypted by the TLS bridge, HTTP/2, or a pipeline with a `WritesResponseBody` plugin — is declined as below. | +| `parse` | The socket is carried: the handshake is completed with the origin, frames are relayed unchanged, and every turn on it — a `response.create` message and the event messages answering it — runs the outbound pipeline as the equivalent HTTPS request and SSE stream would. Each turn records a request/response pair, is priced, and finishes; the socket records an open and a close with `tunnelReason: "websocket"`. A plugin rejecting a turn keeps it from the origin and answers the client with a Responses `error` event; a request-body writer's rewrite is what the origin receives. `permessage-deflate` is stripped from the handshake so frames stay readable. A socket that cannot be carried — not decrypted by the TLS bridge, or a pipeline with a `WritesResponseBody` or `WritesDestination` plugin (inference-router: a turn on a socket already connected to OpenAI cannot be sent elsewhere) — is declined as below, and a reload that adds such a plugin ends an open socket at its next turn with close code 1013 so Codex reconnects and is declined. A message over the HTTPS limits (32 MiB from the client, 10 MB per event from the origin) ends the socket with 1009. Each turn and the handshake run on the pipeline they began on. An HTTP/2 client never reaches the handler: Go's HTTP/2 server refuses a request carrying `Upgrade` with a 400. | | `decline` | Answered **426 Upgrade Required** by the proxy, contacting no origin. Codex reads that as "use HTTP" and switches to HTTPS at once. Recorded as `error.kind: "upgrade_declined"`, which `/v1/usage` does not count as an error. | | `forward` | The behaviour before either existed: forwarded with `Upgrade` stripped (it is hop-by-hop), so the origin sees a bare `GET`, which OpenAI's endpoint answers 405. Codex retries a 405 five times over about six seconds before falling back to HTTPS. | diff --git a/docs/session-api.md b/docs/session-api.md index 8017ae034..5b145e254 100644 --- a/docs/session-api.md +++ b/docs/session-api.md @@ -235,9 +235,14 @@ Every event on `/v1/sessions/{id}` and `/v1/events` carries: is the client hanging up first — since neither proxy times out a slow upstream, that is how a hung one usually ends. A 426 with `error.kind: "upgrade_declined"` is the forward proxy's own answer to a WebSocket upgrade on - OpenAI's Responses API (`chatgpt.com` or `api.openai.com`): nothing is forwarded, - and a client such as Codex reads it as "use HTTPS" and retries at once over - HTTPS. It is not counted as an error in `/v1/usage`. + OpenAI's Responses API (`chatgpt.com` or `api.openai.com`) that it will not carry: + every one under `listener.responses_websockets: decline`, and under the default, + `parse`, one it cannot (not decrypted, or a plugin that rewrites response bodies or + redirects requests). Nothing is forwarded, and a client such as Codex reads it as + "use HTTPS" and retries at once over HTTPS. It is not counted as an error in + `/v1/usage`. A carried socket's turn the client ends early (Codex's Esc) is a 200 + with `error.kind: "turn_interrupted"`, as an HTTPS stream the client leaves part way + keeps its 200, and is not counted as an error either. - `requestedHost` — present only when a plugin redirected the request: the host the client asked for. `host` is where the request actually went, and usage, the cost ledger and pricing all follow `host`. From 944bd00dc3eebd7843d271f717a84ec0a2f89540 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Wed, 28 Oct 2026 16:04:10 -0400 Subject: [PATCH 09/12] fix: Build each WebSocket turn's session view as HTTPS does CodeQL's clear-text-logging query treats every request header as sensitive, and an id cleaned of line breaks still counts. Two changes in the last commit gave it a path from the handshake's headers into plugins' and the store's log lines that main has no counterpart for: wrapping the socket's session id in logValue, and building each turn's session view from that id. Each turn's view now comes from the handshake's resolved session, as hydrateSession builds an HTTPS request's: the same per-turn freshness, from the same source an HTTPS request uses. The logValue wrapper is dropped; it changed nothing a header id can hold, since IDHeaderFrom already refuses control characters. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/server.go | 7 ++---- core/listener/forwardproxy/websocket.go | 32 +++++++++++++++---------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index b4bc183bb..ed312f3b4 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -508,10 +508,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne if !skipped && wantsResponsesWebSocket(r) { switch s.ResponsesWebSockets { case config.ResponsesWebSocketsParse: - // logValue changes nothing a header id can hold — IDHeaderFrom already refused - // control characters — but the socket carries this id into rows and log lines - // for a whole session, so it is cleaned where the socket takes it. - sid := logValue(s.recordingSessionID(sessionID, r.Header, chain)) + sid := s.recordingSessionID(sessionID, r.Header, chain) if why := s.cannotCarryWebSocket(tl, p); why != "" { slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", logValue(r.Host), "reason", why) s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) @@ -520,7 +517,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain, p) return case config.ResponsesWebSocketsDecline: - s.declineResponsesWebSocket(w, r, pctx, tl, logValue(s.recordingSessionID(sessionID, r.Header, chain)), chain) + s.declineResponsesWebSocket(w, r, pctx, tl, s.recordingSessionID(sessionID, r.Header, chain), chain) return } } diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index d8298da43..cc051f0cf 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -807,20 +807,26 @@ func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { headers.Del(h) } headers.Set("Content-Type", "application/json") + // The session as it is when this turn starts, built per turn as HTTPS builds one per + // request: the handshake's view was taken before any turn was recorded, so a plugin + // reading the session (an intent, a request count) would see none of the turns before + // this one. From the handshake's resolved session, as hydrateSession resolves one for + // an HTTPS request, so a turn has no view exactly when that request would have none. + var view *pipeline.SessionView + if hs.Session != nil { + view = ws.s.sessionViewFor(hs.Session.ID) + } pctx := &pipeline.Context{ - Direction: pipeline.Outbound, - Method: http.MethodPost, - Scheme: hs.Scheme, - Host: hs.Host, - Path: hs.Path, - Headers: headers, - Body: body, - Shared: ws.s.Shared, - StartedAt: time.Now(), - // Built per turn, as HTTPS builds one per request: the handshake's was taken - // before any turn was recorded, and a plugin reading the session (an intent, a - // request count) would see none of the turns before this one. - Session: ws.s.sessionViewFor(ws.sid), + Direction: pipeline.Outbound, + Method: http.MethodPost, + Scheme: hs.Scheme, + Host: hs.Host, + Path: hs.Path, + Headers: headers, + Body: body, + Shared: ws.s.Shared, + StartedAt: time.Now(), + Session: view, SessionSource: hs.SessionSource, SessionHeader: hs.SessionHeader, ChainSession: hs.ChainSession, From 431901157187dc213bbb110e11a62f83030b6fac Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Wed, 28 Oct 2026 16:18:59 -0400 Subject: [PATCH 10/12] refactor: Drop the CodeQL-only log wrapper and size check logValue replaced line breaks in seven WebSocket log values, and Encode bounded 14+len(payload) against int overflow, only to clear CodeQL alerts. Neither guarded anything real: slog's text handler, which every binary installs, already quotes a value containing a line break; and 14+len cannot overflow for a slice that fits in memory, since every caller passes a payload already bounded by the 32 MiB message limit. The log lines now pass values as the rest of the proxy does. The alerts these targeted are dismissed as false positives; #1533 tracks fixing the session id's flow at its source. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/server.go | 4 ++-- core/listener/forwardproxy/websocket.go | 19 +++++-------------- core/listener/internal/wsframe/wsframe.go | 9 +-------- 3 files changed, 8 insertions(+), 24 deletions(-) diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index ed312f3b4..aa4bebae6 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -510,7 +510,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne case config.ResponsesWebSocketsParse: sid := s.recordingSessionID(sessionID, r.Header, chain) if why := s.cannotCarryWebSocket(tl, p); why != "" { - slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", logValue(r.Host), "reason", why) + slog.Debug("forward-proxy: declining a Responses WebSocket the proxy cannot carry", "host", r.Host, "reason", why) s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) return } @@ -1060,7 +1060,7 @@ const declinedUpgradeBody = `{"detail":"WebSocket transport is not available thr // does for the 421. func (s *Server) declineResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc) { slog.Debug("forward-proxy: declined a WebSocket upgrade on a Responses path; the client falls back to HTTPS", - "host", logValue(r.Host), "path", logValue(r.URL.Path)) + "host", r.Host, "path", r.URL.Path) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusUpgradeRequired) _, _ = io.WriteString(w, declinedUpgradeBody) diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index cc051f0cf..680200d4c 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -192,7 +192,7 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, if err != nil { cancel() status, fail := pipeline.ExchangeFailure(r.Context(), err) - slog.Warn("forward-proxy: WebSocket handshake with the origin failed", "host", logValue(r.Host), "kind", fail.Kind, "error", logValue(fail.Message)) + slog.Warn("forward-proxy: WebSocket handshake with the origin failed", "host", r.Host, "kind", fail.Kind, "error", fail.Message) http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) s.recordWebSocketRefusal(pctx, tl, sid, chain, status, fail) return @@ -237,7 +237,7 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, clientConn, clientBuf, err := hj.Hijack() if err != nil { _ = upConn.Close() - slog.Warn("forward-proxy: could not take over the client connection for a WebSocket", "host", logValue(r.Host), "error", logValue(err.Error())) + slog.Warn("forward-proxy: could not take over the client connection for a WebSocket", "host", r.Host, "error", err) return } // The bridge's handshake deadline was cleared before serving; nothing sets one @@ -630,10 +630,10 @@ func (ws *wsSocket) logEnd(side string, err error) { return } if errors.Is(err, wsframe.ErrProtocol) { - slog.Warn("forward-proxy: ending a WebSocket on a frame the relay cannot read", "host", logValue(ws.r.Host), "side", side, "error", logValue(err.Error())) + slog.Warn("forward-proxy: ending a WebSocket on a frame the relay cannot read", "host", ws.r.Host, "side", side, "error", err) return } - slog.Debug("forward-proxy: WebSocket ended", "host", logValue(ws.r.Host), "side", side, "error", logValue(err.Error())) + slog.Debug("forward-proxy: WebSocket ended", "host", ws.r.Host, "side", side, "error", err) } // wsEnvelope is the one field every message on the socket carries. @@ -657,7 +657,7 @@ func (ws *wsSocket) clientMessage(payload []byte) (replace []byte, refuse bool) // A reload since the handshake made this socket one the proxy would have // declined. End it with "try again later": Codex reconnects, and the new // handshake is answered 426 and falls back to HTTPS, where the plugin applies. - slog.Info("forward-proxy: ending a carried WebSocket after a reload", "host", logValue(ws.r.Host), "reason", why) + slog.Info("forward-proxy: ending a carried WebSocket after a reload", "host", ws.r.Host, "reason", why) ws.closeBoth(wsframe.CloseTryAgainLater, "the proxy's pipeline changed; reconnect") return nil, true } @@ -945,12 +945,3 @@ func withType(body []byte, prewarm bool) []byte { } return out } - -// logValue is s as it may go in a log line: line breaks replaced, so a value a client -// chose (a Host, an error that quotes one) cannot start a log line of its own. slog's -// text handler already quotes such a value; this keeps the guarantee independent of -// which handler a binary installs. -func logValue(s string) string { - s = strings.ReplaceAll(s, "\n", `\n`) - return strings.ReplaceAll(s, "\r", `\r`) -} diff --git a/core/listener/internal/wsframe/wsframe.go b/core/listener/internal/wsframe/wsframe.go index 71432cd92..08a23e751 100644 --- a/core/listener/internal/wsframe/wsframe.go +++ b/core/listener/internal/wsframe/wsframe.go @@ -19,7 +19,6 @@ import ( "errors" "fmt" "io" - "math" "unicode/utf8" ) @@ -197,13 +196,7 @@ func maskBytes(b []byte, key [4]byte) { // must not (RFC 6455 §5.1), so the caller says which side it is writing as. func Encode(opcode byte, payload []byte, masked bool) []byte { length := len(payload) - // 14 is the largest header; the check keeps the capacity from overflowing an int, - // which only a payload near the int's own limit could do. - capacity := length - if capacity <= math.MaxInt-14 { - capacity += 14 - } - hdr := make([]byte, 0, capacity) + hdr := make([]byte, 0, 14+length) hdr = append(hdr, 0x80|opcode) var maskBit byte if masked { From cb3f4fbb477d0cb324b2d1cbe46d1c634614248d Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Thu, 29 Oct 2026 11:19:38 -0400 Subject: [PATCH 11/12] fix: Don't finish a queued prewarm again when the socket closes From huang195's review of #1467: - A prewarm still waiting for its answer when the socket closed was ended like an unanswered turn: a response row with no request row, and a second RunFinish, which RunFinish drops with a WARN. A prewarm is finished when it is sent, so the close now leaves it alone. - The request row no longer says the upgrade was stripped once the origin answers the forwarded upgrade, refusal or not. A dial that failed, and the local 426, keep the marker. - The over-limit comment and docs say which direction is refused before forwarding: the client's. An origin message is forwarded frame by frame before it is read, so frames before the limit reach Codex. - Comments that still said truncated frames are relayed unread, or that a client hang-up is a 499, now match the code. - wsSocket.hsp, and serveResponsesWebSocket's pipeline parameter, are dropped: nothing read them. Tests: the unanswered prewarm records no turn row and is finished once (no RunFinish-twice WARN); an origin refusal's request row has no upgrade marker. Both fail without their fix. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/server.go | 2 +- core/listener/forwardproxy/websocket.go | 25 +++- .../forwardproxy/websocket_review_test.go | 126 ++++++++++++++++++ core/listener/forwardproxy/websocket_test.go | 7 +- core/listener/internal/wsframe/wsframe.go | 9 +- core/pipeline/context.go | 6 +- docs/agents/codex.md | 6 +- 7 files changed, 160 insertions(+), 21 deletions(-) diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index aa4bebae6..adcc8d3e3 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -514,7 +514,7 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, tl *tunne s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) return } - s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain, p) + s.serveResponsesWebSocket(w, r, pctx, tl, sid, chain) return case config.ResponsesWebSocketsDecline: s.declineResponsesWebSocket(w, r, pctx, tl, s.recordingSessionID(sessionID, r.Header, chain), chain) diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index 680200d4c..525d0088c 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -58,7 +58,10 @@ const ( // conversation, the same body an HTTPS POST would — and maxBodySize for each event // the origin sends back, the SSE path's per-frame cap. A larger message ends the // socket with close code 1009, as the HTTPS path refuses an oversized body with a - // 413: relaying it unread would carry a turn no plugin saw and no row records. + // 413. From the client that happens before any of it is forwarded, since relaying + // it unread would carry a turn no plugin saw and no row records. From the origin + // the frames already sent stay sent: relayDown forwards each frame before reading + // it, so a too-big answer reaches the client up to the frame that crossed the limit. wsClientMessageLimit = maxRequestBodySize wsServerMessageLimit = maxBodySize // wsFrameLimit is the largest single frame relayed. Raw is buffered whole before @@ -164,7 +167,7 @@ func pipelineCannotCarry(p *pipeline.Pipeline) string { // pctx is the handshake's, after the pipeline ran on it. Its rows are the socket's open // (now) and close (when it ends), sharing its RequestID the way an opaque tunnel's do; // the turns record their own pairs between them. -func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc, p *pipeline.Pipeline) { +func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, pctx *pipeline.Context, tl *tunnelLog, sid string, chain []session.Proc) { hj, ok := w.(http.Hijacker) if !ok { s.declineResponsesWebSocket(w, r, pctx, tl, sid, chain) @@ -189,6 +192,12 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, client := s.TLSBridge.Upstream resp, err := client.Do(up) + if err == nil { + // The origin answered the upgrade, so it was forwarded, whatever it answered: + // the request row must not say otherwise. A dial that failed forwarded nothing, + // and keeps the marker. + pctx.UpgradeStripped = "" + } if err != nil { cancel() status, fail := pipeline.ExchangeFailure(r.Context(), err) @@ -269,7 +278,7 @@ func (s *Server) serveResponsesWebSocket(w http.ResponseWriter, r *http.Request, sock := &wsSocket{ s: s, r: r, hs: pctx, tl: tl, sid: sid, chain: chain, client: clientConn, clientR: clientBuf.Reader, up: upConn, - start: time.Now(), ctxv: ctx, stop: stop, hsp: p, + start: time.Now(), ctxv: ctx, stop: stop, } s.recordWebSocketOpen(sock) sock.run() @@ -340,8 +349,6 @@ type wsSocket struct { // waiting on something (a budget pause, a slow script) is released with it. ctxv context.Context stop context.CancelFunc - // hsp is the pipeline the handshake ran on. - hsp *pipeline.Pipeline // endedBy is the side whose relay stopped first, "client" or "origin", which // decides how a turn still open at that moment is recorded. Set once, by the @@ -396,14 +403,18 @@ func (ws *wsSocket) run() { <-done // Both relays have returned, so nothing else touches the queue: every turn still - // in it was sent and never answered in full. + // in it was sent and never answered in full. A prewarm in it is left alone: it was + // finished when it was sent and has no request row, so ending it here would run + // its Finishers twice and record a response row with nothing to pair with. ws.turnMu.Lock() open := ws.turns ws.turns = nil ws.turnMu.Unlock() status, fail := ws.unfinishedTurn() for _, t := range open { - ws.endTurn(t, status, fail) + if !t.prewarm { + ws.endTurn(t, status, fail) + } } ws.recordClose() // Released last: endTurn's teardown contexts derive from it. diff --git a/core/listener/forwardproxy/websocket_review_test.go b/core/listener/forwardproxy/websocket_review_test.go index e87127a15..48bcf472f 100644 --- a/core/listener/forwardproxy/websocket_review_test.go +++ b/core/listener/forwardproxy/websocket_review_test.go @@ -3,6 +3,7 @@ package forwardproxy import ( "bufio" "context" + "log/slog" "net/http" "strings" "sync" @@ -292,3 +293,128 @@ func (r *sessionLenRecorder) lens() []int { defer r.mu.Unlock() return append([]int(nil), r.got...) } + +// A prewarm still waiting for its answer when the socket closes is not a turn: it was +// already finished when it was sent. Closing must not finish it again, nor record a +// response row for a request no row was recorded for (#1467 review, reproduced). +func TestResponsesWebSocket_UnansweredPrewarmIsNotFinishedTwice(t *testing.T) { + count := &finishCounter{} + logs := captureLogs(t) + h := newWSHarness(t, config.ResponsesWebSocketsParse, count) + scriptedOrigin(h, func(read func() *wsframe.Frame, write func(string)) { + read() // the prewarm, never answered + }) + c := h.dial(t, nil) + c.send(`{"type":"response.create","generate":false,"model":"gpt-6-luna","input":[]}`) + // Wait until the origin has the prewarm, so the close lands while it is queued. + deadline := time.Now().Add(10 * time.Second) + for { + h.origin.mu.Lock() + n := len(h.origin.creates) + h.origin.mu.Unlock() + if n > 0 || time.Now().After(deadline) { + break + } + time.Sleep(10 * time.Millisecond) + } + c.close() + events := h.events(t, codexSID, true) + + if rows := turnRows(events); len(rows) != 0 { + t.Errorf("%d turn response rows for a prewarm that was never answered; want none", len(rows)) + } + // The prewarm is the one POST, finished once, when it was sent. Before the fix the + // close finished it again: RunFinish drops a second call with a WARN rather than run + // the Finishers twice, so the Finisher count alone cannot see it, but the WARN can. + // (The CONNECT tunnel and the handshake finish too, at moments the test does not + // control, so they are not counted.) + if got := count.n(); got != 1 { + t.Errorf("the prewarm's Finisher ran %d times; want 1", got) + } + if logs.contains("RunFinish called twice") { + t.Error("the prewarm was finished twice: RunFinish dropped a second call at close") + } +} + +// finishCounter counts Finisher calls for POST requests: a turn, or a prewarm. +type finishCounter struct { + mu sync.Mutex + calls int +} + +func (*finishCounter) Name() string { return "finish-counter" } +func (*finishCounter) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{} +} +func (*finishCounter) OnRequest(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (*finishCounter) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (c *finishCounter) OnFinish(_ context.Context, pctx *pipeline.Context) { + if pctx.Method != http.MethodPost { + return + } + c.mu.Lock() + c.calls++ + c.mu.Unlock() +} +func (c *finishCounter) n() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.calls +} + +// An origin that refuses the forwarded upgrade still received it, so the request row +// must not say the upgrade was stripped (#1467 review). +func TestResponsesWebSocket_OriginRefusalIsNotMarkedStripped(t *testing.T) { + h := newWSHarness(t, config.ResponsesWebSocketsParse) + h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodHead { + return + } + w.WriteHeader(http.StatusForbidden) + }) + c := h.dial(t, nil) + if c.status != http.StatusForbidden { + t.Fatalf("handshake answered %d, want the origin's 403", c.status) + } + v := h.store.View(codexSID) + if v == nil { + t.Fatal("no session recorded") + } + for _, e := range v.Events { + if e.Phase == pipeline.SessionRequest && e.UpgradeStripped != "" { + t.Errorf("request row carries upgradeStripped=%q; the origin received the upgrade and refused it", e.UpgradeStripped) + } + } +} + +// logBuffer is a slog destination a test can search. +type logBuffer struct { + mu sync.Mutex + buf strings.Builder +} + +func (b *logBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *logBuffer) contains(s string) bool { + b.mu.Lock() + defer b.mu.Unlock() + return strings.Contains(b.buf.String(), s) +} + +// captureLogs sends slog's default logger to a buffer for the rest of the test. +func captureLogs(t *testing.T) *logBuffer { + t.Helper() + b := &logBuffer{} + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(b, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + return b +} diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index 7115b8f86..e159ecb60 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -621,8 +621,8 @@ func TestResponsesWebSocket_OriginRefusalIsRelayed(t *testing.T) { // A client that hangs up mid-turn — Codex does this on Esc, closing the socket rather // than interrupting over it — still gets that turn recorded with what had arrived, as a -// 499 client_canceled the way an HTTPS client hanging up is, not as an upstream -// failure: the origin did nothing wrong. +// 200 turn_interrupted the way an HTTPS stream the client leaves keeps its 200, not as an +// upstream failure: the origin did nothing wrong. func TestResponsesWebSocket_ClientHangupFinalizesTheOpenTurn(t *testing.T) { h := newWSHarness(t, config.ResponsesWebSocketsParse) // An origin that answers a create with a partial stream and then stalls. @@ -711,7 +711,8 @@ func TestResponsesWebSocket_OriginDropFinalizesTheOpenTurnAsUpstreamError(t *tes // stops first. Here the client sends a close frame, the origin answers with one of its // own and drops the connection at once, so the relay reading from the origin is the // first to see its connection gone — timing alone would blame the origin. The client -// sent the first close, so the turn is the client's: 499, with the code it closed with. +// sent the first close, so the turn is the client's: 200 turn_interrupted, with the code +// it closed with. func TestResponsesWebSocket_FirstCloseFrameDecidesWhoEndedTheTurn(t *testing.T) { h := newWSHarness(t, config.ResponsesWebSocketsParse) h.origin.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/core/listener/internal/wsframe/wsframe.go b/core/listener/internal/wsframe/wsframe.go index 08a23e751..897605376 100644 --- a/core/listener/internal/wsframe/wsframe.go +++ b/core/listener/internal/wsframe/wsframe.go @@ -54,7 +54,7 @@ type Frame struct { // Forward this, not a re-encoding. Raw []byte // Payload is the unmasked payload, or nil when the frame is over the reader's - // limit: such a frame is still relayed whole, but not read. + // limit, in which case Truncated is set and Message.Add refuses the message. Payload []byte // Truncated reports that Payload is nil because the frame exceeded the limit. Truncated bool @@ -66,10 +66,9 @@ func (f *Frame) IsControl() bool { return f.Opcode&0x8 != 0 } // Reader reads frames from one direction of a connection. type Reader struct { r io.Reader - // limit is the largest payload kept for reading. A larger frame is read and - // relayed in full but its Payload is dropped, so one huge message costs the - // proxy no more than limit bytes of reading memory... beyond the Raw copy the - // relay needs anyway. + // limit is the largest payload kept for reading. A larger frame is read whole + // into Raw, which maxFrame bounds, but its Payload is not kept, and Message.Add + // turns it into ErrMessageTooBig, which ends the socket. limit int64 // maxFrame bounds Raw: a frame over it is refused, since relaying it would mean // buffering it whole first. diff --git a/core/pipeline/context.go b/core/pipeline/context.go index ee13dbbb2..d023cfc6c 100644 --- a/core/pipeline/context.go +++ b/core/pipeline/context.go @@ -140,9 +140,9 @@ type Context struct { // with nothing connecting the two. (Codex no longer reaches that 405 by default: // the forward proxy carries a Responses API WebSocket, or answers it 426 itself, // which Codex reads as "use HTTPS" at once — see listener.responses_websockets. - // Neither sets the marker: a carried socket forwards the upgrade, and its open - // row is a tunnel row; a declined one forwards nothing. It is set only under - // responses_websockets: forward, and for every other WebSocket.) + // A declined one keeps the marker, since its upgrade was not forwarded either; a + // carried one clears it as soon as the origin answers the forwarded upgrade, + // refusal or not, and a dial that failed keeps it.) // // Recorded on the REQUEST row, where the fact is true, and paired with the // response row's status by RequestID — the same split Tunnel/TunnelReason use. diff --git a/docs/agents/codex.md b/docs/agents/codex.md index 5f3ff0f12..fab3003ce 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -256,8 +256,10 @@ that refuses or rewrites content does so there too. It records no turn: it produ output. A turn over the **32 MiB** an HTTPS request body may be ends the socket with close code -1009, "message too big" — the WebSocket counterpart of the 413 HTTPS answers. It is never -relayed unread, since no plugin would see it and no row would record it. +1009, "message too big" — the WebSocket counterpart of the 413 HTTPS answers — before any +of it reaches OpenAI, since no plugin would see it and no row would record it. An answer +from OpenAI over 10 MiB ends the socket the same way; the frames before the limit have +already reached Codex, since the relay forwards each frame before reading it. **Pressing Esc mid-answer closes the socket.** Codex drops the turn's stream, which ends the WebSocket, and opens a new one for the next prompt. The interrupted turn is recorded From 167e1b66fabb5cd44b0a90ae8cf493c51cb05ad2 Mon Sep 17 00:00:00 2001 From: Alan Cha Date: Thu, 29 Oct 2026 14:30:26 -0400 Subject: [PATCH 12/12] fix: Carry the handshake's attribution onto each WebSocket turn #1535 pins a request's process and the session-id header that named its session on its pipeline.Context (SetAttribution, from hydrateSession), and the session store counts a session's procs and owner only from request rows that carry SessionHeader. A WebSocket turn builds a fresh Context, so it had neither. On a session Codex uses only the socket for, no request row named the session, and session.Link had nothing to tie Codex's earlier calls to. Each turn's Context now takes the handshake's attribution, and its request row records Proc, ClientProc and SessionHeader as an HTTPS POST's does. The turn's response and deny rows read Proc from the same Context. TurnsAreParsedLikeHTTPS asserts SessionHeader on each turn's request row. Assisted-By: Claude (Anthropic AI) Signed-off-by: Alan Cha --- core/listener/forwardproxy/websocket.go | 35 ++++++++++++-------- core/listener/forwardproxy/websocket_test.go | 5 +++ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/core/listener/forwardproxy/websocket.go b/core/listener/forwardproxy/websocket.go index 17a029aeb..e3336955c 100644 --- a/core/listener/forwardproxy/websocket.go +++ b/core/listener/forwardproxy/websocket.go @@ -843,6 +843,10 @@ func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { ChainSession: hs.ChainSession, } pctx.ResolveClient() + // The handshake's attribution, which hydrateSession pinned on it: a turn is the same + // process on the same session, and without it no turn row names the header that named + // the session — the rows #1535's fold counts procs and owner from. + pctx.SetAttribution(hs.AttributedProc(), hs.AttributedClientProc(), hs.NamingHeader()) return pctx } @@ -850,20 +854,23 @@ func (ws *wsSocket) newTurnContext(body []byte) *pipeline.Context { // would have recorded. func (ws *wsSocket) turnRequestEvent(pctx *pipeline.Context) pipeline.SessionEvent { return pipeline.SessionEvent{ - At: time.Now(), - Direction: pipeline.Outbound, - Phase: pipeline.SessionRequest, - RequestID: pctx.RequestID(), - BytesUp: int64(len(pctx.Body)), - MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), - Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), - Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), - Plugins: pipeline.SnapshotPlugins(pctx.Extensions.Custom), - Identity: pipeline.SnapshotIdentity(pctx), - Host: pctx.Host, - HTTPMethod: pctx.Method, - HTTPPath: pctx.Path, - Client: pctx.ClientInfo(), + At: time.Now(), + Direction: pipeline.Outbound, + Phase: pipeline.SessionRequest, + RequestID: pctx.RequestID(), + BytesUp: int64(len(pctx.Body)), + MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), + Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), + Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest), + Plugins: pipeline.SnapshotPlugins(pctx.Extensions.Custom), + Identity: pipeline.SnapshotIdentity(pctx), + Host: pctx.Host, + HTTPMethod: pctx.Method, + HTTPPath: pctx.Path, + Client: pctx.ClientInfo(), + Proc: pctx.AttributedProc(), + ClientProc: pctx.AttributedClientProc(), + SessionHeader: pctx.NamingHeader(), } } diff --git a/core/listener/forwardproxy/websocket_test.go b/core/listener/forwardproxy/websocket_test.go index e159ecb60..58077633e 100644 --- a/core/listener/forwardproxy/websocket_test.go +++ b/core/listener/forwardproxy/websocket_test.go @@ -448,6 +448,11 @@ func TestResponsesWebSocket_TurnsAreParsedLikeHTTPS(t *testing.T) { if req.Inference == nil || req.Inference.Model != "gpt-6-luna" || len(req.Inference.Messages) != 1 { t.Errorf("turn %d request parse = %+v, want model gpt-6-luna and its one message", i+1, req.Inference) } + // The handshake's attribution: on a session Codex uses only the socket for, the + // turns are the only request rows that can name the session's header. + if req.SessionHeader != session.CodexSessionHeader { + t.Errorf("turn %d request row sessionHeader = %q, want %q", i+1, req.SessionHeader, session.CodexSessionHeader) + } if res.StatusCode != http.StatusOK || res.Inference == nil { t.Errorf("turn %d response = status %d inference %+v", i+1, res.StatusCode, res.Inference) continue