Python: (AG-UI) Support service-managed thread on AG-UI (#3136)

* added service thread support

* set service_thread_id to only supplied_thread_id

* uses raw_representation to extract the conversation_id

* removed accidental edit

* updated test to use raw_representation

* resolves copilot review feedback

* revert back StubAgent, since not used

* removed relative module import

* removed hasattr check per PR feedback

python/packages/ag-ui/agent_framework_ag_ui/_agent.py

@@ -24,6 +24,7 @@ class AgentConfig:
         self,
         state_schema: Any | None = None,
         predict_state_config: dict[str, dict[str, str]] | None = None,
+        use_service_thread: bool = False,
         require_confirmation: bool = True,
     ):
         """Initialize agent configuration.
@@ -31,10 +32,12 @@ class AgentConfig:
         Args:
             state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
             predict_state_config: Configuration for predictive state updates
+            use_service_thread: Whether the agent thread is service-managed
             require_confirmation: Whether predictive updates require confirmation
         """
         self.state_schema = self._normalize_state_schema(state_schema)
         self.predict_state_config = predict_state_config or {}
+        self.use_service_thread = use_service_thread
         self.require_confirmation = require_confirmation
 
     @staticmethod
@@ -86,6 +89,7 @@ class AgentFrameworkAgent:
         predict_state_config: dict[str, dict[str, str]] | None = None,
         require_confirmation: bool = True,
         orchestrators: list[Orchestrator] | None = None,
+        use_service_thread: bool = False,
         confirmation_strategy: ConfirmationStrategy | None = None,
     ):
         """Initialize the AG-UI compatible agent wrapper.
@@ -101,6 +105,7 @@ class AgentFrameworkAgent:
                 Set to False for agentic generative UI that updates automatically.
             orchestrators: Custom orchestrators (auto-configured if None).
                 Orchestrators are checked in order; first match handles the request.
+            use_service_thread: Whether the agent thread is service-managed.
             confirmation_strategy: Strategy for generating confirmation messages.
                 Defaults to DefaultConfirmationStrategy if None.
         """
@@ -111,6 +116,7 @@ class AgentFrameworkAgent:
         self.config = AgentConfig(
             state_schema=state_schema,
             predict_state_config=predict_state_config,
+            use_service_thread=use_service_thread,
             require_confirmation=require_confirmation,
         )
 

python/packages/ag-ui/agent_framework_ag_ui/_orchestrators.py

@@ -6,7 +6,7 @@ import json
 import logging
 import uuid
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from typing import TYPE_CHECKING, Any
 
 from ag_ui.core import (
@@ -53,11 +53,18 @@ from ._orchestration._tooling import (
     merge_tools,
     register_additional_client_tools,
 )
-from ._utils import convert_agui_tools_to_agent_framework, generate_event_id, get_role_value
+from ._utils import (
+    convert_agui_tools_to_agent_framework,
+    generate_event_id,
+    get_conversation_id_from_update,
+    get_role_value,
+)
 
 if TYPE_CHECKING:
     from ._agent import AgentConfig
     from ._confirmation_strategies import ConfirmationStrategy
+    from ._events import AgentFrameworkEventBridge
+    from ._orchestration._state_manager import StateManager
 
 
 logger = logging.getLogger(__name__)
@@ -92,6 +99,8 @@ class ExecutionContext:
         self._last_message = None
         self._run_id: str | None = None
         self._thread_id: str | None = None
+        self._supplied_run_id: str | None = None
+        self._supplied_thread_id: str | None = None
 
     @property
     def messages(self):
@@ -125,26 +134,66 @@ class ExecutionContext:
             self._last_message = self.messages[-1]
         return self._last_message
 
+    @property
+    def supplied_run_id(self) -> str | None:
+        """Get the supplied run ID, if any."""
+        if self._supplied_run_id is None:
+            self._supplied_run_id = self.input_data.get("run_id") or self.input_data.get("runId")
+        return self._supplied_run_id
+
     @property
     def run_id(self) -> str:
-        """Get or generate run ID."""
+        """Get supplied run ID or generate a new run ID."""
+        if self._run_id:
+            return self._run_id
+
+        if self.supplied_run_id:
+            self._run_id = self.supplied_run_id
+
         if self._run_id is None:
-            self._run_id = self.input_data.get("run_id") or self.input_data.get("runId") or str(uuid.uuid4())
-        # This should never be None after the if block above, but satisfy type checkers
-        if self._run_id is None:  # pragma: no cover
-            raise RuntimeError("Failed to initialize run_id")
+            self._run_id = str(uuid.uuid4())
+
         return self._run_id
 
+    @property
+    def supplied_thread_id(self) -> str | None:
+        """Get the supplied thread ID, if any."""
+        if self._supplied_thread_id is None:
+            self._supplied_thread_id = self.input_data.get("thread_id") or self.input_data.get("threadId")
+        return self._supplied_thread_id
+
     @property
     def thread_id(self) -> str:
-        """Get or generate thread ID."""
+        """Get supplied thread ID or generate a new thread ID."""
+        if self._thread_id:
+            return self._thread_id
+
+        if self.supplied_thread_id:
+            self._thread_id = self.supplied_thread_id
+
         if self._thread_id is None:
-            self._thread_id = self.input_data.get("thread_id") or self.input_data.get("threadId") or str(uuid.uuid4())
-        # This should never be None after the if block above, but satisfy type checkers
-        if self._thread_id is None:  # pragma: no cover
-            raise RuntimeError("Failed to initialize thread_id")
+            self._thread_id = str(uuid.uuid4())
+
         return self._thread_id
 
+    def update_run_id(self, new_run_id: str) -> None:
+        """Update the run ID in the context.
+
+        Args:
+            new_run_id: The new run ID to set
+        """
+        self._supplied_run_id = new_run_id
+        self._run_id = new_run_id
+
+    def update_thread_id(self, new_thread_id: str) -> None:
+        """Update the thread ID in the context.
+
+        Args:
+            new_thread_id: The new thread ID to set
+        """
+        self._supplied_thread_id = new_thread_id
+        self._thread_id = new_thread_id
+
 
 class Orchestrator(ABC):
     """Base orchestrator for agent execution flows."""
@@ -297,6 +346,28 @@ class DefaultOrchestrator(Orchestrator):
         """
         return True
 
+    def _create_initial_events(
+        self, event_bridge: "AgentFrameworkEventBridge", state_manager: "StateManager"
+    ) -> Sequence[BaseEvent]:
+        """Generate initial events for the run.
+
+        Args:
+            event_bridge: Event bridge for creating events
+        Returns:
+            Initial AG-UI events
+        """
+        events: list[BaseEvent] = [event_bridge.create_run_started_event()]
+
+        predict_event = state_manager.predict_state_event()
+        if predict_event:
+            events.append(predict_event)
+
+        snapshot_event = state_manager.initial_snapshot_event(event_bridge)
+        if snapshot_event:
+            events.append(snapshot_event)
+
+        return events
+
     async def run(
         self,
         context: ExecutionContext,
@@ -342,17 +413,11 @@ class DefaultOrchestrator(Orchestrator):
             approval_tool_name=approval_tool_name,
         )
 
-        yield event_bridge.create_run_started_event()
+        if context.config.use_service_thread:
+            thread = AgentThread(service_thread_id=context.supplied_thread_id)
+        else:
+            thread = AgentThread()
 
-        predict_event = state_manager.predict_state_event()
-        if predict_event:
-            yield predict_event
-
-        snapshot_event = state_manager.initial_snapshot_event(event_bridge)
-        if snapshot_event:
-            yield snapshot_event
-
-        thread = AgentThread()
         thread.metadata = {  # type: ignore[attr-defined]
             "ag_ui_thread_id": context.thread_id,
             "ag_ui_run_id": context.run_id,
@@ -363,6 +428,8 @@ class DefaultOrchestrator(Orchestrator):
         provider_messages = context.messages or []
         snapshot_messages = context.snapshot_messages
         if not provider_messages:
+            for event in self._create_initial_events(event_bridge, state_manager):
+                yield event
             logger.warning("No messages provided in AG-UI input")
             yield event_bridge.create_run_finished_event()
             return
@@ -554,13 +621,41 @@ class DefaultOrchestrator(Orchestrator):
                     confirmation_message = strategy.on_state_rejected()
 
             message_id = generate_event_id()
+            for event in self._create_initial_events(event_bridge, state_manager):
+                yield event
             yield TextMessageStartEvent(message_id=message_id, role="assistant")
             yield TextMessageContentEvent(message_id=message_id, delta=confirmation_message)
             yield TextMessageEndEvent(message_id=message_id)
             yield event_bridge.create_run_finished_event()
             return
 
+        should_recreate_event_bridge = False
         async for update in context.agent.run_stream(messages_to_run, **run_kwargs):
+            conv_id = get_conversation_id_from_update(update)
+            if conv_id and conv_id != context.thread_id:
+                context.update_thread_id(conv_id)
+                should_recreate_event_bridge = True
+
+            if update.response_id and update.response_id != context.run_id:
+                context.update_run_id(update.response_id)
+                should_recreate_event_bridge = True
+
+            if should_recreate_event_bridge:
+                event_bridge = AgentFrameworkEventBridge(
+                    run_id=context.run_id,
+                    thread_id=context.thread_id,
+                    predict_state_config=context.config.predict_state_config,
+                    current_state=current_state,
+                    skip_text_content=skip_text_content,
+                    require_confirmation=context.config.require_confirmation,
+                    approval_tool_name=approval_tool_name,
+                )
+                should_recreate_event_bridge = False
+
+            if update_count == 0:
+                for event in self._create_initial_events(event_bridge, state_manager):
+                    yield event
+
             update_count += 1
             logger.info(f"[STREAM] Received update #{update_count} from agent")
             if all_updates is not None:
@@ -672,6 +767,11 @@ class DefaultOrchestrator(Orchestrator):
                     yield TextMessageEndEvent(message_id=message_id)
                     logger.info(f"Emitted conversational message with length={len(response_dict['message'])}")
 
+        if all_updates is not None and len(all_updates) == 0:
+            logger.info("No updates received from agent - emitting initial events")
+            for event in self._create_initial_events(event_bridge, state_manager):
+                yield event
+
         logger.info(f"[FINALIZE] Checking for unclosed message. current_message_id={event_bridge.current_message_id}")
         if event_bridge.current_message_id:
             logger.info(f"[FINALIZE] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")

python/packages/ag-ui/agent_framework_ag_ui/_utils.py

@@ -10,7 +10,7 @@ from dataclasses import asdict, is_dataclass
 from datetime import date, datetime
 from typing import Any
 
-from agent_framework import AIFunction, Role, ToolProtocol
+from agent_framework import AgentResponseUpdate, AIFunction, ChatResponseUpdate, Role, ToolProtocol
 
 # Role mapping constants
 AGUI_TO_FRAMEWORK_ROLE: dict[str, Role] = {
@@ -259,3 +259,17 @@ def convert_tools_to_agui_format(
             continue
 
     return results if results else None
+
+
+def get_conversation_id_from_update(update: AgentResponseUpdate) -> str | None:
+    """Extract conversation ID from AgentResponseUpdate metadata.
+
+    Args:
+        update: AgentRunResponseUpdate instance
+    Returns:
+        Conversation ID if present, else None
+
+    """
+    if isinstance(update.raw_representation, ChatResponseUpdate):
+        return update.raw_representation.conversation_id
+    return None

python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py

@@ -637,6 +637,60 @@ async def test_suppressed_summary_with_document_state():
     assert "written" in full_text.lower() or "document" in full_text.lower()
 
 
+async def test_agent_with_use_service_thread_is_false():
+    """Test that when use_service_thread is False, the AgentThread used to run the agent is NOT set to the service thread ID."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    request_service_thread_id: str | None = None
+
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        nonlocal request_service_thread_id
+        thread = kwargs.get("thread")
+        request_service_thread_id = thread.service_thread_id if thread else None
+        yield ChatResponseUpdate(
+            contents=[TextContent(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
+        )
+
+    agent = ChatAgent(chat_client=StreamingChatClientStub(stream_fn))
+    wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=False)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
+
+    events: list[Any] = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+    assert request_service_thread_id is None  # type: ignore[attr-defined] (service_thread_id should be set)
+
+
+async def test_agent_with_use_service_thread_is_true():
+    """Test that when use_service_thread is True, the AgentThread used to run the agent is set to the service thread ID."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    request_service_thread_id: str | None = None
+
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        nonlocal request_service_thread_id
+        thread = kwargs.get("thread")
+        request_service_thread_id = thread.service_thread_id if thread else None
+        yield ChatResponseUpdate(
+            contents=[TextContent(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
+        )
+
+    agent = ChatAgent(chat_client=StreamingChatClientStub(stream_fn))
+    wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=True)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
+
+    events: list[Any] = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+    assert request_service_thread_id == "conv_123456"  # type: ignore[attr-defined] (service_thread_id should be set)
+
+
 async def test_function_approval_mode_executes_tool():
     """Test that function approval with approval_mode='always_require' sends the correct messages."""
     from agent_framework import FunctionResultContent, ai_function

python/packages/ag-ui/tests/test_orchestrators.py

@@ -6,10 +6,13 @@ from collections.abc import AsyncGenerator
 from typing import Any
 from unittest.mock import MagicMock
 
+from ag_ui.core import BaseEvent, RunFinishedEvent
 from agent_framework import (
     AgentResponseUpdate,
+    AgentThread,
     BaseChatClient,
     ChatAgent,
+    ChatResponseUpdate,
     FunctionInvocationConfiguration,
     TextContent,
     ai_function,
@@ -52,7 +55,22 @@ def _create_mock_chat_agent(
     async def mock_run_stream(
         messages: list[Any],
         *,
-        thread: Any = None,
+        #     thread: AgentThread,
+        #     tools: list[Any] | None = None,
+        #     **kwargs: Any,
+        # ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
+        #     self.seen_tools = tools
+        #     yield AgentRunResponseUpdate(
+        #         contents=[TextContent(text="ok")],
+        #         role="assistant",
+        #         response_id=thread.metadata.get("ag_ui_run_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+        #         raw_representation=ChatResponseUpdate(
+        #             contents=[TextContent(text="ok")],
+        #             conversation_id=thread.metadata.get("ag_ui_thread_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+        #             response_id=thread.metadata.get("ag_ui_run_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+        #         ),
+        #     )
+        thread: AgentThread,
         tools: list[Any] | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[AgentResponseUpdate, None]:
@@ -60,7 +78,16 @@ def _create_mock_chat_agent(
             capture_tools.extend(tools)
         if capture_messages is not None:
             capture_messages.extend(messages)
-        yield AgentResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
+        yield AgentResponseUpdate(
+            contents=[TextContent(text="ok")],
+            role="assistant",
+            response_id=thread.metadata.get("ag_ui_run_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+            raw_representation=ChatResponseUpdate(
+                contents=[TextContent(text="ok")],
+                conversation_id=thread.metadata.get("ag_ui_thread_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+                response_id=thread.metadata.get("ag_ui_run_id"),  # type: ignore[attr-defined] (metadata always created in orchestrator)
+            ),
+        )
 
     # Patch the run_stream method
     agent.run_stream = mock_run_stream  # type: ignore[method-assign]
@@ -139,6 +166,7 @@ async def test_default_orchestrator_with_camel_case_ids() -> None:
         events.append(event)
 
     # assert the last event has the expected run_id and thread_id
+    assert isinstance(events[-1], RunFinishedEvent)
     last_event = events[-1]
     assert last_event.run_id == "test-camelcase-runid"
     assert last_event.thread_id == "test-camelcase-threadid"
@@ -167,11 +195,12 @@ async def test_default_orchestrator_with_snake_case_ids() -> None:
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[BaseEvent] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
     # assert the last event has the expected run_id and thread_id
+    assert isinstance(events[-1], RunFinishedEvent)
     last_event = events[-1]
     assert last_event.run_id == "test-snakecase-runid"
     assert last_event.thread_id == "test-snakecase-threadid"

python/packages/ag-ui/tests/test_service_thread_id.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for service-managed thread IDs, and service-generated response ids."""
+
+import sys
+from pathlib import Path
+from typing import Any
+
+from ag_ui.core import RunFinishedEvent, RunStartedEvent
+from agent_framework import TextContent
+from agent_framework._types import AgentResponseUpdate, ChatResponseUpdate
+
+sys.path.insert(0, str(Path(__file__).parent))
+from utils_test_ag_ui import StubAgent
+
+
+async def test_service_thread_id_when_there_are_updates():
+    """Test that service-managed thread IDs (conversation_id) are correctly set as the thread_id in events."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = [
+        AgentResponseUpdate(
+            contents=[TextContent(text="Hello, user!")],
+            response_id="resp_67890",
+            raw_representation=ChatResponseUpdate(
+                contents=[TextContent(text="Hello, user!")],
+                conversation_id="conv_12345",
+                response_id="resp_67890",
+            ),
+        )
+    ]
+    agent = StubAgent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hi"}],
+    }
+
+    events: list[Any] = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].run_id == "resp_67890"
+    assert events[0].thread_id == "conv_12345"
+    assert isinstance(events[-1], RunFinishedEvent)
+
+
+async def test_service_thread_id_when_no_user_message():
+    """Test when user submits no messages, emitted events still have with a thread_id"""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = []
+    agent = StubAgent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data: dict[str, list[dict[str, str]]] = {
+        "messages": [],
+    }
+
+    events: list[Any] = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    assert len(events) == 2
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].thread_id
+    assert isinstance(events[-1], RunFinishedEvent)
+
+
+async def test_service_thread_id_when_user_supplied_thread_id():
+    """Test that user-supplied thread IDs are preserved in emitted events."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = []
+    agent = StubAgent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}], "threadId": "conv_12345"}
+
+    events: list[Any] = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].thread_id == "conv_12345"
+    assert isinstance(events[-1], RunFinishedEvent)