Skip to content

Commit 273205c

Browse files
ilblackdragonclaude
andcommitted
fix(mcp): add TTL to PendingAuth and clear auth mode on all failure paths
Auth mode (pending_auth on a Thread) had no timeout and several code paths that failed to clear it, causing user messages to be swallowed indefinitely. This adds defense-in-depth: - Add created_at + 5-minute TTL to PendingAuth; auto-clear on next message if expired (safety net for edge cases like user closing browser mid-OAuth) - Clear auth mode on OAuth callback failure paths (unknown/consumed state, expired flow) - Move clear_auth_mode before configure() match in setup_submit so it runs on failure too (addresses Copilot review feedback) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 8b89176 commit 273205c

3 files changed

Lines changed: 75 additions & 26 deletions

File tree

src/agent/agent_loop.rs

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -838,19 +838,31 @@ impl Agent {
838838
};
839839

840840
if let Some(pending) = pending_auth {
841-
match &submission {
842-
Submission::UserInput { content } => {
843-
return self
844-
.process_auth_token(message, &pending, content, session, thread_id)
845-
.await;
841+
if pending.is_expired() {
842+
// TTL exceeded — clear stale auth and fall through to normal handling
843+
tracing::warn!(
844+
extension = %pending.extension_name,
845+
"Auth mode expired after TTL, clearing"
846+
);
847+
let mut sess = session.lock().await;
848+
if let Some(thread) = sess.threads.get_mut(&thread_id) {
849+
thread.pending_auth = None;
846850
}
847-
_ => {
848-
// Any control submission (interrupt, undo, etc.) cancels auth mode
849-
let mut sess = session.lock().await;
850-
if let Some(thread) = sess.threads.get_mut(&thread_id) {
851-
thread.pending_auth = None;
851+
} else {
852+
match &submission {
853+
Submission::UserInput { content } => {
854+
return self
855+
.process_auth_token(message, &pending, content, session, thread_id)
856+
.await;
857+
}
858+
_ => {
859+
// Any control submission (interrupt, undo, etc.) cancels auth mode
860+
let mut sess = session.lock().await;
861+
if let Some(thread) = sess.threads.get_mut(&thread_id) {
862+
thread.pending_auth = None;
863+
}
864+
// Fall through to normal handling
852865
}
853-
// Fall through to normal handling
854866
}
855867
}
856868
}

src/agent/session.rs

Lines changed: 46 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
1313
use std::collections::{HashMap, HashSet};
1414

15-
use chrono::{DateTime, Utc};
15+
use chrono::{DateTime, TimeDelta, Utc};
1616
use serde::{Deserialize, Serialize};
1717
use uuid::Uuid;
1818

@@ -132,6 +132,9 @@ pub enum ThreadState {
132132

133133
/// Pending auth token request.
134134
///
135+
/// Auth mode TTL — matches `OAUTH_FLOW_EXPIRY` (5 minutes).
136+
const AUTH_MODE_TTL: TimeDelta = TimeDelta::minutes(5);
137+
135138
/// When `tool_auth` returns `awaiting_token`, the thread enters auth mode.
136139
/// The next user message is intercepted before entering the normal pipeline
137140
/// (no logging, no turn creation, no history) and routed directly to the
@@ -140,6 +143,16 @@ pub enum ThreadState {
140143
pub struct PendingAuth {
141144
/// Extension name to authenticate.
142145
pub extension_name: String,
146+
/// When this auth mode was entered. Used for TTL expiry.
147+
#[serde(default = "Utc::now")]
148+
pub created_at: DateTime<Utc>,
149+
}
150+
151+
impl PendingAuth {
152+
/// Returns `true` if this auth mode has exceeded the TTL.
153+
pub fn is_expired(&self) -> bool {
154+
Utc::now() - self.created_at > AUTH_MODE_TTL
155+
}
143156
}
144157

145158
/// Pending tool approval request stored on a thread.
@@ -295,7 +308,10 @@ impl Thread {
295308
/// Enter auth mode: next user message will be routed directly to
296309
/// the credential store, bypassing the normal pipeline entirely.
297310
pub fn enter_auth_mode(&mut self, extension_name: String) {
298-
self.pending_auth = Some(PendingAuth { extension_name });
311+
self.pending_auth = Some(PendingAuth {
312+
extension_name,
313+
created_at: Utc::now(),
314+
});
299315
self.updated_at = Utc::now();
300316
}
301317

@@ -684,15 +700,16 @@ mod tests {
684700

685701
#[test]
686702
fn test_enter_auth_mode() {
703+
let before = Utc::now();
687704
let mut thread = Thread::new(Uuid::new_v4());
688705
assert!(thread.pending_auth.is_none());
689706

690707
thread.enter_auth_mode("telegram".to_string());
691-
assert!(thread.pending_auth.is_some());
692-
assert_eq!(
693-
thread.pending_auth.as_ref().unwrap().extension_name,
694-
"telegram"
695-
);
708+
assert!(thread.pending_auth.is_some()); // safety: test assertion
709+
let pending = thread.pending_auth.as_ref().unwrap(); // safety: checked above
710+
assert_eq!(pending.extension_name, "telegram"); // safety: test assertion
711+
assert!(pending.created_at >= before); // safety: test assertion
712+
assert!(!pending.is_expired()); // safety: test assertion
696713
}
697714

698715
#[test]
@@ -701,8 +718,10 @@ mod tests {
701718
thread.enter_auth_mode("notion".to_string());
702719

703720
let pending = thread.take_pending_auth();
704-
assert!(pending.is_some());
705-
assert_eq!(pending.unwrap().extension_name, "notion");
721+
assert!(pending.is_some()); // safety: test assertion
722+
let pending = pending.unwrap(); // safety: checked above
723+
assert_eq!(pending.extension_name, "notion"); // safety: test assertion
724+
assert!(!pending.is_expired()); // safety: test assertion
706725

707726
// Should be cleared after take
708727
assert!(thread.pending_auth.is_none());
@@ -717,10 +736,26 @@ mod tests {
717736
let json = serde_json::to_string(&thread).expect("should serialize");
718737
assert!(json.contains("pending_auth"));
719738
assert!(json.contains("openai"));
739+
assert!(json.contains("created_at")); // safety: test assertion
720740

721741
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
722-
assert!(restored.pending_auth.is_some());
723-
assert_eq!(restored.pending_auth.unwrap().extension_name, "openai");
742+
assert!(restored.pending_auth.is_some()); // safety: test assertion
743+
let pending = restored.pending_auth.unwrap(); // safety: checked above
744+
assert_eq!(pending.extension_name, "openai"); // safety: test assertion
745+
assert!(!pending.is_expired()); // safety: test assertion
746+
}
747+
748+
#[test]
749+
fn test_pending_auth_expiry() {
750+
let mut pending = PendingAuth {
751+
extension_name: "test".to_string(),
752+
created_at: Utc::now(),
753+
};
754+
assert!(!pending.is_expired()); // safety: test assertion
755+
756+
// Backdate beyond the TTL
757+
pending.created_at = Utc::now() - AUTH_MODE_TTL - TimeDelta::seconds(1);
758+
assert!(pending.is_expired()); // safety: test assertion
724759
}
725760

726761
#[test]

src/channels/web/server.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,7 @@ async fn oauth_callback_handler(
563563
lookup_key = %lookup_key,
564564
"OAuth callback received with unknown or expired state"
565565
);
566+
clear_auth_mode(&state).await;
566567
return oauth_error_page("IronClaw");
567568
}
568569
};
@@ -581,6 +582,7 @@ async fn oauth_callback_handler(
581582
message: "OAuth flow expired. Please try again.".to_string(),
582583
});
583584
}
585+
clear_auth_mode(&state).await;
584586
return oauth_error_page(&flow.display_name);
585587
}
586588

@@ -2186,12 +2188,12 @@ async fn extensions_setup_submit_handler(
21862188
"Extension manager not available (secrets store required)".to_string(),
21872189
))?;
21882190

2191+
// Clear auth mode regardless of outcome so the next user message goes
2192+
// through to the LLM instead of being intercepted as a token.
2193+
clear_auth_mode(&state).await;
2194+
21892195
match ext_mgr.configure(&name, &req.secrets).await {
21902196
Ok(result) => {
2191-
// Clear auth mode so the next user message goes through to the LLM
2192-
// instead of being intercepted as a token.
2193-
clear_auth_mode(&state).await;
2194-
21952197
// Broadcast auth_completed so the chat UI can dismiss any in-progress
21962198
// auth card or setup modal that was triggered by tool_auth/tool_activate.
21972199
state.sse.broadcast(SseEvent::AuthCompleted {

0 commit comments

Comments
 (0)