From 34253c8ccba9b76c82018165a974be739db07bed Mon Sep 17 00:00:00 2001 From: u010143241 Date: Sun, 12 Jul 2026 09:13:39 +0800 Subject: [PATCH] fix(harness): drain session mirrors on close --- .../harness/agent/HarnessAgent.java | 13 ++++-- .../agent/memory/session/SessionTree.java | 28 ++++++++++++ .../memory/session/SessionTreeMirrorTest.java | 44 ++++++++++++++++++- 3 files changed, 79 insertions(+), 6 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..ed2cf04eef 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 @@ -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(); + } } } } diff --git a/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/session/SessionTree.java b/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/session/SessionTree.java index 0413860b39..13ecaa50fe 100644 --- a/agentscope-harness/src/main/java/io/agentscope/harness/agent/memory/session/SessionTree.java +++ b/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. + * + *

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; diff --git a/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/session/SessionTreeMirrorTest.java b/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/session/SessionTreeMirrorTest.java index c0f01fa097..f02d560438 100644 --- a/agentscope-harness/src/test/java/io/agentscope/harness/agent/memory/session/SessionTreeMirrorTest.java +++ b/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(); } }