From 9d23167fb0e741849f3229c470cab6486bd0b6c3 Mon Sep 17 00:00:00 2001 From: Asad Date: Thu, 16 Jul 2026 09:59:40 +0530 Subject: [PATCH 1/3] fix(harness): detach memory maintenance from agent call response MemoryMaintenanceMiddleware concatenated the maintenance Mono onto the returned Flux via concatWith, so callers that consume the response to completion (blockLast(), takeLast(1)) waited for the full LLM-based consolidation call to finish -- inflating agent call latency by up to 30s and contradicting documentation that describes this as a background/fire-and-forget step. Detach maintenance via doOnComplete(...).subscribe() instead, so the returned Flux completes as soon as the underlying agent call does, independent of maintenance duration. Also add a timeout to the consolidation LLM call so a hung provider can't tie up a shared boundedElastic thread indefinitely. Fixes #2225 --- .../agent/memory/MemoryConsolidator.java | 12 ++ .../MemoryMaintenanceMiddleware.java | 31 ++-- .../memory/MemoryConsolidatorTimeoutTest.java | 69 ++++++++ ...aintenanceMiddlewareAsyncBehaviorTest.java | 165 ++++++++++++++++++ 4 files changed, 264 insertions(+), 13 deletions(-) create mode 100644 agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/MemoryConsolidatorTimeoutTest.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/memory/MemoryConsolidator.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/MemoryConsolidator.java index b324db42a4..fdc61cbf64 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/MemoryConsolidator.java +++ b/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/MemoryConsolidator.java @@ -19,11 +19,13 @@ import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; import io.agentscope.core.message.TextBlock; +import io.agentscope.core.model.ExecutionConfig; import io.agentscope.core.model.Model; import io.agentscope.harness.agent.filesystem.AbstractFilesystem; import io.agentscope.harness.agent.filesystem.model.FileInfo; import io.agentscope.harness.agent.filesystem.model.GlobResult; import io.agentscope.harness.agent.workspace.WorkspaceManager; +import java.time.Duration; import java.time.Instant; import java.util.ArrayList; import java.util.Comparator; @@ -61,6 +63,15 @@ public class MemoryConsolidator { /** Hidden state file inside {@code memory/} tracking the last consolidation Instant. */ public static final String STATE_FILE = ".consolidation_state"; + /** + * Upper bound on the consolidation LLM call. Consolidation runs on a shared + * {@code boundedElastic} worker thread (see {@code MemoryMaintenanceMiddleware}); without a + * timeout, a hung model provider would tie up that thread indefinitely. Reuses + * {@link ExecutionConfig#MODEL_DEFAULTS}'s timeout so this stays consistent with the + * standard per-call model timeout used elsewhere in the codebase. + */ + static final Duration CONSOLIDATION_TIMEOUT = ExecutionConfig.MODEL_DEFAULTS.getTimeout(); + /** * Default prompt for the consolidation step. Exposed publicly so callers can extend * (e.g. append project-specific guidelines) when constructing @@ -167,6 +178,7 @@ public Mono consolidate(RuntimeContext rc) { .build()); return model.stream(messages, null, null) + .timeout(CONSOLIDATION_TIMEOUT) .reduce( new StringBuilder(), (sb, chatResponse) -> { 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 2a8b75e764..3a37260fb5 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 @@ -41,9 +41,12 @@ /** * 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 on the agent invocation completion (via {@code onAgent doOnComplete}, after + * {@link MemoryFlushMiddleware}) in a genuinely detached, fire-and-forget fashion: the + * maintenance {@code Mono} is subscribed independently of the returned {@code Flux} 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: *

    @@ -140,16 +143,18 @@ public Flux onAgent( Function> next) { final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty(); return next.apply(input) - .concatWith( - Mono.fromRunnable(() -> maybeRunMaintenance(rc)) - .subscribeOn(Schedulers.boundedElastic()) - .onErrorResume( - e -> { - log.warn( - "Memory maintenance failed: {}", - e.getMessage()); - return Mono.empty(); - })); + .doOnComplete( + () -> + Mono.fromRunnable(() -> maybeRunMaintenance(rc)) + .subscribeOn(Schedulers.boundedElastic()) + .onErrorResume( + e -> { + log.warn( + "Memory maintenance failed: {}", + e.getMessage()); + return Mono.empty(); + }) + .subscribe()); } private void maybeRunMaintenance(RuntimeContext rc) { diff --git a/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/MemoryConsolidatorTimeoutTest.java b/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/MemoryConsolidatorTimeoutTest.java new file mode 100644 index 0000000000..dfd51eb3d1 --- /dev/null +++ b/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/MemoryConsolidatorTimeoutTest.java @@ -0,0 +1,69 @@ +/* + * 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.memory; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.agentscope.core.agent.RuntimeContext; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.Model; +import io.agentscope.harness.agent.filesystem.remote.RemoteFilesystem; +import io.agentscope.harness.agent.filesystem.remote.store.InMemoryStore; +import io.agentscope.harness.agent.workspace.WorkspaceManager; +import java.nio.file.Path; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.TimeoutException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import reactor.core.publisher.Flux; +import reactor.test.StepVerifier; + +/** + * Verifies that {@link MemoryConsolidator#consolidate} does not hang forever when the model + * provider never responds. Without a timeout, this would tie up a shared {@code boundedElastic} + * worker thread indefinitely (see {@link io.agentscope.harness.agent.middleware + * .MemoryMaintenanceMiddleware}, which runs consolidation on that scheduler). + */ +class MemoryConsolidatorTimeoutTest { + + @Test + void consolidate_timesOutInsteadOfHangingForever(@TempDir Path tmp) throws Exception { + InMemoryStore store = new InMemoryStore(); + List ns = List.of("test-ns"); + RemoteFilesystem fs = new RemoteFilesystem(store, ns); + + try (WorkspaceManager wsm = new WorkspaceManager(tmp, fs)) { + // Seed a fresh daily entry so consolidate() doesn't short-circuit with Mono.empty() + // before ever reaching the model call. + wsm.writeUtf8WorkspaceRelative( + RuntimeContext.empty(), "memory/2025-06-15.md", "Some daily notes"); + + Model hungModel = mock(Model.class); + when(hungModel.stream(anyList(), any(), any())).thenReturn(Flux.never()); + + MemoryConsolidator consolidator = new MemoryConsolidator(wsm, hungModel); + + StepVerifier.withVirtualTime(() -> consolidator.consolidate(RuntimeContext.empty())) + .thenAwait(MemoryConsolidator.CONSOLIDATION_TIMEOUT.plus(Duration.ofSeconds(1))) + .expectError(TimeoutException.class) + .verify(Duration.ofSeconds(10)); + } + } +} 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..b55d88fab4 --- /dev/null +++ b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java @@ -0,0 +1,165 @@ +/* + * 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.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.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.Test; +import org.junit.jupiter.api.io.TempDir; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Hooks; +import reactor.core.publisher.Mono; + +/** + * Regression tests for the issue where {@code onAgent} used {@code concatWith} to append memory + * maintenance onto the returned {@link reactor.core.publisher.Flux}, forcing callers that consume + * the response to completion (e.g. {@code blockLast()}, {@code takeLast(1)}) to wait for the full + * maintenance duration — including the {@link MemoryConsolidator#consolidate} LLM call. + * + *

    {@code onAgent} must now detach maintenance via {@code doOnComplete(...).subscribe()} so the + * returned Flux completes as soon as the underlying agent call completes, independent of how long + * maintenance takes. + */ +class MemoryMaintenanceMiddlewareAsyncBehaviorTest { + + @BeforeEach + void resetSharedMap() { + MemoryMaintenanceMiddleware.SHARED_LAST_RUN_AT.clear(); + } + + @Test + void onAgent_completesBeforeSlowConsolidationFinishes(@TempDir Path tmp) throws Exception { + // No filesystem configured -> expireDailyFiles/pruneOldSessions are no-ops, isolating + // this test to the consolidation step's timing. + 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(); + } + + @Test + void onAgent_onErrorResume_handlesExceptionFromDetachedMaintenance(@TempDir Path tmp) + throws Exception { + WorkspaceManager wsm = new WorkspaceManager(tmp); + MemoryMaintenanceMiddleware mw = new MemoryMaintenanceMiddleware(wsm, null); + + // getUserId() is called (via timerKeyFor) before maybeRunMaintenance's own try/catch, + // so throwing here reaches onAgent's onErrorResume rather than being swallowed inside + // maybeRunMaintenance itself. + RuntimeContext brokenCtx = mock(RuntimeContext.class); + when(brokenCtx.getUserId()).thenThrow(new RuntimeException("boom")); + + AtomicBoolean errorDropped = new AtomicBoolean(false); + Hooks.onErrorDropped(t -> errorDropped.set(true)); + try { + AgentInput input = new AgentInput(List.of(userMsg("hi"))); + + List events = + mw.onAgent((Agent) null, brokenCtx, input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + + assertTrue(events != null && events.isEmpty()); + + // Give the detached boundedElastic subscription time to run and (not) drop the error. + // Without onErrorResume, subscribing with no error consumer routes the error to + // Hooks.onErrorDropped instead of handling it. + long deadline = System.nanoTime() + Duration.ofSeconds(2).toNanos(); + while (!errorDropped.get() && System.nanoTime() < deadline) { + Thread.sleep(20); + } + assertFalse( + errorDropped.get(), + "onErrorResume should have handled the maintenance error, not dropped it"); + } finally { + Hooks.resetOnErrorDropped(); + } + } + + private static void await(CountDownLatch latch) { + try { + latch.await(3, 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(); + } +} From 0ec27e02cee31002a9565b3b03bc8b548e9a59e9 Mon Sep 17 00:00:00 2001 From: Asad Date: Fri, 17 Jul 2026 18:41:44 +0530 Subject: [PATCH 2/3] fix(harness): drain in-flight memory maintenance on HarnessAgent.close() Detaching MemoryMaintenanceMiddleware's fire-and-forget maintenance run (#2225's fix) left it fully untracked: HarnessAgent.close() returned immediately even while a background maintenance write to the workspace was still in flight. In tests using JUnit @TempDir this raced against directory cleanup and intermittently failed with DirectoryNotEmptyException; in production it could tear down a workspace/sandbox out from under an in-progress write. Track each scheduled maintenance run's Disposable and add MemoryMaintenanceMiddleware.close(), called first from HarnessAgent.close(), which stops scheduling new runs and waits (bounded, 5s) for outstanding ones to finish before disposing stragglers. --- .../harness/agent/HarnessAgent.java | 26 ++++-- .../MemoryMaintenanceMiddleware.java | 88 ++++++++++++++++--- 2 files changed, 94 insertions(+), 20 deletions(-) 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 d669184cb9..8349ffe668 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 @@ -168,6 +168,7 @@ public class HarnessAgent implements Agent, AutoCloseable { private final SkillCurator skillCurator; private final SkillAuditLog skillAuditLog; private final MemoryConfig memoryConfig; + private final MemoryMaintenanceMiddleware memoryMaintenanceMw; /** The subagent middleware (either SubagentsMiddleware or DynamicSubagentsMiddleware). */ private final Object subagentMiddleware; @@ -199,6 +200,7 @@ private HarnessAgent( SkillCurator skillCurator, SkillAuditLog skillAuditLog, MemoryConfig memoryConfig, + MemoryMaintenanceMiddleware memoryMaintenanceMw, Object subagentMiddleware, DistributedStore distributedStore, WorkspacePathNormalizer pathNormalizer) { @@ -217,6 +219,7 @@ private HarnessAgent( this.skillCurator = skillCurator; this.skillAuditLog = skillAuditLog; this.memoryConfig = memoryConfig != null ? memoryConfig : MemoryConfig.defaults(); + this.memoryMaintenanceMw = memoryMaintenanceMw; this.subagentMiddleware = subagentMiddleware; this.distributedStore = distributedStore; this.pathNormalizer = pathNormalizer; @@ -375,14 +378,20 @@ public io.agentscope.core.permission.PermissionMode getPermissionMode( @Override public void close() { try { - shutdownTaskRepository(); + if (memoryMaintenanceMw != null) { + memoryMaintenanceMw.close(); + } } finally { try { - if (ownedWorkspaceIndex != null) { - ownedWorkspaceIndex.close(); - } + shutdownTaskRepository(); } finally { - delegate.close(); + try { + if (ownedWorkspaceIndex != null) { + ownedWorkspaceIndex.close(); + } + } finally { + delegate.close(); + } } } } @@ -2130,6 +2139,7 @@ public HarnessAgent build() { inner.middleware(new AtPathExpansionMiddleware(wsManager)); } Model memoryModel = memoryConfig.model() != null ? memoryConfig.model() : model; + MemoryMaintenanceMiddleware memoryMaintenanceMw = null; if (memoryModel != null && !disableMemoryHooks) { IsolationScope effectiveIsolationScope = fsIsolationScope; @@ -2155,14 +2165,15 @@ public HarnessAgent build() { memoryModel, effectiveConsolidationPrompt, memoryConfig.consolidationMaxTokens()); - inner.middleware( + memoryMaintenanceMw = new MemoryMaintenanceMiddleware( wsManager, consolidator, memoryConfig.dailyFileRetentionDays(), memoryConfig.sessionRetentionDays(), memoryConfig.consolidationMinGap(), - effectiveIsolationScope)); + effectiveIsolationScope); + inner.middleware(memoryMaintenanceMw); } CompactionMiddleware compactionHook = null; if (!disableCompaction && compactionConfig != null) { @@ -2491,6 +2502,7 @@ public HarnessAgent build() { pendingSkillCurator, pendingSkillAuditLog, memoryConfig, + memoryMaintenanceMw, capturedSubagentMw, distributedStore, pathNormalizer); 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 3a37260fb5..3c164265c0 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 @@ -29,11 +29,13 @@ import java.time.Duration; import java.time.Instant; import java.time.LocalDate; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; 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; @@ -80,6 +82,19 @@ public class MemoryMaintenanceMiddleware implements HarnessRuntimeMiddleware { private final Duration minGap; private final IsolationScope isolationScope; + /** 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; + /** * Process-wide per-isolation-key maintenance timestamps. Static so that the throttle window * survives across {@code HarnessAgent.Builder.build()} calls — each rebuild creates a new @@ -142,19 +157,66 @@ public Flux onAgent( AgentInput input, Function> next) { final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty(); - return next.apply(input) - .doOnComplete( - () -> - Mono.fromRunnable(() -> maybeRunMaintenance(rc)) - .subscribeOn(Schedulers.boundedElastic()) - .onErrorResume( - e -> { - log.warn( - "Memory maintenance failed: {}", - e.getMessage()); - return Mono.empty(); - }) - .subscribe()); + return next.apply(input).doOnComplete(() -> scheduleMaintenance(rc)); + } + + /** + * 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 (and disposes immediately) once {@link #close()} has been called, so a call that + * races with shutdown doesn't spawn untracked work. + */ + private void scheduleMaintenance(RuntimeContext rc) { + if (closed) { + return; + } + Disposable[] holder = new Disposable[1]; + Disposable d = + Mono.fromRunnable(() -> maybeRunMaintenance(rc)) + .subscribeOn(Schedulers.boundedElastic()) + .onErrorResume( + e -> { + log.warn("Memory maintenance failed: {}", e.getMessage()); + return Mono.empty(); + }) + .doFinally( + sig -> { + if (holder[0] != null) { + pending.remove(holder[0]); + } + }) + .subscribe(); + holder[0] = d; + pending.add(d); + if (closed) { + pending.remove(d); + d.dispose(); + } + } + + /** + * 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 — see {@link MemoryMaintenanceMiddleware class docs} for why maintenance is + * detached in the first place. + */ + public void close() { + 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 void maybeRunMaintenance(RuntimeContext rc) { From 27a92274de57f39a796e1e185fb69f84edf21bac Mon Sep 17 00:00:00 2001 From: Asad Date: Fri, 17 Jul 2026 18:52:53 +0530 Subject: [PATCH 3/3] test(harness): cover close()'s wait/dispose paths, drop untestable race guard Adds coverage for the branches introduced by the previous commit that jacoco flagged as untested: scheduleMaintenance() no-op'ing after close(), and close()'s interrupted-wait / dispose-stragglers paths. Also drops the extra "dispose immediately if closed flipped true right after pending.add(d)" guard in scheduleMaintenance(). It only mattered for a TOCTOU race too narrow to hit deterministically in a test, and close()'s own wait-then-dispose-all loop already reclaims that disposable in every realistic timing. --- .../MemoryMaintenanceMiddleware.java | 8 +-- ...aintenanceMiddlewareAsyncBehaviorTest.java | 71 +++++++++++++++++++ 2 files changed, 73 insertions(+), 6 deletions(-) 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 3c164265c0..e249ae23f1 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 @@ -163,8 +163,8 @@ public Flux onAgent( /** * 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 (and disposes immediately) once {@link #close()} has been called, so a call that - * races with shutdown doesn't spawn untracked work. + * it. No-ops once {@link #close()} has been called, so a call that races with shutdown + * doesn't spawn new untracked work. */ private void scheduleMaintenance(RuntimeContext rc) { if (closed) { @@ -188,10 +188,6 @@ private void scheduleMaintenance(RuntimeContext rc) { .subscribe(); holder[0] = d; pending.add(d); - if (closed) { - pending.remove(d); - d.dispose(); - } } /** 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 index b55d88fab4..7029d6ef10 100644 --- 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 @@ -19,6 +19,7 @@ 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; @@ -147,6 +148,76 @@ void onAgent_onErrorResume_handlesExceptionFromDetachedMaintenance(@TempDir Path } } + @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()); + + // Give any (incorrectly) scheduled maintenance a chance to run before asserting it + // didn't. + Thread.sleep(200); + verifyNoInteractions(consolidator); + } + + @Test + void close_interruptedWhileWaiting_returnsPromptlyAndDisposesOutstandingRun(@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"))); + mw.onAgent((Agent) null, RuntimeContext.empty(), input, in -> Flux.empty()) + .collectList() + .block(Duration.ofSeconds(5)); + assertTrue( + consolidationStarted.await(2, TimeUnit.SECONDS), + "consolidation should have started on a detached background thread"); + + AtomicBoolean closeReturned = new AtomicBoolean(false); + Thread closer = + new Thread( + () -> { + mw.close(); + closeReturned.set(true); + }); + closer.start(); + Thread.sleep(50); // let close() enter its bounded poll-wait loop + closer.interrupt(); + closer.join(2000); + + assertFalse(closer.isAlive(), "close() should return promptly once interrupted"); + assertTrue(closeReturned.get(), "close() should have completed on the closer thread"); + + releaseConsolidation.countDown(); + } + private static void await(CountDownLatch latch) { try { latch.await(3, TimeUnit.SECONDS);