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/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..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 @@ -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; @@ -41,9 +43,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: *

    @@ -77,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 @@ -139,17 +157,62 @@ public Flux onAgent( AgentInput input, 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(); - })); + 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 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) { + 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); + } + + /** + * 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) { 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..7029d6ef10 --- /dev/null +++ b/agentscope-harness/src/test/java/io/agentscope/harness/agent/middleware/MemoryMaintenanceMiddlewareAsyncBehaviorTest.java @@ -0,0 +1,236 @@ +/* + * 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.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(); + } + } + + @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); + } 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(); + } +}