From 554f18dfc4fa7077e2a13bbd85b9ece5b9a00f90 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 18 Aug 2026 22:31:29 -0700 Subject: [PATCH] fix(agent-core): keep compaction streams under idle watchdog (#126159) --- .../compaction/branch-summarization.test.ts | 33 ++++++++ .../compaction/branch-summarization.ts | 3 +- .../src/harness/compaction/compaction.test.ts | 47 +++++++++++ .../src/harness/compaction/compaction.ts | 3 +- packages/agent-core/src/runtime-deps.ts | 9 +++ .../run/attempt-execution-phase.test.ts | 38 +++++++-- .../run/attempt-execution-phase.ts | 10 ++- .../run/llm-idle-timeout.abort.test.ts | 79 +++++++++++++++++++ .../run/llm-idle-timeout.ts | 22 ++++-- 9 files changed, 228 insertions(+), 16 deletions(-) create mode 100644 src/agents/embedded-agent-runner/run/llm-idle-timeout.abort.test.ts diff --git a/packages/agent-core/src/harness/compaction/branch-summarization.test.ts b/packages/agent-core/src/harness/compaction/branch-summarization.test.ts index abf1ab2c0b2c..60c0024b5fe4 100644 --- a/packages/agent-core/src/harness/compaction/branch-summarization.test.ts +++ b/packages/agent-core/src/harness/compaction/branch-summarization.test.ts @@ -93,6 +93,39 @@ function createLongBranchEntries(count: number): SessionTreeEntry[] { } describe("branch summarization", () => { + it("consumes the decorated stream before reading its result", async () => { + const model = createModel(128_000); + let consumed = false; + const streamFn = vi.fn(() => ({ + [Symbol.asyncIterator]() { + return { + async next() { + consumed = true; + return { done: true as const, value: undefined }; + }, + }; + }, + async result() { + if (!consumed) { + throw new Error("stream result read before iteration"); + } + return createResponse(model); + }, + })); + + await generateBranchSummary( + [createMessageEntry({ role: "user", content: "summarize this branch", timestamp: 1 }, 0)], + { + model, + apiKey: "test-key", + signal: new AbortController().signal, + streamFn, + }, + ); + + expect(consumed).toBe(true); + }); + it.each([ ["empty", []], ["whitespace-only", [{ type: "text" as const, text: " \n\t " }]], diff --git a/packages/agent-core/src/harness/compaction/branch-summarization.ts b/packages/agent-core/src/harness/compaction/branch-summarization.ts index 330c6710e329..ffef0bda22d7 100644 --- a/packages/agent-core/src/harness/compaction/branch-summarization.ts +++ b/packages/agent-core/src/harness/compaction/branch-summarization.ts @@ -2,6 +2,7 @@ import type { Model, StreamFn } from "@openclaw/llm-core"; import { type AgentCoreCompletionRuntimeDeps, + consumeAgentCoreStream, resolveAgentCoreCompleteFn, } from "../../runtime-deps.js"; import type { AgentMessage } from "../../types.js"; @@ -232,7 +233,7 @@ export async function generateBranchSummary( const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }; const streamOptions = { apiKey, headers, signal, maxTokens: maxSummaryOutputTokens }; const response = options.streamFn - ? await (await options.streamFn(model, context, streamOptions)).result() + ? await consumeAgentCoreStream(options.streamFn(model, context, streamOptions)) : await resolveAgentCoreCompleteFn(options.runtime)(model, context, streamOptions); if (response.stopReason === "aborted") { return err( diff --git a/packages/agent-core/src/harness/compaction/compaction.test.ts b/packages/agent-core/src/harness/compaction/compaction.test.ts index d3bbf5695a1b..4a87d0bede04 100644 --- a/packages/agent-core/src/harness/compaction/compaction.test.ts +++ b/packages/agent-core/src/harness/compaction/compaction.test.ts @@ -647,6 +647,53 @@ describe("session-entry compaction budgeting", () => { }); describe("generateSummary thinking options", () => { + it("consumes the decorated stream before reading its result", async () => { + const model: Model = { + id: "summary-model", + name: "Summary Model", + api: "test-api", + provider: "test-provider", + baseUrl: "https://example.test", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 100_000, + maxTokens: 8_000, + }; + let consumed = false; + const streamFn = vi.fn(() => ({ + [Symbol.asyncIterator]() { + return { + async next() { + consumed = true; + return { done: true as const, value: undefined }; + }, + }; + }, + async result() { + if (!consumed) { + throw new Error("stream result read before iteration"); + } + return createAssistant("summary", createUsage(1), 1); + }, + })); + + await generateSummary( + [{ role: "user", content: "hello", timestamp: 1 }], + model, + 1_000, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + streamFn, + ); + + expect(consumed).toBe(true); + }); + it("maps explicit Fable off to low effort for compaction", async () => { const model: Model = { id: "production-fable", diff --git a/packages/agent-core/src/harness/compaction/compaction.ts b/packages/agent-core/src/harness/compaction/compaction.ts index 5f59f318a656..80f6e36e3530 100644 --- a/packages/agent-core/src/harness/compaction/compaction.ts +++ b/packages/agent-core/src/harness/compaction/compaction.ts @@ -14,6 +14,7 @@ import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; import { resolveAgentReasoningOption } from "../../reasoning.js"; import { type AgentCoreCompletionRuntimeDeps, + consumeAgentCoreStream, resolveAgentCoreCompleteFn, } from "../../runtime-deps.js"; import type { AgentMessage, ThinkingLevel } from "../../types.js"; @@ -608,7 +609,7 @@ async function runSummarizationCompletion(params: { params.thinkingLevel, ); const response = params.streamFn - ? await (await params.streamFn(params.model, context, options)).result() + ? await consumeAgentCoreStream(params.streamFn(params.model, context, options)) : await resolveAgentCoreCompleteFn(params.runtime)(params.model, context, options); if (response.stopReason === "aborted") { return err( diff --git a/packages/agent-core/src/runtime-deps.ts b/packages/agent-core/src/runtime-deps.ts index 492eda15bc8a..974737853c13 100644 --- a/packages/agent-core/src/runtime-deps.ts +++ b/packages/agent-core/src/runtime-deps.ts @@ -34,6 +34,15 @@ export function resolveAgentCoreStreamFn( throw missingRuntimeDep("streamSimple"); } +/** Drain a host-decorated stream before reading its final assistant message. */ +export async function consumeAgentCoreStream(stream: ReturnType) { + const response = await stream; + for await (const _ of response) { + // drain + } + return response.result(); +} + /** Resolve the completion function used by non-streaming helper flows. */ export function resolveAgentCoreCompleteFn( runtime: AgentCoreCompletionRuntimeDeps | undefined, diff --git a/src/agents/embedded-agent-runner/run/attempt-execution-phase.test.ts b/src/agents/embedded-agent-runner/run/attempt-execution-phase.test.ts index 5444b226f292..29e950b9426d 100644 --- a/src/agents/embedded-agent-runner/run/attempt-execution-phase.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt-execution-phase.test.ts @@ -44,7 +44,12 @@ import { runEmbeddedAttemptExecutionPhase } from "./attempt-execution-phase.js"; type ExecutionInput = Parameters[0]; -function createFixture(options: { aborted?: boolean } = {}) { +function createFixture( + options: { + aborted?: boolean; + exerciseTerminalMerges?: boolean; + } = {}, +) { const order: string[] = []; const attemptAbortController = new AbortController(); if (options.aborted) { @@ -201,15 +206,19 @@ function createFixture(options: { aborted?: boolean } = {}) { }); mocks.prepareStream.mockImplementation((streamInput) => { order.push("stream"); - const idleError = new Error("idle timeout"); - mocks.installStreamGuards.mock.calls[0]?.[0].onIdleTimeout(idleError); - streamInput.markExternalAbort(); + if (options.exerciseTerminalMerges !== false) { + const idleError = new Error("idle timeout"); + mocks.installStreamGuards.mock.calls[0]?.[0].onIdleTimeout(idleError); + streamInput.markExternalAbort(); + } return streamResult; }); mocks.prepareTimeout.mockImplementation((timeoutInput) => { order.push("timeout"); - timeoutInput.markTimedOutDuringCompaction(); - timeoutInput.markTimedOutByRunBudget(); + if (options.exerciseTerminalMerges !== false) { + timeoutInput.markTimedOutDuringCompaction(); + timeoutInput.markTimedOutByRunBudget(); + } return timeoutResult; }); mocks.runSettledPhase.mockImplementation(async (settledInput) => { @@ -353,6 +362,23 @@ describe("runEmbeddedAttemptExecutionPhase", () => { expect(fixture.activeSession.prompt).not.toHaveBeenCalled(); }); + it("attributes an idle timeout during authoritative compaction to compaction", async () => { + const fixture = createFixture({ exerciseTerminalMerges: false }); + fixture.activeSession.isCompacting = true; + await runEmbeddedAttemptExecutionPhase(fixture.input); + const idleError = new Error("idle timeout"); + const guardInput = mocks.installStreamGuards.mock.calls[0]?.[0]; + + guardInput.onIdleTimeout(idleError); + + expect(fixture.state.terminal).toEqual({ + kind: "timeout", + phase: "compaction", + source: "idle", + }); + expect(fixture.runAbort).toHaveBeenCalledWith(true, idleError); + }); + it("flushes pending tool results and disposes the session when history preparation fails", async () => { const fixture = createFixture({ aborted: true }); const failure = new Error("history failed"); diff --git a/src/agents/embedded-agent-runner/run/attempt-execution-phase.ts b/src/agents/embedded-agent-runner/run/attempt-execution-phase.ts index c1c9c2a9f3d9..bf88aad5e988 100644 --- a/src/agents/embedded-agent-runner/run/attempt-execution-phase.ts +++ b/src/agents/embedded-agent-runner/run/attempt-execution-phase.ts @@ -143,7 +143,15 @@ export async function runEmbeddedAttemptExecutionPhase( }); input.externalAbortController.setRunAbort(abortRun); idleTimeoutTriggerRef.current = (error) => { - mergeTerminal({ kind: "timeout", phase: "prompt", source: "idle" }); + // Caller cancellation owns the terminal outcome when it beats a late watchdog callback. + if (input.runAbortController.signal.aborted) { + return; + } + mergeTerminal({ + kind: "timeout", + phase: activeSession.isCompacting ? "compaction" : "prompt", + source: "idle", + }); abortRun(true, error); }; const abortable = (promise: Promise): Promise => diff --git a/src/agents/embedded-agent-runner/run/llm-idle-timeout.abort.test.ts b/src/agents/embedded-agent-runner/run/llm-idle-timeout.abort.test.ts new file mode 100644 index 000000000000..ccd87e470122 --- /dev/null +++ b/src/agents/embedded-agent-runner/run/llm-idle-timeout.abort.test.ts @@ -0,0 +1,79 @@ +import type { AssistantMessageEventStream } from "openclaw/plugin-sdk/llm"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { streamWithIdleTimeout } from "./llm-idle-timeout.js"; + +function createNeverYieldingStream(): AsyncIterable { + return { + [Symbol.asyncIterator]() { + return { + async next() { + return new Promise>(() => {}); + }, + }; + }, + }; +} + +describe("streamWithIdleTimeout caller cancellation", () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it("preempts a provider iterator that ignores abort", async () => { + vi.useFakeTimers(); + const callerAbortController = new AbortController(); + const callerReason = new Error("caller cancelled"); + const baseFn = vi.fn().mockReturnValue(createNeverYieldingStream()); + const onIdleTimeout = vi.fn(); + const iterator = ( + streamWithIdleTimeout(baseFn, 50, onIdleTimeout)( + {} as Parameters[0], + {} as Parameters[1], + { signal: callerAbortController.signal }, + ) as AsyncIterable + )[Symbol.asyncIterator](); + const outcome = iterator.next().catch((error: unknown) => error); + + callerAbortController.abort(callerReason); + + await expect(outcome).resolves.toMatchObject({ + name: "AbortError", + message: callerReason.message, + cause: callerReason, + }); + await vi.advanceTimersByTimeAsync(50); + const providerSignal = (baseFn.mock.calls.at(0)?.[2] as { signal?: AbortSignal } | undefined) + ?.signal; + expect([providerSignal?.reason, onIdleTimeout.mock.calls.length]).toEqual([callerReason, 0]); + }); + + it("preempts provider stream creation", async () => { + vi.useFakeTimers(); + const callerAbortController = new AbortController(); + const callerReason = new Error("caller cancelled"); + const baseFnMock = vi.fn( + (_model: unknown, _context: unknown, _options?: { signal?: AbortSignal }) => + new Promise(() => {}), + ); + const baseFn = baseFnMock as unknown as Parameters[0]; + const onIdleTimeout = vi.fn(); + const pending = streamWithIdleTimeout(baseFn, 50, onIdleTimeout)( + {} as Parameters[0], + {} as Parameters[1], + { signal: callerAbortController.signal }, + ); + + callerAbortController.abort(callerReason); + + await expect(pending).rejects.toMatchObject({ + name: "AbortError", + message: callerReason.message, + cause: callerReason, + }); + await vi.advanceTimersByTimeAsync(50); + const providerSignal = ( + baseFnMock.mock.calls.at(0)?.[2] as { signal?: AbortSignal } | undefined + )?.signal; + expect([providerSignal?.reason, onIdleTimeout.mock.calls.length]).toEqual([callerReason, 0]); + }); +}); diff --git a/src/agents/embedded-agent-runner/run/llm-idle-timeout.ts b/src/agents/embedded-agent-runner/run/llm-idle-timeout.ts index fe31a87a62b6..348889fe0c3f 100644 --- a/src/agents/embedded-agent-runner/run/llm-idle-timeout.ts +++ b/src/agents/embedded-agent-runner/run/llm-idle-timeout.ts @@ -13,6 +13,7 @@ import { toErrorObject } from "../../../infra/errors.js"; import type { StreamFn } from "../../runtime/index.js"; import type { MutableAssistantMessageEventStream } from "../../stream-compat.js"; import { createStreamIteratorWrapper } from "../../stream-iterator-wrapper.js"; +import { abortable } from "./abortable.js"; import type { EmbeddedRunTrigger } from "./params.js"; import { getLastToolActivityMs, onToolActivity } from "./tool-activity-heartbeat.js"; @@ -451,6 +452,8 @@ export function streamWithIdleTimeout( const cleanupSourceSignal = () => { sourceSignal?.removeEventListener("abort", abortFromSourceSignal); }; + const withSourceAbort = (promise: Promise) => + sourceSignal ? abortable(sourceSignal, promise) : promise; const wrappedOptions = { ...options, signal: streamAbortController.signal, @@ -549,7 +552,11 @@ export function streamWithIdleTimeout( firstArmPending = true; armTimer(); }); - const result = await Promise.race([streamIterator.next(), timeoutPromise]); + // Providers may ignore their mirrored abort signal, so caller + // cancellation must also settle this exact iterator wait. + const result = await withSourceAbort( + Promise.race([streamIterator.next(), timeoutPromise]), + ); if (result.done) { cleanupIterator(); @@ -591,12 +598,13 @@ export function streamWithIdleTimeout( // Some providers return a pending Promise before the stream object exists; // protect that creation phase with the same idle watchdog. - return Promise.race([ - Promise.resolve(maybeStream), - createTimeoutPromise((timer) => { - streamPromiseTimer = timer; - }), - ]).then( + const timeoutPromise = createTimeoutPromise((timer) => { + streamPromiseTimer = timer; + }); + const streamPromise = withSourceAbort( + Promise.race([Promise.resolve(maybeStream), timeoutPromise]), + ); + return streamPromise.then( (stream) => { clearStreamPromiseTimer(); return wrapStream(stream);