feat(hooks): add MessageDisplay hook for mid-turn streaming#2512
feat(hooks): add MessageDisplay hook for mid-turn streaming#2512yanchenko wants to merge 1 commit into
Conversation
Fires repeatedly as the assistant reply streams, before Stop (which only fires once at the end of the turn). Fire-and-forget, cumulative text payload, debounced (~200ms) except for the unconditional final firing. Fires from the single streaming path in KimiSoul._step shared by the shell, print, ACP, wire, and web surfaces.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| def _on_message_part(part: StreamedMessagePart) -> None: | ||
| wire_send(part) | ||
| # Only final-answer text counts as displayed; thinking and tool | ||
| # parts are different part types and never feed the accumulator. | ||
| if message_display is not None and isinstance(part, TextPart): | ||
| message_display.add_chunk(part.text) |
There was a problem hiding this comment.
🟡 Retried assistant replies send duplicated, garbled text to streaming hooks
The reply text handed to hook scripts keeps accumulating (add_chunk at src/kimi_cli/soul/kimisoul.py:1221) across a retry instead of being cleared when an attempt fails, so a retried reply arrives with the failed attempt's leftover partial text stuck in front of it.
Impact: When a step is retried after having already streamed some text, hook scripts receive corrupted, duplicated reply content instead of the real reply.
Accumulator is never reset across retry attempts
The dispatcher is constructed once per _step at src/kimi_cli/soul/kimisoul.py:1208-1214, before the tenacity retry wrapper _kosong_step_with_retry. kosong.generate streams every TextPart delta to on_message_part as it arrives (packages/kosong/src/kosong/_generate.py:61-64), so a retryable failure (e.g. APIConnectionError/APITimeoutError) that occurs after partial text has streamed leaves that partial text folded into MessageDisplayState.displayed_text. On the next attempt the full reply is streamed and appended onto that stale prefix, so the mid-stream and final (is_final: true) payloads carry attempt1_partial + attempt2_full.
The comment at src/kimi_cli/soul/kimisoul.py:1206-1207 justifies spanning the retry wrapper by claiming it matches "how the UI merges the attempts' partial text", but the shell UI actually discards the failed attempt's partial content on StepRetry via discard_retry_attempt (src/kimi_cli/ui/shell/visualize/_live_view.py:432-434,547), setting self._current_content_block = None. So the hook displayed_text diverges from what the UI renders, and the displayed_text cumulative-text contract is violated (text can appear duplicated / non-monotonic).
A fix would reset the dispatcher's accumulator at the start of each attempt (e.g. on StepRetry, mirroring discard_retry_attempt) while keeping the same message_id.
Prompt for agents
The MessageDisplayDispatcher is created once per _step (kimisoul.py:1208-1214), before the tenacity retry wrapper, and its cumulative text accumulator (MessageDisplayState.displayed_text) is never reset between retry attempts. Because kosong streams partial TextPart deltas to on_message_part before a retryable error is raised, a retried step accumulates attempt1's partial text and then appends attempt2's full text, producing garbled/duplicated displayed_text in both mid-stream and is_final hook payloads. The shell UI (src/kimi_cli/ui/shell/visualize/_live_view.py, discard_retry_attempt) discards the failed attempt's partial content on StepRetry, so the hook text diverges from the UI and the comment's parity claim (kimisoul.py:1206-1207) is incorrect. Consider adding a method on MessageDisplayDispatcher to reset the debounce accumulator (displayed_text/last_flushed_text) to empty while keeping message_id, and invoking it when an attempt is retried (e.g. via the retry before_sleep callback _before_step_retry_sleep, which already fires on each retry). Ensure any in-flight mid-stream state remains consistent after reset.
Was this helpful? React with 👍 or 👎 to provide feedback.
1 participant
Closes #2511.
Summary
Adds a fire-and-forget
MessageDisplayhook event that fires repeatedly as the assistant reply streams — beforeStop, which fires once at turn end. Modeled on the same feature in Qwen Code (QwenLM/qwen-code#6488), adapted to kimi-cli's hook engine.Design
message_id(stable per streamed message / model call),displayed_text(cumulative, never a delta),is_final, on top of the common base fields (session_id,cwd,hook_event_name).step_message_displayaccumulator (src/kimi_cli/hooks/message_display.py) — leading-edge ~200ms per chunk, no timer; the unconditional final flush delivers the reply's tail.MessageDisplayDispatcher): at most one mid-stream execution in flight; newer cumulative payloads replace the queued one (lossless).is_finalis dispatched immediately, alongside any stale in-flight execution, and the turn waits for it bounded by a shared 5s drain budget. Suppressed on cancellation and on textless (tool-call-only) messages.KimiSoul._stepwrappingon_message_part— every surface (shell, print/headless, ACP, wire, web) shares it. One dispatcher per model call, spanning the retry wrapper so retried attempts keep onemessage_id(matching how the UI merges partial text).is_finallands before theStophook. OnlyTextPartcounts as displayed; thinking and tool parts are excluded.Tests
tests/hooks/test_message_display.py— pure debounce table plus dispatcher semantics: coalescing to newest pending, final-supersede alongside stale in-flight, idempotentfinish(), abort suppression, drain timeout, delivery-failure swallow (18 tests).tests/core/test_kimisoul_message_display.py— soul-level with scripted providers: final firing with cumulative text beforeStop, thinking parts excluded, no-op without configured hooks, cancellation suppressesis_final(4 tests).supported_eventslists.Full suite:
pytest tests→ 2809 passed. The 65 failures are pre-existing Windows/POSIX environment issues (fcntl/ptyimports, missingpython3, tab-in-filename, editor/symlink tests) — verified byte-identical on unmodifiedmain.ruff check/ruff formatclean; pyright shows only the 6 pre-existing pty-helper errors also present onmain.Docs
docs/{en,zh}/customization/hooks.md: new event table row plus aMessageDisplaysection documenting the full delivery-semantics contract.