fix

Cargo.lock

@@ -217,6 +217,7 @@ dependencies = [
  "axum",
  "axum-core",
  "bytes",
+ "cookie",
  "fastrand",
  "futures-util",
  "headers",
@@ -325,6 +326,7 @@ dependencies = [
  "axum-extra",
  "base64",
  "chrono",
+ "cookie",
  "dashmap",
  "jsonwebtoken",
  "migration",
@@ -486,6 +488,17 @@ 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"

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"] }
+tower-http = { version = "0.5", features = ["cors", "set-header", "limit"] }
 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,4 +22,5 @@ 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"] }
+axum-extra = { version = "0.9", features = ["typed-header", "cookie"] }
+cookie = "0.18"

frontend/src/main.js

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

src/agent.rs

@@ -1,4 +1,5 @@
 use chrono::Utc;
+use std::time::{Duration, Instant};
 
 use crate::api::{ChatRequest, ChatResponse, Message, Tool};
 use crate::tools;
@@ -27,7 +28,6 @@ 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,15 +62,35 @@ 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<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 {
                 model: "kimi-k2.5".to_string(),
                 messages: self.messages.clone(),
@@ -100,7 +120,12 @@ impl Agent {
             }
 
             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());
 
@@ -128,16 +153,16 @@ impl Agent {
 
                     self.messages.push(tool_message);
                     if written {
-                        file_written = true;
+                        finished = true;
                     }
                 }
                 // Continue the loop to send tool results back
                 continue;
             }
 
-            // 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... ---");
+            // 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... ---");
             }
         }
 

src/auth.rs

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

src/scheduler.rs

@@ -98,19 +98,28 @@ impl Scheduler {
         let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;
 
         tokio::spawn(async move {
-            let (logs, answer) = match agent.run().await {
-                Ok(res) => res,
-                Err(e) => (format!("Scheduled run failed: {}", e), None),
+            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 {
                 id: Set(run_id),
-                status: Set("completed".to_string()),
+                status: Set(status),
                 logs: Set(logs),
                 answer: Set(answer),
                 ..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(())

src/server.rs

@@ -1,11 +1,13 @@
 use axum::{
     Json, RequestPartsExt, Router,
     extract::{FromRef, FromRequestParts, Path, Query, State},
-    http::{StatusCode, request::Parts},
+    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;
@@ -14,7 +16,7 @@ use sea_orm::{
 };
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
-use tower_http::cors::{Any, CorsLayer};
+use tower_http::cors::{AllowOrigin, CorsLayer};
 use uuid::Uuid;
 
 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 =
         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(
         crate::auth::Authenticator::new(
             authentik_issuer,
@@ -116,19 +121,31 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
         authenticator,
     });
 
-    let cors = CorsLayer::new()
-        .allow_origin(Any)
-        .allow_methods(Any)
-        .allow_headers(Any);
+    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());
@@ -140,6 +157,44 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
     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(
     _user: AuthenticatedUser,
     State(state): State<Arc<AppState>>,
@@ -249,9 +304,13 @@ async fn execute_agent_run(
     )
     .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    let (logs, answer) = match agent.run().await {
-        Ok((logs, answer)) => (logs, answer),
-        Err(e) => (format!("Execution failed: {}", e), None),
+    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
@@ -267,7 +326,7 @@ async fn execute_agent_run(
 
     run.logs = Set(logs.clone());
     run.answer = Set(answer.clone());
-    run.status = Set("completed".to_string());
+    run.status = Set(status);
 
     run.update(&state.db)
         .await
@@ -319,26 +378,27 @@ where
     async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
         let app_state = Arc::<AppState>::from_ref(state);
 
-        let TypedHeader(Authorization(bearer)) = parts
-            .extract::<TypedHeader<Authorization<Bearer>>>()
-            .await
-            .map_err(|_| {
-                (
-                    StatusCode::UNAUTHORIZED,
-                    "Missing or invalid Authorization header".to_string(),
-                )
-            })?;
+        let token = if let Ok(TypedHeader(Authorization(bearer))) =
+            parts.extract::<TypedHeader<Authorization<Bearer>>>().await
+        {
+            Some(bearer.token().to_string())
+        } else {
+            let jar = parts.extract::<CookieJar>().await.unwrap();
+            jar.get("access_token")
+                .map(|cookie| cookie.value().to_string())
+        };
 
-        let claims = app_state
-            .verifier
-            .verify(bearer.token())
-            .await
-            .map_err(|e| {
-                (
-                    StatusCode::UNAUTHORIZED,
-                    format!("Token verification failed: {}", e),
-                )
-            })?;
+        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))
     }
@@ -352,36 +412,117 @@ pub struct AuthCallbackQuery {
 
 #[derive(Deserialize)]
 struct RefreshRequest {
-    refresh_token: String,
+    refresh_token: Option<String>,
 }
 
 async fn auth_refresh(
     State(state): State<Arc<AppState>>,
+    jar: CookieJar,
     Json(payload): Json<RefreshRequest>,
-) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
-    state
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    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(payload.refresh_token)
+        .refresh_token(refresh_token)
         .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(
     State(state): State<Arc<AppState>>,
+    jar: CookieJar,
     Query(query): Query<AuthCallbackQuery>,
-) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
-    state
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let data = 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<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(

src/tools.rs

@@ -51,11 +51,6 @@ pub async fn handle_tool_call(
         let args: HashMap<String, String> = 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,
@@ -69,8 +64,6 @@ pub async fn handle_tool_call(
         let args: HashMap<String, String> = 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 {