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:
Illia Polosukhin
2026-02-06 13:35:17 -08:00
co-authored by Claude Opus 4.6
parent 9b729795fb
commit 2cdd04a359
7 changed files with 533 additions and 142 deletions
+82
View File
@@ -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());
}
}