Skip to content

Commit 00b7410

Browse files
ZhiXiao-LinZhiXiao-Lin
committed
style: cargo fmt --all
1 parent 72f87dd commit 00b7410

5 files changed

Lines changed: 127 additions & 68 deletions

File tree

‎core/examples/test_mcp_parallel_fix.rs‎

Lines changed: 41 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,10 @@ async fn main() -> Result<()> {
3737
// -------------------------------------------------------------------------
3838
println!("--- Test 1: MCP tool name format ---");
3939
{
40-
use a3s_code_core::mcp::{manager::McpManager, protocol::{McpServerConfig, McpTransportConfig, McpTool}};
40+
use a3s_code_core::mcp::{
41+
manager::McpManager,
42+
protocol::{McpServerConfig, McpTool, McpTransportConfig},
43+
};
4144
use std::sync::Arc;
4245

4346
let manager = Arc::new(McpManager::new());
@@ -60,7 +63,8 @@ async fn main() -> Result<()> {
6063
let tool_name = wrappers[0].name();
6164
assert_eq!(
6265
tool_name, "mcp__github__create_issue",
63-
"Tool name must use double-underscore: got '{}'", tool_name
66+
"Tool name must use double-underscore: got '{}'",
67+
tool_name
6468
);
6569
println!(" tool name: {} ✓", tool_name);
6670
}
@@ -72,9 +76,12 @@ async fn main() -> Result<()> {
7276
// -------------------------------------------------------------------------
7377
println!("\n--- Test 2: add_mcp_server on live session ---");
7478
{
75-
use a3s_code_core::mcp::{manager::McpManager, protocol::{McpServerConfig, McpTransportConfig}};
76-
use std::sync::Arc;
79+
use a3s_code_core::mcp::{
80+
manager::McpManager,
81+
protocol::{McpServerConfig, McpTransportConfig},
82+
};
7783
use std::collections::HashMap;
84+
use std::sync::Arc;
7885

7986
// Check if npx is available
8087
let npx_available = std::process::Command::new("npx")
@@ -84,32 +91,38 @@ async fn main() -> Result<()> {
8491
.unwrap_or(false);
8592

8693
if npx_available {
87-
let session = agent.session(
88-
"/tmp",
89-
Some(SessionOptions::new().with_permissive_policy()),
90-
)?;
94+
let session =
95+
agent.session("/tmp", Some(SessionOptions::new().with_permissive_policy()))?;
9196

9297
let manager = Arc::new(McpManager::new());
93-
manager.register_server(McpServerConfig {
94-
name: "filesystem".to_string(),
95-
transport: McpTransportConfig::Stdio {
96-
command: "npx".to_string(),
97-
args: vec![
98-
"-y".to_string(),
99-
"@modelcontextprotocol/server-filesystem".to_string(),
100-
"/tmp".to_string(),
101-
],
102-
},
103-
enabled: true,
104-
env: HashMap::new(),
105-
oauth: None,
106-
tool_timeout_secs: 30,
107-
}).await;
108-
109-
match session.add_mcp_server(Arc::clone(&manager), "filesystem").await {
98+
manager
99+
.register_server(McpServerConfig {
100+
name: "filesystem".to_string(),
101+
transport: McpTransportConfig::Stdio {
102+
command: "npx".to_string(),
103+
args: vec![
104+
"-y".to_string(),
105+
"@modelcontextprotocol/server-filesystem".to_string(),
106+
"/tmp".to_string(),
107+
],
108+
},
109+
enabled: true,
110+
env: HashMap::new(),
111+
oauth: None,
112+
tool_timeout_secs: 30,
113+
})
114+
.await;
115+
116+
match session
117+
.add_mcp_server(Arc::clone(&manager), "filesystem")
118+
.await
119+
{
110120
Ok(count) => {
111121
println!(" add_mcp_server: registered {} tools ✓", count);
112-
assert!(count > 0, "Expected at least 1 tool from filesystem MCP server");
122+
assert!(
123+
count > 0,
124+
"Expected at least 1 tool from filesystem MCP server"
125+
);
113126
}
114127
Err(e) => {
115128
println!(" add_mcp_server failed (non-fatal): {}", e);
@@ -125,10 +138,8 @@ async fn main() -> Result<()> {
125138
// -------------------------------------------------------------------------
126139
println!("\n--- Test 3: Basic LLM send ---");
127140
{
128-
let session = agent.session(
129-
"/tmp",
130-
Some(SessionOptions::new().with_permissive_policy()),
131-
)?;
141+
let session =
142+
agent.session("/tmp", Some(SessionOptions::new().with_permissive_policy()))?;
132143

133144
let result = session.send("Reply with exactly: OK", None).await?;
134145
println!(" LLM response: '{}' ✓", result.text.trim());

‎core/examples/test_submit_api.rs‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ impl SessionCommand for ShellCommand {
4040
.output()
4141
.await?;
4242
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
43-
Ok(serde_json::json!({ "cmd": self.cmd, "output": stdout, "exit_code": output.status.code() }))
43+
Ok(
44+
serde_json::json!({ "cmd": self.cmd, "output": stdout, "exit_code": output.status.code() }),
45+
)
4446
}
4547
fn command_type(&self) -> &str {
4648
"shell"
@@ -79,14 +81,12 @@ async fn main() -> Result<()> {
7981
println!("{}", "=".repeat(70));
8082

8183
let agent = Agent::new(config_path.to_str().unwrap()).await?;
82-
let opts = SessionOptions::new().with_queue_config(
83-
SessionQueueConfig {
84-
query_max_concurrency: 8,
85-
execute_max_concurrency: 4,
86-
enable_metrics: true,
87-
..Default::default()
88-
},
89-
);
84+
let opts = SessionOptions::new().with_queue_config(SessionQueueConfig {
85+
query_max_concurrency: 8,
86+
execute_max_concurrency: 4,
87+
enable_metrics: true,
88+
..Default::default()
89+
});
9090
let session = agent.session(".", Some(opts))?;
9191

9292
// ── Test 1: single submit() ───────────────────────────────────────────────
@@ -114,9 +114,7 @@ async fn main() -> Result<()> {
114114
.collect();
115115

116116
let start = Instant::now();
117-
let rxs = session
118-
.submit_batch(SessionLane::Execute, commands)
119-
.await?;
117+
let rxs = session.submit_batch(SessionLane::Execute, commands).await?;
120118

121119
// Await all in parallel
122120
let results = futures::future::join_all(rxs).await;
@@ -171,8 +169,10 @@ async fn main() -> Result<()> {
171169
// ── Queue stats ───────────────────────────────────────────────────────────
172170
let stats = session.queue_stats().await;
173171
println!("\nQueue stats:");
174-
println!(" pending={} active={} external_pending={}",
175-
stats.total_pending, stats.total_active, stats.external_pending);
172+
println!(
173+
" pending={} active={} external_pending={}",
174+
stats.total_pending, stats.total_active, stats.external_pending
175+
);
176176

177177
println!("\n{}", "=".repeat(70));
178178
println!("All tests passed.");

‎core/src/agent_api.rs‎

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,7 +1496,8 @@ impl AgentSession {
14961496
&self,
14971497
lane: SessionLane,
14981498
commands: Vec<Box<dyn crate::queue::SessionCommand>>,
1499-
) -> anyhow::Result<Vec<tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>>> {
1499+
) -> anyhow::Result<Vec<tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>>>
1500+
{
15001501
let queue = self
15011502
.command_queue
15021503
.as_ref()
@@ -1529,12 +1530,13 @@ impl AgentSession {
15291530
manager: Arc<crate::mcp::manager::McpManager>,
15301531
server_name: &str,
15311532
) -> crate::error::Result<usize> {
1532-
manager.connect(server_name).await.map_err(|e| {
1533-
crate::error::CodeError::Tool {
1533+
manager
1534+
.connect(server_name)
1535+
.await
1536+
.map_err(|e| crate::error::CodeError::Tool {
15341537
tool: server_name.to_string(),
15351538
message: format!("Failed to connect MCP server: {}", e),
1536-
}
1537-
})?;
1539+
})?;
15381540

15391541
let tools = manager.get_server_tools(server_name).await;
15401542
let count = tools.len();
@@ -1574,9 +1576,13 @@ mod tests {
15741576
async fn execute(&self) -> anyhow::Result<serde_json::Value> {
15751577
Ok(serde_json::json!(null))
15761578
}
1577-
fn command_type(&self) -> &str { "noop" }
1579+
fn command_type(&self) -> &str {
1580+
"noop"
1581+
}
15781582
}
1579-
let result: anyhow::Result<tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>> = session.submit(SessionLane::Query, Box::new(Noop)).await;
1583+
let result: anyhow::Result<
1584+
tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>,
1585+
> = session.submit(SessionLane::Query, Box::new(Noop)).await;
15801586
assert!(result.is_err());
15811587
assert!(result.unwrap_err().to_string().contains("No queue"));
15821588
}
@@ -1591,10 +1597,14 @@ mod tests {
15911597
async fn execute(&self) -> anyhow::Result<serde_json::Value> {
15921598
Ok(serde_json::json!(null))
15931599
}
1594-
fn command_type(&self) -> &str { "noop" }
1600+
fn command_type(&self) -> &str {
1601+
"noop"
1602+
}
15951603
}
15961604
let cmds: Vec<Box<dyn crate::queue::SessionCommand>> = vec![Box::new(Noop)];
1597-
let result: anyhow::Result<Vec<tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>>> = session.submit_batch(SessionLane::Query, cmds).await;
1605+
let result: anyhow::Result<
1606+
Vec<tokio::sync::oneshot::Receiver<anyhow::Result<serde_json::Value>>>,
1607+
> = session.submit_batch(SessionLane::Query, cmds).await;
15981608
assert!(result.is_err());
15991609
assert!(result.unwrap_err().to_string().contains("No queue"));
16001610
}
@@ -2270,11 +2280,16 @@ mod tests {
22702280
async fn execute(&self) -> anyhow::Result<serde_json::Value> {
22712281
Ok(self.0.clone())
22722282
}
2273-
fn command_type(&self) -> &str { "echo" }
2283+
fn command_type(&self) -> &str {
2284+
"echo"
2285+
}
22742286
}
22752287

22762288
let rx = session
2277-
.submit(SessionLane::Query, Box::new(Echo(serde_json::json!({"ok": true}))))
2289+
.submit(
2290+
SessionLane::Query,
2291+
Box::new(Echo(serde_json::json!({"ok": true}))),
2292+
)
22782293
.await
22792294
.expect("submit should succeed with queue configured");
22802295

‎core/src/session/manager.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,11 @@ use super::{ContextUsage, Session, SessionConfig, SessionState};
44
use crate::agent::{AgentConfig, AgentEvent, AgentLoop, AgentResult};
55
use crate::hitl::ConfirmationPolicy;
66
use crate::llm::{self, LlmClient, LlmConfig, Message};
7+
use crate::mcp::McpManager;
78
use crate::memory::AgentMemory;
89
use crate::prompts::SystemPromptSlots;
910
use crate::skills::SkillRegistry;
1011
use crate::store::{FileSessionStore, LlmConfigData, SessionData, SessionStore};
11-
use crate::mcp::McpManager;
1212
use crate::tools::ToolExecutor;
1313
use a3s_memory::MemoryStore;
1414
use anyhow::{Context, Result};
@@ -176,7 +176,9 @@ impl SessionManager {
176176
by_server.entry(server).or_default().push(tool);
177177
}
178178
for (server_name, tools) in by_server {
179-
for tool in crate::mcp::tools::create_mcp_tools(&server_name, tools, Arc::clone(&manager)) {
179+
for tool in
180+
crate::mcp::tools::create_mcp_tools(&server_name, tools, Arc::clone(&manager))
181+
{
180182
self.tool_executor.register_dynamic_tool(tool);
181183
}
182184
}
@@ -218,12 +220,15 @@ impl SessionManager {
218220
if let Some(ref manager) = *guard {
219221
manager.disconnect(name).await?;
220222
}
221-
self.tool_executor.unregister_tools_by_prefix(&format!("mcp__{name}__"));
223+
self.tool_executor
224+
.unregister_tools_by_prefix(&format!("mcp__{name}__"));
222225
Ok(())
223226
}
224227

225228
/// Get status of all connected MCP servers.
226-
pub async fn mcp_status(&self) -> std::collections::HashMap<String, crate::mcp::McpServerStatus> {
229+
pub async fn mcp_status(
230+
&self,
231+
) -> std::collections::HashMap<String, crate::mcp::McpServerStatus> {
227232
let guard = self.mcp_manager.read().await;
228233
match guard.as_ref() {
229234
Some(m) => {

‎core/tests/integration_mcp.rs‎

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,10 @@ fn echo_mcp_config(name: &str) -> McpServerConfig {
5757
async fn test_mcp_status_empty_when_no_manager() {
5858
let manager = make_session_manager();
5959
let status = manager.mcp_status().await;
60-
assert!(status.is_empty(), "status should be empty before any MCP server is added");
60+
assert!(
61+
status.is_empty(),
62+
"status should be empty before any MCP server is added"
63+
);
6164
}
6265

6366
#[tokio::test]
@@ -72,15 +75,22 @@ async fn test_add_mcp_server_registers_in_status() {
7275

7376
// connect() will fail because `cat` doesn't implement MCP initialize.
7477
// The important thing is that the error is returned cleanly.
75-
assert!(result.is_err(), "connecting to a non-MCP process should fail");
78+
assert!(
79+
result.is_err(),
80+
"connecting to a non-MCP process should fail"
81+
);
7682
}
7783

7884
#[tokio::test]
7985
async fn test_remove_mcp_server_nonexistent_is_ok() {
8086
let manager = make_session_manager();
8187
// Removing a server that was never added should not panic.
8288
let result = manager.remove_mcp_server("nonexistent").await;
83-
assert!(result.is_ok(), "removing nonexistent server should be a no-op: {:?}", result.err());
89+
assert!(
90+
result.is_ok(),
91+
"removing nonexistent server should be a no-op: {:?}",
92+
result.err()
93+
);
8494
}
8595

8696
#[tokio::test]
@@ -131,7 +141,9 @@ async fn test_parallel_tasks_via_session() {
131141
use a3s_code_core::{Agent, CodeConfig};
132142

133143
let config = CodeConfig::from_file(&config_path()).expect("failed to load config.hcl");
134-
let agent = Agent::from_config(config).await.expect("failed to create agent");
144+
let agent = Agent::from_config(config)
145+
.await
146+
.expect("failed to create agent");
135147

136148
let tmp = tempfile::tempdir().unwrap();
137149
let session = agent
@@ -198,17 +210,33 @@ async fn test_add_real_mcp_filesystem_server() {
198210
};
199211

200212
let result = manager.add_mcp_server(config).await;
201-
assert!(result.is_ok(), "filesystem MCP server should connect: {:?}", result.err());
213+
assert!(
214+
result.is_ok(),
215+
"filesystem MCP server should connect: {:?}",
216+
result.err()
217+
);
202218

203219
let status = manager.mcp_status().await;
204-
assert!(status.contains_key("filesystem"), "filesystem server should appear in status");
205-
assert!(status["filesystem"].connected, "filesystem server should be connected");
206-
assert!(status["filesystem"].tool_count > 0, "filesystem server should expose tools");
220+
assert!(
221+
status.contains_key("filesystem"),
222+
"filesystem server should appear in status"
223+
);
224+
assert!(
225+
status["filesystem"].connected,
226+
"filesystem server should be connected"
227+
);
228+
assert!(
229+
status["filesystem"].tool_count > 0,
230+
"filesystem server should expose tools"
231+
);
207232

208233
// Clean up
209234
let _ = manager.remove_mcp_server("filesystem").await;
210235
let status_after = manager.mcp_status().await;
211236
if let Some(s) = status_after.get("filesystem") {
212-
assert!(!s.connected, "filesystem server should be disconnected after removal");
237+
assert!(
238+
!s.connected,
239+
"filesystem server should be disconnected after removal"
240+
);
213241
}
214242
}

0 commit comments

Comments
 (0)