From 991632d7fc3bcf8dbaf48776166da9347f5d662d Mon Sep 17 00:00:00 2001 From: webmin Date: Sat, 8 Aug 2026 12:39:30 +0800 Subject: [PATCH] fix(harness): detach memory flush and maintenance from agent call response MemoryFlushMiddleware and MemoryMaintenanceMiddleware appended their LLM-backed work via concatWith onto the returned Flux. Because ReActAgent.callInternal ends with takeLast(1), which can't emit until the upstream signals onComplete, callers consuming the agent response to completion (blockLast, takeLast(1), WebFlux controllers awaiting Mono) ended up waiting for the full memory flush LLM call (19-27s per call) and consolidation LLM call (~44s first run). The original implementation used doOnComplete (fire-and-forget). A later commit (PR #1802, RC4) swapped it for concatWith, introducing this regression. This fix restores the fire-and-forget behavior: both middlewares now subscribe their work independently of the returned Flux via doOnComplete().subscribe(), so the Flux completes as soon as the underlying agent call does. Additional engineering safeguards: - Both middlewares implement AutoCloseable and track pending Disposables in a Set guarded by synchronized(pending); a pre-subscribe closed check prevents new work, a post-subscribe check disposes any subscription that raced with close(). HarnessAgent.close() drains them (bounded by 5s) before workspace teardown to prevent races with temp dir deletion in tests and CLI shutdown. - captureFlushRequest is wrapped in try/catch so exceptions (e.g. from List.copyOf or resolveAgentState) cannot escape into the doOnComplete callback and turn an already-completed Flux into an error. - MemoryFlushMiddleware: flush LLM call is capped by a 5-minute timeout (after subscribeOn, measuring execution not queue wait) so a hung model provider cannot tie up a boundedElastic worker. - MemoryMaintenanceMiddleware: consolidate LLM call is capped by a 5-minute timeout. - A FlushRequest record snapshots messages (List.copyOf) on the complete thread before background execution, avoiding races with mutable AgentState context. New regression tests (TDD): - onAgent_completesBeforeSlowFlushFinishes / completesBeforeSlowConsolidationFinishes: mock a slow LLM, assert Flux completes in <1s while the LLM is still running in the background. - onAgent_flushError_doesNotPropagateToFlux / maintenanceError_doesNotPropagateToFlux: verify errors in detached work never reach the caller. - onAgent_afterClose_doesNotScheduleNewFlush / doesNotScheduleNewMaintenance: verify close() prevents new work; uses polling (not fixed sleep) for deterministic verification on slow CI. - close_waitsForPendingFlush_thenReturns / waitsForPendingMaintenance_thenReturns: verify close() drains in-flight work before returning. - close_disposesHungFlush_andReturnsPromptly: verifies close() disposes a hanging model (Flux.never) and returns within CLOSE_AWAIT_TIMEOUT. Fixes #2276 Fixes #2225 --- .../harness/agent/HarnessAgent.java | 38 +- .../middleware/MemoryFlushMiddleware.java | 154 ++++++- .../MemoryMaintenanceMiddleware.java | 152 ++++++- ...emoryFlushMiddlewareAsyncBehaviorTest.java | 387 ++++++++++++++++++ ...aintenanceMiddlewareAsyncBehaviorTest.java | 315 ++++++++++++++ 5 files changed, 1010 insertions(+), 36 deletions(-) create mode 100644 agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddlewareAsyncBehaviorTest.java create mode 100644 agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java index be1cb74f42..ae132b1360 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java @@ -181,6 +181,9 @@ public class HarnessAgent implements Agent, AutoCloseable { private final SkillAuditLog skillAuditLog; private final MemoryConfig memoryConfig; + /** Closeable middlewares (e.g. memory flush/maintenance) drained during {@link #close()}. */ + private final List closeableMiddlewares; + /** The subagent middleware (either SubagentsMiddleware or DynamicSubagentsMiddleware). */ private final Object subagentMiddleware; @@ -211,6 +214,7 @@ private HarnessAgent( SkillCurator skillCurator, SkillAuditLog skillAuditLog, MemoryConfig memoryConfig, + List closeableMiddlewares, Object subagentMiddleware, DistributedStore distributedStore, WorkspacePathNormalizer pathNormalizer) { @@ -229,6 +233,8 @@ private HarnessAgent( this.skillCurator = skillCurator; this.skillAuditLog = skillAuditLog; this.memoryConfig = memoryConfig != null ? memoryConfig : MemoryConfig.defaults(); + this.closeableMiddlewares = + closeableMiddlewares != null ? List.copyOf(closeableMiddlewares) : List.of(); this.subagentMiddleware = subagentMiddleware; this.distributedStore = distributedStore; this.pathNormalizer = pathNormalizer; @@ -431,11 +437,23 @@ public void close() { shutdownTaskRepository(); } finally { try { - if (ownedWorkspaceIndex != null) { - ownedWorkspaceIndex.close(); + // Drain detached memory flush/maintenance so async writes do not race with + // workspace teardown (e.g., temp workspace deletion in tests). + for (AutoCloseable mw : closeableMiddlewares) { + try { + mw.close(); + } catch (Exception e) { + log.warn("Failed to close middleware", e); + } } } finally { - delegate.close(); + try { + if (ownedWorkspaceIndex != null) { + ownedWorkspaceIndex.close(); + } + } finally { + delegate.close(); + } } } } @@ -2386,6 +2404,7 @@ public HarnessAgent build() { wsManager, effectiveTranscriptStore, transcriptTenant)); } Model memoryModel = memoryConfig.model() != null ? memoryConfig.model() : model; + List pendingCloseableMiddlewares = new java.util.ArrayList<>(); if (memoryModel != null && !disableMemoryHooks) { IsolationScope effectiveIsolationScope = fsIsolationScope; @@ -2393,14 +2412,16 @@ public HarnessAgent build() { memoryConfig.flushPrompt() != null ? memoryConfig.flushPrompt() : MemoryFlushManager.DEFAULT_FLUSH_PROMPT; - inner.middleware( + MemoryFlushMiddleware memoryFlushMw = new MemoryFlushMiddleware( wsManager, memoryModel, effectiveFlushPrompt, memoryConfig.flushTrigger(), effectiveIsolationScope, - periodicGate)); + periodicGate); + inner.middleware(memoryFlushMw); + pendingCloseableMiddlewares.add(memoryFlushMw); String effectiveConsolidationPrompt = memoryConfig.consolidationPrompt() != null @@ -2413,7 +2434,7 @@ public HarnessAgent build() { effectiveConsolidationPrompt, memoryConfig.consolidationMaxTokens(), distributedStore != null ? distributedStore.baseStore() : null); - inner.middleware( + MemoryMaintenanceMiddleware memoryMaintenanceMw = new MemoryMaintenanceMiddleware( wsManager, consolidator, @@ -2421,7 +2442,9 @@ public HarnessAgent build() { memoryConfig.sessionRetentionDays(), memoryConfig.consolidationMinGap(), effectiveIsolationScope, - periodicGate)); + periodicGate); + inner.middleware(memoryMaintenanceMw); + pendingCloseableMiddlewares.add(memoryMaintenanceMw); } CompactionMiddleware compactionHook = null; if (!disableCompaction && compactionConfig != null) { @@ -2785,6 +2808,7 @@ public HarnessAgent build() { pendingSkillCurator, pendingSkillAuditLog, memoryConfig, + pendingCloseableMiddlewares, capturedSubagentMw, distributedStore, pathNormalizer); diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java index 39c6f8c5cb..c5d4664cb8 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java @@ -28,10 +28,14 @@ import io.agentscope.harness.agent.memory.MemoryConfig; import io.agentscope.harness.agent.memory.MemoryFlushManager; import io.agentscope.harness.agent.workspace.WorkspaceManager; +import java.time.Duration; import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -39,10 +43,14 @@ /** * Middleware that triggers memory flush and message offload at the end of each agent call. * - *

Runs in {@link #onAgent}'s {@code doOnComplete} so long-term memories are extracted and - * persisted after every call, even when conversation compaction was not triggered during that - * call. When {@link CompactionMiddleware} is active, it handles flush/offload for the messages - * it summarizes; this middleware covers the remaining tail of messages that were kept verbatim. + *

Runs in a genuinely detached, fire-and-forget fashion: the flush {@code Mono} is subscribed + * independently of the returned {@code Flux} (via {@code doOnComplete}) rather than being + * concatenated onto it, so callers that wait for the response to complete (e.g. + * {@code blockLast()}, {@code takeLast(1)}) are not delayed by flush work. Long-term memories are + * extracted and persisted after every call, even when conversation compaction was not triggered + * during that call. When {@link CompactionMiddleware} is active, it handles flush/offload for the + * messages it summarizes; this middleware covers the remaining tail of messages that were kept + * verbatim. * *

Flush is gated by a {@link MemoryConfig.FlushTrigger}: *

    @@ -66,7 +74,7 @@ * the whole agent instance (prevents concurrent flush races on shared memory files). *
*/ -public class MemoryFlushMiddleware implements HarnessRuntimeMiddleware { +public class MemoryFlushMiddleware implements HarnessRuntimeMiddleware, AutoCloseable { private static final Logger log = LoggerFactory.getLogger(MemoryFlushMiddleware.class); @@ -77,6 +85,21 @@ public class MemoryFlushMiddleware implements HarnessRuntimeMiddleware { private final IsolationScope isolationScope; private final PeriodicGate periodicGate; + /** Upper bound on the flush LLM call, preventing a hung model from tying up a worker thread. */ + static final Duration FLUSH_TIMEOUT = Duration.ofMinutes(5); + + /** Upper bound {@link #close()} waits for outstanding fire-and-forget flushes to drain. */ + static final Duration CLOSE_AWAIT_TIMEOUT = Duration.ofSeconds(5); + + /** + * Tracks the {@link Disposable} of every fire-and-forget flush subscription that has been + * scheduled but not yet finished, so {@link #close()} can wait for them instead of leaving + * them racing against teardown of the workspace resources they read/write. + */ + private final Set pending = ConcurrentHashMap.newKeySet(); + + private volatile boolean closed = false; + public MemoryFlushMiddleware(WorkspaceManager workspaceManager, Model model) { this( workspaceManager, @@ -140,28 +163,89 @@ public Flux onAgent( AgentInput input, Function> next) { final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty(); - return next.apply(input) - .concatWith( - Mono.defer(() -> doFlush(agent, rc)) + return next.apply(input).doOnComplete(() -> scheduleFlush(agent, rc)); + } + + /** + * Fires the fire-and-forget flush {@code Mono} on {@code boundedElastic}, tracking its + * {@link Disposable} in {@link #pending} until it terminates so {@link #close()} can wait for + * it. No-ops once {@link #close()} has been called, so a call that races with shutdown + * doesn't spawn new untracked work. + * + *

The whole method body (snapshot capture of the RuntimeContext shallow copy plus the + * message copy, operator assembly, and subscribe) is wrapped in try/catch so that no + * exception can escape into the {@code doOnComplete} callback and turn an already-completed + * Flux into an error. + * + *

The {@code closed} flag and {@code pending} add are guarded by + * {@code synchronized(pending)}. The {@code pending} set (a ConcurrentHashMap key set) is + * reused as the mutex object to avoid allocating a separate lock; its own concurrency + * features are not relied upon for the closed-check/add atomicity. + */ + private void scheduleFlush(Agent agent, RuntimeContext rc) { + Disposable[] holder = new Disposable[1]; + // The whole body is wrapped in try/catch so that no exception (from capture, operator + // assembly, or subscribe) can escape into the doOnComplete callback and turn an + // already-completed Flux into an error. + try { + FlushRequest request = captureFlushRequest(agent, rc); + if (request == null) { + return; + } + // Subscribe and register inside one critical section: if the subscription were + // started outside the lock, a fast-completing flush could run doFinally (which + // removes the Disposable) before the scheduling thread executed pending.add, + // permanently leaking a terminated Disposable into pending and stalling close()'s + // drain until its timeout. doFinally reads holder[0] and removes under the same + // lock, so the add-before-remove ordering is also guaranteed to be VISIBLE (the + // holder write happens inside the lock, before its release; the read happens after + // acquiring it). subscribeOn(boundedElastic) ensures no callback ever runs on this + // thread, so the lock cannot self-deadlock. + synchronized (pending) { + if (closed) { + return; + } + Disposable d = + Mono.defer(() -> doFlush(request)) .subscribeOn(Schedulers.boundedElastic()) + .timeout(FLUSH_TIMEOUT) .onErrorResume( e -> { - log.warn("Memory flush failed: {}", e.getMessage()); + log.warn("Memory flush failed", e); return Mono.empty(); }) - .then(Mono.empty())); + .doFinally( + sig -> { + synchronized (pending) { + if (holder[0] != null) { + pending.remove(holder[0]); + } + } + }) + .subscribe(); + holder[0] = d; + pending.add(d); + } + } catch (Exception e) { + log.warn("Failed to schedule memory flush", e); + } } - private Mono doFlush(Agent agent, RuntimeContext rc) { + private FlushRequest captureFlushRequest(Agent agent, RuntimeContext rc) { AgentState state = RuntimeContext.resolveAgentState(rc, agent); if (state == null) { - return Mono.empty(); + return null; } List messages = state.getContext(); if (messages.isEmpty()) { - return Mono.empty(); + return null; } + return new FlushRequest(RuntimeContext.builder().from(rc).build(), List.copyOf(messages)); + } + private Mono doFlush(FlushRequest request) { + RuntimeContext rc = request.runtimeContext(); + List messages = request.messages(); MemoryFlushManager flushManager = new MemoryFlushManager(workspaceManager, model, flushPrompt); @@ -174,18 +258,56 @@ private Mono doFlush(Agent agent, RuntimeContext rc) { .doOnSuccess(v -> log.debug("Memory flush completed")) .onErrorResume( e -> { - log.warn("Memory flush failed: {}", e.getMessage()); + log.warn("Memory flush failed", e); return Mono.empty(); }); } else { log.debug("Memory flush skipped (trigger={})", flushTrigger); flushMono = Mono.empty(); } - - // Message offload is owned by TranscriptMiddleware (independent of memory flush). return flushMono; } + /** + * Waits (bounded by {@link #CLOSE_AWAIT_TIMEOUT}) for outstanding fire-and-forget flushes to + * finish, then disposes anything still outstanding. Intended to be called from {@code + * HarnessAgent#close()} so short-lived callers (tests using JUnit {@code @TempDir}, CLI runs, + * etc.) don't tear down the workspace while a detached flush write is still in flight. + * + *

The {@code closed} flag and {@code pending} add are guarded by {@code synchronized(pending)} + * so that a flush scheduled concurrently with close() is either fully tracked (and drained) or + * disposed immediately, but never lost. + */ + @Override + public void close() { + synchronized (pending) { + closed = true; + } + long deadline = System.nanoTime() + CLOSE_AWAIT_TIMEOUT.toNanos(); + while (!pending.isEmpty() && System.nanoTime() < deadline) { + try { + Thread.sleep(20); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + for (Disposable d : pending) { + d.dispose(); + } + pending.clear(); + } + + private record FlushRequest(RuntimeContext runtimeContext, List messages) {} + + /** + * Returns whether any fire-and-forget flush is currently in flight. Package-private, intended + * for tests that need to poll for quiescence instead of relying on a fixed sleep. + */ + boolean hasPendingFlushes() { + return !pending.isEmpty(); + } + /** * Returns whether this call should trigger a flush, applying the configured trigger policy. * For {@link MemoryConfig.FlushMode#THROTTLED}, uses an {@link AtomicReference#compareAndSet} diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddleware.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddleware.java index 2ac5db9788..5b18c2bf20 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddleware.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddleware.java @@ -31,9 +31,12 @@ import java.time.Duration; import java.time.Instant; import java.time.LocalDate; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.scheduler.Schedulers; @@ -41,9 +44,11 @@ /** * Middleware that performs periodic memory maintenance after each agent call. * - *

Fires on the agent invocation completion (via {@code onAgent concatWith}, after - * {@link MemoryFlushMiddleware}) and is throttled by a configurable minimum gap so it - * does not run on every single call. + *

Fires in a genuinely detached, fire-and-forget fashion: the maintenance {@code Mono} is + * subscribed independently of the returned {@code Flux} (via {@code doOnComplete}) rather than + * being concatenated onto it, so callers that wait for the response to complete (e.g. + * {@code blockLast()}, {@code takeLast(1)}) are not delayed by maintenance work. It is also + * throttled by a configurable minimum gap so it does not run on every single call. * *

Maintenance steps executed in order: *

    @@ -63,13 +68,23 @@ * the whole agent instance (prevents concurrent maintenance races on shared memory files). * */ -public class MemoryMaintenanceMiddleware implements HarnessRuntimeMiddleware { +public class MemoryMaintenanceMiddleware implements HarnessRuntimeMiddleware, AutoCloseable { private static final Logger log = LoggerFactory.getLogger(MemoryMaintenanceMiddleware.class); /** Default minimum gap between two maintenance runs. */ public static final Duration DEFAULT_MIN_GAP = Duration.ofMinutes(30); + /** + * Upper bound on the consolidation LLM call. Set slightly below + * {@link #MAINTENANCE_TIMEOUT} so the error log can distinguish "consolidation itself + * timed out" from "the entire maintenance run timed out". + */ + static final Duration CONSOLIDATION_TIMEOUT = Duration.ofMinutes(4).plusSeconds(30); + + /** Upper bound on the entire maintenance run (file IO + LLM), consistent with flush. */ + static final Duration MAINTENANCE_TIMEOUT = Duration.ofMinutes(5); + private final WorkspaceManager workspaceManager; private final MemoryConsolidator consolidator; private final int dailyFileRetentionDays; @@ -78,6 +93,19 @@ public class MemoryMaintenanceMiddleware implements HarnessRuntimeMiddleware { private final IsolationScope isolationScope; private final PeriodicGate periodicGate; + /** Upper bound {@link #close()} waits for outstanding fire-and-forget runs to drain. */ + static final Duration CLOSE_AWAIT_TIMEOUT = Duration.ofSeconds(5); + + /** + * Tracks the {@link Disposable} of every fire-and-forget maintenance subscription that has + * been scheduled but not yet finished, so {@link #close()} can wait for/dispose them instead + * of leaving them racing against teardown of the resources they read/write (e.g. a workspace + * directory being deleted). + */ + private final Set pending = ConcurrentHashMap.newKeySet(); + + private volatile boolean closed = false; + public MemoryMaintenanceMiddleware( WorkspaceManager workspaceManager, MemoryConsolidator consolidator, @@ -140,27 +168,125 @@ public Flux onAgent( AgentInput input, Function> next) { final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty(); + // Snapshot the full RuntimeContext (shallow copy of attribute maps) at completion time — + // not eagerly here — so attributes added during the call are visible, matching the + // original rc-at-completion semantics. The copy makes the background task immune to + // concurrent mutations on the caller's attribute maps after the call completes, while + // preserving all fields that filesystem operations or NamespaceFactory implementations + // may depend on (not just userId/sessionId). return next.apply(input) - .concatWith( - Mono.fromRunnable(() -> maybeRunMaintenance(rc)) + .doOnComplete(() -> scheduleMaintenance(RuntimeContext.builder().from(rc).build())); + } + + /** + * Fires the fire-and-forget maintenance {@code Mono} on {@code boundedElastic}, tracking its + * {@link Disposable} in {@link #pending} until it terminates so {@link #close()} can wait for + * it. No-ops once {@link #close()} has been called, so a call that races with shutdown + * doesn't spawn new untracked work. + * + *

    The whole method body (RuntimeContext copy, operator assembly, and subscribe) is + * wrapped in try/catch so that no exception can escape into the {@code doOnComplete} + * callback and turn an already-completed Flux into an error. + * + *

    The {@code closed} flag and {@code pending} add are guarded by + * {@code synchronized(pending)} so that maintenance scheduled concurrently with close() is + * either fully tracked (and drained) or disposed immediately, but never lost. The + * {@code pending} set (a ConcurrentHashMap key set) is reused as the mutex object to avoid + * allocating a separate lock; its own concurrency features are not relied upon for the + * closed-check/add atomicity. + */ + private void scheduleMaintenance(RuntimeContext snapshot) { + Disposable[] holder = new Disposable[1]; + // The whole body is wrapped in try/catch so that no exception (from the RuntimeContext + // copy, operator assembly, or subscribe) can escape into the doOnComplete callback and + // turn an already-completed Flux into an error. + try { + // Subscribe and register inside one critical section: if the subscription were + // started outside the lock, a fast-completing run (e.g. the throttle gate rejecting + // within its gap) could run doFinally (which removes the Disposable) before the + // scheduling thread executed pending.add, permanently leaking a terminated + // Disposable into pending and stalling close()'s drain until its timeout. + // doFinally reads holder[0] and removes under the same lock, so the + // add-before-remove ordering is also guaranteed to be VISIBLE (the holder write + // happens inside the lock, before its release; the read happens after acquiring + // it). subscribeOn(boundedElastic) ensures no callback ever runs on this thread, so + // the lock cannot self-deadlock. + synchronized (pending) { + if (closed) { + return; + } + Disposable d = + Mono.fromRunnable(() -> maybeRunMaintenance(snapshot)) .subscribeOn(Schedulers.boundedElastic()) + .timeout(MAINTENANCE_TIMEOUT) .onErrorResume( e -> { - log.warn( - "Memory maintenance failed: {}", - e.getMessage()); + log.warn("Memory maintenance failed", e); return Mono.empty(); - })); + }) + .doFinally( + sig -> { + synchronized (pending) { + if (holder[0] != null) { + pending.remove(holder[0]); + } + } + }) + .subscribe(); + holder[0] = d; + pending.add(d); + } + } catch (Exception e) { + log.warn("Failed to schedule memory maintenance", e); + } + } + + /** + * Waits (bounded by {@link #CLOSE_AWAIT_TIMEOUT}) for outstanding fire-and-forget maintenance + * runs to finish, then disposes anything still outstanding. Intended to be called from {@code + * HarnessAgent#close()} so short-lived callers (tests using JUnit {@code @TempDir}, CLI runs, + * etc.) don't tear down the workspace while a detached maintenance write is still in flight. + */ + @Override + public void close() { + synchronized (pending) { + closed = true; + } + long deadline = System.nanoTime() + CLOSE_AWAIT_TIMEOUT.toNanos(); + while (!pending.isEmpty() && System.nanoTime() < deadline) { + try { + Thread.sleep(20); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + for (Disposable d : pending) { + d.dispose(); + } + pending.clear(); + } + + /** + * Returns whether any fire-and-forget maintenance is currently in flight. Package-private, + * intended for tests that need to poll for quiescence instead of relying on a fixed sleep. + */ + boolean hasPendingMaintenance() { + return !pending.isEmpty(); } private void maybeRunMaintenance(RuntimeContext rc) { + // rc is a shallow-copy snapshot taken at doOnComplete time (see onAgent), so attribute + // maps are independent of the caller's original context. Value objects within the maps + // are shared, but all identity fields (userId, sessionId) and attributes needed by + // filesystem namespace resolution are preserved. if (!periodicGate.tryClaim(compositeTimerKey(rc), minGap)) { return; } try { runMaintenance(rc); } catch (Exception e) { - log.warn("Memory maintenance failed: {}", e.getMessage()); + log.warn("Memory maintenance failed", e); } } @@ -243,9 +369,9 @@ private void consolidateMemory(RuntimeContext rc) { return; } try { - consolidator.consolidate(rc).block(); + consolidator.consolidate(rc).timeout(CONSOLIDATION_TIMEOUT).block(); } catch (Exception e) { - log.warn("Memory consolidation failed: {}", e.getMessage()); + log.warn("Memory consolidation failed", e); } } diff --git a/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddlewareAsyncBehaviorTest.java b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddlewareAsyncBehaviorTest.java new file mode 100644 index 0000000000..15cc472201 --- /dev/null +++ b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddlewareAsyncBehaviorTest.java @@ -0,0 +1,387 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.harness.agent.middleware; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.Agent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.event.AgentEvent; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.middleware.AgentInput; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.Model; +import io.agentscope.core.state.AgentState; +import io.agentscope.harness.agent.coordination.LocalPeriodicGate; +import io.agentscope.harness.agent.memory.MemoryConfig; +import io.agentscope.harness.agent.memory.MemoryFlushManager; +import io.agentscope.harness.agent.workspace.WorkspaceManager; +import java.nio.file.Path; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import reactor.core.publisher.Flux; + +/** + * Regression tests for the issue where {@code onAgent} used {@code concatWith} to append + * memory flush onto the returned {@link Flux}, forcing callers that consume the response to + * completion (e.g. {@code blockLast()}, {@code takeLast(1)}) to wait for the full flush LLM + * call duration. + * + *

    {@code onAgent} must detach flush so the returned Flux completes as soon as the + * underlying agent call completes, independent of how long the flush LLM takes. + * + *

    See issue #2276 + * and issue #2225. + */ +@Tag("unit") +class MemoryFlushMiddlewareAsyncBehaviorTest { + + @BeforeEach + void resetSharedTimerMap() { + LocalPeriodicGate.clearForTests(); + } + + /** + * The returned Flux must complete before a slow flush LLM finishes. With {@code concatWith} + * the Flux is held open until the LLM call completes, inflating latency by the full model + * round-trip (reported as 19-27s per call in production). + */ + @Test + void onAgent_completesBeforeSlowFlushFinishes(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + + CountDownLatch flushStarted = new CountDownLatch(1); + CountDownLatch releaseFlush = new CountDownLatch(1); + Model slowModel = mock(Model.class); + when(slowModel.stream(any(), any(), any())) + .thenAnswer( + inv -> + Flux.create( + sink -> { + flushStarted.countDown(); + try { + releaseFlush.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + sink.complete(); + })); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + slowModel, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + long start = System.nanoTime(); + List events = + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue(events != null && events.isEmpty()); + assertTrue( + elapsedMs < 1000, + () -> + "onAgent should complete before flush finishes, but took " + + elapsedMs + + "ms"); + assertTrue( + flushStarted.await(2, TimeUnit.SECONDS), + "flush should have started on a detached background thread"); + + releaseFlush.countDown(); + } + + /** + * A flush error must not propagate to the returned Flux. With detach, the error is caught + * inside the background subscription and logged, never reaching the caller. + */ + @Test + void onAgent_flushError_doesNotPropagateToFlux(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + Model errorModel = mock(Model.class); + when(errorModel.stream(any(), any(), any())) + .thenReturn(Flux.error(new RuntimeException("flush LLM failed"))); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + errorModel, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + // Should complete normally, not error + List events = + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + assertTrue(events != null && events.isEmpty()); + } + + @Test + void onAgent_noMessages_completesImmediately(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + Model model = mock(Model.class); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + model, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(AgentState.builder().sessionId("s1").build()); + + AgentInput input = new AgentInput(List.of()); + + long start = System.nanoTime(); + List events = + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue(events != null && events.isEmpty()); + assertTrue(elapsedMs < 500, "should complete immediately with no messages to flush"); + } + + @Test + void onAgent_afterClose_doesNotScheduleNewFlush(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + Model model = mock(Model.class); + when(model.stream(any(), any(), any())).thenReturn(Flux.empty()); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + model, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + // close() with nothing in flight returns immediately but marks the middleware closed. + mw.close(); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + List events = + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + assertTrue(events != null && events.isEmpty()); + + // Poll for the background subscription to settle, then verify no model interaction. + awaitPendingSettled(mw, Duration.ofSeconds(2)); + verifyNoInteractions(model); + } + + @Test + void close_waitsForPendingFlush_thenReturns(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + + CountDownLatch releaseFlush = new CountDownLatch(1); + AtomicBoolean flushCompleted = new AtomicBoolean(false); + Model slowModel = mock(Model.class); + when(slowModel.stream(any(), any(), any())) + .thenAnswer( + inv -> + Flux.create( + sink -> { + try { + releaseFlush.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + flushCompleted.set(true); + sink.complete(); + })); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + slowModel, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + // Trigger a flush (detached, running in background) + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + // Release the flush so close() can drain it + releaseFlush.countDown(); + + // close() should return within the await timeout (flush completes quickly after release) + long start = System.nanoTime(); + mw.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + elapsedMs < 3000, + "close() should return promptly after flush completes, took " + elapsedMs + "ms"); + assertTrue(flushCompleted.get(), "flush should have completed before close() returned"); + } + + /** + * When the flush LLM hangs indefinitely, close() must still return within CLOSE_AWAIT_TIMEOUT + * by disposing the outstanding subscription. This verifies the timeout+dispose safety net + * works without waiting the full 5-minute FLUSH_TIMEOUT. + */ + @Test + void close_disposesHungFlush_andReturnsPromptly(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + + Model hangingModel = mock(Model.class); + when(hangingModel.stream(any(), any(), any())) + .thenAnswer(inv -> Flux.never()); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + hangingModel, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + // Trigger a flush that will hang (model never returns) + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + assertTrue(mw.hasPendingFlushes(), "flush should be pending (hanging)"); + + // close() should dispose the hung flush and return within the await timeout (5s) + long start = System.nanoTime(); + mw.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + elapsedMs < 6000, + "close() should return within CLOSE_AWAIT_TIMEOUT even with hung flush, took " + + elapsedMs + + "ms"); + assertFalse(mw.hasPendingFlushes(), "no pending flushes should remain after close()"); + } + + /** + * Fast-completing flushes must not leak terminated Disposables into pending. Regression + * guard for the subscribe-then-add race: with subscribe() outside the critical section, a + * flush completing between subscribe() and pending.add() ran doFinally before the entry was + * registered, leaving a dead Disposable in pending that stalled close()'s drain for its full + * timeout. Scheduling many instant flushes makes the old race highly likely to reproduce; + * the fixed single-lock pattern keeps pending consistent on every iteration. + */ + @Test + void rapidFastCompletingFlushes_leaveNoTerminatedEntriesInPending(@TempDir Path tmp) + throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + Model instantModel = mock(Model.class); + when(instantModel.stream(any(), any(), any())).thenReturn(Flux.empty()); + + MemoryFlushMiddleware mw = + new MemoryFlushMiddleware( + wsm, + instantModel, + MemoryFlushManager.DEFAULT_FLUSH_PROMPT, + MemoryConfig.FlushTrigger.always()); + + RuntimeContext rc = RuntimeContext.builder().sessionId("s1").userId("u1").build(); + rc.setAgentState(stateWithMessages(userMsg("hello"))); + AgentInput input = new AgentInput(List.of(userMsg("hello"))); + + for (int i = 0; i < 50; i++) { + mw.onAgent((Agent) null, rc, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + } + + // If any terminated Disposable leaked into pending, this polls to its 2s timeout and + // the assertFalse inside fails. + awaitPendingSettled(mw, Duration.ofSeconds(2)); + } + + // ── helpers ────────────────────────────────────────────────────────────── + + /** + * Polls until the middleware has no pending background subscriptions. Throws + * AssertionError if the timeout expires with pending work still in flight, so a + * slow CI cannot silently mask a scheduling bug. + */ + private static void awaitPendingSettled(MemoryFlushMiddleware mw, Duration timeout) + throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (System.nanoTime() < deadline) { + if (!mw.hasPendingFlushes()) { + return; + } + Thread.sleep(10); + } + assertFalse(mw.hasPendingFlushes(), "pending flushes should have settled"); + } + + private static Msg userMsg(String text) { + return Msg.builder() + .name("user") + .role(MsgRole.USER) + .content(List.of(TextBlock.builder().text(text).build())) + .build(); + } + + private static AgentState stateWithMessages(Msg... msgs) { + AgentState.Builder b = AgentState.builder().sessionId("s1"); + for (Msg m : msgs) { + b.addMessage(m); + } + return b.build(); + } +} diff --git a/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java new file mode 100644 index 0000000000..5a18ac3e77 --- /dev/null +++ b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java @@ -0,0 +1,315 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.agentscope.harness.agent.middleware; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.Agent; +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.event.AgentEvent; +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.middleware.AgentInput; +import io.agentscope.harness.agent.coordination.LocalPeriodicGate; +import io.agentscope.harness.agent.memory.MemoryConsolidator; +import io.agentscope.harness.agent.workspace.WorkspaceManager; +import java.nio.file.Path; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +/** + * Regression tests for the issue where {@code onAgent} used {@code concatWith} to append memory + * maintenance onto the returned {@link Flux}, forcing callers that consume the response to + * completion (e.g. {@code blockLast()}, {@code takeLast(1)}) to wait for the full consolidation + * LLM call inside {@code consolidateMemory()}. + * + *

    {@code onAgent} must detach maintenance so the returned Flux completes as soon as the + * underlying agent call completes, independent of how long maintenance takes. + * + *

    See issue #2225 + * and issue #2276. + */ +@Tag("unit") +class MemoryMaintenanceMiddlewareAsyncBehaviorTest { + + @BeforeEach + void resetSharedTimerMap() { + LocalPeriodicGate.clearForTests(); + } + + /** + * The returned Flux must complete before a slow consolidation finishes. With + * {@code concatWith} the Flux is held open until {@code consolidate().block()} returns, + * inflating latency by the full model round-trip (reported as ~44s in production). + */ + @Test + void onAgent_completesBeforeSlowConsolidationFinishes(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + + CountDownLatch consolidationStarted = new CountDownLatch(1); + CountDownLatch releaseConsolidation = new CountDownLatch(1); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + when(consolidator.consolidate(any())) + .thenAnswer( + invocation -> + Mono.fromRunnable( + () -> { + consolidationStarted.countDown(); + await(releaseConsolidation); + })); + + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + long start = System.nanoTime(); + List events = + mw.onAgent( + (Agent) null, + RuntimeContext.empty(), + input, + in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue(events != null && events.isEmpty()); + assertTrue( + elapsedMs < 1000, + () -> + "onAgent should complete before consolidation finishes, but took " + + elapsedMs + + "ms"); + assertTrue( + consolidationStarted.await(2, TimeUnit.SECONDS), + "consolidation should have started on a detached background thread"); + + releaseConsolidation.countDown(); + } + + /** + * A maintenance error must not propagate to the returned Flux. With detach, the error is + * caught inside the background subscription and logged, never reaching the caller. + */ + @Test + void onAgent_maintenanceError_doesNotPropagateToFlux(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + when(consolidator.consolidate(any())) + .thenReturn(Mono.error(new RuntimeException("consolidation LLM failed"))); + + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + List events = + mw.onAgent( + (Agent) null, + RuntimeContext.empty(), + input, + in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + assertTrue(events != null && events.isEmpty()); + } + + @Test + void onAgent_afterClose_doesNotScheduleNewMaintenance(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + + // Nothing in flight, so close() returns immediately but marks the middleware closed. + mw.close(); + + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + List events = + mw.onAgent( + (Agent) null, + RuntimeContext.empty(), + input, + in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + assertTrue(events != null && events.isEmpty()); + + // Poll for the background subscription to settle, then verify no consolidator + // interaction. + awaitPendingSettled(mw, Duration.ofSeconds(2)); + verifyNoInteractions(consolidator); + } + + @Test + void close_waitsForPendingMaintenance_thenReturns(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + + CountDownLatch releaseConsolidation = new CountDownLatch(1); + AtomicBoolean consolidationCompleted = new AtomicBoolean(false); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + when(consolidator.consolidate(any())) + .thenAnswer( + invocation -> + Mono.fromRunnable( + () -> { + try { + releaseConsolidation.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + consolidationCompleted.set(true); + })); + + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + // Trigger maintenance (detached, running in background) + mw.onAgent((Agent) null, RuntimeContext.empty(), input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + // Release consolidation so close() can drain it + releaseConsolidation.countDown(); + + long start = System.nanoTime(); + mw.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + elapsedMs < 3000, + "close() should return promptly after maintenance completes, took " + + elapsedMs + + "ms"); + assertTrue( + consolidationCompleted.get(), + "consolidation should have completed before close() returned"); + } + + /** + * When consolidation hangs indefinitely, close() must still return within + * CLOSE_AWAIT_TIMEOUT by disposing the outstanding subscription. This verifies the + * timeout+dispose safety net works without waiting the full 5-minute + * CONSOLIDATION_TIMEOUT. + */ + @Test + void close_disposesHungConsolidation_andReturnsPromptly(@TempDir Path tmp) throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + when(consolidator.consolidate(any())).thenReturn(Mono.never()); + + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + // Trigger maintenance that will hang (consolidation never returns) + mw.onAgent((Agent) null, RuntimeContext.empty(), input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + assertTrue(mw.hasPendingMaintenance(), "maintenance should be pending (hanging)"); + + // close() should dispose the hung maintenance and return within the await timeout (5s) + long start = System.nanoTime(); + mw.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + elapsedMs < 6000, + "close() should return within CLOSE_AWAIT_TIMEOUT even with hung consolidation," + + " took " + + elapsedMs + + "ms"); + assertFalse( + mw.hasPendingMaintenance(), "no pending maintenance should remain after close()"); + } + + /** + * Fast-completing maintenance runs must not leak terminated Disposables into pending. + * Regression guard for the subscribe-then-add race: with subscribe() outside the critical + * section, a run rejected by the throttle gate (microsecond completion) could run doFinally + * before the scheduling thread executed pending.add, leaving a dead Disposable in pending + * that stalled close()'s drain for its full timeout. Scheduling many instant runs makes the + * old race highly likely to reproduce; the fixed single-lock pattern keeps pending + * consistent on every iteration. + */ + @Test + void rapidFastCompletingMaintenance_leavesNoTerminatedEntriesInPending(@TempDir Path tmp) + throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + MemoryConsolidator consolidator = mock(MemoryConsolidator.class); + when(consolidator.consolidate(any())).thenReturn(Mono.empty()); + + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, consolidator); + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + for (int i = 0; i < 50; i++) { + mw.onAgent((Agent) null, RuntimeContext.empty(), input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + } + + // If any terminated Disposable leaked into pending, this polls to its 2s timeout and + // the assertFalse inside fails. + awaitPendingSettled(mw, Duration.ofSeconds(2)); + } + + // ── helpers ────────────────────────────────────────────────────────────── + + /** + * Polls until the middleware has no pending background subscriptions. Throws + * AssertionError if the timeout expires with pending work still in flight, so a + * slow CI cannot silently mask a scheduling bug. + */ + private static void awaitPendingSettled(MemoryMaintenanceMiddleware mw, Duration timeout) + throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (System.nanoTime() < deadline) { + if (!mw.hasPendingMaintenance()) { + return; + } + Thread.sleep(10); + } + assertFalse(mw.hasPendingMaintenance(), "pending maintenance should have settled"); + } + + private static void await(CountDownLatch latch) { + try { + latch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + private static Msg userMsg(String text) { + return Msg.builder() + .name("user") + .role(MsgRole.USER) + .content(List.of(TextBlock.builder().text(text).build())) + .build(); + } +}