feat: add SSE passthrough stream action
This commit is contained in:
+166
-4
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user