fix(soul): clear partial UI output when LLM step is retried - #2177
Merged
Conversation
When a streamed LLM call fails after emitting partial content, tenacity retries the step but the aborted attempt's text/think/tool-call output was already shown to the user, so the new attempt appeared concatenated with the failed one. Add a StepRetry wire event emitted from tenacity's before_sleep hook so UI consumers (live shell view, print backends) can discard the partial state and show a retry banner before the new attempt's output.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new StepRetry wire event to support clearing/discarding partial streamed UI output when an LLM step is retried (e.g., after a 429), preventing the next attempt’s stream from appearing concatenated with the failed attempt in shell/print outputs.
Changes:
- Introduces
StepRetryas a first-class wire message type and emits it from the step-level tenacity retry hook. - Updates shell live view and print backends to discard partial assistant/tool-call state on retry and render a retry banner (shell).
- Adds targeted test coverage for serde, shell rendering behavior, print backend buffering behavior, and end-to-end retry recovery.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
src/kimi_cli/wire/types.py |
Adds the StepRetry message and wires it into the event/message unions and exports. |
src/kimi_cli/soul/kimisoul.py |
Emits StepRetry from tenacity’s before_sleep hook during step retries. |
src/kimi_cli/ui/shell/visualize/_live_view.py |
Handles StepRetry by discarding partial streamed state and printing a retry banner. |
src/kimi_cli/ui/print/visualize.py |
Drops buffered assistant output/tool-calls on StepRetry across print backends while preserving notification flushing semantics. |
src/kimi_cli/acp/session.py |
Accepts/ignores StepRetry during ACP streaming to avoid disrupting the session loop. |
tests/core/test_wire_message.py |
Adds serde snapshot coverage for StepRetry. |
tests/core/test_kimisoul_retry_recovery.py |
Adds an integration-style test ensuring partial stream + retry yields correct UI/history behavior and emits StepRetry. |
tests/ui_and_conv/test_empty_think_part_indicator.py |
Verifies live view clears partial content blocks and renders a retry banner on StepRetry. |
tests/ui_and_conv/test_print_notifications.py |
Adds tests ensuring print backends discard partial assistant buffers on StepRetry. |
CHANGELOG.md |
Notes the user-facing retry/partial-output clearing behavior change. |
docs/en/release-notes/changelog.md |
Mirrors the changelog entry in the English docs. |
docs/zh/release-notes/changelog.md |
Adds the corresponding Chinese changelog entry. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This was referenced May 8, 2026
Update wire-mode.md (en/zh) to add the StepRetry event under the Event union with full field documentation, and refresh the protocol version reference. Update kimi-info.md examples to show the bumped wire protocol version.
Replace persistent console.print of StepRetry with a live-status block that updates in place, so consecutive retries no longer pile up as separate lines. Cleared once new content (think/text/tool call) starts or the turn ends. Also tighten the surrounding docstrings: discard_retry_attempt no longer claims the banner is a permanent scrollback boundary, and compose_agent_output's layout section is rewritten to distinguish modal preemption from additive composition rather than implying a strict priority order.
Collaborator
Author
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a5885471dc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
discardRetryAttemptMessages was iterating over every entry in currentToolCallsRef and clearing the map. Because that ref is session- scoped (only ever cleared in resetState), tool calls from earlier successful steps in the same turn would be deleted from the UI when StepRetry fired during a later step. Filter to tool calls without a result (the in-flight ones) and delete only those, mirroring the running-status filter the subagent path already uses.
Collaborator
Author
|
Codex Review: Didn't find any major issues. Can't wait for the next one! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
7Sageer
added a commit
to 7Sageer/kimi-cli
that referenced
this pull request
May 9, 2026
PR MoonshotAI#2177 bumped `WIRE_PROTOCOL_VERSION` from 1.9 to 1.10 in `src/kimi_cli/wire/protocol.py` but didn't update the e2e tests, so `tests_e2e/test_wire_protocol.py::{test_initialize_handshake, test_initialize_external_tool_conflict}` started failing on push to main. Update the three `"1.9"` literals in the e2e snapshot to `"1.10"` to match the new server response.
RealKai42
pushed a commit
that referenced
this pull request
May 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related Issue
N/A — small targeted bug fix.
Description
When a streamed LLM call fails after emitting partial content, tenacity retries the step but the aborted attempt's text/think/tool-call output was already shown to the user, so the new attempt appeared concatenated with the failed one.
This PR adds a
StepRetrywire event emitted from tenacity'sbefore_sleephook so UI consumers can discard the partial state and render a retry banner before the new attempt's output._current_content_block,_tool_call_blocks,_last_tool_call_block, then printsRetrying after <reason> · attempt N/M · <wait>. State that is independent of the LLM stream (status, compaction, MCP loading, notifications, approval queues) is intentionally preserved — see the docstring on_LiveView.discard_retry_attemptfor the rationale.JsonPrinter,FinalOnlyTextPrinter,FinalOnlyJsonPrinter) drop the partial assistant-message buffers; pending notifications captured during the failed attempt still flush so they aren't lost.web/src/hooks/wireTypes.ts,useSessionStream.ts) consumesStepRetryfor both the parent agent and subagents: the partial thinking/text blocks and any in-flight tool-call messages from the failed attempt are removed, a single live status lineRetrying after … · attempt N/M · <wait>is rendered (and updated in place across consecutive retries), and tool calls from earlier successful steps in the same turn are preserved (filtered bytc.result === undefined).Wire protocol version is bumped
1.9 → 1.10to record the additiveStepRetryevent, matching the convention from prior wire-type changes (#1601, #1552, #1743). The change is backward-compatible — older readers ofwire.jsonlskip unknown record types via existing exception handling inWireMessageEnvelope.to_wire_message()and the replay loop, so no client breaks. Public wire docs (docs/{en,zh}/customization/wire-mode.md) are updated with the new event schema.Tests
test_kimisoul_retry_recovery.py::test_step_retry_event_after_partial_stream— partial stream → 429 → retry yields only the new attempt in history.test_wire_message.py— StepRetry serde snapshot.test_empty_think_part_indicator.py::test_step_retry_clears_partial_content_and_updates_live_status— live view clears blocks and renders banner.test_print_notifications.py— three new tests covering each print backend's discard behavior.make gen-changelogfor this PR's own version-bump commit, the upstream model returned 503 — the new retry banner rendered correctly (Retrying after server error · attempt 2/3 · <1s), confirmingerror_type → reasonmapping and attempt/wait formatting on a real failure path.Known limitations / follow-ups
vis/src/features/wire-viewer/*) TypeScript client does not yet handleStepRetry; it will silently ignore the event, so the wire-viewer will continue to render the failed attempt's partial output until a follow-up PR adds the handler. Web is handled in this PR (see the bullet above)._run_with_connection_recoveryperforms first-layer retries (OAuth refresh, single-shot connection recovery) silently, without emittingStepRetry.track("api_retry", ...)telemetry and noretriescounter in the session-stats endpoint yet.Before:
After:
Checklist
make gen-changelogto update the changelog.make gen-docsto update the user documentation.