From 1354c43d1f90f5749114fb5e6cdeb99fe0a6d3ee Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Sun, 19 Jul 2026 23:08:56 +0530 Subject: [PATCH 1/6] fix(foundry_hosting): preserve auth credentials across FoundryToolbox reconnections --- .../_toolbox.py | 18 +++++- .../foundry_hosting/tests/test_toolbox.py | 59 +++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index a3848a75be..91c858bcf7 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -4,7 +4,9 @@ import logging import os +from contextlib import AbstractAsyncContextManager from typing import TYPE_CHECKING +from typing import Any from urllib.parse import urlsplit import httpx @@ -170,6 +172,9 @@ def __init__( auth=_ToolboxAuth(credential, token_scope), timeout=timeout, ) + self._credential = credential + self._token_scope = token_scope + self._timeout = timeout super().__init__( name=tool_name, @@ -179,8 +184,19 @@ def __init__( load_tools=load_tools, ) + def get_mcp_client(self) -> AbstractAsyncContextManager[Any]: + """Get an authenticated MCP HTTP client. + + Recreates the underlying HTTP client if it was previously closed.""" + if self._httpx_client is None: + self._httpx_client = httpx.AsyncClient( + auth=_ToolboxAuth(self._credential, self._token_scope), + timeout=self._timeout, + ) + return super().get_mcp_client() + async def close(self) -> None: - """Close the MCP session and the toolbox-owned HTTP client.""" + """Close the MCP session and toolbox HTTP client while preserving credentials and timeout for reconnection.""" try: await super().close() finally: diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 005eefa576..f133f3f715 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -1,4 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. +# pyright: reportPrivateUsage=false """Unit tests for FoundryToolbox.""" @@ -229,3 +230,61 @@ async def get_skills(self, context: SkillsSourceContext) -> list[str]: assert result == ["skill-a"] assert captured["client"] is sentinel_session + +class TestFoundryToolboxReconnection: + async def test_close_preserves_credential_for_reconnection(self) -> None: + """After close(), get_mcp_client() should recreate an authenticated client.""" + cred = _FakeCredential("reconnect-token") + toolbox = FoundryToolbox( + cred, + url="https://h/toolboxes/recon/mcp", + timeout=60.0, + ) + + assert toolbox._credential is cred + assert toolbox._token_scope == "https://ai.azure.com/.default" + assert toolbox._timeout == 60.0 + + assert toolbox._httpx_client is not None + assert isinstance(toolbox._httpx_client.auth, _ToolboxAuth) + original_auth = toolbox._httpx_client.auth + + client = toolbox._httpx_client + client.aclose = AsyncMock() + await toolbox.close() + + client.aclose.assert_awaited_once() + assert toolbox._httpx_client is None + + assert toolbox._credential is cred + assert toolbox._timeout == 60.0 + + ctx_manager = toolbox.get_mcp_client() + assert toolbox._httpx_client is not None + assert isinstance(toolbox._httpx_client.auth, _ToolboxAuth) + + new_auth = toolbox._httpx_client.auth + assert new_auth is not original_auth + assert new_auth._credential is cred + + assert hasattr(ctx_manager, "__aenter__") + assert hasattr(ctx_manager, "__aexit__") + + await toolbox.close() + + async def test_close_idempotent_with_reconnection(self) -> None: + """Multiple close() calls don't break reconnection.""" + cred = _FakeCredential() + toolbox = FoundryToolbox( + cred, + url="https://h/toolboxes/idem/mcp", + ) + + await toolbox.close() + await toolbox.close() + + toolbox.get_mcp_client() + assert toolbox._httpx_client is not None + assert isinstance(toolbox._httpx_client.auth, _ToolboxAuth) + + await toolbox.close() From 2e13ec5cd0a64728dae702161ff25a11ef16e988 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 21 Jul 2026 14:29:25 -0700 Subject: [PATCH 2/6] Address copilot comments --- .../agent_framework_foundry_hosting/_toolbox.py | 14 ++++++++------ .../packages/foundry_hosting/tests/test_toolbox.py | 1 + 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 91c858bcf7..20fffc3ddf 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -4,9 +4,8 @@ import logging import os -from contextlib import AbstractAsyncContextManager -from typing import TYPE_CHECKING -from typing import Any +from contextlib import _AsyncGeneratorContextManager # pyright: ignore[reportPrivateUsage] +from typing import TYPE_CHECKING, Any, override from urllib.parse import urlsplit import httpx @@ -184,10 +183,12 @@ def __init__( load_tools=load_tools, ) - def get_mcp_client(self) -> AbstractAsyncContextManager[Any]: + @override + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an authenticated MCP HTTP client. - - Recreates the underlying HTTP client if it was previously closed.""" + + Recreates the underlying HTTP client if it was previously closed. + """ if self._httpx_client is None: self._httpx_client = httpx.AsyncClient( auth=_ToolboxAuth(self._credential, self._token_scope), @@ -195,6 +196,7 @@ def get_mcp_client(self) -> AbstractAsyncContextManager[Any]: ) return super().get_mcp_client() + @override async def close(self) -> None: """Close the MCP session and toolbox HTTP client while preserving credentials and timeout for reconnection.""" try: diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index f133f3f715..775f52092b 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -231,6 +231,7 @@ async def get_skills(self, context: SkillsSourceContext) -> list[str]: assert result == ["skill-a"] assert captured["client"] is sentinel_session + class TestFoundryToolboxReconnection: async def test_close_preserves_credential_for_reconnection(self) -> None: """After close(), get_mcp_client() should recreate an authenticated client.""" From 558679ba343081d011e0305653ec514febe4106c Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 22 Jul 2026 11:16:47 -0700 Subject: [PATCH 3/6] fix syntax check --- .../agent_framework_foundry_hosting/_toolbox.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 0a76ffcaef..00f03ef995 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -5,7 +5,7 @@ import logging import os from contextlib import _AsyncGeneratorContextManager # pyright: ignore[reportPrivateUsage] -from typing import TYPE_CHECKING, Any, override +from typing import TYPE_CHECKING, Any from urllib.parse import urlsplit import httpx @@ -19,6 +19,7 @@ SkillsSourceContext, ) from azure.ai.agentserver.core import get_request_context +from typing_extensions import override if TYPE_CHECKING: from collections.abc import Generator From e831da5e9dbfeb1afef55de663fa321c117c8154 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 22 Jul 2026 11:41:27 -0700 Subject: [PATCH 4/6] Fix tests --- python/packages/foundry_hosting/tests/test_toolbox.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 8b4f62ae9f..1f5eee7e9b 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -231,7 +231,6 @@ async def get_skills(self, context: SkillsSourceContext) -> list[str]: result = await _FoundryToolboxSkillsSource(toolbox).get_skills(_source_context()) assert result == ["skill-a"] - assert captured["client"] is sentinel_session # The source hands MCPSkillsSource a provider (not a fixed session) that resolves # the toolbox's current session, so it survives a reconnect that swaps it. provider = captured["session_provider"] From e0454dc592701533b69a8d1eb05c90840c2df85f Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 22 Jul 2026 14:45:21 -0700 Subject: [PATCH 5/6] Fix formatting --- .../foundry_hosting/tests/test_toolbox.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 1f5eee7e9b..69c0b4ecd0 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -21,7 +21,7 @@ ) from agent_framework_foundry_hosting import FoundryToolbox -from agent_framework_foundry_hosting._toolbox import ( # pyright: ignore[reportPrivateUsage] +from agent_framework_foundry_hosting._toolbox import ( _FoundryToolboxSkillsSource, _resolve_toolbox_endpoint, _toolbox_name_from_endpoint, @@ -146,9 +146,9 @@ async def test_close_closes_owned_http_client() -> None: _FakeCredential(), # type: ignore url="https://h/toolboxes/tb/mcp", ) - client = toolbox._httpx_client # pyright: ignore[reportPrivateUsage] + client = toolbox._httpx_client assert client is not None - client.aclose = AsyncMock() # type: ignore[method-assign] + client.aclose = AsyncMock() await toolbox.close() @@ -175,9 +175,9 @@ def test_as_skills_provider_requires_approval_by_default() -> None: ) provider = toolbox.as_skills_provider() # By default every skill tool keeps its approval requirement. - assert provider._disable_load_skill_approval is False # pyright: ignore[reportPrivateUsage] - assert provider._disable_read_skill_resource_approval is False # pyright: ignore[reportPrivateUsage] - assert provider._disable_run_skill_script_approval is False # pyright: ignore[reportPrivateUsage] + assert provider._disable_load_skill_approval is False + assert provider._disable_read_skill_resource_approval is False + assert provider._disable_run_skill_script_approval is False def test_as_skills_provider_forwards_approval_overrides() -> None: @@ -192,9 +192,9 @@ def test_as_skills_provider_forwards_approval_overrides() -> None: ) # Overrides flow through to the underlying SkillsProvider so an unattended # host (no AgentSession) can load skills without an approval round-trip. - assert provider._disable_load_skill_approval is True # pyright: ignore[reportPrivateUsage] - assert provider._disable_read_skill_resource_approval is True # pyright: ignore[reportPrivateUsage] - assert provider._disable_run_skill_script_approval is True # pyright: ignore[reportPrivateUsage] + assert provider._disable_load_skill_approval is True + assert provider._disable_read_skill_resource_approval is True + assert provider._disable_run_skill_script_approval is True async def test_skills_source_requires_connection() -> None: @@ -249,9 +249,9 @@ async def test_skills_source_requires_connection_via_provider() -> None: source = _FoundryToolboxSkillsSource(toolbox) # Discovery captures the bound provider; a later reconnect gap (session is None) # surfaces the same clear error when the provider is resolved. - toolbox.session = None # type: ignore + toolbox.session = None with pytest.raises(RuntimeError, match="not connected"): - source._require_session() # pyright: ignore[reportPrivateUsage] + source._require_session() class _FakeSkill: @@ -292,7 +292,7 @@ async def test_as_skills_provider_caches_by_default(monkeypatch: pytest.MonkeyPa provider = toolbox.as_skills_provider() context = _source_context() for _ in range(3): - await provider._source.get_skills(context) # pyright: ignore[reportPrivateUsage] + await provider._source.get_skills(context) # By default the toolbox index is read once and reused across agent runs. assert read_count[0] == 1 @@ -309,7 +309,7 @@ async def test_as_skills_provider_disable_caching_rereads_every_run(monkeypatch: provider = toolbox.as_skills_provider(disable_caching=True) context = _source_context() for _ in range(3): - await provider._source.get_skills(context) # pyright: ignore[reportPrivateUsage] + await provider._source.get_skills(context) # With caching disabled the index is re-read on every agent run. assert read_count[0] == 3 @@ -330,7 +330,7 @@ async def test_as_skills_provider_cache_refresh_interval_rereads_after_staleness provider = toolbox.as_skills_provider(cache_refresh_interval=timedelta(0)) context = _source_context() for _ in range(3): - await provider._source.get_skills(context) # pyright: ignore[reportPrivateUsage] + await provider._source.get_skills(context) assert read_count[0] == 3 @@ -340,7 +340,7 @@ async def test_close_preserves_credential_for_reconnection(self) -> None: """After close(), get_mcp_client() should recreate an authenticated client.""" cred = _FakeCredential("reconnect-token") toolbox = FoundryToolbox( - cred, + cred, # type: ignore url="https://h/toolboxes/recon/mcp", timeout=60.0, ) @@ -380,7 +380,7 @@ async def test_close_idempotent_with_reconnection(self) -> None: """Multiple close() calls don't break reconnection.""" cred = _FakeCredential() toolbox = FoundryToolbox( - cred, + cred, # type: ignore url="https://h/toolboxes/idem/mcp", ) From 094c2a33e26c80813c386770fbf80a91499eea6f Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 22 Jul 2026 15:33:19 -0700 Subject: [PATCH 6/6] Fix formatting --- python/packages/foundry_hosting/tests/test_toolbox.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 69c0b4ecd0..7d9faedd47 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -148,7 +148,7 @@ async def test_close_closes_owned_http_client() -> None: ) client = toolbox._httpx_client assert client is not None - client.aclose = AsyncMock() + client.aclose = AsyncMock() # zuban: ignore await toolbox.close() @@ -354,7 +354,7 @@ async def test_close_preserves_credential_for_reconnection(self) -> None: original_auth = toolbox._httpx_client.auth client = toolbox._httpx_client - client.aclose = AsyncMock() + client.aclose = AsyncMock() # zuban: ignore await toolbox.close() client.aclose.assert_awaited_once()