feat(ai): added websocket frame forwarding for raw SSE debug events

- Added helpers to synthesize RawSseEvent records for inbound, outbound, and malformed Codex WebSocket traffic.
- Passed onSseEvent through websocket transport and stream setup so frames are forwarded to the raw-SSE debug pipeline during streaming.
- Added a stream test that verifies outbound and inbound websocket frames are emitted with SSE-style raw lines for the debug viewer.

packages/ai/CHANGELOG.md

@@ -9,6 +9,7 @@
 
 - Added `PI_CODEX_WEBSOCKET_FIRST_EVENT_TIMEOUT_MS` and `PI_CODEX_WEBSOCKET_IDLE_TIMEOUT_MS` options to tune Codex WebSocket timeout behavior before fallback
 - Added `AuthStorage.getOAuthAccess` to return a refreshed OAuth access token with identity metadata (`accountId`, `email`, `projectId`, `enterpriseUrl`) for callers that need bearer-token headers together
+- Added Codex WebSocket forwarding to the `onSseEvent` observer so the raw provider-stream debug viewer captures the inbound JSON frames and the outbound request frame from the WS transport using the same synthesized SSE-wire shape (`event:` + `data:` lines, prefixed with a `: ws ← <type>` (inbound) or `: ws → <type>` (outbound) comment).
 
 ### Changed
 

packages/ai/src/providers/openai-codex-responses.ts

@@ -29,6 +29,7 @@ import {
 	type FetchImpl,
 	type Model,
 	type ProviderSessionState,
+	type RawSseEvent,
 	resolveServiceTier,
 	type ServiceTier,
 	type StreamFunction,
@@ -50,6 +51,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream";
 import { finalizeErrorMessage, type RawHttpRequestDump } from "../utils/http-inspector";
 import { parseStreamingJson } from "../utils/json-parse";
 import { adaptSchemaForStrict, NO_STRICT, sanitizeSchemaForOpenAIResponses, toolWireSchema } from "../utils/schema";
+import { notifyRawSseEvent } from "../utils/sse-debug";
 import { compactGrammarDefinition } from "./grammar";
 import { CODEX_BASE_URL, getCodexAccountId, OPENAI_HEADER_VALUES, OPENAI_HEADERS } from "./openai-codex/constants";
 import {
@@ -356,6 +358,52 @@ function extractCodexWebSocketHandshakeHeaders(socket: Bun.WebSocket, openEvent?
 	);
 }
 
+// Synthesizes a `RawSseEvent` for a Codex WebSocket frame so the same debug
+// pipeline used for HTTP SSE (`onSseEvent` → `RawSseDebugBuffer.recordEvent`)
+// also captures WebSocket traffic. The `raw` array mirrors SSE wire format
+// (one line per field) so the existing TUI viewer renders it identically:
+//   : ws ← <type>
+//   event: <type>
+//   data: <json>
+// Outbound (client → server) uses `: ws → <type>`. The viewer pretty-prints
+// `data:` JSON lines, so we keep the wire JSON single-line here and let the
+// renderer expand it.
+function notifyCodexWebSocketInbound(
+	observer: ((event: RawSseEvent) => void) | undefined,
+	parsed: Record<string, unknown>,
+	text: string,
+): void {
+	const type = typeof parsed.type === "string" ? parsed.type : null;
+	const raw: string[] = [`: ws ← ${type ?? "(untyped)"}`];
+	if (type) raw.push(`event: ${type}`);
+	raw.push(`data: ${text}`);
+	notifyRawSseEvent(observer, { event: type, data: text, raw });
+}
+
+function notifyCodexWebSocketOutbound(
+	observer: ((event: RawSseEvent) => void) | undefined,
+	request: Record<string, unknown>,
+	payload: string,
+): void {
+	const type = typeof request.type === "string" ? request.type : null;
+	const raw: string[] = [`: ws → ${type ?? "(untyped)"}`];
+	if (type) raw.push(`event: ${type}`);
+	raw.push(`data: ${payload}`);
+	notifyRawSseEvent(observer, { event: type, data: payload, raw });
+}
+
+function notifyCodexWebSocketMalformed(
+	observer: ((event: RawSseEvent) => void) | undefined,
+	data: unknown,
+	error: unknown,
+): void {
+	const text = typeof data === "string" ? data : "";
+	const reason = error instanceof Error ? error.message : String(error);
+	const raw: string[] = [`: ws ← (parse-error: ${reason})`];
+	if (text) raw.push(`data: ${text}`);
+	notifyRawSseEvent(observer, { event: "parse_error", data: text, raw });
+}
+
 /** @internal Exported for tests. */
 export function normalizeCodexToolChoice(
 	choice: ToolChoice | undefined,
@@ -613,7 +661,13 @@ async function openInitialCodexEventStream(
 		let websocketRetries = 0;
 		while (true) {
 			try {
-				return await openCodexWebSocketTransport(requestContext, requestSetup, websocketState, websocketRetries);
+				return await openCodexWebSocketTransport(
+					requestContext,
+					requestSetup,
+					websocketState,
+					websocketRetries,
+					options ? event => options.onSseEvent?.(event, model) : undefined,
+				);
 			} catch (error) {
 				const websocketError = error instanceof Error ? error : new Error(String(error));
 				const isFatal = isCodexWebSocketFatalError(websocketError);
@@ -644,6 +698,7 @@ async function openCodexWebSocketTransport(
 	requestSetup: CodexRequestSetup,
 	websocketState: CodexWebSocketSessionState,
 	retry: number,
+	onSseEvent?: (event: RawSseEvent) => void,
 ): Promise<{
 	eventStream: AsyncGenerator<Record<string, unknown>>;
 	requestBodyForState: RequestBody;
@@ -676,6 +731,7 @@ async function openCodexWebSocketTransport(
 		websocketRequest,
 		websocketState,
 		requestSetup.requestSignal,
+		onSseEvent,
 	);
 	return { eventStream, requestBodyForState, transport: "websocket" };
 }
@@ -721,6 +777,7 @@ async function reopenCodexWebSocketRuntimeStream(
 			context.requestSetup,
 			state,
 			runtime.websocketStreamRetries,
+			context.options ? event => context.options?.onSseEvent?.(event, context.model) : undefined,
 		);
 		runtime.eventStream = next.eventStream;
 		runtime.requestBodyForState = next.requestBodyForState;
@@ -1865,6 +1922,7 @@ class CodexWebSocketConnection {
 	#waiters: Array<() => void> = [];
 	#connectPromise?: Promise<void>;
 	#activeRequest = false;
+	#streamObserver?: (event: RawSseEvent) => void;
 
 	constructor(url: string, headers: Record<string, string>, options: CodexWebSocketConnectionOptions) {
 		this.#url = url;
@@ -1983,8 +2041,10 @@ class CodexWebSocketConnection {
 						parsed.message = inner.message;
 					}
 				}
+				notifyCodexWebSocketInbound(this.#streamObserver, parsed, text);
 				this.#push(parsed);
 			} catch (error) {
+				notifyCodexWebSocketMalformed(this.#streamObserver, event.data, error);
 				this.#push(createCodexWebSocketTransportError(String(error)));
 			}
 		};
@@ -2000,6 +2060,7 @@ class CodexWebSocketConnection {
 	async *streamRequest(
 		request: Record<string, unknown>,
 		signal?: AbortSignal,
+		onSseEvent?: (event: RawSseEvent) => void,
 	): AsyncGenerator<Record<string, unknown>> {
 		if (!this.#socket || this.#socket.readyState !== WebSocket.OPEN) {
 			throw createCodexWebSocketTransportError("websocket connection is unavailable");
@@ -2008,6 +2069,7 @@ class CodexWebSocketConnection {
 			throw createCodexWebSocketTransportError("websocket request already in progress");
 		}
 		this.#activeRequest = true;
+		this.#streamObserver = onSseEvent;
 		const onAbort = () => {
 			this.close("aborted");
 			this.#push(createCodexWebSocketTransportError("request was aborted"));
@@ -2021,7 +2083,9 @@ class CodexWebSocketConnection {
 		}
 
 		try {
-			this.#socket.send(JSON.stringify(request));
+			const requestPayload = JSON.stringify(request);
+			notifyCodexWebSocketOutbound(onSseEvent, request, requestPayload);
+			this.#socket.send(requestPayload);
 			let sawFirstEvent = false;
 			let lastProgressAt = Date.now();
 			while (true) {
@@ -2060,6 +2124,7 @@ class CodexWebSocketConnection {
 			}
 		} finally {
 			this.#activeRequest = false;
+			this.#streamObserver = undefined;
 			if (signal) {
 				signal.removeEventListener("abort", onAbort);
 			}
@@ -2193,9 +2258,10 @@ async function openCodexWebSocketEventStream(
 	request: Record<string, unknown>,
 	state: CodexWebSocketSessionState,
 	signal?: AbortSignal,
+	onSseEvent?: (event: RawSseEvent) => void,
 ): Promise<AsyncGenerator<Record<string, unknown>>> {
 	const connection = await getOrCreateCodexWebSocketConnection(state, url, headers, signal);
-	return connection.streamRequest(request, signal);
+	return connection.streamRequest(request, signal, onSseEvent);
 }
 
 function createCodexHeaders(

packages/ai/src/types.ts

@@ -336,7 +336,10 @@ export interface StreamOptions {
 	 */
 	onResponse?: (response: ProviderResponseMetadata, model?: Model<Api>) => void | Promise<void>;
 	/**
-	 * Optional callback for raw Server-Sent Events as they arrive from HTTP streaming providers.
+	 * Optional callback for raw Server-Sent Events as they arrive from HTTP streaming providers,
+	 * plus synthesized SSE-shaped frames for the Codex WebSocket transport (one synthetic frame
+	 * per JSON request/response message). WebSocket frames are tagged with a leading
+	 * `: ws → <type>` (outbound) or `: ws ← <type>` (inbound) comment line in `RawSseEvent.raw`.
 	 *
 	 * Diagnostic only: provider implementations must ignore callback failures and must not
 	 * let observers alter stream contents.

packages/ai/test/openai-codex-stream.test.ts

@@ -359,6 +359,63 @@ describe("openai-codex streaming", () => {
 		expect(result.stopReason).toBe("stop");
 	});
 
+	it("forwards websocket frames through onSseEvent for the raw-SSE debug viewer", async () => {
+		const tempDir = TempDir.createSync("@pi-codex-stream-");
+		setAgentDir(tempDir.path());
+		const token = createCodexTestToken();
+
+		class ObservedWebSocket extends MockWebSocket {
+			constructor(url: string, options?: { headers?: WsHeaders }) {
+				super(url, options);
+				this.scheduleOpen();
+			}
+
+			send(): void {
+				this.emitCodexResponse({ messageId: "msg_obs", responseId: "resp_obs", text: "Observed" });
+			}
+		}
+		global.WebSocket = ObservedWebSocket as unknown as typeof WebSocket;
+
+		const observed: Array<{ event: string | null; data: string; raw: string[] }> = [];
+		const result = await streamOpenAICodexResponses(
+			createCodexTestModel("https://chatgpt.com/backend-api"),
+			createCodexTestContext(),
+			{
+				apiKey: token,
+				sessionId: "ws-observer-session",
+				providerSessionState: new Map<string, ProviderSessionState>(),
+				onSseEvent: event => {
+					observed.push({ event: event.event, data: event.data, raw: [...event.raw] });
+				},
+			},
+		).result();
+
+		expect(result.stopReason).toBe("stop");
+
+		// First record is the outbound request frame (the JSON we sent).
+		const [outbound, ...inbound] = observed;
+		expect(outbound).toBeDefined();
+		expect(outbound.raw[0]).toMatch(/^: ws → /);
+		expect(outbound.data.length).toBeGreaterThan(0);
+		expect(() => JSON.parse(outbound.data)).not.toThrow();
+
+		// Inbound frames mirror the Codex response sequence emitted by `emitCodexResponse`.
+		expect(inbound.map(e => e.event)).toEqual([
+			"response.output_item.added",
+			"response.content_part.added",
+			"response.output_text.delta",
+			"response.output_item.done",
+			"response.done",
+		]);
+		for (const event of inbound) {
+			expect(event.raw[0]).toBe(`: ws ← ${event.event}`);
+			// Synthesized SSE wire shape: `event:` line then `data:` line.
+			expect(event.raw[1]).toBe(`event: ${event.event}`);
+			expect(event.raw[2]).toBe(`data: ${event.data}`);
+			expect(JSON.parse(event.data)).toMatchObject({ type: event.event });
+		}
+	});
+
 	it("omits request-body headers and replaces stale beta headers for websocket handshakes", async () => {
 		const tempDir = TempDir.createSync("@pi-codex-stream-");
 		setAgentDir(tempDir.path());

packages/coding-agent/CHANGELOG.md

@@ -10,6 +10,7 @@
 
 - Added `codex` and `gemini` to the web search provider settings so users can configure OpenAI and Gemini web search directly from provider selection
 - Added OpenAI (`codex`) and Gemini web search options with updated setup descriptions for `omp /login openai-codex` and Gemini OAuth login
+- Added pretty-printing for wide JSON `data:` payloads in the raw provider-stream debug viewer so streamed event bodies expand across multiple `data:` lines instead of getting clipped by the per-line truncator, and updated the viewer header to read `raw provider stream (SSE + WS)` now that Codex WebSocket frames also flow through the buffer
 
 ### Changed
 

packages/coding-agent/src/debug/raw-sse.ts

@@ -2,15 +2,56 @@ import { type Component, matchesKey, padding, replaceTabs, truncateToWidth, visi
 import { sanitizeText } from "@oh-my-pi/pi-utils";
 import { theme } from "../modes/theme/theme";
 import { copyToClipboard } from "../utils/clipboard";
-import { formatRawSseIsoTime, type RawSseDebugBuffer, rawSseRecordLines } from "./raw-sse-buffer";
+import {
+	formatRawSseIsoTime,
+	type RawSseDebugBuffer,
+	type RawSseDebugRecord,
+	rawSseRecordLines,
+} from "./raw-sse-buffer";
 
 const MIN_VIEWER_WIDTH = 20;
 const VIEWER_FRAME_LINES = 5;
+// `data:` lines below this width render fine on a single row; anything wider gets pretty-printed
+// across multiple `data:` lines so streamed JSON blobs stop getting clipped by `truncateToWidth`.
+const PRETTY_PRINT_DATA_THRESHOLD = 100;
 
 function sanitizeFrameLine(line: string, width: number): string {
 	return truncateToWidth(replaceTabs(sanitizeText(line)), width);
 }
 
+// Walks the SSE wire lines and replaces single-line `data: <json>` payloads with
+// multi-line `data: <indented-json>` entries when the JSON is wide enough to clip.
+// Multi-line `data:` is still valid SSE (the spec joins lines with `\n`), so the
+// transformed view round-trips back to the same event when copied.
+/** @internal Exported for tests. */
+export function expandPrettyDataLines(raw: readonly string[]): string[] {
+	const out: string[] = [];
+	for (const line of raw) {
+		if (!line.startsWith("data: ") || line.length <= PRETTY_PRINT_DATA_THRESHOLD) {
+			out.push(line);
+			continue;
+		}
+		const body = line.slice("data: ".length);
+		const trimmed = body.trim();
+		if (trimmed.length === 0 || (trimmed[0] !== "{" && trimmed[0] !== "[")) {
+			out.push(line);
+			continue;
+		}
+		let parsed: unknown;
+		try {
+			parsed = JSON.parse(trimmed);
+		} catch {
+			out.push(line);
+			continue;
+		}
+		const pretty = JSON.stringify(parsed, null, 2);
+		for (const prettyLine of pretty.split("\n")) {
+			out.push(`data: ${prettyLine}`);
+		}
+	}
+	return out;
+}
+
 export interface RawSseViewerOptions {
 	buffer: RawSseDebugBuffer;
 	terminalRows: number;
@@ -30,6 +71,12 @@ export class RawSseViewerComponent implements Component {
 	#followTail = true;
 	#lastRenderWidth = MIN_VIEWER_WIDTH;
 	#statusMessage: string | undefined;
+	// Pretty-printed wire lines keyed by `record.sequence`. Pretty-printing is
+	// the JSON.parse + JSON.stringify per `data:` line, so we cache the result —
+	// the render path runs on every keypress and from `#maxScrollOffset()`.
+	// Sequences are monotonic; we prune entries below the oldest live record
+	// after each render so the cache tracks the buffer's eviction window.
+	readonly #prettyLinesCache = new Map<number, string[]>();
 
 	constructor(options: RawSseViewerOptions) {
 		this.#buffer = options.buffer;
@@ -131,8 +178,9 @@ export class RawSseViewerComponent implements Component {
 			);
 			lines.push("");
 		}
+		const firstSequence = snapshot.records[0]?.sequence;
 		for (const record of snapshot.records) {
-			for (const line of rawSseRecordLines(record)) {
+			for (const line of this.#prettyLinesFor(record)) {
 				lines.push(sanitizeFrameLine(line, innerWidth));
 			}
 			if (record.kind === "event" && record.truncated) {
@@ -140,14 +188,31 @@ export class RawSseViewerComponent implements Component {
 			}
 			lines.push("");
 		}
+		if (firstSequence !== undefined) this.#pruneCache(firstSequence);
 		return lines;
 	}
 
+	#prettyLinesFor(record: RawSseDebugRecord): string[] {
+		const cached = this.#prettyLinesCache.get(record.sequence);
+		if (cached) return cached;
+		const expanded = expandPrettyDataLines(rawSseRecordLines(record));
+		this.#prettyLinesCache.set(record.sequence, expanded);
+		return expanded;
+	}
+
+	#pruneCache(firstSequence: number): void {
+		// Bounded by the buffer eviction rate; with `MAX_RAW_SSE_EVENTS = 1000`
+		// this rarely runs and only walks freshly-evicted entries.
+		for (const key of this.#prettyLinesCache.keys()) {
+			if (key < firstSequence) this.#prettyLinesCache.delete(key);
+		}
+	}
+
 	#summaryText(): string {
 		const snapshot = this.#buffer.snapshot();
 		const last = snapshot.lastUpdatedAt ? ` last=${formatRawSseIsoTime(snapshot.lastUpdatedAt)}` : "";
 		const follow = this.#followTail ? "follow:on" : "follow:off";
-		return ` # raw SSE | events=${snapshot.totalEvents} records=${snapshot.records.length}${last} | ${follow} | Esc back Ctrl+C copy End follow`;
+		return ` # raw provider stream (SSE + WS) | events=${snapshot.totalEvents} records=${snapshot.records.length}${last} | ${follow} | Esc back Ctrl+C copy End follow`;
 	}
 
 	#statusText(): string {

packages/coding-agent/test/debug/raw-sse-pretty.test.ts

@@ -0,0 +1,70 @@
+import { describe, expect, it } from "bun:test";
+import { expandPrettyDataLines } from "../../src/debug/raw-sse";
+
+// Wide enough that `truncateToWidth` would clip the payload in the viewer; matches
+// what real Codex `response.output_item.done` frames look like on the wire.
+function wideObjectLine(): string {
+	const payload = {
+		type: "response.output_item.done",
+		item: {
+			type: "message",
+			id: "msg_1234567890",
+			role: "assistant",
+			status: "completed",
+			content: [{ type: "output_text", text: "Hello there, this is a long enough message" }],
+		},
+	};
+	return `data: ${JSON.stringify(payload)}`;
+}
+
+describe("expandPrettyDataLines", () => {
+	it("expands wide JSON `data:` payloads into multi-line indented `data:` entries", () => {
+		const input = [`: ws ← response.output_item.done`, `event: response.output_item.done`, wideObjectLine()];
+
+		const out = expandPrettyDataLines(input);
+
+		// Comment + event: lines pass through.
+		expect(out[0]).toBe(input[0]);
+		expect(out[1]).toBe(input[1]);
+
+		// Original payload spans many lines now, each prefixed with `data: `.
+		const dataLines = out.slice(2);
+		expect(dataLines.length).toBeGreaterThan(1);
+		for (const line of dataLines) {
+			expect(line.startsWith("data: ")).toBe(true);
+		}
+
+		// The expanded JSON, with `data: ` stripped, round-trips to the original payload.
+		const rejoined = dataLines.map(line => line.slice("data: ".length)).join("\n");
+		expect(JSON.parse(rejoined)).toEqual(JSON.parse(wideObjectLine().slice("data: ".length)));
+	});
+
+	it("leaves short single-line payloads alone so small deltas stay compact", () => {
+		const input = [
+			`: ws ← response.output_text.delta`,
+			`event: response.output_text.delta`,
+			`data: {"type":"response.output_text.delta","delta":"hi"}`,
+		];
+		expect(expandPrettyDataLines(input)).toEqual([...input]);
+	});
+
+	it("falls back to the raw line when the payload is wide but not JSON", () => {
+		// Wide enough to cross the threshold but doesn't start with `{` or `[`.
+		const wideNonJson = `data: ${"x".repeat(200)}`;
+		const out = expandPrettyDataLines([wideNonJson]);
+		expect(out).toEqual([wideNonJson]);
+	});
+
+	it("falls back when the payload is wide JSON-looking text but parses as invalid", () => {
+		const wideBrokenJson = `data: {"unterminated":"${"x".repeat(200)}`;
+		const out = expandPrettyDataLines([wideBrokenJson]);
+		expect(out).toEqual([wideBrokenJson]);
+	});
+
+	it("preserves non-`data:` lines (event/comment) verbatim regardless of length", () => {
+		const wideComment = `: ${"x".repeat(300)}`;
+		const wideEvent = `event: ${"x".repeat(300)}`;
+		const out = expandPrettyDataLines([wideComment, wideEvent]);
+		expect(out).toEqual([wideComment, wideEvent]);
+	});
+});