fix(ai): stop streams getting cut off and keep tool steps across messages

This commit is contained in:
LukeGus committed 2026-09-28 21:08:48 -05:00
1 parent c5b2b880b6
commit b4457b6121
17 files changed
+726 -254

No files matched your search

+1
View File
@@ -185,6 +185,7 @@ export function createPluginFetch(deps: Deps): PluginFetch {
? { rejectUnauthorized: false }
: {}),
},
init.timeoutMs,
);
} finally {
clearTimeout(timer);
@@ -25,6 +25,7 @@ const state = vi.hoisted(() => ({
]),
fetchCalls: [] as Array<{ url: string; allowlist: readonly string[] }>,
lastTls: null as Record<string, unknown> | null,
lastIdleTimeout: undefined as number | undefined,
lastSignal: null as AbortSignal | null,
lastRedirect: undefined as RequestRedirect | undefined,
}));
@@ -65,8 +66,10 @@ vi.mock("../../utils/safe-outbound-fetch.js", () => ({
init: RequestInit,
allowlist: readonly string[],
tls: Record<string, unknown> = {},
idleTimeoutMs?: number,
) => {
state.lastTls = tls;
state.lastIdleTimeout = idleTimeoutMs;
state.lastRedirect = init.redirect;
state.lastSignal = init.signal ?? null;
state.fetchCalls.push({ url, allowlist });
@@ -302,6 +305,12 @@ describe("ctx.fetch", () => {
expect(state.lastTls).toEqual({});
});
it("hands the caller's timeout to the guard for the body idle timeout", async () => {
const ctx = contextFor(["network:outbound"]);
await ctx.fetch("https://example.com/stream", { timeoutMs: 600_000 });
expect(state.lastIdleTimeout).toBe(600_000);
});
it("aborts the request, and a streamed body, when the caller's signal fires", async () => {
const ctx = contextFor(["network:outbound"]);
const controller = new AbortController();
@@ -5,6 +5,7 @@ import type { LookupAddress, LookupOptions } from "dns";
import {
createDnsLookupHook,
isBlockedAddress,
outboundIdleTimeout,
readResponseTextLimited,
safeOutboundFetch,
} from "../../utils/safe-outbound-fetch.js";
@@ -282,6 +283,17 @@ describe("safeOutboundFetch", () => {
}, 10_000);
});
describe("outboundIdleTimeout", () => {
it("keeps undici's 5 minute default as the floor", () => {
expect(outboundIdleTimeout()).toBe(300_000);
expect(outboundIdleTimeout(30_000)).toBe(300_000);
});
it("raises the idle timeout for a longer request timeout", () => {
expect(outboundIdleTimeout(600_000)).toBe(600_000);
});
});
describe("redirect handling", () => {
it.each([undefined, "error", "follow", "manual"] as const)(
"never follows redirects with mode %s",
+13
View File
@@ -188,11 +188,19 @@ export async function readResponseTextLimited(
return Buffer.concat(chunks, total).toString("utf8");
}
/** undici's own default for both, kept as the floor. */
const DEFAULT_IDLE_TIMEOUT_MS = 300_000;
export function outboundIdleTimeout(timeoutMs?: number): number {
return Math.max(timeoutMs ?? 0, DEFAULT_IDLE_TIMEOUT_MS);
}
export async function safeOutboundFetch(
rawUrl: string,
options: RequestInit,
allowedPrivateHosts: readonly string[] = [],
tls: OutboundTlsOptions = {},
idleTimeoutMs?: number,
): Promise<Response> {
const url = new URL(rawUrl);
if (
@@ -226,7 +234,12 @@ export async function safeOutboundFetch(
} as never)) as unknown as Response;
}
// A slow model can go quiet for minutes mid-stream. Past bodyTimeout undici
// kills the body with a bare "terminated".
const idle = outboundIdleTimeout(idleTimeoutMs);
const dispatcher = new Agent({
headersTimeout: idle,
bodyTimeout: idle,
connect: {
lookup: createDnsLookupHook(lookup, allowPrivate),
...(tls.ca ? { ca: tls.ca } : {}),