diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index 4109c644d2..09713425bb 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -120,6 +120,7 @@ class AzureAIAgentClient(BaseChatClient): project_endpoint: str | None = None, model_deployment_name: str | None = None, async_credential: AsyncTokenCredential | None = None, + should_cleanup_agent: bool = True, env_file_path: str | None = None, env_file_encoding: str | None = None, **kwargs: Any, @@ -140,6 +141,9 @@ class AzureAIAgentClient(BaseChatClient): model_deployment_name: The model deployment name to use for agent creation. Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME. async_credential: Azure async credential to use for authentication. + should_cleanup_agent: Whether to cleanup (delete) agents created by this client when + the client is closed or context is exited. Defaults to True. Only affects agents + created by this client instance; existing agents passed via agent_id are never deleted. env_file_path: Path to environment file for loading settings. env_file_encoding: Encoding of the environment file. kwargs: Additional keyword arguments passed to the parent class. @@ -212,7 +216,8 @@ class AzureAIAgentClient(BaseChatClient): self.agent_name = agent_name self.model_id = azure_ai_settings.model_deployment_name self.thread_id = thread_id - self._should_delete_agent = False # Track whether we should delete the agent + self.should_cleanup_agent = should_cleanup_agent # Track whether we should delete the agent + self._agent_created = False # Track whether agent was created inside this class self._should_close_client = should_close_client # Track whether we should close client connection self._agent_definition: Agent | None = None # Cached definition for existing agent @@ -245,6 +250,7 @@ class AzureAIAgentClient(BaseChatClient): agent_name=settings.get("agent_name"), credential=settings.get("credential"), env_file_path=settings.get("env_file_path"), + should_cleanup_agent=settings.get("should_cleanup_agent", True), ) async def _inner_get_response( @@ -323,7 +329,7 @@ class AzureAIAgentClient(BaseChatClient): self.agent_id = str(created_agent.id) self._agent_definition = created_agent - self._should_delete_agent = True + self._agent_created = True return self.agent_id @@ -655,10 +661,10 @@ class AzureAIAgentClient(BaseChatClient): async def _cleanup_agent_if_needed(self) -> None: """Clean up the agent if we created it.""" - if self._should_delete_agent and self.agent_id is not None: + if self._agent_created and self.should_cleanup_agent and self.agent_id is not None: await self.agents_client.delete_agent(self.agent_id) self.agent_id = None - self._should_delete_agent = False + self._agent_created = False async def _load_agent_definition_if_needed(self) -> Agent | None: """Load and cache agent details if not already loaded.""" diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py index 00f8aa404d..555d27d560 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py @@ -71,7 +71,7 @@ def create_test_azure_ai_chat_client( agent_id: str | None = None, thread_id: str | None = None, azure_ai_settings: AzureAISettings | None = None, - should_delete_agent: bool = False, + should_cleanup_agent: bool = True, agent_name: str | None = None, ) -> AzureAIAgentClient: """Helper function to create AzureAIAgentClient instances for testing, bypassing normal validation.""" @@ -88,9 +88,10 @@ def create_test_azure_ai_chat_client( client.agent_name = agent_name client.model_id = azure_ai_settings.model_deployment_name client.thread_id = thread_id - client._should_delete_agent = should_delete_agent # type: ignore - client._should_close_client = False # type: ignore - client._agent_definition = None # type: ignore + client.should_cleanup_agent = should_cleanup_agent + client._agent_created = False + client._should_close_client = False + client._agent_definition = None client.additional_properties = {} client.middleware = None @@ -125,7 +126,6 @@ def test_azure_ai_chat_client_init_with_client(mock_agents_client: MagicMock) -> assert chat_client.agents_client is mock_agents_client assert chat_client.agent_id == "existing-agent-id" assert chat_client.thread_id == "test-thread-id" - assert not chat_client._should_delete_agent # type: ignore assert isinstance(chat_client, ChatClientProtocol) @@ -141,7 +141,6 @@ def test_azure_ai_chat_client_init_auto_create_client( chat_client.agents_client = mock_agents_client chat_client.agent_id = None chat_client.thread_id = None - chat_client._should_delete_agent = False # type: ignore chat_client._should_close_client = False # type: ignore chat_client.credential = None chat_client.model_id = azure_ai_settings.model_deployment_name @@ -151,7 +150,6 @@ def test_azure_ai_chat_client_init_auto_create_client( assert chat_client.agents_client is mock_agents_client assert chat_client.agent_id is None - assert not chat_client._should_delete_agent # type: ignore def test_azure_ai_chat_client_init_missing_project_endpoint() -> None: @@ -302,7 +300,7 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_existing_agent( agent_id = await chat_client._get_agent_id_or_create() # type: ignore assert agent_id == "existing-agent-id" - assert not chat_client._should_delete_agent # type: ignore + assert not chat_client._agent_created async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( @@ -316,7 +314,7 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( agent_id = await chat_client._get_agent_id_or_create(run_options={"model": azure_ai_settings.model_deployment_name}) # type: ignore assert agent_id == "test-agent-id" - assert chat_client._should_delete_agent # type: ignore + assert chat_client._agent_created async def test_azure_ai_chat_client_thread_management_through_public_api(mock_agents_client: MagicMock) -> None: @@ -364,74 +362,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_missing_model( await chat_client._get_agent_id_or_create() # type: ignore -async def test_azure_ai_chat_client_cleanup_agent_if_needed_should_delete( - mock_agents_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed when agent should be deleted.""" - chat_client = create_test_azure_ai_chat_client( - mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - await chat_client._cleanup_agent_if_needed() # type: ignore - # Verify agent deletion was called - mock_agents_client.delete_agent.assert_called_once_with("agent-to-delete") - assert not chat_client._should_delete_agent # type: ignore - - -async def test_azure_ai_chat_client_cleanup_agent_if_needed_should_not_delete( - mock_agents_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed when agent should not be deleted.""" - chat_client = create_test_azure_ai_chat_client( - mock_agents_client, agent_id="agent-to-keep", should_delete_agent=False - ) - - await chat_client._cleanup_agent_if_needed() # type: ignore - - # Verify agent deletion was not called - mock_agents_client.delete_agent.assert_not_called() - assert not chat_client._should_delete_agent # type: ignore - - -async def test_azure_ai_chat_client_cleanup_agent_if_needed_exception_handling( - mock_agents_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed propagates exceptions (it doesn't handle them).""" - chat_client = create_test_azure_ai_chat_client( - mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True - ) - mock_agents_client.delete_agent.side_effect = Exception("Deletion failed") - - with pytest.raises(Exception, match="Deletion failed"): - await chat_client._cleanup_agent_if_needed() # type: ignore - - -async def test_azure_ai_chat_client_aclose(mock_agents_client: MagicMock) -> None: - """Test aclose method calls cleanup.""" - chat_client = create_test_azure_ai_chat_client( - mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - await chat_client.close() - - # Verify agent deletion was called - mock_agents_client.delete_agent.assert_called_once_with("agent-to-delete") - - -async def test_azure_ai_chat_client_async_context_manager(mock_agents_client: MagicMock) -> None: - """Test async context manager functionality.""" - chat_client = create_test_azure_ai_chat_client( - mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - # Test context manager - async with chat_client: - pass # Just test that we can enter and exit - - # Verify cleanup was called on exit - mock_agents_client.delete_agent.assert_called_once_with("agent-to-delete") - - async def test_azure_ai_chat_client_create_run_options_basic(mock_agents_client: MagicMock) -> None: """Test _create_run_options with basic ChatOptions.""" chat_client = create_test_azure_ai_chat_client(mock_agents_client) @@ -1347,23 +1277,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_tool_resources( assert call_kwargs["tool_resources"] == {"vector_store_ids": ["vs-123"]} -async def test_azure_ai_chat_client_close_method(mock_agents_client: MagicMock) -> None: - """Test close method.""" - chat_client = create_test_azure_ai_chat_client(mock_agents_client, should_delete_agent=True) - chat_client._should_close_client = True # type: ignore - chat_client.agent_id = "test-agent" - - # Mock cleanup methods - mock_agents_client.delete_agent = AsyncMock() - mock_agents_client.close = AsyncMock() - - await chat_client.close() - - # Verify cleanup was called - mock_agents_client.delete_agent.assert_called_once_with("test-agent") - mock_agents_client.close.assert_called_once() - - async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs( mock_agents_client: MagicMock, ) -> None: @@ -1837,3 +1750,57 @@ async def test_azure_ai_chat_client_agent_chat_options_agent_level() -> None: assert isinstance(response, AgentRunResponse) assert response.text is not None assert len(response.text) > 0 + + +async def test_azure_ai_chat_client_cleanup_agent_when_enabled_and_created( + mock_agents_client: MagicMock, +) -> None: + """Test that agent is cleaned up when should_cleanup_agent=True and agent was created by client.""" + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id=None, should_cleanup_agent=True) + + # Simulate agent creation + chat_client.agent_id = "created-agent-id" + chat_client._agent_created = True # type: ignore + + await chat_client._cleanup_agent_if_needed() # type: ignore + + # Verify agent was deleted + mock_agents_client.delete_agent.assert_called_once_with("created-agent-id") + assert chat_client.agent_id is None + assert chat_client._agent_created is False # type: ignore + + +async def test_azure_ai_chat_client_no_cleanup_when_disabled( + mock_agents_client: MagicMock, +) -> None: + """Test that agent is not cleaned up when should_cleanup_agent=False.""" + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id=None, should_cleanup_agent=False) + + # Simulate agent creation + chat_client.agent_id = "created-agent-id" + chat_client._agent_created = True + + await chat_client._cleanup_agent_if_needed() # type: ignore + + # Verify agent was NOT deleted + mock_agents_client.delete_agent.assert_not_called() + assert chat_client.agent_id == "created-agent-id" + assert chat_client._agent_created is True + + +async def test_azure_ai_chat_client_no_cleanup_when_agent_not_created_by_client( + mock_agents_client: MagicMock, +) -> None: + """Test that agent is not cleaned up when it was not created by this client instance.""" + chat_client = create_test_azure_ai_chat_client( + mock_agents_client, agent_id="existing-agent-id", should_cleanup_agent=True + ) + + # Agent exists but was not created by this client (_agent_created = False) + assert chat_client._agent_created is False # type: ignore + + await chat_client._cleanup_agent_if_needed() # type: ignore + + # Verify agent was NOT deleted + mock_agents_client.delete_agent.assert_not_called() + assert chat_client.agent_id == "existing-agent-id" diff --git a/python/samples/getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py b/python/samples/getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py index c695ea9a90..0ac2ee620c 100644 --- a/python/samples/getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py +++ b/python/samples/getting_started/agents/azure_ai_agent/azure_ai_with_explicit_settings.py @@ -41,6 +41,7 @@ async def main() -> None: model_deployment_name=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], async_credential=credential, agent_name="WeatherAgent", + should_cleanup_agent=True, # Set to False if you want to disable automatic agent cleanup ), instructions="You are a helpful weather agent.", tools=get_weather,