Merge branch 'main' into cynthiawang/agx1-618-2a-task-stream-lifecycle-metrics

Author: cyntwang99

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

@@ -43,6 +43,9 @@
     DAuthorizedResourceIds,
 )
 from src.utils.logging import make_logger
+from src.utils.rpc_metrics import (
+    rpc_request_timing,
+)
 from src.utils.task_authorization import check_task_or_collapse_to_404
 
 logger = make_logger(__name__)
@@ -549,53 +552,62 @@ async def _handle_sync_rpc(
     request_headers: dict[str, str] | None = None,
 ) -> AgentRPCResponse:
     """Handle synchronous JSON-RPC requests."""
-    try:
-        result_entity = await agents_acp_use_case.handle_rpc_request(
-            agent_id=agent_id,
-            agent_name=agent_name,
-            method=request.method,
-            params=request.params,
-            request_headers=request_headers,
-        )
+    with rpc_request_timing(request.method.value, streaming=False) as rpc_call:
+        try:
+            result_entity = await agents_acp_use_case.handle_rpc_request(
+                agent_id=agent_id,
+                agent_name=agent_name,
+                method=request.method,
+                params=request.params,
+                request_headers=request_headers,
+            )
 
-        if isinstance(result_entity, AsyncIterator):
-            raise ValueError(f"Expected non-async iterator, got {type(result_entity)}")
+            if isinstance(result_entity, AsyncIterator):
+                raise ValueError(
+                    f"Expected non-async iterator, got {type(result_entity)}"
+                )
 
-        if isinstance(result_entity, list):
-            serialized_result = [item.model_dump() for item in result_entity]
-        else:
-            serialized_result = result_entity.model_dump()
-
-        # if request.method == AgentRPCMethod.MESSAGE_SEND:
-        #     if isinstance(result_entity, list):
-        #         result = [TaskMessage.model_validate(task_message_entity) for task_message_entity in result_entity]
-        #     else:
-        #         raise ValueError(f"Expected list of TaskMessage entities, got {type(result_entity)}")
-        # elif request.method == AgentRPCMethod.TASK_CREATE:
-        #     result = Task.model_validate(result_entity)
-        # elif request.method == AgentRPCMethod.TASK_CANCEL:
-        #     result = Task.model_validate(result_entity)
-        # elif request.method == AgentRPCMethod.EVENT_SEND:
-        #     result = Event.model_validate(result_entity)
-        # else:
-        #     raise ValueError(f"Unsupported method: {request.method}")
-        # logger.info(f"AgentRPCResponse Result: {result}")
-        return AgentRPCResponse.model_validate(
-            {
-                "id": request.id,
-                "result": serialized_result,
-                "error": None,
-            }
-        )
+            if isinstance(result_entity, list):
+                serialized_result = [item.model_dump() for item in result_entity]
+            else:
+                serialized_result = result_entity.model_dump()
+
+            # if request.method == AgentRPCMethod.MESSAGE_SEND:
+            #     if isinstance(result_entity, list):
+            #         result = [TaskMessage.model_validate(task_message_entity) for task_message_entity in result_entity]
+            #     else:
+            #         raise ValueError(f"Expected list of TaskMessage entities, got {type(result_entity)}")
+            # elif request.method == AgentRPCMethod.TASK_CREATE:
+            #     result = Task.model_validate(result_entity)
+            # elif request.method == AgentRPCMethod.TASK_CANCEL:
+            #     result = Task.model_validate(result_entity)
+            # elif request.method == AgentRPCMethod.EVENT_SEND:
+            #     result = Event.model_validate(result_entity)
+            # else:
+            #     raise ValueError(f"Unsupported method: {request.method}")
+            # logger.info(f"AgentRPCResponse Result: {result}")
+            return AgentRPCResponse.model_validate(
+                {
+                    "id": request.id,
+                    "result": serialized_result,
+                    "error": None,
+                }
+            )
 
-    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}")
-        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))
-        return AgentRPCResponse(id=request.id, error=error.model_dump(), result=None)
+        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}")
+            rpc_call.fail(error.code, e)
+            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))
+            rpc_call.fail(error.code, e)
+            return AgentRPCResponse(
+                id=request.id, error=error.model_dump(), result=None
+            )
 
 
 async def _handle_streaming_rpc(
@@ -608,56 +620,60 @@ async def _handle_streaming_rpc(
     """Handle streaming JSON-RPC requests."""
 
     async def rpc_response_generator():
-        result_entity_async_iterator = None
-        try:
-            result_entity_async_iterator = await agents_acp_use_case.handle_rpc_request(
-                agent_id=agent_id,
-                agent_name=agent_name,
-                method=request.method,
-                params=request.params,
-                request_headers=request_headers,
-            )
-
-            if not isinstance(result_entity_async_iterator, AsyncIterator):
-                raise ValueError(
-                    f"Expected AsyncIterator, got {type(result_entity_async_iterator)}"
+        with rpc_request_timing(request.method.value, streaming=True) as rpc_call:
+            result_entity_async_iterator = None
+            try:
+                result_entity_async_iterator = (
+                    await agents_acp_use_case.handle_rpc_request(
+                        agent_id=agent_id,
+                        agent_name=agent_name,
+                        method=request.method,
+                        params=request.params,
+                        request_headers=request_headers,
+                    )
                 )
 
-            # At this point we know it's an AsyncIterator[TaskMessage]
-            async for task_message_update_entity in result_entity_async_iterator:
-                logger.debug(
-                    f"Streaming message chunk type: {type(task_message_update_entity).__name__}"
-                )
-                rpc_response = AgentRPCResponse.model_validate(
-                    {
-                        "id": request.id,
-                        "result": task_message_update_entity.model_dump(),
-                        "error": None,
-                    }
-                )
-                # Yield JSON bytes with newline for NDJSON format
-                yield rpc_response.model_dump_json().encode() + b"\n"
+                if not isinstance(result_entity_async_iterator, AsyncIterator):
+                    raise ValueError(
+                        f"Expected AsyncIterator, got {type(result_entity_async_iterator)}"
+                    )
 
-        except Exception as e:
-            logger.error(f"Error in streaming RPC response: {e}", exc_info=True)
-            # Yield error response
-            error_response = AgentRPCResponse(
-                id=request.id,
-                result=None,
-                error=JSONRPCError(code=-32603, message=str(e)).model_dump(),
-            )
-            yield error_response.model_dump_json().encode() + b"\n"
-        finally:
-            # 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(
-                result_entity_async_iterator, "aclose"
-            ):
-                try:
-                    await result_entity_async_iterator.aclose()
-                    logger.debug("Closed streaming iterator properly")
-                except Exception as e:
-                    logger.warning(f"Error closing streaming iterator: {e}")
+                # At this point we know it's an AsyncIterator[TaskMessage]
+                async for task_message_update_entity in result_entity_async_iterator:
+                    logger.debug(
+                        f"Streaming message chunk type: {type(task_message_update_entity).__name__}"
+                    )
+                    rpc_response = AgentRPCResponse.model_validate(
+                        {
+                            "id": request.id,
+                            "result": task_message_update_entity.model_dump(),
+                            "error": None,
+                        }
+                    )
+                    # Yield JSON bytes with newline for NDJSON format
+                    yield rpc_response.model_dump_json().encode() + b"\n"
+
+            except Exception as e:
+                logger.error(f"Error in streaming RPC response: {e}", exc_info=True)
+                rpc_call.fail(-32603, e)
+                # Yield error response
+                error_response = AgentRPCResponse(
+                    id=request.id,
+                    result=None,
+                    error=JSONRPCError(code=-32603, message=str(e)).model_dump(),
+                )
+                yield error_response.model_dump_json().encode() + b"\n"
+            finally:
+                # 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(
+                    result_entity_async_iterator, "aclose"
+                ):
+                    try:
+                        await result_entity_async_iterator.aclose()
+                        logger.debug("Closed streaming iterator properly")
+                    except Exception as e:
+                        logger.warning(f"Error closing streaming iterator: {e}")
 
     return StreamingResponse(
         rpc_response_generator(),

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,