fix(harness): avoid blocking stream completion on transcript persistence - #3106
fix(harness): avoid blocking stream completion on transcript persistence#3106ningmao-hlyz wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
oss-maintainer
left a comment
There was a problem hiding this comment.
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
01c83aa to
105959a
Compare
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.
105959a to
a66780c
Compare
oss-maintainer
left a comment
There was a problem hiding this comment.
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 beforeunlock(), so a new caller can publish a secondWriteLockfor 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 remoteTranscriptStore; 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 previousonErrorResumewas dropped fromappendTranscript. - [Info]
TranscriptMiddleware.java:84/:91— state resolution runs inline on the completing thread; failure logging lost itsagent=/session=context. - [Info]
TranscriptMiddlewareAsyncTest.java:242— no test pins down the cancel semantics of thedoOnComplete-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(); |
There was a problem hiding this comment.
[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. | ||
| */ |
There was a problem hiding this comment.
[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(); | ||
| }) |
There was a problem hiding this comment.
[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); |
There was a problem hiding this comment.
[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())); |
There was a problem hiding this comment.
[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() { |
There was a problem hiding this comment.
[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.
|
Follow-up on CI for the current head ( https://github.com/agentscope-ai/agentscope-java/actions/runs/34626731278 The leftover path is a session transcript directory under a 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 |
Summary
Fixes #3059.
MemoryBackgroundTasksso shutdown can wait for them.Problem
TranscriptMiddlewareusedconcatWith(...)for transcript persistence. Because the persistence task was part of the sequential stream, downstream subscribers did not receiveonCompleteuntil 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
doOnComplete(...)without delaying downstream completion.MemoryBackgroundTasks, preserving lifecycle shutdown/quiescence guarantees.MemoryFlushMiddlewareandMemoryMaintenanceMiddlewarealready use tracked fire-and-forget execution and are unchanged by this PR.Testing
mvn -pl agentscope-harness -Dtest=SessionTranscriptWriterTest,TranscriptMiddlewareAsyncTest testmvn -pl agentscope-harness spotless:check