more refactoring

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",

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"

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)
+        // Check if we should follow the latest run
-        let shouldFollowLatest = false;
+        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 we don't have a selected run or the runs changed, we might want to update
-                if (state.selectedRunId === latestRunId) {
+                if (!state.selectedRunId || (state.tasks.find(t => t.id === state.selectedTaskId)?.runs?.length !== currentTask.runs.length)) {
-                    shouldFollowLatest = true;
+                    // 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;
+        updateState({
-        renderTaskList();
+            tasks: newTasks,
-
+            selectedRunId: newSelectedRunId
-        // 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);
-            }
-        }
     } catch (error) {
         console.error('Error fetching tasks:', error);
     }

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<Self, Box<dyn std::error::Error>> {
+    pub fn from_env() -> AppResult<Self> {
-        let database_url = env::var("DATABASE_URL").map_err(|_| "DATABASE_URL must be set")?;
+        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 =
+        let authentik_issuer = env::var("AUTHENTIK_ISSUER")
-            env::var("AUTHENTIK_ISSUER").map_err(|_| "AUTHENTIK_ISSUER must be set")?;
+            .map_err(|_| AppError::Config("AUTHENTIK_ISSUER must be set".into()))?;
-        let authentik_client_id =
+        let authentik_client_id = env::var("AUTHENTIK_CLIENT_ID")
-            env::var("AUTHENTIK_CLIENT_ID").map_err(|_| "AUTHENTIK_CLIENT_ID must be set")?;
+            .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();
 

src/domain/agent/mod.rs

@@ -17,12 +17,14 @@ pub struct Agent {
     answer: Option<String>,
 }
 
+use crate::error::{AppError, AppResult};
+
 impl Agent {
     pub fn new(
         zen_api_key: Option<String>,
         tavily_api_key: Option<String>,
         initial_message: String,
-    ) -> Result<Self, Box<dyn std::error::Error>> {
+    ) -> AppResult<Self> {
         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<String>), Box<dyn std::error::Error>> {
+    ) -> AppResult<(String, Option<String>)> {
         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 {
+            let current_role = self
-                model: "kimi-k2.5".to_string(),
+                .messages
-                messages: self.messages.clone(),
+                .last()
-                tools: self.tools.clone(),
+                .map(|m| m.role.as_str())
-            };
+                .unwrap_or("unknown");
-
             self.log(&format!(
-                "--- Sending request to Zen API (Role: {}) ---",
+                "\n[Turn {}] Sending request (Last role: {})",
-                self.messages.last().unwrap().role
+                turns, current_role
             ));
 
-            let mut request_builder = self.client.post(&self.url).json(&request);
+            let chat_response = self.call_llm().await?;
-
-            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 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) =
+                    self.log(&format!("Calling tool: {}", tool_call.function.name));
-                        tools::handle_tool_call(&tool_call, &self.tavily_api_key).await?;
+
+                    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!(
+                        self.log(&format!("Tool result: {}", content));
-                            "Tool result ({}): {}",
-                            tool_call.function.name, content
-                        ));
                     }
 
                     self.messages.push(tool_message);
-                    if written {
+                    if is_final {
                         finished = true;
                     }
                 }
-                // Continue the loop to send tool results back
+            } else if assistant_message.content.is_some() {
-                continue;
+                // 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.");
-            // No tool calls from assistant, but we only exit if the task was finished
+                // For now we continue unless the assistant explicitly uses a tool to finish,
-            if !finished {
+                // or we could add heuristic here if needed.
-                self.log("--- Assistant didn't finish yet. Waiting for next turn... ---");
             }
         }
 
+        self.log("\n--- Execution Finished ---");
         Ok((self.logs.clone(), self.answer.clone()))
     }
+
+    async fn call_llm(&self) -> AppResult<ChatResponse> {
+        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)))
+    }
 }

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<Tool> {
     vec![
@@ -24,13 +34,13 @@ pub fn get_tools() -> Vec<Tool> {
             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<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let args: GoogleSearchArgs = serde_json::from_str(&tool_call.function.arguments)?;
-        let query = args.get("query").ok_or("Missing query argument")?;
+        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<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?;
-        let result = args.get("result").ok_or("Missing result argument")?;
+        let result = &args.result;
 
         answer = Some(result.clone());
         (result.clone(), true)

src/domain/tasks.rs

@@ -49,13 +49,15 @@ pub struct RecentRunResponse {
     pub created_at: chrono::DateTime<chrono::FixedOffset>,
 }
 
+use crate::error::AppResult;
+
 pub async fn execute_agent_run(
     db: &DatabaseConnection,
     _scheduler: &Arc<Scheduler>,
     config: &Arc<Config>,
     task_id: Uuid,
     goal: String,
-) -> Result<TaskResponse, Box<dyn std::error::Error>> {
+) -> AppResult<TaskResponse> {
     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?
+        .await
-        .ok_or("Run not found after insert")?
+        .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(
+pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult<TaskResponse> {
-    id: Uuid,
-    db: &DatabaseConnection,
-) -> Result<TaskResponse, Box<dyn std::error::Error>> {
     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));
 

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<T> = Result<T, AppError>;

src/main.rs

@@ -1,6 +1,7 @@
 mod config;
 mod domain;
 mod entities;
+mod error;
 mod scheduler;
 mod server;
 

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 crate::domain::auth::Claims);
-pub struct AuthenticatedUser(pub Claims);
 
 #[axum::async_trait]
 impl<S> FromRequestParts<S> for AuthenticatedUser
@@ -23,7 +22,7 @@ where
     Arc<AppState>: axum::extract::FromRef<S>,
     S: Send + Sync,
 {
-    type Rejection = (StatusCode, String);
+    type Rejection = AppError;
 
     async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
         let app_state = Arc::<AppState>::from_ref(state);
@@ -33,22 +32,22 @@ where
         {
             Some(bearer.token().to_string())
         } else {
-            let jar = parts.extract::<CookieJar>().await.unwrap();
+            let jar = parts
+                .extract::<CookieJar>()
+                .await
+                .map_err(|e| AppError::Internal(e.to_string()))?;
             jar.get("access_token")
                 .map(|cookie| cookie.value().to_string())
         };
 
-        let token = token.ok_or((
+        let token = token
-            StatusCode::UNAUTHORIZED,
+            .ok_or_else(|| AppError::Unauthorized("Missing or invalid access token".into()))?;
-            "Missing or invalid access token".to_string(),
-        ))?;
 
-        let claims = app_state.verifier.verify(&token).await.map_err(|e| {
+        let claims = app_state
-            (
+            .verifier
-                StatusCode::UNAUTHORIZED,
+            .verify(&token)
-                format!("Token verification failed: {}", e),
+            .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<Arc<AppState>>,
     jar: CookieJar,
     Json(payload): Json<RefreshRequest>,
-) -> Result<impl IntoResponse, (StatusCode, String)> {
+) -> AppResult<impl IntoResponse> {
     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((
+        .ok_or_else(|| AppError::Unauthorized("Missing refresh token".into()))?;
-            StatusCode::UNAUTHORIZED,
-            "Missing refresh token".to_string(),
-        ))?;
 
     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<Arc<AppState>>,
     jar: CookieJar,
     Query(query): Query<AuthCallbackQuery>,
-) -> Result<impl IntoResponse, (StatusCode, String)> {
+) -> AppResult<impl IntoResponse> {
     let data = state
         .authenticator
         .exchange_code(query.code, query.redirect_uri)
         .await
-        .map_err(|e| {
+        .map_err(|e| AppError::Internal(format!("Token exchange failed: {}", e)))?;
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                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<Arc<AppState>>, 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<serde_json::Value> {

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<crate::domain::auth::Authenticator>,
 }
 
-pub async fn start(config: crate::config::Config) -> Result<(), Box<dyn std::error::Error>> {
+pub async fn start(config: crate::config::Config) -> AppResult<()> {
-    let db = Database::connect(&config.database_url).await?;
+    let db = setup_database(&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?);
+    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(
+    let (verifier, authenticator) = setup_auth(&config).await?;
-        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,
@@ -63,9 +57,61 @@ pub async fn start(config: crate::config::Config) -> Result<(), Box<dyn std::err
         authenticator,
     });
 
-    let cors = build_cors_layer(&config);
+    let app = build_app(state, &config);
 
-    let app = Router::new()
+    let addr = format!("0.0.0.0:{}", config.port);
+    let listener = tokio::net::TcpListener::bind(&addr)
+        .await
+        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
+    println!("Server running on http://localhost:{}", config.port);
+    axum::serve(listener, app)
+        .await
+        .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
+
+    Ok(())
+}
+
+async fn setup_database(database_url: &str) -> AppResult<DatabaseConnection> {
+    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<crate::domain::auth::JwksVerifier>,
+    Arc<crate::domain::auth::Authenticator>,
+)> {
+    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<AppState>, 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<dyn std::err
             HeaderValue::from_static("strict-origin-when-cross-origin"),
         ))
         .layer(tower_http::limit::RequestBodyLimitLayer::new(1024 * 1024)) // 1MB limit
-        .with_state(state);
+        .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 {

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<Arc<AppState>>,
-) -> Result<Json<Vec<TaskResponse>>, (StatusCode, String)> {
+) -> AppResult<Json<Vec<TaskResponse>>> {
     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
+        .await?;
-        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
     let response = tasks
         .into_iter()
@@ -54,7 +54,7 @@ pub async fn create_task(
     _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
     Json(payload): Json<CreateTaskRequest>,
-) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+) -> AppResult<Json<TaskResponse>> {
     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
+    new_task.insert(&state.db).await?;
-        .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;
@@ -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<Arc<AppState>>,
     Path(id): Path<Uuid>,
-) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+) -> AppResult<Json<TaskResponse>> {
     let task = crate::entities::task::Entity::find_by_id(id)
         .one(&state.db)
-        .await
+        .await?
-        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or_else(|| crate::error::AppError::NotFound("Task not found".to_string()))?;
-        .ok_or((StatusCode::NOT_FOUND, "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<Arc<AppState>>,
     Path(id): Path<Uuid>,
     Json(payload): Json<UpdateTaskRequest>,
-) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+) -> AppResult<Json<TaskResponse>> {
-    let mut task: crate::entities::task::ActiveModel =
+    let task: crate::entities::task::ActiveModel = crate::entities::task::Entity::find_by_id(id)
-        crate::entities::task::Entity::find_by_id(id)
+        .one(&state.db)
-            .one(&state.db)
+        .await?
-            .await
+        .ok_or_else(|| crate::error::AppError::NotFound("Task not found".to_string()))?
-            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .into();
-            .ok_or((StatusCode::NOT_FOUND, "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)
+    task.update(&state.db).await?;
-        .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;
@@ -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<Uuid>,
     State(state): State<Arc<AppState>>,
-) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+) -> AppResult<Json<TaskResponse>> {
     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<Arc<AppState>>,
-) -> Result<Json<Vec<RecentRunResponse>>, (StatusCode, String)> {
+) -> AppResult<Json<Vec<RecentRunResponse>>> {
     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
+        .await?;
-        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
     let response = results
         .into_iter()