Python: Fix AG-UI conversation correlation across runs

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

@@ -11,6 +11,7 @@
 from collections import OrderedDict
 from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
 from dataclasses import dataclass, field
+from functools import partial
 from typing import TYPE_CHECKING, Any, TypedDict, cast
 
 from ag_ui.core import (
@@ -50,6 +51,9 @@
 )
 from agent_framework._types import ResponseStream
 from agent_framework.exceptions import AgentInvalidResponseException
+from agent_framework.observability import (
+    _use_telemetry_conversation_id,  # pyright: ignore[reportPrivateUsage]
+)
 
 from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id
 from ._message_adapters import normalize_agui_input_messages
@@ -66,6 +70,7 @@
     _extract_resume_payload,  # type: ignore
     _extract_tool_result_display,  # type: ignore
     _has_only_tool_calls,  # type: ignore
+    _iterate_with_context,  # type: ignore
     _normalize_resume_interrupts,  # type: ignore
     _reconstruct_messages_from_thread_snapshot,  # type: ignore
     _resume_contract_error,  # type: ignore
@@ -2154,8 +2159,10 @@ async def run_agent_stream(
         AG-UI events
     """
     # Parse IDs
-    thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4())
-    run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4())
+    supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId")
+    supplied_run_id = input_data.get("run_id") or input_data.get("runId")
+    thread_id = supplied_thread_id or str(uuid.uuid4())
+    run_id = supplied_run_id or str(uuid.uuid4())
     snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY))
     approval_scope = cast(str | None, input_data.get(_APPROVAL_SCOPE_INPUT_KEY))
     approval_thread_id = approval_state_thread_id(scope=approval_scope, thread_id=thread_id)
@@ -2280,7 +2287,6 @@ async def run_agent_stream(
 
     # Create session (with service session support)
     if config.use_service_session:
-        supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId")
         session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id)
     else:
         session = AgentSession(session_id=thread_id)
@@ -2390,23 +2396,34 @@ async def run_agent_stream(
 
     # Stream from agent - emit RunStarted after first update to get service IDs
     run_started_emitted = False
+    provider_thread_id: str | None = None
     all_updates: list[Any] = []  # Collect for structured output processing
     latest_state_snapshot: dict[str, Any] | None = (
         cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None
     )
-    response_stream = agent.run(messages, stream=True, **run_kwargs)
-    stream = await _normalize_response_stream(response_stream)
-    async for update in stream:
+    # Agent middleware can defer the inner run until streaming begins, so the
+    # telemetry override must cover construction, stream resolution, and every pull.
+    telemetry_conversation_id = str(supplied_thread_id) if supplied_thread_id is not None else None
+    telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id)
+    with telemetry_context():
+        response_stream = agent.run(messages, stream=True, **run_kwargs)
+        stream = await _normalize_response_stream(response_stream)
+
+    async for update in _iterate_with_context(stream, telemetry_context):
         # Collect updates for structured output processing
         if response_format is not None:
             all_updates.append(update)
 
-        # Update IDs from service response on first update and emit RunStarted
+        # Use service-generated IDs only when the AG-UI request omitted them. Client-supplied
+        # IDs remain authoritative for lifecycle correlation and thread-scoped persistence.
         if not run_started_emitted:
             conv_id = get_conversation_id_from_update(update)
             if conv_id:
+                provider_thread_id = conv_id
+            if supplied_thread_id is None and conv_id:
                 thread_id = conv_id
-            if update.response_id:
+                snapshot_session.rebind_thread_id(thread_id)
+            if supplied_run_id is None and update.response_id:
                 run_id = update.response_id
             # NOW emit RunStarted with proper IDs
             yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
@@ -2446,7 +2463,10 @@ async def run_agent_stream(
             if content_type == "function_approval_request" and pending_approvals is not None:
                 if content.id and content.function_call and content.function_call.name:
                     canonical_interrupt_id = content.function_call.call_id or content.id
-                    provider_approval_thread_id = approval_state_thread_id(scope=approval_scope, thread_id=thread_id)
+                    provider_approval_thread_id = approval_state_thread_id(
+                        scope=approval_scope,
+                        thread_id=provider_thread_id or thread_id,
+                    )
                     _register_pending_approval(
                         pending_approvals,
                         [approval_thread_id, provider_approval_thread_id],

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

@@ -7,9 +7,10 @@
 import copy
 import json
 import logging
-from collections.abc import Mapping
+from collections.abc import AsyncGenerator, AsyncIterable, Callable, Mapping
+from contextlib import AbstractContextManager
 from dataclasses import dataclass, field
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 from ag_ui.core import (
     BaseEvent,
@@ -32,18 +33,41 @@
     ToolCallResultEvent,
     ToolCallStartEvent,
 )
-from agent_framework import Content
+from agent_framework import Content, ResponseStream
 
 from ._predictive_state import PredictiveStateHandler
 from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY
 from ._utils import generate_event_id, make_json_safe, normalize_agui_role
 
 logger = logging.getLogger(__name__)
 
+_StreamItemT = TypeVar("_StreamItemT")
+
 # Sentinel for an unset display_result; distinguishes "caller didn't pass" from None/{}/"".
 _UNSET = object()
 
 
+async def _iterate_with_context(
+    stream: AsyncIterable[_StreamItemT],
+    context_factory: Callable[[], AbstractContextManager[Any]],
+) -> AsyncGenerator[_StreamItemT]:
+    """Advance a response stream with a fresh execution context for every pull."""
+    if isinstance(stream, ResponseStream):
+        stream.with_pull_context_manager(context_factory)
+        async for item in stream:
+            yield item
+        return
+
+    stream_iterator = aiter(stream)
+    while True:
+        with context_factory():
+            try:
+                item = await anext(stream_iterator)
+            except StopAsyncIteration:
+                return
+        yield item
+
+
 def _has_only_tool_calls(contents: list[Any]) -> bool:
     """Check if contents have only tool calls (no text)."""
     has_tool_call = any(getattr(c, "type", None) == "function_call" for c in contents)

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

@@ -86,6 +86,14 @@ def stored(self) -> AGUIThreadSnapshot | None:
         """The snapshot loaded at open, or ``None``."""
         return self._stored
 
+    def rebind_thread_id(self, thread_id: str) -> None:
+        """Use a provider-resolved fallback ID for subsequent snapshot operations.
+
+        Runners call this only when the request omitted its AG-UI Thread ID and
+        the provider supplies the lifecycle fallback after the session opened.
+        """
+        self._thread_id = thread_id
+
     async def hydrate_events(self, *, run_id: str) -> AsyncGenerator[BaseEvent]:
         """Replay the stored snapshot as a complete run without invoking the agent."""
         yield RunStartedEvent(run_id=run_id, thread_id=self._thread_id)

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

@@ -9,6 +9,7 @@
 import logging
 import uuid
 from collections.abc import AsyncGenerator
+from functools import partial
 from typing import Any, cast, get_args, get_origin
 
 from ag_ui.core import (
@@ -25,6 +26,9 @@
     ToolCallStartEvent,
 )
 from agent_framework import AgentResponse, AgentResponseUpdate, Content, Message, Workflow, WorkflowRunState
+from agent_framework.observability import (
+    _use_telemetry_conversation_id,  # pyright: ignore[reportPrivateUsage]
+)
 
 from ._message_adapters import normalize_agui_input_messages
 from ._run_common import (
@@ -33,6 +37,7 @@
     _close_reasoning_block,
     _emit_content,
     _extract_resume_payload,
+    _iterate_with_context,
     _normalize_resume_interrupts,
     _resume_contract_error,
 )
@@ -777,7 +782,8 @@ async def run_workflow_stream(
     workflow: Workflow,
 ) -> AsyncGenerator[BaseEvent]:
     """Run a Workflow and emit AG-UI protocol events."""
-    thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4())
+    supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId")
+    thread_id = supplied_thread_id or str(uuid.uuid4())
     run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4())
     available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts")
     if available_interrupts:
@@ -890,12 +896,15 @@ def _drain_open_message() -> list[TextMessageEndEvent]:
             fwd_kwargs = {}
 
     try:
-        if responses:
-            event_stream = workflow.run(responses=responses, stream=True, **fwd_kwargs)
-        else:
-            event_stream = workflow.run(message=messages, stream=True, **fwd_kwargs)
-
-        async for event in event_stream:
+        telemetry_conversation_id = str(supplied_thread_id) if supplied_thread_id is not None else None
+        telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id)
+        with telemetry_context():
+            if responses:
+                event_stream = workflow.run(responses=responses, stream=True, **fwd_kwargs)
+            else:
+                event_stream = workflow.run(message=messages, stream=True, **fwd_kwargs)
+
+        async for event in _iterate_with_context(event_stream, telemetry_context):
             event_type = getattr(event, "type", None)
 
             if event_type == "started":