use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use async_trait::async_trait;
use talos_core::session::{
MAX_SUBMISSION_ITEM_BYTES, SessionOp, StructuredSubmission, SubmissionItem, SubmissionKind,
SubmissionReceipt, SubmissionReceiptDisposition, SubmissionSource,
};
use talos_core::tool::{AgentTool, ToolFamily, ToolNature, ToolResult};
use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
pub(crate) const MIN_DELAY_SECS: u64 = 1;
pub(crate) const MAX_DELAY_SECS: u64 = 86_400;
pub(crate) const MIN_INTERVAL_SECS: u64 = 5;
pub(crate) const MAX_INTERVAL_SECS: u64 = 3_600;
const SCHEDULER_COMMAND_CAPACITY: usize = 64;
const DELIVERY_SEND_TIMEOUT: Duration = Duration::from_millis(250);
const DELIVERY_RECEIPT_TIMEOUT: Duration = Duration::from_secs(1);
const DELIVERY_RETRY_DELAY: Duration = Duration::from_secs(1);
pub(crate) const SCHEDULED_FOLLOWUP_LABEL: &str = "[scheduled-followup]";
static NEXT_TASK_SEQ: AtomicU64 = AtomicU64::new(1);
static NEXT_DELIVERY_SEQ: AtomicU64 = AtomicU64::new(1);
pub(crate) fn validate_delay_secs(delay_secs: u64) -> Result<(), String> {
if delay_secs < MIN_DELAY_SECS {
return Err(format!(
"delay_secs must be at least {MIN_DELAY_SECS}; got {delay_secs}"
));
}
if delay_secs > MAX_DELAY_SECS {
return Err(format!(
"delay_secs must be at most {MAX_DELAY_SECS}; got {delay_secs}"
));
}
Ok(())
}
pub(crate) fn validate_interval_secs(interval_secs: u64) -> Result<(), String> {
if interval_secs < MIN_INTERVAL_SECS {
return Err(format!(
"interval_secs must be at least {MIN_INTERVAL_SECS}; got {interval_secs}"
));
}
if interval_secs > MAX_INTERVAL_SECS {
return Err(format!(
"interval_secs must be at most {MAX_INTERVAL_SECS}; got {interval_secs}"
));
}
Ok(())
}
pub(crate) fn label_scheduled_message(message: &str) -> String {
format!("{SCHEDULED_FOLLOWUP_LABEL} {message}")
}
fn validate_scheduled_message(message: &str) -> Result<String, String> {
if message.is_empty() {
return Err("message must not be empty".into());
}
let labeled = label_scheduled_message(message);
if labeled.len() > MAX_SUBMISSION_ITEM_BYTES {
return Err(format!(
"message exceeds the structured submission item limit of {MAX_SUBMISSION_ITEM_BYTES} UTF-8 bytes"
));
}
Ok(labeled)
}
pub(crate) fn next_task_id() -> String {
format!("sched_{}", NEXT_TASK_SEQ.fetch_add(1, Ordering::Relaxed))
}
fn next_delivery_identity(task_id: &str) -> String {
let registration = NEXT_DELIVERY_SEQ.fetch_add(1, Ordering::Relaxed);
format!("scheduler:{task_id}:registration:{registration}")
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ScheduleKind {
OneShot,
Recurring { interval: Duration },
}
impl fmt::Display for ScheduleKind {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::OneShot => write!(f, "one-shot"),
Self::Recurring { interval } => write!(f, "recurring ({}s)", interval.as_secs()),
}
}
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub(crate) struct ScheduledTaskInfo {
pub id: String,
pub message: String,
pub kind: ScheduleKind,
pub created_at: Instant,
pub fire_at: Instant,
}
impl ScheduledTaskInfo {
#[must_use]
pub fn remaining(&self) -> Duration {
self.fire_at.saturating_duration_since(Instant::now())
}
#[must_use]
fn delivery_state(&self) -> &'static str {
if self.remaining().is_zero() {
"delivery-blocked"
} else {
"scheduled"
}
}
}
#[derive(Debug)]
pub(crate) enum ScheduleRegistrationResult {
Registered { task_id: String },
InvalidDuration { reason: String },
InvalidMessage { reason: String },
}
#[derive(Debug)]
pub(crate) enum CancelResult {
Cancelled,
NotFound,
}
#[derive(Debug)]
#[allow(dead_code)]
pub(crate) enum ScheduleCommand {
RegisterOneShot {
id: Option<String>,
message: String,
delay: Duration,
response_tx: oneshot::Sender<ScheduleRegistrationResult>,
},
RegisterRecurring {
id: Option<String>,
message: String,
interval: Duration,
response_tx: oneshot::Sender<ScheduleRegistrationResult>,
},
Cancel {
id: String,
response_tx: oneshot::Sender<CancelResult>,
},
List {
response_tx: oneshot::Sender<Vec<ScheduledTaskInfo>>,
},
Shutdown,
}
#[derive(Clone)]
pub(crate) struct SchedulerHandle {
cmd_tx: mpsc::Sender<ScheduleCommand>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum SchedulerSendError {
Full,
Closed,
}
impl SchedulerHandle {
pub(crate) fn new(cmd_tx: mpsc::Sender<ScheduleCommand>) -> Self {
Self { cmd_tx }
}
pub(crate) async fn send(&self, command: ScheduleCommand) -> Result<(), SchedulerSendError> {
match self.cmd_tx.try_send(command) {
Ok(()) => Ok(()),
Err(mpsc::error::TrySendError::Full(_)) => Err(SchedulerSendError::Full),
Err(mpsc::error::TrySendError::Closed(_)) => Err(SchedulerSendError::Closed),
}
}
}
impl fmt::Debug for SchedulerHandle {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SchedulerHandle")
.field("cmd_tx", &"mpsc::Sender<ScheduleCommand>")
.finish()
}
}
struct ActiveTask {
info: ScheduledTaskInfo,
handle: JoinHandle<()>,
}
struct AcceptedFire {
task_id: String,
next_fire_at: Option<Instant>,
}
pub(crate) struct SchedulerActor {
cmd_rx: mpsc::Receiver<ScheduleCommand>,
sq_tx: mpsc::Sender<SessionOp>,
session_generation: u64,
cancel_token: CancellationToken,
tasks: HashMap<String, ActiveTask>,
accepted_tx: mpsc::UnboundedSender<AcceptedFire>,
accepted_rx: mpsc::UnboundedReceiver<AcceptedFire>,
}
impl SchedulerActor {
pub(crate) fn new(
cmd_rx: mpsc::Receiver<ScheduleCommand>,
sq_tx: mpsc::Sender<SessionOp>,
session_generation: u64,
cancel_token: CancellationToken,
) -> Self {
let (accepted_tx, accepted_rx) = mpsc::unbounded_channel();
Self {
cmd_rx,
sq_tx,
session_generation,
cancel_token,
tasks: HashMap::new(),
accepted_tx,
accepted_rx,
}
}
pub(crate) async fn run(mut self) {
loop {
tokio::select! {
biased;
_ = self.cancel_token.cancelled() => break,
Some(accepted) = self.accepted_rx.recv() => {
match accepted.next_fire_at {
None => {
self.tasks.remove(&accepted.task_id);
}
Some(next_fire_at) => {
if let Some(task) = self.tasks.get_mut(&accepted.task_id) {
task.info.fire_at = next_fire_at;
}
}
}
}
command = self.cmd_rx.recv() => {
match command {
Some(ScheduleCommand::RegisterOneShot {
id,
message,
delay,
response_tx,
}) => self.handle_register_one_shot(id, message, delay, response_tx),
Some(ScheduleCommand::RegisterRecurring {
id,
message,
interval,
response_tx,
}) => self.handle_register_recurring(
id,
message,
interval,
response_tx,
),
Some(ScheduleCommand::Cancel { id, response_tx }) => {
self.handle_cancel(id, response_tx);
}
Some(ScheduleCommand::List { response_tx }) => {
self.handle_list(response_tx);
}
Some(ScheduleCommand::Shutdown) | None => break,
}
}
}
}
for (_, task) in self.tasks.drain() {
task.handle.abort();
}
}
fn handle_register_one_shot(
&mut self,
id: Option<String>,
message: String,
delay: Duration,
response_tx: oneshot::Sender<ScheduleRegistrationResult>,
) {
if let Err(reason) = validate_delay_secs(delay.as_secs()) {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidDuration { reason });
return;
}
let labeled_message = match validate_scheduled_message(&message) {
Ok(message) => message,
Err(reason) => {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidMessage { reason });
return;
}
};
let task_id = id.unwrap_or_else(next_task_id);
if self.tasks.contains_key(&task_id) {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidMessage {
reason: format!("task ID {task_id} is already active"),
});
return;
}
let delivery_identity = next_delivery_identity(&task_id);
let now = Instant::now();
let sq_tx = self.sq_tx.clone();
let session_generation = self.session_generation;
let accepted_tx = self.accepted_tx.clone();
let task_id_for_fire = task_id.clone();
let message_for_fire = labeled_message.clone();
let handle = tokio::spawn(async move {
tokio::time::sleep(delay).await;
let submission =
scheduled_submission(&delivery_identity, 1, session_generation, message_for_fire);
deliver_until_accepted(&sq_tx, submission).await;
let _ = accepted_tx.send(AcceptedFire {
task_id: task_id_for_fire,
next_fire_at: None,
});
});
self.tasks.insert(
task_id.clone(),
ActiveTask {
info: ScheduledTaskInfo {
id: task_id.clone(),
message: labeled_message,
kind: ScheduleKind::OneShot,
created_at: now,
fire_at: now + delay,
},
handle,
},
);
let _ = response_tx.send(ScheduleRegistrationResult::Registered { task_id });
}
fn handle_register_recurring(
&mut self,
id: Option<String>,
message: String,
interval: Duration,
response_tx: oneshot::Sender<ScheduleRegistrationResult>,
) {
if let Err(reason) = validate_interval_secs(interval.as_secs()) {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidDuration { reason });
return;
}
let labeled_message = match validate_scheduled_message(&message) {
Ok(message) => message,
Err(reason) => {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidMessage { reason });
return;
}
};
let task_id = id.unwrap_or_else(next_task_id);
if self.tasks.contains_key(&task_id) {
let _ = response_tx.send(ScheduleRegistrationResult::InvalidMessage {
reason: format!("task ID {task_id} is already active"),
});
return;
}
let delivery_identity = next_delivery_identity(&task_id);
let now = Instant::now();
let sq_tx = self.sq_tx.clone();
let session_generation = self.session_generation;
let accepted_tx = self.accepted_tx.clone();
let task_id_for_fire = task_id.clone();
let message_for_fire = labeled_message.clone();
let handle = tokio::spawn(async move {
let mut fire_sequence = 1_u64;
tokio::time::sleep(interval).await;
loop {
let submission = scheduled_submission(
&delivery_identity,
fire_sequence,
session_generation,
message_for_fire.clone(),
);
deliver_until_accepted(&sq_tx, submission).await;
let next_fire_at = Instant::now() + interval;
if accepted_tx
.send(AcceptedFire {
task_id: task_id_for_fire.clone(),
next_fire_at: Some(next_fire_at),
})
.is_err()
{
break;
}
let Some(next_sequence) = fire_sequence.checked_add(1) else {
tracing::warn!(
task_id = %task_id_for_fire,
"recurring scheduler exhausted its fire sequence"
);
break;
};
fire_sequence = next_sequence;
tokio::time::sleep(interval).await;
}
});
self.tasks.insert(
task_id.clone(),
ActiveTask {
info: ScheduledTaskInfo {
id: task_id.clone(),
message: labeled_message,
kind: ScheduleKind::Recurring { interval },
created_at: now,
fire_at: now + interval,
},
handle,
},
);
let _ = response_tx.send(ScheduleRegistrationResult::Registered { task_id });
}
fn handle_cancel(&mut self, id: String, response_tx: oneshot::Sender<CancelResult>) {
if let Some(task) = self.tasks.remove(&id) {
task.handle.abort();
let _ = response_tx.send(CancelResult::Cancelled);
} else {
let _ = response_tx.send(CancelResult::NotFound);
}
}
fn handle_list(&self, response_tx: oneshot::Sender<Vec<ScheduledTaskInfo>>) {
let snapshot = self.tasks.values().map(|task| task.info.clone()).collect();
let _ = response_tx.send(snapshot);
}
}
fn scheduled_submission(
delivery_identity: &str,
fire_sequence: u64,
session_generation: u64,
text: String,
) -> StructuredSubmission {
let submission_id = format!("{delivery_identity}:fire:{fire_sequence}");
StructuredSubmission {
id: submission_id.clone(),
source: SubmissionSource::Scheduler,
sender_generation: session_generation,
items: vec![SubmissionItem {
id: format!("{submission_id}:item:1"),
enqueue_sequence: fire_sequence,
kind: SubmissionKind::UserTurn,
text,
attachments: Vec::new(),
}],
}
}
async fn deliver_until_accepted(sq_tx: &mpsc::Sender<SessionOp>, submission: StructuredSubmission) {
let (receipt_tx, mut receipt_rx) = mpsc::unbounded_channel();
let mut submit_required = true;
let mut attempts = 0_u64;
loop {
attempts = attempts.saturating_add(1);
let operation = if submit_required {
SessionOp::SubmitStructuredTracked {
submission: submission.clone(),
receipt_tx: Some(receipt_tx.clone()),
}
} else {
SessionOp::ReconcileStructuredTracked {
submission: submission.clone(),
receipt_tx: Some(receipt_tx.clone()),
}
};
match tokio::time::timeout(DELIVERY_SEND_TIMEOUT, sq_tx.send(operation)).await {
Ok(Ok(())) => {}
Ok(Err(_)) | Err(_) => {
log_blocked_delivery(&submission.id, attempts, "session queue unavailable");
tokio::time::sleep(DELIVERY_RETRY_DELAY).await;
continue;
}
}
let receipt = tokio::time::timeout(
DELIVERY_RECEIPT_TIMEOUT,
next_matching_receipt(&mut receipt_rx, &submission.id),
)
.await;
match receipt {
Ok(Some(receipt)) if receipt.disposition.has_durable_custody() => return,
Ok(Some(receipt)) => {
submit_required = match receipt.disposition {
SubmissionReceiptDisposition::NotAccepted => true,
SubmissionReceiptDisposition::Rejected { reason } => {
if attempts == 1 || attempts.is_power_of_two() {
tracing::warn!(
submission_id = %submission.id,
?reason,
attempts,
"scheduled fire remains blocked before durable custody"
);
}
true
}
SubmissionReceiptDisposition::AcceptedPending
| SubmissionReceiptDisposition::AlreadyAccepted { .. } => false,
};
}
Ok(None) | Err(_) => {
submit_required = false;
log_blocked_delivery(&submission.id, attempts, "durable receipt unavailable");
}
}
tokio::time::sleep(DELIVERY_RETRY_DELAY).await;
}
}
async fn next_matching_receipt(
receipt_rx: &mut mpsc::UnboundedReceiver<SubmissionReceipt>,
submission_id: &str,
) -> Option<SubmissionReceipt> {
while let Some(receipt) = receipt_rx.recv().await {
if receipt.submission_id == submission_id && receipt.source == SubmissionSource::Scheduler {
return Some(receipt);
}
}
None
}
fn log_blocked_delivery(submission_id: &str, attempts: u64, reason: &str) {
if attempts == 1 || attempts.is_power_of_two() {
tracing::warn!(
submission_id,
attempts,
reason,
"scheduled fire retained for bounded retry"
);
}
}
#[allow(dead_code)]
pub(crate) fn spawn_scheduler_actor(
sq_tx: mpsc::Sender<SessionOp>,
session_generation: u64,
cancel_token: CancellationToken,
) -> (SchedulerHandle, JoinHandle<()>) {
let (cmd_tx, cmd_rx) = mpsc::channel(SCHEDULER_COMMAND_CAPACITY);
let handle = SchedulerHandle::new(cmd_tx);
let actor = SchedulerActor::new(cmd_rx, sq_tx, session_generation, cancel_token);
let join = tokio::spawn(async move { actor.run().await });
(handle, join)
}
pub fn create_delay_tool_and_scheduler() -> (Arc<dyn AgentTool>, PendingSchedulerActor) {
let (cmd_tx, cmd_rx) = mpsc::channel(SCHEDULER_COMMAND_CAPACITY);
let handle = SchedulerHandle::new(cmd_tx);
let tool: Arc<dyn AgentTool> = Arc::new(DelayTool::new(handle));
(tool, PendingSchedulerActor { cmd_rx })
}
pub fn create_scheduler_tools() -> (Vec<Arc<dyn AgentTool>>, PendingSchedulerActor) {
let (cmd_tx, cmd_rx) = mpsc::channel(SCHEDULER_COMMAND_CAPACITY);
let handle = SchedulerHandle::new(cmd_tx);
let tools: Vec<Arc<dyn AgentTool>> = vec![
Arc::new(DelayTool::new(handle.clone())),
Arc::new(ScheduleTool::new(handle.clone())),
Arc::new(ListScheduledTasksTool::new(handle.clone())),
Arc::new(CancelScheduledTaskTool::new(handle)),
];
(tools, PendingSchedulerActor { cmd_rx })
}
pub struct PendingSchedulerActor {
cmd_rx: mpsc::Receiver<ScheduleCommand>,
}
impl PendingSchedulerActor {
pub fn spawn(
self,
sq_tx: mpsc::Sender<SessionOp>,
session_generation: u64,
cancel_token: CancellationToken,
) -> JoinHandle<()> {
let actor = SchedulerActor::new(self.cmd_rx, sq_tx, session_generation, cancel_token);
tokio::spawn(async move { actor.run().await })
}
}
impl fmt::Debug for PendingSchedulerActor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PendingSchedulerActor")
.field("cmd_rx", &"mpsc::Receiver<ScheduleCommand>")
.finish()
}
}
pub(crate) struct DelayTool {
handle: SchedulerHandle,
}
impl DelayTool {
pub(crate) fn new(handle: SchedulerHandle) -> Self {
Self { handle }
}
}
#[async_trait]
impl AgentTool for DelayTool {
fn name(&self) -> &str {
"delay"
}
fn description(&self) -> &str {
"Schedule a one-shot delayed follow-up message. Delivery transfers only after the Session Actor durably accepts the exact scheduled fire. Session-scoped; minimum 1 second, maximum 86400 seconds."
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"message": {"type": "string", "description": "Follow-up message."},
"delay_secs": {
"type": "integer",
"minimum": MIN_DELAY_SECS,
"maximum": MAX_DELAY_SECS
}
},
"required": ["message", "delay_secs"]
})
}
async fn execute(&self, input: serde_json::Value) -> ToolResult {
let message = match input.get("message").and_then(|value| value.as_str()) {
Some(message) if !message.is_empty() => message.to_owned(),
_ => return ToolResult::error("missing or empty 'message' field"),
};
let delay_secs = match input.get("delay_secs").and_then(|value| value.as_u64()) {
Some(delay_secs) => delay_secs,
None => {
return ToolResult::error(
"missing or invalid 'delay_secs' field (expected a positive integer)",
);
}
};
if let Err(reason) = validate_delay_secs(delay_secs) {
return ToolResult::error(reason);
}
let (response_tx, response_rx) = oneshot::channel();
if let Err(error) = self
.handle
.send(ScheduleCommand::RegisterOneShot {
id: None,
message,
delay: Duration::from_secs(delay_secs),
response_tx,
})
.await
{
return scheduler_send_error(error);
}
registration_result(response_rx.await, format!("Delay: {delay_secs} second(s)"))
}
fn nature(&self) -> ToolNature {
ToolNature::Execute
}
fn family(&self) -> ToolFamily {
ToolFamily::Extension
}
}
pub(crate) struct ScheduleTool {
handle: SchedulerHandle,
}
impl ScheduleTool {
pub(crate) fn new(handle: SchedulerHandle) -> Self {
Self { handle }
}
}
#[async_trait]
impl AgentTool for ScheduleTool {
fn name(&self) -> &str {
"schedule"
}
fn description(&self) -> &str {
"Schedule a recurring follow-up. Each fire uses a stable structured identity and the next interval does not begin until Actor durable acceptance. Session-scoped; interval 5 to 3600 seconds."
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"message": {"type": "string", "description": "Recurring follow-up message."},
"interval_secs": {
"type": "integer",
"minimum": MIN_INTERVAL_SECS,
"maximum": MAX_INTERVAL_SECS
}
},
"required": ["message", "interval_secs"]
})
}
async fn execute(&self, input: serde_json::Value) -> ToolResult {
let message = match input.get("message").and_then(|value| value.as_str()) {
Some(message) if !message.is_empty() => message.to_owned(),
_ => return ToolResult::error("missing or empty 'message' field"),
};
let interval_secs = match input.get("interval_secs").and_then(|value| value.as_u64()) {
Some(interval_secs) => interval_secs,
None => {
return ToolResult::error(
"missing or invalid 'interval_secs' field (expected a positive integer)",
);
}
};
if let Err(reason) = validate_interval_secs(interval_secs) {
return ToolResult::error(reason);
}
let (response_tx, response_rx) = oneshot::channel();
if let Err(error) = self
.handle
.send(ScheduleCommand::RegisterRecurring {
id: None,
message,
interval: Duration::from_secs(interval_secs),
response_tx,
})
.await
{
return scheduler_send_error(error);
}
registration_result(
response_rx.await,
format!("Interval: {interval_secs} second(s)"),
)
}
fn nature(&self) -> ToolNature {
ToolNature::Execute
}
fn family(&self) -> ToolFamily {
ToolFamily::Extension
}
}
pub(crate) struct ListScheduledTasksTool {
handle: SchedulerHandle,
}
impl ListScheduledTasksTool {
pub(crate) fn new(handle: SchedulerHandle) -> Self {
Self { handle }
}
}
#[async_trait]
impl AgentTool for ListScheduledTasksTool {
fn name(&self) -> &str {
"list_scheduled_tasks"
}
fn description(&self) -> &str {
"List active scheduled tasks, including whether a due fire is blocked awaiting durable Actor custody."
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({"type": "object", "properties": {}})
}
async fn execute(&self, _input: serde_json::Value) -> ToolResult {
let (response_tx, response_rx) = oneshot::channel();
if let Err(error) = self
.handle
.send(ScheduleCommand::List { response_tx })
.await
{
return scheduler_send_error(error);
}
match response_rx.await {
Ok(tasks) if tasks.is_empty() => ToolResult::success("No active scheduled tasks."),
Ok(tasks) => {
const MAX_DISPLAY: usize = 20;
let total = tasks.len();
let mut text = format!("{total} active task(s):\n");
for info in tasks.iter().take(MAX_DISPLAY) {
text.push_str(&format!(
" {} | {} | next: {}s | {}\n",
info.id,
info.kind,
info.remaining().as_secs(),
info.delivery_state(),
));
}
let omitted = total.saturating_sub(MAX_DISPLAY);
if omitted > 0 {
text.push_str(&format!("... and {omitted} more task(s) not shown\n"));
}
ToolResult::success(text.trim_end().to_owned())
}
Err(_) => ToolResult::error("scheduler dropped the request"),
}
}
fn nature(&self) -> ToolNature {
ToolNature::Read
}
fn family(&self) -> ToolFamily {
ToolFamily::Extension
}
}
pub(crate) struct CancelScheduledTaskTool {
handle: SchedulerHandle,
}
impl CancelScheduledTaskTool {
pub(crate) fn new(handle: SchedulerHandle) -> Self {
Self { handle }
}
}
#[async_trait]
impl AgentTool for CancelScheduledTaskTool {
fn name(&self) -> &str {
"cancel_scheduled_task"
}
fn description(&self) -> &str {
"Cancel an active scheduled follow-up by task ID."
}
fn parameters(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {"task_id": {"type": "string"}},
"required": ["task_id"]
})
}
async fn execute(&self, input: serde_json::Value) -> ToolResult {
let task_id = match input.get("task_id").and_then(|value| value.as_str()) {
Some(task_id) if !task_id.is_empty() => task_id.to_owned(),
_ => return ToolResult::error("missing or empty 'task_id' field"),
};
let (response_tx, response_rx) = oneshot::channel();
if let Err(error) = self
.handle
.send(ScheduleCommand::Cancel {
id: task_id.clone(),
response_tx,
})
.await
{
return scheduler_send_error(error);
}
match response_rx.await {
Ok(CancelResult::Cancelled) => {
ToolResult::success(format!("Task {task_id} cancelled."))
}
Ok(CancelResult::NotFound) => {
ToolResult::success(format!("Task {task_id} not found or already completed."))
}
Err(_) => ToolResult::error("scheduler dropped the request"),
}
}
fn nature(&self) -> ToolNature {
ToolNature::Execute
}
fn family(&self) -> ToolFamily {
ToolFamily::Extension
}
fn summary_fields(&self) -> &'static [&'static str] {
&["task_id"]
}
}
fn scheduler_send_error(error: SchedulerSendError) -> ToolResult {
match error {
SchedulerSendError::Full => ToolResult::error("scheduler is busy; try again"),
SchedulerSendError::Closed => ToolResult::error("scheduler is not available"),
}
}
fn registration_result(
result: Result<ScheduleRegistrationResult, oneshot::error::RecvError>,
timing: String,
) -> ToolResult {
match result {
Ok(ScheduleRegistrationResult::Registered { task_id }) => ToolResult::success(format!(
"Scheduled follow-up registered.\nTask ID: {task_id}\n{timing}"
)),
Ok(ScheduleRegistrationResult::InvalidDuration { reason })
| Ok(ScheduleRegistrationResult::InvalidMessage { reason }) => ToolResult::error(reason),
Err(_) => ToolResult::error("scheduler dropped the request"),
}
}
#[cfg(test)]
mod tests {
use super::*;
use talos_core::session::{PendingSubmissionState, SubmissionRejectionReason};
async fn yield_times(times: usize) {
for _ in 0..times {
tokio::task::yield_now().await;
}
}
async fn advance_delivery_retry() {
yield_times(10).await;
tokio::time::advance(DELIVERY_RECEIPT_TIMEOUT).await;
yield_times(10).await;
tokio::time::advance(DELIVERY_RETRY_DELAY).await;
yield_times(20).await;
}
async fn register_one_shot(handle: &SchedulerHandle, id: &str, delay: Duration) {
let (response_tx, response_rx) = oneshot::channel();
handle
.send(ScheduleCommand::RegisterOneShot {
id: Some(id.to_owned()),
message: "check the build".into(),
delay,
response_tx,
})
.await
.expect("operation should succeed");
assert!(matches!(
response_rx.await.expect("operation should succeed"),
ScheduleRegistrationResult::Registered { .. }
));
}
fn split_tracked_operation(
operation: SessionOp,
) -> (
bool,
StructuredSubmission,
mpsc::UnboundedSender<SubmissionReceipt>,
) {
match operation {
SessionOp::SubmitStructuredTracked {
submission,
receipt_tx: Some(receipt_tx),
} => (true, submission, receipt_tx),
SessionOp::ReconcileStructuredTracked {
submission,
receipt_tx: Some(receipt_tx),
} => (false, submission, receipt_tx),
other => panic!("expected tracked structured operation, got {other:?}"),
}
}
fn accept(
submission: &StructuredSubmission,
receipt_tx: &mpsc::UnboundedSender<SubmissionReceipt>,
disposition: SubmissionReceiptDisposition,
) {
receipt_tx
.send(SubmissionReceipt {
session_id: "session-test".into(),
session_generation: submission.sender_generation,
submission_id: submission.id.clone(),
reservation_id: submission.id.clone(),
receipt_id: "receipt-test".into(),
source: submission.source,
item_count: submission.items.len(),
total_text_bytes: submission.total_text_bytes(),
disposition,
})
.expect("operation should succeed");
}
async fn list(handle: &SchedulerHandle) -> Vec<ScheduledTaskInfo> {
let (response_tx, response_rx) = oneshot::channel();
handle
.send(ScheduleCommand::List { response_tx })
.await
.expect("operation should succeed");
response_rx.await.expect("operation should succeed")
}
#[test]
fn validates_bounds_and_label() {
assert!(validate_delay_secs(MIN_DELAY_SECS).is_ok());
assert!(validate_delay_secs(0).is_err());
assert!(validate_interval_secs(MIN_INTERVAL_SECS).is_ok());
assert!(validate_interval_secs(MIN_INTERVAL_SECS - 1).is_err());
assert_eq!(
label_scheduled_message("check"),
"[scheduled-followup] check"
);
}
#[tokio::test(start_paused = true)]
async fn scheduled_fire_uses_the_bound_actor_generation() {
let (sq_tx, mut sq_rx) = mpsc::channel(8);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 7, CancellationToken::new());
register_one_shot(&handle, "generation-bound", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let (is_submit, submission, receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert!(is_submit);
assert_eq!(submission.source, SubmissionSource::Scheduler);
assert_eq!(submission.sender_generation, 7);
accept(
&submission,
&receipt_tx,
SubmissionReceiptDisposition::AcceptedPending,
);
yield_times(10).await;
assert!(list(&handle).await.is_empty());
}
#[tokio::test(start_paused = true)]
async fn one_shot_transfers_only_after_durable_receipt() {
let (sq_tx, mut sq_rx) = mpsc::channel(8);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
register_one_shot(&handle, "one", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let (is_submit, submission, receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert!(is_submit);
assert_eq!(submission.source, SubmissionSource::Scheduler);
assert!(
submission.items[0]
.text
.starts_with(SCHEDULED_FOLLOWUP_LABEL)
);
assert_eq!(list(&handle).await.len(), 1);
accept(
&submission,
&receipt_tx,
SubmissionReceiptDisposition::AcceptedPending,
);
yield_times(10).await;
assert!(list(&handle).await.is_empty());
}
#[tokio::test(start_paused = true)]
async fn lost_ack_reconciles_the_exact_fire_identity() {
let (sq_tx, mut sq_rx) = mpsc::channel(8);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
register_one_shot(&handle, "lost-ack", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let (_, submitted, _lost_receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
advance_delivery_retry().await;
let (is_submit, reconciled, receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert!(!is_submit);
assert_eq!(reconciled, submitted);
accept(
&reconciled,
&receipt_tx,
SubmissionReceiptDisposition::AlreadyAccepted {
state: PendingSubmissionState::Running,
turn_id: Some("turn-1".into()),
},
);
yield_times(10).await;
assert!(list(&handle).await.is_empty());
}
#[tokio::test(start_paused = true)]
async fn closed_session_queue_retains_a_blocked_fire() {
let (sq_tx, sq_rx) = mpsc::channel(1);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
drop(sq_rx);
register_one_shot(&handle, "closed", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(5)).await;
yield_times(20).await;
let tasks = list(&handle).await;
assert_eq!(tasks.len(), 1);
assert_eq!(tasks[0].delivery_state(), "delivery-blocked");
}
#[tokio::test(start_paused = true)]
async fn rejection_retries_without_changing_identity() {
let (sq_tx, mut sq_rx) = mpsc::channel(8);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
register_one_shot(&handle, "rejected", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let (_, first, receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
accept(
&first,
&receipt_tx,
SubmissionReceiptDisposition::Rejected {
reason: SubmissionRejectionReason::LimitExceeded,
},
);
yield_times(10).await;
tokio::time::advance(DELIVERY_RETRY_DELAY).await;
yield_times(20).await;
let (is_submit, retry, _) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert!(is_submit);
assert_eq!(retry, first);
assert_eq!(list(&handle).await.len(), 1);
}
#[tokio::test(start_paused = true)]
async fn recurring_does_not_replace_an_unresolved_fire() {
let (sq_tx, mut sq_rx) = mpsc::channel(32);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
let (response_tx, response_rx) = oneshot::channel();
handle
.send(ScheduleCommand::RegisterRecurring {
id: Some("recurring".into()),
message: "tick".into(),
interval: Duration::from_secs(5),
response_tx,
})
.await
.expect("operation should succeed");
assert!(response_rx.await.is_ok());
tokio::time::advance(Duration::from_secs(6)).await;
yield_times(10).await;
let (_, first, _) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
advance_delivery_retry().await;
let (is_submit, retry, receipt_tx) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert!(!is_submit);
assert_eq!(
retry.id, first.id,
"later interval replaced unresolved fire"
);
assert!(sq_rx.try_recv().is_err());
accept(
&retry,
&receipt_tx,
SubmissionReceiptDisposition::AcceptedPending,
);
yield_times(20).await;
tokio::time::advance(Duration::from_secs(4)).await;
yield_times(10).await;
assert!(sq_rx.try_recv().is_err());
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let (_, second, _) =
split_tracked_operation(sq_rx.try_recv().expect("operation should succeed"));
assert_ne!(second.id, first.id);
assert!(second.id.ends_with(":fire:2"));
}
#[tokio::test(start_paused = true)]
async fn cancellation_aborts_an_unaccepted_fire() {
let (sq_tx, mut sq_rx) = mpsc::channel(8);
let (handle, _join) = spawn_scheduler_actor(sq_tx, 0, CancellationToken::new());
register_one_shot(&handle, "cancel", Duration::from_secs(1)).await;
tokio::time::advance(Duration::from_secs(2)).await;
yield_times(10).await;
let _ = sq_rx.try_recv().expect("operation should succeed");
let (response_tx, response_rx) = oneshot::channel();
handle
.send(ScheduleCommand::Cancel {
id: "cancel".into(),
response_tx,
})
.await
.expect("operation should succeed");
assert!(matches!(
response_rx.await.expect("operation should succeed"),
CancelResult::Cancelled
));
tokio::time::advance(Duration::from_secs(10)).await;
yield_times(20).await;
assert!(sq_rx.try_recv().is_err());
}
#[test]
fn tool_natures_remain_permission_safe() {
let (tools, _pending) = create_scheduler_tools();
assert_eq!(tools[0].nature(), ToolNature::Execute);
assert_eq!(tools[1].nature(), ToolNature::Execute);
assert_eq!(tools[2].nature(), ToolNature::Read);
assert_eq!(tools[3].nature(), ToolNature::Execute);
}
}