From bb48154777bee0914b4361e7de0b87de52cd7a57 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 10 Aug 2026 10:18:47 -0700 Subject: [PATCH] Fix FHA session ID traslation --- python/packages/foundry/README.md | 7 +- .../foundry/agent_framework_foundry/_agent.py | 72 ++++++++----------- .../tests/foundry/test_foundry_agent.py | 54 +++++++------- 3 files changed, 62 insertions(+), 71 deletions(-) diff --git a/python/packages/foundry/README.md b/python/packages/foundry/README.md index bcaa83f616..84a574778f 100644 --- a/python/packages/foundry/README.md +++ b/python/packages/foundry/README.md @@ -126,9 +126,10 @@ session = await agent.create_conversation() response = await agent.run("Help me plan a trip to Seattle.", session=session) ``` -This is separate from hosted-agent `isolation_key` sessions: the created -conversation ID is stored on `AgentSession.service_session_id`, while the local -`session_id` remains available for application/session storage. +This is separate from hosted-agent compute sessions. The conversation ID is +stored on `AgentSession.service_session_id`, while a pre-created Foundry +`agent_session_id` is stored in `AgentSession.state`. Reusing a Foundry session +routes requests to the same hosted compute and does not continue a conversation. ## Publishing an agent as a Foundry prompt agent diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 28a840c9e3..fd69e6d339 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -69,6 +69,8 @@ logger: logging.Logger = logging.getLogger("agent_framework.foundry") +_FOUNDRY_AGENT_SESSION_ID_STATE_KEY = "agent_framework_foundry.agent_session_id" + AzureTokenProvider = Callable[[], str | Awaitable[str]] AzureCredentialTypes = TokenCredential | AsyncTokenCredential @@ -128,16 +130,6 @@ def _merge_extra_body(extra_body: Any | None, *, additions: Mapping[str, Any] | return merged -def _uses_foundry_agent_session(conversation_id: Any) -> bool: - """Return whether a conversation_id should be treated as a Foundry agent session id.""" - return ( - isinstance(conversation_id, str) - and bool(conversation_id) - and not conversation_id.startswith("resp_") - and not conversation_id.startswith("conv_") - ) - - 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. @@ -375,23 +367,14 @@ async def _prepare_options( run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"])) # Merge caller-supplied extra_body with any agent-specific request payload. - conversation_id = options.get("conversation_id") extra_body = _merge_extra_body(run_options.pop("extra_body", None)) - if _uses_foundry_agent_session(conversation_id): - run_options.pop("previous_response_id", None) - run_options.pop("conversation", None) - run_options.pop("model", None) - extra_body["agent_session_id"] = conversation_id # Non-preview Prompt/Hosted Agent calls need agent_reference in the request body to # tell the Responses API which Foundry agent (and version) is in use, since ``model`` # is stripped below. The preview path injects the reference via the OpenAI client kwarg # ``agent_name`` instead, so skip there. See issue #5582. if not self.allow_preview: extra_body.setdefault("agent_reference", _build_agent_reference(self.agent_name, self.agent_version)) - should_strip_model = _uses_foundry_agent_session(conversation_id) or ( - conversation_id is None and not options.get("model") - ) - if should_strip_model: + if not options.get("model"): run_options.pop("model", None) if extra_body: run_options["extra_body"] = extra_body @@ -425,10 +408,7 @@ def _parse_response_from_openai( response: Any, options: dict[str, Any], ) -> Any: - parsed_response = super()._parse_response_from_openai(response, options) - if _uses_foundry_agent_session(options.get("conversation_id")): - parsed_response.conversation_id = None - return parsed_response + return super()._parse_response_from_openai(response, options) @override def _parse_chunk_from_openai( @@ -438,7 +418,7 @@ def _parse_chunk_from_openai( function_call_ids: dict[int, tuple[str, str]], seen_reasoning_delta_item_ids: set[str] | None = None, ) -> ChatResponseUpdate: - """Parse streaming events while preserving hosted-agent session state.""" + """Parse streaming events, including Foundry OAuth consent events.""" update = try_parse_oauth_consent_event(event, self.model) if update is None: update = super()._parse_chunk_from_openai( @@ -447,8 +427,6 @@ def _parse_chunk_from_openai( function_call_ids, seen_reasoning_delta_item_ids, ) - if _uses_foundry_agent_session(options.get("conversation_id")): - update.conversation_id = None return update @override @@ -752,7 +730,7 @@ def __init__( additional_properties=dict(additional_properties) if additional_properties is not None else None, ) - def _resolve_service_session_isolation_key(self, isolation_key: str | None = None) -> str: + def _resolve_agent_session_isolation_key(self, isolation_key: str | None = None) -> str: """Resolve the isolation key from an explicit value or default_options.""" resolved_isolation_key = ( isolation_key if isolation_key is not None else self.default_options.get("isolation_key") @@ -761,20 +739,20 @@ def _resolve_service_session_isolation_key(self, isolation_key: str | None = Non raise ValueError("isolation_key is required. Pass it explicitly or set default_options['isolation_key'].") return resolved_isolation_key - async def _create_service_session_id( + async def _create_agent_session_id( self, *, isolation_key: str | None = None, ) -> str: - """Create a hosted Foundry service session and return the service session ID.""" + """Create a hosted Foundry compute session and return its agent session ID.""" if not isinstance(self.client, RawFoundryAgentChatClient): - raise TypeError("_create_service_session_id requires a RawFoundryAgentChatClient-based client.") + raise TypeError("_create_agent_session_id requires a RawFoundryAgentChatClient-based client.") if not self.client.allow_preview: raise RuntimeError("Hosted Foundry service sessions require allow_preview=True.") create_session_kwargs: dict[str, Any] = { "agent_name": self.client.agent_name, - "isolation_key": self._resolve_service_session_isolation_key(isolation_key), + "isolation_key": self._resolve_agent_session_isolation_key(isolation_key), } if version := await self.client.get_agent_version(): from azure.ai.projects.models import VersionRefIndicator @@ -831,14 +809,20 @@ async def _prepare_run_context( **{key: value for key, value in runtime_options.items() if value is not None}, } - if ( - session is not None - and session.service_session_id is None - and effective_options.get("isolation_key") is not None - ): - session.service_session_id = await self._create_service_session_id( - isolation_key=cast(str | None, effective_options.get("isolation_key")), - ) + if session is not None: + agent_session_id = session.state.get(_FOUNDRY_AGENT_SESSION_ID_STATE_KEY) + if agent_session_id is None and effective_options.get("isolation_key") is not None: + agent_session_id = await self._create_agent_session_id( + isolation_key=cast(str | None, effective_options.get("isolation_key")), + ) + session.state[_FOUNDRY_AGENT_SESSION_ID_STATE_KEY] = agent_session_id + if agent_session_id is not None: + if not isinstance(agent_session_id, str) or not agent_session_id: + raise ValueError("Foundry agent session state must contain a non-empty string ID.") + runtime_options["extra_body"] = _merge_extra_body( + effective_options.get("extra_body"), + additions={"agent_session_id": agent_session_id}, + ) return await super()._prepare_run_context( messages=messages, @@ -999,10 +983,12 @@ def __init__( already be configured for preview APIs before being passed to ``FoundryAgent``. - To lazily create HostedAgent service sessions inside the agent, pass an + To lazily create a HostedAgent compute session inside the agent, pass an ``isolation_key`` through ``default_options`` (or per-run options). The - agent stores the resulting HostedAgent session ID in - ``AgentSession.service_session_id`` and reuses it on subsequent runs. + agent stores the resulting Foundry ``agent_session_id`` in + ``AgentSession.state`` and reuses it on subsequent runs. This is separate + from ``AgentSession.service_session_id``, which tracks conversation + continuation. Keyword Args: project_endpoint: The Foundry project endpoint URL. diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index e6fc1cde13..b5ff3d839d 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -537,8 +537,8 @@ async def test_raw_foundry_agent_chat_client_prepare_options_no_tool_warning_whe assert not any("cannot be sent when an agent is specified" in record.message for record in caplog.records) -async def test_raw_foundry_agent_chat_client_prepare_options_strips_model_for_hosted_session() -> None: - """Test that model is stripped when using a hosted agent session (not a PromptAgent).""" +async def test_raw_foundry_agent_chat_client_preserves_separate_compute_and_conversation_ids() -> None: + """Test that compute-session and conversation IDs use separate request fields.""" mock_project = MagicMock() mock_openai = MagicMock() @@ -555,15 +555,16 @@ async def test_raw_foundry_agent_chat_client_prepare_options_strips_model_for_ho return_value={ "model": "gpt-4.1", "previous_response_id": "resp_abc", + "extra_body": {"agent_session_id": "agent-session-123"}, }, ): result = await client._prepare_options( messages=[Message(role="user", contents="hi")], - options={"conversation_id": "agent-session-123"}, + options={"conversation_id": "resp_abc"}, ) assert "model" not in result - assert "previous_response_id" not in result + assert result["previous_response_id"] == "resp_abc" assert result["extra_body"]["agent_session_id"] == "agent-session-123" assert result["extra_body"]["agent_reference"] == {"name": "test-agent", "type": "agent_reference"} @@ -745,8 +746,8 @@ async def test_raw_foundry_agent_chat_client_prepare_options_preserves_model_for assert "extra_body" not in result or "agent_session_id" not in result.get("extra_body", {}) -async def test_raw_foundry_agent_chat_client_prepare_options_maps_agent_session_id_to_extra_body() -> None: - """Test that service_session_id is forwarded as agent_session_id for hosted sessions.""" +async def test_raw_foundry_agent_chat_client_prepare_options_preserves_conversation_continuation() -> None: + """Test that conversation IDs are not treated as hosted-agent compute session IDs.""" mock_project = MagicMock() mock_openai = MagicMock() @@ -767,21 +768,19 @@ async def test_raw_foundry_agent_chat_client_prepare_options_maps_agent_session_ ): result = await client._prepare_options( messages=[Message(role="user", contents="hi")], - options={"conversation_id": "agent-session-123", "isolation_key": "iso-key"}, + options={"conversation_id": "conversation-123", "isolation_key": "iso-key"}, ) assert result["extra_body"] == { "custom": "value", - "agent_session_id": "agent-session-123", "agent_reference": {"name": "test-agent", "type": "agent_reference"}, } - assert "previous_response_id" not in result - assert "conversation" not in result + assert result["previous_response_id"] == "should-be-removed" assert "isolation_key" not in result -def test_raw_foundry_agent_chat_client_parse_response_suppresses_conversation_id_for_agent_sessions() -> None: - """Test that agent-session continuations do not overwrite session.service_session_id.""" +def test_raw_foundry_agent_chat_client_parse_response_preserves_conversation_id() -> None: + """Test that response conversation IDs remain available for framework continuation.""" mock_project = MagicMock() mock_project.get_openai_client.return_value = MagicMock() @@ -798,14 +797,14 @@ def test_raw_foundry_agent_chat_client_parse_response_suppresses_conversation_id ): result = client._parse_response_from_openai( response=MagicMock(), - options={"conversation_id": "agent-session-123"}, + options={"conversation_id": "conversation-123"}, ) - assert result.conversation_id is None + assert result.conversation_id == "resp_123" -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.""" +def test_raw_foundry_agent_chat_client_parse_chunk_preserves_conversation_id() -> None: + """Test that streaming conversation IDs remain available for framework continuation.""" mock_project = MagicMock() mock_project.get_openai_client.return_value = MagicMock() @@ -822,11 +821,11 @@ def test_raw_foundry_agent_chat_client_parse_chunk_suppresses_conversation_id_fo ): result = client._parse_chunk_from_openai( event=MagicMock(type="response.output_text.delta"), - options={"conversation_id": "agent-session-123"}, + options={"conversation_id": "conversation-123"}, function_call_ids={}, ) - assert result.conversation_id is None + assert result.conversation_id == "resp_123" def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None: @@ -1051,8 +1050,8 @@ def my_func() -> str: assert agent.default_options.get("tools") is not None -async def test_raw_foundry_agent_prepare_run_context_creates_service_session_from_isolation_key() -> None: - """Test that RawFoundryAgent lazily creates a hosted session and stores it on service_session_id.""" +async def test_raw_foundry_agent_prepare_run_context_creates_agent_session_from_isolation_key() -> None: + """Test that a compute session is stored in state without replacing the conversation ID.""" mock_project = MagicMock() mock_project.get_openai_client.return_value = MagicMock() @@ -1069,7 +1068,7 @@ async def test_raw_foundry_agent_prepare_run_context_creates_service_session_fro agent_version="1.0", allow_preview=True, ) - session = AgentSession() + session = AgentSession(service_session_id="conv_123") with patch( "agent_framework._agents.RawAgent._prepare_run_context", @@ -1087,16 +1086,21 @@ async def test_raw_foundry_agent_prepare_run_context_creates_service_session_fro ) assert result == {"ok": True} - assert session.service_session_id == "agent-session-123" + assert session.service_session_id == "conv_123" + assert session.state["agent_framework_foundry.agent_session_id"] == "agent-session-123" mock_project.beta.agents.create_session.assert_awaited_once() create_session_kwargs = mock_project.beta.agents.create_session.await_args.kwargs assert create_session_kwargs["agent_name"] == "test-agent" assert create_session_kwargs["isolation_key"] == "iso-key" assert "version_indicator" in create_session_kwargs + assert mock_prepare_run_context.await_args.kwargs["options"] == { + "isolation_key": "iso-key", + "extra_body": {"agent_session_id": "agent-session-123"}, + } mock_prepare_run_context.assert_awaited_once() -async def test_raw_foundry_agent_create_service_session_uses_stable_agents_operations() -> None: +async def test_raw_foundry_agent_create_agent_session_uses_stable_agents_operations() -> None: """Test that hosted sessions use the stable agents operations when available.""" create_session = AsyncMock(return_value=SimpleNamespace(agent_session_id="agent-session-123")) @@ -1111,7 +1115,7 @@ async def test_raw_foundry_agent_create_service_session_uses_stable_agents_opera allow_preview=True, ) - result = await agent._create_service_session_id(isolation_key="iso-key") + result = await agent._create_agent_session_id(isolation_key="iso-key") assert result == "agent-session-123" create_session.assert_awaited_once() @@ -1122,7 +1126,7 @@ async def test_raw_foundry_agent_create_service_session_uses_stable_agents_opera assert "version_indicator" in create_session_kwargs -async def test_raw_foundry_agent_prepare_run_context_requires_preview_for_hosted_sessions() -> None: +async def test_raw_foundry_agent_prepare_run_context_requires_preview_for_agent_sessions() -> None: """Test that hosted-agent sessions require allow_preview=True.""" mock_project = MagicMock()