Skip to content

refactor(responses): single-pass SSE payload rewrite (#588 follow-up) #602

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Wibias merged 1 commit into from
Jul 28, 2026
Merged

refactor(responses): single-pass SSE payload rewrite (#588 follow-up) #602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 9 additions & 86 deletions src/server/responses-image-gen-repair.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { collectResponsesToolGroups } from "../responses/tool-groups";
import { relaySseWithPayloadRewrite, type SsePayloadRewrite } from "./sse-payload-rewrite";

interface NamespacedTool {
namespace: string;
Expand Down Expand Up @@ -96,45 +97,12 @@ export function restoreImageGenCallsInJson(
return restored.changed ? JSON.stringify(restored.value) : text;
}

/** Split one complete SSE event block while retaining its original blank-line delimiter. */
function nextSseBlock(buffer: string): { block: string; delimiter: string; rest: string } | null {
const match = buffer.match(/\r?\n\r?\n/);
if (!match || match.index === undefined) return null;
return {
block: buffer.slice(0, match.index),
delimiter: match[0],
rest: buffer.slice(match.index + match[0].length),
};
}

/** Join all data lines from one SSE event according to the event-stream field rules. */
function sseDataPayload(block: string): string | null {
const data: string[] = [];
for (const line of block.split(/\r?\n/)) {
if (!line.startsWith("data:")) continue;
const value = line.slice(5);
data.push(value.startsWith(" ") ? value.slice(1) : value);
}
return data.length > 0 ? data.join("\n") : null;
}

/** Replace an SSE event's data field while preserving non-data fields and newline style. */
function replaceSseDataPayload(block: string, payload: string): string {
const newline = block.includes("\r\n") ? "\r\n" : "\n";
const lines = block.split(/\r?\n/);
const rewritten: string[] = [];
let replaced = false;
for (const line of lines) {
if (!line.startsWith("data:")) {
rewritten.push(line);
continue;
}
if (!replaced) {
rewritten.push(`data: ${payload}`);
replaced = true;
}
}
return replaced ? rewritten.join(newline) : block;
/** Payload rewrite for composition with other client-facing SSE transforms. */
export function createImageGenCallRestoreRewrite(
aliases: ReadonlyMap<string, NamespacedTool>,
): SsePayloadRewrite | undefined {
if (aliases.size === 0) return undefined;
return (payload) => restoreImageGenCallsInJson(payload, aliases);
}

/**
Expand All @@ -145,51 +113,6 @@ export function relaySseWithImageGenCallRestore(
body: ReadableStream<Uint8Array>,
aliases: ReadonlyMap<string, NamespacedTool>,
): ReadableStream<Uint8Array> {
if (aliases.size === 0) return body;
const reader = body.getReader();
const decoder = new TextDecoder();
const encoder = new TextEncoder();
let buffer = "";

const emitProcessedBlocks = (
controller: ReadableStreamDefaultController<Uint8Array>,
flushFinal = false,
): void => {
let next: { block: string; delimiter: string; rest: string } | null;
while ((next = nextSseBlock(buffer))) {
buffer = next.rest;
const payload = sseDataPayload(next.block);
const restoredPayload = payload ? restoreImageGenCallsInJson(payload, aliases) : undefined;
const block = payload && restoredPayload !== undefined && restoredPayload !== payload
? replaceSseDataPayload(next.block, restoredPayload)
: next.block;
controller.enqueue(encoder.encode(block + next.delimiter));
}
if (flushFinal && buffer.length > 0) {
const payload = sseDataPayload(buffer);
const restoredPayload = payload ? restoreImageGenCallsInJson(payload, aliases) : undefined;
const block = payload && restoredPayload !== undefined && restoredPayload !== payload
? replaceSseDataPayload(buffer, restoredPayload)
: buffer;
controller.enqueue(encoder.encode(block));
buffer = "";
}
};

return new ReadableStream<Uint8Array>({
async pull(controller) {
const { done, value } = await reader.read();
if (done) {
buffer += decoder.decode();
emitProcessedBlocks(controller, true);
controller.close();
return;
}
buffer += decoder.decode(value, { stream: true });
emitProcessedBlocks(controller);
},
cancel(reason) {
reader.cancel(reason).catch(() => {});
},
});
const rewrite = createImageGenCallRestoreRewrite(aliases);
return rewrite ? relaySseWithPayloadRewrite(body, rewrite) : body;
}
95 changes: 10 additions & 85 deletions src/server/responses-item-id-repair.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { randomUUID } from "node:crypto";
import type { ResponsesItemIdRepairConfig } from "../types";
import { relaySseWithPayloadRewrite, type SsePayloadRewrite } from "./sse-payload-rewrite";

type RepairableItemType = "message" | "reasoning";

Expand Down Expand Up @@ -35,44 +36,6 @@ function isPlainObject(value: unknown): value is Record<string, unknown> {
return !!value && typeof value === "object" && !Array.isArray(value);
}

function nextSseBlock(buffer: string): { block: string; delimiter: string; rest: string } | null {
const match = buffer.match(/\r?\n\r?\n/);
if (!match || match.index === undefined) return null;
return {
block: buffer.slice(0, match.index),
delimiter: match[0],
rest: buffer.slice(match.index + match[0].length),
};
}

function sseDataPayload(block: string): string | null {
const data: string[] = [];
for (const line of block.split(/\r?\n/)) {
if (!line.startsWith("data:")) continue;
const value = line.slice(5);
data.push(value.startsWith(" ") ? value.slice(1) : value);
}
return data.length > 0 ? data.join("\n") : null;
}

function replaceSseDataPayload(block: string, payload: string): string {
const newline = block.includes("\r\n") ? "\r\n" : "\n";
const lines = block.split(/\r?\n/);
const rewritten: string[] = [];
let replaced = false;
for (const line of lines) {
if (!line.startsWith("data:")) {
rewritten.push(line);
continue;
}
if (!replaced) {
rewritten.push(`data: ${payload}`);
replaced = true;
}
}
return replaced ? rewritten.join(newline) : block;
}

function asOutputIndex(value: unknown): number | null {
return typeof value === "number" && Number.isInteger(value) && value >= 0 ? value : null;
}
Expand Down Expand Up @@ -221,57 +184,19 @@ function repairEventPayload(
* type에 재사용해도 function_call id/call_id는 보존된다. opt-in 게이트웨이는 sequential streams에서도
* 고유한 canonical id를 얻지만, 보정이 필요한 경우에만 JS stream 재작성 비용을 지불한다.
*/
/** Stateful payload rewrite for composition with other client-facing SSE transforms. */
export function createResponsesItemIdPayloadRewrite(
config: ResponsesItemIdRepairConfig,
): SsePayloadRewrite {
const state = createRepairState(config);
return (payload) => repairEventPayload(payload, state);
}

export function relaySseWithResponsesItemIdRepair(
body: ReadableStream<Uint8Array>,
config: ResponsesItemIdRepairConfig,
): ReadableStream<Uint8Array> {
const reader = body.getReader();
const decoder = new TextDecoder();
const encoder = new TextEncoder();
const state = createRepairState(config);
let buffer = "";

const emitProcessedBlocks = (
controller: ReadableStreamDefaultController<Uint8Array>,
flushFinal = false,
): void => {
let next: { block: string; delimiter: string; rest: string } | null;
while ((next = nextSseBlock(buffer))) {
buffer = next.rest;
const payload = sseDataPayload(next.block);
const repairedPayload = payload ? repairEventPayload(payload, state) : undefined;
const block = payload && repairedPayload !== undefined && repairedPayload !== payload
? replaceSseDataPayload(next.block, repairedPayload)
: next.block;
controller.enqueue(encoder.encode(block + next.delimiter));
}
if (flushFinal && buffer.length > 0) {
const payload = sseDataPayload(buffer);
const repairedPayload = payload ? repairEventPayload(payload, state) : undefined;
const block = payload && repairedPayload !== undefined && repairedPayload !== payload
? replaceSseDataPayload(buffer, repairedPayload)
: buffer;
controller.enqueue(encoder.encode(block));
buffer = "";
}
};

return new ReadableStream<Uint8Array>({
async pull(controller) {
const { done, value } = await reader.read();
if (done) {
buffer += decoder.decode();
emitProcessedBlocks(controller, true);
controller.close();
return;
}
buffer += decoder.decode(value, { stream: true });
emitProcessedBlocks(controller);
},
cancel(reason) {
reader.cancel(reason).catch(() => {});
},
});
return relaySseWithPayloadRewrite(body, createResponsesItemIdPayloadRewrite(config));
}

export function hasResponsesItemIdRepair(config: ResponsesItemIdRepairConfig | undefined): boolean {
Expand Down
24 changes: 17 additions & 7 deletions src/server/responses/core.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,12 +132,16 @@ import {
import { relaySseEagerBounded } from "../relay-eager";
import { decideEagerRelay } from "../../lib/bun-stream-caps";
import { cancelBodyOnAbort } from "../../lib/abort";
import { hasResponsesItemIdRepair, relaySseWithResponsesItemIdRepair } from "../responses-item-id-repair";
import {
createResponsesItemIdPayloadRewrite,
hasResponsesItemIdRepair,
} from "../responses-item-id-repair";
import {
createImageGenCallRestoreRewrite,
imageGenToolCallAliases,
relaySseWithImageGenCallRestore,
restoreImageGenCallsInJson,
} from "../responses-image-gen-repair";
import { composeSsePayloadRewrites, relaySseWithPayloadRewrite } from "../sse-payload-rewrite";
import type { EffectiveSubagentRoster, SpawnAgentSurface } from "../../codex/catalog";

import { buildToolBridgeMaps, collabSurface, injectDeveloperMessage, multiAgentGuidanceText } from "./collaboration";
Expand Down Expand Up @@ -1663,13 +1667,19 @@ export async function handleResponses(
// win32 must keep the pure native relay (Bun#32111 JS-sink segfault); elsewhere a JS pull
// relay is established practice (relayWithAbort, relaySseWithHeartbeat) and lets a
// mid-stream reset end with a clean response.failed terminal instead of a raw socket error.
const restoredBody = relaySseWithImageGenCallRestore(nativeBody, imageGenCallAliases);
const repairedBody = hasResponsesItemIdRepair(repairConfig)
? relaySseWithResponsesItemIdRepair(restoredBody, repairConfig!)
: restoredBody;
// Compose opt-in payload rewrites into one parse/stringify pass (image-gen restore first).
const payloadRewrites = [
createImageGenCallRestoreRewrite(imageGenCallAliases),
hasResponsesItemIdRepair(repairConfig)
? createResponsesItemIdPayloadRewrite(repairConfig!)
: undefined,
].filter((rewrite): rewrite is NonNullable<typeof rewrite> => rewrite !== undefined);
const rewrittenBody = payloadRewrites.length > 0
? relaySseWithPayloadRewrite(nativeBody, composeSsePayloadRewrites(...payloadRewrites))
: nativeBody;
const clientBody = process.platform === "win32" && !needsClientRewrite
? nativeBody
: relaySseWithFailedTail(repairedBody, upstream);
: relaySseWithFailedTail(rewrittenBody, upstream);
return markNativePassthroughSseResponse(new Response(clientBody, {
status: upstreamResponse.status,
headers,
Expand Down
116 changes: 116 additions & 0 deletions src/server/sse-payload-rewrite.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
/**
* Shared client-facing SSE payload rewrite shell.
*
* Multiple opt-in transforms (image-gen namespace restore, item-id repair, …) compose into one
* parse/stringify pass so a tee'd stream is not re-framed twice per event.
*/

export type SsePayloadRewrite = (payload: string) => string;

/** Split one complete SSE event block while retaining its original blank-line delimiter. */
export function nextSseBlock(buffer: string): { block: string; delimiter: string; rest: string } | null {
const match = buffer.match(/\r?\n\r?\n/);
if (!match || match.index === undefined) return null;
return {
block: buffer.slice(0, match.index),
delimiter: match[0],
rest: buffer.slice(match.index + match[0].length),
};
}

/** Join all data lines from one SSE event according to the event-stream field rules. */
export function sseDataPayload(block: string): string | null {
const data: string[] = [];
for (const line of block.split(/\r?\n/)) {
if (!line.startsWith("data:")) continue;
const value = line.slice(5);
data.push(value.startsWith(" ") ? value.slice(1) : value);
}
return data.length > 0 ? data.join("\n") : null;
}

/** Replace an SSE event's data field while preserving non-data fields and newline style. */
export function replaceSseDataPayload(block: string, payload: string): string {
const newline = block.includes("\r\n") ? "\r\n" : "\n";
const lines = block.split(/\r?\n/);
const rewritten: string[] = [];
let replaced = false;
for (const line of lines) {
if (!line.startsWith("data:")) {
rewritten.push(line);
continue;
}
if (!replaced) {
rewritten.push(`data: ${payload}`);
replaced = true;
}
}
return replaced ? rewritten.join(newline) : block;
}

/** Apply rewrites left-to-right; empty list is identity. */
export function composeSsePayloadRewrites(...rewrites: SsePayloadRewrite[]): SsePayloadRewrite {
if (rewrites.length === 0) return (payload) => payload;
if (rewrites.length === 1) return rewrites[0]!;
return (payload) => {
let next = payload;
for (const rewrite of rewrites) next = rewrite(next);
return next;
};
}

/**
* Relay an SSE body through a single JS pull wrapper, rewriting each event's data payload in place.
* Non-data fields and framing are preserved; invalid JSON payloads are left to the rewrite callback.
*/
export function relaySseWithPayloadRewrite(
body: ReadableStream<Uint8Array>,
rewrite: SsePayloadRewrite,
): ReadableStream<Uint8Array> {
const reader = body.getReader();
const decoder = new TextDecoder();
const encoder = new TextEncoder();
let buffer = "";

const emitProcessedBlocks = (
controller: ReadableStreamDefaultController<Uint8Array>,
flushFinal = false,
): void => {
let next: { block: string; delimiter: string; rest: string } | null;
while ((next = nextSseBlock(buffer))) {
buffer = next.rest;
const payload = sseDataPayload(next.block);
const rewrittenPayload = payload ? rewrite(payload) : undefined;
const block = payload && rewrittenPayload !== undefined && rewrittenPayload !== payload
? replaceSseDataPayload(next.block, rewrittenPayload)
: next.block;
controller.enqueue(encoder.encode(block + next.delimiter));
}
if (flushFinal && buffer.length > 0) {
const payload = sseDataPayload(buffer);
const rewrittenPayload = payload ? rewrite(payload) : undefined;
const block = payload && rewrittenPayload !== undefined && rewrittenPayload !== payload
? replaceSseDataPayload(buffer, rewrittenPayload)
: buffer;
controller.enqueue(encoder.encode(block));
buffer = "";
}
};

return new ReadableStream<Uint8Array>({
async pull(controller) {
const { done, value } = await reader.read();
if (done) {
buffer += decoder.decode();
emitProcessedBlocks(controller, true);
controller.close();
return;
}
buffer += decoder.decode(value, { stream: true });
emitProcessedBlocks(controller);
},
cancel(reason) {
reader.cancel(reason).catch(() => {});
},
});
}
Loading
Loading