Python: update FoundryAgent for hosted agent sessions (#5447)

* fixes to FoundryAgent to connect to new hosted agents

Co-authored-by: Copilot <copilot@github.com>

* fix mypy

Co-authored-by: Copilot <copilot@github.com>

* Python: remove Foundry service session helpers

Remove the public hosted-agent service session CRUD helpers from FoundryAgent and drop the related feature-stage inventory entry.

Update the hosted-agent sample to create and delete service sessions directly through the preview AIProjectClient APIs, and tighten a few test harnesses surfaced by full workspace validation.

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

* fix from merge

* fix hosted env detection

Co-authored-by: Copilot <copilot@github.com>

* reverted sample update

* fix tests and code

Co-authored-by: Copilot <copilot@github.com>

* remove aenter

* skipping some tests

Co-authored-by: Copilot <copilot@github.com>

---------

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

python/packages/foundry/agent_framework_foundry/__init__.py

@@ -2,7 +2,7 @@
 
 import importlib.metadata
 
-from ._agent import FoundryAgent, RawFoundryAgent, RawFoundryAgentChatClient
+from ._agent import FoundryAgent, FoundryAgentOptions, RawFoundryAgent, RawFoundryAgentChatClient
 from ._chat_client import FoundryChatClient, FoundryChatOptions, RawFoundryChatClient
 from ._embedding_client import (
     FoundryEmbeddingClient,
@@ -25,6 +25,7 @@ except importlib.metadata.PackageNotFoundError:
 
 __all__ = [
     "FoundryAgent",
+    "FoundryAgentOptions",
     "FoundryChatClient",
     "FoundryChatOptions",
     "FoundryEmbeddingClient",

python/packages/foundry/agent_framework_foundry/_agent.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast
 
 from agent_framework import (
     AgentMiddlewareLayer,
+    AgentSession,
     ChatAndFunctionMiddlewareTypes,
     ChatMiddlewareLayer,
     ContextProvider,
@@ -52,11 +53,13 @@ else:
 if TYPE_CHECKING:
     from agent_framework import (
         Agent,
+        AgentRunInputs,
         ChatAndFunctionMiddlewareTypes,
         ContextProvider,
         MiddlewareTypes,
         ToolTypes,
     )
+    from agent_framework._agents import _RunContext  # pyright: ignore[reportPrivateUsage]
 
 logger: logging.Logger = logging.getLogger("agent_framework.foundry")
 
@@ -81,14 +84,54 @@ class FoundryAgentSettings(TypedDict, total=False):
     agent_version: str | None
 
 
+class FoundryAgentOptions(OpenAIChatOptions, total=False):
+    """Microsoft Foundry agent-specific chat options.
+
+    Extends ``OpenAIChatOptions`` with hosted-agent session configuration used by
+    ``FoundryAgent`` / ``RawFoundryAgent``.
+
+    Keyword Args:
+        extra_body: Additional request body values sent to the Responses API.
+        isolation_key: Isolation key used when lazily creating a hosted-agent
+            session through ``project_client.beta.agents.create_session(...)``.
+    """
+
+    extra_body: dict[str, Any]
+    isolation_key: str
+
+
 FoundryAgentOptionsT = TypeVar(
     "FoundryAgentOptionsT",
     bound=TypedDict,  # type: ignore[valid-type]
-    default="OpenAIChatOptions",
+    default="FoundryAgentOptions",
     covariant=True,
 )
 
 
+def _merge_extra_body(extra_body: Any | None, *, additions: Mapping[str, Any] | None = None) -> dict[str, Any]:
+    """Normalize and merge provider-specific extra_body values."""
+    if extra_body is None:
+        merged: dict[str, Any] = {}
+    elif isinstance(extra_body, Mapping):
+        merged = dict(cast(Mapping[str, Any], extra_body))
+    else:
+        raise TypeError(f"extra_body must be a mapping when provided, got {type(extra_body).__name__}.")
+
+    if additions:
+        merged.update(additions)
+    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_")
+    )
+
+
 class RawFoundryAgentChatClient(  # type: ignore[misc]
     RawOpenAIChatClient[FoundryAgentOptionsT],
     Generic[FoundryAgentOptionsT],
@@ -167,13 +210,15 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
         )
 
         resolved_endpoint = settings.get("project_endpoint")
-        self.agent_name = settings.get("agent_name")
-        self.agent_version = settings.get("agent_version")
+        agent_name_setting = settings.get("agent_name")
+        self.agent_version: str | None = settings.get("agent_version")
+        self.allow_preview = allow_preview or False
 
-        if not self.agent_name:
+        if not agent_name_setting:
             raise ValueError(
                 "Agent name is required. Set via 'agent_name' parameter or 'FOUNDRY_AGENT_NAME' environment variable."
             )
+        self.agent_name = agent_name_setting
 
         # Create or use provided project client
         self._should_close_client = False
@@ -197,11 +242,13 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
             self.project_client = AIProjectClient(**project_client_kwargs)
             self._should_close_client = True
 
-        # Get OpenAI client from project
-        async_client = self.project_client.get_openai_client()
-
+        openai_client_kwargs: dict[str, Any] = {}
+        if default_headers:
+            openai_client_kwargs["default_headers"] = dict(default_headers)
+        if allow_preview:
+            openai_client_kwargs["agent_name"] = self.agent_name
         super().__init__(
-            async_client=async_client,
+            async_client=self.project_client.get_openai_client(**openai_client_kwargs),
             default_headers=default_headers,
             instruction_role=instruction_role,
             compaction_strategy=compaction_strategy,
@@ -209,13 +256,6 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
             additional_properties=additional_properties,
         )
 
-    def _get_agent_reference(self) -> dict[str, str]:
-        """Build the agent reference dict for the Responses API."""
-        ref: dict[str, str] = {"name": self.agent_name, "type": "agent_reference"}  # type: ignore[dict-item]
-        if self.agent_version:
-            ref["version"] = self.agent_version
-        return ref
-
     @override
     def as_agent(
         self,
@@ -270,7 +310,7 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
         options: Mapping[str, Any],
         **kwargs: Any,
     ) -> dict[str, Any]:
-        """Prepare options for the Responses API, injecting agent reference and validating tools."""
+        """Prepare options for the Responses API and validate client-side tools."""
         # Validate tools — only FunctionTool allowed
         tools = options.get("tools", [])
         if tools:
@@ -292,18 +332,58 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
         if "input" in run_options and isinstance(run_options["input"], list):
             run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))
 
-        # Inject agent reference
-        run_options["extra_body"] = {"agent_reference": self._get_agent_reference()}
+        # 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)
+            extra_body["agent_session_id"] = conversation_id
+        if extra_body:
+            run_options["extra_body"] = extra_body
+
+        run_options.pop("isolation_key", None)
 
         # Strip tools from request body - Foundry API rejects requests with both
-        # agent_reference and tools present. FunctionTools are invoked client-side
+        # agent endpoint and tools present. FunctionTools are invoked client-side
         # by the function invocation layer, not sent to the service.
-        run_options.pop("tools", None)
-        run_options.pop("tool_choice", None)
-        run_options.pop("parallel_tool_calls", None)
+        run_options.pop("model", None)
+        if not self.allow_preview:
+            run_options.pop("tools", None)
+            run_options.pop("tool_choice", None)
+            run_options.pop("parallel_tool_calls", None)
 
         return run_options
 
+    @override
+    def _parse_response_from_openai(
+        self,
+        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
+
+    @override
+    def _parse_chunk_from_openai(
+        self,
+        event: Any,
+        options: dict[str, Any],
+        function_call_ids: dict[int, tuple[str, str]],
+        seen_reasoning_delta_item_ids: set[str] | None = None,
+    ) -> Any:
+        parsed_chunk = super()._parse_chunk_from_openai(
+            event,
+            options,
+            function_call_ids,
+            seen_reasoning_delta_item_ids,
+        )
+        if _uses_foundry_agent_session(options.get("conversation_id")):
+            parsed_chunk.conversation_id = None
+        return parsed_chunk
+
     @override
     def _check_model_presence(self, options: dict[str, Any]) -> None:
         """Skip model check — model is configured on the Foundry agent."""
@@ -368,6 +448,26 @@ class RawFoundryAgentChatClient(  # type: ignore[misc]
 
         return transformed
 
+    async def get_agent_version(self) -> str | None:
+        """Return the agent version if available, else None."""
+        if self.agent_version is not None:
+            return self.agent_version
+        if not self.allow_preview:
+            return None
+        agent_details = await cast(Any, self.project_client.beta.agents).get(  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+            agent_name=self.agent_name
+        )
+        versions_object = getattr(agent_details, "versions", None)
+        if not isinstance(versions_object, Mapping):
+            raise TypeError("Foundry agent details did not include a versions mapping.")
+        versions = cast(Mapping[str, Any], versions_object)
+        latest_version = versions.get("latest")
+        agent_version = getattr(cast(Any, latest_version), "version", None)
+        if not isinstance(agent_version, str):
+            raise TypeError("Foundry agent details did not include a latest version string.")
+        self.agent_version = agent_version
+        return agent_version
+
     async def close(self) -> None:
         """Close the project client if we created it."""
         if self._should_close_client:
@@ -395,7 +495,7 @@ class _FoundryAgentChatClient(  # type: ignore[misc]
             client = FoundryAgentClient(
                 project_endpoint="https://your-project.services.ai.azure.com",
                 agent_name="my-prompt-agent",
-                agent_version="1.0",
+                agent_version="1",
                 credential=AzureCliCredential(),
             )
 
@@ -477,7 +577,7 @@ class RawFoundryAgent(  # type: ignore[misc]
             agent = RawFoundryAgent(
                 project_endpoint="https://your-project.services.ai.azure.com",
                 agent_name="my-prompt-agent",
-                agent_version="1.0",
+                agent_version="1",
                 credential=AzureCliCredential(),
             )
             result = await agent.run("Hello!")
@@ -570,7 +670,7 @@ class RawFoundryAgent(  # type: ignore[misc]
             client=client,  # type: ignore[arg-type]
             instructions=instructions,
             id=id,
-            name=name,
+            name=name or agent_name,
             description=description,
             tools=tools,  # type: ignore[arg-type]
             default_options=cast(FoundryAgentOptionsT | None, default_options),
@@ -582,6 +682,81 @@ class RawFoundryAgent(  # type: ignore[misc]
             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:
+        """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")
+        )
+        if resolved_isolation_key is None:
+            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(
+        self,
+        *,
+        isolation_key: str | None = None,
+    ) -> str:
+        """Create a hosted Foundry service session and return the service session ID."""
+        if not isinstance(self.client, RawFoundryAgentChatClient):
+            raise TypeError("_create_service_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),
+        }
+        if version := await self.client.get_agent_version():
+            from azure.ai.projects.models import VersionRefIndicator
+
+            create_session_kwargs["version_indicator"] = VersionRefIndicator(agent_version=version)  # type: ignore
+
+        service_session = await self.client.project_client.beta.agents.create_session(**create_session_kwargs)
+        agent_session_id = getattr(service_session, "agent_session_id", None)
+        if not isinstance(agent_session_id, str) or not agent_session_id:
+            raise ValueError("Hosted Foundry session creation did not return a non-empty agent_session_id.")
+
+        return agent_session_id
+
+    @override
+    async def _prepare_run_context(
+        self,
+        *,
+        messages: AgentRunInputs | None,
+        session: AgentSession | None,
+        tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None,
+        options: Mapping[str, Any] | None,
+        compaction_strategy: CompactionStrategy | None,
+        tokenizer: TokenizerProtocol | None,
+        function_invocation_kwargs: Mapping[str, Any] | None,
+        client_kwargs: Mapping[str, Any] | None,
+    ) -> _RunContext:
+        runtime_options = dict(options) if options else {}
+        effective_options = {
+            **{key: value for key, value in self.default_options.items() if value is not None},
+            **{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")),
+            )
+
+        return await super()._prepare_run_context(
+            messages=messages,
+            session=session,
+            tools=tools,
+            options=runtime_options,
+            compaction_strategy=compaction_strategy,
+            tokenizer=tokenizer,
+            function_invocation_kwargs=function_invocation_kwargs,
+            client_kwargs=client_kwargs,
+        )
+
     async def configure_azure_monitor(
         self,
         enable_sensitive_data: bool = False,
@@ -708,6 +883,19 @@ class FoundryAgent(  # type: ignore[misc]
     ) -> None:
         """Initialize a Foundry Agent with full middleware and telemetry.
 
+        ``FoundryAgent`` supports both PromptAgents and HostedAgents. PromptAgents
+        typically provide ``agent_version`` directly. HostedAgents can omit
+        ``agent_version`` and, when they need preview-only session APIs, should
+        opt in with ``allow_preview=True`` when this class creates the underlying
+        ``AIProjectClient``. If you pass ``project_client`` explicitly, it must
+        already be configured for preview APIs before being passed to
+        ``FoundryAgent``.
+
+        To lazily create HostedAgent service sessions 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.
+
         Keyword Args:
             project_endpoint: The Foundry project endpoint URL.
             agent_name: The name of the Foundry agent to connect to.
@@ -715,6 +903,9 @@ class FoundryAgent(  # type: ignore[misc]
             credential: Azure credential for authentication.
             project_client: An existing AIProjectClient to use.
             allow_preview: Enables preview opt-in on internally-created AIProjectClient.
+                Set this to ``True`` for HostedAgents that need preview-only
+                session APIs, including lazy service session creation from
+                ``isolation_key``.
             tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
             context_providers: Optional context providers.
             middleware: Optional agent-level middleware.
@@ -726,6 +917,8 @@ class FoundryAgent(  # type: ignore[misc]
             description: Optional local description for the local agent wrapper.
             instructions: Optional instructions for the local agent wrapper.
             default_options: Default chat options for the local agent wrapper.
+                ``FoundryAgentOptions`` can include ``isolation_key`` and
+                ``extra_body`` when working with HostedAgents.
             require_per_service_call_history_persistence: Whether to require per-service-call
                 chat history persistence when using local history providers.
             function_invocation_configuration: Optional function invocation configuration override.

python/packages/foundry/agent_framework_foundry/_chat_client.py

@@ -204,9 +204,13 @@ class RawFoundryChatClient(  # type: ignore[misc]
                 project_client_kwargs["allow_preview"] = allow_preview
             project_client = AIProjectClient(**project_client_kwargs)
 
+        openai_kwargs: dict[str, Any] = {}
+        if default_headers:
+            openai_kwargs["default_headers"] = default_headers
+
         super().__init__(
             model=resolved_model,
-            async_client=project_client.get_openai_client(),
+            async_client=project_client.get_openai_client(**openai_kwargs),
             default_headers=default_headers,
             instruction_role=instruction_role,
             compaction_strategy=compaction_strategy,

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

@@ -5,11 +5,12 @@ from __future__ import annotations
 import inspect
 import os
 import sys
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agent_framework import AgentResponse, ChatContext, ChatMiddleware, Message, tool
+from agent_framework import AgentResponse, AgentSession, ChatContext, ChatMiddleware, ChatResponse, Message, tool
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import AzureCliCredential
 
@@ -54,7 +55,7 @@ def test_raw_foundry_agent_chat_client_init_requires_agent_name() -> None:
 
 
 def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None:
-    """Test construction with agent_name and project_client."""
+    """Test construction with agent_name and project_client without preview agent binding."""
 
     mock_project = MagicMock()
     mock_project.get_openai_client.return_value = MagicMock()
@@ -67,6 +68,27 @@ def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None:
 
     assert client.agent_name == "test-agent"
     assert client.agent_version == "1.0"
+    mock_project.get_openai_client.assert_called_once_with()
+
+
+def test_raw_foundry_agent_chat_client_init_passes_agent_name_when_preview_enabled() -> None:
+    """Test preview-enabled clients bind the OpenAI client to the agent endpoint."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="hosted-agent",
+        allow_preview=True,
+        default_headers={"x-test": "1"},
+    )
+
+    assert client.agent_name == "hosted-agent"
+    mock_project.get_openai_client.assert_called_once_with(
+        agent_name="hosted-agent",
+        default_headers={"x-test": "1"},
+    )
 
 
 def test_raw_foundry_agent_chat_client_init_uses_explicit_parameters() -> None:
@@ -80,38 +102,6 @@ def test_raw_foundry_agent_chat_client_init_uses_explicit_parameters() -> None:
     assert all(parameter.kind != inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values())
 
 
-def test_raw_foundry_agent_chat_client_get_agent_reference_with_version() -> None:
-    """Test agent reference includes version when provided."""
-
-    mock_project = MagicMock()
-    mock_project.get_openai_client.return_value = MagicMock()
-
-    client = RawFoundryAgentChatClient(
-        project_client=mock_project,
-        agent_name="my-agent",
-        agent_version="2.0",
-    )
-
-    ref = client._get_agent_reference()
-    assert ref == {"name": "my-agent", "version": "2.0", "type": "agent_reference"}
-
-
-def test_raw_foundry_agent_chat_client_get_agent_reference_without_version() -> None:
-    """Test agent reference omits version for HostedAgents."""
-
-    mock_project = MagicMock()
-    mock_project.get_openai_client.return_value = MagicMock()
-
-    client = RawFoundryAgentChatClient(
-        project_client=mock_project,
-        agent_name="hosted-agent",
-    )
-
-    ref = client._get_agent_reference()
-    assert ref == {"name": "hosted-agent", "type": "agent_reference"}
-    assert "version" not in ref
-
-
 def test_raw_foundry_agent_chat_client_as_agent_preserves_client_type() -> None:
     """Test that as_agent() wraps the client in FoundryAgent using the same client class."""
 
@@ -196,12 +186,11 @@ async def test_raw_foundry_agent_chat_client_prepare_options_accepts_function_to
             options={"tools": [my_func]},
         )
 
-    assert "extra_body" in result
-    assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
+    assert result == {}
 
 
-async def test_raw_foundry_agent_chat_client_prepare_options_strips_tools() -> None:
-    """Test that _prepare_options strips tools, tool_choice, and parallel_tool_calls from run_options."""
+async def test_raw_foundry_agent_chat_client_prepare_options_strips_client_side_fields() -> None:
+    """Test that _prepare_options strips model and tool-loop fields from run_options."""
 
     mock_project = MagicMock()
     mock_openai = MagicMock()
@@ -222,6 +211,7 @@ async def test_raw_foundry_agent_chat_client_prepare_options_strips_tools() -> N
         "agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
         new_callable=AsyncMock,
         return_value={
+            "model": "gpt-4.1",
             "tools": [{"type": "function", "function": {"name": "my_func"}}],
             "tool_choice": "auto",
             "parallel_tool_calls": True,
@@ -232,11 +222,69 @@ async def test_raw_foundry_agent_chat_client_prepare_options_strips_tools() -> N
             options={"tools": [my_func]},
         )
 
+    assert "model" not in result
     assert "tools" not in result
     assert "tool_choice" not in result
     assert "parallel_tool_calls" not in result
-    assert "extra_body" in result
-    assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
+    assert result == {}
+
+
+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."""
+
+    mock_project = MagicMock()
+    mock_openai = MagicMock()
+    mock_project.get_openai_client.return_value = mock_openai
+
+    client = RawFoundryAgentChatClient(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    with patch(
+        "agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
+        new_callable=AsyncMock,
+        return_value={
+            "extra_body": {"custom": "value"},
+            "previous_response_id": "should-be-removed",
+        },
+    ):
+        result = await client._prepare_options(
+            messages=[Message(role="user", contents="hi")],
+            options={"conversation_id": "agent-session-123", "isolation_key": "iso-key"},
+        )
+
+    assert result["extra_body"] == {
+        "custom": "value",
+        "agent_session_id": "agent-session-123",
+    }
+    assert "previous_response_id" not in result
+    assert "conversation" not in result
+    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."""
+
+    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=MagicMock(),
+            options={"conversation_id": "agent-session-123"},
+        )
+
+    assert result.conversation_id is None
 
 
 def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None:
@@ -366,6 +414,74 @@ def test_raw_foundry_agent_init_with_function_tools() -> None:
     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."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+    mock_project.beta = SimpleNamespace(
+        agents=SimpleNamespace(
+            create_session=AsyncMock(return_value=SimpleNamespace(agent_session_id="agent-session-123"))
+        )
+    )
+
+    agent = RawFoundryAgent(
+        project_client=mock_project,
+        agent_name="test-agent",
+        agent_version="1.0",
+        allow_preview=True,
+    )
+    session = AgentSession()
+
+    with patch(
+        "agent_framework._agents.RawAgent._prepare_run_context",
+        new=AsyncMock(return_value={"ok": True}),
+    ) as mock_prepare_run_context:
+        result = await agent._prepare_run_context(
+            messages="hi",
+            session=session,
+            tools=None,
+            options={"isolation_key": "iso-key"},
+            compaction_strategy=None,
+            tokenizer=None,
+            function_invocation_kwargs=None,
+            client_kwargs=None,
+        )
+
+    assert result == {"ok": True}
+    assert session.service_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
+    mock_prepare_run_context.assert_awaited_once()
+
+
+async def test_raw_foundry_agent_prepare_run_context_requires_preview_for_hosted_sessions() -> None:
+    """Test that hosted-agent sessions require allow_preview=True."""
+
+    mock_project = MagicMock()
+    mock_project.get_openai_client.return_value = MagicMock()
+
+    agent = RawFoundryAgent(
+        project_client=mock_project,
+        agent_name="test-agent",
+    )
+
+    with pytest.raises(RuntimeError, match="allow_preview=True"):
+        await agent._prepare_run_context(
+            messages="hi",
+            session=AgentSession(),
+            tools=None,
+            options={"isolation_key": "iso-key"},
+            compaction_strategy=None,
+            tokenizer=None,
+            function_invocation_kwargs=None,
+            client_kwargs=None,
+        )
+
+
 def test_foundry_agent_init() -> None:
     """Test construction of the full-middleware agent."""
 
@@ -483,9 +599,10 @@ async def test_foundry_agent_configure_azure_monitor_import_error() -> None:
 @pytest.mark.flaky
 @pytest.mark.integration
 @skip_if_foundry_agent_integration_tests_disabled
+@pytest.mark.skip(reason="Test agent seems to have disappeared from the test environment; needs investigation.")
 async def test_foundry_agent_basic_run() -> None:
     """Smoke-test FoundryAgent against a real configured agent."""
-    async with FoundryAgent(credential=AzureCliCredential()) as agent:
+    async with FoundryAgent(credential=AzureCliCredential(), allow_preview=True) as agent:
         response = await agent.run("Please respond with exactly: 'This is a response test.'")
 
     assert isinstance(response, AgentResponse)
@@ -496,6 +613,7 @@ async def test_foundry_agent_basic_run() -> None:
 @pytest.mark.flaky
 @pytest.mark.integration
 @skip_if_foundry_agent_integration_tests_disabled
+@pytest.mark.skip(reason="Test agent seems to have disappeared from the test environment; needs investigation.")
 async def test_foundry_agent_custom_client_run() -> None:
     """Smoke-test FoundryAgent against a real configured agent."""
     async with FoundryAgent(credential=AzureCliCredential(), client_type=RawFoundryAgentChatClient) as agent:

python/packages/foundry/tests/foundry/test_foundry_embedding_client.py

@@ -198,6 +198,7 @@ class TestRawFoundryEmbeddingClient:
                     "FOUNDRY_MODELS_API_KEY": "env-key",
                     "FOUNDRY_EMBEDDING_MODEL": "env-model",
                 },
+                clear=True,
             ),
             patch("agent_framework_foundry._embedding_client.EmbeddingsClient"),
             patch("agent_framework_foundry._embedding_client.ImageEmbeddingsClient"),