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 18709c76aa..42e7d4ef01 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; @@ -460,11 +466,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(); + } } } } @@ -2415,6 +2433,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; @@ -2422,14 +2441,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 @@ -2442,7 +2463,7 @@ public HarnessAgent build() { effectiveConsolidationPrompt, memoryConfig.consolidationMaxTokens(), distributedStore != null ? distributedStore.baseStore() : null); - inner.middleware( + MemoryMaintenanceMiddleware memoryMaintenanceMw = new MemoryMaintenanceMiddleware( wsManager, consolidator, @@ -2450,7 +2471,9 @@ public HarnessAgent build() { memoryConfig.sessionRetentionDays(), memoryConfig.consolidationMinGap(), effectiveIsolationScope, - periodicGate)); + periodicGate); + inner.middleware(memoryMaintenanceMw); + pendingCloseableMiddlewares.add(memoryMaintenanceMw); } CompactionMiddleware compactionHook = null; if (!disableCompaction && compactionConfig != null) { @@ -2814,6 +2837,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(); + } +}