From 0d0fe55519669839c540859e8fcf5089257974b7 Mon Sep 17 00:00:00 2001
From: rock <locartrock@gmail.com>
Date: Sun, 13 Sep 2026 21:37:11 +0900
Subject: [PATCH] feat(O2): instrument query handler with Prometheus metrics

- Track query requests, errors, auth failures, rate limits
- Track embedding failures and embedding call duration
- Track result counts, empty results
- In-flight gauge for concurrent queries
- Timer for query duration histogram
- Metrics: Q1-Q12 (12 metrics instrumented)
---
 crates/mem-cli/src/handlers/unified_query.rs | 31 +++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/crates/mem-cli/src/handlers/unified_query.rs b/crates/mem-cli/src/handlers/unified_query.rs
index 1d8415c..d84a5c9 100644
--- a/crates/mem-cli/src/handlers/unified_query.rs
+++ b/crates/mem-cli/src/handlers/unified_query.rs
@@ -118,17 +118,26 @@ pub async fn unified_query_handler(
     body: web::Json<UnifiedQueryRequest>,
     state: web::Data<AppState>,
 ) -> HttpResponse {
+    use crate::metrics::*;
+    QUERY_REQUESTS_TOTAL.inc();
+    QUERY_IN_FLIGHT.inc();
+    let _timer = Timer::new(&QUERY_DURATION);
     let start_time = std::time::Instant::now();
 
     // 1. Validate JWT + rate limit
     if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
         &req, &state, "query", 500
     ) {
+        QUERY_AUTH_FAILURES.inc();
+        QUERY_ERRORS_TOTAL.inc();
+        QUERY_IN_FLIGHT.dec();
         return response;
     }
 
     // 2. Validate input
     if let Err(response) = validate_unified_request(&body) {
+        QUERY_ERRORS_TOTAL.inc();
+        QUERY_IN_FLIGHT.dec();
         return response;
     }
 
@@ -136,9 +145,16 @@ pub async fn unified_query_handler(
            body.search_type, body.query, body.entity_type, body.relation_type);
 
     // 3. Embed query once (reused for all search types)
+    let embed_start = std::time::Instant::now();
     let query_embedding = match state.embeddings.embed_one(&body.query).await {
-        Ok(emb) => emb.to_vec(),
+        Ok(emb) => {
+            QUERY_EMBEDDING_DURATION.observe(embed_start.elapsed().as_secs_f64());
+            emb.to_vec()
+        }
         Err(e) => {
+            QUERY_EMBEDDING_FAILURES.inc();
+            QUERY_ERRORS_TOTAL.inc();
+            QUERY_IN_FLIGHT.dec();
             error!("Embedding failed: {}", e);
             return crate::handlers::response_builder::internal_error(
                 "Failed to embed query"
@@ -152,12 +168,15 @@ pub async fn unified_query_handler(
         "edges" => search_edges(&body, &state, &query_embedding, start_time).await,
         "hybrid" => search_hybrid(&body, &state, &query_embedding, start_time).await,
         _ => {
+            QUERY_ERRORS_TOTAL.inc();
+            QUERY_IN_FLIGHT.dec();
             return crate::handlers::response_builder::bad_request(
                 "search_type must be 'entities', 'edges', or 'hybrid'"
             );
         }
     };
 
+    QUERY_IN_FLIGHT.dec();
     response
 }
 
@@ -247,6 +266,10 @@ async fn search_entities(
 
     info!("Unified query (entities): {} results in {}ms", count, elapsed);
 
+    // O2: Track result counts
+    crate::metrics::QUERY_RESULTS_TOTAL.inc_by(count as u64);
+    if count == 0 { crate::metrics::QUERY_EMPTY_RESULTS.inc(); }
+
     let response = UnifiedQueryResponse {
         query: req.query.clone(),
         search_type: "entities".to_string(),
@@ -305,6 +328,9 @@ async fn search_edges(
 
     info!("Unified query (edges): {} results in {}ms", count, elapsed);
 
+    crate::metrics::QUERY_RESULTS_TOTAL.inc_by(count as u64);
+    if count == 0 { crate::metrics::QUERY_EMPTY_RESULTS.inc(); }
+
     let response = UnifiedQueryResponse {
         query: req.query.clone(),
         search_type: "edges".to_string(),
@@ -350,6 +376,9 @@ async fn search_hybrid(
 
     info!("Unified query (hybrid): {} results in {}ms", count, elapsed);
 
+    crate::metrics::QUERY_RESULTS_TOTAL.inc_by(count as u64);
+    if count == 0 { crate::metrics::QUERY_EMPTY_RESULTS.inc(); }
+
     let response = UnifiedQueryResponse {
         query: req.query.clone(),
         search_type: "hybrid".to_string(),