mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Add auth mode, fix MCP token handling, and parallelize startup loading
Auth mode: when a tool requires an API key, the thread enters a special mode where the next user message is routed directly to the credential store, bypassing logs, turns, history, and compaction entirely. This prevents tokens from leaking into debug output or persistent storage. Fix MCP auth: auth_mcp now actually uses the token parameter (was ignored as _token) and falls back to manual token entry when OAuth and DCR are both unsupported. Parallel loading: WASM tools, WASM channels, and MCP servers now load concurrently at startup. Within each loader, individual items also load in parallel (join_all for WASM, JoinSet for MCP servers). Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
9b729795fb
commit
2cdd04a359
+233
-12
@@ -42,6 +42,7 @@ use crate::config::{AgentConfig, HeartbeatConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::context::JobContext;
|
||||
use crate::error::Error;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult};
|
||||
use crate::safety::SafetyLayer;
|
||||
@@ -68,6 +69,7 @@ pub struct AgentDeps {
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub workspace: Option<Arc<Workspace>>,
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
}
|
||||
|
||||
/// The main agent that coordinates all components.
|
||||
@@ -374,13 +376,6 @@ impl Agent {
|
||||
}
|
||||
|
||||
async fn handle_message(&self, message: &IncomingMessage) -> Result<Option<String>, Error> {
|
||||
tracing::debug!(
|
||||
"Received message from {} on {}: {}",
|
||||
message.user_id,
|
||||
message.channel,
|
||||
truncate(&message.content, 100)
|
||||
);
|
||||
|
||||
// Parse submission type first
|
||||
let submission = SubmissionParser::parse(&message.content);
|
||||
|
||||
@@ -394,6 +389,41 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
// Auth mode interception: if the thread is awaiting a token, route
|
||||
// the message directly to the credential store. Nothing touches
|
||||
// logs, turns, history, or compaction.
|
||||
let pending_auth = {
|
||||
let sess = session.lock().await;
|
||||
sess.threads
|
||||
.get(&thread_id)
|
||||
.and_then(|t| t.pending_auth.clone())
|
||||
};
|
||||
|
||||
if let Some(pending) = pending_auth {
|
||||
match &submission {
|
||||
Submission::UserInput { content } => {
|
||||
return self
|
||||
.process_auth_token(message, &pending, content, session, thread_id)
|
||||
.await;
|
||||
}
|
||||
_ => {
|
||||
// Any control submission (interrupt, undo, etc.) cancels auth mode
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
// Fall through to normal handling
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"Received message from {} on {} ({} chars)",
|
||||
message.user_id,
|
||||
message.channel,
|
||||
message.content.len()
|
||||
);
|
||||
|
||||
// Process based on submission type
|
||||
let result = match submission {
|
||||
Submission::UserInput { content } => {
|
||||
@@ -867,6 +897,19 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// If tool_auth returned awaiting_token, enter auth mode
|
||||
// and short-circuit: return the instructions directly so
|
||||
// the LLM doesn't get a chance to hallucinate tool calls.
|
||||
if let Some((ext_name, instructions)) =
|
||||
detect_auth_awaiting(&tc.name, &tool_result)
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name);
|
||||
}
|
||||
return Ok(AgenticLoopResult::Response(instructions));
|
||||
}
|
||||
|
||||
// Add tool result to context for next LLM call
|
||||
let result_content = match tool_result {
|
||||
Ok(output) => {
|
||||
@@ -1236,6 +1279,29 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// If tool_auth returned awaiting_token, enter auth mode and
|
||||
// return instructions directly (skip agentic loop continuation).
|
||||
if let Some((ext_name, instructions)) =
|
||||
detect_auth_awaiting(&pending.tool_name, &tool_result)
|
||||
{
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name);
|
||||
thread.complete_turn(&instructions);
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting token".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
return Ok(SubmissionResult::response(instructions));
|
||||
}
|
||||
|
||||
// Add tool result to context
|
||||
let result_content = match tool_result {
|
||||
Ok(output) => {
|
||||
@@ -1336,6 +1402,75 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle an auth token submitted while the thread is in auth mode.
|
||||
///
|
||||
/// The token goes directly to the extension manager's credential store,
|
||||
/// completely bypassing logging, turn creation, history, and compaction.
|
||||
async fn process_auth_token(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
pending: &crate::agent::session::PendingAuth,
|
||||
token: &str,
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
) -> Result<Option<String>, Error> {
|
||||
let token = token.trim();
|
||||
|
||||
// Clear auth mode regardless of outcome
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
}
|
||||
|
||||
let ext_mgr = match self.deps.extension_manager.as_ref() {
|
||||
Some(mgr) => mgr,
|
||||
None => return Ok(Some("Extension manager not available.".to_string())),
|
||||
};
|
||||
|
||||
match ext_mgr.auth(&pending.extension_name, Some(token)).await {
|
||||
Ok(result) if result.status == "authenticated" => {
|
||||
tracing::info!(
|
||||
"Extension '{}' authenticated via auth mode",
|
||||
pending.extension_name
|
||||
);
|
||||
|
||||
// Notify via channel status so the response doesn't echo the token
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Authenticated".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(Some(format!(
|
||||
"{} authenticated successfully.",
|
||||
pending.extension_name
|
||||
)))
|
||||
}
|
||||
Ok(result) => {
|
||||
// Unexpected state, re-enter auth mode
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(pending.extension_name.clone());
|
||||
}
|
||||
}
|
||||
let msg = result
|
||||
.instructions
|
||||
.unwrap_or_else(|| "Invalid token. Please try again.".to_string());
|
||||
Ok(Some(msg))
|
||||
}
|
||||
Err(e) => Ok(Some(format!(
|
||||
"Authentication failed for {}: {}",
|
||||
pending.extension_name, e
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn process_new_thread(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
@@ -1737,10 +1872,96 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate(s: &str, max_len: usize) -> String {
|
||||
if s.len() <= max_len {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}...", &s[..max_len])
|
||||
/// Check if a tool_auth result indicates the extension is awaiting a token.
|
||||
///
|
||||
/// Returns `Some((extension_name, instructions))` if the tool result contains
|
||||
/// `awaiting_token: true`, meaning the thread should enter auth mode.
|
||||
fn detect_auth_awaiting(
|
||||
tool_name: &str,
|
||||
result: &Result<String, Error>,
|
||||
) -> Option<(String, String)> {
|
||||
if tool_name != "tool_auth" {
|
||||
return None;
|
||||
}
|
||||
let output = result.as_ref().ok()?;
|
||||
let parsed: serde_json::Value = serde_json::from_str(output).ok()?;
|
||||
if parsed.get("awaiting_token") != Some(&serde_json::Value::Bool(true)) {
|
||||
return None;
|
||||
}
|
||||
let name = parsed.get("name")?.as_str()?.to_string();
|
||||
let instructions = parsed
|
||||
.get("instructions")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("Please provide your API token/key.")
|
||||
.to_string();
|
||||
Some((name, instructions))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::error::Error;
|
||||
|
||||
use super::detect_auth_awaiting;
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_positive() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "telegram",
|
||||
"kind": "WasmTool",
|
||||
"awaiting_token": true,
|
||||
"status": "awaiting_token",
|
||||
"instructions": "Please provide your Telegram Bot API token."
|
||||
})
|
||||
.to_string());
|
||||
|
||||
let detected = detect_auth_awaiting("tool_auth", &result);
|
||||
assert!(detected.is_some());
|
||||
let (name, instructions) = detected.unwrap();
|
||||
assert_eq!(name, "telegram");
|
||||
assert!(instructions.contains("Telegram Bot API"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_not_awaiting() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "telegram",
|
||||
"kind": "WasmTool",
|
||||
"awaiting_token": false,
|
||||
"status": "authenticated"
|
||||
})
|
||||
.to_string());
|
||||
|
||||
assert!(detect_auth_awaiting("tool_auth", &result).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_wrong_tool() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "telegram",
|
||||
"awaiting_token": true,
|
||||
})
|
||||
.to_string());
|
||||
|
||||
assert!(detect_auth_awaiting("tool_list", &result).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_error_result() {
|
||||
let result: Result<String, Error> =
|
||||
Err(crate::error::ToolError::NotFound { name: "x".into() }.into());
|
||||
assert!(detect_auth_awaiting("tool_auth", &result).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_default_instructions() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "custom_tool",
|
||||
"awaiting_token": true,
|
||||
"status": "awaiting_token"
|
||||
})
|
||||
.to_string());
|
||||
|
||||
let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap();
|
||||
assert_eq!(instructions, "Please provide your API token/key.");
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -30,7 +30,7 @@ pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_hea
|
||||
pub use router::{MessageIntent, Router};
|
||||
pub use scheduler::Scheduler;
|
||||
pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob};
|
||||
pub use session::{PendingApproval, Session, Thread, ThreadState, Turn, TurnState};
|
||||
pub use session::{PendingApproval, PendingAuth, Session, Thread, ThreadState, Turn, TurnState};
|
||||
pub use session_manager::SessionManager;
|
||||
pub use submission::{Submission, SubmissionParser, SubmissionResult};
|
||||
pub use task::{Task, TaskContext, TaskHandler, TaskOutput, TaskStatus};
|
||||
|
||||
@@ -121,6 +121,18 @@ pub enum ThreadState {
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
/// Pending auth token request.
|
||||
///
|
||||
/// When `tool_auth` returns `awaiting_token`, the thread enters auth mode.
|
||||
/// The next user message is intercepted before entering the normal pipeline
|
||||
/// (no logging, no turn creation, no history) and routed directly to the
|
||||
/// credential store.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PendingAuth {
|
||||
/// Extension name to authenticate.
|
||||
pub extension_name: String,
|
||||
}
|
||||
|
||||
/// Pending tool approval request stored on a thread.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PendingApproval {
|
||||
@@ -158,6 +170,9 @@ pub struct Thread {
|
||||
/// Pending approval request (when state is AwaitingApproval).
|
||||
#[serde(default)]
|
||||
pub pending_approval: Option<PendingApproval>,
|
||||
/// Pending auth token request (thread is in auth mode).
|
||||
#[serde(default)]
|
||||
pub pending_auth: Option<PendingAuth>,
|
||||
}
|
||||
|
||||
impl Thread {
|
||||
@@ -173,6 +188,7 @@ impl Thread {
|
||||
updated_at: now,
|
||||
metadata: serde_json::Value::Null,
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -238,6 +254,18 @@ impl Thread {
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
/// Enter auth mode: next user message will be routed directly to
|
||||
/// the credential store, bypassing the normal pipeline entirely.
|
||||
pub fn enter_auth_mode(&mut self, extension_name: String) {
|
||||
self.pending_auth = Some(PendingAuth { extension_name });
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
/// Take the pending auth (clearing auth mode).
|
||||
pub fn take_pending_auth(&mut self) -> Option<PendingAuth> {
|
||||
self.pending_auth.take()
|
||||
}
|
||||
|
||||
/// Interrupt the current turn.
|
||||
pub fn interrupt(&mut self) {
|
||||
if let Some(turn) = self.turns.last_mut() {
|
||||
@@ -511,4 +539,58 @@ mod tests {
|
||||
assert_eq!(thread.turns[1].user_input, "How are you?");
|
||||
assert!(thread.turns[1].response.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enter_auth_mode() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
assert!(thread.pending_auth.is_none());
|
||||
|
||||
thread.enter_auth_mode("telegram".to_string());
|
||||
assert!(thread.pending_auth.is_some());
|
||||
assert_eq!(
|
||||
thread.pending_auth.as_ref().unwrap().extension_name,
|
||||
"telegram"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_take_pending_auth() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
thread.enter_auth_mode("notion".to_string());
|
||||
|
||||
let pending = thread.take_pending_auth();
|
||||
assert!(pending.is_some());
|
||||
assert_eq!(pending.unwrap().extension_name, "notion");
|
||||
|
||||
// Should be cleared after take
|
||||
assert!(thread.pending_auth.is_none());
|
||||
assert!(thread.take_pending_auth().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_auth_serialization() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
thread.enter_auth_mode("openai".to_string());
|
||||
|
||||
let json = serde_json::to_string(&thread).expect("should serialize");
|
||||
assert!(json.contains("pending_auth"));
|
||||
assert!(json.contains("openai"));
|
||||
|
||||
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
|
||||
assert!(restored.pending_auth.is_some());
|
||||
assert_eq!(restored.pending_auth.unwrap().extension_name, "openai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_auth_default_none() {
|
||||
// Deserialization of old data without pending_auth should default to None
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
thread.pending_auth = None;
|
||||
let json = serde_json::to_string(&thread).expect("serialize");
|
||||
|
||||
// Remove the pending_auth field to simulate old data
|
||||
let json = json.replace(",\"pending_auth\":null", "");
|
||||
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
|
||||
assert!(restored.pending_auth.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,17 +151,18 @@ impl WasmChannelLoader {
|
||||
}
|
||||
|
||||
let mut results = LoadResults::default();
|
||||
|
||||
// Collect all .wasm entries first, then load in parallel
|
||||
let mut channel_entries = Vec::new();
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
|
||||
// Only process .wasm files
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("wasm") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extract channel name from filename
|
||||
let name = match path.file_stem().and_then(|s| s.to_str()) {
|
||||
Some(n) => n.to_string(),
|
||||
None => {
|
||||
@@ -173,15 +174,20 @@ impl WasmChannelLoader {
|
||||
}
|
||||
};
|
||||
|
||||
// Look for sidecar capabilities file
|
||||
let cap_path = path.with_extension("capabilities.json");
|
||||
let cap_path_option = if cap_path.exists() {
|
||||
Some(cap_path.as_path())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let has_cap = cap_path.exists();
|
||||
channel_entries.push((name, path, if has_cap { Some(cap_path) } else { None }));
|
||||
}
|
||||
|
||||
match self.load_from_files(&name, &path, cap_path_option).await {
|
||||
// Load all channels in parallel (file I/O + WASM compilation)
|
||||
let load_futures = channel_entries
|
||||
.iter()
|
||||
.map(|(name, path, cap_path)| self.load_from_files(name, path, cap_path.as_deref()));
|
||||
|
||||
let load_results = futures::future::join_all(load_futures).await;
|
||||
|
||||
for ((name, path, _), result) in channel_entries.into_iter().zip(load_results) {
|
||||
match result {
|
||||
Ok(loaded) => {
|
||||
results.loaded.push(loaded);
|
||||
}
|
||||
|
||||
@@ -460,12 +460,35 @@ impl ExtensionManager {
|
||||
async fn auth_mcp(
|
||||
&self,
|
||||
name: &str,
|
||||
_token: Option<&str>,
|
||||
token: Option<&str>,
|
||||
) -> Result<AuthResult, ExtensionError> {
|
||||
let server = get_mcp_server(name)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
|
||||
// If a token was provided directly, store it and we're done.
|
||||
if let Some(token_value) = token {
|
||||
let secret_name = server.token_secret_name();
|
||||
let params =
|
||||
CreateSecretParams::new(&secret_name, token_value).with_provider(name.to_string());
|
||||
self.secrets
|
||||
.create(&self.user_id, params)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||
|
||||
tracing::info!("MCP server '{}' authenticated via manual token", name);
|
||||
return Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
auth_url: None,
|
||||
callback_type: None,
|
||||
instructions: None,
|
||||
setup_url: None,
|
||||
awaiting_token: false,
|
||||
status: "authenticated".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check if already authenticated
|
||||
if is_authenticated(&server, &self.secrets, &self.user_id).await {
|
||||
return Ok(AuthResult {
|
||||
@@ -483,7 +506,7 @@ impl ExtensionManager {
|
||||
// Run the full OAuth flow (opens browser, waits for callback)
|
||||
match authorize_mcp_server(&server, &self.secrets, &self.user_id).await {
|
||||
Ok(_token) => {
|
||||
tracing::info!("MCP server '{}' authenticated successfully", name);
|
||||
tracing::info!("MCP server '{}' authenticated via OAuth", name);
|
||||
Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
@@ -496,10 +519,45 @@ impl ExtensionManager {
|
||||
})
|
||||
}
|
||||
Err(crate::tools::mcp::auth::AuthError::NotSupported) => {
|
||||
// Server doesn't support OAuth at all, try to build a non-interactive auth URL
|
||||
self.auth_mcp_build_url(name, &server).await
|
||||
// Server doesn't support OAuth, try building a URL first
|
||||
match self.auth_mcp_build_url(name, &server).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(_) => {
|
||||
// No OAuth, no DCR: fall back to manual token entry
|
||||
Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
auth_url: None,
|
||||
callback_type: None,
|
||||
instructions: Some(format!(
|
||||
"Server '{}' does not support OAuth. \
|
||||
Please provide an API token/key for this server.",
|
||||
name
|
||||
)),
|
||||
setup_url: None,
|
||||
awaiting_token: true,
|
||||
status: "awaiting_token".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// OAuth failed for some other reason, fall back to manual token
|
||||
Ok(AuthResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
auth_url: None,
|
||||
callback_type: None,
|
||||
instructions: Some(format!(
|
||||
"OAuth failed for '{}': {}. \
|
||||
Please provide an API token/key manually.",
|
||||
name, e
|
||||
)),
|
||||
setup_url: None,
|
||||
awaiting_token: true,
|
||||
status: "awaiting_token".to_string(),
|
||||
})
|
||||
}
|
||||
Err(e) => Err(ExtensionError::AuthFailed(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+124
-106
@@ -301,44 +301,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Builder mode enabled");
|
||||
}
|
||||
|
||||
// Load installed WASM tools (save runtime handle for extension manager)
|
||||
let wasm_tool_runtime: Option<Arc<WasmToolRuntime>> = if config.wasm.enabled
|
||||
&& config.wasm.tools_dir.exists()
|
||||
{
|
||||
match WasmToolRuntime::new(config.wasm.to_runtime_config()) {
|
||||
Ok(runtime) => {
|
||||
let runtime = Arc::new(runtime);
|
||||
let loader = WasmToolLoader::new(Arc::clone(&runtime), Arc::clone(&tools));
|
||||
|
||||
match loader.load_from_dir(&config.wasm.tools_dir).await {
|
||||
Ok(results) => {
|
||||
if !results.loaded.is_empty() {
|
||||
tracing::info!(
|
||||
"Loaded {} WASM tools from {}",
|
||||
results.loaded.len(),
|
||||
config.wasm.tools_dir.display()
|
||||
);
|
||||
}
|
||||
for (path, err) in &results.errors {
|
||||
tracing::warn!("Failed to load WASM tool {}: {}", path.display(), err);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to scan WASM tools directory: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
Some(runtime)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize WASM runtime: {}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Create secrets store if master key is configured (needed for MCP auth and WASM channels)
|
||||
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
|
||||
if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) {
|
||||
@@ -356,91 +318,146 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
// Load configured MCP servers
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
match load_mcp_servers().await {
|
||||
Ok(servers) => {
|
||||
let enabled_count = servers.servers.iter().filter(|s| s.enabled).count();
|
||||
if enabled_count > 0 {
|
||||
tracing::info!("Loading {} configured MCP server(s)...", enabled_count);
|
||||
|
||||
// Create WASM tool runtime (sync, just builds the wasmtime engine)
|
||||
let wasm_tool_runtime: Option<Arc<WasmToolRuntime>> =
|
||||
if config.wasm.enabled && config.wasm.tools_dir.exists() {
|
||||
match WasmToolRuntime::new(config.wasm.to_runtime_config()) {
|
||||
Ok(runtime) => Some(Arc::new(runtime)),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize WASM runtime: {}", e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
for server in servers.enabled_servers() {
|
||||
tracing::debug!(
|
||||
"Checking authentication for MCP server '{}'...",
|
||||
server.name
|
||||
);
|
||||
// Check for stored tokens (from either pre-configured OAuth or DCR)
|
||||
let has_tokens = is_authenticated(server, secrets, "default").await;
|
||||
tracing::debug!("MCP server '{}' has_tokens={}", server.name, has_tokens);
|
||||
// Load WASM tools and MCP servers concurrently.
|
||||
// Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe.
|
||||
let wasm_tools_future = async {
|
||||
if let Some(ref runtime) = wasm_tool_runtime {
|
||||
let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
|
||||
match loader.load_from_dir(&config.wasm.tools_dir).await {
|
||||
Ok(results) => {
|
||||
if !results.loaded.is_empty() {
|
||||
tracing::info!(
|
||||
"Loaded {} WASM tools from {}",
|
||||
results.loaded.len(),
|
||||
config.wasm.tools_dir.display()
|
||||
);
|
||||
}
|
||||
for (path, err) in &results.errors {
|
||||
tracing::warn!("Failed to load WASM tool {}: {}", path.display(), err);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to scan WASM tools directory: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let client = if has_tokens || server.requires_auth() {
|
||||
// Use authenticated client if we have tokens or OAuth is configured
|
||||
McpClient::new_authenticated(
|
||||
server.clone(),
|
||||
Arc::clone(&mcp_session_manager),
|
||||
Arc::clone(secrets),
|
||||
"default",
|
||||
)
|
||||
} else {
|
||||
// No tokens and no OAuth - try unauthenticated
|
||||
McpClient::new_with_name(&server.name, &server.url)
|
||||
};
|
||||
let mcp_servers_future = async {
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
match load_mcp_servers().await {
|
||||
Ok(servers) => {
|
||||
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||
if !enabled.is_empty() {
|
||||
tracing::info!("Loading {} configured MCP server(s)...", enabled.len());
|
||||
}
|
||||
|
||||
tracing::debug!("Fetching tools from MCP server '{}'...", server.name);
|
||||
match client.list_tools().await {
|
||||
Ok(mcp_tools) => {
|
||||
let mut join_set = tokio::task::JoinSet::new();
|
||||
for server in enabled {
|
||||
let mcp_sm = Arc::clone(&mcp_session_manager);
|
||||
let secrets = Arc::clone(secrets);
|
||||
let tools = Arc::clone(&tools);
|
||||
|
||||
join_set.spawn(async move {
|
||||
let server_name = server.name.clone();
|
||||
tracing::debug!(
|
||||
"Got {} tools from MCP server '{}'",
|
||||
mcp_tools.len(),
|
||||
server.name
|
||||
"Checking authentication for MCP server '{}'...",
|
||||
server_name
|
||||
);
|
||||
match client.create_tools().await {
|
||||
Ok(tool_impls) => {
|
||||
for tool in tool_impls {
|
||||
tools.register(tool).await;
|
||||
}
|
||||
tracing::info!(
|
||||
"Loaded {} tools from MCP server '{}'",
|
||||
mcp_tools.len(),
|
||||
server.name
|
||||
let has_tokens = is_authenticated(&server, &secrets, "default").await;
|
||||
tracing::debug!(
|
||||
"MCP server '{}' has_tokens={}",
|
||||
server_name,
|
||||
has_tokens
|
||||
);
|
||||
|
||||
let client = if has_tokens || server.requires_auth() {
|
||||
McpClient::new_authenticated(server, mcp_sm, secrets, "default")
|
||||
} else {
|
||||
McpClient::new_with_name(&server_name, &server.url)
|
||||
};
|
||||
|
||||
tracing::debug!("Fetching tools from MCP server '{}'...", server_name);
|
||||
match client.list_tools().await {
|
||||
Ok(mcp_tools) => {
|
||||
let tool_count = mcp_tools.len();
|
||||
tracing::debug!(
|
||||
"Got {} tools from MCP server '{}'",
|
||||
tool_count,
|
||||
server_name
|
||||
);
|
||||
match client.create_tools().await {
|
||||
Ok(tool_impls) => {
|
||||
for tool in tool_impls {
|
||||
tools.register(tool).await;
|
||||
}
|
||||
tracing::info!(
|
||||
"Loaded {} tools from MCP server '{}'",
|
||||
tool_count,
|
||||
server_name
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to create tools from MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to create tools from MCP server '{}': {}",
|
||||
server.name,
|
||||
e
|
||||
);
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("401") || err_str.contains("authentication")
|
||||
{
|
||||
tracing::warn!(
|
||||
"MCP server '{}' requires authentication. \
|
||||
Run: ironclaw mcp auth {}",
|
||||
server_name,
|
||||
server_name
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Failed to connect to MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Check if it's an auth error
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("401") || err_str.contains("authentication") {
|
||||
tracing::warn!(
|
||||
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
|
||||
server.name,
|
||||
server.name
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Failed to connect to MCP server '{}': {}",
|
||||
server.name,
|
||||
e
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(result) = join_set.join_next().await {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!("MCP server loading task panicked: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("No MCP servers configured ({})", e);
|
||||
Err(e) => {
|
||||
tracing::debug!("No MCP servers configured ({})", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
tokio::join!(wasm_tools_future, mcp_servers_future);
|
||||
|
||||
// Create extension manager for in-chat discovery/install/auth/activate
|
||||
let extension_manager = if let Some(ref secrets) = secrets_store {
|
||||
@@ -728,6 +745,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
safety,
|
||||
tools,
|
||||
workspace,
|
||||
extension_manager,
|
||||
};
|
||||
let agent = Agent::new(
|
||||
config.agent.clone(),
|
||||
|
||||
@@ -165,17 +165,18 @@ impl WasmToolLoader {
|
||||
}
|
||||
|
||||
let mut results = LoadResults::default();
|
||||
|
||||
// Collect all .wasm entries first, then load in parallel
|
||||
let mut tool_entries = Vec::new();
|
||||
let mut entries = fs::read_dir(dir).await?;
|
||||
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let path = entry.path();
|
||||
|
||||
// Only process .wasm files
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("wasm") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extract tool name from filename
|
||||
let name = match path.file_stem().and_then(|s| s.to_str()) {
|
||||
Some(n) => n.to_string(),
|
||||
None => {
|
||||
@@ -187,15 +188,20 @@ impl WasmToolLoader {
|
||||
}
|
||||
};
|
||||
|
||||
// Look for sidecar capabilities file
|
||||
let cap_path = path.with_extension("capabilities.json");
|
||||
let cap_path_option = if cap_path.exists() {
|
||||
Some(cap_path.as_path())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let has_cap = cap_path.exists();
|
||||
tool_entries.push((name, path, if has_cap { Some(cap_path) } else { None }));
|
||||
}
|
||||
|
||||
match self.load_from_files(&name, &path, cap_path_option).await {
|
||||
// Load all tools in parallel (file I/O + WASM compilation + registration)
|
||||
let load_futures = tool_entries
|
||||
.iter()
|
||||
.map(|(name, path, cap_path)| self.load_from_files(name, path, cap_path.as_deref()));
|
||||
|
||||
let load_results = futures::future::join_all(load_futures).await;
|
||||
|
||||
for ((name, path, _), result) in tool_entries.into_iter().zip(load_results) {
|
||||
match result {
|
||||
Ok(()) => {
|
||||
results.loaded.push(name);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user