Skip to content

Commit 4c89207

Browse files
its-mashits-mash
committed
backup
1 parent 93f04c5 commit 4c89207

3 files changed

Lines changed: 335 additions & 27 deletions

File tree

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,3 +83,4 @@ lto = true
8383
codegen-units = 1
8484
strip = true
8585

86+

‎crates/mcpmux-gateway/src/pool/routing.rs‎

Lines changed: 248 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use serde_json::Value;
1717
use tracing::{debug, info, warn};
1818
use uuid::Uuid;
1919

20+
use super::connection::ConnectionResult;
2021
use super::features::FeatureService;
2122
use super::service::PoolService;
2223

@@ -332,17 +333,128 @@ impl RoutingService {
332333
Ok(result) => {
333334
let duration = call_start.elapsed();
334335
if result.is_error {
335-
warn!(
336-
"[RoutingService] Tool execution error: {} (duration: {:?})",
337-
actual_tool_name, duration
338-
);
339-
self.log(
340-
&space_id,
341-
&server_id,
342-
LogLevel::Error,
343-
format!("Tool execution error: {}", actual_tool_name),
344-
Some(serde_json::json!({ "result": result.content, "duration_ms": duration.as_millis() }))
345-
).await;
336+
// Check if this is an auth error embedded in the tool result.
337+
// Some servers (e.g., Atlassian) return 401 as tool results rather than
338+
// HTTP errors. The SDK refreshes the token successfully, but the server's
339+
// internal session may be stale. A fresh MCP connection fixes this.
340+
if Self::content_has_auth_error(&result.content) {
341+
warn!(
342+
"[RoutingService] Auth error in tool result for {}/{}, attempting auto-reconnect",
343+
server_id, actual_tool_name
344+
);
345+
self.log(
346+
&space_id,
347+
&server_id,
348+
LogLevel::Warn,
349+
format!(
350+
"Auth error in tool result for '{}' - auto-reconnecting",
351+
actual_tool_name
352+
),
353+
Some(serde_json::json!({ "result": result.content, "duration_ms": duration.as_millis() })),
354+
)
355+
.await;
356+
357+
match self
358+
.pool_service
359+
.reconnect_instance(space_id, &server_id)
360+
.await
361+
{
362+
ConnectionResult::Connected { .. } => {
363+
info!(
364+
"[RoutingService] Reconnected {}, retrying tool call: {}",
365+
server_id, actual_tool_name
366+
);
367+
368+
let retry_start = std::time::Instant::now();
369+
match execute_call(
370+
self.pool_service.clone(),
371+
space_id,
372+
server_id.clone(),
373+
actual_tool_name.clone(),
374+
arguments.clone(),
375+
)
376+
.await
377+
{
378+
Ok(retry_result) => {
379+
let retry_duration = retry_start.elapsed();
380+
if retry_result.is_error {
381+
warn!(
382+
"[RoutingService] Tool retry still has error: {} (duration: {:?})",
383+
actual_tool_name, retry_duration
384+
);
385+
} else {
386+
info!(
387+
"[RoutingService] Tool retry succeeded after reconnect: {} (duration: {:?})",
388+
actual_tool_name, retry_duration
389+
);
390+
}
391+
self.log(
392+
&space_id,
393+
&server_id,
394+
LogLevel::Info,
395+
format!(
396+
"Tool '{}' retried after auto-reconnect (is_error={})",
397+
actual_tool_name, retry_result.is_error
398+
),
399+
Some(serde_json::json!({ "retry_duration_ms": retry_duration.as_millis() })),
400+
)
401+
.await;
402+
Ok(retry_result)
403+
}
404+
Err(retry_err) => {
405+
warn!(
406+
"[RoutingService] Tool retry transport error: {} - {}",
407+
actual_tool_name, retry_err
408+
);
409+
self.log(
410+
&space_id,
411+
&server_id,
412+
LogLevel::Error,
413+
format!(
414+
"Tool '{}' still failing after reconnect",
415+
actual_tool_name
416+
),
417+
Some(serde_json::json!({ "error": retry_err.to_string() })),
418+
)
419+
.await;
420+
// Return original tool result since it has the error details
421+
Ok(result)
422+
}
423+
}
424+
}
425+
other => {
426+
warn!(
427+
"[RoutingService] Auto-reconnect failed for {}: {:?}",
428+
server_id, other
429+
);
430+
self.log(
431+
&space_id,
432+
&server_id,
433+
LogLevel::Error,
434+
format!(
435+
"Auto-reconnect failed for tool '{}' - manual reconnection required",
436+
actual_tool_name
437+
),
438+
Some(serde_json::json!({ "reconnect_result": format!("{:?}", other) })),
439+
)
440+
.await;
441+
Ok(result)
442+
}
443+
}
444+
} else {
445+
warn!(
446+
"[RoutingService] Tool execution error: {} (duration: {:?})",
447+
actual_tool_name, duration
448+
);
449+
self.log(
450+
&space_id,
451+
&server_id,
452+
LogLevel::Error,
453+
format!("Tool execution error: {}", actual_tool_name),
454+
Some(serde_json::json!({ "result": result.content, "duration_ms": duration.as_millis() }))
455+
).await;
456+
Ok(result)
457+
}
346458
} else {
347459
info!(
348460
"[RoutingService] Tool executed successfully: {} (duration: {:?})",
@@ -356,8 +468,8 @@ impl RoutingService {
356468
Some(serde_json::json!({ "duration_ms": duration.as_millis() })),
357469
)
358470
.await;
471+
Ok(result)
359472
}
360-
Ok(result)
361473
}
362474
Err(e) => {
363475
let duration = call_start.elapsed();
@@ -368,35 +480,126 @@ impl RoutingService {
368480
actual_tool_name, server_id, e, duration
369481
);
370482

371-
// Check if it's an auth error
372-
// NOTE: With RMCP's AuthClient, token refresh happens automatically per-request.
373-
// If we still get an auth error, it means the refresh token is invalid or expired.
374-
// The user needs to reconnect to re-authorize.
375483
let is_auth = Self::is_auth_error(&err_str);
376-
let is_timeout = err_str.contains("timed out");
377484

378-
if is_auth || is_timeout {
485+
if is_auth {
486+
// Auth error detected - attempt auto-reconnect and retry once.
487+
// This handles the case where RMCP's AuthClient failed to refresh
488+
// the token (e.g., stale in-memory state after idle).
489+
// Creating a fresh connection loads latest tokens from the database.
379490
warn!(
380-
"[RoutingService] Auth/timeout error for {}/{} - RMCP auto-refresh likely failed, user needs to reconnect",
491+
"[RoutingService] Auth error for {}/{}, attempting auto-reconnect",
381492
server_id, actual_tool_name
382493
);
383494
self.log(
384495
&space_id,
385496
&server_id,
386-
LogLevel::Error,
387-
format!("Authentication failed for tool '{}' - reconnection required", actual_tool_name),
388-
Some(serde_json::json!({ "error": e.to_string(), "duration_ms": duration.as_millis() }))
389-
).await;
390-
Err(anyhow!("Server '{}' requires reconnection. Token may have expired. Please disconnect and connect again.", server_id))
497+
LogLevel::Warn,
498+
format!(
499+
"Auth error on tool '{}' - auto-reconnecting to refresh credentials",
500+
actual_tool_name
501+
),
502+
Some(serde_json::json!({ "error": e.to_string(), "duration_ms": duration.as_millis() })),
503+
)
504+
.await;
505+
506+
match self
507+
.pool_service
508+
.reconnect_instance(space_id, &server_id)
509+
.await
510+
{
511+
ConnectionResult::Connected { .. } => {
512+
info!(
513+
"[RoutingService] Reconnected {}, retrying tool call: {}",
514+
server_id, actual_tool_name
515+
);
516+
517+
// Retry the call once with the fresh connection
518+
let retry_start = std::time::Instant::now();
519+
match execute_call(
520+
self.pool_service.clone(),
521+
space_id,
522+
server_id.clone(),
523+
actual_tool_name.clone(),
524+
arguments.clone(),
525+
)
526+
.await
527+
{
528+
Ok(result) => {
529+
let retry_duration = retry_start.elapsed();
530+
info!(
531+
"[RoutingService] Tool retry succeeded: {} (duration: {:?})",
532+
actual_tool_name, retry_duration
533+
);
534+
self.log(
535+
&space_id,
536+
&server_id,
537+
LogLevel::Info,
538+
format!(
539+
"Tool '{}' succeeded after auto-reconnect",
540+
actual_tool_name
541+
),
542+
Some(serde_json::json!({ "retry_duration_ms": retry_duration.as_millis() })),
543+
)
544+
.await;
545+
Ok(result)
546+
}
547+
Err(retry_err) => {
548+
warn!(
549+
"[RoutingService] Tool retry also failed: {} - {}",
550+
actual_tool_name, retry_err
551+
);
552+
self.log(
553+
&space_id,
554+
&server_id,
555+
LogLevel::Error,
556+
format!(
557+
"Tool '{}' still failing after reconnect - manual reconnection required",
558+
actual_tool_name
559+
),
560+
Some(serde_json::json!({ "error": retry_err.to_string() })),
561+
)
562+
.await;
563+
Err(anyhow!(
564+
"Server '{}' auth error persists after auto-reconnect. Please disconnect and connect again. Error: {}",
565+
server_id,
566+
retry_err
567+
))
568+
}
569+
}
570+
}
571+
other => {
572+
warn!(
573+
"[RoutingService] Auto-reconnect failed for {}: {:?}",
574+
server_id, other
575+
);
576+
self.log(
577+
&space_id,
578+
&server_id,
579+
LogLevel::Error,
580+
format!(
581+
"Auto-reconnect failed for tool '{}' - manual reconnection required",
582+
actual_tool_name
583+
),
584+
Some(serde_json::json!({ "reconnect_result": format!("{:?}", other) })),
585+
)
586+
.await;
587+
Err(anyhow!(
588+
"Server '{}' requires reconnection. Auto-reconnect failed. Please disconnect and connect again.",
589+
server_id
590+
))
591+
}
592+
}
391593
} else {
392594
// Not an auth error, return original error
393595
self.log(
394596
&space_id,
395597
&server_id,
396598
LogLevel::Error,
397599
format!("Tool call failed: {}", e),
398-
Some(serde_json::json!({ "error": e.to_string(), "duration_ms": duration.as_millis() }))
399-
).await;
600+
Some(serde_json::json!({ "error": e.to_string(), "duration_ms": duration.as_millis() })),
601+
)
602+
.await;
400603
Err(e)
401604
}
402605
}
@@ -426,7 +629,7 @@ impl RoutingService {
426629
}
427630
}
428631

429-
/// Check if an error indicates authentication is needed
632+
/// Check if an error string indicates authentication is needed
430633
fn is_auth_error(error_str: &str) -> bool {
431634
let indicators = [
432635
"401",
@@ -437,4 +640,22 @@ impl RoutingService {
437640
];
438641
indicators.iter().any(|s| error_str.contains(s))
439642
}
643+
644+
/// Check if tool result content contains authentication error indicators.
645+
///
646+
/// Some MCP servers (e.g., Atlassian) return auth errors as tool results
647+
/// (`is_error: true` with 401 in the text) rather than HTTP-level errors.
648+
/// The SDK may have already refreshed the token, but the server's internal
649+
/// session can be stale. A fresh connection (reconnect) fixes this.
650+
fn content_has_auth_error(content: &[Value]) -> bool {
651+
for item in content {
652+
if let Some(text) = item.get("text").and_then(|v| v.as_str()) {
653+
let lower = text.to_lowercase();
654+
if Self::is_auth_error(&lower) {
655+
return true;
656+
}
657+
}
658+
}
659+
false
660+
}
440661
}

0 commit comments

Comments
 (0)