Python: Fix AG-UI conversation correlation across runs

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

@@ -10,6 +10,7 @@
 import uuid
 from collections import OrderedDict
 from collections.abc import AsyncIterable, Awaitable, Mapping
+from functools import partial
 from typing import TYPE_CHECKING, Any, TypedDict, cast
 
 from ag_ui.core import (
@@ -48,6 +49,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
@@ -63,6 +67,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
@@ -1815,8 +1820,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)
@@ -1959,7 +1966,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)
@@ -2062,23 +2068,33 @@ 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:
+            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)
@@ -2118,7 +2134,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 ._orchestration._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/_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":

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -15,6 +15,7 @@
 from ag_ui.core import MessagesSnapshotEvent, RunStartedEvent, StateSnapshotEvent
 from agent_framework import (
     Agent,
+    AgentContext,
     AgentResponseUpdate,
     AgentSession,
     ChatResponseUpdate,
@@ -3724,6 +3725,145 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any):
     assert state_snapshots[0]["snapshot"] == {"recipe": "pasta"}
 
 
+async def test_agent_endpoint_keeps_request_thread_key_when_provider_returns_conversation_id(
+    streaming_chat_client_stub: Any,
+) -> None:
+    """A provider conversation id must not move snapshots away from the requested AG-UI thread."""
+    app = FastAPI()
+    captured_messages: list[list[tuple[str, str]]] = []
+
+    async def stream_fn(messages: Any, options: Any, **kwargs: Any) -> AsyncIterator[ChatResponseUpdate]:
+        del options, kwargs
+        captured_messages.append([(message.role, message.text) for message in messages])
+        yield ChatResponseUpdate(
+            contents=[Content.from_text(text=f"Reply {len(captured_messages)}")],
+            conversation_id="conv_foundry_123",
+            response_id=f"resp_foundry_{len(captured_messages)}",
+        )
+
+    agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn))
+    store = InMemoryAGUIThreadSnapshotStore()
+    add_agent_framework_fastapi_endpoint(
+        app,
+        agent,
+        path="/snapshots",
+        snapshot_store=store,
+        snapshot_scope_resolver=lambda _request: "tenant-a",
+    )
+    client = TestClient(app)
+
+    first_response = client.post(
+        "/snapshots",
+        json={
+            "thread_id": "ag-ui-thread-1",
+            "run_id": "run-1",
+            "messages": [{"id": "user-1", "role": "user", "content": "Remember LANTERN-482"}],
+        },
+    )
+    assert first_response.status_code == 200
+    first_events = _decode_sse_events(first_response)
+    assert (first_events[0]["threadId"], first_events[0]["runId"]) == ("ag-ui-thread-1", "run-1")
+    assert (first_events[-1]["threadId"], first_events[-1]["runId"]) == ("ag-ui-thread-1", "run-1")
+
+    second_response = client.post(
+        "/snapshots",
+        json={
+            "thread_id": "ag-ui-thread-1",
+            "run_id": "run-2",
+            "messages": [{"id": "user-2", "role": "user", "content": "What token?"}],
+        },
+    )
+
+    assert second_response.status_code == 200
+    second_events = _decode_sse_events(second_response)
+    assert (second_events[0]["threadId"], second_events[0]["runId"]) == ("ag-ui-thread-1", "run-2")
+    assert (second_events[-1]["threadId"], second_events[-1]["runId"]) == ("ag-ui-thread-1", "run-2")
+    assert captured_messages[1] == [
+        ("user", "Remember LANTERN-482"),
+        ("assistant", "Reply 1"),
+        ("user", "What token?"),
+    ]
+
+
+async def test_agent_endpoint_correlates_agent_spans_with_supplied_thread_id(
+    streaming_chat_client_stub: Any,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Agent spans use the stable AG-UI thread id as their OTel conversation id."""
+    from types import SimpleNamespace
+
+    import agent_framework.observability as observability
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+    exporter = InMemorySpanExporter()
+    tracer_provider = TracerProvider()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
+    monkeypatch.setattr(
+        observability,
+        "OBSERVABILITY_SETTINGS",
+        SimpleNamespace(ENABLED=True, SENSITIVE_DATA_ENABLED=False),
+    )
+    monkeypatch.setattr(observability, "get_tracer", lambda *args, **kwargs: tracer_provider.get_tracer("test"))
+
+    call_count = 0
+
+    async def stream_fn(messages: Any, options: Any, **kwargs: Any) -> AsyncIterator[ChatResponseUpdate]:
+        nonlocal call_count
+        del messages, options, kwargs
+        call_count += 1
+        yield ChatResponseUpdate(
+            contents=[Content.from_text(text=f"Reply {call_count}")],
+            conversation_id=f"resp_foundry_{call_count}",
+        )
+
+    app = FastAPI()
+
+    async def passthrough_middleware(_context: AgentContext, call_next: Any) -> None:
+        await call_next()
+
+    agent = Agent(
+        name="test",
+        instructions="Test agent",
+        client=streaming_chat_client_stub(stream_fn),
+        middleware=[passthrough_middleware],
+    )
+    add_agent_framework_fastapi_endpoint(app, agent, path="/agent")
+    client = TestClient(app)
+
+    for run_number in (1, 2):
+        response = client.post(
+            "/agent",
+            json={
+                "thread_id": "ag-ui-thread-1",
+                "run_id": f"run-{run_number}",
+                "messages": [{"role": "user", "content": f"Turn {run_number}"}],
+            },
+        )
+        assert response.status_code == 200
+
+    agent_spans = []
+    for span in exporter.get_finished_spans():
+        if span.attributes is not None and span.attributes.get("gen_ai.operation.name") == "invoke_agent":
+            agent_spans.append(span)
+
+    assert len(agent_spans) == 2
+    trace_ids: set[int] = set()
+    conversation_ids = []
+    for span in agent_spans:
+        assert span.context is not None
+        assert span.attributes is not None
+        trace_ids.add(span.context.trace_id)
+        conversation_ids.append(span.attributes.get("gen_ai.conversation.id"))
+
+    assert len(trace_ids) == 2
+    assert conversation_ids == [
+        "ag-ui-thread-1",
+        "ag-ui-thread-1",
+    ]
+
+
 async def test_agent_endpoint_deduplicates_full_history_and_merges_fresh_state(streaming_chat_client_stub):
     """Stored prior history is authoritative while incoming full history and fresh state remain supported."""
     app = FastAPI()