@@ -32,6 +32,7 @@
 from typing import Union
 
 from google.genai import types
+from opentelemetry import context
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
@@ -87,6 +88,21 @@ class BaseAgentState(BaseModel):
 
 
 AgentState = TypeVar('AgentState', bound=BaseAgentState)
+_T = TypeVar('_T')
+
+
+async def _with_caller_context(
+    agen: AsyncGenerator[_T, None],
+    caller_ctx: context.Context,
+) -> AsyncGenerator[_T, None]:
+  """Wraps an async generator to attach caller_ctx around each yield."""
+  async with Aclosing(agen) as a:
+    async for item in a:
+      token = context.attach(caller_ctx)
+      try:
+        yield item
+      finally:
+        context.detach(token)
 
 
 # TODO: drop the explicit abc.ABC base once BaseNode surfaces ABCMeta to
@@ -295,26 +311,33 @@ async def run_async(
       Event: the events generated by the agent.
     """
 
-    ctx = self._create_invocation_context(parent_context)
-    async with _instrumentation.record_agent_invocation(ctx, self):
-      try:
-        if event := await self._handle_before_agent_callback(ctx):
-          yield event
-        if ctx.end_invocation:
-          return
+    caller_ctx = context.get_current()
 
-        async with Aclosing(self._run_async_impl(ctx)) as agen:
-          async for event in agen:
+    async def _run() -> AsyncGenerator[Event, None]:
+      ctx = self._create_invocation_context(parent_context)
+      async with _instrumentation.record_agent_invocation(ctx, self):
+        try:
+          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 ctx.end_invocation:
+            return
 
-        if event := await self._handle_after_agent_callback(ctx):
-          yield event
-      except Exception as e:
-        await self._handle_agent_error_callback(ctx, e)
-        raise
+          if event := await self._handle_after_agent_callback(ctx):
+            yield event
+        except Exception as e:
+          await self._handle_agent_error_callback(ctx, e)
+          raise
+
+    async with Aclosing(_with_caller_context(_run(), caller_ctx)) as agen:
+      async for event in agen:
+        yield event
 
   @override
   async def _run_impl(
@@ -350,23 +373,30 @@ async def run_live(
       Event: the events generated by the agent.
     """
 
-    ctx = self._create_invocation_context(parent_context)
-    async with _instrumentation.record_agent_invocation(ctx, self):
-      try:
-        if event := await self._handle_before_agent_callback(ctx):
-          yield event
-        if ctx.end_invocation:
-          return
+    caller_ctx = context.get_current()
+
+    async def _run() -> AsyncGenerator[Event, None]:
+      ctx = self._create_invocation_context(parent_context)
+      async with _instrumentation.record_agent_invocation(ctx, self):
+        try:
+          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
 
-        async with Aclosing(self._run_live_impl(ctx)) as agen:
-          async for event in agen:
+          if event := await self._handle_after_agent_callback(ctx):
             yield event
+        except Exception as e:
+          await self._handle_agent_error_callback(ctx, e)
+          raise
 
-        if event := await self._handle_after_agent_callback(ctx):
-          yield event
-      except Exception as e:
-        await self._handle_agent_error_callback(ctx, e)
-        raise
+    async with Aclosing(_with_caller_context(_run(), caller_ctx)) as agen:
+      async for event in agen:
+        yield event
 
   async def _run_async_impl(
       self, ctx: InvocationContext