tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Platform Tool for Scheduler: Allow Goose to Manage Its Own Schedule (#2944)

Browse Source
This commit is contained in:
tlongwell-block
2025-06-19 10:04:31 -04:00
committed by GitHub
parent 46a8b9218a
commit f0b627c96a
10 changed files with 2023 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+21 -3
View File
@@ -17,6 +17,7 @@ use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::recipe::{Author, Recipe, Settings};
use crate::scheduler_trait::SchedulerTrait;
use crate::tool_monitor::{ToolCall, ToolMonitor};
use regex::Regex;
use serde_json::Value;
@@ -27,7 +28,8 @@ use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult,
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::platform_tools::{
PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME,
PLATFORM_READ_RESOURCE_TOOL_NAME, PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
PLATFORM_MANAGE_SCHEDULE_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME,
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
};
use crate::agents::prompt_manager::PromptManager;
use crate::agents::router_tool_selector::{
@@ -59,6 +61,7 @@ pub struct Agent {
pub(super) tool_result_rx: ToolResultReceiver,
pub(super) tool_monitor: Mutex<Option<ToolMonitor>>,
pub(super) router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
}
#[derive(Clone, Debug)]
@@ -86,6 +89,7 @@ impl Agent {
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
tool_monitor: Mutex::new(None),
router_tool_selector: Mutex::new(None),
scheduler_service: Mutex::new(None),
}
}
@@ -104,6 +108,12 @@ impl Agent {
monitor.reset();
}
}
/// Set the scheduler service for this agent
pub async fn set_scheduler(&self, scheduler: Arc<dyn SchedulerTrait>) {
let mut scheduler_service = self.scheduler_service.lock().await;
*scheduler_service = Some(scheduler);
}
}
impl Default for Agent {
@@ -185,7 +195,7 @@ impl Agent {
/// Dispatch a single tool call to the appropriate client
#[instrument(skip(self, tool_call, request_id), fields(input, output))]
pub(super) async fn dispatch_tool_call(
pub async fn dispatch_tool_call(
&self,
tool_call: mcp_core::tool::ToolCall,
request_id: String,
@@ -204,6 +214,13 @@ impl Agent {
}
}
if tool_call.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME {
let result = self
.handle_schedule_management(tool_call.arguments, request_id.clone())
.await;
return (request_id, Ok(ToolCallResult::from(result)));
}
if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME {
let extension_name = tool_call
.arguments
@@ -414,7 +431,7 @@ impl Agent {
let mut extension_manager = self.extension_manager.lock().await;
extension_manager.add_extension(extension.clone()).await?;
}
};
}
// If vector tool selection is enabled, index the tools
let selector = self.router_tool_selector.lock().await.clone();
@@ -453,6 +470,7 @@ impl Agent {
// Add platform tools
prefixed_tools.push(platform_tools::search_available_extensions_tool());
prefixed_tools.push(platform_tools::manage_extensions_tool());
prefixed_tools.push(platform_tools::manage_schedule_tool());
// Add resource tools if supported
if extension_manager.supports_resources() {
crates/goose/src/agents/mod.rs
+2
View File
@@ -8,6 +8,8 @@ pub mod prompt_manager;
mod reply_parts;
mod router_tool_selector;
mod router_tools;
mod schedule_tool;
mod tool_execution;
mod tool_router_index_manager;
pub(crate) mod tool_vectordb;
crates/goose/src/agents/platform_tools.rs
+45
View File
@@ -7,6 +7,7 @@ pub const PLATFORM_LIST_RESOURCES_TOOL_NAME: &str = "platform__list_resources";
pub const PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME: &str =
"platform__search_available_extensions";
pub const PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME: &str = "platform__manage_extensions";
pub const PLATFORM_MANAGE_SCHEDULE_TOOL_NAME: &str = "platform__manage_schedule";
pub fn read_resource_tool() -> Tool {
Tool::new(
@@ -112,3 +113,47 @@ pub fn manage_extensions_tool() -> Tool {
}),
)
}
pub fn manage_schedule_tool() -> Tool {
Tool::new(
PLATFORM_MANAGE_SCHEDULE_TOOL_NAME.to_string(),
indoc! {r#"
Manage scheduled recipe execution for this Goose instance.
Actions:
- "list": List all scheduled jobs
- "create": Create a new scheduled job from a recipe file
- "run_now": Execute a scheduled job immediately
- "pause": Pause a scheduled job
- "unpause": Resume a paused job
- "delete": Remove a scheduled job
- "kill": Terminate a currently running job
- "inspect": Get details about a running job
- "sessions": List execution history for a job
- "session_content": Get the full content (messages) of a specific session
"#}
.to_string(),
json!({
"type": "object",
"required": ["action"],
"properties": {
"action": {
"type": "string",
"enum": ["list", "create", "run_now", "pause", "unpause", "delete", "kill", "inspect", "sessions", "session_content"]
},
"job_id": {"type": "string", "description": "Job identifier for operations on existing jobs"},
"recipe_path": {"type": "string", "description": "Path to recipe file for create action"},
"cron_expression": {"type": "string", "description": "A six field cron expression for create action"},
"limit": {"type": "integer", "description": "Limit for sessions list", "default": 50},
"session_id": {"type": "string", "description": "Session identifier for session_content action"}
}
}),
Some(ToolAnnotations {
title: Some("Manage scheduled recipes".to_string()),
read_only_hint: false,
destructive_hint: true, // Can kill jobs
idempotent_hint: false,
open_world_hint: false,
}),
)
}
crates/goose/src/agents/schedule_tool.rs
+420
View File
@@ -0,0 +1,420 @@
//! Schedule tool handlers for the Goose agent
//!
//! This module contains all the handlers for the schedule management platform tool,
//! including job creation, execution, monitoring, and session management.
use std::sync::Arc;
use chrono::Utc;
use mcp_core::{Content, ToolError, ToolResult};
use crate::recipe::Recipe;
use crate::scheduler_trait::SchedulerTrait;
use super::Agent;
impl Agent {
/// Handle schedule management tool calls
pub async fn handle_schedule_management(
&self,
arguments: serde_json::Value,
_request_id: String,
) -> ToolResult<Vec<Content>> {
let scheduler = match self.scheduler_service.lock().await.as_ref() {
Some(s) => s.clone(),
None => {
return Err(ToolError::ExecutionError(
"Scheduler not available. This tool only works in server mode.".to_string(),
))
}
};
let action = arguments
.get("action")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'action' parameter".to_string()))?;
match action {
"list" => self.handle_list_jobs(scheduler).await,
"create" => self.handle_create_job(scheduler, arguments).await,
"run_now" => self.handle_run_now(scheduler, arguments).await,
"pause" => self.handle_pause_job(scheduler, arguments).await,
"unpause" => self.handle_unpause_job(scheduler, arguments).await,
"delete" => self.handle_delete_job(scheduler, arguments).await,
"kill" => self.handle_kill_job(scheduler, arguments).await,
"inspect" => self.handle_inspect_job(scheduler, arguments).await,
"sessions" => self.handle_list_sessions(scheduler, arguments).await,
"session_content" => self.handle_session_content(arguments).await,
_ => Err(ToolError::ExecutionError(format!(
"Unknown action: {}",
action
))),
}
}
/// List all scheduled jobs
async fn handle_list_jobs(
&self,
scheduler: Arc<dyn SchedulerTrait>,
) -> ToolResult<Vec<Content>> {
match scheduler.list_scheduled_jobs().await {
Ok(jobs) => {
let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| {
ToolError::ExecutionError(format!("Failed to serialize jobs: {}", e))
})?;
Ok(vec![Content::text(format!(
"Scheduled Jobs:\n{}",
jobs_json
))])
}
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to list jobs: {}",
e
))),
}
}
/// Create a new scheduled job from a recipe file
async fn handle_create_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let recipe_path = arguments
.get("recipe_path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::ExecutionError("Missing 'recipe_path' parameter".to_string())
})?;
let cron_expression = arguments
.get("cron_expression")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::ExecutionError("Missing 'cron_expression' parameter".to_string())
})?;
// Validate recipe file exists and is readable
if !std::path::Path::new(recipe_path).exists() {
return Err(ToolError::ExecutionError(format!(
"Recipe file not found: {}",
recipe_path
)));
}
// 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| {
ToolError::ExecutionError(format!("Invalid JSON recipe: {}", e))
})?;
} else {
serde_yaml::from_str::<Recipe>(&content).map_err(|e| {
ToolError::ExecutionError(format!("Invalid YAML recipe: {}", e))
})?;
}
}
Err(e) => {
return Err(ToolError::ExecutionError(format!(
"Cannot read recipe file: {}",
e
)))
}
}
// Generate unique job ID
let job_id = format!("agent_created_{}", Utc::now().timestamp());
let job = crate::scheduler::ScheduledJob {
id: job_id.clone(),
source: recipe_path.to_string(),
cron: cron_expression.to_string(),
last_run: None,
currently_running: false,
paused: false,
current_session_id: None,
process_start_time: None,
};
match scheduler.add_scheduled_job(job).await {
Ok(()) => Ok(vec![Content::text(format!(
"Successfully created scheduled job '{}' for recipe '{}' with cron expression '{}'",
job_id, recipe_path, cron_expression
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to create job: {}",
e
))),
}
}
/// Run a scheduled job immediately
async fn handle_run_now(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.run_now(job_id).await {
Ok(session_id) => Ok(vec![Content::text(format!(
"Successfully started job '{}'. Session ID: {}",
job_id, session_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to run job: {}",
e
))),
}
}
/// Pause a scheduled job
async fn handle_pause_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.pause_schedule(job_id).await {
Ok(()) => Ok(vec![Content::text(format!(
"Successfully paused job '{}'",
job_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to pause job: {}",
e
))),
}
}
/// Resume a paused scheduled job
async fn handle_unpause_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.unpause_schedule(job_id).await {
Ok(()) => Ok(vec![Content::text(format!(
"Successfully unpaused job '{}'",
job_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to unpause job: {}",
e
))),
}
}
/// Delete a scheduled job
async fn handle_delete_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.remove_scheduled_job(job_id).await {
Ok(()) => Ok(vec![Content::text(format!(
"Successfully deleted job '{}'",
job_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to delete job: {}",
e
))),
}
}
/// Terminate a currently running job
async fn handle_kill_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.kill_running_job(job_id).await {
Ok(()) => Ok(vec![Content::text(format!(
"Successfully killed running job '{}'",
job_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to kill job: {}",
e
))),
}
}
/// Get information about a running job
async fn handle_inspect_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
match scheduler.get_running_job_info(job_id).await {
Ok(Some((session_id, start_time))) => {
let duration = Utc::now().signed_duration_since(start_time);
Ok(vec![Content::text(format!(
"Job '{}' is currently running:\n- Session ID: {}\n- Started: {}\n- Duration: {} seconds",
job_id, session_id, start_time.to_rfc3339(), duration.num_seconds()
))])
}
Ok(None) => Ok(vec![Content::text(format!(
"Job '{}' is not currently running",
job_id
))]),
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to inspect job: {}",
e
))),
}
}
/// List execution sessions for a job
async fn handle_list_sessions(
&self,
scheduler: Arc<dyn SchedulerTrait>,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let job_id = arguments
.get("job_id")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?;
let limit = arguments
.get("limit")
.and_then(|v| v.as_u64())
.unwrap_or(50) as usize;
match scheduler.sessions(job_id, limit).await {
Ok(sessions) => {
if sessions.is_empty() {
Ok(vec![Content::text(format!(
"No sessions found for job '{}'",
job_id
))])
} else {
let sessions_info: Vec<String> = sessions
.into_iter()
.map(|(session_name, metadata)| {
format!(
"- Session: {} (Messages: {}, Working Dir: {})",
session_name,
metadata.message_count,
metadata.working_dir.display()
)
})
.collect();
Ok(vec![Content::text(format!(
"Sessions for job '{}':\n{}",
job_id,
sessions_info.join("\n")
))])
}
}
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to list sessions: {}",
e
))),
}
}
/// Get the full content (metadata and messages) of a specific session
async fn handle_session_content(
&self,
arguments: serde_json::Value,
) -> ToolResult<Vec<Content>> {
let session_id = arguments
.get("session_id")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::ExecutionError("Missing 'session_id' parameter".to_string())
})?;
// Get the session file path
let session_path = crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_id.to_string()),
);
// Check if session file exists
if !session_path.exists() {
return Err(ToolError::ExecutionError(format!(
"Session '{}' not found",
session_id
)));
}
// Read session metadata
let metadata = match crate::session::storage::read_metadata(&session_path) {
Ok(metadata) => metadata,
Err(e) => {
return Err(ToolError::ExecutionError(format!(
"Failed to read session metadata: {}",
e
)));
}
};
// Read session messages
let messages = match crate::session::storage::read_messages(&session_path) {
Ok(messages) => messages,
Err(e) => {
return Err(ToolError::ExecutionError(format!(
"Failed to read session messages: {}",
e
)));
}
};
// Format the response with metadata and messages
let metadata_json = match serde_json::to_string_pretty(&metadata) {
Ok(json) => json,
Err(e) => {
return Err(ToolError::ExecutionError(format!(
"Failed to serialize metadata: {}",
e
)));
}
};
let messages_json = match serde_json::to_string_pretty(&messages) {
Ok(json) => json,
Err(e) => {
return Err(ToolError::ExecutionError(format!(
"Failed to serialize messages: {}",
e
)));
}
};
Ok(vec![Content::text(format!(
"Session '{}' Content:\n\nMetadata:\n{}\n\nMessages:\n{}",
session_id, metadata_json, messages_json
))])
}
}