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

Add basic cron scheduler to goose-server (#2621)

Browse Source
This commit is contained in:
Max Novich
2025-05-27 10:36:27 -07:00
committed by GitHub
parent c8e3f6ac69
commit c272b5df95
39 changed files with 3554 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/commands/agent.rs
+13 -7
View File
@@ -3,7 +3,10 @@ use std::sync::Arc;
use crate::configuration;
use crate::state;
use anyhow::Result;
use etcetera::{choose_app_strategy, AppStrategy};
use goose::agents::Agent;
use goose::config::APP_STRATEGY;
use goose::scheduler::Scheduler as GooseScheduler;
use tower_http::cors::{Any, CorsLayer};
use tracing::info;
@@ -11,27 +14,30 @@ pub async fn run() -> Result<()> {
// Initialize logging
crate::logging::setup_logging(Some("goosed"))?;
// Load configuration
let settings = configuration::Settings::new()?;
// load secret key from GOOSE_SERVER__SECRET_KEY environment variable
let secret_key =
std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
let new_agent = Agent::new();
let agent_ref = Arc::new(new_agent);
// Create app state with agent
let state = state::AppState::new(Arc::new(new_agent), secret_key.clone()).await;
let app_state = state::AppState::new(agent_ref.clone(), secret_key.clone()).await;
let schedule_file_path = choose_app_strategy(APP_STRATEGY.clone())?
.data_dir()
.join("schedules.json");
let scheduler_instance = GooseScheduler::new(schedule_file_path).await?;
app_state.set_scheduler(scheduler_instance).await;
// Create router with CORS support
let cors = CorsLayer::new()
.allow_origin(Any)
.allow_methods(Any)
.allow_headers(Any);
let app = crate::routes::configure(state).layer(cors);
let app = crate::routes::configure(app_state).layer(cors);
// Run server
let listener = tokio::net::TcpListener::bind(settings.socket_addr()).await?;
info!("listening on {}", listener.local_addr()?);
axum::serve(listener, app).await?;
crates/goose-server/src/logging.rs
+3 -2
View File
@@ -8,6 +8,7 @@ use tracing_subscriber::{
Registry,
};
use goose::config::APP_STRATEGY;
use goose::tracing::langfuse_layer;
/// Returns the directory where log files should be stored.
@@ -17,8 +18,8 @@ fn get_log_directory() -> Result<PathBuf> {
// - macOS/Linux: ~/.local/state/goose/logs/server
// - Windows: ~\AppData\Roaming\Block\goose\data\logs\server
// - Windows has no convention for state_dir, use data_dir instead
let home_dir = choose_app_strategy(crate::APP_STRATEGY.clone())
.context("HOME environment variable not set")?;
let home_dir =
choose_app_strategy(APP_STRATEGY.clone()).context("HOME environment variable not set")?;
let base_log_dir = home_dir
.in_state_dir("logs/server")
crates/goose-server/src/main.rs
-9
View File
@@ -1,12 +1,3 @@
use etcetera::AppStrategyArgs;
use once_cell::sync::Lazy;
pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
top_level_domain: "Block".to_string(),
author: "Block".to_string(),
app_name: "goose".to_string(),
});
mod commands;
mod configuration;
mod error;
crates/goose-server/src/openapi.rs
+12 -1
View File
@@ -37,7 +37,12 @@ use utoipa::OpenApi;
super::routes::reply::confirm_permission,
super::routes::context::manage_context,
super::routes::session::list_sessions,
super::routes::session::get_session_history
super::routes::session::get_session_history,
super::routes::schedule::create_schedule,
super::routes::schedule::list_schedules,
super::routes::schedule::delete_schedule,
super::routes::schedule::run_now_handler,
super::routes::schedule::sessions_handler
),
components(schemas(
super::routes::config_management::UpsertConfigQuery,
@@ -85,6 +90,12 @@ use utoipa::OpenApi;
ModelInfo,
SessionInfo,
SessionMetadata,
super::routes::schedule::CreateScheduleRequest,
goose::scheduler::ScheduledJob,
super::routes::schedule::RunNowResponse,
super::routes::schedule::ListSchedulesResponse,
super::routes::schedule::SessionsQuery,
super::routes::schedule::SessionDisplayInfo,
))
)]
pub struct ApiDoc;
crates/goose-server/src/routes/config_management.rs
+15 -46
View File
@@ -6,8 +6,9 @@ use axum::{
routing::{delete, get, post},
Json, Router,
};
use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs};
use etcetera::{choose_app_strategy, AppStrategy};
use goose::config::Config;
use goose::config::APP_STRATEGY;
use goose::config::{extensions::name_to_key, PermissionManager};
use goose::config::{ExtensionConfigManager, ExtensionEntry};
use goose::model::ModelConfig;
@@ -15,7 +16,6 @@ use goose::providers::base::ProviderMetadata;
use goose::providers::providers as get_providers;
use goose::{agents::ExtensionConfig, config::permission::PermissionLevel};
use http::{HeaderMap, StatusCode};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_yaml;
@@ -52,14 +52,12 @@ pub struct ConfigResponse {
pub config: HashMap<String, Value>,
}
// Define a new structure to encapsulate the provider details along with configuration status
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ProviderDetails {
/// Unique identifier and name of the provider
pub name: String,
/// Metadata about the provider
pub metadata: ProviderMetadata,
/// Indicates whether the provider is fully configured
pub is_configured: bool,
}
@@ -70,7 +68,6 @@ pub struct ProvidersResponse {
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct ToolPermission {
/// Unique identifier and name of the tool, format <extension_name>__<tool_name>
pub tool_name: String,
pub permission: PermissionLevel,
}
@@ -94,7 +91,6 @@ pub async fn upsert_config(
headers: HeaderMap,
Json(query): Json<UpsertConfigQuery>,
) -> Result<Json<Value>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
@@ -121,12 +117,10 @@ pub async fn remove_config(
headers: HeaderMap,
Json(query): Json<ConfigKeyQuery>,
) -> Result<Json<String>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
// Check if the secret flag is true and call the appropriate method
let result = if query.is_secret {
config.delete_secret(&query.key)
} else {
@@ -142,7 +136,7 @@ pub async fn remove_config(
#[utoipa::path(
post,
path = "/config/read",
request_body = ConfigKeyQuery, // Switch back to request_body
request_body = ConfigKeyQuery,
responses(
(status = 200, description = "Configuration value retrieved successfully", body = Value),
(status = 404, description = "Configuration key not found")
@@ -155,7 +149,6 @@ pub async fn read_config(
) -> Result<Json<Value>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Special handling for model-limits
if query.key == "model-limits" {
let limits = ModelConfig::get_all_model_limits();
return Ok(Json(
@@ -166,13 +159,10 @@ pub async fn read_config(
let config = Config::global();
match config.get(&query.key, query.is_secret) {
// Always get the actual value
Ok(value) => {
if query.is_secret {
// If it's marked as secret, return a boolean indicating presence
Ok(Json(Value::Bool(true)))
} else {
// Return the actual value if not secret
Ok(Json(value))
}
}
@@ -197,7 +187,6 @@ pub async fn get_extensions(
match ExtensionConfigManager::get_all() {
Ok(extensions) => Ok(Json(ExtensionResponse { extensions })),
Err(err) => {
// Return UNPROCESSABLE_ENTITY only for DeserializeError, INTERNAL_SERVER_ERROR for everything else
if err
.downcast_ref::<goose::config::base::ConfigError>()
.is_some_and(|e| matches!(e, goose::config::base::ConfigError::DeserializeError(_)))
@@ -228,7 +217,6 @@ pub async fn add_extension(
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Get existing extensions to check if this is an update
let extensions =
ExtensionConfigManager::get_all().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let key = name_to_key(&extension_query.name);
@@ -284,12 +272,10 @@ pub async fn read_all_config(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
) -> Result<Json<ConfigResponse>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
// Load values from config file
let values = config
.load_values()
.map_err(|_| StatusCode::UNPROCESSABLE_ENTITY)?;
@@ -297,7 +283,6 @@ pub async fn read_all_config(
Ok(Json(ConfigResponse { config: values }))
}
// Modified providers function using the new response type
#[utoipa::path(
get,
path = "/config/providers",
@@ -311,14 +296,11 @@ pub async fn providers(
) -> Result<Json<Vec<ProviderDetails>>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Fetch the list of providers, which are likely stored in the AppState or can be retrieved via a function call
let providers_metadata = get_providers();
// Construct the response by checking configuration status for each provider
let providers_response: Vec<ProviderDetails> = providers_metadata
.into_iter()
.map(|metadata| {
// Check if the provider is configured (this will depend on how you track configuration status)
let is_configured = check_provider_configured(&metadata);
ProviderDetails {
@@ -348,21 +330,16 @@ pub async fn init_config(
let config = Config::global();
// 200 if config already exists
if config.exists() {
return Ok(Json("Config already exists".to_string()));
}
// Find the workspace root (where the top-level Cargo.toml with [workspace] is)
let workspace_root = match std::env::current_exe() {
Ok(mut exe_path) => {
// Start from the executable's directory and traverse up
while let Some(parent) = exe_path.parent() {
let cargo_toml = parent.join("Cargo.toml");
if cargo_toml.exists() {
// Read the Cargo.toml file
if let Ok(content) = std::fs::read_to_string(&cargo_toml) {
// Check if it contains [workspace]
if content.contains("[workspace]") {
exe_path = parent.to_path_buf();
break;
@@ -376,7 +353,6 @@ pub async fn init_config(
Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
};
// Check if init-config.yaml exists at workspace root
let init_config_path = workspace_root.join("init-config.yaml");
if !init_config_path.exists() {
return Ok(Json(
@@ -384,7 +360,6 @@ pub async fn init_config(
));
}
// Read init-config.yaml and validate
let init_content = match std::fs::read_to_string(&init_config_path) {
Ok(content) => content,
Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
@@ -394,7 +369,6 @@ pub async fn init_config(
Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
};
// Save init-config.yaml to ~/.config/goose/config.yaml
match config.save_values(init_values) {
Ok(_) => Ok(Json("Config initialized successfully".to_string())),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
@@ -418,7 +392,7 @@ pub async fn upsert_permissions(
verify_secret_key(&headers, &state)?;
let mut permission_manager = PermissionManager::default();
// Iterate over each tool permission and update permissions
for tool_permission in &query.tool_permissions {
permission_manager.update_user_permission(
&tool_permission.tool_name,
@@ -429,12 +403,6 @@ pub async fn upsert_permissions(
Ok(Json("Permissions updated successfully".to_string()))
}
pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
top_level_domain: "Block".to_string(),
author: "Block".to_string(),
app_name: "goose".to_string(),
});
#[utoipa::path(
post,
path = "/config/backup",
@@ -460,11 +428,9 @@ pub async fn backup_config(
.file_name()
.ok_or(StatusCode::INTERNAL_SERVER_ERROR)?;
// Append ".bak" to the file name
let mut backup_name = file_name.to_os_string();
backup_name.push(".bak");
// Construct the new path with the same parent directory
let backup = config_path.with_file_name(backup_name);
match std::fs::rename(&config_path, &backup) {
Ok(_) => Ok(Json(format!("Moved {:?} to {:?}", config_path, backup))),
@@ -483,7 +449,7 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/read", post(read_config))
.route("/config/extensions", get(get_extensions))
.route("/config/extensions", post(add_extension))
.route("/config/extensions/:name", delete(remove_extension))
.route("/config/extensions/{name}", delete(remove_extension))
.route("/config/providers", get(providers))
.route("/config/init", post(init_config))
.route("/config/backup", post(backup_config))
@@ -497,16 +463,22 @@ mod tests {
#[tokio::test]
async fn test_read_model_limits() {
// Create test state and headers
let test_state = AppState::new(
Arc::new(goose::agents::Agent::default()),
"test".to_string(),
)
.await;
let sched_storage_path = choose_app_strategy(APP_STRATEGY.clone())
.unwrap()
.data_dir()
.join("schedules.json");
let sched = goose::scheduler::Scheduler::new(sched_storage_path)
.await
.unwrap();
test_state.set_scheduler(sched).await;
let mut headers = HeaderMap::new();
headers.insert("X-Secret-Key", "test".parse().unwrap());
// Execute
let result = read_config(
State(test_state),
headers,
@@ -517,16 +489,13 @@ mod tests {
)
.await;
// Assert
assert!(result.is_ok());
let response = result.unwrap();
// Parse the response and check the contents
let limits: Vec<goose::model::ModelLimitConfig> =
serde_json::from_value(response.0).unwrap();
assert!(!limits.is_empty());
// Check for some expected patterns
let gpt4_limit = limits.iter().find(|l| l.pattern == "gpt-4o");
assert!(gpt4_limit.is_some());
assert_eq!(gpt4_limit.unwrap().context_limit, 128_000);
crates/goose-server/src/routes/mod.rs
+2
View File
@@ -6,6 +6,7 @@ pub mod extension;
pub mod health;
pub mod recipe;
pub mod reply;
pub mod schedule;
pub mod session;
pub mod utils;
use std::sync::Arc;
@@ -23,4 +24,5 @@ pub fn configure(state: Arc<crate::state::AppState>) -> Router {
.merge(config_management::routes(state.clone()))
.merge(recipe::routes(state.clone()))
.merge(session::routes(state.clone()))
.merge(schedule::routes(state.clone()))
}
crates/goose-server/src/routes/reply.rs
+9 -30
View File
@@ -35,7 +35,6 @@ use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
use utoipa::ToSchema;
// Direct message serialization for the chat request
#[derive(Debug, Deserialize)]
struct ChatRequest {
messages: Vec<Message>,
@@ -43,7 +42,6 @@ struct ChatRequest {
session_working_dir: String,
}
// Custom SSE response type for streaming messages
pub struct SseResponse {
rx: ReceiverStream<String>,
}
@@ -78,7 +76,6 @@ impl IntoResponse for SseResponse {
}
}
// Message event types for SSE streaming
#[derive(Debug, Serialize)]
#[serde(tag = "type")]
enum MessageEvent {
@@ -87,7 +84,6 @@ enum MessageEvent {
Finish { reason: String },
}
// Stream a message as an SSE event
async fn stream_event(
event: MessageEvent,
tx: &mpsc::Sender<String>,
@@ -108,19 +104,16 @@ async fn handler(
) -> Result<SseResponse, StatusCode> {
verify_secret_key(&headers, &state)?;
// Create channel for streaming
let (tx, rx) = mpsc::channel(100);
let stream = ReceiverStream::new(rx);
let messages = request.messages;
let session_working_dir = request.session_working_dir;
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
// Spawn task to handle streaming
tokio::spawn(async move {
let agent = state.get_agent().await;
let agent = match agent {
@@ -166,7 +159,6 @@ async fn handler(
}
};
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
let mut stream = match agent
@@ -175,6 +167,7 @@ async fn handler(
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(session_working_dir),
schedule_id: None,
}),
)
.await
@@ -200,7 +193,6 @@ async fn handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
@@ -221,7 +213,7 @@ async fn handler(
break;
}
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = Arc::clone(provider.as_ref().unwrap());
@@ -255,7 +247,6 @@ async fn handler(
}
}
// Send finish event
let _ = stream_event(
MessageEvent::Finish {
reason: "stop".to_string(),
@@ -280,7 +271,6 @@ struct AskResponse {
response: String,
}
// Simple ask an AI for a response, non streaming
async fn ask_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
@@ -290,7 +280,6 @@ async fn ask_handler(
let session_working_dir = request.session_working_dir;
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
@@ -300,13 +289,10 @@ async fn ask_handler(
.await
.map_err(|_| StatusCode::PRECONDITION_FAILED)?;
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
// Create a single message for the prompt
let messages = vec![Message::user().with_text(request.prompt)];
// Get response from agent
let mut response_text = String::new();
let mut stream = match agent
.reply(
@@ -314,6 +300,7 @@ async fn ask_handler(
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(session_working_dir),
schedule_id: None,
}),
)
.await
@@ -325,7 +312,6 @@ async fn ask_handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let mut response_message = Message::assistant();
@@ -349,15 +335,12 @@ async fn ask_handler(
}
}
// Add the complete response message to the conversation history
if !response_message.content.is_empty() {
all_messages.push(response_message);
}
// Get the session path - file will be created when needed
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = Arc::clone(provider.as_ref().unwrap());
@@ -438,13 +421,11 @@ async fn submit_tool_result(
) -> Result<Json<Value>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Log the raw request for debugging
tracing::info!(
"Received tool result request: {}",
serde_json::to_string_pretty(&raw.0).unwrap()
);
// Try to parse into our struct
let payload: ToolResultRequest = match serde_json::from_value(raw.0.clone()) {
Ok(req) => req,
Err(e) => {
@@ -465,7 +446,6 @@ async fn submit_tool_result(
Ok(Json(json!({"status": "ok"})))
}
// Configure routes for this module
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/reply", post(handler))
@@ -488,7 +468,6 @@ mod tests {
};
use mcp_core::tool::Tool;
// Mock Provider implementation for testing
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
@@ -523,10 +502,8 @@ mod tests {
use std::sync::Arc;
use tower::ServiceExt;
// This test requires tokio runtime
#[tokio::test]
async fn test_ask_endpoint() {
// Create a mock app state with mock provider
let mock_model_config = ModelConfig::new("test-model".to_string());
let mock_provider = Arc::new(MockProvider {
model_config: mock_model_config,
@@ -534,11 +511,15 @@ mod tests {
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let state = AppState::new(Arc::new(agent), "test-secret".to_string()).await;
let scheduler_path = goose::scheduler::get_default_scheduler_storage_path()
.expect("Failed to get default scheduler storage path");
let scheduler = goose::scheduler::Scheduler::new(scheduler_path)
.await
.unwrap();
state.set_scheduler(scheduler).await;
// Build router
let app = routes(state);
// Create request
let request = Request::builder()
.uri("/ask")
.method("POST")
@@ -554,10 +535,8 @@ mod tests {
))
.unwrap();
// Send request
let response = app.oneshot(request).await.unwrap();
// Assert response status
assert_eq!(response.status(), StatusCode::OK);
}
}
crates/goose-server/src/routes/schedule.rs
+270
View File
@@ -0,0 +1,270 @@
use std::sync::Arc;
use axum::{
extract::{Path, Query, State},
http::{HeaderMap, StatusCode},
routing::{delete, get, post},
Json, Router,
};
use serde::{Deserialize, Serialize};
use chrono::NaiveDateTime;
use crate::routes::utils::verify_secret_key;
use crate::state::AppState;
use goose::scheduler::ScheduledJob;
#[derive(Deserialize, Serialize, utoipa::ToSchema)]
pub struct CreateScheduleRequest {
id: String,
recipe_source: String,
cron: String,
}
#[derive(Serialize, utoipa::ToSchema)]
pub struct ListSchedulesResponse {
jobs: Vec<ScheduledJob>,
}
// Response for the run_now endpoint
#[derive(Serialize, utoipa::ToSchema)]
pub struct RunNowResponse {
session_id: String,
}
// Query parameters for the sessions endpoint
#[derive(Deserialize, utoipa::ToSchema, utoipa::IntoParams)]
pub struct SessionsQuery {
#[serde(default = "default_limit")]
limit: u32,
}
fn default_limit() -> u32 {
50 // Default limit for sessions listed
}
// Struct for the frontend session list
#[derive(Serialize, utoipa::ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct SessionDisplayInfo {
id: String, // Derived from session_name (filename)
name: String, // From metadata.description
created_at: String, // Derived from session_name, in ISO 8601 format
working_dir: String, // from metadata.working_dir (as String)
schedule_id: Option<String>,
message_count: usize,
total_tokens: Option<i32>,
input_tokens: Option<i32>,
output_tokens: Option<i32>,
accumulated_total_tokens: Option<i32>,
accumulated_input_tokens: Option<i32>,
accumulated_output_tokens: Option<i32>,
}
fn parse_session_name_to_iso(session_name: &str) -> String {
NaiveDateTime::parse_from_str(session_name, "%Y%m%d_%H%M%S")
.map(|dt| dt.and_utc().to_rfc3339())
.unwrap_or_else(|_| String::new()) // Fallback to empty string if parsing fails
}
#[utoipa::path(
post,
path = "/schedule/create",
request_body = CreateScheduleRequest,
responses(
(status = 200, description = "Scheduled job created successfully", body = ScheduledJob),
(status = 500, description = "Internal server error")
),
tag = "schedule"
)]
#[axum::debug_handler]
async fn create_schedule(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(req): Json<CreateScheduleRequest>,
) -> Result<Json<ScheduledJob>, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let job = ScheduledJob {
id: req.id,
source: req.recipe_source,
cron: req.cron,
last_run: None,
};
scheduler
.add_scheduled_job(job.clone())
.await
.map_err(|e| {
eprintln!("Error creating schedule: {:?}", e); // Log error
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(job))
}
#[utoipa::path(
get,
path = "/schedule/list",
responses(
(status = 200, description = "A list of scheduled jobs", body = ListSchedulesResponse),
(status = 500, description = "Internal server error")
),
tag = "schedule"
)]
#[axum::debug_handler]
async fn list_schedules(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
) -> Result<Json<ListSchedulesResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let jobs = scheduler.list_scheduled_jobs().await;
Ok(Json(ListSchedulesResponse { jobs }))
}
#[utoipa::path(
delete,
path = "/schedule/delete/{id}",
params(
("id" = String, Path, description = "ID of the schedule to delete")
),
responses(
(status = 204, description = "Scheduled job deleted successfully"),
(status = 404, description = "Scheduled job not found"),
(status = 500, description = "Internal server error")
),
tag = "schedule"
)]
#[axum::debug_handler]
async fn delete_schedule(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Path(id): Path<String>,
) -> Result<StatusCode, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
scheduler.remove_scheduled_job(&id).await.map_err(|e| {
eprintln!("Error deleting schedule '{}': {:?}", id, e);
match e {
goose::scheduler::SchedulerError::JobNotFound(_) => StatusCode::NOT_FOUND,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
Ok(StatusCode::NO_CONTENT)
}
#[utoipa::path(
post,
path = "/schedule/{id}/run_now",
params(
("id" = String, Path, description = "ID of the schedule to run")
),
responses(
(status = 200, description = "Scheduled job triggered successfully, returns new session ID", body = RunNowResponse),
(status = 404, description = "Scheduled job not found"),
(status = 500, description = "Internal server error when trying to run the job")
),
tag = "schedule"
)]
#[axum::debug_handler]
async fn run_now_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Path(id): Path<String>,
) -> Result<Json<RunNowResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
match scheduler.run_now(&id).await {
Ok(session_id) => Ok(Json(RunNowResponse { session_id })),
Err(e) => {
eprintln!("Error running schedule '{}' now: {:?}", id, e);
match e {
goose::scheduler::SchedulerError::JobNotFound(_) => Err(StatusCode::NOT_FOUND),
_ => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
}
}
#[utoipa::path(
get,
path = "/schedule/{id}/sessions",
params(
("id" = String, Path, description = "ID of the schedule"),
SessionsQuery // This will automatically pick up 'limit' as a query parameter
),
responses(
(status = 200, description = "A list of session display info", body = Vec<SessionDisplayInfo>),
(status = 500, description = "Internal server error")
),
tag = "schedule"
)]
#[axum::debug_handler]
async fn sessions_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap, // Added this line
Path(schedule_id_param): Path<String>, // Renamed to avoid confusion with session_id
Query(query_params): Query<SessionsQuery>,
) -> Result<Json<Vec<SessionDisplayInfo>>, StatusCode> {
verify_secret_key(&headers, &state)?; // Added this line
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
match scheduler
.sessions(&schedule_id_param, query_params.limit as usize)
.await
{
Ok(session_tuples) => {
// Expecting Vec<(String, goose::session::storage::SessionMetadata)>
let display_infos: Vec<SessionDisplayInfo> = session_tuples
.into_iter()
.map(|(session_name, metadata)| SessionDisplayInfo {
id: session_name.clone(),
name: metadata.description, // Use description as name
created_at: parse_session_name_to_iso(&session_name),
working_dir: metadata.working_dir.to_string_lossy().into_owned(),
schedule_id: metadata.schedule_id, // This is the ID of the schedule itself
message_count: metadata.message_count,
total_tokens: metadata.total_tokens,
input_tokens: metadata.input_tokens,
output_tokens: metadata.output_tokens,
accumulated_total_tokens: metadata.accumulated_total_tokens,
accumulated_input_tokens: metadata.accumulated_input_tokens,
accumulated_output_tokens: metadata.accumulated_output_tokens,
})
.collect();
Ok(Json(display_infos))
}
Err(e) => {
eprintln!(
"Error fetching sessions for schedule '{}': {:?}",
schedule_id_param, e
);
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/schedule/create", post(create_schedule))
.route("/schedule/list", get(list_schedules))
.route("/schedule/delete/{id}", delete(delete_schedule)) // Corrected
.route("/schedule/{id}/run_now", post(run_now_handler)) // Corrected
.route("/schedule/{id}/sessions", get(sessions_handler)) // Corrected
.with_state(state)
}
crates/goose-server/src/routes/session.rs
+1 -1
View File
@@ -108,6 +108,6 @@ async fn get_session_history(
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/sessions", get(list_sessions))
.route("/sessions/:session_id", get(get_session_history))
.route("/sessions/{session_id}", get(get_session_history))
.with_state(state)
}
crates/goose-server/src/state.rs
+17 -9
View File
@@ -1,21 +1,15 @@
use goose::agents::Agent;
use goose::scheduler::Scheduler;
use std::sync::Arc;
use tokio::sync::Mutex;
/// Shared reference to an Agent that can be cloned cheaply
/// without cloning the underlying Agent object
pub type AgentRef = Arc<Agent>;
/// Thread-safe container for an optional Agent reference
/// Outer Arc: Allows multiple route handlers to access the same Mutex
/// - Mutex provides exclusive access for updates
/// - Option allows for the case where no agent exists yet
///
/// Shared application state
#[derive(Clone)]
pub struct AppState {
// agent: SharedAgentStore,
agent: Option<AgentRef>,
pub secret_key: String,
pub scheduler: Arc<Mutex<Option<Arc<Scheduler>>>>,
}
impl AppState {
@@ -23,6 +17,7 @@ impl AppState {
Arc::new(Self {
agent: Some(agent.clone()),
secret_key,
scheduler: Arc::new(Mutex::new(None)),
})
}
@@ -31,4 +26,17 @@ impl AppState {
.clone()
.ok_or_else(|| anyhow::anyhow!("Agent needs to be created first."))
}
pub async fn set_scheduler(&self, sched: Arc<Scheduler>) {
let mut guard = self.scheduler.lock().await;
*guard = Some(sched);
}
pub async fn scheduler(&self) -> Result<Arc<Scheduler>, anyhow::Error> {
self.scheduler
.lock()
.await
.clone()
.ok_or_else(|| anyhow::anyhow!("Scheduler not initialized"))
}
}