Skip to content

fix(harness): avoid blocking stream completion on transcript persistence - #3106

Draft
ningmao-hlyz wants to merge 1 commit into
agentscope-ai:mainfrom
ningmao-hlyz:fix/transcript-async-completion
Draft

fix(harness): avoid blocking stream completion on transcript persistence#3106
ningmao-hlyz wants to merge 1 commit into
agentscope-ai:mainfrom
ningmao-hlyz:fix/transcript-async-completion

Conversation

@ningmao-hlyz

Copy link
Copy Markdown
Contributor

Summary

Fixes #3059.

  • Keep transcript persistence fire-and-forget for regular agent calls.
  • Track pending transcript writes with MemoryBackgroundTasks so shutdown can wait for them.
  • Keep sandbox-backed calls in the stream until transcript persistence finishes.
  • Serialize concurrent transcript appends for the same session.
  • Add regression coverage for stream completion, shutdown quiescence, and sandbox release ordering.

Problem

TranscriptMiddleware used concatWith(...) for transcript persistence. Because the persistence task was part of the sequential stream, downstream subscribers did not receive onComplete until transcript writing finished.

This added user-visible tail latency for streaming callers such as SSE and WebSocket deployments, even after the final assistant event had already been emitted.

Changes

  • Regular calls now schedule transcript persistence from doOnComplete(...) without delaying downstream completion.
  • Pending asynchronous writes are counted by MemoryBackgroundTasks, preserving lifecycle shutdown/quiescence guarantees.
  • The conversation context is copied before the asynchronous task is scheduled.
  • Same-session transcript read/merge/write cycles are serialized across writer instances.
  • Sandbox-backed calls intentionally keep persistence inside the stream so the sandbox is not released before the writer finishes.

MemoryFlushMiddleware and MemoryMaintenanceMiddleware already use tracked fire-and-forget execution and are unchanged by this PR.

Testing

  • mvn -pl agentscope-harness -Dtest=SessionTranscriptWriterTest,TranscriptMiddlewareAsyncTest test
  • mvn -pl agentscope-harness spotless:check

@codecov

codecov Bot commented Sep 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.00000% with 4 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
.../agent/memory/session/SessionTranscriptWriter.java 92.59% 2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

@oss-maintainer oss-maintainer left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

Fixes #3059 by scheduling transcript persistence from doOnComplete instead of concatWith, keeping sandbox-backed calls in-stream, and serialising same-session read/merge/write behind a ref-counted per-path lock. The direction is right, the fire-and-forget task is correctly tracked by MemoryBackgroundTasks, and the regression tests cover completion ordering, sandbox release ordering and the pre-detach snapshot.

Verdict

No blocking issue found — inline comments are latency/robustness polish plus one doc-accuracy note.


Automated review by github-manager-bot

@ningmao-hlyz
ningmao-hlyz force-pushed the fix/transcript-async-completion branch from 01c83aa to 105959a Compare September 11, 2026 17:16
Preserve fire-and-forget transcript writes for regular calls while delaying sandbox release until the transcript append completes. Snapshot messages before asynchronous persistence and add regression coverage for the sandbox lifecycle.
@ningmao-hlyz
ningmao-hlyz force-pushed the fix/transcript-async-completion branch from 105959a to a66780c Compare September 11, 2026 17:58

@oss-maintainer oss-maintainer left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

Moves non-sandbox transcript persistence off the critical path: TranscriptMiddleware now dispatches the append in doOnComplete as a tracked fire-and-forget MemoryBackgroundTasks task (sandbox calls keep it in-stream so the per-call sandbox stays alive), and SessionTranscriptWriter guards the load → syncFromRemote → flush cycle with a reference-counted per-path lock so concurrent appends for one session can't interleave. The direction is right and the test coverage is genuinely good (onComplete fires before the write finishes, sandbox writes stay in-stream, message snapshot taken before detaching, failure paths don't break the stream). One ordering bug in the new lock protocol and a scope question on the guard are worth fixing before this lands.

Findings

  • [Warning] SessionTranscriptWriter.java:108 — the refcount is dropped before unlock(), so a new caller can publish a second WriteLock for the same path while an append is still in flight; the eviction check runs under the map bin lock, not the write lock.
  • [Warning] SessionTranscriptWriter.java:66 — the lock is process-local while the guarded cycle reaches the remote TranscriptStore; two JVMs on a shared workspace are still unprotected. Please confirm the intended scope and document it.
  • [Warning] TranscriptMiddleware.java:88 — the sandbox branch and the background branch now have different error-discipline; the previous onErrorResume was dropped from appendTranscript.
  • [Info] TranscriptMiddleware.java:84 / :91 — state resolution runs inline on the completing thread; failure logging lost its agent=/session= context.
  • [Info] TranscriptMiddlewareAsyncTest.java:242 — no test pins down the cancel semantics of the doOnComplete-only dispatch.

Suggestions

The lock fix is small — take the eviction decision while still holding the monitor (see the inline comment for a snippet), so users == 0 → remove and users++ → reuse can never overlap with an in-flight append. For the error-discipline item, hoisting the existing onErrorResume lambda into a private static Mono<Void> logFailure(Throwable) and applying it in both branches would keep the two paths from drifting again.

Neither item blocks the design: the change is a clear improvement over blocking stream completion on a disk/remote write, and the added tests exercise exactly the races that matter.


Automated review by github-manager-bot

WRITE_LOCKS.compute(
lockKey,
(k, existing) -> {
WriteLock next = existing != null ? existing : new WriteLock();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Warning] The refcount is released before lock.unlock(), so mutual exclusion can be lost. computeIfPresent here only decrements users (it returns a non-null value when the count is still > 0), but the monitor for that key is the map bin lock, not writeLock.lock. If a new caller runs its compute increment while we are still between this line and the unlock() below, the counter is pinned at >= 1 and the eviction branch can never be taken, so a third caller legitimately creates a different WriteLock and enters the read/merge/write cycle concurrently with us — exactly the interleaving this PR is meant to prevent, just much rarer.

Suggested fix: release the reference after unlocking, and let the unlock itself happen inside the map mutation so no new user can be attached to an about-to-be-discarded lock:

try {
    appendMessagesLocked(rc, messages, agentId, sessionId, contextFile);
} finally {
    writeLock.lock.lock(); // never contended, just re-enters the monitor we own
    try {
        WRITE_LOCKS.compute(lockKey, (k, current) -> {
            if (current != writeLock) {
                return current; // defensive: another lock replaced us
            }
            return --current.users == 0 ? null : current;
        });
    } finally {
        writeLock.lock.unlock();
    }
}

With this ordering the eviction and the increment are serialized by the same lock we hold, so a fresh WriteLock can never be published while an append is in flight.


/**
* Active per-session locks for the transcript read/merge/write cycle; idle entries are evicted.
*/

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Warning] WRITE_LOCKS is process-local, but the guarded cycle is tree.load()tree.syncFromRemote()tree.flush() (see appendMessagesLocked). Two JVMs sharing the same workspace/TranscriptStore (agentscope-distribution is explicitly the multi-instance path) will both pass through this lock and can still interleave a remote read-modify-write, losing entries. The Javadoc wording ("serialized across writer instances") could be read as covering that case — could we either state explicitly that this only serializes writers inside one JVM, or push the guard down to the session-index/TranscriptStore layer where a cross-process lease/version check is possible? A follow-up issue is fine if the intended scope is single-process.

e -> {
log.warn("Transcript append failed: {}", e.getMessage());
return Mono.empty();
})

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Warning] Inconsistency with the old behavior: the sandbox path keeps .onErrorResume(e -> { log.warn(...); return Mono.empty(); }), while the new background path swallows the error only via the subscribe(null, consumer) overload at line 105, and the Mono.fromRunnable in appendTranscript(agent, rc, messages) no longer has the previous onErrorResume/doOnSuccess wrappers. That works today because SessionTranscriptWriter.appendMessages has a catch-all catch (Exception e) — but Mono.fromRunnable also turns a cancellation/interrupt of the boundedElastic worker, or any Error (e.g. NoClassDefFoundError on a lazily loaded class), into an onError signal that will now reach the subscriber's error consumer, and a throwing log.warn handler there would surface as ErrorCallbackNotImplemented. Since the two branches of this method now have different error-discipline, could we apply .onErrorResume(e -> { log.warn("Transcript append failed: {}", e.getMessage()); return Mono.empty(); }) to the background append as well, so both paths are identical and the failure policy is explicit at the call site?

}

private void scheduleTranscript(Agent agent, RuntimeContext rc) {
AgentState state = RuntimeContext.resolveAgentState(rc, agent);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Info] scheduleTranscript is invoked from doOnComplete, so it runs on whatever thread emitted the stream terminal signal, and RuntimeContext.resolveAgentState(rc, agent) plus the defensive state.getContext() copy happen inline there. For the goal of this PR (not delaying completion) it is worth checking whether resolveAgentState can touch the state store: if it can, that work should move onto boundedElastic together with the append, e.g. by doing the capture inside the Mono.fromRunnable(...) that is then subscribeOn'd, keeping only the already-obtained state reference on the stream thread.

MemoryBackgroundTasks.begin();
append.subscribeOn(Schedulers.boundedElastic())
.doFinally(signal -> MemoryBackgroundTasks.end())
.subscribe(null, e -> log.warn("Transcript append failed: {}", e.getMessage()));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Info] MemoryBackgroundTasks.begin() is called synchronously before subscribeOn, and end() is paired via doFinally, which is the right shape (no window where the task is uncounted). One residual gap: end() runs on doFinally, and awaitQuiescence is only consulted in HarnessAgent.close() per the class docs of MemoryBackgroundTasks. If a session is evicted / a non-sandbox agent call completes and the harness drops the agent before the write lands, the transcript for that turn is still lost silently (background failures are only logged at WARN). It would help operators if the failure log carried agent=/session= context here so a dropped transcript is traceable — the previous concatWith path had the same weakness, so this is just an observability suggestion, not a blocker.

}

@Test
void sandboxTranscriptFailureDoesNotFailAgentStream() {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Info] sandboxTranscriptWriteStaysInStreamUntilPersistenceFinishes asserts the sandbox path via Flux.using + a blocking collectList(), which is good, but there is no equivalent test that a cancelled non-sandbox stream still persists what it had (the PR title's "completion" signal is doOnComplete only). If cancellation without persistence is the intended contract, a small test asserting doOnCancel does not schedule a write would lock that decision down; if it is not intended, doOnCancel/doFinally would need the same dispatch.

@oss-maintainer

Copy link
Copy Markdown
Collaborator

Follow-up on CI for the current head (105959a6) — the Linux build is failing, and the failure shape looks related to this change rather than flaky:

[ERROR] AgentSpawnToolPlanModeTest.agentSend_reappliesParentPlanModeBeforeInvokingExistingChild » JUnit Failed to close extension context
Caused by: java.io.IOException: Failed to delete temp directory /tmp/junit-…
  The following paths could not be deleted: <root>, user, user/agents, user/agents/worker/sessions
[ERROR] Tests run: 1048, Failures: 0, Errors: 1, Skipped: 9

https://github.com/agentscope-ai/agentscope-java/actions/runs/34626731278

The leftover path is a session transcript directory under a @TempDir workspace, and the error is DirectoryNotEmptyException during JUnit's extension teardown — i.e. something wrote into the temp workspace after the test body finished. With this PR the non-sandbox transcript append moves to doOnComplete + a detached boundedElastic task, so a @TempDir-based test whose agent stream completes can now have an in-flight write outliving the directory. MemoryBackgroundTasks.awaitQuiescence(...) is already the right tool: awaiting it in the affected test's teardown (or in a shared @AfterEach for harness tests that use @TempDir workspaces) should make the ordering explicit instead of timing-dependent.

Not blocking on this, and if you can show the same test failing on an unrelated head please treat this as a red herring — but it is worth ruling the ordering out before the build is re-run a few times and assumed flaky.


Automated review follow-up by github-manager-bot

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants