diff --git a/Cargo.lock b/Cargo.lock index 44e16d9..73d8ffb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -328,6 +328,7 @@ dependencies = [ "chrono", "cookie", "dashmap", + "dotenvy", "jsonwebtoken", "migration", "reqwest", @@ -335,9 +336,12 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", + "thiserror", "tokio", "tokio-cron-scheduler", "tower-http 0.5.2", + "tracing", + "tracing-subscriber", "uuid", ] @@ -1725,6 +1729,15 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.60.2", +] + [[package]] name = "num-bigint" version = "0.4.6" @@ -3503,6 +3516,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", ] [[package]] @@ -3512,12 +3537,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e" dependencies = [ "matchers", + "nu-ansi-term", "once_cell", "regex-automata", "sharded-slab", + "smallvec", "thread_local", "tracing", "tracing-core", + "tracing-log", ] [[package]] @@ -3607,6 +3635,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/Cargo.toml b/Cargo.toml index d0b8b6f..6d3bcb2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,3 +24,7 @@ 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" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/frontend/index.html b/frontend/index.html index 6567fa6..ea8c9bd 100644 --- a/frontend/index.html +++ b/frontend/index.html @@ -175,6 +175,7 @@ +
diff --git a/frontend/src/main.js b/frontend/src/main.js index 91e3375..64ef57e 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,50 @@ 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 showToast(message, type = 'info') { + const container = document.getElementById('toast-container'); + const toast = document.createElement('div'); + toast.className = `toast ${type}`; + + const icons = { + success: '✓', + error: '✕', + info: 'ℹ' + }; + + toast.innerHTML = ` + ${icons[type] || 'ℹ'} + ${message} + `; + + container.appendChild(toast); + + // Auto remove + setTimeout(() => { + toast.style.animation = 'fadeOut 0.3s forwards'; + setTimeout(() => toast.remove(), 300); + }, 4000); +} + +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 +140,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); } @@ -339,9 +368,10 @@ rerunBtn.addEventListener('click', async () => { state.tasks[index] = updatedTask; } selectTask(updatedTask.id); + showToast('Task rerun successfully!', 'success'); } catch (error) { - console.error('Error running task:', error); - alert('Failed to run task.'); + console.error('Failed to rerun task:', error); + showToast('Failed to rerun task.', 'error'); } }); @@ -460,9 +490,10 @@ newTaskForm.addEventListener('submit', async (e) => { state.isEditing = false; selectTask(updatedTask.id); renderTaskList(); + showToast(state.isEditing ? 'Task updated successfully' : 'Task created successfully', 'success'); } catch (error) { - console.error('Error creating task:', error); - alert('Failed to execute task. Check console.'); + console.error('Save task failed:', error); + showToast('Failed to execute task. Check console.', 'error'); } }); @@ -532,8 +563,8 @@ async function handleCallback() { throw new Error('No access token in response'); } } catch (error) { - console.error('Auth callback failed:', error); - alert('Authentication failed.'); + console.error('Callback failed:', error); + showToast('Authentication failed.', 'error'); showLogin(); } } diff --git a/frontend/src/style.css b/frontend/src/style.css index 742bc21..0761f29 100644 --- a/frontend/src/style.css +++ b/frontend/src/style.css @@ -798,6 +798,77 @@ textarea:focus { .btn-sm { padding: 6px 12px; font-size: 12px; + width: 100%; + justify-content: center; +} + +/* Toast System */ +#toast-container { + position: fixed; + bottom: 24px; + right: 24px; + display: flex; + flex-direction: column; + gap: 12px; + z-index: 2000; +} + +.toast { + min-width: 300px; + padding: 16px 20px; + border-radius: 12px; + background: var(--bg-sidebar); + border: 1px solid var(--glass-border); + box-shadow: 0 10px 30px rgba(0, 0, 0, 0.5); + display: flex; + align-items: center; + gap: 12px; + animation: slideIn 0.3s cubic-bezier(0.4, 0, 0.2, 1); +} + +.toast.error { + border-left: 4px solid var(--status-failed); +} + +.toast.success { + border-left: 4px solid var(--status-completed); +} + +.toast.info { + border-left: 4px solid var(--primary); +} + +.toast-icon { + font-size: 18px; +} + +.toast-message { + font-size: 14px; + font-weight: 500; +} + +@keyframes slideIn { + from { + transform: translateX(100%); + opacity: 0; + } + + to { + transform: translateX(0); + opacity: 1; + } +} + +@keyframes fadeOut { + from { + opacity: 1; + transform: scale(1); + } + + to { + opacity: 0; + transform: scale(0.95); + } } .logout-btn { diff --git a/src/agent.rs b/src/agent.rs deleted file mode 100644 index 09c93ce..0000000 --- a/src/agent.rs +++ /dev/null @@ -1,171 +0,0 @@ -use chrono::Utc; -use std::time::{Duration, Instant}; - -use crate::api::{ChatRequest, ChatResponse, Message, Tool}; -use crate::tools; - -pub struct Agent { - client: reqwest::Client, - url: String, - zen_api_key: Option, - tavily_api_key: Option, - messages: Vec, - tools: Option>, - logs: String, - answer: Option, -} - -impl Agent { - pub fn new( - zen_api_key: Option, - tavily_api_key: Option, - initial_message: String, - ) -> Result> { - 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 - so do not ask for clarification and use the - answer tool once you to give your final answer. current date is {}", - Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string() - ); - let messages = vec![ - Message { - role: "system".to_string(), - content: Some(intro), - tool_calls: None, - tool_call_id: None, - }, - Message { - role: "user".to_string(), - content: Some(initial_message), - tool_calls: None, - tool_call_id: None, - }, - ]; - - let tools = Some(tools::get_tools()); - - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(60)) - .build()?; - - Ok(Self { - client, - url: "https://opencode.ai/zen/v1/chat/completions".to_string(), - zen_api_key, - tavily_api_key, - messages, - tools, - logs: String::new(), - answer: None, - }) - } - - fn log(&mut self, message: &str) { - self.logs.push_str(message); - self.logs.push('\n'); - } - - pub async fn run(&mut self) -> Result<(String, Option), Box> { - let mut finished = false; - let start_time = Instant::now(); - let max_duration_secs = std::env::var("AGENT_MAX_DURATION_SECS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(120); - let max_duration = Duration::from_secs(max_duration_secs); - - let max_turns = std::env::var("AGENT_MAX_TURNS") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(20); - let mut turns = 0; - - while !finished { - if start_time.elapsed() > max_duration { - return Err("Agent run timed out".into()); - } - - if turns >= max_turns { - return Err("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(), - }; - - self.log(&format!( - "--- Sending request to Zen API (Role: {}) ---", - self.messages.last().unwrap().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 assistant_message = chat_response - .choices - .get(0) - .ok_or("Missing assistant response")? - .message - .clone(); - - self.messages.push(assistant_message.clone()); - - if let Some(content) = &assistant_message.content { - if !content.is_empty() { - self.log(&format!("\nAssistant response:\n{}\n", 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?; - - if let Some(ans) = tool_answer { - self.answer = Some(ans); - } - - if let Some(content) = &tool_message.content { - self.log(&format!( - "Tool result ({}): {}", - tool_call.function.name, content - )); - } - - self.messages.push(tool_message); - if written { - 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... ---"); - } - } - - Ok((self.logs.clone(), self.answer.clone())) - } -} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..c7a714e --- /dev/null +++ b/src/config.rs @@ -0,0 +1,71 @@ +use crate::error::{AppError, AppResult}; +use std::env; + +#[derive(Clone, Debug)] +pub struct Config { + pub database_url: String, + pub port: u16, + pub zen_api_key: Option, + pub tavily_api_key: Option, + pub authentik_issuer: String, + pub authentik_client_id: String, + pub authentik_client_secret: String, + pub cors_allowed_origins: Option, + pub cookie_secure: bool, + pub agent_max_turns: u32, + pub agent_max_duration_secs: u64, +} + +impl Config { + 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() + .and_then(|s| s.parse().ok()) + .unwrap_or(3000); + + 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(|_| 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(|_| AppError::Config("AUTHENTIK_CLIENT_SECRET must be set".into()))?; + + let cors_allowed_origins = env::var("CORS_ALLOWED_ORIGINS").ok(); + + let cookie_secure = env::var("COOKIE_SECURE") + .map(|v| v == "true") + .unwrap_or(false); + + let agent_max_turns = env::var("AGENT_MAX_TURNS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(20); + + let agent_max_duration_secs = env::var("AGENT_MAX_DURATION_SECS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(120); + + Ok(Config { + database_url, + port, + zen_api_key, + tavily_api_key, + authentik_issuer, + authentik_client_id, + authentik_client_secret, + cors_allowed_origins, + cookie_secure, + agent_max_turns, + agent_max_duration_secs, + }) + } +} 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/domain/agent/mod.rs b/src/domain/agent/mod.rs new file mode 100644 index 0000000..9f6933f --- /dev/null +++ b/src/domain/agent/mod.rs @@ -0,0 +1,190 @@ +pub mod api; +pub mod tools; + +use chrono::Utc; +use std::time::{Duration, Instant}; + +use self::api::{ChatRequest, ChatResponse, Message, Tool}; + +pub struct Agent { + client: reqwest::Client, + url: String, + zen_api_key: Option, + tavily_api_key: Option, + messages: Vec, + tools: Option>, + logs: String, + answer: Option, +} + +use crate::error::{AppError, AppResult}; + +impl Agent { + pub fn new( + zen_api_key: Option, + tavily_api_key: Option, + initial_message: String, + ) -> 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 + so do not ask for clarification and use the + answer tool once you to give your final answer. current date is {}", + Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string() + ); + let messages = vec![ + Message { + role: "system".to_string(), + content: Some(intro), + tool_calls: None, + tool_call_id: None, + }, + Message { + role: "user".to_string(), + content: Some(initial_message), + tool_calls: None, + tool_call_id: None, + }, + ]; + + let tools = Some(tools::get_tools()); + + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(60)) + .build() + .map_err(|e| AppError::Internal(format!("Failed to build HTTP client: {}", e)))?; + + Ok(Self { + client, + url: "https://opencode.ai/zen/v1/chat/completions".to_string(), + zen_api_key, + tavily_api_key, + messages, + tools, + logs: String::new(), + answer: None, + }) + } + + fn log(&mut self, message: &str) { + self.logs.push_str(message); + self.logs.push('\n'); + } + + pub async fn run( + &mut self, + config: &crate::config::Config, + ) -> AppResult<(String, Option)> { + let mut finished = false; + let start_time = Instant::now(); + let max_duration = Duration::from_secs(config.agent_max_duration_secs); + let max_turns = config.agent_max_turns; + let mut turns = 0; + + while !finished { + if start_time.elapsed() > max_duration { + return Err(AppError::Internal("Agent run timed out".into())); + } + + if turns >= max_turns { + return Err(AppError::Internal("Agent run exceeded max turns".into())); + } + + turns += 1; + let current_role = self + .messages + .last() + .map(|m| m.role.as_str()) + .unwrap_or("unknown"); + self.log(&format!( + "\n[Turn {}] Sending request (Last role: {})", + turns, current_role + )); + + let chat_response = self.call_llm().await?; + let assistant_message = chat_response + .choices + .get(0) + .ok_or_else(|| AppError::Internal("Missing assistant response".into()))? + .message + .clone(); + + self.messages.push(assistant_message.clone()); + + if let Some(content) = &assistant_message.content { + if !content.is_empty() { + self.log(&format!("\nAssistant: {}", content)); + } + } + + if let Some(tool_calls) = assistant_message.tool_calls { + for tool_call in tool_calls { + 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: {}", content)); + } + + self.messages.push(tool_message); + if is_final { + finished = true; + } + } + } 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/tools.rs b/src/domain/agent/tools.rs similarity index 78% rename from src/tools.rs rename to src/domain/agent/tools.rs index 37e665c..663b1b1 100644 --- a/src/tools.rs +++ b/src/domain/agent/tools.rs @@ -1,5 +1,15 @@ -use crate::api::{self, FunctionDefinition, Message, Tool, ToolCall}; -use std::collections::HashMap; +use super::api::{self, FunctionDefinition, Message, Tool, ToolCall}; +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/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..5092091 --- /dev/null +++ b/src/domain/tasks.rs @@ -0,0 +1,148 @@ +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, +} + +use crate::error::AppResult; + +pub async fn execute_agent_run( + db: &DatabaseConnection, + _scheduler: &Arc, + config: &Arc, + task_id: Uuid, + goal: String, +) -> AppResult { + let run_id = Uuid::new_v4(); + tracing::info!(%task_id, %run_id, "Starting agent execution run"); + + 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 + .map_err(crate::error::AppError::Database)?; + + 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)) => { + tracing::info!(%task_id, %run_id, "Agent execution completed successfully"); + (logs, answer, "completed".to_string()) + } + Err(e) => { + tracing::error!(%task_id, %run_id, error = %e, "Agent execution failed"); + ( + format!("Execution failed: {}", e), + None, + "failed".to_string(), + ) + } + }; + + let run: task_run::ActiveModel = TaskRun::find_by_id(run_id) + .one(db) + .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 + .map_err(crate::error::AppError::Database)?; + + get_task_inner(task_id, db).await +} + +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 + .map_err(crate::error::AppError::Database)?; + + 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)); + + 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/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 14ad58e..59849ad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,28 @@ -mod agent; -mod api; -mod auth; +mod config; +mod domain; mod entities; +mod error; mod scheduler; mod server; -mod tools; +#[cfg(test)] +mod tests; + +use tracing::info; #[tokio::main] async fn main() -> Result<(), Box> { - let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "bot=info,axum=info".into()), + ) + .init(); - server::start(&db_url).await?; + info!("Starting Antigravity Agent..."); + + let config = config::Config::from_env()?; + + server::start(config).await?; Ok(()) } diff --git a/src/scheduler.rs b/src/scheduler.rs index d2c8d47..2ee8a01 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,8 +1,10 @@ -use crate::agent::Agent; +use crate::domain::agent::Agent; use crate::entities::task::Entity as Task; use crate::entities::task_run; +use crate::error::{AppError, AppResult}; use dashmap::DashMap; use sea_orm::{DatabaseConnection, EntityTrait, Set}; +use std::sync::Arc; use tokio_cron_scheduler::{Job, JobScheduler}; use uuid::Uuid; @@ -10,78 +12,87 @@ pub struct Scheduler { scheduler: JobScheduler, db: DatabaseConnection, tasks_to_jobs: DashMap, - zen_api_key: Option, - tavily_api_key: Option, + config: Arc, } impl Scheduler { pub async fn new( db: DatabaseConnection, - zen_api_key: Option, - tavily_api_key: Option, - ) -> Result> { - let scheduler = JobScheduler::new().await?; - scheduler.start().await?; + config: Arc, + ) -> AppResult { + let scheduler = JobScheduler::new() + .await + .map_err(|e| AppError::Internal(format!("Failed to create scheduler: {}", e)))?; + scheduler + .start() + .await + .map_err(|e| AppError::Internal(format!("Failed to start scheduler: {}", e)))?; Ok(Self { scheduler, db, tasks_to_jobs: DashMap::new(), - zen_api_key, - tavily_api_key, + config, }) } - pub async fn add_task_job( - &self, - task_id: Uuid, - cron_expr: &str, - ) -> Result<(), Box> { + pub async fn add_task_job(&self, task_id: Uuid, cron_expr: &str) -> AppResult<()> { // Remove existing job if any if let Some((_, old_job_id)) = self.tasks_to_jobs.remove(&task_id) { let _ = self.scheduler.remove(&old_job_id).await; } let db = self.db.clone(); - let zen_key = self.zen_api_key.clone(); - let tavily_key = self.tavily_api_key.clone(); + let config = self.config.clone(); let job = Job::new_async(cron_expr, move |_uuid, _l| { let db = db.clone(); - let zen_key = zen_key.clone(); - let tavily_key = tavily_key.clone(); + let config = config.clone(); Box::pin(async move { - if let Err(e) = Self::run_task(db, zen_key, tavily_key, task_id).await { - eprintln!("Error in scheduled task {}: {}", task_id, e); + if let Err(e) = Self::run_task(db, config, task_id).await { + tracing::error!("Error in scheduled task {}: {}", task_id, e); } }) - })?; + }) + .map_err(|e| AppError::Internal(format!("Failed to create job: {}", e)))?; - let job_id = self.scheduler.add(job).await?; + let job_id = self + .scheduler + .add(job) + .await + .map_err(|e| AppError::Internal(format!("Failed to add job: {}", e)))?; self.tasks_to_jobs.insert(task_id, job_id); + tracing::info!(%task_id, %cron_expr, "Added task to scheduler"); + Ok(()) } - pub async fn remove_task_job(&self, task_id: Uuid) -> Result<(), Box> { + pub async fn remove_task_job(&self, task_id: Uuid) -> AppResult<()> { if let Some((_, job_id)) = self.tasks_to_jobs.remove(&task_id) { - self.scheduler.remove(&job_id).await?; + self.scheduler + .remove(&job_id) + .await + .map_err(|e| AppError::Internal(format!("Failed to remove job: {}", e)))?; + tracing::info!(%task_id, "Removed task from scheduler"); } Ok(()) } async fn run_task( db: DatabaseConnection, - zen_key: Option, - tavily_key: Option, + config: Arc, task_id: Uuid, - ) -> Result<(), Box> { + ) -> AppResult<()> { let task = Task::find_by_id(task_id) .one(&db) - .await? - .ok_or("Task not found")?; + .await + .map_err(AppError::Database)? + .ok_or_else(|| AppError::NotFound("Task not found".into()))?; // Create a new run entry let run_id = Uuid::new_v4(); + tracing::info!(task_id = %task_id, run_id = %run_id, "Starting scheduled task execution"); + let run = task_run::ActiveModel { id: Set(run_id), task_id: Set(task_id), @@ -92,19 +103,29 @@ impl Scheduler { }; use sea_orm::ActiveModelTrait; - run.insert(&db).await?; + run.insert(&db).await.map_err(AppError::Database)?; // Start agent in background - let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?; + let mut agent = Agent::new( + config.zen_api_key.clone(), + config.tavily_api_key.clone(), + task.goal.clone(), + )?; tokio::spawn(async move { - let (logs, answer, status) = match agent.run().await { - Ok((logs, answer)) => (logs, answer, "completed".to_string()), - Err(e) => ( - format!("Scheduled run failed: {}", e), - None, - "failed".to_string(), - ), + let (logs, answer, status) = match agent.run(&config).await { + Ok((logs, answer)) => { + tracing::info!(task_id = %task_id, run_id = %run_id, "Scheduled task execution completed successfully"); + (logs, answer, "completed".to_string()) + } + Err(e) => { + tracing::error!(task_id = %task_id, run_id = %run_id, error = %e, "Scheduled task execution failed"); + ( + format!("Scheduled run failed: {}", e), + None, + "failed".to_string(), + ) + } }; let run_complete = task_run::ActiveModel { @@ -115,9 +136,10 @@ impl Scheduler { ..Default::default() }; if let Err(e) = run_complete.update(&db).await { - eprintln!( + tracing::error!( "Failed to update scheduled run status for task {}: {}", - task_id, e + task_id, + e ); } }); diff --git a/src/server.rs b/src/server.rs deleted file mode 100644 index a7e97b1..0000000 --- a/src/server.rs +++ /dev/null @@ -1,594 +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 zen_api_key: Option, - pub tavily_api_key: Option, - 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(db_url: &str) -> Result<(), Box> { - let db = Database::connect(db_url).await?; - Migrator::up(&db, None).await?; - - let zen_api_key = std::env::var("ZEN_API_KEY").ok(); - let tavily_api_key = std::env::var("TAVILY_API_KEY").ok(); - - let scheduler = - Arc::new(Scheduler::new(db.clone(), zen_api_key.clone(), tavily_api_key.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 authentik_issuer = - std::env::var("AUTHENTIK_ISSUER").map_err(|_| "AUTHENTIK_ISSUER not set")?; - let authentik_client_id = - std::env::var("AUTHENTIK_CLIENT_ID").map_err(|_| "AUTHENTIK_CLIENT_ID not set")?; - let authentik_client_secret = - std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?; - - let verifier = Arc::new( - crate::auth::JwksVerifier::new(authentik_issuer.clone(), authentik_client_id.clone()) - .await?, - ); - let authenticator = Arc::new( - crate::auth::Authenticator::new( - authentik_issuer, - authentik_client_id, - authentik_client_secret, - ) - .await?, - ); - - let state = Arc::new(AppState { - db, - scheduler, - zen_api_key, - tavily_api_key, - verifier, - authenticator, - }); - - let cors = build_cors_layer(); - - 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 port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string()); - let addr = format!("0.0.0.0:{}", port); - let listener = tokio::net::TcpListener::bind(&addr).await?; - println!("Server running on http://localhost:{}", port); - axum::serve(listener, app).await?; - - Ok(()) -} - -fn build_cors_layer() -> CorsLayer { - let origins = std::env::var("CORS_ALLOWED_ORIGINS").ok(); - - let allow_origin = if let Some(origins) = 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.zen_api_key.clone(), - state.tavily_api_key.clone(), - goal.clone(), - ) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - - let (logs, answer, status) = match agent.run().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); - 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); - Ok((jar, Json(data))) -} - -async fn auth_logout(jar: CookieJar) -> impl IntoResponse { - let jar = clear_auth_cookies(jar); - (jar, StatusCode::NO_CONTENT) -} - -async fn auth_session(user: AuthenticatedUser) -> Json { - Json(serde_json::json!({ - "authenticated": true, - "user": user.0 - })) -} - -fn secure() -> bool { - std::env::var("COOKIE_SECURE") - .map(|value| value == "true") - .unwrap_or(false) -} - -fn update_auth_cookies(jar: CookieJar, data: &serde_json::Value) -> 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()) - .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()) - .build(); - jar = jar.add(cookie); - } - - jar -} - -fn clear_auth_cookies(jar: CookieJar) -> 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()) - .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..7ff1086 --- /dev/null +++ b/src/server/auth.rs @@ -0,0 +1,169 @@ +use axum::{ + Json, RequestPartsExt, + extract::{FromRef, FromRequestParts, Query, State}, + http::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::error::{AppError, AppResult}; + +pub struct AuthenticatedUser(pub crate::domain::auth::Claims); + +#[axum::async_trait] +impl FromRequestParts for AuthenticatedUser +where + Arc: axum::extract::FromRef, + S: Send + Sync, +{ + type Rejection = AppError; + + 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 + .map_err(|e| AppError::Internal(e.to_string()))?; + jar.get("access_token") + .map(|cookie| cookie.value().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| AppError::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, +) -> AppResult { + 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_else(|| AppError::Unauthorized("Missing refresh token".into()))?; + + let data = state + .authenticator + .refresh_token(refresh_token) + .await + .map_err(|e| AppError::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, +) -> AppResult { + let data = state + .authenticator + .exchange_code(query.code, query.redirect_uri) + .await + .map_err(|e| AppError::Internal(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, axum::http::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..79d28e2 --- /dev/null +++ b/src/server/mod.rs @@ -0,0 +1,174 @@ +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; + +use crate::error::AppResult; + +#[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) -> 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 + .map_err(|e| crate::error::AppError::Internal(e.to_string()))?, + ); + + // Load existing scheduled tasks + 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, authenticator) = setup_auth(&config).await?; + + let state = Arc::new(AppState { + db, + scheduler, + config: config.clone(), + verifier, + authenticator, + }); + + let app = build_app(state, &config); + + 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()))?; + tracing::info!("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 { + 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)) + .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) +} + +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..9782b9c --- /dev/null +++ b/src/server/tasks.rs @@ -0,0 +1,192 @@ +use axum::{ + Json, + extract::{Path, State}, +}; +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, +}; + +use crate::error::AppResult; + +pub async fn list_tasks( + _user: AuthenticatedUser, + State(state): State>, +) -> 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?; + + 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, +) -> AppResult> { + if payload.goal.trim().is_empty() { + return Err(crate::error::AppError::InvalidRequest( + "Goal cannot be empty".into(), + )); + } + + if payload.goal.trim().len() < 5 { + return Err(crate::error::AppError::InvalidRequest( + "Goal is too short (min 5 characters)".into(), + )); + } + + let task_id = Uuid::new_v4(); + tracing::info!(%task_id, goal = %payload.goal, "Creating new task"); + + let new_task = crate::entities::task::ActiveModel { + id: sea_orm::Set(task_id), + goal: sea_orm::Set(payload.goal), + 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?; + + if let Some(cron) = &payload.cron { + let _ = state.scheduler.add_task_job(task_id, cron).await; + } + + tasks::get_task_inner(task_id, &state.db).await.map(Json) +} + +pub async fn rerun_task( + _user: AuthenticatedUser, + State(state): State>, + Path(id): Path, +) -> AppResult> { + let task = crate::entities::task::Entity::find_by_id(id) + .one(&state.db) + .await? + .ok_or_else(|| crate::error::AppError::NotFound("Task not found".to_string()))?; + + tracing::info!(task_id = %task.id, "Manually triggering task rerun"); + + tasks::execute_agent_run( + &state.db, + &state.scheduler, + &state.config, + task.id, + task.goal, + ) + .await + .map(Json) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) +} + +pub async fn update_task( + _user: AuthenticatedUser, + State(state): State>, + Path(id): Path, + Json(payload): Json, +) -> AppResult> { + if payload.goal.trim().is_empty() { + return Err(crate::error::AppError::InvalidRequest( + "Goal cannot be empty".into(), + )); + } + + if payload.goal.trim().len() < 5 { + return Err(crate::error::AppError::InvalidRequest( + "Goal is too short (min 5 characters)".into(), + )); + } + + 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; + tracing::info!(task_id = %id, goal = %payload.goal, "Updating task"); + + task.goal = sea_orm::Set(payload.goal); + task.cron = sea_orm::Set(payload.cron.clone()); + + use sea_orm::ActiveModelTrait; + task.update(&state.db).await?; + + 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) +} + +pub async fn get_task( + _user: AuthenticatedUser, + Path(id): Path, + State(state): State>, +) -> AppResult> { + tasks::get_task_inner(id, &state.db) + .await + .map(Json) + .map_err(|e| crate::error::AppError::Internal(e.to_string())) +} + +pub async fn get_recent_runs( + _user: AuthenticatedUser, + State(state): State>, +) -> 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?; + + 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/tests.rs b/src/tests.rs new file mode 100644 index 0000000..ca3206e --- /dev/null +++ b/src/tests.rs @@ -0,0 +1,39 @@ +use crate::config::Config; +use crate::error::{AppError, AppResult}; +use axum::http::StatusCode; +use axum::response::IntoResponse; + +#[test] +fn test_app_error_into_response() { + let err = AppError::NotFound("Resource not found".into()); + let response = err.into_response(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + let err = AppError::Unauthorized("Invalid token".into()); + let response = err.into_response(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + + let err = AppError::Internal("Server glitch".into()); + let response = err.into_response(); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); +} + +#[tokio::test] +async fn test_config_validation() { + // We can't easily clear all env vars in multi-threaded tests, + // but we can test that it fails if a required one is missing (if we can ensure it's missing) + // However, for this environment, it's safer to test the mapping logic if it was more complex. + + // Instead, let's test a helper if we had one, or just verify AppResult works as expected. + let result: AppResult = Err(AppError::Config("Missing DATABASE_URL".into())); + assert!(result.is_err()); + if let Err(AppError::Config(msg)) = result { + assert_eq!(msg, "Missing DATABASE_URL"); + } +} + +#[test] +fn test_error_variants() { + let err = AppError::InvalidRequest("Bad input".into()); + assert!(err.to_string().contains("Bad input")); +}