fix: bound scheduled recipe validation (#10509)
This commit is contained in:
@@ -240,6 +240,8 @@ keyring = { workspace = true, features = ["apple-native"], optional = true }
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
keyring = { workspace = true, features = ["sync-secret-service"], optional = true }
|
||||
|
||||
[target."cfg(unix)".dependencies]
|
||||
libc = { version = "0.2.182", default-features = false, features = ["std"] }
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
@@ -3,6 +3,8 @@
|
||||
//! This module contains all the handlers for the schedule management platform tool,
|
||||
//! including job creation, execution, monitoring, and session management.
|
||||
|
||||
use std::io::Read;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::mcp_utils::ToolResult;
|
||||
@@ -11,8 +13,55 @@ use rmcp::model::{Content, ErrorCode, ErrorData};
|
||||
|
||||
use super::Agent;
|
||||
use crate::recipe::Recipe;
|
||||
use crate::scheduler::{
|
||||
open_regular_schedule_recipe, ValidatedScheduleRecipe, MAX_SCHEDULE_RECIPE_BYTES,
|
||||
};
|
||||
use crate::scheduler_trait::SchedulerTrait;
|
||||
|
||||
fn recipe_file_error(message: &str) -> ErrorData {
|
||||
ErrorData::new(ErrorCode::INTERNAL_ERROR, message.to_string(), None)
|
||||
}
|
||||
|
||||
fn read_schedule_recipe(path: &Path) -> Result<(String, PathBuf), ErrorData> {
|
||||
let canonical_path = path
|
||||
.canonicalize()
|
||||
.map_err(|_| recipe_file_error("Cannot read recipe file"))?;
|
||||
let file = open_regular_schedule_recipe(&canonical_path).map_err(|error| {
|
||||
if error.kind() == std::io::ErrorKind::InvalidInput {
|
||||
recipe_file_error("Recipe path must reference a regular file")
|
||||
} else {
|
||||
recipe_file_error("Cannot read recipe file")
|
||||
}
|
||||
})?;
|
||||
let opened_metadata = file
|
||||
.metadata()
|
||||
.map_err(|_| recipe_file_error("Cannot read recipe file"))?;
|
||||
if !opened_metadata.is_file() {
|
||||
return Err(recipe_file_error(
|
||||
"Recipe path must reference a regular file",
|
||||
));
|
||||
}
|
||||
if opened_metadata.len() > MAX_SCHEDULE_RECIPE_BYTES {
|
||||
return Err(recipe_file_error(
|
||||
"Recipe file exceeds the 1048576 byte limit",
|
||||
));
|
||||
}
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
file.take(MAX_SCHEDULE_RECIPE_BYTES + 1)
|
||||
.read_to_end(&mut bytes)
|
||||
.map_err(|_| recipe_file_error("Cannot read recipe file"))?;
|
||||
if bytes.len() as u64 > MAX_SCHEDULE_RECIPE_BYTES {
|
||||
return Err(recipe_file_error(
|
||||
"Recipe file exceeds the 1048576 byte limit",
|
||||
));
|
||||
}
|
||||
|
||||
let content = String::from_utf8(bytes)
|
||||
.map_err(|_| recipe_file_error("Recipe file must be valid UTF-8"))?;
|
||||
Ok((content, canonical_path))
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
/// Handle schedule management tool calls
|
||||
pub async fn handle_schedule_management(
|
||||
@@ -109,50 +158,24 @@ impl Agent {
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("background");
|
||||
|
||||
if !std::path::Path::new(recipe_path).exists() {
|
||||
return Err(ErrorData::new(
|
||||
ErrorCode::INTERNAL_ERROR,
|
||||
format!("Recipe file not found: {}", recipe_path),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
// Validate it's a valid recipe by trying to parse it
|
||||
match std::fs::read_to_string(recipe_path) {
|
||||
Ok(content) => {
|
||||
if recipe_path.ends_with(".json") {
|
||||
serde_json::from_str::<Recipe>(&content).map_err(|e| {
|
||||
ErrorData::new(
|
||||
ErrorCode::INTERNAL_ERROR,
|
||||
format!("Invalid JSON recipe: {}", e),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
} else {
|
||||
serde_yaml::from_str::<Recipe>(&content).map_err(|e| {
|
||||
ErrorData::new(
|
||||
ErrorCode::INTERNAL_ERROR,
|
||||
format!("Invalid YAML recipe: {}", e),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(ErrorData::new(
|
||||
ErrorCode::INTERNAL_ERROR,
|
||||
format!("Cannot read recipe file: {}", e),
|
||||
None,
|
||||
))
|
||||
}
|
||||
let (content, canonical_recipe_path) = read_schedule_recipe(Path::new(recipe_path))?;
|
||||
if recipe_path.ends_with(".json") {
|
||||
serde_json::from_str::<Recipe>(&content)
|
||||
.map_err(|_| recipe_file_error("Invalid JSON recipe"))?;
|
||||
} else {
|
||||
serde_yaml::from_str::<Recipe>(&content)
|
||||
.map_err(|_| recipe_file_error("Invalid YAML recipe"))?;
|
||||
}
|
||||
|
||||
// Generate unique job ID
|
||||
let job_id = format!("agent_created_{}", Utc::now().timestamp());
|
||||
|
||||
let recipe_base_dir = canonical_recipe_path
|
||||
.parent()
|
||||
.map(|path| path.to_string_lossy().into_owned());
|
||||
let job = crate::scheduler::ScheduledJob {
|
||||
id: job_id.clone(),
|
||||
source: recipe_path.to_string(),
|
||||
source: canonical_recipe_path.to_string_lossy().into_owned(),
|
||||
cron: cron_expression.to_string(),
|
||||
last_run: None,
|
||||
currently_running: false,
|
||||
@@ -160,10 +183,16 @@ impl Agent {
|
||||
current_session_id: None,
|
||||
process_start_time: None,
|
||||
parameters: vec![],
|
||||
recipe_base_dir: None,
|
||||
recipe_base_dir,
|
||||
};
|
||||
|
||||
match scheduler.add_scheduled_job(job, true).await {
|
||||
match scheduler
|
||||
.add_scheduled_job_with_recipe(
|
||||
job,
|
||||
ValidatedScheduleRecipe::new(content.into_bytes(), canonical_recipe_path),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(()) => Ok(vec![Content::text(format!(
|
||||
"Successfully created scheduled job '{}' for recipe '{}' with cron expression '{}' in {} mode",
|
||||
job_id, recipe_path, cron_expression, execution_mode
|
||||
|
||||
+276
-16
@@ -1,6 +1,6 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::io;
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::{self, Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -30,6 +30,116 @@ use crate::session::{Session, SessionManager};
|
||||
type RunningTasksMap = HashMap<String, CancellationToken>;
|
||||
type JobsMap = HashMap<String, (JobId, ScheduledJob)>;
|
||||
|
||||
pub(crate) const MAX_SCHEDULE_RECIPE_BYTES: u64 = 1024 * 1024;
|
||||
|
||||
pub struct ValidatedScheduleRecipe {
|
||||
bytes: Vec<u8>,
|
||||
source: PathBuf,
|
||||
}
|
||||
|
||||
impl ValidatedScheduleRecipe {
|
||||
pub(crate) fn new(bytes: Vec<u8>, source: PathBuf) -> Self {
|
||||
Self { bytes, source }
|
||||
}
|
||||
|
||||
pub fn bytes(&self) -> &[u8] {
|
||||
&self.bytes
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn open_regular_schedule_recipe(path: &Path) -> io::Result<File> {
|
||||
let metadata = fs::metadata(path)?;
|
||||
if !metadata.is_file() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"Recipe path must reference a regular file",
|
||||
));
|
||||
}
|
||||
|
||||
let mut options = OpenOptions::new();
|
||||
options.read(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.custom_flags(libc::O_NONBLOCK | libc::O_NOFOLLOW);
|
||||
}
|
||||
let file = options.open(path)?;
|
||||
if !file.metadata()?.is_file() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"Recipe path must reference a regular file",
|
||||
));
|
||||
}
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
fn copy_bounded_schedule_recipe(source: &Path, destination: &Path) -> Result<(), SchedulerError> {
|
||||
let source = open_regular_schedule_recipe(source).map_err(|error| {
|
||||
SchedulerError::RecipeLoadError(format!("Cannot read recipe file: {error}"))
|
||||
})?;
|
||||
let metadata = source.metadata().map_err(|error| {
|
||||
SchedulerError::RecipeLoadError(format!("Cannot inspect recipe file: {error}"))
|
||||
})?;
|
||||
if !metadata.is_file() {
|
||||
return Err(SchedulerError::RecipeLoadError(
|
||||
"Recipe path must reference a regular file".to_string(),
|
||||
));
|
||||
}
|
||||
if metadata.len() > MAX_SCHEDULE_RECIPE_BYTES {
|
||||
return Err(SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file exceeds the {MAX_SCHEDULE_RECIPE_BYTES} byte limit"
|
||||
)));
|
||||
}
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
source
|
||||
.take(MAX_SCHEDULE_RECIPE_BYTES + 1)
|
||||
.read_to_end(&mut bytes)
|
||||
.map_err(|error| {
|
||||
SchedulerError::RecipeLoadError(format!("Cannot read recipe file: {error}"))
|
||||
})?;
|
||||
if bytes.len() as u64 > MAX_SCHEDULE_RECIPE_BYTES {
|
||||
return Err(SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file exceeds the {MAX_SCHEDULE_RECIPE_BYTES} byte limit"
|
||||
)));
|
||||
}
|
||||
|
||||
write_schedule_recipe_bytes(destination, &bytes)
|
||||
}
|
||||
|
||||
fn write_schedule_recipe_bytes(destination: &Path, bytes: &[u8]) -> Result<(), SchedulerError> {
|
||||
if bytes.len() as u64 > MAX_SCHEDULE_RECIPE_BYTES {
|
||||
return Err(SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file exceeds the {MAX_SCHEDULE_RECIPE_BYTES} byte limit"
|
||||
)));
|
||||
}
|
||||
|
||||
let result = (|| {
|
||||
let mut options = OpenOptions::new();
|
||||
options.write(true).create(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
options.mode(0o600);
|
||||
}
|
||||
|
||||
let mut file = options.open(destination)?;
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
file.set_permissions(fs::Permissions::from_mode(0o600))?;
|
||||
}
|
||||
file.set_len(0)?;
|
||||
file.write_all(bytes)
|
||||
})();
|
||||
if let Err(error) = result {
|
||||
let _ = fs::remove_file(destination);
|
||||
return Err(SchedulerError::StorageError(error));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_default_scheduler_storage_path() -> Result<PathBuf, io::Error> {
|
||||
let data_dir = Paths::data_dir();
|
||||
fs::create_dir_all(&data_dir)?;
|
||||
@@ -299,6 +409,25 @@ impl Scheduler {
|
||||
&self,
|
||||
original_job_spec: ScheduledJob,
|
||||
make_copy: bool,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.add_scheduled_job_inner(original_job_spec, make_copy, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
original_job_spec: ScheduledJob,
|
||||
validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.add_scheduled_job_inner(original_job_spec, true, Some(validated_recipe))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_inner(
|
||||
&self,
|
||||
original_job_spec: ScheduledJob,
|
||||
make_copy: bool,
|
||||
validated_recipe: Option<ValidatedScheduleRecipe>,
|
||||
) -> Result<(), SchedulerError> {
|
||||
{
|
||||
let jobs_guard = self.jobs.lock().await;
|
||||
@@ -309,19 +438,25 @@ impl Scheduler {
|
||||
|
||||
let mut stored_job = original_job_spec;
|
||||
if make_copy {
|
||||
let original_recipe_path =
|
||||
Path::new(&stored_job.source).canonicalize().map_err(|e| {
|
||||
SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file not found: {}: {}",
|
||||
stored_job.source, e
|
||||
))
|
||||
})?;
|
||||
if !original_recipe_path.is_file() {
|
||||
return Err(SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file not found: {}",
|
||||
stored_job.source
|
||||
)));
|
||||
}
|
||||
let (original_recipe_path, validated_recipe) =
|
||||
if let Some(validated_recipe) = validated_recipe {
|
||||
(validated_recipe.source, Some(validated_recipe.bytes))
|
||||
} else {
|
||||
let original_recipe_path =
|
||||
Path::new(&stored_job.source).canonicalize().map_err(|e| {
|
||||
SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file not found: {}: {}",
|
||||
stored_job.source, e
|
||||
))
|
||||
})?;
|
||||
if !original_recipe_path.is_file() {
|
||||
return Err(SchedulerError::RecipeLoadError(format!(
|
||||
"Recipe file not found: {}",
|
||||
stored_job.source
|
||||
)));
|
||||
}
|
||||
(original_recipe_path, None)
|
||||
};
|
||||
|
||||
let scheduled_recipes_dir = get_default_scheduled_recipes_dir()?;
|
||||
let original_extension = original_recipe_path
|
||||
@@ -332,7 +467,11 @@ impl Scheduler {
|
||||
let destination_filename = format!("{}.{}", stored_job.id, original_extension);
|
||||
let destination_recipe_path = scheduled_recipes_dir.join(destination_filename);
|
||||
|
||||
fs::copy(&original_recipe_path, &destination_recipe_path)?;
|
||||
if let Some(recipe) = validated_recipe.as_deref() {
|
||||
write_schedule_recipe_bytes(&destination_recipe_path, recipe)?;
|
||||
} else {
|
||||
copy_bounded_schedule_recipe(&original_recipe_path, &destination_recipe_path)?;
|
||||
}
|
||||
stored_job.recipe_base_dir = original_recipe_path
|
||||
.parent()
|
||||
.map(|p| p.to_string_lossy().into_owned());
|
||||
@@ -1085,6 +1224,15 @@ impl SchedulerTrait for Scheduler {
|
||||
self.add_scheduled_job(job, make_copy).await
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.add_scheduled_job_with_recipe(job, validated_recipe)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
recipe_path: PathBuf,
|
||||
@@ -1157,6 +1305,118 @@ mod tests {
|
||||
recipe_path
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bounded_recipe_copy_rejects_source_that_grew_after_validation() {
|
||||
let temp_dir = tempdir().unwrap();
|
||||
let source = temp_dir.path().join("source.yaml");
|
||||
let destination = temp_dir.path().join("destination.yaml");
|
||||
fs::write(
|
||||
&source,
|
||||
"title: Valid\ndescription: Initially valid\nprompt: Run safely\n",
|
||||
)
|
||||
.unwrap();
|
||||
let validated = fs::read_to_string(&source).unwrap();
|
||||
serde_yaml::from_str::<Recipe>(&validated).unwrap();
|
||||
File::options()
|
||||
.write(true)
|
||||
.open(&source)
|
||||
.unwrap()
|
||||
.set_len(MAX_SCHEDULE_RECIPE_BYTES + 1)
|
||||
.unwrap();
|
||||
|
||||
let error = copy_bounded_schedule_recipe(&source, &destination).unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("exceeds the 1048576 byte limit"));
|
||||
assert!(!destination.exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validated_recipe_bytes_and_base_are_persisted_after_source_replacement() {
|
||||
let temp_dir = tempdir().unwrap();
|
||||
let _guard =
|
||||
env_lock::lock_env([("GOOSE_PATH_ROOT", Some(temp_dir.path().to_str().unwrap()))]);
|
||||
let trusted_dir = temp_dir.path().join("trusted");
|
||||
let replacement_dir = temp_dir.path().join("replacement");
|
||||
fs::create_dir_all(&trusted_dir).unwrap();
|
||||
fs::create_dir_all(&replacement_dir).unwrap();
|
||||
let trusted_source = trusted_dir.join("source.yaml");
|
||||
let replacement_source = replacement_dir.join("source.yaml");
|
||||
let validated =
|
||||
b"title: Validated\ndescription: Original recipe\nprompt: Run safely\n".to_vec();
|
||||
let replacement =
|
||||
b"title: Replacement\ndescription: Swapped recipe\nprompt: Run something else\n";
|
||||
fs::write(&trusted_source, &validated).unwrap();
|
||||
serde_yaml::from_slice::<Recipe>(&validated).unwrap();
|
||||
fs::write(&replacement_source, replacement).unwrap();
|
||||
|
||||
let storage_path = temp_dir.path().join("schedule.json");
|
||||
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
|
||||
let scheduler = Scheduler::new(storage_path, session_manager).await.unwrap();
|
||||
let job = ScheduledJob {
|
||||
id: "validated_recipe_copy".to_string(),
|
||||
source: replacement_source.to_string_lossy().into_owned(),
|
||||
cron: "0 0 0 1 1 *".to_string(),
|
||||
last_run: None,
|
||||
currently_running: false,
|
||||
paused: false,
|
||||
current_session_id: None,
|
||||
process_start_time: None,
|
||||
parameters: vec![],
|
||||
recipe_base_dir: None,
|
||||
};
|
||||
|
||||
scheduler
|
||||
.add_scheduled_job_with_recipe(
|
||||
job,
|
||||
ValidatedScheduleRecipe::new(validated.clone(), trusted_source.clone()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let jobs = scheduler.list_scheduled_jobs().await;
|
||||
let stored = jobs
|
||||
.iter()
|
||||
.find(|job| job.id == "validated_recipe_copy")
|
||||
.unwrap();
|
||||
assert_eq!(fs::read(&stored.source).unwrap(), validated);
|
||||
assert_eq!(stored.recipe_base_dir.as_deref(), trusted_dir.to_str());
|
||||
assert_ne!(
|
||||
stored.recipe_base_dir.as_deref(),
|
||||
replacement_source.parent().and_then(Path::to_str)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validated_recipe_copy_rejects_oversized_bytes_without_destination() {
|
||||
let temp_dir = tempdir().unwrap();
|
||||
let destination = temp_dir.path().join("destination.yaml");
|
||||
let oversized = vec![0; (MAX_SCHEDULE_RECIPE_BYTES + 1) as usize];
|
||||
|
||||
let error = write_schedule_recipe_bytes(&destination, &oversized).unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("exceeds the 1048576 byte limit"));
|
||||
assert!(!destination.exists());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn validated_recipe_copy_makes_existing_destination_owner_private() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
let temp_dir = tempdir().unwrap();
|
||||
let destination = temp_dir.path().join("destination.yaml");
|
||||
fs::write(&destination, b"old contents").unwrap();
|
||||
fs::set_permissions(&destination, fs::Permissions::from_mode(0o644)).unwrap();
|
||||
|
||||
write_schedule_recipe_bytes(&destination, b"private recipe").unwrap();
|
||||
|
||||
assert_eq!(fs::read(&destination).unwrap(), b"private recipe");
|
||||
assert_eq!(
|
||||
fs::metadata(&destination).unwrap().permissions().mode() & 0o777,
|
||||
0o600
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_runs_on_schedule() {
|
||||
let _guard = env_lock::lock_env([
|
||||
|
||||
@@ -2,7 +2,7 @@ use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::scheduler::{ScheduledJob, SchedulerError};
|
||||
use crate::scheduler::{ScheduledJob, SchedulerError, ValidatedScheduleRecipe};
|
||||
use crate::session::Session;
|
||||
|
||||
#[async_trait]
|
||||
@@ -12,6 +12,11 @@ pub trait SchedulerTrait: Send + Sync {
|
||||
job: ScheduledJob,
|
||||
copy_recipe: bool,
|
||||
) -> Result<(), SchedulerError>;
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError>;
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
recipe_path: PathBuf,
|
||||
|
||||
@@ -20,7 +20,7 @@ use goose::config::{GooseMode, PermissionManager};
|
||||
use goose::providers::api_client::{ApiClient, AuthMethod as ApiAuthMethod};
|
||||
use goose::providers::base::Provider;
|
||||
use goose::providers::openai::OpenAiProvider;
|
||||
use goose::scheduler::{ScheduledJob, SchedulerError};
|
||||
use goose::scheduler::{ScheduledJob, SchedulerError, ValidatedScheduleRecipe};
|
||||
use goose::scheduler_trait::SchedulerTrait;
|
||||
use goose::session::Session as GooseSession;
|
||||
use goose::session_context::SESSION_ID_HEADER;
|
||||
@@ -78,6 +78,14 @@ impl SchedulerTrait for FixtureScheduler {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
_validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.add_scheduled_job(job, false).await
|
||||
}
|
||||
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
recipe_path: PathBuf,
|
||||
|
||||
@@ -18,7 +18,7 @@ mod tests {
|
||||
use goose::agents::AgentConfig;
|
||||
use goose::config::permission::PermissionManager;
|
||||
use goose::config::GooseMode;
|
||||
use goose::scheduler::{ScheduledJob, SchedulerError};
|
||||
use goose::scheduler::{ScheduledJob, SchedulerError, ValidatedScheduleRecipe};
|
||||
use goose::scheduler_trait::SchedulerTrait;
|
||||
use goose::session::{Session, SessionManager};
|
||||
use std::path::PathBuf;
|
||||
@@ -49,6 +49,14 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
_job: ScheduledJob,
|
||||
_validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
_recipe_path: PathBuf,
|
||||
@@ -129,6 +137,16 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
_validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
let mut jobs = self.jobs.lock().await;
|
||||
jobs.push(job);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
_recipe_path: PathBuf,
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use goose::agents::{Agent, AgentConfig, GoosePlatform};
|
||||
use goose::config::permission::PermissionManager;
|
||||
use goose::config::GooseMode;
|
||||
use goose::scheduler::{ScheduledJob, SchedulerError, ValidatedScheduleRecipe};
|
||||
use goose::scheduler_trait::SchedulerTrait;
|
||||
use goose::session::{Session, SessionManager};
|
||||
use tempfile::TempDir;
|
||||
|
||||
struct MockScheduler {
|
||||
jobs: tokio::sync::Mutex<Vec<ScheduledJob>>,
|
||||
validated_recipes: tokio::sync::Mutex<Vec<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl MockScheduler {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
jobs: tokio::sync::Mutex::new(Vec::new()),
|
||||
validated_recipes: tokio::sync::Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SchedulerTrait for MockScheduler {
|
||||
async fn add_scheduled_job(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
_copy: bool,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.jobs.lock().await.push(job);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn add_scheduled_job_with_recipe(
|
||||
&self,
|
||||
job: ScheduledJob,
|
||||
validated_recipe: ValidatedScheduleRecipe,
|
||||
) -> Result<(), SchedulerError> {
|
||||
self.jobs.lock().await.push(job);
|
||||
self.validated_recipes
|
||||
.lock()
|
||||
.await
|
||||
.push(validated_recipe.bytes().to_vec());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn schedule_recipe(
|
||||
&self,
|
||||
_recipe_path: PathBuf,
|
||||
_cron_schedule: Option<String>,
|
||||
) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_scheduled_jobs(&self) -> Vec<ScheduledJob> {
|
||||
self.jobs.lock().await.clone()
|
||||
}
|
||||
|
||||
async fn remove_scheduled_job(&self, _id: &str, _remove: bool) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn pause_schedule(&self, _id: &str) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn unpause_schedule(&self, _id: &str) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_now(&self, _id: &str) -> Result<String, SchedulerError> {
|
||||
Ok("test-session".to_string())
|
||||
}
|
||||
|
||||
async fn sessions(
|
||||
&self,
|
||||
_sched_id: &str,
|
||||
_limit: usize,
|
||||
) -> Result<Vec<(String, Session)>, SchedulerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn update_schedule(
|
||||
&self,
|
||||
_sched_id: &str,
|
||||
_new_cron: String,
|
||||
) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn kill_running_job(&self, _sched_id: &str) -> Result<(), SchedulerError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_running_job_info(
|
||||
&self,
|
||||
_sched_id: &str,
|
||||
) -> Result<Option<(String, DateTime<Utc>)>, SchedulerError> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_with_scheduler(temp_dir: &TempDir, scheduler: Arc<MockScheduler>) -> Agent {
|
||||
let data_dir = temp_dir.path().join("data");
|
||||
let session_manager = Arc::new(SessionManager::new(data_dir.clone()));
|
||||
let permission_manager = Arc::new(PermissionManager::new(data_dir));
|
||||
let config = AgentConfig::new(
|
||||
session_manager,
|
||||
permission_manager,
|
||||
Some(scheduler),
|
||||
GooseMode::Auto,
|
||||
false,
|
||||
GoosePlatform::GooseCli,
|
||||
);
|
||||
Agent::with_config(config)
|
||||
}
|
||||
|
||||
async fn create_schedule(agent: &Agent, recipe_path: &Path) -> Result<(), String> {
|
||||
agent
|
||||
.handle_schedule_management(
|
||||
serde_json::json!({
|
||||
"action": "create",
|
||||
"recipe_path": recipe_path,
|
||||
"cron_expression": "0 * * * *"
|
||||
}),
|
||||
"test-request".to_string(),
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(|error| error.message.to_string())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn parse_errors_do_not_reflect_recipe_contents() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
let cases = [
|
||||
("invalid.yaml", "yaml-secret-242", "Invalid YAML recipe"),
|
||||
("invalid.json", "\"json-secret-242\"", "Invalid JSON recipe"),
|
||||
];
|
||||
|
||||
for (name, secret, expected) in cases {
|
||||
let path = temp_dir.path().join(name);
|
||||
std::fs::write(&path, secret).unwrap();
|
||||
let message = create_schedule(&agent, &path).await.unwrap_err();
|
||||
assert_eq!(message, expected);
|
||||
assert!(!message.contains(secret));
|
||||
}
|
||||
|
||||
assert!(scheduler.jobs.lock().await.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_non_regular_recipe_path() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
|
||||
let message = create_schedule(&agent, temp_dir.path()).await.unwrap_err();
|
||||
|
||||
assert_eq!(message, "Recipe path must reference a regular file");
|
||||
assert!(scheduler.jobs.lock().await.is_empty());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn rejects_fifo_without_blocking() {
|
||||
use std::ffi::CString;
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
use std::time::Duration;
|
||||
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
let path = temp_dir.path().join("recipe.yaml");
|
||||
let fifo_path = CString::new(path.as_os_str().as_bytes()).unwrap();
|
||||
// SAFETY: fifo_path is a valid, NUL-terminated path and mode contains only permission bits.
|
||||
assert_eq!(unsafe { libc::mkfifo(fifo_path.as_ptr(), 0o600) }, 0);
|
||||
|
||||
let (finished_tx, finished_rx) = std::sync::mpsc::channel();
|
||||
let watchdog_path = path.clone();
|
||||
let watchdog = std::thread::spawn(move || {
|
||||
let timed_out = finished_rx.recv_timeout(Duration::from_secs(2)).is_err();
|
||||
if timed_out {
|
||||
let _ = std::fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.open(watchdog_path);
|
||||
}
|
||||
timed_out
|
||||
});
|
||||
|
||||
let message = create_schedule(&agent, &path).await.unwrap_err();
|
||||
let _ = finished_tx.send(());
|
||||
|
||||
assert!(!watchdog.join().unwrap(), "FIFO validation blocked on open");
|
||||
assert_eq!(message, "Recipe path must reference a regular file");
|
||||
assert!(scheduler.jobs.lock().await.is_empty());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn accepts_symlink_to_regular_recipe_with_canonical_provenance() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
let target = temp_dir.path().join("target.yaml");
|
||||
let link = temp_dir.path().join("recipe-link.yaml");
|
||||
std::fs::write(
|
||||
&target,
|
||||
b"title: Valid recipe\ndescription: A small recipe\nprompt: Run safely\n",
|
||||
)
|
||||
.unwrap();
|
||||
std::os::unix::fs::symlink(&target, &link).unwrap();
|
||||
|
||||
create_schedule(&agent, &link).await.unwrap();
|
||||
|
||||
let canonical_target = target.canonicalize().unwrap();
|
||||
let jobs = scheduler.jobs.lock().await;
|
||||
assert_eq!(jobs[0].source, canonical_target.to_string_lossy());
|
||||
assert_eq!(
|
||||
jobs[0].recipe_base_dir.as_deref(),
|
||||
canonical_target.parent().and_then(Path::to_str)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_oversized_recipe() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
let path = temp_dir.path().join("oversized.yaml");
|
||||
std::fs::File::create(&path)
|
||||
.unwrap()
|
||||
.set_len(1_048_577)
|
||||
.unwrap();
|
||||
|
||||
let message = create_schedule(&agent, &path).await.unwrap_err();
|
||||
|
||||
assert_eq!(message, "Recipe file exceeds the 1048576 byte limit");
|
||||
assert!(scheduler.jobs.lock().await.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accepts_valid_regular_recipe() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let scheduler = Arc::new(MockScheduler::new());
|
||||
let agent = agent_with_scheduler(&temp_dir, scheduler.clone());
|
||||
let path = temp_dir.path().join("valid.yaml");
|
||||
let recipe = b"title: Valid recipe\ndescription: A small recipe\nprompt: Run safely\n";
|
||||
std::fs::write(&path, recipe).unwrap();
|
||||
|
||||
create_schedule(&agent, &path).await.unwrap();
|
||||
|
||||
assert_eq!(scheduler.jobs.lock().await.len(), 1);
|
||||
assert_eq!(
|
||||
scheduler.validated_recipes.lock().await.as_slice(),
|
||||
&[recipe.to_vec()]
|
||||
);
|
||||
let canonical_path = path.canonicalize().unwrap();
|
||||
let jobs = scheduler.jobs.lock().await;
|
||||
assert_eq!(jobs[0].source, canonical_path.to_string_lossy());
|
||||
assert_eq!(
|
||||
jobs[0].recipe_base_dir.as_deref(),
|
||||
canonical_path.parent().and_then(Path::to_str)
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user