MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(mcp): enhance MCP server registration with error handling and improve transport process management

Browse Source 
refactor: update various tools for better error handling and path resolution
fix: improve markdown rendering and input display in TUI
This commit is contained in:
asepharyana
2026-07-20 13:12:04 +07:00
parent 89ee213454
commit 148ba4e07b
17 changed files with 242 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/domain/src/core/provider.rs
+29 -2
View File
@@ -218,14 +218,27 @@ impl SseParser {
///
/// Return: all `StreamEvent`s completed by this chunk.
pub fn feed(&mut self, chunk: &str) -> Vec<StreamEvent> {
self.buffer.push_str(chunk);
// Normalize \r\n and bare \r to \n for consistent line ending handling
let chunk = chunk.replace("\r\n", "\n").replace('\r', "\n");
// Prevent unbounded buffer growth for long lines without \n
const MAX_BUFFER_SIZE: usize = 1_048_576; // 1 MB
if self.buffer.len() + chunk.len() > MAX_BUFFER_SIZE {
tracing::warn!("SSE buffer exceeded maximum size, resetting");
self.buffer.clear();
self.event_type = None;
self.data_lines.clear();
}
self.buffer.push_str(&chunk);
let mut events = Vec::new();
while let Some(line_end) = self.buffer.find('\n') {
let line = self.buffer[..line_end].trim_end_matches('\r').to_string();
self.buffer = self.buffer[line_end + 1..].to_string();
if line.is_empty() {
events.extend(self.flush_event());
} else if let Some(ty) = line.strip_prefix("event: ") {
} else if let Some(ty) = line.strip_prefix("event:") {
// Handle both "event:foo" and "event: foo"
self.event_type = Some(ty.trim().to_string());
} else if let Some(data) = line.strip_prefix("data:") {
let data = data.trim_start().to_string();
@@ -372,6 +385,20 @@ impl SseParser {
}
d_events
}
"content_block_delta" => {
let mut d_events = Vec::new();
if let Some(delta) = value.get("delta") {
if let Some(content) = delta.get("text").and_then(|c| c.as_str()) {
d_events.push(StreamEvent::Token(content.to_string()));
}
if let Some(reasoning) =
delta.get("reasoning_content").and_then(|r| r.as_str())
{
d_events.push(StreamEvent::Reasoning(reasoning.to_string()));
}
}
d_events
}
_ => vec![],
};
apps/infrastructure/src/mcp/manager.rs
+11 -3
View File
@@ -32,7 +32,13 @@ impl Default for McpManager {
impl McpManager {
pub fn register(&mut self, name: &str, transport: &str) {
/// Register an MCP server by name and transport string.
///
/// Returns an error if a server with the same name is already registered.
pub fn register(&mut self, name: &str, transport: &str) -> anyhow::Result<()> {
if self.servers.contains_key(name) {
anyhow::bail!("MCP server '{name}' is already registered");
}
self.servers.insert(
name.to_string(),
McpServerHandle {
@@ -40,10 +46,12 @@ impl McpManager {
transport: transport.to_string(),
},
);
Ok(())
}
pub fn unregister(&mut self, name: &str) {
self.servers.remove(name);
/// Remove a registered MCP server and return its handle, if it existed.
pub fn unregister(&mut self, name: &str) -> Option<McpServerHandle> {
self.servers.remove(name)
}
pub fn list(&self) -> Vec<McpServerHandle> {
apps/infrastructure/src/mcp/transport.rs
+69 -16
View File
@@ -1,44 +1,97 @@
//! MCP transport layer — manages child-process and HTTP-based transport
//! for connecting to MCP servers.
use std::process::{Child, Command, Stdio};
use std::{
io::{Read, Write},
process::{Child, ChildStdin, ChildStdout, Command, Stdio},
};
/// A running MCP server process connected via stdio.
///
/// Holds the child process handle plus the piped stdin/stdout streams
/// so callers can send JSON-RPC messages and read responses.
pub struct McpTransport {
process: Option<Child>,
stdin: Option<ChildStdin>,
stdout: Option<ChildStdout>,
}
impl McpTransport {
pub fn start_child_process(command: &str, args: &[String]) -> anyhow::Result<Self> {
let child = Command::new(command)
.args(args)
/// Spawn a child process as an MCP server over stdio.
///
/// The command is passed to `sh -c` so shell syntax (pipes, redirects, etc.)
/// works naturally. Stderr is discarded to avoid corrupting a TUI that may
/// be running in the same terminal.
pub fn start_child_process(name: &str, command: &str) -> anyhow::Result<Self> {
tracing::info!("starting MCP transport '{name}': {command}");
let mut child = Command::new("sh")
.arg("-c")
.arg(command)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.stderr(Stdio::null())
.spawn()?;
let stdin = child.stdin.take();
let stdout = child.stdout.take();
Ok(McpTransport {
process: Some(child),
stdin,
stdout,
})
}
pub fn stop(&mut self) -> anyhow::Result<()> {
if let Some(mut child) = self.process.take() {
if let Err(e) = child.kill() {
tracing::warn!("MCP transport kill error: {e}");
}
/// Write raw bytes to the child's stdin.
pub fn send(&mut self, data: &[u8]) -> anyhow::Result<()> {
if let Some(ref mut stdin) = self.stdin {
stdin.write_all(data)?;
stdin.flush()?;
}
Ok(())
}
/// Read from the child's stdout into the provided buffer.
///
/// Returns `Ok(Some(n))` with the number of bytes read,
/// `Ok(None)` on EOF, or `Err` on I/O errors.
pub fn receive(&mut self, buf: &mut [u8]) -> anyhow::Result<Option<usize>> {
match self.stdout.as_mut() {
Some(stdout) => match stdout.read(buf) {
Ok(0) => Ok(None),
Ok(n) => Ok(Some(n)),
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(e.into()),
},
None => Ok(None),
}
}
/// Gracefully shut down the child by closing stdin (sending EOF) and
/// then killing the process.
pub fn kill(&mut self) {
// Close stdin first to signal EOF to the MCP server.
let _ = self.stdin.take();
if let Some(ref mut child) = self.process {
let _ = child.kill();
let _ = child.wait();
}
}
/// Check whether the child process is still running.
pub fn is_running(&mut self) -> bool {
self.process
.as_mut()
.is_some_and(|c| matches!(c.try_wait(), Ok(None)))
}
/// Stop the child process. This is the public API alias for `kill`.
pub fn stop(&mut self) -> anyhow::Result<()> {
self.kill();
Ok(())
}
}
impl Drop for McpTransport {
fn drop(&mut self) {
if let Some(mut child) = self.process.take() {
if let Err(e) = child.kill() {
tracing::warn!("MCP transport kill error: {e}");
}
let _ = child.wait();
}
self.kill();
}
}
apps/infrastructure/src/tools/plan.rs
+2 -1
View File
@@ -70,7 +70,8 @@ impl Tool for PlanReady {
if std::fs::create_dir_all(&plan_dir).is_ok() {
let filename = format!("plan-{}.md", chrono::Utc::now().format("%Y%m%d_%H%M%S"));
let path = plan_dir.join(&filename);
let _ = std::fs::write(&path, &plan_content);
std::fs::write(&path, &plan_content)
.map_err(|e| anyhow::anyhow!("failed to save plan: {e}"))?;
Ok(format!("Plan saved to {filename}. Starting execution."))
} else {
Ok("Plan is ready. Starting execution.".to_string())
apps/infrastructure/src/tools/search.rs
+8 -1
View File
@@ -55,7 +55,14 @@ impl Tool for Grep {
}
if let Ok(content) = fs::read_to_string(file_path) {
for (i, line) in content.lines().enumerate() {
if line.contains(&pattern) {
let is_match = if let Ok(re) = regex::Regex::new(&pattern) {
re.is_match(line)
} else {
// Fall back to literal substring search when the
// pattern is not a valid regex.
line.contains(&pattern)
};
if is_match {
let rel_path = file_path
.strip_prefix(&path)
.unwrap_or(file_path)
apps/infrastructure/src/tools/utility/cd.rs
+5 -4
View File
@@ -1,6 +1,6 @@
//! Change the working directory for subsequent commands.
use crate::tools::{Tool, ToolCtx};
use crate::tools::{resolve_path, Tool, ToolCtx};
use anyhow::Result;
use serde_json::{json, Value};
@@ -28,9 +28,10 @@ impl Tool for Cd {
})
}
fn run(&self, _ctx: &ToolCtx, args: &Value) -> Result<String> {
fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
let dir = crate::tools::arg_str(args, "directory")?;
std::env::set_current_dir(&dir)?;
Ok(format!("Changed directory to '{dir}'"))
let resolved = resolve_path(&ctx.workspaces, &dir)?;
std::env::set_current_dir(&resolved)?;
Ok(format!("Changed directory to '{}'", resolved.display()))
}
}
apps/infrastructure/src/tools/utility/dir_cache_update.rs
+18 -4
View File
@@ -1,8 +1,10 @@
//! Update the shared directory cache.
//! Update the shared directory cache by resolving each path against
//! workspaces and storing the resolved paths in `ctx.dir_cache`.
use crate::tools::ToolCtx;
use crate::tools::{resolve_path, ToolCtx};
use anyhow::Result;
use serde_json::{json, Value};
use std::path::PathBuf;
pub struct DirCacheUpdate;
@@ -28,7 +30,7 @@ impl crate::tools::Tool for DirCacheUpdate {
})
}
fn run(&self, _ctx: &ToolCtx, args: &Value) -> Result<String> {
fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
let paths: Vec<String> = args
.get("paths")
.and_then(|v| v.as_array())
@@ -39,7 +41,19 @@ impl crate::tools::Tool for DirCacheUpdate {
})
.unwrap_or_default();
let count = paths.len();
let resolved: Vec<PathBuf> = paths
.iter()
.map(|p| resolve_path(&ctx.workspaces, p))
.collect::<Result<Vec<_>>>()?;
let count = resolved.len();
// Persist the resolved paths into the shared DirCache so the TUI
// and other tools can read the cached listing without re-scanning.
let dc = ctx.dir_cache.clone();
let rt = tokio::runtime::Runtime::new()?;
rt.block_on(async { dc.write().await.set(resolved).await });
Ok(format!("Directory cache updated with {} entries", count))
}
}
apps/infrastructure/src/tools/workflow.rs
+24 -5
View File
@@ -9,7 +9,9 @@ use crate::tools::{arg_str, Tool, ToolCtx};
use crate::workflow::engine::execution::execute_workflow;
use crate::workflow::hive_mind::cycle::execute_cycle;
use crate::workflow::hive_mind::synthesis::synthesize_consensus;
use crate::workflow::hive_mind::types::{CognitiveCycle, NodeOutput};
use crate::workflow::hive_mind::types::{
CognitiveCycle, NodeDirective, NodeOutput,
};
use crate::workflow::script::WorkflowScript;
/// Execute a multi-step workflow defined in YAML.
@@ -102,10 +104,16 @@ impl Tool for NoteFinding {
fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
let finding = crate::tools::arg_str(args, "finding")?;
let category = args
.get("category")
.and_then(|v| v.as_str())
.unwrap_or("general");
let tagged = format!("[{category}] {finding}");
if let Some(ref findings) = ctx.workflow_findings {
if let Ok(mut guard) = findings.lock() {
guard.push(finding.clone());
guard.push(tagged);
}
}
@@ -205,13 +213,24 @@ impl Tool for HiveMind {
let mut all_node_outputs = Vec::new();
for (cycle_idx, cycle_val) in cycles_val.iter().enumerate() {
let directives: Vec<String> = cycle_val
let directives: Vec<NodeDirective> = cycle_val
.get("directives")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|d| d.get("directive").and_then(|v| v.as_str()))
.map(String::from)
.filter_map(|d| {
let directive = d
.get("directive")
.and_then(|v| v.as_str())?;
let access = d
.get("access")
.and_then(|v| v.as_str())
.unwrap_or("read");
Some(NodeDirective {
directive: directive.to_string(),
access_tier: access.to_string(),
})
})
.collect()
})
.unwrap_or_default();
apps/infrastructure/src/utils.rs
+31 -20
View File
@@ -69,31 +69,42 @@ impl CastOr<u32> for u128 {
/// Atomically write serializable `data` to `path`.
///
/// Flow: serialize to pretty JSON -> write to `path.tmp` -> fsync -> rename
/// Flow: serialize to pretty JSON -> write to `path.tmp.<uuid>` -> fsync -> rename
/// -> fsync parent. If `mode` is `Some`, set permissions before rename (Unix only).
///
/// A UUID-based temporary filename avoids collisions from concurrent writes.
/// Orphaned tmp files are cleaned up on any error after creation.
pub fn write_json_atomic<T: Serialize>(path: &Path, data: &T, mode: Option<u32>) -> std::io::Result<()> {
let tmp = path.with_extension("tmp");
let bytes = serde_json::to_vec_pretty(data)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
{
let mut f = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&tmp)?;
f.write_all(&bytes)?;
f.sync_all()?;
}
if let Some(m) = mode {
#[cfg(unix)]
let tmp = path.with_extension(format!("tmp.{}", uuid::Uuid::new_v4()));
let result = (|| -> std::io::Result<()> {
let bytes = serde_json::to_vec_pretty(data)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(m))?;
let mut f = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&tmp)?;
f.write_all(&bytes)?;
f.sync_all()?;
}
#[cfg(not(unix))]
{ let _ = m; }
if let Some(m) = mode {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(m))?;
}
#[cfg(not(unix))]
{ let _ = m; }
}
std::fs::rename(&tmp, path)?;
Ok(())
})();
// Clean up orphaned temp file on error after creation
if let Err(e) = result {
let _ = std::fs::remove_file(&tmp);
return Err(e);
}
std::fs::rename(&tmp, path)?;
if let Some(parent) = path.parent() {
let _ = std::fs::File::open(parent).and_then(|d| d.sync_all());
}
apps/infrastructure/src/workflow/hive_mind/cycle.rs
+13 -4
View File
@@ -56,25 +56,34 @@ pub async fn execute_cycle(
let cycle_index = cycle.index;
use crate::workflow::hive_mind::types::NodeDirective;
// Run all directives in this cycle concurrently.
let handles: Vec<_> = cycle
.directives
.iter()
.enumerate()
.map(|(i, directive)| {
let dir = directive.clone();
.map(|(i, node_dir): (usize, &NodeDirective)| {
let dir = node_dir.directive.clone();
let access_tier = node_dir.access_tier.clone();
let ctx = SubagentContext::new(
dir.clone(),
tool_ctx.clone(),
"full".to_string(),
access_tier.clone(),
base_url.clone(),
api_key.clone(),
model.clone(),
);
let tc = tool_ctx.clone();
let access = match access_tier.as_str() {
"write" => AccessTier::Write,
"full" => AccessTier::Full,
_ => AccessTier::Read,
};
async move {
let result = run_agent(ctx, &dir, AccessTier::Full, tc).await?;
let result = run_agent(ctx, &dir, access, tc).await?;
Ok::<NodeOutput, anyhow::Error>(NodeOutput {
id: format!("Node-{}-{}", cycle_index, i),
directive: dir,
apps/infrastructure/src/workflow/hive_mind/types.rs
+1 -1
View File
@@ -19,7 +19,7 @@ pub struct CognitiveCyclePlan {
#[derive(Debug, Clone)]
pub struct CognitiveCycle {
pub index: u32,
pub directives: Vec<String>,
pub directives: Vec<NodeDirective>,
}
/// Output from a single hive-mind processing node after a cycle completes.
apps/interfaces/daemon/src/handler.rs
+9 -3
View File
@@ -427,9 +427,15 @@ fn handle_open_editor(state: &mut AppStateRest, path: String) {
fn handle_mcp_add(state: &mut AppStateRest, name: String, command: String) {
tracing::info!("adding MCP server: {name}");
state.mcp_manager.register(&name, &command);
state.toast_success(format!("MCP server '{name}' added with command: {command}"));
state.dirty = true;
match state.mcp_manager.register(&name, &command) {
Ok(()) => {
state.toast_success(format!("MCP server '{name}' added with command: {command}"));
state.dirty = true;
}
Err(e) => {
state.toast_error(format!("Failed to add MCP server '{name}': {e}"));
}
}
}
fn handle_start_oauth(state: &mut AppStateRest, provider: String) {
apps/interfaces/tui/src/model/msglog/schema.rs
+1 -2
View File
@@ -17,8 +17,7 @@ pub fn init_schema(conn: &Connection) -> Result<()> {
tool_call_id TEXT,
tool_name TEXT,
tool_arguments TEXT,
created_at INTEGER NOT NULL,
FOREIGN KEY (session_id) REFERENCES archives(session_id)
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS archives (
apps/interfaces/tui/src/view/markdown.rs
+2 -2
View File
@@ -237,12 +237,12 @@ pub fn render_markdown(text: &str, width: u16, dim: bool) -> Vec<Span<'static>>
}
if r == 0 {
spans.push(Span::styled(
" |",
" | ",
apply_dim(Style::default().fg(Theme::BORDER), dim),
));
for w in &col_widths {
spans.push(Span::styled(
format!("{}-|", "-".repeat(*w + 2)),
format!("{}| ", "-".repeat(*w + 1)),
apply_dim(Style::default().fg(Theme::BORDER), dim),
));
}
apps/interfaces/tui/src/view/mod.rs
+9 -3
View File
@@ -147,7 +147,13 @@ fn render_input_bar(frame: &mut Frame, area: Rect, state: &AppStateRest) {
} else {
let (before, after) = input_text.split_at(cursor_pos);
spans.push(Span::raw(before.to_string()));
let cursor_char = if after.is_empty() { " " } else { &after[..1] };
let (cursor_char, after_char) = if after.is_empty() {
(" ".to_string(), "")
} else {
let c = after.chars().next().unwrap();
let char_len = c.len_utf8();
(c.to_string(), &after[char_len..])
};
spans.push(Span::styled(
cursor_char,
Style::default()
@@ -155,8 +161,8 @@ fn render_input_bar(frame: &mut Frame, area: Rect, state: &AppStateRest) {
.fg(Theme::BG)
.add_modifier(Modifier::BOLD),
));
if after.len() > 1 {
spans.push(Span::raw(after[1..].to_string()));
if !after_char.is_empty() {
spans.push(Span::raw(after_char.to_string()));
}
}
apps/interfaces/tui/src/view/overlays/key_input.rs
+8 -9
View File
@@ -22,19 +22,18 @@ pub fn render(
))
.border_style(Style::default().fg(Theme::WARNING));
let input_text = &state.input.buffer;
let display = if input_text.is_empty() {
" Type your API key..."
let char_count = input_text.chars().count();
let display: String = if input_text.is_empty() {
" Type your API key...".to_string()
} else if char_count > 8 {
input_text.chars().take(4).collect()
} else {
if input_text.len() > 8 {
&input_text[..4]
} else {
input_text.as_str()
}
input_text.to_string()
};
let masked = if input_text.is_empty() {
display.to_string()
display
} else {
let suffix = if input_text.len() > 8 { "****" } else { "" };
let suffix = if char_count > 8 { "****" } else { "" };
format!("{display}{suffix}")
};
let lines = vec![
apps/interfaces/tui/src/view/overlays/learning.rs
+2 -2
View File
@@ -89,12 +89,12 @@ pub fn render(
let max_lines = h_chunks[0].height.saturating_sub(2) as usize;
let selected = state.misc.selected_index;
let start_idx = if selected >= max_lines {
selected - max_lines + 1
selected.saturating_sub(max_lines).saturating_add(1)
} else {
0
};
let end_idx = (start_idx + max_lines).min(left_lines.len());
let visible_lines = if left_lines.is_empty() {
let visible_lines = if left_lines.is_empty() || start_idx >= left_lines.len() {
Vec::new()
} else {
left_lines[start_idx..end_idx].to_vec()