diff --git a/src/api.rs b/src/domain/agent/api.rs similarity index 100% rename from src/api.rs rename to src/domain/agent/api.rs diff --git a/src/agent.rs b/src/domain/agent/mod.rs similarity index 98% rename from src/agent.rs rename to src/domain/agent/mod.rs index 7b433e5..91f5147 100644 --- a/src/agent.rs +++ b/src/domain/agent/mod.rs @@ -1,8 +1,10 @@ +pub mod api; +pub mod tools; + use chrono::Utc; use std::time::{Duration, Instant}; -use crate::api::{ChatRequest, ChatResponse, Message, Tool}; -use crate::tools; +use self::api::{ChatRequest, ChatResponse, Message, Tool}; pub struct Agent { client: reqwest::Client, diff --git a/src/tools.rs b/src/domain/agent/tools.rs similarity index 97% rename from src/tools.rs rename to src/domain/agent/tools.rs index 37e665c..eb9b2ab 100644 --- a/src/tools.rs +++ b/src/domain/agent/tools.rs @@ -1,4 +1,4 @@ -use crate::api::{self, FunctionDefinition, Message, Tool, ToolCall}; +use super::api::{self, FunctionDefinition, Message, Tool, ToolCall}; use std::collections::HashMap; pub fn get_tools() -> Vec { diff --git a/src/auth.rs b/src/domain/auth.rs similarity index 96% rename from src/auth.rs rename to src/domain/auth.rs index b2f8c8d..c70e0c7 100644 --- a/src/auth.rs +++ b/src/domain/auth.rs @@ -45,12 +45,8 @@ pub struct JwksVerifier { } impl JwksVerifier { - pub async fn new( - issuer: String, - audience: String, - ) -> Result> { + pub async fn new(issuer: String, audience: String) -> Result> { let client = Client::new(); - // Authentik OIDC discovery let discovery_url = format!( "{}/.well-known/openid-configuration", issuer.trim_end_matches('/') diff --git a/src/domain/mod.rs b/src/domain/mod.rs new file mode 100644 index 0000000..7982060 --- /dev/null +++ b/src/domain/mod.rs @@ -0,0 +1,3 @@ +pub mod agent; +pub mod auth; +pub mod tasks; diff --git a/src/domain/tasks.rs b/src/domain/tasks.rs new file mode 100644 index 0000000..9e64d8e --- /dev/null +++ b/src/domain/tasks.rs @@ -0,0 +1,131 @@ +use chrono::Utc; +use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, Set}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use uuid::Uuid; + +use crate::config::Config; +use crate::domain::agent::Agent; +use crate::entities::task::Entity as Task; +use crate::entities::task_run::{self, Entity as TaskRun}; +use crate::scheduler::Scheduler; + +#[derive(Deserialize)] +pub struct CreateTaskRequest { + pub goal: String, + pub cron: Option, +} + +#[derive(Deserialize)] +pub struct UpdateTaskRequest { + pub goal: String, + pub cron: Option, +} + +#[derive(Serialize)] +pub struct TaskResponse { + pub id: Uuid, + pub goal: String, + pub cron: Option, + pub created_at: chrono::DateTime, + pub runs: Vec, +} + +#[derive(Serialize)] +pub struct TaskRunResponse { + pub id: Uuid, + pub status: String, + pub logs: String, + pub answer: Option, + pub created_at: chrono::DateTime, +} + +#[derive(Serialize)] +pub struct RecentRunResponse { + pub id: Uuid, + pub task_id: Uuid, + pub goal: String, + pub status: String, + pub created_at: chrono::DateTime, +} + +pub async fn execute_agent_run( + db: &DatabaseConnection, + _scheduler: &Arc, + config: &Arc, + task_id: Uuid, + goal: String, +) -> Result> { + let run_id = Uuid::new_v4(); + + let new_run = task_run::ActiveModel { + id: Set(run_id), + task_id: Set(task_id), + status: Set("running".to_string()), + logs: Set(String::new()), + answer: Set(None), + created_at: Set(Utc::now().into()), + }; + + new_run.insert(db).await?; + + let mut agent = Agent::new( + config.zen_api_key.clone(), + config.tavily_api_key.clone(), + goal.clone(), + )?; + + let (logs, answer, status) = match agent.run(config).await { + Ok((logs, answer)) => (logs, answer, "completed".to_string()), + Err(e) => ( + format!("Execution failed: {}", e), + None, + "failed".to_string(), + ), + }; + + let mut run: task_run::ActiveModel = TaskRun::find_by_id(run_id) + .one(db) + .await? + .ok_or("Run not found after insert")? + .into(); + + run.logs = Set(logs.clone()); + run.answer = Set(answer.clone()); + run.status = Set(status); + + run.update(db).await?; + + get_task_inner(task_id, db).await +} + +pub async fn get_task_inner( + id: Uuid, + db: &DatabaseConnection, +) -> Result> { + let results = Task::find_by_id(id) + .find_with_related(TaskRun) + .all(db) + .await?; + + let (t, mut runs) = results.into_iter().next().ok_or("Task not found")?; + + runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); + + Ok(TaskResponse { + id: t.id, + goal: t.goal, + cron: t.cron, + created_at: t.created_at, + runs: runs + .into_iter() + .map(|r| TaskRunResponse { + id: r.id, + status: r.status, + logs: r.logs, + answer: r.answer, + created_at: r.created_at, + }) + .collect(), + }) +} diff --git a/src/main.rs b/src/main.rs index efc614d..36c5a8d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,8 @@ -mod agent; -mod api; -mod auth; mod config; +mod domain; mod entities; mod scheduler; mod server; -mod tools; #[tokio::main] async fn main() -> Result<(), Box> { diff --git a/src/scheduler.rs b/src/scheduler.rs index 1c4aeaf..106a047 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,4 +1,4 @@ -use crate::agent::Agent; +use crate::domain::agent::Agent; use crate::entities::task::Entity as Task; use crate::entities::task_run; use dashmap::DashMap; diff --git a/src/server.rs b/src/server.rs deleted file mode 100644 index 443ee04..0000000 --- a/src/server.rs +++ /dev/null @@ -1,585 +0,0 @@ -use axum::{ - Json, RequestPartsExt, Router, - extract::{FromRef, FromRequestParts, Path, Query, State}, - http::{HeaderValue, StatusCode, request::Parts}, - response::IntoResponse, - routing::{get, post}, -}; -use axum_extra::{ - TypedHeader, - extract::cookie::{Cookie, CookieJar, SameSite}, - headers::{Authorization, authorization::Bearer}, -}; -use chrono::Utc; -use sea_orm::{ - ActiveModelTrait, Database, DatabaseConnection, EntityTrait, QueryOrder, QuerySelect, Set, -}; -use serde::{Deserialize, Serialize}; -use std::sync::Arc; -use tower_http::cors::{AllowOrigin, CorsLayer}; -use uuid::Uuid; - -use crate::agent::Agent; -use crate::entities::task::{self, Entity as Task}; -use crate::entities::task_run::{self, Entity as TaskRun}; -use crate::scheduler::Scheduler; -use migration::{Migrator, MigratorTrait}; - -#[derive(Clone)] -pub struct AppState { - pub db: DatabaseConnection, - pub scheduler: Arc, - pub config: Arc, - pub verifier: Arc, - pub authenticator: Arc, -} - -#[derive(Deserialize)] -pub struct CreateTaskRequest { - pub goal: String, - pub cron: Option, -} - -#[derive(Deserialize)] -pub struct UpdateTaskRequest { - pub goal: String, - pub cron: Option, -} - -#[derive(Serialize)] -pub struct TaskResponse { - pub id: Uuid, - pub goal: String, - pub cron: Option, - pub created_at: chrono::DateTime, - pub runs: Vec, -} - -#[derive(Serialize)] -pub struct TaskRunResponse { - pub id: Uuid, - pub status: String, - pub logs: String, - pub answer: Option, - pub created_at: chrono::DateTime, -} - -#[derive(Serialize)] -pub struct RecentRunResponse { - pub id: Uuid, - pub task_id: Uuid, - pub goal: String, - pub status: String, - pub created_at: chrono::DateTime, -} - -pub async fn start(config: crate::config::Config) -> Result<(), Box> { - let db = Database::connect(&config.database_url).await?; - Migrator::up(&db, None).await?; - - let config = Arc::new(config); - - let scheduler = Arc::new(Scheduler::new(db.clone(), config.clone()).await?); - - // Load existing scheduled tasks - let existing_tasks = Task::find().all(&db).await?; - for task in existing_tasks { - if let Some(cron) = task.cron { - let _ = scheduler.add_task_job(task.id, &cron).await; - } - } - - let verifier = Arc::new( - crate::auth::JwksVerifier::new( - config.authentik_issuer.clone(), - config.authentik_client_id.clone(), - ) - .await?, - ); - let authenticator = Arc::new( - crate::auth::Authenticator::new( - config.authentik_issuer.clone(), - config.authentik_client_id.clone(), - config.authentik_client_secret.clone(), - ) - .await?, - ); - - let state = Arc::new(AppState { - db, - scheduler, - config: config.clone(), - verifier, - authenticator, - }); - - let cors = build_cors_layer(&config); - - let app = Router::new() - .route("/api/tasks", post(create_task).get(list_tasks)) - .route("/api/tasks/:id", get(get_task).put(update_task)) - .route("/api/tasks/:id/runs", post(rerun_task)) - .route("/api/runs/recent", get(get_recent_runs)) - .route("/api/auth/session", get(auth_session)) - .route("/api/auth/callback", get(auth_callback)) - .route("/api/auth/refresh", post(auth_refresh)) - .route("/api/auth/logout", post(auth_logout)) - .layer(cors) - .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( - axum::http::header::CONTENT_SECURITY_POLICY, - HeaderValue::from_static("default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data:;"), - )) - .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( - axum::http::header::X_CONTENT_TYPE_OPTIONS, - HeaderValue::from_static("nosniff"), - )) - .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( - axum::http::header::REFERRER_POLICY, - HeaderValue::from_static("strict-origin-when-cross-origin"), - )) - .layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit - .with_state(state); - - let addr = format!("0.0.0.0:{}", config.port); - let listener = tokio::net::TcpListener::bind(&addr).await?; - println!("Server running on http://localhost:{}", config.port); - axum::serve(listener, app).await?; - - Ok(()) -} - -fn build_cors_layer(config: &crate::config::Config) -> CorsLayer { - let allow_origin = if let Some(origins) = &config.cors_allowed_origins { - let values: Vec = origins - .split(',') - .map(|origin| origin.trim()) - .filter(|origin| !origin.is_empty()) - .filter_map(|origin| HeaderValue::from_str(origin).ok()) - .collect(); - - if values.is_empty() { - AllowOrigin::mirror_request() - } else { - AllowOrigin::list(values) - } - } else { - AllowOrigin::mirror_request() - }; - - CorsLayer::new() - .allow_origin(allow_origin) - .allow_methods([ - axum::http::Method::GET, - axum::http::Method::POST, - axum::http::Method::PUT, - axum::http::Method::PATCH, - axum::http::Method::DELETE, - axum::http::Method::OPTIONS, - ]) - .allow_headers([ - axum::http::header::CONTENT_TYPE, - axum::http::header::AUTHORIZATION, - axum::http::header::ACCEPT, - ]) - .allow_credentials(true) -} - -async fn list_tasks( - _user: AuthenticatedUser, - State(state): State>, -) -> Result>, (StatusCode, String)> { - let tasks = Task::find() - .find_with_related(TaskRun) - .order_by_desc(task::Column::CreatedAt) - .all(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let response = tasks - .into_iter() - .map(|(t, mut runs)| { - runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); - TaskResponse { - id: t.id, - goal: t.goal, - cron: t.cron, - created_at: t.created_at, - runs: runs - .into_iter() - .map(|r| TaskRunResponse { - id: r.id, - status: r.status, - logs: r.logs, - answer: r.answer, - created_at: r.created_at, - }) - .collect(), - } - }) - .collect(); - - Ok(Json(response)) -} - -async fn create_task( - _user: AuthenticatedUser, - State(state): State>, - Json(payload): Json, -) -> Result, (StatusCode, String)> { - let task_id = Uuid::new_v4(); - - // Initial task save - let new_task = task::ActiveModel { - id: Set(task_id), - goal: Set(payload.goal.clone()), - cron: Set(payload.cron.clone()), - created_at: Set(Utc::now().into()), - }; - - new_task - .insert(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - if let Some(cron) = &payload.cron { - let _ = state.scheduler.add_task_job(task_id, cron).await; - } else { - let _ = state.scheduler.remove_task_job(task_id).await; - } - - get_task_inner(task_id, &state).await.map(Json) -} - -async fn rerun_task( - _user: AuthenticatedUser, - State(state): State>, - Path(id): Path, -) -> Result, (StatusCode, String)> { - let task = Task::find_by_id(id) - .one(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?; - - execute_agent_run(state, task.id, task.goal).await -} - -async fn execute_agent_run( - state: Arc, - task_id: Uuid, - goal: String, -) -> Result, (StatusCode, String)> { - let run_id = Uuid::new_v4(); - - // Initial run save - let new_run = task_run::ActiveModel { - id: Set(run_id), - task_id: Set(task_id), - status: Set("running".to_string()), - logs: Set(String::new()), - answer: Set(None), - created_at: Set(Utc::now().into()), - }; - - new_run - .insert(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let mut agent = Agent::new( - state.config.zen_api_key.clone(), - state.config.tavily_api_key.clone(), - goal.clone(), - ) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let (logs, answer, status) = match agent.run(&state.config).await { - Ok((logs, answer)) => (logs, answer, "completed".to_string()), - Err(e) => ( - format!("Execution failed: {}", e), - None, - "failed".to_string(), - ), - }; - - // Update with final logs and status - let mut run: task_run::ActiveModel = TaskRun::find_by_id(run_id) - .one(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .ok_or(( - StatusCode::NOT_FOUND, - "Run not found after insert".to_string(), - ))? - .into(); - - run.logs = Set(logs.clone()); - run.answer = Set(answer.clone()); - run.status = Set(status); - - run.update(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - get_task_inner(task_id, &state).await.map(Json) -} - -async fn update_task( - _user: AuthenticatedUser, - State(state): State>, - Path(id): Path, - Json(payload): Json, -) -> Result, (StatusCode, String)> { - let mut task: task::ActiveModel = Task::find_by_id(id) - .one(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))? - .into(); - - task.goal = Set(payload.goal.clone()); - task.cron = Set(payload.cron.clone()); - - task.update(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - if let Some(cron) = &payload.cron { - let _ = state.scheduler.add_task_job(id, cron).await; - } else { - let _ = state.scheduler.remove_task_job(id).await; - } - - get_task_inner(id, &state).await.map(Json) -} - -#[allow(dead_code)] -pub struct AuthenticatedUser(pub crate::auth::Claims); - -#[axum::async_trait] -impl FromRequestParts for AuthenticatedUser -where - Arc: axum::extract::FromRef, - S: Send + Sync, -{ - type Rejection = (StatusCode, String); - - async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { - let app_state = Arc::::from_ref(state); - - let token = if let Ok(TypedHeader(Authorization(bearer))) = - parts.extract::>>().await - { - Some(bearer.token().to_string()) - } else { - let jar = parts.extract::().await.unwrap(); - jar.get("access_token") - .map(|cookie| cookie.value().to_string()) - }; - - let token = token.ok_or(( - StatusCode::UNAUTHORIZED, - "Missing or invalid access token".to_string(), - ))?; - - let claims = app_state.verifier.verify(&token).await.map_err(|e| { - ( - StatusCode::UNAUTHORIZED, - format!("Token verification failed: {}", e), - ) - })?; - - Ok(AuthenticatedUser(claims)) - } -} - -#[derive(Deserialize)] -pub struct AuthCallbackQuery { - pub code: String, - pub redirect_uri: String, -} - -#[derive(Deserialize)] -struct RefreshRequest { - refresh_token: Option, -} - -async fn auth_refresh( - State(state): State>, - jar: CookieJar, - Json(payload): Json, -) -> Result { - let refresh_token = payload - .refresh_token - .filter(|token| !token.is_empty()) - .or_else(|| { - jar.get("refresh_token") - .map(|cookie| cookie.value().to_string()) - }) - .ok_or(( - StatusCode::UNAUTHORIZED, - "Missing refresh token".to_string(), - ))?; - - let data = state - .authenticator - .refresh_token(refresh_token) - .await - .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?; - - let jar = update_auth_cookies(jar, &data, &state.config); - Ok((jar, Json(data))) -} - -async fn auth_callback( - State(state): State>, - jar: CookieJar, - Query(query): Query, -) -> Result { - let data = state - .authenticator - .exchange_code(query.code, query.redirect_uri) - .await - .map_err(|e| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Token exchange failed: {}", e), - ) - })?; - - let jar = update_auth_cookies(jar, &data, &state.config); - Ok((jar, Json(data))) -} - -async fn auth_logout(State(state): State>, jar: CookieJar) -> impl IntoResponse { - let jar = clear_auth_cookies(jar, &state.config); - (jar, StatusCode::NO_CONTENT) -} - -async fn auth_session(user: AuthenticatedUser) -> Json { - Json(serde_json::json!({ - "authenticated": true, - "user": user.0 - })) -} - -fn secure(config: &crate::config::Config) -> bool { - config.cookie_secure -} - -fn update_auth_cookies( - jar: CookieJar, - data: &serde_json::Value, - config: &crate::config::Config, -) -> CookieJar { - let access_token = data.get("access_token"); - let refresh_token = data.get("refresh_token"); - - let mut jar = jar; - - if let Some(token) = access_token.and_then(|t| t.as_str()) { - let cookie = Cookie::build(("access_token", token.to_owned())) - .path("/") - .http_only(true) - .same_site(SameSite::Lax) - .secure(secure(config)) - .build(); - jar = jar.add(cookie); - } - - if let Some(token) = refresh_token.and_then(|t| t.as_str()) { - let cookie = Cookie::build(("refresh_token", token.to_owned())) - .path("/") - .http_only(true) - .same_site(SameSite::Lax) - .secure(secure(config)) - .build(); - jar = jar.add(cookie); - } - - jar -} - -fn clear_auth_cookies(jar: CookieJar, config: &crate::config::Config) -> CookieJar { - let mut jar = jar; - for name in ["access_token", "refresh_token"] { - let cookie = Cookie::build((name, "")) - .path("/") - .http_only(true) - .same_site(SameSite::Lax) - .secure(secure(config)) - .max_age(cookie::time::Duration::seconds(0)) - .build(); - jar = jar.add(cookie); - } - - jar -} - -async fn get_task( - _user: AuthenticatedUser, - Path(id): Path, - State(state): State>, -) -> Result, (StatusCode, String)> { - get_task_inner(id, &state).await.map(Json) -} - -async fn get_task_inner(id: Uuid, state: &AppState) -> Result { - let results = Task::find_by_id(id) - .find_with_related(TaskRun) - .all(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let (t, mut runs) = results - .into_iter() - .next() - .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?; - - runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); - - Ok(TaskResponse { - id: t.id, - goal: t.goal, - cron: t.cron, - created_at: t.created_at, - runs: runs - .into_iter() - .map(|r| TaskRunResponse { - id: r.id, - status: r.status, - logs: r.logs, - answer: r.answer, - created_at: r.created_at, - }) - .collect(), - }) -} - -async fn get_recent_runs( - _user: AuthenticatedUser, - State(state): State>, -) -> Result>, (StatusCode, String)> { - let results = TaskRun::find() - .find_also_related(Task) - .order_by_desc(task_run::Column::CreatedAt) - .limit(50) - .all(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let response = results - .into_iter() - .filter_map(|(run, task_opt)| { - task_opt.map(|task| RecentRunResponse { - id: run.id, - task_id: run.task_id, - goal: task.goal, - status: run.status, - created_at: run.created_at, - }) - }) - .collect(); - - Ok(Json(response)) -} diff --git a/src/server/auth.rs b/src/server/auth.rs new file mode 100644 index 0000000..edc7465 --- /dev/null +++ b/src/server/auth.rs @@ -0,0 +1,178 @@ +use axum::{ + Json, RequestPartsExt, + extract::{FromRef, FromRequestParts, Query, State}, + http::{StatusCode, request::Parts}, + response::IntoResponse, +}; +use axum_extra::{ + TypedHeader, + extract::cookie::{Cookie, CookieJar, SameSite}, + headers::{Authorization, authorization::Bearer}, +}; +use std::sync::Arc; + +use super::AppState; +use crate::domain::auth::Claims; + +#[allow(dead_code)] +pub struct AuthenticatedUser(pub Claims); + +#[axum::async_trait] +impl FromRequestParts for AuthenticatedUser +where + Arc: axum::extract::FromRef, + S: Send + Sync, +{ + type Rejection = (StatusCode, String); + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let app_state = Arc::::from_ref(state); + + let token = if let Ok(TypedHeader(Authorization(bearer))) = + parts.extract::>>().await + { + Some(bearer.token().to_string()) + } else { + let jar = parts.extract::().await.unwrap(); + jar.get("access_token") + .map(|cookie| cookie.value().to_string()) + }; + + let token = token.ok_or(( + StatusCode::UNAUTHORIZED, + "Missing or invalid access token".to_string(), + ))?; + + let claims = app_state.verifier.verify(&token).await.map_err(|e| { + ( + StatusCode::UNAUTHORIZED, + format!("Token verification failed: {}", e), + ) + })?; + + Ok(AuthenticatedUser(claims)) + } +} + +#[derive(serde::Deserialize)] +pub struct AuthCallbackQuery { + pub code: String, + pub redirect_uri: String, +} + +#[derive(serde::Deserialize)] +pub struct RefreshRequest { + pub refresh_token: Option, +} + +pub async fn auth_refresh( + State(state): State>, + jar: CookieJar, + Json(payload): Json, +) -> Result { + let refresh_token = payload + .refresh_token + .filter(|token| !token.is_empty()) + .or_else(|| { + jar.get("refresh_token") + .map(|cookie| cookie.value().to_string()) + }) + .ok_or(( + StatusCode::UNAUTHORIZED, + "Missing refresh token".to_string(), + ))?; + + let data = state + .authenticator + .refresh_token(refresh_token) + .await + .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?; + + let jar = update_auth_cookies(jar, &data, &state.config); + Ok((jar, Json(data))) +} + +pub async fn auth_callback( + State(state): State>, + jar: CookieJar, + Query(query): Query, +) -> Result { + let data = state + .authenticator + .exchange_code(query.code, query.redirect_uri) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Token exchange failed: {}", e), + ) + })?; + + let jar = update_auth_cookies(jar, &data, &state.config); + Ok((jar, Json(data))) +} + +pub async fn auth_logout(State(state): State>, jar: CookieJar) -> impl IntoResponse { + let jar = clear_auth_cookies(jar, &state.config); + (jar, StatusCode::NO_CONTENT) +} + +pub async fn auth_session(user: AuthenticatedUser) -> Json { + Json(serde_json::json!({ + "authenticated": true, + "user": user.0 + })) +} + +fn secure(config: &crate::config::Config) -> bool { + config.cookie_secure +} + +pub fn update_auth_cookies( + jar: CookieJar, + data: &serde_json::Value, + config: &crate::config::Config, +) -> CookieJar { + let access_token = data.get("access_token"); + let refresh_token = data.get("refresh_token"); + + let mut jar = jar; + + if let Some(token) = access_token.and_then(|t| t.as_str()) { + let cookie = Cookie::build(("access_token", token.to_owned())) + .path("/") + .http_only(true) + .same_site(SameSite::Lax) + .secure(secure(config)) + .build(); + jar = jar.add(cookie); + } + + if let Some(token) = refresh_token.and_then(|t| t.as_str()) { + let cookie = Cookie::build(("refresh_token", token.to_owned())) + .path("/") + .http_only(true) + .same_site(SameSite::Lax) + .secure(secure(config)) + .build(); + jar = jar.add(cookie); + } + + jar +} + +pub fn clear_auth_cookies(jar: CookieJar, config: &crate::config::Config) -> CookieJar { + let mut jar = jar; + for name in ["access_token", "refresh_token"] { + let cookie = Cookie::build((name, "")) + .path("/") + .http_only(true) + .same_site(SameSite::Lax) + .secure(secure(config)) + .max_age(cookie::time::Duration::seconds(0)) + .build(); + jar = jar.add(cookie); + } + + jar +} diff --git a/src/server/mod.rs b/src/server/mod.rs new file mode 100644 index 0000000..8f88b7d --- /dev/null +++ b/src/server/mod.rs @@ -0,0 +1,135 @@ +pub mod auth; +pub mod tasks; + +use axum::{ + Router, + http::HeaderValue, + routing::{get, post}, +}; +use migration::{Migrator, MigratorTrait}; +use sea_orm::{Database, DatabaseConnection, EntityTrait}; +use std::sync::Arc; +use tower_http::cors::{AllowOrigin, CorsLayer}; + +use crate::entities::task::Entity as Task; +use crate::scheduler::Scheduler; + +#[derive(Clone)] +pub struct AppState { + pub db: DatabaseConnection, + pub scheduler: Arc, + pub config: Arc, + pub verifier: Arc, + pub authenticator: Arc, +} + +pub async fn start(config: crate::config::Config) -> Result<(), Box> { + let db = Database::connect(&config.database_url).await?; + Migrator::up(&db, None).await?; + + let config = Arc::new(config); + + let scheduler = Arc::new(Scheduler::new(db.clone(), config.clone()).await?); + + // Load existing scheduled tasks + let existing_tasks = Task::find().all(&db).await?; + for task in existing_tasks { + if let Some(cron) = task.cron { + let _ = scheduler.add_task_job(task.id, &cron).await; + } + } + + let verifier = Arc::new( + crate::domain::auth::JwksVerifier::new( + config.authentik_issuer.clone(), + config.authentik_client_id.clone(), + ) + .await?, + ); + let authenticator = Arc::new( + crate::domain::auth::Authenticator::new( + config.authentik_issuer.clone(), + config.authentik_client_id.clone(), + config.authentik_client_secret.clone(), + ) + .await?, + ); + + let state = Arc::new(AppState { + db, + scheduler, + config: config.clone(), + verifier, + authenticator, + }); + + let cors = build_cors_layer(&config); + + let app = Router::new() + .route("/api/tasks", post(tasks::create_task).get(tasks::list_tasks)) + .route("/api/tasks/:id", get(tasks::get_task).put(tasks::update_task)) + .route("/api/tasks/:id/runs", post(tasks::rerun_task)) + .route("/api/runs/recent", get(tasks::get_recent_runs)) + .route("/api/auth/session", get(auth::auth_session)) + .route("/api/auth/callback", get(auth::auth_callback)) + .route("/api/auth/refresh", post(auth::auth_refresh)) + .route("/api/auth/logout", post(auth::auth_logout)) + .layer(cors) + .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( + axum::http::header::CONTENT_SECURITY_POLICY, + HeaderValue::from_static("default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; img-src 'self' data:;"), + )) + .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( + axum::http::header::X_CONTENT_TYPE_OPTIONS, + HeaderValue::from_static("nosniff"), + )) + .layer(tower_http::set_header::SetResponseHeaderLayer::overriding( + axum::http::header::REFERRER_POLICY, + HeaderValue::from_static("strict-origin-when-cross-origin"), + )) + .layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit + .with_state(state); + + let addr = format!("0.0.0.0:{}", config.port); + let listener = tokio::net::TcpListener::bind(&addr).await?; + println!("Server running on http://localhost:{}", config.port); + axum::serve(listener, app).await?; + + Ok(()) +} + +fn build_cors_layer(config: &crate::config::Config) -> CorsLayer { + let allow_origin = if let Some(origins) = &config.cors_allowed_origins { + let values: Vec = origins + .split(',') + .map(|origin| origin.trim()) + .filter(|origin| !origin.is_empty()) + .filter_map(|origin| HeaderValue::from_str(origin).ok()) + .collect(); + + if values.is_empty() { + AllowOrigin::mirror_request() + } else { + AllowOrigin::list(values) + } + } else { + AllowOrigin::mirror_request() + }; + + CorsLayer::new() + .allow_origin(allow_origin) + .allow_methods([ + axum::http::Method::GET, + axum::http::Method::POST, + axum::http::Method::PUT, + axum::http::Method::PATCH, + axum::http::Method::DELETE, + axum::http::Method::OPTIONS, + ]) + .allow_headers([ + axum::http::header::CONTENT_TYPE, + axum::http::header::AUTHORIZATION, + axum::http::header::ACCEPT, + ]) + .allow_credentials(true) +} diff --git a/src/server/tasks.rs b/src/server/tasks.rs new file mode 100644 index 0000000..f9ff855 --- /dev/null +++ b/src/server/tasks.rs @@ -0,0 +1,179 @@ +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; +use sea_orm::{EntityTrait, QueryOrder, QuerySelect}; +use std::sync::Arc; +use uuid::Uuid; + +use super::AppState; +use super::auth::AuthenticatedUser; +use crate::domain::tasks::{ + self, CreateTaskRequest, RecentRunResponse, TaskResponse, UpdateTaskRequest, +}; + +pub async fn list_tasks( + _user: AuthenticatedUser, + State(state): State>, +) -> Result>, (StatusCode, String)> { + let tasks = crate::entities::task::Entity::find() + .find_with_related(crate::entities::task_run::Entity) + .order_by_desc(crate::entities::task::Column::CreatedAt) + .all(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let response = tasks + .into_iter() + .map(|(t, mut runs)| { + runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); + TaskResponse { + id: t.id, + goal: t.goal, + cron: t.cron, + created_at: t.created_at, + runs: runs + .into_iter() + .map(|r| tasks::TaskRunResponse { + id: r.id, + status: r.status, + logs: r.logs, + answer: r.answer, + created_at: r.created_at, + }) + .collect(), + } + }) + .collect(); + + Ok(Json(response)) +} + +pub async fn create_task( + _user: AuthenticatedUser, + State(state): State>, + Json(payload): Json, +) -> Result, (StatusCode, String)> { + let task_id = Uuid::new_v4(); + + let new_task = crate::entities::task::ActiveModel { + id: sea_orm::Set(task_id), + goal: sea_orm::Set(payload.goal.clone()), + cron: sea_orm::Set(payload.cron.clone()), + created_at: sea_orm::Set(chrono::Utc::now().into()), + }; + + use sea_orm::ActiveModelTrait; + new_task + .insert(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if let Some(cron) = &payload.cron { + let _ = state.scheduler.add_task_job(task_id, cron).await; + } else { + let _ = state.scheduler.remove_task_job(task_id).await; + } + + tasks::get_task_inner(task_id, &state.db) + .await + .map(Json) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) +} + +pub async fn rerun_task( + _user: AuthenticatedUser, + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let task = crate::entities::task::Entity::find_by_id(id) + .one(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?; + + tasks::execute_agent_run( + &state.db, + &state.scheduler, + &state.config, + task.id, + task.goal, + ) + .await + .map(Json) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) +} + +pub async fn update_task( + _user: AuthenticatedUser, + State(state): State>, + Path(id): Path, + Json(payload): Json, +) -> Result, (StatusCode, String)> { + let mut task: crate::entities::task::ActiveModel = + crate::entities::task::Entity::find_by_id(id) + .one(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))? + .into(); + + task.goal = sea_orm::Set(payload.goal.clone()); + task.cron = sea_orm::Set(payload.cron.clone()); + + use sea_orm::ActiveModelTrait; + task.update(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if let Some(cron) = &payload.cron { + let _ = state.scheduler.add_task_job(id, cron).await; + } else { + let _ = state.scheduler.remove_task_job(id).await; + } + + tasks::get_task_inner(id, &state.db) + .await + .map(Json) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) +} + +pub async fn get_task( + _user: AuthenticatedUser, + Path(id): Path, + State(state): State>, +) -> Result, (StatusCode, String)> { + tasks::get_task_inner(id, &state.db) + .await + .map(Json) + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) +} + +pub async fn get_recent_runs( + _user: AuthenticatedUser, + State(state): State>, +) -> Result>, (StatusCode, String)> { + let results = crate::entities::task_run::Entity::find() + .find_also_related(crate::entities::task::Entity) + .order_by_desc(crate::entities::task_run::Column::CreatedAt) + .limit(50) + .all(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let response = results + .into_iter() + .filter_map(|(run, task_opt)| { + task_opt.map(|task| RecentRunResponse { + id: run.id, + task_id: run.task_id, + goal: task.goal, + status: run.status, + created_at: run.created_at, + }) + }) + .collect(); + + Ok(Json(response)) +}