From ef2039ba2af0a8d71f0088a29d025c245090d7dd Mon Sep 17 00:00:00 2001 From: Mux Date: Sun, 23 Aug 2026 10:44:41 -0500 Subject: [PATCH] =?UTF-8?q?[stream-manager]=20=F0=9F=A4=96=20fix:=20abort?= =?UTF-8?q?=20stalled=20provider=20streams?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Abort a provider stream after 15 minutes without SDK progress. Pause the deadline while a local tool executes. Preserve partial output and report a retryable truncation error. --- _Generated with `mux` • Model: `openai:gpt-5.6-sol` • Thinking: `xhigh` • Cost: `$0.00`_ Co-authored-by: Mux --- src/constants/streaming.ts | 7 ++ src/node/services/providerModelFactory.ts | 4 +- src/node/services/streamManager.test.ts | 86 ++++++++++++++ src/node/services/streamManager.ts | 109 ++++++++++++++++-- .../tools/withSequentialExecution.test.ts | 58 ++++++++++ .../services/tools/withSequentialExecution.ts | 16 ++- 6 files changed, 266 insertions(+), 14 deletions(-) diff --git a/src/constants/streaming.ts b/src/constants/streaming.ts index 8745aa8d616..7947760924d 100644 --- a/src/constants/streaming.ts +++ b/src/constants/streaming.ts @@ -13,6 +13,13 @@ export const WORKSPACE_STREAMING_STATUS_TRANSITION_MS = 150; */ export const APPROX_CHARS_PER_TOKEN = 4; +/** + * Abort a provider stream that emits no SDK progress for this interval. + * A half-open connection otherwise leaves the workspace busy forever. + * Local tool execution pauses this deadline because tools have their own bounds. + */ +export const PROVIDER_STREAM_IDLE_TIMEOUT_MS = 15 * 60 * 1000; + export const STREAM_SMOOTHING = { /** Baseline reveal speed in characters per second when no live model rate is known yet. */ BASE_CHARS_PER_SEC: 72, diff --git a/src/node/services/providerModelFactory.ts b/src/node/services/providerModelFactory.ts index 50dad21c6a1..658aef28875 100644 --- a/src/node/services/providerModelFactory.ts +++ b/src/node/services/providerModelFactory.ts @@ -83,8 +83,8 @@ import { EnvHttpProxyAgent, type Dispatcher } from "undici"; import packageJson from "../../../package.json"; // --------------------------------------------------------------------------- -// Undici agent with unlimited timeouts for AI streaming requests. -// Safe because users control cancellation via AbortSignal from the UI. +// Undici agent with unlimited transport timeouts for AI streaming requests. +// StreamManager applies a semantic progress deadline and users can also cancel. // Uses EnvHttpProxyAgent to automatically respect HTTP_PROXY, HTTPS_PROXY, // and NO_PROXY environment variables for debugging/corporate network support. // --------------------------------------------------------------------------- diff --git a/src/node/services/streamManager.test.ts b/src/node/services/streamManager.test.ts index b6ea466da96..a21c65863d5 100644 --- a/src/node/services/streamManager.test.ts +++ b/src/node/services/streamManager.test.ts @@ -214,6 +214,7 @@ function createStreamInfoForTests( toolCompletionTimestamps: new Map(), pendingWorkflowRunAttachments: new Map(), pendingToolExecutionStarts: new Map(), + activeLocalToolExecutions: new Set(), model, metadataModel: overrides.metadataModel ?? model, historySequence: 1, @@ -406,6 +407,13 @@ describe("StreamManager - tool execution start timing", () => { expect(events[0].timestamp).toBeGreaterThan(timestamp); const parts = streamInfo.parts as Array>; + expect((streamInfo.activeLocalToolExecutions as Set).has("tool-call-1")).toBe(true); + + const handleToolExecutionEnd = getPrivateMethodForTests< + (workspaceId: string, messageId: string, toolCallId: string) => void + >(streamManager, "handleToolExecutionEnd"); + handleToolExecutionEnd.call(streamManager, workspaceId, messageId, "tool-call-1"); + expect((streamInfo.activeLocalToolExecutions as Set).size).toBe(0); expect(parts[0].executionStartedAt).toBe(events[0].timestamp); }); @@ -2328,6 +2336,83 @@ describe("StreamManager - empty stream completions", () => { expect(partial?.parts).toMatchObject([{ type: "text", text: "partial answer" }]); }); + test("aborts and persists a retryable error when a provider stream stops making progress", async () => { + const streamManager = new StreamManager(historyService); + expect(Reflect.set(streamManager, "providerStreamIdleTimeoutMs", 20)).toBe(true); + + const errorEvents: Array<{ messageId: string; error: string; errorType?: string }> = []; + const streamEndEvents: unknown[] = []; + streamManager.on("error", (data) => { + errorEvents.push(data as { messageId: string; error: string; errorType?: string }); + }); + streamManager.on("stream-end", (data) => { + streamEndEvents.push(data); + }); + expect( + Reflect.set(streamManager, "tokenTracker", { + setModel: () => Promise.resolve(undefined), + countTokens: () => Promise.resolve(0), + }) + ).toBe(true); + + const workspaceId = "idle-timeout-workspace"; + const messageId = "idle-timeout-message"; + const historySequence = 1; + const abortController = new AbortController(); + let abortObserved = false; + abortController.signal.addEventListener( + "abort", + () => { + abortObserved = true; + }, + { once: true } + ); + + await appendPartialAssistantForTests(workspaceId, messageId, historySequence); + const processStreamWithCleanup = getProcessStreamWithCleanupForTests(streamManager); + const startTime = Date.now() - 250; + const streamInfo = createStreamInfoForTests({ + abortController, + streamResult: createStreamResultForTests( + (async function* () { + yield { type: "text-delta", text: "Ratified — spawning the implementation child now." }; + await new Promise((resolve) => { + abortController.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + })(), + { inputTokens: 3, outputTokens: 2, totalTokens: 5 } + ), + messageId, + startTime, + lastPartTimestamp: startTime, + model: KNOWN_MODELS.SONNET.id, + metadataModel: KNOWN_MODELS.SONNET.id, + historySequence, + initialMetadata: { agentId: "exec" }, + runtime, + }); + getWorkspaceStreamsForTests(streamManager).set(workspaceId, streamInfo); + + await processStreamWithCleanup.call(streamManager, workspaceId, streamInfo, historySequence); + + expect(abortObserved).toBe(true); + expect(streamEndEvents).toHaveLength(0); + expect(errorEvents).toHaveLength(1); + expect(errorEvents[0]).toMatchObject({ + messageId, + errorType: "stream_truncated", + }); + expect(errorEvents[0]?.error).toContain("produced no stream progress for 20ms"); + expect(getWorkspaceStreamsForTests(streamManager).has(workspaceId)).toBe(false); + + const partial = await historyService.readPartial(workspaceId); + expect(partial?.metadata?.errorType).toBe("stream_truncated"); + expect(partial?.metadata?.error).toContain("produced no stream progress for 20ms"); + expect(partial?.parts).toMatchObject([ + { type: "text", text: "Ratified — spawning the implementation child now." }, + ]); + }); + test("treats streamText's synthesized (other, undefined) finish part as a truncated stream", async () => { // streamText's runStep initializes stepFinishReason="other" / // stepRawFinishReason=undefined and unconditionally emits those from its @@ -5731,6 +5816,7 @@ describe("StreamManager - mid-turn thinking override", () => { undefined, undefined, undefined, // onToolExecutionStart + undefined, // onToolExecutionEnd state, rebuild ); diff --git a/src/node/services/streamManager.ts b/src/node/services/streamManager.ts index 7aad3f6c590..c0556c1f51c 100644 --- a/src/node/services/streamManager.ts +++ b/src/node/services/streamManager.ts @@ -102,6 +102,10 @@ import { PROVIDER_DEFINITIONS } from "@/common/constants/providers"; import { isRefusalFinishReason } from "@/common/utils/messages/refusalFinishReason"; import { getOpenAIResponsesBaseUrlHint } from "@/node/services/utils/openAIResponsesBaseUrlHint"; +import { PROVIDER_STREAM_IDLE_TIMEOUT_MS } from "@/constants/streaming"; +import { formatDuration } from "@/common/utils/formatDuration"; +import { raceWithAbortAndTimeout } from "@/node/utils/concurrency/withTimeout"; + // Disable noisy AI SDK warning logging. globalThis.AI_SDK_LOG_WARNINGS = false; @@ -141,13 +145,26 @@ class ModelRefusalError extends Error { class StreamTruncatedError extends Error { readonly providerDisplayName: string; - constructor(providerDisplayName: string) { - super(`${providerDisplayName} ${STREAM_TRUNCATED_MESSAGE_SUFFIX}`); + constructor(providerDisplayName: string, message?: string) { + super(message ?? `${providerDisplayName} ${STREAM_TRUNCATED_MESSAGE_SUFFIX}`); this.name = "StreamTruncatedError"; this.providerDisplayName = providerDisplayName; } } +class StreamIdleTimeoutError extends StreamTruncatedError { + readonly timeoutMs: number; + + constructor(providerDisplayName: string, timeoutMs: number) { + super( + providerDisplayName, + `${providerDisplayName} produced no stream progress for ${formatDuration(timeoutMs)}. Xum aborted the stalled stream. Retry the turn or switch models.` + ); + this.name = "StreamIdleTimeoutError"; + this.timeoutMs = timeoutMs; + } +} + // Type definitions for stream parts with extended properties interface ReasoningDeltaPart { type: "reasoning-delta"; @@ -570,6 +587,10 @@ interface WorkspaceStreamInfo { // and apply it as soon as the part lands. pendingToolExecutionStarts: Map; + // Local tool execution can legitimately produce no provider events for a long time. + // Pause the provider idle deadline until each execute() call returns. + activeLocalToolExecutions: Set; + model: string; /** Metadata model resolved from provider mapping for cost/token metadata lookups. */ metadataModel: string; @@ -672,6 +693,7 @@ export class StreamManager extends EventEmitter { private workspaceStreams = new Map(); private streamLocks = new Map(); private readonly PARTIAL_WRITE_THROTTLE_MS = 500; + private readonly providerStreamIdleTimeoutMs = PROVIDER_STREAM_IDLE_TIMEOUT_MS; private readonly historyService: HistoryService; private mcpServerManager?: MCPServerManager; private readonly sessionUsageService?: SessionUsageService; @@ -855,6 +877,8 @@ export class StreamManager extends EventEmitter { return; } + streamInfo.activeLocalToolExecutions.add(toolCallId); + // Use the stream's monotonic clock, not raw Date.now(): the tool-call part timestamp // was monotonicized by nextPartTimestamp(), so a same-millisecond raw reading could be // <= it. Reconnect replay repairs missed execution starts only when @@ -868,6 +892,19 @@ export class StreamManager extends EventEmitter { } } + private handleToolExecutionEnd( + workspaceId: WorkspaceId, + messageId: string, + toolCallId: string + ): void { + const streamInfo = this.workspaceStreams.get(workspaceId); + if (streamInfo?.messageId !== messageId) { + return; + } + + streamInfo.activeLocalToolExecutions.delete(toolCallId); + } + /** * Write the current partial message to disk (throttled by mtime) * Ensures writes happen during rapid streaming (crash-resilient) @@ -1676,6 +1713,7 @@ export class StreamManager extends EventEmitter { onStepMessages?: (messages: ModelMessage[]) => void, toolSearchState?: ToolSearchStreamState, onToolExecutionStart?: (toolCallId: string) => void, + onToolExecutionEnd?: (toolCallId: string) => void, thinkingOverrideState?: ActiveTurnThinkingOverride, rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel, forcedFirstStepToolNames?: string[], @@ -1764,7 +1802,7 @@ export class StreamManager extends EventEmitter { system: finalSystem, // Keep provider-level parallel tool planning enabled, but serialize sibling // execute() handlers inside this stream so shared mutable state cannot race. - tools: withSequentialExecution(finalTools, onToolExecutionStart), + tools: withSequentialExecution(finalTools, onToolExecutionStart, onToolExecutionEnd), providerOptions: finalProviderOptions, headers, maxOutputTokens: effectiveMaxOutputTokens, @@ -2069,6 +2107,7 @@ export class StreamManager extends EventEmitter { onStepMessages, toolSearchState, (toolCallId) => this.handleToolExecutionStart(workspaceId, messageId, toolCallId), + (toolCallId) => this.handleToolExecutionEnd(workspaceId, messageId, toolCallId), thinkingOverrideState, rebuildProviderOptionsForThinkingLevel, forcedFirstStepToolNames, @@ -2100,6 +2139,7 @@ export class StreamManager extends EventEmitter { toolCompletionTimestamps: new Map(), pendingWorkflowRunAttachments: new Map(), pendingToolExecutionStarts: new Map(), + activeLocalToolExecutions: new Set(), model: modelString, metadataModel, thinkingLevel, @@ -2801,6 +2841,7 @@ export class StreamManager extends EventEmitter { // against the fallback toolset, so prepareStep keeps reading live state. streamInfo.request.toolSearchState, (toolCallId) => this.handleToolExecutionStart(workspaceId, streamInfo.messageId, toolCallId), + (toolCallId) => this.handleToolExecutionEnd(workspaceId, streamInfo.messageId, toolCallId), // Same holder object (the session's setter keeps working across the // hop) with a closure bound to the FALLBACK model. Attached before // createStreamResult below in case the SDK eagerly prepares step 1. @@ -2962,6 +3003,55 @@ export class StreamManager extends EventEmitter { return true; } + private async readNextStreamPart( + iterator: AsyncIterator, + streamInfo: WorkspaceStreamInfo, + workspaceLog: Logger + ): Promise | null> { + const waitStartedAt = Date.now(); + const nextPartPromise = iterator.next(); + + while (true) { + const outcome = await raceWithAbortAndTimeout(nextPartPromise, { + signal: streamInfo.abortController.signal, + timeoutMs: this.providerStreamIdleTimeoutMs, + }); + if (outcome.kind === "ok") { + return outcome.value; + } + if (outcome.kind === "aborted") { + return null; + } + + // Tool execute() calls have their own bounds. They can legitimately keep + // fullStream.next() pending while no provider event exists to observe. + if (streamInfo.activeLocalToolExecutions.size > 0) { + continue; + } + + const timeoutError = new StreamIdleTimeoutError( + getStreamProviderDisplayName(streamInfo.model), + this.providerStreamIdleTimeoutMs + ); + workspaceLog.warn("Aborting provider stream after idle timeout", { + messageId: streamInfo.messageId, + model: streamInfo.model, + timeoutMs: this.providerStreamIdleTimeoutMs, + idleDurationMs: Date.now() - waitStartedAt, + lastProgressAt: new Date(waitStartedAt).toISOString(), + }); + + streamInfo.abortController.abort(); + const returnPromise = iterator.return?.(); + if (returnPromise != null) { + void Promise.resolve(returnPromise).catch(() => { + // The provider abort is authoritative. Iterator cleanup is best-effort. + }); + } + throw timeoutError; + } + } + /** * Processes a stream with guaranteed cleanup, regardless of success or failure */ @@ -2992,12 +3082,17 @@ export class StreamManager extends EventEmitter { const toolCalls: ToolCallMap = new Map(); try { - for await (const part of streamInfo.streamResult.fullStream) { - // Check if stream was cancelled BEFORE processing any parts - // This improves interruption responsiveness by catching aborts earlier - if (streamInfo.abortController.signal.aborted) { + const streamIterator = streamInfo.streamResult.fullStream[Symbol.asyncIterator](); + while (true) { + const nextPart = await this.readNextStreamPart( + streamIterator, + streamInfo, + workspaceLog + ); + if (nextPart == null || nextPart.done) { break; } + const part = nextPart.value; // Log all stream parts to debug reasoning (commented out - too spammy) // console.log("[DEBUG streamManager]: Stream part", { diff --git a/src/node/services/tools/withSequentialExecution.test.ts b/src/node/services/tools/withSequentialExecution.test.ts index 12cefcf0183..52a6a8db0e4 100644 --- a/src/node/services/tools/withSequentialExecution.test.ts +++ b/src/node/services/tools/withSequentialExecution.test.ts @@ -168,6 +168,64 @@ describe("withSequentialExecution", () => { expect(events).toEqual(["execution-start call-a", "run A", "execution-start call-b", "run B"]); }); + test("reports execution end after success and failure", async () => { + const events: string[] = []; + const tools = { + success: tool({ + description: "Successful tool", + inputSchema: z.object({}), + execute: () => { + events.push("run success"); + return Promise.resolve({ ok: true }); + }, + }), + failure: tool({ + description: "Failing tool", + inputSchema: z.object({}), + execute: (): Promise<{ ok: boolean }> => { + events.push("run failure"); + return Promise.reject(new Error("failed")); + }, + }), + }; + const wrappedTools = withSequentialExecution( + tools, + (toolCallId) => events.push(`start ${toolCallId}`), + (toolCallId) => events.push(`end ${toolCallId}`) + ); + + await callWrappedExecute( + wrappedTools!.success as Record, + {}, + { + toolCallId: "success-call", + } + ); + let failure: unknown; + try { + await callWrappedExecute( + wrappedTools!.failure as Record, + {}, + { + toolCallId: "failure-call", + } + ); + } catch (error) { + failure = error; + } + expect(failure).toBeInstanceOf(Error); + expect((failure as Error).message).toBe("failed"); + + expect(events).toEqual([ + "start success-call", + "run success", + "end success-call", + "start failure-call", + "run failure", + "end failure-call", + ]); + }); + test("does not execute queued siblings after stream abort", async () => { const executionLog: string[] = []; const startedA = createDeferred(); diff --git a/src/node/services/tools/withSequentialExecution.ts b/src/node/services/tools/withSequentialExecution.ts index d0fb72cbd10..cb71b661ed7 100644 --- a/src/node/services/tools/withSequentialExecution.ts +++ b/src/node/services/tools/withSequentialExecution.ts @@ -209,13 +209,13 @@ class SharedExecutionLock { * tasks share the read side so they can overlap with each other, while every * other tool call stays exclusive. * - * `onExecutionStart` fires right after the execution lock is acquired (i.e. - * when the tool actually starts running, not when the model emitted the call), - * so queued siblings don't count wait time as execution time. + * `onExecutionStart` fires after the execution lock is acquired. `onExecutionEnd` + * fires after execute() settles. These callbacks exclude queued lock wait time. */ export function withSequentialExecution( tools: Record | undefined, - onExecutionStart?: (toolCallId: string) => void + onExecutionStart?: (toolCallId: string) => void, + onExecutionEnd?: (toolCallId: string) => void ): Record | undefined { if (!tools) { return tools; @@ -250,7 +250,13 @@ export function withSequentialExecution( if (onExecutionStart && toolCallId !== undefined) { onExecutionStart(toolCallId); } - return await executeFn.call(baseTool, args, options); + try { + return await executeFn.call(baseTool, args, options); + } finally { + if (onExecutionEnd && toolCallId !== undefined) { + onExecutionEnd(toolCallId); + } + } }; wrappedTools[toolName] = wrappedTool;