Skip to content

🦌 Fix multi-step agent context and display merging#54

Open
masudahiroto wants to merge 13 commits intomainfrom
masudahiroto/fix-per-step-text-message-context
Open

🦌 Fix multi-step agent context and display merging#54
masudahiroto wants to merge 13 commits intomainfrom
masudahiroto/fix-per-step-text-message-context

Conversation

@masudahiroto
Copy link
Copy Markdown
Contributor

@masudahiroto masudahiroto commented Apr 3, 2026 •
edited
Loading

Summary

Fixes a bug where multi-step agent runs (producing multiple assistant steps with text+tool per step) corrupted the conversation context sent to the LLM on subsequent runs. Tool calls from all steps were being merged into a single assistant message, and all intermediate text was being concatenated into a single final message instead of being preserved per-step.

Also adds a display-layer merge (mergeAssistantMessagesForDisplay) so the UI shows a single unified bubble per turn while keeping the correct per-step structure for LLM context.

Changes

  • packages/client/src/client.ts — Handle STEP_FINISHED event to flush per-step assistant+toolCall messages and tool results, preserving correct API message order
  • packages/client/src/components/UseAIChatPanel.tsx — Replace inline filter with mergeAssistantMessagesForDisplay to combine intermediate assistant text into a single display bubble
  • packages/client/src/utils/mergeAssistantMessages.ts — New utility to merge consecutive assistant messages within a turn for display, using deterministic IDs
  • packages/client/src/utils/mergeAssistantMessages.test.ts — Tests for merge utility including edge cases
  • packages/client/src/utils/messageConversion.ts — Moved transformMessagesToClientFormat here as a shared utility; fixed extractTurnMessages to preserve content on intermediate assistant messages
  • packages/client/src/utils/messageConversion.test.ts — Tests using real transformMessagesToClientFormat and extractTurnMessages
  • packages/client/src/hooks/useChatManagement.ts — Import transformMessagesToClientFormat from shared utility instead of local copy
  • packages/client/src/hooks/useServerEvents.ts — Remove local copy of extractTurnMessages, use shared utility
  • packages/client/src/client.multiStepContext.test.ts — Integration tests validating per-step message structure and 2nd-run context correctness
  • packages/server/src/agents/AISDKAgent.ts — Server-side per-step event emission fixes
  • packages/server/src/agents/AISDKAgent.perStepEvents.test.ts — Tests for per-step event emission
  • packages/client/src/providers/chatRepository/toolCallPersistence.test.ts — Updated persistence tests

Created by deer — review carefully.

deer-agent and others added 10 commits April 2, 2026 22:47
@claude
fix: preserve per-step text+tool association in multi-step agent runs
c99b869 
When a multi-step agent run produces multiple steps each with text AND
tool calls, the conversation context was losing per-step boundaries.
After page reload, all tool calls were merged into one assistant message
and all text was concatenated into a separate message, causing the LLM
to misunderstand previous turns.

Server: Move hasEmittedTextStart/messageId to StepContext so
TEXT_MESSAGE_START/END is emitted per step instead of per run.
Client: Add STEP_FINISHED handler to flush per-step messages while
maintaining backward compatibility with servers without step events.
useServerEvents: Fix extractTurnMessages to preserve assistant content.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
fix: combine multi-step text into single bubble and prevent streaming…
bfb8c60 
… flash

Two UI issues with per-step TEXT_MESSAGE emission:

1. Streaming text cleared at TEXT_MESSAGE_END caused flash between steps.
   Fix: accumulate text across steps via runTextRef, only clear at
   RUN_FINISHED.

2. Intermediate assistant messages with text+toolCalls displayed as
   separate bubbles. Fix: hide all assistant messages with toolCalls
   from display, combine all step text into the final saved message.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
style: remove bubble background from assistant messages
25e5bdb 
Assistant messages now render with transparent background and minimal
padding for a cleaner, less cluttered chat appearance. User messages
retain the existing bubble style.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Revert "style: remove bubble background from assistant messages"
4d2d47d 
This reverts commit 25e5bdb.
@claude
fix: prevent text duplication in final message across multi-step runs
9296cd3 
saveAIResponse was using runTextRef (accumulated text from all steps)
instead of client.currentMessageContent (last step only). Since
intermediate steps' text is already preserved in turnMessages via
extractTurnMessages, the final message duplicated that text.

Use client.currentMessageContent for the saved message content.
runTextRef continues to be used for streaming UI display only.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
refactor: use real extractTurnMessages in tests instead of re-impleme…
b94aa2b 
…ntation

Tests were re-implementing extractTurnMessages locally, so bugs in the
real function would not be caught. Export extractTurnMessages from
useServerEvents.ts and import it in tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
refactor: use real transformMessagesToClientFormat in tests
a8716aa 
Tests were re-implementing transformMessagesToClientFormat locally.
Export from useChatManagement.ts and import in both persistence test
files so bugs in the real function are caught by tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
fix: merge multi-step assistant text into single display bubble
48c8831 
Per-step messages are preserved for correct LLM context, but the UI
now combines consecutive assistant messages within each turn into a
single bubble. Intermediate assistant messages (with toolCalls) have
their text merged with the final text-only message using paragraph
separators.

Add mergeAssistantMessagesForDisplay utility with 8 unit tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
test: consolidate multi-step context tests into 2 focused files
6827999 
Replace 5 scattered test files (~1850 lines) with 2 focused test files
(~350 lines) that directly verify the core bug fix: per-step text+tool
association is preserved in conversation context sent to LLM on 2nd run.

- client.multiStepContext.test.ts: integration test covering message
  assembly, 2nd run_agent context correctness, and text deduplication
- AISDKAgent.perStepEvents.test.ts: server unit test for per-step
  TEXT_MESSAGE_START/END emission and event ordering

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@claude
refactor: replace runTextRef with boolean flag, add mergeAssistantMes…
238e87e 
…sages tests

- Replace runTextRef (accumulated string never read) with
  hasTextFromPriorStepRef boolean for step separator logic
- Add unit tests for mergeAssistantMessagesForDisplay (7 cases:
  simple exchange, multi-step merge, tool filtering, tool-only step,
  multiple turns, property preservation, ContentPart arrays)
- Strengthen extractTurnMessages assertions in integration test

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@masudahiroto
Copy link
Copy Markdown
Contributor Author

masudahiroto commented Apr 3, 2026

still in draft. i will check this later and open it if there is no issue.

@claude
fix: use deterministic IDs in mergeAssistantMessages, add edge-case t…
0e8249c 
…ests

Replace Date.now() with constituent message IDs to produce stable React
keys. Update stale JSDoc in extractTurnMessages to reflect STEP_FINISHED
flushing. Add tests for empty input, trailing pending text, and ID
determinism.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@masudahiroto
Copy link
Copy Markdown
Contributor Author

masudahiroto commented Apr 6, 2026

This PR is ready for review.

@masudahiroto masudahiroto marked this pull request as ready for review April 6, 2026 04:09
@masudahiroto masudahiroto mentioned this pull request Apr 7, 2026
Draft
@claude
Merge origin/main: resolve conflicts with PR #50 token limit recovery
4e7b7f7 
Both branches added fields to StepContext and RUN_FINISHED handler.
Kept all additions from both sides:
- This branch: messageId, hasEmittedTextStart, streamingText/chatId cleanup
- PR #50: stepFinishReason, executingTool cleanup on truncation

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
mm-zacharydavison
mm-zacharydavison requested changes Apr 7, 2026
View reviewed changes
packages/client/src/hooks/useChatManagement.ts Outdated
* so the server can reconstruct valid API messages.
*/
function transformMessagesToClientFormat(persistedMessages: PersistedMessage[]): AGUIMessage[] {
export function transformMessagesToClientFormat(persistedMessages: PersistedMessage[]): AGUIMessage[] {
Copy link
Copy Markdown
Contributor

@mm-zacharydavison mm-zacharydavison Apr 7, 2026

Choose a reason for hiding this comment

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

i dont think you should export functions from hooks files, it's a leaking abstraction.
If they are common functions, they should probably be somewhere else.

Copy link
Copy Markdown
Contributor Author

@masudahiroto masudahiroto Apr 8, 2026

Choose a reason for hiding this comment

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

Thank you for pointing this out. I've addressed the review in 4deed3d.

I refactored the code to move these two functions (transformMessagesToClientFormat and extractTurnMessages) to packages/client/src/utils/messageConversion.ts.

packages/client/src/hooks/useServerEvents.ts Outdated
* backward compatibility when the server does not emit step events).
*/
function extractTurnMessages(messages: Message[], startIndex: number): PersistedMessage[] {
export function extractTurnMessages(messages: Message[], startIndex: number): PersistedMessage[] {
Copy link
Copy Markdown
Contributor

@mm-zacharydavison mm-zacharydavison Apr 7, 2026

Choose a reason for hiding this comment

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

same here

Copy link
Copy Markdown
Contributor Author

@masudahiroto masudahiroto Apr 8, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

I've addressed this review too, as noted in this comment.

@masudahiroto masudahiroto changed the title fix: preserve per-step text+tool_call association in multi-step agent runs 🦌 Fix multi-step agent context and display merging Apr 8, 2026
@masudahiroto masudahiroto force-pushed the masudahiroto/fix-per-step-text-message-context branch from 6191e6c to 4deed3d Compare April 8, 2026 03:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mm-zacharydavison mm-zacharydavison Awaiting requested review from mm-zacharydavison mm-zacharydavison is a code owner

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@masudahiroto @mm-zacharydavison