PicoBot/src/tools/complete_scheduled_run.rs

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<ToolResult> {
        self.execute_with_context(&ToolExecutionContext::default(), args)
            .await
            .map(|output| output.result)
    }

    async fn execute_with_context(
        &self,
        context: &ToolExecutionContext,
        args: Value,
    ) -> anyhow::Result<ToolOutput> {
        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<String>) -> 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
        );
    }
}