use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionLimits {
pub max_turns: usize,
pub max_total_tokens: usize,
pub max_duration: std::time::Duration,
#[serde(default)]
pub max_cost: Option<f64>,
}
impl Default for ExecutionLimits {
fn default() -> Self {
Self {
max_turns: 50,
max_total_tokens: 1_000_000,
max_duration: std::time::Duration::from_secs(600),
max_cost: None,
}
}
}
pub struct ExecutionTracker {
pub limits: ExecutionLimits,
pub turns: usize,
pub tokens_used: usize,
pub cost_accumulated: f64,
pub started_at: std::time::Instant,
}
#[derive(Debug, Clone)]
pub struct CurrentToolExecution {
pub name: String,
pub timeout: Option<std::time::Duration>,
}
impl ExecutionTracker {
pub fn new(limits: ExecutionLimits) -> Self {
Self {
limits,
turns: 0,
tokens_used: 0,
cost_accumulated: 0.0,
started_at: std::time::Instant::now(),
}
}
pub fn record_turn(&mut self, tokens: usize) {
self.turns += 1;
self.tokens_used += tokens;
}
pub fn record_cost(&mut self, cost: f64) {
self.cost_accumulated += cost;
}
pub fn check_limits(&self) -> Option<String> {
if self.turns >= self.limits.max_turns {
return Some(format!(
"Max turns reached ({}/{})",
self.turns, self.limits.max_turns
));
}
if self.tokens_used >= self.limits.max_total_tokens {
return Some(format!(
"Max tokens reached ({}/{})",
self.tokens_used, self.limits.max_total_tokens
));
}
let elapsed = self.started_at.elapsed(); if elapsed >= self.limits.max_duration {
return Some(format!(
"Max duration reached ({:.0}s/{:.0}s)", elapsed.as_secs_f64(),
self.limits.max_duration.as_secs_f64()
));
}
if let Some(max) = self.limits.max_cost {
if self.cost_accumulated >= max {
return Some(format!(
"Max cost reached (${:.4}/${:.4})",
self.cost_accumulated, max
));
}
}
None }
}