Wait for cancellation

Author: tconley1428

temporalio/contrib/openai_agents/_mcp.py

@@ -17,7 +17,12 @@
     TIMEOUT_TYPE_HEARTBEAT,
     TIMEOUT_TYPE_SCHEDULE_TO_START,
 )
-from temporalio.exceptions import ActivityError, ApplicationError
+from temporalio.exceptions import (
+    ActivityError,
+    ApplicationError,
+    CancelledError,
+    is_cancelled_exception,
+)
 from temporalio.worker import PollerBehaviorSimpleMaximum, Worker
 from temporalio.workflow import ActivityConfig, ActivityHandle
 
@@ -213,6 +218,13 @@ async def connect(self) -> None:
     async def cleanup(self) -> None:
         if self._connect_handle:
             self._connect_handle.cancel()
+            try:
+                await self._connect_handle
+            except Exception as e:
+                if is_cancelled_exception(e):
+                    pass
+                else:
+                    raise
 
     async def __aenter__(self):
         await self.connect()

tests/contrib/openai_agents/test_openai.py

@@ -12,6 +12,7 @@
     Optional,
     Sequence,
     Union,
+    cast,
     no_type_check,
 )
 
@@ -2466,6 +2467,7 @@ async def get_prompt(
                 "list_tools",
                 "cleanup",
             ]
+            assert len(cast(StatefulMCPServerProvider, server)._servers) == 0
         else:
             assert tracking_server.calls == [
                 "connect",