diff --git a/Cargo.lock b/Cargo.lock index 44e16d9..1a9f8e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,7 +217,6 @@ dependencies = [ "axum", "axum-core", "bytes", - "cookie", "fastrand", "futures-util", "headers", @@ -326,7 +325,6 @@ dependencies = [ "axum-extra", "base64", "chrono", - "cookie", "dashmap", "jsonwebtoken", "migration", @@ -488,17 +486,6 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" -[[package]] -name = "cookie" -version = "0.18.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" -dependencies = [ - "percent-encoding", - "time", - "version_check", -] - [[package]] name = "core-foundation" version = "0.9.4" diff --git a/Cargo.toml b/Cargo.toml index d0b8b6f..e995879 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ reqwest = { version = "0.12", features = ["json"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" axum = "0.7" -tower-http = { version = "0.5", features = ["cors", "set-header", "limit"] } +tower-http = { version = "0.5", features = ["cors"] } sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] } sea-orm-migration = "1.1" uuid = { version = "1.8", features = ["v4", "serde"] } @@ -22,5 +22,4 @@ tokio-cron-scheduler = "0.15.1" dashmap = "6.1.0" 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" +axum-extra = { version = "0.9", features = ["typed-header"] } diff --git a/frontend/src/main.js b/frontend/src/main.js index 91e3375..3e5100f 100644 --- a/frontend/src/main.js +++ b/frontend/src/main.js @@ -16,7 +16,8 @@ const state = { selectedRunId: null, currentView: 'dashboard', // 'dashboard' or 'task' isEditing: false, - isAuthenticated: false + token: localStorage.getItem('auth_token'), + refreshToken: localStorage.getItem('refresh_token') }; // DOM elements @@ -57,15 +58,29 @@ const presetBtns = document.querySelectorAll('.btn-preset'); // Wrapper for fetch to include Authorization header async function fetchWithAuth(url, options = {}) { - let response = await fetch(url, { ...options, credentials: 'include' }); + if (!state.token) { + showLogin(); + throw new Error('Not authenticated'); + } - if (response.status === 401) { + const headers = { + ...options.headers, + 'Authorization': `Bearer ${state.token}` + }; + + let response = await fetch(url, { ...options, headers }); + + if (response.status === 401 && state.refreshToken) { // Try to refresh token try { const success = await attemptTokenRefresh(); if (success) { - // Retry original request - response = await fetch(url, { ...options, credentials: 'include' }); + // Retry original request with new token + const newHeaders = { + ...options.headers, + 'Authorization': `Bearer ${state.token}` + }; + response = await fetch(url, { ...options, headers: newHeaders }); } } catch (error) { console.error('Token refresh failed:', error); @@ -85,12 +100,20 @@ async function attemptTokenRefresh() { const response = await fetch(`${API_URL}/auth/refresh`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - credentials: 'include', - body: JSON.stringify({ refresh_token: '' }) + body: JSON.stringify({ refresh_token: state.refreshToken }) }); if (response.ok) { - return true; + const data = await response.json(); + if (data.access_token) { + state.token = data.access_token; + localStorage.setItem('auth_token', data.access_token); + if (data.refresh_token) { + state.refreshToken = data.refresh_token; + localStorage.setItem('refresh_token', data.refresh_token); + } + return true; + } } } catch (error) { console.error('Error during token refresh:', error); @@ -130,7 +153,7 @@ async function fetchTasks() { 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; + state.selectedRunId = task.runs[task.runs.length - 1].id; } renderRunHistory(task); @@ -154,7 +177,13 @@ async function fetchRecentRuns() { function renderTaskList() { const sortedTasks = [...state.tasks].sort((a, b) => { - return new Date(b.created_at) - new Date(a.created_at); + const aDate = a.runs && a.runs.length > 0 + ? new Date(a.runs[0].created_at) + : new Date(a.created_at); + const bDate = b.runs && b.runs.length > 0 + ? new Date(b.runs[0].created_at) + : new Date(b.created_at); + return bDate - aDate; }); taskListEl.innerHTML = sortedTasks @@ -194,8 +223,8 @@ function selectTask(id, runId = null) { if (runId) { state.selectedRunId = runId; } else if (task.runs && task.runs.length > 0) { - // Default to latest run if not specified (index 0 is newest) - state.selectedRunId = task.runs[0].id; + // Default to latest run if not specified + state.selectedRunId = task.runs[task.runs.length - 1].id; } else { state.selectedRunId = null; } @@ -466,20 +495,6 @@ newTaskForm.addEventListener('submit', async (e) => { } }); -async function checkSession() { - try { - const response = await fetch(`${API_URL}/auth/session`, { credentials: 'include' }); - if (response.ok) { - state.isAuthenticated = true; - return true; - } - } catch (error) { - console.error('Session check failed:', error); - } - state.isAuthenticated = false; - return false; -} - let isPolling = false; async function startAutoRefresh() { setInterval(async () => { @@ -499,12 +514,10 @@ async function showLogin() { } async function logout() { - try { - await fetch(`${API_URL}/auth/logout`, { method: 'POST', credentials: 'include' }); - } catch (error) { - console.error('Logout failed:', error); - } - state.isAuthenticated = false; + state.token = null; + state.refreshToken = null; + localStorage.removeItem('auth_token'); + localStorage.removeItem('refresh_token'); showLogin(); } @@ -518,13 +531,16 @@ async function handleCallback() { loginOverlay.classList.add('hidden'); try { - const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`, { - credentials: 'include' - }); + const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`); const data = await response.json(); - if (response.ok) { - state.isAuthenticated = true; + if (data.access_token) { + state.token = data.access_token; + localStorage.setItem('auth_token', data.access_token); + if (data.refresh_token) { + state.refreshToken = data.refresh_token; + localStorage.setItem('refresh_token', data.refresh_token); + } callbackOverlay.classList.add('hidden'); appEl.classList.remove('hidden'); initializeApp(); @@ -548,16 +564,16 @@ logoutBtn.addEventListener('click', () => { }); async function initializeApp() { - const hasSession = await checkSession(); - - if (hasSession) { - appEl.classList.remove('hidden'); - loginOverlay.classList.add('hidden'); - await fetchTasks(); - startAutoRefresh(); - } else { + if (!state.token) { showLogin(); + return; } + + appEl.classList.remove('hidden'); + loginOverlay.classList.add('hidden'); + + await fetchTasks(); + startAutoRefresh(); } // Check for callback on load diff --git a/src/agent.rs b/src/agent.rs index 09c93ce..2bbc042 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -1,5 +1,4 @@ use chrono::Utc; -use std::time::{Duration, Instant}; use crate::api::{ChatRequest, ChatResponse, Message, Tool}; use crate::tools; @@ -28,6 +27,7 @@ impl Agent { answer tool once you to give your final answer. current date is {}", Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string() ); + println!("initial_message: {}", intro); let messages = vec![ Message { role: "system".to_string(), @@ -62,35 +62,15 @@ impl Agent { } fn log(&mut self, message: &str) { + println!("{}", message); 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 mut file_written = false; - 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; + while !file_written { let request = ChatRequest { model: "kimi-k2.5".to_string(), messages: self.messages.clone(), @@ -120,12 +100,7 @@ impl Agent { } let chat_response: ChatResponse = response.json().await?; - let assistant_message = chat_response - .choices - .get(0) - .ok_or("Missing assistant response")? - .message - .clone(); + let assistant_message = chat_response.choices.get(0).unwrap().message.clone(); self.messages.push(assistant_message.clone()); @@ -153,16 +128,16 @@ impl Agent { self.messages.push(tool_message); if written { - finished = true; + file_written = 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... ---"); + // No more tool calls from assistant, but we only exit if file was written + if !file_written { + self.log("--- Assistant didn't use write_file yet. Waiting for next turn... ---"); } } diff --git a/src/auth.rs b/src/auth.rs index b2f8c8d..0cf0f67 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -10,17 +10,10 @@ pub struct Claims { pub exp: usize, pub iat: usize, pub iss: String, - pub aud: Audience, + pub aud: String, } -#[derive(Debug, Serialize, Deserialize)] -#[serde(untagged)] -pub enum Audience { - Single(String), - Multiple(Vec), -} - -#[derive(Debug, Deserialize, Clone)] +#[derive(Debug, Deserialize)] struct Jwk { #[serde(rename = "kty")] _kty: String, @@ -38,17 +31,13 @@ struct Jwks { pub struct JwksVerifier { issuer: String, - audience: String, jwks_uri: String, keys: Arc>>, client: Client, } impl JwksVerifier { - pub async fn new( - issuer: String, - audience: String, - ) -> Result> { + pub async fn new(issuer: String) -> Result> { let client = Client::new(); // Authentik OIDC discovery let discovery_url = format!( @@ -64,7 +53,6 @@ impl JwksVerifier { let verifier = Self { issuer, - audience, jwks_uri, keys: Arc::new(RwLock::new(Vec::new())), client, @@ -85,29 +73,18 @@ impl JwksVerifier { let header = decode_header(token)?; let kid = header.kid.ok_or("Missing kid in token header")?; - let jwk = { - let keys = self.keys.read().await; - keys.iter().find(|k| k.kid == kid).cloned() - }; - - let jwk = match jwk { - Some(key) => key, - None => { - self.refresh_keys().await?; - let keys = self.keys.read().await; - keys.iter() - .find(|k| k.kid == kid) - .cloned() - .ok_or("Key not found in JWKS")? - } - }; + let keys = self.keys.read().await; + let jwk = keys + .iter() + .find(|k| k.kid == kid) + .ok_or("Key not found in JWKS")?; let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?; let mut validation = Validation::new(Algorithm::RS256); validation.set_issuer(&[self.issuer.clone()]); - validation.set_audience(&[self.audience.clone()]); - validation.validate_aud = true; + // Aud validation might need careful config, usually it's the client_id + validation.validate_aud = false; let token_data = decode::(token, &decoding_key, &validation)?; Ok(token_data.claims) diff --git a/src/scheduler.rs b/src/scheduler.rs index d2c8d47..d20f257 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -98,28 +98,19 @@ impl Scheduler { let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?; tokio::spawn(async move { - let (logs, answer, status) = match agent.run().await { - Ok((logs, answer)) => (logs, answer, "completed".to_string()), - Err(e) => ( - format!("Scheduled run failed: {}", e), - None, - "failed".to_string(), - ), + let (logs, answer) = match agent.run().await { + Ok(res) => res, + Err(e) => (format!("Scheduled run failed: {}", e), None), }; let run_complete = task_run::ActiveModel { id: Set(run_id), - status: Set(status), + status: Set("completed".to_string()), logs: Set(logs), answer: Set(answer), ..Default::default() }; - if let Err(e) = run_complete.update(&db).await { - eprintln!( - "Failed to update scheduled run status for task {}: {}", - task_id, e - ); - } + let _ = run_complete.update(&db).await; }); Ok(()) diff --git a/src/server.rs b/src/server.rs index a7e97b1..e7d53b5 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,13 +1,11 @@ use axum::{ Json, RequestPartsExt, Router, extract::{FromRef, FromRequestParts, Path, Query, State}, - http::{HeaderValue, StatusCode, request::Parts}, - response::IntoResponse, + http::{StatusCode, request::Parts}, routing::{get, post}, }; use axum_extra::{ TypedHeader, - extract::cookie::{Cookie, CookieJar, SameSite}, headers::{Authorization, authorization::Bearer}, }; use chrono::Utc; @@ -16,7 +14,7 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; use std::sync::Arc; -use tower_http::cors::{AllowOrigin, CorsLayer}; +use tower_http::cors::{Any, CorsLayer}; use uuid::Uuid; use crate::agent::Agent; @@ -99,10 +97,7 @@ pub async fn start(db_url: &str) -> Result<(), Box> { 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 verifier = Arc::new(crate::auth::JwksVerifier::new(authentik_issuer.clone()).await?); let authenticator = Arc::new( crate::auth::Authenticator::new( authentik_issuer, @@ -121,31 +116,19 @@ pub async fn start(db_url: &str) -> Result<(), Box> { authenticator, }); - let cors = build_cors_layer(); + let cors = CorsLayer::new() + .allow_origin(Any) + .allow_methods(Any) + .allow_headers(Any); 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()); @@ -157,44 +140,6 @@ pub async fn start(db_url: &str) -> Result<(), Box> { 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>, @@ -304,13 +249,9 @@ async fn execute_agent_run( ) .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(), - ), + let (logs, answer) = match agent.run().await { + Ok((logs, answer)) => (logs, answer), + Err(e) => (format!("Execution failed: {}", e), None), }; // Update with final logs and status @@ -326,7 +267,7 @@ async fn execute_agent_run( run.logs = Set(logs.clone()); run.answer = Set(answer.clone()); - run.status = Set(status); + run.status = Set("completed".to_string()); run.update(&state.db) .await @@ -378,27 +319,26 @@ where 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 TypedHeader(Authorization(bearer)) = parts + .extract::>>() + .await + .map_err(|_| { + ( + StatusCode::UNAUTHORIZED, + "Missing or invalid Authorization header".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), - ) - })?; + let claims = app_state + .verifier + .verify(bearer.token()) + .await + .map_err(|e| { + ( + StatusCode::UNAUTHORIZED, + format!("Token verification failed: {}", e), + ) + })?; Ok(AuthenticatedUser(claims)) } @@ -412,117 +352,36 @@ pub struct AuthCallbackQuery { #[derive(Deserialize)] struct RefreshRequest { - refresh_token: Option, + refresh_token: String, } 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 +) -> Result, (StatusCode, String)> { + state .authenticator - .refresh_token(refresh_token) + .refresh_token(payload.refresh_token) .await - .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?; - - let jar = update_auth_cookies(jar, &data); - Ok((jar, Json(data))) + .map(Json) + .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string())) } async fn auth_callback( State(state): State>, - jar: CookieJar, Query(query): Query, -) -> Result { - let data = state +) -> Result, (StatusCode, String)> { + state .authenticator .exchange_code(query.code, query.redirect_uri) .await + .map(Json) .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( diff --git a/src/tools.rs b/src/tools.rs index 37e665c..ef1c094 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -51,6 +51,11 @@ pub async fn handle_tool_call( let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; let query = args.get("query").ok_or("Missing query argument")?; + println!( + "--- Executing tool: google_search(query: \"{}\") ---", + query + ); + let search_result = if let Some(key) = tavily_api_key { match api::perform_search(query, key).await { Ok(results) => results, @@ -64,6 +69,8 @@ pub async fn handle_tool_call( let args: HashMap = serde_json::from_str(&tool_call.function.arguments)?; let result = args.get("result").ok_or("Missing result argument")?; + println!("--- Finishing task: {}", result); + answer = Some(result.clone()); (result.clone(), true) } else {