已合并
fix(mcp): retry rate-limited RAM-A initialization [AI: Codex GPT-5] #293
DoraA_Mengjie创建于 7月30日
fix(mcp): retry rate-limited RAM-A initialization [AI: Codex GPT-5] #293
已合并
共 38 个文件变更+6335-260
| @@ -1898,10 +1898,12 @@ name = "mcp" | |||
| 1898 | version = "0.1.0" | 1898 | version = "0.1.0" |
| 1899 | dependencies = [ | 1899 | dependencies = [ |
| 1900 | "async-trait", | 1900 | "async-trait", |
| 1901 | + "axum", | ||
| 1901 | "futures-util", | 1902 | "futures-util", |
| 1902 | "reqwest", | 1903 | "reqwest", |
| 1903 | "serde", | 1904 | "serde", |
| 1904 | "serde_json", | 1905 | "serde_json", |
| 1906 | + "tempfile", | ||
| 1905 | "thiserror 1.0.69", | 1907 | "thiserror 1.0.69", |
| 1906 | "tokio", | 1908 | "tokio", |
| 1907 | "toml", | 1909 | "toml", |
| @@ -5188,6 +5190,7 @@ dependencies = [ | |||
| 5188 | "anyhow", | 5190 | "anyhow", |
| 5189 | "async-trait", | 5191 | "async-trait", |
| 5190 | "base64", | 5192 | "base64", |
| 5193 | + "chrono", | ||
| 5191 | "compact", | 5194 | "compact", |
| 5192 | "dirs 5.0.1", | 5195 | "dirs 5.0.1", |
| 5193 | "futures", | 5196 | "futures", |
| @@ -6,6 +6,7 @@ use serde::Deserialize; | |||
| 6 | use serde_json::Value; | 6 | use serde_json::Value; |
| 7 | use std::collections::BTreeMap; | 7 | use std::collections::BTreeMap; |
| 8 | use std::path::{Path, PathBuf}; | 8 | use std::path::{Path, PathBuf}; |
| 9 | +use xiaoo_shared::gateway::MemoryAutomationConfig; | ||
| 9 | 10 | ||
| 10 | const CONFIG_ENV_VAR: &str = "XIAOO_CONFIG"; | 11 | const CONFIG_ENV_VAR: &str = "XIAOO_CONFIG"; |
| 11 | 12 | ||
| @@ -24,6 +25,8 @@ pub struct FileConfig { | |||
| 24 | pub subagent: BTreeMap<String, SubagentRoleConfig>, | 25 | pub subagent: BTreeMap<String, SubagentRoleConfig>, |
| 25 | 26 | ||
| 26 | pub mcp: McpSection, | 27 | pub mcp: McpSection, |
| 28 | + | ||
| 29 | + pub memory_automation: MemoryAutomationConfig, | ||
| 27 | } | 30 | } |
| 28 | 31 | ||
| 29 | 32 | ||
| @@ -115,6 +118,13 @@ impl FileConfig { | |||
| 115 | subagent: parse_optional_section(&root, "subagent", &path, debug) | 118 | subagent: parse_optional_section(&root, "subagent", &path, debug) |
| 116 | .unwrap_or_default(), | 119 | .unwrap_or_default(), |
| 117 | mcp: parse_optional_section(&root, "mcp", &path, debug).unwrap_or_default(), | 120 | mcp: parse_optional_section(&root, "mcp", &path, debug).unwrap_or_default(), |
| 121 | + memory_automation: parse_optional_section( | ||
| 122 | + &root, | ||
| 123 | + "memory_automation", | ||
| 124 | + &path, | ||
| 125 | + debug, | ||
| 126 | + ) | ||
| 127 | + .unwrap_or_default(), | ||
| 118 | } | 128 | } |
| 119 | } | 129 | } |
| 120 | Err(e) => { | 130 | Err(e) => { |
| @@ -125,6 +135,22 @@ impl FileConfig { | |||
| 125 | Err(_) => Self::default(), | 135 | Err(_) => Self::default(), |
| 126 | } | 136 | } |
| 127 | } | 137 | } |
| 138 | + | ||
| 139 | + pub fn resolve_mcp_servers( | ||
| 140 | + &self, | ||
| 141 | + explicit_path: Option<&Path>, | ||
| 142 | + workspace: &Path, | ||
| 143 | + home: Option<&Path>, | ||
| 144 | + toml_source: &Path, | ||
| 145 | + ) -> Result<Vec<mcp::McpServerConfig>, mcp::McpConfigError> { | ||
| 146 | + crate::support::config::load_merged_mcp_servers( | ||
| 147 | + &self.mcp.servers, | ||
| 148 | + explicit_path, | ||
| 149 | + workspace, | ||
| 150 | + home, | ||
| 151 | + toml_source, | ||
| 152 | + ) | ||
| 153 | + } | ||
| 128 | } | 154 | } |
| 129 | 155 | ||
| 130 | fn parse_optional_section<T>( | 156 | fn parse_optional_section<T>( |
| @@ -242,4 +268,40 @@ provider = "anthropic" | |||
| 242 | let config = FileConfig::load_from_path(temp_file.path(), false); | 268 | let config = FileConfig::load_from_path(temp_file.path(), false); |
| 243 | assert_eq!(config.subagent.len(), 0); | 269 | assert_eq!(config.subagent.len(), 0); |
| 244 | } | 270 | } |
| 271 | + | ||
| 272 | + | ||
| 273 | + fn test_loads_memory_automation_config() { | ||
| 274 | + let config_content = r#" | ||
| 275 | +[memory_automation] | ||
| 276 | +enabled = true | ||
| 277 | +server = "ram-a" | ||
| 278 | +recall_top_k = 3 | ||
| 279 | +recall_token_budget = 128 | ||
| 280 | +context_messages = 2 | ||
| 281 | +queue_path = "/tmp/xiaoo-memory-queue.jsonl" | ||
| 282 | +queue_capacity = 32 | ||
| 283 | +max_retries = 4 | ||
| 284 | +retry_backoff_ms = 50 | ||
| 285 | +allowed_agent_roles = ["main", "researcher"] | ||
| 286 | +"#; | ||
| 287 | + | ||
| 288 | + let mut temp_file = NamedTempFile::new().unwrap(); | ||
| 289 | + temp_file.write_all(config_content.as_bytes()).unwrap(); | ||
| 290 | + temp_file.flush().unwrap(); | ||
| 291 | + | ||
| 292 | + let config = FileConfig::load_from_path(temp_file.path(), false); | ||
| 293 | + | ||
| 294 | + assert!(config.memory_automation.enabled); | ||
| 295 | + assert_eq!(config.memory_automation.server, "ram-a"); | ||
| 296 | + assert_eq!(config.memory_automation.recall_top_k, 3); | ||
| 297 | + assert_eq!(config.memory_automation.recall_token_budget, 128); | ||
| 298 | + assert_eq!(config.memory_automation.context_messages, 2); | ||
| 299 | + assert_eq!(config.memory_automation.queue_capacity, 32); | ||
| 300 | + assert_eq!(config.memory_automation.max_retries, 4); | ||
| 301 | + assert_eq!(config.memory_automation.retry_backoff_ms, 50); | ||
| 302 | + assert_eq!( | ||
| 303 | + config.memory_automation.allowed_agent_roles, | ||
| 304 | + vec!["main".to_string(), "researcher".to_string()] | ||
| 305 | + ); | ||
| 306 | + } | ||
| 245 | } | 307 | } |
| @@ -1,5 +1,5 @@ | |||
| 1 | use std::io::Write; | 1 | use std::io::Write; |
| 2 | -use std::path::PathBuf; | 2 | +use std::path::{Path, PathBuf}; |
| 3 | use std::sync::Arc; | 3 | use std::sync::Arc; |
| 4 | 4 | ||
| 5 | use crate::cli::config::FileConfig; | 5 | use crate::cli::config::FileConfig; |
| @@ -18,7 +18,7 @@ use skill::types::config::SkillsConfig; | |||
| 18 | use xiaoo_shared::gateway::{ | 18 | use xiaoo_shared::gateway::{ |
| 19 | session_record::SubagentRoleRecord, AppBootstrap, AppTurnRequest, GatewayEntryContext, | 19 | session_record::SubagentRoleRecord, AppBootstrap, AppTurnRequest, GatewayEntryContext, |
| 20 | HostedSessionRuntimeConfig, HostedSessionRuntimeResolver, InMemorySessionStore, | 20 | HostedSessionRuntimeConfig, HostedSessionRuntimeResolver, InMemorySessionStore, |
| 21 | - LlmRuntimeConfig, SessionDetachRequest, SessionOpenRequest, | 21 | + LlmRuntimeConfig, McpMemoryAutomation, SessionDetachRequest, SessionOpenRequest, |
| 22 | SessionRuntimeBindings, SessionRuntimeDescriptor, SessionRuntimeResolver, SessionStore, | 22 | SessionRuntimeBindings, SessionRuntimeDescriptor, SessionRuntimeResolver, SessionStore, |
| 23 | }; | 23 | }; |
| 24 | 24 | ||
| @@ -36,6 +36,10 @@ struct Args { | |||
| 36 | 36 | ||
| 37 | config: Option<String>, | 37 | config: Option<String>, |
| 38 | 38 | ||
| 39 | + /// Path to standard MCP JSON config (default discovery uses .mcp.json) | ||
| 40 | + | ||
| 41 | + mcp_config: Option<PathBuf>, | ||
| 42 | + | ||
| 39 | /// Show intermediate results (turns, tool calls, tokens) | 43 | /// Show intermediate results (turns, tool calls, tokens) |
| 40 | 44 | ||
| 41 | debug: bool, | 45 | debug: bool, |
| @@ -191,6 +195,7 @@ where | |||
| 191 | let args = Args::parse_from(args); | 195 | let args = Args::parse_from(args); |
| 192 | let debug = args.debug; | 196 | let debug = args.debug; |
| 193 | let config_path = FileConfig::resolve_path(args.config.as_deref()); | 197 | let config_path = FileConfig::resolve_path(args.config.as_deref()); |
| 198 | + let mcp_config_path = args.mcp_config; | ||
| 194 | 199 | ||
| 195 | if args.version { | 200 | if args.version { |
| 196 | println!("{}", env!("CARGO_PKG_VERSION")); | 201 | println!("{}", env!("CARGO_PKG_VERSION")); |
| @@ -255,6 +260,20 @@ where | |||
| 255 | let reasoning_effort = reasoning_effort.unwrap_or_default(); | 260 | let reasoning_effort = reasoning_effort.unwrap_or_default(); |
| 256 | 261 | ||
| 257 | let skills_config = resolve_skills_config_from_file(&file_cfg); | 262 | let skills_config = resolve_skills_config_from_file(&file_cfg); |
| 263 | + let workspace = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); | ||
| 264 | + let default_toml_source = Path::new("config.toml"); | ||
| 265 | + let mcp_servers = match file_cfg.resolve_mcp_servers( | ||
| 266 | + mcp_config_path.as_deref(), | ||
| 267 | + &workspace, | ||
| 268 | + dirs::home_dir().as_deref(), | ||
| 269 | + config_path.as_deref().unwrap_or(default_toml_source), | ||
| 270 | + ) { | ||
| 271 | + Ok(servers) => servers, | ||
| 272 | + Err(error) => { | ||
| 273 | + eprintln!("Failed to load MCP config: {error}"); | ||
| 274 | + std::process::exit(1); | ||
| 275 | + } | ||
| 276 | + }; | ||
| 258 | 277 | ||
| 259 | let config = CliConfig { | 278 | let config = CliConfig { |
| 260 | provider, | 279 | provider, |
| @@ -281,11 +300,12 @@ where | |||
| 281 | operation_backend: file_cfg.operation_backend.clone(), | 300 | operation_backend: file_cfg.operation_backend.clone(), |
| 282 | skills_config, | 301 | skills_config, |
| 283 | subagent: file_cfg.subagent.clone(), | 302 | subagent: file_cfg.subagent.clone(), |
| 284 | - mcp_servers: file_cfg.mcp.servers.clone(), | 303 | + mcp_servers, |
| 304 | + memory_automation: file_cfg.memory_automation.clone(), | ||
| 285 | }; | 305 | }; |
| 286 | 306 | ||
| 287 | -let session_title = title.or_else(|| generate_title_from_prompt(&prompt)); | 307 | + let session_title = title.or_else(|| generate_title_from_prompt(&prompt)); |
| 288 | - | 308 | + |
| 289 | run_once( | 309 | run_once( |
| 290 | config, | 310 | config, |
| 291 | prompt, | 311 | prompt, |
| @@ -295,12 +315,17 @@ let session_title = title.or_else(|| generate_title_from_prompt(&prompt)); | |||
| 295 | session, | 315 | session, |
| 296 | agent, | 316 | agent, |
| 297 | attach, | 317 | attach, |
| 298 | - ).await; | 318 | + ) |
| 319 | + .await; | ||
| 299 | } | 320 | } |
| 300 | Some(Command::Serve { port, hostname }) => { | 321 | Some(Command::Serve { port, hostname }) => { |
| 301 | handle_serve_command(port, hostname).await; | 322 | handle_serve_command(port, hostname).await; |
| 302 | } | 323 | } |
| 303 | - Some(Command::Export { session_id, port, client_id }) => { | 324 | + Some(Command::Export { |
| 325 | + session_id, | ||
| 326 | + port, | ||
| 327 | + client_id, | ||
| 328 | + }) => { | ||
| 304 | handle_export_command(session_id, port, client_id).await; | 329 | handle_export_command(session_id, port, client_id).await; |
| 305 | } | 330 | } |
| 306 | Some(Command::Debug { command }) => { | 331 | Some(Command::Debug { command }) => { |
| @@ -897,7 +922,7 @@ async fn run_once( | |||
| 897 | } | 922 | } |
| 898 | } | 923 | } |
| 899 | 924 | ||
| 900 | -if let Some(attach_url) = &attach { | 925 | + if let Some(attach_url) = &attach { |
| 901 | run_with_attach(attach_url, prompt, format, title, session, agent, debug).await; | 926 | run_with_attach(attach_url, prompt, format, title, session, agent, debug).await; |
| 902 | return; | 927 | return; |
| 903 | } | 928 | } |
| @@ -1009,6 +1034,7 @@ if let Some(attach_url) = &attach { | |||
| 1009 | }) | 1034 | }) |
| 1010 | .collect(), | 1035 | .collect(), |
| 1011 | mcp_servers: config.mcp_servers.clone(), | 1036 | mcp_servers: config.mcp_servers.clone(), |
| 1037 | + memory_automation: config.memory_automation.clone(), | ||
| 1012 | }; | 1038 | }; |
| 1013 | 1039 | ||
| 1014 | // 4. Bindings (CliEventSink for debug output) | 1040 | // 4. Bindings (CliEventSink for debug output) |
| @@ -1025,15 +1051,44 @@ if let Some(attach_url) = &attach { | |||
| 1025 | 1051 | ||
| 1026 | // 5. Bootstrap gateway | 1052 | // 5. Bootstrap gateway |
| 1027 | let store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default()); | 1053 | let store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default()); |
| 1054 | + let memory_automation = match McpMemoryAutomation::connect( | ||
| 1055 | + config.memory_automation.clone(), | ||
| 1056 | + &config.mcp_servers, | ||
| 1057 | + ) | ||
| 1058 | + .await | ||
| 1059 | + { | ||
| 1060 | + Ok(automation) => automation, | ||
| 1061 | + Err(error) => { | ||
| 1062 | + tracing::warn!(error = %error, "memory automation disabled after CLI startup error"); | ||
| 1063 | + None | ||
| 1064 | + } | ||
| 1065 | + }; | ||
| 1066 | + let memory_automation_for_shutdown = memory_automation.clone(); | ||
| 1028 | let resolver: Arc<dyn SessionRuntimeResolver> = | 1067 | let resolver: Arc<dyn SessionRuntimeResolver> = |
| 1029 | Arc::new(HostedSessionRuntimeResolver::new(runtime_config, bindings)); | 1068 | Arc::new(HostedSessionRuntimeResolver::new(runtime_config, bindings)); |
| 1030 | - let deps = match AppBootstrap::from_session_components_with_hooks( | 1069 | + let deps = match AppBootstrap::from_session_components_with_hooks_and_backend_manager_and_memory_automation( |
| 1031 | store, | 1070 | store, |
| 1032 | resolver, | 1071 | resolver, |
| 1033 | config.hooker.clone(), | 1072 | config.hooker.clone(), |
| 1073 | + Arc::new(xiaoo_shared::backend::BackendManager::new()), | ||
| 1074 | + memory_automation, | ||
| 1034 | ) { | 1075 | ) { |
| 1035 | Ok(d) => d, | 1076 | Ok(d) => d, |
| 1036 | Err(e) => { | 1077 | Err(e) => { |
| 1078 | + if let Some(automation) = memory_automation_for_shutdown { | ||
| 1079 | + match tokio::time::timeout( | ||
| 1080 | + std::time::Duration::from_secs(5), | ||
| 1081 | + automation.close(), | ||
| 1082 | + ) | ||
| 1083 | + .await | ||
| 1084 | + { | ||
| 1085 | + Ok(Ok(())) => {} | ||
| 1086 | + Ok(Err(error)) => { | ||
| 1087 | + eprintln!("[warn] failed to close MCP memory automation: {error}") | ||
| 1088 | + } | ||
| 1089 | + Err(_) => eprintln!("[warn] MCP memory automation close timed out after 5 seconds"), | ||
| 1090 | + } | ||
| 1091 | + } | ||
| 1037 | eprintln!("Failed to bootstrap session: {}", e); | 1092 | eprintln!("Failed to bootstrap session: {}", e); |
| 1038 | std::process::exit(1); | 1093 | std::process::exit(1); |
| 1039 | } | 1094 | } |
| @@ -1074,13 +1129,20 @@ if let Some(attach_url) = &attach { | |||
| 1074 | "agent": agent, | 1129 | "agent": agent, |
| 1075 | }); | 1130 | }); |
| 1076 | if format == OutputFormat::Json { | 1131 | if format == OutputFormat::Json { |
| 1077 | - println!("{}", serde_json::to_string(&serde_json::json!({ | 1132 | + println!( |
| 1078 | - "type": "session_start", | 1133 | + "{}", |
| 1079 | - "data": session_info | 1134 | + serde_json::to_string(&serde_json::json!({ |
| 1080 | - })).unwrap()); | 1135 | + "type": "session_start", |
| 1136 | + "data": session_info | ||
| 1137 | + })) | ||
| 1138 | + .unwrap() | ||
| 1139 | + ); | ||
| 1081 | let _ = std::io::stdout().flush(); | 1140 | let _ = std::io::stdout().flush(); |
| 1082 | } else if debug { | 1141 | } else if debug { |
| 1083 | - eprintln!("[session] {}", serde_json::to_string_pretty(&session_info).unwrap()); | 1142 | + eprintln!( |
| 1143 | + "[session] {}", | ||
| 1144 | + serde_json::to_string_pretty(&session_info).unwrap() | ||
| 1145 | + ); | ||
| 1084 | } | 1146 | } |
| 1085 | } | 1147 | } |
| 1086 | 1148 | ||
| @@ -1093,28 +1155,43 @@ if let Some(attach_url) = &attach { | |||
| 1093 | .await | 1155 | .await |
| 1094 | { | 1156 | { |
| 1095 | if format == OutputFormat::Json { | 1157 | if format == OutputFormat::Json { |
| 1096 | - println!("{}", serde_json::to_string(&serde_json::json!({ | 1158 | + println!( |
| 1097 | - "type": "error", | 1159 | + "{}", |
| 1098 | - "data": { | 1160 | + serde_json::to_string(&serde_json::json!({ |
| 1099 | - "message": format!("failed to close session: {}", err) | 1161 | + "type": "error", |
| 1100 | - } | 1162 | + "data": { |
| 1101 | - })).unwrap()); | 1163 | + "message": format!("failed to close session: {}", err) |
| 1164 | + } | ||
| 1165 | + })) | ||
| 1166 | + .unwrap() | ||
| 1167 | + ); | ||
| 1102 | let _ = std::io::stdout().flush(); | 1168 | let _ = std::io::stdout().flush(); |
| 1103 | } else { | 1169 | } else { |
| 1104 | eprintln!("[warn] failed to close session: {}", err); | 1170 | eprintln!("[warn] failed to close session: {}", err); |
| 1105 | } | 1171 | } |
| 1106 | } | 1172 | } |
| 1173 | + if let Some(automation) = memory_automation_for_shutdown { | ||
| 1174 | + match tokio::time::timeout(std::time::Duration::from_secs(5), automation.close()).await { | ||
| 1175 | + Ok(Ok(())) => {} | ||
| 1176 | + Ok(Err(error)) => eprintln!("[warn] failed to close MCP memory automation: {error}"), | ||
| 1177 | + Err(_) => eprintln!("[warn] MCP memory automation close timed out after 5 seconds"), | ||
| 1178 | + } | ||
| 1179 | + } | ||
| 1107 | 1180 | ||
| 1108 | match turn_result { | 1181 | match turn_result { |
| 1109 | Ok(result) => { | 1182 | Ok(result) => { |
| 1110 | if format == OutputFormat::Json { | 1183 | if format == OutputFormat::Json { |
| 1111 | - println!("{}", serde_json::to_string(&serde_json::json!({ | 1184 | + println!( |
| 1112 | - "type": "response", | 1185 | + "{}", |
| 1113 | - "data": { | 1186 | + serde_json::to_string(&serde_json::json!({ |
| 1114 | - "raw_reply": result.raw_reply, | 1187 | + "type": "response", |
| 1115 | - "session_id": session_id, | 1188 | + "data": { |
| 1116 | - } | 1189 | + "raw_reply": result.raw_reply, |
| 1117 | - })).unwrap()); | 1190 | + "session_id": session_id, |
| 1191 | + } | ||
| 1192 | + })) | ||
| 1193 | + .unwrap() | ||
| 1194 | + ); | ||
| 1118 | let _ = std::io::stdout().flush(); | 1195 | let _ = std::io::stdout().flush(); |
| 1119 | } else { | 1196 | } else { |
| 1120 | if !result.raw_reply.is_empty() { | 1197 | if !result.raw_reply.is_empty() { |
| @@ -1124,12 +1201,16 @@ if let Some(attach_url) = &attach { | |||
| 1124 | } | 1201 | } |
| 1125 | Err(e) => { | 1202 | Err(e) => { |
| 1126 | if format == OutputFormat::Json { | 1203 | if format == OutputFormat::Json { |
| 1127 | - println!("{}", serde_json::to_string(&serde_json::json!({ | 1204 | + println!( |
| 1128 | - "type": "error", | 1205 | + "{}", |
| 1129 | - "data": { | 1206 | + serde_json::to_string(&serde_json::json!({ |
| 1130 | - "message": e.to_string() | 1207 | + "type": "error", |
| 1131 | - } | 1208 | + "data": { |
| 1132 | - })).unwrap()); | 1209 | + "message": e.to_string() |
| 1210 | + } | ||
| 1211 | + })) | ||
| 1212 | + .unwrap() | ||
| 1213 | + ); | ||
| 1133 | let _ = std::io::stdout().flush(); | 1214 | let _ = std::io::stdout().flush(); |
| 1134 | } else { | 1215 | } else { |
| 1135 | eprintln!("[error] {}", e); | 1216 | eprintln!("[error] {}", e); |
| @@ -1152,12 +1233,9 @@ async fn handle_serve_command(port: u16, hostname: String) { | |||
| 1152 | eprintln!("Starting xiaoo daemon server on {}:{}", hostname, port); | 1233 | eprintln!("Starting xiaoo daemon server on {}:{}", hostname, port); |
| 1153 | eprintln!("Use 'xiaoo-daemon' binary directly for full daemon functionality"); | 1234 | eprintln!("Use 'xiaoo-daemon' binary directly for full daemon functionality"); |
| 1154 | let status = std::process::Command::new("xiaoo-daemon") | 1235 | let status = std::process::Command::new("xiaoo-daemon") |
| 1155 | - .args([ | 1236 | + .args(["--port", &port.to_string(), "--host", &hostname]) |
| 1156 | - "--port", &port.to_string(), | ||
| 1157 | - "--host", &hostname, | ||
| 1158 | - ]) | ||
| 1159 | .status(); | 1237 | .status(); |
| 1160 | - | 1238 | + |
| 1161 | match status { | 1239 | match status { |
| 1162 | Ok(s) if s.success() => std::process::exit(0), | 1240 | Ok(s) if s.success() => std::process::exit(0), |
| 1163 | Ok(s) => { | 1241 | Ok(s) => { |
| @@ -1244,7 +1322,10 @@ async fn run_with_attach( | |||
| 1244 | Ok(resp) if !resp.status().is_success() => { | 1322 | Ok(resp) if !resp.status().is_success() => { |
| 1245 | let status = resp.status(); | 1323 | let status = resp.status(); |
| 1246 | let body = resp.text().await.unwrap_or_default(); | 1324 | let body = resp.text().await.unwrap_or_default(); |
| 1247 | - attach_fail(format!("session open failed: HTTP {status} {body}"), is_json); | 1325 | + attach_fail( |
| 1326 | + format!("session open failed: HTTP {status} {body}"), | ||
| 1327 | + is_json, | ||
| 1328 | + ); | ||
| 1248 | } | 1329 | } |
| 1249 | Ok(_) => {} | 1330 | Ok(_) => {} |
| 1250 | Err(error) => attach_fail(format!("failed to connect to daemon: {error}"), is_json), | 1331 | Err(error) => attach_fail(format!("failed to connect to daemon: {error}"), is_json), |
| @@ -1284,7 +1365,10 @@ async fn run_with_attach( | |||
| 1284 | if !response.status().is_success() { | 1365 | if !response.status().is_success() { |
| 1285 | let status = response.status(); | 1366 | let status = response.status(); |
| 1286 | let body = response.text().await.unwrap_or_default(); | 1367 | let body = response.text().await.unwrap_or_default(); |
| 1287 | - attach_fail(format!("turn submission failed: HTTP {status} {body}"), is_json); | 1368 | + attach_fail( |
| 1369 | + format!("turn submission failed: HTTP {status} {body}"), | ||
| 1370 | + is_json, | ||
| 1371 | + ); | ||
| 1288 | } | 1372 | } |
| 1289 | 1373 | ||
| 1290 | // 3. Consume the SSE event stream emitted by /api/v1/runtimes/input. | 1374 | // 3. Consume the SSE event stream emitted by /api/v1/runtimes/input. |
| @@ -1345,7 +1429,10 @@ async fn run_with_attach( | |||
| 1345 | } | 1429 | } |
| 1346 | 1430 | ||
| 1347 | if !saw_done { | 1431 | if !saw_done { |
| 1348 | - attach_fail("daemon stream ended without a completion event".to_string(), is_json); | 1432 | + attach_fail( |
| 1433 | + "daemon stream ended without a completion event".to_string(), | ||
| 1434 | + is_json, | ||
| 1435 | + ); | ||
| 1349 | } | 1436 | } |
| 1350 | 1437 | ||
| 1351 | // 4. Best-effort detach so the daemon releases this process's lease | 1438 | // 4. Best-effort detach so the daemon releases this process's lease |
| @@ -1409,30 +1496,40 @@ fn parse_sse_event(frame: &str) -> Option<Value> { | |||
| 1409 | serde_json::from_str(&data).ok() | 1496 | serde_json::from_str(&data).ok() |
| 1410 | } | 1497 | } |
| 1411 | 1498 | ||
| 1412 | - | ||
| 1413 | fn handle_debug_command(command: DebugCommands, config_path: Option<&PathBuf>, debug: bool) { | 1499 | fn handle_debug_command(command: DebugCommands, config_path: Option<&PathBuf>, debug: bool) { |
| 1414 | match command { | 1500 | match command { |
| 1415 | DebugCommands::Config => { | 1501 | DebugCommands::Config => { |
| 1416 | let file_cfg = config_path | 1502 | let file_cfg = config_path |
| 1417 | .map(|path| FileConfig::load_from_path(path, debug)) | 1503 | .map(|path| FileConfig::load_from_path(path, debug)) |
| 1418 | .unwrap_or_default(); | 1504 | .unwrap_or_default(); |
| 1419 | - | 1505 | + |
| 1420 | let mut config_json = serde_json::Map::new(); | 1506 | let mut config_json = serde_json::Map::new(); |
| 1421 | - config_json.insert("$schema".to_string(), Value::String("https://xiaoo.ai/config.json".to_string())); | 1507 | + config_json.insert( |
| 1508 | + "$schema".to_string(), | ||
| 1509 | + Value::String("https://xiaoo.ai/config.json".to_string()), | ||
| 1510 | + ); | ||
| 1422 | 1511 | ||
| 1423 | if let Some(llm) = &file_cfg.llm { | 1512 | if let Some(llm) = &file_cfg.llm { |
| 1424 | let provider = llm.provider.as_deref().unwrap_or("openai"); | 1513 | let provider = llm.provider.as_deref().unwrap_or("openai"); |
| 1425 | let model = llm.model.as_deref().unwrap_or(""); | 1514 | let model = llm.model.as_deref().unwrap_or(""); |
| 1426 | - config_json.insert("model".to_string(), Value::String(format!("{}/{}", provider, model))); | 1515 | + config_json.insert( |
| 1516 | + "model".to_string(), | ||
| 1517 | + Value::String(format!("{}/{}", provider, model)), | ||
| 1518 | + ); | ||
| 1427 | } | 1519 | } |
| 1428 | 1520 | ||
| 1429 | - println!("{}", serde_json::to_string_pretty(&Value::Object(config_json)).unwrap()); | 1521 | + println!( |
| 1522 | + "{}", | ||
| 1523 | + serde_json::to_string_pretty(&Value::Object(config_json)).unwrap() | ||
| 1524 | + ); | ||
| 1430 | } | 1525 | } |
| 1431 | } | 1526 | } |
| 1432 | - | ||
| 1433 | } | 1527 | } |
| 1434 | async fn handle_export_command(session_id: String, port: u16, client_id: Option<String>) { | 1528 | async fn handle_export_command(session_id: String, port: u16, client_id: Option<String>) { |
| 1435 | - let url = format!("http://127.0.0.1:{}/api/v1/runtimes/export/{}", port, session_id); | 1529 | + let url = format!( |
| 1530 | + "http://127.0.0.1:{}/api/v1/runtimes/export/{}", | ||
| 1531 | + port, session_id | ||
| 1532 | + ); | ||
| 1436 | 1533 | ||
| 1437 | let client = reqwest::Client::new(); | 1534 | let client = reqwest::Client::new(); |
| 1438 | let mut req = client.get(&url); | 1535 | let mut req = client.get(&url); |
| @@ -1483,6 +1580,24 @@ mod tests { | |||
| 1483 | use std::fs; | 1580 | use std::fs; |
| 1484 | use tempfile::tempdir; | 1581 | use tempfile::tempdir; |
| 1485 | 1582 | ||
| 1583 | + | ||
| 1584 | + fn parses_explicit_mcp_config_path() { | ||
| 1585 | + let args = Args::try_parse_from([ | ||
| 1586 | + "xiaoo", | ||
| 1587 | + "--mcp-config", | ||
| 1588 | + "/tmp/mcp.json", | ||
| 1589 | + "run", | ||
| 1590 | + "--prompt", | ||
| 1591 | + "hello", | ||
| 1592 | + ]) | ||
| 1593 | + .expect("CLI should accept --mcp-config"); | ||
| 1594 | + | ||
| 1595 | + assert_eq!( | ||
| 1596 | + args.mcp_config.as_deref(), | ||
| 1597 | + Some(std::path::Path::new("/tmp/mcp.json")) | ||
| 1598 | + ); | ||
| 1599 | + } | ||
| 1600 | + | ||
| 1486 | 1601 | ||
| 1487 | fn copy_dir_rejects_destination_inside_source() { | 1602 | fn copy_dir_rejects_destination_inside_source() { |
| 1488 | let temp = tempdir().unwrap(); | 1603 | let temp = tempdir().unwrap(); |
| @@ -1527,7 +1642,6 @@ mod tests { | |||
| 1527 | } | 1642 | } |
| 1528 | } | 1643 | } |
| 1529 | 1644 | ||
| 1530 | - | ||
| 1531 | 1645 | ||
| 1532 | mod attach_sse_tests { | 1646 | mod attach_sse_tests { |
| 1533 | use super::{parse_sse_event, take_sse_frame}; | 1647 | use super::{parse_sse_event, take_sse_frame}; |
| @@ -133,6 +133,7 @@ pub struct CliConfig { | |||
| 133 | pub skills_config: skill::SkillsConfig, | 133 | pub skills_config: skill::SkillsConfig, |
| 134 | pub subagent: std::collections::BTreeMap<String, config::SubagentRoleConfig>, | 134 | pub subagent: std::collections::BTreeMap<String, config::SubagentRoleConfig>, |
| 135 | pub mcp_servers: Vec<mcp::McpServerConfig>, | 135 | pub mcp_servers: Vec<mcp::McpServerConfig>, |
| 136 | + pub memory_automation: xiaoo_shared::gateway::MemoryAutomationConfig, | ||
| 136 | } | 137 | } |
| 137 | 138 | ||
| 138 | // --------------------------------------------------------------------------- | 139 | // --------------------------------------------------------------------------- |
| @@ -308,6 +309,7 @@ mod tests { | |||
| 308 | operation_backend: None, | 309 | operation_backend: None, |
| 309 | subagent: Default::default(), | 310 | subagent: Default::default(), |
| 310 | mcp_servers: Vec::new(), | 311 | mcp_servers: Vec::new(), |
| 312 | + memory_automation: Default::default(), | ||
| 311 | } | 313 | } |
| 312 | } | 314 | } |
| 313 | 315 | ||
| @@ -205,6 +205,9 @@ impl GatewayRuntime { | |||
| 205 | .status_panel | 205 | .status_panel |
| 206 | .set_backend(format!("Remote: {base_url}")); | 206 | .set_backend(format!("Remote: {base_url}")); |
| 207 | state.status_panel.set_remote_workspace(&base_url); | 207 | state.status_panel.set_remote_workspace(&base_url); |
| 208 | + // The remote SSE protocol does not yet publish RAM-A health. Do not | ||
| 209 | + // imply that memory is disabled merely because this TUI cannot see it. | ||
| 210 | + state.status_panel.memory_status = crate::status_panel::MemoryStatus::Unknown; | ||
| 208 | } | 211 | } |
| 209 | 212 | ||
| 210 | pub async fn connect_remote( | 213 | pub async fn connect_remote( |
| @@ -252,6 +255,7 @@ impl GatewayRuntime { | |||
| 252 | .status_panel | 255 | .status_panel |
| 253 | .set_backend(sandbox_display_name(&state.agent_config.operation_backend)); | 256 | .set_backend(sandbox_display_name(&state.agent_config.operation_backend)); |
| 254 | state.status_panel.set_workspace(&state.workspace); | 257 | state.status_panel.set_workspace(&state.workspace); |
| 258 | + state.status_panel.memory_status = self.session_gateway.current_memory_status(); | ||
| 255 | Ok(()) | 259 | Ok(()) |
| 256 | } | 260 | } |
| 257 | 261 | ||
| @@ -1282,7 +1286,46 @@ fn default_interaction_response(request: &InteractionRequest) -> InteractionResp | |||
| 1282 | 1286 | ||
| 1283 | 1287 | ||
| 1284 | mod tests { | 1288 | mod tests { |
| 1285 | - use super::{parse_sse_frame, take_sse_frame, RemoteSseEvent}; | 1289 | + use std::path::PathBuf; |
| 1290 | + use tokio::sync::watch; | ||
| 1291 | + | ||
| 1292 | + use crate::app_state::AppState; | ||
| 1293 | + use crate::gateway::MemoryAutomationHealth; | ||
| 1294 | + use crate::status_panel::MemoryStatus; | ||
| 1295 | + | ||
| 1296 | + use super::{parse_sse_frame, take_sse_frame, GatewayRuntime, RemoteSseEvent}; | ||
| 1297 | + | ||
| 1298 | + | ||
| 1299 | + fn configuring_remote_marks_memory_state_unknown() { | ||
| 1300 | + let mut state = AppState::new(PathBuf::from("config.toml"), PathBuf::from(".")) | ||
| 1301 | + .expect("test app state should initialize"); | ||
| 1302 | + let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); | ||
| 1303 | + | ||
| 1304 | + runtime.configure_remote(&mut state, "http://daemon.example".to_string(), None); | ||
| 1305 | + | ||
| 1306 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Unknown); | ||
| 1307 | + } | ||
| 1308 | + | ||
| 1309 | + | ||
| 1310 | + async fn disconnecting_remote_restores_cached_local_memory_health() { | ||
| 1311 | + let mut state = AppState::new(PathBuf::from("config.toml"), PathBuf::from(".")) | ||
| 1312 | + .expect("test app state should initialize"); | ||
| 1313 | + let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); | ||
| 1314 | + let (_health_tx, health_rx) = watch::channel(MemoryAutomationHealth::Healthy); | ||
| 1315 | + *runtime | ||
| 1316 | + .session_gateway | ||
| 1317 | + .memory_health | ||
| 1318 | + .lock() | ||
| 1319 | + .expect("memory health lock should not be poisoned") = Some(health_rx); | ||
| 1320 | + runtime.configure_remote(&mut state, "http://daemon.example".to_string(), None); | ||
| 1321 | + | ||
| 1322 | + runtime | ||
| 1323 | + .disconnect_remote(&mut state) | ||
| 1324 | + .await | ||
| 1325 | + .expect("remote disconnect should succeed"); | ||
| 1326 | + | ||
| 1327 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Connected); | ||
| 1328 | + } | ||
| 1286 | 1329 | ||
| 1287 | 1330 | ||
| 1288 | fn parses_sse_frame_from_split_buffer() { | 1331 | fn parses_sse_frame_from_split_buffer() { |
| @@ -321,7 +321,8 @@ impl GatewayRuntime { | |||
| 321 | ) | 321 | ) |
| 322 | }) | 322 | }) |
| 323 | .collect(), | 323 | .collect(), |
| 324 | - mcp_servers: state.agent_config.mcp.servers.clone(), | 324 | + mcp_servers: state.agent_config.mcp_servers().to_vec(), |
| 325 | + memory_automation: state.agent_config.memory_automation.clone(), | ||
| 325 | }) | 326 | }) |
| 326 | } | 327 | } |
| 327 | 328 | ||
| @@ -10,6 +10,19 @@ use super::runtime::{GatewayRuntime, PendingStreamDone, STREAM_REVEAL_CHARS_PER_ | |||
| 10 | impl GatewayRuntime { | 10 | impl GatewayRuntime { |
| 11 | pub fn poll_stream_updates(&mut self, state: &mut AppState) -> bool { | 11 | pub fn poll_stream_updates(&mut self, state: &mut AppState) -> bool { |
| 12 | let mut changed = false; | 12 | let mut changed = false; |
| 13 | + if self.remote.is_none() { | ||
| 14 | + if let Some(health) = self.session_gateway.take_memory_health_update() { | ||
| 15 | + state.status_panel.memory_status = match health { | ||
| 16 | + crate::gateway::MemoryAutomationHealth::Healthy => { | ||
| 17 | + crate::status_panel::MemoryStatus::Connected | ||
| 18 | + } | ||
| 19 | + crate::gateway::MemoryAutomationHealth::Degraded => { | ||
| 20 | + crate::status_panel::MemoryStatus::Degraded | ||
| 21 | + } | ||
| 22 | + }; | ||
| 23 | + changed = true; | ||
| 24 | + } | ||
| 25 | + } | ||
| 13 | while let Some(receiver) = &mut self.stream_rx { | 26 | while let Some(receiver) = &mut self.stream_rx { |
| 14 | let update = match receiver.try_recv() { | 27 | let update = match receiver.try_recv() { |
| 15 | Ok(update) => update, | 28 | Ok(update) => update, |
| @@ -140,6 +153,9 @@ impl GatewayRuntime { | |||
| 140 | } | 153 | } |
| 141 | state.chat_state.stick_to_bottom = true; | 154 | state.chat_state.stick_to_bottom = true; |
| 142 | } | 155 | } |
| 156 | + SessionTurnUpdate::MemoryStatus(memory_status) => { | ||
| 157 | + state.status_panel.memory_status = memory_status; | ||
| 158 | + } | ||
| 143 | SessionTurnUpdate::Done { | 159 | SessionTurnUpdate::Done { |
| 144 | prompt_tokens, | 160 | prompt_tokens, |
| 145 | completion_tokens, | 161 | completion_tokens, |
| @@ -835,13 +851,15 @@ mod tests { | |||
| 835 | use std::time::{Duration, Instant}; | 851 | use std::time::{Duration, Instant}; |
| 836 | 852 | ||
| 837 | use agent_types::common::ids::AgentId; | 853 | use agent_types::common::ids::AgentId; |
| 838 | - use tokio::sync::mpsc; | 854 | + use tokio::sync::{mpsc, watch}; |
| 839 | 855 | ||
| 840 | use crate::app_state::AppState; | 856 | use crate::app_state::AppState; |
| 841 | use crate::chat::{ | 857 | use crate::chat::{ |
| 842 | Message, MessageRole, TodoDisplayStatus, ToolExecutionStatus, ToolExecutionUpdate, | 858 | Message, MessageRole, TodoDisplayStatus, ToolExecutionStatus, ToolExecutionUpdate, |
| 843 | }; | 859 | }; |
| 860 | + use crate::gateway::MemoryAutomationHealth; | ||
| 844 | use crate::session_gateway::SessionTurnUpdate; | 861 | use crate::session_gateway::SessionTurnUpdate; |
| 862 | + use crate::status_panel::MemoryStatus; | ||
| 845 | 863 | ||
| 846 | use super::{GatewayRuntime, PendingStreamDone}; | 864 | use super::{GatewayRuntime, PendingStreamDone}; |
| 847 | 865 | ||
| @@ -999,6 +1017,49 @@ mod tests { | |||
| 999 | assert_eq!(state.chat_state.messages[0].content, "answer"); | 1017 | assert_eq!(state.chat_state.messages[0].content, "answer"); |
| 1000 | } | 1018 | } |
| 1001 | 1019 | ||
| 1020 | + | ||
| 1021 | + fn stream_updates_surface_memory_state_transitions() { | ||
| 1022 | + let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); | ||
| 1023 | + let mut state = test_state(); | ||
| 1024 | + | ||
| 1025 | + let (tx, rx) = mpsc::unbounded_channel(); | ||
| 1026 | + runtime.stream_rx = Some(rx); | ||
| 1027 | + tx.send(SessionTurnUpdate::MemoryStatus(MemoryStatus::Disabled)) | ||
| 1028 | + .expect("memory status update should send"); | ||
| 1029 | + assert!(runtime.poll_stream_updates(&mut state)); | ||
| 1030 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Disabled); | ||
| 1031 | + | ||
| 1032 | + tx.send(SessionTurnUpdate::MemoryStatus(MemoryStatus::Degraded)) | ||
| 1033 | + .expect("memory status update should send"); | ||
| 1034 | + assert!(runtime.poll_stream_updates(&mut state)); | ||
| 1035 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Degraded); | ||
| 1036 | + | ||
| 1037 | + tx.send(SessionTurnUpdate::MemoryStatus(MemoryStatus::Connected)) | ||
| 1038 | + .expect("memory status update should send"); | ||
| 1039 | + | ||
| 1040 | + assert!(runtime.poll_stream_updates(&mut state)); | ||
| 1041 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Connected); | ||
| 1042 | + } | ||
| 1043 | + | ||
| 1044 | + | ||
| 1045 | + fn background_memory_health_change_updates_status_without_a_turn() { | ||
| 1046 | + let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); | ||
| 1047 | + let mut state = test_state(); | ||
| 1048 | + let (health_tx, health_rx) = watch::channel(MemoryAutomationHealth::Healthy); | ||
| 1049 | + *runtime | ||
| 1050 | + .session_gateway | ||
| 1051 | + .memory_health | ||
| 1052 | + .lock() | ||
| 1053 | + .expect("memory health lock should not be poisoned") = Some(health_rx); | ||
| 1054 | + | ||
| 1055 | + health_tx | ||
| 1056 | + .send(MemoryAutomationHealth::Degraded) | ||
| 1057 | + .expect("memory health receiver should be present"); | ||
| 1058 | + | ||
| 1059 | + assert!(runtime.poll_stream_updates(&mut state)); | ||
| 1060 | + assert_eq!(state.status_panel.memory_status, MemoryStatus::Degraded); | ||
| 1061 | + } | ||
| 1062 | + | ||
| 1002 | 1063 | ||
| 1003 | fn child_stream_updates_create_subagent_lane_without_touching_root_messages() { | 1064 | fn child_stream_updates_create_subagent_lane_without_touching_root_messages() { |
| 1004 | let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); | 1065 | let mut runtime = GatewayRuntime::new(uuid::Uuid::new_v4().to_string()); |
| @@ -2,12 +2,18 @@ use std::collections::{HashSet, VecDeque}; | |||
| 2 | use std::sync::{Arc, Mutex}; | 2 | use std::sync::{Arc, Mutex}; |
| 3 | 3 | ||
| 4 | use async_trait::async_trait; | 4 | use async_trait::async_trait; |
| 5 | -use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; | 5 | +use tokio::sync::{ |
| 6 | + mpsc::{UnboundedReceiver, UnboundedSender}, | ||
| 7 | + watch, | ||
| 8 | +}; | ||
| 6 | 9 | ||
| 7 | use crate::backend::BackendManager; | 10 | use crate::backend::BackendManager; |
| 8 | use crate::chat::{FileChangeDelta, ToolExecutionUpdate}; | 11 | use crate::chat::{FileChangeDelta, ToolExecutionUpdate}; |
| 9 | -use crate::gateway::{InMemorySessionStore, SessionControlPlane, SessionStore}; | 12 | +use crate::gateway::{ |
| 13 | + InMemorySessionStore, SessionControlPlane, SessionStore, TurnMemoryAutomation, | ||
| 14 | +}; | ||
| 10 | use crate::interaction_prompt::PromptRequest; | 15 | use crate::interaction_prompt::PromptRequest; |
| 16 | +use crate::status_panel::MemoryStatus; | ||
| 11 | 17 | ||
| 12 | use agent_types::common::ids::AgentId; | 18 | use agent_types::common::ids::AgentId; |
| 13 | use agent_types::events::LoopEndSummary; | 19 | use agent_types::events::LoopEndSummary; |
| @@ -60,6 +66,7 @@ pub enum SessionTurnUpdate { | |||
| 60 | PendingUserMessagesConsumed { | 66 | PendingUserMessagesConsumed { |
| 61 | prompts: Vec<String>, | 67 | prompts: Vec<String>, |
| 62 | }, | 68 | }, |
| 69 | + MemoryStatus(MemoryStatus), | ||
| 63 | Done { | 70 | Done { |
| 64 | prompt_tokens: u64, | 71 | prompt_tokens: u64, |
| 65 | completion_tokens: u64, | 72 | completion_tokens: u64, |
| @@ -80,6 +87,15 @@ pub struct SessionGateway { | |||
| 80 | Arc<tokio::sync::Mutex<Option<Arc<dyn SessionControlPlane>>>>, | 87 | Arc<tokio::sync::Mutex<Option<Arc<dyn SessionControlPlane>>>>, |
| 81 | /// Session IDs that have been opened and not yet closed. | 88 | /// Session IDs that have been opened and not yet closed. |
| 82 | pub(super) active_session_ids: Arc<tokio::sync::Mutex<HashSet<String>>>, | 89 | pub(super) active_session_ids: Arc<tokio::sync::Mutex<HashSet<String>>>, |
| 90 | + /// One MCP memory client is shared by all local TUI turns. The nested | ||
| 91 | + /// option distinguishes not-yet-initialized from a disabled/failed setup. | ||
| 92 | + pub(super) memory_automation: | ||
| 93 | + Arc<tokio::sync::Mutex<Option<Option<Arc<dyn TurnMemoryAutomation>>>>>, | ||
| 94 | + /// Latest RAM-A health receiver. Unlike a turn's stream receiver, this | ||
| 95 | + /// remains available after a turn finishes so background ingest failures | ||
| 96 | + /// can update the TUI immediately. | ||
| 97 | + pub(super) memory_health: | ||
| 98 | + Arc<Mutex<Option<watch::Receiver<crate::gateway::MemoryAutomationHealth>>>>, | ||
| 83 | pub(super) backend_manager: Arc<BackendManager>, | 99 | pub(super) backend_manager: Arc<BackendManager>, |
| 84 | } | 100 | } |
| 85 | 101 | ||
| @@ -105,6 +121,8 @@ impl Default for SessionGateway { | |||
| 105 | session_store, | 121 | session_store, |
| 106 | lifecycle_control_plane: Arc::new(tokio::sync::Mutex::new(None)), | 122 | lifecycle_control_plane: Arc::new(tokio::sync::Mutex::new(None)), |
| 107 | active_session_ids: Arc::new(tokio::sync::Mutex::new(HashSet::new())), | 123 | active_session_ids: Arc::new(tokio::sync::Mutex::new(HashSet::new())), |
| 124 | + memory_automation: Arc::new(tokio::sync::Mutex::new(None)), | ||
| 125 | + memory_health: Arc::new(Mutex::new(None)), | ||
| 108 | backend_manager, | 126 | backend_manager, |
| 109 | } | 127 | } |
| 110 | } | 128 | } |
| @@ -142,6 +160,59 @@ impl ChannelPendingUserMessages { | |||
| 142 | } | 160 | } |
| 143 | } | 161 | } |
| 144 | 162 | ||
| 163 | + | ||
| 164 | +mod tests { | ||
| 165 | + use super::SessionGateway; | ||
| 166 | + use async_trait::async_trait; | ||
| 167 | + use std::sync::atomic::{AtomicBool, Ordering}; | ||
| 168 | + use std::sync::Arc; | ||
| 169 | + use xiaoo_shared::gateway::memory_automation::{ | ||
| 170 | + CompletedTurnIngest, MemoryAutomationError, RecallMemory, TurnMemoryContext, | ||
| 171 | + }; | ||
| 172 | + use xiaoo_shared::gateway::TurnMemoryAutomation; | ||
| 173 | + | ||
| 174 | + struct ClosingAutomation(AtomicBool); | ||
| 175 | + | ||
| 176 | + | ||
| 177 | + impl TurnMemoryAutomation for ClosingAutomation { | ||
| 178 | + async fn recall( | ||
| 179 | + &self, | ||
| 180 | + _context: &TurnMemoryContext, | ||
| 181 | + ) -> Result<Vec<RecallMemory>, MemoryAutomationError> { | ||
| 182 | + Ok(Vec::new()) | ||
| 183 | + } | ||
| 184 | + | ||
| 185 | + async fn enqueue_ingest( | ||
| 186 | + &self, | ||
| 187 | + _ingest: CompletedTurnIngest, | ||
| 188 | + ) -> Result<(), MemoryAutomationError> { | ||
| 189 | + Ok(()) | ||
| 190 | + } | ||
| 191 | + | ||
| 192 | + fn recall_token_budget(&self) -> usize { | ||
| 193 | + 0 | ||
| 194 | + } | ||
| 195 | + | ||
| 196 | + async fn close(&self) -> Result<(), MemoryAutomationError> { | ||
| 197 | + self.0.store(true, Ordering::SeqCst); | ||
| 198 | + Ok(()) | ||
| 199 | + } | ||
| 200 | + } | ||
| 201 | + | ||
| 202 | + | ||
| 203 | + async fn close_all_sessions_closes_cached_memory_automation() { | ||
| 204 | + let gateway = SessionGateway::new(); | ||
| 205 | + let automation = Arc::new(ClosingAutomation(AtomicBool::new(false))); | ||
| 206 | + *gateway.memory_automation.lock().await = | ||
| 207 | + Some(Some(automation.clone() as Arc<dyn TurnMemoryAutomation>)); | ||
| 208 | + | ||
| 209 | + gateway.close_all_sessions().await; | ||
| 210 | + | ||
| 211 | + assert!(automation.0.load(Ordering::SeqCst)); | ||
| 212 | + assert!(gateway.memory_automation.lock().await.is_none()); | ||
| 213 | + } | ||
| 214 | +} | ||
| 215 | + | ||
| 145 | 216 | ||
| 146 | impl xiaoo_core::PendingUserMessageSource for ChannelPendingUserMessages { | 217 | impl xiaoo_core::PendingUserMessageSource for ChannelPendingUserMessages { |
| 147 | async fn drain_pending_user_messages(&self) -> Vec<String> { | 218 | async fn drain_pending_user_messages(&self) -> Vec<String> { |
| @@ -3,8 +3,8 @@ use std::sync::{Arc, Mutex}; | |||
| 3 | 3 | ||
| 4 | use crate::gateway::{ | 4 | use crate::gateway::{ |
| 5 | AppBootstrap, AppDependencies, AppTurnRequest, AppTurnResult, HostedSessionRuntimeConfig, | 5 | AppBootstrap, AppDependencies, AppTurnRequest, AppTurnResult, HostedSessionRuntimeConfig, |
| 6 | - HostedSessionRuntimeResolver, SessionControlPlane, SessionOpenRequest, SessionRecord, | 6 | + HostedSessionRuntimeResolver, McpMemoryAutomation, MemoryAutomationHealth, SessionControlPlane, |
| 7 | - SessionRuntimeBindings, SessionStore, | 7 | + SessionOpenRequest, SessionRecord, SessionRuntimeBindings, SessionStore, |
| 8 | }; | 8 | }; |
| 9 | use crate::interaction_prompt::UserPromptResult; | 9 | use crate::interaction_prompt::UserPromptResult; |
| 10 | 10 | ||
| @@ -14,6 +14,13 @@ use super::session::{ | |||
| 14 | }; | 14 | }; |
| 15 | use xiaoo_core::spawn_prefetch; | 15 | use xiaoo_core::spawn_prefetch; |
| 16 | 16 | ||
| 17 | +fn memory_status_from_health(health: MemoryAutomationHealth) -> crate::status_panel::MemoryStatus { | ||
| 18 | + match health { | ||
| 19 | + MemoryAutomationHealth::Healthy => crate::status_panel::MemoryStatus::Connected, | ||
| 20 | + MemoryAutomationHealth::Degraded => crate::status_panel::MemoryStatus::Degraded, | ||
| 21 | + } | ||
| 22 | +} | ||
| 23 | + | ||
| 17 | impl SessionGateway { | 24 | impl SessionGateway { |
| 18 | pub fn new() -> Self { | 25 | pub fn new() -> Self { |
| 19 | Self::default() | 26 | Self::default() |
| @@ -67,6 +74,77 @@ impl SessionGateway { | |||
| 67 | self.session_store.load(session_id).await | 74 | self.session_store.load(session_id).await |
| 68 | } | 75 | } |
| 69 | 76 | ||
| 77 | + async fn get_or_init_memory_automation( | ||
| 78 | + memory_automation: &tokio::sync::Mutex< | ||
| 79 | + Option<Option<Arc<dyn crate::gateway::TurnMemoryAutomation>>>, | ||
| 80 | + >, | ||
| 81 | + config: &crate::gateway::HostedSessionRuntimeConfig, | ||
| 82 | + ) -> ( | ||
| 83 | + Option<Arc<dyn crate::gateway::TurnMemoryAutomation>>, | ||
| 84 | + Option<crate::status_panel::MemoryStatus>, | ||
| 85 | + ) { | ||
| 86 | + let mut state = memory_automation.lock().await; | ||
| 87 | + if let Some(automation) = state.as_ref() { | ||
| 88 | + return (automation.clone(), None); | ||
| 89 | + } | ||
| 90 | + if !config.memory_automation.enabled { | ||
| 91 | + *state = Some(None); | ||
| 92 | + return (None, Some(crate::status_panel::MemoryStatus::Disabled)); | ||
| 93 | + } | ||
| 94 | + match McpMemoryAutomation::connect(config.memory_automation.clone(), &config.mcp_servers) | ||
| 95 | + .await | ||
| 96 | + { | ||
| 97 | + Ok(automation) => { | ||
| 98 | + *state = Some(automation.clone()); | ||
| 99 | + ( | ||
| 100 | + automation, | ||
| 101 | + Some(crate::status_panel::MemoryStatus::Connected), | ||
| 102 | + ) | ||
| 103 | + } | ||
| 104 | + Err(error) => { | ||
| 105 | + tracing::warn!(error = %error, "memory automation disabled after TUI startup error"); | ||
| 106 | + // Keep the state uninitialized: an optional MCP server can | ||
| 107 | + // recover while the TUI is still open, so the next turn gets | ||
| 108 | + // a fresh connection attempt instead of permanent fail-open. | ||
| 109 | + (None, Some(crate::status_panel::MemoryStatus::Degraded)) | ||
| 110 | + } | ||
| 111 | + } | ||
| 112 | + } | ||
| 113 | + | ||
| 114 | + fn register_memory_health_receiver( | ||
| 115 | + memory_health: &Mutex<Option<tokio::sync::watch::Receiver<MemoryAutomationHealth>>>, | ||
| 116 | + automation: Option<&Arc<dyn crate::gateway::TurnMemoryAutomation>>, | ||
| 117 | + ) { | ||
| 118 | + let Some(automation) = automation else { | ||
| 119 | + return; | ||
| 120 | + }; | ||
| 121 | + let Ok(mut receiver) = memory_health.lock() else { | ||
| 122 | + tracing::warn!("memory health receiver lock poisoned"); | ||
| 123 | + return; | ||
| 124 | + }; | ||
| 125 | + if receiver.is_none() { | ||
| 126 | + *receiver = automation.subscribe_health(); | ||
| 127 | + } | ||
| 128 | + } | ||
| 129 | + | ||
| 130 | + pub(super) fn take_memory_health_update(&self) -> Option<MemoryAutomationHealth> { | ||
| 131 | + let mut receiver = self.memory_health.lock().ok()?; | ||
| 132 | + let receiver = receiver.as_mut()?; | ||
| 133 | + match receiver.has_changed() { | ||
| 134 | + Ok(true) => Some(*receiver.borrow_and_update()), | ||
| 135 | + Ok(false) | Err(_) => None, | ||
| 136 | + } | ||
| 137 | + } | ||
| 138 | + | ||
| 139 | + pub(super) fn current_memory_status(&self) -> crate::status_panel::MemoryStatus { | ||
| 140 | + self.memory_health | ||
| 141 | + .lock() | ||
| 142 | + .ok() | ||
| 143 | + .and_then(|receiver| receiver.as_ref().map(|receiver| *receiver.borrow())) | ||
| 144 | + .map(memory_status_from_health) | ||
| 145 | + .unwrap_or(crate::status_panel::MemoryStatus::Disabled) | ||
| 146 | + } | ||
| 147 | + | ||
| 70 | pub async fn import_session_snapshot(&self, record: SessionRecord) { | 148 | pub async fn import_session_snapshot(&self, record: SessionRecord) { |
| 71 | let session_id = record.session_id.clone(); | 149 | let session_id = record.session_id.clone(); |
| 72 | let kvcache_enabled = record.runtime.feature_flags.kvcache_enabled; | 150 | let kvcache_enabled = record.runtime.feature_flags.kvcache_enabled; |
| @@ -99,6 +177,8 @@ impl SessionGateway { | |||
| 99 | ) { | 177 | ) { |
| 100 | let session_store: Arc<dyn SessionStore> = self.session_store.clone(); | 178 | let session_store: Arc<dyn SessionStore> = self.session_store.clone(); |
| 101 | let active_session_ids = Arc::clone(&self.active_session_ids); | 179 | let active_session_ids = Arc::clone(&self.active_session_ids); |
| 180 | + let memory_automation_state = Arc::clone(&self.memory_automation); | ||
| 181 | + let memory_health_state = Arc::clone(&self.memory_health); | ||
| 102 | let backend_manager = self.backend_manager.clone(); | 182 | let backend_manager = self.backend_manager.clone(); |
| 103 | tokio::spawn(async move { | 183 | tokio::spawn(async move { |
| 104 | active_session_ids | 184 | active_session_ids |
| @@ -126,13 +206,22 @@ impl SessionGateway { | |||
| 126 | }; | 206 | }; |
| 127 | 207 | ||
| 128 | let hooker_config = runtime_config.hooker.clone(); | 208 | let hooker_config = runtime_config.hooker.clone(); |
| 209 | + let (memory_automation, memory_status) = | ||
| 210 | + Self::get_or_init_memory_automation(&memory_automation_state, &runtime_config) | ||
| 211 | + .await; | ||
| 212 | + Self::register_memory_health_receiver(&memory_health_state, memory_automation.as_ref()); | ||
| 213 | + if let Some(memory_status) = memory_status { | ||
| 214 | + let _ = updates_tx.send(SessionTurnUpdate::MemoryStatus(memory_status)); | ||
| 215 | + } | ||
| 129 | let resolver = Arc::new(HostedSessionRuntimeResolver::new(runtime_config, bindings)); | 216 | let resolver = Arc::new(HostedSessionRuntimeResolver::new(runtime_config, bindings)); |
| 217 | + let memory_automation_for_status = memory_automation.clone(); | ||
| 130 | let dependencies = | 218 | let dependencies = |
| 131 | - match AppBootstrap::from_session_components_with_hooks_and_backend_manager( | 219 | + match AppBootstrap::from_session_components_with_hooks_and_backend_manager_and_memory_automation( |
| 132 | session_store, | 220 | session_store, |
| 133 | resolver, | 221 | resolver, |
| 134 | hooker_config, | 222 | hooker_config, |
| 135 | backend_manager, | 223 | backend_manager, |
| 224 | + memory_automation, | ||
| 136 | ) { | 225 | ) { |
| 137 | Ok(dependencies) => dependencies, | 226 | Ok(dependencies) => dependencies, |
| 138 | Err(error) => { | 227 | Err(error) => { |
| @@ -142,6 +231,11 @@ impl SessionGateway { | |||
| 142 | }; | 231 | }; |
| 143 | 232 | ||
| 144 | let result = dependencies.session_service.run_turn(request).await; | 233 | let result = dependencies.session_service.run_turn(request).await; |
| 234 | + if let Some(automation) = memory_automation_for_status.as_ref() { | ||
| 235 | + let _ = updates_tx.send(SessionTurnUpdate::MemoryStatus( | ||
| 236 | + memory_status_from_health(automation.health()), | ||
| 237 | + )); | ||
| 238 | + } | ||
| 145 | match result { | 239 | match result { |
| 146 | Ok(AppTurnResult { | 240 | Ok(AppTurnResult { |
| 147 | messages, | 241 | messages, |
| @@ -222,21 +316,52 @@ impl SessionGateway { | |||
| 222 | lock.clear(); | 316 | lock.clear(); |
| 223 | ids | 317 | ids |
| 224 | }; | 318 | }; |
| 225 | - if ids.is_empty() { | ||
| 226 | - return; | ||
| 227 | - } | ||
| 228 | let cp = self.lifecycle_control_plane.lock().await.clone(); | 319 | let cp = self.lifecycle_control_plane.lock().await.clone(); |
| 229 | - let Some(control_plane) = cp else { | 320 | + if let Some(control_plane) = cp { |
| 230 | - return; | 321 | + for session_id in ids { |
| 231 | - }; | 322 | + if let Err(err) = control_plane.force_close_session(&session_id).await { |
| 232 | - for session_id in ids { | 323 | + tracing::warn!( |
| 233 | - if let Err(err) = control_plane.force_close_session(&session_id).await { | 324 | + session_id = %session_id, |
| 234 | - tracing::warn!( | 325 | + error = %err, |
| 235 | - session_id = %session_id, | 326 | + "failed to close session on exit" |
| 236 | - error = %err, | 327 | + ); |
| 237 | - "failed to close session on exit" | 328 | + } |
| 238 | - ); | 329 | + } |
| 330 | + } | ||
| 331 | + let automation = self.memory_automation.lock().await.take().flatten(); | ||
| 332 | + if let Ok(mut receiver) = self.memory_health.lock() { | ||
| 333 | + *receiver = None; | ||
| 334 | + } | ||
| 335 | + if let Some(automation) = automation { | ||
| 336 | + let close = | ||
| 337 | + tokio::time::timeout(std::time::Duration::from_secs(5), automation.close()).await; | ||
| 338 | + match close { | ||
| 339 | + Ok(Ok(())) => {} | ||
| 340 | + Ok(Err(error)) => { | ||
| 341 | + tracing::warn!(error = %error, "failed to close MCP memory automation") | ||
| 342 | + } | ||
| 343 | + Err(_) => tracing::warn!("MCP memory automation close timed out after 5 seconds"), | ||
| 239 | } | 344 | } |
| 240 | } | 345 | } |
| 241 | } | 346 | } |
| 242 | } | 347 | } |
| 348 | + | ||
| 349 | + | ||
| 350 | +mod tests { | ||
| 351 | + use crate::gateway::MemoryAutomationHealth; | ||
| 352 | + use crate::status_panel::MemoryStatus; | ||
| 353 | + | ||
| 354 | + use super::memory_status_from_health; | ||
| 355 | + | ||
| 356 | + | ||
| 357 | + fn memory_health_maps_to_operator_status() { | ||
| 358 | + assert_eq!( | ||
| 359 | + memory_status_from_health(MemoryAutomationHealth::Healthy), | ||
| 360 | + MemoryStatus::Connected | ||
| 361 | + ); | ||
| 362 | + assert_eq!( | ||
| 363 | + memory_status_from_health(MemoryAutomationHealth::Degraded), | ||
| 364 | + MemoryStatus::Degraded | ||
| 365 | + ); | ||
| 366 | + } | ||
| 367 | +} | ||
| @@ -78,8 +78,15 @@ where | |||
| 78 | config_arg.path.display() | 78 | config_arg.path.display() |
| 79 | ) | 79 | ) |
| 80 | })?; | 80 | })?; |
| 81 | - let config = config::require_tui_bootstrap_config(config, &config_arg.path)?; | 81 | + let mut config = config::require_tui_bootstrap_config(config, &config_arg.path)?; |
| 82 | - run_tui(config, config_arg.path).await | 82 | + let workspace = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); |
| 83 | + config.load_runtime_mcp_servers( | ||
| 84 | + config_arg.mcp_config.as_deref(), | ||
| 85 | + &workspace, | ||
| 86 | + dirs::home_dir().as_deref(), | ||
| 87 | + &config_arg.path, | ||
| 88 | + )?; | ||
| 89 | + run_tui(config, config_arg.path, workspace).await | ||
| 83 | } | 90 | } |
| 84 | 91 | ||
| 85 | 92 | ||
| @@ -112,7 +119,7 @@ fn os_str_eq(value: &OsStr, expected: &str) -> bool { | |||
| 112 | 119 | ||
| 113 | fn print_end_side_usage(program: &OsStr) { | 120 | fn print_end_side_usage(program: &OsStr) { |
| 114 | eprintln!( | 121 | eprintln!( |
| 115 | - "Usage: {} [--config <path>]\n {} --cli <command>\n\nDefault: launch the TUI.\nCLI: pass --cli before existing CLI commands, for example `{} --cli run -p \"hello\"`.", | 122 | + "Usage: {} [--config <path>] [--mcp-config <path>]\n {} --cli <command>\n\nDefault: launch the TUI.\nCLI: pass --cli before existing CLI commands, for example `{} --cli run -p \"hello\"`.", |
| 116 | PathBuf::from(program).display(), | 123 | PathBuf::from(program).display(), |
| 117 | PathBuf::from(program).display(), | 124 | PathBuf::from(program).display(), |
| 118 | PathBuf::from(program).display() | 125 | PathBuf::from(program).display() |
| @@ -122,6 +129,7 @@ fn print_end_side_usage(program: &OsStr) { | |||
| 122 | struct ConfigArg { | 129 | struct ConfigArg { |
| 123 | path: PathBuf, | 130 | path: PathBuf, |
| 124 | explicit: bool, | 131 | explicit: bool, |
| 132 | + mcp_config: Option<PathBuf>, | ||
| 125 | } | 133 | } |
| 126 | 134 | ||
| 127 | fn parse_config_path_from<I, T>(args: I) -> Result<ConfigArg> | 135 | fn parse_config_path_from<I, T>(args: I) -> Result<ConfigArg> |
| @@ -132,30 +140,36 @@ where | |||
| 132 | let mut args = args.into_iter().map(Into::into); | 140 | let mut args = args.into_iter().map(Into::into); |
| 133 | let program = args.next().unwrap_or_else(|| OsString::from("xiaoo")); | 141 | let program = args.next().unwrap_or_else(|| OsString::from("xiaoo")); |
| 134 | 142 | ||
| 135 | - let cli_path = match args.next() { | 143 | + let mut cli_path = None; |
| 136 | - None => None, | 144 | + let mut mcp_config = None; |
| 137 | - Some(first) if first == "--help" || first == "-h" => { | 145 | + while let Some(argument) = args.next() { |
| 146 | + if argument == "--help" || argument == "-h" { | ||
| 138 | print_usage(&program); | 147 | print_usage(&program); |
| 139 | std::process::exit(0); | 148 | std::process::exit(0); |
| 140 | - } | 149 | + } else if argument == "--config" || argument == "-c" { |
| 141 | - Some(first) if first == "--config" || first == "-c" => { | ||
| 142 | let Some(path) = args.next() else { | 150 | let Some(path) = args.next() else { |
| 143 | bail!("missing value for --config"); | 151 | bail!("missing value for --config"); |
| 144 | }; | 152 | }; |
| 145 | - if args.next().is_some() { | 153 | + if cli_path.replace(PathBuf::from(path)).is_some() { |
| 146 | - bail!("unexpected extra arguments after --config"); | 154 | + bail!("--config may only be specified once"); |
| 147 | } | 155 | } |
| 148 | - Some(PathBuf::from(path)) | 156 | + } else if argument == "--mcp-config" { |
| 157 | + let Some(path) = args.next() else { | ||
| 158 | + bail!("missing value for --mcp-config"); | ||
| 159 | + }; | ||
| 160 | + if mcp_config.replace(PathBuf::from(path)).is_some() { | ||
| 161 | + bail!("--mcp-config may only be specified once"); | ||
| 162 | + } | ||
| 163 | + } else { | ||
| 164 | + bail!("unsupported argument {:?}. use --help for usage", argument); | ||
| 149 | } | 165 | } |
| 150 | - Some(_) => { | 166 | + } |
| 151 | - bail!("unsupported arguments. use --help for usage, or pass only --config <path>") | ||
| 152 | - } | ||
| 153 | - }; | ||
| 154 | 167 | ||
| 155 | if let Some(path) = cli_path { | 168 | if let Some(path) = cli_path { |
| 156 | return Ok(ConfigArg { | 169 | return Ok(ConfigArg { |
| 157 | path, | 170 | path, |
| 158 | explicit: true, | 171 | explicit: true, |
| 172 | + mcp_config, | ||
| 159 | }); | 173 | }); |
| 160 | } | 174 | } |
| 161 | 175 | ||
| @@ -166,12 +180,14 @@ where | |||
| 166 | return Ok(ConfigArg { | 180 | return Ok(ConfigArg { |
| 167 | path, | 181 | path, |
| 168 | explicit: true, | 182 | explicit: true, |
| 183 | + mcp_config, | ||
| 169 | }); | 184 | }); |
| 170 | } | 185 | } |
| 171 | 186 | ||
| 172 | Ok(ConfigArg { | 187 | Ok(ConfigArg { |
| 173 | path: default_config_path()?, | 188 | path: default_config_path()?, |
| 174 | explicit: false, | 189 | explicit: false, |
| 190 | + mcp_config, | ||
| 175 | }) | 191 | }) |
| 176 | } | 192 | } |
| 177 | 193 | ||
| @@ -187,7 +203,7 @@ fn load_tui_config(config_arg: &ConfigArg) -> Result<Option<config::Config>> { | |||
| 187 | 203 | ||
| 188 | fn print_usage(program: &std::ffi::OsStr) { | 204 | fn print_usage(program: &std::ffi::OsStr) { |
| 189 | eprintln!( | 205 | eprintln!( |
| 190 | - "Usage: {} [--config <path>]\n\nConfig lookup order: --config > XIAOO_CONFIG > platform default.\nLaunch the TUI binary directly.", | 206 | + "Usage: {} [--config <path>] [--mcp-config <path>]\n\nConfig lookup order: --config > XIAOO_CONFIG > platform default.\nMCP lookup order: --mcp-config > XIAOO_MCP_CONFIG > workspace .mcp.json > ~/.config/xiaoo/mcp.json.\nLaunch the TUI binary directly.", |
| 191 | PathBuf::from(program).display() | 207 | PathBuf::from(program).display() |
| 192 | ); | 208 | ); |
| 193 | } | 209 | } |
| @@ -216,7 +232,7 @@ fn default_config_path() -> Result<PathBuf> { | |||
| 216 | } | 232 | } |
| 217 | } | 233 | } |
| 218 | 234 | ||
| 219 | -async fn run_tui(config: config::Config, config_path: PathBuf) -> Result<()> { | 235 | +async fn run_tui(config: config::Config, config_path: PathBuf, workspace: PathBuf) -> Result<()> { |
| 220 | let (validation_errors, validation_warnings) = validate_config_for_tui(&config, &config_path); | 236 | let (validation_errors, validation_warnings) = validate_config_for_tui(&config, &config_path); |
| 221 | 237 | ||
| 222 | for warning in &validation_warnings { | 238 | for warning in &validation_warnings { |
| @@ -241,7 +257,6 @@ async fn run_tui(config: config::Config, config_path: PathBuf) -> Result<()> { | |||
| 241 | let _ = execute!(io::stdout(), EnableMouseCapture); | 257 | let _ = execute!(io::stdout(), EnableMouseCapture); |
| 242 | let _ = execute!(io::stdout(), EnableBracketedPaste); | 258 | let _ = execute!(io::stdout(), EnableBracketedPaste); |
| 243 | 259 | ||
| 244 | - let workspace = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); | ||
| 245 | let mut app = app::App::new_with_config(&config, config_path.clone(), workspace) | 260 | let mut app = app::App::new_with_config(&config, config_path.clone(), workspace) |
| 246 | .context("failed to initialize TUI app state")?; | 261 | .context("failed to initialize TUI app state")?; |
| 247 | 262 | ||
| @@ -339,9 +354,27 @@ fn validate_config_for_tui( | |||
| 339 | 354 | ||
| 340 | 355 | ||
| 341 | mod tests { | 356 | mod tests { |
| 342 | - use super::{classify_args, EntryInvocation}; | 357 | + use super::{classify_args, parse_config_path_from, EntryInvocation}; |
| 343 | use std::ffi::OsString; | 358 | use std::ffi::OsString; |
| 344 | 359 | ||
| 360 | + | ||
| 361 | + fn tui_accepts_mcp_config_alongside_toml_config() { | ||
| 362 | + let parsed = parse_config_path_from([ | ||
| 363 | + "xiaoo", | ||
| 364 | + "--config", | ||
| 365 | + "/tmp/config.toml", | ||
| 366 | + "--mcp-config", | ||
| 367 | + "/tmp/mcp.json", | ||
| 368 | + ]) | ||
| 369 | + .expect("TUI should accept both config paths"); | ||
| 370 | + | ||
| 371 | + assert_eq!(parsed.path, std::path::PathBuf::from("/tmp/config.toml")); | ||
| 372 | + assert_eq!( | ||
| 373 | + parsed.mcp_config, | ||
| 374 | + Some(std::path::PathBuf::from("/tmp/mcp.json")) | ||
| 375 | + ); | ||
| 376 | + } | ||
| 377 | + | ||
| 345 | 378 | ||
| 346 | fn no_args_dispatches_to_tui() { | 379 | fn no_args_dispatches_to_tui() { |
| 347 | assert_eq!( | 380 | assert_eq!( |
| @@ -281,6 +281,13 @@ impl App { | |||
| 281 | ), | 281 | ), |
| 282 | Span::styled(" WS ", Style::default().fg(self.state.theme.muted)), | 282 | Span::styled(" WS ", Style::default().fg(self.state.theme.muted)), |
| 283 | Span::styled(workspace, Style::default().fg(self.state.theme.foreground)), | 283 | Span::styled(workspace, Style::default().fg(self.state.theme.foreground)), |
| 284 | + Span::styled(" Mem ", Style::default().fg(self.state.theme.muted)), | ||
| 285 | + Span::styled( | ||
| 286 | + self.state.status_panel.memory_status.label(), | ||
| 287 | + Style::default() | ||
| 288 | + .fg(self.state.theme.primary) | ||
| 289 | + .add_modifier(Modifier::BOLD), | ||
| 290 | + ), | ||
| 284 | Span::styled(" Tok ", Style::default().fg(self.state.theme.muted)), | 291 | Span::styled(" Tok ", Style::default().fg(self.state.theme.muted)), |
| 285 | Span::styled( | 292 | Span::styled( |
| 286 | StatusPanel::format_token_count(self.state.status_panel.total_tokens), | 293 | StatusPanel::format_token_count(self.state.status_panel.total_tokens), |
| @@ -2,6 +2,26 @@ use crate::render::utils::sanitize_terminal_text; | |||
| 2 | 2 | ||
| 3 | use std::path::Path; | 3 | use std::path::Path; |
| 4 | 4 | ||
| 5 | + | ||
| 6 | +pub enum MemoryStatus { | ||
| 7 | + | ||
| 8 | + Disabled, | ||
| 9 | + Unknown, | ||
| 10 | + Connected, | ||
| 11 | + Degraded, | ||
| 12 | +} | ||
| 13 | + | ||
| 14 | +impl MemoryStatus { | ||
| 15 | + pub(crate) fn label(self) -> &'static str { | ||
| 16 | + match self { | ||
| 17 | + Self::Disabled => "OFF", | ||
| 18 | + Self::Unknown => "UNKNOWN", | ||
| 19 | + Self::Connected => "OK", | ||
| 20 | + Self::Degraded => "DEGRADED", | ||
| 21 | + } | ||
| 22 | + } | ||
| 23 | +} | ||
| 24 | + | ||
| 5 | pub struct StatusPanel { | 25 | pub struct StatusPanel { |
| 6 | pub model_name: String, | 26 | pub model_name: String, |
| 7 | pub provider_name: String, | 27 | pub provider_name: String, |
| @@ -15,6 +35,7 @@ pub struct StatusPanel { | |||
| 15 | pub is_connected: bool, | 35 | pub is_connected: bool, |
| 16 | pub input_context_tokens: u64, | 36 | pub input_context_tokens: u64, |
| 17 | pub input_context_tokens_estimated: bool, | 37 | pub input_context_tokens_estimated: bool, |
| 38 | + pub memory_status: MemoryStatus, | ||
| 18 | } | 39 | } |
| 19 | 40 | ||
| 20 | impl Default for StatusPanel { | 41 | impl Default for StatusPanel { |
| @@ -31,6 +52,7 @@ impl Default for StatusPanel { | |||
| 31 | is_connected: false, | 52 | is_connected: false, |
| 32 | input_context_tokens: 0, | 53 | input_context_tokens: 0, |
| 33 | input_context_tokens_estimated: false, | 54 | input_context_tokens_estimated: false, |
| 55 | + memory_status: MemoryStatus::Disabled, | ||
| 34 | } | 56 | } |
| 35 | } | 57 | } |
| 36 | } | 58 | } |
| @@ -108,7 +130,14 @@ fn shorten_path_display(path: &Path, max_chars: usize) -> String { | |||
| 108 | 130 | ||
| 109 | 131 | ||
| 110 | mod tests { | 132 | mod tests { |
| 111 | - use super::StatusPanel; | 133 | + use super::{MemoryStatus, StatusPanel}; |
| 134 | + | ||
| 135 | + | ||
| 136 | + fn memory_status_labels_are_operator_readable() { | ||
| 137 | + assert_eq!(MemoryStatus::Unknown.label(), "UNKNOWN"); | ||
| 138 | + assert_eq!(MemoryStatus::Connected.label(), "OK"); | ||
| 139 | + assert_eq!(MemoryStatus::Degraded.label(), "DEGRADED"); | ||
| 140 | + } | ||
| 112 | 141 | ||
| 113 | 142 | ||
| 114 | fn format_context_usage_marks_estimated_values() { | 143 | fn format_context_usage_marks_estimated_values() { |
| @@ -13,6 +13,7 @@ use std::fs; | |||
| 13 | use std::path::{Path, PathBuf}; | 13 | use std::path::{Path, PathBuf}; |
| 14 | use std::sync::Arc; | 14 | use std::sync::Arc; |
| 15 | use xiaoo_shared::builtin_agent_roles::{PLAN_AGENT_DESCRIPTION, PLAN_AGENT_ID, PLAN_AGENT_PROMPT}; | 15 | use xiaoo_shared::builtin_agent_roles::{PLAN_AGENT_DESCRIPTION, PLAN_AGENT_ID, PLAN_AGENT_PROMPT}; |
| 16 | +use xiaoo_shared::gateway::MemoryAutomationConfig; | ||
| 16 | 17 | ||
| 17 | const DEFAULT_AGENT_ID: &str = "main"; | 18 | const DEFAULT_AGENT_ID: &str = "main"; |
| 18 | const DEFAULT_LLM_MAX_TOKENS: u32 = 16384; | 19 | const DEFAULT_LLM_MAX_TOKENS: u32 = 16384; |
| @@ -64,6 +65,13 @@ pub struct Config { | |||
| 64 | pub tui: TuiConfig, | 65 | pub tui: TuiConfig, |
| 65 | 66 | ||
| 66 | pub mcp: McpSection, | 67 | pub mcp: McpSection, |
| 68 | + | ||
| 69 | + pub memory_automation: MemoryAutomationConfig, | ||
| 70 | + /// Effective MCP servers after adding optional `.mcp.json` entries. This | ||
| 71 | + /// is runtime-only so TUI config saves never copy imported servers into | ||
| 72 | + /// `config.toml`. | ||
| 73 | + | ||
| 74 | + runtime_mcp_servers: Option<Vec<mcp::McpServerConfig>>, | ||
| 67 | /// Preserve daemon- or plugin-owned top-level sections when the TUI | 75 | /// Preserve daemon- or plugin-owned top-level sections when the TUI |
| 68 | /// rewrites config.toml after changing providers or other UI settings. | 76 | /// rewrites config.toml after changing providers or other UI settings. |
| 69 | 77 | ||
| @@ -224,6 +232,29 @@ impl Config { | |||
| 224 | Ok(()) | 232 | Ok(()) |
| 225 | } | 233 | } |
| 226 | 234 | ||
| 235 | + pub fn load_runtime_mcp_servers( | ||
| 236 | + &mut self, | ||
| 237 | + explicit_path: Option<&Path>, | ||
| 238 | + workspace: &Path, | ||
| 239 | + home: Option<&Path>, | ||
| 240 | + toml_source: &Path, | ||
| 241 | + ) -> Result<(), mcp::McpConfigError> { | ||
| 242 | + self.runtime_mcp_servers = Some(load_merged_mcp_servers( | ||
| 243 | + &self.mcp.servers, | ||
| 244 | + explicit_path, | ||
| 245 | + workspace, | ||
| 246 | + home, | ||
| 247 | + toml_source, | ||
| 248 | + )?); | ||
| 249 | + Ok(()) | ||
| 250 | + } | ||
| 251 | + | ||
| 252 | + pub fn mcp_servers(&self) -> &[mcp::McpServerConfig] { | ||
| 253 | + self.runtime_mcp_servers | ||
| 254 | + .as_deref() | ||
| 255 | + .unwrap_or(&self.mcp.servers) | ||
| 256 | + } | ||
| 257 | + | ||
| 227 | pub fn list_agent_ids(&self) -> Vec<String> { | 258 | pub fn list_agent_ids(&self) -> Vec<String> { |
| 228 | self.agents | 259 | self.agents |
| 229 | .list | 260 | .list |
| @@ -351,6 +382,24 @@ impl Config { | |||
| 351 | } | 382 | } |
| 352 | } | 383 | } |
| 353 | 384 | ||
| 385 | +pub(crate) fn load_merged_mcp_servers( | ||
| 386 | + toml_servers: &[mcp::McpServerConfig], | ||
| 387 | + explicit_path: Option<&Path>, | ||
| 388 | + workspace: &Path, | ||
| 389 | + home: Option<&Path>, | ||
| 390 | + toml_source: &Path, | ||
| 391 | +) -> Result<Vec<mcp::McpServerConfig>, mcp::McpConfigError> { | ||
| 392 | + let json_source = mcp::resolve_json_config_path(explicit_path, workspace, home); | ||
| 393 | + let json_servers = mcp::load_json_servers(explicit_path, workspace, home)?; | ||
| 394 | + let fallback_json_source = workspace.join(".mcp.json"); | ||
| 395 | + mcp::merge_server_configs( | ||
| 396 | + toml_servers.to_vec(), | ||
| 397 | + json_servers, | ||
| 398 | + toml_source, | ||
| 399 | + json_source.as_deref().unwrap_or(&fallback_json_source), | ||
| 400 | + ) | ||
| 401 | +} | ||
| 402 | + | ||
| 354 | pub fn require_tui_bootstrap_config(config: Option<Config>, config_path: &Path) -> Result<Config> { | 403 | pub fn require_tui_bootstrap_config(config: Option<Config>, config_path: &Path) -> Result<Config> { |
| 355 | let mut config = config | 404 | let mut config = config |
| 356 | .ok_or_else(|| anyhow::anyhow!("config file not found: {}", config_path.display()))?; | 405 | .ok_or_else(|| anyhow::anyhow!("config file not found: {}", config_path.display()))?; |
| @@ -554,6 +603,39 @@ mod tests { | |||
| 554 | assert_eq!(resolve_context_window(&config), Some(200_000)); | 603 | assert_eq!(resolve_context_window(&config), Some(200_000)); |
| 555 | } | 604 | } |
| 556 | 605 | ||
| 606 | + | ||
| 607 | + fn parses_memory_automation_config() { | ||
| 608 | + let config: Config = toml::from_str( | ||
| 609 | + r#" | ||
| 610 | +[memory_automation] | ||
| 611 | +enabled = true | ||
| 612 | +server = "ram-a" | ||
| 613 | +recall_top_k = 3 | ||
| 614 | +recall_token_budget = 128 | ||
| 615 | +context_messages = 2 | ||
| 616 | +queue_path = "/tmp/xiaoo-memory-queue.jsonl" | ||
| 617 | +queue_capacity = 32 | ||
| 618 | +max_retries = 4 | ||
| 619 | +retry_backoff_ms = 50 | ||
| 620 | +allowed_agent_roles = ["main", "researcher"] | ||
| 621 | +"#, | ||
| 622 | + ) | ||
| 623 | + .expect("config should parse"); | ||
| 624 | + | ||
| 625 | + assert!(config.memory_automation.enabled); | ||
| 626 | + assert_eq!(config.memory_automation.server, "ram-a"); | ||
| 627 | + assert_eq!(config.memory_automation.recall_top_k, 3); | ||
| 628 | + assert_eq!(config.memory_automation.recall_token_budget, 128); | ||
| 629 | + assert_eq!(config.memory_automation.context_messages, 2); | ||
| 630 | + assert_eq!(config.memory_automation.queue_capacity, 32); | ||
| 631 | + assert_eq!(config.memory_automation.max_retries, 4); | ||
| 632 | + assert_eq!(config.memory_automation.retry_backoff_ms, 50); | ||
| 633 | + assert_eq!( | ||
| 634 | + config.memory_automation.allowed_agent_roles, | ||
| 635 | + vec!["main".to_string(), "researcher".to_string()] | ||
| 636 | + ); | ||
| 637 | + } | ||
| 638 | + | ||
| 557 | 639 | ||
| 558 | fn secret_provider_initializes_when_existing_secrets_file_is_present() { | 640 | fn secret_provider_initializes_when_existing_secrets_file_is_present() { |
| 559 | let mut config = valid_config(); | 641 | let mut config = valid_config(); |
| @@ -665,6 +747,56 @@ kind = "local" | |||
| 665 | ); | 747 | ); |
| 666 | } | 748 | } |
| 667 | 749 | ||
| 750 | + | ||
| 751 | + fn imported_json_mcp_servers_are_runtime_only_when_tui_saves() { | ||
| 752 | + let temp = tempdir().expect("tempdir"); | ||
| 753 | + let config_path = temp.path().join("config.toml"); | ||
| 754 | + let json_path = temp.path().join("mcp.json"); | ||
| 755 | + std::fs::write( | ||
| 756 | + &config_path, | ||
| 757 | + r#" | ||
| 758 | +[llm] | ||
| 759 | +provider = "openai" | ||
| 760 | +model = "gpt-4o" | ||
| 761 | + | ||
| 762 | +[[mcp.servers]] | ||
| 763 | +name = "toml-server" | ||
| 764 | +transport = "stdio" | ||
| 765 | +command = "toml-server" | ||
| 766 | +"#, | ||
| 767 | + ) | ||
| 768 | + .expect("write TOML config"); | ||
| 769 | + std::fs::write( | ||
| 770 | + &json_path, | ||
| 771 | + r#"{"mcpServers":{"json-server":{"transport":"stdio","command":"json-server"}}}"#, | ||
| 772 | + ) | ||
| 773 | + .expect("write JSON config"); | ||
| 774 | + | ||
| 775 | + let mut config = Config::load_from(&config_path).expect("load TOML config"); | ||
| 776 | + config | ||
| 777 | + .load_runtime_mcp_servers( | ||
| 778 | + Some(&json_path), | ||
| 779 | + temp.path(), | ||
| 780 | + Some(temp.path()), | ||
| 781 | + &config_path, | ||
| 782 | + ) | ||
| 783 | + .expect("merge JSON MCP servers"); | ||
| 784 | + assert_eq!( | ||
| 785 | + config | ||
| 786 | + .mcp_servers() | ||
| 787 | + .iter() | ||
| 788 | + .map(|server| server.name.as_str()) | ||
| 789 | + .collect::<Vec<_>>(), | ||
| 790 | + vec!["toml-server", "json-server"] | ||
| 791 | + ); | ||
| 792 | + | ||
| 793 | + config.llm.model = "gpt-4.1".to_string(); | ||
| 794 | + config.save_to(&config_path).expect("save TUI config"); | ||
| 795 | + let persisted = Config::load_from(&config_path).expect("reload TOML config"); | ||
| 796 | + assert_eq!(persisted.mcp.servers.len(), 1); | ||
| 797 | + assert_eq!(persisted.mcp.servers[0].name, "toml-server"); | ||
| 798 | + } | ||
| 799 | + | ||
| 668 | 800 | ||
| 669 | fn tui_bootstrap_adds_builtin_plan_agent_role() { | 801 | fn tui_bootstrap_adds_builtin_plan_agent_role() { |
| 670 | let config = | 802 | let config = |
| @@ -18,6 +18,7 @@ use std::path::{Path, PathBuf}; | |||
| 18 | use std::sync::Arc; | 18 | use std::sync::Arc; |
| 19 | use xiaoo_shared::backend::GatewayBackendConfig; | 19 | use xiaoo_shared::backend::GatewayBackendConfig; |
| 20 | use xiaoo_shared::builtin_agent_roles::{PLAN_AGENT_DESCRIPTION, PLAN_AGENT_ID, PLAN_AGENT_PROMPT}; | 20 | use xiaoo_shared::builtin_agent_roles::{PLAN_AGENT_DESCRIPTION, PLAN_AGENT_ID, PLAN_AGENT_PROMPT}; |
| 21 | +use xiaoo_shared::gateway::MemoryAutomationConfig; | ||
| 21 | 22 | ||
| 22 | const DEFAULT_OUTPUT_TOKENS: usize = 16384; | 23 | const DEFAULT_OUTPUT_TOKENS: usize = 16384; |
| 23 | const DEFAULT_SYSTEM_PROMPT: &str = include_str!("prompts/default_system_prompt.txt"); | 24 | const DEFAULT_SYSTEM_PROMPT: &str = include_str!("prompts/default_system_prompt.txt"); |
| @@ -54,6 +55,8 @@ pub struct AppConfig { | |||
| 54 | 55 | ||
| 55 | pub mcp: McpSection, | 56 | pub mcp: McpSection, |
| 56 | 57 | ||
| 58 | + pub memory_automation: MemoryAutomationConfig, | ||
| 59 | + | ||
| 57 | pub mcp_server: McpServerConfig, | 60 | pub mcp_server: McpServerConfig, |
| 58 | } | 61 | } |
| 59 | 62 | ||
| @@ -376,6 +379,27 @@ impl DaemonConfig { | |||
| 376 | Ok(Self { app, config_path }) | 379 | Ok(Self { app, config_path }) |
| 377 | } | 380 | } |
| 378 | 381 | ||
| 382 | + pub fn load_with_mcp_config( | ||
| 383 | + path: impl AsRef<Path>, | ||
| 384 | + explicit_path: Option<&Path>, | ||
| 385 | + workspace: &Path, | ||
| 386 | + home: Option<&Path>, | ||
| 387 | + ) -> Result<Self> { | ||
| 388 | + let mut config = Self::load_from(path)?; | ||
| 389 | + let json_source = mcp::resolve_json_config_path(explicit_path, workspace, home); | ||
| 390 | + let json_servers = mcp::load_json_servers(explicit_path, workspace, home) | ||
| 391 | + .context("failed to load MCP JSON config")?; | ||
| 392 | + let fallback_json_source = workspace.join(".mcp.json"); | ||
| 393 | + config.app.mcp.servers = mcp::merge_server_configs( | ||
| 394 | + std::mem::take(&mut config.app.mcp.servers), | ||
| 395 | + json_servers, | ||
| 396 | + &config.config_path, | ||
| 397 | + json_source.as_deref().unwrap_or(&fallback_json_source), | ||
| 398 | + ) | ||
| 399 | + .context("failed to merge MCP server configs")?; | ||
| 400 | + Ok(config) | ||
| 401 | + } | ||
| 402 | + | ||
| 379 | pub fn resolve_agent(&self) -> Result<ResolvedAgentConfig> { | 403 | pub fn resolve_agent(&self) -> Result<ResolvedAgentConfig> { |
| 380 | let default_agent_id = self | 404 | let default_agent_id = self |
| 381 | .app | 405 | .app |
| @@ -1013,6 +1037,87 @@ mod tests { | |||
| 1013 | use super::{resolve_config_path, AppConfig, DaemonConfig}; | 1037 | use super::{resolve_config_path, AppConfig, DaemonConfig}; |
| 1014 | use tempfile::TempDir; | 1038 | use tempfile::TempDir; |
| 1015 | 1039 | ||
| 1040 | + | ||
| 1041 | + fn parses_memory_automation_config() { | ||
| 1042 | + let content = r#" | ||
| 1043 | +[llm] | ||
| 1044 | +provider = "openai" | ||
| 1045 | +model = "gpt-4o" | ||
| 1046 | + | ||
| 1047 | +[memory_automation] | ||
| 1048 | +enabled = true | ||
| 1049 | +server = "ram-a" | ||
| 1050 | +recall_top_k = 3 | ||
| 1051 | +recall_token_budget = 128 | ||
| 1052 | +context_messages = 2 | ||
| 1053 | +queue_path = "/tmp/xiaoo-memory-queue.jsonl" | ||
| 1054 | +queue_capacity = 32 | ||
| 1055 | +max_retries = 4 | ||
| 1056 | +retry_backoff_ms = 50 | ||
| 1057 | +allowed_agent_roles = ["main", "researcher"] | ||
| 1058 | +"#; | ||
| 1059 | + | ||
| 1060 | + let config: AppConfig = toml::from_str(content).expect("config should parse"); | ||
| 1061 | + | ||
| 1062 | + assert!(config.memory_automation.enabled); | ||
| 1063 | + assert_eq!(config.memory_automation.server, "ram-a"); | ||
| 1064 | + assert_eq!(config.memory_automation.recall_top_k, 3); | ||
| 1065 | + assert_eq!(config.memory_automation.recall_token_budget, 128); | ||
| 1066 | + assert_eq!(config.memory_automation.context_messages, 2); | ||
| 1067 | + assert_eq!(config.memory_automation.queue_capacity, 32); | ||
| 1068 | + assert_eq!(config.memory_automation.max_retries, 4); | ||
| 1069 | + assert_eq!(config.memory_automation.retry_backoff_ms, 50); | ||
| 1070 | + assert_eq!( | ||
| 1071 | + config.memory_automation.allowed_agent_roles, | ||
| 1072 | + vec!["main".to_string(), "researcher".to_string()] | ||
| 1073 | + ); | ||
| 1074 | + } | ||
| 1075 | + | ||
| 1076 | + | ||
| 1077 | + fn daemon_load_merges_runtime_json_mcp_servers() { | ||
| 1078 | + let temp = TempDir::new().expect("tempdir"); | ||
| 1079 | + let config_path = temp.path().join("config.toml"); | ||
| 1080 | + let json_path = temp.path().join("mcp.json"); | ||
| 1081 | + std::fs::write( | ||
| 1082 | + &config_path, | ||
| 1083 | + r#" | ||
| 1084 | +[llm] | ||
| 1085 | +provider = "openrouter" | ||
| 1086 | +model = "z-ai/glm-5" | ||
| 1087 | + | ||
| 1088 | +[[mcp.servers]] | ||
| 1089 | +name = "toml-server" | ||
| 1090 | +transport = "stdio" | ||
| 1091 | +command = "toml-server" | ||
| 1092 | +"#, | ||
| 1093 | + ) | ||
| 1094 | + .expect("write TOML config"); | ||
| 1095 | + std::fs::write( | ||
| 1096 | + &json_path, | ||
| 1097 | + r#"{"mcpServers":{"json-server":{"transport":"stdio","command":"json-server"}}}"#, | ||
| 1098 | + ) | ||
| 1099 | + .expect("write JSON config"); | ||
| 1100 | + | ||
| 1101 | + let daemon = DaemonConfig::load_with_mcp_config( | ||
| 1102 | + &config_path, | ||
| 1103 | + Some(&json_path), | ||
| 1104 | + temp.path(), | ||
| 1105 | + Some(temp.path()), | ||
| 1106 | + ) | ||
| 1107 | + .expect("load merged daemon config"); | ||
| 1108 | + | ||
| 1109 | + assert_eq!( | ||
| 1110 | + daemon | ||
| 1111 | + .app | ||
| 1112 | + .mcp | ||
| 1113 | + .servers | ||
| 1114 | + .iter() | ||
| 1115 | + .map(|server| server.name.as_str()) | ||
| 1116 | + .collect::<Vec<_>>(), | ||
| 1117 | + vec!["toml-server", "json-server"] | ||
| 1118 | + ); | ||
| 1119 | + } | ||
| 1120 | + | ||
| 1016 | 1121 | ||
| 1017 | fn parses_feishu_channel_config() { | 1122 | fn parses_feishu_channel_config() { |
| 1018 | let content = r#" | 1123 | let content = r#" |
| @@ -452,7 +452,10 @@ fn create_router_from_state( | |||
| 452 | "/api/v1/runtimes/write-file", | 452 | "/api/v1/runtimes/write-file", |
| 453 | post(handle_runtime_write_file), | 453 | post(handle_runtime_write_file), |
| 454 | ) | 454 | ) |
| 455 | - .route("/api/v1/runtimes/export/:session_id", get(handle_session_export)), | 455 | + .route( |
| 456 | + "/api/v1/runtimes/export/:session_id", | ||
| 457 | + get(handle_session_export), | ||
| 458 | + ), | ||
| 456 | bearer_auth.clone(), | 459 | bearer_auth.clone(), |
| 457 | ); | 460 | ); |
| 458 | 461 | ||
| @@ -28,7 +28,9 @@ use std::path::PathBuf; | |||
| 28 | use std::sync::Arc; | 28 | use std::sync::Arc; |
| 29 | use tracing_subscriber::EnvFilter; | 29 | use tracing_subscriber::EnvFilter; |
| 30 | use xiaoo_shared::backend::BackendManager; | 30 | use xiaoo_shared::backend::BackendManager; |
| 31 | -use xiaoo_shared::gateway::{AppBootstrap, InMemorySessionStore, SessionStore}; | 31 | +use xiaoo_shared::gateway::{ |
| 32 | + AppBootstrap, InMemorySessionStore, McpMemoryAutomation, SessionStore, | ||
| 33 | +}; | ||
| 32 | 34 | ||
| 33 | 35 | ||
| 34 | async fn main() -> Result<()> { | 36 | async fn main() -> Result<()> { |
| @@ -42,6 +44,7 @@ async fn main() -> Result<()> { | |||
| 42 | } | 44 | } |
| 43 | run_daemon( | 45 | run_daemon( |
| 44 | cli.config, | 46 | cli.config, |
| 47 | + cli.mcp_config, | ||
| 45 | cli.host, | 48 | cli.host, |
| 46 | cli.port, | 49 | cli.port, |
| 47 | cli.dashboard_host, | 50 | cli.dashboard_host, |
| @@ -52,6 +55,7 @@ async fn main() -> Result<()> { | |||
| 52 | 55 | ||
| 53 | async fn run_daemon( | 56 | async fn run_daemon( |
| 54 | config_path: Option<PathBuf>, | 57 | config_path: Option<PathBuf>, |
| 58 | + mcp_config_path: Option<PathBuf>, | ||
| 55 | host: String, | 59 | host: String, |
| 56 | port: u16, | 60 | port: u16, |
| 57 | dashboard_cli_host: Option<String>, | 61 | dashboard_cli_host: Option<String>, |
| @@ -64,7 +68,13 @@ async fn run_daemon( | |||
| 64 | config_path.display() | 68 | config_path.display() |
| 65 | ) | 69 | ) |
| 66 | })?; | 70 | })?; |
| 67 | - let config = DaemonConfig::load_from(&config_path)?; | 71 | + let workspace = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); |
| 72 | + let config = DaemonConfig::load_with_mcp_config( | ||
| 73 | + &config_path, | ||
| 74 | + mcp_config_path.as_deref(), | ||
| 75 | + &workspace, | ||
| 76 | + dirs::home_dir().as_deref(), | ||
| 77 | + )?; | ||
| 68 | let mcp_server_config = config.resolve_mcp_server_config()?; | 78 | let mcp_server_config = config.resolve_mcp_server_config()?; |
| 69 | let hooker_config = config.app.hooker.clone(); | 79 | let hooker_config = config.app.hooker.clone(); |
| 70 | let bearer_auth = config.http_bearer_token()?.map(HttpBearerAuthConfig::new); | 80 | let bearer_auth = config.http_bearer_token()?.map(HttpBearerAuthConfig::new); |
| @@ -72,162 +82,192 @@ async fn run_daemon( | |||
| 72 | let resolver = Arc::new(ConfiguredRuntimeResolver::from_config(&config).await?); | 82 | let resolver = Arc::new(ConfiguredRuntimeResolver::from_config(&config).await?); |
| 73 | let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default()); | 83 | let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default()); |
| 74 | let backend_manager = Arc::new(BackendManager::new()); | 84 | let backend_manager = Arc::new(BackendManager::new()); |
| 75 | - // Start the cross-process signal handler so backends owned by this | 85 | + let memory_automation = match McpMemoryAutomation::connect( |
| 76 | - // daemon that another process has marked for eviction get evicted | 86 | + config.app.memory_automation.clone(), |
| 77 | - // immediately upon receiving SIGUSR1. | 87 | + &config.app.mcp.servers, |
| 78 | - let handler_handle = backend_manager | 88 | + ) |
| 79 | - .clone() | 89 | + .await |
| 80 | - .start_signal_handler(session_store.clone()); | 90 | + { |
| 81 | - tokio::spawn(async move { | 91 | + Ok(automation) => automation, |
| 82 | - handler_handle.await.ok(); | 92 | + Err(error) => { |
| 83 | - }); | 93 | + tracing::warn!(error = %error, "memory automation disabled after startup error"); |
| 84 | - let app = AppBootstrap::from_session_components_with_hooks_and_backend_manager( | 94 | + None |
| 85 | - session_store.clone(), | 95 | + } |
| 86 | - resolver, | 96 | + }; |
| 87 | - hooker_config, | 97 | + let memory_automation_for_shutdown = memory_automation.clone(); |
| 88 | - backend_manager.clone(), | 98 | + let serve_result = async { |
| 89 | - )?; | 99 | + // Start the cross-process signal handler so backends owned by this |
| 90 | - let interaction_timeout_secs = config.interaction_timeout_secs(); | 100 | + // daemon that another process has marked for eviction get evicted |
| 91 | - let session_service = app.session_service.clone(); | 101 | + // immediately upon receiving SIGUSR1. |
| 92 | - let session_control_plane = app.session_control_plane.clone(); | 102 | + let handler_handle = backend_manager |
| 103 | + .clone() | ||
| 104 | + .start_signal_handler(session_store.clone()); | ||
| 105 | + tokio::spawn(async move { | ||
| 106 | + handler_handle.await.ok(); | ||
| 107 | + }); | ||
| 108 | + let app = | ||
| 109 | + AppBootstrap::from_session_components_with_hooks_and_backend_manager_and_memory_automation( | ||
| 110 | + session_store.clone(), | ||
| 111 | + resolver, | ||
| 112 | + hooker_config, | ||
| 113 | + backend_manager.clone(), | ||
| 114 | + memory_automation, | ||
| 115 | + )?; | ||
| 116 | + let interaction_timeout_secs = config.interaction_timeout_secs(); | ||
| 117 | + let session_service = app.session_service.clone(); | ||
| 118 | + let session_control_plane = app.session_control_plane.clone(); | ||
| 93 | 119 | ||
| 94 | - if let Some(telegram_config) = config.telegram_polling_config()? { | 120 | + if let Some(telegram_config) = config.telegram_polling_config()? { |
| 95 | - spawn_telegram_polling_service( | 121 | + spawn_telegram_polling_service( |
| 96 | - telegram_config, | 122 | + telegram_config, |
| 97 | - session_service.clone(), | ||
| 98 | - interaction_timeout_secs, | ||
| 99 | - ) | ||
| 100 | - .context("failed to start telegram polling service")?; | ||
| 101 | - } | ||
| 102 | - | ||
| 103 | - if let Some(feishu_config) = config.feishu_config()? { | ||
| 104 | - if feishu_config.event_transport == FeishuEventTransport::Websocket { | ||
| 105 | - spawn_feishu_websocket_service( | ||
| 106 | - feishu_config, | ||
| 107 | session_service.clone(), | 123 | session_service.clone(), |
| 108 | interaction_timeout_secs, | 124 | interaction_timeout_secs, |
| 109 | ) | 125 | ) |
| 110 | - .context("failed to start Feishu websocket service")?; | 126 | + .context("failed to start telegram polling service")?; |
| 111 | } | 127 | } |
| 112 | - } | ||
| 113 | 128 | ||
| 114 | - // ── Cron scheduler ────────────────────────────────────────── | 129 | + if let Some(feishu_config) = config.feishu_config()? { |
| 115 | - let cron_enabled = config.cron_section().is_some(); | 130 | + if feishu_config.event_transport == FeishuEventTransport::Websocket { |
| 116 | - let cron_scheduler = match config.resolve_cron_jobs() { | 131 | + spawn_feishu_websocket_service( |
| 117 | - Ok(jobs) if !jobs.is_empty() => { | 132 | + feishu_config, |
| 118 | - let global = config | ||
| 119 | - .cron_section() | ||
| 120 | - .expect("cron section must exist when jobs loaded"); | ||
| 121 | - let total = jobs.len(); | ||
| 122 | - let enabled_count = jobs.iter().filter(|j| j.enabled).count(); | ||
| 123 | - if enabled_count > 0 { | ||
| 124 | - Some(Arc::new(CronScheduler::new( | ||
| 125 | - jobs, | ||
| 126 | - global.max_concurrent_jobs, | ||
| 127 | session_service.clone(), | 133 | session_service.clone(), |
| 128 | - ))) | 134 | + interaction_timeout_secs, |
| 129 | - } else { | 135 | + ) |
| 130 | - tracing::info!(total, "no enabled cron jobs"); | 136 | + .context("failed to start Feishu websocket service")?; |
| 137 | + } | ||
| 138 | + } | ||
| 139 | + | ||
| 140 | + // ── Cron scheduler ────────────────────────────────────────── | ||
| 141 | + let cron_enabled = config.cron_section().is_some(); | ||
| 142 | + let cron_scheduler = match config.resolve_cron_jobs() { | ||
| 143 | + Ok(jobs) if !jobs.is_empty() => { | ||
| 144 | + let global = config | ||
| 145 | + .cron_section() | ||
| 146 | + .expect("cron section must exist when jobs loaded"); | ||
| 147 | + let total = jobs.len(); | ||
| 148 | + let enabled_count = jobs.iter().filter(|j| j.enabled).count(); | ||
| 149 | + if enabled_count > 0 { | ||
| 150 | + Some(Arc::new(CronScheduler::new( | ||
| 151 | + jobs, | ||
| 152 | + global.max_concurrent_jobs, | ||
| 153 | + session_service.clone(), | ||
| 154 | + ))) | ||
| 155 | + } else { | ||
| 156 | + tracing::info!(total, "no enabled cron jobs"); | ||
| 157 | + None | ||
| 158 | + } | ||
| 159 | + } | ||
| 160 | + Ok(_) => { | ||
| 161 | + if cron_enabled { | ||
| 162 | + tracing::info!("cron section present but no jobs configured"); | ||
| 163 | + } | ||
| 131 | None | 164 | None |
| 132 | } | 165 | } |
| 133 | - } | 166 | + Err(error) => { |
| 134 | - Ok(_) => { | 167 | + tracing::error!(%error, "failed to load cron jobs, cron disabled"); |
| 135 | - if cron_enabled { | 168 | + None |
| 136 | - tracing::info!("cron section present but no jobs configured"); | ||
| 137 | } | 169 | } |
| 138 | - None | 170 | + }; |
| 139 | - } | ||
| 140 | - Err(error) => { | ||
| 141 | - tracing::error!(%error, "failed to load cron jobs, cron disabled"); | ||
| 142 | - None | ||
| 143 | - } | ||
| 144 | - }; | ||
| 145 | 171 | ||
| 146 | - let channel_runtimes = config.channel_runtimes()?; | 172 | + let channel_runtimes = config.channel_runtimes()?; |
| 147 | - let mut router = if channel_runtimes.is_empty() { | 173 | + let mut router = if channel_runtimes.is_empty() { |
| 148 | - create_router_with_control_plane_and_auth( | 174 | + create_router_with_control_plane_and_auth( |
| 149 | - session_service.clone(), | 175 | + session_service.clone(), |
| 150 | - session_control_plane.clone(), | 176 | + session_control_plane.clone(), |
| 151 | - bearer_auth, | 177 | + bearer_auth, |
| 152 | - rate_limit.clone(), | 178 | + rate_limit.clone(), |
| 153 | - ) | 179 | + ) |
| 154 | - } else { | 180 | + } else { |
| 155 | - create_router_with_channel_runtimes_control_plane_and_timeout_and_auth( | 181 | + create_router_with_channel_runtimes_control_plane_and_timeout_and_auth( |
| 156 | - session_service.clone(), | 182 | + session_service.clone(), |
| 157 | - session_control_plane.clone(), | 183 | + session_control_plane.clone(), |
| 158 | - channel_runtimes, | 184 | + channel_runtimes, |
| 159 | - interaction_timeout_secs, | 185 | + interaction_timeout_secs, |
| 160 | - bearer_auth, | 186 | + bearer_auth, |
| 161 | - rate_limit.clone(), | 187 | + rate_limit.clone(), |
| 162 | - ) | 188 | + ) |
| 163 | - .map_err(anyhow::Error::new) | 189 | + .map_err(anyhow::Error::new) |
| 164 | - .context("failed to create router with channel runtimes")? | 190 | + .context("failed to create router with channel runtimes")? |
| 165 | - }; | 191 | + }; |
| 166 | - if let Some(mcp_server_config) = mcp_server_config { | 192 | + if let Some(mcp_server_config) = mcp_server_config { |
| 167 | - router = router.merge(create_mcp_router( | 193 | + router = router.merge(create_mcp_router( |
| 168 | - mcp_server_config, | 194 | + mcp_server_config, |
| 169 | - session_service.clone(), | 195 | + session_service.clone(), |
| 170 | - session_control_plane.clone(), | 196 | + session_control_plane.clone(), |
| 197 | + session_store.clone(), | ||
| 198 | + rate_limit.clone(), | ||
| 199 | + )); | ||
| 200 | + } | ||
| 201 | + | ||
| 202 | + // Dashboard runs on its own listener so it never shares the runtime | ||
| 203 | + // API port (and its bearer auth). When `[http.dashboard].enabled = false` | ||
| 204 | + // is set in the config, `dashboard_port` resolves to `None` and no | ||
| 205 | + // dashboard server is started. | ||
| 206 | + if let Some(dash_addr) = spawn_dashboard_server( | ||
| 207 | + &config, | ||
| 208 | + dashboard_cli_host, | ||
| 209 | + dashboard_cli_port, | ||
| 171 | session_store.clone(), | 210 | session_store.clone(), |
| 172 | - rate_limit.clone(), | 211 | + backend_manager.clone(), |
| 173 | - )); | 212 | + ) |
| 174 | - } | 213 | + .await? |
| 175 | - | 214 | + { |
| 176 | - // Dashboard runs on its own listener so it never shares the runtime | 215 | + tracing::info!(%dash_addr, "dashboard ready at http://{dash_addr}"); |
| 177 | - // API port (and its bearer auth). When `[http.dashboard].enabled = false` | 216 | + eprintln!("dashboard ready at http://{dash_addr}"); |
| 178 | - // is set in the config, `dashboard_port` resolves to `None` and no | 217 | + } else { |
| 179 | - // dashboard server is started. | 218 | + tracing::info!("dashboard disabled by config ([http.dashboard].enabled = false)"); |
| 180 | - if let Some(dash_addr) = spawn_dashboard_server( | ||
| 181 | - &config, | ||
| 182 | - dashboard_cli_host, | ||
| 183 | - dashboard_cli_port, | ||
| 184 | - session_store.clone(), | ||
| 185 | - backend_manager.clone(), | ||
| 186 | - ) | ||
| 187 | - .await? | ||
| 188 | - { | ||
| 189 | - tracing::info!(%dash_addr, "dashboard ready at http://{dash_addr}"); | ||
| 190 | - eprintln!("dashboard ready at http://{dash_addr}"); | ||
| 191 | - } else { | ||
| 192 | - tracing::info!("dashboard disabled by config ([http.dashboard].enabled = false)"); | ||
| 193 | - } | ||
| 194 | - | ||
| 195 | - let addr: SocketAddr = format!("{host}:{port}") | ||
| 196 | - .parse() | ||
| 197 | - .with_context(|| format!("invalid listen address {host}:{port}"))?; | ||
| 198 | - let listener = tokio::net::TcpListener::bind(addr) | ||
| 199 | - .await | ||
| 200 | - .with_context(|| format!("failed to bind {addr}"))?; | ||
| 201 | - tracing::info!(config = %config_path.display(), %addr, "starting rebuild daemon"); | ||
| 202 | - let serve_result = axum::serve(listener, router) | ||
| 203 | - .with_graceful_shutdown(shutdown_signal()) | ||
| 204 | - .await | ||
| 205 | - .context("axum server exited unexpectedly"); | ||
| 206 | - | ||
| 207 | - // Gracefully shutdown cron scheduler | ||
| 208 | - if let Some(scheduler) = cron_scheduler { | ||
| 209 | - scheduler.stop().await; | ||
| 210 | - } | ||
| 211 | - | ||
| 212 | - // Best-effort sandbox cleanup with a bounded timeout so a slow/stuck | ||
| 213 | - // provider delete call cannot keep the daemon alive indefinitely after a | ||
| 214 | - // shutdown signal. Matches the TUI exit path which also bounds remote | ||
| 215 | - // close to a few seconds. | ||
| 216 | - let shutdown_result = tokio::time::timeout( | ||
| 217 | - std::time::Duration::from_secs(10), | ||
| 218 | - backend_manager.shutdown_all(), | ||
| 219 | - ) | ||
| 220 | - .await; | ||
| 221 | - match shutdown_result { | ||
| 222 | - Ok(Ok(())) => {} | ||
| 223 | - Ok(Err(error)) => { | ||
| 224 | - tracing::warn!(error = %error, "failed to shutdown daemon backend manager"); | ||
| 225 | } | 219 | } |
| 226 | - Err(_) => { | 220 | + |
| 227 | - tracing::warn!( | 221 | + let addr: SocketAddr = format!("{host}:{port}") |
| 228 | - "daemon backend manager shutdown timed out after 10s; \ | 222 | + .parse() |
| 223 | + .with_context(|| format!("invalid listen address {host}:{port}"))?; | ||
| 224 | + let listener = tokio::net::TcpListener::bind(addr) | ||
| 225 | + .await | ||
| 226 | + .with_context(|| format!("failed to bind {addr}"))?; | ||
| 227 | + tracing::info!(config = %config_path.display(), %addr, "starting rebuild daemon"); | ||
| 228 | + let serve_result = axum::serve(listener, router) | ||
| 229 | + .with_graceful_shutdown(shutdown_signal()) | ||
| 230 | + .await | ||
| 231 | + .context("axum server exited unexpectedly"); | ||
| 232 | + | ||
| 233 | + // Gracefully shutdown cron scheduler | ||
| 234 | + if let Some(scheduler) = cron_scheduler { | ||
| 235 | + scheduler.stop().await; | ||
| 236 | + } | ||
| 237 | + | ||
| 238 | + // Best-effort sandbox cleanup with a bounded timeout so a slow/stuck | ||
| 239 | + // provider delete call cannot keep the daemon alive indefinitely after a | ||
| 240 | + // shutdown signal. Matches the TUI exit path which also bounds remote | ||
| 241 | + // close to a few seconds. | ||
| 242 | + let shutdown_result = tokio::time::timeout( | ||
| 243 | + std::time::Duration::from_secs(10), | ||
| 244 | + backend_manager.shutdown_all(), | ||
| 245 | + ) | ||
| 246 | + .await; | ||
| 247 | + match shutdown_result { | ||
| 248 | + Ok(Ok(())) => {} | ||
| 249 | + Ok(Err(error)) => { | ||
| 250 | + tracing::warn!(error = %error, "failed to shutdown daemon backend manager"); | ||
| 251 | + } | ||
| 252 | + Err(_) => { | ||
| 253 | + tracing::warn!( | ||
| 254 | + "daemon backend manager shutdown timed out after 10s; \ | ||
| 229 | some sandboxes may linger and will be reclaimed lazily" | 255 | some sandboxes may linger and will be reclaimed lazily" |
| 230 | - ); | 256 | + ); |
| 257 | + } | ||
| 258 | + } | ||
| 259 | + serve_result | ||
| 260 | + } | ||
| 261 | + .await; | ||
| 262 | + if let Some(automation) = memory_automation_for_shutdown { | ||
| 263 | + match tokio::time::timeout(std::time::Duration::from_secs(5), automation.close()).await { | ||
| 264 | + Ok(Ok(())) => {} | ||
| 265 | + Ok(Err(error)) => { | ||
| 266 | + tracing::warn!(error = %error, "failed to close daemon MCP memory automation"); | ||
| 267 | + } | ||
| 268 | + Err(_) => { | ||
| 269 | + tracing::warn!("daemon MCP memory automation close timed out after 5 seconds"); | ||
| 270 | + } | ||
| 231 | } | 271 | } |
| 232 | } | 272 | } |
| 233 | serve_result | 273 | serve_result |
| @@ -421,6 +461,7 @@ fn init_tracing() { | |||
| 421 | 461 | ||
| 422 | struct Cli { | 462 | struct Cli { |
| 423 | config: Option<PathBuf>, | 463 | config: Option<PathBuf>, |
| 464 | + mcp_config: Option<PathBuf>, | ||
| 424 | host: String, | 465 | host: String, |
| 425 | port: u16, | 466 | port: u16, |
| 426 | dashboard_host: Option<String>, | 467 | dashboard_host: Option<String>, |
| @@ -434,6 +475,7 @@ impl Cli { | |||
| 434 | I: IntoIterator<Item = String>, | 475 | I: IntoIterator<Item = String>, |
| 435 | { | 476 | { |
| 436 | let mut config = None; | 477 | let mut config = None; |
| 478 | + let mut mcp_config = None; | ||
| 437 | let mut host = "0.0.0.0".to_string(); | 479 | let mut host = "0.0.0.0".to_string(); |
| 438 | let mut port = 18080_u16; | 480 | let mut port = 18080_u16; |
| 439 | let mut dashboard_host: Option<String> = None; | 481 | let mut dashboard_host: Option<String> = None; |
| @@ -445,6 +487,7 @@ impl Cli { | |||
| 445 | "--help" | "-h" => { | 487 | "--help" | "-h" => { |
| 446 | return Ok(Self { | 488 | return Ok(Self { |
| 447 | config, | 489 | config, |
| 490 | + mcp_config, | ||
| 448 | host, | 491 | host, |
| 449 | port, | 492 | port, |
| 450 | dashboard_host, | 493 | dashboard_host, |
| @@ -457,6 +500,13 @@ impl Cli { | |||
| 457 | let value = remaining.get(index).context("missing value for --config")?; | 500 | let value = remaining.get(index).context("missing value for --config")?; |
| 458 | config = Some(PathBuf::from(value)); | 501 | config = Some(PathBuf::from(value)); |
| 459 | } | 502 | } |
| 503 | + "--mcp-config" => { | ||
| 504 | + index += 1; | ||
| 505 | + let value = remaining | ||
| 506 | + .get(index) | ||
| 507 | + .context("missing value for --mcp-config")?; | ||
| 508 | + mcp_config = Some(PathBuf::from(value)); | ||
| 509 | + } | ||
| 460 | "--host" => { | 510 | "--host" => { |
| 461 | index += 1; | 511 | index += 1; |
| 462 | let value = remaining.get(index).context("missing value for --host")?; | 512 | let value = remaining.get(index).context("missing value for --host")?; |
| @@ -493,6 +543,7 @@ impl Cli { | |||
| 493 | } | 543 | } |
| 494 | Ok(Self { | 544 | Ok(Self { |
| 495 | config, | 545 | config, |
| 546 | + mcp_config, | ||
| 496 | host, | 547 | host, |
| 497 | port, | 548 | port, |
| 498 | dashboard_host, | 549 | dashboard_host, |
| @@ -504,7 +555,7 @@ impl Cli { | |||
| 504 | 555 | ||
| 505 | fn print_usage() { | 556 | fn print_usage() { |
| 506 | eprintln!( | 557 | eprintln!( |
| 507 | - "Usage: xiaoo-daemon [--config <path>] [--host <host>] [--port <port>]\n\ | 558 | + "Usage: xiaoo-daemon [--config <path>] [--mcp-config <path>] [--host <host>] [--port <port>]\n\ |
| 508 | \x20 [--dashboard-host <host>] [--dashboard-port <port>]\n\n\ | 559 | \x20 [--dashboard-host <host>] [--dashboard-port <port>]\n\n\ |
| 509 | Defaults: --host 0.0.0.0 --port 18080\n\ | 560 | Defaults: --host 0.0.0.0 --port 18080\n\ |
| 510 | \x20 --dashboard-host 127.0.0.1 --dashboard-port 28081\n\n\ | 561 | \x20 --dashboard-host 127.0.0.1 --dashboard-port 28081\n\n\ |
| @@ -539,6 +590,18 @@ mod tests { | |||
| 539 | assert!(!cli.help); | 590 | assert!(!cli.help); |
| 540 | } | 591 | } |
| 541 | 592 | ||
| 593 | + | ||
| 594 | + fn parses_daemon_mcp_config_argument() { | ||
| 595 | + let cli = Cli::parse( | ||
| 596 | + ["--mcp-config", "/tmp/mcp.json"] | ||
| 597 | + .into_iter() | ||
| 598 | + .map(str::to_string), | ||
| 599 | + ) | ||
| 600 | + .expect("daemon should accept --mcp-config"); | ||
| 601 | + | ||
| 602 | + assert_eq!(cli.mcp_config, Some(PathBuf::from("/tmp/mcp.json"))); | ||
| 603 | + } | ||
| 604 | + | ||
| 542 | 605 | ||
| 543 | fn daemon_defaults_to_port_18080() { | 606 | fn daemon_defaults_to_port_18080() { |
| 544 | let cli = Cli::parse(std::iter::empty::<String>()).expect("cli should parse with defaults"); | 607 | let cli = Cli::parse(std::iter::empty::<String>()).expect("cli should parse with defaults"); |
| @@ -6,6 +6,9 @@ use std::path::PathBuf; | |||
| 6 | 6 | ||
| 7 | pub fn build_base_url_candidates(original_base: &str) -> Vec<String> { | 7 | pub fn build_base_url_candidates(original_base: &str) -> Vec<String> { |
| 8 | let base = original_base.trim_end_matches('/'); | 8 | let base = original_base.trim_end_matches('/'); |
| 9 | + if is_chat_completions_endpoint(base) { | ||
| 10 | + return vec![base.to_string()]; | ||
| 11 | + } | ||
| 9 | let mut candidates = Vec::new(); | 12 | let mut candidates = Vec::new(); |
| 10 | 13 | ||
| 11 | // B1: 用户原始配置(最高优先级) | 14 | // B1: 用户原始配置(最高优先级) |
| @@ -50,6 +53,9 @@ pub fn build_final_candidates(base_candidates: &[String]) -> Vec<String> { | |||
| 50 | } | 53 | } |
| 51 | 54 | ||
| 52 | for base in base_candidates { | 55 | for base in base_candidates { |
| 56 | + if is_chat_completions_endpoint(base) { | ||
| 57 | + continue; | ||
| 58 | + } | ||
| 53 | let paths = build_endpoint_paths(base); | 59 | let paths = build_endpoint_paths(base); |
| 54 | for path in paths { | 60 | for path in paths { |
| 55 | let url = format!("{}{}", base.trim_end_matches('/'), path); | 61 | let url = format!("{}{}", base.trim_end_matches('/'), path); |
| @@ -65,6 +71,14 @@ pub fn build_final_candidates(base_candidates: &[String]) -> Vec<String> { | |||
| 65 | final_urls | 71 | final_urls |
| 66 | } | 72 | } |
| 67 | 73 | ||
| 74 | +fn is_chat_completions_endpoint(base: &str) -> bool { | ||
| 75 | + base.split(['?', '#']) | ||
| 76 | + .next() | ||
| 77 | + .unwrap_or(base) | ||
| 78 | + .trim_end_matches('/') | ||
| 79 | + .ends_with("/chat/completions") | ||
| 80 | +} | ||
| 81 | + | ||
| 68 | fn has_version_path(base: &str) -> bool { | 82 | fn has_version_path(base: &str) -> bool { |
| 69 | base.ends_with("/v1") | 83 | base.ends_with("/v1") |
| 70 | || base.ends_with("/v4") | 84 | || base.ends_with("/v4") |
| @@ -451,6 +465,22 @@ mod tests { | |||
| 451 | assert_eq!(final_urls[1], "http://example.com/v1/chat/completions"); | 465 | assert_eq!(final_urls[1], "http://example.com/v1/chat/completions"); |
| 452 | } | 466 | } |
| 453 | 467 | ||
| 468 | + | ||
| 469 | + fn full_chat_completions_endpoint_is_not_extended_again() { | ||
| 470 | + for endpoint in [ | ||
| 471 | + "https://api.example.com/v1/chat/completions", | ||
| 472 | + "https://api.example.com/v1/chat/completions/", | ||
| 473 | + "https://api.example.com/v1/chat/completions?api-version=2026-01-01", | ||
| 474 | + ] { | ||
| 475 | + let expected = endpoint.trim_end_matches('/'); | ||
| 476 | + let bases = build_base_url_candidates(endpoint); | ||
| 477 | + assert_eq!(bases, vec![expected]); | ||
| 478 | + | ||
| 479 | + let final_urls = build_final_candidates(&bases); | ||
| 480 | + assert_eq!(final_urls, vec![expected]); | ||
| 481 | + } | ||
| 482 | + } | ||
| 483 | + | ||
| 454 | 484 | ||
| 455 | fn test_endpoint_path_error_detection() { | 485 | fn test_endpoint_path_error_detection() { |
| 456 | let http_error = LlmError::HttpError("Connection failed".to_string()); | 486 | let http_error = LlmError::HttpError("Connection failed".to_string()); |
| @@ -16,4 +16,7 @@ futures-util.workspace = true | |||
| 16 | url.workspace = true | 16 | url.workspace = true |
| 17 | 17 | ||
| 18 | [dev-dependencies] | 18 | [dev-dependencies] |
| 19 | +axum.workspace = true | ||
| 20 | +tempfile.workspace = true | ||
| 19 | toml.workspace = true | 21 | toml.workspace = true |
| 22 | +tokio = { workspace = true, features = ["net", "rt-multi-thread"] } | ||
| @@ -4,7 +4,7 @@ use std::sync::Arc; | |||
| 4 | 4 | ||
| 5 | use crate::config::{McpServerConfig, Transport}; | 5 | use crate::config::{McpServerConfig, Transport}; |
| 6 | use crate::error::McpError; | 6 | use crate::error::McpError; |
| 7 | -use crate::transport::{McpTransport, SseTransport, StdioTransport}; | 7 | +use crate::transport::{McpTransport, SseTransport, StdioTransport, StreamableHttpTransport}; |
| 8 | use crate::types::{ | 8 | use crate::types::{ |
| 9 | CallToolParams, CallToolResult, ClientInfo, ContentBlock, InitializeParams, InitializeResult, | 9 | CallToolParams, CallToolResult, ClientInfo, ContentBlock, InitializeParams, InitializeResult, |
| 10 | ListToolsResult, McpToolDef, | 10 | ListToolsResult, McpToolDef, |
| @@ -15,6 +15,7 @@ use crate::types::{ | |||
| 15 | pub struct McpCallResult { | 15 | pub struct McpCallResult { |
| 16 | pub content: Vec<ContentBlock>, | 16 | pub content: Vec<ContentBlock>, |
| 17 | pub is_error: bool, | 17 | pub is_error: bool, |
| 18 | + pub structured_content: Option<Value>, | ||
| 18 | } | 19 | } |
| 19 | 20 | ||
| 20 | /// A connected, initialised MCP client. Cloning shares the underlying | 21 | /// A connected, initialised MCP client. Cloning shares the underlying |
| @@ -46,6 +47,7 @@ impl McpClient { | |||
| 46 | })?; | 47 | })?; |
| 47 | Arc::new(SseTransport::connect(&url, config.timeout_ms).await?) | 48 | Arc::new(SseTransport::connect(&url, config.timeout_ms).await?) |
| 48 | } | 49 | } |
| 50 | + Transport::StreamableHttp => Arc::new(StreamableHttpTransport::connect(config).await?), | ||
| 49 | }; | 51 | }; |
| 50 | 52 | ||
| 51 | Ok(Self { | 53 | Ok(Self { |
| @@ -66,7 +68,7 @@ impl McpClient { | |||
| 66 | /// Perform the MCP `initialize` handshake. | 68 | /// Perform the MCP `initialize` handshake. |
| 67 | pub async fn initialize(&self) -> Result<InitializeResult, McpError> { | 69 | pub async fn initialize(&self) -> Result<InitializeResult, McpError> { |
| 68 | let params = InitializeParams { | 70 | let params = InitializeParams { |
| 69 | - protocol_version: "2024-11-05".to_string(), | 71 | + protocol_version: self.transport.initialize_protocol_version().to_string(), |
| 70 | capabilities: serde_json::json!({}), | 72 | capabilities: serde_json::json!({}), |
| 71 | client_info: ClientInfo { | 73 | client_info: ClientInfo { |
| 72 | name: "xiaoo".to_string(), | 74 | name: "xiaoo".to_string(), |
| @@ -81,6 +83,11 @@ impl McpClient { | |||
| 81 | .await?; | 83 | .await?; |
| 82 | let init: InitializeResult = | 84 | let init: InitializeResult = |
| 83 | serde_json::from_value(result).map_err(|e| McpError::HandshakeFailed(e.to_string()))?; | 85 | serde_json::from_value(result).map_err(|e| McpError::HandshakeFailed(e.to_string()))?; |
| 86 | + self.transport | ||
| 87 | + .validate_negotiated_protocol_version(&init.protocol_version)?; | ||
| 88 | + self.transport | ||
| 89 | + .set_protocol_version(&init.protocol_version) | ||
| 90 | + .await; | ||
| 84 | // Notify the server that initialisation is complete. | 91 | // Notify the server that initialisation is complete. |
| 85 | self.transport | 92 | self.transport |
| 86 | .send_notification("notifications/initialized", None) | 93 | .send_notification("notifications/initialized", None) |
| @@ -88,6 +95,13 @@ impl McpClient { | |||
| 88 | Ok(init) | 95 | Ok(init) |
| 89 | } | 96 | } |
| 90 | 97 | ||
| 98 | + /// Release the connected transport. For Streamable HTTP this sends the | ||
| 99 | + /// MCP session-termination DELETE request; older transports retain their | ||
| 100 | + /// existing no-op close behaviour. | ||
| 101 | + pub async fn close(&self) -> Result<(), McpError> { | ||
| 102 | + self.transport.close().await | ||
| 103 | + } | ||
| 104 | + | ||
| 91 | /// List all tools exposed by the server, following pagination cursors. | 105 | /// List all tools exposed by the server, following pagination cursors. |
| 92 | pub async fn list_tools(&self) -> Result<Vec<McpToolDef>, McpError> { | 106 | pub async fn list_tools(&self) -> Result<Vec<McpToolDef>, McpError> { |
| 93 | let mut tools = Vec::new(); | 107 | let mut tools = Vec::new(); |
| @@ -129,6 +143,7 @@ impl McpClient { | |||
| 129 | Ok(McpCallResult { | 143 | Ok(McpCallResult { |
| 130 | content: call.content, | 144 | content: call.content, |
| 131 | is_error: call.is_error, | 145 | is_error: call.is_error, |
| 146 | + structured_content: call.structured_content, | ||
| 132 | }) | 147 | }) |
| 133 | } | 148 | } |
| 134 | } | 149 | } |
| @@ -158,13 +173,18 @@ impl McpCallResult { | |||
| 158 | } | 173 | } |
| 159 | out.push_str(&segment); | 174 | out.push_str(&segment); |
| 160 | } | 175 | } |
| 176 | + if out.is_empty() { | ||
| 177 | + if let Some(structured_content) = &self.structured_content { | ||
| 178 | + return serde_json::to_string(structured_content).unwrap_or_default(); | ||
| 179 | + } | ||
| 180 | + } | ||
| 161 | out | 181 | out |
| 162 | } | 182 | } |
| 163 | } | 183 | } |
| 164 | 184 | ||
| 165 | fn base64_decoded_len(s: &str) -> usize { | 185 | fn base64_decoded_len(s: &str) -> usize { |
| 166 | let len = s.len(); | 186 | let len = s.len(); |
| 167 | - if len == 0 || len % 4 != 0 { | 187 | + if len == 0 || !len.is_multiple_of(4) { |
| 168 | return len; | 188 | return len; |
| 169 | } | 189 | } |
| 170 | let padding = s.bytes().filter(|b| *b == b'=').count(); | 190 | let padding = s.bytes().filter(|b| *b == b'=').count(); |
| @@ -1,5 +1,5 @@ | |||
| 1 | use serde::{Deserialize, Serialize}; | 1 | use serde::{Deserialize, Serialize}; |
| 2 | -use std::collections::HashMap; | 2 | +use std::collections::{BTreeMap, HashMap}; |
| 3 | 3 | ||
| 4 | /// Top-level `[mcp]` configuration section. | 4 | /// Top-level `[mcp]` configuration section. |
| 5 | 5 | ||
| @@ -30,6 +30,19 @@ pub struct McpServerConfig { | |||
| 30 | 30 | ||
| 31 | pub url: Option<String>, | 31 | pub url: Option<String>, |
| 32 | 32 | ||
| 33 | + /// Name of the environment variable containing the bearer token. The | ||
| 34 | + /// secret itself is resolved by the HTTP transport at runtime. | ||
| 35 | + | ||
| 36 | + pub bearer_token_env: Option<String>, | ||
| 37 | + | ||
| 38 | + /// Optional agent selector sent by Streamable HTTP transports. | ||
| 39 | + | ||
| 40 | + pub agent_id: Option<String>, | ||
| 41 | + | ||
| 42 | + /// Non-sensitive, fixed headers for HTTP transports. | ||
| 43 | + | ||
| 44 | + pub headers: BTreeMap<String, String>, | ||
| 45 | + | ||
| 33 | /// Override the enabled flag (defaults to true). | 46 | /// Override the enabled flag (defaults to true). |
| 34 | 47 | ||
| 35 | pub enabled: Option<bool>, | 48 | pub enabled: Option<bool>, |
| @@ -52,10 +65,48 @@ impl McpServerConfig { | |||
| 52 | } | 65 | } |
| 53 | } | 66 | } |
| 54 | 67 | ||
| 55 | -fn default_timeout_ms() -> u64 { | 68 | +pub(crate) fn default_timeout_ms() -> u64 { |
| 56 | 30_000 | 69 | 30_000 |
| 57 | } | 70 | } |
| 58 | 71 | ||
| 72 | +pub(crate) fn validate_fixed_headers(headers: &BTreeMap<String, String>) -> Result<(), String> { | ||
| 73 | + for (name, value) in headers { | ||
| 74 | + let parsed_name = reqwest::header::HeaderName::from_bytes(name.as_bytes()) | ||
| 75 | + .map_err(|_| format!("invalid header name `{name}`"))?; | ||
| 76 | + let normalized = parsed_name.as_str(); | ||
| 77 | + if is_sensitive_header_name(normalized) { | ||
| 78 | + return Err(format!( | ||
| 79 | + "sensitive header `{name}` is not allowed; use bearer_token_env" | ||
| 80 | + )); | ||
| 81 | + } | ||
| 82 | + if matches!( | ||
| 83 | + normalized, | ||
| 84 | + "origin" | ||
| 85 | + | "mcp-session-id" | ||
| 86 | + | "mcp-protocol-version" | ||
| 87 | + | "accept" | ||
| 88 | + | "content-type" | ||
| 89 | + | "x-agent-id" | ||
| 90 | + | "last-event-id" | ||
| 91 | + | "mcp-method" | ||
| 92 | + | "mcp-name" | ||
| 93 | + ) { | ||
| 94 | + return Err(format!("transport-managed header `{name}` is not allowed")); | ||
| 95 | + } | ||
| 96 | + reqwest::header::HeaderValue::from_str(value) | ||
| 97 | + .map_err(|_| format!("invalid value for header `{name}`"))?; | ||
| 98 | + } | ||
| 99 | + Ok(()) | ||
| 100 | +} | ||
| 101 | + | ||
| 102 | +fn is_sensitive_header_name(name: &str) -> bool { | ||
| 103 | + matches!( | ||
| 104 | + name, | ||
| 105 | + "authorization" | "proxy-authorization" | "cookie" | "set-cookie" | ||
| 106 | + ) || name.contains("token") | ||
| 107 | + || name.contains("api-key") | ||
| 108 | +} | ||
| 109 | + | ||
| 59 | /// Effect profile for an MCP server's tools. The MCP protocol does not expose | 110 | /// Effect profile for an MCP server's tools. The MCP protocol does not expose |
| 60 | /// effect metadata in tool definitions, so this is a per-server declaration. | 111 | /// effect metadata in tool definitions, so this is a per-server declaration. |
| 61 | /// Defaults to the most conservative assumption (all effects present); users | 112 | /// Defaults to the most conservative assumption (all effects present); users |
| @@ -96,17 +147,13 @@ fn default_true() -> bool { | |||
| 96 | true | 147 | true |
| 97 | } | 148 | } |
| 98 | 149 | ||
| 99 | -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] | 150 | +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] |
| 100 | -#[serde(rename_all = "lowercase")] | 151 | +#[serde(rename_all = "snake_case")] |
| 101 | pub enum Transport { | 152 | pub enum Transport { |
| 153 | + | ||
| 102 | Stdio, | 154 | Stdio, |
| 103 | Sse, | 155 | Sse, |
| 104 | -} | 156 | + StreamableHttp, |
| 105 | - | ||
| 106 | -impl Default for Transport { | ||
| 107 | - fn default() -> Self { | ||
| 108 | - Self::Stdio | ||
| 109 | - } | ||
| 110 | } | 157 | } |
| 111 | 158 | ||
| 112 | 159 | ||
| @@ -1,6 +1,6 @@ | |||
| 1 | use thiserror::Error; | 1 | use thiserror::Error; |
| 2 | 2 | ||
| 3 | -#[derive(Debug, Error)] | 3 | +#[derive(Clone, Debug, Error)] |
| 4 | pub enum McpError { | 4 | pub enum McpError { |
| 5 | 5 | ||
| 6 | SpawnFailed { command: String, error: String }, | 6 | SpawnFailed { command: String, error: String }, |
| @@ -26,6 +26,9 @@ pub enum McpError { | |||
| 26 | 26 | ||
| 27 | Http(String), | 27 | Http(String), |
| 28 | 28 | ||
| 29 | + | ||
| 30 | + BearerTokenUnavailable { env_var: String }, | ||
| 31 | + | ||
| 29 | 32 | ||
| 30 | Disconnected, | 33 | Disconnected, |
| 31 | } | 34 | } |