feat(agents): make openai responses websocket-first with fallback

This commit is contained in:
Peter Steinberger
2026-03-01 21:50:33 +00:00
parent 38da2d076c
commit 7ced38b5ef
7 changed files with 1260 additions and 6 deletions

View File

@@ -319,6 +319,15 @@ function createCodexDefaultTransportWrapper(baseStreamFn: StreamFn | undefined):
});
}
function createOpenAIDefaultTransportWrapper(baseStreamFn: StreamFn | undefined): StreamFn {
const underlying = baseStreamFn ?? streamSimple;
return (model, context, options) =>
underlying(model, context, {
...options,
transport: options?.transport ?? "auto",
});
}
function isAnthropic1MModel(modelId: string): boolean {
const normalized = modelId.trim().toLowerCase();
return ANTHROPIC_1M_MODEL_PREFIXES.some((prefix) => normalized.startsWith(prefix));
@@ -740,6 +749,9 @@ export function applyExtraParamsToAgent(
if (provider === "openai-codex") {
// Default Codex to WebSocket-first when nothing else specifies transport.
agent.streamFn = createCodexDefaultTransportWrapper(agent.streamFn);
} else if (provider === "openai") {
// Default OpenAI Responses to WebSocket-first with transparent SSE fallback.
agent.streamFn = createOpenAIDefaultTransportWrapper(agent.streamFn);
}
const override =
extraParamsOverride && Object.keys(extraParamsOverride).length > 0

View File

@@ -42,6 +42,7 @@ import { resolveImageSanitizationLimits } from "../../image-sanitization.js";
import { resolveModelAuthMode } from "../../model-auth.js";
import { normalizeProviderId, resolveDefaultModelForAgent } from "../../model-selection.js";
import { createOllamaStreamFn, OLLAMA_NATIVE_BASE_URL } from "../../ollama-stream.js";
import { createOpenAIWebSocketStreamFn, releaseWsSession } from "../../openai-ws-stream.js";
import { resolveOwnerDisplaySetting } from "../../owner-display.js";
import {
isCloudCodeAssistFormatError,
@@ -866,6 +867,16 @@ export async function runEmbeddedAttempt(
typeof providerConfig?.baseUrl === "string" ? providerConfig.baseUrl.trim() : "";
const ollamaBaseUrl = modelBaseUrl || providerBaseUrl || OLLAMA_NATIVE_BASE_URL;
activeSession.agent.streamFn = createOllamaStreamFn(ollamaBaseUrl);
} else if (params.model.api === "openai-responses" && params.provider === "openai") {
const wsApiKey = await params.authStorage.getApiKey(params.provider);
if (wsApiKey) {
activeSession.agent.streamFn = createOpenAIWebSocketStreamFn(wsApiKey, params.sessionId, {
signal: runAbortController.signal,
});
} else {
log.warn(`[ws-stream] no API key for provider=${params.provider}; using HTTP transport`);
activeSession.agent.streamFn = streamSimple;
}
} else {
// Force a stable streamFn reference so vitest can reliably mock @mariozechner/pi-ai.
activeSession.agent.streamFn = streamSimple;
@@ -1548,6 +1559,7 @@ export async function runEmbeddedAttempt(
sessionManager,
});
session?.dispose();
releaseWsSession(params.sessionId);
await sessionLock.release();
}
} finally {