Compare commits

...

7 commits

Author SHA1 Message Date
d90052cfd1 fix 2026-02-11 01:04:46 +01:00
39143e1998 fix: stuff 2026-02-11 00:56:33 +01:00
6e1251b136 fix: use cookies 2026-02-11 00:51:43 +01:00
4dc3d84d38 fix: improve task status 2026-02-11 00:26:47 +01:00
886d408c3b fix: terminate stuck agent 2026-02-11 00:23:57 +01:00
fcad49bf7a fix: cors 2026-02-11 00:18:03 +01:00
844677cc27 fix: token validation 2026-02-11 00:14:27 +01:00
8 changed files with 324 additions and 135 deletions

13
Cargo.lock generated
View file

@ -217,6 +217,7 @@ dependencies = [
"axum", "axum",
"axum-core", "axum-core",
"bytes", "bytes",
"cookie",
"fastrand", "fastrand",
"futures-util", "futures-util",
"headers", "headers",
@ -325,6 +326,7 @@ dependencies = [
"axum-extra", "axum-extra",
"base64", "base64",
"chrono", "chrono",
"cookie",
"dashmap", "dashmap",
"jsonwebtoken", "jsonwebtoken",
"migration", "migration",
@ -486,6 +488,17 @@ version = "0.9.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" 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]] [[package]]
name = "core-foundation" name = "core-foundation"
version = "0.9.4" version = "0.9.4"

View file

@ -13,7 +13,7 @@ reqwest = { version = "0.12", features = ["json"] }
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"
axum = "0.7" axum = "0.7"
tower-http = { version = "0.5", features = ["cors"] } tower-http = { version = "0.5", features = ["cors", "set-header", "limit"] }
sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] } sea-orm = { version = "1.1", features = ["sqlx-postgres", "runtime-tokio-rustls", "macros"] }
sea-orm-migration = "1.1" sea-orm-migration = "1.1"
uuid = { version = "1.8", features = ["v4", "serde"] } uuid = { version = "1.8", features = ["v4", "serde"] }
@ -22,4 +22,5 @@ tokio-cron-scheduler = "0.15.1"
dashmap = "6.1.0" dashmap = "6.1.0"
jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] } jsonwebtoken = { version = "10.3.0", features = ["rsa", "rust_crypto"] }
base64 = "0.22.1" base64 = "0.22.1"
axum-extra = { version = "0.9", features = ["typed-header"] } axum-extra = { version = "0.9", features = ["typed-header", "cookie"] }
cookie = "0.18"

View file

@ -16,8 +16,7 @@ const state = {
selectedRunId: null, selectedRunId: null,
currentView: 'dashboard', // 'dashboard' or 'task' currentView: 'dashboard', // 'dashboard' or 'task'
isEditing: false, isEditing: false,
token: localStorage.getItem('auth_token'), isAuthenticated: false
refreshToken: localStorage.getItem('refresh_token')
}; };
// DOM elements // DOM elements
@ -58,29 +57,15 @@ const presetBtns = document.querySelectorAll('.btn-preset');
// Wrapper for fetch to include Authorization header // Wrapper for fetch to include Authorization header
async function fetchWithAuth(url, options = {}) { async function fetchWithAuth(url, options = {}) {
if (!state.token) { let response = await fetch(url, { ...options, credentials: 'include' });
showLogin();
throw new Error('Not authenticated');
}
const headers = { if (response.status === 401) {
...options.headers,
'Authorization': `Bearer ${state.token}`
};
let response = await fetch(url, { ...options, headers });
if (response.status === 401 && state.refreshToken) {
// Try to refresh token // Try to refresh token
try { try {
const success = await attemptTokenRefresh(); const success = await attemptTokenRefresh();
if (success) { if (success) {
// Retry original request with new token // Retry original request
const newHeaders = { response = await fetch(url, { ...options, credentials: 'include' });
...options.headers,
'Authorization': `Bearer ${state.token}`
};
response = await fetch(url, { ...options, headers: newHeaders });
} }
} catch (error) { } catch (error) {
console.error('Token refresh failed:', error); console.error('Token refresh failed:', error);
@ -100,21 +85,13 @@ async function attemptTokenRefresh() {
const response = await fetch(`${API_URL}/auth/refresh`, { const response = await fetch(`${API_URL}/auth/refresh`, {
method: 'POST', method: 'POST',
headers: { 'Content-Type': 'application/json' }, headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ refresh_token: state.refreshToken }) credentials: 'include',
body: JSON.stringify({ refresh_token: '' })
}); });
if (response.ok) { if (response.ok) {
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; return true;
} }
}
} catch (error) { } catch (error) {
console.error('Error during token refresh:', error); console.error('Error during token refresh:', error);
} }
@ -153,7 +130,7 @@ async function fetchTasks() {
const task = state.tasks.find((t) => t.id === state.selectedTaskId); const task = state.tasks.find((t) => t.id === state.selectedTaskId);
if (task) { if (task) {
if (shouldFollowLatest && task.runs && task.runs.length > 0) { if (shouldFollowLatest && task.runs && task.runs.length > 0) {
state.selectedRunId = task.runs[task.runs.length - 1].id; state.selectedRunId = task.runs[0].id;
} }
renderRunHistory(task); renderRunHistory(task);
@ -177,13 +154,7 @@ async function fetchRecentRuns() {
function renderTaskList() { function renderTaskList() {
const sortedTasks = [...state.tasks].sort((a, b) => { const sortedTasks = [...state.tasks].sort((a, b) => {
const aDate = a.runs && a.runs.length > 0 return new Date(b.created_at) - new Date(a.created_at);
? 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 taskListEl.innerHTML = sortedTasks
@ -223,8 +194,8 @@ function selectTask(id, runId = null) {
if (runId) { if (runId) {
state.selectedRunId = runId; state.selectedRunId = runId;
} else if (task.runs && task.runs.length > 0) { } else if (task.runs && task.runs.length > 0) {
// Default to latest run if not specified // Default to latest run if not specified (index 0 is newest)
state.selectedRunId = task.runs[task.runs.length - 1].id; state.selectedRunId = task.runs[0].id;
} else { } else {
state.selectedRunId = null; state.selectedRunId = null;
} }
@ -495,6 +466,20 @@ 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; let isPolling = false;
async function startAutoRefresh() { async function startAutoRefresh() {
setInterval(async () => { setInterval(async () => {
@ -514,10 +499,12 @@ async function showLogin() {
} }
async function logout() { async function logout() {
state.token = null; try {
state.refreshToken = null; await fetch(`${API_URL}/auth/logout`, { method: 'POST', credentials: 'include' });
localStorage.removeItem('auth_token'); } catch (error) {
localStorage.removeItem('refresh_token'); console.error('Logout failed:', error);
}
state.isAuthenticated = false;
showLogin(); showLogin();
} }
@ -531,16 +518,13 @@ async function handleCallback() {
loginOverlay.classList.add('hidden'); loginOverlay.classList.add('hidden');
try { try {
const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`); const response = await fetch(`${API_URL}/auth/callback?code=${code}&redirect_uri=${encodeURIComponent(AUTH_CONFIG.redirectUri)}`, {
credentials: 'include'
});
const data = await response.json(); const data = await response.json();
if (data.access_token) { if (response.ok) {
state.token = data.access_token; state.isAuthenticated = true;
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'); callbackOverlay.classList.add('hidden');
appEl.classList.remove('hidden'); appEl.classList.remove('hidden');
initializeApp(); initializeApp();
@ -564,16 +548,16 @@ logoutBtn.addEventListener('click', () => {
}); });
async function initializeApp() { async function initializeApp() {
if (!state.token) { const hasSession = await checkSession();
showLogin();
return;
}
if (hasSession) {
appEl.classList.remove('hidden'); appEl.classList.remove('hidden');
loginOverlay.classList.add('hidden'); loginOverlay.classList.add('hidden');
await fetchTasks(); await fetchTasks();
startAutoRefresh(); startAutoRefresh();
} else {
showLogin();
}
} }
// Check for callback on load // Check for callback on load

View file

@ -1,4 +1,5 @@
use chrono::Utc; use chrono::Utc;
use std::time::{Duration, Instant};
use crate::api::{ChatRequest, ChatResponse, Message, Tool}; use crate::api::{ChatRequest, ChatResponse, Message, Tool};
use crate::tools; use crate::tools;
@ -27,7 +28,6 @@ impl Agent {
answer tool once you to give your final answer. current date is {}", answer tool once you to give your final answer. current date is {}",
Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string() Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string()
); );
println!("initial_message: {}", intro);
let messages = vec![ let messages = vec![
Message { Message {
role: "system".to_string(), role: "system".to_string(),
@ -62,15 +62,35 @@ impl Agent {
} }
fn log(&mut self, message: &str) { fn log(&mut self, message: &str) {
println!("{}", message);
self.logs.push_str(message); self.logs.push_str(message);
self.logs.push('\n'); self.logs.push('\n');
} }
pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> { pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
let mut file_written = false; 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);
while !file_written { 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 { let request = ChatRequest {
model: "kimi-k2.5".to_string(), model: "kimi-k2.5".to_string(),
messages: self.messages.clone(), messages: self.messages.clone(),
@ -100,7 +120,12 @@ impl Agent {
} }
let chat_response: ChatResponse = response.json().await?; let chat_response: ChatResponse = response.json().await?;
let assistant_message = chat_response.choices.get(0).unwrap().message.clone(); let assistant_message = chat_response
.choices
.get(0)
.ok_or("Missing assistant response")?
.message
.clone();
self.messages.push(assistant_message.clone()); self.messages.push(assistant_message.clone());
@ -128,16 +153,16 @@ impl Agent {
self.messages.push(tool_message); self.messages.push(tool_message);
if written { if written {
file_written = true; finished = true;
} }
} }
// Continue the loop to send tool results back // Continue the loop to send tool results back
continue; continue;
} }
// No more tool calls from assistant, but we only exit if file was written // No tool calls from assistant, but we only exit if the task was finished
if !file_written { if !finished {
self.log("--- Assistant didn't use write_file yet. Waiting for next turn... ---"); self.log("--- Assistant didn't finish yet. Waiting for next turn... ---");
} }
} }

View file

@ -10,10 +10,17 @@ pub struct Claims {
pub exp: usize, pub exp: usize,
pub iat: usize, pub iat: usize,
pub iss: String, pub iss: String,
pub aud: String, pub aud: Audience,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Audience {
Single(String),
Multiple(Vec<String>),
}
#[derive(Debug, Deserialize, Clone)]
struct Jwk { struct Jwk {
#[serde(rename = "kty")] #[serde(rename = "kty")]
_kty: String, _kty: String,
@ -31,13 +38,17 @@ struct Jwks {
pub struct JwksVerifier { pub struct JwksVerifier {
issuer: String, issuer: String,
audience: String,
jwks_uri: String, jwks_uri: String,
keys: Arc<RwLock<Vec<Jwk>>>, keys: Arc<RwLock<Vec<Jwk>>>,
client: Client, client: Client,
} }
impl JwksVerifier { impl JwksVerifier {
pub async fn new(issuer: String) -> Result<Self, Box<dyn std::error::Error>> { pub async fn new(
issuer: String,
audience: String,
) -> Result<Self, Box<dyn std::error::Error>> {
let client = Client::new(); let client = Client::new();
// Authentik OIDC discovery // Authentik OIDC discovery
let discovery_url = format!( let discovery_url = format!(
@ -53,6 +64,7 @@ impl JwksVerifier {
let verifier = Self { let verifier = Self {
issuer, issuer,
audience,
jwks_uri, jwks_uri,
keys: Arc::new(RwLock::new(Vec::new())), keys: Arc::new(RwLock::new(Vec::new())),
client, client,
@ -73,18 +85,29 @@ impl JwksVerifier {
let header = decode_header(token)?; let header = decode_header(token)?;
let kid = header.kid.ok_or("Missing kid in token header")?; let kid = header.kid.ok_or("Missing kid in token header")?;
let jwk = {
let keys = self.keys.read().await; let keys = self.keys.read().await;
let jwk = keys keys.iter().find(|k| k.kid == kid).cloned()
.iter() };
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) .find(|k| k.kid == kid)
.ok_or("Key not found in JWKS")?; .cloned()
.ok_or("Key not found in JWKS")?
}
};
let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?; let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
let mut validation = Validation::new(Algorithm::RS256); let mut validation = Validation::new(Algorithm::RS256);
validation.set_issuer(&[self.issuer.clone()]); validation.set_issuer(&[self.issuer.clone()]);
// Aud validation might need careful config, usually it's the client_id validation.set_audience(&[self.audience.clone()]);
validation.validate_aud = false; validation.validate_aud = true;
let token_data = decode::<Claims>(token, &decoding_key, &validation)?; let token_data = decode::<Claims>(token, &decoding_key, &validation)?;
Ok(token_data.claims) Ok(token_data.claims)

View file

@ -98,19 +98,28 @@ impl Scheduler {
let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?; let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;
tokio::spawn(async move { tokio::spawn(async move {
let (logs, answer) = match agent.run().await { let (logs, answer, status) = match agent.run().await {
Ok(res) => res, Ok((logs, answer)) => (logs, answer, "completed".to_string()),
Err(e) => (format!("Scheduled run failed: {}", e), None), Err(e) => (
format!("Scheduled run failed: {}", e),
None,
"failed".to_string(),
),
}; };
let run_complete = task_run::ActiveModel { let run_complete = task_run::ActiveModel {
id: Set(run_id), id: Set(run_id),
status: Set("completed".to_string()), status: Set(status),
logs: Set(logs), logs: Set(logs),
answer: Set(answer), answer: Set(answer),
..Default::default() ..Default::default()
}; };
let _ = run_complete.update(&db).await; if let Err(e) = run_complete.update(&db).await {
eprintln!(
"Failed to update scheduled run status for task {}: {}",
task_id, e
);
}
}); });
Ok(()) Ok(())

View file

@ -1,11 +1,13 @@
use axum::{ use axum::{
Json, RequestPartsExt, Router, Json, RequestPartsExt, Router,
extract::{FromRef, FromRequestParts, Path, Query, State}, extract::{FromRef, FromRequestParts, Path, Query, State},
http::{StatusCode, request::Parts}, http::{HeaderValue, StatusCode, request::Parts},
response::IntoResponse,
routing::{get, post}, routing::{get, post},
}; };
use axum_extra::{ use axum_extra::{
TypedHeader, TypedHeader,
extract::cookie::{Cookie, CookieJar, SameSite},
headers::{Authorization, authorization::Bearer}, headers::{Authorization, authorization::Bearer},
}; };
use chrono::Utc; use chrono::Utc;
@ -14,7 +16,7 @@ use sea_orm::{
}; };
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::sync::Arc; use std::sync::Arc;
use tower_http::cors::{Any, CorsLayer}; use tower_http::cors::{AllowOrigin, CorsLayer};
use uuid::Uuid; use uuid::Uuid;
use crate::agent::Agent; use crate::agent::Agent;
@ -97,7 +99,10 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
let authentik_client_secret = let authentik_client_secret =
std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?; std::env::var("AUTHENTIK_CLIENT_SECRET").map_err(|_| "AUTHENTIK_CLIENT_SECRET not set")?;
let verifier = Arc::new(crate::auth::JwksVerifier::new(authentik_issuer.clone()).await?); let verifier = Arc::new(
crate::auth::JwksVerifier::new(authentik_issuer.clone(), authentik_client_id.clone())
.await?,
);
let authenticator = Arc::new( let authenticator = Arc::new(
crate::auth::Authenticator::new( crate::auth::Authenticator::new(
authentik_issuer, authentik_issuer,
@ -116,19 +121,31 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
authenticator, authenticator,
}); });
let cors = CorsLayer::new() let cors = build_cors_layer();
.allow_origin(Any)
.allow_methods(Any)
.allow_headers(Any);
let app = Router::new() let app = Router::new()
.route("/api/tasks", post(create_task).get(list_tasks)) .route("/api/tasks", post(create_task).get(list_tasks))
.route("/api/tasks/:id", get(get_task).put(update_task)) .route("/api/tasks/:id", get(get_task).put(update_task))
.route("/api/tasks/:id/runs", post(rerun_task)) .route("/api/tasks/:id/runs", post(rerun_task))
.route("/api/runs/recent", get(get_recent_runs)) .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/callback", get(auth_callback))
.route("/api/auth/refresh", post(auth_refresh)) .route("/api/auth/refresh", post(auth_refresh))
.route("/api/auth/logout", post(auth_logout))
.layer(cors) .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); .with_state(state);
let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string()); let port = std::env::var("PORT").unwrap_or_else(|_| "3000".to_string());
@ -140,6 +157,44 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
Ok(()) 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<HeaderValue> = 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( async fn list_tasks(
_user: AuthenticatedUser, _user: AuthenticatedUser,
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
@ -249,9 +304,13 @@ async fn execute_agent_run(
) )
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let (logs, answer) = match agent.run().await { let (logs, answer, status) = match agent.run().await {
Ok((logs, answer)) => (logs, answer), Ok((logs, answer)) => (logs, answer, "completed".to_string()),
Err(e) => (format!("Execution failed: {}", e), None), Err(e) => (
format!("Execution failed: {}", e),
None,
"failed".to_string(),
),
}; };
// Update with final logs and status // Update with final logs and status
@ -267,7 +326,7 @@ async fn execute_agent_run(
run.logs = Set(logs.clone()); run.logs = Set(logs.clone());
run.answer = Set(answer.clone()); run.answer = Set(answer.clone());
run.status = Set("completed".to_string()); run.status = Set(status);
run.update(&state.db) run.update(&state.db)
.await .await
@ -319,21 +378,22 @@ where
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> { async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
let app_state = Arc::<AppState>::from_ref(state); let app_state = Arc::<AppState>::from_ref(state);
let TypedHeader(Authorization(bearer)) = parts let token = if let Ok(TypedHeader(Authorization(bearer))) =
.extract::<TypedHeader<Authorization<Bearer>>>() parts.extract::<TypedHeader<Authorization<Bearer>>>().await
.await {
.map_err(|_| { Some(bearer.token().to_string())
( } else {
StatusCode::UNAUTHORIZED, let jar = parts.extract::<CookieJar>().await.unwrap();
"Missing or invalid Authorization header".to_string(), jar.get("access_token")
) .map(|cookie| cookie.value().to_string())
})?; };
let claims = app_state let token = token.ok_or((
.verifier StatusCode::UNAUTHORIZED,
.verify(bearer.token()) "Missing or invalid access token".to_string(),
.await ))?;
.map_err(|e| {
let claims = app_state.verifier.verify(&token).await.map_err(|e| {
( (
StatusCode::UNAUTHORIZED, StatusCode::UNAUTHORIZED,
format!("Token verification failed: {}", e), format!("Token verification failed: {}", e),
@ -352,36 +412,117 @@ pub struct AuthCallbackQuery {
#[derive(Deserialize)] #[derive(Deserialize)]
struct RefreshRequest { struct RefreshRequest {
refresh_token: String, refresh_token: Option<String>,
} }
async fn auth_refresh( async fn auth_refresh(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
jar: CookieJar,
Json(payload): Json<RefreshRequest>, Json(payload): Json<RefreshRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> { ) -> Result<impl IntoResponse, (StatusCode, String)> {
state 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 .authenticator
.refresh_token(payload.refresh_token) .refresh_token(refresh_token)
.await .await
.map(Json) .map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))?;
.map_err(|e| (StatusCode::UNAUTHORIZED, e.to_string()))
let jar = update_auth_cookies(jar, &data);
Ok((jar, Json(data)))
} }
async fn auth_callback( async fn auth_callback(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
jar: CookieJar,
Query(query): Query<AuthCallbackQuery>, Query(query): Query<AuthCallbackQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> { ) -> Result<impl IntoResponse, (StatusCode, String)> {
state let data = state
.authenticator .authenticator
.exchange_code(query.code, query.redirect_uri) .exchange_code(query.code, query.redirect_uri)
.await .await
.map(Json)
.map_err(|e| { .map_err(|e| {
( (
StatusCode::INTERNAL_SERVER_ERROR, StatusCode::INTERNAL_SERVER_ERROR,
format!("Token exchange failed: {}", e), 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<serde_json::Value> {
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( async fn get_task(

View file

@ -51,11 +51,6 @@ pub async fn handle_tool_call(
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?; let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
let query = args.get("query").ok_or("Missing query argument")?; 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 { let search_result = if let Some(key) = tavily_api_key {
match api::perform_search(query, key).await { match api::perform_search(query, key).await {
Ok(results) => results, Ok(results) => results,
@ -69,8 +64,6 @@ pub async fn handle_tool_call(
let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?; let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
let result = args.get("result").ok_or("Missing result argument")?; let result = args.get("result").ok_or("Missing result argument")?;
println!("--- Finishing task: {}", result);
answer = Some(result.clone()); answer = Some(result.clone());
(result.clone(), true) (result.clone(), true)
} else { } else {