fix(model): streaming retry re-delivers already-sent chunks downstream - #2479
fix(model): streaming retry re-delivers already-sent chunks downstream#2479gxgeek-n wants to merge 3 commits into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
oss-maintainer
left a comment
oss-maintainer
left a comment
There was a problem hiding this comment.
Review — PR #2479
Streaming Retry Fix (commit 2) — ✅ Good
The core fix is clean and well-targeted:
AtomicBoolean emittedcorrectly tracks whether any chunk has been delivered.filter(error -> !emitted.get() && retryableError.test(error))narrows retry to pre-first-chunk failures only.doOnNext(r -> emitted.set(true))placed upstream ofretryWhen— correct ordering- 14 lines in production code — minimal surface area
- 4 test cases in
ModelUtilsStreamingRetryTestcover the bug scenario + regression cases (pre-first-chunk retry still works, non-retryable errors pass through, clean stream unaffected)
The fix preserves the valuable retry cases (connection setup, HTTP 429 before first token) while preventing the duplication bug. Well done.
Concerns
1. PR mixes two independent fixes
This PR contains two commits from different authors:
- Commit 1 (waterWang):
fix(#2475): store userId/sessionId directly on CallExecution— already submitted as standalone PR #2476 - Commit 2 (gxgeek-n):
fix(model): stop retrying a streaming response once chunks were delivered— the actual #2478 fix
Mixing them creates merge/review complexity. If #2476 is merged first, this PR will have a conflict on ReActAgent.java. Consider rebasing this PR onto main after #2476 lands, or splitting into a standalone PR for just the streaming fix.
2. CI failures — Spotless formatting
Both Linux and Windows builds fail with Spotless violations in ReActAgent.java:
- Extra blank lines before
initialStateLoadJavadoc (line 352-353) CallExecutionconstructor parameters need one-per-line formatting
Fix with:
mvn spotless:apply -pl agentscope-coreSummary
The streaming retry fix itself is solid. Address the Spotless formatting and consider the PR composition concern.
There was a problem hiding this comment.
Pull request overview
This PR improves correctness in two areas of AgentScope’s streaming and session state handling: (1) it prevents retryWhen from re-subscribing a streaming model response after chunks have already been emitted (avoiding duplicate delivery downstream), and (2) it fixes AgentState persistence when sessionId contains / by avoiding round-tripping through a concatenated slot key.
Changes:
- Update
ModelUtils.applyTimeoutAndRetryto stop retrying once any streaming chunk has been emitted. - Add regression tests covering streaming retry semantics (mid-stream retryable failure, pre-first-chunk retry, non-retryable failure, clean stream).
- Fix
ReActAgentAgentState persistence for sessionIds containing/and add end-to-end regression tests for the slot integrity issue.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| agentscope-core/src/main/java/io/agentscope/core/model/ModelUtils.java | Stops retries after the first emitted streaming element to prevent duplicate downstream chunks. |
| agentscope-core/src/test/java/io/agentscope/core/model/ModelUtilsStreamingRetryTest.java | Adds tests pinning streaming retry behavior and preventing regression. |
| agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java | Persists AgentState using original (userId, sessionId) instead of parsing a concatenated slot key. |
| agentscope-core/src/test/java/io/agentscope/core/state/AgentStateSlotKeySessionIdSlashTest.java | Adds end-to-end regression coverage for sessionIds containing / to ensure state lands in the correct slot. |
| java.util.concurrent.atomic.AtomicBoolean emitted = | ||
| new java.util.concurrent.atomic.AtomicBoolean(false); | ||
| final Predicate<Throwable> retryableError = retryOn; | ||
|
|
||
| Retry retrySpec = |
| private Mono<Void> saveStateToSession(CallExecution scope) { | ||
| if (stateStore == null) { | ||
| return Mono.empty(); | ||
| } | ||
| syncToolkitToState(scope.state); | ||
| SlotRef ref = SlotRef.parse(scope.slotKey); | ||
| AgentState toSave = scope.state; | ||
| return Mono.<Void>fromRunnable( | ||
| () -> stateStore.save(ref.userId, ref.sessionId, "agent_state", toSave)) | ||
| () -> stateStore.save(scope.userId, scope.sessionId, "agent_state", toSave)) | ||
| .subscribeOn(Schedulers.boundedElastic()); |
…ered
ModelUtils.applyTimeoutAndRetry attaches retryWhen directly to the streaming
Flux. retryWhen re-subscribes its upstream, so when a retryable error arrives
mid-stream the model regenerates the whole response and downstream consumers
(middlewares, memory, UI) receive the already-delivered chunks a second time.
With MODEL_DEFAULTS (maxAttempts=3) the same content can be delivered three
times. It also inflates latency by a full regeneration, which for long or
thinking-mode responses is tens of seconds of apparent silence.
Reproduced without a live model — upstream emits two chunks then throws
IOException (retryable per ExecutionConfig.RETRYABLE_ERRORS):
before: upstream subscribed 3x, downstream got
[Hello, " world", Hello, " world", Hello, " world"]
after: upstream subscribed 1x, downstream got [Hello, " world"]
Fix: track whether anything has been emitted and fold that into the retry
filter, so retries stop after the first chunk reaches downstream. Failures
*before* the first chunk (connection setup, HTTP 429 ahead of the first token)
still retry — that is where retrying is both safe and valuable.
Affects the streaming path of every provider routed through this helper:
OpenAI, DashScope, Anthropic and Ollama.
Adds ModelUtilsStreamingRetryTest covering four cases: mid-stream retryable
error, failure before the first chunk, non-retryable mid-stream error, and a
clean stream. Only the first fails without this change; the other three pass
both ways, pinning that retries are not over-disabled.
ba74fb1 to
178a982
Compare
|
Apologies — I force-pushed to correct a mistake on my side. The branch was accidentally cut from a local branch based on #2476's head rather than from
That explains @copilot-pull-request-reviewer describing this PR as covering "two areas" and @oss-maintainer referring to the streaming fix as "commit 2". The branch is now rebuilt directly on
Re-verified after the rebuild: Thanks @oss-maintainer for the review of the streaming portion — the assessment of the approach matches the intent: |
|
Correction to my previous comment: the author of #2476 / commit Responding to both concerns in the review — both are resolved by the force-push: 1. PR mixed two independent fixes — fixed. The branch is rebuilt on
2. Spotless failures in Thanks for the precise review — the composition problem was real and entirely my error in cutting the branch. |
Review feedback: the emitted flag and the Retry instance were created once per applyTimeoutAndRetry call and captured by the returned Flux. A Flux may be subscribed more than once, so a second subscriber could inherit emitted=true from an earlier subscription and skip retries it should have performed. Moved the retry wiring inside Flux.defer so both the flag and the Retry spec are allocated per subscription. Adds emittedFlagIsPerSubscription: the same wrapped Flux is subscribed twice, each subscription emitting a chunk before a retryable failure. Shared state would make the second subscription diverge from the first. ModelUtilsStreamingRetryTest: 5 tests, 0 failures.
|
Thanks @copilot — that catch is correct and I've fixed it. The Fixed by moving the retry wiring inside responseFlux = Flux.defer(() -> {
AtomicBoolean emitted = new AtomicBoolean(false);
Retry retrySpec = Retry.backoff(retryMaxAttempts - 1, retryInitialBackoff)
.maxBackoff(retryMaxBackoff)
.jitter(0.5)
.filter(error -> !emitted.get() && retryableError.test(error))
.doBeforeRetry(...);
return retrySource.doOnNext(response -> emitted.set(true)).retryWhen(retrySpec);
});Added
Your second comment (about this PR also modifying |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Codecov flagged one partial branch on the retry filter `!emitted.get() && retryableError.test(error)`. The existing non-retryable case emits a chunk first, so the emitted guard short-circuits and the predicate's false path was never taken. Adds a case that fails with a non-retryable error before emitting anything, so both operands of the condition are exercised in both directions. 6 tests, 0 failures.
|
Addressed the Codecov partial-branch warning (96.15%, The partial was on the retry filter's compound condition: .filter(error -> !emitted.get() && retryableError.test(error))The existing non-retryable case emits a chunk before failing, so
This PR now stands at: 3 commits, 2 files ( |
gxgeek-n
changed the title
Fixes #2478.
Problem
ModelUtils.applyTimeoutAndRetryattachesretryWhendirectly to the streamingFlux.retryWhenre-subscribes its upstream, so a retryable error arriving mid-stream makes the model regenerate the whole response — and the chunks already handed to middlewares, memory and the UI get delivered again.ExecutionConfig.RETRYABLE_ERRORSclassifiesIOExceptionandTimeoutExceptionas retryable, which is exactly what a dropped streaming connection looks like. WithMODEL_DEFAULTS(maxAttempts=3) the same content can land downstream three times.Affects the streaming path of every provider routed through this helper:
OpenAIChatModel,DashScopeChatModel,AnthropicChatModel,OllamaChatModel.Change
Track whether the stream has emitted anything and fold that into the retry filter:
doOnNextsits upstream ofretryWhenso every emitted element flips the flag. Failures before the first chunk still retry — connection setup failures and HTTP 429 ahead of the first token — which is where retrying is safe and valuable. Non-streaming calls emit a single element and are unaffected when they fail before producing it.14 lines changed in production code.
Verification
ModelUtilsStreamingRetryTest, four cases:[Hello, world]expected, got[Hello, world, Hello, world, Hello, world](upstream subscribed 3x)Only the first case fails without the change; the other three pass both ways, which pins that retries are not over-disabled — the fix is narrow rather than "turn retry off for streaming".
Reproduced with no live model: upstream emits two chunks then throws
IOException.Relationship to #2279
This likely explains the 30 s of silence reported there.
doFinallyis attached to the lifecycleMono, so it cannot fire until thatMonoterminates; a retry keeps it alive across backoff plus a full regeneration. The timeline in that report — middlewaresonCompleteat00:15:33,POST_CALLat00:16:03— matches "the Mono finished 30 s later" rather than "the completion signal took 30 s to propagate".Worth noting the fix proposed in #2279,
sink.tryEmitComplete(), cannot compile:FluxSinkhas notryEmit*method in any Reactor 3.x release (checked withjavapagainst 3.1.7.RELEASE / 3.4.16 / 3.6.18 / 3.7.6) —tryEmitComplete()belongs toSinks.Many. I have not closed #2279 since I cannot reproduce its exact environment; if maintainers confirm this is the same root cause it can be linked or closed as a duplicate.spotless:applyhas been run.