diff --git a/src/google/adk/agents/base_agent.py b/src/google/adk/agents/base_agent.py
index 88cd05f2..4e441a03 100644
--- a/src/google/adk/agents/base_agent.py
+++ b/src/google/adk/agents/base_agent.py
@@ -282,27 +282,22 @@ class BaseAgent(BaseModel):
       Event: the events generated by the agent.
     """
 
-    async def _run_with_trace() -> AsyncGenerator[Event, None]:
-      with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
-        ctx = self._create_invocation_context(parent_context)
-        tracing.trace_agent_invocation(span, self, ctx)
-        if event := await self._handle_before_agent_callback(ctx):
-          yield event
-        if ctx.end_invocation:
-          return
+    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
+      ctx = self._create_invocation_context(parent_context)
+      tracing.trace_agent_invocation(span, self, ctx)
+      if event := await self._handle_before_agent_callback(ctx):
+        yield event
+      if ctx.end_invocation:
+        return
 
-        async with Aclosing(self._run_async_impl(ctx)) as agen:
-          async for event in agen:
-            yield event
-
-        if ctx.end_invocation:
-          return
-
-        if event := await self._handle_after_agent_callback(ctx):
+      async with Aclosing(self._run_async_impl(ctx)) as agen:
+        async for event in agen:
           yield event
 
-    async with Aclosing(_run_with_trace()) as agen:
-      async for event in agen:
+      if ctx.end_invocation:
+        return
+
+      if event := await self._handle_after_agent_callback(ctx):
         yield event
 
   @final
@@ -320,24 +315,19 @@ class BaseAgent(BaseModel):
       Event: the events generated by the agent.
     """
 
-    async def _run_with_trace() -> AsyncGenerator[Event, None]:
-      with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
-        ctx = self._create_invocation_context(parent_context)
-        tracing.trace_agent_invocation(span, self, ctx)
-        if event := await self._handle_before_agent_callback(ctx):
-          yield event
-        if ctx.end_invocation:
-          return
+    with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
+      ctx = self._create_invocation_context(parent_context)
+      tracing.trace_agent_invocation(span, self, ctx)
+      if event := await self._handle_before_agent_callback(ctx):
+        yield event
+      if ctx.end_invocation:
+        return
 
-        async with Aclosing(self._run_live_impl(ctx)) as agen:
-          async for event in agen:
-            yield event
-
-        if event := await self._handle_after_agent_callback(ctx):
+      async with Aclosing(self._run_live_impl(ctx)) as agen:
+        async for event in agen:
           yield event
 
-    async with Aclosing(_run_with_trace()) as agen:
-      async for event in agen:
+      if event := await self._handle_after_agent_callback(ctx):
         yield event
 
   async def _run_async_impl(