use crate::utils::clock::Instant;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use super::client::ClientId;
use super::server::ServerId;
use crate::server::connection::ConnectionId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TaskId(u64);
impl TaskId {
pub fn new(id: u64) -> Self {
Self(id)
}
pub fn as_u64(&self) -> u64 {
self.0
}
pub fn from_string(s: &str) -> Option<Self> {
s.parse::<u64>().ok().map(Self)
}
}
impl std::fmt::Display for TaskId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone)]
pub enum TaskScope {
Global,
Server(ServerId),
Connection(ServerId, ConnectionId),
Client(ClientId),
}
#[derive(Debug, Clone)]
pub enum TaskType {
OneShot { delay_secs: u64 },
Recurring {
interval_secs: u64,
max_executions: Option<u64>,
executions_count: u64,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaskStatus {
Scheduled,
Executing,
Completed,
Failed(String),
}
#[derive(Debug, Clone)]
pub struct ScheduledTask {
pub id: TaskId,
pub name: String,
pub scope: TaskScope,
pub task_type: TaskType,
pub instruction: String,
pub context: Option<serde_json::Value>,
pub status: TaskStatus,
pub created_at: Instant,
pub next_execution: Instant,
pub last_error: Option<String>,
pub failure_count: u64,
}
impl ScheduledTask {
pub fn new_one_shot(
id: TaskId,
name: String,
scope: TaskScope,
delay_secs: u64,
instruction: String,
context: Option<serde_json::Value>,
) -> anyhow::Result<Self> {
let now = Instant::now();
let next_execution = checked_task_deadline(delay_secs)?;
Ok(Self {
id,
name,
scope,
task_type: TaskType::OneShot { delay_secs },
instruction,
context,
status: TaskStatus::Scheduled,
created_at: now,
next_execution,
last_error: None,
failure_count: 0,
})
}
pub fn new_recurring(
id: TaskId,
name: String,
scope: TaskScope,
interval_secs: u64,
max_executions: Option<u64>,
instruction: String,
context: Option<serde_json::Value>,
) -> anyhow::Result<Self> {
let now = Instant::now();
anyhow::ensure!(
interval_secs > 0,
"recurring interval_secs must be greater than zero"
);
anyhow::ensure!(
max_executions != Some(0),
"max_executions must be greater than zero"
);
let next_execution = checked_task_deadline(interval_secs)?;
Ok(Self {
id,
name,
scope,
task_type: TaskType::Recurring {
interval_secs,
max_executions,
executions_count: 0,
},
instruction,
context,
status: TaskStatus::Scheduled,
created_at: now,
next_execution,
last_error: None,
failure_count: 0,
})
}
pub fn new_connection_one_shot(
id: TaskId,
name: String,
server_id: ServerId,
connection_id: ConnectionId,
delay_secs: u64,
instruction: String,
context: Option<serde_json::Value>,
) -> anyhow::Result<Self> {
Self::new_one_shot(
id,
name,
TaskScope::Connection(server_id, connection_id),
delay_secs,
instruction,
context,
)
}
#[allow(clippy::too_many_arguments)]
pub fn new_connection_recurring(
id: TaskId,
name: String,
server_id: ServerId,
connection_id: ConnectionId,
interval_secs: u64,
max_executions: Option<u64>,
instruction: String,
context: Option<serde_json::Value>,
) -> anyhow::Result<Self> {
Self::new_recurring(
id,
name,
TaskScope::Connection(server_id, connection_id),
interval_secs,
max_executions,
instruction,
context,
)
}
pub fn from_definition(
definition: &crate::llm::actions::common::ServerTaskDefinition,
scope: TaskScope,
) -> anyhow::Result<Self> {
let next_execution = checked_task_deadline(definition.delay_secs.unwrap_or(0))?;
let mut task = if definition.recurring {
Self::new_recurring(
TaskId::new(0),
definition.task_id.clone(),
scope,
definition.interval_secs.unwrap_or(60),
definition.max_executions,
definition.instruction.clone(),
definition.context.clone(),
)?
} else {
Self::new_one_shot(
TaskId::new(0),
definition.task_id.clone(),
scope,
definition.delay_secs.unwrap_or(0),
definition.instruction.clone(),
definition.context.clone(),
)?
};
task.next_execution = next_execution;
Ok(task)
}
pub fn interval_secs(&self) -> Option<u64> {
match &self.task_type {
TaskType::Recurring { interval_secs, .. } => Some(*interval_secs),
_ => None,
}
}
pub fn executions_count(&self) -> u64 {
match &self.task_type {
TaskType::Recurring {
executions_count, ..
} => *executions_count,
_ => 0,
}
}
pub fn is_max_executions_reached(&self) -> bool {
match &self.task_type {
TaskType::Recurring {
max_executions: Some(max),
executions_count,
..
} => executions_count >= max,
_ => false,
}
}
pub fn format_for_prompt(&self) -> String {
let scope_str = match &self.scope {
TaskScope::Global => "global".to_string(),
TaskScope::Server(sid) => format!("server #{}", sid.as_u32()),
TaskScope::Connection(sid, cid) => {
format!("connection {} on server #{}", cid, sid.as_u32())
}
TaskScope::Client(cid) => format!("client #{}", cid.as_u32()),
};
let type_str = match &self.task_type {
TaskType::OneShot { delay_secs } => {
format!("one-shot ({}s delay)", delay_secs)
}
TaskType::Recurring {
interval_secs,
max_executions,
executions_count,
} => {
if let Some(max) = max_executions {
format!(
"recurring ({}s interval, {}/{} executions)",
interval_secs, executions_count, max
)
} else {
format!(
"recurring ({}s interval, {} executions)",
interval_secs, executions_count
)
}
}
};
let next_in = if self.next_execution > Instant::now() {
let secs = (self.next_execution - Instant::now()).as_secs();
format!("{}s", secs)
} else {
"now".to_string()
};
format!(
"- Task '{}' (ID: {}): {} scope, {}, next in {}, status: {:?}",
self.name, self.id, scope_str, type_str, next_in, self.status
)
}
}
pub struct TaskExecutionResult {
pub success: bool,
pub actions: Vec<serde_json::Value>,
pub error: Option<String>,
}
pub fn format_duration(duration: Duration) -> String {
let secs = duration.as_secs();
if secs < 60 {
format!("{}s", secs)
} else if secs < 3600 {
format!("{}m", secs / 60)
} else if secs < 86400 {
format!("{}h", secs / 3600)
} else {
format!("{}d", secs / 86400)
}
}
pub fn checked_task_deadline(seconds: u64) -> anyhow::Result<Instant> {
Instant::now()
.checked_add(Duration::from_secs(seconds))
.ok_or_else(|| {
anyhow::anyhow!(
"scheduled task delay/interval of {} seconds exceeds the supported clock range",
seconds
)
})
}
pub fn prepare_tasks(
definitions: Option<&[crate::llm::actions::common::ServerTaskDefinition]>,
scope: TaskScope,
) -> anyhow::Result<Vec<ScheduledTask>> {
definitions
.unwrap_or_default()
.iter()
.map(|definition| ScheduledTask::from_definition(definition, scope.clone()))
.collect()
}