diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 2013d0401a..7e10bd4998 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/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 diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index e6fc1cde13..27831f7901 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/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)."""