Conversation
+ +Start the conversation with a prompt:
+{message.text}
+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 cd976cad9c..a19320c31e 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 @@ -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], diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index fc887ee919..777450a686 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/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,7 +33,7 @@ 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 @@ -40,10 +41,33 @@ 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) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py index da4360832a..ea291853b0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py +++ b/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) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 3f6e76e5fd..bd3089775e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/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": diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index d378ac0fa4..ab7e609af1 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/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,226 @@ 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_uses_provider_thread_key_when_request_omits_thread_id( + streaming_chat_client_stub: Any, +) -> None: + """A provider fallback ID becomes the lifecycle and snapshot key when AG-UI omits one.""" + app = FastAPI() + 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="Stored reply")], + conversation_id="conv_foundry_123", + response_id="resp_foundry_1", + ) + + 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={"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"]) == ( + "conv_foundry_123", + "resp_foundry_1", + ) + assert (first_events[-1]["threadId"], first_events[-1]["runId"]) == ( + "conv_foundry_123", + "resp_foundry_1", + ) + + hydrate_response = client.post( + "/snapshots", + json={"thread_id": "conv_foundry_123", "run_id": "hydrate-run", "messages": []}, + ) + + assert hydrate_response.status_code == 200 + assert call_count == 1 + hydrated_messages = _latest_messages_snapshot(hydrate_response) + assert any( + message.get("role") == "user" and message.get("content") == "Remember LANTERN-482" + for message in hydrated_messages + ) + assert any( + message.get("role") == "assistant" and message.get("content") == "Stored reply" for message in hydrated_messages + ) + + +async def test_agent_endpoint_correlates_gen_ai_spans_with_supplied_thread_id( + streaming_chat_client_stub: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Agent and chat 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 + provider_conversation_ids: list[str | None] = [] + + async def stream_fn(messages: Any, options: Any, **kwargs: Any) -> AsyncIterator[ChatResponseUpdate]: + nonlocal call_count + del messages, kwargs + call_count += 1 + provider_conversation_ids.append(options.get("conversation_id")) + 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 + + spans_by_operation: dict[str, list[Any]] = {"invoke_agent": [], "chat": []} + for span in exporter.get_finished_spans(): + if span.attributes is None: + continue + operation = span.attributes.get("gen_ai.operation.name") + if isinstance(operation, str) and operation in spans_by_operation: + spans_by_operation[operation].append(span) + + trace_ids_by_operation: dict[str, set[int]] = {} + for operation, spans in spans_by_operation.items(): + assert len(spans) == 2 + trace_ids: set[int] = set() + conversation_ids = [] + for span in 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")) + trace_ids_by_operation[operation] = trace_ids + assert conversation_ids == [ + "ag-ui-thread-1", + "ag-ui-thread-1", + ] + + assert len(trace_ids_by_operation["invoke_agent"]) == 2 + assert trace_ids_by_operation["chat"] == trace_ids_by_operation["invoke_agent"] + for chat_span in spans_by_operation["chat"]: + assert chat_span.context is not None + assert chat_span.parent is not None + matching_agent_span = next( + span + for span in spans_by_operation["invoke_agent"] + if span.context is not None and span.context.trace_id == chat_span.context.trace_id + ) + assert matching_agent_span.context is not None + assert chat_span.parent.span_id == matching_agent_span.context.span_id + assert provider_conversation_ids == [None, None] + + 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() diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py index 413da380d7..d0f7522e1a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -123,6 +123,25 @@ async def test_snapshot_without_state_or_interrupts_replays_messages_only(self) ] +class TestRebindThreadId: + """A late provider fallback becomes the key for subsequent writes.""" + + async def test_save_uses_rebound_thread_id(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="generated-thread") + + session.rebind_thread_id("provider-thread") + await session.save( + messages=[{"id": "m1", "role": "user", "content": "hi"}], + state=None, + interrupt=None, + session_state=None, + ) + + assert await store.get(scope="user-1", thread_id="generated-thread") is None + assert await store.get(scope="user-1", thread_id="provider-thread") is not None + + class TestEffectiveState: """Request values overlay stored values; defaults never reset either.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 8af9a33be2..6eac56864c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -8,9 +8,11 @@ from types import SimpleNamespace from typing import Any, cast -from ag_ui.core import EventType, StateSnapshotEvent +import pytest +from ag_ui.core import EventType, RunFinishedEvent, RunStartedEvent, StateSnapshotEvent from agent_framework import ( Agent, + AgentContext, AgentResponse, AgentResponseUpdate, ChatResponseUpdate, @@ -114,6 +116,121 @@ async def start(message: Any, ctx: WorkflowContext[Any, str]) -> None: assert custom_events[0].value == {"progress": 10} # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] +async def test_workflow_and_agent_spans_use_supplied_agui_thread_id(monkeypatch: pytest.MonkeyPatch) -> None: + """Workflow spans use supplied AG-UI threads without replacing provider fallback behavior.""" + 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 scripted_stream( + 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="provider-conversation", + ) + + async def passthrough_middleware(_context: AgentContext, call_next: Any) -> None: + await call_next() + + participant = Agent( + client=StreamingChatClientStub(scripted_stream), + name="workflow-agent", + middleware=[passthrough_middleware], + ) + workflow = WorkflowBuilder(start_executor=participant, output_from="all").build() + + for run_number in (1, 2): + events = [ + event + async for event in run_workflow_stream( + { + "thread_id": "ag-ui-workflow-thread", + "run_id": f"run-{run_number}", + "messages": [{"role": "user", "content": f"Turn {run_number}"}], + }, + workflow, + ) + ] + run_started = next(event for event in events if isinstance(event, RunStartedEvent)) + run_finished = next(event for event in events if isinstance(event, RunFinishedEvent)) + assert (run_started.thread_id, run_started.run_id) == ( + "ag-ui-workflow-thread", + f"run-{run_number}", + ) + assert (run_finished.thread_id, run_finished.run_id) == ( + "ag-ui-workflow-thread", + f"run-{run_number}", + ) + + spans_by_operation: dict[str, list[Any]] = {"invoke_agent": [], "chat": []} + for span in exporter.get_finished_spans(): + if span.attributes is None: + continue + operation = span.attributes.get("gen_ai.operation.name") + if isinstance(operation, str) and operation in spans_by_operation: + spans_by_operation[operation].append(span) + + for spans in spans_by_operation.values(): + assert len(spans) == 2 + assert [span.attributes.get("gen_ai.conversation.id") for span in spans if span.attributes is not None] == [ + "ag-ui-workflow-thread", + "ag-ui-workflow-thread", + ] + + workflow_spans = [span for span in exporter.get_finished_spans() if span.name == "workflow.run"] + assert len(workflow_spans) == 2 + workflow_conversation_ids = [] + for span in workflow_spans: + assert span.attributes is not None + workflow_conversation_ids.append(span.attributes.get("gen_ai.conversation.id")) + + assert workflow_conversation_ids == [ + "ag-ui-workflow-thread", + "ag-ui-workflow-thread", + ] + + exporter.clear() + events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "Provider fallback"}]}, + workflow, + ) + ] + assert any(isinstance(event, RunFinishedEvent) for event in events) + + fallback_agent_span = next( + span + for span in exporter.get_finished_spans() + if span.attributes is not None and span.attributes.get("gen_ai.operation.name") == "invoke_agent" + ) + assert fallback_agent_span.attributes is not None + assert fallback_agent_span.attributes.get("gen_ai.conversation.id") == "provider-conversation" + + fallback_workflow_span = next(span for span in exporter.get_finished_spans() if span.name == "workflow.run") + assert fallback_workflow_span.attributes is not None + assert "gen_ai.conversation.id" not in fallback_workflow_span.attributes + + async def test_workflow_run_request_info_emits_interrupt_and_resume_works(): """request_info should emit interrupt metadata and resume should continue run.""" diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 20629b76b8..063947d365 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -127,6 +127,29 @@ "inner_accumulated_usage", default=None ) +# Allows protocol adapters to supply an application-managed conversation identity for one execution +# without putting that value into a service-owned continuation field. +_TELEMETRY_CONVERSATION_ID: Final[contextvars.ContextVar[str | None]] = contextvars.ContextVar( + "telemetry_conversation_id", default=None +) + + +@contextlib.contextmanager +def _use_telemetry_conversation_id( # pyright: ignore[reportUnusedFunction] + conversation_id: str | None, +) -> Generator[None]: + """Set an application-managed OTel conversation id for the current execution.""" + if conversation_id is None: + yield + return + + token = _TELEMETRY_CONVERSATION_ID.set(conversation_id) + try: + yield + finally: + _TELEMETRY_CONVERSATION_ID.reset(token) + + OTEL_METRICS: Final[str] = "__otel_metrics__" TOKEN_USAGE_BUCKET_BOUNDARIES: Final[tuple[float, ...]] = ( 1, @@ -1528,6 +1551,10 @@ def get_response( service_url=service_url, **merged_client_kwargs, ) + if (telemetry_conversation_id := _TELEMETRY_CONVERSATION_ID.get()) is not None: + # Keep application-managed telemetry correlation separate from the + # provider-owned conversation_id forwarded through chat options. + attributes[OtelAttr.CONVERSATION_ID] = telemetry_conversation_id if stream: agent_span = trace.get_current_span() @@ -1824,11 +1851,13 @@ def _trace_agent_invocation( "Callable[[AgentSession | None], str | None] | None", getattr(self, "_get_otel_conversation_id", None), ) - conversation_id = ( - get_otel_conversation_id(session) - if callable(get_otel_conversation_id) - else (session.service_session_id if (session and isinstance(session.service_session_id, str)) else None) - ) + conversation_id = _TELEMETRY_CONVERSATION_ID.get() + if conversation_id is None: + conversation_id = ( + get_otel_conversation_id(session) + if callable(get_otel_conversation_id) + else (session.service_session_id if (session and isinstance(session.service_session_id, str)) else None) + ) attributes = _get_span_attributes( operation_name=OtelAttr.AGENT_INVOKE_OPERATION, provider_name=provider_name, @@ -2903,7 +2932,11 @@ def create_workflow_span( kind: trace.SpanKind = trace.SpanKind.INTERNAL, ) -> _AgnosticContextManager[trace.Span]: """Create a generic workflow span.""" - return workflow_tracer().start_as_current_span(name, kind=kind, attributes=attributes) + span_attributes = dict(attributes) if attributes is not None else {} + conversation_id = _TELEMETRY_CONVERSATION_ID.get() + if name == OtelAttr.WORKFLOW_RUN_SPAN and conversation_id is not None: + span_attributes.setdefault(OtelAttr.CONVERSATION_ID, conversation_id) + return workflow_tracer().start_as_current_span(name, kind=kind, attributes=span_attributes or None) def create_processing_span( diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 31f75c78e0..4256e62b2d 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -164,6 +164,10 @@ def mock_chat_client(): """Create a mock chat client for testing.""" class MockChatClient(ChatTelemetryLayer, BaseChatClient[Any]): + def __init__(self) -> None: + super().__init__() + self.observed_options: list[dict[str, Any]] = [] + def service_url(self): return "https://test.example.com" @@ -175,6 +179,7 @@ def _inner_get_response( # pyrefly: ignore[bad-override] options: Mapping[str, Any], **kwargs: Any, # type: ignore[override] ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.observed_options.append(dict(options)) if stream: return self._get_streaming_response(messages=messages, options=options, **kwargs) @@ -207,6 +212,61 @@ def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: return MockChatClient +@pytest.mark.parametrize("stream", [False, True]) +async def test_chat_telemetry_conversation_override_is_scoped_and_telemetry_only( + mock_chat_client: Any, + span_exporter: InMemorySpanExporter, + stream: bool, +) -> None: + """An application conversation id changes telemetry without changing provider options.""" + from agent_framework.observability import ( + _use_telemetry_conversation_id, # pyright: ignore[reportPrivateUsage] + ) + + client = mock_chat_client() + messages = [Message(role="user", contents=["Test message"])] + provider_options = { + "model": "Test", + "conversation_id": "provider-conversation", + "metadata": {"sentinel": "unchanged"}, + } + expected_options = { + "model": "Test", + "conversation_id": "provider-conversation", + "metadata": {"sentinel": "unchanged"}, + } + + async def invoke() -> None: + if stream: + response_stream = client.get_response(messages=messages, stream=True, options=provider_options) + async for _ in response_stream: + pass + await response_stream.get_final_response() + return + await client.get_response(messages=messages, stream=False, options=provider_options) + + span_exporter.clear() + with _use_telemetry_conversation_id("application-thread"): + await invoke() + + assert provider_options == expected_options + assert client.observed_options == [expected_options] + scoped_spans = span_exporter.get_finished_spans() + assert len(scoped_spans) == 1 + assert scoped_spans[0].attributes is not None + assert scoped_spans[0].attributes.get(OtelAttr.CONVERSATION_ID) == "application-thread" + + span_exporter.clear() + await invoke() + + assert provider_options == expected_options + assert client.observed_options == [expected_options, expected_options] + unscoped_spans = span_exporter.get_finished_spans() + assert len(unscoped_spans) == 1 + assert unscoped_spans[0].attributes is not None + assert unscoped_spans[0].attributes.get(OtelAttr.CONVERSATION_ID) != "application-thread" + + @pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True) async def test_chat_client_observability(mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data): """Test that when diagnostics are enabled, telemetry is applied.""" @@ -608,6 +668,36 @@ class MockChatClientAgent(AgentTelemetryLayer, _MockChatClientAgent): # type: i return MockChatClientAgent +async def test_agent_telemetry_conversation_override_is_scoped( + mock_chat_agent: SupportsAgentRun, + span_exporter: InMemorySpanExporter, +) -> None: + """An application-managed conversation id overrides provider continuation for one run only.""" + from agent_framework import AgentSession + from agent_framework.observability import ( + _use_telemetry_conversation_id, # pyright: ignore[reportPrivateUsage] + ) + + agent = mock_chat_agent() # type: ignore[operator] # pyrefly: ignore[not-callable] # ty: ignore[call-non-callable] + session = AgentSession(service_session_id="provider-conversation") + span_exporter.clear() + + with _use_telemetry_conversation_id("application-thread"): + await agent.run("First turn", session=session) + await agent.run("Second turn", session=session) + + spans = span_exporter.get_finished_spans() + conversation_ids = [] + for span in spans: + assert span.attributes is not None + conversation_ids.append(span.attributes.get(OtelAttr.CONVERSATION_ID)) + + assert conversation_ids == [ + "application-thread", + "provider-conversation", + ] + + @pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True) async def test_agent_span_captures_response_telemetry_without_inner_chat_span( mock_chat_agent: SupportsAgentRun, span_exporter: InMemorySpanExporter, enable_sensitive_data @@ -2081,6 +2171,46 @@ def test_create_workflow_span(span_exporter): assert spans[0].attributes["key"] == "value" +def test_create_workflow_span_uses_scoped_conversation_id(span_exporter: InMemorySpanExporter) -> None: + """An ambient conversation id is applied only within its workflow execution scope.""" + from agent_framework.observability import ( + OtelAttr, + _use_telemetry_conversation_id, # pyright: ignore[reportPrivateUsage] + create_workflow_span, + ) + + span_exporter.clear() # type: ignore[attr-defined] + with _use_telemetry_conversation_id("application-thread"): + with create_workflow_span(OtelAttr.WORKFLOW_RUN_SPAN): + pass + with create_workflow_span( + OtelAttr.WORKFLOW_RUN_SPAN, + attributes={OtelAttr.CONVERSATION_ID: "explicit-thread"}, + ): + pass + with create_workflow_span(OtelAttr.MESSAGE_SEND_SPAN): + pass + with create_workflow_span(OtelAttr.WORKFLOW_RUN_SPAN): + pass + + spans = span_exporter.get_finished_spans() # type: ignore[attr-defined] + workflow_spans = [span for span in spans if span.name == OtelAttr.WORKFLOW_RUN_SPAN] + assert len(workflow_spans) == 3 + ambient_attributes = workflow_spans[0].attributes + explicit_attributes = workflow_spans[1].attributes + unscoped_attributes = workflow_spans[2].attributes + assert ambient_attributes is not None + assert explicit_attributes is not None + assert unscoped_attributes is not None + assert ambient_attributes[OtelAttr.CONVERSATION_ID] == "application-thread" + assert explicit_attributes[OtelAttr.CONVERSATION_ID] == "explicit-thread" + assert OtelAttr.CONVERSATION_ID not in unscoped_attributes + message_send_span = next(span for span in spans if span.name == OtelAttr.MESSAGE_SEND_SPAN) + message_send_attributes = message_send_span.attributes + assert message_send_attributes is not None + assert OtelAttr.CONVERSATION_ID not in message_send_attributes + + def test_create_processing_span(span_exporter): """Test create_processing_span creates a span with correct attributes.""" from agent_framework.observability import OtelAttr, create_processing_span diff --git a/python/samples/05-end-to-end/ag_ui_single_agent/README.md b/python/samples/05-end-to-end/ag_ui_single_agent/README.md new file mode 100644 index 0000000000..b1d2bfb1ff --- /dev/null +++ b/python/samples/05-end-to-end/ag_ui_single_agent/README.md @@ -0,0 +1,96 @@ +# AG-UI Single Agent Demo + +The simplest possible AG-UI integration: a **single chat agent** with **no tools** and **no context providers**, +served over the AG-UI protocol and consumed by a small React client. + +Use this sample as the starting point for AG-UI. For a richer, multi-agent example with tool-approval checkpoints +and human-in-the-loop resumes, see [`../ag_ui_workflow_handoff`](../ag_ui_workflow_handoff/README.md). + +## Folder Layout + +- `backend/server.py` - FastAPI + AG-UI endpoint wrapping a single `Agent` +- `frontend/` - Vite + React AG-UI client UI + +## Prerequisites + +- Python 3.10+ +- Node.js 20.19+ or 22.12+ +- npm 9+ +- Azure AI project + model deployment configured in environment variables: + - `FOUNDRY_PROJECT_ENDPOINT` + - `FOUNDRY_MODEL` +- Azure CLI authenticated with `az login` + +## 1) Run Backend + +From the repository root: + +```bash +cd python +uv sync +uv run python samples/05-end-to-end/ag_ui_single_agent/backend/server.py +``` + +Backend default URL: + +- `http://127.0.0.1:8892` +- AG-UI endpoint: `POST http://127.0.0.1:8892/agent` + +To export traces to the Application Insights resource connected to the Foundry project, run the backend with: + +```bash +ENABLE_AZURE_MONITOR=true uv run python samples/05-end-to-end/ag_ui_single_agent/backend/server.py +``` + +Each user turn is a separate run and trace. The stable AG-UI `thread_id` is recorded as +`gen_ai.conversation.id`, which lets Foundry group those turns into one conversation. + +## 2) Install Frontend Packages (npm) + +From the `python/` directory (where Step 1 left you): + +```bash +cd samples/05-end-to-end/ag_ui_single_agent/frontend +npm install +``` + +## 3) Run Frontend Locally + +```bash +npm run dev +``` + +Frontend default URL: + +- `http://127.0.0.1:5173` + +If you changed backend host/port, run with: + +```bash +VITE_BACKEND_URL=http://127.0.0.1:8892 npm run dev +``` + +## 4) Demo Flow to Verify + +1. Click one of the starter prompts (or type your own message). +2. Watch the assistant response stream in token by token. +3. Send a follow-up that depends on the previous turn (for example: "summarize what you just told me"). + The client only sends the newest message plus the `thread_id`; the server replays the stored history. +4. Click **New Thread** to start a fresh conversation (a new `thread_id`). + +## Conversation History + +The client only ever sends the **newest message** plus a `thread_id`. The backend retains history **server-side**, +keyed by that `thread_id`, using an `InMemoryAGUIThreadSnapshotStore`. Because an AG-UI thread id is not an +authorization boundary, a `snapshot_scope_resolver` is required whenever a snapshot store is configured; this +single-tenant demo maps every request to one shared `"demo"` scope. + +The in-memory store is process-local and not durable. Swap in your own `AGUIThreadSnapshotStore` implementation +(and a real scope resolver) for production. + +## What This Validates + +- `add_agent_framework_fastapi_endpoint(...)` with a plain `Agent` (no `AgentFrameworkWorkflow` wrapper) +- Streaming assistant text via `TEXT_MESSAGE_START` / `TEXT_MESSAGE_CONTENT` / `TEXT_MESSAGE_END` AG-UI events +- Server-side conversation history keyed by `thread_id` via a snapshot store +- Foundry trace correlation across runs using the stable AG-UI `thread_id` diff --git a/python/samples/05-end-to-end/ag_ui_single_agent/backend/server.py b/python/samples/05-end-to-end/ag_ui_single_agent/backend/server.py new file mode 100644 index 0000000000..a399a0372b --- /dev/null +++ b/python/samples/05-end-to-end/ag_ui_single_agent/backend/server.py @@ -0,0 +1,148 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "agent-framework-ag-ui", +# "agent-framework-foundry", +# "azure-identity", +# "azure-monitor-opentelemetry", +# "fastapi", +# "python-dotenv", +# "uvicorn", +# ] +# /// + +# Copyright (c) Microsoft. All rights reserved. + +"""AG-UI single-agent demo backend. + +This sample exposes one Foundry-backed Agent over AG-UI and pairs it with the +React frontend in `../frontend`. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT: Microsoft Foundry project endpoint. + FOUNDRY_MODEL: Model deployment name. + ENABLE_AZURE_MONITOR: Set to true to export traces to the project's Application Insights resource. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +import uvicorn +from agent_framework import Agent +from agent_framework.ag_ui import ( + InMemoryAGUIThreadSnapshotStore, + add_agent_framework_fastapi_endpoint, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware + +load_dotenv() + +logger = logging.getLogger(__name__) + + +# 1. Create one Foundry-backed agent with no tools or context providers. +def create_client() -> FoundryChatClient: + """Create the Foundry chat client used by the sample.""" + + return FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["FOUNDRY_MODEL"], + credential=AzureCliCredential(), + ) + + +def create_agent(client: FoundryChatClient) -> Agent: + """Create a single chat agent with no tools and no context providers.""" + + return Agent( + id="assistant", + name="assistant", + instructions="You are a helpful, concise assistant. Answer the user's questions directly.", + client=client, + ) + + +# 2. Configure the AG-UI endpoint, thread history, and optional trace export. +def create_app() -> FastAPI: + """Create and configure the FastAPI application.""" + + client = create_client() + agent = create_agent(client) + + @asynccontextmanager + async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + if os.getenv("ENABLE_AZURE_MONITOR", "false").casefold() in {"1", "true", "yes", "on"}: + await client.configure_azure_monitor() + logger.info("Azure Monitor telemetry export is enabled") + yield + + app = FastAPI(title="AG-UI Single Agent Demo", lifespan=lifespan) + + cors_origins = [ + origin.strip() for origin in os.getenv("CORS_ORIGINS", "http://127.0.0.1:5173").split(",") if origin.strip() + ] + app.add_middleware( + CORSMiddleware, + allow_origins=cors_origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + add_agent_framework_fastapi_endpoint( + app=app, + agent=agent, + path="/agent", + # Persist conversation history server-side, keyed by thread_id, so the + # client only ever sends the newest message plus its thread_id. + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + # AG-UI thread ids are not an authorization boundary, so a scope is required + # when a snapshot store is configured. This demo is single-tenant, so every + # request maps to one shared scope. + snapshot_scope_resolver=lambda _request: "demo", + ) + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + return app + + +app = create_app() + + +# 3. Run the backend for the React frontend. +async def main() -> None: + """Run the AG-UI single-agent demo backend.""" + + logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") + + host = os.getenv("HOST", "127.0.0.1") + port = int(os.getenv("PORT", "8892")) + + print(f"AG-UI single-agent demo backend running at http://{host}:{port}") + print("AG-UI endpoint: POST /agent") + + server = uvicorn.Server(uvicorn.Config(app, host=host, port=port)) + await server.serve() + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +AG-UI single-agent demo backend running at http://127.0.0.1:8892 +AG-UI endpoint: POST /agent +""" diff --git a/python/samples/05-end-to-end/ag_ui_single_agent/frontend/.gitignore b/python/samples/05-end-to-end/ag_ui_single_agent/frontend/.gitignore new file mode 100644 index 0000000000..16c69217c0 --- /dev/null +++ b/python/samples/05-end-to-end/ag_ui_single_agent/frontend/.gitignore @@ -0,0 +1,7 @@ +# dependencies +/node_modules + +# build artifacts +*.tsbuildinfo +vite.config.js +vite.config.d.ts diff --git a/python/samples/05-end-to-end/ag_ui_single_agent/frontend/index.html b/python/samples/05-end-to-end/ag_ui_single_agent/frontend/index.html new file mode 100644 index 0000000000..fa3da04d3e --- /dev/null +++ b/python/samples/05-end-to-end/ag_ui_single_agent/frontend/index.html @@ -0,0 +1,13 @@ + + + +
+ + +Agent Framework ยท AG-UI
++ The simplest AG-UI integration: one chat agent with no tools and no context providers, streamed to a React + client over Server-Sent Events. +
+Start the conversation with a prompt:
+{message.text}
+