fix(agents): preserve embedded auth on HTTP fallback

This commit is contained in:
Ayaan Zaidi
2026-03-27 21:13:38 +05:30
parent 9098e948ac
commit 5e8db468ff
2 changed files with 14 additions and 10 deletions

View File

@@ -299,7 +299,7 @@ export function createOpenAIWebSocketStreamFn(
const run = async () => {
const transport = resolveWsTransport(options);
if (transport === "sse") {
return fallbackToHttp(model, context, options, eventStream, opts.signal);
return fallbackToHttp(model, context, options, apiKey, eventStream, opts.signal);
}
// ── 1. Get or create session state ──────────────────────────────────
@@ -339,7 +339,7 @@ export function createOpenAIWebSocketStreamFn(
`[ws-stream] WebSocket connect failed for session=${sessionId}; falling back to HTTP. error=${String(connErr)}`,
);
// Fall back to HTTP immediately
return fallbackToHttp(model, context, options, eventStream, opts.signal);
return fallbackToHttp(model, context, options, apiKey, eventStream, opts.signal);
}
}
@@ -356,7 +356,7 @@ export function createOpenAIWebSocketStreamFn(
/* ignore */
}
wsRegistry.delete(sessionId);
return fallbackToHttp(model, context, options, eventStream, opts.signal);
return fallbackToHttp(model, context, options, apiKey, eventStream, opts.signal);
}
const signal = opts.signal ?? (options as WsOptions | undefined)?.signal;
@@ -401,7 +401,7 @@ export function createOpenAIWebSocketStreamFn(
log.warn(
`[ws-stream] reconnect after warm-up failed for session=${sessionId}; falling back to HTTP. error=${String(reconnectErr)}`,
);
return fallbackToHttp(model, context, options, eventStream, opts.signal);
return fallbackToHttp(model, context, options, apiKey, eventStream, opts.signal);
}
}
}
@@ -506,7 +506,7 @@ export function createOpenAIWebSocketStreamFn(
/* ignore */
}
wsRegistry.delete(sessionId);
return fallbackToHttp(model, context, options, eventStream, opts.signal);
return fallbackToHttp(model, context, options, apiKey, eventStream, opts.signal);
}
eventStream.push({
@@ -616,10 +616,15 @@ async function fallbackToHttp(
model: Parameters<StreamFn>[0],
context: Parameters<StreamFn>[1],
options: Parameters<StreamFn>[2],
apiKey: string,
eventStream: AssistantMessageEventStreamLike,
signal?: AbortSignal,
): Promise<void> {
const mergedOptions = signal ? { ...options, signal } : options;
const mergedOptions = {
...options,
apiKey,
...(signal ? { signal } : {}),
};
const httpStream = openAIWsStreamDeps.streamSimple(model, context, mergedOptions);
for await (const event of httpStream) {
eventStream.push(event);

View File

@@ -1,7 +1,6 @@
import fs from "node:fs/promises";
import os from "node:os";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { streamSimple } from "@mariozechner/pi-ai";
import {
createAgentSession,
DefaultResourceLoader,
@@ -844,6 +843,7 @@ export async function runEmbeddedAttempt(
workspaceDir: params.workspaceDir,
});
const defaultSessionStreamFn = activeSession.agent.streamFn;
const providerStreamFn = registerProviderStreamForModel({
model: params.model,
cfg: params.config,
@@ -865,15 +865,14 @@ export async function runEmbeddedAttempt(
});
} else {
log.warn(`[ws-stream] no API key for provider=${params.provider}; using HTTP transport`);
activeSession.agent.streamFn = streamSimple;
activeSession.agent.streamFn = defaultSessionStreamFn;
}
} else if (params.model.provider === "anthropic-vertex") {
// Anthropic Vertex AI: inject AnthropicVertex client into pi-ai's
// streamAnthropic for GCP IAM auth instead of Anthropic API keys.
activeSession.agent.streamFn = createAnthropicVertexStreamFnForModel(params.model);
} else {
// Force a stable streamFn reference so vitest can reliably mock @mariozechner/pi-ai.
activeSession.agent.streamFn = streamSimple;
activeSession.agent.streamFn = defaultSessionStreamFn;
}
const { effectiveExtraParams } = applyExtraParamsToAgent(