diff --git a/src/agents/pi-embedded-runner/run/attempt-http-runtime.ts b/src/agents/pi-embedded-runner/run/attempt-http-runtime.ts new file mode 100644 index 00000000000..ad2d2f38b06 --- /dev/null +++ b/src/agents/pi-embedded-runner/run/attempt-http-runtime.ts @@ -0,0 +1,11 @@ +import { + ensureGlobalUndiciEnvProxyDispatcher, + ensureGlobalUndiciStreamTimeouts, +} from "../../../infra/net/undici-global-dispatcher.js"; + +export function configureEmbeddedAttemptHttpRuntime(params: { timeoutMs: number }): void { + // Proxy bootstrap must happen before timeout tuning so the timeouts wrap the + // active EnvHttpProxyAgent instead of being replaced by a bare proxy dispatcher. + ensureGlobalUndiciEnvProxyDispatcher(); + ensureGlobalUndiciStreamTimeouts({ timeoutMs: params.timeoutMs }); +} diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.timeout.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.timeout.test.ts index 6cefe21c187..625b2f0603f 100644 --- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.timeout.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.timeout.test.ts @@ -1,41 +1,28 @@ -import { afterEach, beforeEach, describe, expect, it } from "vitest"; -import { - cleanupTempPaths, - createContextEngineAttemptRunner, - getHoisted, - resetEmbeddedAttemptHarness, -} from "./attempt.spawn-workspace.test-support.js"; +import { beforeEach, describe, expect, it, vi } from "vitest"; -const hoisted = getHoisted(); +const mocks = vi.hoisted(() => ({ + ensureGlobalUndiciEnvProxyDispatcher: vi.fn(), + ensureGlobalUndiciStreamTimeouts: vi.fn(), +})); + +vi.mock("../../../infra/net/undici-global-dispatcher.js", () => ({ + ensureGlobalUndiciEnvProxyDispatcher: mocks.ensureGlobalUndiciEnvProxyDispatcher, + ensureGlobalUndiciStreamTimeouts: mocks.ensureGlobalUndiciStreamTimeouts, +})); + +import { configureEmbeddedAttemptHttpRuntime } from "./attempt-http-runtime.js"; describe("runEmbeddedAttempt undici timeout wiring", () => { - const tempPaths: string[] = []; - beforeEach(() => { - resetEmbeddedAttemptHarness(); + mocks.ensureGlobalUndiciEnvProxyDispatcher.mockReset(); + mocks.ensureGlobalUndiciStreamTimeouts.mockReset(); }); - afterEach(async () => { - await cleanupTempPaths(tempPaths); - }); + it("forwards the configured run timeout into global undici stream tuning", () => { + configureEmbeddedAttemptHttpRuntime({ timeoutMs: 123_456 }); - it("forwards the configured run timeout into global undici stream tuning", async () => { - await createContextEngineAttemptRunner({ - sessionKey: "agent:main:ollama-timeout-test", - tempPaths, - contextEngine: { - assemble: async ({ messages }) => ({ - messages, - estimatedTokens: 1, - }), - }, - attemptOverrides: { - timeoutMs: 123_456, - }, - }); - - expect(hoisted.ensureGlobalUndiciEnvProxyDispatcherMock).toHaveBeenCalledOnce(); - expect(hoisted.ensureGlobalUndiciStreamTimeoutsMock).toHaveBeenCalledWith({ + expect(mocks.ensureGlobalUndiciEnvProxyDispatcher).toHaveBeenCalledOnce(); + expect(mocks.ensureGlobalUndiciStreamTimeouts).toHaveBeenCalledWith({ timeoutMs: 123_456, }); }); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index c05d3a155cf..5b77dbf1884 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -12,10 +12,6 @@ import { resolveChannelCapabilities } from "../../../config/channel-capabilities import { formatErrorMessage } from "../../../infra/errors.js"; import { resolveHeartbeatSummaryForAgent } from "../../../infra/heartbeat-summary.js"; import { getMachineDisplayName } from "../../../infra/machine-name.js"; -import { - ensureGlobalUndiciEnvProxyDispatcher, - ensureGlobalUndiciStreamTimeouts, -} from "../../../infra/net/undici-global-dispatcher.js"; import { MAX_IMAGE_BYTES } from "../../../media/constants.js"; import { isOllamaCompatProvider, @@ -186,6 +182,7 @@ import { splitSdkTools } from "../tool-split.js"; import { mapThinkingLevel } from "../utils.js"; import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js"; export { buildContextEnginePromptCacheInfo } from "./attempt.context-engine-helpers.js"; +import { configureEmbeddedAttemptHttpRuntime } from "./attempt-http-runtime.js"; import { assembleAttemptContextEngine, buildLoopPromptCacheInfo, @@ -423,10 +420,7 @@ export async function runEmbeddedAttempt( ): Promise { const resolvedWorkspace = resolveUserPath(params.workspaceDir); const runAbortController = new AbortController(); - // Proxy bootstrap must happen before timeout tuning so the timeouts wrap the - // active EnvHttpProxyAgent instead of being replaced by a bare proxy dispatcher. - ensureGlobalUndiciEnvProxyDispatcher(); - ensureGlobalUndiciStreamTimeouts({ timeoutMs: params.timeoutMs }); + configureEmbeddedAttemptHttpRuntime({ timeoutMs: params.timeoutMs }); log.debug( `embedded run start: runId=${params.runId} sessionId=${params.sessionId} provider=${params.provider} model=${params.modelId} thinking=${params.thinkLevel} messageChannel=${params.messageChannel ?? params.messageProvider ?? "unknown"}`,