refactor: adjust computation of workflow.steps metric and add new unit tests

PiperOrigin-RevId: 910601469

Author: google-genai-bot

src/google/adk/telemetry/_instrumentation.py

@@ -34,6 +34,7 @@
 if TYPE_CHECKING:
   from ..agents.base_agent import BaseAgent
   from ..agents.invocation_context import InvocationContext
+  from ..tools.base_tool import BaseTool
 
 
 def _get_elapsed_ms(span: trace.Span | None, fallback_start: float) -> float:
@@ -97,7 +98,7 @@ async def record_agent_invocation(
       )
       _metrics.record_agent_request_size(agent.name, ctx.user_content)
       _metrics.record_agent_response_size(agent.name, ctx.session.events)
-      _metrics.record_agent_workflow_steps(agent.name, len(ctx.session.events))
+      _metrics.record_agent_workflow_steps(agent.name, ctx.session.events)
     except Exception:  # pylint: disable=broad-exception-caught
       logger.exception(
           "Failed to record agent metrics for agent %s", agent.name

src/google/adk/telemetry/_metrics.py

@@ -98,12 +98,11 @@ def record_agent_response_size(agent_name: str, events: list[Event]):
   _agent_response_size.record(size, attributes=attrs)
 
 
-def record_agent_workflow_steps(agent_name: str, steps_count: int):
-  """Records the number of steps in the agent workflow."""
-  attrs = {
-      gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name,
-  }
-  _agent_workflow_steps.record(steps_count, attributes=attrs)
+def record_agent_workflow_steps(agent_name: str, events: list[Event]):
+  """Records the number of steps in the agent workflow by counting the number of events."""
+  attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name}
+  count = sum(1 for event in events if event.author == agent_name)
+  _agent_workflow_steps.record(count, attributes=attrs)
 
 
 def record_tool_execution_duration(

tests/unittests/telemetry/test_functional.py

@@ -349,7 +349,8 @@ async def generate_random_number():
   got_steps = _extract_metrics(metrics_list, "gen_ai.agent.workflow.steps")
   assert len(got_steps) == 1
   want_steps = [
-      MetricPoint(attributes={"gen_ai.agent.name": "complex_agent"}, value=6)
+      # (tool call + result) x 2 + text response = 5 steps
+      MetricPoint(attributes={"gen_ai.agent.name": "complex_agent"}, value=5)
   ]
   assert got_steps == want_steps
 

tests/unittests/telemetry/test_instrumentation.py

@@ -0,0 +1,82 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=protected-access
+
+import time
+from unittest import mock
+
+from google.adk.telemetry import _instrumentation
+from opentelemetry import trace
+
+
+def test_get_elapsed_ms_span_none():
+  """Tests fallback when span is None."""
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _instrumentation._get_elapsed_ms(None, start_time)
+  assert elapsed == 2000.0  # (12 - 10) * 1000
+
+
+def test_get_elapsed_ms_span_valid():
+  """Tests duration calculation with valid span times."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000  # 1s in ns
+  mock_span.end_time = 2000000000  # 2s in ns
+  elapsed = _instrumentation._get_elapsed_ms(mock_span, time.monotonic())
+  assert elapsed == 1000.0  # (2 - 1) * 1000 ms
+
+
+def test_get_elapsed_ms_span_missing_start():
+  """Tests fallback when start_time is missing."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  del mock_span.start_time
+  mock_span.end_time = 2000000000
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _instrumentation._get_elapsed_ms(mock_span, start_time)
+  assert elapsed == 2000.0
+
+
+def test_get_elapsed_ms_span_missing_end():
+  """Tests fallback when end_time is missing."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000
+  del mock_span.end_time
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _instrumentation._get_elapsed_ms(mock_span, start_time)
+  assert elapsed == 2000.0
+
+
+def test_get_elapsed_ms_span_non_int_start():
+  """Tests fallback when start_time is not an integer."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000.0
+  mock_span.end_time = 2000000000
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _instrumentation._get_elapsed_ms(mock_span, start_time)
+  assert elapsed == 2000.0
+
+
+def test_get_elapsed_ms_span_non_int_end():
+  """Tests fallback when end_time is not an integer."""
+  mock_span = mock.MagicMock(spec=trace.Span)
+  mock_span.start_time = 1000000000
+  mock_span.end_time = 2000000000.0
+  start_time = 10.0
+  with mock.patch("time.monotonic", return_value=12.0):
+    elapsed = _instrumentation._get_elapsed_ms(mock_span, start_time)
+  assert elapsed == 2000.0

tests/unittests/telemetry/test_metrics.py

@@ -75,13 +75,14 @@ def create_histogram_side_effect(name, **_kwargs):
 
 def test_record_agent_request_size(mock_meter_setup):
   """Tests record_agent_request_size records correctly."""
+  user_content = "hello"
   _metrics.record_agent_request_size(
-      "test_agent", types.Content(parts=[types.Part(text="hello")])
+      "test_agent", types.Content(parts=[types.Part(text=user_content)])
   )
   request_size_hist = mock_meter_setup["request_size"]
   request_size_hist.record.assert_called_once()
   args, kwargs = request_size_hist.record.call_args
-  assert args[0] == 5  # len('hello')
+  assert args[0] == len(user_content)
   want_attributes = {
       "gen_ai.agent.name": "test_agent",
   }
@@ -121,7 +122,7 @@ def test_record_agent_response_size(mock_meter_setup):
   response_text = "response"
   event = mock.MagicMock(
       author="test_agent",
-      content=types.Content(parts=[types.Part(text="response")]),
+      content=types.Content(parts=[types.Part(text=response_text)]),
   )
   _metrics.record_agent_response_size("test_agent", [event])
   response_size_hist = mock_meter_setup["response_size"]
@@ -134,14 +135,19 @@ def test_record_agent_response_size(mock_meter_setup):
 
 def test_record_agent_workflow_steps(mock_meter_setup):
   """Tests record_agent_workflow_steps records correctly."""
-  _metrics.record_agent_workflow_steps("test_agent", 5)
+  _metrics.record_agent_workflow_steps(
+      "test_agent",
+      [
+          mock.MagicMock(author="test_agent"),
+          mock.MagicMock(author="test_agent"),
+          mock.MagicMock(author="other_agent"),
+      ],
+  )
   steps_hist = mock_meter_setup["steps"]
   steps_hist.record.assert_called_once()
   args, kwargs = steps_hist.record.call_args
-  assert args[0] == 5
-  want_attributes = {
-      "gen_ai.agent.name": "test_agent",
-  }
+  assert args[0] == 2
+  want_attributes = {"gen_ai.agent.name": "test_agent"}
   assert kwargs["attributes"] == want_attributes