fix(harness): drain session mirrors on close

agentscope-harness/src/main/java/io/agentscope/harness/agent/HarnessAgent.java

@@ -59,6 +59,7 @@
 import io.agentscope.harness.agent.memory.compaction.CompactionConfig;
 import io.agentscope.harness.agent.memory.compaction.ConversationCompactor;
 import io.agentscope.harness.agent.memory.compaction.ToolResultEvictionConfig;
+import io.agentscope.harness.agent.memory.session.SessionTree;
 import io.agentscope.harness.agent.middleware.AgentTraceMiddleware;
 import io.agentscope.harness.agent.middleware.AsyncToolMiddleware;
 import io.agentscope.harness.agent.middleware.AtPathExpansionMiddleware;
@@ -378,11 +379,15 @@ public void close() {
             shutdownTaskRepository();
         } finally {
             try {
-                if (ownedWorkspaceIndex != null) {
-                    ownedWorkspaceIndex.close();
-                }
+                SessionTree.awaitPendingMirrors();
             } finally {
-                delegate.close();
+                try {
+                    if (ownedWorkspaceIndex != null) {
+                        ownedWorkspaceIndex.close();
+                    }
+                } finally {
+                    delegate.close();
+                }
             }
         }
     }

agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/session/SessionTree.java

@@ -34,6 +34,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
@@ -88,6 +89,33 @@ public class SessionTree {
                         return t;
                     });
 
+    /**
+     * Waits until all remote mirror work submitted before this call has completed.
+     *
+     * <p>Harness lifecycle owners should call this before closing filesystem or index resources.
+     * The wait is deliberately uninterruptible so that resource shutdown cannot race an in-flight
+     * upload; the interrupted status is restored before returning.
+     */
+    public static void awaitPendingMirrors() {
+        boolean interrupted = false;
+        try {
+            while (true) {
+                try {
+                    MIRROR_EXECUTOR.submit(() -> {}).get();
+                    return;
+                } catch (InterruptedException e) {
+                    interrupted = true;
+                } catch (ExecutionException e) {
+                    throw new IllegalStateException("Failed to await session tree mirrors", e);
+                }
+            }
+        } finally {
+            if (interrupted) {
+                Thread.currentThread().interrupt();
+            }
+        }
+    }
+
     private final Path contextFile;
     private final Path logFile;
     private final Path workspaceRoot;

agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/session/SessionTreeMirrorTest.java

@@ -16,7 +16,11 @@
 package io.agentscope.harness.agent.memory.session;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+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.doAnswer;
+import static org.mockito.Mockito.mock;
 
 import io.agentscope.core.agent.RuntimeContext;
 import io.agentscope.harness.agent.filesystem.AbstractFilesystem;
@@ -26,6 +30,10 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -200,11 +208,43 @@ void flush_localWriteCompletesImmediately() throws Exception {
                 "local log file must exist immediately after flush()");
     }
 
+    @Test
+    void awaitPendingMirrors_blocksUntilInFlightUploadCompletes() throws Exception {
+        AbstractFilesystem fs = mock(AbstractFilesystem.class);
+        CountDownLatch uploadStarted = new CountDownLatch(1);
+        CountDownLatch releaseUpload = new CountDownLatch(1);
+        doAnswer(
+                        invocation -> {
+                            uploadStarted.countDown();
+                            releaseUpload.await();
+                            return null;
+                        })
+                .when(fs)
+                .uploadFiles(any(), any());
+        Path context = workspace.resolve("agents/agent-a/sessions/blocking.jsonl");
+        SessionTree tree = new SessionTree(context, workspace, fs);
+        tree.append(new SessionEntry.MessageEntry(null, null, null, "USER", "hi", null));
+        tree.flush();
+        assertTrue(uploadStarted.await(5, TimeUnit.SECONDS));
+
+        ExecutorService closer = Executors.newSingleThreadExecutor();
+        try {
+            Future<?> barrier = closer.submit(SessionTree::awaitPendingMirrors);
+            assertFalse(barrier.isDone(), "barrier must wait for the in-flight upload");
+
+            releaseUpload.countDown();
+            barrier.get(5, TimeUnit.SECONDS);
+        } finally {
+            releaseUpload.countDown();
+            closer.shutdownNow();
+        }
+    }
+
     // -----------------------------------------------------------------------
     //  Helper
     // -----------------------------------------------------------------------
 
-    private static void awaitMirror() throws InterruptedException {
-        TimeUnit.MILLISECONDS.sleep(300);
+    private static void awaitMirror() {
+        SessionTree.awaitPendingMirrors();
     }
 }