From 87391462cc90586f468b4571ee2d8514f8358f82 Mon Sep 17 00:00:00 2001 From: mjlbach Date: Sat, 29 Aug 2026 09:43:24 -0700 Subject: [PATCH 1/2] Add an opt-in WebSocket transport for Codex Responses Upgrade the Codex Responses endpoint to a WebSocket (OpenAI-Beta: responses_websockets=2026-02-06), send the request body as one response.create text frame, and feed each received frame into the same Responses reducer the SSE transport uses. Connections are cached per (session, account) with a five-minute idle and 55-minute age limit, so later turns skip DNS, TCP, and TLS setup; a busy or mismatched entry gets a one-off connection, and a reused connection that fails before any model output is replaced by a fresh one once. Enabled with FX_OPENAI_CODEX_TRANSPORT=websocket; the default transport is unchanged. A WebSocket transport failure before any model output falls back to SSE under the original invocation admission and latches the process to SSE; after output has been emitted the failure is surfaced instead, so a replay can never duplicate tool calls. Verified against the production Codex endpoint, including cross-turn connection reuse in a saved session. Co-Authored-By: Claude Fable 5 --- src/gateway/openai_codex.zig | 41 +- src/gateway/openai_codex_websocket.zig | 991 +++++++++++++++++++++++++ src/gateway/websocket.zig | 256 +++++++ src/main.zig | 2 + 4 files changed, 1288 insertions(+), 2 deletions(-) create mode 100644 src/gateway/openai_codex_websocket.zig create mode 100644 src/gateway/websocket.zig diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index 95e3356ed..ab9b9a2d4 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -6,6 +6,7 @@ const stream_provider = @import("../core/agent/stream_provider.zig"); const io_mod = @import("../core/shared/io.zig"); const types = @import("../core/shared/types.zig"); const gateway_client = @import("client.zig"); +const codex_websocket = @import("openai_codex_websocket.zig"); const responses_protocol = @import("responses_protocol.zig"); const model_tool_schema = @import("../core/tooling/model_tool_schema.zig"); @@ -140,7 +141,32 @@ fn streamCompletion( try validateModel(request.model); const payload = try buildRequest(alloc, request.data()); defer alloc.free(payload); - return streamPrepared(alloc, request, payload) catch |err| { + + var admission_policy: AdmissionPolicy = .admit; + if (codex_websocket.webSocketAvailable()) { + var output_emitted = false; + if (codex_websocket.streamViaWebSocket(alloc, request, payload, &output_emitted)) |result| { + return result; + } else |err| { + if (request.cancel_flag.load(.seq_cst)) return stream_provider.failResult(error.Cancelled); + if (err == error.OutOfMemory) return err; + // Once the model produced output, replaying the request over SSE + // could duplicate tool calls; surface the failure instead. + if (output_emitted or !codex_websocket.errorAllowsSseFallback(err)) { + request.attempt_evidence.network_failure = gateway_client.networkFailureEvidence(err, request.delivery.load()); + return err; + } + codex_websocket.armSseFallback(err); + // The WebSocket attempt may have consumed this invocation's + // admission; the SSE fallback continues under it. + admission_policy = if (request.attempt_evidence.provider_admitted) + .already_admitted + else + .admit; + } + } + + return streamPreparedWithAdmission(alloc, request, payload, admission_policy) 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; @@ -183,10 +209,21 @@ const OpenRequestOperation = struct { } }; +const AdmissionPolicy = enum { admit, already_admitted }; + pub fn streamPrepared( alloc: Allocator, request: stream_provider.ModelRequest, payload: []const u8, +) !stream_provider.Result { + return streamPreparedWithAdmission(alloc, request, payload, .admit); +} + +fn streamPreparedWithAdmission( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, + admission_policy: AdmissionPolicy, ) !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); @@ -230,7 +267,7 @@ pub fn streamPrepared( .clock = .awake, .raw = .fromMilliseconds(connect_timeout_ms), }); - try request.admission.admit(); + if (admission_policy == .admit) try request.admission.admit(); var opened = try gateway_client.runBoundedHttpOperation( OpenedRequest, alloc, diff --git a/src/gateway/openai_codex_websocket.zig b/src/gateway/openai_codex_websocket.zig new file mode 100644 index 000000000..63a36b6b5 --- /dev/null +++ b/src/gateway/openai_codex_websocket.zig @@ -0,0 +1,991 @@ +//! WebSocket-first transport for the OpenAI Codex Responses endpoint. +//! +//! The wire contract mirrors the Codex CLI and pi clients: the HTTP Responses +//! endpoint is upgraded to a WebSocket (`OpenAI-Beta: responses_websockets`), +//! the request body is sent as one text frame with `"type":"response.create"` +//! prepended, and each received text frame carries one Responses stream event +//! — the same JSON payloads the SSE transport delivers in `data:` lines, so +//! both transports share `responses_protocol.Reducer`. +//! +//! Transport policy, mirrored from the reference clients: +//! - opt-in via `FX_OPENAI_CODEX_TRANSPORT=websocket`; +//! - one process-wide fallback latch: after a WebSocket transport failure the +//! process stops attempting WebSocket so a broken proxy is paid for once; +//! - a connection is cached per (session, account) and reused across turns; +//! a busy or mismatched cache entry yields a fresh uncached connection; +//! - a reused connection that fails before any model output is replaced by a +//! fresh connection once before the caller falls back to SSE; +//! - fallback decisions after model output has been emitted are forbidden — +//! the caller surfaces the error instead of replaying the request. + +const std = @import("std"); +const chatgpt_oauth = @import("../core/auth/chatgpt_oauth.zig"); +const debug_trace = @import("../core/shared/debug_trace.zig"); +const io_mod = @import("../core/shared/io.zig"); +const secret = @import("../core/auth/secret.zig"); +const stream_provider = @import("../core/agent/stream_provider.zig"); +const types = @import("../core/shared/types.zig"); +const gateway_client = @import("client.zig"); +const responses_protocol = @import("responses_protocol.zig"); +const websocket = @import("websocket.zig"); + +const Allocator = std.mem.Allocator; + +const endpoint = "https://chatgpt.com/backend-api/codex/responses"; +const e2e_endpoint_env = "FX_E2E_OPENAI_CODEX_RESPONSES_URL"; +const transport_env = "FX_OPENAI_CODEX_TRANSPORT"; +const beta_header_value = "responses_websockets=2026-02-06"; + +const connect_timeout_ms: i64 = 15_000; +const idle_ttl_ms: i64 = 5 * std.time.ms_per_min; +const max_age_ms: i64 = 55 * std.time.ms_per_min; + +// Stream limits mirror the SSE transport in openai_codex.zig so switching +// transports never changes what the client accepts. +const max_ws_message_bytes: usize = 32 * 1024 * 1024; +const max_sse_aggregate_bytes: usize = 64 * 1024 * 1024; +const max_events: usize = 100_000; +const max_tool_calls: usize = 128; +const max_tool_identity_bytes: usize = 1024; +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 stream_limits = responses_protocol.StreamLimits{ + .aggregate_bytes = max_sse_aggregate_bytes, + .events = max_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, +}; + +pub const TransportMode = enum { sse, websocket }; + +pub fn transportMode() TransportMode { + const value = io_mod.getenv(transport_env) orelse return .sse; + if (std.ascii.eqlIgnoreCase(value, "websocket") or std.ascii.eqlIgnoreCase(value, "ws")) { + return .websocket; + } + return .sse; +} + +/// Once a WebSocket attempt fails in a way the caller decided to fall back +/// from, the rest of the process sticks to SSE: a deterministically broken +/// path (proxy, firewall) should cost one failed handshake, not one per turn. +var sse_fallback_active = std.atomic.Value(bool).init(false); + +pub fn webSocketAvailable() bool { + return transportMode() == .websocket and !sse_fallback_active.load(.seq_cst); +} + +pub 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)}, + ); +} + +/// Errors that indicate the provider processed (and refused) the request, or +/// that a resource limit fired locally, must not trigger an SSE replay: the +/// replay would repeat the same failure at full request cost. +pub fn errorAllowsSseFallback(err: anyerror) bool { + return switch (err) { + error.Cancelled, + error.OutOfMemory, + error.OpenAICodexResponseFailed, + error.OpenAICodexToolCallLimitExceeded, + error.OpenAICodexToolArgumentsTooLarge, + error.OpenAICodexResourceLimitExceeded, + error.ProviderAdmissionMissing, + error.ProviderAdmissionRepeated, + => false, + else => true, + }; +} + +pub fn streamViaWebSocket( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, + output_emitted: *bool, +) !stream_provider.Result { + const url = if (io_mod.getenv(e2e_endpoint_env)) |override| url: { + if (!gateway_client.isLoopbackHttpUrl(override)) return error.InvalidE2EOpenAICodexEndpoint; + break :url override; + } else endpoint; + return streamViaWebSocketAtUrl(alloc, request, payload, url, output_emitted); +} + +pub fn streamViaWebSocketAtUrl( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, + url: []const u8, + output_emitted: *bool, +) !stream_provider.Result { + output_emitted.* = false; + 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); + + // One admission covers the whole WebSocket attempt cycle, including the + // single fresh-connection retry below and the caller's SSE fallback. + try request.admission.admit(); + + var attempt: usize = 0; + while (true) : (attempt += 1) { + const acquired = try acquireConnection(alloc, request, url, auth_header, account_id); + debug_trace.eventf( + "stream", + "codex_ws_connection", + request.trace_ctx, + "reused={any} cacheable={any}", + .{ acquired.reused, acquired.cacheable }, + ); + if (runOnConnection(alloc, request, payload, acquired.conn, output_emitted)) |result| { + releaseConnection(acquired, true); + return result; + } else |err| { + releaseConnection(acquired, false); + const retry_with_fresh = acquired.reused and attempt == 0 and + !output_emitted.* and + err != error.Cancelled and err != error.OutOfMemory; + if (retry_with_fresh) { + debug_trace.eventf( + "stream", + "codex_ws_reused_connection_replaced", + request.trace_ctx, + "error={s}", + .{@errorName(err)}, + ); + continue; + } + return err; + } + } +} + +// --------------------------------------------------------------------------- +// Connection lifecycle +// --------------------------------------------------------------------------- + +const WsConnection = struct { + arena_state: std.heap.ArenaAllocator, + client: std.http.Client, + request: std.http.Client.Request, + response: std.http.Client.Response, + reader: *std.Io.Reader, + session_key: ?[]u8, + created_at_ms: i64, + last_used_ms: i64, + + fn connectionRef(self: *WsConnection) *std.http.Client.Connection { + return self.request.connection.?; + } + + fn netStream(self: *WsConnection) std.Io.net.Stream { + return self.connectionRef().stream_writer.stream; + } + + fn destroy(self: *WsConnection) void { + // Shut the socket down first so releasing the request can never block + // draining a live stream. + self.netStream().shutdown(io_mod.getIo(), .both) catch {}; + self.request.deinit(); + self.client.deinit(); + var arena = self.arena_state; + arena.deinit(); + } + + /// `runBoundedHttpOperation` cleanup contract for late arrivals. + pub fn deinit(self: *WsConnection, _: Allocator) void { + self.destroy(); + } +}; + +const cache_mutex_state = struct { + var mutex: std.Io.Mutex = .init; + var entry: ?*WsConnection = null; + var entry_busy: bool = false; + + fn lock() void { + mutex.lockUncancelable(io_mod.getIo()); + } + + fn unlock() void { + mutex.unlock(io_mod.getIo()); + } +}; + +const Acquired = struct { + conn: *WsConnection, + reused: bool, + cacheable: bool, +}; + +fn sessionKey(alloc: Allocator, session_id: []const u8, account_id: []const u8) ![]u8 { + return std.fmt.allocPrint(alloc, "{s}\x00{s}", .{ session_id, account_id }); +} + +fn acquireConnection( + alloc: Allocator, + request: stream_provider.ModelRequest, + url: []const u8, + auth_header: []const u8, + account_id: []const u8, +) !Acquired { + const session_id: ?[]const u8 = if (request.session_id) |sid| + (if (sid.len > 0) sid else null) + else + null; + const key: ?[]u8 = if (session_id) |sid| try sessionKey(alloc, sid, account_id) else null; + defer if (key) |value| alloc.free(value); + + if (key) |wanted| { + var evicted: ?*WsConnection = null; + var reusable: ?*WsConnection = null; + { + cache_mutex_state.lock(); + defer cache_mutex_state.unlock(); + if (cache_mutex_state.entry) |conn| { + if (!cache_mutex_state.entry_busy) { + const now = io_mod.milliTimestamp(); + const matches = conn.session_key != null and + std.mem.eql(u8, conn.session_key.?, wanted); + const expired = now - conn.created_at_ms >= max_age_ms or + now - conn.last_used_ms >= idle_ttl_ms; + if (matches and !expired) { + cache_mutex_state.entry_busy = true; + reusable = conn; + } else { + cache_mutex_state.entry = null; + evicted = conn; + } + } + // A busy entry stays untouched; this request uses a one-off + // connection so concurrent turns never share one socket. + } + } + if (evicted) |conn| conn.destroy(); + if (reusable) |conn| return .{ .conn = conn, .reused = true, .cacheable = true }; + } + + const conn = try connectBounded(alloc, request, url, auth_header, account_id, key); + return .{ .conn = conn, .reused = false, .cacheable = key != null }; +} + +fn releaseConnection(acquired: Acquired, keep: bool) void { + const conn = acquired.conn; + if (!keep or !acquired.cacheable or conn.session_key == null) { + removeFromCache(conn); + conn.destroy(); + return; + } + var displaced: ?*WsConnection = null; + { + cache_mutex_state.lock(); + defer cache_mutex_state.unlock(); + conn.last_used_ms = io_mod.milliTimestamp(); + if (cache_mutex_state.entry == conn) { + cache_mutex_state.entry_busy = false; + } else if (cache_mutex_state.entry == null) { + cache_mutex_state.entry = conn; + cache_mutex_state.entry_busy = false; + } else { + // Another connection claimed the slot while this one was in + // flight; the newest one wins and this one is closed. + displaced = conn; + } + } + if (displaced) |value| value.destroy(); +} + +fn removeFromCache(conn: *WsConnection) void { + cache_mutex_state.lock(); + defer cache_mutex_state.unlock(); + if (cache_mutex_state.entry == conn) { + cache_mutex_state.entry = null; + cache_mutex_state.entry_busy = false; + } +} + +/// Closes the cached connection, if any. Intended for tests and shutdown. +pub fn closeCachedConnection() void { + var evicted: ?*WsConnection = null; + { + cache_mutex_state.lock(); + defer cache_mutex_state.unlock(); + if (cache_mutex_state.entry) |conn| { + if (!cache_mutex_state.entry_busy) { + cache_mutex_state.entry = null; + evicted = conn; + } + } + } + if (evicted) |conn| conn.destroy(); +} + +const ConnectOperation = struct { + url: []const u8, + auth_header: []const u8, + account_id: []const u8, + session_id: ?[]const u8, + session_key: ?[]const u8, + + pub fn run(self: *@This()) !*WsConnection { + return createConnection( + self.url, + self.auth_header, + self.account_id, + self.session_id, + self.session_key, + ); + } +}; + +fn connectBounded( + alloc: Allocator, + request: stream_provider.ModelRequest, + url: []const u8, + auth_header: []const u8, + account_id: []const u8, + session_key: ?[]const u8, +) !*WsConnection { + var operation = ConnectOperation{ + .url = url, + .auth_header = auth_header, + .account_id = account_id, + .session_id = request.session_id, + .session_key = session_key, + }; + return gateway_client.runBoundedHttpOperation( + *WsConnection, + alloc, + request.cancel_flag, + std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ + .clock = .awake, + .raw = .fromMilliseconds(connect_timeout_ms), + }), + &operation, + ); +} + +fn createConnection( + url: []const u8, + auth_header: []const u8, + account_id: []const u8, + session_id: ?[]const u8, + session_key: ?[]const u8, +) !*WsConnection { + var boot = std.heap.ArenaAllocator.init(std.heap.page_allocator); + const self = boot.allocator().create(WsConnection) catch |err| { + boot.deinit(); + return err; + }; + self.* = undefined; + self.arena_state = boot; + initConnection(self, url, auth_header, account_id, session_id, session_key) catch |err| { + var arena = self.arena_state; + arena.deinit(); + return err; + }; + return self; +} + +fn initConnection( + self: *WsConnection, + url: []const u8, + auth_header: []const u8, + account_id: []const u8, + session_id: ?[]const u8, + session_key: ?[]const u8, +) !void { + const arena = self.arena_state.allocator(); + const now = io_mod.milliTimestamp(); + self.created_at_ms = now; + self.last_used_ms = now; + self.session_key = if (session_key) |key| try arena.dupe(u8, key) else null; + + // `Request.uri` borrows this memory for the request's lifetime. + const url_copy = try arena.dupe(u8, url); + const uri = try std.Uri.parse(url_copy); + + var random_key: [16]u8 = undefined; + io_mod.getIo().random(&random_key); + const sec_key = websocket.secKey(random_key); + const expected_accept = websocket.acceptKey(&sec_key); + + var extra_headers_buf: [8]std.http.Header = undefined; + var extra_count: usize = 0; + extra_headers_buf[extra_count] = .{ .name = "upgrade", .value = "websocket" }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "sec-websocket-version", .value = "13" }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "sec-websocket-key", .value = &sec_key }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = account_id }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "OpenAI-Beta", .value = beta_header_value }; + extra_count += 1; + if (session_id) |sid| if (sid.len > 0) { + extra_headers_buf[extra_count] = .{ .name = "session-id", .value = sid }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "x-client-request-id", .value = sid }; + extra_count += 1; + }; + + self.client = .{ .allocator = arena, .io = io_mod.getIo() }; + errdefer self.client.deinit(); + self.request = try self.client.request(.GET, uri, .{ + .headers = .{ + .authorization = .{ .override = auth_header }, + .user_agent = .{ .override = gateway_client.user_agent }, + .accept_encoding = .omit, + .connection = .{ .override = "Upgrade" }, + }, + .extra_headers = extra_headers_buf[0..extra_count], + .keep_alive = true, + .redirect_behavior = .unhandled, + }); + errdefer self.request.deinit(); + try self.request.sendBodiless(); + self.response = try self.request.receiveHead(&.{}); + if (self.response.head.status != .switching_protocols) { + return error.WebSocketHandshakeRejected; + } + try verifyHandshakeHeaders(self.response.head.bytes, &expected_accept); + + const transfer_buffer = try arena.alloc(u8, transfer_buffer_bytes); + self.reader = self.response.reader(transfer_buffer); +} + +fn verifyHandshakeHeaders(head_bytes: []const u8, expected_accept: []const u8) !void { + var saw_upgrade = false; + var saw_accept = false; + var it = std.http.HeaderIterator.init(head_bytes); + while (it.next()) |header| { + if (std.ascii.eqlIgnoreCase(header.name, "upgrade")) { + if (!std.ascii.eqlIgnoreCase(header.value, "websocket")) { + return error.WebSocketHandshakeInvalid; + } + saw_upgrade = true; + } else if (std.ascii.eqlIgnoreCase(header.name, "sec-websocket-accept")) { + if (!std.mem.eql(u8, header.value, expected_accept)) { + return error.WebSocketHandshakeInvalid; + } + saw_accept = true; + } + } + if (!saw_upgrade or !saw_accept) return error.WebSocketHandshakeInvalid; +} + +// --------------------------------------------------------------------------- +// Streaming +// --------------------------------------------------------------------------- + +const TrackingSink = struct { + inner: stream_provider.EventSink, + emitted: *bool, + + fn emit(raw: *anyopaque, event: stream_provider.Event) void { + const self: *TrackingSink = @ptrCast(@alignCast(raw)); + self.emitted.* = true; + self.inner.emit(event); + } + + fn content(raw: *anyopaque, chunk: []const u8) void { + const self: *TrackingSink = @ptrCast(@alignCast(raw)); + self.emitted.* = true; + self.inner.emit(.{ .content_delta = chunk }); + } + + fn reasoning(raw: *anyopaque, chunk: []const u8) void { + const self: *TrackingSink = @ptrCast(@alignCast(raw)); + self.emitted.* = true; + self.inner.emit(.{ .reasoning_delta = chunk }); + } + + fn toolInput(raw: *anyopaque, chunk: []const u8) void { + const self: *TrackingSink = @ptrCast(@alignCast(raw)); + self.emitted.* = true; + self.inner.emit(.{ .tool_input_delta = chunk }); + } + + fn toolStart(raw: *anyopaque, id: []const u8, name: []const u8, label: ?[]const u8) void { + const self: *TrackingSink = @ptrCast(@alignCast(raw)); + self.emitted.* = true; + self.inner.emit(.{ .tool_started = .{ .id = id, .name = name, .label = label } }); + } +}; + +fn runOnConnection( + alloc: Allocator, + request: stream_provider.ModelRequest, + payload: []const u8, + conn: *WsConnection, + output_emitted: *bool, +) !stream_provider.Result { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + + var cancel_watch_done = std.atomic.Value(bool).init(false); + const cancel_watcher = try gateway_client.spawnHttpCancelWatcher( + &cancel_watch_done, + request.cancel_flag, + conn.netStream(), + ); + defer { + cancel_watch_done.store(true, .seq_cst); + cancel_watcher.join(); + } + + // The SSE body is a complete JSON object; the WebSocket request is the + // same object with the message type prepended. + if (payload.len < 2 or payload[0] != '{') return error.InvalidOpenAICodexRequestPayload; + const ws_payload = try std.fmt.allocPrint( + alloc, + "{{\"type\":\"response.create\",{s}", + .{payload[1..]}, + ); + defer secret.zeroAndFree(alloc, ws_payload); + + request.delivery.markPossiblySent(); + sendFrame(conn, .text, ws_payload) catch |err| return wsIoError(request, err); + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + + var tracking = TrackingSink{ .inner = request.events, .emitted = output_emitted }; + var completion = try consumeWebSocket(alloc, request, conn, &tracking); + errdefer { + var owned = stream_provider.Result{ .completed = .{ + .completion = completion, + .ownership = .owned, + } }; + owned.deinit(alloc); + } + const usage_outcome: stream_provider.UsageOutcome = usage: { + if (completion.generation_id == null) { + break :usage .{ .unavailable = .possibly_billed }; + } + completion.billing = try responses_protocol.buildSubscriptionBilling( + alloc, + .codex, + request.model, + @max(io_mod.milliTimestamp(), 0), + completion.usage, + ) orelse break :usage .{ .unavailable = .possibly_billed }; + break :usage .{ .exact = .codex }; + }; + return .{ .completed = .{ + .completion = completion, + .usage = usage_outcome, + .ownership = .owned, + } }; +} + +fn wsIoError(request: stream_provider.ModelRequest, err: anyerror) anyerror { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + return err; +} + +fn sendFrame(conn: *WsConnection, opcode: websocket.Opcode, payload: []const u8) !void { + var mask: [4]u8 = undefined; + io_mod.getIo().random(&mask); + const connection = conn.connectionRef(); + try websocket.writeClientFrame(connection.writer(), opcode, payload, mask); + try connection.flush(); +} + +fn consumeWebSocket( + alloc: Allocator, + request: stream_provider.ModelRequest, + conn: *WsConnection, + tracking: *TrackingSink, +) !types.ModelCompletion { + var reducer = responses_protocol.Reducer.init(alloc); + defer reducer.deinit(alloc); + var scratch: std.ArrayList(u8) = .empty; + defer scratch.deinit(alloc); + const callbacks = responses_protocol.StreamCallbacks{ + .context = tracking, + .on_content = TrackingSink.content, + .on_tool_start = TrackingSink.toolStart, + .on_reasoning = TrackingSink.reasoning, + .on_tool_input = TrackingSink.toolInput, + }; + while (true) { + const message = nextTextMessage(alloc, request, conn, &scratch) catch |err| + return wsIoError(request, err); + const json_text = message orelse break; + if (reducer.applyJson( + alloc, + json_text, + callbacks, + request.cancel_flag, + request.content_capture_limit, + stream_limits, + ) catch |err| return mapReducerError(err)) break; + } + return reducer.finish(alloc, request.cancel_flag, stream_limits) catch |err| + return mapReducerError(err); +} + +/// Reads the next complete text message, transparently answering pings and +/// treating a close frame or clean socket end as end-of-stream (`null`). +/// The reducer decides whether an end-of-stream at that point is an error. +fn nextTextMessage( + alloc: Allocator, + request: stream_provider.ModelRequest, + conn: *WsConnection, + scratch: *std.ArrayList(u8), +) !?[]const u8 { + scratch.clearRetainingCapacity(); + var assembling = false; + while (true) { + if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + const head = websocket.readFrameHead(conn.reader, max_ws_message_bytes) catch |err| switch (err) { + error.EndOfStream => return null, + else => return err, + }; + switch (head.opcode) { + .ping => { + var control: std.ArrayList(u8) = .empty; + defer control.deinit(alloc); + try websocket.readPayloadInto(alloc, conn.reader, head, &control); + try sendFrame(conn, .pong, control.items); + }, + .pong => { + var control: std.ArrayList(u8) = .empty; + defer control.deinit(alloc); + try websocket.readPayloadInto(alloc, conn.reader, head, &control); + }, + .close => return null, + .text, .binary => { + if (assembling) return error.WebSocketProtocolError; + try websocket.readPayloadInto(alloc, conn.reader, head, scratch); + if (head.fin) return scratch.items; + assembling = true; + }, + .continuation => { + if (!assembling) return error.WebSocketProtocolError; + if (head.len > max_ws_message_bytes - scratch.items.len) { + return error.WebSocketFrameTooLarge; + } + try websocket.readPayloadInto(alloc, conn.reader, head, scratch); + if (head.fin) return scratch.items; + }, + _ => return error.WebSocketProtocolError, + } + } +} + +fn mapReducerError(err: anyerror) anyerror { + return switch (err) { + error.InvalidEvent => error.InvalidOpenAICodexSseEvent, + error.ResponseFailed => error.OpenAICodexResponseFailed, + error.StreamIncomplete => error.OpenAICodexStreamIncomplete, + error.ToolCallLimitExceeded => error.OpenAICodexToolCallLimitExceeded, + error.ToolArgumentsTooLarge => error.OpenAICodexToolArgumentsTooLarge, + error.ResourceLimitExceeded => error.OpenAICodexResourceLimitExceeded, + else => err, + }; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +const test_ws_events = [_][]const u8{ + "{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"message\"}}", + "{\"type\":\"response.output_text.delta\",\"output_index\":0,\"delta\":\"hello\"}", + "{\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":4}}}", +}; + +fn testChatGptJwt(alloc: Allocator) ![]u8 { + const claims = "{\"https://api.openai.com/auth\":{\"chatgpt_account_id\":\"acct-test\"}}"; + const encoded_len = std.base64.url_safe_no_pad.Encoder.calcSize(claims.len); + const encoded = try alloc.alloc(u8, encoded_len); + defer alloc.free(encoded); + _ = std.base64.url_safe_no_pad.Encoder.encode(encoded, claims); + return std.fmt.allocPrint(alloc, "h.{s}.s", .{encoded}); +} + +const WsLoopbackMode = enum { serve, reject_handshake }; + +const WsLoopbackFixture = struct { + io_backend: std.Io.Threaded = .init_single_threaded, + server: std.Io.net.Server, + mode: WsLoopbackMode, + thread: ?std.Thread = null, + server_open: bool = true, + stopping: std.atomic.Value(bool) = .init(false), + accept_count: std.atomic.Value(usize) = .init(0), + request_count: std.atomic.Value(usize) = .init(0), + saw_response_create: std.atomic.Value(bool) = .init(false), + failure: ?anyerror = null, + + fn init(mode: WsLoopbackMode) !@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 start(self: *@This()) !void { + std.debug.assert(self.thread == null); + self.thread = try std.Thread.spawn(.{}, run, .{self}); + } + + fn deinit(self: *@This()) void { + if (!self.server_open) return; + const zio = self.io(); + self.stopping.store(true, .seq_cst); + if (self.thread) |thread| { + const listener = std.Io.net.Stream{ .socket = self.server.socket }; + listener.shutdown(zio, .both) catch {}; + self.wakeAccept(); + thread.join(); + self.thread = null; + } + self.server.deinit(zio); + self.server_open = false; + } + + fn io(self: *@This()) std.Io { + return self.io_backend.io(); + } + + fn port(self: *@This()) u16 { + return self.server.socket.address.getPort(); + } + + fn url(self: *@This(), alloc: Allocator) ![]u8 { + return std.fmt.allocPrint(alloc, "http://127.0.0.1:{d}", .{self.port()}); + } + + fn wakeAccept(self: *@This()) void { + var wake_io_backend: std.Io.Threaded = .init_single_threaded; + const zio = wake_io_backend.io(); + const address = std.Io.net.IpAddress{ .ip4 = .loopback(self.port()) }; + var stream = address.connect(zio, .{ .mode = .stream }) catch return; + stream.close(zio); + } + + fn run(self: *@This()) void { + self.runFallible() catch |err| { + if (self.stopping.load(.seq_cst)) return; + self.failure = err; + }; + } + + fn runFallible(self: *@This()) !void { + const zio = self.io(); + while (!self.stopping.load(.seq_cst)) { + var stream = self.server.accept(zio) catch |err| { + if (self.stopping.load(.seq_cst)) return; + return err; + }; + defer stream.close(zio); + if (self.stopping.load(.seq_cst)) return; + _ = self.accept_count.fetchAdd(1, .seq_cst); + self.handleConnection(zio, stream) catch |err| switch (err) { + error.EndOfStream => continue, + else => return err, + }; + } + } + + fn handleConnection(self: *@This(), zio: std.Io, stream: std.Io.net.Stream) !void { + var read_buffer: [16 * 1024]u8 = undefined; + var reader = stream.reader(zio, &read_buffer); + var write_buffer: [4096]u8 = undefined; + var writer = stream.writer(zio, &write_buffer); + + var head: [16 * 1024]u8 = undefined; + var head_len: usize = 0; + while (head_len < head.len) { + head[head_len] = reader.interface.takeByte() catch return error.EndOfStream; + head_len += 1; + if (std.mem.endsWith(u8, head[0..head_len], "\r\n\r\n")) break; + } else return error.TestRequestTooLarge; + + if (self.mode == .reject_handshake) { + try writer.interface.writeAll("HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"); + try writer.interface.flush(); + return; + } + + const key = headerValue(head[0..head_len], "sec-websocket-key") orelse + return error.TestHandshakeKeyMissing; + const accept = websocket.acceptKey(key); + try writer.interface.print( + "HTTP/1.1 101 Switching Protocols\r\n" ++ + "upgrade: websocket\r\n" ++ + "connection: Upgrade\r\n" ++ + "sec-websocket-accept: {s}\r\n\r\n", + .{accept}, + ); + try writer.interface.flush(); + + var scratch: std.ArrayList(u8) = .empty; + defer scratch.deinit(std.heap.page_allocator); + while (!self.stopping.load(.seq_cst)) { + scratch.clearRetainingCapacity(); + const frame_head = websocket.readFrameHead(&reader.interface, 1 << 20) catch return error.EndOfStream; + if (frame_head.opcode == .close) return error.EndOfStream; + if (frame_head.opcode != .text) return error.TestUnexpectedFrame; + try websocket.readPayloadInto(std.heap.page_allocator, &reader.interface, frame_head, &scratch); + if (std.mem.startsWith(u8, scratch.items, "{\"type\":\"response.create\",")) { + self.saw_response_create.store(true, .seq_cst); + } + _ = self.request_count.fetchAdd(1, .seq_cst); + for (test_ws_events) |event| { + try writeServerTextFrame(&writer.interface, event); + } + try writer.interface.flush(); + } + } +}; + +fn headerValue(head: []const u8, name: []const u8) ?[]const u8 { + var lines = std.mem.splitSequence(u8, head, "\r\n"); + while (lines.next()) |line| { + const colon = std.mem.indexOfScalar(u8, line, ':') orelse continue; + if (!std.ascii.eqlIgnoreCase(std.mem.trim(u8, line[0..colon], " \t"), name)) continue; + return std.mem.trim(u8, line[colon + 1 ..], " \t"); + } + return null; +} + +fn writeServerTextFrame(writer: *std.Io.Writer, payload: []const u8) !void { + std.debug.assert(payload.len <= std.math.maxInt(u16)); + if (payload.len <= 125) { + try writer.writeAll(&.{ 0x81, @as(u8, @intCast(payload.len)) }); + } else { + var head: [4]u8 = .{ 0x81, 126, 0, 0 }; + std.mem.writeInt(u16, head[2..4], @intCast(payload.len), .big); + try writer.writeAll(&head); + } + try writer.writeAll(payload); +} + +const WsTestHarness = struct { + content: std.ArrayList(u8) = .empty, + admissions: usize = 0, + + fn emit(raw: *anyopaque, event: stream_provider.Event) void { + const self: *@This() = @ptrCast(@alignCast(raw)); + switch (event) { + .content_delta => |chunk| self.content.appendSlice(std.testing.allocator, chunk) catch unreachable, + else => {}, + } + } + + fn admit(raw: *anyopaque) !void { + const self: *@This() = @ptrCast(@alignCast(raw)); + self.admissions += 1; + } +}; + +fn runWsTestRequest( + alloc: Allocator, + harness: *WsTestHarness, + token: []const u8, + url: []const u8, + output_emitted: *bool, +) !stream_provider.Result { + var cancelled = std.atomic.Value(bool).init(false); + var delivery = stream_provider.DeliveryCertainty.init(); + var evidence: stream_provider.AttemptEvidence = .{}; + return streamViaWebSocketAtUrl(alloc, .{ + .credential = .{ .secret = token, .source = .chatgpt_subscription }, + .session_id = "sess-ws-test", + .model = "gpt-test", + .retry_count = 1, + .messages = &.{}, + .tool_choice = .none, + .provider_options = .{}, + .trace_ctx = .{}, + .content_capture_limit = null, + .delivery = &delivery, + .attempt_evidence = &evidence, + .events = .{ .context = harness, .emit_fn = WsTestHarness.emit }, + .admission = .{ .context = harness, .admit_fn = WsTestHarness.admit }, + .cancel_flag = &cancelled, + }, "{\"model\":\"gpt-test\",\"stream\":true}", url, output_emitted); +} + +test "Codex WebSocket transport streams a completion and reuses the connection" { + const alloc = std.testing.allocator; + var fixture = try WsLoopbackFixture.init(.serve); + defer fixture.deinit(); + try fixture.start(); + defer closeCachedConnection(); + + const token = try testChatGptJwt(alloc); + defer alloc.free(token); + const url = try fixture.url(alloc); + defer alloc.free(url); + + var harness = WsTestHarness{}; + defer harness.content.deinit(alloc); + + for (0..2) |round| { + var output_emitted = false; + var result = try runWsTestRequest(alloc, &harness, token, url, &output_emitted); + defer result.deinit(alloc); + try std.testing.expect(output_emitted); + switch (result) { + .completed => |completed| { + try std.testing.expect(completed.completion.content != null); + try std.testing.expectEqualStrings("hello", completed.completion.content.?); + try std.testing.expectEqual(@as(?u64, 10), completed.completion.usage.input_tokens); + }, + else => return error.TestExpectedCompletion, + } + try std.testing.expectEqual(round + 1, harness.admissions); + } + + try std.testing.expectEqualStrings("hellohello", harness.content.items); + try std.testing.expectEqual(@as(usize, 1), fixture.accept_count.load(.seq_cst)); + try std.testing.expectEqual(@as(usize, 2), fixture.request_count.load(.seq_cst)); + try std.testing.expect(fixture.saw_response_create.load(.seq_cst)); + if (fixture.failure) |err| return err; +} + +test "Codex WebSocket handshake rejection is fallback safe" { + const alloc = std.testing.allocator; + var fixture = try WsLoopbackFixture.init(.reject_handshake); + defer fixture.deinit(); + try fixture.start(); + defer closeCachedConnection(); + + const token = try testChatGptJwt(alloc); + defer alloc.free(token); + const url = try fixture.url(alloc); + defer alloc.free(url); + + var harness = WsTestHarness{}; + defer harness.content.deinit(alloc); + + var output_emitted = true; + const result = runWsTestRequest(alloc, &harness, token, url, &output_emitted); + try std.testing.expectError(error.WebSocketHandshakeRejected, result); + try std.testing.expect(!output_emitted); + try std.testing.expect(errorAllowsSseFallback(error.WebSocketHandshakeRejected)); + try std.testing.expectEqual(@as(usize, 1), harness.admissions); +} + +test "Codex WebSocket fallback classification blocks provider refusals" { + try std.testing.expect(!errorAllowsSseFallback(error.OpenAICodexResponseFailed)); + try std.testing.expect(!errorAllowsSseFallback(error.Cancelled)); + try std.testing.expect(!errorAllowsSseFallback(error.OutOfMemory)); + try std.testing.expect(errorAllowsSseFallback(error.WebSocketHandshakeInvalid)); + try std.testing.expect(errorAllowsSseFallback(error.Timeout)); + try std.testing.expect(errorAllowsSseFallback(error.EndOfStream)); +} diff --git a/src/gateway/websocket.zig b/src/gateway/websocket.zig new file mode 100644 index 000000000..9914dbf27 --- /dev/null +++ b/src/gateway/websocket.zig @@ -0,0 +1,256 @@ +//! Minimal RFC 6455 client-side WebSocket codec. +//! +//! This module is transport-policy free: it encodes and decodes frames over +//! caller-provided `std.Io` reader/writer streams and computes handshake keys. +//! Connection management, control-frame policy, and message semantics belong +//! to the caller. + +const std = @import("std"); + +pub const handshake_guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; +pub const sec_key_len = 24; +pub const accept_key_len = 28; +/// Control frames are never allowed a payload above 125 bytes (RFC 6455 5.5). +pub const max_control_payload_bytes: usize = 125; + +/// Encodes 16 random bytes as the `Sec-WebSocket-Key` request header value. +pub fn secKey(random_bytes: [16]u8) [sec_key_len]u8 { + var out: [sec_key_len]u8 = undefined; + const written = std.base64.standard.Encoder.encode(&out, &random_bytes); + std.debug.assert(written.len == sec_key_len); + return out; +} + +/// Computes the `Sec-WebSocket-Accept` value the server must echo for `key`. +pub fn acceptKey(key: []const u8) [accept_key_len]u8 { + var sha1 = std.crypto.hash.Sha1.init(.{}); + sha1.update(key); + sha1.update(handshake_guid); + var digest: [std.crypto.hash.Sha1.digest_length]u8 = undefined; + sha1.final(&digest); + var out: [accept_key_len]u8 = undefined; + const written = std.base64.standard.Encoder.encode(&out, &digest); + std.debug.assert(written.len == accept_key_len); + return out; +} + +pub const Opcode = enum(u4) { + continuation = 0x0, + text = 0x1, + binary = 0x2, + close = 0x8, + ping = 0x9, + pong = 0xA, + _, + + pub fn isControl(self: Opcode) bool { + return @intFromEnum(self) >= 0x8; + } +}; + +pub const FrameHead = struct { + fin: bool, + opcode: Opcode, + masked: bool, + mask: [4]u8, + len: u64, +}; + +pub const ReadFrameError = error{ + WebSocketProtocolError, + WebSocketFrameTooLarge, + ReadFailed, + EndOfStream, +}; + +/// Reads one frame header. `max_payload_bytes` bounds the declared payload +/// length before the caller allocates anything for it. +pub fn readFrameHead(reader: *std.Io.Reader, max_payload_bytes: u64) ReadFrameError!FrameHead { + const b0 = takeByte(reader) catch |err| return err; + const b1 = takeByte(reader) catch |err| return err; + // RSV bits must be zero: no extension was negotiated. + if (b0 & 0x70 != 0) return error.WebSocketProtocolError; + const opcode: Opcode = @enumFromInt(@as(u4, @truncate(b0 & 0x0f))); + const fin = b0 & 0x80 != 0; + const masked = b1 & 0x80 != 0; + var len: u64 = b1 & 0x7f; + if (opcode.isControl() and (!fin or len > max_control_payload_bytes)) { + return error.WebSocketProtocolError; + } + if (len == 126) { + var ext: [2]u8 = undefined; + readAll(reader, &ext) catch |err| return err; + len = std.mem.readInt(u16, &ext, .big); + } else if (len == 127) { + var ext: [8]u8 = undefined; + readAll(reader, &ext) catch |err| return err; + len = std.mem.readInt(u64, &ext, .big); + if (len & (1 << 63) != 0) return error.WebSocketProtocolError; + } + if (len > max_payload_bytes) return error.WebSocketFrameTooLarge; + var mask: [4]u8 = .{ 0, 0, 0, 0 }; + if (masked) readAll(reader, &mask) catch |err| return err; + return .{ .fin = fin, .opcode = opcode, .masked = masked, .mask = mask, .len = len }; +} + +/// Appends the frame payload described by `head` to `out`, unmasking when the +/// sender masked it. The caller has already bounded `head.len`. +pub fn readPayloadInto( + alloc: std.mem.Allocator, + reader: *std.Io.Reader, + head: FrameHead, + out: *std.ArrayList(u8), +) !void { + const len = std.math.cast(usize, head.len) orelse return error.WebSocketFrameTooLarge; + if (len == 0) return; + const start = out.items.len; + try out.resize(alloc, start + len); + errdefer out.shrinkRetainingCapacity(start); + reader.readSliceAll(out.items[start..]) catch |err| switch (err) { + error.EndOfStream => return error.EndOfStream, + error.ReadFailed => return error.ReadFailed, + }; + if (head.masked) { + for (out.items[start..], 0..) |*byte, index| { + byte.* ^= head.mask[index % 4]; + } + } +} + +/// Writes one complete client frame. Client frames are always masked +/// (RFC 6455 5.1); the caller supplies the mask so the codec stays +/// deterministic under test. +pub fn writeClientFrame( + writer: *std.Io.Writer, + opcode: Opcode, + payload: []const u8, + mask: [4]u8, +) !void { + if (opcode.isControl() and payload.len > max_control_payload_bytes) { + return error.WebSocketProtocolError; + } + var head: [14]u8 = undefined; + var head_len: usize = 2; + head[0] = 0x80 | @as(u8, @intFromEnum(opcode)); + if (payload.len <= 125) { + head[1] = 0x80 | @as(u8, @intCast(payload.len)); + } else if (payload.len <= std.math.maxInt(u16)) { + head[1] = 0x80 | 126; + std.mem.writeInt(u16, head[2..4], @intCast(payload.len), .big); + head_len += 2; + } else { + head[1] = 0x80 | 127; + std.mem.writeInt(u64, head[2..10], payload.len, .big); + head_len += 8; + } + @memcpy(head[head_len..][0..4], &mask); + head_len += 4; + try writer.writeAll(head[0..head_len]); + + var chunk: [4096]u8 = undefined; + var offset: usize = 0; + while (offset < payload.len) { + const take = @min(chunk.len, payload.len - offset); + for (payload[offset..][0..take], 0..) |byte, index| { + chunk[index] = byte ^ mask[(offset + index) % 4]; + } + try writer.writeAll(chunk[0..take]); + offset += take; + } +} + +fn takeByte(reader: *std.Io.Reader) ReadFrameError!u8 { + return reader.takeByte() catch |err| switch (err) { + error.EndOfStream => error.EndOfStream, + error.ReadFailed => error.ReadFailed, + }; +} + +fn readAll(reader: *std.Io.Reader, buffer: []u8) ReadFrameError!void { + reader.readSliceAll(buffer) catch |err| switch (err) { + error.EndOfStream => return error.EndOfStream, + error.ReadFailed => return error.ReadFailed, + }; +} + +test "WebSocket accept key matches the RFC 6455 vector" { + const accept = acceptKey("dGhlIHNhbXBsZSBub25jZQ=="); + try std.testing.expectEqualStrings("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=", &accept); +} + +test "WebSocket sec key is 24 base64 bytes" { + const key = secKey(.{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 }); + try std.testing.expectEqual(@as(usize, sec_key_len), key.len); + var decoded: [16]u8 = undefined; + try std.base64.standard.Decoder.decode(&decoded, &key); + try std.testing.expectEqual(@as(u8, 16), decoded[15]); +} + +test "WebSocket client frames round-trip through the frame reader" { + const alloc = std.testing.allocator; + const cases = [_][]const u8{ + "", + "short", + "x" ** 126, + "y" ** 70_000, + }; + for (cases) |payload| { + var out: std.Io.Writer.Allocating = .init(alloc); + defer out.deinit(); + try writeClientFrame(&out.writer, .text, payload, .{ 0x11, 0x22, 0x33, 0x44 }); + + var reader: std.Io.Reader = .fixed(out.written()); + const head = try readFrameHead(&reader, 1 << 20); + try std.testing.expect(head.fin); + try std.testing.expect(head.masked); + try std.testing.expectEqual(Opcode.text, head.opcode); + try std.testing.expectEqual(@as(u64, payload.len), head.len); + + var collected: std.ArrayList(u8) = .empty; + defer collected.deinit(alloc); + try readPayloadInto(alloc, &reader, head, &collected); + try std.testing.expectEqualStrings(payload, collected.items); + } +} + +test "WebSocket frame reader decodes unmasked server frames" { + const alloc = std.testing.allocator; + // FIN text frame, no mask, payload "ok". + const bytes = [_]u8{ 0x81, 0x02, 'o', 'k' }; + var reader: std.Io.Reader = .fixed(&bytes); + const head = try readFrameHead(&reader, 1024); + try std.testing.expect(head.fin); + try std.testing.expect(!head.masked); + try std.testing.expectEqual(@as(u64, 2), head.len); + var collected: std.ArrayList(u8) = .empty; + defer collected.deinit(alloc); + try readPayloadInto(alloc, &reader, head, &collected); + try std.testing.expectEqualStrings("ok", collected.items); +} + +test "WebSocket frame reader rejects reserved bits oversized declarations and long control frames" { + // RSV1 set. + { + const bytes = [_]u8{ 0xC1, 0x00 }; + var reader: std.Io.Reader = .fixed(&bytes); + try std.testing.expectError(error.WebSocketProtocolError, readFrameHead(&reader, 1024)); + } + // Declared length above the caller's bound. + { + const bytes = [_]u8{ 0x81, 0x7E, 0xFF, 0xFF }; + var reader: std.Io.Reader = .fixed(&bytes); + try std.testing.expectError(error.WebSocketFrameTooLarge, readFrameHead(&reader, 1024)); + } + // Fragmented ping. + { + const bytes = [_]u8{ 0x09, 0x00 }; + var reader: std.Io.Reader = .fixed(&bytes); + try std.testing.expectError(error.WebSocketProtocolError, readFrameHead(&reader, 1024)); + } + // Control frame with an oversized payload declaration. + { + const bytes = [_]u8{ 0x89, 0x7E, 0x00, 0x80 }; + var reader: std.Io.Reader = .fixed(&bytes); + try std.testing.expectError(error.WebSocketProtocolError, readFrameHead(&reader, 1024)); + } +} diff --git a/src/main.zig b/src/main.zig index 47163a8d2..c00b14e6b 100644 --- a/src/main.zig +++ b/src/main.zig @@ -4054,6 +4054,8 @@ test { _ = @import("core/auth/provider_catalog.zig"); _ = @import("gateway/openai_codex_models.zig"); _ = @import("gateway/openai_codex.zig"); + _ = @import("gateway/openai_codex_websocket.zig"); + _ = @import("gateway/websocket.zig"); _ = @import("gateway/openai_codex_permission_reviewer.zig"); _ = @import("core/auth/grok_session.zig"); _ = @import("core/auth/grok_oauth.zig"); From f486bb209043efa4025152532e9dbbb665177344 Mon Sep 17 00:00:00 2001 From: mjlbach Date: Sat, 29 Aug 2026 10:39:20 -0700 Subject: [PATCH 2/2] Propagate deadlines, reject masked server frames, add WebSocket e2e Respect request.deadline across the WebSocket path: it bounds the handshake alongside the connect timeout and bounds the stream phase via the bounded cancel watcher, surfacing error.Timeout. Reject masked server frames (RFC 6455 5.1: only clients mask). Add two e2e tests that drive the built binary against a Bun WebSocket loopback server: a completion streamed over one upgraded connection with zero SSE requests, and pre-output fallback to SSE when the upgrade is rejected. Co-Authored-By: Claude Fable 5 --- src/gateway/openai_codex_websocket.zig | 41 +++++-- tests/e2e/tui-auth-source-selection.test.ts | 126 +++++++++++++++++++- 2 files changed, 157 insertions(+), 10 deletions(-) diff --git a/src/gateway/openai_codex_websocket.zig b/src/gateway/openai_codex_websocket.zig index 63a36b6b5..c7940918e 100644 --- a/src/gateway/openai_codex_websocket.zig +++ b/src/gateway/openai_codex_websocket.zig @@ -363,14 +363,20 @@ fn connectBounded( .session_id = request.session_id, .session_key = session_key, }; + 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; + } + } return gateway_client.runBoundedHttpOperation( *WsConnection, alloc, request.cancel_flag, - std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{ - .clock = .awake, - .raw = .fromMilliseconds(connect_timeout_ms), - }), + connect_deadline, &operation, ); } @@ -535,11 +541,19 @@ fn runOnConnection( if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; var cancel_watch_done = std.atomic.Value(bool).init(false); - const cancel_watcher = try gateway_client.spawnHttpCancelWatcher( - &cancel_watch_done, - request.cancel_flag, - conn.netStream(), - ); + const cancel_watcher = if (request.deadline) |deadline| + try gateway_client.spawnHttpCancelWatcherBounded( + &cancel_watch_done, + request.cancel_flag, + deadline, + conn.netStream(), + ) + else + try gateway_client.spawnHttpCancelWatcher( + &cancel_watch_done, + request.cancel_flag, + conn.netStream(), + ); defer { cancel_watch_done.store(true, .seq_cst); cancel_watcher.join(); @@ -590,9 +604,16 @@ fn runOnConnection( fn wsIoError(request: stream_provider.ModelRequest, err: anyerror) anyerror { if (request.cancel_flag.load(.seq_cst)) return error.Cancelled; + if (deadlineExpired(request.deadline)) return error.Timeout; return err; } +fn deadlineExpired(deadline: ?std.Io.Clock.Timestamp) bool { + const limit = deadline orelse return false; + const now = std.Io.Clock.Timestamp.now(io_mod.getIo(), .awake); + return !std.Io.Clock.Timestamp.compare(now, .lt, limit); +} + fn sendFrame(conn: *WsConnection, opcode: websocket.Opcode, payload: []const u8) !void { var mask: [4]u8 = undefined; io_mod.getIo().random(&mask); @@ -652,6 +673,8 @@ fn nextTextMessage( error.EndOfStream => return null, else => return err, }; + // Only clients mask frames (RFC 6455 5.1); a masking server is broken. + if (head.masked) return error.WebSocketProtocolError; switch (head.opcode) { .ping => { var control: std.ArrayList(u8) = .empty; diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index bac87d5b9..9377b9e05 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -444,10 +444,13 @@ function startFakeChatGptOAuth( tokenDelayMs?: number; responseDelayMs?: number; unauthorizedResponses?: number; + websocketMode?: "serve" | "reject"; } = {}, ) { const accessToken = chatgptAccessToken(); let responseCount = 0; + let websocketUpgrades = 0; + let websocketMessages = 0; let models = [ { slug: "gpt-5.6-sol", visibility: "list", supported_in_api: true, supported_reasoning_levels: [{ effort: "max" }, { effort: "high" }], additional_speed_tiers: ["fast"], input_modalities: ["text", "image"], context_window: 272000 }, { slug: "gpt-5.4-mini", visibility: "list", supported_in_api: true, supported_reasoning_levels: [{ effort: "low" }], additional_speed_tiers: [], input_modalities: ["text"], context_window: 128000 }, @@ -461,8 +464,40 @@ function startFakeChatGptOAuth( const server = Bun.serve({ hostname: "127.0.0.1", port: 0, - async fetch(request) { + websocket: { + message(ws: unknown, message: string | Uint8Array) { + websocketMessages += 1; + const text = typeof message === "string" ? message : Buffer.from(message).toString("utf8"); + const socket = ws as { send(data: string): void; close(code?: number, reason?: string): void }; + if (!text.startsWith('{"type":"response.create",')) { + socket.close(1008, "expected response.create"); + return; + } + socket.send('{"type":"response.output_text.delta","delta":"CHATGPT_WEBSOCKET_RESPONSE"}'); + socket.send( + '{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":4,"output_tokens":2}}}', + ); + }, + }, + async fetch(request, bunServer) { const url = new URL(request.url); + if ( + url.pathname === "/chatgpt/responses" && + request.headers.get("upgrade")?.toLowerCase() === "websocket" + ) { + websocketUpgrades += 1; + requests.push({ + method: request.method, + path: url.pathname, + authorization: request.headers.get("authorization"), + body: null, + }); + if (options.websocketMode !== "serve") { + return new Response("websocket upgrade rejected", { status: 403 }); + } + if (bunServer.upgrade(request)) return undefined as unknown as Response; + return new Response("websocket upgrade failed", { status: 400 }); + } const body = url.pathname === "/chatgpt/responses" || url.pathname === "/chatgpt/token" ? await request.text() : null; @@ -524,6 +559,12 @@ function startFakeChatGptOAuth( setModels(next: typeof models) { models = next; }, + get websocketUpgrades() { + return websocketUpgrades; + }, + get websocketMessages() { + return websocketMessages; + }, stop() { server.stop(true); }, @@ -2182,6 +2223,89 @@ test( 60_000, ); +test( + "Codex WebSocket transport streams a response over one upgraded connection", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-ws-")); + gateway = startFakeGateway([]); + chatgptOauth = startFakeChatGptOAuth({ websocketMode: "serve" }); + const env = { + HOME: home, + AI_GATEWAY_API_KEY: ENV_TOKEN, + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_SKIP_ONBOARDING: "1", + FX_AUTO_UPGRADE: "0", + FX_NO_OPEN_BROWSER: "1", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_OPENAI_CODEX_TRANSPORT: "websocket", + ...chatgptOauth.env, + }; + + const login = await runCodexLoginWithBrowser(env); + expect(login.code, `stdout: ${login.stdout}\nstderr: ${login.stderr}`).toBe(0); + + const ask = await runFx(["ask", "--json", "--auto", "--no-save", "Answer directly."], { + env, + timeoutMs: TIMEOUT, + }); + expect(ask.code, `stdout: ${ask.stdout}\nstderr: ${ask.stderr}`).toBe(0); + expect(ask.stdout).toContain("CHATGPT_WEBSOCKET_RESPONSE"); + expect(ask.stdout).not.toContain("CHATGPT_DIRECT_RESPONSE"); + expect(chatgptOauth.websocketUpgrades).toBe(1); + expect(chatgptOauth.websocketMessages).toBe(1); + const ssePosts = chatgptOauth.requests.filter( + (request) => request.path === "/chatgpt/responses" && request.method === "POST", + ); + expect(ssePosts).toHaveLength(0); + const upgradeRequest = chatgptOauth.requests.find( + (request) => request.path === "/chatgpt/responses" && request.method === "GET", + ); + expect(upgradeRequest?.authorization).toBe(`Bearer ${chatgptOauth.accessToken}`); + }, + TIMEOUT, +); + +test( + "Codex WebSocket upgrade rejection falls back to SSE before any output", + async () => { + home = mkdtempSync(join(tmpdir(), "fx-codex-ws-fallback-")); + gateway = startFakeGateway([]); + chatgptOauth = startFakeChatGptOAuth({ websocketMode: "reject" }); + const env = { + HOME: home, + AI_GATEWAY_API_KEY: ENV_TOKEN, + VERCEL_OIDC_TOKEN: undefined, + FX_DISABLE_KEYCHAIN: "1", + FX_SKIP_ONBOARDING: "1", + FX_AUTO_UPGRADE: "0", + FX_NO_OPEN_BROWSER: "1", + FX_GATEWAY_BASE_URL: gateway.baseUrl, + FX_E2E_GATEWAY_MODELS_URL: `${gateway.baseUrl}/coding-agent/v1/models`, + FX_OPENAI_CODEX_TRANSPORT: "websocket", + ...chatgptOauth.env, + }; + + const login = await runCodexLoginWithBrowser(env); + expect(login.code, `stdout: ${login.stdout}\nstderr: ${login.stderr}`).toBe(0); + + const ask = await runFx(["ask", "--json", "--auto", "--no-save", "Answer directly."], { + env, + timeoutMs: TIMEOUT, + }); + expect(ask.code, `stdout: ${ask.stdout}\nstderr: ${ask.stderr}`).toBe(0); + expect(ask.stdout).toContain("CHATGPT_DIRECT_RESPONSE"); + expect(chatgptOauth.websocketUpgrades).toBe(1); + expect(chatgptOauth.websocketMessages).toBe(0); + const ssePosts = chatgptOauth.requests.filter( + (request) => request.path === "/chatgpt/responses" && request.method === "POST", + ); + expect(ssePosts).toHaveLength(1); + }, + TIMEOUT, +); + test( "Codex CLI browser login fetches raw models and replays one 401 without Gateway leakage", async () => {