feat(observability): RPC + workflow + DB-pool metrics for the agent JSON-RPC endpoint

agentex/otel/otel-collector-config.yaml

@@ -21,10 +21,14 @@ exporters:
     sampling_initial: 5
     sampling_thereafter: 200
 
-  # Expose Prometheus endpoint for scraping
+  # Expose Prometheus endpoint for scraping.
+  # No `namespace:` on purpose — the cluster's otel-operator -> daemonset
+  # pipeline adds no prefix, so metric names there are exactly their instrument
+  # names (e.g. agentex_rpc_requests_total, db_client_connection_wait_time_seconds).
+  # A namespace here would double-prefix agentex.* metrics (agentex_agentex_...)
+  # and make locally-tested PromQL diverge from Mimir.
   prometheus:
     endpoint: 0.0.0.0:8889
-    namespace: agentex
     send_timestamps: true
     metric_expiration: 5m
 

agentex/src/api/routes/agents.py

@@ -1,4 +1,5 @@
 import secrets
+import time
 from collections.abc import AsyncIterator
 
 from fastapi import APIRouter, HTTPException, Query, Request
@@ -43,6 +44,9 @@
     DAuthorizedResourceIds,
 )
 from src.utils.logging import make_logger
+from src.utils.rpc_metrics import (
+    record_rpc_request,
+)
 from src.utils.task_authorization import check_task_or_collapse_to_404
 
 logger = make_logger(__name__)
@@ -549,6 +553,7 @@ async def _handle_sync_rpc(
     request_headers: dict[str, str] | None = None,
 ) -> AgentRPCResponse:
     """Handle synchronous JSON-RPC requests."""
+    start = time.perf_counter()
     try:
         result_entity = await agents_acp_use_case.handle_rpc_request(
             agent_id=agent_id,
@@ -580,6 +585,11 @@ async def _handle_sync_rpc(
         # else:
         #     raise ValueError(f"Unsupported method: {request.method}")
         # logger.info(f"AgentRPCResponse Result: {result}")
+        record_rpc_request(
+            method=request.method.value,
+            streaming=False,
+            duration_s=time.perf_counter() - start,
+        )
         return AgentRPCResponse.model_validate(
             {
                 "id": request.id,
@@ -591,10 +601,24 @@ async def _handle_sync_rpc(
     except ValidationError as e:
         logger.error(f"Validation error in RPC request: {e}", exc_info=True)
         error = JSONRPCError(code=-32602, message=f"Invalid parameters: {e}")
+        record_rpc_request(
+            method=request.method.value,
+            streaming=False,
+            duration_s=time.perf_counter() - start,
+            error_code=error.code,
+            error_type=type(e).__name__,
+        )
         return AgentRPCResponse(id=request.id, error=error.model_dump(), result=None)
     except Exception as e:
         logger.error(f"Error handling JSON-RPC request: {e}", exc_info=True)
         error = JSONRPCError(code=-32603, message=str(e))
+        record_rpc_request(
+            method=request.method.value,
+            streaming=False,
+            duration_s=time.perf_counter() - start,
+            error_code=error.code,
+            error_type=type(e).__name__,
+        )
         return AgentRPCResponse(id=request.id, error=error.model_dump(), result=None)
 
 
@@ -608,6 +632,8 @@ async def _handle_streaming_rpc(
     """Handle streaming JSON-RPC requests."""
 
     async def rpc_response_generator():
+        start = time.perf_counter()
+        error_type: str | None = None
         result_entity_async_iterator = None
         try:
             result_entity_async_iterator = await agents_acp_use_case.handle_rpc_request(
@@ -640,6 +666,7 @@ async def rpc_response_generator():
 
         except Exception as e:
             logger.error(f"Error in streaming RPC response: {e}", exc_info=True)
+            error_type = type(e).__name__
             # Yield error response
             error_response = AgentRPCResponse(
                 id=request.id,
@@ -648,6 +675,13 @@ async def rpc_response_generator():
             )
             yield error_response.model_dump_json().encode() + b"\n"
         finally:
+            record_rpc_request(
+                method=request.method.value,
+                streaming=True,
+                duration_s=time.perf_counter() - start,
+                error_code=None if error_type is None else -32603,
+                error_type=error_type,
+            )
             # CRITICAL: Ensure the async iterator is properly closed
             # This ensures HTTP connections are released back to the pool
             if result_entity_async_iterator is not None and hasattr(

agentex/src/config/dependencies.py

@@ -19,7 +19,10 @@
 
 from src.config.environment_variables import Environment, EnvironmentVariables
 from src.utils.database import async_db_engine_creator
-from src.utils.db_metrics import PostgresMetricsCollector
+from src.utils.db_metrics import (
+    InstrumentedAsyncAdaptedQueuePool,
+    PostgresMetricsCollector,
+)
 from src.utils.logging import make_logger
 
 logger = make_logger(__name__)
@@ -97,6 +100,7 @@ async def load(self):
                 self.environment_variables.DATABASE_URL,
             ),
             echo=echo_db_engine,
+            poolclass=InstrumentedAsyncAdaptedQueuePool,  # emits pool wait_time/pending_requests/timeouts
             pool_size=async_db_pool_size,
             max_overflow=20,  # Allow 20 additional connections beyond pool_size when needed
             pool_pre_ping=True,
@@ -109,6 +113,7 @@ async def load(self):
                 self.environment_variables.DATABASE_URL,
             ),
             echo=echo_db_engine,
+            poolclass=InstrumentedAsyncAdaptedQueuePool,  # emits pool wait_time/pending_requests/timeouts
             pool_size=middleware_db_pool_size,
             max_overflow=10,  # Allow 10 additional connections for middleware
             pool_pre_ping=True,
@@ -188,6 +193,7 @@ async def load(self):
                 "postgresql+asyncpg://",
                 async_creator=async_db_engine_creator(read_only_db_url),
                 echo=echo_db_engine,
+                poolclass=InstrumentedAsyncAdaptedQueuePool,  # emits pool wait_time/pending_requests/timeouts
                 pool_size=async_db_pool_size,
                 max_overflow=20,
                 pool_pre_ping=True,

agentex/src/utils/db_metrics.py

@@ -19,6 +19,8 @@
 from datadog import statsd
 from sqlalchemy import event
 from sqlalchemy.engine import ExecutionContext
+from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError
+from sqlalchemy.pool import AsyncAdaptedQueuePool
 
 from src.utils.logging import make_logger
 from src.utils.otel_metrics import get_meter
@@ -39,6 +41,28 @@
 # StatsD is only enabled if DD_AGENT_HOST is configured
 _STATSD_ENABLED = bool(os.environ.get("DD_AGENT_HOST"))
 
+# Bucket boundaries (seconds) for the connection-acquisition wait histogram.
+# The SDK's default boundaries are millisecond-magnitude integers [0, 5, 10, ...]
+# that assume a millisecond unit; against a seconds-unit metric almost every
+# pool wait (µs–ms normally, up to the ~30s pool_timeout under saturation) falls
+# in the first bucket and quantiles are meaningless. These span instant-to-
+# timeout so p95/p99 stay usable.
+_POOL_WAIT_BUCKET_BOUNDARIES_S = (
+    0.001,
+    0.005,
+    0.01,
+    0.025,
+    0.05,
+    0.1,
+    0.25,
+    0.5,
+    1.0,
+    2.5,
+    5.0,
+    10.0,
+    30.0,
+)
+
 
 def _format_statsd_tags(attributes: dict) -> list[str]:
     """Convert OTel attributes dict to Datadog StatsD tags list."""
@@ -70,6 +94,84 @@ def _parse_db_url(url: str) -> tuple[str, int, str]:
     return host, port, db_name
 
 
+class _PoolWaitInstruments:
+    """OTel instruments + attributes for connection-acquisition metrics.
+
+    Held on the pool instance so ``connect()`` can emit without looking back
+    through the collector. A ``None`` slot on the pool means OTel is not
+    configured and the override degrades to a straight passthrough.
+    """
+
+    __slots__ = ("wait_time", "pending_requests", "timeouts", "attributes")
+
+    def __init__(self, wait_time, pending_requests, timeouts, attributes):
+        self.wait_time = wait_time
+        self.pending_requests = pending_requests
+        self.timeouts = timeouts
+        self.attributes = attributes
+
+
+class InstrumentedAsyncAdaptedQueuePool(AsyncAdaptedQueuePool):
+    """Async connection pool that measures connection acquisition.
+
+    Sources the three OTel DB connection-pool metrics that the checkout/checkin
+    events in ``PostgresPoolMetrics`` cannot: by the time ``checkout`` fires the
+    wait is already over, SQLAlchemy exposes no "started waiting" event, and a
+    pool timeout raises before any connection is handed out.
+
+    ``Pool.connect()`` is the single, non-recursive entry point for obtaining a
+    connection (``connect() -> _ConnectionFairy._checkout() -> _do_get()``) and
+    runs inside the acquiring greenlet, so wall-clock timing around it captures
+    the real wait — including time blocked on a saturated pool. Overriding
+    ``_do_get`` directly would double-count, because ``QueuePool._do_get``
+    recurses on itself on the overflow-retry path.
+
+    Emits (OTel-only; no StatsD dual-emit, these are new series with no Datadog
+    consumer yet):
+      * ``db.client.connection.wait_time``        — time to obtain a connection
+      * ``db.client.connection.pending_requests`` — requests waiting right now
+      * ``db.client.connection.timeouts``         — pool acquisition timeouts
+
+    The pool is constructed inside ``create_async_engine`` before the collector
+    has meters, so instruments are attached post-construction; until then, and
+    whenever OTel is disabled, ``connect()`` is a plain passthrough.
+    """
+
+    # None => passthrough (OTel disabled, or instruments not yet attached).
+    _wait_instruments: _PoolWaitInstruments | None = None
+
+    def attach_wait_instruments(self, instruments: _PoolWaitInstruments) -> None:
+        self._wait_instruments = instruments
+
+    def recreate(self) -> InstrumentedAsyncAdaptedQueuePool:
+        # dispose()/reset recreates the pool object; carry instrumentation
+        # forward so metrics don't silently go dark after a pool recreate.
+        new_pool = super().recreate()
+        new_pool._wait_instruments = self._wait_instruments
+        return new_pool
+
+    def connect(self):
+        instruments = self._wait_instruments
+        if instruments is None:
+            return super().connect()
+
+        attributes = instruments.attributes
+        instruments.pending_requests.add(1, attributes)
+        start = time.monotonic()
+        try:
+            connection = super().connect()
+        except SQLAlchemyTimeoutError:
+            instruments.timeouts.add(1, attributes)
+            raise
+        finally:
+            instruments.pending_requests.add(-1, attributes)
+
+        # Only reached on success: a timed-out (or otherwise failed) acquisition
+        # obtained no connection, so it must not land in the wait_time histogram.
+        instruments.wait_time.record(time.monotonic() - start, attributes)
+        return connection
+
+
 class PostgresPoolMetrics:
     """
     Collects and emits PostgreSQL connection pool metrics.
@@ -151,6 +253,52 @@ def __init__(
             unit="s",
         )
 
+        # Connection-acquisition metrics (OTel DB semconv). Sourced from the
+        # InstrumentedAsyncAdaptedQueuePool.connect() seam rather than pool
+        # events, which can't see the wait, the current waiters, or a timeout.
+        self._connection_wait_time: Histogram = meter.create_histogram(
+            name="db.client.connection.wait_time",
+            description="The time it took to obtain an open connection from the pool",
+            unit="s",
+            explicit_bucket_boundaries_advisory=_POOL_WAIT_BUCKET_BOUNDARIES_S,
+        )
+
+        self._connection_pending_requests: UpDownCounter = meter.create_up_down_counter(
+            name="db.client.connection.pending_requests",
+            description="The number of current pending requests for an open connection",
+            unit="{request}",
+        )
+
+        self._connection_timeouts: Counter = meter.create_counter(
+            name="db.client.connection.timeouts",
+            description=(
+                "The number of connection timeouts that have occurred trying to "
+                "obtain a connection from the pool"
+            ),
+            unit="{timeout}",
+        )
+
+        # Hand the instruments to the pool so its connect() override can emit.
+        # Only InstrumentedAsyncAdaptedQueuePool carries the hook; a differently
+        # configured engine (e.g. NullPool in tests) is left as a passthrough.
+        pool = self.engine.sync_engine.pool
+        if isinstance(pool, InstrumentedAsyncAdaptedQueuePool):
+            pool.attach_wait_instruments(
+                _PoolWaitInstruments(
+                    wait_time=self._connection_wait_time,
+                    pending_requests=self._connection_pending_requests,
+                    timeouts=self._connection_timeouts,
+                    attributes=self.base_attributes,
+                )
+            )
+        else:
+            logger.warning(
+                "Pool for %s is %s, not InstrumentedAsyncAdaptedQueuePool; "
+                "wait_time/pending_requests/timeouts will not be emitted",
+                pool_name,
+                type(pool).__name__,
+            )
+
         # Track last reported values for delta calculation
         self._last_idle = 0
         self._last_used = 0