From 3ddc4a0a8de065791a2d9ae10fb5e54a33279fad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexander=20Schn=C3=B6rch?= Date: Mon, 9 Nov 2026 21:31:16 +0100 Subject: [PATCH 1/2] fix(mcp): every tool shares one persistent session per endpoint (#451) Most aimdb-mcp tools opened a new AimX connection per call. Over serial that either failed (the tty is already open) or sent a new hello, which reset the board's session and discarded the drain cursor drain_record reads. - All tools and the records resource go through ConnectionPool::get_client (renamed from get_drain_client) and hold the per-client lock for one call. - get_connection, invalidate_connection, invalidate_drain_client and the per-tool fallbacks are gone, so no path opens a second connection. - Each endpoint has one connect in flight: parallel first calls wait for it instead of racing, where the loser found the tty busy. Tests: serial_session calls every tool between opening a ring's drain cursor and draining it, and checks the queued events and exactly one hello. New tests cover parallel first calls and the records resource. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 + Makefile | 2 +- tools/aimdb-mcp/CHANGELOG.md | 5 + tools/aimdb-mcp/src/connection.rs | 153 +++++--------------- tools/aimdb-mcp/src/resources/records.rs | 7 +- tools/aimdb-mcp/src/tools/buffer_metrics.rs | 19 +-- tools/aimdb-mcp/src/tools/graph.rs | 40 +---- tools/aimdb-mcp/src/tools/instance.rs | 11 +- tools/aimdb-mcp/src/tools/mod.rs | 14 ++ tools/aimdb-mcp/src/tools/profiling.rs | 19 +-- tools/aimdb-mcp/src/tools/record.rs | 57 ++------ tools/aimdb-mcp/tests/persistent_client.rs | 21 ++- tools/aimdb-mcp/tests/serial_session.rs | 148 ++++++++++++++++++- 13 files changed, 259 insertions(+), 243 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 231d2921..b184a13c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed — `aimdb-mcp` tools share one session per endpoint (#451) + +- **No tool opens its own connection any more**, so over serial a tool call no + longer fails on a busy tty or resets the board's session and drain cursor. See + [aimdb-mcp](tools/aimdb-mcp/CHANGELOG.md). + ### Fixed — `record.get` answers at once and no longer drains (#444) - **`record.get` is a single, non-destructive read.** `aimdb-mcp`'s diff --git a/Makefile b/Makefile index 7f42589c..9bc11c2c 100644 --- a/Makefile +++ b/Makefile @@ -243,7 +243,7 @@ test: test-operational-events test-build-info-link cargo test --package aimdb-cli --features yaml @printf "$(YELLOW) → Testing MCP server$(NC)\n" cargo test --package aimdb-mcp - @printf "$(YELLOW) → Testing MCP server (serial: one session for get and drain)$(NC)\n" + @printf "$(YELLOW) → Testing MCP server (serial: one session for every tool)$(NC)\n" cargo test --package aimdb-mcp --features "transport-serial" --test serial_session @printf "$(YELLOW) → Testing persistence backend$(NC)\n" cargo test --package aimdb-persistence diff --git a/tools/aimdb-mcp/CHANGELOG.md b/tools/aimdb-mcp/CHANGELOG.md index 7cb2e3a2..4fe0b264 100644 --- a/tools/aimdb-mcp/CHANGELOG.md +++ b/tools/aimdb-mcp/CHANGELOG.md @@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **Every tool shares one session per endpoint (#451).** Tools opened a new + connection per call, which over serial failed (the tty was already open) or + sent a new `hello` that reset the board's session and discarded its drain + cursor. All tools and the records resource now use the persistent client, one + call at a time, and parallel first calls wait for a single connect. - **`get_record` answers at once (#444).** It no longer retries `not_found` for up to 3 s, so an empty record or a misspelled name returns immediately instead of after 3 s. A ring returns `unsupported` (use `drain_record`), and diff --git a/tools/aimdb-mcp/src/connection.rs b/tools/aimdb-mcp/src/connection.rs index fe69a065..f93742f2 100644 --- a/tools/aimdb-mcp/src/connection.rs +++ b/tools/aimdb-mcp/src/connection.rs @@ -1,7 +1,7 @@ -//! Connection pool management for AimDB instances +//! Persistent AimX clients, one per endpoint. //! -//! Manages persistent connections to AimDB instances to avoid -//! reconnecting on every tool call. Includes auto-reconnect logic. +//! Every tool and resource reaches an AimDB instance through this pool, so an +//! endpoint only ever has one connection. use aimdb_client::AimxConnection; use aimdb_client::ClientError; @@ -10,28 +10,24 @@ use std::sync::Arc; use tokio::sync::Mutex; use tracing::debug; -/// Connection entry with metadata -#[derive(Debug)] -struct ConnectionEntry { - /// Last successful connection time (for staleness detection) - last_used: std::time::Instant, -} +/// A persistent client. Lock it for the duration of one call. +pub type SharedClient = Arc>; + +/// An endpoint's client, or `None` before the first connect. Its lock is held +/// across `connect` so an endpoint only ever has one attempt in flight. +type ClientSlot = Arc>>; -/// Connection pool for managing AimDB connections +/// Pool of persistent AimDB clients #[derive(Clone)] pub struct ConnectionPool { - /// Track which connections we've attempted (for logging/metrics) - connections: Arc>>, - /// Persistent drain clients — kept alive so drain readers accumulate values - /// Key: endpoint, Value: shared AimxConnection - drain_clients: Arc>>>>, + /// Key: endpoint + clients: Arc>>, } impl std::fmt::Debug for ConnectionPool { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ConnectionPool") - .field("connections", &"<...>") - .field("drain_clients", &"<...>") + .field("clients", &"<...>") .finish() } } @@ -40,70 +36,34 @@ impl ConnectionPool { /// Create a new connection pool pub fn new() -> Self { Self { - connections: Arc::new(Mutex::new(HashMap::new())), - drain_clients: Arc::new(Mutex::new(HashMap::new())), - } - } - - /// Get or create a connection to an AimDB instance - /// - /// Note: since `AimxConnection` does not implement `Clone`, a fresh - /// connection is created each time. The pool tracks connection metadata for - /// monitoring and future optimization (e.g. persistent connections via - /// `Arc>` if `AimxConnection` becomes `Sync`). - pub async fn get_connection(&self, endpoint: &str) -> Result { - let mut pool = self.connections.lock().await; - - // Update or insert connection metadata - let now = std::time::Instant::now(); - - if let Some(entry) = pool.get_mut(endpoint) { - debug!( - "♻️ Connection metadata exists for {}, reconnecting", - endpoint - ); - entry.last_used = now; - } else { - debug!("🔌 First connection to {}", endpoint); - pool.insert(endpoint.to_string(), ConnectionEntry { last_used: now }); - } - - // Always create a new connection (until AimxConnection supports cloning/sharing) - drop(pool); // Release lock before async operation - AimxConnection::connect(endpoint).await - } - - /// Remove a connection from the pool (called when operations fail) - pub async fn invalidate_connection(&self, endpoint: &str) { - let mut pool = self.connections.lock().await; - if pool.remove(endpoint).is_some() { - debug!("❌ Invalidated connection metadata for {}", endpoint); + clients: Arc::new(Mutex::new(HashMap::new())), } } /// Get or create the persistent client for an endpoint. /// - /// One connection per endpoint, kept across calls and shared by - /// `get_record` and `drain_record`; callers take its lock for each call. - /// Keeping it open matters twice: the server's drain cursor lives in the - /// session, so values accumulate between drains, and on serial a second - /// connection either fails (the tty is open) or sends a new `hello` that - /// resets the board's session and discards that cursor. + /// One connection per endpoint, kept across calls and shared by every tool; + /// callers take its lock for each call. Keeping it open matters twice: the + /// server's drain cursor lives in the session, so values accumulate between + /// drains, and on serial a second connection either fails (the tty is open) + /// or sends a new `hello` that resets the board's session and discards that + /// cursor. For the same reason concurrent first calls wait for one connect + /// instead of racing their own. /// /// The client is replaced only once its engine has stopped /// ([`AimxConnection::is_closed`]), never because a call returned an error: - /// `not_found`, `unsupported` and even `internal` are replies on a healthy + /// `not_found`, `denied` and even `internal` are replies on a healthy /// session. A client busy with a call is returned as is; if its engine has /// stopped, the next call replaces it. - pub async fn get_drain_client( - &self, - endpoint: &str, - ) -> Result>, ClientError> { - let mut drain_map = self.drain_clients.lock().await; - - if let Some(client) = drain_map.get(endpoint) { - let closed = client.try_lock().is_ok_and(|c| c.is_closed()); - if !closed { + pub async fn get_client(&self, endpoint: &str) -> Result { + let slot = { + let mut clients = self.clients.lock().await; + Arc::clone(clients.entry(endpoint.to_string()).or_default()) + }; + let mut slot = slot.lock().await; + + if let Some(client) = slot.as_ref() { + if !client.try_lock().is_ok_and(|c| c.is_closed()) { debug!("♻️ Reusing persistent client for {}", endpoint); return Ok(Arc::clone(client)); } @@ -111,49 +71,24 @@ impl ConnectionPool { "🔁 Persistent client for {} has stopped, reconnecting", endpoint ); - drain_map.remove(endpoint); - } - - debug!("🔌 Creating persistent drain client for {}", endpoint); - - // Drop lock before async connect - drop(drain_map); - - let client = AimxConnection::connect(endpoint).await?; - let shared = Arc::new(tokio::sync::Mutex::new(client)); - - let mut drain_map = self.drain_clients.lock().await; - // Double-check: another task may have inserted while we were connecting - if let Some(existing) = drain_map.get(endpoint) { - if !existing.try_lock().is_ok_and(|c| c.is_closed()) { - return Ok(Arc::clone(existing)); - } + *slot = None; } - drain_map.insert(endpoint.to_string(), Arc::clone(&shared)); - Ok(shared) - } - /// Invalidate (remove) a persistent drain client, e.g. after connection error - pub async fn invalidate_drain_client(&self, endpoint: &str) { - let mut drain_map = self.drain_clients.lock().await; - if drain_map.remove(endpoint).is_some() { - debug!("❌ Invalidated drain client for {}", endpoint); - } + debug!("🔌 Creating persistent client for {}", endpoint); + let client = Arc::new(Mutex::new(AimxConnection::connect(endpoint).await?)); + *slot = Some(Arc::clone(&client)); + Ok(client) } - /// Clear all connections in the pool + /// Forget every client pub async fn clear(&self) { - let mut pool = self.connections.lock().await; - pool.clear(); - let mut drain_map = self.drain_clients.lock().await; - drain_map.clear(); + self.clients.lock().await.clear(); debug!("🧹 Cleared connection pool"); } - /// Get the number of tracked connections + /// Number of endpoints the pool has been asked to connect to pub async fn connection_count(&self) -> usize { - let pool = self.connections.lock().await; - pool.len() + self.clients.lock().await.len() } } @@ -179,12 +114,4 @@ mod tests { pool.clear().await; assert_eq!(pool.connection_count().await, 0); } - - #[tokio::test] - async fn test_invalidate_nonexistent_connection() { - let pool = ConnectionPool::new(); - // Should not panic - pool.invalidate_connection("/tmp/nonexistent.sock").await; - assert_eq!(pool.connection_count().await, 0); - } } diff --git a/tools/aimdb-mcp/src/resources/records.rs b/tools/aimdb-mcp/src/resources/records.rs index 846ec58d..78cb9dfc 100644 --- a/tools/aimdb-mcp/src/resources/records.rs +++ b/tools/aimdb-mcp/src/resources/records.rs @@ -5,7 +5,6 @@ use crate::error::{McpError, McpResult}; use crate::protocol::{Resource, ResourceContent, ResourceReadResult}; -use aimdb_client::AimxConnection; use serde_json::json; use std::path::Path; use tracing::debug; @@ -81,8 +80,9 @@ pub async fn read_records_resource(socket_path: &str) -> McpResult, } -async fn connect(endpoint: &str) -> McpResult { - if let Some(pool) = super::connection_pool() { - pool.get_connection(endpoint) - .await - .map_err(McpError::Client) - } else { - AimxConnection::connect(endpoint) - .await - .map_err(McpError::Client) - } -} - /// Returns live buffer metrics for records whose name contains `record_key`. /// /// Session-agnostic: works in any phase. @@ -45,7 +32,8 @@ pub async fn get_buffer_metrics(args: Option) -> McpResult { .map_err(|e| McpError::InvalidParams(format!("get_buffer_metrics: {e}")))?; let endpoint = super::resolve_endpoint(params.endpoint)?; - let client = connect(&endpoint).await?; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; let raw = client.list_records().await.map_err(McpError::Client)?; let matching: Vec<_> = raw @@ -75,7 +63,8 @@ pub async fn reset_buffer_metrics(args: Option) -> McpResult { .map_err(|e| McpError::InvalidParams(format!("reset_buffer_metrics: {e}")))?; let endpoint = super::resolve_endpoint(params.endpoint)?; - let client = connect(&endpoint).await?; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; match client.reset_buffer_metrics().await { Ok(_) => Ok(json!({ "reset": true, diff --git a/tools/aimdb-mcp/src/tools/graph.rs b/tools/aimdb-mcp/src/tools/graph.rs index 04c7ad81..cef35ade 100644 --- a/tools/aimdb-mcp/src/tools/graph.rs +++ b/tools/aimdb-mcp/src/tools/graph.rs @@ -1,7 +1,6 @@ //! Graph introspection tools (graph_nodes, graph_edges, graph_topo_order) use crate::error::{McpError, McpResult}; -use aimdb_client::AimxConnection; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::debug; @@ -56,17 +55,8 @@ pub async fn graph_nodes(args: Option) -> McpResult { debug!("🔌 Connecting to {}", endpoint); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint) - .await - .map_err(McpError::Client)? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint) - .await - .map_err(McpError::Client)? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // Get graph nodes let nodes = client.graph_nodes().await.map_err(McpError::Client)?; @@ -101,17 +91,8 @@ pub async fn graph_edges(args: Option) -> McpResult { debug!("🔌 Connecting to {}", endpoint); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint) - .await - .map_err(McpError::Client)? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint) - .await - .map_err(McpError::Client)? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // Get graph edges let edges = client.graph_edges().await.map_err(McpError::Client)?; @@ -144,17 +125,8 @@ pub async fn graph_topo_order(args: Option) -> McpResult { debug!("🔌 Connecting to {}", endpoint); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint) - .await - .map_err(McpError::Client)? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint) - .await - .map_err(McpError::Client)? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // Get topological order let order = client.graph_topo_order().await.map_err(McpError::Client)?; diff --git a/tools/aimdb-mcp/src/tools/instance.rs b/tools/aimdb-mcp/src/tools/instance.rs index dfb45c82..34e45c8a 100644 --- a/tools/aimdb-mcp/src/tools/instance.rs +++ b/tools/aimdb-mcp/src/tools/instance.rs @@ -2,7 +2,6 @@ use crate::error::McpResult; use aimdb_client::protocol::InstanceInfoReply; -use aimdb_client::{self, AimxConnection}; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::debug; @@ -111,13 +110,8 @@ pub async fn get_instance_info(args: Option) -> McpResult { debug!("🔍 Getting instance info for: {}", endpoint); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint).await? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint).await? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // Get server info from the welcome message let server_info = client.server_info(); @@ -231,6 +225,7 @@ mod tests { #[tokio::test] async fn test_get_instance_info_invalid_socket() { + crate::tools::init_connection_pool(crate::connection::ConnectionPool::new()); let params = serde_json::json!({ "endpoint": "/tmp/nonexistent.sock" }); diff --git a/tools/aimdb-mcp/src/tools/mod.rs b/tools/aimdb-mcp/src/tools/mod.rs index 942d937a..9b355764 100644 --- a/tools/aimdb-mcp/src/tools/mod.rs +++ b/tools/aimdb-mcp/src/tools/mod.rs @@ -33,6 +33,20 @@ pub(crate) fn connection_pool() -> Option<&'static ConnectionPool> { CONNECTION_POOL.get() } +/// The persistent client for `endpoint`, the only way tools reach an instance. +/// Callers hold its lock for the duration of one call, so requests on a session +/// never overlap. +pub(crate) async fn shared_client( + endpoint: &str, +) -> crate::error::McpResult { + let pool = connection_pool().ok_or_else(|| { + crate::error::McpError::Internal("Connection pool not initialized".to_string()) + })?; + pool.get_client(endpoint) + .await + .map_err(crate::error::McpError::Client) +} + /// Resolve the endpoint from an explicit argument, the `--connect` flag, or the /// `AIMDB_CONNECT` env var (checked in that order). The value is a `scheme://` /// URL (`unix://PATH`, `serial://DEVICE?baud=N`) or a bare path. diff --git a/tools/aimdb-mcp/src/tools/profiling.rs b/tools/aimdb-mcp/src/tools/profiling.rs index cc69d84d..433b094b 100644 --- a/tools/aimdb-mcp/src/tools/profiling.rs +++ b/tools/aimdb-mcp/src/tools/profiling.rs @@ -8,7 +8,6 @@ //! feature; without it, records simply carry no `stage_profiling` data. use crate::error::{McpError, McpResult}; -use aimdb_client::AimxConnection; use serde::Deserialize; use serde_json::{json, Value}; use tracing::debug; @@ -25,18 +24,6 @@ struct ResetStageProfilingParams { endpoint: Option, } -async fn connect(endpoint: &str) -> McpResult { - if let Some(pool) = super::connection_pool() { - pool.get_connection(endpoint) - .await - .map_err(McpError::Client) - } else { - AimxConnection::connect(endpoint) - .await - .map_err(McpError::Client) - } -} - fn ms(ns: u64) -> f64 { ns as f64 / 1_000_000.0 } @@ -49,7 +36,8 @@ pub async fn get_stage_profiling(args: Option) -> McpResult { .map_err(|e| McpError::InvalidParams(format!("get_stage_profiling: {e}")))?; let endpoint = super::resolve_endpoint(params.endpoint)?; - let client = connect(&endpoint).await?; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; let records = client.list_records().await.map_err(McpError::Client)?; let mut out = Vec::new(); @@ -110,7 +98,8 @@ pub async fn reset_stage_profiling(args: Option) -> McpResult { .map_err(|e| McpError::InvalidParams(format!("reset_stage_profiling: {e}")))?; let endpoint = super::resolve_endpoint(params.endpoint)?; - let client = connect(&endpoint).await?; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; match client.reset_stage_profiling().await { Ok(_) => Ok(json!({ "reset": true, diff --git a/tools/aimdb-mcp/src/tools/record.rs b/tools/aimdb-mcp/src/tools/record.rs index 1cb64cd0..0875a4bd 100644 --- a/tools/aimdb-mcp/src/tools/record.rs +++ b/tools/aimdb-mcp/src/tools/record.rs @@ -1,7 +1,6 @@ //! Record-related tools (list_records, get_record, set_record) use crate::error::{McpError, McpResult}; -use aimdb_client::AimxConnection; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use tracing::debug; @@ -97,17 +96,8 @@ pub async fn list_records(args: Option) -> McpResult { let endpoint = super::resolve_endpoint(params.endpoint)?; debug!("🔌 Connecting to {}", endpoint); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint) - .await - .map_err(McpError::Client)? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint) - .await - .map_err(McpError::Client)? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // List records let records = client.list_records().await.map_err(McpError::Client)?; @@ -165,17 +155,12 @@ pub async fn get_record(args: Option) -> McpResult { // fails (the tty is open) or says `hello` again, which resets the board's // session and its drain cursors. Sharing is safe because `record.get` never // touches the drain cursor (the server keeps them apart). - let pool = super::connection_pool() - .ok_or_else(|| McpError::Internal("Connection pool not initialized".to_string()))?; - let client_arc = pool - .get_drain_client(&endpoint) - .await - .map_err(McpError::Client)?; - let client = client_arc.lock().await; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // One `record.get`, no retry: `not_found` (no such record, or no value yet) // and `unsupported` (a ring) come back at once and leave the session as it - // is. A stopped engine is replaced by `get_drain_client` on the next call. + // is. A stopped engine is replaced by `get_client` on the next call. let sample = client .get_record(¶ms.record_name) .await @@ -208,17 +193,8 @@ pub async fn set_record(args: Option) -> McpResult { endpoint, params.record_name ); - // Get or create connection from pool (if available) - let client = if let Some(pool) = super::connection_pool() { - pool.get_connection(&endpoint) - .await - .map_err(McpError::Client)? - } else { - // Fallback to direct connection if pool not initialized - AimxConnection::connect(&endpoint) - .await - .map_err(McpError::Client)? - }; + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // Set record value let result = client @@ -257,20 +233,13 @@ pub async fn drain_record(args: Option) -> McpResult { endpoint, params.record_name ); - // The persistent client, shared with `get_record`: the server-side drain - // cursor lives in its session, so values accumulate between calls. - let pool = super::connection_pool() - .ok_or_else(|| McpError::Internal("Connection pool not initialized".to_string()))?; - - let client_arc = pool - .get_drain_client(&endpoint) - .await - .map_err(McpError::Client)?; - - let client = client_arc.lock().await; + // The server-side drain cursor lives in the persistent client's session, + // so values accumulate between calls. + let client = super::shared_client(&endpoint).await?; + let client = client.lock().await; // An error reply keeps the session (and its cursors); a stopped engine is - // replaced by `get_drain_client` on the next call. + // replaced by `get_client` on the next call. let response = match params.limit { Some(limit) => client .drain_record_with_limit(¶ms.record_name, limit) @@ -364,6 +333,7 @@ mod tests { #[tokio::test] async fn test_list_records_invalid_socket() { + crate::tools::init_connection_pool(crate::connection::ConnectionPool::new()); // Should fail with non-existent socket let params = json!({ "endpoint": "/tmp/nonexistent.sock" @@ -412,6 +382,7 @@ mod tests { #[tokio::test] async fn test_set_record_invalid_socket() { + crate::tools::init_connection_pool(crate::connection::ConnectionPool::new()); // Should fail with non-existent socket let params = json!({ "endpoint": "/tmp/nonexistent.sock", diff --git a/tools/aimdb-mcp/tests/persistent_client.rs b/tools/aimdb-mcp/tests/persistent_client.rs index 79ad6ff3..7f511d32 100644 --- a/tools/aimdb-mcp/tests/persistent_client.rs +++ b/tools/aimdb-mcp/tests/persistent_client.rs @@ -79,12 +79,12 @@ async fn the_client_is_kept_across_error_replies_and_replaced_once_stopped() { let pool = ConnectionPool::new(); let server = Server::start(&sock); - let first = pool.get_drain_client(&endpoint).await.expect("connect"); + let first = pool.get_client(&endpoint).await.expect("connect"); // Error replies are answers on a healthy session: the client stays. for missing in ["missing", "also-missing"] { assert!(first.lock().await.get_record(missing).await.is_err()); - let again = pool.get_drain_client(&endpoint).await.unwrap(); + let again = pool.get_client(&endpoint).await.unwrap(); assert!(Arc::ptr_eq(&first, &again), "kept after an error reply"); } assert!(!first.lock().await.is_closed()); @@ -103,8 +103,23 @@ async fn the_client_is_kept_across_error_replies_and_replaced_once_stopped() { // The instance is back: the next call replaces the stopped client. let _server = Server::start(&sock); - let next = pool.get_drain_client(&endpoint).await.expect("reconnect"); + let next = pool.get_client(&endpoint).await.expect("reconnect"); assert!(!Arc::ptr_eq(&first, &next), "a stopped client is replaced"); let got = next.lock().await.get_record("state").await.expect("get"); assert_eq!(got.value["n"], 1); } + +#[tokio::test] +async fn the_records_resource_uses_the_pool() { + let dir = tempfile::tempdir().unwrap(); + let sock = dir.path().join("aimdb.sock"); + let endpoint = sock.to_str().unwrap().to_string(); + let pool = ConnectionPool::new(); + aimdb_mcp::tools::init_connection_pool(pool.clone()); + + let _server = Server::start(&sock); + aimdb_mcp::resources::records::read_records_resource(&endpoint) + .await + .expect("read the records resource"); + assert_eq!(pool.connection_count().await, 1, "the pool made the client"); +} diff --git a/tools/aimdb-mcp/tests/serial_session.rs b/tools/aimdb-mcp/tests/serial_session.rs index 80a5c68f..7c0c7a34 100644 --- a/tools/aimdb-mcp/tests/serial_session.rs +++ b/tools/aimdb-mcp/tests/serial_session.rs @@ -1,9 +1,9 @@ -//! `get_record` and `drain_record` share one serial session (#450). +//! Every `aimdb-mcp` tool shares one serial session (#451). //! //! A UART never closes, so a board learns of a new client only from a `hello`, -//! and a `hello` resets its session, drain cursors included. If `get_record` -//! opened a second connection, it would either fail (the tty is already open) -//! or say `hello` again and lose the cursor `drain_record` reads. +//! and a `hello` resets its session, drain cursors included. If a tool opened +//! its own connection, it would either fail (the tty is already open) or say +//! `hello` again and lose the cursor `drain_record` reads. //! //! The board is a `SerialServer` behind a pty; `aimdb-mcp` opens the other //! side as `serial://…`. A relay between the pty and the board counts the @@ -19,6 +19,7 @@ use std::sync::Arc; use std::time::Duration; use aimdb_core::buffer::BufferCfg; +use aimdb_core::remote::SecurityPolicy; use aimdb_core::AimDbBuilder; use aimdb_mcp::tools; use aimdb_serial_connector::framing::FrameAccumulator; @@ -130,13 +131,18 @@ fn relay(master: File, board_line: tokio::io::DuplexStream, hellos: Arc (Board, String, Arc) { let (board_end, line_end) = tokio::io::duplex(64 * 1024); let mut builder = AimDbBuilder::new() .runtime(Arc::new(TokioAdapter)) - .with_connector(SerialServer::new(TokioByteStream(board_end))); + .with_connector( + SerialServer::new(TokioByteStream(board_end)) + .security_policy(SecurityPolicy::read_write().with_writable_key("config")), + ); builder.configure::("events", |reg| { reg.buffer(BufferCfg::SpmcRing { capacity: 16 }) .with_remote_access(); @@ -144,6 +150,9 @@ async fn get_record_shares_the_drain_session_over_serial() { builder.configure::("state", |reg| { reg.buffer(BufferCfg::SingleLatest).with_remote_access(); }); + builder.configure::("config", |reg| { + reg.buffer(BufferCfg::SingleLatest).with_remote_access(); + }); let (board, runner) = builder.build().await.expect("build board"); tokio::spawn(runner.run()); board @@ -154,9 +163,17 @@ async fn get_record_shares_the_drain_session_over_serial() { let pty = openpty(); let hellos = Arc::new(AtomicUsize::new(0)); relay(pty.master, line_end, hellos.clone()); + // The relay threads keep their ends; the slave must outlive the test. + std::mem::forget(pty._slave); tools::init_connection_pool(aimdb_mcp::connection::ConnectionPool::new()); let endpoint = format!("serial://{}?baud=115200", pty.path); + (board, endpoint, hellos) +} + +#[tokio::test] +async fn get_record_shares_the_drain_session_over_serial() { + let (board, endpoint, hellos) = start_board().await; let drain = || { tools::drain_record(Some( json!({ "endpoint": endpoint, "record_name": "events" }), @@ -223,3 +240,118 @@ async fn get_record_shares_the_drain_session_over_serial() { "the board received exactly one hello" ); } + +#[tokio::test] +async fn every_tool_shares_the_drain_session_over_serial() { + let (board, endpoint, hellos) = start_board().await; + let at = |extra: Value| { + let mut args = json!({ "endpoint": endpoint }); + args.as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + Some(args) + }; + let drain = || tools::drain_record(at(json!({ "record_name": "events" }))); + + // 1. Open the ring's drain cursor. + let first = drain().await.expect("first drain"); + assert_eq!(first["count"], 0); + + // 2. The board publishes events. + let events = board.producer::("events").unwrap(); + for n in 1..=3 { + events.produce(Reading { n }); + } + + // 3. Call each tool once on the same session. + let listed = tools::list_records(at(json!({}))) + .await + .expect("list_records"); + assert_eq!(listed.as_array().unwrap().len(), 3); + tools::set_record(at(json!({ "record_name": "config", "value": { "n": 7 } }))) + .await + .expect("set_record"); + let info = tools::get_instance_info(at(json!({}))) + .await + .expect("get_instance_info"); + assert_eq!(info["endpoint"], json!(endpoint)); + tools::graph_nodes(at(json!({}))) + .await + .expect("graph_nodes"); + tools::graph_edges(at(json!({}))) + .await + .expect("graph_edges"); + tools::graph_topo_order(at(json!({}))) + .await + .expect("graph_topo_order"); + // This tool matches the record type name, not its key. + let metrics = tools::get_buffer_metrics(at(json!({ "record_key": "Reading" }))) + .await + .expect("get_buffer_metrics"); + assert_eq!(metrics["found"], true); + tools::reset_buffer_metrics(at(json!({}))) + .await + .expect("reset_buffer_metrics"); + let profiling = tools::get_stage_profiling(at(json!({ "record_key": "events" }))) + .await + .expect("get_stage_profiling"); + assert_eq!(profiling["found"], true); + tools::reset_stage_profiling(at(json!({}))) + .await + .expect("reset_stage_profiling"); + let got = tools::get_record(at(json!({ "record_name": "state" }))) + .await + .expect("get_record"); + assert_eq!(got["value"], json!({ "n": 42 })); + + // Error replies keep the session too. + let err = tools::drain_record(at(json!({ "record_name": "missing" }))) + .await + .expect_err("unknown record") + .to_string(); + assert!(err.contains("not_found"), "{err}"); + tools::set_record(at(json!({ "record_name": "state", "value": { "n": 0 } }))) + .await + .expect_err("a record that is not writable"); + // Let any deferred cleanup run before the session is used again. + tokio::time::sleep(Duration::from_millis(50)).await; + + // 4. The cursor survived: the drain returns the queued events. + let mut ns = Vec::new(); + for _ in 0..100 { + let page = drain().await.expect("drain"); + ns.extend( + page["values"] + .as_array() + .unwrap() + .iter() + .map(|v| v["value"]["n"].clone()), + ); + if ns.len() >= 3 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(ns, [1, 2, 3]); + + // 5. One session the whole time. + assert_eq!( + hellos.load(SeqCst), + 1, + "the board received exactly one hello" + ); +} + +/// Parallel first calls wait for one connect: racing their own would find the +/// tty busy. +#[tokio::test] +async fn concurrent_first_calls_open_one_session() { + let (_board, endpoint, hellos) = start_board().await; + let list = || tools::list_records(Some(json!({ "endpoint": endpoint }))); + + let (a, b, c, d) = tokio::join!(list(), list(), list(), list()); + for result in [a, b, c, d] { + result.expect("list_records"); + } + assert_eq!(hellos.load(SeqCst), 1, "one connect, one hello"); +} From 4212ee83e5e9ee9fdd8cb8898fb1bda1f9e2390f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexander=20Schn=C3=B6rch?= Date: Mon, 9 Nov 2026 21:48:55 +0100 Subject: [PATCH 2/2] feat(uds): warn at startup when max_connections is 1 (#451) A persistent client such as aimdb-mcp can occupy the only slot, and instance discovery's separate connection then fails. The default limit of 16 is not affected. Co-Authored-By: Claude Opus 5.5 --- aimdb-uds-connector/CHANGELOG.md | 6 ++++++ aimdb-uds-connector/src/lib.rs | 14 +++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/aimdb-uds-connector/CHANGELOG.md b/aimdb-uds-connector/CHANGELOG.md index 1eb032c0..c2663937 100644 --- a/aimdb-uds-connector/CHANGELOG.md +++ b/aimdb-uds-connector/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **Startup warning for `max_connections(1)`.** A persistent client such as + `aimdb-mcp` can occupy the only slot, and instance discovery's separate + connection then fails. + ## [0.2.0] - 2026-11-03 ### Changed (breaking, wire) diff --git a/aimdb-uds-connector/src/lib.rs b/aimdb-uds-connector/src/lib.rs index 13149de5..13d3dae4 100644 --- a/aimdb-uds-connector/src/lib.rs +++ b/aimdb-uds-connector/src/lib.rs @@ -56,12 +56,12 @@ use std::pin::Pin; use std::sync::Arc; use aimdb_core::connector::ConnectorBuilder; -use aimdb_core::log_info; use aimdb_core::remote::AimxConfig; use aimdb_core::session::aimx::{AimxCodec, AimxDispatch}; use aimdb_core::session::{ Dispatch, SessionClientConnector, SessionConfig, SessionLimits, SessionServerConnector, }; +use aimdb_core::{log_info, log_warn}; use aimdb_core::{AimDb, DbError, DbResult}; type BoxFuture = Pin + Send + 'static>>; @@ -133,6 +133,10 @@ impl UdsServer { } /// Maximum concurrently served connections. + /// + /// With `1`, a persistent client such as `aimdb-mcp` can occupy the only + /// slot, and instance discovery's separate connection then fails. The + /// server logs a warning at startup in that case. pub fn max_connections(mut self, max: usize) -> Self { self.config = self.config.max_connections(max); self @@ -162,6 +166,14 @@ impl ConnectorBuilder for UdsServer { let config = self.config.clone(); let scheme = self.scheme.clone(); Box::pin(async move { + if config.max_connections == 1 { + log_warn!( + "AimX UDS server at {} has max_connections=1: a persistent client such as \ + aimdb-mcp can occupy the only slot, and instance discovery's separate \ + connection will then fail", + config.socket_path + ); + } let session_config = SessionConfig { limits: SessionLimits { max_connections: config.max_connections,