From edb0fd2d8796443bd141a610725497a5f8829af3 Mon Sep 17 00:00:00 2001
From: Max Ind <maxind@google.com>
Date: Wed, 1 Jul 2026 14:42:47 -0700
Subject: [PATCH] refactor(telemetry): extract invocation span into
 _instrumentation.record_invocation

Co-authored-by: Max Ind <maxind@google.com>
PiperOrigin-RevId: 941303834
---
 src/google/adk/runners.py                    | 13 ++++++++++--
 src/google/adk/telemetry/_instrumentation.py | 21 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py
index 55ab638085..9802801dca 100644
--- a/src/google/adk/runners.py
+++ b/src/google/adk/runners.py
@@ -56,6 +56,7 @@
 from .sessions.base_session_service import BaseSessionService
 from .sessions.base_session_service import GetSessionConfig
 from .sessions.session import Session
+from .telemetry import _instrumentation
 from .telemetry.tracing import tracer
 from .tools.base_toolset import BaseToolset
 from .utils._debug_output import print_event
@@ -66,6 +67,10 @@
 
 logger = logging.getLogger('google_adk.' + __name__)
 
+# Silence unused warning.
+# tracer is imported for backwards compatibility, to avoid breaking change in the API.
+_ = tracer
+
 
 def _find_active_task_isolation_scope(session) -> Optional[str]:
   """Walk session backwards; find the active paused task agent's scope.
@@ -454,7 +459,9 @@ async def _run_node_async(
     Events flow through ic._event_queue via NodeRunner.
     """
 
-    with tracer.start_as_current_span('invocation'):
+    with _instrumentation.record_invocation(
+        entrypoint_node=node or self.agent, conversation_id=session_id
+    ):
       # 1. Setup
       session = await self._get_or_create_session(
           user_id=user_id, session_id=session_id
@@ -1040,7 +1047,9 @@ async def _run_with_trace(
         new_message: Optional[types.Content] = None,
         invocation_id: Optional[str] = None,
     ) -> AsyncGenerator[Event, None]:
-      with tracer.start_as_current_span('invocation'):
+      with _instrumentation.record_invocation(
+          entrypoint_node=self.agent, conversation_id=session_id
+      ):
         session = await self._get_or_create_session(
             user_id=user_id,
             session_id=session_id,
diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py
index 2284fb839d..1f8cc31feb 100644
--- a/src/google/adk/telemetry/_instrumentation.py
+++ b/src/google/adk/telemetry/_instrumentation.py
@@ -20,6 +20,7 @@
 import sys
 import time
 from typing import AsyncIterator
+from typing import Iterator
 from typing import TYPE_CHECKING
 
 from opentelemetry import trace
@@ -35,6 +36,7 @@
   from ..models.llm_request import LlmRequest
   from ..models.llm_response import LlmResponse
   from ..tools.base_tool import BaseTool
+  from ..workflow._base_node import BaseNode
 
 logger = logging.getLogger("google_adk." + __name__)
 
@@ -69,6 +71,25 @@ def _get_elapsed_s(
   return time.monotonic() - fallback_start
 
 
+@contextlib.contextmanager
+def record_invocation(
+    entrypoint_node: BaseNode | None,
+    conversation_id: str,
+) -> Iterator[None]:
+  """Top-level ``invocation`` span for a runner invocation.
+
+  Args:
+    entrypoint_node: The runner's root agent/node.
+    conversation_id: Session/conversation id.
+
+  Yields:
+    Nothing; the span is active for the duration of the block.
+  """
+  del entrypoint_node, conversation_id  # Unused until schema v2 lands.
+  with tracing.tracer.start_as_current_span("invocation"):
+    yield
+
+
 @dataclasses.dataclass
 class TelemetryContext:
   """Stores all telemetry related state."""