diff --git a/CLAUDE.md b/CLAUDE.md index 789cfc8e3..25281a891 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -604,16 +604,22 @@ 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. + 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, 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 403 with `error.kind: "destination_private"` is the egress guard's answer, not the upstream's** (`listener.egress_guard`, off by default): the destination resolved into loopback, link-local, private or this host's own address space and no permit covered it. The Finisher sees an allow with status 403, like the 426, and `/v1/usage` does not count it as an error. - **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) + **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 @@ -641,7 +647,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. A response relayed straight through is counted, not zero — both proxies record the row after the copy (#1251's fix), so the bytes they relayed through their own process are on it. The one arm that still reports ↓0 is extproc with a body-free pipeline: it asked Envoy for body mode NONE, so the bytes it never saw are bytes it asked not to see — `core/listener/parity`'s `TestDivergence_InboundResponseBodyCountWhenNoPluginReadsIt` pins that disagreement with its reason. 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. Everywhere else the 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. `tunnelReason: "dest-private"` with a 403 close and `error.kind: "destination_private"` is the egress guard refusing the destination before any dial — never confused with `dial-failed`, which attempted. +- `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. `tunnelReason: "dest-private"` with a 403 close and `error.kind: "destination_private"` is the egress guard refusing the destination before any dial — never confused with `dial-failed`, which attempted. - `dialAddr` — the address the egress guard checked and the proxy dialed, `"ip:port"`, on a tunnel's open row when `listener.egress_guard` is on. With `host` it is the audit trail against DNS rebinding: under the guard the proxy resolves once, checks that answer, and dials exactly it, so the two fields together say "asked for the name, reached this address". Absent when the guard is off and on rows that dialed nothing. - `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. - `subagent` — present on rows from a Claude Code subagent or agent-team teammate: the `X-Claude-Code-Agent-Id` the request carried. For a subagent that is also the `` in `~/.claude/projects///subagents/agent-.meta.json`; for a teammate the join is unverified. Resolved by `Context.Subagent`, pinned with the User-Agent by `ResolveClient` so no plugin can re-file a row, and **refused rather than sanitised** when unusable (control runes, over 128 bytes), because it is a join key and a cleaned-up id joins nothing. It does not move the row: a subagent's requests stay in their parent's session, so the header must not be added to `session.id_headers`. **It names an agent, it does not classify the row:** `inference.agentRole` stays the field that says whether a row is a subagent's, and a teammate's row is `agentRole: main` with `subagent` set — Claude Code writes `cc_is_subagent=true` only for `agentType: subagent`. diff --git a/cmd/agentop/tui/detail_pane.go b/cmd/agentop/tui/detail_pane.go index 0dbd31d80..668bd8a7b 100644 --- a/cmd/agentop/tui/detail_pane.go +++ b/cmd/agentop/tui/detail_pane.go @@ -150,7 +150,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 994c78f01..b73a16571 100644 --- a/cmd/cortex/main.go +++ b/cmd/cortex/main.go @@ -1016,10 +1016,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 1d4e45ef2..594d041c0 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -935,15 +935,18 @@ 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"` + + // 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:"-"` // EgressGuard arms the SSRF dial guard on the forward proxy's every // outbound dial (#1484): each destination is resolved once, refused with @@ -970,7 +973,7 @@ type ListenerConfig struct { // protect. // // A pointer so "absent" (off) stays distinguishable from an explicit - // enabled: false, mirroring DeclineResponsesWebSockets. Listener.* field: + // enabled: false, as SessionConfig.Enabled is. Listener.* field: // changing it requires a restart, like every other listener knob. EgressGuard *EgressGuardConfig `yaml:"egress_guard,omitempty" json:"egress_guard,omitempty"` } @@ -989,18 +992,37 @@ type EgressGuardConfig struct { PermitCIDRs []string `yaml:"permit_cidrs,omitempty" json:"permit_cidrs,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 + // 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. + 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 e44b18fb8..0a2a34baa 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,59 @@ 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) + } +} + +// 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) + } + } +} + // TestListenerConfig_EgressGuard pins the default (off) and the block that arms // it. Unset must mean off with or without a listener block: a sidecar's ordinary // upstreams live in exactly the refused space, so a config predating the key diff --git a/core/config/validate.go b/core/config/validate.go index bc24959c3..0b8f5a6e7 100644 --- a/core/config/validate.go +++ b/core/config/validate.go @@ -112,6 +112,22 @@ 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 + 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/cost/usage/usage.go b/core/cost/usage/usage.go index 88a2b3fe0..bfb53f947 100644 --- a/core/cost/usage/usage.go +++ b/core/cost/usage/usage.go @@ -1278,6 +1278,7 @@ func (a *Aggregator) evictColdestLocked() { // failing. 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 || + e.Error.Kind == pipeline.ErrorKindTurnInterrupted || e.Error.Kind == pipeline.ErrorKindDestinationPrivate) { return false } diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 889e57e0d..bf4cc5ddc 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -24,6 +24,7 @@ import ( "errors" "github.com/rossoctl/cortex/core/attribution" + "github.com/rossoctl/cortex/core/config" "github.com/rossoctl/cortex/core/listener/egressguard" "github.com/rossoctl/cortex/core/listener/httpx" "github.com/rossoctl/cortex/core/listener/internal/bodycount" @@ -134,12 +135,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 // EgressGuard, when non-nil, checks every destination this proxy dials — // CONNECT tunnels, plain and bridged requests, transparent redirects, and @@ -161,6 +163,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 @@ -473,6 +480,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 @@ -480,7 +494,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)) }() } @@ -518,7 +532,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 @@ -563,11 +577,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(tl, p); 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 { @@ -1131,59 +1158,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. // @@ -2181,6 +2155,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. @@ -2251,9 +2233,12 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { var sessionID string var chain []attribution.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 @@ -2271,7 +2256,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 @@ -2372,7 +2357,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // address" — moca's egress conformance reads it off the tunnel open row. pctx.DialedAddr = checked } - upstream, err := net.DialTimeout("tcp", dialTarget, connectDialTimeout) + upstream, err := s.dialUpstream(dialTarget) if err != nil { slog.Warn("forward-proxy: CONNECT upstream dial failed", "host", r.Host, "error", err) http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) @@ -2429,7 +2414,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // like any other opaque one. dialTarget, not r.Host: the address was // already checked once, and re-resolving here would reopen the rebinding // window the check closed. - up2, derr := net.DialTimeout("tcp", dialTarget, connectDialTimeout) + up2, derr := s.dialUpstream(dialTarget) 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..e3336955c --- /dev/null +++ b/core/listener/forwardproxy/websocket.go @@ -0,0 +1,963 @@ +package forwardproxy + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/rossoctl/cortex/core/attribution" + "github.com/rossoctl/cortex/core/listener/httpx" + "github.com/rossoctl/cortex/core/listener/internal/wsframe" + "github.com/rossoctl/cortex/core/pipeline" +) + +// 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 ( + // 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 ends the + // socket with close code 1009, as the HTTPS path refuses an oversized body with a + // 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 + // 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. +// +// 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" + } + 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" + } + 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 []attribution.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) + 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) + 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. + // 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 { + 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 + } + // 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() + 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 + 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(), ctxv: ctx, stop: stop, + } + 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 []attribution.Proc, status int, fail *pipeline.EventError) { + if s.Sessions == nil { + return + } + pctx.OutboundSessionID = sid + if chain != nil { + defer s.touch(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 []attribution.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 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 + 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 + + // 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 + + // 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. +func (ws *wsSocket) wsEnd(side string) { + ws.endedBy.CompareAndSwap(nil, side) +} + +// 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 + // 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() { + done := make(chan struct{}, 2) + 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() + _ = 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. 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 { + if !t.prewarm { + 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 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 +// 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) { + 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 http.StatusOK, &pipeline.EventError{ + Kind: pipeline.ErrorKindTurnInterrupted, + 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" + 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. +// 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 + // 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 "client" + } + 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" + } + continue + } + 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 + } + 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 "origin" + } + } + } +} + +// relayDown carries origin frames to the client, feeding each data message to the turn +// 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 "origin" + } + 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 "client" + } + if f.IsControl() { + if f.Opcode == wsframe.OpClose { + ws.sawClose("origin", f) + } + continue + } + 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" + } + if complete && op == wsframe.OpText && readable { + ws.serverMessage(payload) + } + } +} + +// 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() + _, 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) + + 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", 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: 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) + p.RunFinish(ws.ctx(), pctx, pipeline.OutcomeFromContext(pctx)) + return nil, true + } + 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.turns = append(ws.turns, t) + ws.turnMu.Unlock() + + if pctx.BodyMutated() { + return withType(pctx.Body, !generate), 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() + var t *wsTurn + if len(ws.turns) > 0 { + t = ws.turns[0] + } + ws.turnMu.Unlock() + if t == nil { + // Nothing is waiting for an answer. Nothing to parse into. + return + } + 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 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) + message := "the origin answered the turn with an error event" + 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 + // 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, 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() + defer ws.turnMu.Unlock() + if len(ws.turns) == 0 || ws.turns[0] != t { + return false + } + ws.turns[0] = nil + ws.turns = ws.turns[1:] + return true +} + +// 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 || fail.Kind == pipeline.ErrorKindTurnInterrupted { + t.pctx.StatusCode = status + } + t.p.RunResponseFrame(finalCtx, t.pctx, nil, true) + ws.s.recordOutboundResponseEvent(t.pctx, status, fail) + 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 +// (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 +// 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() + // 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") + // 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(), + Session: view, + SessionSource: hs.SessionSource, + SessionHeader: hs.SessionHeader, + 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 +} + +// 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(), + Proc: pctx.AttributedProc(), + ClientProc: pctx.AttributedClientProc(), + SessionHeader: pctx.NamingHeader(), + } +} + +// 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) + ws.s.touch(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, 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 +} diff --git a/core/listener/forwardproxy/websocket_review_test.go b/core/listener/forwardproxy/websocket_review_test.go new file mode 100644 index 000000000..48bcf472f --- /dev/null +++ b/core/listener/forwardproxy/websocket_review_test.go @@ -0,0 +1,420 @@ +package forwardproxy + +import ( + "bufio" + "context" + "log/slog" + "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...) +} + +// 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 new file mode 100644 index 000000000..58077633e --- /dev/null +++ b/core/listener/forwardproxy/websocket_test.go @@ -0,0 +1,1096 @@ +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, 64<<20, 64<<20) + msg := wsframe.NewMessage(64 << 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 + // 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 { + 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) + } + h.holder = pipeline.NewHolder(p) + srv := &Server{ + OutboundPipeline: h.holder, + 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(10 * 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) + } + // 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) + } +} + +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) + } + // 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 + } + 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() + + // 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() + if n != 0 { + t.Errorf("origin received %d response.create; a rejected turn must not be forwarded", n) + } + var denied bool + for _, e := range events { + 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 — Codex does this on Esc, closing the socket rather +// than interrupting over it — still gets that turn recorded with what had arrived, as a +// 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. + 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.Inference == nil || res.Inference.Completion != "par" { + t.Errorf("interrupted turn completion = %+v, want the partial text", res.Inference) + } + // 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) + } +} + +// 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() +} + +// 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: 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) { + 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 != 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) + } +} + +// 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 +// response.interrupt.accepted carrying that id, which it rejects as "unexpected" or +// "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. +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 +// 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") + } + 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") + } +} + +// 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) + _ = 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() + + 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("an 11 MB first turn was not parsed: %+v", req) + } +} + +// --- 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..897605376 --- /dev/null +++ b/core/listener/internal/wsframe/wsframe.go @@ -0,0 +1,359 @@ +// 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 ( + "bytes" + "crypto/rand" + "encoding/binary" + "errors" + "fmt" + "io" + "unicode/utf8" +) + +// 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, 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 +} + +// 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 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. + 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) + } + + // 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 + } + // 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 + } + 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 { + // 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...) +} + +// 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 +} + +// 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. 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 + 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. 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 + } + 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.wire, m.fragments, m.inFlight = f.Opcode, nil, 0, 0, true + } + 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 + 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 new file mode 100644 index 000000000..c8957eeb8 --- /dev/null +++ b/core/listener/internal/wsframe/wsframe_test.go @@ -0,0 +1,340 @@ +package wsframe + +import ( + "bytes" + "errors" + "io" + "runtime" + "strings" + "testing" + "time" + "unicode/utf8" +) + +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 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"), + ) + 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) + } + 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) + } +} + +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]) + } +} + +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) + } + } +} + +// 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 27f504ff5..e4cd5f9e1 100644 --- a/core/listener/parity/dispatch_audit_test.go +++ b/core/listener/parity/dispatch_audit_test.go @@ -93,6 +93,16 @@ 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. 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 // 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 f13756258..a311577f5 100644 --- a/core/pipeline/context.go +++ b/core/pipeline/context.go @@ -138,9 +138,11 @@ 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. + // 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/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 58e101429..52ef92f68 100644 --- a/core/pipeline/session.go +++ b/core/pipeline/session.go @@ -277,7 +277,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 @@ -476,6 +476,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" // TunnelDestPrivate — the egress guard refused the destination before any // dial: it resolved to loopback, link-local (cloud metadata), RFC 1918 / // unique-local space, or one of this host's own addresses, and no permit @@ -719,12 +724,18 @@ 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. 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" + // ErrorKindDestinationPrivate is the egress guard refusing a destination that // resolved into loopback, link-local, private or this host's own address space // (listener.egress_guard), rendered as 403 destination_private. Like diff --git a/docs/README.md b/docs/README.md index 8a5205f19..2164cb772 100644 --- a/docs/README.md +++ b/docs/README.md @@ -74,7 +74,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 2f61c30ea..923dbc0de 100644 --- a/docs/agents/codex.md +++ b/docs/agents/codex.md @@ -259,42 +259,87 @@ 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 (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 — 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 +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; 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. + +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; 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. + +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. @@ -384,22 +429,50 @@ 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`; +- **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`; +- **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; +- **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. -- **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. +- **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 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. diff --git a/docs/framework-architecture.md b/docs/framework-architecture.md index 9e4bfcabc..5fb4633c5 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, 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. | + +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. #### Refusing private destinations (`listener.egress_guard`) diff --git a/docs/session-api.md b/docs/session-api.md index c0231aa9e..4fca4112e 100644 --- a/docs/session-api.md +++ b/docs/session-api.md @@ -273,9 +273,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`. A **403** with + 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. A **403** with `error.kind: "destination_private"` is the egress guard's refusal (`listener.egress_guard`): the destination resolved into loopback, link-local, private or this host's own address space and no permit covered @@ -319,7 +324,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. `tunnelReason: "dest-private"` is the + 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. `tunnelReason: "dest-private"` is the egress guard's refusal: the close row carries 403 and `error.kind: "destination_private"`, and no tunnel was ever opened to the destination — the same refusal a plain request gets, one layer down.