feat: add SSE passthrough stream action

This commit is contained in:
2026-07-20 00:40:14 +08:00
parent ca8935d73e
commit 31e673f777
4 changed files with 282 additions and 4 deletions
+166 -4
View File
@@ -80,6 +80,7 @@ const DEFAULT_CONFIG = {
log_match: true,
health_path: "/__codex_retry_gateway/health",
};
const STREAM_ACTIONS = ["strict_502", "disconnect", "passthrough"];
const REASONING_POINTERS = [
"/usage/output_tokens_details/reasoning_tokens",
@@ -935,6 +936,15 @@ function normalizeReasoningMatchMode(value) {
return DEFAULT_CONFIG.reasoning_match_mode;
}
function normalizeStreamAction(value) {
const action = `${value || DEFAULT_CONFIG.stream_action}`.trim().toLowerCase();
return STREAM_ACTIONS.includes(action) ? action : DEFAULT_CONFIG.stream_action;
}
function streamPassthroughEnabled(config) {
return normalizeStreamAction(config?.stream_action) === "passthrough";
}
function normalizeReasoningEquals(values, fallback = DEFAULT_CONFIG.reasoning_equals) {
const normalized = normalizeIntegerList(values, fallback);
return normalized.length > 0 ? normalized : [...fallback];
@@ -1646,6 +1656,7 @@ function buildProfileFormModel(env) {
env.CODEX_RETRY_GATEWAY_UPSTREAM_FETCH_RETRY_ATTEMPTS || `${DEFAULT_CONFIG.upstream_fetch_retry_attempts}`,
upstream_fetch_retry_backoff_ms:
env.CODEX_RETRY_GATEWAY_UPSTREAM_FETCH_RETRY_BACKOFF_MS || `${DEFAULT_CONFIG.upstream_fetch_retry_backoff_ms}`,
stream_action: normalizeStreamAction(env.CODEX_RETRY_GATEWAY_STREAM_ACTION),
endpoints: normalizeStringList(env.CODEX_RETRY_GATEWAY_ENDPOINTS || DEFAULT_CONFIG.endpoints, DEFAULT_CONFIG.endpoints),
};
}
@@ -1758,7 +1769,7 @@ function buildConfigFromProfileEnv(profileName, env, imageConfig = {}) {
non_stream_status_code: env.CODEX_RETRY_GATEWAY_NON_STREAM_STATUS_CODE
? Number.parseInt(`${env.CODEX_RETRY_GATEWAY_NON_STREAM_STATUS_CODE}`, 10)
: DEFAULT_CONFIG.non_stream_status_code,
stream_action: env.CODEX_RETRY_GATEWAY_STREAM_ACTION || DEFAULT_CONFIG.stream_action,
stream_action: normalizeStreamAction(env.CODEX_RETRY_GATEWAY_STREAM_ACTION),
log_match: env.CODEX_RETRY_GATEWAY_LOG_MATCH === undefined
? DEFAULT_CONFIG.log_match
: ["1", "true", "yes", "on"].includes(`${env.CODEX_RETRY_GATEWAY_LOG_MATCH}`.trim().toLowerCase()),
@@ -1895,6 +1906,7 @@ async function buildProfileEnvText(payload) {
if (upstreamFetchRetryBackoffMs < 0) {
throw new Error("upstream_fetch_retry_backoff_ms 不能为负数");
}
const streamAction = normalizeStreamAction(payload.stream_action);
const envPairs = [
["CODEX_RETRY_GATEWAY_LISTEN_HOST", `${payload.listen_host || DEFAULT_CONFIG.listen_host}`.trim()],
@@ -1907,6 +1919,7 @@ async function buildProfileEnvText(payload) {
["CODEX_RETRY_GATEWAY_RETRYABLE_ERROR_MESSAGES", retryableErrorMessages.join("\n")],
["CODEX_RETRY_GATEWAY_UPSTREAM_FETCH_RETRY_ATTEMPTS", `${upstreamFetchRetryAttempts}`],
["CODEX_RETRY_GATEWAY_UPSTREAM_FETCH_RETRY_BACKOFF_MS", `${upstreamFetchRetryBackoffMs}`],
["CODEX_RETRY_GATEWAY_STREAM_ACTION", streamAction],
["CODEX_RETRY_GATEWAY_REQUEST_HISTORY_LIMIT", `${requestHistoryLimit}`],
["CODEX_RETRY_GATEWAY_ENDPOINTS", endpoints.join(",")],
];
@@ -2132,6 +2145,7 @@ async function buildTextProfileExportBundle(runtime, payload) {
retryable_error_messages: config.retryable_error_messages,
upstream_fetch_retry_attempts: config.upstream_fetch_retry_attempts,
upstream_fetch_retry_backoff_ms: config.upstream_fetch_retry_backoff_ms,
stream_action: config.stream_action,
endpoints: config.endpoints,
},
};
@@ -2242,6 +2256,7 @@ async function importTextProfileBundle(runtime, payload) {
retryable_error_messages: imported.profile.retryable_error_messages,
upstream_fetch_retry_attempts: imported.profile.upstream_fetch_retry_attempts,
upstream_fetch_retry_backoff_ms: imported.profile.upstream_fetch_retry_backoff_ms,
stream_action: imported.profile.stream_action,
endpoints: imported.profile.endpoints,
});
return {
@@ -2544,6 +2559,7 @@ async function loadConfig(configPath) {
config.retryable_error_messages,
DEFAULT_CONFIG.retryable_error_messages,
);
config.stream_action = normalizeStreamAction(config.stream_action);
config.management_access_key = normalizeManagementAccessKey(config.management_access_key);
config.upstream_fetch_retry_attempts = normalizePositiveInteger(
config.upstream_fetch_retry_attempts,
@@ -2941,6 +2957,7 @@ async function listProfiles(runtime) {
auth_source: summarizeProfileAuthSource(env),
reasoning_match_mode: form.reasoning_match_mode,
reasoning_equals: form.reasoning_equals,
stream_action: form.stream_action,
},
form,
});
@@ -3703,6 +3720,9 @@ function buildEditableConfig(currentConfig, payload) {
payload.non_stream_status_code === undefined
? currentConfig.non_stream_status_code
: Number.parseInt(`${payload.non_stream_status_code}`, 10);
const nextStreamAction = normalizeStreamAction(
payload.stream_action === undefined ? currentConfig.stream_action : payload.stream_action,
);
if (nextRetryableStatusCodes.length === 0) {
throw new Error("retryable_status_codes 不能为空");
@@ -3737,6 +3757,7 @@ function buildEditableConfig(currentConfig, payload) {
upstream_fetch_retry_attempts: nextUpstreamFetchRetryAttempts,
upstream_fetch_retry_backoff_ms: nextUpstreamFetchRetryBackoffMs,
non_stream_status_code: nextStatusCode,
stream_action: nextStreamAction,
log_match: payload.log_match === undefined ? currentConfig.log_match : Boolean(payload.log_match),
};
}
@@ -4322,7 +4343,7 @@ async function handleManagementRequest(runtime, req, res, requestUrl) {
await writeConfig(runtime.configPath, nextConfig);
runtime.config = nextConfig;
runtime.logger(
`[config] updated reasoning_match_mode=${nextConfig.reasoning_match_mode} reasoning_equals=${nextConfig.reasoning_equals.join(",")} retryable_status_codes=${nextConfig.retryable_status_codes.join(",")} endpoints=${nextConfig.endpoints.join(",")}`,
`[config] updated reasoning_match_mode=${nextConfig.reasoning_match_mode} reasoning_equals=${nextConfig.reasoning_equals.join(",")} retryable_status_codes=${nextConfig.retryable_status_codes.join(",")} stream_action=${nextConfig.stream_action} endpoints=${nextConfig.endpoints.join(",")}`,
);
const state = await readRuntimeState(runtime);
jsonResponse(req, res, 200, {
@@ -4538,6 +4559,34 @@ async function writeBufferedStreamChunks(res, chunks) {
}
}
async function writeResponseChunk(res, chunk) {
if (!chunk || chunk.length === 0 || res.destroyed || res.writableEnded) {
return false;
}
if (res.write(chunk)) {
return true;
}
await new Promise((resolve) => {
let settled = false;
const finish = () => {
if (settled) {
return;
}
settled = true;
res.off("drain", onDrain);
res.off("close", onClose);
res.off("error", onClose);
resolve();
};
const onDrain = () => finish();
const onClose = () => finish();
res.once("drain", onDrain);
res.once("close", onClose);
res.once("error", onClose);
});
return !res.destroyed && !res.writableEnded;
}
async function writeCapturedResponse(res, delivery) {
if (!delivery) {
throw new Error("missing captured response delivery");
@@ -5281,7 +5330,7 @@ async function handleStreaming({
captureOnly = false,
persistEntry = true,
}) {
const strict502Mode = captureOnly || config.stream_action !== "disconnect";
const strict502Mode = captureOnly || normalizeStreamAction(config.stream_action) !== "disconnect";
const reader = upstreamResponse.body.getReader();
const sseState = {
decoder: new TextDecoder("utf8"),
@@ -5659,6 +5708,97 @@ async function handleStreaming({
}
}
async function handlePassthroughStreaming({
runtime,
upstreamResponse,
res,
requestEntry,
requestAbortSignal = null,
persistEntry = true,
}) {
const reader = upstreamResponse.body?.getReader();
if (!reader) {
throw new Error("upstream SSE response is missing a readable body");
}
requestEntry.inspected = false;
copyHeadersToClient(upstreamResponse.headers, res);
res.writeHead(upstreamResponse.status);
res.socket?.setNoDelay(true);
res.flushHeaders?.();
while (true) {
let readResult;
try {
readResult = await reader.read();
} catch (error) {
if (requestAbortSignal?.aborted) {
throw requestAbortSignal.reason || error;
}
if (persistEntry) {
persistStreamingProgress(runtime, requestEntry, { force: true }, new Date());
}
if (!res.destroyed && !res.writableEnded) {
res.destroy(error);
}
return {
inspected: false,
matched: false,
status_code: upstreamResponse.status,
upstream_status_code: upstreamResponse.status,
response_id: requestEntry.response_id,
thread_id: requestEntry.thread_id,
response_bytes_received: requestEntry.response_bytes_received,
stream_chunk_count: requestEntry.stream_chunk_count,
error: `upstream SSE stream terminated: ${error?.message || error}`,
};
}
if (readResult.done) {
if (persistEntry) {
persistStreamingProgress(runtime, requestEntry, { force: true }, new Date());
}
if (!res.destroyed && !res.writableEnded) {
res.end();
}
return {
inspected: false,
matched: false,
status_code: upstreamResponse.status,
upstream_status_code: upstreamResponse.status,
response_id: requestEntry.response_id,
thread_id: requestEntry.thread_id,
response_bytes_received: requestEntry.response_bytes_received,
stream_chunk_count: requestEntry.stream_chunk_count,
};
}
const chunk = Buffer.from(readResult.value);
const now = new Date();
markRequestEntryFirstResponse(runtime, requestEntry, { persistEntry, at: now });
updateStreamingProgress(requestEntry, { chunkBytes: chunk.length, at: now });
if (persistEntry) {
persistStreamingProgress(
runtime,
requestEntry,
{ force: requestEntry.stream_chunk_count === 1 },
now,
);
}
await writeResponseChunk(res, chunk);
if (res.destroyed || res.writableEnded) {
return {
cancelled: true,
cancel_reason: REASONING_RETRY_ABORT_CLIENT,
response_id: requestEntry.response_id,
thread_id: requestEntry.thread_id,
response_bytes_received: requestEntry.response_bytes_received,
stream_chunk_count: requestEntry.stream_chunk_count,
};
}
}
}
function buildAttemptRequestEntry(baseEntry) {
return {
...baseEntry,
@@ -5778,6 +5918,22 @@ async function executeGatewayQuery({
);
}
if (isSseContentType(responseContentType) && !captureOnly && streamPassthroughEnabled(config)) {
const result = await handlePassthroughStreaming({
runtime,
upstreamResponse,
res,
requestEntry,
requestAbortSignal: queryAbortLink.controller.signal,
persistEntry,
});
return {
response_stream: true,
...result,
total_upstream_attempts: totalUpstreamAttempts,
};
}
if (!shouldInspect) {
markRequestEntryFirstResponse(runtime, requestEntry, { persistEntry });
const body = Buffer.from(await upstreamResponse.arrayBuffer());
@@ -6303,6 +6459,12 @@ async function proxyRequest(runtime, req, res) {
requestEntry.reasoning_summary = extractRequestReasoningSummary(requestJson);
requestEntry.request_stream = requestIsStream;
applyThreadReasoningState(runtime, requestEntry, pathname);
const streamPassthroughRequest = streamPassthroughEnabled(config) && requestIsStream;
if (streamPassthroughRequest && isResponsesReasoningRetryPath(pathname)) {
requestEntry.reasoning_retry_enabled = false;
requestEntry.reasoning_retry_thread_mode = "stream_passthrough";
requestEntry.reasoning_retry_stop_reason = "stream_passthrough";
}
upsertRequestEntry(runtime, requestEntry);
const upstreamRoute = selectUpstreamRoute(config, pathname);
@@ -6323,7 +6485,7 @@ async function proxyRequest(runtime, req, res) {
}
const clientAbortContext = createClientAbortContext(req, res);
try {
const result = isResponsesReasoningRetryEligible(pathname, requestEntry)
const result = !streamPassthroughRequest && isResponsesReasoningRetryEligible(pathname, requestEntry)
? await runResponsesReasoningRetry({
runtime,
config,