From 7268d49b4a6c11d2349bb538d78feeaf74bc8f73 Mon Sep 17 00:00:00 2001 From: pavel Date: Wed, 11 Feb 2026 01:41:07 +0100 Subject: [PATCH] more refactoring --- Cargo.lock | 2 + Cargo.toml | 2 + frontend/src/main.js | 64 +++++++++++---------- src/config.rs | 18 +++--- src/domain/agent/mod.rs | 116 +++++++++++++++++++++++--------------- src/domain/agent/tools.rs | 24 +++++--- src/domain/tasks.rs | 34 +++++++---- src/error.rs | 53 +++++++++++++++++ src/main.rs | 1 + src/server/auth.rs | 51 +++++++---------- src/server/mod.rs | 99 ++++++++++++++++++++++---------- src/server/tasks.rs | 56 ++++++++---------- 12 files changed, 325 insertions(+), 195 deletions(-) create mode 100644 src/error.rs diff --git a/Cargo.lock b/Cargo.lock index 44e16d9..6be1feb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -328,6 +328,7 @@ dependencies = [ "chrono", "cookie", "dashmap", + "dotenvy", "jsonwebtoken", "migration", "reqwest", @@ -335,6 +336,7 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", + "thiserror", "tokio", "tokio-cron-scheduler", "tower-http 0.5.2", diff --git a/Cargo.toml b/Cargo.toml index d0b8b6f..1e6e013 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,3 +24,5 @@ jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] } base64 = "0.22.1" axum-extra = { version = "0.9", features = ["typed-header", "cookie"] } cookie = "0.18" +thiserror = "2.0.18" +dotenvy = "0.15.7" diff --git a/frontend/src/main.js b/frontend/src/main.js index 91e3375..dd7c913 100644 --- a/frontend/src/main.js +++ b/frontend/src/main.js @@ -18,7 +18,6 @@ const state = { isEditing: false, isAuthenticated: false }; - // DOM elements const loginOverlay = document.getElementById('login-overlay'); const callbackOverlay = document.getElementById('callback-overlay'); @@ -55,6 +54,25 @@ const toggleCustomCronBtn = document.getElementById('toggle-custom-cron'); const customCronContainer = document.getElementById('custom-cron-container'); const presetBtns = document.querySelectorAll('.btn-preset'); +function updateState(newState) { + Object.assign(state, newState); + renderApp(); +} + +function renderApp() { + renderTaskList(); + + if (state.currentView === 'dashboard') { + fetchRecentRuns(); + } else if (state.selectedTaskId) { + const task = state.tasks.find(t => t.id === state.selectedTaskId); + if (task) { + renderRunHistory(task); + showTaskView(task); + } + } +} + // Wrapper for fetch to include Authorization header async function fetchWithAuth(url, options = {}) { let response = await fetch(url, { ...options, credentials: 'include' }); @@ -97,46 +115,32 @@ async function attemptTokenRefresh() { } return false; } - +// I'll replace the fetchTasks function and add updateState async function fetchTasks() { try { const response = await fetchWithAuth(`${API_URL}/tasks`); const newTasks = await response.json(); - // Check if we should follow the latest run (if we were already watching it) - let shouldFollowLatest = false; + // Check if we should follow the latest run + let newSelectedRunId = state.selectedRunId; if (state.selectedTaskId) { - const currentTask = state.tasks.find(t => t.id === state.selectedTaskId); + const currentTask = newTasks.find(t => t.id === state.selectedTaskId); if (currentTask && currentTask.runs && currentTask.runs.length > 0) { - const latestRunId = currentTask.runs[currentTask.runs.length - 1].id; - if (state.selectedRunId === latestRunId) { - shouldFollowLatest = true; + // If we don't have a selected run or the runs changed, we might want to update + if (!state.selectedRunId || (state.tasks.find(t => t.id === state.selectedTaskId)?.runs?.length !== currentTask.runs.length)) { + // Only auto-switch if we are "following" the latest + const wasFollowingLatest = state.tasks.find(t => t.id === state.selectedTaskId)?.runs?.[0]?.id === state.selectedRunId; + if (wasFollowingLatest || !state.selectedRunId) { + newSelectedRunId = currentTask.runs[0].id; + } } - } else if (!state.selectedRunId) { - shouldFollowLatest = true; } } - state.tasks = newTasks; - renderTaskList(); - - // If we are on the dashboard, refresh it too - if (state.currentView === 'dashboard') { - fetchRecentRuns(); - } - - // If a task is selected, update it - if (state.selectedTaskId) { - const task = state.tasks.find((t) => t.id === state.selectedTaskId); - if (task) { - if (shouldFollowLatest && task.runs && task.runs.length > 0) { - state.selectedRunId = task.runs[0].id; - } - - renderRunHistory(task); - showTaskView(task); - } - } + updateState({ + tasks: newTasks, + selectedRunId: newSelectedRunId + }); } catch (error) { console.error('Error fetching tasks:', error); } diff --git a/src/config.rs b/src/config.rs index 8da6394..c7a714e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,3 +1,4 @@ +use crate::error::{AppError, AppResult}; use std::env; #[derive(Clone, Debug)] @@ -16,8 +17,11 @@ pub struct Config { } impl Config { - pub fn from_env() -> Result> { - let database_url = env::var("DATABASE_URL").map_err(|_| "DATABASE_URL must be set")?; + pub fn from_env() -> AppResult { + dotenvy::dotenv().ok(); + + let database_url = env::var("DATABASE_URL") + .map_err(|_| AppError::Config("DATABASE_URL must be set".into()))?; let port = env::var("PORT") .ok() @@ -27,12 +31,12 @@ impl Config { let zen_api_key = env::var("ZEN_API_KEY").ok(); let tavily_api_key = env::var("TAVILY_API_KEY").ok(); - let authentik_issuer = - env::var("AUTHENTIK_ISSUER").map_err(|_| "AUTHENTIK_ISSUER must be set")?; - let authentik_client_id = - env::var("AUTHENTIK_CLIENT_ID").map_err(|_| "AUTHENTIK_CLIENT_ID must be set")?; + let authentik_issuer = env::var("AUTHENTIK_ISSUER") + .map_err(|_| AppError::Config("AUTHENTIK_ISSUER must be set".into()))?; + let authentik_client_id = env::var("AUTHENTIK_CLIENT_ID") + .map_err(|_| AppError::Config("AUTHENTIK_CLIENT_ID must be set".into()))?; let authentik_client_secret = env::var("AUTHENTIK_CLIENT_SECRET") - .map_err(|_| "AUTHENTIK_CLIENT_SECRET must be set")?; + .map_err(|_| AppError::Config("AUTHENTIK_CLIENT_SECRET must be set".into()))?; let cors_allowed_origins = env::var("CORS_ALLOWED_ORIGINS").ok(); diff --git a/src/domain/agent/mod.rs b/src/domain/agent/mod.rs index 91f5147..9f6933f 100644 --- a/src/domain/agent/mod.rs +++ b/src/domain/agent/mod.rs @@ -17,12 +17,14 @@ pub struct Agent { answer: Option, } +use crate::error::{AppError, AppResult}; + impl Agent { pub fn new( zen_api_key: Option, tavily_api_key: Option, initial_message: String, - ) -> Result> { + ) -> AppResult { let intro = format!( "You are an autonomous agent. You have access to tools that can help you achieve your goals. Use them wisely. The user is unable to respond to you @@ -49,7 +51,8 @@ impl Agent { let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(60)) - .build()?; + .build() + .map_err(|e| AppError::Internal(format!("Failed to build HTTP client: {}", e)))?; Ok(Self { client, @@ -71,7 +74,7 @@ impl Agent { pub async fn run( &mut self, config: &crate::config::Config, - ) -> Result<(String, Option), Box> { + ) -> AppResult<(String, Option)> { let mut finished = false; let start_time = Instant::now(); let max_duration = Duration::from_secs(config.agent_max_duration_secs); @@ -80,47 +83,29 @@ impl Agent { while !finished { if start_time.elapsed() > max_duration { - return Err("Agent run timed out".into()); + return Err(AppError::Internal("Agent run timed out".into())); } if turns >= max_turns { - return Err("Agent run exceeded max turns".into()); + return Err(AppError::Internal("Agent run exceeded max turns".into())); } turns += 1; - let request = ChatRequest { - model: "kimi-k2.5".to_string(), - messages: self.messages.clone(), - tools: self.tools.clone(), - }; - + let current_role = self + .messages + .last() + .map(|m| m.role.as_str()) + .unwrap_or("unknown"); self.log(&format!( - "--- Sending request to Zen API (Role: {}) ---", - self.messages.last().unwrap().role + "\n[Turn {}] Sending request (Last role: {})", + turns, current_role )); - let mut request_builder = self.client.post(&self.url).json(&request); - - if let Some(key) = &self.zen_api_key { - request_builder = - request_builder.header("Authorization", format!("Bearer {}", key)); - } - - let response = request_builder.send().await?; - - if !response.status().is_success() { - let status = response.status(); - let error_text = response.text().await?; - self.log(&format!("Error: API request failed with status {}", status)); - self.log(&format!("Error details: {}", error_text)); - return Err(format!("API request failed: {}", status).into()); - } - - let chat_response: ChatResponse = response.json().await?; + let chat_response = self.call_llm().await?; let assistant_message = chat_response .choices .get(0) - .ok_or("Missing assistant response")? + .ok_or_else(|| AppError::Internal("Missing assistant response".into()))? .message .clone(); @@ -128,41 +113,78 @@ impl Agent { if let Some(content) = &assistant_message.content { if !content.is_empty() { - self.log(&format!("\nAssistant response:\n{}\n", content)); + self.log(&format!("\nAssistant: {}", content)); } } if let Some(tool_calls) = assistant_message.tool_calls { for tool_call in tool_calls { - let (tool_message, written, tool_answer) = - tools::handle_tool_call(&tool_call, &self.tavily_api_key).await?; + self.log(&format!("Calling tool: {}", tool_call.function.name)); + + let (tool_message, is_final, tool_answer) = + tools::handle_tool_call(&tool_call, &self.tavily_api_key) + .await + .map_err(|e| { + AppError::Internal(format!("Tool execution failed: {}", e)) + })?; if let Some(ans) = tool_answer { self.answer = Some(ans); + self.log("Task marked as finished by tool."); } if let Some(content) = &tool_message.content { - self.log(&format!( - "Tool result ({}): {}", - tool_call.function.name, content - )); + self.log(&format!("Tool result: {}", content)); } self.messages.push(tool_message); - if written { + if is_final { finished = true; } } - // Continue the loop to send tool results back - continue; - } - - // No tool calls from assistant, but we only exit if the task was finished - if !finished { - self.log("--- Assistant didn't finish yet. Waiting for next turn... ---"); + } else if assistant_message.content.is_some() { + // If assistant just talked without tools, we might be stuck or finished. + // But typically we expect a 'finish' tool call. + self.log("Assistant responded without tool calls."); + // For now we continue unless the assistant explicitly uses a tool to finish, + // or we could add heuristic here if needed. } } + self.log("\n--- Execution Finished ---"); Ok((self.logs.clone(), self.answer.clone())) } + + async fn call_llm(&self) -> AppResult { + let request = ChatRequest { + model: "kimi-k2.5".to_string(), + messages: self.messages.clone(), + tools: self.tools.clone(), + }; + + let mut request_builder = self.client.post(&self.url).json(&request); + + if let Some(key) = &self.zen_api_key { + request_builder = request_builder.header("Authorization", format!("Bearer {}", key)); + } + + let response = request_builder.send().await.map_err(AppError::Network)?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response + .text() + .await + .unwrap_or_else(|_| "Unknown error".into()); + return Err(AppError::Internal(format!( + "API request failed: {} - {}", + status, error_text + ))); + } + + response + .json() + .await + .map_err(|e| AppError::Internal(format!("Failed to parse LLM response: {}", e))) + } } diff --git a/src/domain/agent/tools.rs b/src/domain/agent/tools.rs index eb9b2ab..663b1b1 100644 --- a/src/domain/agent/tools.rs +++ b/src/domain/agent/tools.rs @@ -1,5 +1,15 @@ use super::api::{self, FunctionDefinition, Message, Tool, ToolCall}; -use std::collections::HashMap; +use serde::Deserialize; + +#[derive(Deserialize)] +struct GoogleSearchArgs { + query: String, +} + +#[derive(Deserialize)] +struct FinishArgs { + result: String, +} pub fn get_tools() -> Vec { vec![ @@ -24,13 +34,13 @@ pub fn get_tools() -> Vec { tool_type: "function".to_string(), function: FunctionDefinition { name: "finish".to_string(), - description: "Finish the task".to_string(), + description: "Finish the task and provide a final answer".to_string(), parameters: serde_json::json!({ "type": "object", "properties": { "result": { "type": "string", - "description": "The result of the task" + "description": "The final detailed answer to the task" } }, "required": ["result"] @@ -48,8 +58,8 @@ pub async fn handle_tool_call( let name = &tool_call.function.name; let (content, written) = if name == "google_search" { - let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; - let query = args.get("query").ok_or("Missing query argument")?; + let args: GoogleSearchArgs = serde_json::from_str(&tool_call.function.arguments)?; + let query = &args.query; let search_result = if let Some(key) = tavily_api_key { match api::perform_search(query, key).await { @@ -61,8 +71,8 @@ pub async fn handle_tool_call( }; (search_result, false) } else if name == "finish" { - let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; - let result = args.get("result").ok_or("Missing result argument")?; + let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?; + let result = &args.result; answer = Some(result.clone()); (result.clone(), true) diff --git a/src/domain/tasks.rs b/src/domain/tasks.rs index 9e64d8e..0968cdf 100644 --- a/src/domain/tasks.rs +++ b/src/domain/tasks.rs @@ -49,13 +49,15 @@ pub struct RecentRunResponse { pub created_at: chrono::DateTime, } +use crate::error::AppResult; + pub async fn execute_agent_run( db: &DatabaseConnection, _scheduler: &Arc, config: &Arc, task_id: Uuid, goal: String, -) -> Result> { +) -> AppResult { let run_id = Uuid::new_v4(); let new_run = task_run::ActiveModel { @@ -67,7 +69,10 @@ pub async fn execute_agent_run( created_at: Set(Utc::now().into()), }; - new_run.insert(db).await?; + new_run + .insert(db) + .await + .map_err(crate::error::AppError::Database)?; let mut agent = Agent::new( config.zen_api_key.clone(), @@ -84,31 +89,36 @@ pub async fn execute_agent_run( ), }; - let mut run: task_run::ActiveModel = TaskRun::find_by_id(run_id) + let run: task_run::ActiveModel = TaskRun::find_by_id(run_id) .one(db) - .await? - .ok_or("Run not found after insert")? + .await + .map_err(crate::error::AppError::Database)? + .ok_or_else(|| crate::error::AppError::NotFound("Run not found after insert".into()))? .into(); + let mut run = run; run.logs = Set(logs.clone()); run.answer = Set(answer.clone()); run.status = Set(status); - run.update(db).await?; + run.update(db) + .await + .map_err(crate::error::AppError::Database)?; get_task_inner(task_id, db).await } -pub async fn get_task_inner( - id: Uuid, - db: &DatabaseConnection, -) -> Result> { +pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult { let results = Task::find_by_id(id) .find_with_related(TaskRun) .all(db) - .await?; + .await + .map_err(crate::error::AppError::Database)?; - let (t, mut runs) = results.into_iter().next().ok_or("Task not found")?; + let (t, mut runs) = results + .into_iter() + .next() + .ok_or_else(|| crate::error::AppError::NotFound("Task not found".into()))?; runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..01943d5 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,53 @@ +use axum::{ + Json, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use serde_json::json; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum AppError { + #[error("Database error: {0}")] + Database(#[from] sea_orm::DbErr), + + #[error("Configuration error: {0}")] + Config(String), + + #[error("Not found: {0}")] + NotFound(String), + + #[error("Unauthorized: {0}")] + Unauthorized(String), + + #[error("Internal server error: {0}")] + Internal(String), + + #[error("Network error: {0}")] + Network(#[from] reqwest::Error), + + #[error("Invalid request: {0}")] + InvalidRequest(String), +} + +impl IntoResponse for AppError { + fn into_response(self) -> Response { + let (status, error_message) = match self { + AppError::Database(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()), + AppError::Config(err) => (StatusCode::INTERNAL_SERVER_ERROR, err), + AppError::NotFound(err) => (StatusCode::NOT_FOUND, err), + AppError::Unauthorized(err) => (StatusCode::UNAUTHORIZED, err), + AppError::Internal(err) => (StatusCode::INTERNAL_SERVER_ERROR, err), + AppError::Network(err) => (StatusCode::BAD_GATEWAY, err.to_string()), + AppError::InvalidRequest(err) => (StatusCode::BAD_REQUEST, err), + }; + + let body = Json(json!({ + "error": error_message, + })); + + (status, body).into_response() + } +} + +pub type AppResult = Result; diff --git a/src/main.rs b/src/main.rs index 36c5a8d..199a393 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ mod config; mod domain; mod entities; +mod error; mod scheduler; mod server; diff --git a/src/server/auth.rs b/src/server/auth.rs index edc7465..7ff1086 100644 --- a/src/server/auth.rs +++ b/src/server/auth.rs @@ -1,7 +1,7 @@ use axum::{ Json, RequestPartsExt, extract::{FromRef, FromRequestParts, Query, State}, - http::{StatusCode, request::Parts}, + http::request::Parts, response::IntoResponse, }; use axum_extra::{ @@ -12,10 +12,9 @@ use axum_extra::{ use std::sync::Arc; use super::AppState; -use crate::domain::auth::Claims; +use crate::error::{AppError, AppResult}; -#[allow(dead_code)] -pub struct AuthenticatedUser(pub Claims); +pub struct AuthenticatedUser(pub crate::domain::auth::Claims); #[axum::async_trait] impl FromRequestParts for AuthenticatedUser @@ -23,7 +22,7 @@ where Arc: axum::extract::FromRef, S: Send + Sync, { - type Rejection = (StatusCode, String); + type Rejection = AppError; async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let app_state = Arc::::from_ref(state); @@ -33,22 +32,22 @@ where { Some(bearer.token().to_string()) } else { - let jar = parts.extract::().await.unwrap(); + let jar = parts + .extract::() + .await + .map_err(|e| AppError::Internal(e.to_string()))?; 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 token = token + .ok_or_else(|| AppError::Unauthorized("Missing or invalid access token".into()))?; - let claims = app_state.verifier.verify(&token).await.map_err(|e| { - ( - StatusCode::UNAUTHORIZED, - format!("Token verification failed: {}", e), - ) - })?; + let claims = app_state + .verifier + .verify(&token) + .await + .map_err(|e| AppError::Unauthorized(format!("Token verification failed: {}", e)))?; Ok(AuthenticatedUser(claims)) } @@ -69,7 +68,7 @@ pub async fn auth_refresh( State(state): State>, jar: CookieJar, Json(payload): Json, -) -> Result { +) -> AppResult { let refresh_token = payload .refresh_token .filter(|token| !token.is_empty()) @@ -77,16 +76,13 @@ pub async fn auth_refresh( jar.get("refresh_token") .map(|cookie| cookie.value().to_string()) }) - .ok_or(( - StatusCode::UNAUTHORIZED, - "Missing refresh token".to_string(), - ))?; + .ok_or_else(|| AppError::Unauthorized("Missing refresh token".into()))?; let data = state .authenticator .refresh_token(refresh_token) .await - .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?; + .map_err(|e| AppError::Unauthorized(e.to_string()))?; let jar = update_auth_cookies(jar, &data, &state.config); Ok((jar, Json(data))) @@ -96,17 +92,12 @@ pub async fn auth_callback( State(state): State>, jar: CookieJar, Query(query): Query, -) -> Result { +) -> AppResult { let data = state .authenticator .exchange_code(query.code, query.redirect_uri) .await - .map_err(|e| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Token exchange failed: {}", e), - ) - })?; + .map_err(|e| AppError::Internal(format!("Token exchange failed: {}", e)))?; let jar = update_auth_cookies(jar, &data, &state.config); Ok((jar, Json(data))) @@ -114,7 +105,7 @@ pub async fn auth_callback( pub async fn auth_logout(State(state): State>, jar: CookieJar) -> impl IntoResponse { let jar = clear_auth_cookies(jar, &state.config); - (jar, StatusCode::NO_CONTENT) + (jar, axum::http::StatusCode::NO_CONTENT) } pub async fn auth_session(user: AuthenticatedUser) -> Json { diff --git a/src/server/mod.rs b/src/server/mod.rs index 8f88b7d..2fe3485 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -14,6 +14,8 @@ use tower_http::cors::{AllowOrigin, CorsLayer}; use crate::entities::task::Entity as Task; use crate::scheduler::Scheduler; +use crate::error::AppResult; + #[derive(Clone)] pub struct AppState { pub db: DatabaseConnection, @@ -23,37 +25,29 @@ pub struct AppState { 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?; +pub async fn start(config: crate::config::Config) -> AppResult<()> { + let db = setup_database(&config.database_url).await?; let config = Arc::new(config); - let scheduler = Arc::new(Scheduler::new(db.clone(), config.clone()).await?); + let scheduler = Arc::new( + Scheduler::new(db.clone(), config.clone()) + .await + .map_err(|e| crate::error::AppError::Internal(e.to_string()))?, + ); // Load existing scheduled tasks - let existing_tasks = Task::find().all(&db).await?; + let existing_tasks = Task::find() + .all(&db) + .await + .map_err(crate::error::AppError::Database)?; 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 (verifier, authenticator) = setup_auth(&config).await?; let state = Arc::new(AppState { db, @@ -63,9 +57,61 @@ pub async fn start(config: crate::config::Config) -> Result<(), Box AppResult { + let db = Database::connect(database_url) + .await + .map_err(crate::error::AppError::Database)?; + Migrator::up(&db, None) + .await + .map_err(crate::error::AppError::Database)?; + Ok(db) +} + +async fn setup_auth( + config: &crate::config::Config, +) -> AppResult<( + Arc, + Arc, +)> { + let verifier = Arc::new( + crate::domain::auth::JwksVerifier::new( + config.authentik_issuer.clone(), + config.authentik_client_id.clone(), + ) + .await + .map_err(|e| crate::error::AppError::Internal(e.to_string()))?, + ); + let authenticator = Arc::new( + crate::domain::auth::Authenticator::new( + config.authentik_issuer.clone(), + config.authentik_client_id.clone(), + config.authentik_client_secret.clone(), + ) + .await + .map_err(|e| crate::error::AppError::Internal(e.to_string()))?, + ); + + Ok((verifier, authenticator)) +} + +fn build_app(state: Arc, config: &crate::config::Config) -> Router { + let cors = build_cors_layer(config); + + 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)) @@ -88,14 +134,7 @@ pub async fn start(config: crate::config::Config) -> Result<(), Box CorsLayer { diff --git a/src/server/tasks.rs b/src/server/tasks.rs index f9ff855..138c474 100644 --- a/src/server/tasks.rs +++ b/src/server/tasks.rs @@ -1,7 +1,6 @@ use axum::{ Json, extract::{Path, State}, - http::StatusCode, }; use sea_orm::{EntityTrait, QueryOrder, QuerySelect}; use std::sync::Arc; @@ -13,16 +12,17 @@ use crate::domain::tasks::{ self, CreateTaskRequest, RecentRunResponse, TaskResponse, UpdateTaskRequest, }; +use crate::error::AppResult; + pub async fn list_tasks( _user: AuthenticatedUser, State(state): State>, -) -> Result>, (StatusCode, String)> { +) -> AppResult>> { 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()))?; + .await?; let response = tasks .into_iter() @@ -54,7 +54,7 @@ pub async fn create_task( _user: AuthenticatedUser, State(state): State>, Json(payload): Json, -) -> Result, (StatusCode, String)> { +) -> AppResult> { let task_id = Uuid::new_v4(); let new_task = crate::entities::task::ActiveModel { @@ -65,10 +65,7 @@ pub async fn create_task( }; use sea_orm::ActiveModelTrait; - new_task - .insert(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + new_task.insert(&state.db).await?; if let Some(cron) = &payload.cron { let _ = state.scheduler.add_task_job(task_id, cron).await; @@ -79,19 +76,18 @@ pub async fn create_task( tasks::get_task_inner(task_id, &state.db) .await .map(Json) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) } pub async fn rerun_task( _user: AuthenticatedUser, State(state): State>, Path(id): Path, -) -> Result, (StatusCode, String)> { +) -> AppResult> { 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()))?; + .await? + .ok_or_else(|| crate::error::AppError::NotFound("Task not found".to_string()))?; tasks::execute_agent_run( &state.db, @@ -102,7 +98,7 @@ pub async fn rerun_task( ) .await .map(Json) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) } pub async fn update_task( @@ -110,22 +106,19 @@ pub async fn update_task( 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(); +) -> AppResult> { + let task: crate::entities::task::ActiveModel = crate::entities::task::Entity::find_by_id(id) + .one(&state.db) + .await? + .ok_or_else(|| crate::error::AppError::NotFound("Task not found".to_string()))? + .into(); + let mut task = task; 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()))?; + task.update(&state.db).await?; if let Some(cron) = &payload.cron { let _ = state.scheduler.add_task_job(id, cron).await; @@ -136,31 +129,30 @@ pub async fn update_task( tasks::get_task_inner(id, &state.db) .await .map(Json) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) } pub async fn get_task( _user: AuthenticatedUser, Path(id): Path, State(state): State>, -) -> Result, (StatusCode, String)> { +) -> AppResult> { tasks::get_task_inner(id, &state.db) .await .map(Json) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) } pub async fn get_recent_runs( _user: AuthenticatedUser, State(state): State>, -) -> Result>, (StatusCode, String)> { +) -> AppResult>> { 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()))?; + .await?; let response = results .into_iter()