Python: Persist Foundry hosted agent session IDs

python/packages/foundry/agent_framework_foundry/_agent.py

@@ -138,6 +138,20 @@ def _uses_foundry_agent_session(conversation_id: Any) -> bool:
     )
 
 
+def _get_foundry_agent_session_id(response: Any) -> str | None:
+    """Return a hosted agent session id from a Foundry response payload."""
+    agent_session_id = getattr(response, "agent_session_id", None)
+    if isinstance(agent_session_id, str) and agent_session_id:
+        return agent_session_id
+
+    session = getattr(response, "session", None)
+    if isinstance(session, Mapping):
+        session_id = cast(Mapping[str, Any], session).get("id")
+    else:
+        session_id = getattr(session, "id", None)
+    return session_id if isinstance(session_id, str) and session_id else None
+
+
 def _build_agent_reference(agent_name: str, agent_version: str | None) -> dict[str, str]:
     """Build the Responses API ``agent_reference`` payload for non-preview Foundry agent calls.
 
@@ -428,6 +442,8 @@ def _parse_response_from_openai(
         parsed_response = super()._parse_response_from_openai(response, options)
         if _uses_foundry_agent_session(options.get("conversation_id")):
             parsed_response.conversation_id = None
+        elif options.get("store") is not False and (agent_session_id := _get_foundry_agent_session_id(response)):
+            parsed_response.conversation_id = agent_session_id
         return parsed_response
 
     @override
@@ -449,6 +465,10 @@ def _parse_chunk_from_openai(
             )
         if _uses_foundry_agent_session(options.get("conversation_id")):
             update.conversation_id = None
+        elif options.get("store") is not False and (
+            agent_session_id := _get_foundry_agent_session_id(getattr(event, "response", None))
+        ):
+            update.conversation_id = agent_session_id
         return update
 
     @override

python/packages/foundry/tests/foundry/test_foundry_agent.py

@@ -37,6 +37,7 @@
 from azure.identity import AzureCliCredential
 from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential
 from openai import AsyncOpenAI
+from openai.types.responses.response import Response as OpenAIResponse
 
 from agent_framework_foundry._agent import (
     FoundryAgent,
@@ -804,6 +805,74 @@ def test_raw_foundry_agent_chat_client_parse_response_suppresses_conversation_id
     assert result.conversation_id is None
 
 
+def _openai_response_with_session_id(session_id: str) -> OpenAIResponse:
+    return OpenAIResponse.model_validate({
+        "id": "resp_123",
+        "created_at": 0,
+        "model": "test-model",
+        "object": "response",
+        "output": [],
+        "parallel_tool_calls": False,
+        "tool_choice": "auto",
+        "tools": [],
+        "session": {"id": session_id},
+    })
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        SimpleNamespace(agent_session_id="agent-session-123"),
+        SimpleNamespace(session=SimpleNamespace(id="agent-session-123")),
+        _openai_response_with_session_id("agent-session-123"),
+        SimpleNamespace(agent_session_id="agent-session-123", session=SimpleNamespace(id="session-id-should-not-win")),
+    ],
+)
+def test_raw_foundry_agent_chat_client_parse_response_uses_hosted_agent_session_id(response: Any) -> None:
+    """Test that first-turn hosted agent responses persist the service session id."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    parsed = ChatResponse(conversation_id="resp_123")
+    with patch(
+        "agent_framework_openai._chat_client.RawOpenAIChatClient._parse_response_from_openai",
+        return_value=parsed,
+    ):
+        result = client._parse_response_from_openai(response=response, options={})
+
+    assert result.conversation_id == "agent-session-123"
+
+
+def test_raw_foundry_agent_chat_client_parse_response_respects_store_false_for_agent_session_id() -> None:
+    """Test that hosted agent session ids are not promoted when storage is disabled."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    parsed = ChatResponse(conversation_id=None)
+    with patch(
+        "agent_framework_openai._chat_client.RawOpenAIChatClient._parse_response_from_openai",
+        return_value=parsed,
+    ):
+        result = client._parse_response_from_openai(
+            response=SimpleNamespace(agent_session_id="agent-session-123"),
+            options={"store": False},
+        )
+
+    assert result.conversation_id is None
+
+
 def test_raw_foundry_agent_chat_client_parse_chunk_suppresses_conversation_id_for_agent_sessions() -> None:
     """Test that agent-session stream updates do not overwrite session.service_session_id."""
 
@@ -829,6 +898,69 @@ def test_raw_foundry_agent_chat_client_parse_chunk_suppresses_conversation_id_fo
     assert result.conversation_id is None
 
 
+@pytest.mark.parametrize(
+    "response",
+    [
+        SimpleNamespace(agent_session_id="agent-session-123"),
+        SimpleNamespace(session=SimpleNamespace(id="agent-session-123")),
+        _openai_response_with_session_id("agent-session-123"),
+        SimpleNamespace(agent_session_id="agent-session-123", session=SimpleNamespace(id="session-id-should-not-win")),
+    ],
+)
+def test_raw_foundry_agent_chat_client_parse_chunk_uses_hosted_agent_session_id(response: Any) -> None:
+    """Test that first-turn hosted agent stream updates persist the service session id."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    event = MagicMock(type="response.created")
+    event.response = response
+    parsed = ChatResponseUpdate(conversation_id="resp_123")
+    with patch(
+        "agent_framework_openai._chat_client.RawOpenAIChatClient._parse_chunk_from_openai",
+        return_value=parsed,
+    ):
+        result = client._parse_chunk_from_openai(
+            event=event,
+            options={},
+            function_call_ids={},
+        )
+
+    assert result.conversation_id == "agent-session-123"
+
+
+def test_raw_foundry_agent_chat_client_parse_chunk_respects_store_false_for_agent_session_id() -> None:
+    """Test that hosted agent session stream ids are not promoted when storage is disabled."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    event = MagicMock(type="response.created")
+    event.response = SimpleNamespace(agent_session_id="agent-session-123")
+    parsed = ChatResponseUpdate(conversation_id=None)
+    with patch(
+        "agent_framework_openai._chat_client.RawOpenAIChatClient._parse_chunk_from_openai",
+        return_value=parsed,
+    ):
+        result = client._parse_chunk_from_openai(
+            event=event,
+            options={"store": False},
+            function_call_ids={},
+        )
+
+    assert result.conversation_id is None
+
+
 def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None:
     """Test that _check_model_presence does nothing (model is on service)."""