Skip to content
18 changes: 17 additions & 1 deletion python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
Comment thread
ktz03 marked this conversation as resolved.
Original file line number Diff line number Diff line change
Expand Up @@ -2483,7 +2483,16 @@ 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)
# Empty approval resumes prepend stored history; non-empty/replayed
# transcripts overlap-merge so clients do not double-write (#8140).
if raw_messages:
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
stored_interrupt=stored_snapshot.interrupt,
)
else:
raw_messages = snapshot_session.resume_seeded_messages(raw_messages)
else:
provider_suffix, snapshot_seed_messages = _split_service_session_input(
stored_snapshot_messages=stored_snapshot.messages,
Expand All @@ -2492,6 +2501,13 @@ async def run_agent_stream(
)
raw_messages = provider_suffix
elif not config.use_service_session:
if resume_payload is not None and raw_messages:
# Client-replayed transcript on predictive/generic resume: overlap-merge
# and mark seeded so save-time resume_seeded_messages does not prepend
# again (#8140). Empty interrupt-only resumes (e.g. confirm_changes)
# must stay empty so synthesized resume tool messages are the only
# turn input; history is restored at save when this flag stays false.
seeded_resume_from_snapshot = True
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -134,11 +134,17 @@ def effective_state(
def resume_seeded_messages(self, incoming: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Prepend copies of stored thread history to a resume request's messages.

Resume requests carry only the synthesized interrupt response; seeding
Resume requests often carry only the synthesized interrupt response; seeding
with stored history keeps the persisted thread from being truncated.

For non-empty client-replayed transcripts that already overlap stored
history, callers should use ``_reconstruct_messages_from_thread_snapshot``
instead so messages are not double-persisted (#8140).
"""
if self._stored is None:
return incoming
if not incoming:
return [copy.deepcopy(message) for message in self._stored.messages]
return [copy.deepcopy(message) for message in self._stored.messages] + incoming

async def save(
Expand Down
14 changes: 10 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,16 @@ 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. Empty input prepends;
# non-empty/replayed transcripts overlap-merge (#8140).
if builder_seed_messages:
builder_seed_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages if stored_snapshot is not None else [],
incoming_messages=builder_seed_messages,
stored_interrupt=stored_snapshot.interrupt if stored_snapshot is not None else None,
)
else:
builder_seed_messages = snapshot_session.resume_seeded_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
34 changes: 34 additions & 0 deletions python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
that interface only; runner integration is covered by the existing suite.
"""

from typing import Any

import pytest
from ag_ui.core import (
EventType,
Expand Down Expand Up @@ -209,6 +211,38 @@ async def test_seeded_copies_do_not_alias_stored_snapshot(self) -> None:
assert session.stored is not None
assert session.stored.messages[0]["content"] == "hi"

async def test_replayed_transcript_uses_reconstructor_not_blind_prepend(self) -> None:
"""Client-replayed history must not be naively prepended again (#8140)."""
from agent_framework_ag_ui._run_common import _reconstruct_messages_from_thread_snapshot

stored: list[dict[str, Any]] = [
{"id": "u1", "role": "user", "content": "please run the tool"},
{
"id": "a1",
"role": "assistant",
"content": "",
"toolCalls": [
{"id": "c1", "type": "function", "function": {"name": "needs_approval", "arguments": "{}"}}
],
},
]
incoming: list[dict[str, Any]] = [
*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 [message["id"] for message in reconstructed] == ["u1", "a1", "u2"]
# resume_seeded_messages remains blind prepend for save-time / empty seeds.
snapshot = AGUIThreadSnapshot(messages=stored, interrupt=[{"interruptId": "int-1"}])
store = await make_store_with("user-1", "t-replay", snapshot)
session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t-replay")
seeded = session.resume_seeded_messages([{"id": "u2", "role": "user", "content": "approved"}])
assert [message["id"] for message in seeded] == ["u1", "a1", "u2"]

async def test_without_stored_snapshot_returns_incoming_unchanged(self) -> None:
session = await ThreadSnapshotSession.open(store=None, scope=None, thread_id="t1")
incoming = [{"id": "m2", "role": "user", "content": "hello"}]
Expand Down
Loading