fix(openai): add async context manager support

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -9,6 +9,7 @@
 from asyncio import Task
 from dataclasses import dataclass
 from importlib.metadata import PackageNotFoundError, version
+from types import TracebackType
 from typing import (
     Any,
     AsyncGenerator,
@@ -1142,6 +1143,17 @@ async def _create_stream_chunks_beta_client(
     async def close(self) -> None:
         await self._client.close()
 
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        await self.close()
+
     def actual_usage(self) -> RequestUsage:
         return self._actual_usage
 

python/packages/autogen-ext/tests/models/test_openai_model_client.py

@@ -3378,3 +3378,28 @@ async def test_reasoning_effort_validation() -> None:
         }
 
         ChatCompletionClient.load_component(config)
+
+
+@pytest.mark.asyncio
+async def test_openai_client_async_context_manager() -> None:
+    """Test that OpenAIChatCompletionClient supports async context manager protocol."""
+    client = OpenAIChatCompletionClient(model="gpt-4o", api_key="fake_key")
+    client._client.close = AsyncMock()  # pyright: ignore[reportPrivateUsage]
+
+    async with client as c:
+        assert c is client
+
+    client._client.close.assert_awaited_once()  # pyright: ignore[reportPrivateUsage]
+
+
+@pytest.mark.asyncio
+async def test_openai_client_async_context_manager_exception_propagation() -> None:
+    """Test that async context manager cleans up resources when an exception is raised inside the block."""
+    client = OpenAIChatCompletionClient(model="gpt-4o", api_key="fake_key")
+    client._client.close = AsyncMock()  # pyright: ignore[reportPrivateUsage]
+
+    with pytest.raises(ValueError, match="Test exception"):
+        async with client:
+            raise ValueError("Test exception")
+
+    client._client.close.assert_awaited_once()  # pyright: ignore[reportPrivateUsage]