Python: Pass client thread_id as session_id when constructing AgentSession in AG-UI (#5384)

* Pass thread_id as session_id when constructing AgentSession in AG-UI

run_agent_stream() was constructing AgentSession without passing the
client's thread_id as session_id, causing every request to receive a
random UUID. This broke session continuity for HistoryProvider
implementations that rely on session_id matching the client's thread_id.

Pass session_id=thread_id in both the service-session and non-service
code paths so the session identity is consistent with the AG-UI client.

Fixes #5357

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add test for service_session with no thread_id edge case (#5357)

When use_service_session=True but no thread_id/threadId is in the payload,
verify session_id is a generated UUID and service_session_id is None.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

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

@@ -790,9 +790,9 @@ 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(service_session_id=supplied_thread_id)
+        session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id)
     else:
-        session = AgentSession()
+        session = AgentSession(session_id=thread_id)
 
     # Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation)
     base_metadata: dict[str, Any] = {

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

@@ -183,6 +183,7 @@ class StubAgent(SupportsAgentRun):
         self.client = client or SimpleNamespace(function_invocation_configuration=None)
         self.messages_received: list[Any] = []
         self.tools_received: list[Any] | None = None
+        self.last_session: AgentSession | None = None
 
     @overload
     def run(
@@ -216,6 +217,7 @@ class StubAgent(SupportsAgentRun):
 
             async def _stream() -> AsyncIterator[AgentResponseUpdate]:
                 self.messages_received = [] if messages is None else list(messages)  # type: ignore[arg-type]
+                self.last_session = session
                 self.tools_received = kwargs.get("tools")
                 for update in self.updates:
                     yield update

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

@@ -1640,3 +1640,115 @@ class TestReasoningInSnapshot:
         # close: MsgEnd(block2) + End(block2)
         assert isinstance(close[0], ReasoningMessageEndEvent)
         assert close[0].message_id == "block2"
+
+
+async def test_session_id_matches_thread_id():
+    """Session created by run_agent_stream uses the client thread_id as session_id."""
+    from conftest import StubAgent
+
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    stub = StubAgent()
+    agent = AgentFrameworkAgent(agent=stub)
+
+    payload = {
+        "thread_id": "my-thread-123",
+        "run_id": "run-1",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    _ = [event async for event in agent.run(payload)]
+
+    assert stub.last_session is not None
+    assert stub.last_session.session_id == "my-thread-123"
+
+
+async def test_session_id_matches_camel_case_thread_id():
+    """Session uses threadId (camelCase) as session_id when snake_case is absent."""
+    from conftest import StubAgent
+
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    stub = StubAgent()
+    agent = AgentFrameworkAgent(agent=stub)
+
+    payload = {
+        "threadId": "camel-thread-456",
+        "run_id": "run-2",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    _ = [event async for event in agent.run(payload)]
+
+    assert stub.last_session is not None
+    assert stub.last_session.session_id == "camel-thread-456"
+
+
+async def test_session_id_matches_thread_id_with_service_session():
+    """Session uses thread_id as session_id even when use_service_session is enabled."""
+    from conftest import StubAgent
+
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    stub = StubAgent()
+    agent = AgentFrameworkAgent(agent=stub, use_service_session=True)
+
+    payload = {
+        "thread_id": "service-thread-789",
+        "run_id": "run-3",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    _ = [event async for event in agent.run(payload)]
+
+    assert stub.last_session is not None
+    assert stub.last_session.session_id == "service-thread-789"
+    assert stub.last_session.service_session_id == "service-thread-789"
+
+
+async def test_session_id_generated_when_no_thread_id():
+    """Session gets a generated UUID as session_id when no thread_id is provided."""
+    import uuid
+
+    from conftest import StubAgent
+
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    stub = StubAgent()
+    agent = AgentFrameworkAgent(agent=stub)
+
+    payload = {
+        "run_id": "run-4",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    _ = [event async for event in agent.run(payload)]
+
+    assert stub.last_session is not None
+    # Should be a valid UUID (auto-generated)
+    uuid.UUID(stub.last_session.session_id)
+
+
+async def test_service_session_no_thread_id_generates_uuid():
+    """With use_service_session=True and no thread_id, session_id is a UUID and service_session_id is None."""
+    import uuid
+
+    from conftest import StubAgent
+
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    stub = StubAgent()
+    agent = AgentFrameworkAgent(agent=stub, use_service_session=True)
+
+    payload = {
+        "run_id": "run-5",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    _ = [event async for event in agent.run(payload)]
+
+    assert stub.last_session is not None
+    # session_id should be a valid auto-generated UUID
+    uuid.UUID(stub.last_session.session_id)
+    # service_session_id should be None since no thread_id was supplied
+    assert stub.last_session.service_session_id is None