Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
Comment thread
ktz03 marked this conversation as resolved.
Resolved
Original file line number Diff line number Diff line change
Expand Up @@ -2483,7 +2483,8 @@ async def run_agent_stream(
seeded_resume_from_snapshot = True

if not config.use_service_session:
raw_messages = snapshot_session.resume_seeded_messages(raw_messages)
# One session-owned reconcile for empty vs replayed resume shapes (#8140).
raw_messages = snapshot_session.reconcile_resume_messages(raw_messages)
else:
provider_suffix, snapshot_seed_messages = _split_service_session_input(
stored_snapshot_messages=stored_snapshot.messages,
Expand All @@ -2492,11 +2493,17 @@ async def run_agent_stream(
)
raw_messages = provider_suffix
elif not config.use_service_session:
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
stored_interrupt=stored_snapshot.interrupt,
)
if resume_payload is not None:
# Predictive-state / generic resumes also merge history here; mark seeded
# so save-time resume_seeded_messages does not prepend again (#8140).
seeded_resume_from_snapshot = True
raw_messages = snapshot_session.reconcile_resume_messages(raw_messages)
Comment thread
ktz03 marked this conversation as resolved.
Resolved
else:
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
stored_interrupt=stored_snapshot.interrupt,
)
else:
provider_suffix, snapshot_seed_messages = _split_service_session_input(
stored_snapshot_messages=stored_snapshot.messages,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
StateSnapshotEvent,
)

from ._run_common import _build_run_finished_event
from ._run_common import _build_run_finished_event, _reconstruct_messages_from_thread_snapshot
from ._snapshots import (
AGUIThreadSnapshot,
AGUIThreadSnapshotStore,
Expand Down Expand Up @@ -141,6 +141,23 @@ def resume_seeded_messages(self, incoming: list[dict[str, Any]]) -> list[dict[st
return incoming
return [copy.deepcopy(message) for message in self._stored.messages] + incoming

def reconcile_resume_messages(self, incoming: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Merge stored history with resume messages without double-persisting.

Empty incoming (interrupt-only / checkpoint resume) prepends stored
history. Non-empty incoming is overlapped via the thread-snapshot
reconstructor so a client that replays its transcript is not duplicated.
"""
if self._stored is None:
return incoming
if not incoming:
return self.resume_seeded_messages(incoming)
return _reconstruct_messages_from_thread_snapshot(
stored_messages=self._stored.messages,
incoming_messages=incoming,
stored_interrupt=self._stored.interrupt,
)

async def save(
self,
*,
Expand Down
7 changes: 3 additions & 4 deletions python/packages/ag-ui/agent_framework_ag_ui/_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -446,10 +446,9 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]:
run_checkpoint_storage = _OwnedWorkflowCheckpointStorage(checkpoint_storage, request_owner)
builder_seed_messages = raw_messages
if resume_payload is not None or (checkpoint_id is not None and not raw_messages):
# Resume requests carry only the synthesized interrupt response, and a
# checkpoint-only resume carries no new messages at all; in both cases seed
# the builder with stored history to avoid persisting a truncated thread.
builder_seed_messages = snapshot_session.resume_seeded_messages(builder_seed_messages)
# Resume / checkpoint-only requests need stored history. Session-owned
# reconcile covers empty vs client-replayed transcripts (#8140).
builder_seed_messages = snapshot_session.reconcile_resume_messages(builder_seed_messages)
snapshot_builder = _WorkflowSnapshotBuilder(builder_seed_messages) if snapshot_session.enabled else None
if snapshot_builder is not None and effective_state:
# Seed builder state so a run that emits no StateSnapshotEvent still
Expand Down
92 changes: 92 additions & 0 deletions python/packages/ag-ui/tests/ag_ui/test_resume_transcript_dedupe.py
Comment thread
moonbox3 marked this conversation as resolved.
Resolved
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
# Copyright (c) Microsoft. All rights reserved.

"""Regression for AG-UI resume with client-replayed transcript (#8140)."""

from agent_framework_ag_ui._run_common import _reconstruct_messages_from_thread_snapshot
from agent_framework_ag_ui._snapshot_session import ThreadSnapshotSession
from agent_framework_ag_ui._snapshots import AGUIThreadSnapshot


def _stored_history() -> list[dict]:
return [
{"id": "u1", "role": "user", "content": "please run the tool"},
{
"id": "a1",
"role": "assistant",
"content": "",
"toolCalls": [{"id": "c1", "type": "function", "function": {"name": "needs_approval", "arguments": "{}"}}],
},
]


def test_resume_with_replayed_transcript_does_not_duplicate_history():
"""Clients that re-send the full transcript on resume must not double-write."""
stored = _stored_history()
incoming = [
*stored,
{"id": "u2", "role": "user", "content": "approved"},
]

reconstructed = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored,
incoming_messages=incoming,
stored_interrupt=[{"interruptId": "int-1"}],
)

assert [m.get("id") for m in reconstructed] == ["u1", "a1", "u2"]
assert len(reconstructed) == 3


def test_resume_with_partial_new_turn_still_seeds_history():
"""Non-empty resume suffix is overlapped onto stored history without duplication."""
stored = _stored_history()
incoming = [{"id": "u2", "role": "user", "content": "approved"}]
Comment thread
ktz03 marked this conversation as resolved.
Resolved

reconstructed = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored,
incoming_messages=incoming,
stored_interrupt=[{"interruptId": "int-1"}],
)

assert [m.get("id") for m in reconstructed] == ["u1", "a1", "u2"]


def test_reconstruct_with_empty_incoming_returns_empty():
"""Empty incoming is unchanged; call sites must use resume seeding instead."""
stored = _stored_history()

assert (
_reconstruct_messages_from_thread_snapshot(
stored_messages=stored,
incoming_messages=[],
stored_interrupt=[{"interruptId": "int-1"}],
)
== []
)


def test_reconcile_resume_messages_empty_vs_replayed():
"""Session-owned reconcile covers interrupt-only and replayed transcript shapes."""
stored = _stored_history()
session = ThreadSnapshotSession(
store=None,
scope=None,
thread_id="t1",
stored=AGUIThreadSnapshot(
messages=stored,
state=None,
interrupt=[{"interruptId": "int-1"}],
session_state=None,
),
)

empty = session.reconcile_resume_messages([])
assert [m.get("id") for m in empty] == ["u1", "a1"]

replayed = session.reconcile_resume_messages(
[
*stored,
{"id": "u2", "role": "user", "content": "approved"},
]
)
assert [m.get("id") for m in replayed] == ["u1", "a1", "u2"]
Loading