Skip to content

Commit 224cb9c

Browse files
lxsaahlxsaah
committed
refactor: update socket_path parameter to be optional and resolve from AIMDB_SOCKET env var
1 parent 68d5d88 commit 224cb9c

9 files changed

Lines changed: 103 additions & 75 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎tools/aimdb-mcp/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "aimdb-mcp"
3-
version = "0.6.1"
3+
version = "0.7.0"
44
edition = "2021"
55
authors.workspace = true
66
license.workspace = true

‎tools/aimdb-mcp/src/server.rs‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -140,10 +140,10 @@ impl McpServer {
140140
"properties": {
141141
"socket_path": {
142142
"type": "string",
143-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
143+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
144144
}
145145
},
146-
"required": ["socket_path"],
146+
"required": [],
147147
"additionalProperties": false
148148
}),
149149
},
@@ -155,14 +155,14 @@ impl McpServer {
155155
"properties": {
156156
"socket_path": {
157157
"type": "string",
158-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
158+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
159159
},
160160
"record_name": {
161161
"type": "string",
162162
"description": "Name of the record to retrieve (e.g., server::Temperature)"
163163
}
164164
},
165-
"required": ["socket_path", "record_name"],
165+
"required": ["record_name"],
166166
"additionalProperties": false
167167
}),
168168
},
@@ -174,7 +174,7 @@ impl McpServer {
174174
"properties": {
175175
"socket_path": {
176176
"type": "string",
177-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
177+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
178178
},
179179
"record_name": {
180180
"type": "string",
@@ -184,7 +184,7 @@ impl McpServer {
184184
"description": "New value for the record (must match record's type schema)"
185185
}
186186
},
187-
"required": ["socket_path", "record_name", "value"],
187+
"required": ["record_name", "value"],
188188
"additionalProperties": false
189189
}),
190190
},
@@ -196,10 +196,10 @@ impl McpServer {
196196
"properties": {
197197
"socket_path": {
198198
"type": "string",
199-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
199+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
200200
}
201201
},
202-
"required": ["socket_path"],
202+
"required": [],
203203
"additionalProperties": false
204204
}),
205205
},
@@ -216,7 +216,7 @@ impl McpServer {
216216
"properties": {
217217
"socket_path": {
218218
"type": "string",
219-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
219+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
220220
},
221221
"record_name": {
222222
"type": "string",
@@ -228,7 +228,7 @@ impl McpServer {
228228
"default": true
229229
}
230230
},
231-
"required": ["socket_path", "record_name"],
231+
"required": ["record_name"],
232232
"additionalProperties": false
233233
}),
234234
},
@@ -245,7 +245,7 @@ impl McpServer {
245245
"properties": {
246246
"socket_path": {
247247
"type": "string",
248-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
248+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
249249
},
250250
"record_name": {
251251
"type": "string",
@@ -257,7 +257,7 @@ impl McpServer {
257257
"minimum": 1
258258
}
259259
},
260-
"required": ["socket_path", "record_name"],
260+
"required": ["record_name"],
261261
"additionalProperties": false
262262
}),
263263
},
@@ -269,10 +269,10 @@ impl McpServer {
269269
"properties": {
270270
"socket_path": {
271271
"type": "string",
272-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
272+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
273273
}
274274
},
275-
"required": ["socket_path"],
275+
"required": [],
276276
"additionalProperties": false
277277
}),
278278
},
@@ -284,10 +284,10 @@ impl McpServer {
284284
"properties": {
285285
"socket_path": {
286286
"type": "string",
287-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
287+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
288288
}
289289
},
290-
"required": ["socket_path"],
290+
"required": [],
291291
"additionalProperties": false
292292
}),
293293
},
@@ -299,10 +299,10 @@ impl McpServer {
299299
"properties": {
300300
"socket_path": {
301301
"type": "string",
302-
"description": "Unix socket path to the AimDB instance (e.g., /tmp/aimdb-demo.sock)"
302+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
303303
}
304304
},
305-
"required": ["socket_path"],
305+
"required": [],
306306
"additionalProperties": false
307307
}),
308308
},
@@ -698,14 +698,14 @@ impl McpServer {
698698
"properties": {
699699
"socket_path": {
700700
"type": "string",
701-
"description": "Unix socket path to the running AimDB instance (e.g., /tmp/aimdb-demo.sock)"
701+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
702702
},
703703
"state_path": {
704704
"type": "string",
705705
"description": "Path to state.toml (default: .aimdb/state.toml)"
706706
}
707707
},
708-
"required": ["socket_path"],
708+
"required": [],
709709
"additionalProperties": false
710710
}),
711711
},
@@ -717,14 +717,14 @@ impl McpServer {
717717
"properties": {
718718
"socket_path": {
719719
"type": "string",
720-
"description": "Unix socket path to the AimDB instance"
720+
"description": "Unix socket path to the AimDB instance. Falls back to AIMDB_SOCKET env var if omitted."
721721
},
722722
"record_key": {
723723
"type": "string",
724724
"description": "Substring to match against record names (e.g., 'Temperature')"
725725
}
726726
},
727-
"required": ["socket_path", "record_key"],
727+
"required": ["record_key"],
728728
"additionalProperties": false
729729
}),
730730
},

‎tools/aimdb-mcp/src/tools/architecture.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -773,7 +773,8 @@ pub async fn rename_record(args: Option<Value>) -> McpResult<Value> {
773773

774774
#[derive(Debug, Deserialize)]
775775
struct ValidateInstanceParams {
776-
socket_path: String,
776+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
777+
socket_path: Option<String>,
777778
#[serde(default)]
778779
state_path: Option<String>,
779780
}
@@ -785,6 +786,7 @@ pub async fn validate_against_instance(args: Option<Value>) -> McpResult<Value>
785786
debug!("validate_against_instance called");
786787
let params: ValidateInstanceParams = serde_json::from_value(args.unwrap_or(Value::Null))
787788
.map_err(|e| McpError::InvalidParams(format!("validate_against_instance: {e}")))?;
789+
let socket_path = super::resolve_socket_path(params.socket_path)?;
788790

789791
let state_path = params
790792
.state_path
@@ -801,7 +803,7 @@ pub async fn validate_against_instance(args: Option<Value>) -> McpResult<Value>
801803
})?;
802804

803805
// Connect and list live records
804-
let mut client = AimxClient::connect(&params.socket_path)
806+
let mut client = AimxClient::connect(&socket_path)
805807
.await
806808
.map_err(McpError::Client)?;
807809

@@ -826,7 +828,8 @@ pub async fn validate_against_instance(args: Option<Value>) -> McpResult<Value>
826828

827829
#[derive(Debug, Deserialize)]
828830
struct GetBufferMetricsParams {
829-
socket_path: String,
831+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
832+
socket_path: Option<String>,
830833
record_key: String,
831834
}
832835

@@ -837,8 +840,9 @@ pub async fn get_buffer_metrics(args: Option<Value>) -> McpResult<Value> {
837840
debug!("get_buffer_metrics called");
838841
let params: GetBufferMetricsParams = serde_json::from_value(args.unwrap_or(Value::Null))
839842
.map_err(|e| McpError::InvalidParams(format!("get_buffer_metrics: {e}")))?;
843+
let socket_path = super::resolve_socket_path(params.socket_path)?;
840844

841-
let mut client = AimxClient::connect(&params.socket_path)
845+
let mut client = AimxClient::connect(&socket_path)
842846
.await
843847
.map_err(McpError::Client)?;
844848

‎tools/aimdb-mcp/src/tools/graph.rs‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,22 +9,22 @@ use tracing::debug;
99
/// Parameters for graph_nodes tool
1010
#[derive(Debug, Clone, Serialize, Deserialize)]
1111
struct GraphNodesParams {
12-
/// Unix socket path to the AimDB instance
13-
socket_path: String,
12+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
13+
socket_path: Option<String>,
1414
}
1515

1616
/// Parameters for graph_edges tool
1717
#[derive(Debug, Clone, Serialize, Deserialize)]
1818
struct GraphEdgesParams {
19-
/// Unix socket path to the AimDB instance
20-
socket_path: String,
19+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
20+
socket_path: Option<String>,
2121
}
2222

2323
/// Parameters for graph_topo_order tool
2424
#[derive(Debug, Clone, Serialize, Deserialize)]
2525
struct GraphTopoOrderParams {
26-
/// Unix socket path to the AimDB instance
27-
socket_path: String,
26+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
27+
socket_path: Option<String>,
2828
}
2929

3030
/// Get all nodes in the dependency graph
@@ -50,17 +50,18 @@ pub async fn graph_nodes(args: Option<Value>) -> McpResult<Value> {
5050
// Parse parameters
5151
let params: GraphNodesParams = serde_json::from_value(args.unwrap_or(Value::Null))
5252
.map_err(|e| McpError::InvalidParams(format!("Invalid parameters: {}", e)))?;
53+
let socket_path = super::resolve_socket_path(params.socket_path)?;
5354

54-
debug!("🔌 Connecting to {}", params.socket_path);
55+
debug!("🔌 Connecting to {}", socket_path);
5556

5657
// Get or create connection from pool (if available)
5758
let mut client = if let Some(pool) = super::connection_pool() {
58-
pool.get_connection(&params.socket_path)
59+
pool.get_connection(&socket_path)
5960
.await
6061
.map_err(McpError::Client)?
6162
} else {
6263
// Fallback to direct connection if pool not initialized
63-
AimxClient::connect(&params.socket_path)
64+
AimxClient::connect(&socket_path)
6465
.await
6566
.map_err(McpError::Client)?
6667
};
@@ -94,17 +95,18 @@ pub async fn graph_edges(args: Option<Value>) -> McpResult<Value> {
9495
// Parse parameters
9596
let params: GraphEdgesParams = serde_json::from_value(args.unwrap_or(Value::Null))
9697
.map_err(|e| McpError::InvalidParams(format!("Invalid parameters: {}", e)))?;
98+
let socket_path = super::resolve_socket_path(params.socket_path)?;
9799

98-
debug!("🔌 Connecting to {}", params.socket_path);
100+
debug!("🔌 Connecting to {}", socket_path);
99101

100102
// Get or create connection from pool (if available)
101103
let mut client = if let Some(pool) = super::connection_pool() {
102-
pool.get_connection(&params.socket_path)
104+
pool.get_connection(&socket_path)
103105
.await
104106
.map_err(McpError::Client)?
105107
} else {
106108
// Fallback to direct connection if pool not initialized
107-
AimxClient::connect(&params.socket_path)
109+
AimxClient::connect(&socket_path)
108110
.await
109111
.map_err(McpError::Client)?
110112
};
@@ -136,17 +138,18 @@ pub async fn graph_topo_order(args: Option<Value>) -> McpResult<Value> {
136138
// Parse parameters
137139
let params: GraphTopoOrderParams = serde_json::from_value(args.unwrap_or(Value::Null))
138140
.map_err(|e| McpError::InvalidParams(format!("Invalid parameters: {}", e)))?;
141+
let socket_path = super::resolve_socket_path(params.socket_path)?;
139142

140-
debug!("🔌 Connecting to {}", params.socket_path);
143+
debug!("🔌 Connecting to {}", socket_path);
141144

142145
// Get or create connection from pool (if available)
143146
let mut client = if let Some(pool) = super::connection_pool() {
144-
pool.get_connection(&params.socket_path)
147+
pool.get_connection(&socket_path)
145148
.await
146149
.map_err(McpError::Client)?
147150
} else {
148151
// Fallback to direct connection if pool not initialized
149-
AimxClient::connect(&params.socket_path)
152+
AimxClient::connect(&socket_path)
150153
.await
151154
.map_err(McpError::Client)?
152155
};

‎tools/aimdb-mcp/src/tools/instance.rs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ pub async fn discover_instances(_args: Option<Value>) -> McpResult<Value> {
5959
/// Parameters for get_instance_info tool
6060
#[derive(Debug, Deserialize)]
6161
struct GetInstanceInfoParams {
62-
socket_path: String,
62+
/// Unix socket path to the AimDB instance (falls back to AIMDB_SOCKET env)
63+
socket_path: Option<String>,
6364
}
6465

6566
/// Result from get_instance_info tool
@@ -93,23 +94,24 @@ pub async fn get_instance_info(args: Option<Value>) -> McpResult<Value> {
9394
))
9495
}
9596
};
97+
let socket_path = super::resolve_socket_path(params.socket_path)?;
9698

97-
debug!("🔍 Getting instance info for: {}", params.socket_path);
99+
debug!("🔍 Getting instance info for: {}", socket_path);
98100

99101
// Get or create connection from pool (if available)
100102
let client = if let Some(pool) = super::connection_pool() {
101-
pool.get_connection(&params.socket_path).await?
103+
pool.get_connection(&socket_path).await?
102104
} else {
103105
// Fallback to direct connection if pool not initialized
104-
AimxClient::connect(&params.socket_path).await?
106+
AimxClient::connect(&socket_path).await?
105107
};
106108

107109
// Get server info from the welcome message
108110
let server_info = client.server_info();
109111

110112
// Convert to result format
111113
let result = InstanceInfoResult {
112-
socket_path: params.socket_path.clone(),
114+
socket_path: socket_path.clone(),
113115
server_version: server_info.server.clone(),
114116
protocol_version: server_info.version.clone(),
115117
permissions: server_info.permissions.clone(),

‎tools/aimdb-mcp/src/tools/mod.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,20 @@ pub(crate) fn connection_pool() -> Option<&'static ConnectionPool> {
2424
CONNECTION_POOL.get()
2525
}
2626

27+
/// Resolve the socket path from an explicit argument or the `AIMDB_SOCKET` env var.
28+
///
29+
/// Returns an error if neither is set.
30+
pub(crate) fn resolve_socket_path(explicit: Option<String>) -> crate::error::McpResult<String> {
31+
explicit
32+
.or_else(|| std::env::var("AIMDB_SOCKET").ok())
33+
.filter(|s| !s.is_empty())
34+
.ok_or_else(|| {
35+
crate::error::McpError::InvalidParams(
36+
"Missing socket_path (pass it explicitly or set AIMDB_SOCKET env var)".into(),
37+
)
38+
})
39+
}
40+
2741
// Re-export tool functions
2842
pub use architecture::{
2943
get_architecture, get_buffer_metrics, propose_add_binary, propose_add_connector,

0 commit comments

Comments
 (0)