From 44f85beb4a5bc07ac8c05d83a6b6da5bc3a4ce42 Mon Sep 17 00:00:00 2001 From: can1357 Date: Tue, 26 May 2026 10:46:00 +0200 Subject: [PATCH] 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 | 1 + .../src/providers/openai-codex-responses.ts | 72 ++++++++++++++++++- packages/ai/src/types.ts | 5 +- packages/ai/test/openai-codex-stream.test.ts | 57 +++++++++++++++ packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/debug/raw-sse.ts | 71 +++++++++++++++++- .../test/debug/raw-sse-pretty.test.ts | 70 ++++++++++++++++++ 7 files changed, 270 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/test/debug/raw-sse-pretty.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index ea61a96fc..2e5553395 100644 --- a/packages/ai/CHANGELOG.md +++ b/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 ← ` (inbound) or `: ws → ` (outbound) comment). ### Changed diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 328d2a497..b6eda8ece 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/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 ← +// event: +// data: +// Outbound (client → server) uses `: ws → `. 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, + 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, + 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>; 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; #activeRequest = false; + #streamObserver?: (event: RawSseEvent) => void; constructor(url: string, headers: Record, 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, signal?: AbortSignal, + onSseEvent?: (event: RawSseEvent) => void, ): AsyncGenerator> { 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, state: CodexWebSocketSessionState, signal?: AbortSignal, + onSseEvent?: (event: RawSseEvent) => void, ): Promise>> { const connection = await getOrCreateCodexWebSocketConnection(state, url, headers, signal); - return connection.streamRequest(request, signal); + return connection.streamRequest(request, signal, onSseEvent); } function createCodexHeaders( diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 60db6d081..d2c256781 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -336,7 +336,10 @@ export interface StreamOptions { */ onResponse?: (response: ProviderResponseMetadata, model?: Model) => void | Promise; /** - * 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 → ` (outbound) or `: ws ← ` (inbound) comment line in `RawSseEvent.raw`. * * Diagnostic only: provider implementations must ignore callback failures and must not * let observers alter stream contents. diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 4e5c42fd6..6b565ca1c 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/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(), + 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()); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 7a6fab6e3..eee47d3b5 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/src/debug/raw-sse.ts b/packages/coding-agent/src/debug/raw-sse.ts index c3d812264..a369fb246 100644 --- a/packages/coding-agent/src/debug/raw-sse.ts +++ b/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: ` payloads with +// multi-line `data: ` 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(); 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 { diff --git a/packages/coding-agent/test/debug/raw-sse-pretty.test.ts b/packages/coding-agent/test/debug/raw-sse-pretty.test.ts new file mode 100644 index 000000000..2dc12e11e --- /dev/null +++ b/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]); + }); +});