From 0c1d709137db545bfc9ae8ce4069819b3dc09f8b Mon Sep 17 00:00:00 2001 From: Manideep Malyala Date: Fri, 21 Aug 2026 22:23:53 +0530 Subject: [PATCH 1/2] feat(ag-ui): add emit_messages_snapshot config to suppress terminal snapshot --- .../ag-ui/agent_framework_ag_ui/_agent.py | 9 +++++ .../ag-ui/agent_framework_ag_ui/_agent_run.py | 2 +- python/packages/ag-ui/tests/ag_ui/test_run.py | 38 +++++++++++++++++++ 3 files changed, 48 insertions(+), 1 deletion(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 9ae73c3790..87a8d4b25e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -26,6 +26,7 @@ def __init__( require_confirmation: bool = True, snapshot_store: AGUIThreadSnapshotStore | None = None, a2ui_config: dict[str, Any] | None = None, + emit_messages_snapshot: bool = True, ): """Initialize agent configuration. @@ -42,6 +43,9 @@ def __init__( require_confirmation: Whether predictive updates require user confirmation before applying a2ui_config: Optional backend A2UI config consumed by auto-injection (``forwardedProps.injectA2UITool``). See ``plan_a2ui_injection``. + emit_messages_snapshot: Whether to emit a terminal MessagesSnapshotEvent at the end of runs. + Defaults to True for backward compatibility. Set to False when using HistoryProvider + to prevent redundant full-transcript rewrites on the client. """ self.state_schema = self._normalize_state_schema(state_schema) self.predict_state_config = predict_state_config or {} @@ -49,6 +53,7 @@ def __init__( self.require_confirmation = require_confirmation self.snapshot_store = snapshot_store self.a2ui_config = a2ui_config + self.emit_messages_snapshot = emit_messages_snapshot @staticmethod def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]: @@ -96,6 +101,7 @@ def __init__( use_service_session: bool = False, snapshot_store: AGUIThreadSnapshotStore | None = None, a2ui_config: dict[str, Any] | None = None, + emit_messages_snapshot: bool = True, ): """Initialize the AG-UI compatible agent wrapper. @@ -113,6 +119,8 @@ def __init__( snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless endpoint setup also provides an explicit Snapshot Scope resolver. a2ui_config: Optional backend A2UI config consumed by auto-injection. + emit_messages_snapshot: Whether to emit a terminal MessagesSnapshotEvent at the end of runs. + Defaults to True. Set to False when using HistoryProvider. """ self.agent = agent self.name = name or getattr(agent, "name", "agent") @@ -125,6 +133,7 @@ def __init__( require_confirmation=require_confirmation, snapshot_store=snapshot_store, a2ui_config=a2ui_config, + emit_messages_snapshot=emit_messages_snapshot, ) # Server-side Approval State. Populated when approval requests are emitted diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 8b99e106dc..141d2610f0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -2952,7 +2952,7 @@ async def run_agent_stream( # so an unrelated user tool named "generate_a2ui" keeps its snapshot. if a2ui_active: logger.info("Suppressing terminal MessagesSnapshotEvent for A2UI run to preserve streamed message order.") - if not a2ui_active and not _should_suppress_intermediate_snapshot( + if config.emit_messages_snapshot and not a2ui_active and not _should_suppress_intermediate_snapshot( last_tool_name, predict_state_config, config.require_confirmation ): yield snapshot_event diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 68db77e428..6334ade427 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -1823,6 +1823,44 @@ async def test_run_agent_stream_accumulates_multiple_confirm_interrupts(): assert interrupt_tool_names == {"generate_tasks", "generate_notes"} +async def test_run_agent_stream_suppresses_messages_snapshot_if_configured(): + """When emit_messages_snapshot=False, no terminal MessagesSnapshotEvent is yielded.""" + from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + + from agent_framework_ag_ui import AgentFrameworkAgent + + updates = [ + AgentResponseUpdate(contents=[Content.from_text("Hello")], role="assistant"), + ] + + stub = StubAgent(updates=updates) + agent = AgentFrameworkAgent( + agent=stub, + emit_messages_snapshot=False, + ) + + payload = { + "thread_id": "thread-1", + "run_id": "run-1", + "messages": [{"role": "user", "content": "Hi"}], + } + + events = [event async for event in agent.run(payload)] + + # We should have TextMessageStart/Delta/End, but no MessagesSnapshot + snapshot_events = [e for e in events if getattr(e, "type", None) == "MESSAGES_SNAPSHOT"] + assert len(snapshot_events) == 0, "MessagesSnapshotEvent should be suppressed" + + # Still finishes normally + finished_events = [ + e + for e in events + if getattr(e, "type", None) == "RUN_FINISHED" + or getattr(getattr(e, "type", None), "value", None) == "RUN_FINISHED" + ] + assert len(finished_events) == 1 + + def test_emit_oauth_consent_request(): """Test that oauth_consent_request content emits a CustomEvent.""" content = Content.from_oauth_consent_request( From 9549a9e193af3a05335d22b1fbace6f5b1b5b661 Mon Sep 17 00:00:00 2001 From: Manideep Malyala Date: Fri, 21 Aug 2026 22:31:37 +0530 Subject: [PATCH 2/2] style: remove trailing whitespace in test_run.py to pass ruff --- python/packages/ag-ui/tests/ag_ui/test_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 6334ade427..490493f95b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -1846,11 +1846,11 @@ async def test_run_agent_stream_suppresses_messages_snapshot_if_configured(): } events = [event async for event in agent.run(payload)] - + # We should have TextMessageStart/Delta/End, but no MessagesSnapshot snapshot_events = [e for e in events if getattr(e, "type", None) == "MESSAGES_SNAPSHOT"] assert len(snapshot_events) == 0, "MessagesSnapshotEvent should be suppressed" - + # Still finishes normally finished_events = [ e