refactor(telemetry): extract invocation span into _instrumentation.record_invocation

Co-authored-by: Max Ind <maxind@google.com>
PiperOrigin-RevId: 939850648

Author: RKest

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,

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."""