use async_trait::async_trait; use serde_json::{Value, json}; use crate::storage::ScheduledOutcomeKind; use crate::tools::{ScheduledOutcome, Tool, ToolExecutionContext, ToolOutput, ToolResult}; pub struct CompleteScheduledRunTool; impl CompleteScheduledRunTool { pub fn new() -> Self { Self } } impl Default for CompleteScheduledRunTool { fn default() -> Self { Self::new() } } #[async_trait] impl Tool for CompleteScheduledRunTool { fn name(&self) -> &str { "complete_scheduled_run" } fn description(&self) -> &str { "Submit the single structured final outcome (ok, alert, failed, or refused) of an unattended scheduled run and end the run immediately." } fn parameters_schema(&self) -> Value { json!({ "type": "object", "properties": { "outcome": { "type": "string", "enum": ["ok", "alert", "failed", "refused"], "description": "ok: completed with nothing needing attention; alert: completed with actionable findings; failed: did not complete reliably; refused: denied for permission or safety reasons" }, "message": { "type": "string", "minLength": 1, "maxLength": 16384, "description": "The user-facing result or notification body" } }, "required": ["outcome", "message"], "additionalProperties": false }) } fn runtime_injected(&self) -> bool { true } fn exclusive(&self) -> bool { true } async fn execute(&self, args: Value) -> anyhow::Result { self.execute_with_context(&ToolExecutionContext::default(), args) .await .map(|output| output.result) } async fn execute_with_context( &self, context: &ToolExecutionContext, args: Value, ) -> anyhow::Result { let Some(sink) = context.scheduled_completion.as_ref() else { return Ok(failure( "complete_scheduled_run is only available to the top-level scheduled Agent", ) .into()); }; if !context.execution_origin.is_scheduled() { return Ok(failure("scheduled completion context is invalid").into()); } let Some(object) = args.as_object() else { return Ok(failure("arguments must be an object").into()); }; if object .keys() .any(|key| key != "outcome" && key != "message") { return Ok(failure("unknown complete_scheduled_run argument").into()); } let kind = match args.get("outcome").and_then(Value::as_str) { Some("ok") => ScheduledOutcomeKind::Ok, Some("alert") => ScheduledOutcomeKind::Alert, Some("failed") => ScheduledOutcomeKind::Failed, Some("refused") => ScheduledOutcomeKind::Refused, Some(other) => return Ok(failure(format!("invalid scheduled outcome: {other}")).into()), None => return Ok(failure("outcome is required").into()), }; let message = args .get("message") .and_then(Value::as_str) .unwrap_or_default() .trim(); if message.is_empty() { return Ok(failure("message must not be empty").into()); } if message.chars().count() > 16_384 { return Ok(failure("message exceeds 16384 characters").into()); } let message = message.to_string(); if let Err(error) = sink.submit(ScheduledOutcome { kind, message: message.clone(), }) { return Ok(failure(error).into()); } Ok(ToolResult { success: true, output: format!("scheduled run completed with outcome={}", kind.as_str()), error: None, } .into()) } } fn failure(error: impl Into) -> ToolResult { ToolResult { success: false, output: String::new(), error: Some(error.into()), } } #[cfg(test)] mod tests { use super::*; use crate::tools::{ExecutionOrigin, ScheduledCompletionSink}; use std::sync::Arc; #[tokio::test] async fn submits_exactly_once_in_scheduled_context() { let sink = Arc::new(ScheduledCompletionSink::default()); let context = ToolExecutionContext::for_session("scheduled-run:1") .with_execution_origin(ExecutionOrigin::Scheduled { job_run_id: 1 }) .with_scheduled_completion(sink.clone()); let tool = CompleteScheduledRunTool::new(); let first = tool .execute_with_context(&context, json!({"outcome":"ok","message":"healthy"})) .await .unwrap(); assert!(first.result.success); assert_eq!(sink.outcome().unwrap().message, "healthy"); let second = tool .execute_with_context(&context, json!({"outcome":"alert","message":"again"})) .await .unwrap(); assert!(!second.result.success); } #[tokio::test] async fn rejects_empty_extra_and_oversized_arguments() { let sink = Arc::new(ScheduledCompletionSink::default()); let context = ToolExecutionContext::for_session("scheduled-run:1") .with_execution_origin(ExecutionOrigin::Scheduled { job_run_id: 1 }) .with_scheduled_completion(sink); let tool = CompleteScheduledRunTool::new(); for args in [ json!({"outcome":"ok","message":" "}), json!({"outcome":"ok","message":"fine","notify":false}), json!({"outcome":"unknown","message":"fine"}), ] { assert!( !tool .execute_with_context(&context, args) .await .unwrap() .result .success ); } assert!( !tool .execute_with_context( &context, json!({"outcome":"ok","message":"x".repeat(16_385)}), ) .await .unwrap() .result .success ); } }