use super::super::elicitation::AdkClientHandler;
use super::super::toolset::McpToolset;
use super::config::{McpServerConfig, RestartPolicy};
use super::status::ServerStatus;
#[derive(Debug, Clone)]
#[allow(dead_code)] pub(crate) struct BackoffState {
pub consecutive_failures: u32,
pub current_delay_ms: u64,
}
#[allow(dead_code)] impl BackoffState {
pub fn new(policy: &Option<RestartPolicy>) -> Self {
Self {
consecutive_failures: 0,
current_delay_ms: policy.as_ref().map_or(1000, |p| p.initial_delay_ms),
}
}
pub fn next_delay(&mut self, policy: &RestartPolicy) -> u64 {
let delay = self.current_delay_ms;
self.consecutive_failures += 1;
self.current_delay_ms = ((self.current_delay_ms as f64 * policy.backoff_multiplier) as u64)
.min(policy.max_delay_ms);
delay
}
pub fn reset(&mut self, policy: &RestartPolicy) {
self.consecutive_failures = 0;
self.current_delay_ms = policy.initial_delay_ms;
}
pub fn exceeded_max_attempts(&self, policy: &RestartPolicy) -> bool {
self.consecutive_failures >= policy.max_restart_attempts
}
}
#[allow(dead_code)] pub(crate) struct McpServerEntry {
pub config: McpServerConfig,
pub status: ServerStatus,
pub toolset: Option<McpToolset<AdkClientHandler>>,
pub child: Option<tokio::process::Child>,
pub backoff: BackoffState,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_backoff_state_new_with_policy() {
let policy = Some(RestartPolicy {
initial_delay_ms: 500,
max_delay_ms: 10000,
backoff_multiplier: 2.0,
max_restart_attempts: 5,
});
let state = BackoffState::new(&policy);
assert_eq!(state.consecutive_failures, 0);
assert_eq!(state.current_delay_ms, 500);
}
#[test]
fn test_backoff_state_new_without_policy() {
let state = BackoffState::new(&None);
assert_eq!(state.consecutive_failures, 0);
assert_eq!(state.current_delay_ms, 1000);
}
#[test]
fn test_backoff_next_delay() {
let policy = RestartPolicy {
initial_delay_ms: 1000,
max_delay_ms: 30000,
backoff_multiplier: 2.0,
max_restart_attempts: 10,
};
let mut state = BackoffState::new(&Some(policy.clone()));
let delay = state.next_delay(&policy);
assert_eq!(delay, 1000);
assert_eq!(state.consecutive_failures, 1);
assert_eq!(state.current_delay_ms, 2000);
let delay = state.next_delay(&policy);
assert_eq!(delay, 2000);
assert_eq!(state.consecutive_failures, 2);
assert_eq!(state.current_delay_ms, 4000);
let delay = state.next_delay(&policy);
assert_eq!(delay, 4000);
assert_eq!(state.consecutive_failures, 3);
assert_eq!(state.current_delay_ms, 8000);
}
#[test]
fn test_backoff_delay_capped_at_max() {
let policy = RestartPolicy {
initial_delay_ms: 10000,
max_delay_ms: 15000,
backoff_multiplier: 2.0,
max_restart_attempts: 10,
};
let mut state = BackoffState::new(&Some(policy.clone()));
let delay = state.next_delay(&policy);
assert_eq!(delay, 10000);
assert_eq!(state.current_delay_ms, 15000);
let delay = state.next_delay(&policy);
assert_eq!(delay, 15000);
assert_eq!(state.current_delay_ms, 15000);
}
#[test]
fn test_backoff_reset() {
let policy = RestartPolicy {
initial_delay_ms: 1000,
max_delay_ms: 30000,
backoff_multiplier: 2.0,
max_restart_attempts: 10,
};
let mut state = BackoffState::new(&Some(policy.clone()));
state.next_delay(&policy);
state.next_delay(&policy);
assert_eq!(state.consecutive_failures, 2);
assert_eq!(state.current_delay_ms, 4000);
state.reset(&policy);
assert_eq!(state.consecutive_failures, 0);
assert_eq!(state.current_delay_ms, 1000);
}
#[test]
fn test_backoff_exceeded_max_attempts() {
let policy = RestartPolicy {
initial_delay_ms: 1000,
max_delay_ms: 30000,
backoff_multiplier: 2.0,
max_restart_attempts: 3,
};
let mut state = BackoffState::new(&Some(policy.clone()));
assert!(!state.exceeded_max_attempts(&policy));
state.next_delay(&policy); assert!(!state.exceeded_max_attempts(&policy));
state.next_delay(&policy); assert!(!state.exceeded_max_attempts(&policy));
state.next_delay(&policy); assert!(state.exceeded_max_attempts(&policy));
}
#[test]
fn test_mcp_server_entry_creation() {
let config = McpServerConfig {
command: "echo".to_string(),
args: vec![],
env: std::collections::HashMap::new(),
disabled: false,
auto_approve: vec![],
restart_policy: None,
};
let entry = McpServerEntry {
backoff: BackoffState::new(&config.restart_policy),
config,
status: ServerStatus::Stopped,
toolset: None,
child: None,
};
assert_eq!(entry.status, ServerStatus::Stopped);
assert!(entry.toolset.is_none());
assert!(entry.child.is_none());
}
}