From e3b4710e19b3ee7af8af2ba021541e2b088a9843 Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 15:28:20 +0530 Subject: [PATCH 01/21] Add experimental Codex WebSocket transport --- src/gateway/openai_codex.zig | 139 ++++++++++++++++- src/gateway/websocket_transport.zig | 229 ++++++++++++++++++++++++++++ 2 files changed, 367 insertions(+), 1 deletion(-) create mode 100644 src/gateway/websocket_transport.zig diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index 95e3356ed..6cfad9ba8 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -7,6 +7,7 @@ const io_mod = @import("../core/shared/io.zig"); const types = @import("../core/shared/types.zig"); const gateway_client = @import("client.zig"); const responses_protocol = @import("responses_protocol.zig"); +const websocket_transport = @import("websocket_transport.zig"); const model_tool_schema = @import("../core/tooling/model_tool_schema.zig"); const Allocator = std.mem.Allocator; @@ -22,6 +23,16 @@ const max_tool_arguments_bytes: usize = 4 * 1024 * 1024; const max_provider_state_bytes: usize = 4 * 1024 * 1024; const transfer_buffer_bytes: usize = 256 * 1024; const connect_timeout_ms: i64 = 30_000; +const transport_env = "FX_CODEX_TRANSPORT"; + +const Transport = enum { sse, websocket }; + +fn selectedTransport() !Transport { + const value = io_mod.getenv(transport_env) orelse return .sse; + if (std.mem.eql(u8, value, "sse") or std.mem.eql(u8, value, "auto")) return .sse; + if (std.mem.eql(u8, value, "websocket")) return .websocket; + return error.InvalidOpenAICodexTransport; +} const CodexLimits = struct { aggregate_bytes: usize = max_sse_aggregate_bytes, @@ -140,7 +151,14 @@ fn streamCompletion( try validateModel(request.model); const payload = try buildRequest(alloc, request.data()); defer alloc.free(payload); - return streamPrepared(alloc, request, payload) catch |err| { + return switch (try selectedTransport()) { + .sse => streamPrepared(alloc, request, payload), + .websocket => blk: { + const websocket_payload = try buildWebSocketRequest(alloc, payload); + defer alloc.free(websocket_payload); + break :blk streamWebSocketPrepared(alloc, request, websocket_payload); + }, + } catch |err| { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); request.attempt_evidence.network_failure = gateway_client.networkFailureEvidence(err, request.delivery.load()); return err; @@ -321,6 +339,100 @@ pub fn streamPrepared( } }; } +fn buildWebSocketRequest(alloc: Allocator, sse_payload: []const u8) ![]u8 { + const stream_fields = ",\"store\":false,\"stream\":true"; + if (sse_payload.len < 2 or sse_payload[0] != '{') return error.InvalidOpenAICodexWebSocketRequest; + const index = std.mem.find(u8, sse_payload, stream_fields) orelse return error.InvalidOpenAICodexWebSocketRequest; + var output: std.Io.Writer.Allocating = .init(alloc); + errdefer output.deinit(); + try output.writer.writeAll("{\"type\":\"response.create\","); + try output.writer.writeAll(sse_payload[1..index]); + try output.writer.writeAll(sse_payload[index + stream_fields.len ..]); + return output.toOwnedSlice(); +} + +fn streamWebSocketPrepared( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, +) !stream_provider.Result { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + const account_id = try chatgpt_oauth.extractAccountId(alloc, request.credential.secret); + defer alloc.free(account_id); + const auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.credential.secret}); + defer secret.zeroAndFree(alloc, auth_header); + const request_endpoint = if (io_mod.getenv(e2e_endpoint_env)) |override| endpoint: { + if (!gateway_client.isLoopbackHttpUrl(override)) return error.InvalidE2EOpenAICodexEndpoint; + break :endpoint override; + } else endpoint; + + var reducer = responses_protocol.Reducer.init(alloc); + defer reducer.deinit(alloc); + var bridge = WebSocketBridge{ + .alloc = alloc, + .reducer = &reducer, + .events = request.events, + .cancel_flag = request.cancel_flag, + .content_capture_limit = request.content_capture_limit, + }; + // The WebSocket transport does not replay after this point. Marking before + // the upgrade remains conservative if an intermediary accepts then drops it. + request.delivery.markPossiblySent(); + try websocket_transport.stream(alloc, .{ + .endpoint = request_endpoint, + .authorization = auth_header, + .account_id = account_id, + .session_id = request.session_id, + .payload = payload, + .cancel_flag = request.cancel_flag, + }, &bridge, WebSocketBridge.event); + const completion = reducer.finish(alloc, request.cancel_flag, bridge.streamLimits()) catch |err| + return mapReducerError(err); + return .{ .completed = .{ + .completion = completion, + .usage = .{ .immediate = null }, + .ownership = .owned, + } }; +} + +const WebSocketBridge = struct { + alloc: Allocator, + reducer: *responses_protocol.Reducer, + events: stream_provider.EventSink, + cancel_flag: *std.atomic.Value(bool), + content_capture_limit: ?usize, + + fn streamLimits(self: @This()) responses_protocol.StreamLimits { + _ = self; + return .{ + .aggregate_bytes = max_sse_aggregate_bytes, + .events = max_sse_events, + .tool_calls = max_tool_calls, + .tool_identity_bytes = max_tool_identity_bytes, + .tool_arguments_bytes = max_tool_arguments_bytes, + .provider_state_bytes = max_provider_state_bytes, + }; + } + + fn event(raw: *anyopaque, json_text: []const u8) !bool { + const self: *@This() = @ptrCast(@alignCast(raw)); + return self.reducer.applyJson( + self.alloc, + json_text, + .{ + .context = &self.events, + .on_content = EventBridge.content, + .on_tool_start = EventBridge.toolStart, + .on_reasoning = EventBridge.reasoning, + .on_tool_input = EventBridge.toolInput, + }, + self.cancel_flag, + self.content_capture_limit, + self.streamLimits(), + ) catch |err| return mapReducerError(err); + } +}; + const EventBridge = struct { fn sink(raw: *anyopaque) *stream_provider.EventSink { return @ptrCast(@alignCast(raw)); @@ -473,6 +585,31 @@ fn mapReducerError(err: anyerror) anyerror { }; } +test "OpenAI Codex WebSocket request uses response create framing" { + const sse_payload = "{\"model\":\"gpt-5.4\",\"store\":false,\"stream\":true,\"input\":[]}"; + const websocket_payload = try buildWebSocketRequest(std.testing.allocator, sse_payload); + defer std.testing.allocator.free(websocket_payload); + try std.testing.expectEqualStrings( + "{\"type\":\"response.create\",\"model\":\"gpt-5.4\",\"input\":[]}", + websocket_payload, + ); +} + +test "OpenAI Codex transport policy keeps auto on SSE during Phase 1" { + // Environment-dependent selection is covered by integration launch tests. + // This assertion records the Phase 1 default when no override is present. + if (io_mod.getenv(transport_env) == null) { + try std.testing.expectEqual(Transport.sse, try selectedTransport()); + } +} + +test "OpenAI Codex WebSocket request rejects a non-SSE payload" { + try std.testing.expectError( + error.InvalidOpenAICodexWebSocketRequest, + buildWebSocketRequest(std.testing.allocator, "{\"model\":\"gpt-5.4\"}"), + ); +} + test "OpenAI Codex request uses Responses input and converts AI SDK tool schemas" { const read_file_schema = model_tool_schema.FunctionSchema{ .name = "read_file", diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig new file mode 100644 index 000000000..7336f05d2 --- /dev/null +++ b/src/gateway/websocket_transport.zig @@ -0,0 +1,229 @@ +const std = @import("std"); +const io_mod = @import("../core/shared/io.zig"); + +const Allocator = std.mem.Allocator; + +pub const max_frame_bytes: usize = 4 * 1024 * 1024; +pub const max_message_bytes: usize = 64 * 1024 * 1024; + +pub const Error = error{ + WebSocketUpgradeRejected, + WebSocketAcceptInvalid, + WebSocketProtocolViolation, + WebSocketUnexpectedBinary, + WebSocketMessageTooLarge, + WebSocketClosedBeforeCompletion, +}; + +pub const EventHandler = *const fn (context: *anyopaque, json: []const u8) anyerror!bool; + +pub const Request = struct { + endpoint: []const u8, + authorization: []const u8, + account_id: []const u8, + session_id: ?[]const u8, + payload: []const u8, + cancel_flag: *std.atomic.Value(bool), +}; + +/// Opens one socket, sends one request, and consumes one terminal response. +/// The caller owns delivery certainty: this function returns an error after a +/// frame write without retrying the request. +pub fn stream( + alloc: Allocator, + request: Request, + context: *anyopaque, + on_event: EventHandler, +) !void { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + const uri = try std.Uri.parse(request.endpoint); + var nonce: [16]u8 = undefined; + try io_mod.getIo().randomSecure(&nonce); + var key_buffer: [std.base64.standard.Encoder.calcSize(nonce.len)]u8 = undefined; + _ = std.base64.standard.Encoder.encode(&key_buffer, &nonce); + var accept_buffer: [std.base64.standard.Encoder.calcSize(std.crypto.hash.Sha1.digest_length)]u8 = undefined; + const expected_accept = websocketAccept(&key_buffer, &accept_buffer); + + var extra_headers: [7]std.http.Header = undefined; + var count: usize = 0; + extra_headers[count] = .{ .name = "chatgpt-account-id", .value = request.account_id }; + count += 1; + extra_headers[count] = .{ .name = "originator", .value = "fx" }; + count += 1; + extra_headers[count] = .{ .name = "OpenAI-Beta", .value = "responses_websockets=v2" }; + count += 1; + extra_headers[count] = .{ .name = "Upgrade", .value = "websocket" }; + count += 1; + extra_headers[count] = .{ .name = "Sec-WebSocket-Version", .value = "13" }; + count += 1; + extra_headers[count] = .{ .name = "Sec-WebSocket-Key", .value = &key_buffer }; + count += 1; + if (request.session_id) |session_id| if (session_id.len > 0) { + extra_headers[count] = .{ .name = "session-id", .value = session_id }; + count += 1; + }; + + var client: std.http.Client = .{ .allocator = alloc, .io = io_mod.getIo() }; + defer client.deinit(); + var http_request = try client.request(.GET, uri, .{ + .headers = .{ + .authorization = .{ .override = request.authorization }, + .connection = .{ .override = "Upgrade" }, + .accept_encoding = .omit, + }, + .extra_headers = extra_headers[0..count], + .keep_alive = false, + .redirect_behavior = .unhandled, + }); + defer { + // An upgraded connection must never return to the HTTP pool. + if (http_request.connection) |connection| connection.closing = true; + http_request.deinit(); + } + try http_request.sendBodiless(); + const response = try http_request.receiveHead(&.{}); + if (response.head.status != .switching_protocols) return error.WebSocketUpgradeRejected; + if (!hasHeader(response.head, "upgrade", "websocket") or + !hasTokenHeader(response.head, "connection", "upgrade") or + !hasHeader(response.head, "sec-websocket-accept", expected_accept)) + { + return error.WebSocketAcceptInvalid; + } + + // `receiveHead` leaves any already-buffered WebSocket bytes on this reader. + const reader = http_request.reader.in; + const writer = http_request.connection.?.writer(); + try writeFrame(writer, .text, request.payload); + try http_request.connection.?.flush(); + + var message: std.ArrayList(u8) = .empty; + defer message.deinit(alloc); + var fragmented_opcode: ?Opcode = null; + while (true) { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + const frame = try readFrame(alloc, reader); + defer alloc.free(frame.payload); + switch (frame.opcode) { + .ping => { + try writeFrame(writer, .pong, frame.payload); + try http_request.connection.?.flush(); + }, + .pong => {}, + .close => return error.WebSocketClosedBeforeCompletion, + .binary => return error.WebSocketUnexpectedBinary, + .continuation => { + if (fragmented_opcode == null) return error.WebSocketProtocolViolation; + try appendMessage(&message, alloc, frame.payload); + if (!frame.fin) continue; + const opcode = fragmented_opcode.?; + fragmented_opcode = null; + if (opcode != .text) return error.WebSocketUnexpectedBinary; + if (try on_event(context, message.items)) return; + message.clearRetainingCapacity(); + }, + .text => { + if (fragmented_opcode != null) return error.WebSocketProtocolViolation; + try appendMessage(&message, alloc, frame.payload); + if (!frame.fin) { + fragmented_opcode = .text; + continue; + } + if (try on_event(context, message.items)) return; + message.clearRetainingCapacity(); + }, + } + } +} + +const Opcode = enum(u4) { continuation = 0, text = 1, binary = 2, close = 8, ping = 9, pong = 10 }; +const Frame = struct { fin: bool, opcode: Opcode, payload: []u8 }; + +fn websocketAccept(key: []const u8, output: []u8) []const u8 { + var hash = std.crypto.hash.Sha1.init(.{}); + hash.update(key); + hash.update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); + var digest: [std.crypto.hash.Sha1.digest_length]u8 = undefined; + hash.final(&digest); + _ = std.base64.standard.Encoder.encode(output, &digest); + return output; +} + +fn hasHeader(head: std.http.Client.Response.Head, name: []const u8, expected: []const u8) bool { + var it = head.iterateHeaders(); + while (it.next()) |header| { + if (std.ascii.eqlIgnoreCase(header.name, name) and std.mem.eql(u8, header.value, expected)) return true; + } + return false; +} + +fn hasTokenHeader(head: std.http.Client.Response.Head, name: []const u8, token: []const u8) bool { + var it = head.iterateHeaders(); + while (it.next()) |header| { + if (!std.ascii.eqlIgnoreCase(header.name, name)) continue; + var tokens = std.mem.splitScalar(u8, header.value, ','); + while (tokens.next()) |candidate| if (std.ascii.eqlIgnoreCase(std.mem.trim(u8, candidate, " \t"), token)) return true; + } + return false; +} + +fn appendMessage(message: *std.ArrayList(u8), alloc: Allocator, payload: []const u8) !void { + if (payload.len > max_message_bytes -| message.items.len) return error.WebSocketMessageTooLarge; + try message.appendSlice(alloc, payload); +} + +fn writeFrame(writer: *std.Io.Writer, opcode: Opcode, payload: []const u8) !void { + if (payload.len > max_message_bytes) return error.WebSocketMessageTooLarge; + var mask: [4]u8 = undefined; + try io_mod.getIo().randomSecure(&mask); + try writer.writeByte(0x80 | @as(u8, @intFromEnum(opcode))); + if (payload.len < 126) { + try writer.writeByte(0x80 | @as(u8, @intCast(payload.len))); + } else if (payload.len <= std.math.maxInt(u16)) { + try writer.writeByte(0x80 | 126); + try writer.writeInt(u16, @intCast(payload.len), .big); + } else { + try writer.writeByte(0x80 | 127); + try writer.writeInt(u64, @intCast(payload.len), .big); + } + try writer.writeAll(&mask); + var chunk: [4096]u8 = undefined; + var offset: usize = 0; + while (offset < payload.len) { + const length = @min(chunk.len, payload.len - offset); + for (payload[offset..][0..length], 0..) |byte, index| chunk[index] = byte ^ mask[(offset + index) % mask.len]; + try writer.writeAll(chunk[0..length]); + offset += length; + } +} + +fn readFrame(alloc: Allocator, reader: *std.Io.Reader) !Frame { + const first = try reader.takeByte(); + const second = try reader.takeByte(); + if (second & 0x80 != 0 or first & 0x70 != 0) return error.WebSocketProtocolViolation; + const fin = first & 0x80 != 0; + const opcode = std.enums.fromInt(Opcode, first & 0x0f) orelse return error.WebSocketProtocolViolation; + var length: u64 = second & 0x7f; + if (length == 126) length = try reader.takeInt(u16, .big); + if (length == 127) { + length = try reader.takeInt(u64, .big); + if (length & (@as(u64, 1) << 63) != 0) return error.WebSocketProtocolViolation; + } + if (length > max_frame_bytes) return error.WebSocketMessageTooLarge; + if (@intFromEnum(opcode) >= @intFromEnum(Opcode.close) and (!fin or length > 125)) return error.WebSocketProtocolViolation; + const payload = try alloc.alloc(u8, @intCast(length)); + errdefer alloc.free(payload); + try reader.readSliceAll(payload); + return .{ .fin = fin, .opcode = opcode, .payload = payload }; +} + +test "WebSocket accept matches RFC 6455" { + var output: [28]u8 = undefined; + try std.testing.expectEqualStrings("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=", websocketAccept("dGhlIHNhbXBsZSBub25jZQ==", &output)); +} + +test "fragment aggregation limits message size" { + var message: std.ArrayList(u8) = .empty; + defer message.deinit(std.testing.allocator); + try appendMessage(&message, std.testing.allocator, "hello"); + try std.testing.expectEqualStrings("hello", message.items); +} From 394c36ac4426a5203fda211f208860d30c343595 Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 15:55:40 +0530 Subject: [PATCH 02/21] Harden WebSocket frame validation --- docs/codex-websocket-transport.md | 331 ++++++++++++++++++++++++++++ src/gateway/websocket_transport.zig | 73 +++++- 2 files changed, 396 insertions(+), 8 deletions(-) create mode 100644 docs/codex-websocket-transport.md diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md new file mode 100644 index 000000000..f86a46c01 --- /dev/null +++ b/docs/codex-websocket-transport.md @@ -0,0 +1,331 @@ +# OpenAI Codex Responses WebSocket transport + +## Scope and evidence status + +This document concerns the ChatGPT-subscription Codex route at `chatgpt.com/backend-api/codex/responses`, not fx's Vercel AI Gateway route. Both use a Responses WebSocket protocol family, but they are separate backends. Public Responses API, Azure, and AI Gateway documentation must not be treated as a specification for the private ChatGPT backend. + +Claims are labelled as follows: + +- **Confirmed:** supported by the fx checkout or current upstream `openai/codex` source. +- **Likely, verify:** supported by public protocol documentation or related implementations, but not yet measured against the ChatGPT subscription backend. +- **Undocumented:** a private-backend behavior that requires live, authenticated probing before fx relies on it. + +This is an implementation design, not a claim that fx already supports WebSockets. + +## Current fx behavior + +**Confirmed.** The Codex provider currently uses HTTPS with Server-Sent Events (SSE): + +1. fx reads or refreshes the local ChatGPT OAuth session, extracts the ChatGPT account ID, and constructs a bearer token. +2. It serializes a complete Responses request containing the model, system instructions, conversation history, tools, tool outputs, images, and retained encrypted reasoning state. +3. It sends that JSON with `POST https://chatgpt.com/backend-api/codex/responses` and asks for `text/event-stream`. +4. The server streams `data:` records over that HTTP request. +5. `src/gateway/responses_protocol.zig` reduces each JSON event into text, reasoning, tool-call, and final-completion events. + +The implementation has a 30-second connection deadline, cancellation handling, and limits for aggregate stream data, events, tool calls, tool identities, tool arguments, and preserved provider state. Each model request opens a new HTTP connection. + +## What WebSockets change + +**Likely, verify for this backend.** A secure WebSocket begins with an HTTPS Upgrade request. After a valid `101 Switching Protocols` response, both sides retain an encrypted, bidirectional connection. + +```text +fx Codex service + HTTPS WebSocket Upgrade → + 101 Switching Protocols ← + response.create frame → + response event frames ← + response.completed ← +``` + +The Responses data remains JSON. The client sends a text message resembling: + +```json +{ + "type": "response.create", + "model": "gpt-5...", + "instructions": "...", + "input": [], + "tools": [] +} +``` + +The server returns text messages such as `response.output_text.delta`, `response.function_call_arguments.delta`, `response.completed`, and `response.failed`. A WebSocket event source should feed decoded JSON text directly into fx's existing Responses reducer. It must not create a second model-event implementation. + +## Upstream Codex CLI behavior + +**Confirmed for upstream Codex, not automatically for the private backend.** Upstream Codex has a dedicated `responses_websocket` transport, gated by provider capability, and retains a healthy connection for its client session. + +### Handshake and metadata + +**Confirmed.** Codex constructs a WebSocket URL, attaches provider and authentication headers during the HTTP upgrade, requests a versioned Responses WebSocket beta protocol, and validates the upgrade. It can record selected server model, reasoning inclusion, Codex turn state, rate-limit data, model ETags, moderation metadata, and timing information. + +It also supports a handshake probe that upgrades without sending a prompt and briefly waits for an immediate close. This distinguishes a usable connection from one accepted by the edge then rejected by policy. + +### One request at a time + +**Confirmed.** Codex serializes one response stream per socket. It sends `response.create`, reads until a terminal event, then permits the next request. It does not interleave independent generations on the same connection. + +This is the correct initial model for fx. Concurrent requests require request identity, event demultiplexing, flow control, and independent recovery. + +### Event failures and reuse + +**Confirmed.** Codex applies an idle timeout to each server-event wait. It treats idle timeout, EOF before completion, close before completion, unexpected binary messages, I/O failure, and structured service errors as stream failures. It handles Ping and Pong as transport housekeeping and discards a connection after a terminal stream error. + +**Confirmed for upstream source; likely, verify for this backend.** Upstream recognizes `websocket_connection_limit_reached` with a 60-minute connection limit and recognizes `previous_response_not_found` as recoverable by a full-context request. + +**Likely, verify.** `response.failed` must invalidate its continuation chain. A later request must not reuse the failed response as `previous_response_id`. + +### Incremental continuation + +**Confirmed for upstream client behavior; undocumented for the ChatGPT backend.** A retained connection can send a new `response.create` with `previous_response_id` and only newly added input items. This reduces repeated upload of long transcripts and tool history. + +Until Phase 0 verifies the ChatGPT backend, fx must regard continuation state as **possibly connection-scoped**, not guaranteed connection-scoped. Saved fx sessions must always retain enough history to reconstruct a full request; remote state is never the sole source of truth. + +## The central reliability rule: delivery certainty + +A network failure does not simply mean a request failed: + +- **Request has not begun writing:** automatic retry is safe. +- **Request may have been written but no acknowledgement arrived:** do not retry blindly. +- **The server acknowledged a response ID:** recover only through supported server state or a known-safe replay. + +If fx sends `response.create` then loses the network before `response.created`, the service may still generate a response and issue tool calls. Blindly retrying could duplicate shell commands or other actions. + +fx already models this distinction for HTTP with `DeliveryCertainty`. A WebSocket transport must preserve it. Upstream Codex may retry an established stream before falling back, but fx deliberately diverges: its tool-capable turns require a stricter no-blind-replay policy after delivery becomes ambiguous. + +## Phase 0: authenticated live-traffic probe + +Before transport implementation, run the isolated maintainer probe at `scripts/codex_websocket_probe.py` against `chatgpt.com/backend-api/codex/responses`. It uses only Python's standard library, accepts credentials only through explicit environment variables, never reads fx credential files, and emits one redacted JSON report. It never writes credentials, account IDs, response IDs, prompts, or response content. + +```bash +export FX_CODEX_PROBE_ACCESS_TOKEN='...' +export FX_CODEX_PROBE_ACCOUNT_ID='...' +export FX_CODEX_PROBE_MODEL='...' +python3 scripts/codex_websocket_probe.py +python3 scripts/codex_websocket_probe.py --execute --continuation +``` + +The first invocation performs only an authenticated upgrade. `--execute` sends a fixed minimal prompt and consumes subscription usage; `--continuation` sends a second request using the first response ID without printing that ID. Run `python3 scripts/codex_websocket_probe.py --self-test` for deterministic, credential-free checks. + +Record only privacy-safe protocol evidence: + +- handshake status, selected response headers, and immediate close behavior; +- event type sequence, terminal event, structured error code, and close code; +- whether WebSocket accepts `previous_response_id` continuation; +- whether HTTP/SSE accepts `previous_response_id` continuation; +- behavior after a failed response in a continuation chain; +- whether model changes on a retained connection produce a policy close such as 1008; +- connection-age behavior and any limit/error code; +- behavior of `store` and related persistence fields, if accepted. + +Do not hard-code public API assumptions until this probe confirms them. In particular, the 60-minute limit, `previous_response_not_found`, model pinning, `store` semantics, and SSE continuation support are undocumented for this private backend. + +## Proposed implementation plan + +### Phase 1: per-turn WebSocket transport + +Implement a fresh socket per model request, with no retained connection or cached continuation. + +- Keep SSE as the permanent compatibility baseline. +- Reuse existing Codex request generation and Responses reduction. +- Send one full `response.create` request on a fresh socket. +- Fall back to SSE immediately for failed upgrades, connection timeouts, or other failures definitely before delivery. +- After a request may have been transmitted, never silently replay through SSE. +- Treat server close before `response.completed` as an explicit failure class. +- Once Phase 0 confirms behavior, treat policy close 1008 as a wrong-model-on-connection failure: poison the connection and recreate it with the correct model, rather than retrying it as a network failure. +- Maintain a strict retry policy for established streams. Fewer retries, including zero, are safer than duplicating a possible tool call. + +### Phase 2: retained socket, full context per turn + +After Phase 1 has stable production evidence: + +- Retain at most one idle connection per active fx session, provider, account identity, and compatible model. +- Serialize one full request at a time over it. +- Preconnect only as an optimization. Never send prompt data during preconnect. +- Discard and recreate the connection after a close, protocol error, cancellation during a stream, timeout, failed write, authentication transition, or model incompatibility. +- Continue sending full context after reconnect. +- Proactively reconnect before a confirmed connection-age limit, with enough margin that the limit cannot interrupt an in-flight turn. +- Reset a connection-health or fallback budget only after confirmed `response.completed`, not merely after a successful handshake. +- Segment telemetry by authentication mode from the start. + +Connect and event-idle timeouts must be configurable and measured per platform. Do not assume Linux and macOS behavior predicts Windows behavior. + +### Phase 3: incremental continuation + +Only begin after Phase 0 verifies continuation semantics and Phase 2 has stable evidence. Maintain a continuation record containing: + +```text +connection identity +credential and account fingerprint +endpoint and protocol version +model and request-shaping options +last fully completed response ID +full request required for recovery +continuation expiry and validity +``` + +Invalidate it when the connection reconnects, model or relevant options change, authentication changes, the referenced turn fails, cancellation interrupts a stream, a tool call or result is uncertain, the server rejects the previous response ID, or the connection reaches its lifetime. + +### Phase 4: optional stream lanes + +**Likely, verify.** If the protocol supports a `stream_id` field with tagged events, named lanes could eventually allow concurrent subagent turns to share a retained connection. Requests in one lane must remain ordered; separate lanes require event demultiplexing and independent continuation state. + +Do not implement this before Phases 1 through 3 are stable. It substantially changes connection ownership, resource accounting, and failure recovery. + +## Transport requirements + +Keep WebSocket mechanics separate from Codex request serialization: + +```text +src/gateway/websocket_transport.zig RFC 6455 handshake, frames, cancellation, limits +src/gateway/responses_stream.zig JSON event source shared by SSE and WebSocket +src/gateway/openai_codex.zig Codex auth, payload, and transport selection +``` + +The transport must implement and test: + +- TLS certificate and hostname validation; +- HTTP `101 Switching Protocols` and `Sec-WebSocket-Accept` validation; +- mandatory masking of client frames; +- text-message fragmentation and continuation-frame reassembly; +- UTF-8 validation for text messages; +- Ping to Pong handling; +- close codes, bounded close reasons, and a finite close deadline; +- cancellation that unblocks connection, read, and write operations; +- separate connection, write, and event-idle deadlines; +- bounded outbound requests, inbound frames, reassembled messages, and buffered unread bytes; +- strict rejection of unexpected message kinds and protocol violations. + +Do not enable per-message WebSocket compression in the first release. It adds decompression limits, compatibility cases, CPU cost, and security surface. + +## Connection lifecycle and fallback policy + +Represent a connection as a state machine: + +```text +disconnected → connecting → open-idle → open-streaming → closing + ↘ poisoned → disconnected +``` + +One component owns socket I/O. UI, retry, and shutdown paths communicate through controlled cancellation or commands; they must never concurrently read from or write to the socket. + +A socket becomes poisoned and its continuation state is discarded after a timeout, protocol violation, unexpected binary event, early EOF, failed write, failed close, failed response, or server error that leaves state uncertain. + +Expose a prominent user-facing policy: + +```text +auto Prefer a known-good WebSocket, otherwise use SSE. +sse Always use HTTP and SSE. +websocket Require WebSocket and return an actionable error if unavailable. +``` + +In `auto`, only pre-delivery WebSocket failures may fall back automatically to SSE. After a request may have been delivered, report uncertainty or use verified server recovery. Do not silently replay. + +If fx later adds remote Responses compaction for this provider, route it over HTTP unless Phase 0 or subsequent probes confirm WebSocket support for that endpoint. + +## Observability and rollout + +Ship with a local kill switch and measured rollout. Record only privacy-safe operational data: + +- selected transport, protocol version, platform, fx version, and authentication mode; +- handshake status and duration; +- new versus reused connection; +- time to first event and terminal completion; +- error, close, retry, circuit-breaker, and fallback classification; +- bounded byte and event counts; +- continuation hit, invalidation, and recovery reason. + +Never log prompts, responses, OAuth tokens, account IDs, raw tool arguments, or unredacted headers. + +Track handshake success, terminal completion rate, fallback and ambiguous-delivery rates, latency by transport and auth mode, errors by platform/network, and long-session memory and file-descriptor stability. + +## Verification plan + +Before broad enablement, require: + +1. Zig unit tests for handshake validation, masking, fragmentation, control frames, malformed frames, UTF-8, and resource limits. +2. A scripted loopback fixture for delayed events, oversized frames, invalid events, disconnects at write/read boundaries, and server close before `response.completed`. +3. Reducer-parity tests that feed identical Responses JSON through SSE and WebSocket and assert identical completion data and callback ordering. +4. A delivery-certainty retry matrix, including a request that may have been sent but was never acknowledged and a duplicate-tool-call prevention case. +5. Model-mismatch coverage, after Phase 0 verifies the behavior, proving a policy close is not retried as a generic transient network error. +6. Continuation coverage proving that a failed referenced turn invalidates the chain and forces full context. +7. Circuit-breaker coverage proving that failures deplete the budget and only a confirmed completion resets it. +8. Cancellation and shutdown tests while connecting, writing, waiting for output, receiving tool arguments, and closing, including leak checks. +9. Platform-matrix connect-timeout tests and long-running soak tests with forced reconnects and connection-age expiry. +10. A real-binary smoke test using freshly built `./zig-out/bin/fx` against a loopback fixture and an interactive terminal path. + +## Phase 1 implementation review and required fixes + +This review applies to the Phase 1 implementation. It is not production-ready and must retain SSE as the default until every release blocker below is resolved and verified. + +### Implementation progress + +- [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target. Close-payload validation is implemented; the bounded close handshake remains below. +- [ ] Provider admission and shared Codex preparation. +- [ ] Bounded upgrade, write, and event-idle I/O with cancellation unblocking. +- [ ] Bounded close handshake and close-code reporting. +- [ ] Loopback WebSocket fixture, reducer-parity, delivery-certainty, and real-binary smoke coverage. +- [ ] Authenticated Phase 0 evidence: requires an explicit maintainer-run probe and is not satisfied by automated tests. + +### Release blockers + +1. **Extended-length frame parsing corrupts a valid 127-byte frame.** + + `src/gateway/websocket_transport.zig` reads the seven-bit length discriminator, then uses two independent `if` statements. If that discriminator is `126` and the following 16-bit length is exactly `127`, the first branch correctly reads `127`, then the second branch incorrectly treats it as the 64-bit discriminator and consumes eight payload bytes as a length. This desynchronizes the stream. + + Fix: make the 64-bit branch `else if (length == 127)`. Add a regression test for a 127-byte text frame followed by another frame, proving the second frame remains aligned. + +2. **The WebSocket path does not participate in provider admission.** + + The SSE path calls `request.admission.admit()` before opening transport. The WebSocket path currently does not, which produces the user-visible `provider admission missing` failure and bypasses the normal provider-attempt lifecycle. + + Fix: preserve the admission boundary for every transport before opening its request. Add a focused test that runs the WebSocket path through the same admission fixture as SSE. + +3. **The WebSocket path has no bounded connection deadline or cancellation unblock.** + + SSE opens through `runBoundedHttpOperation` with the 30-second connect deadline and installs `spawnHttpCancelWatcher` to interrupt blocked connection I/O. The WebSocket path invokes the HTTP client directly and only checks `cancel_flag` between completed frame reads. A blocked `reader.takeByte()` cannot observe cancellation until the peer sends bytes or closes. + + Fix: use an operation that bounds connection setup, writes, and event-idle reads, and that closes or interrupts the underlying connection when cancellation occurs. Verify cancellation while connecting, while writing, and while waiting for a frame. + +4. **Close handling does not meet the documented transport contract.** + + The current code marks the HTTP connection closing and deinitializes it without sending a WebSocket close frame on normal terminal, cancellation, and error paths. It also does not parse or report close code and bounded reason. + + Fix: send one bounded close frame when the socket is open, await a peer close for a finite deadline where appropriate, then forcibly release the connection. Record and classify peer close code and bounded reason. Cover normal completion, cancellation, malformed frames, and early peer close. + +5. **Text frames are not explicitly validated as UTF-8.** + + The reducer eventually parses JSON, but WebSocket text-message validity must be enforced at the frame-message boundary so malformed text is classified as a transport protocol error rather than an incidental JSON failure. + + Fix: validate completed text messages before calling the event handler, return a dedicated protocol error, and add malformed UTF-8 coverage. + +### Required cleanup before broad enablement + +1. **Replace fragile WebSocket request string surgery.** + + `buildWebSocketRequest` finds the literal `,"store":false,"stream":true` in an SSE payload and splices around it. This couples WebSocket behavior to exact field ordering and formatting in `buildRequest`. The current test uses a hardcoded SSE payload, so it cannot catch drift in the real request builder. + + Minimum fix: add an end-to-end unit assertion for `buildWebSocketRequest(buildRequest(...))` using representative system messages, tools, reasoning, images, and structured output. Preferred fix: split request construction into a transport-neutral Responses request model or shared field writer, then add transport-specific envelopes without parsing a serialized request. + +2. **Extract shared Codex authentication and endpoint preparation.** + + `streamPrepared` and `streamWebSocketPrepared` duplicate account-ID extraction, authorization-header allocation, and loopback-only E2E endpoint selection. + + Fix: introduce a small helper whose result owns or scopes the prepared account ID, authorization header, and endpoint. Keep credential material zeroed and free it at the same boundary as today. + +3. **Create one stream-limits builder.** + + The six stream-limit values are repeated in `CodexLimits`, `WebSocketBridge.streamLimits`, and the SSE `consumeSse` conversion. The WebSocket method ignores `self`, which is evidence it should be a shared pure builder. + + Fix: use one `codexStreamLimits(CodexLimits)` helper returning `responses_protocol.StreamLimits`; use it for both SSE and WebSocket reducers. + +### Lower-priority hardening and explicit decisions + +- Guard `http_request.connection` after a successful upgrade instead of using `connection.?`. A `101` should have a connection, but treating an impossible state as an error is safer than trapping. +- Document the asymmetric limits: outbound frames permit a full 64 MiB request, while inbound frames are limited to 4 MiB and reassembled messages to 64 MiB. This is reasonable if large outbound contexts are intentional and response events are expected to be fragmented, but it needs rationale and tests at each boundary. +- Keep `auto` mapped to SSE in Phase 1. This is intentional and compatible with the rollout plan, not a bug. Do not change `auto` to WebSocket until pre-delivery fallback, health evidence, and the release blockers are implemented. +- The Phase 0 Python probe is standard-library-only, reads only explicitly named environment credentials, emits redacted output, and is suitable to retain as an untracked diagnostic script. + +## Current status + +fx has robust Codex HTTPS/SSE request generation and Responses reduction. The uncommitted Phase 1 branch adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. The experiment is blocked from broad use by the required fixes above. diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index 7336f05d2..97f6959c7 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -12,6 +12,7 @@ pub const Error = error{ WebSocketProtocolViolation, WebSocketUnexpectedBinary, WebSocketMessageTooLarge, + WebSocketInvalidUtf8, WebSocketClosedBeforeCompletion, }; @@ -92,9 +93,10 @@ pub fn stream( // `receiveHead` leaves any already-buffered WebSocket bytes on this reader. const reader = http_request.reader.in; - const writer = http_request.connection.?.writer(); + const connection = http_request.connection orelse return error.WebSocketConnectionMissing; + const writer = connection.writer(); try writeFrame(writer, .text, request.payload); - try http_request.connection.?.flush(); + try connection.flush(); var message: std.ArrayList(u8) = .empty; defer message.deinit(alloc); @@ -106,10 +108,13 @@ pub fn stream( switch (frame.opcode) { .ping => { try writeFrame(writer, .pong, frame.payload); - try http_request.connection.?.flush(); + try connection.flush(); }, .pong => {}, - .close => return error.WebSocketClosedBeforeCompletion, + .close => { + try validateClosePayload(frame.payload); + return error.WebSocketClosedBeforeCompletion; + }, .binary => return error.WebSocketUnexpectedBinary, .continuation => { if (fragmented_opcode == null) return error.WebSocketProtocolViolation; @@ -118,7 +123,7 @@ pub fn stream( const opcode = fragmented_opcode.?; fragmented_opcode = null; if (opcode != .text) return error.WebSocketUnexpectedBinary; - if (try on_event(context, message.items)) return; + if (try dispatchTextMessage(context, on_event, message.items)) return; message.clearRetainingCapacity(); }, .text => { @@ -128,7 +133,7 @@ pub fn stream( fragmented_opcode = .text; continue; } - if (try on_event(context, message.items)) return; + if (try dispatchTextMessage(context, on_event, message.items)) return; message.clearRetainingCapacity(); }, } @@ -171,6 +176,21 @@ fn appendMessage(message: *std.ArrayList(u8), alloc: Allocator, payload: []const try message.appendSlice(alloc, payload); } +fn dispatchTextMessage(context: *anyopaque, on_event: EventHandler, message: []const u8) !bool { + if (!std.unicode.utf8ValidateSlice(message)) return error.WebSocketInvalidUtf8; + return on_event(context, message); +} + +fn validateClosePayload(payload: []const u8) !void { + if (payload.len == 1) return error.WebSocketProtocolViolation; + if (payload.len < 2) return; + const code = std.mem.readInt(u16, payload[0..2], .big); + if (code < 1000 or code >= 5000 or code == 1004 or code == 1005 or code == 1006 or code == 1015) { + return error.WebSocketProtocolViolation; + } + if (!std.unicode.utf8ValidateSlice(payload[2..])) return error.WebSocketInvalidUtf8; +} + fn writeFrame(writer: *std.Io.Writer, opcode: Opcode, payload: []const u8) !void { if (payload.len > max_message_bytes) return error.WebSocketMessageTooLarge; var mask: [4]u8 = undefined; @@ -203,8 +223,7 @@ fn readFrame(alloc: Allocator, reader: *std.Io.Reader) !Frame { const fin = first & 0x80 != 0; const opcode = std.enums.fromInt(Opcode, first & 0x0f) orelse return error.WebSocketProtocolViolation; var length: u64 = second & 0x7f; - if (length == 126) length = try reader.takeInt(u16, .big); - if (length == 127) { + if (length == 126) length = try reader.takeInt(u16, .big) else if (length == 127) { length = try reader.takeInt(u64, .big); if (length & (@as(u64, 1) << 63) != 0) return error.WebSocketProtocolViolation; } @@ -227,3 +246,41 @@ test "fragment aggregation limits message size" { try appendMessage(&message, std.testing.allocator, "hello"); try std.testing.expectEqualStrings("hello", message.items); } + +test "extended 127-byte frame keeps the following frame aligned" { + var encoded: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer encoded.deinit(); + try encoded.writer.writeAll(&.{ 0x81, 126, 0, 127 }); + try encoded.writer.splatByteAll('a', 127); + try encoded.writer.writeAll(&.{ 0x81, 2, 'o', 'k' }); + + var reader = std.Io.Reader.fixed(encoded.written()); + const first = try readFrame(std.testing.allocator, &reader); + defer std.testing.allocator.free(first.payload); + try std.testing.expectEqual(Opcode.text, first.opcode); + try std.testing.expectEqual(@as(usize, 127), first.payload.len); + + const second = try readFrame(std.testing.allocator, &reader); + defer std.testing.allocator.free(second.payload); + try std.testing.expectEqual(Opcode.text, second.opcode); + try std.testing.expectEqualStrings("ok", second.payload); +} + +test "text messages reject malformed UTF-8 before event dispatch" { + const Handler = struct { + fn handle(_: *anyopaque, _: []const u8) !bool { + return false; + } + }; + var context: u8 = 0; + try std.testing.expectError( + error.WebSocketInvalidUtf8, + dispatchTextMessage(@ptrCast(&context), Handler.handle, &.{ 0xc3, 0x28 }), + ); +} + +test "close payload rejects reserved codes and malformed UTF-8 reasons" { + try std.testing.expectError(error.WebSocketProtocolViolation, validateClosePayload(&.{ 0x03, 0xed })); + try std.testing.expectError(error.WebSocketInvalidUtf8, validateClosePayload(&.{ 0x03, 0xe8, 0xc3, 0x28 })); + try validateClosePayload(&.{ 0x03, 0xe8, 'o', 'k' }); +} From 5ac5be4c161801956ea602ecd199c21d236cd36d Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 16:07:43 +0530 Subject: [PATCH 03/21] Unify Codex WebSocket request setup --- docs/codex-websocket-transport.md | 2 +- src/gateway/openai_codex.zig | 242 +++++++++++++++++++----------- 2 files changed, 154 insertions(+), 90 deletions(-) diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index f86a46c01..366950f1b 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -261,7 +261,7 @@ This review applies to the Phase 1 implementation. It is not production-ready an ### Implementation progress - [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target. Close-payload validation is implemented; the bounded close handshake remains below. -- [ ] Provider admission and shared Codex preparation. +- [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. - [ ] Bounded upgrade, write, and event-idle I/O with cancellation unblocking. - [ ] Bounded close handshake and close-code reporting. - [ ] Loopback WebSocket fixture, reducer-parity, delivery-certainty, and real-binary smoke coverage. diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index 6cfad9ba8..fb6953b95 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -43,6 +43,17 @@ const CodexLimits = struct { provider_state_bytes: usize = max_provider_state_bytes, }; +fn codexStreamLimits(limits: CodexLimits) responses_protocol.StreamLimits { + return .{ + .aggregate_bytes = limits.aggregate_bytes, + .events = limits.events, + .tool_calls = limits.tool_calls, + .tool_identity_bytes = limits.tool_identity_bytes, + .tool_arguments_bytes = limits.tool_arguments_bytes, + .provider_state_bytes = limits.provider_state_bytes, + }; +} + pub const agent_stream_provider = stream_provider.Provider{ .stream_fn = streamCompletion, }; @@ -58,12 +69,48 @@ pub fn buildRequest( alloc: Allocator, request: stream_provider.RequestData, ) ![]u8 { + var out: std.Io.Writer.Allocating = .init(alloc); + errdefer out.deinit(); + try writeResponseRequestStart(&out.writer, alloc, request, null); + try out.writer.writeAll(",\"store\":false,\"stream\":true"); + try out.writer.writeByte('}'); + return out.toOwnedSlice(); +} + +fn buildWebSocketRequest( + alloc: Allocator, + request: stream_provider.RequestData, +) ![]u8 { + var out: std.Io.Writer.Allocating = .init(alloc); + errdefer out.deinit(); + try writeResponseRequestStart(&out.writer, alloc, request, "response.create"); + try out.writer.writeByte('}'); + return out.toOwnedSlice(); +} + +/// Writes the fields shared by SSE and WebSocket Responses envelopes. The +/// caller owns the opening and closing JSON object delimiters. +fn writeResponseRequestStart( + writer: *std.Io.Writer, + alloc: Allocator, + request: stream_provider.RequestData, + websocket_type: ?[]const u8, +) !void { try validateModel(request.model); if (request.budget) |budget| { if (budget.cancel_flag) |flag| if (flag.load(.seq_cst)) return error.Cancelled; _ = budget.deadline; } + try writer.writeByte('{'); + if (websocket_type) |value| { + try writer.writeAll("\"type\":"); + try std.json.Stringify.value(value, .{}, writer); + try writer.writeByte(','); + } + try writer.writeAll("\"model\":"); + try std.json.Stringify.value(request.model, .{}, writer); + var instructions: std.Io.Writer.Allocating = .init(alloc); defer instructions.deinit(); for (request.messages) |message| { @@ -75,12 +122,7 @@ pub fn buildRequest( } if (instructions.written().len == 0) try instructions.writer.writeAll("You are a helpful assistant."); - var out: std.Io.Writer.Allocating = .init(alloc); - errdefer out.deinit(); - const writer = &out.writer; - try writer.writeAll("{\"model\":"); - try std.json.Stringify.value(request.model, .{}, writer); - try writer.writeAll(",\"store\":false,\"stream\":true,\"instructions\":"); + try writer.writeAll(",\"instructions\":"); try std.json.Stringify.value(instructions.written(), .{}, writer); try writer.writeAll(",\"input\":["); try writeResponsesInput(writer, alloc, request.messages, request.verified_images); @@ -115,8 +157,6 @@ pub fn buildRequest( } // The ChatGPT Codex endpoint chooses the model's output limit and rejects // the public Responses API max_output_tokens parameter. - try writer.writeByte('}'); - return out.toOwnedSlice(); } fn writeResponsesInput( @@ -149,14 +189,16 @@ fn streamCompletion( return stream_provider.failResult(error.CodexSubscriptionCredentialRequired); } try validateModel(request.model); - const payload = try buildRequest(alloc, request.data()); - defer alloc.free(payload); return switch (try selectedTransport()) { - .sse => streamPrepared(alloc, request, payload), + .sse => blk: { + const payload = try buildRequest(alloc, request.data()); + defer alloc.free(payload); + break :blk streamPrepared(alloc, request, payload); + }, .websocket => blk: { - const websocket_payload = try buildWebSocketRequest(alloc, payload); - defer alloc.free(websocket_payload); - break :blk streamWebSocketPrepared(alloc, request, websocket_payload); + const payload = try buildWebSocketRequest(alloc, request.data()); + defer alloc.free(payload); + break :blk streamWebSocketPrepared(alloc, request, payload); }, } catch |err| { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); @@ -165,6 +207,10 @@ fn streamCompletion( }; } +fn admitCodexTransport(admission: stream_provider.Admission) !void { + try admission.admit(); +} + const OpenedRequest = struct { request: ?std.http.Client.Request, @@ -207,21 +253,13 @@ pub fn streamPrepared( payload: []const u8, ) !stream_provider.Result { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); - const account_id = try chatgpt_oauth.extractAccountId(alloc, request.credential.secret); - defer alloc.free(account_id); - const auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.credential.secret}); - defer secret.zeroAndFree(alloc, auth_header); - const request_endpoint = if (io_mod.getenv(e2e_endpoint_env)) |override| endpoint: { - if (!gateway_client.isLoopbackHttpUrl(override)) { - return stream_provider.failResult(error.InvalidE2EOpenAICodexEndpoint); - } - break :endpoint override; - } else endpoint; - const uri = try std.Uri.parse(request_endpoint); + var prepared = try prepareCodexTransport(alloc, request); + defer prepared.deinit(alloc); + const uri = try std.Uri.parse(prepared.endpoint); var extra_headers_buf: [7]std.http.Header = undefined; var extra_count: usize = 0; - extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = account_id }; + extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = prepared.account_id }; extra_count += 1; extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" }; extra_count += 1; @@ -241,14 +279,14 @@ pub fn streamPrepared( var open_operation = OpenRequestOperation{ .client = &client, .uri = uri, - .auth_header = auth_header, + .auth_header = prepared.authorization, .extra_headers = extra_headers_buf[0..extra_count], }; const connect_deadline = std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ .clock = .awake, .raw = .fromMilliseconds(connect_timeout_ms), }); - try request.admission.admit(); + try admitCodexTransport(request.admission); var opened = try gateway_client.runBoundedHttpOperation( OpenedRequest, alloc, @@ -339,32 +377,45 @@ pub fn streamPrepared( } }; } -fn buildWebSocketRequest(alloc: Allocator, sse_payload: []const u8) ![]u8 { - const stream_fields = ",\"store\":false,\"stream\":true"; - if (sse_payload.len < 2 or sse_payload[0] != '{') return error.InvalidOpenAICodexWebSocketRequest; - const index = std.mem.find(u8, sse_payload, stream_fields) orelse return error.InvalidOpenAICodexWebSocketRequest; - var output: std.Io.Writer.Allocating = .init(alloc); - errdefer output.deinit(); - try output.writer.writeAll("{\"type\":\"response.create\","); - try output.writer.writeAll(sse_payload[1..index]); - try output.writer.writeAll(sse_payload[index + stream_fields.len ..]); - return output.toOwnedSlice(); -} +const PreparedCodexTransport = struct { + account_id: []u8, + authorization: []u8, + endpoint: []const u8, -fn streamWebSocketPrepared( + fn deinit(self: *PreparedCodexTransport, alloc: Allocator) void { + alloc.free(self.account_id); + secret.zeroAndFree(alloc, self.authorization); + self.* = undefined; + } +}; + +fn prepareCodexTransport( alloc: Allocator, request: stream_provider.ModelRequest, - payload: []const u8, -) !stream_provider.Result { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; +) !PreparedCodexTransport { const account_id = try chatgpt_oauth.extractAccountId(alloc, request.credential.secret); - defer alloc.free(account_id); - const auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.credential.secret}); - defer secret.zeroAndFree(alloc, auth_header); + errdefer alloc.free(account_id); + const authorization = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.credential.secret}); + errdefer secret.zeroAndFree(alloc, authorization); const request_endpoint = if (io_mod.getenv(e2e_endpoint_env)) |override| endpoint: { if (!gateway_client.isLoopbackHttpUrl(override)) return error.InvalidE2EOpenAICodexEndpoint; break :endpoint override; } else endpoint; + return .{ + .account_id = account_id, + .authorization = authorization, + .endpoint = request_endpoint, + }; +} + +fn streamWebSocketPrepared( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, +) !stream_provider.Result { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + var prepared = try prepareCodexTransport(alloc, request); + defer prepared.deinit(alloc); var reducer = responses_protocol.Reducer.init(alloc); defer reducer.deinit(alloc); @@ -374,19 +425,21 @@ fn streamWebSocketPrepared( .events = request.events, .cancel_flag = request.cancel_flag, .content_capture_limit = request.content_capture_limit, + .stream_limits = codexStreamLimits(.{}), }; - // The WebSocket transport does not replay after this point. Marking before - // the upgrade remains conservative if an intermediary accepts then drops it. + // Admission is shared with SSE and happens before the upgrade can make + // delivery possible. Once frame writing begins, this transport never replays. + try admitCodexTransport(request.admission); request.delivery.markPossiblySent(); try websocket_transport.stream(alloc, .{ - .endpoint = request_endpoint, - .authorization = auth_header, - .account_id = account_id, + .endpoint = prepared.endpoint, + .authorization = prepared.authorization, + .account_id = prepared.account_id, .session_id = request.session_id, .payload = payload, .cancel_flag = request.cancel_flag, }, &bridge, WebSocketBridge.event); - const completion = reducer.finish(alloc, request.cancel_flag, bridge.streamLimits()) catch |err| + const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| return mapReducerError(err); return .{ .completed = .{ .completion = completion, @@ -401,18 +454,7 @@ const WebSocketBridge = struct { events: stream_provider.EventSink, cancel_flag: *std.atomic.Value(bool), content_capture_limit: ?usize, - - fn streamLimits(self: @This()) responses_protocol.StreamLimits { - _ = self; - return .{ - .aggregate_bytes = max_sse_aggregate_bytes, - .events = max_sse_events, - .tool_calls = max_tool_calls, - .tool_identity_bytes = max_tool_identity_bytes, - .tool_arguments_bytes = max_tool_arguments_bytes, - .provider_state_bytes = max_provider_state_bytes, - }; - } + stream_limits: responses_protocol.StreamLimits, fn event(raw: *anyopaque, json_text: []const u8) !bool { const self: *@This() = @ptrCast(@alignCast(raw)); @@ -428,7 +470,7 @@ const WebSocketBridge = struct { }, self.cancel_flag, self.content_capture_limit, - self.streamLimits(), + self.stream_limits, ) catch |err| return mapReducerError(err); } }; @@ -550,14 +592,7 @@ fn consumeSse( .on_reasoning = on_reasoning_chunk, .on_tool_input = on_tool_input_chunk, }; - const stream_limits = responses_protocol.StreamLimits{ - .aggregate_bytes = limits.aggregate_bytes, - .events = limits.events, - .tool_calls = limits.tool_calls, - .tool_identity_bytes = limits.tool_identity_bytes, - .tool_arguments_bytes = limits.tool_arguments_bytes, - .provider_state_bytes = limits.provider_state_bytes, - }; + const stream_limits = codexStreamLimits(limits); while (try sse.next(alloc, reader)) |json_text| { defer sse.release(); if (reducer.applyJson( @@ -585,14 +620,43 @@ fn mapReducerError(err: anyerror) anyerror { }; } -test "OpenAI Codex WebSocket request uses response create framing" { - const sse_payload = "{\"model\":\"gpt-5.4\",\"store\":false,\"stream\":true,\"input\":[]}"; - const websocket_payload = try buildWebSocketRequest(std.testing.allocator, sse_payload); +test "OpenAI Codex WebSocket request uses the shared response request fields" { + const messages = [_]types.ChatMessage{ + .{ .role = .system, .content = "Be concise." }, + .{ .role = .user, .content = "Read it." }, + }; + const request: stream_provider.RequestData = .{ + .model = "gpt-5.4", + .messages = &messages, + .tool_choice = .auto, + .provider_options = .{}, + }; + const sse_payload = try buildRequest(std.testing.allocator, request); + defer std.testing.allocator.free(sse_payload); + const websocket_payload = try buildWebSocketRequest(std.testing.allocator, request); defer std.testing.allocator.free(websocket_payload); - try std.testing.expectEqualStrings( - "{\"type\":\"response.create\",\"model\":\"gpt-5.4\",\"input\":[]}", - websocket_payload, - ); + + try std.testing.expect(std.mem.find(u8, websocket_payload, "\"type\":\"response.create\"") != null); + try std.testing.expect(std.mem.find(u8, websocket_payload, "\"model\":\"gpt-5.4\"") != null); + try std.testing.expect(std.mem.find(u8, websocket_payload, "\"instructions\":\"Be concise.\"") != null); + try std.testing.expect(std.mem.find(u8, websocket_payload, "\"stream\"") == null); + try std.testing.expect(std.mem.find(u8, websocket_payload, "\"store\"") == null); + try std.testing.expect(std.mem.find(u8, sse_payload, "\"stream\":true") != null); + try std.testing.expect(std.mem.find(u8, sse_payload, "\"store\":false") != null); +} + +test "OpenAI Codex transport admission invokes the shared admission boundary" { + const Capture = struct { + called: bool = false, + + fn admit(raw: *anyopaque) !void { + const self: *@This() = @ptrCast(@alignCast(raw)); + self.called = true; + } + }; + var capture: Capture = .{}; + try admitCodexTransport(.{ .context = &capture, .admit_fn = Capture.admit }); + try std.testing.expect(capture.called); } test "OpenAI Codex transport policy keeps auto on SSE during Phase 1" { @@ -603,13 +667,6 @@ test "OpenAI Codex transport policy keeps auto on SSE during Phase 1" { } } -test "OpenAI Codex WebSocket request rejects a non-SSE payload" { - try std.testing.expectError( - error.InvalidOpenAICodexWebSocketRequest, - buildWebSocketRequest(std.testing.allocator, "{\"model\":\"gpt-5.4\"}"), - ); -} - test "OpenAI Codex request uses Responses input and converts AI SDK tool schemas" { const read_file_schema = model_tool_schema.FunctionSchema{ .name = "read_file", @@ -626,14 +683,17 @@ test "OpenAI Codex request uses Responses input and converts AI SDK tool schemas }, .{ .role = .tool, .tool_call_id = "call_1", .tool_name = "read_file", .content = "contents" }, }; - const body = try buildRequest(std.testing.allocator, .{ + const request: stream_provider.RequestData = .{ .model = "gpt-5.4", .messages = &messages, .tools = .{ .additional_functions = &.{read_file_schema} }, .tool_choice = .auto, .provider_options = .{ .reasoning = types.ReasoningEffort.literal("high"), .fast = true }, - }); + }; + const body = try buildRequest(std.testing.allocator, request); defer std.testing.allocator.free(body); + const websocket_body = try buildWebSocketRequest(std.testing.allocator, request); + defer std.testing.allocator.free(websocket_body); try std.testing.expect(std.mem.find(u8, body, "\"model\":\"gpt-5.4\"") != null); try std.testing.expect(std.mem.find(u8, body, "\"instructions\":\"Be concise.\"") != null); @@ -643,6 +703,10 @@ test "OpenAI Codex request uses Responses input and converts AI SDK tool schemas try std.testing.expect(std.mem.find(u8, body, "\"reasoning\":{\"effort\":\"high\"") != null); try std.testing.expect(std.mem.find(u8, body, "\"service_tier\":\"priority\"") != null); try std.testing.expect(std.mem.find(u8, body, "\"max_output_tokens\"") == null); + try std.testing.expect(std.mem.find(u8, websocket_body, "\"type\":\"response.create\"") != null); + try std.testing.expect(std.mem.find(u8, websocket_body, "\"encrypted_content\":\"opaque\"") != null); + try std.testing.expect(std.mem.find(u8, websocket_body, "\"parameters\":{\"type\":\"object\",\"properties\":{}}") != null); + try std.testing.expect(std.mem.find(u8, websocket_body, "\"reasoning\":{\"effort\":\"high\"") != null); } fn makeSizedProviderState(alloc: Allocator, size: usize) ![]u8 { From 9592e2eaf61c1e266e4325918bd86ec016d4ec87 Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 16:20:49 +0530 Subject: [PATCH 04/21] Bound Codex WebSocket I/O --- docs/codex-websocket-transport.md | 4 +- src/gateway/openai_codex.zig | 1 + src/gateway/websocket_transport.zig | 201 +++++++++++++++++++++++++--- 3 files changed, 183 insertions(+), 23 deletions(-) diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index 366950f1b..82d5bcb3c 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -262,8 +262,8 @@ This review applies to the Phase 1 implementation. It is not production-ready an - [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target. Close-payload validation is implemented; the bounded close handshake remains below. - [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. -- [ ] Bounded upgrade, write, and event-idle I/O with cancellation unblocking. -- [ ] Bounded close handshake and close-code reporting. +- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented; loopback cancellation coverage remains. +- [~] Bounded close handshake and close-code reporting: client close emission, peer close validation, and policy-close classification are implemented; loopback close-sequence coverage remains. - [ ] Loopback WebSocket fixture, reducer-parity, delivery-certainty, and real-binary smoke coverage. - [ ] Authenticated Phase 0 evidence: requires an explicit maintainer-run probe and is not satisfied by automated tests. diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index fb6953b95..aad2d5a98 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -437,6 +437,7 @@ fn streamWebSocketPrepared( .account_id = prepared.account_id, .session_id = request.session_id, .payload = payload, + .deadline = request.deadline, .cancel_flag = request.cancel_flag, }, &bridge, WebSocketBridge.event); const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index 97f6959c7..c72173ab4 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -1,5 +1,6 @@ const std = @import("std"); const io_mod = @import("../core/shared/io.zig"); +const gateway_client = @import("client.zig"); const Allocator = std.mem.Allocator; @@ -13,20 +14,60 @@ pub const Error = error{ WebSocketUnexpectedBinary, WebSocketMessageTooLarge, WebSocketInvalidUtf8, + WebSocketPolicyClosed, WebSocketClosedBeforeCompletion, }; pub const EventHandler = *const fn (context: *anyopaque, json: []const u8) anyerror!bool; +const connect_timeout_ms: i64 = 30_000; +const event_idle_timeout_ms: i64 = 30_000; + pub const Request = struct { endpoint: []const u8, authorization: []const u8, account_id: []const u8, session_id: ?[]const u8, payload: []const u8, + deadline: ?std.Io.Clock.Timestamp, cancel_flag: *std.atomic.Value(bool), }; +const OpenedRequest = struct { + request: ?std.http.Client.Request, + + pub fn deinit(self: *OpenedRequest, _: Allocator) void { + if (self.request) |*request| request.deinit(); + self.request = null; + } + + fn take(self: *OpenedRequest) std.http.Client.Request { + const request = self.request.?; + self.request = null; + return request; + } +}; + +const OpenWebSocketOperation = struct { + client: *std.http.Client, + uri: std.Uri, + authorization: []const u8, + headers: []const std.http.Header, + + pub fn run(self: *@This()) !OpenedRequest { + return .{ .request = try self.client.request(.GET, self.uri, .{ + .headers = .{ + .authorization = .{ .override = self.authorization }, + .connection = .{ .override = "Upgrade" }, + .accept_encoding = .omit, + }, + .extra_headers = self.headers, + .keep_alive = false, + .redirect_behavior = .unhandled, + }) }; + } +}; + /// Opens one socket, sends one request, and consumes one terminal response. /// The caller owns delivery certainty: this function returns an error after a /// frame write without retrying the request. @@ -66,23 +107,60 @@ pub fn stream( var client: std.http.Client = .{ .allocator = alloc, .io = io_mod.getIo() }; defer client.deinit(); - var http_request = try client.request(.GET, uri, .{ - .headers = .{ - .authorization = .{ .override = request.authorization }, - .connection = .{ .override = "Upgrade" }, - .accept_encoding = .omit, - }, - .extra_headers = extra_headers[0..count], - .keep_alive = false, - .redirect_behavior = .unhandled, + var open_operation = OpenWebSocketOperation{ + .client = &client, + .uri = uri, + .authorization = request.authorization, + .headers = extra_headers[0..count], + }; + var connect_deadline = std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ + .clock = .awake, + .raw = .fromMilliseconds(connect_timeout_ms), }); + if (request.deadline) |deadline| { + if (std.Io.Clock.Timestamp.compare(deadline, .lt, connect_deadline)) { + connect_deadline = deadline; + } + } + var opened = try gateway_client.runBoundedHttpOperation( + OpenedRequest, + alloc, + request.cancel_flag, + connect_deadline, + &open_operation, + ); + var http_request = opened.take(); defer { // An upgraded connection must never return to the HTTP pool. if (http_request.connection) |connection| connection.closing = true; http_request.deinit(); } - try http_request.sendBodiless(); - const response = try http_request.receiveHead(&.{}); + var watcher_done = std.atomic.Value(bool).init(false); + var timeout_fired = std.atomic.Value(bool).init(false); + var last_progress_ms = std.atomic.Value(i64).init(io_mod.milliTimestamp()); + const watcher = if (http_request.connection) |connection| + try spawnConnectionWatcher( + &watcher_done, + request.cancel_flag, + request.deadline, + &timeout_fired, + &last_progress_ms, + connection.stream_writer.stream, + ) + else + null; + defer { + watcher_done.store(true, .seq_cst); + if (watcher) |thread| thread.join(); + } + http_request.sendBodiless() catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; + const response = http_request.receiveHead(&.{}) catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; if (response.head.status != .switching_protocols) return error.WebSocketUpgradeRejected; if (!hasHeader(response.head, "upgrade", "websocket") or !hasTokenHeader(response.head, "connection", "upgrade") or @@ -95,25 +173,47 @@ pub fn stream( const reader = http_request.reader.in; const connection = http_request.connection orelse return error.WebSocketConnectionMissing; const writer = connection.writer(); - try writeFrame(writer, .text, request.payload); - try connection.flush(); + defer { + // The watcher bounds this write and forces release if the peer does not + // cooperate. A fresh Phase 1 socket is never returned to the HTTP pool. + writeFrame(writer, .close, &.{ 0x03, 0xe8 }) catch {}; + connection.flush() catch {}; + } + writeFrame(writer, .text, request.payload) catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; + connection.flush() catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; + last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); var message: std.ArrayList(u8) = .empty; defer message.deinit(alloc); var fragmented_opcode: ?Opcode = null; while (true) { if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - const frame = try readFrame(alloc, reader); + const frame = readFrame(alloc, reader) catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; + last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); defer alloc.free(frame.payload); switch (frame.opcode) { .ping => { - try writeFrame(writer, .pong, frame.payload); - try connection.flush(); + writeFrame(writer, .pong, frame.payload) catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; + connection.flush() catch |err| { + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; }, .pong => {}, .close => { - try validateClosePayload(frame.payload); - return error.WebSocketClosedBeforeCompletion; + return closeError(try validateClosePayload(frame.payload)); }, .binary => return error.WebSocketUnexpectedBinary, .continuation => { @@ -181,14 +281,71 @@ fn dispatchTextMessage(context: *anyopaque, on_event: EventHandler, message: []c return on_event(context, message); } -fn validateClosePayload(payload: []const u8) !void { +fn closeError(code: ?u16) anyerror { + if (code == 1008) return error.WebSocketPolicyClosed; + return error.WebSocketClosedBeforeCompletion; +} + +fn validateClosePayload(payload: []const u8) !?u16 { if (payload.len == 1) return error.WebSocketProtocolViolation; - if (payload.len < 2) return; + if (payload.len < 2) return null; const code = std.mem.readInt(u16, payload[0..2], .big); if (code < 1000 or code >= 5000 or code == 1004 or code == 1005 or code == 1006 or code == 1015) { return error.WebSocketProtocolViolation; } if (!std.unicode.utf8ValidateSlice(payload[2..])) return error.WebSocketInvalidUtf8; + return code; +} + +const ConnectionWatcher = struct { + fn run( + done: *std.atomic.Value(bool), + cancel_flag: *std.atomic.Value(bool), + deadline: ?std.Io.Clock.Timestamp, + timeout_fired: *std.atomic.Value(bool), + last_progress_ms: *std.atomic.Value(i64), + socket: std.Io.net.Stream, + ) void { + while (!done.load(.seq_cst)) { + if (cancel_flag.load(.seq_cst)) { + socket.shutdown(io_mod.getIo(), .both) catch {}; + return; + } + if (deadline) |limit| { + const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + if (!std.Io.Clock.Timestamp.compare(now, .lt, limit)) { + timeout_fired.store(true, .seq_cst); + socket.shutdown(io_mod.getIo(), .both) catch {}; + return; + } + } + const elapsed_ms = io_mod.milliTimestamp() - last_progress_ms.load(.seq_cst); + if (elapsed_ms >= event_idle_timeout_ms) { + timeout_fired.store(true, .seq_cst); + socket.shutdown(io_mod.getIo(), .both) catch {}; + return; + } + io_mod.sleep(10 * std.time.ns_per_ms); + } + } +}; + +fn spawnConnectionWatcher( + done: *std.atomic.Value(bool), + cancel_flag: *std.atomic.Value(bool), + deadline: ?std.Io.Clock.Timestamp, + timeout_fired: *std.atomic.Value(bool), + last_progress_ms: *std.atomic.Value(i64), + socket: std.Io.net.Stream, +) !std.Thread { + return std.Thread.spawn(.{}, ConnectionWatcher.run, .{ + done, + cancel_flag, + deadline, + timeout_fired, + last_progress_ms, + socket, + }); } fn writeFrame(writer: *std.Io.Writer, opcode: Opcode, payload: []const u8) !void { @@ -282,5 +439,7 @@ test "text messages reject malformed UTF-8 before event dispatch" { test "close payload rejects reserved codes and malformed UTF-8 reasons" { try std.testing.expectError(error.WebSocketProtocolViolation, validateClosePayload(&.{ 0x03, 0xed })); try std.testing.expectError(error.WebSocketInvalidUtf8, validateClosePayload(&.{ 0x03, 0xe8, 0xc3, 0x28 })); - try validateClosePayload(&.{ 0x03, 0xe8, 'o', 'k' }); + try std.testing.expectEqual(@as(?u16, 1000), try validateClosePayload(&.{ 0x03, 0xe8, 'o', 'k' })); + try std.testing.expectEqual(error.WebSocketPolicyClosed, closeError(1008)); + try std.testing.expectEqual(error.WebSocketClosedBeforeCompletion, closeError(1000)); } From 7b06750cb6e61ff57ba7e66f9e9d579db52db05f Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 16:22:44 +0530 Subject: [PATCH 05/21] Complete WebSocket close handshake --- src/gateway/websocket_transport.zig | 33 +++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index c72173ab4..b81718ad5 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -173,12 +173,13 @@ pub fn stream( const reader = http_request.reader.in; const connection = http_request.connection orelse return error.WebSocketConnectionMissing; const writer = connection.writer(); - defer { + var close_sent = false; + defer if (!close_sent) { // The watcher bounds this write and forces release if the peer does not // cooperate. A fresh Phase 1 socket is never returned to the HTTP pool. writeFrame(writer, .close, &.{ 0x03, 0xe8 }) catch {}; connection.flush() catch {}; - } + }; writeFrame(writer, .text, request.payload) catch |err| { if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; @@ -223,7 +224,11 @@ pub fn stream( const opcode = fragmented_opcode.?; fragmented_opcode = null; if (opcode != .text) return error.WebSocketUnexpectedBinary; - if (try dispatchTextMessage(context, on_event, message.items)) return; + if (try dispatchTextMessage(context, on_event, message.items)) { + try closeAfterCompletion(alloc, reader, writer, connection); + close_sent = true; + return; + } message.clearRetainingCapacity(); }, .text => { @@ -233,7 +238,11 @@ pub fn stream( fragmented_opcode = .text; continue; } - if (try dispatchTextMessage(context, on_event, message.items)) return; + if (try dispatchTextMessage(context, on_event, message.items)) { + try closeAfterCompletion(alloc, reader, writer, connection); + close_sent = true; + return; + } message.clearRetainingCapacity(); }, } @@ -281,6 +290,22 @@ fn dispatchTextMessage(context: *anyopaque, on_event: EventHandler, message: []c return on_event(context, message); } +fn closeAfterCompletion( + alloc: Allocator, + reader: *std.Io.Reader, + writer: *std.Io.Writer, + connection: anytype, +) !void { + try writeFrame(writer, .close, &.{ 0x03, 0xe8 }); + try connection.flush(); + const frame = try readFrame(alloc, reader); + defer alloc.free(frame.payload); + switch (frame.opcode) { + .close => _ = try validateClosePayload(frame.payload), + else => return error.WebSocketProtocolViolation, + } +} + fn closeError(code: ?u16) anyerror { if (code == 1008) return error.WebSocketPolicyClosed; return error.WebSocketClosedBeforeCompletion; From 3e03e89b9e01848bd17d169d59dc297be31e4a97 Mon Sep 17 00:00:00 2001 From: thinkter Date: Wed, 26 Aug 2026 16:57:49 +0530 Subject: [PATCH 06/21] Cover Codex WebSocket loopback behavior --- docs/codex-websocket-transport.md | 10 +- tests/e2e/tui-auth-source-selection.test.ts | 173 ++++++++++++++++++++ 2 files changed, 178 insertions(+), 5 deletions(-) diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index 82d5bcb3c..e3424c065 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -260,11 +260,11 @@ This review applies to the Phase 1 implementation. It is not production-ready an ### Implementation progress -- [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target. Close-payload validation is implemented; the bounded close handshake remains below. +- [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target, including a 127-byte extended frame followed by another frame. Close payloads and malformed text are also validated at the frame boundary. - [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. -- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented; loopback cancellation coverage remains. -- [~] Bounded close handshake and close-code reporting: client close emission, peer close validation, and policy-close classification are implemented; loopback close-sequence coverage remains. -- [ ] Loopback WebSocket fixture, reducer-parity, delivery-certainty, and real-binary smoke coverage. +- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented. A tmux loopback test proves cancellation unblocks an idle frame read; connect- and write-block cancellation remain untested. +- [~] Bounded close handshake and close-code reporting: normal completion sends and receives a close frame in the loopback smoke test; malformed close payloads are unit-tested, and an early policy close is classified without retry. Cancellation-close sequencing remains untested. +- [~] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for a completion, early policy close, and idle cancellation. Full SSE/WebSocket reducer-parity and delivery-certainty matrix coverage remain. - [ ] Authenticated Phase 0 evidence: requires an explicit maintainer-run probe and is not satisfied by automated tests. ### Release blockers @@ -328,4 +328,4 @@ This review applies to the Phase 1 implementation. It is not production-ready an ## Current status -fx has robust Codex HTTPS/SSE request generation and Responses reduction. The uncommitted Phase 1 branch adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. The experiment is blocked from broad use by the required fixes above. +fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. The experiment is intentionally blocked from broad use until the remaining coverage and authenticated Phase 0 evidence are complete. diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index bac87d5b9..f6a73662f 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -837,6 +837,52 @@ function startFakeCodexToolLoop(options: { }; } +function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: number } = {}) { + const requests: string[] = []; + const closeCodes: number[] = []; + const accessToken = chatgptAccessToken("acct_websocket"); + const server = Bun.serve<{ opened: boolean }>({ + hostname: "127.0.0.1", + port: 0, + fetch(request, server) { + const path = new URL(request.url).pathname; + if (path === "/models") { + return Response.json({ models: [ + { slug: "gpt-5.6-sol", visibility: "list", supported_in_api: true, supported_reasoning_levels: [{ effort: "high" }], additional_speed_tiers: [], input_modalities: ["text"], context_window: 272000 }, + ] }); + } + if (path === "/responses" && server.upgrade(request, { data: { opened: true } })) return; + return new Response("not found", { status: 404 }); + }, + websocket: { + open(ws) { + if (options.closeOnOpen !== undefined) ws.close(options.closeOnOpen, "fixture close"); + }, + message(ws, message) { + const payload = String(message); + requests.push(payload); + if (options.holdOpen) return; + ws.send(JSON.stringify({ type: "response.output_text.delta", delta: "CODEX_WEBSOCKET_OK" })); + ws.send(JSON.stringify({ + type: "response.completed", + response: { id: "resp_websocket", status: "completed", usage: { input_tokens: 5, output_tokens: 2 } }, + })); + }, + close(_ws, code) { + closeCodes.push(code); + }, + }, + }); + return { + accessToken, + requests, + closeCodes, + responsesUrl: `http://127.0.0.1:${server.port}/responses`, + modelsUrl: `http://127.0.0.1:${server.port}/models`, + stop() { server.stop(true); }, + }; +} + function startFakeCodexCapacityLoop() { const bodies: string[] = []; const accessToken = chatgptAccessToken("acct_capacity_loop"); @@ -2762,6 +2808,133 @@ tmuxTest( 60_000, ); +test( + "Codex WebSocket streams a completion through the freshly built binary", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-")); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket(); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + const result = await runFx( + ["ask", "--json", "--auto", "--no-save", "Use the WebSocket transport."], + { + env: { + HOME: home, + AI_GATEWAY_API_KEY: "gateway-websocket-sentinel", + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_AUTO_UPGRADE: "0", + FX_CODEX_TRANSPORT: "websocket", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }, + timeoutMs: TIMEOUT, + }, + ); + expect(result.code, `stdout: ${result.stdout}\nstderr: ${result.stderr}`).toBe(0); + expect(result.stdout).toContain("CODEX_WEBSOCKET_OK"); + expect(codex.requests).toHaveLength(1); + expect(codex.requests[0]).toContain('"type":"response.create"'); + expect(codex.requests[0]).not.toContain('"stream"'); + expect(codex.closeCodes).toContain(1000); + expect(gateway.requests).toHaveLength(0); + } finally { + codex.stop(); + } + }, + 60_000, +); + +test( + "Codex WebSocket policy close is not retried as a transport failure", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-policy-close-")); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ closeOnOpen: 1008 }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + const result = await runFx( + ["ask", "--json", "--auto", "--no-save", "Reject this request by policy."], + { + env: { + HOME: home, + AI_GATEWAY_API_KEY: "gateway-websocket-policy-sentinel", + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_AUTO_UPGRADE: "0", + FX_CODEX_TRANSPORT: "websocket", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }, + timeoutMs: TIMEOUT, + }, + ); + expect(result.code).toBe(1); + expect(`${result.stdout}\n${result.stderr}`).toContain("WebSocketPolicyClosed"); + expect(codex.requests).toHaveLength(0); + expect(gateway.requests).toHaveLength(0); + } finally { + codex.stop(); + } + }, + 60_000, +); + +tmuxTest( + "Codex WebSocket cancellation unblocks an idle response", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-cancel-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ holdOpen: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Wait for a held WebSocket response."); + const requestDeadline = Date.now() + TIMEOUT; + while (codex.requests.length === 0) { + if (Date.now() >= requestDeadline) throw new Error("Codex WebSocket request did not arrive"); + await Bun.sleep(25); + } + await session.sendKeys("C-c"); + await session.waitForComposer(TIMEOUT); + expect(session.isAlive()).toBe(true); + expect(codex.requests).toHaveLength(1); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + test( "ChatGPT tool loops round-trip encrypted reasoning without Gateway leakage", async () => { From 8485e618d74d15b789351f678e128c3abac0ab5f Mon Sep 17 00:00:00 2001 From: thinkter Date: Thu, 27 Aug 2026 22:05:55 +0530 Subject: [PATCH 07/21] Record Codex WebSocket Phase 0 evidence --- docs/codex-websocket-transport.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index e3424c065..462a255eb 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -265,7 +265,7 @@ This review applies to the Phase 1 implementation. It is not production-ready an - [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented. A tmux loopback test proves cancellation unblocks an idle frame read; connect- and write-block cancellation remain untested. - [~] Bounded close handshake and close-code reporting: normal completion sends and receives a close frame in the loopback smoke test; malformed close payloads are unit-tested, and an early policy close is classified without retry. Cancellation-close sequencing remains untested. - [~] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for a completion, early policy close, and idle cancellation. Full SSE/WebSocket reducer-parity and delivery-certainty matrix coverage remain. -- [ ] Authenticated Phase 0 evidence: requires an explicit maintainer-run probe and is not satisfied by automated tests. +- [x] Authenticated Phase 0 evidence: an explicit maintainer probe on 2026-08-27 returned `101`, completed a request, and accepted a same-socket `previous_response_id` continuation. The redacted report recorded only protocol labels, a 1,749 ms handshake, and the `x-models-etag` header name. ### Release blockers @@ -328,4 +328,4 @@ This review applies to the Phase 1 implementation. It is not production-ready an ## Current status -fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. The experiment is intentionally blocked from broad use until the remaining coverage and authenticated Phase 0 evidence are complete. +fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. A live authenticated probe confirms the Phase 0 handshake and continuation assumptions, but the experiment remains intentionally blocked from broad use until the remaining local coverage is complete. From 5cba97ca8e3ac2ba2da21e1e16cf4bc942ef4b6f Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 00:15:38 +0530 Subject: [PATCH 08/21] Add Codex WebSocket protocol probe --- scripts/codex_websocket_probe.py | 391 +++++++++++++++++++++++++++++++ 1 file changed, 391 insertions(+) create mode 100644 scripts/codex_websocket_probe.py diff --git a/scripts/codex_websocket_probe.py b/scripts/codex_websocket_probe.py new file mode 100644 index 000000000..98a042ad0 --- /dev/null +++ b/scripts/codex_websocket_probe.py @@ -0,0 +1,391 @@ +#!/usr/bin/env python3 +"""Isolated Phase 0 probe for the ChatGPT Codex Responses WebSocket endpoint. + +This is intentionally not part of fx's runtime transport. It uses only Python's +standard library, reads credentials only from explicitly named environment +variables, never writes credentials or response content, and prints a redacted +JSON report. + +Required for network use: + FX_CODEX_PROBE_ACCESS_TOKEN + FX_CODEX_PROBE_ACCOUNT_ID + FX_CODEX_PROBE_MODEL + +Examples: + python3 scripts/codex_websocket_probe.py + python3 scripts/codex_websocket_probe.py --execute --continuation + +The default performs only the authenticated WebSocket upgrade. --execute sends +a fixed minimal prompt and consumes subscription usage. --continuation sends a +second request using the first response ID, but does not print that ID. +""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import os +import secrets +import socket +import ssl +import sys +import time +from dataclasses import dataclass, field +from typing import Any + +HOST = "chatgpt.com" +PATH = "/backend-api/codex/responses" +ORIGIN = "https://chatgpt.com" +PROTOCOL_HEADER = "responses_websockets=v2" +CONNECT_TIMEOUT_SECONDS = 15.0 +EVENT_IDLE_TIMEOUT_SECONDS = 45.0 +MAX_FRAME_BYTES = 1 << 20 +MAX_MESSAGE_BYTES = 4 << 20 +MAX_EVENTS = 256 + + +class ProbeError(Exception): + pass + + +@dataclass +class Report: + handshake_status: int | None = None + handshake_elapsed_ms: int | None = None + selected_header_names: list[str] = field(default_factory=list) + immediate_close_code: int | None = None + event_types: list[str] = field(default_factory=list) + terminal_event: str | None = None + close_code: int | None = None + close_reason_length: int | None = None + continuation_attempted: bool = False + continuation_accepted: bool | None = None + error_code: str | None = None + error: str | None = None + + def emit(self) -> None: + print(json.dumps(self.__dict__, separators=(",", ":"), sort_keys=True)) + + +def required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise ProbeError(f"missing required environment variable {name}") + return value + + +def websocket_accept(key: str) -> str: + digest = hashlib.sha1( + (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii") + ).digest() + return base64.b64encode(digest).decode("ascii") + + +def read_exact(sock: ssl.SSLSocket, size: int) -> bytes: + chunks: list[bytes] = [] + remaining = size + while remaining: + chunk = sock.recv(remaining) + if not chunk: + raise ProbeError("socket closed while reading frame") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def read_http_head(sock: ssl.SSLSocket) -> tuple[int, dict[str, str], bytes]: + data = bytearray() + while b"\r\n\r\n" not in data: + if len(data) >= 64 * 1024: + raise ProbeError("HTTP upgrade headers exceed local limit") + chunk = sock.recv(4096) + if not chunk: + raise ProbeError("socket closed during HTTP upgrade") + data.extend(chunk) + raw_head, remainder = bytes(data).split(b"\r\n\r\n", 1) + lines = raw_head.decode("iso-8859-1").split("\r\n") + parts = lines[0].split(" ", 2) + if len(parts) < 2 or not parts[1].isdigit(): + raise ProbeError("malformed HTTP upgrade status") + headers: dict[str, str] = {} + for line in lines[1:]: + if not line or ":" not in line: + raise ProbeError("malformed HTTP upgrade header") + name, value = line.split(":", 1) + headers[name.strip().lower()] = value.strip() + return int(parts[1]), headers, remainder + + +class WebSocket: + def __init__(self, sock: ssl.SSLSocket, buffered: bytes = b"") -> None: + self.sock = sock + self.buffered = bytearray(buffered) + + def _read_exact(self, size: int) -> bytes: + if len(self.buffered) >= size: + data = bytes(self.buffered[:size]) + del self.buffered[:size] + return data + prefix = bytes(self.buffered) + self.buffered.clear() + return prefix + read_exact(self.sock, size - len(prefix)) + + def send_text(self, value: str) -> None: + self._send_frame(0x1, value.encode("utf-8")) + + def send_pong(self, payload: bytes) -> None: + self._send_frame(0xA, payload) + + def close(self) -> None: + try: + self._send_frame(0x8, b"\x03\xe8") + except OSError: + pass + + def _send_frame(self, opcode: int, payload: bytes) -> None: + if len(payload) > MAX_MESSAGE_BYTES: + raise ProbeError("outbound frame exceeds local limit") + mask = secrets.token_bytes(4) + header = bytearray([0x80 | opcode]) + if len(payload) < 126: + header.append(0x80 | len(payload)) + elif len(payload) <= 0xFFFF: + header.append(0x80 | 126) + header.extend(len(payload).to_bytes(2, "big")) + else: + header.append(0x80 | 127) + header.extend(len(payload).to_bytes(8, "big")) + masked = bytes(byte ^ mask[index % 4] for index, byte in enumerate(payload)) + self.sock.sendall(bytes(header) + mask + masked) + + def read_message(self) -> tuple[int, bytes]: + fragments: list[bytes] = [] + initial_opcode: int | None = None + while True: + first, second = self._read_exact(2) + fin = (first & 0x80) != 0 + opcode = first & 0x0F + masked = (second & 0x80) != 0 + length = second & 0x7F + if masked: + raise ProbeError("server sent a masked WebSocket frame") + if length == 126: + length = int.from_bytes(self._read_exact(2), "big") + elif length == 127: + length = int.from_bytes(self._read_exact(8), "big") + if length & (1 << 63): + raise ProbeError("invalid WebSocket frame length") + if length > MAX_FRAME_BYTES: + raise ProbeError("inbound frame exceeds local limit") + if opcode >= 0x8 and (not fin or length > 125): + raise ProbeError("invalid WebSocket control frame") + payload = self._read_exact(length) + if opcode == 0x9: + self.send_pong(payload) + continue + if opcode == 0xA: + continue + if opcode == 0x8: + return opcode, payload + if opcode == 0x0: + if initial_opcode is None: + raise ProbeError("unexpected continuation frame") + elif opcode in (0x1, 0x2): + if initial_opcode is not None: + raise ProbeError("new data frame before fragmented message completed") + initial_opcode = opcode + else: + raise ProbeError("unsupported WebSocket opcode") + fragments.append(payload) + if sum(len(fragment) for fragment in fragments) > MAX_MESSAGE_BYTES: + raise ProbeError("reassembled WebSocket message exceeds local limit") + if fin: + return initial_opcode or opcode, b"".join(fragments) + + +def connect(token: str, account_id: str, report: Report) -> WebSocket: + key = base64.b64encode(secrets.token_bytes(16)).decode("ascii") + context = ssl.create_default_context() + started = time.monotonic() + raw = socket.create_connection((HOST, 443), CONNECT_TIMEOUT_SECONDS) + sock = context.wrap_socket(raw, server_hostname=HOST) + sock.settimeout(EVENT_IDLE_TIMEOUT_SECONDS) + request = "\r\n".join( + [ + f"GET {PATH} HTTP/1.1", + f"Host: {HOST}", + "Connection: Upgrade", + "Upgrade: websocket", + "Sec-WebSocket-Version: 13", + f"Sec-WebSocket-Key: {key}", + f"Authorization: Bearer {token}", + f"chatgpt-account-id: {account_id}", + "originator: fx-phase-0-probe", + f"OpenAI-Beta: {PROTOCOL_HEADER}", + f"Origin: {ORIGIN}", + "\r\n", + ] + ).encode("ascii") + sock.sendall(request) + status, headers, remainder = read_http_head(sock) + report.handshake_elapsed_ms = round((time.monotonic() - started) * 1000) + report.handshake_status = status + report.selected_header_names = sorted( + name + for name in headers + if name in {"openai-model", "x-codex-turn-state", "x-reasoning-included", "x-models-etag"} + ) + if status != 101: + raise ProbeError(f"WebSocket upgrade returned HTTP {status}") + if headers.get("sec-websocket-accept") != websocket_accept(key): + raise ProbeError("invalid Sec-WebSocket-Accept response") + if "upgrade" not in headers.get("connection", "").lower(): + raise ProbeError("upgrade response does not retain Connection: Upgrade") + if headers.get("upgrade", "").lower() != "websocket": + raise ProbeError("upgrade response does not select websocket") + return WebSocket(sock, remainder) + + +def safe_protocol_label(value: Any) -> str | None: + if not isinstance(value, str) or len(value) > 128: + return None + if not value.isascii() or any(not (char.isalnum() or char in "._-") for char in value): + return None + return value + + +def event_type(payload: bytes) -> tuple[str | None, dict[str, Any] | None]: + try: + value = json.loads(payload) + except (UnicodeDecodeError, json.JSONDecodeError): + return None, None + if not isinstance(value, dict): + return None, None + return safe_protocol_label(value.get("type")), value + + +def response_id(value: dict[str, Any]) -> str | None: + response = value.get("response") + if not isinstance(response, dict): + return None + identifier = response.get("id") + return identifier if isinstance(identifier, str) else None + + +def structured_error_code(value: dict[str, Any]) -> str | None: + error = value.get("error") + if not isinstance(error, dict): + return None + return safe_protocol_label(error.get("code")) + + +def wait_for_terminal(ws: WebSocket, report: Report) -> str | None: + identifier: str | None = None + for _ in range(MAX_EVENTS): + opcode, payload = ws.read_message() + if opcode == 0x8: + report.close_code = int.from_bytes(payload[:2], "big") if len(payload) >= 2 else None + report.close_reason_length = max(len(payload) - 2, 0) + raise ProbeError("server closed before terminal response event") + if opcode != 0x1: + raise ProbeError("server sent an unexpected binary message") + kind, value = event_type(payload) + if kind is None or value is None: + report.event_types.append("invalid_json") + continue + report.event_types.append(kind) + if kind == "response.created": + identifier = response_id(value) + if kind in {"response.completed", "response.done", "response.incomplete", "response.failed", "error"}: + report.terminal_event = kind + report.error_code = structured_error_code(value) + return identifier + raise ProbeError("event count exceeds local limit before terminal response") + + +def request_body(model: str, previous_response_id: str | None = None) -> dict[str, Any]: + body: dict[str, Any] = { + "type": "response.create", + "model": model, + "input": [{"role": "user", "content": [{"type": "input_text", "text": "Reply with exactly: probe"}]}], + } + if previous_response_id is not None: + body["previous_response_id"] = previous_response_id + return body + + +def run(args: argparse.Namespace) -> Report: + report = Report() + ws: WebSocket | None = None + try: + token = required_env("FX_CODEX_PROBE_ACCESS_TOKEN") + account_id = required_env("FX_CODEX_PROBE_ACCOUNT_ID") + model = required_env("FX_CODEX_PROBE_MODEL") + ws = connect(token, account_id, report) + if not args.execute: + ws.sock.settimeout(0.2) + try: + opcode, payload = ws.read_message() + if opcode == 0x8: + report.immediate_close_code = int.from_bytes(payload[:2], "big") if len(payload) >= 2 else None + except (socket.timeout, ssl.SSLWantReadError): + pass + return report + ws.send_text(json.dumps(request_body(model), separators=(",", ":"))) + first_id = wait_for_terminal(ws, report) + if args.continuation: + report.continuation_attempted = True + if not first_id or report.terminal_event != "response.completed": + report.continuation_accepted = False + return report + report.event_types = [] + report.terminal_event = None + ws.send_text(json.dumps(request_body(model, first_id), separators=(",", ":"))) + wait_for_terminal(ws, report) + report.continuation_accepted = report.terminal_event == "response.completed" + return report + except (OSError, ssl.SSLError, ProbeError) as error: + report.error = type(error).__name__ + if report.handshake_status is None: + report.handshake_status = 0 + return report + finally: + if ws is not None: + ws.close() + ws.sock.close() + + +def self_test() -> None: + key = "dGhlIHNhbXBsZSBub25jZQ==" + assert websocket_accept(key) == "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" + body = request_body("model", "response") + assert body["previous_response_id"] == "response" + assert "stream" not in body and "background" not in body + code, value = event_type(b'{"type":"response.created","response":{"id":"r"}}') + assert code == "response.created" and response_id(value or {}) == "r" + assert safe_protocol_label("response.completed") == "response.completed" + assert safe_protocol_label("prompt content") is None + print("codex websocket probe self-test passed") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the isolated fx Codex WebSocket Phase 0 probe") + parser.add_argument("--execute", action="store_true", help="send a fixed minimal model request after the upgrade") + parser.add_argument("--continuation", action="store_true", help="test previous_response_id after --execute") + parser.add_argument("--self-test", action="store_true", help="run deterministic local checks without credentials or network") + args = parser.parse_args() + if args.continuation and not args.execute: + parser.error("--continuation requires --execute") + return args + + +if __name__ == "__main__": + arguments = parse_args() + if arguments.self_test: + self_test() + else: + result = run(arguments) + result.emit() + sys.exit(0 if result.error is None else 1) From 11bc046e8b284b0fea44822bf36251e76717cdd4 Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 00:15:44 +0530 Subject: [PATCH 09/21] Harden Codex WebSocket delivery handling --- src/gateway/openai_codex.zig | 114 +++++++++++++++++++++- src/gateway/websocket_transport.zig | 144 +++++++++++++++++++++++++++- 2 files changed, 251 insertions(+), 7 deletions(-) diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index aad2d5a98..6645cc473 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -428,9 +428,9 @@ fn streamWebSocketPrepared( .stream_limits = codexStreamLimits(.{}), }; // Admission is shared with SSE and happens before the upgrade can make - // delivery possible. Once frame writing begins, this transport never replays. + // delivery possible. The transport marks delivery only immediately before + // it begins writing the masked response.create frame. try admitCodexTransport(request.admission); - request.delivery.markPossiblySent(); try websocket_transport.stream(alloc, .{ .endpoint = prepared.endpoint, .authorization = prepared.authorization, @@ -439,12 +439,13 @@ fn streamWebSocketPrepared( .payload = payload, .deadline = request.deadline, .cancel_flag = request.cancel_flag, + .delivery = request.delivery, }, &bridge, WebSocketBridge.event); const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| return mapReducerError(err); return .{ .completed = .{ .completion = completion, - .usage = .{ .immediate = null }, + .usage = .{ .unavailable = .possibly_billed }, .ownership = .owned, } }; } @@ -945,6 +946,113 @@ test "OpenAI Codex SSE maps text reasoning tools and usage" { try std.testing.expectEqual(types.ProviderFinishReason.tool_calls, completion.finish_reason.?); } +test "OpenAI Codex SSE and WebSocket reducers preserve callback order and completion data" { + const raw_events = [_][]const u8{ + "{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\"}}", + "{\"type\":\"response.reasoning_summary_text.delta\",\"output_index\":0,\"delta\":\"thinking\"}", + "{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"rs_1\",\"type\":\"reasoning\",\"summary\":[],\"encrypted_content\":\"opaque\"}}", + "{\"type\":\"response.output_text.delta\",\"output_index\":1,\"delta\":\"hello\"}", + "{\"type\":\"response.output_item.added\",\"output_index\":2,\"item\":{\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"read_file\"}}", + "{\"type\":\"response.function_call_arguments.delta\",\"output_index\":2,\"delta\":\"{\\\"path\\\":\\\"README.md\\\"}\"}", + "{\"type\":\"response.completed\",\"response\":{\"id\":\"response_1\",\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":4}}}", + }; + const Capture = struct { + events: std.Io.Writer.Allocating = .init(std.testing.allocator), + failed: bool = false, + + fn deinit(self: *@This()) void { + self.events.deinit(); + } + + fn emit(raw: *anyopaque, event: stream_provider.Event) void { + const self: *@This() = @ptrCast(@alignCast(raw)); + const writer = &self.events.writer; + switch (event) { + .content_delta => |value| { + writer.print("content:{s}|", .{value}) catch { + self.failed = true; + }; + }, + .reasoning_delta => |value| { + writer.print("reasoning:{s}|", .{value}) catch { + self.failed = true; + }; + }, + .tool_started => |tool| { + writer.print("tool:{s}:{s}|", .{ tool.id, tool.name }) catch { + self.failed = true; + }; + }, + .tool_input_delta => |value| { + writer.print("input:{s}|", .{value}) catch { + self.failed = true; + }; + }, + } + } + }; + + var sse_body: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer sse_body.deinit(); + for (raw_events) |raw_event| try sse_body.writer.print("data: {s}\n\n", .{raw_event}); + var sse_capture: Capture = .{}; + defer sse_capture.deinit(); + var sse_events = stream_provider.EventSink{ .context = &sse_capture, .emit_fn = Capture.emit }; + var sse_reader: std.Io.Reader = .fixed(sse_body.written()); + var sse_cancelled = std.atomic.Value(bool).init(false); + const sse_completion = try consumeSse( + std.testing.allocator, + &sse_reader, + &sse_events, + EventBridge.content, + EventBridge.toolStart, + EventBridge.reasoning, + EventBridge.toolInput, + &sse_cancelled, + null, + .{}, + ); + defer freeOpenAICodexTestCompletion(sse_completion); + + var websocket_capture: Capture = .{}; + defer websocket_capture.deinit(); + const websocket_events = stream_provider.EventSink{ .context = &websocket_capture, .emit_fn = Capture.emit }; + var websocket_cancelled = std.atomic.Value(bool).init(false); + var websocket_reducer = responses_protocol.Reducer.init(std.testing.allocator); + defer websocket_reducer.deinit(std.testing.allocator); + var bridge = WebSocketBridge{ + .alloc = std.testing.allocator, + .reducer = &websocket_reducer, + .events = websocket_events, + .cancel_flag = &websocket_cancelled, + .content_capture_limit = null, + .stream_limits = codexStreamLimits(.{}), + }; + for (raw_events) |raw_event| { + if (try WebSocketBridge.event(&bridge, raw_event)) break; + } + const websocket_completion = try websocket_reducer.finish( + std.testing.allocator, + &websocket_cancelled, + bridge.stream_limits, + ); + defer freeOpenAICodexTestCompletion(websocket_completion); + + try std.testing.expect(!sse_capture.failed); + try std.testing.expect(!websocket_capture.failed); + try std.testing.expectEqualStrings(sse_capture.events.written(), websocket_capture.events.written()); + try std.testing.expectEqualStrings(sse_completion.content.?, websocket_completion.content.?); + try std.testing.expectEqualStrings(sse_completion.generation_id.?, websocket_completion.generation_id.?); + try std.testing.expectEqualStrings(sse_completion.provider_state_json.?, websocket_completion.provider_state_json.?); + try std.testing.expectEqual(sse_completion.usage.input_tokens, websocket_completion.usage.input_tokens); + try std.testing.expectEqual(sse_completion.usage.output_tokens, websocket_completion.usage.output_tokens); + try std.testing.expectEqual(sse_completion.finish_reason, websocket_completion.finish_reason); + try std.testing.expectEqual(@as(usize, 1), websocket_completion.tool_calls.len); + try std.testing.expectEqualStrings(sse_completion.tool_calls[0].id, websocket_completion.tool_calls[0].id); + try std.testing.expectEqualStrings(sse_completion.tool_calls[0].name, websocket_completion.tool_calls[0].name); + try std.testing.expectEqualStrings(sse_completion.tool_calls[0].arguments_json, websocket_completion.tool_calls[0].arguments_json); +} + fn consumeOpenAICodexTestSse(sse_text: []const u8, limits: CodexLimits) !types.ModelCompletion { var reader: std.Io.Reader = .fixed(sse_text); var cancelled = std.atomic.Value(bool).init(false); diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index b81718ad5..e18b79753 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -31,8 +31,8 @@ pub const Request = struct { payload: []const u8, deadline: ?std.Io.Clock.Timestamp, cancel_flag: *std.atomic.Value(bool), + delivery: *gateway_client.DeliveryCertainty, }; - const OpenedRequest = struct { request: ?std.http.Client.Request, @@ -174,17 +174,21 @@ pub fn stream( const connection = http_request.connection orelse return error.WebSocketConnectionMissing; const writer = connection.writer(); var close_sent = false; - defer if (!close_sent) { - // The watcher bounds this write and forces release if the peer does not - // cooperate. A fresh Phase 1 socket is never returned to the HTTP pool. + defer if (!close_sent and !request.cancel_flag.load(.seq_cst)) { + // Cancellation force-releases the socket from the watcher. Do not race + // that release with a best-effort close frame from this I/O owner. + // A fresh Phase 1 socket is never returned to the HTTP pool. writeFrame(writer, .close, &.{ 0x03, 0xe8 }) catch {}; connection.flush() catch {}; }; + request.delivery.markPossiblySent(); writeFrame(writer, .text, request.payload) catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; connection.flush() catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; @@ -468,3 +472,135 @@ test "close payload rejects reserved codes and malformed UTF-8 reasons" { try std.testing.expectEqual(error.WebSocketPolicyClosed, closeError(1008)); try std.testing.expectEqual(error.WebSocketClosedBeforeCompletion, closeError(1000)); } + +const StalledWriteFixture = struct { + io_backend: std.Io.Threaded = .init_single_threaded, + server: std.Io.net.Server, + thread: ?std.Thread = null, + stopping: std.atomic.Value(bool) = .init(false), + upgraded: std.atomic.Value(bool) = .init(false), + failure: ?anyerror = null, + + fn init() !@This() { + var fixture: @This() = .{ .server = undefined }; + var address = try std.Io.net.IpAddress.parse("127.0.0.1", 0); + fixture.server = try address.listen(fixture.io(), .{ .reuse_address = true }); + return fixture; + } + + fn io(self: *@This()) std.Io { + return self.io_backend.io(); + } + + fn endpoint(self: *@This(), buffer: []u8) ![]const u8 { + return std.fmt.bufPrint(buffer, "http://127.0.0.1:{d}/responses", .{self.server.socket.address.getPort()}); + } + + fn start(self: *@This()) !void { + self.thread = try std.Thread.spawn(.{}, run, .{self}); + } + + fn deinit(self: *@This()) void { + self.stopping.store(true, .seq_cst); + if (self.thread) |thread| { + const listener = std.Io.net.Stream{ .socket = self.server.socket }; + listener.shutdown(self.io(), .both) catch {}; + thread.join(); + self.thread = null; + } + self.server.deinit(self.io()); + } + + fn run(self: *@This()) void { + self.runFallible() catch |err| { + if (!self.stopping.load(.seq_cst)) self.failure = err; + }; + } + + fn runFallible(self: *@This()) !void { + const zio = self.io(); + var client_stream = try self.server.accept(zio); + defer client_stream.close(zio); + if (self.stopping.load(.seq_cst)) return; + const receive_buffer: c_int = 1024; + std.posix.setsockopt(client_stream.socket.handle, std.posix.SOL.SOCKET, std.posix.SO.RCVBUF, std.mem.asBytes(&receive_buffer)) catch {}; + + var socket_buffer: [4096]u8 = undefined; + var reader = client_stream.reader(zio, &socket_buffer); + var request: [16 * 1024]u8 = undefined; + var request_len: usize = 0; + while (request_len < request.len) { + request[request_len] = try reader.interface.takeByte(); + request_len += 1; + if (std.mem.endsWith(u8, request[0..request_len], "\r\n\r\n")) break; + } else return error.TestRequestTooLarge; + const key = headerValue(request[0 .. request_len - 4], "sec-websocket-key") orelse return error.TestMissingWebSocketKey; + var accept_buffer: [28]u8 = undefined; + const accept = websocketAccept(key, &accept_buffer); + var write_buffer: [4096]u8 = undefined; + var writer = client_stream.writer(zio, &write_buffer); + try writer.interface.print( + "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: {s}\r\n\r\n", + .{accept}, + ); + try writer.interface.flush(); + self.upgraded.store(true, .seq_cst); + while (!self.stopping.load(.seq_cst)) { + var sleep_io: std.Io.Threaded = .init_single_threaded; + sleep_io.io().sleep(.fromMilliseconds(1), .real) catch {}; + } + } +}; + +fn headerValue(headers: []const u8, name: []const u8) ?[]const u8 { + var lines = std.mem.splitSequence(u8, headers, "\r\n"); + _ = lines.next(); + while (lines.next()) |line| { + const colon = std.mem.findScalar(u8, line, ':') orelse continue; + if (std.ascii.eqlIgnoreCase(std.mem.trim(u8, line[0..colon], " \t"), name)) { + return std.mem.trim(u8, line[colon + 1 ..], " \t"); + } + } + return null; +} + +test "WebSocket cancellation interrupts a backpressured response.create write" { + var fixture = try StalledWriteFixture.init(); + defer fixture.deinit(); + try fixture.start(); + + var endpoint_buffer: [128]u8 = undefined; + const payload = try std.testing.allocator.alloc(u8, 4 * 1024 * 1024); + defer std.testing.allocator.free(payload); + @memset(payload, 'x'); + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const Canceller = struct { + fn run(server: *StalledWriteFixture, flag: *std.atomic.Value(bool)) void { + while (!server.upgraded.load(.seq_cst)) { + var sleep_io: std.Io.Threaded = .init_single_threaded; + sleep_io.io().sleep(.fromMilliseconds(1), .real) catch {}; + } + flag.store(true, .seq_cst); + } + }; + const canceller = try std.Thread.spawn(.{}, Canceller.run, .{ &fixture, &cancelled }); + defer canceller.join(); + const result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = payload, + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + try std.testing.expectError(error.Cancelled, result); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.possibly_sent, delivery.load()); + if (fixture.failure) |err| return err; +} From 05d1ba0a4f2077fb03298f2376be09763ba41c57 Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 00:15:48 +0530 Subject: [PATCH 10/21] Cover Codex WebSocket cancellation paths --- tests/e2e/tui-auth-source-selection.test.ts | 59 ++++++++++++++++++++- 1 file changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index f6a73662f..27110e969 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -837,9 +837,14 @@ function startFakeCodexToolLoop(options: { }; } -function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: number } = {}) { +function startFakeCodexWebSocket(options: { + holdOpen?: boolean; + closeOnOpen?: number; + stallUpgrade?: boolean; +} = {}) { const requests: string[] = []; const closeCodes: number[] = []; + let upgradeRequests = 0; const accessToken = chatgptAccessToken("acct_websocket"); const server = Bun.serve<{ opened: boolean }>({ hostname: "127.0.0.1", @@ -851,7 +856,11 @@ function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: nu { slug: "gpt-5.6-sol", visibility: "list", supported_in_api: true, supported_reasoning_levels: [{ effort: "high" }], additional_speed_tiers: [], input_modalities: ["text"], context_window: 272000 }, ] }); } - if (path === "/responses" && server.upgrade(request, { data: { opened: true } })) return; + if (path === "/responses") { + upgradeRequests += 1; + if (options.stallUpgrade) return new Promise(() => {}); + if (server.upgrade(request, { data: { opened: true } })) return; + } return new Response("not found", { status: 404 }); }, websocket: { @@ -877,12 +886,14 @@ function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: nu accessToken, requests, closeCodes, + get upgradeRequests() { return upgradeRequests; }, responsesUrl: `http://127.0.0.1:${server.port}/responses`, modelsUrl: `http://127.0.0.1:${server.port}/models`, stop() { server.stop(true); }, }; } + function startFakeCodexCapacityLoop() { const bodies: string[] = []; const accessToken = chatgptAccessToken("acct_capacity_loop"); @@ -2887,6 +2898,7 @@ test( expect(result.code).toBe(1); expect(`${result.stdout}\n${result.stderr}`).toContain("WebSocketPolicyClosed"); expect(codex.requests).toHaveLength(0); + expect(codex.upgradeRequests).toBe(1); expect(gateway.requests).toHaveLength(0); } finally { codex.stop(); @@ -2895,6 +2907,48 @@ test( 60_000, ); +tmuxTest( + "Codex WebSocket cancellation unblocks a stalled upgrade", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-upgrade-cancel-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ stallUpgrade: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Cancel a stalled WebSocket upgrade."); + const upgradeDeadline = Date.now() + TIMEOUT; + while (codex.upgradeRequests === 0) { + if (Date.now() >= upgradeDeadline) throw new Error("Codex WebSocket upgrade did not arrive"); + await Bun.sleep(25); + } + await session.sendKeys("C-c"); + await session.waitForComposer(TIMEOUT); + expect(session.isAlive()).toBe(true); + expect(codex.upgradeRequests).toBe(1); + expect(codex.requests).toHaveLength(0); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + + tmuxTest( "Codex WebSocket cancellation unblocks an idle response", async () => { @@ -2927,6 +2981,7 @@ tmuxTest( await session.waitForComposer(TIMEOUT); expect(session.isAlive()).toBe(true); expect(codex.requests).toHaveLength(1); + expect(codex.upgradeRequests).toBe(1); expect(readFileSync(stderrPath, "utf8")).toBe(""); } finally { codex.stop(); From 4e103201e1e1dd0c285c5eb8ae6c1c63b5ba915f Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 00:15:53 +0530 Subject: [PATCH 11/21] Record Codex WebSocket Phase One evidence --- docs/codex-websocket-transport.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index 462a255eb..b8a2715a3 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -262,9 +262,9 @@ This review applies to the Phase 1 implementation. It is not production-ready an - [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target, including a 127-byte extended frame followed by another frame. Close payloads and malformed text are also validated at the frame boundary. - [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. -- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented. A tmux loopback test proves cancellation unblocks an idle frame read; connect- and write-block cancellation remain untested. +- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented. Tmux loopback tests prove cancellation unblocks both a stalled upgrade and an idle frame read; connect- and write-block cancellation remain untested. - [~] Bounded close handshake and close-code reporting: normal completion sends and receives a close frame in the loopback smoke test; malformed close payloads are unit-tested, and an early policy close is classified without retry. Cancellation-close sequencing remains untested. -- [~] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for a completion, early policy close, and idle cancellation. Full SSE/WebSocket reducer-parity and delivery-certainty matrix coverage remain. +- [~] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for a completion, early policy close, stalled-upgrade cancellation, and idle cancellation. It proves policy-close and post-send cancellation each create exactly one connection, so they cannot silently replay a turn. Delivery certainty remains definitely unsent until immediately before the masked `response.create` frame write. A unit test proves the SSE and WebSocket reducers preserve callback order and completion data for the same Responses event sequence. Full delivery-certainty matrix coverage remains. - [x] Authenticated Phase 0 evidence: an explicit maintainer probe on 2026-08-27 returned `101`, completed a request, and accepted a same-socket `previous_response_id` continuation. The redacted report recorded only protocol labels, a 1,749 ms handshake, and the `x-models-etag` header name. ### Release blockers From 6fef080e63a8e542633df0afcba6cdea6f10d135 Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 00:30:12 +0530 Subject: [PATCH 12/21] Prevent Codex WebSocket replay after delivery --- tests/e2e/tui-auth-source-selection.test.ts | 49 +++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index 27110e969..d8c4cf818 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -840,6 +840,7 @@ function startFakeCodexToolLoop(options: { function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: number; + closeAfterMessage?: number; stallUpgrade?: boolean; } = {}) { const requests: string[] = []; @@ -870,6 +871,10 @@ function startFakeCodexWebSocket(options: { message(ws, message) { const payload = String(message); requests.push(payload); + if (options.closeAfterMessage !== undefined) { + ws.close(options.closeAfterMessage, "fixture close"); + return; + } if (options.holdOpen) return; ws.send(JSON.stringify({ type: "response.output_text.delta", delta: "CODEX_WEBSOCKET_OK" })); ws.send(JSON.stringify({ @@ -2907,6 +2912,50 @@ test( 60_000, ); +test( + "Codex WebSocket close after response.create never replays the turn", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-post-send-close-")); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ closeAfterMessage: 1011 }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + const result = await runFx( + ["ask", "--json", "--auto", "--no-save", "Do not replay this request."], + { + env: { + HOME: home, + AI_GATEWAY_API_KEY: "gateway-websocket-post-send-close-sentinel", + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_AUTO_UPGRADE: "0", + FX_CODEX_TRANSPORT: "websocket", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }, + timeoutMs: TIMEOUT, + }, + ); + expect(result.code).toBe(1); + expect(`${result.stdout}\n${result.stderr}`).toContain("WebSocketClosedBeforeCompletion"); + expect(codex.requests).toHaveLength(1); + expect(codex.upgradeRequests).toBe(1); + expect(gateway.requests).toHaveLength(0); + } finally { + codex.stop(); + } + }, + 60_000, +); + + tmuxTest( "Codex WebSocket cancellation unblocks a stalled upgrade", async () => { From f0272b757b2d94a8cbfd20f6ec80bb424fb0a945 Mon Sep 17 00:00:00 2001 From: thinkter Date: Fri, 28 Aug 2026 21:58:13 +0530 Subject: [PATCH 13/21] Harden Codex WebSocket transport --- src/gateway/websocket_transport.zig | 415 +++++++++++++++++++- tests/e2e/tui-auth-source-selection.test.ts | 58 +++ 2 files changed, 470 insertions(+), 3 deletions(-) diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index e18b79753..22c991fb4 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -154,10 +154,12 @@ pub fn stream( if (watcher) |thread| thread.join(); } http_request.sendBodiless() catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; const response = http_request.receiveHead(&.{}) catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; @@ -200,6 +202,7 @@ pub fn stream( while (true) { if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; const frame = readFrame(alloc, reader) catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; @@ -208,10 +211,12 @@ pub fn stream( switch (frame.opcode) { .ping => { writeFrame(writer, .pong, frame.payload) catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; connection.flush() catch |err| { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; if (timeout_fired.load(.seq_cst)) return error.Timeout; return err; }; @@ -229,7 +234,10 @@ pub fn stream( fragmented_opcode = null; if (opcode != .text) return error.WebSocketUnexpectedBinary; if (try dispatchTextMessage(context, on_event, message.items)) { - try closeAfterCompletion(alloc, reader, writer, connection); + closeAfterCompletion(alloc, reader, writer, connection, request.cancel_flag, &timeout_fired) catch |err| { + if (err == error.Cancelled or err == error.Timeout) close_sent = true; + return err; + }; close_sent = true; return; } @@ -243,7 +251,10 @@ pub fn stream( continue; } if (try dispatchTextMessage(context, on_event, message.items)) { - try closeAfterCompletion(alloc, reader, writer, connection); + closeAfterCompletion(alloc, reader, writer, connection, request.cancel_flag, &timeout_fired) catch |err| { + if (err == error.Cancelled or err == error.Timeout) close_sent = true; + return err; + }; close_sent = true; return; } @@ -299,10 +310,16 @@ fn closeAfterCompletion( reader: *std.Io.Reader, writer: *std.Io.Writer, connection: anytype, + cancel_flag: *std.atomic.Value(bool), + timeout_fired: *std.atomic.Value(bool), ) !void { try writeFrame(writer, .close, &.{ 0x03, 0xe8 }); try connection.flush(); - const frame = try readFrame(alloc, reader); + const frame = readFrame(alloc, reader) catch |err| { + if (cancel_flag.load(.seq_cst)) return error.Cancelled; + if (timeout_fired.load(.seq_cst)) return error.Timeout; + return err; + }; defer alloc.free(frame.payload); switch (frame.opcode) { .close => _ = try validateClosePayload(frame.payload), @@ -551,6 +568,167 @@ const StalledWriteFixture = struct { } } }; +const LoopbackMode = enum { + never_accept, + hang_after_upgrade, + reset_after_upgrade, + complete_then_hang_close, + binary_then_close, + ping_then_complete, + oversized_frame, +}; + +const LoopbackWebSocketFixture = struct { + io_backend: std.Io.Threaded = .init_single_threaded, + server: std.Io.net.Server, + mode: LoopbackMode, + thread: ?std.Thread = null, + stopping: std.atomic.Value(bool) = .init(false), + upgraded: std.atomic.Value(bool) = .init(false), + failure: ?anyerror = null, + + fn init(mode: LoopbackMode) !@This() { + var fixture: @This() = .{ .server = undefined, .mode = mode }; + var address = try std.Io.net.IpAddress.parse("127.0.0.1", 0); + fixture.server = try address.listen(fixture.io(), .{ .reuse_address = true }); + return fixture; + } + + fn io(self: *@This()) std.Io { + return self.io_backend.io(); + } + + fn endpoint(self: *@This(), buffer: []u8) ![]const u8 { + return std.fmt.bufPrint(buffer, "http://127.0.0.1:{d}/responses", .{self.server.socket.address.getPort()}); + } + + fn start(self: *@This()) !void { + self.thread = try std.Thread.spawn(.{}, run, .{self}); + } + + fn deinit(self: *@This()) void { + self.stopping.store(true, .seq_cst); + if (self.thread) |thread| { + const listener = std.Io.net.Stream{ .socket = self.server.socket }; + listener.shutdown(self.io(), .both) catch {}; + thread.join(); + self.thread = null; + } + self.server.deinit(self.io()); + } + + fn hold(self: *@This()) void { + while (!self.stopping.load(.seq_cst)) { + self.io().sleep(.fromMilliseconds(1), .real) catch {}; + } + } + + fn run(self: *@This()) void { + self.runFallible() catch |err| { + if (!self.stopping.load(.seq_cst)) self.failure = err; + }; + } + + fn runFallible(self: *@This()) !void { + if (self.mode == .never_accept) return self.hold(); + const zio = self.io(); + var client_stream = try self.server.accept(zio); + defer client_stream.close(zio); + if (self.stopping.load(.seq_cst)) return; + + var socket_buffer: [4096]u8 = undefined; + var reader = client_stream.reader(zio, &socket_buffer); + var request: [16 * 1024]u8 = undefined; + var request_len: usize = 0; + while (request_len < request.len) { + request[request_len] = try reader.interface.takeByte(); + request_len += 1; + if (std.mem.endsWith(u8, request[0..request_len], "\r\n\r\n")) break; + } else return error.TestRequestTooLarge; + const key = headerValue(request[0 .. request_len - 4], "sec-websocket-key") orelse return error.TestMissingWebSocketKey; + var accept_buffer: [28]u8 = undefined; + const accept = websocketAccept(key, &accept_buffer); + var write_buffer: [4096]u8 = undefined; + var writer = client_stream.writer(zio, &write_buffer); + try writer.interface.print( + "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: {s}\r\n\r\n", + .{accept}, + ); + try writer.interface.flush(); + self.upgraded.store(true, .seq_cst); + + if (self.mode == .reset_after_upgrade) { + const reset_on_close: std.posix.linger = .{ .onoff = 1, .linger = 0 }; + try std.posix.setsockopt( + client_stream.socket.handle, + std.posix.SOL.SOCKET, + std.posix.SO.LINGER, + std.mem.asBytes(&reset_on_close), + ); + return; + } + if (self.mode == .hang_after_upgrade) return self.hold(); + + try discardClientFrame(&reader.interface); + switch (self.mode) { + .complete_then_hang_close => { + try writeServerFrame(&writer.interface, .text, "{\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}"); + try writeServerFrame(&writer.interface, .text, "{\"type\":\"response.completed\",\"response\":{\"id\":\"r1\",\"status\":\"completed\"}}"); + try writer.interface.flush(); + self.hold(); + }, + .binary_then_close => { + try writeServerFrame(&writer.interface, .binary, &.{0}); + try writer.interface.flush(); + }, + .ping_then_complete => { + try writeServerFrame(&writer.interface, .ping, "hi"); + try writeServerFrame(&writer.interface, .text, "{\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}"); + try writeServerFrame(&writer.interface, .text, "{\"type\":\"response.completed\",\"response\":{\"id\":\"r1\",\"status\":\"completed\"}}"); + try writeServerFrame(&writer.interface, .close, &.{ 0x03, 0xe8 }); + try writer.interface.flush(); + try discardClientFrame(&reader.interface); + }, + .oversized_frame => { + try writer.interface.writeAll(&.{ 0x81, 127 }); + try writer.interface.writeInt(u64, max_frame_bytes + 1, .big); + try writer.interface.flush(); + }, + else => unreachable, + } + } +}; + +fn writeServerFrame(writer: *std.Io.Writer, opcode: Opcode, payload: []const u8) !void { + try writer.writeByte(0x80 | @as(u8, @intFromEnum(opcode))); + if (payload.len < 126) { + try writer.writeByte(@intCast(payload.len)); + } else if (payload.len <= std.math.maxInt(u16)) { + try writer.writeByte(126); + try writer.writeInt(u16, @intCast(payload.len), .big); + } else { + try writer.writeByte(127); + try writer.writeInt(u64, @intCast(payload.len), .big); + } + try writer.writeAll(payload); +} + +fn discardClientFrame(reader: *std.Io.Reader) !void { + _ = try reader.takeByte(); + const second = try reader.takeByte(); + if (second & 0x80 == 0) return error.WebSocketProtocolViolation; + var length: u64 = second & 0x7f; + if (length == 126) length = try reader.takeInt(u16, .big) else if (length == 127) length = try reader.takeInt(u64, .big); + var mask: [4]u8 = undefined; + try reader.readSliceAll(&mask); + var discarded: [4096]u8 = undefined; + var remaining = length; + while (remaining > 0) { + const chunk_len: usize = @intCast(@min(remaining, discarded.len)); + try reader.readSliceAll(discarded[0..chunk_len]); + remaining -= chunk_len; + } +} fn headerValue(headers: []const u8, name: []const u8) ?[]const u8 { var lines = std.mem.splitSequence(u8, headers, "\r\n"); @@ -604,3 +782,234 @@ test "WebSocket cancellation interrupts a backpressured response.create write" { try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.possibly_sent, delivery.load()); if (fixture.failure) |err| return err; } + +test "WebSocket cancellation interrupts a stalled connect" { + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const Canceller = struct { + fn run(flag: *std.atomic.Value(bool)) void { + io_mod.sleep(20 * std.time.ns_per_ms); + flag.store(true, .seq_cst); + } + }; + const canceller = try std.Thread.spawn(.{}, Canceller.run, .{&cancelled}); + const started = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + const result = stream(std.testing.allocator, .{ + .endpoint = "http://192.0.2.1:9/responses", + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + const elapsed_ms = started.durationTo(std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake)).raw.toMilliseconds(); + canceller.join(); + if (result) |_| { + return error.TestExpectedError; + } else |err| { + if (err == error.Cancelled) { + try std.testing.expect(elapsed_ms < 2_000); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.definitely_unsent, delivery.load()); + return; + } + if (elapsed_ms >= 5) return err; + } + + var fixture = try LoopbackWebSocketFixture.init(.never_accept); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + cancelled.store(false, .seq_cst); + delivery = gateway_client.DeliveryCertainty.init(); + const fallback_canceller = try std.Thread.spawn(.{}, Canceller.run, .{&cancelled}); + const fallback_started = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + const fallback_result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + const fallback_elapsed_ms = fallback_started.durationTo(std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake)).raw.toMilliseconds(); + fallback_canceller.join(); + try std.testing.expectError(error.Cancelled, fallback_result); + try std.testing.expect(fallback_elapsed_ms < 2_000); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.definitely_unsent, delivery.load()); + if (fixture.failure) |err| return err; +} + +test "WebSocket cancellation interrupts a hung close handshake" { + var fixture = try LoopbackWebSocketFixture.init(.complete_then_hang_close); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const Canceller = struct { + fn run(server: *LoopbackWebSocketFixture, flag: *std.atomic.Value(bool)) void { + while (!server.upgraded.load(.seq_cst)) io_mod.sleep(std.time.ns_per_ms); + io_mod.sleep(50 * std.time.ns_per_ms); + flag.store(true, .seq_cst); + } + }; + const canceller = try std.Thread.spawn(.{}, Canceller.run, .{ &fixture, &cancelled }); + defer canceller.join(); + const result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn completed(_: *anyopaque, json: []const u8) !bool { + return std.mem.find(u8, json, "\"response.completed\"") != null; + } + }.completed); + try std.testing.expectError(error.Cancelled, result); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.possibly_sent, delivery.load()); + if (fixture.failure) |err| return err; +} + +test "WebSocket peer reset after upgrade leaves delivery possibly sent" { + var fixture = try LoopbackWebSocketFixture.init(.reset_after_upgrade); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + const payload = try std.testing.allocator.alloc(u8, 4 * 1024); + defer std.testing.allocator.free(payload); + @memset(payload, 'x'); + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = payload, + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + if (result) |_| return error.TestExpectedError else |err| try std.testing.expect(err != error.Cancelled); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.possibly_sent, delivery.load()); + if (fixture.failure) |err| return err; +} + +test "WebSocket rejects unexpected binary frames" { + var fixture = try LoopbackWebSocketFixture.init(.binary_then_close); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + try std.testing.expectError(error.WebSocketUnexpectedBinary, result); + try std.testing.expectEqual(gateway_client.DeliveryCertainty.State.possibly_sent, delivery.load()); + if (fixture.failure) |err| return err; +} + +test "WebSocket answers ping then completes" { + var fixture = try LoopbackWebSocketFixture.init(.ping_then_complete); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + var completed = false; + try stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&completed), struct { + fn handle(context: *anyopaque, json: []const u8) !bool { + const done: *bool = @ptrCast(@alignCast(context)); + if (std.mem.find(u8, json, "\"response.completed\"") != null) { + done.* = true; + return true; + } + return false; + } + }.handle); + try std.testing.expect(completed); + if (fixture.failure) |err| return err; +} + +test "WebSocket rejects inbound frames over max_frame_bytes" { + var fixture = try LoopbackWebSocketFixture.init(.oversized_frame); + defer fixture.deinit(); + try fixture.start(); + var endpoint_buffer: [128]u8 = undefined; + var cancelled = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const result = stream(std.testing.allocator, .{ + .endpoint = try fixture.endpoint(&endpoint_buffer), + .authorization = "Bearer test", + .account_id = "test", + .session_id = null, + .payload = "{}", + .deadline = null, + .cancel_flag = &cancelled, + .delivery = &delivery, + }, @ptrCast(&cancelled), struct { + fn ignore(_: *anyopaque, _: []const u8) !bool { + return false; + } + }.ignore); + try std.testing.expectError(error.WebSocketMessageTooLarge, result); + if (fixture.failure) |err| return err; +} + +test "appendMessage rejects a 64 MiB overflow" { + var message: std.ArrayList(u8) = .empty; + defer message.deinit(std.testing.allocator); + const payload = try std.testing.allocator.alloc(u8, max_message_bytes); + defer std.testing.allocator.free(payload); + try appendMessage(&message, std.testing.allocator, payload); + try std.testing.expectError(error.WebSocketMessageTooLarge, appendMessage(&message, std.testing.allocator, "x")); +} + +test "writeFrame rejects payloads over max_message_bytes" { + var encoded: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer encoded.deinit(); + const payload = try std.testing.allocator.alloc(u8, max_message_bytes + 1); + defer std.testing.allocator.free(payload); + try std.testing.expectError(error.WebSocketMessageTooLarge, writeFrame(&encoded.writer, .text, payload)); +} diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index d8c4cf818..86804197a 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -841,6 +841,7 @@ function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: number; closeAfterMessage?: number; + toolThenClose?: boolean; stallUpgrade?: boolean; } = {}) { const requests: string[] = []; @@ -871,6 +872,20 @@ function startFakeCodexWebSocket(options: { message(ws, message) { const payload = String(message); requests.push(payload); + if (options.toolThenClose) { + ws.send(JSON.stringify({ + type: "response.output_item.added", + output_index: 0, + item: { type: "function_call", call_id: "call_1", name: "read_file" }, + })); + ws.send(JSON.stringify({ + type: "response.function_call_arguments.delta", + output_index: 0, + delta: '{"path":"README.md"}', + })); + ws.close(1011, "fixture close"); + return; + } if (options.closeAfterMessage !== undefined) { ws.close(options.closeAfterMessage, "fixture close"); return; @@ -2956,6 +2971,49 @@ test( ); +test( + "Codex WebSocket tool stream close never replays the turn", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-tool-close-")); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ toolThenClose: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + const result = await runFx( + ["ask", "--json", "--auto", "--no-save", "Do not replay this tool request."], + { + env: { + HOME: home, + AI_GATEWAY_API_KEY: "gateway-websocket-tool-close-sentinel", + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_AUTO_UPGRADE: "0", + FX_CODEX_TRANSPORT: "websocket", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }, + timeoutMs: TIMEOUT, + }, + ); + expect(result.code).toBe(1); + expect(`${result.stdout}\n${result.stderr}`).toContain("WebSocketClosedBeforeCompletion"); + expect(codex.requests).toHaveLength(1); + expect(codex.upgradeRequests).toBe(1); + expect(gateway.requests).toHaveLength(0); + } finally { + codex.stop(); + } + }, + 60_000, +); + tmuxTest( "Codex WebSocket cancellation unblocks a stalled upgrade", async () => { From 0576b276497780ceb8ee36c0126997e34ae0e6a5 Mon Sep 17 00:00:00 2001 From: thinkter Date: Sat, 29 Aug 2026 13:55:52 +0530 Subject: [PATCH 14/21] Retain Codex WebSocket sessions --- docs/codex-websocket-transport.md | 26 +- src/gateway/codex_websocket_session.zig | 260 +++++++++++++++ src/gateway/openai_codex.zig | 88 +++-- src/gateway/websocket_transport.zig | 346 ++++++++++++++------ src/main.zig | 13 +- tests/e2e/tui-auth-source-selection.test.ts | 170 +++++++++- 6 files changed, 757 insertions(+), 146 deletions(-) create mode 100644 src/gateway/codex_websocket_session.zig diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md index b8a2715a3..6bd40528b 100644 --- a/docs/codex-websocket-transport.md +++ b/docs/codex-websocket-transport.md @@ -137,18 +137,16 @@ Implement a fresh socket per model request, with no retained connection or cache ### Phase 2: retained socket, full context per turn -After Phase 1 has stable production evidence: +- [x] Retain one idle native WebSocket per session ID, account, model, and endpoint. A slot serializes one stream at a time and pings before reuse. +- [x] Send a complete `response.create` body for every turn. Phase 2 does not send `previous_response_id` or retain continuation state. +- [x] Keep `FX_CODEX_TRANSPORT=websocket` opt-in only. SSE remains the default, `auto` remains SSE, and no SSE fallback occurs from the required WebSocket path. +- [x] Poison and close sockets after cancellation, timeout, protocol or binary-frame failures, writes, policy close, unhealthy ping, identity changes, and age eviction. +- [x] Evict idle connections at 55 minutes by default, with `FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS=0` disabling age eviction. +- [x] Reset the health budget only after a completed response, never after an upgrade alone. -- Retain at most one idle connection per active fx session, provider, account identity, and compatible model. -- Serialize one full request at a time over it. -- Preconnect only as an optimization. Never send prompt data during preconnect. -- Discard and recreate the connection after a close, protocol error, cancellation during a stream, timeout, failed write, authentication transition, or model incompatibility. -- Continue sending full context after reconnect. -- Proactively reconnect before a confirmed connection-age limit, with enough margin that the limit cannot interrupt an in-flight turn. -- Reset a connection-health or fallback budget only after confirmed `response.completed`, not merely after a successful handshake. -- Segment telemetry by authentication mode from the start. +Connect and event-idle timeouts default to 30 seconds and may be overridden with `FX_CODEX_WEBSOCKET_CONNECT_TIMEOUT_MS` and `FX_CODEX_WEBSOCKET_EVENT_IDLE_TIMEOUT_MS`. Values must be positive integer milliseconds. -Connect and event-idle timeouts must be configurable and measured per platform. Do not assume Linux and macOS behavior predicts Windows behavior. +Local Phase 2 verification exercises the freshly built binary against a loopback WebSocket service. It covers two full-context turns on one socket, absence of `previous_response_id`, pre-delivery reconnect after a peer close, poisoning and recovery after an in-flight failure, maximum-age eviction, policy and post-delivery close handling, and cancellation during upgrade and idle reads. The retained identity includes a SHA-256 fingerprint of the authorization value so credential rotation invalidates a socket without retaining the credential itself. ### Phase 3: incremental continuation @@ -262,9 +260,9 @@ This review applies to the Phase 1 implementation. It is not production-ready an - [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target, including a 127-byte extended frame followed by another frame. Close payloads and malformed text are also validated at the frame boundary. - [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. -- [~] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: bounded upgrade plus a socket-shutdown watcher are implemented. Tmux loopback tests prove cancellation unblocks both a stalled upgrade and an idle frame read; connect- and write-block cancellation remain untested. -- [~] Bounded close handshake and close-code reporting: normal completion sends and receives a close frame in the loopback smoke test; malformed close payloads are unit-tested, and an early policy close is classified without retry. Cancellation-close sequencing remains untested. -- [~] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for a completion, early policy close, stalled-upgrade cancellation, and idle cancellation. It proves policy-close and post-send cancellation each create exactly one connection, so they cannot silently replay a turn. Delivery certainty remains definitely unsent until immediately before the masked `response.create` frame write. A unit test proves the SSE and WebSocket reducers preserve callback order and completion data for the same Responses event sequence. Full delivery-certainty matrix coverage remains. +- [x] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: unit and tmux loopback tests cover stalled connect, backpressured write, stalled upgrade, idle frame read, and a hung close handshake. +- [x] Bounded close handshake and close-code reporting: malformed close payloads are unit-tested, an early policy close is classified without retry, and shutdown uses independent bounded teardown state rather than retaining pointers into a completed turn. +- [x] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for completion, close failures, retained reuse, reconnect, poisoning recovery, age eviction, and cancellation. Delivery certainty remains definitely unsent until immediately before the masked `response.create` frame write. A unit test proves the SSE and WebSocket reducers preserve callback order and completion data for the same Responses event sequence. - [x] Authenticated Phase 0 evidence: an explicit maintainer probe on 2026-08-27 returned `101`, completed a request, and accepted a same-socket `previous_response_id` continuation. The redacted report recorded only protocol labels, a 1,749 ms handshake, and the `x-models-etag` header name. ### Release blockers @@ -328,4 +326,4 @@ This review applies to the Phase 1 implementation. It is not production-ready an ## Current status -fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. A live authenticated probe confirms the Phase 0 handshake and continuation assumptions, but the experiment remains intentionally blocked from broad use until the remaining local coverage is complete. +fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. Phase 2 retains a healthy socket for later turns in the same process while continuing to send full request context. It has no continuation protocol, no `previous_response_id`, and no WebSocket-to-SSE fallback. Local Phase 2 verification is green; exact-commit Full CI remains the release-readiness authority. diff --git a/src/gateway/codex_websocket_session.zig b/src/gateway/codex_websocket_session.zig new file mode 100644 index 000000000..5094d2473 --- /dev/null +++ b/src/gateway/codex_websocket_session.zig @@ -0,0 +1,260 @@ +const std = @import("std"); +const io_mod = @import("../core/shared/io.zig"); +const gateway_client = @import("client.zig"); +const websocket_transport = @import("websocket_transport.zig"); + +const Allocator = std.mem.Allocator; +const pool_alloc = std.heap.c_allocator; + +pub const health_budget: u8 = 3; +pub const default_max_connection_age_ms: i64 = 55 * 60 * 1000; +const max_connection_age_env = "FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS"; + +const Slot = struct { + session_id: []u8, + account_id: []u8, + model: []u8, + endpoint: []u8, + authorization_fingerprint: [std.crypto.hash.sha2.Sha256.digest_length]u8, + connection: ?*websocket_transport.Connection, + busy: bool, + health_failures: u8, + opened_at_ms: i64, + + fn deinit(self: *Slot) void { + if (self.connection) |connection| websocket_transport.close(connection, pool_alloc); + pool_alloc.free(self.session_id); + pool_alloc.free(self.account_id); + pool_alloc.free(self.model); + pool_alloc.free(self.endpoint); + self.* = undefined; + } +}; + +var pool_mutex: std.Io.Mutex = .init; +var slots: std.ArrayList(Slot) = .empty; + +pub const AcquireArgs = struct { + session_id: ?[]const u8, + account_id: []const u8, + model: []const u8, + endpoint: []const u8, + authorization: []const u8, + deadline: ?std.Io.Clock.Timestamp, + cancel_flag: *std.atomic.Value(bool), + delivery: *gateway_client.DeliveryCertainty, +}; + +pub const Checkout = struct { + slot: usize, + connection: *websocket_transport.Connection, + reused: bool, + handshake_ms: i64, + health_failures: u8, +}; + +pub const Outcome = enum { completed, failed }; + +fn sessionKey(session_id: ?[]const u8) []const u8 { + const value = session_id orelse return ""; + return if (value.len == 0) "" else value; +} + +fn authorizationFingerprint(authorization: []const u8) [std.crypto.hash.sha2.Sha256.digest_length]u8 { + var digest: [std.crypto.hash.sha2.Sha256.digest_length]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(authorization, &digest, .{}); + return digest; +} + +fn matches(slot: *const Slot, args: AcquireArgs) bool { + const fingerprint = authorizationFingerprint(args.authorization); + return std.mem.eql(u8, slot.session_id, sessionKey(args.session_id)) and + std.mem.eql(u8, slot.account_id, args.account_id) and + std.mem.eql(u8, slot.model, args.model) and + std.mem.eql(u8, slot.endpoint, args.endpoint) and + std.mem.eql(u8, &slot.authorization_fingerprint, &fingerprint); +} + +fn maxConnectionAgeMs() !i64 { + const value = io_mod.getenv(max_connection_age_env) orelse return default_max_connection_age_ms; + const parsed = std.fmt.parseInt(i64, value, 10) catch return error.InvalidOpenAICodexTransport; + if (parsed < 0) return error.InvalidOpenAICodexTransport; + return parsed; +} + +fn incompatibleIdentity(slot: *const Slot, args: AcquireArgs) bool { + const fingerprint = authorizationFingerprint(args.authorization); + return std.mem.eql(u8, slot.session_id, sessionKey(args.session_id)) and + (!std.mem.eql(u8, slot.account_id, args.account_id) or + !std.mem.eql(u8, slot.model, args.model) or + !std.mem.eql(u8, slot.endpoint, args.endpoint) or + !std.mem.eql(u8, &slot.authorization_fingerprint, &fingerprint)); +} + +fn appendSlot(args: AcquireArgs) !usize { + const session_id = try pool_alloc.dupe(u8, sessionKey(args.session_id)); + errdefer pool_alloc.free(session_id); + const account_id = try pool_alloc.dupe(u8, args.account_id); + errdefer pool_alloc.free(account_id); + const model = try pool_alloc.dupe(u8, args.model); + errdefer pool_alloc.free(model); + const endpoint = try pool_alloc.dupe(u8, args.endpoint); + errdefer pool_alloc.free(endpoint); + try slots.append(pool_alloc, .{ + .session_id = session_id, + .account_id = account_id, + .model = model, + .endpoint = endpoint, + .authorization_fingerprint = authorizationFingerprint(args.authorization), + .connection = null, + .busy = false, + .health_failures = 0, + .opened_at_ms = 0, + }); + return slots.items.len - 1; +} + +fn findSlot(args: AcquireArgs) ?usize { + for (slots.items, 0..) |*slot, index| if (matches(slot, args)) return index; + return null; +} + +fn incrementFailure(slot: *Slot) void { + slot.health_failures = std.math.add(u8, slot.health_failures, 1) catch std.math.maxInt(u8); +} + +pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { + while (true) { + if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; + pool_mutex.lockUncancelable(io_mod.getIo()); + var locked = true; + errdefer if (locked) pool_mutex.unlock(io_mod.getIo()); + for (slots.items) |*existing| { + if (!incompatibleIdentity(existing, args) or existing.busy) continue; + if (existing.connection) |connection| websocket_transport.close(connection, pool_alloc); + existing.connection = null; + } + const index = findSlot(args) orelse try appendSlot(args); + const slot = &slots.items[index]; + if (slot.busy) { + pool_mutex.unlock(io_mod.getIo()); + locked = false; + io_mod.sleep(10 * std.time.ns_per_ms); + continue; + } + + const age_limit = try maxConnectionAgeMs(); + if (slot.connection != null and slot.health_failures >= health_budget) { + websocket_transport.close(slot.connection.?, pool_alloc); + slot.connection = null; + } + if (slot.connection != null and age_limit != 0 and io_mod.milliTimestamp() - slot.opened_at_ms > age_limit) { + websocket_transport.close(slot.connection.?, pool_alloc); + slot.connection = null; + } + if (slot.connection) |connection| { + const ping_result = websocket_transport.ping(connection, args.cancel_flag, args.deadline, args.delivery); + if (ping_result) |_| { + slot.busy = true; + const checkout = Checkout{ + .slot = index, + .connection = connection, + .reused = true, + .handshake_ms = 0, + .health_failures = slot.health_failures, + }; + pool_mutex.unlock(io_mod.getIo()); + return checkout; + } else |err| { + websocket_transport.close(connection, pool_alloc); + slot.connection = null; + incrementFailure(slot); + if (err == error.Cancelled) return err; + } + } + + const started_at_ms = io_mod.milliTimestamp(); + const connection = websocket_transport.connect(pool_alloc, .{ + .endpoint = args.endpoint, + .authorization = args.authorization, + .account_id = args.account_id, + .session_id = args.session_id, + .deadline = args.deadline, + .cancel_flag = args.cancel_flag, + .delivery = args.delivery, + }) catch |err| return err; + slot.connection = connection; + slot.opened_at_ms = connection.opened_at_ms; + slot.busy = true; + const checkout = Checkout{ + .slot = index, + .connection = connection, + .reused = false, + .handshake_ms = @max(io_mod.milliTimestamp() - started_at_ms, 0), + .health_failures = slot.health_failures, + }; + pool_mutex.unlock(io_mod.getIo()); + return checkout; + } +} + +pub fn release(index: usize, outcome: Outcome) void { + pool_mutex.lockUncancelable(io_mod.getIo()); + defer pool_mutex.unlock(io_mod.getIo()); + if (index >= slots.items.len) return; + const slot = &slots.items[index]; + slot.busy = false; + switch (outcome) { + .completed => slot.health_failures = 0, + .failed => { + incrementFailure(slot); + if (slot.connection) |connection| websocket_transport.close(connection, pool_alloc); + slot.connection = null; + }, + } +} + +pub fn shutdown() void { + pool_mutex.lockUncancelable(io_mod.getIo()); + defer pool_mutex.unlock(io_mod.getIo()); + for (slots.items) |*slot| slot.deinit(); + slots.deinit(pool_alloc); + slots = .empty; +} + +test "retained Codex WebSocket identity includes authorization without storing it" { + const slot = Slot{ + .session_id = @constCast("session-a"), + .account_id = @constCast("account-a"), + .model = @constCast("gpt-5.6-sol"), + .endpoint = @constCast("http://127.0.0.1/responses"), + .authorization_fingerprint = authorizationFingerprint("Bearer token-a"), + .connection = null, + .busy = false, + .health_failures = 0, + .opened_at_ms = 0, + }; + const base = AcquireArgs{ + .session_id = "session-a", + .account_id = "account-a", + .model = "gpt-5.6-sol", + .endpoint = "http://127.0.0.1/responses", + .authorization = "Bearer token-a", + .deadline = null, + .cancel_flag = undefined, + .delivery = undefined, + }; + + try std.testing.expect(matches(&slot, base)); + try std.testing.expect(!incompatibleIdentity(&slot, base)); + + var rotated = base; + rotated.authorization = "Bearer token-b"; + try std.testing.expect(!matches(&slot, rotated)); + try std.testing.expect(incompatibleIdentity(&slot, rotated)); + + var changed_model = base; + changed_model.model = "gpt-5.4"; + try std.testing.expect(!matches(&slot, changed_model)); + try std.testing.expect(incompatibleIdentity(&slot, changed_model)); +} diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index 6645cc473..0ba98bb3a 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -8,6 +8,8 @@ const types = @import("../core/shared/types.zig"); const gateway_client = @import("client.zig"); const responses_protocol = @import("responses_protocol.zig"); const websocket_transport = @import("websocket_transport.zig"); +const codex_websocket_session = @import("codex_websocket_session.zig"); +const debug_trace = @import("../core/shared/debug_trace.zig"); const model_tool_schema = @import("../core/tooling/model_tool_schema.zig"); const Allocator = std.mem.Allocator; @@ -58,6 +60,10 @@ pub const agent_stream_provider = stream_provider.Provider{ .stream_fn = streamCompletion, }; +pub fn shutdownWebSockets() void { + codex_websocket_session.shutdown(); +} + fn validateModel(model: []const u8) !void { if (model.len == 0 or model.len > 1024) return error.InvalidOpenAICodexModel; for (model) |byte| { @@ -427,27 +433,69 @@ fn streamWebSocketPrepared( .content_capture_limit = request.content_capture_limit, .stream_limits = codexStreamLimits(.{}), }; - // Admission is shared with SSE and happens before the upgrade can make - // delivery possible. The transport marks delivery only immediately before - // it begins writing the masked response.create frame. try admitCodexTransport(request.admission); - try websocket_transport.stream(alloc, .{ - .endpoint = prepared.endpoint, - .authorization = prepared.authorization, - .account_id = prepared.account_id, - .session_id = request.session_id, - .payload = payload, - .deadline = request.deadline, - .cancel_flag = request.cancel_flag, - .delivery = request.delivery, - }, &bridge, WebSocketBridge.event); - const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| - return mapReducerError(err); - return .{ .completed = .{ - .completion = completion, - .usage = .{ .unavailable = .possibly_billed }, - .ownership = .owned, - } }; + + var acquisition_attempt: u8 = 0; + while (true) { + const checkout = codex_websocket_session.acquire(alloc, .{ + .session_id = request.session_id, + .account_id = prepared.account_id, + .model = request.model, + .endpoint = prepared.endpoint, + .authorization = prepared.authorization, + .deadline = request.deadline, + .cancel_flag = request.cancel_flag, + .delivery = request.delivery, + }) catch |err| { + if (acquisition_attempt == 0 and request.delivery.load() == .definitely_unsent) { + acquisition_attempt += 1; + continue; + } + return err; + }; + debug_trace.eventf("codex.ws", "turn", request.trace_ctx, "reused={d} handshake_ms={d} health={d} auth=chatgpt_subscription", .{ + @as(u8, @intFromBool(checkout.reused)), + checkout.handshake_ms, + checkout.health_failures, + }); + websocket_transport.streamOn(checkout.connection, alloc, .{ + .endpoint = prepared.endpoint, + .authorization = prepared.authorization, + .account_id = prepared.account_id, + .session_id = request.session_id, + .payload = payload, + .deadline = request.deadline, + .cancel_flag = request.cancel_flag, + .delivery = request.delivery, + }, &bridge, WebSocketBridge.event) catch |err| { + codex_websocket_session.release(checkout.slot, .failed); + debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ websocketFailureReason(err), @as(u16, 0) }); + return err; + }; + const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| { + codex_websocket_session.release(checkout.slot, .failed); + debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ "protocol", @as(u16, 0) }); + return mapReducerError(err); + }; + codex_websocket_session.release(checkout.slot, .completed); + return .{ .completed = .{ + .completion = completion, + .usage = .{ .unavailable = .possibly_billed }, + .ownership = .owned, + } }; + } +} + +fn websocketFailureReason(err: anyerror) []const u8 { + return switch (err) { + error.Cancelled => "cancel", + error.Timeout => "timeout", + error.WebSocketPolicyClosed => "policy", + error.WebSocketUnexpectedBinary => "binary", + error.WebSocketProtocolViolation, error.WebSocketInvalidUtf8 => "protocol", + error.WebSocketUpgradeRejected, error.WebSocketAcceptInvalid => "auth", + else => "close", + }; } const WebSocketBridge = struct { diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index 22c991fb4..0d212a060 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -20,10 +20,22 @@ pub const Error = error{ pub const EventHandler = *const fn (context: *anyopaque, json: []const u8) anyerror!bool; -const connect_timeout_ms: i64 = 30_000; -const event_idle_timeout_ms: i64 = 30_000; +const default_connect_timeout_ms: i64 = 30_000; +const default_event_idle_timeout_ms: i64 = 30_000; +const connect_timeout_env = "FX_CODEX_WEBSOCKET_CONNECT_TIMEOUT_MS"; +const event_idle_timeout_env = "FX_CODEX_WEBSOCKET_EVENT_IDLE_TIMEOUT_MS"; -pub const Request = struct { +pub const ConnectArgs = struct { + endpoint: []const u8, + authorization: []const u8, + account_id: []const u8, + session_id: ?[]const u8, + deadline: ?std.Io.Clock.Timestamp, + cancel_flag: *std.atomic.Value(bool), + delivery: *gateway_client.DeliveryCertainty, +}; + +pub const StreamArgs = struct { endpoint: []const u8, authorization: []const u8, account_id: []const u8, @@ -33,6 +45,47 @@ pub const Request = struct { cancel_flag: *std.atomic.Value(bool), delivery: *gateway_client.DeliveryCertainty, }; + +pub const Request = StreamArgs; + +pub const Connection = struct { + alloc: Allocator, + client: std.http.Client, + request: std.http.Client.Request, + opened_at_ms: i64, + close_sent: bool = false, + watcher_done: std.atomic.Value(bool) = std.atomic.Value(bool).init(true), + timeout_fired: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), + last_progress_ms: std.atomic.Value(i64), + watcher: ?std.Thread = null, + + fn socket(self: *Connection) !*std.http.Client.Connection { + return self.request.connection orelse error.WebSocketConnectionMissing; + } + + fn startWatcher(self: *Connection, cancel_flag: *std.atomic.Value(bool), deadline: ?std.Io.Clock.Timestamp) !void { + self.watcher_done.store(false, .seq_cst); + self.timeout_fired.store(false, .seq_cst); + self.last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); + const http_connection = try self.socket(); + self.watcher = try spawnConnectionWatcher( + &self.watcher_done, + cancel_flag, + deadline, + &self.timeout_fired, + &self.last_progress_ms, + try eventIdleTimeoutMs(), + http_connection.stream_writer.stream, + ); + } + + fn stopWatcher(self: *Connection) void { + self.watcher_done.store(true, .seq_cst); + if (self.watcher) |thread| thread.join(); + self.watcher = null; + } +}; + const OpenedRequest = struct { request: ?std.http.Client.Request, @@ -68,17 +121,24 @@ const OpenWebSocketOperation = struct { } }; -/// Opens one socket, sends one request, and consumes one terminal response. -/// The caller owns delivery certainty: this function returns an error after a -/// frame write without retrying the request. -pub fn stream( - alloc: Allocator, - request: Request, - context: *anyopaque, - on_event: EventHandler, -) !void { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - const uri = try std.Uri.parse(request.endpoint); +fn positiveTimeoutFromEnv(name: []const u8, fallback: i64) !i64 { + const value = io_mod.getenv(name) orelse return fallback; + const parsed = std.fmt.parseInt(i64, value, 10) catch return error.InvalidOpenAICodexTransport; + if (parsed <= 0) return error.InvalidOpenAICodexTransport; + return parsed; +} + +fn connectTimeoutMs() !i64 { + return positiveTimeoutFromEnv(connect_timeout_env, default_connect_timeout_ms); +} + +fn eventIdleTimeoutMs() !i64 { + return positiveTimeoutFromEnv(event_idle_timeout_env, default_event_idle_timeout_ms); +} + +pub fn connect(alloc: Allocator, args: ConnectArgs) !*Connection { + if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; + const uri = try std.Uri.parse(args.endpoint); var nonce: [16]u8 = undefined; try io_mod.getIo().randomSecure(&nonce); var key_buffer: [std.base64.standard.Encoder.calcSize(nonce.len)]u8 = undefined; @@ -88,7 +148,7 @@ pub fn stream( var extra_headers: [7]std.http.Header = undefined; var count: usize = 0; - extra_headers[count] = .{ .name = "chatgpt-account-id", .value = request.account_id }; + extra_headers[count] = .{ .name = "chatgpt-account-id", .value = args.account_id }; count += 1; extra_headers[count] = .{ .name = "originator", .value = "fx" }; count += 1; @@ -100,67 +160,53 @@ pub fn stream( count += 1; extra_headers[count] = .{ .name = "Sec-WebSocket-Key", .value = &key_buffer }; count += 1; - if (request.session_id) |session_id| if (session_id.len > 0) { + if (args.session_id) |session_id| if (session_id.len > 0) { extra_headers[count] = .{ .name = "session-id", .value = session_id }; count += 1; }; - var client: std.http.Client = .{ .allocator = alloc, .io = io_mod.getIo() }; - defer client.deinit(); + const connection = try alloc.create(Connection); + errdefer alloc.destroy(connection); + connection.* = undefined; + connection.alloc = alloc; + connection.client = .{ .allocator = alloc, .io = io_mod.getIo() }; + errdefer connection.client.deinit(); var open_operation = OpenWebSocketOperation{ - .client = &client, + .client = &connection.client, .uri = uri, - .authorization = request.authorization, + .authorization = args.authorization, .headers = extra_headers[0..count], }; var connect_deadline = std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ .clock = .awake, - .raw = .fromMilliseconds(connect_timeout_ms), + .raw = .fromMilliseconds(try connectTimeoutMs()), }); - if (request.deadline) |deadline| { - if (std.Io.Clock.Timestamp.compare(deadline, .lt, connect_deadline)) { - connect_deadline = deadline; - } + if (args.deadline) |deadline| { + if (std.Io.Clock.Timestamp.compare(deadline, .lt, connect_deadline)) connect_deadline = deadline; } var opened = try gateway_client.runBoundedHttpOperation( OpenedRequest, alloc, - request.cancel_flag, + args.cancel_flag, connect_deadline, &open_operation, ); - var http_request = opened.take(); - defer { - // An upgraded connection must never return to the HTTP pool. - if (http_request.connection) |connection| connection.closing = true; - http_request.deinit(); - } - var watcher_done = std.atomic.Value(bool).init(false); - var timeout_fired = std.atomic.Value(bool).init(false); - var last_progress_ms = std.atomic.Value(i64).init(io_mod.milliTimestamp()); - const watcher = if (http_request.connection) |connection| - try spawnConnectionWatcher( - &watcher_done, - request.cancel_flag, - request.deadline, - &timeout_fired, - &last_progress_ms, - connection.stream_writer.stream, - ) - else - null; - defer { - watcher_done.store(true, .seq_cst); - if (watcher) |thread| thread.join(); - } - http_request.sendBodiless() catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; + errdefer opened.deinit(alloc); + connection.request = opened.take(); + errdefer connection.request.deinit(); + connection.opened_at_ms = io_mod.milliTimestamp(); + connection.close_sent = false; + connection.watcher_done = std.atomic.Value(bool).init(true); + connection.timeout_fired = std.atomic.Value(bool).init(false); + connection.last_progress_ms = std.atomic.Value(i64).init(connection.opened_at_ms); + connection.watcher = null; + + connection.request.sendBodiless() catch |err| { + if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; return err; }; - const response = http_request.receiveHead(&.{}) catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; + const response = connection.request.receiveHead(&.{}) catch |err| { + if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; return err; }; if (response.head.status != .switching_protocols) return error.WebSocketUpgradeRejected; @@ -170,61 +216,81 @@ pub fn stream( { return error.WebSocketAcceptInvalid; } + _ = try connection.socket(); + return connection; +} - // `receiveHead` leaves any already-buffered WebSocket bytes on this reader. - const reader = http_request.reader.in; - const connection = http_request.connection orelse return error.WebSocketConnectionMissing; - const writer = connection.writer(); - var close_sent = false; - defer if (!close_sent and !request.cancel_flag.load(.seq_cst)) { - // Cancellation force-releases the socket from the watcher. Do not race - // that release with a best-effort close frame from this I/O owner. - // A fresh Phase 1 socket is never returned to the HTTP pool. - writeFrame(writer, .close, &.{ 0x03, 0xe8 }) catch {}; - connection.flush() catch {}; - }; +fn operationError(connection: *Connection, cancel_flag: *std.atomic.Value(bool), err: anyerror) anyerror { + if (cancel_flag.load(.seq_cst)) return error.Cancelled; + if (connection.timeout_fired.load(.seq_cst)) return error.Timeout; + return err; +} + +pub fn ping( + connection: *Connection, + cancel_flag: *std.atomic.Value(bool), + deadline: ?std.Io.Clock.Timestamp, + _: *gateway_client.DeliveryCertainty, +) !void { + if (cancel_flag.load(.seq_cst)) return error.Cancelled; + try connection.startWatcher(cancel_flag, deadline); + defer connection.stopWatcher(); + const socket = try connection.socket(); + const writer = socket.writer(); + writeFrame(writer, .ping, &.{}) catch |err| return operationError(connection, cancel_flag, err); + socket.flush() catch |err| return operationError(connection, cancel_flag, err); + while (true) { + const frame = readFrame(connection.alloc, connection.request.reader.in) catch |err| return operationError(connection, cancel_flag, err); + defer connection.alloc.free(frame.payload); + connection.last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); + switch (frame.opcode) { + .pong => return, + .ping => { + writeFrame(writer, .pong, frame.payload) catch |err| return operationError(connection, cancel_flag, err); + socket.flush() catch |err| return operationError(connection, cancel_flag, err); + }, + .close => return closeError(try validateClosePayload(frame.payload)), + .binary => return error.WebSocketUnexpectedBinary, + else => return error.WebSocketProtocolViolation, + } + } +} + +pub fn streamOn( + connection: *Connection, + alloc: Allocator, + request: StreamArgs, + context: *anyopaque, + on_event: EventHandler, +) !void { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + try connection.startWatcher(request.cancel_flag, request.deadline); + defer connection.stopWatcher(); + const socket = try connection.socket(); + const reader = connection.request.reader.in; + const writer = socket.writer(); + var succeeded = false; + defer if (!succeeded) closeAfterFailure(connection, request.cancel_flag); request.delivery.markPossiblySent(); - writeFrame(writer, .text, request.payload) catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; - return err; - }; - connection.flush() catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; - return err; - }; - last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); + writeFrame(writer, .text, request.payload) catch |err| return operationError(connection, request.cancel_flag, err); + socket.flush() catch |err| return operationError(connection, request.cancel_flag, err); + connection.last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); var message: std.ArrayList(u8) = .empty; defer message.deinit(alloc); var fragmented_opcode: ?Opcode = null; while (true) { if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - const frame = readFrame(alloc, reader) catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; - return err; - }; - last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); + const frame = readFrame(alloc, reader) catch |err| return operationError(connection, request.cancel_flag, err); + connection.last_progress_ms.store(io_mod.milliTimestamp(), .seq_cst); defer alloc.free(frame.payload); switch (frame.opcode) { .ping => { - writeFrame(writer, .pong, frame.payload) catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; - return err; - }; - connection.flush() catch |err| { - if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (timeout_fired.load(.seq_cst)) return error.Timeout; - return err; - }; + writeFrame(writer, .pong, frame.payload) catch |err| return operationError(connection, request.cancel_flag, err); + socket.flush() catch |err| return operationError(connection, request.cancel_flag, err); }, .pong => {}, - .close => { - return closeError(try validateClosePayload(frame.payload)); - }, + .close => return closeError(try validateClosePayload(frame.payload)), .binary => return error.WebSocketUnexpectedBinary, .continuation => { if (fragmented_opcode == null) return error.WebSocketProtocolViolation; @@ -234,11 +300,7 @@ pub fn stream( fragmented_opcode = null; if (opcode != .text) return error.WebSocketUnexpectedBinary; if (try dispatchTextMessage(context, on_event, message.items)) { - closeAfterCompletion(alloc, reader, writer, connection, request.cancel_flag, &timeout_fired) catch |err| { - if (err == error.Cancelled or err == error.Timeout) close_sent = true; - return err; - }; - close_sent = true; + succeeded = true; return; } message.clearRetainingCapacity(); @@ -251,11 +313,7 @@ pub fn stream( continue; } if (try dispatchTextMessage(context, on_event, message.items)) { - closeAfterCompletion(alloc, reader, writer, connection, request.cancel_flag, &timeout_fired) catch |err| { - if (err == error.Cancelled or err == error.Timeout) close_sent = true; - return err; - }; - close_sent = true; + succeeded = true; return; } message.clearRetainingCapacity(); @@ -264,6 +322,75 @@ pub fn stream( } } +fn closeAfterFailure(connection: *Connection, cancel_flag: *std.atomic.Value(bool)) void { + if (connection.close_sent or cancel_flag.load(.seq_cst)) return; + const socket = connection.socket() catch return; + writeFrame(socket.writer(), .close, &.{ 0x03, 0xe8 }) catch return; + socket.flush() catch {}; + connection.close_sent = true; +} + +fn closeChecked( + connection: *Connection, + alloc: Allocator, + cancel_flag: *std.atomic.Value(bool), + deadline: ?std.Io.Clock.Timestamp, +) !void { + connection.stopWatcher(); + defer { + if (connection.request.connection) |http_connection| http_connection.closing = true; + connection.request.deinit(); + connection.client.deinit(); + alloc.destroy(connection); + } + if (!connection.close_sent) { + try connection.startWatcher(cancel_flag, deadline); + defer connection.stopWatcher(); + const http_connection = try connection.socket(); + try closeAfterCompletion( + alloc, + connection.request.reader.in, + http_connection.writer(), + http_connection, + cancel_flag, + &connection.timeout_fired, + ); + connection.close_sent = true; + } +} + +pub fn close(connection: *Connection, alloc: Allocator) void { + var cancelled = std.atomic.Value(bool).init(false); + const deadline = std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ + .clock = .awake, + .raw = .fromMilliseconds(1_000), + }); + closeChecked(connection, alloc, &cancelled, deadline) catch {}; +} + +/// Opens one socket, sends one request, and consumes one terminal response. +pub fn stream( + alloc: Allocator, + request: Request, + context: *anyopaque, + on_event: EventHandler, +) !void { + const connection = try connect(alloc, .{ + .endpoint = request.endpoint, + .authorization = request.authorization, + .account_id = request.account_id, + .session_id = request.session_id, + .deadline = request.deadline, + .cancel_flag = request.cancel_flag, + .delivery = request.delivery, + }); + var owned = true; + defer if (owned) close(connection, alloc); + try streamOn(connection, alloc, request, context, on_event); + owned = false; + return closeChecked(connection, alloc, request.cancel_flag, request.deadline); +} + const Opcode = enum(u4) { continuation = 0, text = 1, binary = 2, close = 8, ping = 9, pong = 10 }; const Frame = struct { fin: bool, opcode: Opcode, payload: []u8 }; @@ -350,6 +477,7 @@ const ConnectionWatcher = struct { deadline: ?std.Io.Clock.Timestamp, timeout_fired: *std.atomic.Value(bool), last_progress_ms: *std.atomic.Value(i64), + event_idle_timeout_ms: i64, socket: std.Io.net.Stream, ) void { while (!done.load(.seq_cst)) { @@ -382,6 +510,7 @@ fn spawnConnectionWatcher( deadline: ?std.Io.Clock.Timestamp, timeout_fired: *std.atomic.Value(bool), last_progress_ms: *std.atomic.Value(i64), + event_idle_timeout_ms: i64, socket: std.Io.net.Stream, ) !std.Thread { return std.Thread.spawn(.{}, ConnectionWatcher.run, .{ @@ -390,6 +519,7 @@ fn spawnConnectionWatcher( deadline, timeout_fired, last_progress_ms, + event_idle_timeout_ms, socket, }); } diff --git a/src/main.zig b/src/main.zig index bfcedb050..a68c21b99 100644 --- a/src/main.zig +++ b/src/main.zig @@ -70,6 +70,7 @@ const js_host_workspace = @import("core/hosts/js_host_workspace.zig"); const host_target = @import("core/hosts/target.zig"); const native_host = @import("core/hosts/native.zig"); const debug_trace = @import("core/shared/debug_trace.zig"); +const openai_codex = @import("gateway/openai_codex.zig"); const display_width = @import("core/shared/display_width.zig"); const file_index_mod = @import("core/workspace/file_index.zig"); const mcp_command_provider = @import("core/mcp/command_provider.zig"); @@ -3064,7 +3065,10 @@ fn mainC(c_argc: c_int, c_argv: [*][*:0]c_char, c_envp: [*:null]?[*:0]c_char) !v }); defer threaded.deinit(); io_mod.setIo(threaded.io()); - defer debug_trace.shutdown(); + defer { + debug_trace.shutdown(); + openai_codex.shutdownWebSockets(); + } debug_trace.configureFromEnv(processAllocator(), "."); try terminal_host.run( processAllocator(), @@ -3145,6 +3149,7 @@ fn runNonBenchmark(raw_args: []const [*:0]const u8, raw_env: RawEnviron, cli_arg }); if (early_threaded) |*threaded| io_mod.setIo(threaded.io()); } + defer openai_codex.shutdownWebSockets(); const before = try app_entry_runtime.runBeforeInteractive(alloc, cli_args, cfg); switch (before) { @@ -3160,7 +3165,10 @@ fn runNonBenchmark(raw_args: []const [*:0]const u8, raw_env: RawEnviron, cli_arg var owned_launch = launch; defer owned_launch.deinit(alloc); - defer debug_trace.shutdown(); + defer { + debug_trace.shutdown(); + openai_codex.shutdownWebSockets(); + } const outcome = try app_entry_runtime.runInteractive(App, alloc, &owned_launch); switch (outcome) { @@ -3646,6 +3654,7 @@ test "session reset traces and clears active paste state" { try std.testing.expectEqual(@as(usize, 0), app.input_runtime.paste.decision_bytes); try std.testing.expectEqual(@as(usize, 0), app.input_runtime.edit_state.input.items.len); debug_trace.shutdown(); + openai_codex.shutdownWebSockets(); var trace_file = try std.Io.Dir.openFileAbsolute(io_mod.getIo(), trace_path, .{}); defer trace_file.close(io_mod.getIo()); diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index 86804197a..805508002 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -841,6 +841,8 @@ function startFakeCodexWebSocket(options: { holdOpen?: boolean; closeOnOpen?: number; closeAfterMessage?: number; + closeAfterFirstMessage?: number; + closeAfterCompletion?: boolean; toolThenClose?: boolean; stallUpgrade?: boolean; } = {}) { @@ -890,12 +892,17 @@ function startFakeCodexWebSocket(options: { ws.close(options.closeAfterMessage, "fixture close"); return; } + if (options.closeAfterFirstMessage !== undefined && requests.length === 1) { + ws.close(options.closeAfterFirstMessage, "fixture first-message close"); + return; + } if (options.holdOpen) return; - ws.send(JSON.stringify({ type: "response.output_text.delta", delta: "CODEX_WEBSOCKET_OK" })); + ws.send(JSON.stringify({ type: "response.output_text.delta", delta: `CODEX_WEBSOCKET_OK_${requests.length}` })); ws.send(JSON.stringify({ type: "response.completed", response: { id: "resp_websocket", status: "completed", usage: { input_tokens: 5, output_tokens: 2 } }, })); + if (options.closeAfterCompletion) ws.close(1000, "fixture completed"); }, close(_ws, code) { closeCodes.push(code); @@ -2875,7 +2882,6 @@ test( expect(codex.requests).toHaveLength(1); expect(codex.requests[0]).toContain('"type":"response.create"'); expect(codex.requests[0]).not.toContain('"stream"'); - expect(codex.closeCodes).toContain(1000); expect(gateway.requests).toHaveLength(0); } finally { codex.stop(); @@ -3014,6 +3020,166 @@ test( 60_000, ); +tmuxTest( + "Codex WebSocket reuses one socket for sequential interactive turns", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-reuse-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket(); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Complete the first retained-socket turn."); + await session.waitForText("CODEX_WEBSOCKET_OK_1", TIMEOUT); + await session.sendText("Complete the second retained-socket turn."); + const secondDeadline = Date.now() + TIMEOUT; + while (codex.requests.length < 2) { + if (Date.now() >= secondDeadline) throw new Error("Second Codex WebSocket request did not arrive"); + await Bun.sleep(25); + } + await session.waitForText("CODEX_WEBSOCKET_OK_2", TIMEOUT); + expect(codex.upgradeRequests).toBe(1); + expect(codex.requests).toHaveLength(2); + expect(codex.requests.every((request) => request.includes('"type":"response.create"'))).toBe(true); + expect(codex.requests[1]).toContain("Complete the first retained-socket turn."); + expect(codex.requests[1]).toContain("Complete the second retained-socket turn."); + expect(codex.requests[1]).toContain("CODEX_WEBSOCKET_OK_1"); + expect(codex.requests[1]).not.toContain("previous_response_id"); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + +tmuxTest( + "Codex WebSocket reconnects before delivery when the retained socket closes", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-reconnect-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ closeAfterCompletion: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Complete before the retained socket closes."); + await session.waitForText("CODEX_WEBSOCKET_OK_1", TIMEOUT); + await session.sendText("Reconnect without replaying either turn."); + await session.waitForText("CODEX_WEBSOCKET_OK_2", TIMEOUT); + expect(codex.upgradeRequests).toBe(2); + expect(codex.requests).toHaveLength(2); + expect(codex.requests[1]).toContain("Complete before the retained socket closes."); + expect(codex.requests[1]).toContain("Reconnect without replaying either turn."); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + +tmuxTest( + "Codex WebSocket poisons a failed stream and recovers on the next turn", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-poison-recovery-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ closeAfterFirstMessage: 1011 }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Fail this retained-socket turn once."); + await session.waitForText("WebSocketClosedBeforeCompletion", TIMEOUT); + await session.sendText("Recover on a fresh socket."); + await session.waitForText("CODEX_WEBSOCKET_OK_2", TIMEOUT); + expect(codex.upgradeRequests).toBe(2); + expect(codex.requests).toHaveLength(2); + expect(codex.requests[0]).toContain("Fail this retained-socket turn once."); + expect(codex.requests[1]).toContain("Recover on a fresh socket."); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + +tmuxTest( + "Codex WebSocket evicts a retained socket after its configured maximum age", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-age-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket(); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS: "1", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Complete on the initial short-lived socket."); + await session.waitForText("CODEX_WEBSOCKET_OK_1", TIMEOUT); + await Bun.sleep(20); + await session.sendText("Complete after connection-age eviction."); + await session.waitForText("CODEX_WEBSOCKET_OK_2", TIMEOUT); + expect(codex.upgradeRequests).toBe(2); + expect(codex.requests).toHaveLength(2); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + tmuxTest( "Codex WebSocket cancellation unblocks a stalled upgrade", async () => { From 7dfd7c9d0796b5bb07c3fef99e54bf14421165a1 Mon Sep 17 00:00:00 2001 From: Ashman Singh <36335693+thinkter@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:40:57 +0530 Subject: [PATCH 15/21] fix --- docs/codex-websocket-transport.md | 329 ------------------------------ 1 file changed, 329 deletions(-) delete mode 100644 docs/codex-websocket-transport.md diff --git a/docs/codex-websocket-transport.md b/docs/codex-websocket-transport.md deleted file mode 100644 index 6bd40528b..000000000 --- a/docs/codex-websocket-transport.md +++ /dev/null @@ -1,329 +0,0 @@ -# OpenAI Codex Responses WebSocket transport - -## Scope and evidence status - -This document concerns the ChatGPT-subscription Codex route at `chatgpt.com/backend-api/codex/responses`, not fx's Vercel AI Gateway route. Both use a Responses WebSocket protocol family, but they are separate backends. Public Responses API, Azure, and AI Gateway documentation must not be treated as a specification for the private ChatGPT backend. - -Claims are labelled as follows: - -- **Confirmed:** supported by the fx checkout or current upstream `openai/codex` source. -- **Likely, verify:** supported by public protocol documentation or related implementations, but not yet measured against the ChatGPT subscription backend. -- **Undocumented:** a private-backend behavior that requires live, authenticated probing before fx relies on it. - -This is an implementation design, not a claim that fx already supports WebSockets. - -## Current fx behavior - -**Confirmed.** The Codex provider currently uses HTTPS with Server-Sent Events (SSE): - -1. fx reads or refreshes the local ChatGPT OAuth session, extracts the ChatGPT account ID, and constructs a bearer token. -2. It serializes a complete Responses request containing the model, system instructions, conversation history, tools, tool outputs, images, and retained encrypted reasoning state. -3. It sends that JSON with `POST https://chatgpt.com/backend-api/codex/responses` and asks for `text/event-stream`. -4. The server streams `data:` records over that HTTP request. -5. `src/gateway/responses_protocol.zig` reduces each JSON event into text, reasoning, tool-call, and final-completion events. - -The implementation has a 30-second connection deadline, cancellation handling, and limits for aggregate stream data, events, tool calls, tool identities, tool arguments, and preserved provider state. Each model request opens a new HTTP connection. - -## What WebSockets change - -**Likely, verify for this backend.** A secure WebSocket begins with an HTTPS Upgrade request. After a valid `101 Switching Protocols` response, both sides retain an encrypted, bidirectional connection. - -```text -fx Codex service - HTTPS WebSocket Upgrade → - 101 Switching Protocols ← - response.create frame → - response event frames ← - response.completed ← -``` - -The Responses data remains JSON. The client sends a text message resembling: - -```json -{ - "type": "response.create", - "model": "gpt-5...", - "instructions": "...", - "input": [], - "tools": [] -} -``` - -The server returns text messages such as `response.output_text.delta`, `response.function_call_arguments.delta`, `response.completed`, and `response.failed`. A WebSocket event source should feed decoded JSON text directly into fx's existing Responses reducer. It must not create a second model-event implementation. - -## Upstream Codex CLI behavior - -**Confirmed for upstream Codex, not automatically for the private backend.** Upstream Codex has a dedicated `responses_websocket` transport, gated by provider capability, and retains a healthy connection for its client session. - -### Handshake and metadata - -**Confirmed.** Codex constructs a WebSocket URL, attaches provider and authentication headers during the HTTP upgrade, requests a versioned Responses WebSocket beta protocol, and validates the upgrade. It can record selected server model, reasoning inclusion, Codex turn state, rate-limit data, model ETags, moderation metadata, and timing information. - -It also supports a handshake probe that upgrades without sending a prompt and briefly waits for an immediate close. This distinguishes a usable connection from one accepted by the edge then rejected by policy. - -### One request at a time - -**Confirmed.** Codex serializes one response stream per socket. It sends `response.create`, reads until a terminal event, then permits the next request. It does not interleave independent generations on the same connection. - -This is the correct initial model for fx. Concurrent requests require request identity, event demultiplexing, flow control, and independent recovery. - -### Event failures and reuse - -**Confirmed.** Codex applies an idle timeout to each server-event wait. It treats idle timeout, EOF before completion, close before completion, unexpected binary messages, I/O failure, and structured service errors as stream failures. It handles Ping and Pong as transport housekeeping and discards a connection after a terminal stream error. - -**Confirmed for upstream source; likely, verify for this backend.** Upstream recognizes `websocket_connection_limit_reached` with a 60-minute connection limit and recognizes `previous_response_not_found` as recoverable by a full-context request. - -**Likely, verify.** `response.failed` must invalidate its continuation chain. A later request must not reuse the failed response as `previous_response_id`. - -### Incremental continuation - -**Confirmed for upstream client behavior; undocumented for the ChatGPT backend.** A retained connection can send a new `response.create` with `previous_response_id` and only newly added input items. This reduces repeated upload of long transcripts and tool history. - -Until Phase 0 verifies the ChatGPT backend, fx must regard continuation state as **possibly connection-scoped**, not guaranteed connection-scoped. Saved fx sessions must always retain enough history to reconstruct a full request; remote state is never the sole source of truth. - -## The central reliability rule: delivery certainty - -A network failure does not simply mean a request failed: - -- **Request has not begun writing:** automatic retry is safe. -- **Request may have been written but no acknowledgement arrived:** do not retry blindly. -- **The server acknowledged a response ID:** recover only through supported server state or a known-safe replay. - -If fx sends `response.create` then loses the network before `response.created`, the service may still generate a response and issue tool calls. Blindly retrying could duplicate shell commands or other actions. - -fx already models this distinction for HTTP with `DeliveryCertainty`. A WebSocket transport must preserve it. Upstream Codex may retry an established stream before falling back, but fx deliberately diverges: its tool-capable turns require a stricter no-blind-replay policy after delivery becomes ambiguous. - -## Phase 0: authenticated live-traffic probe - -Before transport implementation, run the isolated maintainer probe at `scripts/codex_websocket_probe.py` against `chatgpt.com/backend-api/codex/responses`. It uses only Python's standard library, accepts credentials only through explicit environment variables, never reads fx credential files, and emits one redacted JSON report. It never writes credentials, account IDs, response IDs, prompts, or response content. - -```bash -export FX_CODEX_PROBE_ACCESS_TOKEN='...' -export FX_CODEX_PROBE_ACCOUNT_ID='...' -export FX_CODEX_PROBE_MODEL='...' -python3 scripts/codex_websocket_probe.py -python3 scripts/codex_websocket_probe.py --execute --continuation -``` - -The first invocation performs only an authenticated upgrade. `--execute` sends a fixed minimal prompt and consumes subscription usage; `--continuation` sends a second request using the first response ID without printing that ID. Run `python3 scripts/codex_websocket_probe.py --self-test` for deterministic, credential-free checks. - -Record only privacy-safe protocol evidence: - -- handshake status, selected response headers, and immediate close behavior; -- event type sequence, terminal event, structured error code, and close code; -- whether WebSocket accepts `previous_response_id` continuation; -- whether HTTP/SSE accepts `previous_response_id` continuation; -- behavior after a failed response in a continuation chain; -- whether model changes on a retained connection produce a policy close such as 1008; -- connection-age behavior and any limit/error code; -- behavior of `store` and related persistence fields, if accepted. - -Do not hard-code public API assumptions until this probe confirms them. In particular, the 60-minute limit, `previous_response_not_found`, model pinning, `store` semantics, and SSE continuation support are undocumented for this private backend. - -## Proposed implementation plan - -### Phase 1: per-turn WebSocket transport - -Implement a fresh socket per model request, with no retained connection or cached continuation. - -- Keep SSE as the permanent compatibility baseline. -- Reuse existing Codex request generation and Responses reduction. -- Send one full `response.create` request on a fresh socket. -- Fall back to SSE immediately for failed upgrades, connection timeouts, or other failures definitely before delivery. -- After a request may have been transmitted, never silently replay through SSE. -- Treat server close before `response.completed` as an explicit failure class. -- Once Phase 0 confirms behavior, treat policy close 1008 as a wrong-model-on-connection failure: poison the connection and recreate it with the correct model, rather than retrying it as a network failure. -- Maintain a strict retry policy for established streams. Fewer retries, including zero, are safer than duplicating a possible tool call. - -### Phase 2: retained socket, full context per turn - -- [x] Retain one idle native WebSocket per session ID, account, model, and endpoint. A slot serializes one stream at a time and pings before reuse. -- [x] Send a complete `response.create` body for every turn. Phase 2 does not send `previous_response_id` or retain continuation state. -- [x] Keep `FX_CODEX_TRANSPORT=websocket` opt-in only. SSE remains the default, `auto` remains SSE, and no SSE fallback occurs from the required WebSocket path. -- [x] Poison and close sockets after cancellation, timeout, protocol or binary-frame failures, writes, policy close, unhealthy ping, identity changes, and age eviction. -- [x] Evict idle connections at 55 minutes by default, with `FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS=0` disabling age eviction. -- [x] Reset the health budget only after a completed response, never after an upgrade alone. - -Connect and event-idle timeouts default to 30 seconds and may be overridden with `FX_CODEX_WEBSOCKET_CONNECT_TIMEOUT_MS` and `FX_CODEX_WEBSOCKET_EVENT_IDLE_TIMEOUT_MS`. Values must be positive integer milliseconds. - -Local Phase 2 verification exercises the freshly built binary against a loopback WebSocket service. It covers two full-context turns on one socket, absence of `previous_response_id`, pre-delivery reconnect after a peer close, poisoning and recovery after an in-flight failure, maximum-age eviction, policy and post-delivery close handling, and cancellation during upgrade and idle reads. The retained identity includes a SHA-256 fingerprint of the authorization value so credential rotation invalidates a socket without retaining the credential itself. - -### Phase 3: incremental continuation - -Only begin after Phase 0 verifies continuation semantics and Phase 2 has stable evidence. Maintain a continuation record containing: - -```text -connection identity -credential and account fingerprint -endpoint and protocol version -model and request-shaping options -last fully completed response ID -full request required for recovery -continuation expiry and validity -``` - -Invalidate it when the connection reconnects, model or relevant options change, authentication changes, the referenced turn fails, cancellation interrupts a stream, a tool call or result is uncertain, the server rejects the previous response ID, or the connection reaches its lifetime. - -### Phase 4: optional stream lanes - -**Likely, verify.** If the protocol supports a `stream_id` field with tagged events, named lanes could eventually allow concurrent subagent turns to share a retained connection. Requests in one lane must remain ordered; separate lanes require event demultiplexing and independent continuation state. - -Do not implement this before Phases 1 through 3 are stable. It substantially changes connection ownership, resource accounting, and failure recovery. - -## Transport requirements - -Keep WebSocket mechanics separate from Codex request serialization: - -```text -src/gateway/websocket_transport.zig RFC 6455 handshake, frames, cancellation, limits -src/gateway/responses_stream.zig JSON event source shared by SSE and WebSocket -src/gateway/openai_codex.zig Codex auth, payload, and transport selection -``` - -The transport must implement and test: - -- TLS certificate and hostname validation; -- HTTP `101 Switching Protocols` and `Sec-WebSocket-Accept` validation; -- mandatory masking of client frames; -- text-message fragmentation and continuation-frame reassembly; -- UTF-8 validation for text messages; -- Ping to Pong handling; -- close codes, bounded close reasons, and a finite close deadline; -- cancellation that unblocks connection, read, and write operations; -- separate connection, write, and event-idle deadlines; -- bounded outbound requests, inbound frames, reassembled messages, and buffered unread bytes; -- strict rejection of unexpected message kinds and protocol violations. - -Do not enable per-message WebSocket compression in the first release. It adds decompression limits, compatibility cases, CPU cost, and security surface. - -## Connection lifecycle and fallback policy - -Represent a connection as a state machine: - -```text -disconnected → connecting → open-idle → open-streaming → closing - ↘ poisoned → disconnected -``` - -One component owns socket I/O. UI, retry, and shutdown paths communicate through controlled cancellation or commands; they must never concurrently read from or write to the socket. - -A socket becomes poisoned and its continuation state is discarded after a timeout, protocol violation, unexpected binary event, early EOF, failed write, failed close, failed response, or server error that leaves state uncertain. - -Expose a prominent user-facing policy: - -```text -auto Prefer a known-good WebSocket, otherwise use SSE. -sse Always use HTTP and SSE. -websocket Require WebSocket and return an actionable error if unavailable. -``` - -In `auto`, only pre-delivery WebSocket failures may fall back automatically to SSE. After a request may have been delivered, report uncertainty or use verified server recovery. Do not silently replay. - -If fx later adds remote Responses compaction for this provider, route it over HTTP unless Phase 0 or subsequent probes confirm WebSocket support for that endpoint. - -## Observability and rollout - -Ship with a local kill switch and measured rollout. Record only privacy-safe operational data: - -- selected transport, protocol version, platform, fx version, and authentication mode; -- handshake status and duration; -- new versus reused connection; -- time to first event and terminal completion; -- error, close, retry, circuit-breaker, and fallback classification; -- bounded byte and event counts; -- continuation hit, invalidation, and recovery reason. - -Never log prompts, responses, OAuth tokens, account IDs, raw tool arguments, or unredacted headers. - -Track handshake success, terminal completion rate, fallback and ambiguous-delivery rates, latency by transport and auth mode, errors by platform/network, and long-session memory and file-descriptor stability. - -## Verification plan - -Before broad enablement, require: - -1. Zig unit tests for handshake validation, masking, fragmentation, control frames, malformed frames, UTF-8, and resource limits. -2. A scripted loopback fixture for delayed events, oversized frames, invalid events, disconnects at write/read boundaries, and server close before `response.completed`. -3. Reducer-parity tests that feed identical Responses JSON through SSE and WebSocket and assert identical completion data and callback ordering. -4. A delivery-certainty retry matrix, including a request that may have been sent but was never acknowledged and a duplicate-tool-call prevention case. -5. Model-mismatch coverage, after Phase 0 verifies the behavior, proving a policy close is not retried as a generic transient network error. -6. Continuation coverage proving that a failed referenced turn invalidates the chain and forces full context. -7. Circuit-breaker coverage proving that failures deplete the budget and only a confirmed completion resets it. -8. Cancellation and shutdown tests while connecting, writing, waiting for output, receiving tool arguments, and closing, including leak checks. -9. Platform-matrix connect-timeout tests and long-running soak tests with forced reconnects and connection-age expiry. -10. A real-binary smoke test using freshly built `./zig-out/bin/fx` against a loopback fixture and an interactive terminal path. - -## Phase 1 implementation review and required fixes - -This review applies to the Phase 1 implementation. It is not production-ready and must retain SSE as the default until every release blocker below is resolved and verified. - -### Implementation progress - -- [x] Frame parsing and UTF-8 validation: regression tests pass in the repository test target, including a 127-byte extended frame followed by another frame. Close payloads and malformed text are also validated at the frame boundary. -- [x] Provider admission, transport-neutral request construction, shared Codex preparation, and shared stream limits: covered by unit assertions in the repository test target. -- [x] Bounded upgrade, write, and event-idle I/O with cancellation unblocking: unit and tmux loopback tests cover stalled connect, backpressured write, stalled upgrade, idle frame read, and a hung close handshake. -- [x] Bounded close handshake and close-code reporting: malformed close payloads are unit-tested, an early policy close is classified without retry, and shutdown uses independent bounded teardown state rather than retaining pointers into a completed turn. -- [x] Loopback WebSocket fixture and real-binary smoke coverage: `tests/e2e/tui-auth-source-selection.test.ts` runs `./zig-out/bin/fx` against Bun's local WebSocket server for completion, close failures, retained reuse, reconnect, poisoning recovery, age eviction, and cancellation. Delivery certainty remains definitely unsent until immediately before the masked `response.create` frame write. A unit test proves the SSE and WebSocket reducers preserve callback order and completion data for the same Responses event sequence. -- [x] Authenticated Phase 0 evidence: an explicit maintainer probe on 2026-08-27 returned `101`, completed a request, and accepted a same-socket `previous_response_id` continuation. The redacted report recorded only protocol labels, a 1,749 ms handshake, and the `x-models-etag` header name. - -### Release blockers - -1. **Extended-length frame parsing corrupts a valid 127-byte frame.** - - `src/gateway/websocket_transport.zig` reads the seven-bit length discriminator, then uses two independent `if` statements. If that discriminator is `126` and the following 16-bit length is exactly `127`, the first branch correctly reads `127`, then the second branch incorrectly treats it as the 64-bit discriminator and consumes eight payload bytes as a length. This desynchronizes the stream. - - Fix: make the 64-bit branch `else if (length == 127)`. Add a regression test for a 127-byte text frame followed by another frame, proving the second frame remains aligned. - -2. **The WebSocket path does not participate in provider admission.** - - The SSE path calls `request.admission.admit()` before opening transport. The WebSocket path currently does not, which produces the user-visible `provider admission missing` failure and bypasses the normal provider-attempt lifecycle. - - Fix: preserve the admission boundary for every transport before opening its request. Add a focused test that runs the WebSocket path through the same admission fixture as SSE. - -3. **The WebSocket path has no bounded connection deadline or cancellation unblock.** - - SSE opens through `runBoundedHttpOperation` with the 30-second connect deadline and installs `spawnHttpCancelWatcher` to interrupt blocked connection I/O. The WebSocket path invokes the HTTP client directly and only checks `cancel_flag` between completed frame reads. A blocked `reader.takeByte()` cannot observe cancellation until the peer sends bytes or closes. - - Fix: use an operation that bounds connection setup, writes, and event-idle reads, and that closes or interrupts the underlying connection when cancellation occurs. Verify cancellation while connecting, while writing, and while waiting for a frame. - -4. **Close handling does not meet the documented transport contract.** - - The current code marks the HTTP connection closing and deinitializes it without sending a WebSocket close frame on normal terminal, cancellation, and error paths. It also does not parse or report close code and bounded reason. - - Fix: send one bounded close frame when the socket is open, await a peer close for a finite deadline where appropriate, then forcibly release the connection. Record and classify peer close code and bounded reason. Cover normal completion, cancellation, malformed frames, and early peer close. - -5. **Text frames are not explicitly validated as UTF-8.** - - The reducer eventually parses JSON, but WebSocket text-message validity must be enforced at the frame-message boundary so malformed text is classified as a transport protocol error rather than an incidental JSON failure. - - Fix: validate completed text messages before calling the event handler, return a dedicated protocol error, and add malformed UTF-8 coverage. - -### Required cleanup before broad enablement - -1. **Replace fragile WebSocket request string surgery.** - - `buildWebSocketRequest` finds the literal `,"store":false,"stream":true` in an SSE payload and splices around it. This couples WebSocket behavior to exact field ordering and formatting in `buildRequest`. The current test uses a hardcoded SSE payload, so it cannot catch drift in the real request builder. - - Minimum fix: add an end-to-end unit assertion for `buildWebSocketRequest(buildRequest(...))` using representative system messages, tools, reasoning, images, and structured output. Preferred fix: split request construction into a transport-neutral Responses request model or shared field writer, then add transport-specific envelopes without parsing a serialized request. - -2. **Extract shared Codex authentication and endpoint preparation.** - - `streamPrepared` and `streamWebSocketPrepared` duplicate account-ID extraction, authorization-header allocation, and loopback-only E2E endpoint selection. - - Fix: introduce a small helper whose result owns or scopes the prepared account ID, authorization header, and endpoint. Keep credential material zeroed and free it at the same boundary as today. - -3. **Create one stream-limits builder.** - - The six stream-limit values are repeated in `CodexLimits`, `WebSocketBridge.streamLimits`, and the SSE `consumeSse` conversion. The WebSocket method ignores `self`, which is evidence it should be a shared pure builder. - - Fix: use one `codexStreamLimits(CodexLimits)` helper returning `responses_protocol.StreamLimits`; use it for both SSE and WebSocket reducers. - -### Lower-priority hardening and explicit decisions - -- Guard `http_request.connection` after a successful upgrade instead of using `connection.?`. A `101` should have a connection, but treating an impossible state as an error is safer than trapping. -- Document the asymmetric limits: outbound frames permit a full 64 MiB request, while inbound frames are limited to 4 MiB and reassembled messages to 64 MiB. This is reasonable if large outbound contexts are intentional and response events are expected to be fragmented, but it needs rationale and tests at each boundary. -- Keep `auto` mapped to SSE in Phase 1. This is intentional and compatible with the rollout plan, not a bug. Do not change `auto` to WebSocket until pre-delivery fallback, health evidence, and the release blockers are implemented. -- The Phase 0 Python probe is standard-library-only, reads only explicitly named environment credentials, emits redacted output, and is suitable to retain as an untracked diagnostic script. - -## Current status - -fx has robust Codex HTTPS/SSE request generation and Responses reduction. Phase 1 adds a fresh-socket WebSocket experiment behind `FX_CODEX_TRANSPORT=websocket`; SSE remains the default and `auto` maps to SSE. Phase 2 retains a healthy socket for later turns in the same process while continuing to send full request context. It has no continuation protocol, no `previous_response_id`, and no WebSocket-to-SSE fallback. Local Phase 2 verification is green; exact-commit Full CI remains the release-readiness authority. From a3058d64a2d0ddd1ba09bbbff048aca67c87369f Mon Sep 17 00:00:00 2001 From: thinkter Date: Sat, 29 Aug 2026 15:45:12 +0530 Subject: [PATCH 16/21] Add Codex WebSocket continuation --- src/gateway/codex_websocket_session.zig | 121 +++++++++++++ src/gateway/openai_codex.zig | 133 ++++++++++++++- src/gateway/responses_protocol.zig | 48 ++++++ tests/e2e/tui-auth-source-selection.test.ts | 180 +++++++++++++++++++- 4 files changed, 470 insertions(+), 12 deletions(-) diff --git a/src/gateway/codex_websocket_session.zig b/src/gateway/codex_websocket_session.zig index 5094d2473..03e1ba0ed 100644 --- a/src/gateway/codex_websocket_session.zig +++ b/src/gateway/codex_websocket_session.zig @@ -20,9 +20,25 @@ const Slot = struct { busy: bool, health_failures: u8, opened_at_ms: i64, + continuation_response_id: ?[]u8, + continuation_baseline: ?[]u8, + continuation_durable_baseline: ?[]u8, + continuation_shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, + continuation_valid: bool, + + fn clearContinuation(self: *Slot) void { + if (self.continuation_response_id) |value| pool_alloc.free(value); + if (self.continuation_baseline) |value| pool_alloc.free(value); + if (self.continuation_durable_baseline) |value| pool_alloc.free(value); + self.continuation_response_id = null; + self.continuation_baseline = null; + self.continuation_durable_baseline = null; + self.continuation_valid = false; + } fn deinit(self: *Slot) void { if (self.connection) |connection| websocket_transport.close(connection, pool_alloc); + self.clearContinuation(); pool_alloc.free(self.session_id); pool_alloc.free(self.account_id); pool_alloc.free(self.model); @@ -53,6 +69,18 @@ pub const Checkout = struct { health_failures: u8, }; +pub const Continuation = struct { + previous_response_id: []const u8, + delta_input: []const u8, +}; + +fn continuationDelta(full_input: []const u8, baseline: []const u8) ?[]const u8 { + if (!std.mem.startsWith(u8, full_input, baseline)) return null; + if (full_input.len == baseline.len) return ""; + if (full_input[baseline.len] != ',') return null; + return full_input[baseline.len + 1 ..]; +} + pub const Outcome = enum { completed, failed }; fn sessionKey(session_id: ?[]const u8) []const u8 { @@ -110,6 +138,11 @@ fn appendSlot(args: AcquireArgs) !usize { .busy = false, .health_failures = 0, .opened_at_ms = 0, + .continuation_response_id = null, + .continuation_baseline = null, + .continuation_durable_baseline = null, + .continuation_shape = undefined, + .continuation_valid = false, }); return slots.items.len - 1; } @@ -133,6 +166,7 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { if (!incompatibleIdentity(existing, args) or existing.busy) continue; if (existing.connection) |connection| websocket_transport.close(connection, pool_alloc); existing.connection = null; + existing.clearContinuation(); } const index = findSlot(args) orelse try appendSlot(args); const slot = &slots.items[index]; @@ -147,10 +181,12 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { if (slot.connection != null and slot.health_failures >= health_budget) { websocket_transport.close(slot.connection.?, pool_alloc); slot.connection = null; + slot.clearContinuation(); } if (slot.connection != null and age_limit != 0 and io_mod.milliTimestamp() - slot.opened_at_ms > age_limit) { websocket_transport.close(slot.connection.?, pool_alloc); slot.connection = null; + slot.clearContinuation(); } if (slot.connection) |connection| { const ping_result = websocket_transport.ping(connection, args.cancel_flag, args.deadline, args.delivery); @@ -168,6 +204,7 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { } else |err| { websocket_transport.close(connection, pool_alloc); slot.connection = null; + slot.clearContinuation(); incrementFailure(slot); if (err == error.Cancelled) return err; } @@ -184,6 +221,7 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { .delivery = args.delivery, }) catch |err| return err; slot.connection = connection; + slot.clearContinuation(); slot.opened_at_ms = connection.opened_at_ms; slot.busy = true; const checkout = Checkout{ @@ -198,6 +236,67 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { } } +pub fn continuation( + index: usize, + full_input: []const u8, + shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, +) ?Continuation { + pool_mutex.lockUncancelable(io_mod.getIo()); + defer pool_mutex.unlock(io_mod.getIo()); + if (index >= slots.items.len) return null; + const slot = &slots.items[index]; + if (!slot.busy or !slot.continuation_valid) return null; + if (!std.mem.eql(u8, &slot.continuation_shape, &shape)) { + slot.clearContinuation(); + return null; + } + const response_id = slot.continuation_response_id orelse return null; + const baseline = slot.continuation_baseline orelse return null; + const delta = continuationDelta(full_input, baseline) orelse durable: { + const durable_baseline = slot.continuation_durable_baseline orelse { + slot.clearContinuation(); + return null; + }; + break :durable continuationDelta(full_input, durable_baseline) orelse { + slot.clearContinuation(); + return null; + }; + }; + return .{ + .previous_response_id = response_id, + .delta_input = delta, + }; +} + +pub fn recordCompletion( + index: usize, + response_id: []const u8, + baseline: []const u8, + durable_baseline: []const u8, + shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, +) void { + pool_mutex.lockUncancelable(io_mod.getIo()); + defer pool_mutex.unlock(io_mod.getIo()); + if (index >= slots.items.len) return; + const slot = &slots.items[index]; + slot.clearContinuation(); + const owned_id = pool_alloc.dupe(u8, response_id) catch return; + const owned_baseline = pool_alloc.dupe(u8, baseline) catch { + pool_alloc.free(owned_id); + return; + }; + const owned_durable_baseline = pool_alloc.dupe(u8, durable_baseline) catch { + pool_alloc.free(owned_id); + pool_alloc.free(owned_baseline); + return; + }; + slot.continuation_response_id = owned_id; + slot.continuation_baseline = owned_baseline; + slot.continuation_durable_baseline = owned_durable_baseline; + slot.continuation_shape = shape; + slot.continuation_valid = true; +} + pub fn release(index: usize, outcome: Outcome) void { pool_mutex.lockUncancelable(io_mod.getIo()); defer pool_mutex.unlock(io_mod.getIo()); @@ -210,6 +309,7 @@ pub fn release(index: usize, outcome: Outcome) void { incrementFailure(slot); if (slot.connection) |connection| websocket_transport.close(connection, pool_alloc); slot.connection = null; + slot.clearContinuation(); }, } } @@ -233,6 +333,11 @@ test "retained Codex WebSocket identity includes authorization without storing i .busy = false, .health_failures = 0, .opened_at_ms = 0, + .continuation_response_id = null, + .continuation_baseline = null, + .continuation_durable_baseline = null, + .continuation_shape = undefined, + .continuation_valid = false, }; const base = AcquireArgs{ .session_id = "session-a", @@ -258,3 +363,19 @@ test "retained Codex WebSocket identity includes authorization without storing i try std.testing.expect(!matches(&slot, changed_model)); try std.testing.expect(incompatibleIdentity(&slot, changed_model)); } + +test "Codex WebSocket continuation requires an exact item boundary prefix" { + try std.testing.expectEqualStrings( + "{\"role\":\"user\",\"content\":[]}", + continuationDelta( + "{\"type\":\"message\"},{\"role\":\"user\",\"content\":[]}", + "{\"type\":\"message\"}", + ).?, + ); + try std.testing.expectEqualStrings( + "", + continuationDelta("{\"type\":\"message\"}", "{\"type\":\"message\"}").?, + ); + try std.testing.expect(continuationDelta("{\"type\":\"message\"}suffix", "{\"type\":\"message\"}") == null); + try std.testing.expect(continuationDelta("{\"type\":\"other\"}", "{\"type\":\"message\"}") == null); +} diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index 0ba98bb3a..c04e91dfb 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -86,10 +86,25 @@ pub fn buildRequest( fn buildWebSocketRequest( alloc: Allocator, request: stream_provider.RequestData, +) ![]u8 { + const input = try buildResponseInput(alloc, request.messages, request.verified_images); + defer alloc.free(input); + return buildWebSocketRequestWithInput(alloc, request, input, null); +} + +fn buildWebSocketRequestWithInput( + alloc: Allocator, + request: stream_provider.RequestData, + input: []const u8, + previous_response_id: ?[]const u8, ) ![]u8 { var out: std.Io.Writer.Allocating = .init(alloc); errdefer out.deinit(); - try writeResponseRequestStart(&out.writer, alloc, request, "response.create"); + try writeResponseRequestStartWithInput(&out.writer, alloc, request, "response.create", input); + if (previous_response_id) |response_id| { + try out.writer.writeAll(",\"previous_response_id\":"); + try std.json.Stringify.value(response_id, .{}, &out.writer); + } try out.writer.writeByte('}'); return out.toOwnedSlice(); } @@ -101,6 +116,18 @@ fn writeResponseRequestStart( alloc: Allocator, request: stream_provider.RequestData, websocket_type: ?[]const u8, +) !void { + const input = try buildResponseInput(alloc, request.messages, request.verified_images); + defer alloc.free(input); + return writeResponseRequestStartWithInput(writer, alloc, request, websocket_type, input); +} + +fn writeResponseRequestStartWithInput( + writer: *std.Io.Writer, + alloc: Allocator, + request: stream_provider.RequestData, + websocket_type: ?[]const u8, + input: []const u8, ) !void { try validateModel(request.model); if (request.budget) |budget| { @@ -131,7 +158,7 @@ fn writeResponseRequestStart( try writer.writeAll(",\"instructions\":"); try std.json.Stringify.value(instructions.written(), .{}, writer); try writer.writeAll(",\"input\":["); - try writeResponsesInput(writer, alloc, request.messages, request.verified_images); + try writer.writeAll(input); try writer.writeByte(']'); _ = try responses_protocol.writeTools(writer, alloc, request.tools); @@ -165,6 +192,42 @@ fn writeResponseRequestStart( // the public Responses API max_output_tokens parameter. } +fn buildResponseInput( + alloc: Allocator, + messages: []const types.ChatMessage, + images: ?[]const image_attachments.VerifiedSnapshot, +) ![]u8 { + var out: std.Io.Writer.Allocating = .init(alloc); + errdefer out.deinit(); + try writeResponsesInput(&out.writer, alloc, messages, images); + return out.toOwnedSlice(); +} + +fn buildContinuationBaseline( + alloc: Allocator, + full_input: []const u8, + response_input: []const u8, +) ![]u8 { + var out: std.Io.Writer.Allocating = .init(alloc); + errdefer out.deinit(); + try out.writer.writeAll(full_input); + if (response_input.len > 0) { + if (out.written().len > 0) try out.writer.writeByte(','); + try out.writer.writeAll(response_input); + } + return out.toOwnedSlice(); +} + +fn buildDurableResponseInput( + alloc: Allocator, + messages: []const types.ChatMessage, +) ![]u8 { + const projected = try alloc.dupe(types.ChatMessage, messages); + defer alloc.free(projected); + for (projected) |*message| message.provider_state_json = null; + return buildResponseInput(alloc, projected, null); +} + fn writeResponsesInput( writer: *std.Io.Writer, alloc: Allocator, @@ -202,9 +265,7 @@ fn streamCompletion( break :blk streamPrepared(alloc, request, payload); }, .websocket => blk: { - const payload = try buildWebSocketRequest(alloc, request.data()); - defer alloc.free(payload); - break :blk streamWebSocketPrepared(alloc, request, payload); + break :blk streamWebSocketPrepared(alloc, request); }, } catch |err| { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); @@ -417,12 +478,18 @@ fn prepareCodexTransport( fn streamWebSocketPrepared( alloc: Allocator, request: stream_provider.ModelRequest, - payload: []const u8, ) !stream_provider.Result { if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; var prepared = try prepareCodexTransport(alloc, request); defer prepared.deinit(alloc); + const full_input = try buildResponseInput(alloc, request.messages, request.verified_images); + defer alloc.free(full_input); + const shape_payload = try buildWebSocketRequestWithInput(alloc, request.data(), "", null); + defer alloc.free(shape_payload); + var shape: [std.crypto.hash.sha2.Sha256.digest_length]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(shape_payload, &shape, .{}); + var reducer = responses_protocol.Reducer.init(alloc); defer reducer.deinit(alloc); var bridge = WebSocketBridge{ @@ -436,6 +503,7 @@ fn streamWebSocketPrepared( try admitCodexTransport(request.admission); var acquisition_attempt: u8 = 0; + var continuation_recovery_attempted = false; while (true) { const checkout = codex_websocket_session.acquire(alloc, .{ .session_id = request.session_id, @@ -458,6 +526,22 @@ fn streamWebSocketPrepared( checkout.handshake_ms, checkout.health_failures, }); + const continued = if (continuation_recovery_attempted) + null + else + codex_websocket_session.continuation(checkout.slot, full_input, shape); + const payload = try buildWebSocketRequestWithInput( + alloc, + request.data(), + if (continued) |value| value.delta_input else full_input, + if (continued) |value| value.previous_response_id else null, + ); + defer alloc.free(payload); + debug_trace.eventf("codex.ws", "continuation", request.trace_ctx, "used={d} delta_bytes={d} recovery={d}", .{ + @as(u8, @intFromBool(continued != null)), + if (continued) |value| value.delta_input.len else full_input.len, + @as(u8, @intFromBool(continuation_recovery_attempted)), + }); websocket_transport.streamOn(checkout.connection, alloc, .{ .endpoint = prepared.endpoint, .authorization = prepared.authorization, @@ -469,6 +553,12 @@ fn streamWebSocketPrepared( .delivery = request.delivery, }, &bridge, WebSocketBridge.event) catch |err| { codex_websocket_session.release(checkout.slot, .failed); + if (err == error.PreviousResponseNotFound and continued != null and !continuation_recovery_attempted) { + continuation_recovery_attempted = true; + reducer.deinit(alloc); + reducer = responses_protocol.Reducer.init(alloc); + continue; + } debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ websocketFailureReason(err), @as(u16, 0) }); return err; }; @@ -477,6 +567,36 @@ fn streamWebSocketPrepared( debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ "protocol", @as(u16, 0) }); return mapReducerError(err); }; + if (completion.generation_id) |response_id| { + const response_message = [_]types.ChatMessage{.{ + .role = .assistant, + .content = completion.content, + .tool_calls = completion.tool_calls, + .provider_state_json = completion.provider_state_json, + }}; + if (buildResponseInput(alloc, &response_message, null)) |response_input| { + defer alloc.free(response_input); + if (buildContinuationBaseline(alloc, full_input, response_input)) |baseline| { + defer alloc.free(baseline); + if (buildDurableResponseInput(alloc, request.messages)) |durable_full_input| { + defer alloc.free(durable_full_input); + if (buildDurableResponseInput(alloc, &response_message)) |durable_response_input| { + defer alloc.free(durable_response_input); + if (buildContinuationBaseline(alloc, durable_full_input, durable_response_input)) |durable_baseline| { + defer alloc.free(durable_baseline); + codex_websocket_session.recordCompletion( + checkout.slot, + response_id, + baseline, + durable_baseline, + shape, + ); + } else |_| {} + } else |_| {} + } else |_| {} + } else |_| {} + } else |_| {} + } codex_websocket_session.release(checkout.slot, .completed); return .{ .completed = .{ .completion = completion, @@ -661,6 +781,7 @@ fn consumeSse( fn mapReducerError(err: anyerror) anyerror { return switch (err) { error.InvalidEvent => error.InvalidOpenAICodexSseEvent, + error.PreviousResponseNotFound => error.PreviousResponseNotFound, error.ResponseFailed => error.OpenAICodexResponseFailed, error.StreamIncomplete => error.OpenAICodexStreamIncomplete, error.ToolCallLimitExceeded => error.OpenAICodexToolCallLimitExceeded, diff --git a/src/gateway/responses_protocol.zig b/src/gateway/responses_protocol.zig index 4707a13d7..186371d5a 100644 --- a/src/gateway/responses_protocol.zig +++ b/src/gateway/responses_protocol.zig @@ -378,6 +378,21 @@ pub const Reducer = struct { } else if (std.mem.eql(u8, event_type, "response.failed") or std.mem.eql(u8, event_type, "error")) { + const error_value = parsed.value.object.get("error"); + const response_value = parsed.value.object.get("response"); + const response_error = if (response_value != null and response_value.? == .object) + response_value.?.object.get("error") + else + null; + const code = if (error_value != null and error_value.? == .object) + stringField(error_value.?.object, "code") + else if (response_error != null and response_error.? == .object) + stringField(response_error.?.object, "code") + else + stringField(parsed.value.object, "code"); + if (code) |value| if (std.mem.eql(u8, value, "previous_response_not_found")) { + return error.PreviousResponseNotFound; + }; return error.ResponseFailed; } return false; @@ -738,6 +753,39 @@ test "Responses usage projection retains optional cached and reasoning detail" { try std.testing.expectEqual(@as(?u64, 3), usage.reasoning_tokens); } +test "Responses protocol distinguishes missing WebSocket continuation state" { + const Capture = struct { + fn content(_: *anyopaque, _: []const u8) void {} + fn toolStart(_: *anyopaque, _: []const u8, _: []const u8, _: ?[]const u8) void {} + }; + var reducer = Reducer.init(std.testing.allocator); + defer reducer.deinit(std.testing.allocator); + var cancelled = std.atomic.Value(bool).init(false); + var context: u8 = 0; + try std.testing.expectError( + error.PreviousResponseNotFound, + reducer.applyJson( + std.testing.allocator, + "{\"type\":\"error\",\"error\":{\"code\":\"previous_response_not_found\",\"message\":\"expired\"}}", + .{ + .context = &context, + .on_content = Capture.content, + .on_tool_start = Capture.toolStart, + }, + &cancelled, + null, + .{ + .aggregate_bytes = 4096, + .events = 8, + .tool_calls = 8, + .tool_identity_bytes = 1024, + .tool_arguments_bytes = 4096, + .provider_state_bytes = 4096, + }, + ), + ); +} + test "Responses protocol owns one subscription billing projection" { const alloc = std.testing.allocator; const billing = (try buildSubscriptionBilling( diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index 805508002..8ed277b5d 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -845,10 +845,14 @@ function startFakeCodexWebSocket(options: { closeAfterCompletion?: boolean; toolThenClose?: boolean; stallUpgrade?: boolean; + rejectPreviousOnce?: boolean; + toolRoundTrip?: boolean; + reasoningState?: boolean; } = {}) { const requests: string[] = []; const closeCodes: number[] = []; let upgradeRequests = 0; + let rejectedPrevious = false; const accessToken = chatgptAccessToken("acct_websocket"); const server = Bun.serve<{ opened: boolean }>({ hostname: "127.0.0.1", @@ -874,6 +878,32 @@ function startFakeCodexWebSocket(options: { message(ws, message) { const payload = String(message); requests.push(payload); + const parsed = JSON.parse(payload) as { previous_response_id?: string }; + if (options.rejectPreviousOnce && parsed.previous_response_id && !rejectedPrevious) { + rejectedPrevious = true; + ws.send(JSON.stringify({ + type: "error", + error: { code: "previous_response_not_found", message: "Previous response was not found." }, + })); + return; + } + if (options.toolRoundTrip && requests.length === 1) { + ws.send(JSON.stringify({ + type: "response.output_item.added", + output_index: 0, + item: { type: "function_call", call_id: "call_phase3", name: "read_file" }, + })); + ws.send(JSON.stringify({ + type: "response.function_call_arguments.done", + output_index: 0, + arguments: '{"path":"README.md"}', + })); + ws.send(JSON.stringify({ + type: "response.completed", + response: { id: "resp_websocket_1", status: "completed", usage: { input_tokens: 5, output_tokens: 2 } }, + })); + return; + } if (options.toolThenClose) { ws.send(JSON.stringify({ type: "response.output_item.added", @@ -897,10 +927,17 @@ function startFakeCodexWebSocket(options: { return; } if (options.holdOpen) return; + if (options.reasoningState) { + ws.send(JSON.stringify({ + type: "response.output_item.done", + output_index: 0, + item: { type: "reasoning", id: `reasoning_${requests.length}`, encrypted_content: "opaque" }, + })); + } ws.send(JSON.stringify({ type: "response.output_text.delta", delta: `CODEX_WEBSOCKET_OK_${requests.length}` })); ws.send(JSON.stringify({ type: "response.completed", - response: { id: "resp_websocket", status: "completed", usage: { input_tokens: 5, output_tokens: 2 } }, + response: { id: `resp_websocket_${requests.length}`, status: "completed", usage: { input_tokens: 5, output_tokens: 2 } }, })); if (options.closeAfterCompletion) ws.close(1000, "fixture completed"); }, @@ -2851,7 +2888,7 @@ test( async () => { home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-")); gateway = startFakeGateway([]); - const codex = startFakeCodexWebSocket(); + const codex = startFakeCodexWebSocket({ reasoningState: true }); try { writeSeededChatGptLogin(home, codex.accessToken); writeFileSync( @@ -2933,6 +2970,51 @@ test( 60_000, ); +test( + "Codex WebSocket continues a completed tool call with only its result", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-tool-continuation-")); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ toolRoundTrip: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol", permission_mode: "yolo" }) + "\n", + { mode: 0o600 }, + ); + const result = await runFx( + ["ask", "--json", "--auto", "--no-save", "Read README.md, then report success."], + { + env: { + HOME: home, + AI_GATEWAY_API_KEY: "gateway-websocket-tool-continuation-sentinel", + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_AUTO_UPGRADE: "0", + FX_CODEX_TRANSPORT: "websocket", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }, + timeoutMs: TIMEOUT, + }, + ); + expect(result.code, `stdout: ${result.stdout}\nstderr: ${result.stderr}`).toBe(0); + expect(codex.upgradeRequests).toBe(1); + expect(codex.requests).toHaveLength(2); + expect(codex.requests[1]).toContain('"previous_response_id":"resp_websocket_1"'); + expect(codex.requests[1]).toContain('"type":"function_call_output"'); + expect(codex.requests[1]).not.toContain("Read README.md, then report success."); + expect(gateway.requests).toHaveLength(0); + } finally { + codex.stop(); + } + }, + 60_000, +); + test( "Codex WebSocket close after response.create never replays the turn", async () => { @@ -2976,6 +3058,49 @@ test( 60_000, ); +tmuxTest( + "Codex WebSocket uses full context when persisted tool evidence changes history", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-tool-turn-continuation-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ toolRoundTrip: true, reasoningState: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol", permission_mode: "yolo" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Read README.md for the first Phase 3 turn."); + await session.waitForText("CODEX_WEBSOCKET_OK_2", TIMEOUT); + await session.sendText("Complete the next Phase 3 turn without replay."); + await session.waitForText("CODEX_WEBSOCKET_OK_3", TIMEOUT); + + expect(codex.upgradeRequests).toBe(1); + expect(codex.requests).toHaveLength(3); + expect(codex.requests[1]).toContain('"previous_response_id":"resp_websocket_1"'); + expect(codex.requests[1]).toContain('"type":"function_call_output"'); + expect(codex.requests[2]).not.toContain("previous_response_id"); + expect(codex.requests[2]).toContain("Complete the next Phase 3 turn without replay."); + expect(codex.requests[2]).toContain("Read README.md for the first Phase 3 turn."); + expect(codex.requests[2]).toContain('"type":"function_call_output"'); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + test( "Codex WebSocket tool stream close never replays the turn", @@ -3027,7 +3152,7 @@ tmuxTest( stderrPath = join(home, "stderr.log"); writeFileSync(stderrPath, ""); gateway = startFakeGateway([]); - const codex = startFakeCodexWebSocket(); + const codex = startFakeCodexWebSocket({ reasoningState: true }); try { writeSeededChatGptLogin(home, codex.accessToken); writeFileSync( @@ -3054,10 +3179,10 @@ tmuxTest( expect(codex.upgradeRequests).toBe(1); expect(codex.requests).toHaveLength(2); expect(codex.requests.every((request) => request.includes('"type":"response.create"'))).toBe(true); - expect(codex.requests[1]).toContain("Complete the first retained-socket turn."); expect(codex.requests[1]).toContain("Complete the second retained-socket turn."); - expect(codex.requests[1]).toContain("CODEX_WEBSOCKET_OK_1"); - expect(codex.requests[1]).not.toContain("previous_response_id"); + expect(codex.requests[1]).not.toContain("Complete the first retained-socket turn."); + expect(codex.requests[1]).not.toContain("CODEX_WEBSOCKET_OK_1"); + expect(codex.requests[1]).toContain('"previous_response_id":"resp_websocket_1"'); expect(readFileSync(stderrPath, "utf8")).toBe(""); } finally { codex.stop(); @@ -3104,6 +3229,49 @@ tmuxTest( 60_000, ); +tmuxTest( + "Codex WebSocket retries full context when continuation state is missing", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-continuation-recovery-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ rejectPreviousOnce: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Complete the first continuation-recovery turn."); + await session.waitForText("CODEX_WEBSOCKET_OK_1", TIMEOUT); + await session.sendText("Recover the second continuation-recovery turn."); + await session.waitForText("CODEX_WEBSOCKET_OK_3", TIMEOUT); + + expect(codex.upgradeRequests).toBe(2); + expect(codex.requests).toHaveLength(3); + expect(codex.requests[1]).toContain('"previous_response_id":"resp_websocket_1"'); + expect(codex.requests[1]).not.toContain("Complete the first continuation-recovery turn."); + expect(codex.requests[2]).not.toContain("previous_response_id"); + expect(codex.requests[2]).toContain("Complete the first continuation-recovery turn."); + expect(codex.requests[2]).toContain("CODEX_WEBSOCKET_OK_1"); + expect(codex.requests[2]).toContain("Recover the second continuation-recovery turn."); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + tmuxTest( "Codex WebSocket poisons a failed stream and recovers on the next turn", async () => { From e7ce834fdb80047aa938a25f4bfbfab28b2550bf Mon Sep 17 00:00:00 2001 From: thinkter Date: Sat, 29 Aug 2026 18:59:49 +0530 Subject: [PATCH 17/21] Add bounded Codex WebSocket lanes --- src/gateway/codex_websocket_session.zig | 101 ++++++++++++++++++++++-- src/gateway/openai_codex.zig | 5 +- 2 files changed, 98 insertions(+), 8 deletions(-) diff --git a/src/gateway/codex_websocket_session.zig b/src/gateway/codex_websocket_session.zig index 03e1ba0ed..874cf83fe 100644 --- a/src/gateway/codex_websocket_session.zig +++ b/src/gateway/codex_websocket_session.zig @@ -8,7 +8,9 @@ const pool_alloc = std.heap.c_allocator; pub const health_budget: u8 = 3; pub const default_max_connection_age_ms: i64 = 55 * 60 * 1000; +const default_max_lanes: usize = 4; const max_connection_age_env = "FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS"; +const max_lanes_env = "FX_CODEX_WEBSOCKET_MAX_LANES"; const Slot = struct { session_id: []u8, @@ -59,6 +61,8 @@ pub const AcquireArgs = struct { deadline: ?std.Io.Clock.Timestamp, cancel_flag: *std.atomic.Value(bool), delivery: *gateway_client.DeliveryCertainty, + continuation_input: ?[]const u8 = null, + continuation_shape: ?[std.crypto.hash.sha2.Sha256.digest_length]u8 = null, }; pub const Checkout = struct { @@ -110,6 +114,13 @@ fn maxConnectionAgeMs() !i64 { return parsed; } +fn maxLanes() !usize { + const value = io_mod.getenv(max_lanes_env) orelse return default_max_lanes; + const parsed = std.fmt.parseInt(usize, value, 10) catch return error.InvalidOpenAICodexTransport; + if (parsed == 0) return error.InvalidOpenAICodexTransport; + return parsed; +} + fn incompatibleIdentity(slot: *const Slot, args: AcquireArgs) bool { const fingerprint = authorizationFingerprint(args.authorization); return std.mem.eql(u8, slot.session_id, sessionKey(args.session_id)) and @@ -147,9 +158,41 @@ fn appendSlot(args: AcquireArgs) !usize { return slots.items.len - 1; } -fn findSlot(args: AcquireArgs) ?usize { - for (slots.items, 0..) |*slot, index| if (matches(slot, args)) return index; - return null; +fn continuationMatches(slot: *const Slot, full_input: []const u8, shape: [std.crypto.hash.sha2.Sha256.digest_length]u8) bool { + if (!slot.continuation_valid or !std.mem.eql(u8, &slot.continuation_shape, &shape)) return false; + if (slot.continuation_baseline) |baseline| { + if (continuationDelta(full_input, baseline) != null) return true; + } + if (slot.continuation_durable_baseline) |baseline| { + if (continuationDelta(full_input, baseline) != null) return true; + } + return false; +} + +const LaneSelection = struct { + index: ?usize, + matching_count: usize, +}; + +fn selectIdleLane(args: AcquireArgs) LaneSelection { + var first_idle: ?usize = null; + var continuation_idle: ?usize = null; + var matching_count: usize = 0; + for (slots.items, 0..) |*slot, index| { + if (!matches(slot, args)) continue; + matching_count += 1; + if (slot.busy) continue; + if (first_idle == null) first_idle = index; + if (args.continuation_input) |full_input| { + if (args.continuation_shape) |shape| { + if (continuationMatches(slot, full_input, shape)) { + continuation_idle = index; + break; + } + } + } + } + return .{ .index = continuation_idle orelse first_idle, .matching_count = matching_count }; } fn incrementFailure(slot: *Slot) void { @@ -168,14 +211,20 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { existing.connection = null; existing.clearContinuation(); } - const index = findSlot(args) orelse try appendSlot(args); - const slot = &slots.items[index]; - if (slot.busy) { + // Responses on one WebSocket are exclusive and ordered. Preserve a + // compatible continuation lane when it is idle; otherwise another + // retained socket provides bounded parallelism without multiplexing + // unrelated response events on the same wire. + const selection = selectIdleLane(args); + const index = selection.index orelse if (selection.matching_count < try maxLanes()) + try appendSlot(args) + else { pool_mutex.unlock(io_mod.getIo()); locked = false; io_mod.sleep(10 * std.time.ns_per_ms); continue; - } + }; + const slot = &slots.items[index]; const age_limit = try maxConnectionAgeMs(); if (slot.connection != null and slot.health_failures >= health_budget) { @@ -379,3 +428,41 @@ test "Codex WebSocket continuation requires an exact item boundary prefix" { try std.testing.expect(continuationDelta("{\"type\":\"message\"}suffix", "{\"type\":\"message\"}") == null); try std.testing.expect(continuationDelta("{\"type\":\"other\"}", "{\"type\":\"message\"}") == null); } + +test "Codex WebSocket lane selection preserves continuation affinity" { + shutdown(); + defer shutdown(); + + var cancel_flag = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + var shape: [std.crypto.hash.sha2.Sha256.digest_length]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash("shape", &shape, .{}); + const args = AcquireArgs{ + .session_id = "session-a", + .account_id = "account-a", + .model = "gpt-5.6-sol", + .endpoint = "http://127.0.0.1/responses", + .authorization = "Bearer token-a", + .deadline = null, + .cancel_flag = &cancel_flag, + .delivery = &delivery, + .continuation_input = "{\"type\":\"message\"},{\"role\":\"user\"}", + .continuation_shape = shape, + }; + + const first = try appendSlot(args); + slots.items[first].busy = true; + const second = try appendSlot(args); + slots.items[second].busy = true; + recordCompletion(second, "response-2", "{\"type\":\"message\"}", "{\"type\":\"message\"}", shape); + slots.items[second].busy = false; + + const selection = selectIdleLane(args); + try std.testing.expectEqual(@as(?usize, second), selection.index); + try std.testing.expectEqual(@as(usize, 2), selection.matching_count); + + slots.items[second].busy = true; + const saturated = selectIdleLane(args); + try std.testing.expectEqual(@as(?usize, null), saturated.index); + try std.testing.expectEqual(@as(usize, 2), saturated.matching_count); +} diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index c04e91dfb..f794f56c7 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -514,6 +514,8 @@ fn streamWebSocketPrepared( .deadline = request.deadline, .cancel_flag = request.cancel_flag, .delivery = request.delivery, + .continuation_input = if (continuation_recovery_attempted) null else full_input, + .continuation_shape = if (continuation_recovery_attempted) null else shape, }) catch |err| { if (acquisition_attempt == 0 and request.delivery.load() == .definitely_unsent) { acquisition_attempt += 1; @@ -521,7 +523,8 @@ fn streamWebSocketPrepared( } return err; }; - debug_trace.eventf("codex.ws", "turn", request.trace_ctx, "reused={d} handshake_ms={d} health={d} auth=chatgpt_subscription", .{ + debug_trace.eventf("codex.ws", "turn", request.trace_ctx, "lane={d} reused={d} handshake_ms={d} health={d} auth=chatgpt_subscription", .{ + checkout.slot, @as(u8, @intFromBool(checkout.reused)), checkout.handshake_ms, checkout.health_failures, From ee7aee87169f11722458119013da3ce361173325 Mon Sep 17 00:00:00 2001 From: thinkter Date: Sat, 29 Aug 2026 19:03:40 +0530 Subject: [PATCH 18/21] Harden Codex WebSocket lane saturation --- src/gateway/codex_websocket_session.zig | 48 ++++++++++++++++--------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/src/gateway/codex_websocket_session.zig b/src/gateway/codex_websocket_session.zig index 874cf83fe..0328e1b16 100644 --- a/src/gateway/codex_websocket_session.zig +++ b/src/gateway/codex_websocket_session.zig @@ -174,11 +174,11 @@ const LaneSelection = struct { matching_count: usize, }; -fn selectIdleLane(args: AcquireArgs) LaneSelection { +fn selectIdleLane(slot_items: []Slot, args: AcquireArgs) LaneSelection { var first_idle: ?usize = null; var continuation_idle: ?usize = null; var matching_count: usize = 0; - for (slots.items, 0..) |*slot, index| { + for (slot_items, 0..) |*slot, index| { if (!matches(slot, args)) continue; matching_count += 1; if (slot.busy) continue; @@ -195,6 +195,19 @@ fn selectIdleLane(args: AcquireArgs) LaneSelection { return .{ .index = continuation_idle orelse first_idle, .matching_count = matching_count }; } +const LaneChoice = union(enum) { + existing: usize, + append, + wait, +}; + +fn chooseLane(slot_items: []Slot, args: AcquireArgs, lane_limit: usize) LaneChoice { + const selection = selectIdleLane(slot_items, args); + if (selection.index) |index| return .{ .existing = index }; + if (selection.matching_count < lane_limit) return .append; + return .wait; +} + fn incrementFailure(slot: *Slot) void { slot.health_failures = std.math.add(u8, slot.health_failures, 1) catch std.math.maxInt(u8); } @@ -202,6 +215,10 @@ fn incrementFailure(slot: *Slot) void { pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { while (true) { if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; + if (args.deadline) |deadline| { + const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + if (!std.Io.Clock.Timestamp.compare(now, .lt, deadline)) return error.Timeout; + } pool_mutex.lockUncancelable(io_mod.getIo()); var locked = true; errdefer if (locked) pool_mutex.unlock(io_mod.getIo()); @@ -215,14 +232,15 @@ pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { // compatible continuation lane when it is idle; otherwise another // retained socket provides bounded parallelism without multiplexing // unrelated response events on the same wire. - const selection = selectIdleLane(args); - const index = selection.index orelse if (selection.matching_count < try maxLanes()) - try appendSlot(args) - else { - pool_mutex.unlock(io_mod.getIo()); - locked = false; - io_mod.sleep(10 * std.time.ns_per_ms); - continue; + const index = switch (chooseLane(slots.items, args, try maxLanes())) { + .existing => |existing| existing, + .append => try appendSlot(args), + .wait => { + pool_mutex.unlock(io_mod.getIo()); + locked = false; + io_mod.sleep(10 * std.time.ns_per_ms); + continue; + }, }; const slot = &slots.items[index]; @@ -457,12 +475,10 @@ test "Codex WebSocket lane selection preserves continuation affinity" { recordCompletion(second, "response-2", "{\"type\":\"message\"}", "{\"type\":\"message\"}", shape); slots.items[second].busy = false; - const selection = selectIdleLane(args); - try std.testing.expectEqual(@as(?usize, second), selection.index); - try std.testing.expectEqual(@as(usize, 2), selection.matching_count); + const selection = chooseLane(slots.items, args, 2); + try std.testing.expectEqual(second, selection.existing); slots.items[second].busy = true; - const saturated = selectIdleLane(args); - try std.testing.expectEqual(@as(?usize, null), saturated.index); - try std.testing.expectEqual(@as(usize, 2), saturated.matching_count); + try std.testing.expect(chooseLane(slots.items, args, 2) == .wait); + try std.testing.expect(chooseLane(slots.items, args, 3) == .append); } From 86f8ddd6eac114d1de9283802ed71278113724fc Mon Sep 17 00:00:00 2001 From: thinkter Date: Sat, 29 Aug 2026 19:45:49 +0530 Subject: [PATCH 19/21] Document Codex WebSocket lanes --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index da65a1f3d..e6b0c6fc8 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,8 @@ fx The OpenAI Codex route uses ChatGPT subscription access directly and never sends its OAuth token to Vercel AI Gateway. The session is stored privately at `~/.fx/chatgpt-auth.json` and refreshed when needed. On supported Codex models, `/fast` requests OpenAI's priority service tier and consumes ChatGPT credits at the higher Fast mode rate. +Codex uses HTTP streaming by default. To require the WebSocket transport, start fx with `FX_CODEX_TRANSPORT=websocket`. WebSocket sessions retain compatible connections and continuation state; concurrent requests use separate ordered lanes rather than sharing one response stream. Each session identity retains at most four lanes by default. Set `FX_CODEX_WEBSOCKET_MAX_LANES` to a positive integer to choose a different limit. + The Grok route uses subscription access directly at xAI and never sends its OAuth token to Vercel AI Gateway or OpenAI. Its session is stored privately at `~/.fx/grok-auth.json`, refreshed when needed, and used only with the authenticated xAI catalog and Responses API. To use an AI Gateway API key instead: From 18229c8f720980e836b8a873300c1d7bae733756 Mon Sep 17 00:00:00 2001 From: thinkter Date: Sun, 30 Aug 2026 01:57:56 +0530 Subject: [PATCH 20/21] Harden Codex WebSocket transport --- README.md | 2 +- docs/codex-websocket-plan.md | 128 +++++++ scripts/codex_websocket_probe.py | 391 -------------------- src/gateway/codex_websocket_session.zig | 367 ++++++++++++------ src/gateway/openai_codex.zig | 299 +++++---------- src/gateway/openai_codex_websocket.zig | 268 ++++++++++++++ src/gateway/websocket_transport.zig | 2 +- tests/e2e/tui-auth-source-selection.test.ts | 73 +++- 8 files changed, 821 insertions(+), 709 deletions(-) create mode 100644 docs/codex-websocket-plan.md delete mode 100644 scripts/codex_websocket_probe.py create mode 100644 src/gateway/openai_codex_websocket.zig diff --git a/README.md b/README.md index e6b0c6fc8..400ba6eb1 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ fx The OpenAI Codex route uses ChatGPT subscription access directly and never sends its OAuth token to Vercel AI Gateway. The session is stored privately at `~/.fx/chatgpt-auth.json` and refreshed when needed. On supported Codex models, `/fast` requests OpenAI's priority service tier and consumes ChatGPT credits at the higher Fast mode rate. -Codex uses HTTP streaming by default. To require the WebSocket transport, start fx with `FX_CODEX_TRANSPORT=websocket`. WebSocket sessions retain compatible connections and continuation state; concurrent requests use separate ordered lanes rather than sharing one response stream. Each session identity retains at most four lanes by default. Set `FX_CODEX_WEBSOCKET_MAX_LANES` to a positive integer to choose a different limit. +Codex uses HTTP streaming by default. To prefer the WebSocket transport, start fx with `FX_CODEX_TRANSPORT=websocket`. If connection setup fails before the request can be delivered, fx completes that turn over HTTP and keeps using HTTP for the rest of the process. It never replays a request whose delivery is uncertain. WebSocket sessions retain compatible connections and continuation state; concurrent requests use separate ordered lanes rather than sharing one response stream. Each session identity retains at most four lanes by default, and the process retains at most 32 lanes across identities. Set `FX_CODEX_WEBSOCKET_MAX_LANES` or `FX_CODEX_WEBSOCKET_MAX_SLOTS` to a positive integer to choose different limits. The Grok route uses subscription access directly at xAI and never sends its OAuth token to Vercel AI Gateway or OpenAI. Its session is stored privately at `~/.fx/grok-auth.json`, refreshed when needed, and used only with the authenticated xAI catalog and Responses API. diff --git a/docs/codex-websocket-plan.md b/docs/codex-websocket-plan.md new file mode 100644 index 000000000..700ece23e --- /dev/null +++ b/docs/codex-websocket-plan.md @@ -0,0 +1,128 @@ +# Codex WebSocket transport: plan to finish PR #521 + +This document replaces the earlier `docs/codex-websocket-transport.md` design +notes. It records the verified comparison between PR #521 +(`feat/codex-websocket-phase-2`, head `05bf5a4`) and PR #523 +(`perf/codex-websocket-transport`, head `f486bb2`), and the work needed to make +#521 the clear landing candidate. + +## Implementation status + +The plan is implemented on the rebased feature branch: + +- Codex WebSocket orchestration now has a separate adapter module. +- The transport uses the dated beta header and strict frame handling. +- Pre-delivery connection failures fall back to SSE and arm a process latch; + ambiguous delivery never falls back or replays. +- Retained lanes perform network I/O outside the pool mutex, use temporary + connections under saturation, and have a global LRU storage bound. +- Focused tests cover fallback latching, ambiguous-delivery no-replay, + continuation recovery, retained connection reuse, expiration, and identity + churn. + +## Verified findings + +Each claim below was checked directly against both PR heads and current `main` +(`bb2dc7d`). + +- **Stale beta header.** `src/gateway/websocket_transport.zig` sends + `OpenAI-Beta: responses_websockets=v2`. #523 sends the dated + `responses_websockets=2026-02-06`, which matches the current upstream Codex + client. +- **Network I/O under the pool mutex.** `acquire` in + `src/gateway/codex_websocket_session.zig` holds `pool_mutex` through `ping`, + `close`, and `connect`. A slow handshake on one lane stalls every other + lane. When all matching lanes are busy it also busy-polls with a 10ms sleep. +- **Unbounded slot array.** `appendSlot` grows the global slot list per new + identity. Incompatible slots get their connection closed but the slot entry + itself is never evicted. +- **Merge state.** Only `src/gateway/openai_codex.zig` conflicts with `main`, + caused by `cd1a5d3` switching error returns to + `stream_provider.failResult(...)`. +- **#523's real advantages.** A clean adapter boundary (about 41 lines touched + in `openai_codex.zig`), a process-wide SSE fallback latch, and the dated + header. +- **#523's real defects.** Binary frames are accepted as text, close frames + are not consumed or validated, there is no idle-event timeout, its cache is + one global entry (alternating sessions displace each other), and its + fallback and fresh-retry trigger is "no output yet" rather than "provably + unsent", which permits duplicate provider requests and billing after an + ambiguous delivery. + +## Strategy + +Keep #521's functional advantages (safe-prefix continuation, delivery-aware +retries, strict frame validation, retained lanes, deep test coverage). Adopt +every legitimate advantage of #523. Fix the three real defects in this branch. +After that, #523's weaknesses have no counterpart here. + +## Phase 1: rebase and restructure + +1. **Rebase onto current `main`.** Only `openai_codex.zig` conflicts. While + resolving, adapt the WebSocket error paths to the new + `stream_provider.failResult(...)` convention. +2. **Extract the added logic from `openai_codex.zig` into the adapter.** Move + `buildWebSocketRequest*`, the WebSocket consume loop, and continuation + orchestration into `codex_websocket_session.zig` (or a new + `openai_codex_websocket.zig`). Target: the provider file gains only a small + transport dispatch, comparable to #523's footprint. This removes the + integration-boundary critique and shrinks the future conflict surface. +3. **Drop `scripts/codex_websocket_probe.py`** from this PR. It is a dev + probe, not product code. Land it separately if it is worth keeping. + +## Phase 2: adopt #523's genuine wins, done better + +4. **Update the beta header** to `responses_websockets=2026-02-06`. Re-check + the current upstream constant when making the change and cite it in the + commit message. +5. **Add a delivery-aware process-wide SSE fallback latch.** Latch to SSE only + when the failure is provably pre-send (handshake or upgrade rejection, + connect timeout, pre-write errors), using the `DeliveryCertainty` already + threaded through `AcquireArgs`. An ambiguously sent request surfaces an + error instead of being silently re-issued. This keeps #523's "a broken + proxy costs one failed handshake, not one per turn" resilience without its + duplicate-billing hole. +6. **Keep strict transport-env validation.** Erroring on an invalid + `FX_CODEX_TRANSPORT` value is deliberate and better than silently coercing + to SSE. Say so in the PR description. + +## Phase 3: fix the pool's real defects + +7. **Move network I/O outside `pool_mutex`.** Restructure `acquire` to: lock, + select and reserve a slot (`busy = true`), unlock, then ping, connect, or + close outside the lock, and relock only to commit or roll back slot state. + Nothing slow ever runs under the lock. +8. **Delete the busy-wait path.** When all matching lanes are busy at the lane + limit, hand out a temporary unpooled connection instead of sleeping in 10ms + slices. Simpler, lower tail latency, and one less polling state. +9. **Bound and evict slots globally.** Add a global slot cap (for example + `max_lanes` times a small constant). Evict idle slots LRU when appending + past the cap, and fully remove, not just disconnect, slots whose identity + is incompatible. Add a unit test proving the bound holds under identity + churn. + +## Phase 4: prove it + +10. **Add two E2E tests for the new behavior:** + - a handshake failure latches the process to SSE and the turn still + completes; + - an ambiguous post-send failure does not fall back or re-send, and errors + with delivery evidence. + Both PRs added tests to `tests/e2e/tui-auth-source-selection.test.ts`; + inherit its corpus classification but confirm it still fits per AGENTS.md. +11. **Run the full ready gate:** `zig fmt --check src/`, focused Zig tests, + build, drive `./zig-out/bin/fx` with `FX_CODEX_TRANSPORT=websocket` + against the loopback E2E mock, push, and require Full CI green on all four + runners for the exact head before marking the PR ready. + +## Optional: split into a two-PR stack + +The strongest structural critique of #521 is that continuation plus pooling is +a lot for a first landing. If reviewers push back, stack the branch: + +- **PR A:** strict codec, adapter boundary, single retained connection, + delivery-aware fallback latch. A superset of #523 with none of its bugs. +- **PR B:** safe-prefix continuation and the multi-lane pool on top. + +If a single PR is preferred, Phases 1 through 4 alone make #521 dominate the +comparison on every row except raw diff size. diff --git a/scripts/codex_websocket_probe.py b/scripts/codex_websocket_probe.py deleted file mode 100644 index 98a042ad0..000000000 --- a/scripts/codex_websocket_probe.py +++ /dev/null @@ -1,391 +0,0 @@ -#!/usr/bin/env python3 -"""Isolated Phase 0 probe for the ChatGPT Codex Responses WebSocket endpoint. - -This is intentionally not part of fx's runtime transport. It uses only Python's -standard library, reads credentials only from explicitly named environment -variables, never writes credentials or response content, and prints a redacted -JSON report. - -Required for network use: - FX_CODEX_PROBE_ACCESS_TOKEN - FX_CODEX_PROBE_ACCOUNT_ID - FX_CODEX_PROBE_MODEL - -Examples: - python3 scripts/codex_websocket_probe.py - python3 scripts/codex_websocket_probe.py --execute --continuation - -The default performs only the authenticated WebSocket upgrade. --execute sends -a fixed minimal prompt and consumes subscription usage. --continuation sends a -second request using the first response ID, but does not print that ID. -""" - -from __future__ import annotations - -import argparse -import base64 -import hashlib -import json -import os -import secrets -import socket -import ssl -import sys -import time -from dataclasses import dataclass, field -from typing import Any - -HOST = "chatgpt.com" -PATH = "/backend-api/codex/responses" -ORIGIN = "https://chatgpt.com" -PROTOCOL_HEADER = "responses_websockets=v2" -CONNECT_TIMEOUT_SECONDS = 15.0 -EVENT_IDLE_TIMEOUT_SECONDS = 45.0 -MAX_FRAME_BYTES = 1 << 20 -MAX_MESSAGE_BYTES = 4 << 20 -MAX_EVENTS = 256 - - -class ProbeError(Exception): - pass - - -@dataclass -class Report: - handshake_status: int | None = None - handshake_elapsed_ms: int | None = None - selected_header_names: list[str] = field(default_factory=list) - immediate_close_code: int | None = None - event_types: list[str] = field(default_factory=list) - terminal_event: str | None = None - close_code: int | None = None - close_reason_length: int | None = None - continuation_attempted: bool = False - continuation_accepted: bool | None = None - error_code: str | None = None - error: str | None = None - - def emit(self) -> None: - print(json.dumps(self.__dict__, separators=(",", ":"), sort_keys=True)) - - -def required_env(name: str) -> str: - value = os.environ.get(name) - if not value: - raise ProbeError(f"missing required environment variable {name}") - return value - - -def websocket_accept(key: str) -> str: - digest = hashlib.sha1( - (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii") - ).digest() - return base64.b64encode(digest).decode("ascii") - - -def read_exact(sock: ssl.SSLSocket, size: int) -> bytes: - chunks: list[bytes] = [] - remaining = size - while remaining: - chunk = sock.recv(remaining) - if not chunk: - raise ProbeError("socket closed while reading frame") - chunks.append(chunk) - remaining -= len(chunk) - return b"".join(chunks) - - -def read_http_head(sock: ssl.SSLSocket) -> tuple[int, dict[str, str], bytes]: - data = bytearray() - while b"\r\n\r\n" not in data: - if len(data) >= 64 * 1024: - raise ProbeError("HTTP upgrade headers exceed local limit") - chunk = sock.recv(4096) - if not chunk: - raise ProbeError("socket closed during HTTP upgrade") - data.extend(chunk) - raw_head, remainder = bytes(data).split(b"\r\n\r\n", 1) - lines = raw_head.decode("iso-8859-1").split("\r\n") - parts = lines[0].split(" ", 2) - if len(parts) < 2 or not parts[1].isdigit(): - raise ProbeError("malformed HTTP upgrade status") - headers: dict[str, str] = {} - for line in lines[1:]: - if not line or ":" not in line: - raise ProbeError("malformed HTTP upgrade header") - name, value = line.split(":", 1) - headers[name.strip().lower()] = value.strip() - return int(parts[1]), headers, remainder - - -class WebSocket: - def __init__(self, sock: ssl.SSLSocket, buffered: bytes = b"") -> None: - self.sock = sock - self.buffered = bytearray(buffered) - - def _read_exact(self, size: int) -> bytes: - if len(self.buffered) >= size: - data = bytes(self.buffered[:size]) - del self.buffered[:size] - return data - prefix = bytes(self.buffered) - self.buffered.clear() - return prefix + read_exact(self.sock, size - len(prefix)) - - def send_text(self, value: str) -> None: - self._send_frame(0x1, value.encode("utf-8")) - - def send_pong(self, payload: bytes) -> None: - self._send_frame(0xA, payload) - - def close(self) -> None: - try: - self._send_frame(0x8, b"\x03\xe8") - except OSError: - pass - - def _send_frame(self, opcode: int, payload: bytes) -> None: - if len(payload) > MAX_MESSAGE_BYTES: - raise ProbeError("outbound frame exceeds local limit") - mask = secrets.token_bytes(4) - header = bytearray([0x80 | opcode]) - if len(payload) < 126: - header.append(0x80 | len(payload)) - elif len(payload) <= 0xFFFF: - header.append(0x80 | 126) - header.extend(len(payload).to_bytes(2, "big")) - else: - header.append(0x80 | 127) - header.extend(len(payload).to_bytes(8, "big")) - masked = bytes(byte ^ mask[index % 4] for index, byte in enumerate(payload)) - self.sock.sendall(bytes(header) + mask + masked) - - def read_message(self) -> tuple[int, bytes]: - fragments: list[bytes] = [] - initial_opcode: int | None = None - while True: - first, second = self._read_exact(2) - fin = (first & 0x80) != 0 - opcode = first & 0x0F - masked = (second & 0x80) != 0 - length = second & 0x7F - if masked: - raise ProbeError("server sent a masked WebSocket frame") - if length == 126: - length = int.from_bytes(self._read_exact(2), "big") - elif length == 127: - length = int.from_bytes(self._read_exact(8), "big") - if length & (1 << 63): - raise ProbeError("invalid WebSocket frame length") - if length > MAX_FRAME_BYTES: - raise ProbeError("inbound frame exceeds local limit") - if opcode >= 0x8 and (not fin or length > 125): - raise ProbeError("invalid WebSocket control frame") - payload = self._read_exact(length) - if opcode == 0x9: - self.send_pong(payload) - continue - if opcode == 0xA: - continue - if opcode == 0x8: - return opcode, payload - if opcode == 0x0: - if initial_opcode is None: - raise ProbeError("unexpected continuation frame") - elif opcode in (0x1, 0x2): - if initial_opcode is not None: - raise ProbeError("new data frame before fragmented message completed") - initial_opcode = opcode - else: - raise ProbeError("unsupported WebSocket opcode") - fragments.append(payload) - if sum(len(fragment) for fragment in fragments) > MAX_MESSAGE_BYTES: - raise ProbeError("reassembled WebSocket message exceeds local limit") - if fin: - return initial_opcode or opcode, b"".join(fragments) - - -def connect(token: str, account_id: str, report: Report) -> WebSocket: - key = base64.b64encode(secrets.token_bytes(16)).decode("ascii") - context = ssl.create_default_context() - started = time.monotonic() - raw = socket.create_connection((HOST, 443), CONNECT_TIMEOUT_SECONDS) - sock = context.wrap_socket(raw, server_hostname=HOST) - sock.settimeout(EVENT_IDLE_TIMEOUT_SECONDS) - request = "\r\n".join( - [ - f"GET {PATH} HTTP/1.1", - f"Host: {HOST}", - "Connection: Upgrade", - "Upgrade: websocket", - "Sec-WebSocket-Version: 13", - f"Sec-WebSocket-Key: {key}", - f"Authorization: Bearer {token}", - f"chatgpt-account-id: {account_id}", - "originator: fx-phase-0-probe", - f"OpenAI-Beta: {PROTOCOL_HEADER}", - f"Origin: {ORIGIN}", - "\r\n", - ] - ).encode("ascii") - sock.sendall(request) - status, headers, remainder = read_http_head(sock) - report.handshake_elapsed_ms = round((time.monotonic() - started) * 1000) - report.handshake_status = status - report.selected_header_names = sorted( - name - for name in headers - if name in {"openai-model", "x-codex-turn-state", "x-reasoning-included", "x-models-etag"} - ) - if status != 101: - raise ProbeError(f"WebSocket upgrade returned HTTP {status}") - if headers.get("sec-websocket-accept") != websocket_accept(key): - raise ProbeError("invalid Sec-WebSocket-Accept response") - if "upgrade" not in headers.get("connection", "").lower(): - raise ProbeError("upgrade response does not retain Connection: Upgrade") - if headers.get("upgrade", "").lower() != "websocket": - raise ProbeError("upgrade response does not select websocket") - return WebSocket(sock, remainder) - - -def safe_protocol_label(value: Any) -> str | None: - if not isinstance(value, str) or len(value) > 128: - return None - if not value.isascii() or any(not (char.isalnum() or char in "._-") for char in value): - return None - return value - - -def event_type(payload: bytes) -> tuple[str | None, dict[str, Any] | None]: - try: - value = json.loads(payload) - except (UnicodeDecodeError, json.JSONDecodeError): - return None, None - if not isinstance(value, dict): - return None, None - return safe_protocol_label(value.get("type")), value - - -def response_id(value: dict[str, Any]) -> str | None: - response = value.get("response") - if not isinstance(response, dict): - return None - identifier = response.get("id") - return identifier if isinstance(identifier, str) else None - - -def structured_error_code(value: dict[str, Any]) -> str | None: - error = value.get("error") - if not isinstance(error, dict): - return None - return safe_protocol_label(error.get("code")) - - -def wait_for_terminal(ws: WebSocket, report: Report) -> str | None: - identifier: str | None = None - for _ in range(MAX_EVENTS): - opcode, payload = ws.read_message() - if opcode == 0x8: - report.close_code = int.from_bytes(payload[:2], "big") if len(payload) >= 2 else None - report.close_reason_length = max(len(payload) - 2, 0) - raise ProbeError("server closed before terminal response event") - if opcode != 0x1: - raise ProbeError("server sent an unexpected binary message") - kind, value = event_type(payload) - if kind is None or value is None: - report.event_types.append("invalid_json") - continue - report.event_types.append(kind) - if kind == "response.created": - identifier = response_id(value) - if kind in {"response.completed", "response.done", "response.incomplete", "response.failed", "error"}: - report.terminal_event = kind - report.error_code = structured_error_code(value) - return identifier - raise ProbeError("event count exceeds local limit before terminal response") - - -def request_body(model: str, previous_response_id: str | None = None) -> dict[str, Any]: - body: dict[str, Any] = { - "type": "response.create", - "model": model, - "input": [{"role": "user", "content": [{"type": "input_text", "text": "Reply with exactly: probe"}]}], - } - if previous_response_id is not None: - body["previous_response_id"] = previous_response_id - return body - - -def run(args: argparse.Namespace) -> Report: - report = Report() - ws: WebSocket | None = None - try: - token = required_env("FX_CODEX_PROBE_ACCESS_TOKEN") - account_id = required_env("FX_CODEX_PROBE_ACCOUNT_ID") - model = required_env("FX_CODEX_PROBE_MODEL") - ws = connect(token, account_id, report) - if not args.execute: - ws.sock.settimeout(0.2) - try: - opcode, payload = ws.read_message() - if opcode == 0x8: - report.immediate_close_code = int.from_bytes(payload[:2], "big") if len(payload) >= 2 else None - except (socket.timeout, ssl.SSLWantReadError): - pass - return report - ws.send_text(json.dumps(request_body(model), separators=(",", ":"))) - first_id = wait_for_terminal(ws, report) - if args.continuation: - report.continuation_attempted = True - if not first_id or report.terminal_event != "response.completed": - report.continuation_accepted = False - return report - report.event_types = [] - report.terminal_event = None - ws.send_text(json.dumps(request_body(model, first_id), separators=(",", ":"))) - wait_for_terminal(ws, report) - report.continuation_accepted = report.terminal_event == "response.completed" - return report - except (OSError, ssl.SSLError, ProbeError) as error: - report.error = type(error).__name__ - if report.handshake_status is None: - report.handshake_status = 0 - return report - finally: - if ws is not None: - ws.close() - ws.sock.close() - - -def self_test() -> None: - key = "dGhlIHNhbXBsZSBub25jZQ==" - assert websocket_accept(key) == "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" - body = request_body("model", "response") - assert body["previous_response_id"] == "response" - assert "stream" not in body and "background" not in body - code, value = event_type(b'{"type":"response.created","response":{"id":"r"}}') - assert code == "response.created" and response_id(value or {}) == "r" - assert safe_protocol_label("response.completed") == "response.completed" - assert safe_protocol_label("prompt content") is None - print("codex websocket probe self-test passed") - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Run the isolated fx Codex WebSocket Phase 0 probe") - parser.add_argument("--execute", action="store_true", help="send a fixed minimal model request after the upgrade") - parser.add_argument("--continuation", action="store_true", help="test previous_response_id after --execute") - parser.add_argument("--self-test", action="store_true", help="run deterministic local checks without credentials or network") - args = parser.parse_args() - if args.continuation and not args.execute: - parser.error("--continuation requires --execute") - return args - - -if __name__ == "__main__": - arguments = parse_args() - if arguments.self_test: - self_test() - else: - result = run(arguments) - result.emit() - sys.exit(0 if result.error is None else 1) diff --git a/src/gateway/codex_websocket_session.zig b/src/gateway/codex_websocket_session.zig index 0328e1b16..8a17d2c42 100644 --- a/src/gateway/codex_websocket_session.zig +++ b/src/gateway/codex_websocket_session.zig @@ -9,8 +9,10 @@ const pool_alloc = std.heap.c_allocator; pub const health_budget: u8 = 3; pub const default_max_connection_age_ms: i64 = 55 * 60 * 1000; const default_max_lanes: usize = 4; +const default_max_slots: usize = 32; const max_connection_age_env = "FX_CODEX_WEBSOCKET_MAX_CONNECTION_AGE_MS"; const max_lanes_env = "FX_CODEX_WEBSOCKET_MAX_LANES"; +const max_slots_env = "FX_CODEX_WEBSOCKET_MAX_SLOTS"; const Slot = struct { session_id: []u8, @@ -22,6 +24,7 @@ const Slot = struct { busy: bool, health_failures: u8, opened_at_ms: i64, + last_used_at_ms: i64, continuation_response_id: ?[]u8, continuation_baseline: ?[]u8, continuation_durable_baseline: ?[]u8, @@ -66,9 +69,10 @@ pub const AcquireArgs = struct { }; pub const Checkout = struct { - slot: usize, + slot: ?usize, connection: *websocket_transport.Connection, reused: bool, + retained: bool, handshake_ms: i64, health_failures: u8, }; @@ -121,16 +125,14 @@ fn maxLanes() !usize { return parsed; } -fn incompatibleIdentity(slot: *const Slot, args: AcquireArgs) bool { - const fingerprint = authorizationFingerprint(args.authorization); - return std.mem.eql(u8, slot.session_id, sessionKey(args.session_id)) and - (!std.mem.eql(u8, slot.account_id, args.account_id) or - !std.mem.eql(u8, slot.model, args.model) or - !std.mem.eql(u8, slot.endpoint, args.endpoint) or - !std.mem.eql(u8, &slot.authorization_fingerprint, &fingerprint)); +fn maxSlots() !usize { + const value = io_mod.getenv(max_slots_env) orelse return default_max_slots; + const parsed = std.fmt.parseInt(usize, value, 10) catch return error.InvalidOpenAICodexTransport; + if (parsed == 0) return error.InvalidOpenAICodexTransport; + return parsed; } -fn appendSlot(args: AcquireArgs) !usize { +fn initSlot(args: AcquireArgs, busy: bool) !Slot { const session_id = try pool_alloc.dupe(u8, sessionKey(args.session_id)); errdefer pool_alloc.free(session_id); const account_id = try pool_alloc.dupe(u8, args.account_id); @@ -139,22 +141,27 @@ fn appendSlot(args: AcquireArgs) !usize { errdefer pool_alloc.free(model); const endpoint = try pool_alloc.dupe(u8, args.endpoint); errdefer pool_alloc.free(endpoint); - try slots.append(pool_alloc, .{ + return .{ .session_id = session_id, .account_id = account_id, .model = model, .endpoint = endpoint, .authorization_fingerprint = authorizationFingerprint(args.authorization), .connection = null, - .busy = false, + .busy = busy, .health_failures = 0, .opened_at_ms = 0, + .last_used_at_ms = io_mod.milliTimestamp(), .continuation_response_id = null, .continuation_baseline = null, .continuation_durable_baseline = null, .continuation_shape = undefined, .continuation_valid = false, - }); + }; +} + +fn appendSlot(args: AcquireArgs, busy: bool) !usize { + try slots.append(pool_alloc, try initSlot(args, busy)); return slots.items.len - 1; } @@ -198,120 +205,189 @@ fn selectIdleLane(slot_items: []Slot, args: AcquireArgs) LaneSelection { const LaneChoice = union(enum) { existing: usize, append, - wait, + temporary, }; fn chooseLane(slot_items: []Slot, args: AcquireArgs, lane_limit: usize) LaneChoice { const selection = selectIdleLane(slot_items, args); if (selection.index) |index| return .{ .existing = index }; if (selection.matching_count < lane_limit) return .append; - return .wait; + return .temporary; } fn incrementFailure(slot: *Slot) void { slot.health_failures = std.math.add(u8, slot.health_failures, 1) catch std.math.maxInt(u8); } -pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { - while (true) { - if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; - if (args.deadline) |deadline| { - const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); - if (!std.Io.Clock.Timestamp.compare(now, .lt, deadline)) return error.Timeout; - } - pool_mutex.lockUncancelable(io_mod.getIo()); - var locked = true; - errdefer if (locked) pool_mutex.unlock(io_mod.getIo()); - for (slots.items) |*existing| { - if (!incompatibleIdentity(existing, args) or existing.busy) continue; - if (existing.connection) |connection| websocket_transport.close(connection, pool_alloc); - existing.connection = null; - existing.clearContinuation(); +fn leastRecentlyUsedIdle(slot_items: []Slot) ?usize { + var selected: ?usize = null; + for (slot_items, 0..) |*slot, index| { + if (slot.busy) continue; + if (selected == null or slot.last_used_at_ms < slot_items[selected.?].last_used_at_ms) { + selected = index; } - // Responses on one WebSocket are exclusive and ordered. Preserve a - // compatible continuation lane when it is idle; otherwise another - // retained socket provides bounded parallelism without multiplexing - // unrelated response events on the same wire. - const index = switch (chooseLane(slots.items, args, try maxLanes())) { - .existing => |existing| existing, - .append => try appendSlot(args), - .wait => { - pool_mutex.unlock(io_mod.getIo()); - locked = false; - io_mod.sleep(10 * std.time.ns_per_ms); - continue; - }, - }; - const slot = &slots.items[index]; + } + return selected; +} - const age_limit = try maxConnectionAgeMs(); - if (slot.connection != null and slot.health_failures >= health_budget) { - websocket_transport.close(slot.connection.?, pool_alloc); - slot.connection = null; - slot.clearContinuation(); - } - if (slot.connection != null and age_limit != 0 and io_mod.milliTimestamp() - slot.opened_at_ms > age_limit) { - websocket_transport.close(slot.connection.?, pool_alloc); - slot.connection = null; - slot.clearContinuation(); - } - if (slot.connection) |connection| { - const ping_result = websocket_transport.ping(connection, args.cancel_flag, args.deadline, args.delivery); - if (ping_result) |_| { - slot.busy = true; - const checkout = Checkout{ - .slot = index, - .connection = connection, - .reused = true, - .handshake_ms = 0, - .health_failures = slot.health_failures, - }; - pool_mutex.unlock(io_mod.getIo()); - return checkout; - } else |err| { - websocket_transport.close(connection, pool_alloc); +fn incompatibleIdle(slot_items: []Slot, args: AcquireArgs) ?usize { + for (slot_items, 0..) |*slot, index| { + if (slot.busy or matches(slot, args)) continue; + if (std.mem.eql(u8, slot.session_id, sessionKey(args.session_id))) return index; + } + return null; +} + +fn replaceSlot(index: usize, args: AcquireArgs) !?*websocket_transport.Connection { + const replacement = try initSlot(args, true); + const displaced = slots.items[index].connection; + slots.items[index].connection = null; + slots.items[index].deinit(); + slots.items[index] = replacement; + return displaced; +} + +pub fn acquire(_: Allocator, args: AcquireArgs) !Checkout { + if (args.cancel_flag.load(.seq_cst)) return error.Cancelled; + if (args.deadline) |deadline| { + const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + if (!std.Io.Clock.Timestamp.compare(now, .lt, deadline)) return error.Timeout; + } + + const lane_limit = try maxLanes(); + const slot_limit = try maxSlots(); + const age_limit = try maxConnectionAgeMs(); + var index: ?usize = null; + var retained = true; + var reusable: ?*websocket_transport.Connection = null; + var displaced: ?*websocket_transport.Connection = null; + var prior_health: u8 = 0; + + pool_mutex.lockUncancelable(io_mod.getIo()); + switch (chooseLane(slots.items, args, lane_limit)) { + .existing => |existing| { + index = existing; + const slot = &slots.items[existing]; + slot.busy = true; + prior_health = slot.health_failures; + const expired = age_limit != 0 and + io_mod.milliTimestamp() - slot.opened_at_ms > age_limit; + if (slot.connection != null and slot.health_failures < health_budget and !expired) { + reusable = slot.connection; + } else { + displaced = slot.connection; slot.connection = null; slot.clearContinuation(); + } + }, + .append => { + if (incompatibleIdle(slots.items, args)) |victim| { + index = victim; + displaced = replaceSlot(victim, args) catch |err| { + pool_mutex.unlock(io_mod.getIo()); + return err; + }; + } else if (slots.items.len < slot_limit) { + index = appendSlot(args, true) catch |err| { + pool_mutex.unlock(io_mod.getIo()); + return err; + }; + } else if (leastRecentlyUsedIdle(slots.items)) |victim| { + index = victim; + displaced = replaceSlot(victim, args) catch |err| { + pool_mutex.unlock(io_mod.getIo()); + return err; + }; + } else { + retained = false; + } + }, + .temporary => retained = false, + } + pool_mutex.unlock(io_mod.getIo()); + + // Socket close, health checks, and connection establishment are all + // deliberately outside the global pool mutex. + if (displaced) |connection| websocket_transport.close(connection, pool_alloc); + if (reusable) |connection| { + websocket_transport.ping(connection, args.cancel_flag, args.deadline, args.delivery) catch |err| { + websocket_transport.close(connection, pool_alloc); + pool_mutex.lockUncancelable(io_mod.getIo()); + if (index) |slot_index| { + const slot = &slots.items[slot_index]; + if (slot.connection == connection) slot.connection = null; + slot.clearContinuation(); incrementFailure(slot); - if (err == error.Cancelled) return err; + prior_health = slot.health_failures; } + pool_mutex.unlock(io_mod.getIo()); + if (err == error.Cancelled) { + rollbackReservation(index); + return err; + } + reusable = null; + }; + if (reusable != null) { + return .{ + .slot = index, + .connection = connection, + .reused = true, + .retained = retained, + .handshake_ms = 0, + .health_failures = prior_health, + }; } + } - const started_at_ms = io_mod.milliTimestamp(); - const connection = websocket_transport.connect(pool_alloc, .{ - .endpoint = args.endpoint, - .authorization = args.authorization, - .account_id = args.account_id, - .session_id = args.session_id, - .deadline = args.deadline, - .cancel_flag = args.cancel_flag, - .delivery = args.delivery, - }) catch |err| return err; + const started_at_ms = io_mod.milliTimestamp(); + const connection = websocket_transport.connect(pool_alloc, .{ + .endpoint = args.endpoint, + .authorization = args.authorization, + .account_id = args.account_id, + .session_id = args.session_id, + .deadline = args.deadline, + .cancel_flag = args.cancel_flag, + .delivery = args.delivery, + }) catch |err| { + rollbackReservation(index); + return err; + }; + if (index) |slot_index| { + pool_mutex.lockUncancelable(io_mod.getIo()); + const slot = &slots.items[slot_index]; slot.connection = connection; slot.clearContinuation(); slot.opened_at_ms = connection.opened_at_ms; - slot.busy = true; - const checkout = Checkout{ - .slot = index, - .connection = connection, - .reused = false, - .handshake_ms = @max(io_mod.milliTimestamp() - started_at_ms, 0), - .health_failures = slot.health_failures, - }; + slot.last_used_at_ms = io_mod.milliTimestamp(); pool_mutex.unlock(io_mod.getIo()); - return checkout; } + return .{ + .slot = index, + .connection = connection, + .reused = false, + .retained = retained, + .handshake_ms = @max(io_mod.milliTimestamp() - started_at_ms, 0), + .health_failures = prior_health, + }; +} + +fn rollbackReservation(index: ?usize) void { + const slot_index = index orelse return; + pool_mutex.lockUncancelable(io_mod.getIo()); + if (slot_index < slots.items.len) slots.items[slot_index].busy = false; + pool_mutex.unlock(io_mod.getIo()); } pub fn continuation( - index: usize, + index: ?usize, full_input: []const u8, shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, ) ?Continuation { + const slot_index = index orelse return null; pool_mutex.lockUncancelable(io_mod.getIo()); defer pool_mutex.unlock(io_mod.getIo()); - if (index >= slots.items.len) return null; - const slot = &slots.items[index]; + if (slot_index >= slots.items.len) return null; + const slot = &slots.items[slot_index]; if (!slot.busy or !slot.continuation_valid) return null; if (!std.mem.eql(u8, &slot.continuation_shape, &shape)) { slot.clearContinuation(); @@ -336,16 +412,17 @@ pub fn continuation( } pub fn recordCompletion( - index: usize, + index: ?usize, response_id: []const u8, baseline: []const u8, durable_baseline: []const u8, shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, ) void { + const slot_index = index orelse return; pool_mutex.lockUncancelable(io_mod.getIo()); defer pool_mutex.unlock(io_mod.getIo()); - if (index >= slots.items.len) return; - const slot = &slots.items[index]; + if (slot_index >= slots.items.len) return; + const slot = &slots.items[slot_index]; slot.clearContinuation(); const owned_id = pool_alloc.dupe(u8, response_id) catch return; const owned_baseline = pool_alloc.dupe(u8, baseline) catch { @@ -364,29 +441,42 @@ pub fn recordCompletion( slot.continuation_valid = true; } -pub fn release(index: usize, outcome: Outcome) void { +pub fn release(checkout: Checkout, outcome: Outcome) void { + const index = checkout.slot orelse { + websocket_transport.close(checkout.connection, pool_alloc); + return; + }; + var discarded: ?*websocket_transport.Connection = null; pool_mutex.lockUncancelable(io_mod.getIo()); - defer pool_mutex.unlock(io_mod.getIo()); - if (index >= slots.items.len) return; + if (index >= slots.items.len) { + pool_mutex.unlock(io_mod.getIo()); + websocket_transport.close(checkout.connection, pool_alloc); + return; + } const slot = &slots.items[index]; slot.busy = false; + slot.last_used_at_ms = io_mod.milliTimestamp(); switch (outcome) { .completed => slot.health_failures = 0, .failed => { incrementFailure(slot); - if (slot.connection) |connection| websocket_transport.close(connection, pool_alloc); + discarded = slot.connection; slot.connection = null; slot.clearContinuation(); }, } + pool_mutex.unlock(io_mod.getIo()); + if (discarded) |connection| websocket_transport.close(connection, pool_alloc); } pub fn shutdown() void { pool_mutex.lockUncancelable(io_mod.getIo()); - defer pool_mutex.unlock(io_mod.getIo()); - for (slots.items) |*slot| slot.deinit(); - slots.deinit(pool_alloc); + const retired = slots; slots = .empty; + pool_mutex.unlock(io_mod.getIo()); + var owned = retired; + for (owned.items) |*slot| slot.deinit(); + owned.deinit(pool_alloc); } test "retained Codex WebSocket identity includes authorization without storing it" { @@ -400,6 +490,7 @@ test "retained Codex WebSocket identity includes authorization without storing i .busy = false, .health_failures = 0, .opened_at_ms = 0, + .last_used_at_ms = 0, .continuation_response_id = null, .continuation_baseline = null, .continuation_durable_baseline = null, @@ -418,17 +509,14 @@ test "retained Codex WebSocket identity includes authorization without storing i }; try std.testing.expect(matches(&slot, base)); - try std.testing.expect(!incompatibleIdentity(&slot, base)); var rotated = base; rotated.authorization = "Bearer token-b"; try std.testing.expect(!matches(&slot, rotated)); - try std.testing.expect(incompatibleIdentity(&slot, rotated)); var changed_model = base; changed_model.model = "gpt-5.4"; try std.testing.expect(!matches(&slot, changed_model)); - try std.testing.expect(incompatibleIdentity(&slot, changed_model)); } test "Codex WebSocket continuation requires an exact item boundary prefix" { @@ -468,9 +556,9 @@ test "Codex WebSocket lane selection preserves continuation affinity" { .continuation_shape = shape, }; - const first = try appendSlot(args); + const first = try appendSlot(args, false); slots.items[first].busy = true; - const second = try appendSlot(args); + const second = try appendSlot(args, false); slots.items[second].busy = true; recordCompletion(second, "response-2", "{\"type\":\"message\"}", "{\"type\":\"message\"}", shape); slots.items[second].busy = false; @@ -479,6 +567,67 @@ test "Codex WebSocket lane selection preserves continuation affinity" { try std.testing.expectEqual(second, selection.existing); slots.items[second].busy = true; - try std.testing.expect(chooseLane(slots.items, args, 2) == .wait); + try std.testing.expect(chooseLane(slots.items, args, 2) == .temporary); try std.testing.expect(chooseLane(slots.items, args, 3) == .append); } + +test "Codex WebSocket global slot eviction selects the least recently used idle lane" { + shutdown(); + defer shutdown(); + + var cancel_flag = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + const args = AcquireArgs{ + .session_id = "session-a", + .account_id = "account-a", + .model = "gpt-5.6-sol", + .endpoint = "http://127.0.0.1/responses", + .authorization = "Bearer token-a", + .deadline = null, + .cancel_flag = &cancel_flag, + .delivery = &delivery, + }; + _ = try appendSlot(args, false); + _ = try appendSlot(args, false); + _ = try appendSlot(args, false); + slots.items[0].last_used_at_ms = 30; + slots.items[1].last_used_at_ms = 10; + slots.items[2].last_used_at_ms = 20; + slots.items[1].busy = true; + + try std.testing.expectEqual(@as(?usize, 2), leastRecentlyUsedIdle(slots.items)); + try std.testing.expectEqual(@as(usize, 3), slots.items.len); +} + +test "Codex WebSocket slot storage remains bounded under identity churn" { + shutdown(); + defer shutdown(); + + var cancel_flag = std.atomic.Value(bool).init(false); + var delivery = gateway_client.DeliveryCertainty.init(); + var session_buffer: [32]u8 = undefined; + var args = AcquireArgs{ + .session_id = "", + .account_id = "account-a", + .model = "gpt-5.6-sol", + .endpoint = "http://127.0.0.1/responses", + .authorization = "Bearer token-a", + .deadline = null, + .cancel_flag = &cancel_flag, + .delivery = &delivery, + }; + const limit: usize = 3; + for (0..64) |identity| { + args.session_id = try std.fmt.bufPrint(&session_buffer, "session-{d}", .{identity}); + if (slots.items.len < limit) { + _ = try appendSlot(args, false); + } else { + const victim = leastRecentlyUsedIdle(slots.items).?; + _ = try replaceSlot(victim, args); + slots.items[victim].busy = false; + slots.items[victim].last_used_at_ms = @intCast(identity); + } + try std.testing.expect(slots.items.len <= limit); + } + try std.testing.expectEqual(limit, slots.items.len); +} diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index f794f56c7..fedf1e211 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -7,8 +7,7 @@ const io_mod = @import("../core/shared/io.zig"); const types = @import("../core/shared/types.zig"); const gateway_client = @import("client.zig"); const responses_protocol = @import("responses_protocol.zig"); -const websocket_transport = @import("websocket_transport.zig"); -const codex_websocket_session = @import("codex_websocket_session.zig"); +const codex_websocket = @import("openai_codex_websocket.zig"); const debug_trace = @import("../core/shared/debug_trace.zig"); const model_tool_schema = @import("../core/tooling/model_tool_schema.zig"); @@ -28,14 +27,39 @@ const connect_timeout_ms: i64 = 30_000; const transport_env = "FX_CODEX_TRANSPORT"; const Transport = enum { sse, websocket }; +var sse_fallback_active = std.atomic.Value(bool).init(false); fn selectedTransport() !Transport { const value = io_mod.getenv(transport_env) orelse return .sse; if (std.mem.eql(u8, value, "sse") or std.mem.eql(u8, value, "auto")) return .sse; - if (std.mem.eql(u8, value, "websocket")) return .websocket; + if (std.mem.eql(u8, value, "websocket")) { + return if (sse_fallback_active.load(.seq_cst)) .sse else .websocket; + } return error.InvalidOpenAICodexTransport; } +fn allowsSseFallback(err: anyerror, delivery: gateway_client.DeliveryCertainty.State) bool { + if (delivery != .definitely_unsent) return false; + return switch (err) { + error.Cancelled, + error.OutOfMemory, + error.ProviderAdmissionMissing, + error.ProviderAdmissionRepeated, + error.InvalidOpenAICodexTransport, + => false, + else => true, + }; +} + +fn armSseFallback(err: anyerror) void { + sse_fallback_active.store(true, .seq_cst); + debug_trace.logf( + "stream", + "Codex WebSocket transport disabled for this process error={s}", + .{@errorName(err)}, + ); +} + const CodexLimits = struct { aggregate_bytes: usize = max_sse_aggregate_bytes, events: usize = max_sse_events, @@ -61,7 +85,7 @@ pub const agent_stream_provider = stream_provider.Provider{ }; pub fn shutdownWebSockets() void { - codex_websocket_session.shutdown(); + codex_websocket.shutdown(); } fn validateModel(model: []const u8) !void { @@ -203,21 +227,6 @@ fn buildResponseInput( return out.toOwnedSlice(); } -fn buildContinuationBaseline( - alloc: Allocator, - full_input: []const u8, - response_input: []const u8, -) ![]u8 { - var out: std.Io.Writer.Allocating = .init(alloc); - errdefer out.deinit(); - try out.writer.writeAll(full_input); - if (response_input.len > 0) { - if (out.written().len > 0) try out.writer.writeByte(','); - try out.writer.writeAll(response_input); - } - return out.toOwnedSlice(); -} - fn buildDurableResponseInput( alloc: Allocator, messages: []const types.ChatMessage, @@ -258,16 +267,27 @@ fn streamCompletion( return stream_provider.failResult(error.CodexSubscriptionCredentialRequired); } try validateModel(request.model); - return switch (try selectedTransport()) { - .sse => blk: { - const payload = try buildRequest(alloc, request.data()); - defer alloc.free(payload); - break :blk streamPrepared(alloc, request, payload); - }, - .websocket => blk: { - break :blk streamWebSocketPrepared(alloc, request); - }, - } catch |err| { + const transport = try selectedTransport(); + if (transport == .websocket) { + if (streamWebSocketPrepared(alloc, request)) |result| return result else |err| { + if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); + const delivery = request.delivery.load(); + if (!allowsSseFallback(err, delivery)) { + request.attempt_evidence.network_failure = gateway_client.networkFailureEvidence(err, delivery); + return err; + } + armSseFallback(err); + } + } + + const payload = try buildRequest(alloc, request.data()); + defer alloc.free(payload); + return streamPreparedWithAdmission( + alloc, + request, + payload, + !request.attempt_evidence.provider_admitted, + ) catch |err| { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); request.attempt_evidence.network_failure = gateway_client.networkFailureEvidence(err, request.delivery.load()); return err; @@ -318,6 +338,15 @@ pub fn streamPrepared( alloc: Allocator, request: stream_provider.ModelRequest, payload: []const u8, +) !stream_provider.Result { + return streamPreparedWithAdmission(alloc, request, payload, true); +} + +fn streamPreparedWithAdmission( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, + should_admit: bool, ) !stream_provider.Result { if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); var prepared = try prepareCodexTransport(alloc, request); @@ -353,7 +382,7 @@ pub fn streamPrepared( .clock = .awake, .raw = .fromMilliseconds(connect_timeout_ms), }); - try admitCodexTransport(request.admission); + if (should_admit) try admitCodexTransport(request.admission); var opened = try gateway_client.runBoundedHttpOperation( OpenedRequest, alloc, @@ -482,172 +511,23 @@ fn streamWebSocketPrepared( if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; var prepared = try prepareCodexTransport(alloc, request); defer prepared.deinit(alloc); - - const full_input = try buildResponseInput(alloc, request.messages, request.verified_images); - defer alloc.free(full_input); - const shape_payload = try buildWebSocketRequestWithInput(alloc, request.data(), "", null); - defer alloc.free(shape_payload); - var shape: [std.crypto.hash.sha2.Sha256.digest_length]u8 = undefined; - std.crypto.hash.sha2.Sha256.hash(shape_payload, &shape, .{}); - - var reducer = responses_protocol.Reducer.init(alloc); - defer reducer.deinit(alloc); - var bridge = WebSocketBridge{ - .alloc = alloc, - .reducer = &reducer, - .events = request.events, - .cancel_flag = request.cancel_flag, - .content_capture_limit = request.content_capture_limit, - .stream_limits = codexStreamLimits(.{}), - }; - try admitCodexTransport(request.admission); - - var acquisition_attempt: u8 = 0; - var continuation_recovery_attempted = false; - while (true) { - const checkout = codex_websocket_session.acquire(alloc, .{ - .session_id = request.session_id, - .account_id = prepared.account_id, - .model = request.model, - .endpoint = prepared.endpoint, - .authorization = prepared.authorization, - .deadline = request.deadline, - .cancel_flag = request.cancel_flag, - .delivery = request.delivery, - .continuation_input = if (continuation_recovery_attempted) null else full_input, - .continuation_shape = if (continuation_recovery_attempted) null else shape, - }) catch |err| { - if (acquisition_attempt == 0 and request.delivery.load() == .definitely_unsent) { - acquisition_attempt += 1; - continue; - } - return err; - }; - debug_trace.eventf("codex.ws", "turn", request.trace_ctx, "lane={d} reused={d} handshake_ms={d} health={d} auth=chatgpt_subscription", .{ - checkout.slot, - @as(u8, @intFromBool(checkout.reused)), - checkout.handshake_ms, - checkout.health_failures, - }); - const continued = if (continuation_recovery_attempted) - null - else - codex_websocket_session.continuation(checkout.slot, full_input, shape); - const payload = try buildWebSocketRequestWithInput( - alloc, - request.data(), - if (continued) |value| value.delta_input else full_input, - if (continued) |value| value.previous_response_id else null, - ); - defer alloc.free(payload); - debug_trace.eventf("codex.ws", "continuation", request.trace_ctx, "used={d} delta_bytes={d} recovery={d}", .{ - @as(u8, @intFromBool(continued != null)), - if (continued) |value| value.delta_input.len else full_input.len, - @as(u8, @intFromBool(continuation_recovery_attempted)), - }); - websocket_transport.streamOn(checkout.connection, alloc, .{ + return codex_websocket.stream( + alloc, + request, + .{ .endpoint = prepared.endpoint, .authorization = prepared.authorization, .account_id = prepared.account_id, - .session_id = request.session_id, - .payload = payload, - .deadline = request.deadline, - .cancel_flag = request.cancel_flag, - .delivery = request.delivery, - }, &bridge, WebSocketBridge.event) catch |err| { - codex_websocket_session.release(checkout.slot, .failed); - if (err == error.PreviousResponseNotFound and continued != null and !continuation_recovery_attempted) { - continuation_recovery_attempted = true; - reducer.deinit(alloc); - reducer = responses_protocol.Reducer.init(alloc); - continue; - } - debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ websocketFailureReason(err), @as(u16, 0) }); - return err; - }; - const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| { - codex_websocket_session.release(checkout.slot, .failed); - debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ "protocol", @as(u16, 0) }); - return mapReducerError(err); - }; - if (completion.generation_id) |response_id| { - const response_message = [_]types.ChatMessage{.{ - .role = .assistant, - .content = completion.content, - .tool_calls = completion.tool_calls, - .provider_state_json = completion.provider_state_json, - }}; - if (buildResponseInput(alloc, &response_message, null)) |response_input| { - defer alloc.free(response_input); - if (buildContinuationBaseline(alloc, full_input, response_input)) |baseline| { - defer alloc.free(baseline); - if (buildDurableResponseInput(alloc, request.messages)) |durable_full_input| { - defer alloc.free(durable_full_input); - if (buildDurableResponseInput(alloc, &response_message)) |durable_response_input| { - defer alloc.free(durable_response_input); - if (buildContinuationBaseline(alloc, durable_full_input, durable_response_input)) |durable_baseline| { - defer alloc.free(durable_baseline); - codex_websocket_session.recordCompletion( - checkout.slot, - response_id, - baseline, - durable_baseline, - shape, - ); - } else |_| {} - } else |_| {} - } else |_| {} - } else |_| {} - } else |_| {} - } - codex_websocket_session.release(checkout.slot, .completed); - return .{ .completed = .{ - .completion = completion, - .usage = .{ .unavailable = .possibly_billed }, - .ownership = .owned, - } }; - } -} - -fn websocketFailureReason(err: anyerror) []const u8 { - return switch (err) { - error.Cancelled => "cancel", - error.Timeout => "timeout", - error.WebSocketPolicyClosed => "policy", - error.WebSocketUnexpectedBinary => "binary", - error.WebSocketProtocolViolation, error.WebSocketInvalidUtf8 => "protocol", - error.WebSocketUpgradeRejected, error.WebSocketAcceptInvalid => "auth", - else => "close", - }; + }, + .{ + .build_input = buildResponseInput, + .build_request = buildWebSocketRequestWithInput, + .build_durable_input = buildDurableResponseInput, + }, + codexStreamLimits(.{}), + ); } -const WebSocketBridge = struct { - alloc: Allocator, - reducer: *responses_protocol.Reducer, - events: stream_provider.EventSink, - cancel_flag: *std.atomic.Value(bool), - content_capture_limit: ?usize, - stream_limits: responses_protocol.StreamLimits, - - fn event(raw: *anyopaque, json_text: []const u8) !bool { - const self: *@This() = @ptrCast(@alignCast(raw)); - return self.reducer.applyJson( - self.alloc, - json_text, - .{ - .context = &self.events, - .on_content = EventBridge.content, - .on_tool_start = EventBridge.toolStart, - .on_reasoning = EventBridge.reasoning, - .on_tool_input = EventBridge.toolInput, - }, - self.cancel_flag, - self.content_capture_limit, - self.stream_limits, - ) catch |err| return mapReducerError(err); - } -}; - const EventBridge = struct { fn sink(raw: *anyopaque) *stream_provider.EventSink { return @ptrCast(@alignCast(raw)); @@ -833,14 +713,20 @@ test "OpenAI Codex transport admission invokes the shared admission boundary" { try std.testing.expect(capture.called); } -test "OpenAI Codex transport policy keeps auto on SSE during Phase 1" { +test "OpenAI Codex transport policy keeps auto on SSE" { // Environment-dependent selection is covered by integration launch tests. - // This assertion records the Phase 1 default when no override is present. if (io_mod.getenv(transport_env) == null) { try std.testing.expectEqual(Transport.sse, try selectedTransport()); } } +test "OpenAI Codex SSE fallback requires definitely unsent delivery" { + try std.testing.expect(allowsSseFallback(error.WebSocketUpgradeRejected, .definitely_unsent)); + try std.testing.expect(!allowsSseFallback(error.WebSocketUpgradeRejected, .possibly_sent)); + try std.testing.expect(!allowsSseFallback(error.Cancelled, .definitely_unsent)); + try std.testing.expect(!allowsSseFallback(error.OutOfMemory, .definitely_unsent)); +} + test "OpenAI Codex request uses Responses input and converts AI SDK tool schemas" { const read_file_schema = model_tool_schema.FunctionSchema{ .name = "read_file", @@ -1192,21 +1078,28 @@ test "OpenAI Codex SSE and WebSocket reducers preserve callback order and comple var websocket_cancelled = std.atomic.Value(bool).init(false); var websocket_reducer = responses_protocol.Reducer.init(std.testing.allocator); defer websocket_reducer.deinit(std.testing.allocator); - var bridge = WebSocketBridge{ - .alloc = std.testing.allocator, - .reducer = &websocket_reducer, - .events = websocket_events, - .cancel_flag = &websocket_cancelled, - .content_capture_limit = null, - .stream_limits = codexStreamLimits(.{}), - }; + var mutable_websocket_events = websocket_events; + const websocket_limits = codexStreamLimits(.{}); for (raw_events) |raw_event| { - if (try WebSocketBridge.event(&bridge, raw_event)) break; + if (try websocket_reducer.applyJson( + std.testing.allocator, + raw_event, + .{ + .context = &mutable_websocket_events, + .on_content = EventBridge.content, + .on_tool_start = EventBridge.toolStart, + .on_reasoning = EventBridge.reasoning, + .on_tool_input = EventBridge.toolInput, + }, + &websocket_cancelled, + null, + websocket_limits, + )) break; } const websocket_completion = try websocket_reducer.finish( std.testing.allocator, &websocket_cancelled, - bridge.stream_limits, + websocket_limits, ); defer freeOpenAICodexTestCompletion(websocket_completion); diff --git a/src/gateway/openai_codex_websocket.zig b/src/gateway/openai_codex_websocket.zig new file mode 100644 index 000000000..e63a9920e --- /dev/null +++ b/src/gateway/openai_codex_websocket.zig @@ -0,0 +1,268 @@ +//! Codex-specific WebSocket adapter. +//! +//! This module owns connection reuse, continuation, event reduction, and the +//! WebSocket request lifecycle. The base provider supplies only the shared +//! Responses request serializers and connection credentials. + +const std = @import("std"); +const image_attachments = @import("../core/images/image_attachments.zig"); +const stream_provider = @import("../core/agent/stream_provider.zig"); +const types = @import("../core/shared/types.zig"); +const debug_trace = @import("../core/shared/debug_trace.zig"); +const responses_protocol = @import("responses_protocol.zig"); +const websocket_transport = @import("websocket_transport.zig"); +const codex_websocket_session = @import("codex_websocket_session.zig"); + +const Allocator = std.mem.Allocator; + +pub const ConnectionConfig = struct { + endpoint: []const u8, + authorization: []const u8, + account_id: []const u8, +}; + +pub const Serializer = struct { + build_input: *const fn ( + alloc: Allocator, + messages: []const types.ChatMessage, + images: ?[]const image_attachments.VerifiedSnapshot, + ) anyerror![]u8, + build_request: *const fn ( + alloc: Allocator, + request: stream_provider.RequestData, + input: []const u8, + previous_response_id: ?[]const u8, + ) anyerror![]u8, + build_durable_input: *const fn ( + alloc: Allocator, + messages: []const types.ChatMessage, + ) anyerror![]u8, +}; + +pub fn shutdown() void { + codex_websocket_session.shutdown(); +} + +pub fn stream( + alloc: Allocator, + request: stream_provider.ModelRequest, + config: ConnectionConfig, + serializer: Serializer, + stream_limits: responses_protocol.StreamLimits, +) !stream_provider.Result { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + + const full_input = try serializer.build_input(alloc, request.messages, request.verified_images); + defer alloc.free(full_input); + const shape_payload = try serializer.build_request(alloc, request.data(), "", null); + defer alloc.free(shape_payload); + var shape: [std.crypto.hash.sha2.Sha256.digest_length]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(shape_payload, &shape, .{}); + + var reducer = responses_protocol.Reducer.init(alloc); + defer reducer.deinit(alloc); + var bridge = WebSocketBridge{ + .alloc = alloc, + .reducer = &reducer, + .events = request.events, + .cancel_flag = request.cancel_flag, + .content_capture_limit = request.content_capture_limit, + .stream_limits = stream_limits, + }; + try request.admission.admit(); + + var continuation_recovery_attempted = false; + while (true) { + const checkout = try codex_websocket_session.acquire(alloc, .{ + .session_id = request.session_id, + .account_id = config.account_id, + .model = request.model, + .endpoint = config.endpoint, + .authorization = config.authorization, + .deadline = request.deadline, + .cancel_flag = request.cancel_flag, + .delivery = request.delivery, + .continuation_input = if (continuation_recovery_attempted) null else full_input, + .continuation_shape = if (continuation_recovery_attempted) null else shape, + }); + debug_trace.eventf("codex.ws", "turn", request.trace_ctx, "lane={d} retained={d} reused={d} handshake_ms={d} health={d} auth=chatgpt_subscription", .{ + checkout.slot orelse std.math.maxInt(usize), + @as(u8, @intFromBool(checkout.retained)), + @as(u8, @intFromBool(checkout.reused)), + checkout.handshake_ms, + checkout.health_failures, + }); + const continued = if (continuation_recovery_attempted) + null + else + codex_websocket_session.continuation(checkout.slot, full_input, shape); + const payload = try serializer.build_request( + alloc, + request.data(), + if (continued) |value| value.delta_input else full_input, + if (continued) |value| value.previous_response_id else null, + ); + defer alloc.free(payload); + debug_trace.eventf("codex.ws", "continuation", request.trace_ctx, "used={d} delta_bytes={d} recovery={d}", .{ + @as(u8, @intFromBool(continued != null)), + if (continued) |value| value.delta_input.len else full_input.len, + @as(u8, @intFromBool(continuation_recovery_attempted)), + }); + websocket_transport.streamOn(checkout.connection, alloc, .{ + .endpoint = config.endpoint, + .authorization = config.authorization, + .account_id = config.account_id, + .session_id = request.session_id, + .payload = payload, + .deadline = request.deadline, + .cancel_flag = request.cancel_flag, + .delivery = request.delivery, + }, &bridge, WebSocketBridge.event) catch |err| { + codex_websocket_session.release(checkout, .failed); + if (err == error.PreviousResponseNotFound and continued != null and !continuation_recovery_attempted) { + continuation_recovery_attempted = true; + reducer.deinit(alloc); + reducer = responses_protocol.Reducer.init(alloc); + continue; + } + debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason={s} close={d}", .{ failureReason(err), @as(u16, 0) }); + return err; + }; + const completion = reducer.finish(alloc, request.cancel_flag, bridge.stream_limits) catch |err| { + codex_websocket_session.release(checkout, .failed); + debug_trace.eventf("codex.ws", "poison", request.trace_ctx, "reason=protocol close={d}", .{@as(u16, 0)}); + return mapReducerError(err); + }; + recordContinuation(alloc, checkout.slot, request, completion, full_input, shape, serializer); + codex_websocket_session.release(checkout, .completed); + return .{ .completed = .{ + .completion = completion, + .usage = .{ .unavailable = .possibly_billed }, + .ownership = .owned, + } }; + } +} + +fn recordContinuation( + alloc: Allocator, + slot: ?usize, + request: stream_provider.ModelRequest, + completion: types.ModelCompletion, + full_input: []const u8, + shape: [std.crypto.hash.sha2.Sha256.digest_length]u8, + serializer: Serializer, +) void { + const response_id = completion.generation_id orelse return; + const response_message = [_]types.ChatMessage{.{ + .role = .assistant, + .content = completion.content, + .tool_calls = completion.tool_calls, + .provider_state_json = completion.provider_state_json, + }}; + const response_input = serializer.build_input(alloc, &response_message, null) catch return; + defer alloc.free(response_input); + const baseline = buildContinuationBaseline(alloc, full_input, response_input) catch return; + defer alloc.free(baseline); + const durable_full_input = serializer.build_durable_input(alloc, request.messages) catch return; + defer alloc.free(durable_full_input); + const durable_response_input = serializer.build_durable_input(alloc, &response_message) catch return; + defer alloc.free(durable_response_input); + const durable_baseline = buildContinuationBaseline(alloc, durable_full_input, durable_response_input) catch return; + defer alloc.free(durable_baseline); + codex_websocket_session.recordCompletion( + slot, + response_id, + baseline, + durable_baseline, + shape, + ); +} + +fn buildContinuationBaseline( + alloc: Allocator, + full_input: []const u8, + response_input: []const u8, +) ![]u8 { + var out: std.Io.Writer.Allocating = .init(alloc); + errdefer out.deinit(); + try out.writer.writeAll(full_input); + if (response_input.len > 0) { + if (out.written().len > 0) try out.writer.writeByte(','); + try out.writer.writeAll(response_input); + } + return out.toOwnedSlice(); +} + +fn failureReason(err: anyerror) []const u8 { + return switch (err) { + error.Cancelled => "cancel", + error.Timeout => "timeout", + error.WebSocketPolicyClosed => "policy", + error.WebSocketUnexpectedBinary => "binary", + error.WebSocketProtocolViolation, error.WebSocketInvalidUtf8 => "protocol", + error.WebSocketUpgradeRejected, error.WebSocketAcceptInvalid => "auth", + else => "close", + }; +} + +const WebSocketBridge = struct { + alloc: Allocator, + reducer: *responses_protocol.Reducer, + events: stream_provider.EventSink, + cancel_flag: *std.atomic.Value(bool), + content_capture_limit: ?usize, + stream_limits: responses_protocol.StreamLimits, + + fn event(raw: *anyopaque, json_text: []const u8) !bool { + const self: *@This() = @ptrCast(@alignCast(raw)); + return self.reducer.applyJson( + self.alloc, + json_text, + .{ + .context = &self.events, + .on_content = EventBridge.content, + .on_tool_start = EventBridge.toolStart, + .on_reasoning = EventBridge.reasoning, + .on_tool_input = EventBridge.toolInput, + }, + self.cancel_flag, + self.content_capture_limit, + self.stream_limits, + ) catch |err| return mapReducerError(err); + } +}; + +const EventBridge = struct { + fn sink(raw: *anyopaque) *stream_provider.EventSink { + return @ptrCast(@alignCast(raw)); + } + + fn content(raw: *anyopaque, chunk: []const u8) void { + sink(raw).emit(.{ .content_delta = chunk }); + } + + fn reasoning(raw: *anyopaque, chunk: []const u8) void { + sink(raw).emit(.{ .reasoning_delta = chunk }); + } + + fn toolInput(raw: *anyopaque, chunk: []const u8) void { + sink(raw).emit(.{ .tool_input_delta = chunk }); + } + + fn toolStart(raw: *anyopaque, id: []const u8, name: []const u8, label: ?[]const u8) void { + sink(raw).emit(.{ .tool_started = .{ .id = id, .name = name, .label = label } }); + } +}; + +fn mapReducerError(err: anyerror) anyerror { + return switch (err) { + error.InvalidEvent => error.InvalidOpenAICodexSseEvent, + error.PreviousResponseNotFound => error.PreviousResponseNotFound, + error.ResponseFailed => error.OpenAICodexResponseFailed, + error.StreamIncomplete => error.OpenAICodexStreamIncomplete, + error.ToolCallLimitExceeded => error.OpenAICodexToolCallLimitExceeded, + error.ToolArgumentsTooLarge => error.OpenAICodexToolArgumentsTooLarge, + error.ResourceLimitExceeded => error.OpenAICodexResourceLimitExceeded, + else => err, + }; +} diff --git a/src/gateway/websocket_transport.zig b/src/gateway/websocket_transport.zig index 0d212a060..96b3e8d5b 100644 --- a/src/gateway/websocket_transport.zig +++ b/src/gateway/websocket_transport.zig @@ -152,7 +152,7 @@ pub fn connect(alloc: Allocator, args: ConnectArgs) !*Connection { count += 1; extra_headers[count] = .{ .name = "originator", .value = "fx" }; count += 1; - extra_headers[count] = .{ .name = "OpenAI-Beta", .value = "responses_websockets=v2" }; + extra_headers[count] = .{ .name = "OpenAI-Beta", .value = "responses_websockets=2026-02-06" }; count += 1; extra_headers[count] = .{ .name = "Upgrade", .value = "websocket" }; count += 1; diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index 8ed277b5d..0ebbac434 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -848,16 +848,19 @@ function startFakeCodexWebSocket(options: { rejectPreviousOnce?: boolean; toolRoundTrip?: boolean; reasoningState?: boolean; + rejectUpgradeWithSse?: boolean; } = {}) { const requests: string[] = []; const closeCodes: number[] = []; let upgradeRequests = 0; + let sseRequests = 0; + let httpResponseRequests = 0; let rejectedPrevious = false; const accessToken = chatgptAccessToken("acct_websocket"); const server = Bun.serve<{ opened: boolean }>({ hostname: "127.0.0.1", port: 0, - fetch(request, server) { + async fetch(request, server) { const path = new URL(request.url).pathname; if (path === "/models") { return Response.json({ models: [ @@ -865,9 +868,31 @@ function startFakeCodexWebSocket(options: { ] }); } if (path === "/responses") { - upgradeRequests += 1; - if (options.stallUpgrade) return new Promise(() => {}); - if (server.upgrade(request, { data: { opened: true } })) return; + if (request.headers.get("upgrade")?.toLowerCase() === "websocket") { + upgradeRequests += 1; + if (options.rejectUpgradeWithSse) return new Response("upgrade unavailable", { status: 426 }); + if (options.stallUpgrade) return new Promise(() => {}); + if (server.upgrade(request, { data: { opened: true } })) return; + } else { + httpResponseRequests += 1; + if (options.rejectUpgradeWithSse) { + requests.push(await request.text()); + sseRequests += 1; + const completed = { + type: "response.completed", + response: { + id: `resp_sse_${sseRequests}`, + status: "completed", + usage: { input_tokens: 5, output_tokens: 2 }, + }, + }; + return new Response( + `data: ${JSON.stringify({ type: "response.output_text.delta", delta: `CODEX_SSE_FALLBACK_${sseRequests}` })}\n\n` + + `data: ${JSON.stringify(completed)}\n\n`, + { headers: { "content-type": "text/event-stream" } }, + ); + } + } } return new Response("not found", { status: 404 }); }, @@ -951,6 +976,8 @@ function startFakeCodexWebSocket(options: { requests, closeCodes, get upgradeRequests() { return upgradeRequests; }, + get sseRequests() { return sseRequests; }, + get httpResponseRequests() { return httpResponseRequests; }, responsesUrl: `http://127.0.0.1:${server.port}/responses`, modelsUrl: `http://127.0.0.1:${server.port}/models`, stop() { server.stop(true); }, @@ -3050,6 +3077,7 @@ test( expect(`${result.stdout}\n${result.stderr}`).toContain("WebSocketClosedBeforeCompletion"); expect(codex.requests).toHaveLength(1); expect(codex.upgradeRequests).toBe(1); + expect(codex.httpResponseRequests).toBe(0); expect(gateway.requests).toHaveLength(0); } finally { codex.stop(); @@ -3191,6 +3219,43 @@ tmuxTest( 60_000, ); +tmuxTest( + "Codex WebSocket handshake failure falls back once and latches SSE", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-websocket-fallback-")); + stderrPath = join(home, "stderr.log"); + writeFileSync(stderrPath, ""); + gateway = startFakeGateway([]); + const codex = startFakeCodexWebSocket({ rejectUpgradeWithSse: true }); + try { + writeSeededChatGptLogin(home, codex.accessToken); + writeFileSync( + join(home, ".fx", "settings.json"), + JSON.stringify({ provider: "codex", codex_model: "gpt-5.6-sol" }) + "\n", + { mode: 0o600 }, + ); + session = await startFx(home, stderrPath, gateway, undefined, undefined, { + FX_MODEL: undefined, + FX_CODEX_TRANSPORT: "websocket", + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + }); + await session.waitForComposer(TIMEOUT); + await session.sendText("Fall back after the rejected WebSocket upgrade."); + await session.waitForText("CODEX_SSE_FALLBACK_1", TIMEOUT); + await session.sendText("Keep using SSE after the fallback latch is armed."); + await session.waitForText("CODEX_SSE_FALLBACK_2", TIMEOUT); + expect(codex.upgradeRequests).toBe(1); + expect(codex.sseRequests).toBe(2); + expect(codex.requests).toHaveLength(2); + expect(readFileSync(stderrPath, "utf8")).toBe(""); + } finally { + codex.stop(); + } + }, + 60_000, +); + tmuxTest( "Codex WebSocket reconnects before delivery when the retained socket closes", async () => { From 77e921de90b5d5ec71082f6bd81d7d12660b0bb1 Mon Sep 17 00:00:00 2001 From: Ashman Singh <36335693+thinkter@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:13:18 +0530 Subject: [PATCH 21/21] Delete docs/codex-websocket-plan.md --- docs/codex-websocket-plan.md | 128 ----------------------------------- 1 file changed, 128 deletions(-) delete mode 100644 docs/codex-websocket-plan.md diff --git a/docs/codex-websocket-plan.md b/docs/codex-websocket-plan.md deleted file mode 100644 index 700ece23e..000000000 --- a/docs/codex-websocket-plan.md +++ /dev/null @@ -1,128 +0,0 @@ -# Codex WebSocket transport: plan to finish PR #521 - -This document replaces the earlier `docs/codex-websocket-transport.md` design -notes. It records the verified comparison between PR #521 -(`feat/codex-websocket-phase-2`, head `05bf5a4`) and PR #523 -(`perf/codex-websocket-transport`, head `f486bb2`), and the work needed to make -#521 the clear landing candidate. - -## Implementation status - -The plan is implemented on the rebased feature branch: - -- Codex WebSocket orchestration now has a separate adapter module. -- The transport uses the dated beta header and strict frame handling. -- Pre-delivery connection failures fall back to SSE and arm a process latch; - ambiguous delivery never falls back or replays. -- Retained lanes perform network I/O outside the pool mutex, use temporary - connections under saturation, and have a global LRU storage bound. -- Focused tests cover fallback latching, ambiguous-delivery no-replay, - continuation recovery, retained connection reuse, expiration, and identity - churn. - -## Verified findings - -Each claim below was checked directly against both PR heads and current `main` -(`bb2dc7d`). - -- **Stale beta header.** `src/gateway/websocket_transport.zig` sends - `OpenAI-Beta: responses_websockets=v2`. #523 sends the dated - `responses_websockets=2026-02-06`, which matches the current upstream Codex - client. -- **Network I/O under the pool mutex.** `acquire` in - `src/gateway/codex_websocket_session.zig` holds `pool_mutex` through `ping`, - `close`, and `connect`. A slow handshake on one lane stalls every other - lane. When all matching lanes are busy it also busy-polls with a 10ms sleep. -- **Unbounded slot array.** `appendSlot` grows the global slot list per new - identity. Incompatible slots get their connection closed but the slot entry - itself is never evicted. -- **Merge state.** Only `src/gateway/openai_codex.zig` conflicts with `main`, - caused by `cd1a5d3` switching error returns to - `stream_provider.failResult(...)`. -- **#523's real advantages.** A clean adapter boundary (about 41 lines touched - in `openai_codex.zig`), a process-wide SSE fallback latch, and the dated - header. -- **#523's real defects.** Binary frames are accepted as text, close frames - are not consumed or validated, there is no idle-event timeout, its cache is - one global entry (alternating sessions displace each other), and its - fallback and fresh-retry trigger is "no output yet" rather than "provably - unsent", which permits duplicate provider requests and billing after an - ambiguous delivery. - -## Strategy - -Keep #521's functional advantages (safe-prefix continuation, delivery-aware -retries, strict frame validation, retained lanes, deep test coverage). Adopt -every legitimate advantage of #523. Fix the three real defects in this branch. -After that, #523's weaknesses have no counterpart here. - -## Phase 1: rebase and restructure - -1. **Rebase onto current `main`.** Only `openai_codex.zig` conflicts. While - resolving, adapt the WebSocket error paths to the new - `stream_provider.failResult(...)` convention. -2. **Extract the added logic from `openai_codex.zig` into the adapter.** Move - `buildWebSocketRequest*`, the WebSocket consume loop, and continuation - orchestration into `codex_websocket_session.zig` (or a new - `openai_codex_websocket.zig`). Target: the provider file gains only a small - transport dispatch, comparable to #523's footprint. This removes the - integration-boundary critique and shrinks the future conflict surface. -3. **Drop `scripts/codex_websocket_probe.py`** from this PR. It is a dev - probe, not product code. Land it separately if it is worth keeping. - -## Phase 2: adopt #523's genuine wins, done better - -4. **Update the beta header** to `responses_websockets=2026-02-06`. Re-check - the current upstream constant when making the change and cite it in the - commit message. -5. **Add a delivery-aware process-wide SSE fallback latch.** Latch to SSE only - when the failure is provably pre-send (handshake or upgrade rejection, - connect timeout, pre-write errors), using the `DeliveryCertainty` already - threaded through `AcquireArgs`. An ambiguously sent request surfaces an - error instead of being silently re-issued. This keeps #523's "a broken - proxy costs one failed handshake, not one per turn" resilience without its - duplicate-billing hole. -6. **Keep strict transport-env validation.** Erroring on an invalid - `FX_CODEX_TRANSPORT` value is deliberate and better than silently coercing - to SSE. Say so in the PR description. - -## Phase 3: fix the pool's real defects - -7. **Move network I/O outside `pool_mutex`.** Restructure `acquire` to: lock, - select and reserve a slot (`busy = true`), unlock, then ping, connect, or - close outside the lock, and relock only to commit or roll back slot state. - Nothing slow ever runs under the lock. -8. **Delete the busy-wait path.** When all matching lanes are busy at the lane - limit, hand out a temporary unpooled connection instead of sleeping in 10ms - slices. Simpler, lower tail latency, and one less polling state. -9. **Bound and evict slots globally.** Add a global slot cap (for example - `max_lanes` times a small constant). Evict idle slots LRU when appending - past the cap, and fully remove, not just disconnect, slots whose identity - is incompatible. Add a unit test proving the bound holds under identity - churn. - -## Phase 4: prove it - -10. **Add two E2E tests for the new behavior:** - - a handshake failure latches the process to SSE and the turn still - completes; - - an ambiguous post-send failure does not fall back or re-send, and errors - with delivery evidence. - Both PRs added tests to `tests/e2e/tui-auth-source-selection.test.ts`; - inherit its corpus classification but confirm it still fits per AGENTS.md. -11. **Run the full ready gate:** `zig fmt --check src/`, focused Zig tests, - build, drive `./zig-out/bin/fx` with `FX_CODEX_TRANSPORT=websocket` - against the loopback E2E mock, push, and require Full CI green on all four - runners for the exact head before marking the PR ready. - -## Optional: split into a two-PR stack - -The strongest structural critique of #521 is that continuation plus pooling is -a lot for a first landing. If reviewers push back, stack the branch: - -- **PR A:** strict codec, adapter boundary, single retained connection, - delivery-aware fallback latch. A superset of #523 with none of its bugs. -- **PR B:** safe-prefix continuation and the multi-lane pool on top. - -If a single PR is preferred, Phases 1 through 4 alone make #521 dominate the -comparison on every row except raw diff size.