diff --git a/Cargo.lock b/Cargo.lock index 73d8ffb..44e16d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -328,7 +328,6 @@ dependencies = [ "chrono", "cookie", "dashmap", - "dotenvy", "jsonwebtoken", "migration", "reqwest", @@ -336,12 +335,9 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", - "thiserror", "tokio", "tokio-cron-scheduler", "tower-http 0.5.2", - "tracing", - "tracing-subscriber", "uuid", ] @@ -1729,15 +1725,6 @@ 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" @@ -3516,18 +3503,6 @@ 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]] @@ -3537,15 +3512,12 @@ 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]] @@ -3635,12 +3607,6 @@ 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 6d3bcb2..d0b8b6f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,3 @@ 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 ea8c9bd..6567fa6 100644 --- a/frontend/index.html +++ b/frontend/index.html @@ -175,7 +175,6 @@ -
diff --git a/frontend/src/main.js b/frontend/src/main.js index 64ef57e..91e3375 100644 --- a/frontend/src/main.js +++ b/frontend/src/main.js @@ -18,6 +18,7 @@ const state = { isEditing: false, isAuthenticated: false }; + // DOM elements const loginOverlay = document.getElementById('login-overlay'); const callbackOverlay = document.getElementById('callback-overlay'); @@ -54,50 +55,6 @@ 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' }); @@ -140,32 +97,46 @@ 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 - let newSelectedRunId = state.selectedRunId; + // Check if we should follow the latest run (if we were already watching it) + let shouldFollowLatest = false; if (state.selectedTaskId) { - const currentTask = newTasks.find(t => t.id === state.selectedTaskId); + const currentTask = state.tasks.find(t => t.id === state.selectedTaskId); if (currentTask && currentTask.runs && currentTask.runs.length > 0) { - // 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; - } + const latestRunId = currentTask.runs[currentTask.runs.length - 1].id; + if (state.selectedRunId === latestRunId) { + shouldFollowLatest = true; } + } else if (!state.selectedRunId) { + shouldFollowLatest = true; } } - updateState({ - tasks: newTasks, - selectedRunId: newSelectedRunId - }); + 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); + } + } } catch (error) { console.error('Error fetching tasks:', error); } @@ -368,10 +339,9 @@ rerunBtn.addEventListener('click', async () => { state.tasks[index] = updatedTask; } selectTask(updatedTask.id); - showToast('Task rerun successfully!', 'success'); } catch (error) { - console.error('Failed to rerun task:', error); - showToast('Failed to rerun task.', 'error'); + console.error('Error running task:', error); + alert('Failed to run task.'); } }); @@ -490,10 +460,9 @@ 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('Save task failed:', error); - showToast('Failed to execute task. Check console.', 'error'); + console.error('Error creating task:', error); + alert('Failed to execute task. Check console.'); } }); @@ -563,8 +532,8 @@ async function handleCallback() { throw new Error('No access token in response'); } } catch (error) { - console.error('Callback failed:', error); - showToast('Authentication failed.', 'error'); + console.error('Auth callback failed:', error); + alert('Authentication failed.'); showLogin(); } } diff --git a/frontend/src/style.css b/frontend/src/style.css index 0761f29..742bc21 100644 --- a/frontend/src/style.css +++ b/frontend/src/style.css @@ -798,77 +798,6 @@ 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 new file mode 100644 index 0000000..09c93ce --- /dev/null +++ b/src/agent.rs @@ -0,0 +1,171 @@ +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/domain/agent/api.rs b/src/api.rs similarity index 100% rename from src/domain/agent/api.rs rename to src/api.rs diff --git a/src/domain/auth.rs b/src/auth.rs similarity index 96% rename from src/domain/auth.rs rename to src/auth.rs index c70e0c7..b2f8c8d 100644 --- a/src/domain/auth.rs +++ b/src/auth.rs @@ -45,8 +45,12 @@ 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/config.rs b/src/config.rs deleted file mode 100644 index c7a714e..0000000 --- a/src/config.rs +++ /dev/null @@ -1,71 +0,0 @@ -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/domain/agent/mod.rs b/src/domain/agent/mod.rs deleted file mode 100644 index 9f6933f..0000000 --- a/src/domain/agent/mod.rs +++ /dev/null @@ -1,190 +0,0 @@ -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/domain/mod.rs b/src/domain/mod.rs deleted file mode 100644 index 7982060..0000000 --- a/src/domain/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod agent; -pub mod auth; -pub mod tasks; diff --git a/src/domain/tasks.rs b/src/domain/tasks.rs deleted file mode 100644 index 5092091..0000000 --- a/src/domain/tasks.rs +++ /dev/null @@ -1,148 +0,0 @@ -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 deleted file mode 100644 index 01943d5..0000000 --- a/src/error.rs +++ /dev/null @@ -1,53 +0,0 @@ -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 59849ad..14ad58e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,28 +1,16 @@ -mod config; -mod domain; +mod agent; +mod api; +mod auth; mod entities; -mod error; mod scheduler; mod server; -#[cfg(test)] -mod tests; - -use tracing::info; +mod tools; #[tokio::main] async fn main() -> Result<(), Box> { - tracing_subscriber::fmt() - .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "bot=info,axum=info".into()), - ) - .init(); + let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - info!("Starting Antigravity Agent..."); - - let config = config::Config::from_env()?; - - server::start(config).await?; + server::start(&db_url).await?; Ok(()) } diff --git a/src/scheduler.rs b/src/scheduler.rs index 2ee8a01..d2c8d47 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -1,10 +1,8 @@ -use crate::domain::agent::Agent; +use crate::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; @@ -12,87 +10,78 @@ pub struct Scheduler { scheduler: JobScheduler, db: DatabaseConnection, tasks_to_jobs: DashMap, - config: Arc, + zen_api_key: Option, + tavily_api_key: Option, } impl Scheduler { pub async fn new( db: DatabaseConnection, - 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)))?; + zen_api_key: Option, + tavily_api_key: Option, + ) -> Result> { + let scheduler = JobScheduler::new().await?; + scheduler.start().await?; Ok(Self { scheduler, db, tasks_to_jobs: DashMap::new(), - config, + zen_api_key, + tavily_api_key, }) } - pub async fn add_task_job(&self, task_id: Uuid, cron_expr: &str) -> AppResult<()> { + pub async fn add_task_job( + &self, + task_id: Uuid, + cron_expr: &str, + ) -> Result<(), Box> { // 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 config = self.config.clone(); + let zen_key = self.zen_api_key.clone(); + let tavily_key = self.tavily_api_key.clone(); let job = Job::new_async(cron_expr, move |_uuid, _l| { let db = db.clone(); - let config = config.clone(); + let zen_key = zen_key.clone(); + let tavily_key = tavily_key.clone(); Box::pin(async move { - if let Err(e) = Self::run_task(db, config, task_id).await { - tracing::error!("Error in scheduled task {}: {}", task_id, e); + if let Err(e) = Self::run_task(db, zen_key, tavily_key, task_id).await { + eprintln!("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 - .map_err(|e| AppError::Internal(format!("Failed to add job: {}", e)))?; + let job_id = self.scheduler.add(job).await?; 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) -> AppResult<()> { + pub async fn remove_task_job(&self, task_id: Uuid) -> Result<(), Box> { if let Some((_, job_id)) = self.tasks_to_jobs.remove(&task_id) { - 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"); + self.scheduler.remove(&job_id).await?; } Ok(()) } async fn run_task( db: DatabaseConnection, - config: Arc, + zen_key: Option, + tavily_key: Option, task_id: Uuid, - ) -> AppResult<()> { + ) -> Result<(), Box> { let task = Task::find_by_id(task_id) .one(&db) - .await - .map_err(AppError::Database)? - .ok_or_else(|| AppError::NotFound("Task not found".into()))?; + .await? + .ok_or("Task not found")?; // 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), @@ -103,29 +92,19 @@ impl Scheduler { }; use sea_orm::ActiveModelTrait; - run.insert(&db).await.map_err(AppError::Database)?; + run.insert(&db).await?; // Start agent in background - let mut agent = Agent::new( - config.zen_api_key.clone(), - config.tavily_api_key.clone(), - task.goal.clone(), - )?; + let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?; tokio::spawn(async move { - 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 (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 run_complete = task_run::ActiveModel { @@ -136,10 +115,9 @@ impl Scheduler { ..Default::default() }; if let Err(e) = run_complete.update(&db).await { - tracing::error!( + eprintln!( "Failed to update scheduled run status for task {}: {}", - task_id, - e + task_id, e ); } }); diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..a7e97b1 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,594 @@ +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 deleted file mode 100644 index 7ff1086..0000000 --- a/src/server/auth.rs +++ /dev/null @@ -1,169 +0,0 @@ -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 deleted file mode 100644 index 79d28e2..0000000 --- a/src/server/mod.rs +++ /dev/null @@ -1,174 +0,0 @@ -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 deleted file mode 100644 index 9782b9c..0000000 --- a/src/server/tasks.rs +++ /dev/null @@ -1,192 +0,0 @@ -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 deleted file mode 100644 index ca3206e..0000000 --- a/src/tests.rs +++ /dev/null @@ -1,39 +0,0 @@ -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")); -} diff --git a/src/domain/agent/tools.rs b/src/tools.rs similarity index 78% rename from src/domain/agent/tools.rs rename to src/tools.rs index 663b1b1..37e665c 100644 --- a/src/domain/agent/tools.rs +++ b/src/tools.rs @@ -1,15 +1,5 @@ -use super::api::{self, FunctionDefinition, Message, Tool, ToolCall}; -use serde::Deserialize; - -#[derive(Deserialize)] -struct GoogleSearchArgs { - query: String, -} - -#[derive(Deserialize)] -struct FinishArgs { - result: String, -} +use crate::api::{self, FunctionDefinition, Message, Tool, ToolCall}; +use std::collections::HashMap; pub fn get_tools() -> Vec { vec![ @@ -34,13 +24,13 @@ pub fn get_tools() -> Vec { tool_type: "function".to_string(), function: FunctionDefinition { name: "finish".to_string(), - description: "Finish the task and provide a final answer".to_string(), + description: "Finish the task".to_string(), parameters: serde_json::json!({ "type": "object", "properties": { "result": { "type": "string", - "description": "The final detailed answer to the task" + "description": "The result of the task" } }, "required": ["result"] @@ -58,8 +48,8 @@ pub async fn handle_tool_call( let name = &tool_call.function.name; let (content, written) = if name == "google_search" { - let args: GoogleSearchArgs = serde_json::from_str(&tool_call.function.arguments)?; - let query = &args.query; + let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; + let query = args.get("query").ok_or("Missing query argument")?; let search_result = if let Some(key) = tavily_api_key { match api::perform_search(query, key).await { @@ -71,8 +61,8 @@ pub async fn handle_tool_call( }; (search_result, false) } else if name == "finish" { - let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?; - let result = &args.result; + let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; + let result = args.get("result").ok_or("Missing result argument")?; answer = Some(result.clone()); (result.clone(), true)