improve agent metrics integration test and cleanup fixtures

- simplified test to use telemetry.query_metrics for verification
- test now validates actual queryable metrics data
- verified by query metrics functionality added in llamastack#3074

Author: skamenan7

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -63,7 +63,7 @@
     UserMessage,
 )
 from llama_stack.apis.safety import Safety
-from llama_stack.apis.telemetry import MetricEvent, Telemetry
+from llama_stack.apis.telemetry import MetricEvent, MetricType, Telemetry
 from llama_stack.apis.tools import ToolGroups, ToolInvocationResult, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
 from llama_stack.core.datatypes import AccessRule
@@ -124,6 +124,9 @@ def __init__(
             output_shields=agent_config.output_shields,
         )
 
+        # Initialize workflow start time to None
+        self._workflow_start_time: float | None = None
+
     def turn_to_messages(self, turn: Turn) -> list[Message]:
         messages = []
 
@@ -174,14 +177,23 @@ async def create_session(self, name: str) -> str:
         return await self.storage.create_session(name)
 
     def _emit_metric(
-        self, metric_name: str, value: int | float, unit: str, attributes: dict[str, str] | None = None
+        self,
+        metric_name: str,
+        value: int | float,
+        unit: str,
+        attributes: dict[str, str] | None = None,
+        metric_type: MetricType | None = None,
     ) -> None:
         """Emit a single metric event"""
+        logger.info(f"_emit_metric called: {metric_name} = {value} {unit}")
+
         if not self.telemetry_api:
+            logger.warning(f"No telemetry_api available for metric {metric_name}")
             return
 
         span = get_current_span()
         if not span:
+            logger.warning(f"No current span available for metric {metric_name}")
             return
 
         context = span.get_span_context()
@@ -193,22 +205,42 @@ def _emit_metric(
             timestamp=time.time(),
             unit=unit,
             attributes={"agent_id": self.agent_id, **(attributes or {})},
+            metric_type=metric_type,
         )
 
-        # Create task with name for better debugging and potential cleanup
+        # Create task with name for better debugging and capture any async errors
         task_name = f"metric-{metric_name}-{self.agent_id}"
-        asyncio.create_task(self.telemetry_api.log_event(metric), name=task_name)
+        logger.info(f"Creating telemetry task: {task_name}")
+        task = asyncio.create_task(self.telemetry_api.log_event(metric), name=task_name)
+
+        def _on_metric_task_done(t: asyncio.Task) -> None:
+            try:
+                exc = t.exception()
+            except asyncio.CancelledError:
+                logger.debug("Metric task %s was cancelled", task_name)
+                return
+            if exc is not None:
+                logger.warning("Metric task %s failed: %s", task_name, exc)
+
+        # Only add callback if task creation succeeded (not None from mocking)
+        if task is not None:
+            task.add_done_callback(_on_metric_task_done)
 
     def _track_step(self):
-        self._emit_metric("llama_stack_agent_steps_total", 1, "1")
+        logger.info("_track_step called")
+        self._emit_metric("llama_stack_agent_steps_total", 1, "1", metric_type=MetricType.COUNTER)
 
     def _track_workflow(self, status: str, duration: float):
-        self._emit_metric("llama_stack_agent_workflows_total", 1, "1", {"status": status})
-        self._emit_metric("llama_stack_agent_workflow_duration_seconds", duration, "s")
+        logger.info(f"_track_workflow called: status={status}, duration={duration:.2f}s")
+        self._emit_metric("llama_stack_agent_workflows_total", 1, "1", {"status": status}, MetricType.COUNTER)
+        self._emit_metric(
+            "llama_stack_agent_workflow_duration_seconds", duration, "s", metric_type=MetricType.HISTOGRAM
+        )
 
     def _track_tool(self, tool_name: str):
+        logger.info(f"_track_tool called: {tool_name}")
         normalized_name = "rag" if tool_name == "knowledge_search" else tool_name
-        self._emit_metric("llama_stack_agent_tool_calls_total", 1, "1", {"tool": normalized_name})
+        self._emit_metric("llama_stack_agent_tool_calls_total", 1, "1", {"tool": normalized_name}, MetricType.COUNTER)
 
     async def get_messages_from_turns(self, turns: list[Turn]) -> list[Message]:
         messages = []
@@ -244,6 +276,9 @@ async def resume_turn(self, request: AgentTurnResumeRequest) -> AsyncGenerator:
             if self.agent_config.name:
                 span.set_attribute("agent_name", self.agent_config.name)
 
+        # Set workflow start time for resume operations
+        self._workflow_start_time = time.time()
+
         await self._initialize_tools()
         async for chunk in self._run_turn(request):
             yield chunk
@@ -255,6 +290,9 @@ async def _run_turn(
     ) -> AsyncGenerator:
         assert request.stream is True, "Non-streaming not supported"
 
+        # Track workflow start time for metrics
+        self._workflow_start_time = time.time()
+
         is_resume = isinstance(request, AgentTurnResumeRequest)
         session_info = await self.storage.get_session_info(request.session_id)
         if session_info is None:
@@ -356,6 +394,10 @@ async def _run_turn(
                 )
             )
         else:
+            # Track workflow completion when turn is actually complete
+            workflow_duration = time.time() - (self._workflow_start_time or time.time())
+            self._track_workflow("completed", workflow_duration)
+
             chunk = AgentTurnResponseStreamChunk(
                 event=AgentTurnResponseEvent(
                     payload=AgentTurnResponseTurnCompletePayload(
@@ -771,6 +813,7 @@ async def _run(
 
                         # Track step execution metric
                         self._track_step()
+                        self._track_tool(tool_call.tool_name)
 
                         # Add the result message to input_messages for the next iteration
                         input_messages.append(result_message)

llama_stack/providers/inline/telemetry/meta_reference/telemetry.py

@@ -12,7 +12,9 @@
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics._internal.aggregation import ExplicitBucketHistogramAggregation
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+from opentelemetry.sdk.metrics.view import View
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -110,7 +112,17 @@ def __init__(self, config: TelemetryConfig, deps: dict[Api, Any]) -> None:
 
                 if TelemetrySink.OTEL_METRIC in self.config.sinks:
                     metric_reader = PeriodicExportingMetricReader(OTLPMetricExporter())
-                    metric_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
+
+                    # decent default buckets for agent workflow timings
+                    hist_buckets = [0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0, 50.0, 100.0]
+                    views = [
+                        View(
+                            instrument_type=metrics.Histogram,
+                            aggregation=ExplicitBucketHistogramAggregation(boundaries=hist_buckets),
+                        )
+                    ]
+
+                    metric_provider = MeterProvider(resource=resource, metric_readers=[metric_reader], views=views)
                     metrics.set_meter_provider(metric_provider)
 
             if TelemetrySink.SQLITE in self.config.sinks:
@@ -140,8 +152,6 @@ async def log_event(self, event: Event, ttl_seconds: int = 604800) -> None:
             self._log_metric(event)
         elif isinstance(event, StructuredLogEvent):
             self._log_structured(event, ttl_seconds)
-        else:
-            raise ValueError(f"Unknown event type: {event}")
 
     async def query_metrics(
         self,
@@ -211,7 +221,7 @@ def _get_or_create_counter(self, name: str, unit: str) -> metrics.Counter:
             _GLOBAL_STORAGE["counters"][name] = self.meter.create_counter(
                 name=name,
                 unit=unit,
-                description=f"Counter for {name}",
+                description=name.replace("_", " "),
             )
         return _GLOBAL_STORAGE["counters"][name]
 
@@ -221,7 +231,7 @@ def _get_or_create_gauge(self, name: str, unit: str) -> metrics.ObservableGauge:
             _GLOBAL_STORAGE["gauges"][name] = self.meter.create_gauge(
                 name=name,
                 unit=unit,
-                description=f"Gauge for {name}",
+                description=name.replace("_", " "),
             )
         return _GLOBAL_STORAGE["gauges"][name]
 
@@ -265,7 +275,6 @@ def _log_metric(self, event: MetricEvent) -> None:
             histogram = self._get_or_create_histogram(
                 event.metric,
                 event.unit,
-                [0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0, 50.0, 100.0],
             )
             histogram.record(event.value, attributes=event.attributes)
         elif event.metric_type == MetricType.UP_DOWN_COUNTER:
@@ -281,17 +290,17 @@ def _get_or_create_up_down_counter(self, name: str, unit: str) -> metrics.UpDown
             _GLOBAL_STORAGE["up_down_counters"][name] = self.meter.create_up_down_counter(
                 name=name,
                 unit=unit,
-                description=f"UpDownCounter for {name}",
+                description=name.replace("_", " "),
             )
         return _GLOBAL_STORAGE["up_down_counters"][name]
 
-    def _get_or_create_histogram(self, name: str, unit: str, buckets: list[float] | None = None) -> metrics.Histogram:
+    def _get_or_create_histogram(self, name: str, unit: str) -> metrics.Histogram:
         assert self.meter is not None
         if name not in _GLOBAL_STORAGE["histograms"]:
             _GLOBAL_STORAGE["histograms"][name] = self.meter.create_histogram(
                 name=name,
                 unit=unit,
-                description=f"Histogram for {name}",
+                description=name.replace("_", " "),
             )
         return _GLOBAL_STORAGE["histograms"][name]
 

tests/integration/agents/conftest.py

@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import AsyncGenerator, Callable
+from pathlib import Path
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from llama_stack.apis.inference import ToolDefinition
+from llama_stack.apis.tools import ToolInvocationResult
+from llama_stack.providers.inline.agents.meta_reference.agent_instance import ChatAgent
+from llama_stack.providers.inline.telemetry.meta_reference.config import (
+    TelemetryConfig,
+    TelemetrySink,
+)
+from llama_stack.providers.inline.telemetry.meta_reference.telemetry import (
+    TelemetryAdapter,
+)
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite.sqlite import SqliteKVStoreImpl
+from llama_stack.providers.utils.telemetry import tracing as telemetry_tracing
+
+
+@pytest.fixture
+def make_agent_fixture():
+    def _make(telemetry, kvstore) -> ChatAgent:
+        agent = ChatAgent(
+            agent_id="test-agent",
+            agent_config=Mock(),
+            inference_api=Mock(),
+            safety_api=Mock(),
+            tool_runtime_api=Mock(),
+            tool_groups_api=Mock(),
+            vector_io_api=Mock(),
+            telemetry_api=telemetry,
+            persistence_store=kvstore,
+            created_at="2025-01-01T00:00:00Z",
+            policy=[],
+        )
+        agent.agent_config.client_tools = []
+        agent.agent_config.max_infer_iters = 5
+        agent.input_shields = []
+        agent.output_shields = []
+        agent.tool_defs = [
+            ToolDefinition(tool_name="web_search", description="", parameters={}),
+            ToolDefinition(tool_name="knowledge_search", description="", parameters={}),
+        ]
+        agent.tool_name_to_args = {}
+
+        # Stub tool runtime invoke_tool
+        async def _mock_invoke_tool(
+            *args: Any,
+            tool_name: str | None = None,
+            kwargs: dict | None = None,
+            **extra: Any,
+        ):
+            return ToolInvocationResult(content="Tool execution result")
+
+        agent.tool_runtime_api.invoke_tool = _mock_invoke_tool
+        return agent
+
+    return _make
+
+
+def _chat_stream(tool_name: str | None, content: str = ""):
+    from llama_stack.apis.common.content_types import (
+        TextDelta,
+        ToolCallDelta,
+        ToolCallParseStatus,
+    )
+    from llama_stack.apis.inference import (
+        ChatCompletionResponseEvent,
+        ChatCompletionResponseEventType,
+        ChatCompletionResponseStreamChunk,
+        StopReason,
+    )
+    from llama_stack.models.llama.datatypes import ToolCall
+
+    async def gen():
+        # Start
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.start,
+                delta=TextDelta(text=""),
+            )
+        )
+
+        # Content
+        if content:
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=TextDelta(text=content),
+                )
+            )
+
+        # Tool call if specified
+        if tool_name:
+            yield ChatCompletionResponseStreamChunk(
+                event=ChatCompletionResponseEvent(
+                    event_type=ChatCompletionResponseEventType.progress,
+                    delta=ToolCallDelta(
+                        tool_call=ToolCall(call_id="call_0", tool_name=tool_name, arguments={}),
+                        parse_status=ToolCallParseStatus.succeeded,
+                    ),
+                )
+            )
+
+        # Complete
+        yield ChatCompletionResponseStreamChunk(
+            event=ChatCompletionResponseEvent(
+                event_type=ChatCompletionResponseEventType.complete,
+                delta=TextDelta(text=""),
+                stop_reason=StopReason.end_of_turn,
+            )
+        )
+
+    return gen()
+
+
+@pytest.fixture
+async def telemetry(tmp_path: Path) -> AsyncGenerator[TelemetryAdapter, None]:
+    db_path = tmp_path / "trace_store.db"
+    cfg = TelemetryConfig(
+        sinks=[TelemetrySink.CONSOLE, TelemetrySink.SQLITE],
+        sqlite_db_path=str(db_path),
+    )
+    telemetry = TelemetryAdapter(cfg, deps={})
+    telemetry_tracing.setup_logger(telemetry)
+    try:
+        yield telemetry
+    finally:
+        await telemetry.shutdown()
+
+
+@pytest.fixture
+async def kvstore(tmp_path: Path) -> SqliteKVStoreImpl:
+    kv_path = tmp_path / "agent_kvstore.db"
+    kv = SqliteKVStoreImpl(SqliteKVStoreConfig(db_path=str(kv_path)))
+    await kv.initialize()
+    return kv
+
+
+@pytest.fixture
+def span_patch():
+    with (
+        patch("llama_stack.providers.inline.agents.meta_reference.agent_instance.get_current_span") as mock_span,
+        patch(
+            "llama_stack.providers.utils.telemetry.tracing.generate_span_id",
+            return_value="0000000000000abc",
+        ),
+    ):
+        mock_span.return_value = Mock(get_span_context=Mock(return_value=Mock(trace_id=0x123, span_id=0xABC)))
+        yield
+
+
+@pytest.fixture
+def make_completion_fn() -> Callable[[str | None, str], Callable]:
+    def _factory(tool_name: str | None = None, content: str = "") -> Callable:
+        async def chat_completion(*args: Any, **kwargs: Any):
+            return _chat_stream(tool_name, content)
+
+        return chat_completion
+
+    return _factory