chore: import upstream snapshot with attribution

libs/python/agent/cua_agent/callbacks/otel.py

@@ -0,0 +1,212 @@
+"""
+OpenTelemetry callback handler for Computer-Use Agent (cua-agent).
+
+Instruments agent operations for the Four Golden Signals:
+- Latency: Operation duration
+- Traffic: Operation counts
+- Errors: Error counts
+- Saturation: Concurrent operations
+"""
+
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import AsyncCallbackHandler
+
+# Import OTEL functions - these are available when cua-core[telemetry] is installed
+try:
+    from cua_core.telemetry import (
+        create_span,
+        is_otel_enabled,
+        record_error,
+        record_operation,
+        record_tokens,
+        track_concurrent,
+    )
+
+    OTEL_AVAILABLE = True
+except ImportError:
+    OTEL_AVAILABLE = False
+
+    def is_otel_enabled() -> bool:
+        return False
+
+
+class OtelCallback(AsyncCallbackHandler):
+    """
+    OpenTelemetry callback handler for instrumentation.
+
+    Tracks:
+    - Agent session lifecycle (start/end)
+    - Agent run lifecycle (start/end with duration)
+    - Individual steps (with duration)
+    - Computer actions (with duration)
+    - Token usage
+    - Errors
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize OTEL callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+        # Timing state
+        self.run_start_time: Optional[float] = None
+        self.step_start_time: Optional[float] = None
+        self.step_count = 0
+
+        # Span management
+        self._session_span: Optional[Any] = None
+        self._run_span: Optional[Any] = None
+
+        # Track concurrent sessions
+        self._concurrent_tracker: Optional[Any] = None
+
+    def _get_agent_type(self) -> str:
+        """Get the agent loop type name."""
+        if hasattr(self.agent, "agent_loop") and self.agent.agent_loop is not None:
+            return type(self.agent.agent_loop).__name__
+        return "unknown"
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Called at the start of an agent run loop."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.run_start_time = time.perf_counter()
+        self.step_start_time = self.run_start_time
+        self.step_count = 0
+
+    async def on_run_end(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> None:
+        """Called at the end of an agent run loop."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        if self.run_start_time is not None:
+            duration = time.perf_counter() - self.run_start_time
+
+            # Record run metrics
+            record_operation(
+                operation="agent.run",
+                duration_seconds=duration,
+                status="success",
+                model=self.model,
+                steps=self.step_count,
+            )
+
+        self.run_start_time = None
+
+    async def on_responses(self, kwargs: Dict[str, Any], responses: Dict[str, Any]) -> None:
+        """Called when responses are received (each step)."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        self.step_count += 1
+        current_time = time.perf_counter()
+
+        # Calculate step duration if we have a start time
+        if self.step_start_time is not None:
+            step_duration = current_time - self.step_start_time
+            record_operation(
+                operation="agent.step",
+                duration_seconds=step_duration,
+                status="success",
+                model=self.model,
+                step_number=self.step_count,
+            )
+
+        # Start timing next step
+        self.step_start_time = current_time
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Called when usage information is received."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+
+        if prompt_tokens > 0 or completion_tokens > 0:
+            record_tokens(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                model=self.model,
+            )
+
+    async def on_computer_call_start(self, item: Dict[str, Any]) -> None:
+        """Called when a computer call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+    async def on_computer_call_end(
+        self, item: Dict[str, Any], result: List[Dict[str, Any]]
+    ) -> None:
+        """Called when a computer call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        action = item.get("action", {})
+        action_type = action.get("type", "unknown")
+
+        # Record computer action metric
+        # Note: We don't have precise timing here, so we record with 0 duration
+        # The actual timing should be done in the computer module
+        record_operation(
+            operation=f"computer.action.{action_type}",
+            duration_seconds=0,  # Timing handled elsewhere
+            status="success",
+            model=self.model,
+        )
+
+    async def on_api_start(self, kwargs: Dict[str, Any]) -> None:
+        """Called when an LLM API call is about to start."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+    async def on_api_end(self, kwargs: Dict[str, Any], result: Any) -> None:
+        """Called when an LLM API call has completed."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+
+class OtelErrorCallback(AsyncCallbackHandler):
+    """
+    Callback that captures errors and sends them to OTEL.
+
+    Should be added early in the callback chain to catch all errors.
+    """
+
+    def __init__(self, agent: Any):
+        """
+        Initialize error callback.
+
+        Args:
+            agent: The ComputerAgent instance
+        """
+        self.agent = agent
+        self.model = getattr(agent, "model", "unknown")
+
+    async def on_error(self, error: Exception, context: Dict[str, Any]) -> None:
+        """Called when an error occurs during agent execution."""
+        if not OTEL_AVAILABLE or not is_otel_enabled():
+            return
+
+        error_type = type(error).__name__
+        operation = context.get("operation", "unknown")
+
+        # Record error metric
+        record_error(
+            error_type=error_type,
+            operation=operation,
+            model=self.model,
+        )