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