//! Chat handlers: send, approval, auth, SSE events, WebSocket, history, threads. use std::sync::Arc; use axum::{ Json, extract::{Query, State, WebSocketUpgrade}, http::StatusCode, response::IntoResponse, }; use serde::Deserialize; use uuid::Uuid; use crate::channels::IncomingMessage; use crate::channels::web::server::GatewayState; use crate::channels::web::types::*; pub async fn chat_send_handler( State(state): State>, Json(req): Json, ) -> Result<(StatusCode, Json), (StatusCode, String)> { if !state.chat_rate_limiter.check() { return Err(( StatusCode::TOO_MANY_REQUESTS, "Rate limit exceeded. Try again shortly.".to_string(), )); } let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content); if let Some(ref thread_id) = req.thread_id { msg = msg.with_thread(thread_id); msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id})); } let msg_id = msg.id; let tx_guard = state.msg_tx.read().await; let tx = tx_guard.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Channel not started".to_string(), ))?; tx.send(msg).await.map_err(|_| { ( StatusCode::INTERNAL_SERVER_ERROR, "Channel closed".to_string(), ) })?; Ok(( StatusCode::ACCEPTED, Json(SendMessageResponse { message_id: msg_id, status: "accepted", }), )) } pub async fn chat_approval_handler( State(state): State>, Json(req): Json, ) -> Result<(StatusCode, Json), (StatusCode, String)> { let (approved, always) = match req.action.as_str() { "approve" => (true, false), "always" => (true, true), "deny" => (false, false), other => { return Err(( StatusCode::BAD_REQUEST, format!("Unknown action: {}", other), )); } }; let request_id = Uuid::parse_str(&req.request_id).map_err(|_| { ( StatusCode::BAD_REQUEST, "Invalid request_id (expected UUID)".to_string(), ) })?; // Build a structured ExecApproval submission as JSON, sent through the // existing message pipeline so the agent loop picks it up. let approval = crate::agent::submission::Submission::ExecApproval { request_id, approved, always, }; let content = serde_json::to_string(&approval).map_err(|e| { ( StatusCode::INTERNAL_SERVER_ERROR, format!("Failed to serialize approval: {}", e), ) })?; let mut msg = IncomingMessage::new("gateway", &state.user_id, content); if let Some(ref thread_id) = req.thread_id { msg = msg.with_thread(thread_id); } let msg_id = msg.id; let tx_guard = state.msg_tx.read().await; let tx = tx_guard.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Channel not started".to_string(), ))?; tx.send(msg).await.map_err(|_| { ( StatusCode::INTERNAL_SERVER_ERROR, "Channel closed".to_string(), ) })?; Ok(( StatusCode::ACCEPTED, Json(SendMessageResponse { message_id: msg_id, status: "accepted", }), )) } /// Submit an auth token directly to the extension manager, bypassing the message pipeline. /// /// The token never touches the LLM, chat history, or SSE stream. pub async fn chat_auth_token_handler( State(state): State>, Json(req): Json, ) -> Result, (StatusCode, String)> { let ext_mgr = state.extension_manager.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Extension manager not available".to_string(), ))?; let result = ext_mgr .auth(&req.extension_name, Some(&req.token)) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; if result.status == "authenticated" { // Auto-activate so tools are available immediately let msg = match ext_mgr.activate(&req.extension_name).await { Ok(r) => format!( "{} authenticated ({} tools loaded)", req.extension_name, r.tools_loaded.len() ), Err(e) => format!( "{} authenticated but activation failed: {}", req.extension_name, e ), }; // Clear auth mode on the active thread clear_auth_mode(&state).await; state.sse.broadcast(SseEvent::AuthCompleted { extension_name: req.extension_name, success: true, message: msg.clone(), }); Ok(Json(ActionResponse::ok(msg))) } else { // Re-emit auth_required for retry state.sse.broadcast(SseEvent::AuthRequired { extension_name: req.extension_name.clone(), instructions: result.instructions.clone(), auth_url: result.auth_url.clone(), setup_url: result.setup_url.clone(), }); Ok(Json(ActionResponse::fail( result .instructions .unwrap_or_else(|| "Invalid token".to_string()), ))) } } /// Cancel an in-progress auth flow. pub async fn chat_auth_cancel_handler( State(state): State>, Json(_req): Json, ) -> Result, (StatusCode, String)> { clear_auth_mode(&state).await; Ok(Json(ActionResponse::ok("Auth cancelled"))) } /// Clear pending auth mode on the active thread. pub async fn clear_auth_mode(state: &GatewayState) { if let Some(ref sm) = state.session_manager { let session = sm.get_or_create_session(&state.user_id).await; let mut sess = session.lock().await; if let Some(thread_id) = sess.active_thread && let Some(thread) = sess.threads.get_mut(&thread_id) { thread.pending_auth = None; } } } pub async fn chat_events_handler( State(state): State>, ) -> Result { state.sse.subscribe().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Too many connections".to_string(), )) } pub async fn chat_ws_handler( headers: axum::http::HeaderMap, ws: WebSocketUpgrade, State(state): State>, ) -> Result { // Validate Origin header to prevent cross-site WebSocket hijacking. let origin = headers .get("origin") .and_then(|v| v.to_str().ok()) .ok_or_else(|| { ( StatusCode::FORBIDDEN, "WebSocket Origin header required".to_string(), ) })?; let host = origin .strip_prefix("http://") .or_else(|| origin.strip_prefix("https://")) .and_then(|rest| rest.split(':').next()?.split('/').next()) .unwrap_or(""); let is_local = matches!(host, "localhost" | "127.0.0.1" | "[::1]"); if !is_local { return Err(( StatusCode::FORBIDDEN, "WebSocket origin not allowed".to_string(), )); } Ok(ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state))) } #[derive(Deserialize)] pub struct HistoryQuery { pub thread_id: Option, pub limit: Option, pub before: Option, } pub async fn chat_history_handler( State(state): State>, Query(query): Query, ) -> Result, (StatusCode, String)> { let session_manager = state.session_manager.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Session manager not available".to_string(), ))?; let session = session_manager.get_or_create_session(&state.user_id).await; let sess = session.lock().await; let limit = query.limit.unwrap_or(50); let before_cursor = query .before .as_deref() .map(|s| { chrono::DateTime::parse_from_rfc3339(s) .map(|dt| dt.with_timezone(&chrono::Utc)) .map_err(|_| { ( StatusCode::BAD_REQUEST, "Invalid 'before' timestamp".to_string(), ) }) }) .transpose()?; // Find the thread let thread_id = if let Some(ref tid) = query.thread_id { Uuid::parse_str(tid) .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid thread_id".to_string()))? } else { sess.active_thread .ok_or((StatusCode::NOT_FOUND, "No active thread".to_string()))? }; // Verify the thread belongs to the authenticated user before returning any data. if query.thread_id.is_some() && let Some(ref store) = state.store { let owned = store .conversation_belongs_to_user(thread_id, &state.user_id) .await .unwrap_or(false); if !owned && !sess.threads.contains_key(&thread_id) { return Err((StatusCode::NOT_FOUND, "Thread not found".to_string())); } } // For paginated requests (before cursor set), always go to DB if before_cursor.is_some() && let Some(ref store) = state.store { let (messages, has_more) = store .list_conversation_messages_paginated(thread_id, before_cursor, limit as i64) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339()); let turns = build_turns_from_db_messages(&messages); return Ok(Json(HistoryResponse { thread_id, turns, has_more, oldest_timestamp, })); } // Try in-memory first (freshest data for active threads) if let Some(thread) = sess.threads.get(&thread_id) && !thread.turns.is_empty() { let turns: Vec = thread .turns .iter() .map(|t| TurnInfo { turn_number: t.turn_number, user_input: t.user_input.clone(), response: t.response.clone(), state: format!("{:?}", t.state), started_at: t.started_at.to_rfc3339(), completed_at: t.completed_at.map(|dt| dt.to_rfc3339()), tool_calls: t .tool_calls .iter() .map(|tc| ToolCallInfo { name: tc.name.clone(), has_result: tc.result.is_some(), has_error: tc.error.is_some(), }) .collect(), }) .collect(); return Ok(Json(HistoryResponse { thread_id, turns, has_more: false, oldest_timestamp: None, })); } // Fall back to DB for historical threads not in memory (paginated) if let Some(ref store) = state.store { let (messages, has_more) = store .list_conversation_messages_paginated(thread_id, None, limit as i64) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; if !messages.is_empty() { let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339()); let turns = build_turns_from_db_messages(&messages); return Ok(Json(HistoryResponse { thread_id, turns, has_more, oldest_timestamp, })); } } // Empty thread (just created, no messages yet) Ok(Json(HistoryResponse { thread_id, turns: Vec::new(), has_more: false, oldest_timestamp: None, })) } /// Build TurnInfo pairs from flat DB messages (alternating user/assistant). pub fn build_turns_from_db_messages( messages: &[crate::history::ConversationMessage], ) -> Vec { let mut turns = Vec::new(); let mut turn_number = 0; let mut iter = messages.iter().peekable(); while let Some(msg) = iter.next() { if msg.role == "user" { let mut turn = TurnInfo { turn_number, user_input: msg.content.clone(), response: None, state: "Completed".to_string(), started_at: msg.created_at.to_rfc3339(), completed_at: None, tool_calls: Vec::new(), }; // Check if next message is an assistant response if let Some(next) = iter.peek() && next.role == "assistant" { let assistant_msg = iter.next().expect("peeked"); turn.response = Some(assistant_msg.content.clone()); turn.completed_at = Some(assistant_msg.created_at.to_rfc3339()); } // Incomplete turn (user message without response) if turn.response.is_none() { turn.state = "Failed".to_string(); } turns.push(turn); turn_number += 1; } } turns } pub async fn chat_threads_handler( State(state): State>, ) -> Result, (StatusCode, String)> { let session_manager = state.session_manager.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Session manager not available".to_string(), ))?; let session = session_manager.get_or_create_session(&state.user_id).await; let sess = session.lock().await; // Try DB first for persistent thread list if let Some(ref store) = state.store { // Auto-create assistant thread if it doesn't exist let assistant_id = store .get_or_create_assistant_conversation(&state.user_id, "gateway") .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; if let Ok(summaries) = store .list_conversations_with_preview(&state.user_id, "gateway", 50) .await { let mut assistant_thread = None; let mut threads = Vec::new(); for s in &summaries { let info = ThreadInfo { id: s.id, state: "Idle".to_string(), turn_count: (s.message_count / 2).max(0) as usize, created_at: s.started_at.to_rfc3339(), updated_at: s.last_activity.to_rfc3339(), title: s.title.clone(), thread_type: s.thread_type.clone(), }; if s.id == assistant_id { assistant_thread = Some(info); } else { threads.push(info); } } // If assistant wasn't in the list (0 messages), synthesize it if assistant_thread.is_none() { assistant_thread = Some(ThreadInfo { id: assistant_id, state: "Idle".to_string(), turn_count: 0, created_at: chrono::Utc::now().to_rfc3339(), updated_at: chrono::Utc::now().to_rfc3339(), title: None, thread_type: Some("assistant".to_string()), }); } return Ok(Json(ThreadListResponse { assistant_thread, threads, active_thread: sess.active_thread, })); } } // Fallback: in-memory only (no assistant thread without DB) let threads: Vec = sess .threads .values() .map(|t| ThreadInfo { id: t.id, state: format!("{:?}", t.state), turn_count: t.turns.len(), created_at: t.created_at.to_rfc3339(), updated_at: t.updated_at.to_rfc3339(), title: None, thread_type: None, }) .collect(); Ok(Json(ThreadListResponse { assistant_thread: None, threads, active_thread: sess.active_thread, })) } pub async fn chat_new_thread_handler( State(state): State>, ) -> Result, (StatusCode, String)> { let session_manager = state.session_manager.as_ref().ok_or(( StatusCode::SERVICE_UNAVAILABLE, "Session manager not available".to_string(), ))?; let session = session_manager.get_or_create_session(&state.user_id).await; let mut sess = session.lock().await; let thread = sess.create_thread(); let thread_id = thread.id; let info = ThreadInfo { id: thread.id, state: format!("{:?}", thread.state), turn_count: thread.turns.len(), created_at: thread.created_at.to_rfc3339(), updated_at: thread.updated_at.to_rfc3339(), title: None, thread_type: Some("thread".to_string()), }; // Persist the empty conversation row with thread_type metadata if let Some(ref store) = state.store { let store = Arc::clone(store); let user_id = state.user_id.clone(); tokio::spawn(async move { if let Err(e) = store .ensure_conversation(thread_id, "gateway", &user_id, None) .await { tracing::warn!("Failed to persist new thread: {}", e); } let metadata_val = serde_json::json!("thread"); if let Err(e) = store .update_conversation_metadata_field(thread_id, "thread_type", &metadata_val) .await { tracing::warn!("Failed to set thread_type metadata: {}", e); } }); } Ok(Json(info)) } #[cfg(test)] mod tests { use super::*; #[test] fn test_build_turns_from_db_messages_complete() { let now = chrono::Utc::now(); let messages = vec![ crate::history::ConversationMessage { id: Uuid::new_v4(), role: "user".to_string(), content: "Hello".to_string(), created_at: now, }, crate::history::ConversationMessage { id: Uuid::new_v4(), role: "assistant".to_string(), content: "Hi there!".to_string(), created_at: now + chrono::TimeDelta::seconds(1), }, crate::history::ConversationMessage { id: Uuid::new_v4(), role: "user".to_string(), content: "How are you?".to_string(), created_at: now + chrono::TimeDelta::seconds(2), }, crate::history::ConversationMessage { id: Uuid::new_v4(), role: "assistant".to_string(), content: "Doing well!".to_string(), created_at: now + chrono::TimeDelta::seconds(3), }, ]; let turns = build_turns_from_db_messages(&messages); assert_eq!(turns.len(), 2); assert_eq!(turns[0].user_input, "Hello"); assert_eq!(turns[0].response.as_deref(), Some("Hi there!")); assert_eq!(turns[0].state, "Completed"); assert_eq!(turns[1].user_input, "How are you?"); assert_eq!(turns[1].response.as_deref(), Some("Doing well!")); } #[test] fn test_build_turns_from_db_messages_incomplete_last() { let now = chrono::Utc::now(); let messages = vec![ crate::history::ConversationMessage { id: Uuid::new_v4(), role: "user".to_string(), content: "Hello".to_string(), created_at: now, }, crate::history::ConversationMessage { id: Uuid::new_v4(), role: "assistant".to_string(), content: "Hi!".to_string(), created_at: now + chrono::TimeDelta::seconds(1), }, crate::history::ConversationMessage { id: Uuid::new_v4(), role: "user".to_string(), content: "Lost message".to_string(), created_at: now + chrono::TimeDelta::seconds(2), }, ]; let turns = build_turns_from_db_messages(&messages); assert_eq!(turns.len(), 2); assert_eq!(turns[1].user_input, "Lost message"); assert!(turns[1].response.is_none()); assert_eq!(turns[1].state, "Failed"); } }