From d80ab19a36e14b3266fd5d813acc030c3ca05dc6 Mon Sep 17 00:00:00 2001
From: arpanjot-coding <50767931+arpanjot-coding@users.noreply.github.com>
Date: Wed, 11 Nov 2026 20:19:21 +0000
Subject: [PATCH] Match configured record keys in buffer metrics

---
 tools/aimdb-mcp/src/tools/buffer_metrics.rs |  8 ++--
 tools/aimdb-mcp/tests/buffer_metrics.rs     | 47 +++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)
 create mode 100644 tools/aimdb-mcp/tests/buffer_metrics.rs

diff --git a/tools/aimdb-mcp/src/tools/buffer_metrics.rs b/tools/aimdb-mcp/src/tools/buffer_metrics.rs
index df98c675..94a98e38 100644
--- a/tools/aimdb-mcp/src/tools/buffer_metrics.rs
+++ b/tools/aimdb-mcp/src/tools/buffer_metrics.rs
@@ -14,7 +14,7 @@ use tracing::debug;
 struct GetBufferMetricsParams {
     /// Endpoint of the AimDB instance: a `scheme://` URL or bare path (falls back to AIMDB_CONNECT env).
     endpoint: Option<String>,
-    /// Substring matched against record names (e.g. `"Temperature"`).
+    /// Substring matched against record keys and names (e.g. `"Temperature"`).
     record_key: String,
 }
 
@@ -23,7 +23,7 @@ struct ResetBufferMetricsParams {
     endpoint: Option<String>,
 }
 
-/// Returns live buffer metrics for records whose name contains `record_key`.
+/// Returns live buffer metrics for records whose key or name contains `record_key`.
 ///
 /// Session-agnostic: works in any phase.
 pub async fn get_buffer_metrics(args: Option<Value>) -> McpResult<Value> {
@@ -38,7 +38,9 @@ pub async fn get_buffer_metrics(args: Option<Value>) -> McpResult<Value> {
 
     let matching: Vec<_> = raw
         .into_iter()
-        .filter(|r| r.name.contains(&params.record_key))
+        .filter(|r| {
+            r.record_key.contains(&params.record_key) || r.name.contains(&params.record_key)
+        })
         .collect();
 
     if matching.is_empty() {
diff --git a/tools/aimdb-mcp/tests/buffer_metrics.rs b/tools/aimdb-mcp/tests/buffer_metrics.rs
new file mode 100644
index 00000000..aa6bc585
--- /dev/null
+++ b/tools/aimdb-mcp/tests/buffer_metrics.rs
@@ -0,0 +1,47 @@
+//! Buffer metrics look up records by their configured keys.
+
+use std::sync::Arc;
+
+use aimdb_core::buffer::BufferCfg;
+use aimdb_core::remote::AimxConfig;
+use aimdb_core::AimDbBuilder;
+use aimdb_mcp::{connection::ConnectionPool, tools};
+use aimdb_tokio_adapter::{TokioAdapter, TokioRecordRegistrarExt};
+use aimdb_uds_connector::UdsServer;
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Reading {
+    value: u64,
+}
+
+#[tokio::test]
+async fn get_buffer_metrics_finds_the_configured_record_key() {
+    let dir = tempfile::tempdir().unwrap();
+    let sock = dir.path().join("aimdb.sock");
+    let endpoint = sock.to_str().unwrap();
+    let config = AimxConfig::uds_default().socket_path(endpoint);
+    let mut builder = AimDbBuilder::new()
+        .runtime(Arc::new(TokioAdapter))
+        .with_connector(UdsServer::from_config(config));
+    for key in ["sensor.indoor", "sensor.outdoor"] {
+        builder.configure::<Reading>(key, |reg| {
+            reg.buffer(BufferCfg::SingleLatest).with_remote_access();
+        });
+    }
+    let (_db, runner) = builder.build().await.expect("build db");
+    tokio::spawn(runner.run());
+    tools::init_connection_pool(ConnectionPool::new());
+
+    let result = tools::get_buffer_metrics(Some(json!({
+        "endpoint": endpoint,
+        "record_key": "sensor.indoor",
+    })))
+    .await
+    .expect("get buffer metrics");
+    assert_eq!(result["found"], true);
+    let records = result["records"].as_array().expect("matching records");
+    assert_eq!(records.len(), 1);
+    assert_eq!(records[0]["record_key"], "sensor.indoor");
+}