.

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/__init__.py

@@ -41,6 +41,8 @@
 if TYPE_CHECKING:
     from typing import Any
 
+    from agents.run_internal.run_steps import SingleStepResult
+
 
 def _patch_runner() -> None:
     # Create the root span for one full agent run (including eventual handoffs)
@@ -119,7 +121,9 @@ def new_wrapped_get_model(
             agents.run_internal.run_loop.get_model = new_wrapped_get_model
 
             @wraps(run_loop.run_single_turn)
-            async def patched_run_single_turn(*args: "Any", **kwargs: "Any") -> "Any":
+            async def patched_run_single_turn(
+                *args: "Any", **kwargs: "Any"
+            ) -> "SingleStepResult":
                 return await _run_single_turn(run_loop.run_single_turn, *args, **kwargs)
 
             agents.run.run_single_turn = patched_run_single_turn
@@ -150,10 +154,10 @@ def old_wrapped_get_model(
 
         original_run_single_turn = AgentRunner._run_single_turn
 
-        @wraps(AgentRunner._run_single_turn)
+        @wraps(AgentRunner._run_single_turn.__func__)
         async def old_wrapped_run_single_turn(
             cls: "agents.Runner", *args: "Any", **kwargs: "Any"
-        ) -> "Any":
+        ) -> "SingleStepResult":
             return await _run_single_turn(original_run_single_turn, *args, **kwargs)
 
         agents.run.AgentRunner._run_single_turn = classmethod(

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -84,7 +84,7 @@ async def _run_single_turn(
     original_run_single_turn: "Callable[..., Awaitable[SingleStepResult]]",
     *args: "Any",
     **kwargs: "Any",
-) -> "Any":
+) -> "SingleStepResult":
     """
     Patched _run_single_turn that
     - creates agent invocation spans if there is no already active agent invocation span.