fix(harness): detach memory flush and maintenance from agent call response

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<AutoCloseable> 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<AutoCloseable> closeableMiddlewares,
             Object subagentMiddleware,
             DistributedStore distributedStore,
             WorkspacePathNormalizer pathNormalizer) {
@@ -229,6 +233,8 @@ private HarnessAgent(
         this.skillCurator = skillCurator;
         this.skillAuditLog = skillAuditLog;
         this.memoryConfig = memoryConfig != null ? memoryConfig : MemoryConfig.defaults();
+        this.closeableMiddlewares =
+                closeableMiddlewares != null ? List.copyOf(closeableMiddlewares) : List.of();
         this.subagentMiddleware = subagentMiddleware;
         this.distributedStore = distributedStore;
         this.pathNormalizer = pathNormalizer;
@@ -431,11 +437,23 @@ public void close() {
             shutdownTaskRepository();
         } finally {
             try {
-                if (ownedWorkspaceIndex != null) {
-                    ownedWorkspaceIndex.close();
+                // Drain detached memory flush/maintenance so async writes do not race with
+                // workspace teardown (e.g., temp workspace deletion in tests).
+                for (AutoCloseable mw : closeableMiddlewares) {
+                    try {
+                        mw.close();
+                    } catch (Exception e) {
+                        log.warn("Failed to close middleware: {}", e.getMessage());
+                    }
                 }
             } finally {
-                delegate.close();
+                try {
+                    if (ownedWorkspaceIndex != null) {
+                        ownedWorkspaceIndex.close();
+                    }
+                } finally {
+                    delegate.close();
+                }
             }
         }
     }
@@ -2386,21 +2404,24 @@ public HarnessAgent build() {
                                 wsManager, effectiveTranscriptStore, transcriptTenant));
             }
             Model memoryModel = memoryConfig.model() != null ? memoryConfig.model() : model;
+            List<AutoCloseable> pendingCloseableMiddlewares = new java.util.ArrayList<>();
             if (memoryModel != null && !disableMemoryHooks) {
                 IsolationScope effectiveIsolationScope = fsIsolationScope;
 
                 String effectiveFlushPrompt =
                         memoryConfig.flushPrompt() != null
                                 ? memoryConfig.flushPrompt()
                                 : MemoryFlushManager.DEFAULT_FLUSH_PROMPT;
-                inner.middleware(
+                MemoryFlushMiddleware memoryFlushMw =
                         new MemoryFlushMiddleware(
                                 wsManager,
                                 memoryModel,
                                 effectiveFlushPrompt,
                                 memoryConfig.flushTrigger(),
                                 effectiveIsolationScope,
-                                periodicGate));
+                                periodicGate);
+                inner.middleware(memoryFlushMw);
+                pendingCloseableMiddlewares.add(memoryFlushMw);
 
                 String effectiveConsolidationPrompt =
                         memoryConfig.consolidationPrompt() != null
@@ -2413,15 +2434,17 @@ public HarnessAgent build() {
                                 effectiveConsolidationPrompt,
                                 memoryConfig.consolidationMaxTokens(),
                                 distributedStore != null ? distributedStore.baseStore() : null);
-                inner.middleware(
+                MemoryMaintenanceMiddleware memoryMaintenanceMw =
                         new MemoryMaintenanceMiddleware(
                                 wsManager,
                                 consolidator,
                                 memoryConfig.dailyFileRetentionDays(),
                                 memoryConfig.sessionRetentionDays(),
                                 memoryConfig.consolidationMinGap(),
                                 effectiveIsolationScope,
-                                periodicGate));
+                                periodicGate);
+                inner.middleware(memoryMaintenanceMw);
+                pendingCloseableMiddlewares.add(memoryMaintenanceMw);
             }
             CompactionMiddleware compactionHook = null;
             if (!disableCompaction && compactionConfig != null) {
@@ -2785,6 +2808,7 @@ public HarnessAgent build() {
                     pendingSkillCurator,
                     pendingSkillAuditLog,
                     memoryConfig,
+                    pendingCloseableMiddlewares,
                     capturedSubagentMw,
                     distributedStore,
                     pathNormalizer);

agentscope-harness/src/main/java/io/agentscope/harness/agent/middleware/MemoryFlushMiddleware.java

@@ -28,21 +28,29 @@
 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;
 
 /**
  * Middleware that triggers memory flush and message offload at the end of each agent call.
  *
- * <p>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.
+ * <p>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.
  *
  * <p>Flush is gated by a {@link MemoryConfig.FlushTrigger}:
  * <ul>
@@ -66,7 +74,7 @@
  *       the whole agent instance (prevents concurrent flush races on shared memory files).</li>
  * </ul>
  */
-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<Disposable> pending = ConcurrentHashMap.newKeySet();
+
+    private volatile boolean closed = false;
+
     public MemoryFlushMiddleware(WorkspaceManager workspaceManager, Model model) {
         this(
                 workspaceManager,
@@ -140,28 +163,83 @@ public Flux<AgentEvent> onAgent(
             AgentInput input,
             Function<AgentInput, Flux<AgentEvent>> next) {
         final RuntimeContext rc = ctx != null ? ctx : RuntimeContext.empty();
-        return next.apply(input)
-                .concatWith(
-                        Mono.defer(() -> doFlush(agent, rc))
-                                .subscribeOn(Schedulers.boundedElastic())
-                                .onErrorResume(
-                                        e -> {
-                                            log.warn("Memory flush failed: {}", e.getMessage());
-                                            return Mono.empty();
-                                        })
-                                .then(Mono.<AgentEvent>empty()));
+        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.
+     *
+     * <p>The snapshot capture (including a shallow copy of the RuntimeContext and a copy of the
+     * messages) is wrapped in try/catch so that any exception (e.g. from
+     * {@code List.copyOf} or {@code resolveAgentState}) does not escape into the {@code
+     * doOnComplete} callback and turn an already-completed Flux into an error.
+     *
+     * <p>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) {
+        FlushRequest request;
+        try {
+            request = captureFlushRequest(agent, rc);
+        } catch (Exception e) {
+            log.warn("Failed to capture flush request: {}", e.getMessage());
+            return;
+        }
+        if (request == null) {
+            return;
+        }
+        synchronized (pending) {
+            if (closed) {
+                return;
+            }
+        }
+        Disposable[] holder = new Disposable[1];
+        Disposable d =
+                Mono.defer(() -> doFlush(request))
+                        .subscribeOn(Schedulers.boundedElastic())
+                        .timeout(FLUSH_TIMEOUT)
+                        .onErrorResume(
+                                e -> {
+                                    log.warn("Memory flush failed: {}", e.getMessage());
+                                    return Mono.empty();
+                                })
+                        .doFinally(
+                                sig -> {
+                                    if (holder[0] != null) {
+                                        pending.remove(holder[0]);
+                                    }
+                                })
+                        .subscribe();
+        holder[0] = d;
+        synchronized (pending) {
+            if (closed) {
+                d.dispose();
+                return;
+            }
+            pending.add(d);
+        }
     }
 
-    private Mono<Void> 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<Msg> messages = state.getContext();
         if (messages.isEmpty()) {
-            return Mono.empty();
+            return null;
         }
+        return new FlushRequest(RuntimeContext.builder().from(rc).build(), List.copyOf(messages));
+    }
 
+    private Mono<Void> doFlush(FlushRequest request) {
+        RuntimeContext rc = request.runtimeContext();
+        List<Msg> messages = request.messages();
         MemoryFlushManager flushManager =
                 new MemoryFlushManager(workspaceManager, model, flushPrompt);
 
@@ -181,11 +259,49 @@ private Mono<Void> doFlush(Agent agent, RuntimeContext rc) {
             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.
+     *
+     * <p>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<Msg> 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}