Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions aimdb-uds-connector/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
14 changes: 13 additions & 1 deletion aimdb-uds-connector/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Box<dyn Future<Output = ()> + Send + 'static>>;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 5 additions & 0 deletions tools/aimdb-mcp/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
153 changes: 40 additions & 113 deletions tools/aimdb-mcp/src/connection.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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<Mutex<AimxConnection>>;

/// 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<Mutex<Option<SharedClient>>>;

/// 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<Mutex<HashMap<String, ConnectionEntry>>>,
/// Persistent drain clients — kept alive so drain readers accumulate values
/// Key: endpoint, Value: shared AimxConnection
drain_clients: Arc<Mutex<HashMap<String, Arc<tokio::sync::Mutex<AimxConnection>>>>>,
/// Key: endpoint
clients: Arc<Mutex<HashMap<String, ClientSlot>>>,
}

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()
}
}
Expand All @@ -40,120 +36,59 @@ 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<Mutex<AimxConnection>>` if `AimxConnection` becomes `Sync`).
pub async fn get_connection(&self, endpoint: &str) -> Result<AimxConnection, ClientError> {
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<Arc<tokio::sync::Mutex<AimxConnection>>, 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<SharedClient, ClientError> {
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));
}
debug!(
"🔁 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()
}
}

Expand All @@ -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);
}
}
7 changes: 4 additions & 3 deletions tools/aimdb-mcp/src/resources/records.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -81,8 +80,9 @@ pub async fn read_records_resource(socket_path: &str) -> McpResult<ResourceReadR
)));
}

// Connect to the instance
let client = AimxConnection::connect(socket_path).await?;
// The instance's persistent client, shared with the tools
let client = crate::tools::shared_client(socket_path).await?;
let client = client.lock().await;

// List records
let records = client.list_records().await?;
Expand Down Expand Up @@ -188,6 +188,7 @@ mod tests {

#[tokio::test]
async fn test_read_records_resource_invalid_socket() {
crate::tools::init_connection_pool(crate::connection::ConnectionPool::new());
let result = read_records_resource("/tmp/nonexistent.sock").await;
assert!(result.is_err());
}
Expand Down
19 changes: 4 additions & 15 deletions tools/aimdb-mcp/src/tools/buffer_metrics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
//! on the server).

use crate::error::{McpError, McpResult};
use aimdb_client::AimxConnection;
use serde::Deserialize;
use serde_json::{json, Value};
use tracing::debug;
Expand All @@ -24,18 +23,6 @@ struct ResetBufferMetricsParams {
endpoint: Option<String>,
}

async fn connect(endpoint: &str) -> McpResult<AimxConnection> {
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.
Expand All @@ -45,7 +32,8 @@ pub async fn get_buffer_metrics(args: Option<Value>) -> McpResult<Value> {
.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
Expand Down Expand Up @@ -75,7 +63,8 @@ pub async fn reset_buffer_metrics(args: Option<Value>) -> McpResult<Value> {
.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,
Expand Down
Loading
Loading