use 3.11 create_task when possible to avoid an extra context copy

Author: VegetarianOrc

temporalio/contrib/openai_agents/_trace_interceptor.py

@@ -5,6 +5,7 @@
 import asyncio
 import contextvars
 import random
+import sys
 import uuid
 from contextlib import contextmanager
 from typing import Any, Mapping, Optional, Protocol, Type
@@ -424,9 +425,14 @@ async def start_child_workflow(
             data={"workflow": input.workflow},
             input=input,
         )
-        handle = await ctx.run(
-            asyncio.create_task, self.next.start_child_workflow(input)
-        )
+        if sys.version_info >= (3, 11):
+            handle = await asyncio.create_task(
+                self.next.start_child_workflow(input), context=ctx
+            )
+        else:
+            handle = await ctx.run(
+                asyncio.create_task, self.next.start_child_workflow(input)
+            )
         if span:
             handle.add_done_callback(lambda _: ctx.run(span.finish))  # type: ignore
         return handle

temporalio/contrib/opentelemetry.py

@@ -9,6 +9,7 @@
     contextmanager,
 )
 from dataclasses import dataclass
+import sys
 from types import TracebackType
 from typing import (
     Any,
@@ -394,7 +395,12 @@ async def execute_workflow(
         :py:meth:`temporalio.worker.WorkflowInboundInterceptor.execute_workflow`.
         """
         with self._top_level_workflow_context(success_is_complete=True) as ctx:
-            return await ctx.run(asyncio.create_task, self._execute_workflow(input))
+            if sys.version_info >= (3, 11):
+                return await asyncio.create_task(
+                    self._execute_workflow(input), context=ctx
+                )
+            else:
+                return await ctx.run(asyncio.create_task, self._execute_workflow(input))
 
     async def _execute_workflow(
         self, input: temporalio.worker.ExecuteWorkflowInput
@@ -414,7 +420,12 @@ async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> Non
         # Create a span in the current context for the signal and link any
         # header given
         with self._top_level_workflow_context(success_is_complete=False) as ctx:
-            return await ctx.run(asyncio.create_task, self._handle_signal(input))
+            if sys.version_info >= (3, 11):
+                return await asyncio.create_task(
+                    self._handle_signal(input), context=ctx
+                )
+            else:
+                return await ctx.run(asyncio.create_task, self._handle_signal(input))
 
     async def _handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None:
         """Implementation of
@@ -511,9 +522,14 @@ async def handle_update_handler(
         :py:meth:`temporalio.worker.WorkflowInboundInterceptor.handle_update_handler`.
         """
         with self._top_level_workflow_context(success_is_complete=False) as ctx:
-            return await ctx.run(
-                asyncio.create_task, self._handle_update_handler(input)
-            )
+            if sys.version_info >= (3, 11):
+                return await asyncio.create_task(
+                    self._handle_update_handler(input), context=ctx
+                )
+            else:
+                return await ctx.run(
+                    asyncio.create_task, self._handle_update_handler(input)
+                )
 
     async def _handle_update_handler(
         self, input: temporalio.worker.HandleUpdateInput

tests/nexus/test_handler.py

@@ -20,7 +20,7 @@
 import uuid
 from collections.abc import Mapping
 from concurrent.futures.thread import ThreadPoolExecutor
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from types import MappingProxyType
 from typing import Any, Callable, Optional, Union
 
@@ -313,7 +313,9 @@ async def non_serializable_output(
 class SuccessfulResponse:
     status_code: int
     body_json: Optional[Union[dict[str, Any], Callable[[dict[str, Any]], bool]]] = None
-    headers: Mapping[str, str] = SUCCESSFUL_RESPONSE_HEADERS
+    headers: Mapping[str, str] = field(
+        default_factory=lambda: SUCCESSFUL_RESPONSE_HEADERS
+    )
 
 
 @dataclass
@@ -325,7 +327,9 @@ class UnsuccessfulResponse:
     # Expected value of inverse of non_retryable attribute of exception.
     retryable_exception: bool = True
     body_json: Optional[Callable[[dict[str, Any]], bool]] = None
-    headers: Mapping[str, str] = UNSUCCESSFUL_RESPONSE_HEADERS
+    headers: Mapping[str, str] = field(
+        default_factory=lambda: UNSUCCESSFUL_RESPONSE_HEADERS
+    )
 
 
 class _TestCase: