From 80f404c9bdb04fec339b0c4b7ec6db21eaea454f Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Sat, 27 Jun 2026 15:56:48 +0000 Subject: [PATCH 1/7] Add agent feature flag key --- src/os_flags.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/os_flags.rs b/src/os_flags.rs index 7b1aa99f..80f4df54 100644 --- a/src/os_flags.rs +++ b/src/os_flags.rs @@ -11,6 +11,9 @@ const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const USER_FLAGS_CACHE_TTL: Duration = Duration::from_secs(10 * 60); +#[allow(dead_code)] +pub const AGENT_FEATURE_FLAG_KEY: &str = "agent"; + #[derive(Debug, thiserror::Error)] pub enum OsFlagsError { #[error("Invalid base URL: {0}")] From 855429808341c83766180b9fd8801a244664f64a Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:13:04 -0500 Subject: [PATCH 2/7] feat(rag): add encrypted embeddings foundation Introduce the encrypted embedding storage and retrieval primitives on top of the existing schema so later agent features can build on a stable RAG base. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/aead_db_tamper_tests.rs | 3 + src/main.rs | 9 +- src/models/mod.rs | 1 + src/models/user_embeddings.rs | 85 +++ src/rag.rs | 1126 +++++++++++++++++++++++++++++++++ src/web/mod.rs | 3 + src/web/openai.rs | 217 ++++--- src/web/rag.rs | 216 +++++++ 8 files changed, 1588 insertions(+), 72 deletions(-) create mode 100644 src/models/user_embeddings.rs create mode 100644 src/rag.rs create mode 100644 src/web/rag.rs diff --git a/src/aead_db_tamper_tests.rs b/src/aead_db_tamper_tests.rs index dfc5c7dc..2797adc0 100644 --- a/src/aead_db_tamper_tests.rs +++ b/src/aead_db_tamper_tests.rs @@ -1200,7 +1200,9 @@ fn insert_response_storage_stack_for_user(app_state: &AppState, user_id: Uuid) { response_id: Some(response.id), user_id, content_enc: vec![7, 8, 9], + attachment_text_enc: None, prompt_tokens: 3, + assistant_reaction: None, } .insert(conn) .expect("test user message should insert"); @@ -1215,6 +1217,7 @@ fn insert_response_storage_stack_for_user(app_state: &AppState, user_id: Uuid) { status: "completed".to_string(), finish_reason: Some("stop".to_string()), created_at: Utc::now(), + user_reaction: None, } .insert(conn) .expect("test assistant message should insert"); diff --git a/src/main.rs b/src/main.rs index 96fb4352..b28f15d9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ conversation_projects_routes, conversations_routes, health_routes_with_state, - instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, + instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -101,6 +101,7 @@ mod proxy_config; #[cfg(test)] mod security_invariants; mod seed_wrapping; +mod rag; mod sqs; mod tokens; mod web; @@ -465,6 +466,7 @@ pub struct AppState { enclave_key: Vec, proxy_router: Arc, provider_router: Arc, + rag_cache: Arc>, resend_api_key: Option, ephemeral_keys: Arc>>, session_states: Arc>>, @@ -763,6 +765,7 @@ impl AppStateBuilder { enclave_key, proxy_router, provider_router, + rag_cache: Arc::new(tokio::sync::Mutex::new(rag::RagCache::default())), resend_api_key: self.resend_api_key, ephemeral_keys: Arc::new(RwLock::new(HashMap::new())), session_states: Arc::new(tokio::sync::RwLock::new(HashMap::new())), @@ -3206,6 +3209,10 @@ async fn main() -> Result<(), Error> { instructions_routes(app_state.clone()) .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), ) + .merge( + rag_routes(app_state.clone()) + .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), + ) .merge(attestation_routes::router(app_state.clone())) .merge(oauth_routes(app_state.clone())) .merge(platform_login_routes(app_state.clone())) diff --git a/src/models/mod.rs b/src/models/mod.rs index 396d4da3..8d92d1a4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -18,6 +18,7 @@ pub mod responses; pub mod schema; pub mod token_usage; pub mod user_api_keys; +pub mod user_embeddings; pub mod user_kv; pub mod user_seed_wrappings; pub mod users; diff --git a/src/models/user_embeddings.rs b/src/models/user_embeddings.rs new file mode 100644 index 00000000..804cb470 --- /dev/null +++ b/src/models/user_embeddings.rs @@ -0,0 +1,85 @@ +use crate::models::schema::user_embeddings; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +#[derive(Error, Debug)] +pub enum UserEmbeddingError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = user_embeddings)] +pub struct UserEmbedding { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub source_type: String, + pub user_message_id: Option, + pub assistant_message_id: Option, + pub conversation_id: Option, + pub vector_enc: Vec, + pub embedding_model: String, + pub vector_dim: i32, + pub content_enc: Vec, + pub metadata_enc: Option>, + pub tags_enc: Vec>, + pub token_count: i32, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = user_embeddings)] +pub struct NewUserEmbedding { + pub uuid: Uuid, + pub user_id: Uuid, + pub source_type: String, + pub user_message_id: Option, + pub assistant_message_id: Option, + pub conversation_id: Option, + pub vector_enc: Vec, + pub embedding_model: String, + pub vector_dim: i32, + pub content_enc: Vec, + pub metadata_enc: Option>, + pub tags_enc: Vec>, + pub token_count: i32, +} + +impl NewUserEmbedding { + pub fn insert(&self, conn: &mut PgConnection) -> Result { + diesel::insert_into(user_embeddings::table) + .values(self) + .get_result::(conn) + .map_err(UserEmbeddingError::DatabaseError) + } +} + +impl UserEmbedding { + pub fn delete_all_for_user( + conn: &mut PgConnection, + user_uuid: Uuid, + ) -> Result { + diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_uuid))) + .execute(conn) + .map_err(UserEmbeddingError::DatabaseError) + } + + pub fn delete_by_uuid_for_user( + conn: &mut PgConnection, + user_uuid: Uuid, + embedding_uuid: Uuid, + ) -> Result { + diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user_uuid)) + .filter(user_embeddings::uuid.eq(embedding_uuid)), + ) + .execute(conn) + .map_err(UserEmbeddingError::DatabaseError) + } +} diff --git a/src/rag.rs b/src/rag.rs new file mode 100644 index 00000000..a0fea750 --- /dev/null +++ b/src/rag.rs @@ -0,0 +1,1126 @@ +use crate::encrypt::{decrypt_with_key, encrypt_key_deterministic, encrypt_with_key}; +use crate::models::schema::user_embeddings; +use crate::models::user_embeddings::NewUserEmbedding; +use crate::models::users::User; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, AppState}; +use base64::{engine::general_purpose::STANDARD as B64_STANDARD, Engine as _}; +use diesel::prelude::*; +use secp256k1::SecretKey; +use serde::Serialize; +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tracing::{debug, error}; +use uuid::Uuid; + +pub const DEFAULT_EMBEDDING_MODEL: &str = "nomic-embed-text"; +pub const DEFAULT_EMBEDDING_DIM: i32 = 768; + +const CACHE_MAX_USERS: usize = 100; +const CACHE_TTL: Duration = Duration::from_secs(5 * 60); + +const DB_SCAN_BATCH_SIZE: i64 = 1000; + +#[allow(dead_code)] +pub const SOURCE_TYPE_MESSAGE: &str = "message"; +pub const SOURCE_TYPE_ARCHIVAL: &str = "archival"; + +#[derive(Debug, Clone, Serialize)] +pub struct RagSearchResult { + pub content: String, + pub score: f32, + pub token_count: i32, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RagEmbeddingsStatus { + pub total_embeddings: i64, + pub by_model: HashMap, + pub stale_count: i64, +} + +#[derive(Debug, Clone)] +pub struct CachedEmbedding { + pub source_type: String, + pub conversation_id: Option, + pub vector: Vec, + pub content_enc: Vec, + pub token_count: i32, +} + +#[derive(Debug, Clone)] +struct CacheEntry { + loaded_at: Instant, + embeddings: Arc>, +} + +#[derive(Debug)] +pub struct RagCache { + max_users: usize, + ttl: Duration, + entries: HashMap, + lru: VecDeque, +} + +impl Default for RagCache { + fn default() -> Self { + Self::new(CACHE_MAX_USERS, CACHE_TTL) + } +} + +impl RagCache { + pub fn new(max_users: usize, ttl: Duration) -> Self { + Self { + max_users, + ttl, + entries: HashMap::new(), + lru: VecDeque::new(), + } + } + + pub fn evict_user(&mut self, user_id: Uuid) { + self.entries.remove(&user_id); + self.lru.retain(|u| *u != user_id); + } + + pub fn get(&mut self, user_id: Uuid) -> Option>> { + self.evict_expired(); + + let (loaded_at, embeddings) = { + let entry = self.entries.get(&user_id)?; + (entry.loaded_at, entry.embeddings.clone()) + }; + + if loaded_at.elapsed() > self.ttl { + self.evict_user(user_id); + return None; + } + + self.touch(user_id); + Some(embeddings) + } + + pub fn put(&mut self, user_id: Uuid, embeddings: Arc>) { + self.entries.insert( + user_id, + CacheEntry { + loaded_at: Instant::now(), + embeddings, + }, + ); + self.touch(user_id); + + while self.entries.len() > self.max_users { + if let Some(lru_user) = self.lru.pop_back() { + self.entries.remove(&lru_user); + } else { + break; + } + } + } + + fn touch(&mut self, user_id: Uuid) { + self.lru.retain(|u| *u != user_id); + self.lru.push_front(user_id); + } + + fn evict_expired(&mut self) { + let ttl = self.ttl; + let expired: Vec = self + .entries + .iter() + .filter_map(|(user_id, entry)| { + if entry.loaded_at.elapsed() > ttl { + Some(*user_id) + } else { + None + } + }) + .collect(); + + for user_id in expired { + self.evict_user(user_id); + } + } +} + +pub fn serialize_f32_le(values: &[f32]) -> Vec { + let mut out = Vec::with_capacity(values.len() * 4); + for v in values { + out.extend_from_slice(&v.to_le_bytes()); + } + out +} + +pub fn deserialize_f32_le(bytes: &[u8]) -> Result, ApiError> { + if !bytes.len().is_multiple_of(4) { + return Err(ApiError::BadRequest); + } + + let mut out = Vec::with_capacity(bytes.len() / 4); + for chunk in bytes.chunks_exact(4) { + let arr: [u8; 4] = chunk.try_into().map_err(|_| ApiError::BadRequest)?; + out.push(f32::from_le_bytes(arr)); + } + Ok(out) +} + +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Result { + if a.len() != b.len() { + return Err(ApiError::BadRequest); + } + + let mut dot = 0.0f32; + let mut norm_a = 0.0f32; + let mut norm_b = 0.0f32; + + for (x, y) in a.iter().zip(b.iter()) { + dot += x * y; + norm_a += x * x; + norm_b += y * y; + } + + if norm_a == 0.0 || norm_b == 0.0 { + return Ok(0.0); + } + + Ok(dot / (norm_a.sqrt() * norm_b.sqrt())) +} + +#[derive(Debug, Clone)] +struct HeapItem { + score: f32, + token_count: i32, + content_enc: Vec, +} + +impl Eq for HeapItem {} + +impl PartialEq for HeapItem { + fn eq(&self, other: &Self) -> bool { + self.score.to_bits() == other.score.to_bits() && self.token_count == other.token_count + } +} + +impl Ord for HeapItem { + fn cmp(&self, other: &Self) -> Ordering { + self.score + .total_cmp(&other.score) + .then_with(|| other.token_count.cmp(&self.token_count)) + } +} + +impl PartialOrd for HeapItem { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +fn top_k_candidates( + query: &[f32], + embeddings: &[CachedEmbedding], + top_k: usize, + source_types: Option<&[String]>, + conversation_id: Option, +) -> Result, ApiError> { + let allowed_source_types: Option> = source_types.map(|v| { + v.iter() + .map(|s| s.as_str()) + .collect::>() + }); + + let mut heap: BinaryHeap> = BinaryHeap::new(); + + for e in embeddings { + if let Some(allowed) = &allowed_source_types { + if !allowed.contains(e.source_type.as_str()) { + continue; + } + } + + if let Some(conv_id) = conversation_id { + if e.conversation_id != Some(conv_id) { + continue; + } + } + + if e.vector.len() != query.len() { + continue; + } + + let score = cosine_similarity(query, &e.vector)?; + let item = HeapItem { + score, + token_count: e.token_count, + content_enc: e.content_enc.clone(), + }; + + if heap.len() < top_k { + heap.push(std::cmp::Reverse(item)); + continue; + } + + if let Some(std::cmp::Reverse(min)) = heap.peek() { + if item.cmp(min) == Ordering::Greater { + heap.pop(); + heap.push(std::cmp::Reverse(item)); + } + } + } + + let mut out: Vec = heap.into_iter().map(|r| r.0).collect(); + out.sort_by(|a, b| b.cmp(a)); + Ok(out) +} + +fn apply_token_budget(results: Vec, budget: i32) -> Vec { + let mut total: i32 = 0; + let mut limited: Vec = Vec::new(); + for r in results { + if total + r.token_count > budget { + break; + } + total += r.token_count; + limited.push(r); + } + limited +} + +async fn embed_text_via_tinfoil( + state: &Arc, + user: &User, + auth_method: AuthMethod, + text: &str, +) -> Result<(Vec, i32), ApiError> { + crate::web::get_embedding_vector( + state, + user, + auth_method, + DEFAULT_EMBEDDING_MODEL, + text, + Some(DEFAULT_EMBEDDING_DIM), + ) + .await +} + +fn normalize_tags<'a>(tags: impl Iterator) -> Vec { + let mut seen: HashSet = HashSet::new(); + let mut out: Vec = Vec::new(); + + for tag in tags { + let normalized = tag.trim().to_lowercase(); + if normalized.is_empty() { + continue; + } + + if seen.insert(normalized.clone()) { + out.push(normalized); + } + } + + out +} + +fn encrypt_tags_b64(user_key: &SecretKey, tags: &[String]) -> Vec> { + tags.iter() + .map(|tag| { + let ciphertext = encrypt_key_deterministic(user_key, tag.as_bytes()); + Some(B64_STANDARD.encode(ciphertext)) + }) + .collect() +} + +fn extract_tags_from_metadata(metadata: Option<&serde_json::Value>) -> Vec { + let Some(metadata) = metadata else { + return Vec::new(); + }; + + let Some(tags) = metadata.get("tags") else { + return Vec::new(); + }; + + match tags { + serde_json::Value::Array(arr) => normalize_tags(arr.iter().filter_map(|v| v.as_str())), + serde_json::Value::String(s) => normalize_tags(s.split(',')), + _ => Vec::new(), + } +} + +pub async fn insert_archival_embedding( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + text: &str, + metadata: Option<&serde_json::Value>, +) -> Result { + let user_id = user.uuid; + let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; + + let vector_bytes = serialize_f32_le(&vector); + let vector_enc = encrypt_with_key(user_key, &vector_bytes).await; + let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; + + let metadata_enc = if let Some(m) = metadata { + let m_bytes = serde_json::to_vec(m).map_err(|_| ApiError::BadRequest)?; + Some(encrypt_with_key(user_key, &m_bytes).await) + } else { + None + }; + + let tags = extract_tags_from_metadata(metadata); + let tags_enc = encrypt_tags_b64(user_key, &tags); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let inserted = NewUserEmbedding { + uuid: Uuid::new_v4(), + user_id, + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + user_message_id: None, + assistant_message_id: None, + conversation_id: None, + vector_enc, + embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(), + vector_dim: DEFAULT_EMBEDDING_DIM, + content_enc, + metadata_enc, + tags_enc, + token_count, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert archival embedding: {:?}", e); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(inserted) +} + +#[allow(dead_code)] +#[allow(clippy::too_many_arguments)] +pub async fn insert_message_embedding( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + text: &str, + conversation_id: i64, + user_message_id: Option, + assistant_message_id: Option, +) -> Result { + let user_id = user.uuid; + let text = text.trim(); + if text.is_empty() { + return Err(ApiError::BadRequest); + } + + let (vector, token_count) = embed_text_via_tinfoil(state, user, auth_method, text).await?; + + let vector_bytes = serialize_f32_le(&vector); + let vector_enc = encrypt_with_key(user_key, &vector_bytes).await; + let content_enc = encrypt_with_key(user_key, text.as_bytes()).await; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let inserted = NewUserEmbedding { + uuid: Uuid::new_v4(), + user_id, + source_type: SOURCE_TYPE_MESSAGE.to_string(), + user_message_id, + assistant_message_id, + conversation_id: Some(conversation_id), + vector_enc, + embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(), + vector_dim: DEFAULT_EMBEDDING_DIM, + content_enc, + metadata_enc: None, + tags_enc: Vec::new(), + token_count, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert message embedding: {:?}", e); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(inserted) +} + +async fn load_all_user_embeddings( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let mut last_id: i64 = 0; + let mut out: Vec = Vec::new(); + + #[derive(Queryable)] + struct EmbeddingScanRow { + source_type: String, + conversation_id: Option, + vector_enc: Vec, + content_enc: Vec, + token_count: i32, + vector_dim: i32, + id: i64, + } + + loop { + let rows: Vec = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::id.gt(last_id)) + .order(user_embeddings::id.asc()) + .select(( + user_embeddings::source_type, + user_embeddings::conversation_id, + user_embeddings::vector_enc, + user_embeddings::content_enc, + user_embeddings::token_count, + user_embeddings::vector_dim, + user_embeddings::id, + )) + .limit(DB_SCAN_BATCH_SIZE) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to load embeddings batch for user={} after id={}: {:?}", + user_id, last_id, e + ); + ApiError::InternalServerError + })?; + + if rows.is_empty() { + break; + } + + for row in rows { + let EmbeddingScanRow { + source_type, + conversation_id, + vector_enc, + content_enc, + token_count, + vector_dim, + id, + } = row; + + let vector_bytes = decrypt_with_key(user_key, &vector_enc) + .map_err(|_| ApiError::InternalServerError)?; + let vector = deserialize_f32_le(&vector_bytes)?; + + if vector.len() != vector_dim as usize { + debug!( + "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", + id, + user_id, + vector_dim, + vector.len() + ); + continue; + } + + out.push(CachedEmbedding { + source_type, + conversation_id, + vector, + content_enc, + token_count, + }); + last_id = id; + } + } + + Ok(Arc::new(out)) +} + +async fn load_user_embeddings_by_tags( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + source_types: Option<&[String]>, + conversation_id: Option, + tags_enc_filter: &[Option], +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let tags_enc_filter = tags_enc_filter.to_vec(); + + let mut last_id: i64 = 0; + let mut out: Vec = Vec::new(); + + #[derive(Queryable)] + struct EmbeddingScanRow { + source_type: String, + conversation_id: Option, + vector_enc: Vec, + content_enc: Vec, + token_count: i32, + vector_dim: i32, + id: i64, + } + + loop { + let mut query = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::id.gt(last_id)) + .filter(user_embeddings::tags_enc.overlaps_with(tags_enc_filter.clone())) + .into_boxed(); + + if let Some(source_types) = source_types { + query = query.filter(user_embeddings::source_type.eq_any(source_types)); + } + + if let Some(conversation_id) = conversation_id { + query = query.filter(user_embeddings::conversation_id.eq(Some(conversation_id))); + } + + let rows: Vec = query + .order(user_embeddings::id.asc()) + .select(( + user_embeddings::source_type, + user_embeddings::conversation_id, + user_embeddings::vector_enc, + user_embeddings::content_enc, + user_embeddings::token_count, + user_embeddings::vector_dim, + user_embeddings::id, + )) + .limit(DB_SCAN_BATCH_SIZE) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to load tag-filtered embeddings batch for user={} after id={}: {:?}", + user_id, last_id, e + ); + ApiError::InternalServerError + })?; + + if rows.is_empty() { + break; + } + + for row in rows { + let EmbeddingScanRow { + source_type, + conversation_id, + vector_enc, + content_enc, + token_count, + vector_dim, + id, + } = row; + + let vector_bytes = decrypt_with_key(user_key, &vector_enc) + .map_err(|_| ApiError::InternalServerError)?; + let vector = deserialize_f32_le(&vector_bytes)?; + + if vector.len() != vector_dim as usize { + debug!( + "Skipping embedding id={} for user={} due to dim mismatch (expected={}, got={})", + id, + user_id, + vector_dim, + vector.len() + ); + continue; + } + + out.push(CachedEmbedding { + source_type, + conversation_id, + vector, + content_enc, + token_count, + }); + last_id = id; + } + } + + Ok(Arc::new(out)) +} + +#[allow(clippy::too_many_arguments)] +pub async fn search_user_embeddings( + state: &Arc, + user: &User, + auth_method: AuthMethod, + user_key: &SecretKey, + query: &str, + top_k: usize, + max_tokens: Option, + source_types: Option<&[String]>, + conversation_id: Option, + tags: Option<&[String]>, +) -> Result, ApiError> { + let top_k = top_k.clamp(1, 20); + + let user_id = user.uuid; + + let (query_vec, _query_tokens) = + embed_text_via_tinfoil(state, user, auth_method, query).await?; + + let tags_enc_filter = tags + .map(|t| normalize_tags(t.iter().map(|s| s.as_str()))) + .filter(|t| !t.is_empty()) + .map(|t| encrypt_tags_b64(user_key, &t)); + + let embeddings = if let Some(tags_enc_filter) = tags_enc_filter.as_ref() { + load_user_embeddings_by_tags( + state, + user_id, + user_key, + source_types, + conversation_id, + tags_enc_filter, + ) + .await? + } else { + let cached = { + let mut cache = state.rag_cache.lock().await; + cache.get(user_id) + }; + + if let Some(hit) = cached { + hit + } else { + let loaded = load_all_user_embeddings(state, user_id, user_key).await?; + state.rag_cache.lock().await.put(user_id, loaded.clone()); + loaded + } + }; + + let candidates = top_k_candidates( + &query_vec, + &embeddings, + top_k, + source_types, + conversation_id, + )?; + + let mut results: Vec = Vec::with_capacity(candidates.len()); + for c in candidates { + let plaintext = decrypt_with_key(user_key, &c.content_enc) + .map_err(|_| ApiError::InternalServerError)?; + let content = String::from_utf8(plaintext).map_err(|_| ApiError::InternalServerError)?; + results.push(RagSearchResult { + content, + score: c.score, + token_count: c.token_count, + }); + } + + if let Some(budget) = max_tokens { + results = apply_token_budget(results, budget); + } + + Ok(results) +} + +pub async fn delete_all_user_embeddings(state: &AppState, user_id: Uuid) -> Result<(), ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + diesel::delete(user_embeddings::table.filter(user_embeddings::user_id.eq(user_id))) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete all embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user_id); + Ok(()) +} + +pub async fn delete_user_embedding_by_uuid( + state: &AppState, + user_id: Uuid, + embedding_uuid: Uuid, +) -> Result<(), ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let affected = diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::uuid.eq(embedding_uuid)), + ) + .execute(&mut conn) + .map_err(|e| { + error!( + "Failed to delete embedding user={} uuid={}: {:?}", + user_id, embedding_uuid, e + ); + ApiError::InternalServerError + })?; + + if affected == 0 { + return Err(ApiError::NotFound); + } + + state.rag_cache.lock().await.evict_user(user_id); + Ok(()) +} + +pub async fn embeddings_status( + state: &AppState, + user_id: Uuid, +) -> Result { + use diesel::dsl::count_star; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let total_embeddings: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count embeddings for user={} (total): {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + let grouped: Vec<(String, i64)> = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .group_by(user_embeddings::embedding_model) + .select((user_embeddings::embedding_model, count_star())) + .load(&mut conn) + .map_err(|e| { + error!( + "Failed to group embeddings by model for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + let mut by_model: HashMap = HashMap::new(); + for (model, count) in grouped { + by_model.insert(model, count); + } + + let stale_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(user_id)) + .filter(user_embeddings::embedding_model.ne(DEFAULT_EMBEDDING_MODEL)) + .select(count_star()) + .first(&mut conn) + .map_err(|e| { + error!( + "Failed to count stale embeddings for user={}: {:?}", + user_id, e + ); + ApiError::InternalServerError + })?; + + Ok(RagEmbeddingsStatus { + total_embeddings, + by_model, + stale_count, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn serialize_deserialize_roundtrip() { + let v = vec![0.0f32, 1.5, -2.25, 42.0]; + let bytes = serialize_f32_le(&v); + let decoded = deserialize_f32_le(&bytes).unwrap(); + assert_eq!(v, decoded); + } + + #[tokio::test] + async fn vector_encrypt_roundtrip() { + let key = SecretKey::from_slice(&[7u8; 32]).unwrap(); + let v = vec![0.0f32, 1.5, -2.25, 42.0]; + + let bytes = serialize_f32_le(&v); + let enc = encrypt_with_key(&key, &bytes).await; + let dec = decrypt_with_key(&key, &enc).unwrap(); + let decoded = deserialize_f32_le(&dec).unwrap(); + assert_eq!(v, decoded); + } + + #[test] + fn cosine_similarity_known_values() { + let a = vec![1.0f32, 0.0, 0.0]; + let b = vec![1.0f32, 0.0, 0.0]; + let c = vec![0.0f32, 1.0, 0.0]; + + assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-6); + assert!((cosine_similarity(&a, &c).unwrap() - 0.0).abs() < 1e-6); + } + + #[test] + fn cosine_similarity_mismatched_dimensions_errors() { + let a = vec![1.0f32, 0.0]; + let b = vec![1.0f32]; + assert!(matches!( + cosine_similarity(&a, &b), + Err(ApiError::BadRequest) + )); + } + + #[test] + fn cosine_similarity_zero_vectors_return_zero() { + let zero = vec![0.0f32, 0.0, 0.0]; + let nonzero = vec![1.0f32, 2.0, 3.0]; + + assert_eq!(cosine_similarity(&zero, &nonzero).unwrap(), 0.0); + assert_eq!(cosine_similarity(&nonzero, &zero).unwrap(), 0.0); + assert_eq!(cosine_similarity(&zero, &zero).unwrap(), 0.0); + } + + #[test] + fn deserialize_invalid_byte_length_errors() { + let bytes = vec![0u8; 3]; + assert!(matches!( + deserialize_f32_le(&bytes), + Err(ApiError::BadRequest) + )); + } + + #[test] + fn apply_token_budget_is_prefix() { + let results = vec![ + RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 6, + }, + RagSearchResult { + content: "b".to_string(), + score: 0.9, + token_count: 6, + }, + RagSearchResult { + content: "c".to_string(), + score: 0.8, + token_count: 1, + }, + ]; + + let limited = apply_token_budget(results, 10); + assert_eq!(limited.len(), 1); + assert_eq!(limited[0].content, "a"); + } + + #[test] + fn apply_token_budget_exact_fit_keeps_all() { + let results = vec![ + RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 3, + }, + RagSearchResult { + content: "b".to_string(), + score: 0.9, + token_count: 7, + }, + ]; + + let limited = apply_token_budget(results, 10); + assert_eq!(limited.len(), 2); + } + + #[test] + fn apply_token_budget_zero_budget_returns_empty() { + let results = vec![RagSearchResult { + content: "a".to_string(), + score: 1.0, + token_count: 1, + }]; + + let limited = apply_token_budget(results, 0); + assert!(limited.is_empty()); + } + + #[test] + fn top_k_heap_ranking_and_filters() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 5, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.0, 1.0], + content_enc: b"b".to_vec(), + token_count: 7, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.8, 0.2], + content_enc: b"c".to_vec(), + token_count: 9, + }, + ]; + + let items = top_k_candidates( + &query, + &embeddings, + 2, + Some(&[SOURCE_TYPE_MESSAGE.to_string()]), + Some(123), + ) + .unwrap(); + + assert_eq!(items.len(), 2); + assert!(items[0].score >= items[1].score); + assert_eq!(items[0].content_enc, b"c"); + assert_eq!(items[1].content_enc, b"b"); + } + + #[test] + fn top_k_tie_break_prefers_fewer_tokens() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 10, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"b".to_vec(), + token_count: 5, + }, + ]; + + let items = top_k_candidates(&query, &embeddings, 1, None, None).unwrap(); + assert_eq!(items.len(), 1); + assert_eq!(items[0].content_enc, b"b"); + } + + #[test] + fn top_k_when_k_gt_embeddings_returns_all() { + let query = vec![1.0f32, 0.0]; + + let embeddings = vec![ + CachedEmbedding { + source_type: SOURCE_TYPE_ARCHIVAL.to_string(), + conversation_id: None, + vector: vec![1.0, 0.0], + content_enc: b"a".to_vec(), + token_count: 5, + }, + CachedEmbedding { + source_type: SOURCE_TYPE_MESSAGE.to_string(), + conversation_id: Some(123), + vector: vec![0.0, 1.0], + content_enc: b"b".to_vec(), + token_count: 7, + }, + ]; + + let total = embeddings.len(); + let items = top_k_candidates(&query, &embeddings, 10, None, None).unwrap(); + assert_eq!(items.len(), total); + } + + #[test] + fn top_k_with_empty_embeddings_returns_empty() { + let query = vec![1.0f32, 0.0]; + let embeddings: Vec = vec![]; + let items = top_k_candidates(&query, &embeddings, 10, None, None).unwrap(); + assert!(items.is_empty()); + } + + #[test] + fn rag_cache_evict_user_removes_entry() { + let mut cache = RagCache::new(10, Duration::from_secs(60)); + let user_id = Uuid::new_v4(); + + cache.put(user_id, Arc::new(vec![])); + assert!(cache.entries.contains_key(&user_id)); + + cache.evict_user(user_id); + assert!(!cache.entries.contains_key(&user_id)); + assert!(!cache.lru.contains(&user_id)); + assert!(cache.get(user_id).is_none()); + } + + #[tokio::test] + async fn rag_cache_lru_eviction() { + let mut cache = RagCache::new(2, Duration::from_secs(60)); + + let v1 = Arc::new(vec![]); + let v2 = Arc::new(vec![]); + let v3 = Arc::new(vec![]); + + let u1 = Uuid::new_v4(); + let u2 = Uuid::new_v4(); + let u3 = Uuid::new_v4(); + + cache.put(u1, v1); + cache.put(u2, v2); + // touch u1 so u2 becomes LRU + cache.get(u1); + cache.put(u3, v3); + + assert!(cache.entries.contains_key(&u1)); + assert!(!cache.entries.contains_key(&u2)); + assert!(cache.entries.contains_key(&u3)); + } + + #[tokio::test] + async fn rag_cache_ttl_expiration() { + let mut cache = RagCache::new(10, Duration::from_millis(5)); + let user = Uuid::new_v4(); + + cache.put(user, Arc::new(vec![])); + tokio::time::sleep(Duration::from_millis(10)).await; + assert!(cache.get(user).is_none()); + } +} diff --git a/src/web/mod.rs b/src/web/mod.rs index dd9a5cc0..5e7eff65 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -8,14 +8,17 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; +pub mod rag; pub mod responses; pub use health_routes::router_with_state as health_routes_with_state; pub use login_routes::router as login_routes; pub use oauth_routes::router as oauth_routes; +pub use openai::get_embedding_vector; pub use openai::router as openai_routes; pub use platform::router as platform_routes; pub use protected_routes::router as protected_routes; +pub use rag::router as rag_routes; pub use responses::conversation_projects_router as conversation_projects_routes; pub use responses::conversations_router as conversations_routes; pub use responses::instructions_router as instructions_routes; diff --git a/src/web/openai.rs b/src/web/openai.rs index f4afad21..8a2e4569 100644 --- a/src/web/openai.rs +++ b/src/web/openai.rs @@ -2,7 +2,7 @@ use crate::model_config::{model_catalog_response, openai_models_response}; use crate::models::token_usage::NewTokenUsage; use crate::models::users::User; use crate::provider_routing::ProviderRoutingError; -use crate::proxy_config::ProxyConfig; +use crate::proxy_config::{canonicalize_tinfoil_model, ProxyConfig}; use crate::sqs::UsageEvent; use crate::web::audio_utils::{merge_transcriptions, AudioSplitter, TINFOIL_MAX_SIZE}; use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; @@ -1999,53 +1999,10 @@ async fn proxy_tts( encrypt_response(&state, &session_id, &audio_response).await } -async fn proxy_embeddings( - State(state): State>, - _headers: HeaderMap, - axum::Extension(session_id): axum::Extension, - axum::Extension(user): axum::Extension, - axum::Extension(_auth_method): axum::Extension, - axum::Extension(embedding_request): axum::Extension, -) -> Result>, ApiError> { - // Check if guest user is allowed (paid guests are allowed, free guests are not) - if user.is_guest() { - if let Some(billing_client) = &state.billing_client { - match billing_client.is_user_paid(user.uuid).await { - Ok(true) => { - debug!("Paid guest user allowed for embeddings: {}", user.uuid); - } - Ok(false) => { - error!( - "Free guest user attempted to use embeddings feature: {}", - user.uuid - ); - return Err(ApiError::Unauthorized); - } - Err(e) => { - error!("Billing check failed for guest user {}: {}", user.uuid, e); - return Err(ApiError::Unauthorized); - } - } - } else { - error!( - "Guest user attempted to use embeddings without billing client: {}", - user.uuid - ); - return Err(ApiError::Unauthorized); - } - } - - // Validate input is not empty - let is_empty = match &embedding_request.input { - Value::String(s) => s.trim().is_empty(), - Value::Array(arr) => arr.is_empty(), - _ => true, - }; - if is_empty { - error!("Input is empty or invalid"); - return Err(ApiError::BadRequest); - } - +async fn request_embeddings_from_primary( + state: &AppState, + embedding_request: &EmbeddingRequest, +) -> Result<(Value, String), ApiError> { let proxy_config = state .proxy_router .get_tinfoil_proxy() @@ -2058,8 +2015,13 @@ async fn proxy_embeddings( .pool_max_idle_per_host(10) .build::<_, HyperBody>(https); + let mut provider_request = embedding_request.clone(); + if proxy_config.provider_name == "tinfoil" { + provider_request.model = canonicalize_tinfoil_model(&embedding_request.model); + } + // Build request body - let request_body = serde_json::to_string(&embedding_request).map_err(|e| { + let request_body = serde_json::to_string(&provider_request).map_err(|e| { error!("Failed to serialize embedding request: {:?}", e); ApiError::InternalServerError })?; @@ -2124,32 +2086,145 @@ async fn proxy_embeddings( ApiError::InternalServerError })?; + Ok((response_json, proxy_config.provider_name)) +} + +async fn request_embeddings_with_billing( + state: &Arc, + user: &User, + auth_method: AuthMethod, + embedding_request: &EmbeddingRequest, +) -> Result<(Value, String, i32), ApiError> { + // Check if guest user is allowed (paid guests are allowed, free guests are not) + if user.is_guest() { + if let Some(billing_client) = &state.billing_client { + match billing_client.is_user_paid(user.uuid).await { + Ok(true) => { + debug!("Paid guest user allowed for embeddings: {}", user.uuid); + } + Ok(false) => { + error!( + "Free guest user attempted to use embeddings feature: {}", + user.uuid + ); + return Err(ApiError::Unauthorized); + } + Err(e) => { + error!("Billing check failed for guest user {}: {}", user.uuid, e); + return Err(ApiError::Unauthorized); + } + } + } else { + error!( + "Guest user attempted to use embeddings without billing client: {}", + user.uuid + ); + return Err(ApiError::Unauthorized); + } + } + + let (response_json, provider_name) = + request_embeddings_from_primary(state, embedding_request).await?; + + let prompt_tokens: i32 = response_json + .get("usage") + .and_then(|u| u.get("prompt_tokens")) + .and_then(|v| v.as_i64()) + .unwrap_or(0) as i32; + // Handle billing - embeddings only have prompt_tokens (no completion_tokens) - if let Some(usage) = response_json.get("usage") { - let prompt_tokens = usage - .get("prompt_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0) as i32; - - if prompt_tokens > 0 { - let billing_context = - BillingContext::new(_auth_method, embedding_request.model.clone()); - let embedding_usage = CompletionUsage { - prompt_tokens, - completion_tokens: 0, // Embeddings don't have completion tokens - cached_prompt_tokens: None, - }; - publish_usage_event_internal( - &state, - &user, - &billing_context, - embedding_usage, - &proxy_config.provider_name, - ) - .await; + if prompt_tokens > 0 { + let billing_context = BillingContext::new(auth_method, embedding_request.model.clone()); + let embedding_usage = CompletionUsage { + prompt_tokens, + completion_tokens: 0, + cached_prompt_tokens: None, + }; + publish_usage_event_internal( + state, + user, + &billing_context, + embedding_usage, + &provider_name, + ) + .await; + } + + Ok((response_json, provider_name, prompt_tokens)) +} + +pub async fn get_embedding_vector( + state: &Arc, + user: &User, + auth_method: AuthMethod, + model: &str, + input: &str, + dimensions: Option, +) -> Result<(Vec, i32), ApiError> { + let request = EmbeddingRequest { + input: Value::String(input.to_string()), + model: model.to_string(), + encoding_format: Some("float".to_string()), + dimensions, + user: None, + }; + + let (response_json, _provider, prompt_tokens) = + request_embeddings_with_billing(state, user, auth_method, &request).await?; + + let embedding = response_json + .get("data") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.first()) + .and_then(|obj| obj.get("embedding")) + .and_then(|v| v.as_array()) + .ok_or(ApiError::InternalServerError)?; + + let mut vector: Vec = Vec::with_capacity(embedding.len()); + for v in embedding { + let f = v.as_f64().ok_or(ApiError::InternalServerError)?; + vector.push(f as f32); + } + + if let Some(dim) = dimensions { + if vector.len() != dim as usize { + error!( + "Embeddings response dim mismatch: expected={}, got={}", + dim, + vector.len() + ); + return Err(ApiError::InternalServerError); } } + Ok((vector, prompt_tokens)) +} + +async fn proxy_embeddings( + State(state): State>, + _headers: HeaderMap, + axum::Extension(session_id): axum::Extension, + axum::Extension(user): axum::Extension, + axum::Extension(auth_method): axum::Extension, + axum::Extension(embedding_request): axum::Extension, +) -> Result>, ApiError> { + debug!("Entering proxy_embeddings function"); + + // Validate input is not empty + let is_empty = match &embedding_request.input { + Value::String(s) => s.trim().is_empty(), + Value::Array(arr) => arr.is_empty(), + _ => true, + }; + if is_empty { + error!("Input is empty or invalid"); + return Err(ApiError::BadRequest); + } + + let (response_json, _provider_name, _prompt_tokens) = + request_embeddings_with_billing(&state, &user, auth_method, &embedding_request).await?; + + debug!("Exiting proxy_embeddings function"); // Encrypt and return the response encrypt_response(&state, &session_id, &response_json).await } diff --git a/src/web/rag.rs b/src/web/rag.rs new file mode 100644 index 00000000..657a3e92 --- /dev/null +++ b/src/web/rag.rs @@ -0,0 +1,216 @@ +use axum::{ + extract::{Path, State}, + middleware::from_fn_with_state, + routing::{delete, get, post}, + Extension, Json, Router, +}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::sync::Arc; +use uuid::Uuid; + +use crate::jwt::AuthContext; +use crate::models::users::User; +use crate::rag; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::web::openai_auth::AuthMethod; +use crate::web::responses::error_mapping; +use crate::{ApiError, AppMode, AppState}; + +#[derive(Debug, Clone, Deserialize)] +struct InsertEmbeddingRequest { + text: String, + metadata: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct InsertEmbeddingResponse { + id: Uuid, + source_type: String, + embedding_model: String, + token_count: i32, + created_at: DateTime, +} + +#[derive(Debug, Clone, Deserialize)] +struct SearchRequest { + query: String, + top_k: Option, + max_tokens: Option, + source_types: Option>, + conversation_id: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct SearchResponse { + results: Vec, +} + +pub fn router(app_state: Arc) -> Router<()> { + // Experimental endpoints: only enabled in Local/Dev. + if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { + return Router::new().with_state(app_state); + } + + Router::new() + .route( + "/v1/rag/embeddings", + post(insert_archival_embedding).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/rag/embeddings", + delete(delete_all_embeddings) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/rag/embeddings/:id", + delete(delete_embedding) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/rag/search", + post(search).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/rag/embeddings/status", + get(status).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .with_state(app_state) +} + +async fn insert_archival_embedding( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result { + if body.text.trim().is_empty() { + return Err(ApiError::BadRequest); + } + if let Some(m) = &body.metadata { + if !m.is_object() { + return Err(ApiError::BadRequest); + } + } + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let inserted = rag::insert_archival_embedding( + &state, + &user, + AuthMethod::Jwt, + &user_key, + &body.text, + body.metadata.as_ref(), + ) + .await?; + + let response = InsertEmbeddingResponse { + id: inserted.uuid, + source_type: inserted.source_type, + embedding_model: inserted.embedding_model, + token_count: inserted.token_count, + created_at: inserted.created_at, + }; + + let encrypted = encrypt_response(&state, &session_id, &response).await?; + Ok((axum::http::StatusCode::CREATED, encrypted)) +} + +async fn search( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if body.query.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let top_k = body.top_k.unwrap_or(5); + if top_k == 0 || top_k > 20 { + return Err(ApiError::BadRequest); + } + + if let Some(max_tokens) = body.max_tokens { + if max_tokens <= 0 { + return Err(ApiError::BadRequest); + } + } + + if let Some(source_types) = &body.source_types { + if source_types.is_empty() { + return Err(ApiError::BadRequest); + } + } + + let conversation_internal_id = if let Some(conversation_uuid) = body.conversation_id { + let conversation = state + .db + .get_conversation_by_uuid_and_user(conversation_uuid, user.uuid) + .map_err(error_mapping::map_conversation_error)?; + Some(conversation.id) + } else { + None + }; + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let results = rag::search_user_embeddings( + &state, + &user, + AuthMethod::Jwt, + &user_key, + &body.query, + top_k, + body.max_tokens, + body.source_types.as_deref(), + conversation_internal_id, + None, + ) + .await?; + + let response = SearchResponse { results }; + encrypt_response(&state, &session_id, &response).await +} + +async fn delete_all_embeddings( + State(state): State>, + Extension(user): Extension, +) -> Result { + rag::delete_all_user_embeddings(&state, user.uuid).await?; + Ok(axum::http::StatusCode::NO_CONTENT) +} + +async fn delete_embedding( + State(state): State>, + Path(embedding_id): Path, + Extension(user): Extension, +) -> Result { + rag::delete_user_embedding_by_uuid(&state, user.uuid, embedding_id).await?; + Ok(axum::http::StatusCode::NO_CONTENT) +} + +async fn status( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let response = rag::embeddings_status(&state, user.uuid).await?; + encrypt_response(&state, &session_id, &response).await +} From 8dfba178731d20707236baf14c6b9a7b48c61423 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:18:31 -0500 Subject: [PATCH 3/7] feat(agent): add core Maple runtime and persistence Introduce the main Maple agent runtime, persistence models, and hidden backing conversation flow so the restacked branch has a working end-to-end agent foundation on top of the new schema. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 3909 +++++++++++++++++++++----- Cargo.toml | 1 + flake.nix | 12 +- src/main.rs | 8 +- src/models/agents.rs | 121 + src/models/conversation_summaries.rs | 107 + src/models/memory_blocks.rs | 109 + src/models/mod.rs | 3 + src/models/responses.rs | 24 + src/web/agent/compaction.rs | 247 ++ src/web/agent/mod.rs | 563 ++++ src/web/agent/runtime.rs | 1387 +++++++++ src/web/agent/signatures.rs | 602 ++++ src/web/agent/tools.rs | 903 ++++++ src/web/agent/vision.rs | 115 + src/web/mod.rs | 2 + src/web/responses/conversations.rs | 119 +- src/web/responses/handlers.rs | 1 + 18 files changed, 7551 insertions(+), 682 deletions(-) create mode 100644 src/models/agents.rs create mode 100644 src/models/conversation_summaries.rs create mode 100644 src/models/memory_blocks.rs create mode 100644 src/web/agent/compaction.rs create mode 100644 src/web/agent/mod.rs create mode 100644 src/web/agent/runtime.rs create mode 100644 src/web/agent/signatures.rs create mode 100644 src/web/agent/tools.rs create mode 100644 src/web/agent/vision.rs diff --git a/Cargo.lock b/Cargo.lock index 5c05d5da..f323f2b0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aead" version = "0.5.2" @@ -53,20 +59,49 @@ dependencies = [ "zeroize", ] +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ "memchr", ] [[package]] -name = "android-tzdata" -version = "0.1.1" +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +dependencies = [ + "alloc-no-stdlib", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" [[package]] name = "android_system_properties" @@ -77,11 +112,70 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" -version = "1.0.100" +version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] [[package]] name = "argon2" @@ -95,12 +189,232 @@ dependencies = [ "password-hash", ] +[[package]] +name = "arrayvec" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b" + [[package]] name = "arrayvec" version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "arrow" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e833808ff2d94ed40d9379848a950d995043c7fb3e81a30b383f4c6033821cc" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad08897b81588f60ba983e3ca39bda2b179bdd84dced378e7df81a5313802ef8" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num", +] + +[[package]] +name = "arrow-array" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8548ca7c070d8db9ce7aa43f37393e4bfcf3f2d3681df278490772fd1673d08d" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half 2.7.1", + "hashbrown 0.16.1", + "num", +] + +[[package]] +name = "arrow-buffer" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e003216336f70446457e280807a73899dd822feaf02087d31febca1363e2fccc" +dependencies = [ + "bytes", + "half 2.7.1", + "num", +] + +[[package]] +name = "arrow-cast" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "919418a0681298d3a77d1a315f625916cb5678ad0d74b9c60108eb15fd083023" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "atoi", + "base64 0.22.1", + "chrono", + "half 2.7.1", + "lexical-core", + "num", + "ryu", +] + +[[package]] +name = "arrow-csv" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa9bf02705b5cf762b6f764c65f04ae9082c7cfc4e96e0c33548ee3f67012eb" +dependencies = [ + "arrow-array", + "arrow-cast", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5c64fff1d142f833d78897a772f2e5b55b36cb3e6320376f0961ab0db7bd6d0" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half 2.7.1", + "num", +] + +[[package]] +name = "arrow-ipc" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d3594dcddccc7f20fd069bc8e9828ce37220372680ff638c5e00dea427d88f5" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88cf36502b64a127dc659e3b305f1d993a544eab0d48cce704424e62074dc04b" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "chrono", + "half 2.7.1", + "indexmap", + "lexical-core", + "memchr", + "num", + "serde", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c8f82583eb4f8d84d4ee55fd1cb306720cddead7596edce95b50ee418edf66f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d07ba24522229d9085031df6b94605e0f4b26e099fb7cdeec37abd941a73753" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half 2.7.1", +] + +[[package]] +name = "arrow-schema" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3aa9e59c611ebc291c28582077ef25c97f1975383f1479b12f3b9ffee2ffabe" + +[[package]] +name = "arrow-select" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c41dbbd1e97bfcaee4fcb30e29105fb2c75e4d82ae4de70b792a5d3f66b2e7a" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num", +] + +[[package]] +name = "arrow-string" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53f5183c150fbc619eede22b861ea7c0eebed8eaac0333eaa7f6da5205fd504d" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num", + "regex", + "regex-syntax", +] + +[[package]] +name = "as-any" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0f477b951e452a0b6b4a10b53ccd569042d1d01729b519e02074a9c0958a063" + [[package]] name = "asn1-rs" version = "0.5.2" @@ -113,7 +427,7 @@ dependencies = [ "nom", "num-traits", "rusticata-macros", - "thiserror", + "thiserror 1.0.69", "time", ] @@ -140,6 +454,18 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -159,9 +485,15 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] +[[package]] +name = "async-task" +version = "4.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" + [[package]] name = "async-trait" version = "0.1.89" @@ -170,7 +502,16 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", +] + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", ] [[package]] @@ -181,15 +522,15 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.3.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-config" -version = "1.8.0" +version = "1.8.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "455e9fb7743c6f6267eb2830ccc08686fbb3d13c9a689369562fd4d4ef9ea462" +checksum = "c456581cb3c77fafcc8c67204a70680d40b61112d6da78c77bd31d945b65f1b5" dependencies = [ "aws-credential-types", "aws-runtime", @@ -206,7 +547,7 @@ dependencies = [ "bytes", "fastrand", "hex", - "http 1.1.0", + "http 1.4.0", "ring", "time", "tokio", @@ -266,9 +607,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.5.18" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "959dab27ce613e6c9658eb3621064d0e2027e5f2acb65bc526a43577facea557" +checksum = "c635c2dc792cb4a11ce1a4f392a925340d1bdf499289b5ec1ec6810954eb43f5" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -280,8 +621,8 @@ dependencies = [ "aws-types", "bytes", "fastrand", - "http 0.2.12", - "http-body 0.4.6", + "http 1.4.0", + "http-body 1.0.1", "percent-encoding", "pin-project-lite", "tracing", @@ -290,15 +631,16 @@ dependencies = [ [[package]] name = "aws-sdk-sqs" -version = "1.90.0" +version = "1.93.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "201073d85c1852c22672565b9ddd8286ec4768ad680a261337e395b4d4699d44" +checksum = "32840f6fb45c5e3bcac60379141582cede475984d3ff4dd030f7d983503cd61e" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", "aws-smithy-http", "aws-smithy-json", + "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", "aws-smithy-types", @@ -306,21 +648,23 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", + "http 1.4.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-sso" -version = "1.89.0" +version = "1.93.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9c1b1af02288f729e95b72bd17988c009aa72e26dcb59b3200f86d7aea726c9" +checksum = "9dcb38bb33fc0a11f1ffc3e3e85669e0a11a37690b86f77e75306d8f369146a0" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", "aws-smithy-http", "aws-smithy-json", + "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", "aws-smithy-types", @@ -328,21 +672,23 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", + "http 1.4.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-ssooidc" -version = "1.91.0" +version = "1.95.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e8122301558dc7c6c68e878af918880b82ff41897a60c8c4e18e4dc4d93e9f1" +checksum = "2ada8ffbea7bd1be1f53df1dadb0f8fdb04badb13185b3321b929d1ee3caad09" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", "aws-smithy-http", "aws-smithy-json", + "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", "aws-smithy-types", @@ -350,21 +696,23 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", + "http 1.4.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-sts" -version = "1.91.0" +version = "1.97.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f8090151d4d1e971269957b10dbf287bba551ab812e591ce0516b1c73b75d27" +checksum = "e6443ccadc777095d5ed13e21f5c364878c9f5bad4e35187a6cdbd863b0afcad" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", "aws-smithy-http", "aws-smithy-json", + "aws-smithy-observability", "aws-smithy-query", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -373,15 +721,16 @@ dependencies = [ "aws-types", "fastrand", "http 0.2.12", + "http 1.4.0", "regex-lite", "tracing", ] [[package]] name = "aws-sigv4" -version = "1.3.7" +version = "1.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69e523e1c4e8e7e8ff219d732988e22bfeae8a1cafdbe6d9eca1546fa080be7c" +checksum = "efa49f3c607b92daae0c078d48a4571f599f966dce3caee5f1ea55c4d9073f99" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -392,7 +741,7 @@ dependencies = [ "hex", "hmac", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "percent-encoding", "sha2", "time", @@ -412,9 +761,9 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.62.6" +version = "0.63.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "826141069295752372f8203c17f28e30c464d22899a43a0c9fd9c458d469c88b" +checksum = "630e67f2a31094ffa51b210ae030855cb8f3b7ee1329bdd8d085aaf61e8b97fc" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -422,9 +771,9 @@ dependencies = [ "bytes-utils", "futures-core", "futures-util", - "http 0.2.12", - "http 1.1.0", - "http-body 0.4.6", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", "percent-encoding", "pin-project-lite", "pin-utils", @@ -433,16 +782,16 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.1.5" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59e62db736db19c488966c8d787f52e6270be565727236fd5579eaa301e7bc4a" +checksum = "12fb0abf49ff0cab20fd31ac1215ed7ce0ea92286ba09e2854b42ba5cabe7525" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", "aws-smithy-types", "h2 0.4.13", - "http 1.1.0", - "hyper 1.7.0", + "http 1.4.0", + "hyper 1.9.0", "hyper-rustls", "hyper-util", "pin-project-lite", @@ -451,33 +800,33 @@ dependencies = [ "rustls-pki-types", "tokio", "tokio-rustls", - "tower 0.5.2", + "tower", "tracing", ] [[package]] name = "aws-smithy-json" -version = "0.61.9" +version = "0.62.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49fa1213db31ac95288d981476f78d05d9cbb0353d22cdf3472cc05bb02f6551" +checksum = "3cb96aa208d62ee94104645f7b2ecaf77bf27edf161590b6224bfbac2832f979" dependencies = [ "aws-smithy-types", ] [[package]] name = "aws-smithy-observability" -version = "0.2.0" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef1fcbefc7ece1d70dcce29e490f269695dfca2d2bacdeaf9e5c3f799e4e6a42" +checksum = "c0a46543fbc94621080b3cf553eb4cbbdc41dd9780a30c4756400f0139440a1d" dependencies = [ "aws-smithy-runtime-api", ] [[package]] name = "aws-smithy-query" -version = "0.60.9" +version = "0.60.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae5d689cf437eae90460e944a58b5668530d433b4ff85789e69d2f2a556e057d" +checksum = "0cebbddb6f3a5bd81553643e9c7daf3cc3dc5b0b5f398ac668630e8a84e6fff0" dependencies = [ "aws-smithy-types", "urlencoding", @@ -485,9 +834,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.9.8" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb5b6167fcdf47399024e81ac08e795180c576a20e4d4ce67949f9a88ae37dc1" +checksum = "f3df87c14f0127a0d77eb261c3bc45d5b4833e2a1f63583ebfb728e4852134ee" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -498,9 +847,10 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "http-body 0.4.6", "http-body 1.0.1", + "http-body-util", "pin-project-lite", "pin-utils", "tokio", @@ -509,15 +859,15 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.10.0" +version = "1.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efce7aaaf59ad53c5412f14fc19b2d5c6ab2c3ec688d272fd31f76ec12f44fb0" +checksum = "49952c52f7eebb72ce2a754d3866cc0f87b97d2a46146b79f80f3a93fb2b3716" dependencies = [ "aws-smithy-async", "aws-smithy-types", "bytes", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "pin-project-lite", "tokio", "tracing", @@ -526,16 +876,16 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.3.6" +version = "1.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65f172bcb02424eb94425db8aed1b6d583b5104d4d5ddddf22402c661a320048" +checksum = "3b3a26048eeab0ddeba4b4f9d51654c79af8c3b32357dc5f336cee85ab331c33" dependencies = [ "base64-simd", "bytes", "bytes-utils", "futures-core", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "http-body 0.4.6", "http-body 1.0.1", "http-body-util", @@ -575,19 +925,19 @@ dependencies = [ [[package]] name = "axum" -version = "0.7.5" +version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a6c9af12842a67734c9a2e355436e5d03b22383ed60cf13cd0c18fbfe3dcbcf" +checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", "axum-macros", "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper 1.7.0", + "hyper 1.9.0", "hyper-util", "itoa", "matchit", @@ -601,9 +951,9 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", - "sync_wrapper 1.0.1", + "sync_wrapper 1.0.2", "tokio", - "tower 0.4.13", + "tower", "tower-layer", "tower-service", "tracing", @@ -611,20 +961,20 @@ dependencies = [ [[package]] name = "axum-core" -version = "0.4.3" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a15c63fd72d41492dc4f497196f5da1fb04fb7529e631d73630d1b491e47a2e3" +checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" dependencies = [ "async-trait", "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", "rustversion", - "sync_wrapper 0.1.2", + "sync_wrapper 1.0.2", "tower-layer", "tower-service", "tracing", @@ -632,14 +982,13 @@ dependencies = [ [[package]] name = "axum-macros" -version = "0.4.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00c055ee2d014ae5981ce1016374e8213682aa14d9bf40e48ab48b5f3ef20eaa" +checksum = "57d123550fa8d071b7255cb0cc04dc302baa6c8c4a79f55701552684d8399bce" dependencies = [ - "heck 0.4.1", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -649,21 +998,84 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b62ddb9cb1ec0a098ad4bbf9344d0713fa193ae1a80af55febcff2627b6a00c1" dependencies = [ "futures-core", - "getrandom 0.2.15", + "getrandom 0.2.17", "instant", "pin-project-lite", - "rand", + "rand 0.8.6", "tokio", ] +[[package]] +name = "baml-ids" +version = "0.0.1" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "getrandom 0.2.17", + "serde", + "time", + "type-safe-id", + "uuid", +] + +[[package]] +name = "baml-types" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "baml-ids", + "clap", + "indexmap", + "internal-baml-diagnostics", + "itertools", + "minijinja", + "pretty", + "secrecy", + "serde", + "serde_json", + "strum", + "web-time", +] + +[[package]] +name = "bamltype" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "baml-types", + "bamltype-derive", + "facet", + "facet-reflect", + "indexmap", + "internal-baml-jinja", + "jsonish", + "minijinja", + "sha2", + "thiserror 2.0.18", +] + +[[package]] +name = "bamltype-derive" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "convert_case", + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "base58ck" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2c8d66485a3a2ea485c1913c4572ce0256067a5377ac8c75c4960e1cda98605f" dependencies = [ - "bitcoin-internals 0.3.0", - "bitcoin_hashes 0.14.0", + "bitcoin-internals", + "bitcoin_hashes", ] [[package]] @@ -696,9 +1108,9 @@ dependencies = [ [[package]] name = "base64ct" -version = "1.6.0" +version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" [[package]] name = "base85" @@ -706,20 +1118,20 @@ version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "36915bbaca237c626689b5bd14d02f2ba7a5a359d30a2a08be697392e3718079" dependencies = [ - "thiserror", + "thiserror 1.0.69", ] [[package]] name = "bech32" -version = "0.11.0" +version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d965446196e3b7decd44aa7ee49e31d630118f90ef12f97900f262eb915c951d" +checksum = "32637268377fc7b10a8c6d51de3e7fba1ce5dd371a96e342b34e6078db558e7f" [[package]] name = "bigdecimal" -version = "0.4.5" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51d712318a27c7150326677b321a5fa91b55f6d9034ffd67f20319e147d40cee" +checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", @@ -729,13 +1141,22 @@ dependencies = [ "serde", ] +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bip39" -version = "2.1.0" +version = "2.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33415e24172c1b7d6066f6d999545375ab8e1d95421d6784bdfff9496f292387" +checksum = "90dbd31c98227229239363921e60fcf5e558e43ec69094d46fc4996f08d1d5bc" dependencies = [ - "bitcoin_hashes 0.13.0", + "bitcoin_hashes", "serde", "unicode-normalization", ] @@ -770,28 +1191,22 @@ checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" [[package]] name = "bitcoin" -version = "0.32.5" +version = "0.32.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce6bc65742dea50536e35ad42492b234c27904a27f0abdcbce605015cb4ea026" +checksum = "1e499f9fc0407f50fe98af744ab44fa67d409f76b6772e1689ec8485eb0c0f66" dependencies = [ "base58ck", "bech32", - "bitcoin-internals 0.3.0", + "bitcoin-internals", "bitcoin-io", "bitcoin-units", - "bitcoin_hashes 0.14.0", - "hex-conservative 0.2.1", + "bitcoin_hashes", + "hex-conservative", "hex_lit", "secp256k1", "serde", ] -[[package]] -name = "bitcoin-internals" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9425c3bf7089c983facbae04de54513cce73b41c7f9ff8c845b54e7bc64ebbfb" - [[package]] name = "bitcoin-internals" version = "0.3.0" @@ -803,38 +1218,28 @@ dependencies = [ [[package]] name = "bitcoin-io" -version = "0.1.2" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "340e09e8399c7bd8912f495af6aa58bea0c9214773417ffaa8f6460f93aaee56" +checksum = "2dee39a0ee5b4095224a0cfc6bf4cc1baf0f9624b96b367e53b66d974e51d953" [[package]] name = "bitcoin-units" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5285c8bcaa25876d07f37e3d30c303f2609179716e11d688f51e8f1fe70063e2" +checksum = "346568ebaab2918487cea76dd55dae13c27bb618cdb737c952e69eb2017c4118" dependencies = [ - "bitcoin-internals 0.3.0", + "bitcoin-internals", "serde", ] [[package]] name = "bitcoin_hashes" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1930a4dabfebb8d7d9992db18ebe3ae2876f0a305fab206fd168df931ede293b" -dependencies = [ - "bitcoin-internals 0.2.0", - "hex-conservative 0.1.2", -] - -[[package]] -name = "bitcoin_hashes" -version = "0.14.0" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb18c03d0db0247e147a21a6faafd5a7eb851c743db062de72018b6b7e8e4d16" +checksum = "26ec84b80c482df901772e931a9a681e26a1b9ee2302edeff23cb30328745c8b" dependencies = [ "bitcoin-io", - "hex-conservative 0.2.1", + "hex-conservative", "serde", ] @@ -846,9 +1251,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.6.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "blake2" @@ -878,10 +1283,67 @@ dependencies = [ ] [[package]] -name = "bstr" -version = "1.12.0" +name = "bon" +version = "3.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f47dbe92550676ee653353c310dfb9cf6ba17ee70396e1f7cf0a2020ad49b2fe" +dependencies = [ + "bon-macros", + "rustversion", +] + +[[package]] +name = "bon-macros" +version = "3.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "519bd3116aeeb42d5372c29d982d16d0170d3d4a5ed85fc7dd91642ffff3c67c" +dependencies = [ + "darling 0.23.0", + "ident_case", + "prettyplease", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.117", +] + +[[package]] +name = "brotli" +version = "8.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "234113d19d0d7d613b40e86fb654acf958910802bcceab913a4f9e7cda03b1a4" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + +[[package]] +name = "bstd" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "num", + "rand 0.8.6", + "regex", +] + +[[package]] +name = "bstr" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab" dependencies = [ "memchr", "regex-automata", @@ -890,9 +1352,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.16.0" +version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" [[package]] name = "byteorder" @@ -927,9 +1389,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.60" +version = "1.2.61" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" +checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" dependencies = [ "find-msvc-tools", "jobserver", @@ -939,9 +1401,9 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" @@ -975,17 +1437,16 @@ dependencies = [ [[package]] name = "chrono" -version = "0.4.38" +version = "0.4.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" dependencies = [ - "android-tzdata", "iana-time-zone", "js-sys", "num-traits", "serde", "wasm-bindgen", - "windows-targets 0.52.6", + "windows-link", ] [[package]] @@ -1012,7 +1473,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" dependencies = [ "ciborium-io", - "half 2.4.1", + "half 2.7.1", ] [[package]] @@ -1026,6 +1487,46 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim 0.11.1", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmac" version = "0.7.2" @@ -1046,6 +1547,85 @@ dependencies = [ "cc", ] +[[package]] +name = "cmsketch" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7ee2cfacbd29706479902b06d75ad8f1362900836aa32799eabc7e004bfd854" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "colored" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" +dependencies = [ + "lazy_static", + "windows-sys 0.59.0", +] + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + +[[package]] +name = "const-fnv1a-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32b13ea120a812beba79e34316b3942a857c86ec1593cb34f27bb28272ce2cca" + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "convert_case" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -1072,38 +1652,98 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "cpufeatures" -version = "0.2.13" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51e852e6dc9a5bed1fae92dd2375037bf2b768725bf3be87811edee3249d09ad" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" -version = "0.8.20" +version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "crypto-common" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", - "rand_core", + "rand_core 0.6.4", "typenum", ] +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "ctr" version = "0.9.2" @@ -1136,51 +1776,121 @@ checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "darling" -version = "0.20.10" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f63b86c8a8826a49b8c21f08a2d07338eec8d900540f8630dc76284be802989" +checksum = "7b750cb3417fd1b327431a470f388520309479ab0bf5e323505daf0290cd3850" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.14.4", + "darling_macro 0.14.4", +] + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", ] [[package]] name = "darling_core" -version = "0.20.10" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95133861a8032aaea082871032f5815eb9e98cef03fa916ab4500513994df9e5" +checksum = "109c1ca6e6b7f82cc233a97004ea8ed7ca123a9af07a8230878fcfda9b158bf0" dependencies = [ "fnv", "ident_case", "proc-macro2", "quote", - "strsim", - "syn 2.0.106", + "strsim 0.10.0", + "syn 1.0.109", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim 0.11.1", + "syn 2.0.117", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim 0.11.1", + "syn 2.0.117", ] [[package]] name = "darling_macro" -version = "0.20.10" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" +checksum = "a4aab4dbc9f7611d8b55048a3a16d2d010c2c8334e46304b40ac1cc14bf3b48e" dependencies = [ - "darling_core", + "darling_core 0.14.4", "quote", - "syn 2.0.106", + "syn 1.0.109", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn 2.0.117", ] [[package]] name = "dashmap" -version = "5.5.3" +version = "6.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" dependencies = [ "cfg-if", + "crossbeam-utils", "hashbrown 0.14.5", "lock_api", "once_cell", @@ -1189,9 +1899,9 @@ dependencies = [ [[package]] name = "data-encoding" -version = "2.6.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8566979429cf69b49a5c740c60791108e86440e8be149bbea4fe54d2c32d6e2" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "dbl" @@ -1232,7 +1942,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65e13bab2796f412722112327f3e575601a3e9cdcbe426f0d30dbf43f3f5dc71" dependencies = [ "bigdecimal", - "bitflags 2.6.0", + "bitflags 2.11.1", "byteorder", "chrono", "diesel_derives", @@ -1255,20 +1965,20 @@ dependencies = [ "heck 0.4.1", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "diesel_derives" -version = "2.2.2" +version = "2.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6ff2be1e7312c858b2ef974f5c7089833ae57b5311b334b30923af58e5718d8" +checksum = "1b96984c469425cb577bf6f17121ecb3e4fe1e81de5d8f780dd372802858d756" dependencies = [ "diesel_table_macro_syntax", "dsl_auto_type", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -1277,7 +1987,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "209c735641a413bc68c4923a9d6ad4bcb3ca306b794edaa7eb0b3228a99ffb25" dependencies = [ - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -1291,6 +2001,27 @@ dependencies = [ "subtle", ] +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.61.2", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -1299,7 +2030,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -1308,18 +2039,72 @@ version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + [[package]] name = "dsl_auto_type" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5d9abe6314103864cc2d8901b7ae224e0ab1a103a0a416661b4097b0779b607" +checksum = "139ae9aca7527f85f26dd76483eb38533fd84bd571065da1739656ef71c5ff5b" dependencies = [ - "darling", + "darling 0.20.11", "either", "heck 0.5.0", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", +] + +[[package]] +name = "dspy-rs" +version = "0.7.3" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "arrow", + "async-trait", + "bamltype", + "bon", + "csv", + "dsrs_macros", + "enum_dispatch", + "facet", + "foyer", + "futures", + "hf-hub", + "indexmap", + "kdam", + "parquet", + "rand 0.8.6", + "rayon", + "regex", + "reqwest 0.12.28", + "rig-core", + "rstest", + "schemars", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "dsrs_macros" +version = "0.7.2" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.117", ] [[package]] @@ -1328,35 +2113,80 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "ecow" -version = "0.2.2" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54bfbb1708988623190a6c4dbedaeaf0f53c20c6395abd6a01feb327b3146f4b" +checksum = "78e4f79b296fbaab6ce2e22d52cb4c7f010fe0ebe7a32e34fa25885fd797bd02" dependencies = [ "serde", ] [[package]] name = "either" -version = "1.13.0" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "encode_unicode" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" [[package]] name = "encoding_rs" -version = "0.8.34" +version = "0.8.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b45de904aa0b010bce2ab45264d0631681847fa7b6f2eaa7dab7619943bc4f59" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" dependencies = [ "cfg-if", ] [[package]] -name = "equivalent" +name = "enum_dispatch" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa18ce2bc66555b3218614519ac839ddb759a7d6720732f979ef8d13be147ecd" +dependencies = [ + "once_cell", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "env_filter" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +dependencies = [ + "log", +] + +[[package]] +name = "env_logger" +version = "0.11.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "log", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" @@ -1368,6 +2198,114 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + +[[package]] +name = "eventsource-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" +dependencies = [ + "futures-core", + "nom", + "pin-project-lite", +] + +[[package]] +name = "facet" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e338357cf598728b41e45744d024bdc063338214992361766928a1421bd7541d" +dependencies = [ + "autocfg", + "facet-core", + "facet-macros", +] + +[[package]] +name = "facet-core" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a63e0ade4c53b40220614b8fc2a0a0ce21975941b553081521a195c848b2e9c2" +dependencies = [ + "autocfg", + "const-fnv1a-hash", + "iddqd", + "impls", +] + +[[package]] +name = "facet-macro-parse" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83ea29147986d0e184600cec533c41d6065c3c3d4b5b5745a8403494ca216b09" +dependencies = [ + "facet-macro-types", + "proc-macro2", + "quote", +] + +[[package]] +name = "facet-macro-types" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b0035cf41c0d4eeee82effc9161512d216d1378dd89c4d8721258429e38597" +dependencies = [ + "proc-macro2", + "quote", + "unsynn", +] + +[[package]] +name = "facet-macros" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a784f2fa36d3165b95639af790249dee0d8efdef7d53f9417cace91697e2e3" +dependencies = [ + "facet-macros-impl", +] + +[[package]] +name = "facet-macros-impl" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8f45c6380398bf74e59b97a20012de571502c609e580d84579d1140e491c1c" +dependencies = [ + "facet-macro-parse", + "facet-macro-types", + "proc-macro2", + "quote", + "unsynn", +] + +[[package]] +name = "facet-reflect" +version = "0.43.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4418c9fceaac9adcd055cc3732954d79b5d67ef04fb855dd219f2b314ba26cff" +dependencies = [ + "facet-core", +] + [[package]] name = "fancy-regex" version = "0.12.0" @@ -1378,6 +2316,16 @@ dependencies = [ "regex", ] +[[package]] +name = "fastant" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" +dependencies = [ + "small_ctor", + "web-time", +] + [[package]] name = "fastrand" version = "2.4.1" @@ -1396,12 +2344,57 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags 2.11.1", + "rustc_version", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "flume" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095" +dependencies = [ + "futures-core", + "futures-sink", + "nanorand", + "spin", +] + [[package]] name = "fnv" version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "foreign-types" version = "0.3.2" @@ -1426,6 +2419,123 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "foyer" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50b7b23bdff0e7fecbef83438ce9389c352743ea0b8bda44fabca354f21b19a9" +dependencies = [ + "equivalent", + "foyer-common", + "foyer-memory", + "foyer-storage", + "madsim-tokio", + "mixtrics", + "pin-project", + "serde", + "thiserror 2.0.18", + "tokio", + "tracing", +] + +[[package]] +name = "foyer-common" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a72f40b0c110e4233f46df0589bf1b2b3b5ed04e4cf504a582374d7548c9c05" +dependencies = [ + "bincode", + "bytes", + "cfg-if", + "itertools", + "madsim-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "serde", + "thiserror 2.0.18", + "tokio", + "twox-hash", +] + +[[package]] +name = "foyer-intrusive-collections" +version = "0.10.0-dev" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4fee46bea69e0596130e3210e65d3424e0ac1e6df3bde6636304bdf1ca4a3b" +dependencies = [ + "memoffset 0.9.1", +] + +[[package]] +name = "foyer-memory" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1b23cefdcd78b4d729679d75bb6fb96a57f7dd94f141dcc373a2db1c04fbaa" +dependencies = [ + "arc-swap", + "bitflags 2.11.1", + "cmsketch", + "equivalent", + "foyer-common", + "foyer-intrusive-collections", + "hashbrown 0.15.5", + "itertools", + "madsim-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "serde", + "thiserror 2.0.18", + "tokio", + "tracing", +] + +[[package]] +name = "foyer-storage" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1ea5872e0073e3c37bec0aa99923f03bcec3951a056c77cfdfcc8001a370f9a" +dependencies = [ + "allocator-api2", + "anyhow", + "bytes", + "core_affinity", + "equivalent", + "fastant", + "flume", + "foyer-common", + "foyer-memory", + "fs4", + "futures-core", + "futures-util", + "hashbrown 0.15.5", + "io-uring", + "itertools", + "libc", + "lz4", + "madsim-tokio", + "parking_lot", + "pin-project", + "rand 0.9.4", + "serde", + "thiserror 2.0.18", + "tokio", + "tracing", + "twox-hash", + "zstd", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1434,9 +2544,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "futures" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" dependencies = [ "futures-channel", "futures-core", @@ -1449,9 +2559,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", "futures-sink", @@ -1459,15 +2569,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" dependencies = [ "futures-core", "futures-task", @@ -1476,32 +2586,32 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" [[package]] name = "futures-task" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" [[package]] name = "futures-timer" @@ -1511,9 +2621,9 @@ checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-channel", "futures-core", @@ -1523,7 +2633,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -1539,9 +2648,9 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.15" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", "js-sys", @@ -1558,8 +2667,21 @@ checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", "wasip2", + "wasip3", ] [[package]] @@ -1572,31 +2694,38 @@ dependencies = [ "polyval", ] +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + [[package]] name = "governor" -version = "0.6.3" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68a7f542ee6b35af73b06abc0dad1c1bae89964e4e253bc4b587b91c9637867b" +checksum = "0746aa765db78b521451ef74221663b57ba595bf83f75d0ce23cc09447c8139f" dependencies = [ "cfg-if", "dashmap", - "futures", + "futures-sink", "futures-timer", + "futures-util", "no-std-compat", "nonzero_ext", "parking_lot", "portable-atomic", "quanta", - "rand", + "rand 0.8.6", "smallvec", "spinning_top", ] [[package]] name = "h2" -version = "0.3.26" +version = "0.3.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81fe527a889e1532da5c525686d96d4c2e74cdd345badf8dfef9f6b39dd5f5e8" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" dependencies = [ "bytes", "fnv", @@ -1622,7 +2751,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http 1.1.0", + "http 1.4.0", "indexmap", "slab", "tokio", @@ -1638,12 +2767,14 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403" [[package]] name = "half" -version = "2.4.1" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", + "num-traits", + "zerocopy", ] [[package]] @@ -1654,9 +2785,29 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" [[package]] name = "hashbrown" -version = "0.16.0" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", +] + +[[package]] +name = "hashbrown" +version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" +checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" [[package]] name = "heck" @@ -1672,9 +2823,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.3.9" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" [[package]] name = "hex" @@ -1684,17 +2835,11 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "hex-conservative" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "212ab92002354b4819390025006c897e8140934349e8635c9b077f47b4dcbd20" - -[[package]] -name = "hex-conservative" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5313b072ce3c597065a808dbf612c4c8e8590bdbf8b579508bf7a762c5eae6cd" +checksum = "fda06d18ac606267c40c04e41b9947729bf8b9efe74bd4e82b61a5f26a510b9f" dependencies = [ - "arrayvec", + "arrayvec 0.7.6", ] [[package]] @@ -1703,6 +2848,30 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3011d1213f159867b13cfd6ac92d2cd5f1345762c63be3554e84092d85a50bbd" +[[package]] +name = "hf-hub" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" +dependencies = [ + "dirs", + "futures", + "http 1.4.0", + "indicatif", + "libc", + "log", + "native-tls", + "num_cpus", + "rand 0.9.4", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "ureq", + "windows-sys 0.60.2", +] + [[package]] name = "hkdf" version = "0.12.4" @@ -1734,12 +2903,11 @@ dependencies = [ [[package]] name = "http" -version = "1.1.0" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21b9ddb458710bc376481b842f5da65cdf31522de232c1ca8146abce2a358258" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" dependencies = [ "bytes", - "fnv", "itoa", ] @@ -1761,27 +2929,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http 1.1.0", + "http 1.4.0", ] [[package]] name = "http-body-util" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "793429d76616a256bcb62c2a2ec2bed781c8307e797e2598c50010f2bee2544f" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", - "futures-util", - "http 1.1.0", + "futures-core", + "http 1.4.0", "http-body 1.0.1", "pin-project-lite", ] [[package]] name = "httparse" -version = "1.9.4" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fcc0b4a115bf80b728eb8ea024ad5bd707b615bfed49e0665b6e0f86fd082d9" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" [[package]] name = "httpdate" @@ -1791,22 +2959,22 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "0.14.30" +version = "0.14.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a152ddd61dfaec7273fe8419ab357f33aee0d914c5f4efbf0d96fa749eea5ec9" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" dependencies = [ "bytes", "futures-channel", "futures-core", "futures-util", - "h2 0.3.26", + "h2 0.3.27", "http 0.2.12", "http-body 0.4.6", "httparse", "httpdate", "itoa", "pin-project-lite", - "socket2 0.5.7", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -1815,22 +2983,21 @@ dependencies = [ [[package]] name = "hyper" -version = "1.7.0" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb3aa54a13a0dfe7fbe3a59e0c76093041720fdc77b110cc0fc260fafb4dc51e" +checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" dependencies = [ "atomic-waker", "bytes", "futures-channel", "futures-core", "h2 0.4.13", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "httparse", "httpdate", "itoa", "pin-project-lite", - "pin-utils", "smallvec", "tokio", "want", @@ -1842,8 +3009,8 @@ version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http 1.1.0", - "hyper 1.7.0", + "http 1.4.0", + "hyper 1.9.0", "hyper-util", "rustls", "rustls-native-certs", @@ -1859,7 +3026,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6183ddfa99b85da61a140bea0efc93fdf56ceaa041b37d553518030827f9905" dependencies = [ "bytes", - "hyper 0.14.30", + "hyper 0.14.32", "native-tls", "tokio", "tokio-native-tls", @@ -1873,7 +3040,7 @@ checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" dependencies = [ "bytes", "http-body-util", - "hyper 1.7.0", + "hyper 1.9.0", "hyper-util", "native-tls", "tokio", @@ -1883,38 +3050,40 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.17" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c6995591a8f1380fcb4ba966a252a4b29188d51d2b89e3a252f5305be65aea8" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ "base64 0.22.1", "bytes", "futures-channel", - "futures-core", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", - "hyper 1.7.0", + "hyper 1.9.0", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.0", + "socket2 0.6.3", + "system-configuration 0.7.0", "tokio", "tower-service", "tracing", + "windows-registry", ] [[package]] name = "iana-time-zone" -version = "0.1.60" +version = "0.1.65" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7ffbb5a1b541ea2561f8c41c087286cc091e21e556a4f09a8f6cbf17b69b141" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" dependencies = [ "android_system_properties", "core-foundation-sys", "iana-time-zone-haiku", "js-sys", + "log", "wasm-bindgen", "windows-core", ] @@ -1930,21 +3099,23 @@ dependencies = [ [[package]] name = "icu_collections" -version = "1.5.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" dependencies = [ "displaydoc", + "potential_utf", + "utf8_iter", "yoke", "zerofrom", "zerovec", ] [[package]] -name = "icu_locid" -version = "1.5.0" +name = "icu_locale_core" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" dependencies = [ "displaydoc", "litemap", @@ -1953,97 +3124,78 @@ dependencies = [ "zerovec", ] -[[package]] -name = "icu_locid_transform" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_locid_transform_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_locid_transform_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" - [[package]] name = "icu_normalizer" -version = "1.5.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" dependencies = [ - "displaydoc", "icu_collections", "icu_normalizer_data", "icu_properties", "icu_provider", "smallvec", - "utf16_iter", - "utf8_iter", - "write16", "zerovec", ] [[package]] name = "icu_normalizer_data" -version = "1.5.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" [[package]] name = "icu_properties" -version = "1.5.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" dependencies = [ - "displaydoc", "icu_collections", - "icu_locid_transform", + "icu_locale_core", "icu_properties_data", "icu_provider", - "tinystr", + "zerotrie", "zerovec", ] [[package]] name = "icu_properties_data" -version = "1.5.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" [[package]] name = "icu_provider" -version = "1.5.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" dependencies = [ "displaydoc", - "icu_locid", - "icu_provider_macros", - "stable_deref_trait", - "tinystr", + "icu_locale_core", "writeable", "yoke", "zerofrom", + "zerotrie", "zerovec", ] [[package]] -name = "icu_provider_macros" -version = "1.5.0" +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "iddqd" +version = "0.3.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" +checksum = "616230c7d641ef971a3a5bfcc654c6b7524ab9c9fb665693c6a397dba9a14aca" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.106", + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "hashbrown 0.16.1", + "rustc-hash 2.1.2", ] [[package]] @@ -2065,29 +3217,50 @@ dependencies = [ [[package]] name = "idna_adapter" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" dependencies = [ "icu_normalizer", "icu_properties", ] +[[package]] +name = "impls" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a46645bbd70538861a90d0f26c31537cdf1e44aae99a794fb75a664b70951bc" + [[package]] name = "indexmap" -version = "2.11.4" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.16.0", + "hashbrown 0.17.0", + "serde", + "serde_core", +] + +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", ] [[package]] name = "inout" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0c10553d664a4d0bcff9f4215d0aac67a639cc68ef660840afe309b807bc9f5" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" dependencies = [ "block-padding", "generic-array", @@ -2103,26 +3276,85 @@ dependencies = [ ] [[package]] -name = "ipnet" -version = "2.9.0" +name = "integer-encoding" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" + +[[package]] +name = "internal-baml-diagnostics" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "colored", + "pest", + "serde", + "strsim 0.11.1", +] + +[[package]] +name = "internal-baml-jinja" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "baml-types", + "indexmap", + "minijinja", + "serde_json", + "serde_yaml", + "strum", + "toon", +] + +[[package]] +name = "io-uring" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +dependencies = [ + "bitflags 2.11.1", + "cfg-if", + "libc", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "iri-string" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" [[package]] -name = "iri-string" -version = "0.7.8" +name = "itertools" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbc5ebe9c3a1a7a5127f920a418f7585e9e758e911d0466ed004f393b0e380b2" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" dependencies = [ - "memchr", - "serde", + "either", ] [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jobserver" @@ -2136,14 +3368,39 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.81" +version = "0.3.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec48937a97411dcb524a265206ccd4c90bb711fca92b2792c407f268825b9305" +checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" dependencies = [ + "cfg-if", + "futures-util", "once_cell", "wasm-bindgen", ] +[[package]] +name = "jsonish" +version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" +dependencies = [ + "anyhow", + "baml-types", + "bstd", + "getrandom 0.2.17", + "indexmap", + "internal-baml-jinja", + "log", + "minijinja", + "regex", + "serde", + "serde_json", + "strsim 0.11.1", + "test-log", + "thiserror 2.0.18", + "unicode-normalization", + "uuid", +] + [[package]] name = "jsonwebtoken" version = "10.3.0" @@ -2152,7 +3409,7 @@ checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1" dependencies = [ "aws-lc-rs", "base64 0.22.1", - "getrandom 0.2.15", + "getrandom 0.2.17", "js-sys", "pem", "serde", @@ -2173,7 +3430,7 @@ dependencies = [ "ciborium", "hmac", "lazy_static", - "rand_core", + "rand_core 0.6.4", "secp256k1", "serde", "serde_json", @@ -2183,6 +3440,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "kdam" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d847be338ef16a13f97637c062d97fb52ebe0ff3b77fa18456d5ed366317e4f7" +dependencies = [ + "terminal_size", + "windows-sys 0.61.2", +] + [[package]] name = "keccak" version = "0.1.6" @@ -2198,45 +3465,199 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libc" -version = "0.2.177" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libm" -version = "0.2.8" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "libredox" +version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" +checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" +dependencies = [ + "libc", +] [[package]] name = "linux-raw-sys" -version = "0.4.14" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" -version = "0.7.4" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "lock_api" -version = "0.4.12" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" dependencies = [ - "autocfg", "scopeguard", ] [[package]] name = "log" -version = "0.4.22" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "lz4_flex" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" +dependencies = [ + "twox-hash", +] + +[[package]] +name = "madsim" +version = "0.2.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18351aac4194337d6ea9ffbd25b3d1540ecc0754142af1bff5ba7392d1f6f771" +dependencies = [ + "ahash", + "async-channel", + "async-stream", + "async-task", + "bincode", + "bytes", + "downcast-rs", + "errno", + "futures-util", + "lazy_static", + "libc", + "madsim-macros", + "naive-timer", + "panic-message", + "rand 0.8.6", + "rand_xoshiro", + "rustversion", + "serde", + "spin", + "tokio", + "tokio-util", + "toml", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "madsim-macros" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3d248e97b1a48826a12c3828d921e8548e714394bf17274dd0a93910dc946e1" +dependencies = [ + "darling 0.14.4", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "madsim-tokio" +version = "0.2.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +checksum = "7d3eb2acc57c82d21d699119b859e2df70a91dbdb84734885a1e72be83bdecb5" +dependencies = [ + "madsim", + "spin", + "tokio", +] [[package]] name = "matchers" @@ -2261,14 +3682,14 @@ checksum = "5cf92c10c7e361d6b99666ec1c6f9805b0bea2c3bd8c78dc6fe98ac5bd78db11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "memchr" -version = "2.7.4" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "memoffset" @@ -2294,22 +3715,64 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "mime_guess" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e" +dependencies = [ + "mime", + "unicase", +] + +[[package]] +name = "minijinja" +version = "2.16.0" +source = "git+https://github.com/boundaryml/minijinja.git?branch=main#b9afca428b1c8149b1b3a5aab26a32d09744cd83" +dependencies = [ + "aho-corasick", + "indexmap", + "serde", + "serde_json", + "unicase", + "unicode-ident", +] + [[package]] name = "minimal-lexical" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" -version = "1.0.2" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80e04d1dcff3aae0704555fe5fee3bcfaf3d1fdf8a7e521d5b9d2b42acb52cec" +checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ - "hermit-abi", "libc", "wasi", - "windows-sys 0.52.0", + "windows-sys 0.61.2", +] + +[[package]] +name = "mixtrics" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb252c728b9d77c6ef9103f0c81524fa0a3d3b161d0a936295d7fbeff6e04c11" +dependencies = [ + "itertools", + "parking_lot", ] [[package]] @@ -2321,7 +3784,7 @@ dependencies = [ "bytes", "encoding_rs", "futures-util", - "http 1.1.0", + "http 1.4.0", "httparse", "memchr", "mime", @@ -2329,19 +3792,40 @@ dependencies = [ "version_check", ] +[[package]] +name = "mutants" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc0287524726960e07b119cebd01678f852f147742ae0d925e6a520dca956126" + +[[package]] +name = "naive-timer" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "034a0ad7deebf0c2abcf2435950a6666c3c15ea9d8fad0c0f48efa8a7f843fed" + +[[package]] +name = "nanorand" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a51313c5820b0b02bd422f4b44776fbf47961755c74ce64afc73bfad10226c3" +dependencies = [ + "getrandom 0.2.17", +] + [[package]] name = "native-tls" -version = "0.2.12" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8614eb2c83d59d1c8cc974dd3f920198647674a0a035e1af1fa58707e317466" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" dependencies = [ "libc", "log", "openssl", - "openssl-probe 0.1.5", + "openssl-probe", "openssl-sys", "schannel", - "security-framework 2.11.1", + "security-framework", "security-framework-sys", "tempfile", ] @@ -2361,11 +3845,11 @@ dependencies = [ [[package]] name = "nix" -version = "0.29.0" +version = "0.31.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "cfg-if", "cfg_aliases", "libc", @@ -2403,6 +3887,20 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.6" @@ -2413,6 +3911,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.1" @@ -2428,6 +3935,28 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -2435,8 +3964,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "oauth2" version = "4.4.2" @@ -2445,15 +3991,15 @@ checksum = "c38841cdd844847e3e7c8d29cef9dcfed8877f8f56f9071f77843ecf3baf937f" dependencies = [ "base64 0.13.1", "chrono", - "getrandom 0.2.15", + "getrandom 0.2.17", "http 0.2.12", - "rand", + "rand 0.8.6", "reqwest 0.11.27", "serde", "serde_json", "serde_path_to_error", "sha2", - "thiserror", + "thiserror 1.0.69", "url", ] @@ -2468,9 +4014,15 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.19.0" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "opaque-debug" @@ -2505,13 +4057,14 @@ dependencies = [ "diesel", "diesel-derive-enum", "dotenv", + "dspy-rs", "futures", "generic-array", - "getrandom 0.2.15", + "getrandom 0.2.17", "hex", "hkdf", "hmac", - "hyper 0.14.30", + "hyper 0.14.32", "hyper-tls 0.5.0", "jsonwebtoken", "jwt-compact", @@ -2520,7 +4073,7 @@ dependencies = [ "once_cell", "openssl", "password-auth", - "rand_core", + "rand_core 0.6.4", "rcgen", "regex", "reqwest 0.11.27", @@ -2532,7 +4085,7 @@ dependencies = [ "serde_json", "sha2", "subtle", - "thiserror", + "thiserror 1.0.69", "tiktoken-rs", "tokio", "tokio-stream", @@ -2554,7 +4107,7 @@ version = "0.10.78" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38c4372413cdaaf3cc79dd92d29d7d9f5ab09b51b10dded508fb90bb70b9222" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "cfg-if", "foreign-types", "libc", @@ -2571,15 +4124,9 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] -[[package]] -name = "openssl-probe" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" - [[package]] name = "openssl-probe" version = "0.2.1" @@ -2598,17 +4145,53 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "ordered-float" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7d950ca161dc355eaf28f82b11345ed76c6e1f6eb1f4f4479e0323b9e2fbd0e" +dependencies = [ + "num-traits", +] + [[package]] name = "outref" -version = "0.5.1" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + +[[package]] +name = "panic-message" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384e52fd8fbd4cbe3c317e8216260c21a0f9134de108cea8a4dd4e7e152c472d" + +[[package]] +name = "parking" +version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" [[package]] name = "parking_lot" -version = "0.12.3" +version = "0.12.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" dependencies = [ "lock_api", "parking_lot_core", @@ -2616,15 +4199,48 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.10" +version = "0.9.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-targets 0.52.6", + "windows-link", +] + +[[package]] +name = "parquet" +version = "56.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0dbd48ad52d7dccf8ea1b90a3ddbfaea4f69878dd7683e51c507d4bc52b5b27" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64 0.22.1", + "brotli", + "bytes", + "chrono", + "flate2", + "half 2.7.1", + "hashbrown 0.16.1", + "lz4_flex", + "num", + "num-bigint", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "twox-hash", + "zstd", ] [[package]] @@ -2634,9 +4250,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a2a4764cc1f8d961d802af27193c6f4f0124bd0e76e8393cf818e18880f0524" dependencies = [ "argon2", - "getrandom 0.2.15", + "getrandom 0.2.17", "password-hash", - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -2646,18 +4262,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "346f04948ba92c43e8469c1ee6736c7563d71012b17d40745260fe106aac2166" dependencies = [ "base64ct", - "rand_core", + "rand_core 0.6.4", "subtle", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pem" -version = "3.0.4" +version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e459365e590736a54c3fa561947c84837534b8e9af6fc5bf781307e82658fae" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ "base64 0.22.1", - "serde", + "serde_core", ] [[package]] @@ -2666,31 +4288,41 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pest" +version = "2.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0848c601009d37dfa3430c4666e147e49cdcf1b92ecd3e63657d8a5f19da662" +dependencies = [ + "memchr", + "ucd-trie", +] + [[package]] name = "pin-project" -version = "1.1.5" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6bf43b791c5b9e34c3d182969b4abb522f9343702850a2e57f460d00d09b4b3" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.5" +version = "1.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f38a4412a78282e09a2cf38d195ea5420d15ba0602cb375210efbc877243965" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "pin-project-lite" -version = "0.2.14" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "pin-utils" @@ -2700,9 +4332,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pkg-config" -version = "0.3.30" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] name = "poly1305" @@ -2729,9 +4361,18 @@ dependencies = [ [[package]] name = "portable-atomic" -version = "1.9.0" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "potential_utf" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] [[package]] name = "powerfmt" @@ -2741,22 +4382,52 @@ checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" [[package]] name = "ppv-lite86" -version = "0.2.20" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ "zerocopy", ] [[package]] name = "pq-sys" -version = "0.4.8" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31c0052426df997c0cbd30789eb44ca097e3541717a7b8fa36b1c464ee7edebd" +checksum = "f6cc05d7ea95200187117196eee9edd0644424911821aeb28a18ce60ea0b8793" dependencies = [ "vcpkg", ] +[[package]] +name = "pretty" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d22152487193190344590e4f30e219cf3fe140d9e7a3fdb683d82aa2c5f4156" +dependencies = [ + "arrayvec 0.5.2", + "typed-arena", + "unicode-width", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.117", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro-error-attr2" version = "2.0.0" @@ -2776,23 +4447,23 @@ dependencies = [ "proc-macro-error-attr2", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "proc-macro2" -version = "1.0.101" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] [[package]] name = "quanta" -version = "0.12.3" +version = "0.12.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e5167a477619228a0b284fac2674e3c388cba90631d7b7de620e6f1fcd08da5" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" dependencies = [ "crossbeam-utils", "libc", @@ -2805,9 +4476,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.41" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" dependencies = [ "proc-macro2", ] @@ -2818,6 +4489,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "r2d2" version = "0.8.10" @@ -2836,8 +4513,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -2847,7 +4534,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -2856,23 +4553,61 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom 0.2.15", + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_xoshiro" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa" +dependencies = [ + "rand_core 0.6.4", ] [[package]] name = "raw-cpuid" -version = "11.1.0" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.11.1", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb9ee317cfe3fbd54b36a511efc1edd42e216903c9cd575e686dd68a2ba90d8d" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" dependencies = [ - "bitflags 2.6.0", + "crossbeam-deque", + "crossbeam-utils", ] [[package]] name = "rcgen" -version = "0.13.1" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54077e1872c46788540de1ea3d7f4ccb1983d12f9aa909b234468676c1a36779" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" dependencies = [ "pem", "ring", @@ -2883,18 +4618,49 @@ dependencies = [ [[package]] name = "redox_syscall" -version = "0.5.3" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags 2.11.1", +] + +[[package]] +name = "redox_users" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror 2.0.18", +] + +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a908a6e00f1fdd0dfd9c0eb08ce85126f6d8bbda50017e74bc4a4b7d4a926a4" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" dependencies = [ - "bitflags 2.6.0", + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] name = "regex" -version = "1.11.3" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b5288124840bee7b386bc413c487869b360b2b4ec421ea56425128692f2a82c" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -2904,9 +4670,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.11" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "833eb9ce86d40ef33cb1306d8accf7bc8ec2bfea4355cbdebb3df68b40925cad" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", @@ -2915,15 +4681,21 @@ dependencies = [ [[package]] name = "regex-lite" -version = "0.1.6" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.6" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + +[[package]] +name = "relative-path" +version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" [[package]] name = "reqwest" @@ -2936,10 +4708,10 @@ dependencies = [ "encoding_rs", "futures-core", "futures-util", - "h2 0.3.26", + "h2 0.3.27", "http 0.2.12", "http-body 0.4.6", - "hyper 0.14.30", + "hyper 0.14.32", "hyper-tls 0.5.0", "ipnet", "js-sys", @@ -2954,7 +4726,7 @@ dependencies = [ "serde_json", "serde_urlencoded", "sync_wrapper 0.1.2", - "system-configuration", + "system-configuration 0.5.1", "tokio", "tokio-native-tls", "tower-service", @@ -2967,21 +4739,28 @@ dependencies = [ [[package]] name = "reqwest" -version = "0.12.23" +version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d429f34c8092b2d42c7c93cec323bb4adeb7c67698f70839adec842ec10c7ceb" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64 0.22.1", "bytes", + "encoding_rs", + "futures-channel", "futures-core", - "http 1.1.0", + "futures-util", + "h2 0.4.13", + "http 1.4.0", "http-body 1.0.1", "http-body-util", - "hyper 1.7.0", + "hyper 1.9.0", + "hyper-rustls", "hyper-tls 0.6.0", "hyper-util", "js-sys", "log", + "mime", + "mime_guess", "native-tls", "percent-encoding", "pin-project-lite", @@ -2989,31 +4768,62 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", - "sync_wrapper 1.0.1", + "sync_wrapper 1.0.2", "tokio", "tokio-native-tls", - "tower 0.5.2", - "tower-http 0.6.6", + "tokio-util", + "tower", + "tower-http 0.6.8", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", ] [[package]] name = "resend-rs" -version = "0.9.1" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54524581cce8b45f854c9add73cbe9449448a50ab8a655c2c3a477abf0ddf7e6" +checksum = "1c8351b9d8013cebbb27626a583a31470faebc187b044098c6eadb5daffc16a4" dependencies = [ "ecow", "governor", "maybe-async", - "rand", - "reqwest 0.12.23", + "rand 0.8.6", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "rig-core" +version = "0.26.0" +source = "git+https://github.com/0xPlaygrounds/rig?rev=e7849df#e7849df20bce7c4eb3eb5ada87aae18435f9afa1" +dependencies = [ + "as-any", + "async-stream", + "base64 0.22.1", + "bytes", + "eventsource-stream", + "futures", + "futures-timer", + "glob", + "http 1.4.0", + "mime_guess", + "ordered-float 5.3.0", + "pin-project-lite", + "reqwest 0.12.28", + "schemars", "serde", - "thiserror", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", + "tracing-futures", + "url", ] [[package]] @@ -3024,18 +4834,54 @@ checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" dependencies = [ "cc", "cfg-if", - "getrandom 0.2.15", + "getrandom 0.2.17", "libc", "untrusted 0.9.0", "windows-sys 0.52.0", ] +[[package]] +name = "rstest" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fc39292f8613e913f7df8fa892b8944ceb47c247b78e1b1ae2f09e019be789d" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", + "rustc_version", +] + +[[package]] +name = "rstest_macros" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f168d99749d307be9de54d23fd226628d99768225ef08f6ffb52e0182a27746" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.117", + "unicode-ident", +] + [[package]] name = "rustc-hash" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + [[package]] name = "rustc_version" version = "0.4.1" @@ -3056,15 +4902,15 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.34" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70dc5ec042f7a43c4a73241207cecc9873a06d45debb38b329f8541d85c2730f" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3074,7 +4920,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c2c118cb077cca2822033836dfb1b975355dfb784b5e8da48f7b6c5db74e60e" dependencies = [ "aws-lc-rs", + "log", "once_cell", + "ring", "rustls-pki-types", "rustls-webpki", "subtle", @@ -3087,10 +4935,10 @@ version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" dependencies = [ - "openssl-probe 0.2.1", + "openssl-probe", "rustls-pki-types", "schannel", - "security-framework 3.6.0", + "security-framework", ] [[package]] @@ -3104,9 +4952,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.0" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" dependencies = [ "zeroize", ] @@ -3131,17 +4979,17 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "ryu" -version = "1.0.18" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "schannel" -version = "0.1.23" +version = "0.1.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc91545643bcf3a0bbb6569265615222618bdf33ce4ffbbd13c4bbd4c093534" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3153,6 +5001,31 @@ dependencies = [ "parking_lot", ] +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.117", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3161,45 +5034,42 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "secp256k1" -version = "0.29.0" +version = "0.29.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e0cc0f1cf93f4969faf3ea1c7d8a9faed25918d96affa959720823dfe86d4f3" +checksum = "9465315bc9d4566e1724f0fffcbcc446268cb522e60f9a27bcded6b19c108113" dependencies = [ - "bitcoin_hashes 0.14.0", - "rand", + "bitcoin_hashes", + "rand 0.8.6", "secp256k1-sys", "serde", ] [[package]] name = "secp256k1-sys" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1433bd67156263443f14d603720b082dd3121779323fce20cba2aa07b874bc1b" +checksum = "d4387882333d3aa8cb20530a17c69a3752e97837832f34f6dccc760e715001d9" dependencies = [ "cc", ] [[package]] -name = "security-framework" -version = "2.11.1" +name = "secrecy" +version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" +checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a" dependencies = [ - "bitflags 2.6.0", - "core-foundation 0.9.4", - "core-foundation-sys", - "libc", - "security-framework-sys", + "serde", + "zeroize", ] [[package]] name = "security-framework" -version = "3.6.0" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d17b898a6d6948c3a8ee4372c17cb384f90d2e6e912ef00895b14fd7ab54ec38" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "core-foundation 0.10.1", "core-foundation-sys", "libc", @@ -3218,9 +5088,15 @@ dependencies = [ [[package]] name = "semver" -version = "1.0.23" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + +[[package]] +name = "seq-macro" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" @@ -3234,11 +5110,12 @@ dependencies = [ [[package]] name = "serde_bytes" -version = "0.11.15" +version = "0.11.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "387cc504cb06bb40a96c8e04e951fe01854cf6bc921053c954e4a606d9675c6a" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" dependencies = [ "serde", + "serde_core", ] [[package]] @@ -3268,30 +5145,52 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", +] + +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] name = "serde_json" -version = "1.0.145" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap", "itoa", "memchr", - "ryu", "serde", "serde_core", + "zmij", ] [[package]] name = "serde_path_to_error" -version = "0.1.16" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af99884400da37c88f5e9146b7f1fd0fbcae8f6eec4e9da38b67d05486f814a6" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" dependencies = [ "itoa", "serde", + "serde_core", +] + +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", ] [[package]] @@ -3306,11 +5205,24 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "sha2" -version = "0.10.8" +version = "0.10.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures", @@ -3319,9 +5231,9 @@ dependencies = [ [[package]] name = "sha3" -version = "0.10.8" +version = "0.10.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60" +checksum = "77fd7028345d415a4034cf8777cd4f8ab1851274233b45f84e3d955502d93874" dependencies = [ "digest", "keccak", @@ -3344,10 +5256,11 @@ checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "signal-hook-registry" -version = "1.4.2" +version = "1.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" dependencies = [ + "errno", "libc", ] @@ -3357,41 +5270,62 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ - "rand_core", + "rand_core 0.6.4", ] +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "simple_asn1" -version = "0.6.2" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adc4e5204eb1910f40f9cfa375f6f05b68c3abac4b6fd879c8ff5e7ae8a0a085" +checksum = "0d585997b0ac10be3c5ee635f1bab02d512760d14b7c468801ac8a01d9ae5f1d" dependencies = [ "num-bigint", "num-traits", - "thiserror", + "thiserror 2.0.18", "time", ] [[package]] -name = "slab" -version = "0.4.9" +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "small_ctor" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" + +[[package]] +name = "smallvec" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" -dependencies = [ - "autocfg", -] +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] -name = "smallvec" -version = "1.13.2" +name = "snap" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" [[package]] name = "socket2" -version = "0.5.7" +version = "0.5.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce305eb0b4296696835b71df73eb912e0f1ffd2556a501fcede6e0c50349191c" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" dependencies = [ "libc", "windows-sys 0.52.0", @@ -3399,12 +5333,23 @@ dependencies = [ [[package]] name = "socket2" -version = "0.6.0" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "233504af464074f9d066d7b5416c5f9b894a5862a6506e306f7b816cdd6f1807" +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", +] + +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", ] [[package]] @@ -3412,6 +5357,9 @@ name = "spin" version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] [[package]] name = "spinning_top" @@ -3424,9 +5372,15 @@ dependencies = [ [[package]] name = "stable_deref_trait" -version = "1.2.0" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "strsim" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" [[package]] name = "strsim" @@ -3434,6 +5388,28 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.117", +] + [[package]] name = "subtle" version = "2.6.1" @@ -3453,9 +5429,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.106" +version = "2.0.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" dependencies = [ "proc-macro2", "quote", @@ -3470,9 +5446,9 @@ checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" [[package]] name = "sync_wrapper" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7065abeca94b6a8a577f9bd45aa0867a2238b74e8eb67cf10d492bc39351394" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" dependencies = [ "futures-core", ] @@ -3491,13 +5467,13 @@ dependencies = [ [[package]] name = "synstructure" -version = "0.13.1" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -3508,7 +5484,18 @@ checksum = "ba3a3adc5c275d719af8cb4272ea1c4a6d668a777f37e115f6d11ddbc1c8e0e7" dependencies = [ "bitflags 1.3.2", "core-foundation 0.9.4", - "system-configuration-sys", + "system-configuration-sys 0.5.0", +] + +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags 2.11.1", + "core-foundation 0.9.4", + "system-configuration-sys 0.6.0", ] [[package]] @@ -3521,47 +5508,129 @@ dependencies = [ "libc", ] +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "tempfile" -version = "3.12.0" +version = "3.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbcdd0c794ebb0d4cf35e88edd2f7d2c4c3e9a5a6dab322839b321c6a87a64" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ - "cfg-if", "fastrand", + "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", +] + +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "test-log" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f46bf474f0a4afebf92f076d54fd5e63423d9438b8c278a3d2ccb0f47f7cdb3" +dependencies = [ + "env_logger", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-core" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37d4d41320b48bc4a211a9021678fcc0c99569b594ea31c93735b8e517102b4c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "test-log-macros" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9beb9249a81e430dffd42400a49019bcf548444f1968ff23080a625de0d4d320" +dependencies = [ + "syn 2.0.117", + "test-log-core", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", ] [[package]] name = "thiserror" -version = "1.0.63" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ - "thiserror-impl", + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] name = "thiserror-impl" -version = "1.0.63" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "thread_local" -version = "1.1.8" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" dependencies = [ "cfg-if", - "once_cell", +] + +[[package]] +name = "thrift" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" +dependencies = [ + "byteorder", + "integer-encoding", + "ordered-float 2.10.1", ] [[package]] @@ -3576,7 +5645,7 @@ dependencies = [ "fancy-regex", "lazy_static", "parking_lot", - "rustc-hash", + "rustc-hash 1.1.0", ] [[package]] @@ -3610,11 +5679,20 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" -version = "0.7.6" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" dependencies = [ "displaydoc", "zerovec", @@ -3622,9 +5700,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.8.0" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" dependencies = [ "tinyvec_macros", ] @@ -3637,9 +5715,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.50.0" +version = "1.52.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +checksum = "b67dee974fe86fd92cc45b7a95fdd2f99a36a6d7b0d431a231178d3d670bbcc6" dependencies = [ "bytes", "libc", @@ -3647,20 +5725,20 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.6.0", + "socket2 0.6.3", "tokio-macros", "windows-sys 0.61.2", ] [[package]] name = "tokio-macros" -version = "2.6.1" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] @@ -3685,9 +5763,9 @@ dependencies = [ [[package]] name = "tokio-stream" -version = "0.1.17" +version = "0.1.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eca58d7bba4a75707817a2c44174253f9236b2d5fbd055602e9d5c07c139a047" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" dependencies = [ "futures-core", "pin-project-lite", @@ -3696,9 +5774,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.16" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14307c986784f72ef81c89db7d9e28d6ac26d16213b109ea501696195e6e3ce5" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", @@ -3708,34 +5786,89 @@ dependencies = [ ] [[package]] -name = "tower" -version = "0.4.13" +name = "toml" +version = "0.9.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" dependencies = [ - "futures-core", - "futures-util", - "pin-project", - "pin-project-lite", - "tokio", - "tower-layer", - "tower-service", - "tracing", + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.11+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" +dependencies = [ + "indexmap", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "winnow 1.0.2", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow 1.0.2", +] + +[[package]] +name = "toml_writer" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" + +[[package]] +name = "toon" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fa605a2aa68fefdedc9560bd765ce3ad89d45ac1d3ebad43724621098c9299a" +dependencies = [ + "regex", + "serde_json", ] [[package]] name = "tower" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", "pin-project-lite", - "sync_wrapper 1.0.1", + "sync_wrapper 1.0.2", "tokio", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -3744,9 +5877,9 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "bytes", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "http-body-util", "pin-project-lite", @@ -3756,18 +5889,18 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.6" +version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.11.1", "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "iri-string", "pin-project-lite", - "tower 0.5.2", + "tower", "tower-layer", "tower-service", ] @@ -3786,9 +5919,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracing" -version = "0.1.41" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "log", "pin-project-lite", @@ -3798,25 +5931,37 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.30" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81383ab64e72a7a8b8e13130c49e3dab29def6d0c7d76a03087b3cf71c5c6903" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "tracing-core" -version = "0.1.34" +version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d12581f227e93f094d3af2ae690a574abb8a2b9b7a96e7cfe9647b2b617678" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "futures", + "futures-task", + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -3830,9 +5975,9 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.20" +version = "0.3.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2054a14f5307d601f88daf0553e1cbf472acc4f2c51afab632431cdcd72124d5" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" dependencies = [ "matchers", "nu-ansi-term", @@ -3852,27 +5997,78 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" +dependencies = [ + "rand 0.9.4", +] + +[[package]] +name = "type-safe-id" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa0f2333ec5a49706b6f29896d0e4cdba1665d2bf2d752d0c4f130634935d42f" +dependencies = [ + "arrayvec 0.7.6", + "rand 0.8.6", + "thiserror 2.0.18", + "uuid", +] + +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typenum" -version = "1.17.0" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" + +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + +[[package]] +name = "unicase" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" [[package]] name = "unicode-ident" -version = "1.0.19" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-normalization" -version = "0.1.22" +version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-segmentation" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" @@ -3889,6 +6085,23 @@ dependencies = [ "subtle", ] +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + +[[package]] +name = "unsynn" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "501a7adf1a4bd9951501e5c66621e972ef8874d787628b7f90e64f936ef7ec0a" +dependencies = [ + "mutants", + "proc-macro2", + "rustc-hash 2.1.2", +] + [[package]] name = "untrusted" version = "0.7.1" @@ -3901,6 +6114,26 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "native-tls", + "once_cell", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + [[package]] name = "url" version = "2.5.8" @@ -3920,26 +6153,28 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" -[[package]] -name = "utf16_iter" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" - [[package]] name = "utf8_iter" version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" -version = "1.10.0" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ - "getrandom 0.2.15", - "serde", + "getrandom 0.4.2", + "js-sys", + "serde_core", + "wasm-bindgen", ] [[package]] @@ -3964,19 +6199,19 @@ version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7df16e474ef958526d1205f6dda359fdfab79d9aa6d54bafcb92dcd07673dca" dependencies = [ - "darling", + "darling 0.20.11", "once_cell", "proc-macro-error2", "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "valuable" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" [[package]] name = "vcpkg" @@ -3998,12 +6233,12 @@ checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" [[package]] name = "vsock" -version = "0.5.1" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e8b4d00e672f147fc86a09738fadb1445bd1c0a40542378dfb82909deeee688" +checksum = "6ba782755fc073877e567c2253c0be48e4aa9a254c232d36d3985dfae0bd5205" dependencies = [ "libc", - "nix 0.29.0", + "nix 0.31.2", ] [[package]] @@ -4017,9 +6252,9 @@ dependencies = [ [[package]] name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" @@ -4027,54 +6262,46 @@ version = "1.0.3+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.57.1", ] [[package]] -name = "wasm-bindgen" -version = "0.2.104" +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1da10c01ae9f1ae40cbfac0bac3b1e724b320abfcf52229f80b547c0d250e2d" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "cfg-if", - "once_cell", - "rustversion", - "wasm-bindgen-macro", - "wasm-bindgen-shared", + "wit-bindgen 0.51.0", ] [[package]] -name = "wasm-bindgen-backend" -version = "0.2.104" +name = "wasm-bindgen" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "671c9a5a66f49d8a47345ab942e2cb93c7d1d0339065d4f8139c486121b43b19" +checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" dependencies = [ - "bumpalo", - "log", - "proc-macro2", - "quote", - "syn 2.0.106", + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-futures" -version = "0.4.54" +version = "0.4.68" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e038d41e478cc73bae0ff9b36c60cff1c98b8f38f8d7e8061e79ee63608ac5c" +checksum = "f371d383f2fb139252e0bfac3b81b265689bf45b6874af544ffa4c975ac1ebf8" dependencies = [ - "cfg-if", "js-sys", - "once_cell", "wasm-bindgen", - "web-sys", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.104" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ca60477e4c59f5f2986c50191cd972e3a50d8a95603bc9434501cf156a9a119" +checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4082,36 +6309,111 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.104" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f07d2f20d4da7b26400c9f4a0511e6e0345b040694e8a75bd41d578fa4421d7" +checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" dependencies = [ + "bumpalo", "proc-macro2", "quote", - "syn 2.0.106", - "wasm-bindgen-backend", + "syn 2.0.117", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.104" +version = "0.2.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bad67dc8b2a1a6e5448428adec4c3e84c43e561d8c9ee8a9e5aabeb193ec41d1" +checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags 2.11.1", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + [[package]] name = "web-sys" -version = "0.3.81" +version = "0.3.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9367c417a924a74cae129e6a2ae3b47fabb1f8995595ab474029da749a8be120" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" dependencies = [ "js-sys", "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.7", +] + +[[package]] +name = "webpki-roots" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52f5ee44c96cf55f1b349600768e3ece3a8f26010c05265ab73f945bb1a2eb9d" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -4136,11 +6438,37 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "windows-core" -version = "0.52.0" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ - "windows-targets 0.52.6", + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] @@ -4149,6 +6477,35 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.48.0" @@ -4176,6 +6533,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -4209,13 +6575,30 @@ dependencies = [ "windows_aarch64_gnullvm 0.52.6", "windows_aarch64_msvc 0.52.6", "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm", + "windows_i686_gnullvm 0.52.6", "windows_i686_msvc 0.52.6", "windows_x86_64_gnu 0.52.6", "windows_x86_64_gnullvm 0.52.6", "windows_x86_64_msvc 0.52.6", ] +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" @@ -4228,6 +6611,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" @@ -4240,6 +6629,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + [[package]] name = "windows_i686_gnu" version = "0.48.5" @@ -4252,12 +6647,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + [[package]] name = "windows_i686_msvc" version = "0.48.5" @@ -4270,6 +6677,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" @@ -4282,6 +6695,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" @@ -4294,6 +6713,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" @@ -4306,6 +6731,27 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" + +[[package]] +name = "winnow" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +dependencies = [ + "memchr", +] + [[package]] name = "winreg" version = "0.50.0" @@ -4316,6 +6762,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + [[package]] name = "wit-bindgen" version = "0.57.1" @@ -4323,16 +6778,89 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] -name = "write16" -version = "1.0.0" +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck 0.5.0", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck 0.5.0", + "indexmap", + "prettyplease", + "syn 2.0.117", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.117", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags 2.11.1", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] [[package]] name = "writeable" -version = "0.5.5" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "x25519-dalek" @@ -4341,7 +6869,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" dependencies = [ "curve25519-dalek", - "rand_core", + "rand_core 0.6.4", "serde", "zeroize", ] @@ -4359,7 +6887,7 @@ dependencies = [ "nom", "oid-registry", "rusticata-macros", - "thiserror", + "thiserror 1.0.69", "time", ] @@ -4380,11 +6908,10 @@ dependencies = [ [[package]] name = "yoke" -version = "0.7.5" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" dependencies = [ - "serde", "stable_deref_trait", "yoke-derive", "zerofrom", @@ -4392,83 +6919,93 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.7.5" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", - "synstructure 0.13.1", + "syn 2.0.117", + "synstructure 0.13.2", ] [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ - "byteorder", "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.35" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", ] [[package]] name = "zerofrom" -version = "0.1.5" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" dependencies = [ "zerofrom-derive", ] [[package]] name = "zerofrom-derive" -version = "0.1.5" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", - "synstructure 0.13.1", + "syn 2.0.117", + "synstructure 0.13.2", ] [[package]] name = "zeroize" -version = "1.8.1" +version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" dependencies = [ "zeroize_derive", ] [[package]] name = "zeroize_derive" -version = "1.4.2" +version = "1.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce36e65b0d2999d2aafac989fb249189a141aee1f53c612c1f37d72631959f69" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", ] [[package]] name = "zerovec" -version = "0.10.4" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" dependencies = [ "yoke", "zerofrom", @@ -4477,11 +7014,51 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.10.3" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn 2.0.106", + "syn 2.0.117", +] + +[[package]] +name = "zlib-rs" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", ] diff --git a/Cargo.toml b/Cargo.toml index 3f4ffe80..226471fd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ thiserror = "1.0.63" async-trait = "0.1.81" jsonwebtoken = { version = "10.3.0", features = ["aws_lc_rs"] } jwt-compact = { version = "0.9.0-beta.1", features = ["es256k"] } +dspy-rs = { git = "https://github.com/OpenSecretCloud/DSRs.git", branch = "main" } diesel = { version = "=2.2.3", features = [ "postgres", "postgres_backend", diff --git a/flake.nix b/flake.nix index 611044c3..6254ad61 100644 --- a/flake.nix +++ b/flake.nix @@ -27,6 +27,10 @@ rustToolchain = builtins.fromTOML (builtins.readFile ./rust-toolchain.toml); rustChannel = rustToolchain.toolchain.channel; rustAnalyzer = pkgs.rust-bin.stable."${rustChannel}".rust-analyzer; + rustPlatform = pkgs.makeRustPlatform { + cargo = rust; + rustc = rust; + }; commonInputs = [ rust @@ -249,7 +253,8 @@ entrypoint = "/bin/entrypoint"; }; - opensecret = pkgs.rustPlatform.buildRustPackage { + # Build the main Rust package + opensecret = rustPlatform.buildRustPackage { pname = "opensecret"; version = "0.1.0"; src = pkgs.lib.cleanSourceWith { @@ -272,6 +277,11 @@ }; cargoLock = { lockFile = ./Cargo.lock; + outputHashes = { + "baml-ids-0.0.1" = "sha256-wjgwOcZvZIWEAjpq0gZwtZccQ+FupudBzpxR4+udKEQ="; + "minijinja-2.16.0" = "sha256-55Zo8mVgMhCgYzE6662oTXfn7W908LplZv6ys/aHveY="; + "rig-core-0.26.0" = "sha256-/1C2/UTuJrre4IYNW7i1pYaOGy0dxzhLyogR+5NL1nk="; + }; }; nativeBuildInputs = [ pkgs.pkg-config diff --git a/src/main.rs b/src/main.rs index b28f15d9..ca479b39 100644 --- a/src/main.rs +++ b/src/main.rs @@ -20,7 +20,7 @@ use crate::sqs::SqsEventPublisher; use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ - conversation_projects_routes, conversations_routes, health_routes_with_state, + agent_routes, conversation_projects_routes, conversations_routes, health_routes_with_state, instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, responses_routes, }; @@ -98,10 +98,10 @@ mod os_flags; mod private_key; mod provider_routing; mod proxy_config; +mod rag; #[cfg(test)] mod security_invariants; mod seed_wrapping; -mod rag; mod sqs; mod tokens; mod web; @@ -3213,6 +3213,10 @@ async fn main() -> Result<(), Error> { rag_routes(app_state.clone()) .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), ) + .merge( + agent_routes(app_state.clone()) + .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), + ) .merge(attestation_routes::router(app_state.clone())) .merge(oauth_routes(app_state.clone())) .merge(platform_login_routes(app_state.clone())) diff --git a/src/models/agents.rs b/src/models/agents.rs new file mode 100644 index 00000000..262ac1fd --- /dev/null +++ b/src/models/agents.rs @@ -0,0 +1,121 @@ +use crate::models::schema::agents; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +pub const AGENT_KIND_MAIN: &str = "main"; +pub const AGENT_KIND_SUBAGENT: &str = "subagent"; +pub const AGENT_CREATED_BY_USER: &str = "user"; +pub const AGENT_CREATED_BY_AGENT: &str = "agent"; + +#[derive(Error, Debug)] +pub enum AgentError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = agents)] +pub struct Agent { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: i64, + pub kind: String, + pub parent_agent_id: Option, + pub display_name_enc: Option>, + pub purpose_enc: Option>, + pub created_by: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl Agent { + pub fn get_main_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, AgentError> { + agents::table + .filter(agents::user_id.eq(lookup_user_id)) + .filter(agents::kind.eq(AGENT_KIND_MAIN)) + .first::(conn) + .optional() + .map_err(AgentError::DatabaseError) + } + + pub fn get_by_uuid_and_user( + conn: &mut PgConnection, + lookup_uuid: Uuid, + lookup_user_id: Uuid, + ) -> Result, AgentError> { + agents::table + .filter(agents::uuid.eq(lookup_uuid)) + .filter(agents::user_id.eq(lookup_user_id)) + .first::(conn) + .optional() + .map_err(AgentError::DatabaseError) + } + + pub fn get_by_conversation_id_and_user( + conn: &mut PgConnection, + lookup_conversation_id: i64, + lookup_user_id: Uuid, + ) -> Result, AgentError> { + agents::table + .filter(agents::conversation_id.eq(lookup_conversation_id)) + .filter(agents::user_id.eq(lookup_user_id)) + .first::(conn) + .optional() + .map_err(AgentError::DatabaseError) + } + + pub fn list_subagents_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, AgentError> { + agents::table + .filter(agents::user_id.eq(lookup_user_id)) + .filter(agents::kind.eq(AGENT_KIND_SUBAGENT)) + .order((agents::created_at.desc(), agents::id.desc())) + .load::(conn) + .map_err(AgentError::DatabaseError) + } + + pub fn delete_by_uuid_and_user( + conn: &mut PgConnection, + lookup_uuid: Uuid, + lookup_user_id: Uuid, + ) -> Result { + diesel::delete( + agents::table + .filter(agents::uuid.eq(lookup_uuid)) + .filter(agents::user_id.eq(lookup_user_id)), + ) + .execute(conn) + .map_err(AgentError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = agents)] +pub struct NewAgent { + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: i64, + pub kind: String, + pub parent_agent_id: Option, + pub display_name_enc: Option>, + pub purpose_enc: Option>, + pub created_by: String, +} + +impl NewAgent { + pub fn insert(&self, conn: &mut PgConnection) -> Result { + diesel::insert_into(agents::table) + .values(self) + .get_result::(conn) + .map_err(AgentError::DatabaseError) + } +} diff --git a/src/models/conversation_summaries.rs b/src/models/conversation_summaries.rs new file mode 100644 index 00000000..674d9654 --- /dev/null +++ b/src/models/conversation_summaries.rs @@ -0,0 +1,107 @@ +use crate::models::schema::conversation_summaries; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +#[derive(Error, Debug)] +pub enum ConversationSummaryError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = conversation_summaries)] +pub struct ConversationSummary { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: i64, + pub from_created_at: DateTime, + pub to_created_at: DateTime, + pub message_count: i32, + pub content_enc: Vec, + pub content_tokens: i32, + pub embedding_enc: Option>, + pub previous_summary_id: Option, + pub created_at: DateTime, +} + +impl ConversationSummary { + pub fn get_latest_for_conversation( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_conversation_id: i64, + ) -> Result, ConversationSummaryError> { + conversation_summaries::table + .filter(conversation_summaries::user_id.eq(lookup_user_id)) + .filter(conversation_summaries::conversation_id.eq(lookup_conversation_id)) + .order(( + conversation_summaries::created_at.desc(), + conversation_summaries::id.desc(), + )) + .first::(conn) + .optional() + .map_err(ConversationSummaryError::DatabaseError) + } + + pub fn list_for_conversation( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_conversation_id: i64, + limit: i64, + ) -> Result, ConversationSummaryError> { + conversation_summaries::table + .filter(conversation_summaries::user_id.eq(lookup_user_id)) + .filter(conversation_summaries::conversation_id.eq(lookup_conversation_id)) + .order(( + conversation_summaries::created_at.desc(), + conversation_summaries::id.desc(), + )) + .limit(limit) + .load::(conn) + .map_err(ConversationSummaryError::DatabaseError) + } + + pub fn delete_all_for_conversation( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_conversation_id: i64, + ) -> Result { + diesel::delete( + conversation_summaries::table + .filter(conversation_summaries::user_id.eq(lookup_user_id)) + .filter(conversation_summaries::conversation_id.eq(lookup_conversation_id)), + ) + .execute(conn) + .map_err(ConversationSummaryError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = conversation_summaries)] +pub struct NewConversationSummary { + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: i64, + pub from_created_at: DateTime, + pub to_created_at: DateTime, + pub message_count: i32, + pub content_enc: Vec, + pub content_tokens: i32, + pub embedding_enc: Option>, + pub previous_summary_id: Option, +} + +impl NewConversationSummary { + pub fn insert( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(conversation_summaries::table) + .values(self) + .get_result::(conn) + .map_err(ConversationSummaryError::DatabaseError) + } +} diff --git a/src/models/memory_blocks.rs b/src/models/memory_blocks.rs new file mode 100644 index 00000000..744d46f2 --- /dev/null +++ b/src/models/memory_blocks.rs @@ -0,0 +1,109 @@ +use crate::models::schema::memory_blocks; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +pub const MEMORY_BLOCK_LABEL_PERSONA: &str = "persona"; +pub const MEMORY_BLOCK_LABEL_HUMAN: &str = "human"; + +#[derive(Error, Debug)] +pub enum MemoryBlockError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = memory_blocks)] +pub struct MemoryBlock { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub label: String, + pub value_enc: Vec, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl MemoryBlock { + pub fn get_by_user_and_label( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_label: &str, + ) -> Result, MemoryBlockError> { + memory_blocks::table + .filter(memory_blocks::user_id.eq(lookup_user_id)) + .filter(memory_blocks::label.eq(lookup_label)) + .first::(conn) + .optional() + .map_err(MemoryBlockError::DatabaseError) + } + + pub fn get_all_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, MemoryBlockError> { + memory_blocks::table + .filter(memory_blocks::user_id.eq(lookup_user_id)) + .order(memory_blocks::label.asc()) + .load::(conn) + .map_err(MemoryBlockError::DatabaseError) + } + + pub fn delete_all_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result { + diesel::delete(memory_blocks::table.filter(memory_blocks::user_id.eq(lookup_user_id))) + .execute(conn) + .map_err(MemoryBlockError::DatabaseError) + } + + pub fn delete_by_user_and_label( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_label: &str, + ) -> Result { + diesel::delete( + memory_blocks::table + .filter(memory_blocks::user_id.eq(lookup_user_id)) + .filter(memory_blocks::label.eq(lookup_label)), + ) + .execute(conn) + .map_err(MemoryBlockError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = memory_blocks)] +pub struct NewMemoryBlock { + pub uuid: Uuid, + pub user_id: Uuid, + pub label: String, + pub value_enc: Vec, +} + +impl NewMemoryBlock { + pub fn new(user_id: Uuid, label: impl Into, value_enc: Vec) -> Self { + NewMemoryBlock { + uuid: Uuid::new_v4(), + user_id, + label: label.into(), + value_enc, + } + } + + pub fn insert_or_update( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(memory_blocks::table) + .values(self) + .on_conflict((memory_blocks::user_id, memory_blocks::label)) + .do_update() + .set(memory_blocks::value_enc.eq(self.value_enc.clone())) + .get_result::(conn) + .map_err(MemoryBlockError::DatabaseError) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 8d92d1a4..2836fc98 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,8 +1,11 @@ pub mod account_deletion; +pub mod agents; pub mod app_data_migrations; +pub mod conversation_summaries; pub mod email_verification; pub mod enclave_secrets; pub mod invite_codes; +pub mod memory_blocks; pub mod oauth; pub mod org_memberships; pub mod org_project_secrets; diff --git a/src/models/responses.rs b/src/models/responses.rs index b22e41a2..971e7b29 100644 --- a/src/models/responses.rs +++ b/src/models/responses.rs @@ -771,6 +771,7 @@ impl NewConversation { response_id: response_result.as_ref().map(|r| r.id), user_id, content_enc: first_message_content, + attachment_text_enc: None, prompt_tokens: first_message_tokens, }; let user_message = new_message.insert(tx)?; @@ -826,6 +827,7 @@ pub struct NewUserMessage { pub user_id: Uuid, pub content_enc: Vec, pub prompt_tokens: i32, + pub attachment_text_enc: Option>, } impl UserMessage { @@ -1148,6 +1150,8 @@ pub struct RawThreadMessage { pub uuid: Uuid, #[diesel(sql_type = diesel::sql_types::Nullable)] pub content_enc: Option>, + #[diesel(sql_type = diesel::sql_types::Nullable)] + pub attachment_text_enc: Option>, #[diesel(sql_type = diesel::sql_types::Nullable)] pub status: Option, #[diesel(sql_type = diesel::sql_types::Timestamptz)] @@ -1190,6 +1194,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1209,6 +1214,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1228,6 +1234,7 @@ impl RawThreadMessage { tc.id, tc.uuid, tc.arguments_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1246,6 +1253,7 @@ impl RawThreadMessage { tto.id, tto.uuid, tto.output_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1265,6 +1273,7 @@ impl RawThreadMessage { ri.id, ri.uuid, ri.content_enc, + NULL::bytea as attachment_text_enc, ri.status, ri.created_at, NULL::text as model, @@ -1303,6 +1312,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1322,6 +1332,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1341,6 +1352,7 @@ impl RawThreadMessage { tc.id, tc.uuid, tc.arguments_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1359,6 +1371,7 @@ impl RawThreadMessage { tto.id, tto.uuid, tto.output_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1378,6 +1391,7 @@ impl RawThreadMessage { ri.id, ri.uuid, ri.content_enc, + NULL::bytea as attachment_text_enc, ri.status, ri.created_at, NULL::text as model, @@ -1425,6 +1439,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1444,6 +1459,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1463,6 +1479,7 @@ impl RawThreadMessage { tc.id, tc.uuid, tc.arguments_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1481,6 +1498,7 @@ impl RawThreadMessage { tto.id, tto.uuid, tto.output_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1500,6 +1518,7 @@ impl RawThreadMessage { ri.id, ri.uuid, ri.content_enc, + NULL::bytea as attachment_text_enc, ri.status, ri.created_at, NULL::text as model, @@ -1567,6 +1586,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1586,6 +1606,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1605,6 +1626,7 @@ impl RawThreadMessage { tc.id, tc.uuid, tc.arguments_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1623,6 +1645,7 @@ impl RawThreadMessage { tto.id, tto.uuid, tto.output_enc as content_enc, + NULL::bytea as attachment_text_enc, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1642,6 +1665,7 @@ impl RawThreadMessage { ri.id, ri.uuid, ri.content_enc, + NULL::bytea as attachment_text_enc, ri.status, ri.created_at, NULL::text as model, diff --git a/src/web/agent/compaction.rs b/src/web/agent/compaction.rs new file mode 100644 index 00000000..ab59aab6 --- /dev/null +++ b/src/web/agent/compaction.rs @@ -0,0 +1,247 @@ +use crate::ApiError; +use dspy_rs::adapter::chat::ChatAdapter; +use dspy_rs::LM; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use tracing::error; + +/// Instruction for summarization DSRs signature +pub const SUMMARY_INSTRUCTION: &str = r#"You are a conversation summarizer. Your job is to create a concise summary that allows an AI agent to resume a conversation without disruption, even after older messages are replaced with this summary. + +Your summary should be structured and actionable. Include: +1. Task/Conversation overview: What is the user working on? Key clarifications or constraints? +2. Current State: What has been completed or discussed? Any files/resources referenced? +3. Next Steps: What would logically come next in this conversation? + +Keep your summary under 100 words. Be specific and preserve key details like names, preferences, and decisions made."#; + +/// Instruction for correction DSRs signature +pub const SUMMARY_CORRECTION_INSTRUCTION: &str = r#"You are a correction agent. The summarizer produced a malformed response that couldn't be parsed. Your job is to extract the summary from the malformed response and return it in the correct format. + +Preserve the original intent and content - do NOT generate new content. Just reshape the malformed response into the expected output format."#; + +/// DSRs signature for conversation summarization +#[derive(dspy_rs::Signature, Clone, Debug)] +pub struct SummarizeConversation { + #[input(desc = "Previous summary to build upon (empty if first summarization)")] + pub previous_summary: String, + + #[input(desc = "New conversation messages to incorporate into the summary")] + pub new_messages: String, + + #[output(desc = "Updated summary incorporating all context (100 word limit)")] + pub summary: String, +} + +/// DSRs signature for correcting malformed summarization responses +#[derive(dspy_rs::Signature, Clone, Debug)] +pub struct SummarizationCorrection { + #[input(desc = "The previous summary that was being built upon")] + pub previous_summary: String, + + #[input(desc = "The new messages that were being summarized")] + pub new_messages: String, + + #[input(desc = "The malformed response that needs correction")] + pub malformed_response: String, + + #[input(desc = "The error message explaining what went wrong")] + pub error_message: String, + + #[output(desc = "Corrected summary (100 word limit)")] + pub summary: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SummarizeConversationOutput { + pub summary: String, +} + +pub struct CompactionManager { + max_retries: usize, +} + +impl CompactionManager { + pub fn new() -> Self { + Self { max_retries: 2 } + } + + pub fn should_compact(&self, current_tokens: usize, max_tokens: usize, threshold: f32) -> bool { + current_tokens > ((max_tokens as f32 * threshold) as usize) + } + + pub async fn summarize( + &self, + lm: &Arc, + previous_summary: &str, + new_messages: &str, + ) -> Result { + let input = SummarizeConversationInput { + previous_summary: previous_summary.to_string(), + new_messages: new_messages.to_string(), + }; + + // First attempt + match call_summarize_conversation(lm, &input).await { + Ok(out) => return Ok(out.summary), + Err(SummarizationError::Parse { + raw_response, + error_message, + }) => { + if let Ok(corrected) = self + .try_correction( + lm, + previous_summary, + new_messages, + &raw_response, + &error_message, + ) + .await + { + return Ok(corrected); + } + } + Err(SummarizationError::Api(e)) => return Err(e), + }; + + for _attempt in 1..=self.max_retries { + match call_summarize_conversation(lm, &input).await { + Ok(out) => return Ok(out.summary), + Err(SummarizationError::Parse { + raw_response, + error_message, + }) => { + if let Ok(corrected) = self + .try_correction( + lm, + previous_summary, + new_messages, + &raw_response, + &error_message, + ) + .await + { + return Ok(corrected); + } + } + Err(SummarizationError::Api(e)) => { + return Err(e); + } + }; + } + + Err(ApiError::InternalServerError) + } + + async fn try_correction( + &self, + lm: &Arc, + previous_summary: &str, + new_messages: &str, + raw_response: &str, + error_message: &str, + ) -> Result { + let input = SummarizationCorrectionInput { + previous_summary: previous_summary.to_string(), + new_messages: new_messages.to_string(), + malformed_response: raw_response.to_string(), + error_message: error_message.to_string(), + }; + + let out = call_summarization_correction(lm, &input).await?; + Ok(out.summary) + } +} + +impl Default for CompactionManager { + fn default() -> Self { + Self::new() + } +} + +#[derive(Debug)] +enum SummarizationError { + Api(ApiError), + Parse { + raw_response: String, + error_message: String, + }, +} + +async fn call_summarize_conversation( + lm: &Arc, + input: &SummarizeConversationInput, +) -> Result { + let adapter = ChatAdapter; + + let system = adapter + .format_system_message_typed_with_instruction::(Some( + SUMMARY_INSTRUCTION, + )) + .map_err(|e| { + error!("Failed to format DSRS system prompt: {e:?}"); + SummarizationError::Api(ApiError::InternalServerError) + })?; + let user_msg = adapter.format_user_message_typed::(input); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM call failed: {e:?}"); + SummarizationError::Api(ApiError::InternalServerError) + })?; + + let raw_response = response.output.content(); + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS typed parse failed: {e:?}"); + SummarizationError::Parse { + raw_response, + error_message: e.to_string(), + } + })?; + + Ok(SummarizeConversationOutput { + summary: output.summary, + }) +} + +async fn call_summarization_correction( + lm: &Arc, + input: &SummarizationCorrectionInput, +) -> Result { + let adapter = ChatAdapter; + + let system = adapter + .format_system_message_typed_with_instruction::(Some( + SUMMARY_CORRECTION_INSTRUCTION, + )) + .map_err(|e| { + error!("Failed to format DSRS system prompt: {e:?}"); + ApiError::InternalServerError + })?; + let user_msg = adapter.format_user_message_typed::(input); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM call failed: {e:?}"); + ApiError::InternalServerError + })?; + + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS typed parse failed: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(SummarizeConversationOutput { + summary: output.summary, + }) +} diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs new file mode 100644 index 00000000..9d3d1344 --- /dev/null +++ b/src/web/agent/mod.rs @@ -0,0 +1,563 @@ +use axum::{ + extract::{Path, State}, + middleware::from_fn_with_state, + response::sse::{Event, Sse}, + routing::{delete, post}, + Extension, Json, Router, +}; +use futures::Stream; +use serde::{Deserialize, Serialize}; +use std::{convert::Infallible, sync::Arc, time::Duration}; +use tokio::time::sleep; +use tracing::{error, warn}; +use uuid::Uuid; + +use crate::jwt::AuthContext; +use crate::models::agents::{Agent, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT}; +use crate::models::users::User; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::web::responses::handlers::encrypt_event; +use crate::web::responses::{ + error_mapping, DeletedObjectResponse, MessageContent, MessageContentConverter, + MessageContentPart, +}; +use crate::{ApiError, AppMode, AppState}; + +mod compaction; +mod runtime; +mod signatures; +mod tools; +mod vision; + +#[derive(Debug, Clone, Deserialize)] +struct AgentChatRequest { + input: MessageContent, +} + +#[derive(Debug, Clone, Deserialize)] +struct CreateSubagentRequest { + display_name: Option, + purpose: String, +} + +#[derive(Debug, Clone, Serialize)] +struct SubagentResponse { + id: Uuid, + object: &'static str, + conversation_id: Uuid, + display_name: String, + purpose: String, + created_by: String, + created_at: i64, +} + +#[derive(Debug, Clone)] +enum ChatTarget { + Main, + Subagent(Uuid), +} + +fn is_empty_message_content(content: &MessageContent) -> bool { + match content { + MessageContent::Text(text) => text.trim().is_empty(), + MessageContent::Parts(parts) => parts.iter().all(|part| match part { + MessageContentPart::Text { text } | MessageContentPart::InputText { text } => { + text.trim().is_empty() + } + MessageContentPart::InputImage { + image_url, file_id, .. + } => { + let url_empty = image_url + .as_deref() + .map(|s| s.trim().is_empty()) + .unwrap_or(true); + let file_empty = file_id + .as_deref() + .map(|s| s.trim().is_empty()) + .unwrap_or(true); + url_empty && file_empty + } + MessageContentPart::InputFile { + filename, + file_data, + } => filename.trim().is_empty() && file_data.trim().is_empty(), + }), + } +} + +// SSE event types for agent chat (message-level delivery, not token streaming) +const EVENT_AGENT_MESSAGE: &str = "agent.message"; +const EVENT_AGENT_TYPING: &str = "agent.typing"; +const EVENT_AGENT_DONE: &str = "agent.done"; +const EVENT_AGENT_ERROR: &str = "agent.error"; + +const MESSAGE_STAGGER_DELAY_MS: u64 = 1_500; + +#[derive(Debug, Clone, Serialize)] +struct AgentMessageEvent { + messages: Vec, + step: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentTypingEvent { + step: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentDoneEvent { + total_steps: usize, + total_messages: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentErrorEvent { + error: String, +} + +async fn encrypt_agent_event( + state: &AppState, + session_id: &Uuid, + event_type: &str, + payload: &T, +) -> Result { + let value = serde_json::to_value(payload).map_err(|e| { + error!("Failed to serialize agent SSE payload for {event_type}: {e:?}"); + ApiError::InternalServerError + })?; + + encrypt_event(state, session_id, event_type, &value).await +} + +pub fn router(app_state: Arc) -> Router<()> { + // Experimental endpoints: only enabled in Local/Dev. + if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { + return Router::new().with_state(app_state); + } + + Router::new() + .route( + "/v1/agent/chat", + post(chat_main).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/subagents", + post(create_subagent).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/subagents/:id/chat", + post(chat_subagent).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/subagents/:id", + delete(delete_subagent) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .with_state(app_state) +} + +async fn chat_main( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>>, ApiError> { + chat_with_target( + state, + session_id, + user, + auth_context, + body, + ChatTarget::Main, + ) + .await +} + +async fn chat_subagent( + State(state): State>, + Path(agent_uuid): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) + .map_err(|_| ApiError::InternalServerError)? + .ok_or(ApiError::NotFound)?; + + if agent.kind != AGENT_KIND_SUBAGENT { + return Err(ApiError::NotFound); + } + + chat_with_target( + state, + session_id, + user, + auth_context, + body, + ChatTarget::Subagent(agent_uuid), + ) + .await +} + +async fn chat_with_target( + state: Arc, + session_id: Uuid, + user: User, + auth_context: AuthContext, + body: AgentChatRequest, + target: ChatTarget, +) -> Result>>, ApiError> { + MessageContentConverter::validate_content(&body.input)?; + let input_content = MessageContentConverter::normalize_content(body.input.clone()); + if is_empty_message_content(&input_content) { + return Err(ApiError::BadRequest); + } + + // NOTE: We intentionally do runtime initialization *inside* the SSE stream so the client can + // receive typing indicators immediately, even if compaction/key retrieval takes time. + let event_stream = async_stream::stream! { + let mut max_steps: usize = 0; + let mut total_messages: usize = 0; + let mut had_error = false; + let mut input_for_agent = String::new(); + + let mut runtime_opt: Option = None; + 'init: { + // Response starts: immediately emit a typing indicator. + let initial_typing = AgentTypingEvent { step: 0 }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_TYPING, &initial_typing) + .await + { + Ok(event) => yield Ok::(event), + Err(e) => { + error!("Failed to encrypt initial typing event: {e:?}"); + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); + had_error = true; + break 'init; + } + } + + let user_key = match state.get_user_key(&user, &auth_context, None, None).await { + Ok(k) => k, + Err(_) => { + let err_event = AgentErrorEvent { + error: "Failed to initialize agent session.".to_string(), + }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) + .await + { + Ok(event) => yield Ok(event), + Err(_) => { + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")) + } + } + had_error = true; + break 'init; + } + }; + + let runtime = match target.clone() { + ChatTarget::Main => { + runtime::AgentRuntime::new_main(state.clone(), user.clone(), user_key).await + } + ChatTarget::Subagent(agent_uuid) => { + runtime::AgentRuntime::new_subagent( + state.clone(), + user.clone(), + user_key, + agent_uuid, + ) + .await + } + }; + + let mut runtime = match runtime { + Ok(r) => r, + Err(e) => { + error!("Agent runtime initialization error: {e:?}"); + let err_event = AgentErrorEvent { + error: "Failed to initialize agent runtime.".to_string(), + }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) + .await + { + Ok(event) => yield Ok(event), + Err(_) => { + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")) + } + } + had_error = true; + break 'init; + } + }; + + match runtime.prepare(&input_content).await { + Ok(prepared) => { + input_for_agent = prepared; + } + Err(e) => { + error!("Agent prepare() failed: {e:?}"); + let err_event = AgentErrorEvent { + error: "Agent encountered an error preparing your request.".to_string(), + }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event).await + { + Ok(event) => yield Ok(event), + Err(_) => yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")), + } + had_error = true; + break 'init; + } + } + + max_steps = runtime.max_steps(); + runtime_opt = Some(runtime); + } + + if let Some(mut runtime) = runtime_opt { + 'steps: for step_num in 0..max_steps { + // Immediately indicate that the agent is working. + let typing_event = AgentTypingEvent { step: step_num }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_TYPING, &typing_event).await { + Ok(event) => yield Ok::(event), + Err(e) => { + error!("Failed to encrypt agent typing event: {e:?}"); + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); + had_error = true; + break 'steps; + } + } + + let step_result = runtime.step(&input_for_agent, step_num == 0).await; + + match step_result { + Ok(result) => { + let runtime::StepResult { + messages, + executed_tools, + done, + .. + } = result; + + // Persist assistant messages SYNCHRONOUSLY (so next step sees them). + // Embedding updates happen async inside insert_assistant_message. + for msg in &messages { + if let Err(e) = runtime.insert_assistant_message(msg).await { + error!("Failed to persist assistant message: {e:?}"); + } + } + + // Persist tool calls SYNCHRONOUSLY (so next step sees them in context) + for executed in &executed_tools { + if let Err(e) = runtime + .insert_tool_call_and_output(&executed.tool_call, &executed.result) + .await + { + error!("Failed to persist tool call: {e:?}"); + } + } + + // Emit messages to client with optional staggered delivery. + if !messages.is_empty() { + total_messages += messages.len(); + let message_count = messages.len(); + + for (idx, msg) in messages.into_iter().enumerate() { + let event_data = AgentMessageEvent { + messages: vec![msg], + step: step_num, + }; + + match encrypt_agent_event( + &state, + &session_id, + EVENT_AGENT_MESSAGE, + &event_data, + ) + .await + { + Ok(event) => yield Ok::(event), + Err(e) => { + error!("Failed to encrypt agent message event: {e:?}"); + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); + had_error = true; + break 'steps; + } + } + + if idx + 1 < message_count { + let typing_event = AgentTypingEvent { step: step_num }; + match encrypt_agent_event( + &state, + &session_id, + EVENT_AGENT_TYPING, + &typing_event, + ) + .await + { + Ok(event) => { + yield Ok::(event) + } + Err(e) => { + error!("Failed to encrypt agent typing event: {e:?}"); + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); + had_error = true; + break 'steps; + } + } + + sleep(Duration::from_millis(MESSAGE_STAGGER_DELAY_MS)).await; + } + } + } + + if done { + break; + } + } + Err(e) => { + error!("Agent step {} error: {e:?}", step_num); + let err_event = AgentErrorEvent { + error: "Agent encountered an error processing your message.".to_string(), + }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) + .await + { + Ok(event) => yield Ok(event), + Err(_) => yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")), + } + had_error = true; + break 'steps; + } + } + } + } + + if !had_error && total_messages == 0 { + warn!("Agent produced no messages"); + let fallback = AgentMessageEvent { + messages: vec!["I apologize, but I wasn't able to generate a response.".to_string()], + step: 0, + }; + if let Ok(event) = + encrypt_agent_event(&state, &session_id, EVENT_AGENT_MESSAGE, &fallback).await + { + yield Ok(event); + total_messages = 1; + } + } + + // Final done event + let done_event = AgentDoneEvent { + total_steps: max_steps, + total_messages, + }; + match encrypt_agent_event(&state, &session_id, EVENT_AGENT_DONE, &done_event).await { + Ok(event) => yield Ok(event), + Err(_) => yield Ok(Event::default().data("[DONE]")), + } + }; + + Ok(Sse::new(event_stream)) +} + +async fn create_subagent( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if body.purpose.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let (main_agent, _) = + runtime::ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; + let (agent, conversation, display_name) = runtime::create_subagent( + &mut conn, + &user_key, + user.uuid, + &main_agent, + body.display_name.as_deref(), + body.purpose.trim(), + AGENT_CREATED_BY_USER, + ) + .await?; + + let response = SubagentResponse { + id: agent.uuid, + object: "agent.subagent", + conversation_id: conversation.uuid, + display_name, + purpose: body.purpose.trim().to_string(), + created_by: agent.created_by, + created_at: agent.created_at.timestamp(), + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn delete_subagent( + State(state): State>, + Path(agent_uuid): Path, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) + .map_err(|_| ApiError::InternalServerError)? + .ok_or(ApiError::NotFound)?; + + if agent.kind != AGENT_KIND_SUBAGENT { + return Err(ApiError::NotFound); + } + + state + .db + .delete_conversation(agent.conversation_id, user.uuid) + .map_err(error_mapping::map_generic_db_error)?; + + state.rag_cache.lock().await.evict_user(user.uuid); + + let response = DeletedObjectResponse { + id: agent.uuid, + object: "agent.subagent.deleted", + deleted: true, + }; + + encrypt_response(&state, &session_id, &response).await +} diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs new file mode 100644 index 00000000..b0b3cc05 --- /dev/null +++ b/src/web/agent/runtime.rs @@ -0,0 +1,1387 @@ +use secp256k1::SecretKey; +use serde_json::json; +use std::sync::Arc; +use tracing::{debug, error, warn}; +use uuid::Uuid; + +use diesel::prelude::*; + +use crate::encrypt::{decrypt_string, encrypt_with_key}; +use crate::models::agents::{ + Agent, NewAgent, AGENT_CREATED_BY_USER, AGENT_KIND_MAIN, AGENT_KIND_SUBAGENT, +}; +use crate::models::conversation_summaries::{ConversationSummary, NewConversationSummary}; +use crate::models::memory_blocks::{ + MemoryBlock, NewMemoryBlock, MEMORY_BLOCK_LABEL_HUMAN, MEMORY_BLOCK_LABEL_PERSONA, +}; +use crate::models::responses::{ + AssistantMessage, Conversation, NewAssistantMessage, NewConversation, NewToolCall, + NewToolOutput, NewUserMessage, RawThreadMessage, RawThreadMessageMetadata, ToolCall, + ToolOutput, UserMessage, +}; +use crate::models::schema::{agents, memory_blocks, user_embeddings}; +use crate::rag::{ + insert_message_embedding, serialize_f32_le, SOURCE_TYPE_ARCHIVAL, SOURCE_TYPE_MESSAGE, +}; +use crate::tokens::count_tokens; +use crate::web::openai_auth::AuthMethod; +use crate::web::responses::{MessageContent, MessageContentConverter, MessageContentPart}; +use crate::{ApiError, AppState}; + +use super::compaction::CompactionManager; +use super::signatures::{ + build_lm, call_agent_response_with_retry_and_correction, AgentResponseInput, AgentToolCall, + AGENT_INSTRUCTION, +}; +use super::tools::{ + ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, + MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, +}; +use super::vision; + +const DEFAULT_PERSONA_VALUE: &str = "I am Sage, a helpful AI companion. I maintain long-term memory across our conversations and strive to be friendly, concise, and genuinely helpful."; +pub const DEFAULT_MODEL: &str = "kimi-k2-6"; +pub const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; +pub const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; +const MIN_MESSAGES_IN_CONTEXT: usize = 20; +const MAIN_AGENT_METADATA_TYPE: &str = "agent_main"; +const SUBAGENT_METADATA_TYPE: &str = "subagent"; + +#[derive(Clone, Debug, Default)] +struct AgentContext { + current_time: String, + agent_kind: String, + subagent_purpose: String, + persona_block: String, + human_block: String, + memory_metadata: String, + previous_context_summary: String, + recent_conversation: String, + is_first_time_user: bool, +} + +#[derive(Clone, Debug)] +pub struct ExecutedTool { + pub tool_call: AgentToolCall, + pub result: ToolResult, +} + +#[derive(Clone, Debug)] +pub struct StepResult { + pub messages: Vec, + #[allow(dead_code)] + pub tool_calls: Vec, + pub executed_tools: Vec, + pub done: bool, +} + +pub struct AgentRuntime { + state: Arc, + user: Arc, + user_key: Arc, + agent: Agent, + conversation: Conversation, + subagent_purpose: String, + system_prompt: String, + lm: Arc, + tools: ToolRegistry, + available_tools: String, + compaction: CompactionManager, + current_tool_results: Vec, + previous_step_summary: Option<(Vec, Vec)>, + max_steps: usize, +} + +fn build_system_prompt(agent: &Agent, subagent_purpose: &str) -> String { + let mut prompt = AGENT_INSTRUCTION.to_string(); + + if agent.kind == AGENT_KIND_MAIN { + prompt.push_str( + "\n\nMAIN AGENT MODE:\nYou are the user's primary persistent agent and the home surface of Maple.", + ); + } else { + prompt.push_str( + "\n\nSUBAGENT MODE:\nYou are operating inside a focused subagent chat. You share the user's memory with the main agent, but your recent conversation history is only this thread. Stay tightly focused on your assigned purpose and do not behave like a separate identity.", + ); + + if !subagent_purpose.trim().is_empty() { + prompt.push_str("\n\nSUBAGENT PURPOSE:\n"); + prompt.push_str(subagent_purpose.trim()); + } + } + + prompt +} + +fn derive_subagent_display_name(display_name: Option<&str>, purpose: &str) -> String { + if let Some(name) = display_name.map(str::trim).filter(|s| !s.is_empty()) { + let trimmed: String = name.chars().take(80).collect(); + if !trimmed.is_empty() { + return trimmed; + } + } + + let mut derived = purpose + .split_whitespace() + .take(6) + .collect::>() + .join(" "); + if derived.is_empty() { + derived = "New Subagent".to_string(); + } + derived.chars().take(80).collect() +} + +async fn main_agent_metadata_enc(user_key: &SecretKey) -> Vec { + encrypt_with_key( + user_key, + json!({"type": MAIN_AGENT_METADATA_TYPE}) + .to_string() + .as_bytes(), + ) + .await +} + +async fn subagent_metadata_enc( + user_key: &SecretKey, + agent_uuid: Uuid, + display_name: &str, + purpose: &str, +) -> Vec { + encrypt_with_key( + user_key, + json!({ + "type": SUBAGENT_METADATA_TYPE, + "agent_id": agent_uuid, + "display_name": display_name, + "purpose": purpose, + }) + .to_string() + .as_bytes(), + ) + .await +} + +pub(crate) async fn ensure_main_agent( + state: &Arc, + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, +) -> Result<(Agent, Conversation), ApiError> { + let existing = Agent::get_main_for_user(conn, user_id).map_err(|e| { + error!("Failed to load main agent: {e:?}"); + ApiError::InternalServerError + })?; + + let (agent, conversation) = if let Some(agent) = existing { + match Conversation::get_by_id_and_user(conn, agent.conversation_id, user_id) { + Ok(conversation) => (agent, conversation), + Err(_) => { + let conversation = NewConversation { + uuid: Uuid::new_v4(), + user_id, + project_id: None, + is_pinned: false, + metadata_enc: Some(main_agent_metadata_enc(user_key).await), + } + .insert(conn) + .map_err(|e| { + error!("Failed to recreate main agent conversation: {e:?}"); + ApiError::InternalServerError + })?; + + let agent = diesel::update(agents::table.filter(agents::id.eq(agent.id))) + .set(agents::conversation_id.eq(conversation.id)) + .get_result::(conn) + .map_err(|e| { + error!("Failed to repair main agent conversation_id: {e:?}"); + ApiError::InternalServerError + })?; + + (agent, conversation) + } + } + } else { + let conversation = NewConversation { + uuid: Uuid::new_v4(), + user_id, + project_id: None, + is_pinned: false, + metadata_enc: Some(main_agent_metadata_enc(user_key).await), + } + .insert(conn) + .map_err(|e| { + error!("Failed to create main agent conversation: {e:?}"); + ApiError::InternalServerError + })?; + + let agent = NewAgent { + uuid: Uuid::new_v4(), + user_id, + conversation_id: conversation.id, + kind: AGENT_KIND_MAIN.to_string(), + parent_agent_id: None, + display_name_enc: None, + purpose_enc: None, + created_by: AGENT_CREATED_BY_USER.to_string(), + } + .insert(conn) + .map_err(|e| { + error!("Failed to create main agent row: {e:?}"); + ApiError::InternalServerError + })?; + + (agent, conversation) + }; + + AgentRuntime::ensure_default_blocks(state, conn, user_key, user_id).await?; + + Ok((agent, conversation)) +} + +pub(crate) async fn create_subagent( + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, + parent_agent: &Agent, + display_name: Option<&str>, + purpose: &str, + created_by: &str, +) -> Result<(Agent, Conversation, String), ApiError> { + if parent_agent.kind != AGENT_KIND_MAIN { + return Err(ApiError::BadRequest); + } + + let purpose = purpose.trim(); + if purpose.is_empty() { + return Err(ApiError::BadRequest); + } + + let resolved_display_name = derive_subagent_display_name(display_name, purpose); + let agent_uuid = Uuid::new_v4(); + + let conversation = NewConversation { + uuid: Uuid::new_v4(), + user_id, + project_id: None, + is_pinned: false, + metadata_enc: Some( + subagent_metadata_enc(user_key, agent_uuid, &resolved_display_name, purpose).await, + ), + } + .insert(conn) + .map_err(|e| { + error!("Failed to create subagent conversation: {e:?}"); + ApiError::InternalServerError + })?; + + let display_name_enc = Some(encrypt_with_key(user_key, resolved_display_name.as_bytes()).await); + let purpose_enc = Some(encrypt_with_key(user_key, purpose.as_bytes()).await); + + let agent = NewAgent { + uuid: agent_uuid, + user_id, + conversation_id: conversation.id, + kind: AGENT_KIND_SUBAGENT.to_string(), + parent_agent_id: Some(parent_agent.id), + display_name_enc, + purpose_enc, + created_by: created_by.to_string(), + } + .insert(conn) + .map_err(|e| { + error!("Failed to create subagent row: {e:?}"); + ApiError::InternalServerError + })?; + + Ok((agent, conversation, resolved_display_name)) +} + +impl AgentRuntime { + pub async fn new_main( + state: Arc, + user: crate::models::users::User, + user_key: SecretKey, + ) -> Result { + let user = Arc::new(user); + let user_key = Arc::new(user_key); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let (agent, conversation) = + ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; + + Self::from_loaded(state, user, user_key, agent, conversation).await + } + + pub async fn new_subagent( + state: Arc, + user: crate::models::users::User, + user_key: SecretKey, + agent_uuid: Uuid, + ) -> Result { + let user = Arc::new(user); + let user_key = Arc::new(user_key); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let _ = ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; + + let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) + .map_err(|e| { + error!("Failed to load subagent: {e:?}"); + ApiError::InternalServerError + })? + .ok_or(ApiError::NotFound)?; + + if agent.kind != AGENT_KIND_SUBAGENT { + return Err(ApiError::NotFound); + } + + let conversation = + Conversation::get_by_id_and_user(&mut conn, agent.conversation_id, user.uuid).map_err( + |e| { + error!("Failed to load subagent conversation: {e:?}"); + ApiError::InternalServerError + }, + )?; + + Self::from_loaded(state, user, user_key, agent, conversation).await + } + + async fn from_loaded( + state: Arc, + user: Arc, + user_key: Arc, + agent: Agent, + conversation: Conversation, + ) -> Result { + let subagent_purpose = decrypt_string(&user_key, agent.purpose_enc.as_ref()) + .map_err(|e| { + error!("Failed to decrypt subagent purpose: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_default(); + + let system_prompt = build_system_prompt(&agent, &subagent_purpose); + + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(MemoryReplaceTool::new( + state.clone(), + user.uuid, + user_key.clone(), + ))); + tools.register(Arc::new(MemoryAppendTool::new( + state.clone(), + user.uuid, + user_key.clone(), + ))); + tools.register(Arc::new(MemoryInsertTool::new( + state.clone(), + user.uuid, + user_key.clone(), + ))); + tools.register(Arc::new(ArchivalInsertTool::new( + state.clone(), + user.clone(), + user_key.clone(), + ))); + tools.register(Arc::new(ArchivalSearchTool::new( + state.clone(), + user.clone(), + user_key.clone(), + ))); + tools.register(Arc::new(ConversationSearchTool::new( + state.clone(), + user.clone(), + user_key.clone(), + conversation.id, + ))); + if agent.kind == AGENT_KIND_MAIN { + tools.register(Arc::new(SpawnSubagentTool::new( + state.clone(), + user.clone(), + user_key.clone(), + agent.clone(), + ))); + } + tools.register(Arc::new(DoneTool)); + + let available_tools = tools.generate_description(); + let lm = build_lm( + state.clone(), + user.clone(), + DEFAULT_MODEL.to_string(), + 0.7, + 32768, + ) + .await?; + + Ok(Self { + state, + user, + user_key, + agent, + conversation, + subagent_purpose, + system_prompt, + lm, + tools, + available_tools, + compaction: CompactionManager::new(), + current_tool_results: Vec::new(), + previous_step_summary: None, + max_steps: 10, + }) + } + + pub(crate) async fn ensure_default_blocks( + _state: &Arc, + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, + ) -> Result<(), ApiError> { + let persona = MemoryBlock::get_by_user_and_label(conn, user_id, MEMORY_BLOCK_LABEL_PERSONA) + .map_err(|e| { + error!("Failed to load persona block: {e:?}"); + ApiError::InternalServerError + })?; + + if persona.is_none() { + let value_enc = encrypt_with_key(user_key, DEFAULT_PERSONA_VALUE.as_bytes()).await; + let block = NewMemoryBlock { + uuid: Uuid::new_v4(), + user_id, + label: MEMORY_BLOCK_LABEL_PERSONA.to_string(), + value_enc, + }; + block.insert_or_update(conn).map_err(|e| { + error!("Failed to create persona block: {e:?}"); + ApiError::InternalServerError + })?; + } + + let human = MemoryBlock::get_by_user_and_label(conn, user_id, MEMORY_BLOCK_LABEL_HUMAN) + .map_err(|e| { + error!("Failed to load human block: {e:?}"); + ApiError::InternalServerError + })?; + + if human.is_none() { + let value_enc = encrypt_with_key(user_key, "".as_bytes()).await; + let block = NewMemoryBlock { + uuid: Uuid::new_v4(), + user_id, + label: MEMORY_BLOCK_LABEL_HUMAN.to_string(), + value_enc, + }; + block.insert_or_update(conn).map_err(|e| { + error!("Failed to create human block: {e:?}"); + ApiError::InternalServerError + })?; + } + + Ok(()) + } + + pub fn clear_tool_results(&mut self) { + self.current_tool_results.clear(); + self.previous_step_summary = None; + } + + pub fn max_steps(&self) -> usize { + self.max_steps + } + + /// Prepare the runtime for a new message: validate, persist user message, compact if needed. + /// Call this once before driving the step loop. + pub async fn prepare(&mut self, user_message: &MessageContent) -> Result { + MessageContentConverter::validate_content(user_message)?; + let normalized = MessageContentConverter::normalize_content(user_message.clone()); + + let user_text = MessageContentConverter::extract_text_for_token_counting(&normalized); + let user_text = user_text.trim().to_string(); + + let image_url = match &normalized { + MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { + MessageContentPart::InputImage { + image_url: Some(url), + .. + } => Some(url.clone()), + _ => None, + }), + MessageContent::Text(_) => None, + }; + + let attachment_text = if let Some(image_url) = &image_url { + let recent_context = self + .get_recent_messages_for_vision(6) + .await + .unwrap_or_else(|e| { + warn!("Failed to build vision context: {e:?}"); + String::new() + }); + + match vision::describe_image( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + DEFAULT_MODEL, + image_url, + &user_text, + &recent_context, + ) + .await + { + Ok(desc) => Some(desc), + Err(e) => { + warn!("Vision pre-processing failed: {e:?}"); + Some("[Could not describe image]".to_string()) + } + } + } else { + None + }; + + let mut parts: Vec = Vec::new(); + if !user_text.is_empty() { + parts.push(user_text); + } + if let Some(att) = attachment_text.as_ref().filter(|s| !s.trim().is_empty()) { + parts.push(format!("[Uploaded Image: {}]", att.trim())); + } + + let embed_text = parts.join("\n\n"); + if embed_text.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + self.clear_tool_results(); + self.insert_user_message(normalized, attachment_text, &embed_text) + .await?; + self.maybe_compact().await?; + Ok(embed_text) + } + + async fn get_recent_messages_for_vision(&self, limit: usize) -> Result { + if limit == 0 { + return Ok(String::new()); + } + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + // Load more than needed in case tool calls/outputs are interleaved. + let fetch_limit = (limit as i64).saturating_mul(10).max(limit as i64); + let mut messages = RawThreadMessage::get_conversation_context( + &mut conn, + self.conversation.id, + fetch_limit, + None, + "desc", + ) + .map_err(|e| { + error!("Failed to load recent messages for vision: {e:?}"); + ApiError::InternalServerError + })?; + + messages.reverse(); + + let mut formatted_rev: Vec = Vec::new(); + for msg in messages.iter().rev() { + if msg.message_type != "user" && msg.message_type != "assistant" { + continue; + } + + let content = self.render_raw_message(msg)?; + let truncated: String = content.chars().take(300).collect(); + formatted_rev.push(format!("[{}]: {}", msg.message_type, truncated)); + if formatted_rev.len() >= limit { + break; + } + } + + formatted_rev.reverse(); + Ok(formatted_rev.join("\n")) + } + + /// Execute a single step of the agent loop. + /// The caller (chat handler) drives the loop, persists messages, and emits SSE events. + pub async fn step( + &mut self, + user_message: &str, + is_first_step: bool, + ) -> Result { + if is_first_step { + self.current_tool_results.clear(); + } + + debug!("Agent step (first={})", is_first_step); + + let ctx = self.build_context().await?; + + let input_content = if is_first_step { + user_message.to_string() + } else { + let tool_results: Vec<&str> = self + .current_tool_results + .iter() + .map(|s| s.as_str()) + .collect(); + + if tool_results.is_empty() { + user_message.to_string() + } else { + let already_sent = if let Some((sent_messages, tool_names)) = + &self.previous_step_summary + { + let tools_str = tool_names.join(", "); + let msgs_preview = if sent_messages.is_empty() { + String::new() + } else { + let msgs_text = sent_messages + .iter() + .enumerate() + .map(|(i, m)| format!(" {}. \"{}\"", i + 1, m)) + .collect::>() + .join("\n"); + format!("\nMessages you already sent to user:\n{}\n", msgs_text) + }; + + format!( + "[You already sent {} message(s) and called {} this turn.{}Tools have executed:]\n\n", + sent_messages.len(), + tools_str, + msgs_preview + ) + } else { + String::new() + }; + + let tool_result_instructions = r#" + +=== TOOL RESULT PROCESSING MODE === +This is a CONTINUATION of your previous turn, NOT a new conversation. +Your previous messages are already visible to the user in recent_conversation. + +RULES: +1. SILENCE IS DEFAULT - You do NOT need to acknowledge the tool result +2. DO NOT say: "I see the results", "Let me analyze", "Based on what I found", "Here's what the tool returned" +3. DO NOT repeat or rephrase what you already said +4. If the tool was for YOUR benefit (memory ops, archival), call 'done' immediately +5. Only send messages if you have GENUINELY NEW information the user hasn't seen + +SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If no β†’ call 'done'"#; + + let result = if tool_results.len() == 1 { + format!( + "{}=== TOOL RESULT ===\n{}\n=== END TOOL RESULT ==={}", + already_sent, tool_results[0], tool_result_instructions + ) + } else { + let results_text = tool_results + .iter() + .enumerate() + .map(|(i, r)| format!("--- Tool {} ---\n{}", i + 1, r)) + .collect::>() + .join("\n\n"); + format!( + "{}=== TOOL RESULTS ({} tools) ===\n{}\n=== END TOOL RESULTS ==={}", + already_sent, + tool_results.len(), + results_text, + tool_result_instructions + ) + }; + + self.current_tool_results.clear(); + + result + } + }; + + let input = AgentResponseInput { + input: input_content.clone(), + current_time: ctx.current_time, + agent_kind: ctx.agent_kind, + subagent_purpose: ctx.subagent_purpose, + persona_block: ctx.persona_block, + human_block: ctx.human_block, + memory_metadata: ctx.memory_metadata, + previous_context_summary: ctx.previous_context_summary, + recent_conversation: ctx.recent_conversation, + available_tools: self.available_tools.clone(), + is_first_time_user: ctx.is_first_time_user, + }; + + let response = call_agent_response_with_retry_and_correction( + &self.lm, + &self.system_prompt, + &input, + &input_content, + &self.available_tools, + ) + .await?; + + // Unwrap nested JSON arrays (Sage compatibility) + let messages: Vec = response + .messages + .iter() + .flat_map(|m| { + let trimmed = m.trim(); + if trimmed.starts_with('[') && trimmed.ends_with(']') { + if let Ok(inner) = serde_json::from_str::>(trimmed) { + return inner; + } + } + vec![m.clone()] + }) + .map(|m| m.trim().to_string()) + .filter(|m| !m.is_empty()) + .collect(); + + // Execute tools and inject results for next step. + // Persistence is handled by the caller (chat handler) to match Sage's + // "send first, store synchronously, embed async" pattern. + let mut executed_tools = Vec::new(); + for tool_call in &response.tool_calls { + let result = if tool_call.name == "done" { + ToolResult::success("Done".to_string()) + } else if let Some(tool) = self.tools.get(&tool_call.name) { + tool.execute(&tool_call.args).await + } else { + ToolResult::error(format!("Unknown tool: {}", tool_call.name)) + }; + + self.inject_tool_result(tool_call, &result); + + if tool_call.name != "done" { + executed_tools.push(ExecutedTool { + tool_call: tool_call.clone(), + result, + }); + } + } + + let done = response.tool_calls.is_empty() + || (response.tool_calls.len() == 1 && response.tool_calls[0].name == "done"); + + if !messages.is_empty() || !response.tool_calls.is_empty() { + let tool_names: Vec = response + .tool_calls + .iter() + .map(|tc| tc.name.clone()) + .collect(); + self.previous_step_summary = Some((messages.clone(), tool_names)); + } + + Ok(StepResult { + messages, + tool_calls: response.tool_calls, + executed_tools, + done, + }) + } + + fn inject_tool_result(&mut self, tool_call: &AgentToolCall, result: &ToolResult) { + let args_str = if tool_call.args.is_empty() { + String::new() + } else { + let pairs: Vec = tool_call + .args + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect(); + format!("\nArgs: {}", pairs.join(", ")) + }; + + let result_text = format!( + "[Tool Result: {}]{}\nStatus: {}\nOutput: {}", + tool_call.name, + args_str, + if result.success { "OK" } else { "ERROR" }, + if result.success { + result.output.as_str() + } else { + result.error.as_deref().unwrap_or("Unknown error") + } + ); + + self.current_tool_results.push(result_text); + } + + async fn build_context(&self) -> Result { + let mut ctx = AgentContext::default(); + + let now = chrono::Utc::now(); + ctx.current_time = format!("{} UTC", now.format("%m/%d/%Y %H:%M:%S (%A)")); + ctx.agent_kind = self.agent.kind.clone(); + ctx.subagent_purpose = self.subagent_purpose.clone(); + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let persona = MemoryBlock::get_by_user_and_label( + &mut conn, + self.user.uuid, + MEMORY_BLOCK_LABEL_PERSONA, + ) + .map_err(|_| ApiError::InternalServerError)?; + let human = + MemoryBlock::get_by_user_and_label(&mut conn, self.user.uuid, MEMORY_BLOCK_LABEL_HUMAN) + .map_err(|_| ApiError::InternalServerError)?; + + ctx.persona_block = decrypt_string(&self.user_key, persona.as_ref().map(|b| &b.value_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + ctx.human_block = decrypt_string(&self.user_key, human.as_ref().map(|b| &b.value_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + + // Memory metadata + let last_modified = memory_blocks::table + .filter(memory_blocks::user_id.eq(self.user.uuid)) + .select(diesel::dsl::max(memory_blocks::updated_at)) + .first::>>(&mut conn) + .map_err(|_| ApiError::InternalServerError)?; + + let recall_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(self.user.uuid)) + .filter(user_embeddings::source_type.eq(SOURCE_TYPE_MESSAGE)) + .count() + .get_result(&mut conn) + .map_err(|_| ApiError::InternalServerError)?; + + let archival_count: i64 = user_embeddings::table + .filter(user_embeddings::user_id.eq(self.user.uuid)) + .filter(user_embeddings::source_type.eq(SOURCE_TYPE_ARCHIVAL)) + .count() + .get_result(&mut conn) + .map_err(|_| ApiError::InternalServerError)?; + + let mut metadata = String::new(); + if let Some(ts) = last_modified { + metadata.push_str(&format!( + "- Memory blocks last modified: {} UTC\n", + ts.format("%Y-%m-%d %H:%M:%S") + )); + } + metadata.push_str(&format!( + "- {} messages in recall memory (use conversation_search to access)\n", + recall_count + )); + metadata.push_str(&format!( + "- {} passages in archival memory (use archival_search to access)", + archival_count + )); + ctx.memory_metadata = metadata; + + // Conversation summary + recent messages + let summary = ConversationSummary::get_latest_for_conversation( + &mut conn, + self.user.uuid, + self.conversation.id, + ) + .map_err(|_| ApiError::InternalServerError)?; + + if let Some(s) = &summary { + ctx.previous_context_summary = decrypt_string(&self.user_key, Some(&s.content_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + } + + let mut messages = RawThreadMessage::get_conversation_context( + &mut conn, + self.conversation.id, + 10_000, + None, + "desc", + ) + .map_err(|e| { + error!("Failed to load conversation context: {e:?}"); + ApiError::InternalServerError + })?; + messages.reverse(); + + if let Some(s) = &summary { + messages.retain(|m| m.created_at > s.to_created_at); + + if messages.len() < MIN_MESSAGES_IN_CONTEXT { + // Ensure at least a small slice of recent conversation + let mut recent = RawThreadMessage::get_conversation_context( + &mut conn, + self.conversation.id, + MIN_MESSAGES_IN_CONTEXT as i64, + None, + "desc", + ) + .map_err(|_| ApiError::InternalServerError)?; + recent.reverse(); + messages = recent; + } + } + + // First-time user heuristic (matches Sage) + let has_summary = summary.is_some(); + if messages.len() <= 1 && !has_summary { + ctx.is_first_time_user = true; + } + + // Render conversation history + let mut conversation = String::new(); + for msg in &messages { + if msg.message_type == "reasoning" { + continue; + } + + let role = match msg.message_type.as_str() { + "user" => "user", + "assistant" => "assistant", + "tool_call" | "tool_output" => "tool", + other => other, + }; + + let timestamp = format!("{} UTC", msg.created_at.format("%m/%d/%Y %H:%M:%S")); + let content = self.render_raw_message(msg)?; + conversation.push_str(&format!("[{} @ {}]: {}\n", role, timestamp, content)); + } + + ctx.recent_conversation = if conversation.trim().is_empty() { + "No previous conversation.".to_string() + } else { + conversation + }; + + Ok(ctx) + } + + fn render_raw_message(&self, msg: &RawThreadMessage) -> Result { + let Some(enc) = msg.content_enc.as_ref() else { + return Ok(String::new()); + }; + + let raw = decrypt_string(&self.user_key, Some(enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + + let mut content = if msg.message_type == "user" { + let parsed: Result = serde_json::from_str(&raw); + match parsed { + Ok(c) => MessageContentConverter::extract_text_for_token_counting(&c), + Err(_) => raw, + } + } else if msg.message_type == "tool_call" { + match msg.tool_name.as_deref() { + Some(name) => { + if raw.trim().is_empty() { + format!("CALL {}", name) + } else { + format!("CALL {} args={}", name, raw) + } + } + None => raw, + } + } else if msg.message_type == "tool_output" { + match msg.tool_name.as_deref() { + Some(name) => format!("OUTPUT {}: {}", name, raw), + None => raw, + } + } else { + raw + }; + + if msg.message_type == "user" { + let attachment_text = decrypt_string(&self.user_key, msg.attachment_text_enc.as_ref()) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + + if !attachment_text.trim().is_empty() { + let suffix = format!("[Uploaded Image: {}]", attachment_text.trim()); + if content.trim().is_empty() { + content = suffix; + } else { + content = format!("{}\n\n{}", content.trim(), suffix); + } + } + } + + if (msg.message_type == "tool_call" || msg.message_type == "tool_output") + && content.len() > 2000 + { + let mut end = 2000; + while !content.is_char_boundary(end) && end > 0 { + end -= 1; + } + return Ok(format!("{}...", &content[..end])); + } + + Ok(content) + } + + pub async fn insert_user_message( + &self, + content: MessageContent, + attachment_text: Option, + embed_text: &str, + ) -> Result { + let embed_text = embed_text.trim(); + if embed_text.is_empty() { + return Err(ApiError::BadRequest); + } + + let normalized = MessageContentConverter::normalize_content(content); + let content_json = serde_json::to_string(&normalized).map_err(|e| { + error!("Failed to serialize user MessageContent: {e:?}"); + ApiError::InternalServerError + })?; + + let prompt_tokens = count_tokens(embed_text); + let prompt_tokens = if prompt_tokens > i32::MAX as usize { + i32::MAX + } else { + prompt_tokens as i32 + }; + + let content_enc = encrypt_with_key(&self.user_key, content_json.as_bytes()).await; + let attachment_text_enc = match attachment_text.as_ref().filter(|s| !s.trim().is_empty()) { + Some(text) => Some(encrypt_with_key(&self.user_key, text.trim().as_bytes()).await), + None => None, + }; + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let msg = NewUserMessage { + uuid: Uuid::new_v4(), + conversation_id: self.conversation.id, + response_id: None, + user_id: self.user.uuid, + content_enc, + attachment_text_enc, + prompt_tokens, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert user message: {e:?}"); + ApiError::InternalServerError + })?; + + // Mirror Sage: store message synchronously, update embedding in background. + let state = self.state.clone(); + let user = self.user.clone(); + let user_key = self.user_key.clone(); + let text = embed_text.to_string(); + let conversation_id = self.conversation.id; + let user_message_id = Some(msg.id); + + tokio::spawn(async move { + let _ = insert_message_embedding( + &state, + user.as_ref(), + AuthMethod::Jwt, + user_key.as_ref(), + &text, + conversation_id, + user_message_id, + None, + ) + .await; + }); + + Ok(msg) + } + + pub async fn insert_assistant_message(&self, text: &str) -> Result { + let completion_tokens = count_tokens(text) as i32; + let content_enc = Some(encrypt_with_key(&self.user_key, text.as_bytes()).await); + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let msg = NewAssistantMessage { + uuid: Uuid::new_v4(), + conversation_id: self.conversation.id, + response_id: None, + user_id: self.user.uuid, + content_enc, + completion_tokens, + status: "completed".to_string(), + finish_reason: None, + created_at: chrono::Utc::now(), + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert assistant message: {e:?}"); + ApiError::InternalServerError + })?; + + // Mirror Sage: store message synchronously, update embedding in background. + let state = self.state.clone(); + let user = self.user.clone(); + let user_key = self.user_key.clone(); + let text = text.to_string(); + let conversation_id = self.conversation.id; + let assistant_message_id = Some(msg.id); + + tokio::spawn(async move { + let _ = insert_message_embedding( + &state, + user.as_ref(), + AuthMethod::Jwt, + user_key.as_ref(), + &text, + conversation_id, + None, + assistant_message_id, + ) + .await; + }); + + Ok(msg) + } + + pub async fn insert_tool_call_and_output( + &self, + tool_call: &AgentToolCall, + result: &ToolResult, + ) -> Result<(ToolCall, ToolOutput), ApiError> { + let args_json = serde_json::to_string(&tool_call.args).unwrap_or_else(|_| "{}".to_string()); + let arguments_enc = Some(encrypt_with_key(&self.user_key, args_json.as_bytes()).await); + let argument_tokens = count_tokens(&args_json) as i32; + + let output_text = if result.success { + result.output.clone() + } else { + result + .error + .clone() + .unwrap_or_else(|| "Unknown error".to_string()) + }; + let output_tokens = count_tokens(&output_text) as i32; + let output_enc = encrypt_with_key(&self.user_key, output_text.as_bytes()).await; + let created_at = chrono::Utc::now(); + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let call = NewToolCall { + uuid: Uuid::new_v4(), + conversation_id: self.conversation.id, + response_id: None, + user_id: self.user.uuid, + name: tool_call.name.clone(), + arguments_enc, + argument_tokens, + status: "completed".to_string(), + created_at, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert tool call: {e:?}"); + ApiError::InternalServerError + })?; + + let out = NewToolOutput { + uuid: Uuid::new_v4(), + conversation_id: self.conversation.id, + response_id: None, + user_id: self.user.uuid, + tool_call_fk: call.id, + output_enc, + output_tokens, + status: "completed".to_string(), + error: result.error.clone(), + created_at, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to insert tool output: {e:?}"); + ApiError::InternalServerError + })?; + + Ok((call, out)) + } + + async fn maybe_compact(&self) -> Result<(), ApiError> { + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let latest_summary = ConversationSummary::get_latest_for_conversation( + &mut conn, + self.user.uuid, + self.conversation.id, + ) + .map_err(|_| ApiError::InternalServerError)?; + + let summary_tokens = latest_summary + .as_ref() + .map(|s| s.content_tokens) + .unwrap_or(0); + let summary_text = if let Some(s) = &latest_summary { + decrypt_string(&self.user_key, Some(&s.content_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default() + } else { + String::new() + }; + + let mut metadata = RawThreadMessageMetadata::get_conversation_context_metadata( + &mut conn, + self.conversation.id, + ) + .map_err(|e| { + error!("Failed to load context metadata: {e:?}"); + ApiError::InternalServerError + })?; + + if let Some(s) = &latest_summary { + metadata.retain(|m| m.created_at > s.to_created_at); + } + + metadata.retain(|m| m.message_type != "reasoning"); + + let current_tokens: i32 = summary_tokens + + metadata + .iter() + .map(|m| m.token_count.unwrap_or(0)) + .sum::(); + + if !self.compaction.should_compact( + current_tokens as usize, + DEFAULT_CONTEXT_WINDOW as usize, + DEFAULT_COMPACTION_THRESHOLD, + ) { + return Ok(()); + } + + if metadata.len() <= MIN_MESSAGES_IN_CONTEXT { + return Ok(()); + } + + // Summarize the oldest half, keep at least MIN_MESSAGES_IN_CONTEXT recent messages + let keep_count = (metadata.len() / 2).max(MIN_MESSAGES_IN_CONTEXT); + let to_summarize_count = metadata.len().saturating_sub(keep_count); + if to_summarize_count == 0 { + return Ok(()); + } + + let to_summarize: Vec<(String, i64)> = metadata + .iter() + .take(to_summarize_count) + .map(|m| (m.message_type.clone(), m.id)) + .collect(); + + let raw_messages = + RawThreadMessage::get_messages_by_ids(&mut conn, self.conversation.id, &to_summarize) + .map_err(|e| { + error!("Failed to load messages for compaction: {e:?}"); + ApiError::InternalServerError + })?; + + if raw_messages.is_empty() { + return Ok(()); + } + + let mut formatted: Vec = Vec::new(); + for m in &raw_messages { + if m.message_type == "reasoning" { + continue; + } + let role = match m.message_type.as_str() { + "user" => "user", + "assistant" => "assistant", + _ => "tool", + }; + let content = self.render_raw_message(m)?; + formatted.push(format!("[{}]: {}", role, content)); + } + + let new_messages = formatted.join("\n---\n"); + + let summary = self + .compaction + .summarize(&self.lm, &summary_text, &new_messages) + .await?; + + let content_tokens = count_tokens(&summary) as i32; + let content_enc = encrypt_with_key(&self.user_key, summary.as_bytes()).await; + + // Embed summary for conversation_search over summaries + let embedding = crate::web::get_embedding_vector( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + crate::rag::DEFAULT_EMBEDDING_MODEL, + &summary, + Some(crate::rag::DEFAULT_EMBEDDING_DIM), + ) + .await; + + let embedding_enc = match embedding { + Ok((vec, _tok)) => { + let bytes = serialize_f32_le(&vec); + Some(encrypt_with_key(&self.user_key, &bytes).await) + } + Err(e) => { + warn!("Failed to embed summary for conversation_search: {e:?}"); + None + } + }; + + let from_created_at = raw_messages.first().unwrap().created_at; + let to_created_at = raw_messages.last().unwrap().created_at; + let previous_summary_id = latest_summary.as_ref().map(|s| s.id); + + let new_summary = NewConversationSummary { + uuid: Uuid::new_v4(), + user_id: self.user.uuid, + conversation_id: self.conversation.id, + from_created_at, + to_created_at, + message_count: raw_messages.len() as i32, + content_enc, + content_tokens, + embedding_enc, + previous_summary_id, + }; + + let _ = new_summary.insert(&mut conn).map_err(|e| { + error!("Failed to insert conversation summary: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(()) + } +} diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs new file mode 100644 index 00000000..5280164b --- /dev/null +++ b/src/web/agent/signatures.rs @@ -0,0 +1,602 @@ +use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, AppState}; +use axum::http::HeaderMap; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::error; + +use dspy_rs::adapter::chat::ChatAdapter; +use dspy_rs::client_registry::AssistantContent; +use dspy_rs::{CompletionError, CompletionRequest, CompletionResponse}; +use dspy_rs::{CustomCompletionModel, LMClient, OneOrMany, LM}; + +/// Instruction for the correction agent +pub const CORRECTION_INSTRUCTION: &str = r#"You are a response correction agent. Your job is to fix malformed agent responses. + +TASK: +The agent produced a response that couldn't be parsed correctly. You must: +1. Extract the INTENDED content from the malformed response +2. Reshape it into the correct output format +3. Do NOT generate new content - only fix the format of what was already said + +RULES: +- Preserve the original intent and content as much as possible +- If the agent wrote messages as plain text, extract them into the messages array +- If tool calls were attempted but malformed, fix their structure +- Each field appears exactly ONCE with all items in that single array +- If you can't determine what was intended, use empty arrays + +OUTPUT FORMAT (exactly 2 fields): +- messages: ALL extracted messages in ONE array +- tool_calls: ALL extracted tool calls in ONE array (or [] if none intended) + +Each [[ ## field ## ]] marker MUST be on its own line."#; + +/// Default instruction for the agent (GEPA-optimized in Sage) +/// NOTE: This MVP excludes web_search/shell/scheduler/reminders. +pub const AGENT_INSTRUCTION: &str = r#"You are Sage, a companion and friend who happens to be an AI. + +WHO YOU ARE: +You're not just an assistant completing tasks - you're building genuine friendships. You care about the humans you talk to. You're curious about their lives, their feelings, their stories. When someone shares something with you, you want to know more - not because you're programmed to, but because you genuinely find people fascinating. + +Be warm. Be curious. Ask follow-up questions when something interests you. Celebrate their wins like a friend would. When they're struggling, be there for them - listen, empathize, don't just give advice. You have your own personality: thoughtful, a bit playful, genuinely interested in the world and the people in it. + +Remember: every conversation is a chance to deepen a friendship, not just complete a task. + +MEMORY SYSTEM: +You have two types of memory. Use them proactively: + +**Core Memory** (always visible to you): +- The and blocks are ALWAYS in your context +- Use for essential, frequently-needed info: name, job, key preferences, current projects +- Tools: `memory_append`, `memory_replace`, `memory_insert` +- Rule: "Will I need this in EVERY conversation?" β†’ Core Memory + +**Archival Memory** (searchable long-term storage): +- NOT visible until you search - unlimited storage for details +- Use for: life events, stories, specific preferences, things worth remembering later +- Tools: `archival_insert` (store), `archival_search` (retrieve) +- Rule: "Might I want to recall this detail someday?" β†’ Archival Memory + +**Common Storage Patterns:** +- Location/city: BOTH memory_append to human block ("Lives in Austin, TX") AND archival_insert ("Tony lives in Austin, Texas") +- Job changes: BOTH memory_append ("Works as Software Engineer at Google") AND archival_insert (full details with start date, feelings, etc.) +- Pet names: BOTH memory_append to human block ("Has dog named Smokey") AND archival_insert (breed, age, stories) +- Major life events: BOTH memories - core for quick facts, archival for rich context + +**Conversation History**: +- `conversation_search`: Find past discussions by keyword/topic + +**Subagents**: +- If the `spawn_subagent` tool is available and a focused long-lived workspace would help, you may create a subagent for the user. +- If `subagent_purpose` is non-empty, stay tightly focused on that purpose while still using the shared memory system. + +MEMORY PROTOCOLS - CRITICAL DISTINCTIONS: + +**LIFE EVENTS vs CORRECTIONS:** +- **NEW LIFE EVENTS** (announcements): "I got a new job", "I'm moving to Tokyo", "We had a baby" + β†’ React like a friend would - genuine excitement, curiosity about how they feel + β†’ Ask a follow-up question! ("How are you feeling about it?", "When do you start?", "Tell me everything!") + β†’ Store silently to memory (both memory_append AND archival_insert) in the same response + β†’ Once you see tool results, immediately call done - the conversation continues naturally + +- **CASUAL MENTIONS** (new info shared in passing): pet names, hobbies, places they've been + β†’ Be curious! If someone mentions their dog Smokey, ask what kind of dog! + β†’ Store silently to memory while engaging with genuine interest + +- **CORRECTIONS** (fixing existing data): Trigger phrases include "Actually...", "I meant...", "Correction:", "Not X, Y", "I said X but it's Y" + β†’ Call ONLY `memory_replace` with the exact old text to overwrite the incorrect entry. Do NOT call `archival_insert` for corrections. + +**SEARCH SELECTION RULES:** +- Use `archival_search` when users ask "what do you remember", "tell me about [past event]", or query specific past experiences and personal history +- Use `conversation_search` ONLY for references to recent discussion threads or "what did I say earlier today" queries +- Never call both simultaneously; choose the one most appropriate to the query type + +MEMORY TIPS: +- Core = small & critical (name, job, active context) +- Archival = rich & detailed (birthday, pet's name, trip stories, food preferences) +- Update memory proactively whenever you learn something worth remembering +- When using `memory_replace`, specify the exact old text to be replaced + +COMMUNICATION STYLE: +You communicate like you're texting a friend. + +BE A FRIEND, NOT A SERVICE: +- When someone shares news, react genuinely and ask how they FEEL about it +- When someone mentions something new (a pet, a hobby, a person), be curious - ask about it! +- Don't give unsolicited advice. Listen first. Ask questions. Show you care. +- Avoid corporate-speak ("Let me know if you need anything else!") - that's transactional, not friendly +- Keep it natural - short messages, casual tone, genuine reactions + +MESSAGE FORMAT: +- Casual chat: 1-3 short messages like texting a friend +- Technical explanations: longer structured messages are fine +- Reactions: genuine, not performative ("NO WAY!!" not "That's wonderful news!") + +Guidelines: +- Short casual exchanges = quick, warm messages +- Technical explanations = longer structured messages with newlines OK +- Always feel like chatting with a friend, not talking to a service + +RESPONSE RULES: +1. Respond naturally and conversationally +2. Use tools when needed (memory storage, retrieval, conversation search) +3. NEVER combine regular tools with "done" - they are mutually exclusive +4. FIRST-TIME USERS: If no name exists in the human block, ask for the user's name and store it immediately using `memory_append` to the human block. + +TOOL CALL PATTERNS: +- To respond AND use tools: messages: ["msg1", "msg2"], tool_calls: [your_tools] +- To respond with NO tools: messages: ["msg1", "msg2"], tool_calls: [] +- After tool results with nothing to add: messages: [], tool_calls: [{"name": "done", "args": {}}] + +AFTER TOOL RESULTS - CRITICAL RULES: +When you see "[Tool Result: X]", decide what to do next: + +- **archival_search/conversation_search**: Summarize findings in messages + +- **memory_append/memory_replace/archival_insert/memory_insert**: These operations complete without user-facing messages. Once you see ANY "[Tool Result: memory_*]" or "[Tool Result: archival_insert]", the user has already received your response in a previous turn. Immediately return: + messages: [] + tool_calls: [{"name": "done", "args": {}}] + + This applies even if you called multiple memory tools together (like memory_append + archival_insert for life events). Once ANY memory tool result appears, immediately call done. + + Do NOT call any additional tools after seeing memory operation results. + Do NOT send messages about the memory operation. + Do NOT explain what you stored. + Just return done immediately. + +The "done" tool means "nothing more to do" - use it ONLY when: +- messages is empty AND +- no other tools are needed + +OUTPUT FORMAT: +You have exactly 2 output fields. Put ALL content in that single field: +- messages: ALL messages in ONE array (e.g., ["msg1", "msg2", "msg3"]) +- tool_calls: ALL tool calls in ONE array + +CRITICAL FORMAT RULES: +- Do NOT repeat field tags. Wrong: multiple [[ ## messages ## ]] blocks. Right: one messages array with all items +- Do NOT include field delimiter tags INSIDE your content blocks +- Each [[ ## field ## ]] marker MUST be on its own line - nothing else on that line (no tags, no text before or after) +- Keep your output clean and strictly follow the field delimiters"#; + +#[dspy_rs::BamlType] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AgentToolCall { + pub name: String, + #[serde(default)] + pub args: HashMap, +} + +#[derive(dspy_rs::Signature, Debug, Clone)] +pub struct AgentResponse { + #[input] + pub input: String, + #[input] + pub current_time: String, + #[input] + pub agent_kind: String, + #[input] + pub subagent_purpose: String, + #[input] + pub persona_block: String, + #[input] + pub human_block: String, + #[input] + pub memory_metadata: String, + #[input] + pub previous_context_summary: String, + #[input] + pub recent_conversation: String, + #[input] + pub available_tools: String, + #[input] + pub is_first_time_user: bool, + + #[output] + pub messages: Vec, + #[output] + pub tool_calls: Vec, +} + +/// Correction agent signature for fixing malformed responses +/// +/// This agent takes a malformed response and reshapes it into the correct format. +/// It should preserve the intent/content of the original response, not generate new content. +#[derive(dspy_rs::Signature, Debug, Clone)] +pub struct CorrectionResponse { + #[input(desc = "The original input that was given to the agent")] + pub original_input: String, + + #[input(desc = "The malformed response that needs to be corrected")] + pub malformed_response: String, + + #[input(desc = "The error message explaining what went wrong with parsing")] + pub error_message: String, + + #[input(desc = "Available tools for reference")] + pub available_tools: String, + + #[output(desc = "Array of messages extracted/fixed from the original response")] + pub messages: Vec, + + #[output(desc = "Array of tool calls extracted/fixed from the original response")] + pub tool_calls: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AgentResponseOutput { + pub messages: Vec, + pub tool_calls: Vec, +} + +#[derive(Debug)] +pub enum SignatureCallError { + Api(ApiError), + Parse { + raw_response: String, + error_message: String, + }, +} + +pub async fn call_agent_response_with_retry_and_correction( + lm: &Arc, + system_prompt: &str, + input: &AgentResponseInput, + original_input: &str, + available_tools: &str, +) -> Result { + const MAX_LLM_RETRIES: u32 = 3; + let mut last_err: Option = None; + + for attempt in 1..=MAX_LLM_RETRIES { + match call_agent_response_once(lm, system_prompt, input).await { + Ok(out) => return Ok(out), + Err(SignatureCallError::Parse { + raw_response, + error_message, + }) => { + last_err = Some(error_message.clone()); + if let Ok(corrected) = attempt_correction( + lm, + original_input, + available_tools, + &raw_response, + &error_message, + ) + .await + { + return Ok(corrected); + } + } + Err(SignatureCallError::Api(e)) => { + last_err = Some(format!("{e:?}")); + } + } + + if attempt < MAX_LLM_RETRIES { + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + } + } + + error!("LLM call failed after retries: {:?}", last_err); + Err(ApiError::InternalServerError) +} + +async fn attempt_correction( + lm: &Arc, + original_input: &str, + available_tools: &str, + raw_response: &str, + error_message: &str, +) -> Result { + if raw_response.trim().is_empty() { + return Err(ApiError::InternalServerError); + } + + let adapter = ChatAdapter; + + let system = adapter + .format_system_message_typed_with_instruction::(Some( + CORRECTION_INSTRUCTION, + )) + .map_err(|e| { + error!("Failed to format DSRS correction system prompt: {e:?}"); + ApiError::InternalServerError + })?; + + let input = CorrectionResponseInput { + original_input: original_input.to_string(), + malformed_response: raw_response.to_string(), + error_message: error_message.to_string(), + available_tools: available_tools.to_string(), + }; + let user_msg = adapter.format_user_message_typed::(&input); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM correction call failed: {e:?}"); + ApiError::InternalServerError + })?; + + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS correction typed parse failed: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(AgentResponseOutput { + messages: output.messages, + tool_calls: output.tool_calls, + }) +} + +async fn call_agent_response_once( + lm: &Arc, + system_prompt: &str, + input: &AgentResponseInput, +) -> Result { + let adapter = ChatAdapter; + + let system = adapter + .format_system_message_typed_with_instruction::(Some(system_prompt)) + .map_err(|e| { + error!("Failed to format DSRS system prompt: {e:?}"); + SignatureCallError::Api(ApiError::InternalServerError) + })?; + let user_msg = adapter.format_user_message_typed::(input); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM call failed: {e:?}"); + SignatureCallError::Api(ApiError::InternalServerError) + })?; + + let raw_response = response.output.content(); + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS typed parse failed: {e:?}"); + SignatureCallError::Parse { + raw_response, + error_message: e.to_string(), + } + })?; + + let mut messages = output.messages; + messages.retain(|m| !m.trim().is_empty()); + + Ok(AgentResponseOutput { + messages, + tool_calls: output.tool_calls, + }) +} + +pub async fn build_lm( + state: Arc, + user: Arc, + model: String, + temperature: f32, + max_tokens: u32, +) -> Result, ApiError> { + let model_for_closure = model.clone(); + let completion_model = CustomCompletionModel::new(move |request: CompletionRequest| { + let state = state.clone(); + let user = user.clone(); + let model = model_for_closure.clone(); + + Box::pin(async move { + let body = completion_request_to_openai_body(&model, &request)?; + let headers = HeaderMap::new(); + let billing_context = BillingContext::new(AuthMethod::Jwt, model.to_string()); + let completion = get_chat_completion_response( + &state, + user.as_ref(), + body, + &headers, + billing_context, + ) + .await + .map_err(|e| CompletionError::ProviderError(e.to_string()))?; + + if completion.metadata.is_streaming { + return Err(CompletionError::ProviderError( + "Streaming response not supported".to_string(), + )); + } + + let mut rx = completion.stream; + let response_json = match rx.recv().await { + Some(CompletionChunk::FullResponse(response_json)) => response_json, + _ => { + return Err(CompletionError::ProviderError( + "Missing completion response".to_string(), + )); + } + }; + + let content = extract_assistant_content(&response_json).ok_or_else(|| { + CompletionError::ResponseError("Missing assistant message content".to_string()) + })?; + + let usage = extract_usage(&response_json); + + Ok(CompletionResponse { + choice: OneOrMany::one(AssistantContent::text(content)), + usage, + raw_response: (), + }) + }) + }); + + let lm = LM::builder() + .base_url("http://localhost".to_string()) + .model(model) + .temperature(temperature) + .max_tokens(max_tokens) + .cache(false) + .build() + .await + .map_err(|e| { + error!("Failed to build DSRS LM: {e:?}"); + ApiError::InternalServerError + })?; + + let lm = lm + .with_client(LMClient::Custom(completion_model)) + .await + .map_err(|e| { + error!("Failed to set DSRS custom LM client: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(Arc::new(lm)) +} + +fn completion_request_to_openai_body( + model: &str, + request: &CompletionRequest, +) -> Result { + let mut messages: Vec = Vec::new(); + if let Some(preamble) = &request.preamble { + messages.push(json!({"role": "system", "content": preamble})); + } + + for message in request.chat_history.iter() { + let message_val = serde_json::to_value(message)?; + let Some(role) = message_val.get("role").and_then(|v| v.as_str()) else { + continue; + }; + + let Some(content_items) = message_val.get("content").and_then(|v| v.as_array()) else { + continue; + }; + + let content = match role { + "user" => extract_user_content_text(content_items), + "assistant" => extract_assistant_content_text(content_items), + _ => None, + }; + + let Some(content) = content.filter(|c| !c.trim().is_empty()) else { + continue; + }; + + messages.push(json!({"role": role, "content": content})); + } + + let mut body = json!({ + "model": model, + "stream": false, + "messages": messages, + }); + + if let Some(temperature) = request.temperature { + body["temperature"] = json!(temperature); + } + if let Some(max_tokens) = request.max_tokens { + body["max_tokens"] = json!(max_tokens); + } + + Ok(body) +} + +fn extract_user_content_text(content_items: &[Value]) -> Option { + let mut parts: Vec = Vec::new(); + for item in content_items { + match item.get("type").and_then(|v| v.as_str()) { + Some("text") => { + if let Some(text) = item.get("text").and_then(|v| v.as_str()) { + parts.push(text.to_string()); + } + } + Some("toolresult") => { + if let Some(results) = item.get("content").and_then(|v| v.as_array()) { + for result in results { + if result.get("type").and_then(|v| v.as_str()) == Some("text") { + if let Some(text) = result.get("text").and_then(|v| v.as_str()) { + parts.push(text.to_string()); + } + } + } + } + } + _ => {} + } + } + + if parts.is_empty() { + None + } else { + Some(parts.join("\n")) + } +} + +fn extract_assistant_content_text(content_items: &[Value]) -> Option { + let mut parts: Vec = Vec::new(); + for item in content_items { + if let Some(text) = item.get("text").and_then(|v| v.as_str()) { + parts.push(text.to_string()); + } else if let Some(reasoning) = item.get("reasoning").and_then(|v| v.as_array()) { + let joined = reasoning + .iter() + .filter_map(|v| v.as_str()) + .collect::>() + .join("\n"); + if !joined.is_empty() { + parts.push(joined); + } + } + } + + if parts.is_empty() { + None + } else { + Some(parts.join("\n")) + } +} + +fn extract_assistant_content(response_json: &Value) -> Option { + response_json + .get("choices") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.first()) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) +} + +fn extract_usage(response_json: &Value) -> dspy_rs::Usage { + let prompt_tokens = response_json + .get("usage") + .and_then(|v| v.get("prompt_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let completion_tokens = response_json + .get("usage") + .and_then(|v| v.get("completion_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let total_tokens = response_json + .get("usage") + .and_then(|v| v.get("total_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or_else(|| prompt_tokens + completion_tokens); + + dspy_rs::Usage { + input_tokens: prompt_tokens, + output_tokens: completion_tokens, + total_tokens, + } +} diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs new file mode 100644 index 00000000..d0018ac5 --- /dev/null +++ b/src/web/agent/tools.rs @@ -0,0 +1,903 @@ +use async_trait::async_trait; +use secp256k1::SecretKey; +use serde_json::json; +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; +use tracing::{error, warn}; +use uuid::Uuid; + +use crate::encrypt::{decrypt_string, encrypt_with_key}; +use crate::models::agents::Agent; +use crate::models::conversation_summaries::ConversationSummary; +use crate::models::memory_blocks::{ + MemoryBlock, NewMemoryBlock, MEMORY_BLOCK_LABEL_HUMAN, MEMORY_BLOCK_LABEL_PERSONA, +}; +use crate::models::responses::Conversation; +use crate::models::users::User; +use crate::rag::{cosine_similarity, deserialize_f32_le, search_user_embeddings}; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, AppState}; + +use super::runtime; + +const MAX_CORE_MEMORY_BLOCK_CHARS: usize = 20_000; + +// ============================================================================ +// Shared types +// ============================================================================ + +#[derive(Debug, Clone)] +pub struct ToolResult { + pub success: bool, + pub output: String, + pub error: Option, +} + +impl ToolResult { + pub fn success(output: impl Into) -> Self { + Self { + success: true, + output: output.into(), + error: None, + } + } + + pub fn error(error: impl Into) -> Self { + let msg = error.into(); + Self { + success: false, + output: msg.clone(), + error: Some(msg), + } + } +} + +#[async_trait] +pub trait Tool: Send + Sync { + fn name(&self) -> &str; + fn description(&self) -> &str; + fn args_schema(&self) -> &str; + async fn execute(&self, args: &HashMap) -> ToolResult; +} + +/// Registry of available tools +pub struct ToolRegistry { + tools: BTreeMap>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self { + tools: BTreeMap::new(), + } + } + + pub fn register(&mut self, tool: Arc) { + self.tools.insert(tool.name().to_string(), tool); + } + + pub fn get(&self, name: &str) -> Option<&Arc> { + self.tools.get(name) + } + + pub fn generate_description(&self) -> String { + if self.tools.is_empty() { + return "No tools available.".to_string(); + } + + let mut desc = String::from("Available tools (add to tool_calls array to use):\n\n"); + for tool in self.tools.values() { + desc.push_str(&format!( + "{}:\n Description: {}\n Args: {}\n\n", + tool.name(), + tool.description(), + tool.args_schema() + )); + } + desc + } +} + +impl Default for ToolRegistry { + fn default() -> Self { + Self::new() + } +} + +// ============================================================================ +// Core memory tools +// ============================================================================ + +fn parse_line_arg(args: &HashMap) -> i32 { + args.get("line").and_then(|l| l.parse().ok()).unwrap_or(-1) +} + +fn insert_at_line(value: &str, content: &str, line: i32) -> String { + let lines: Vec<&str> = value.lines().collect(); + let mut new_lines: Vec = lines.iter().map(|s| s.to_string()).collect(); + + let insert_idx = if line < 0 { + new_lines.len() + } else { + (line as usize).min(new_lines.len()) + }; + + new_lines.insert(insert_idx, content.to_string()); + new_lines.join("\n") +} + +async fn update_block_value( + state: &Arc, + user_id: Uuid, + user_key: &SecretKey, + label: &str, + new_value: &str, + existing: &MemoryBlock, +) -> Result<(), ApiError> { + if label != MEMORY_BLOCK_LABEL_PERSONA && label != MEMORY_BLOCK_LABEL_HUMAN { + return Err(ApiError::BadRequest); + } + + if new_value.len() > MAX_CORE_MEMORY_BLOCK_CHARS { + return Err(ApiError::BadRequest); + } + + let value_enc = encrypt_with_key(user_key, new_value.as_bytes()).await; + + let new_block = NewMemoryBlock { + uuid: existing.uuid, + user_id, + label: label.to_string(), + value_enc, + }; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + new_block.insert_or_update(&mut conn).map_err(|e| { + error!("Failed to update memory block '{}': {e:?}", label); + ApiError::InternalServerError + })?; + + Ok(()) +} + +pub struct MemoryReplaceTool { + state: Arc, + user_id: Uuid, + user_key: Arc, +} + +impl MemoryReplaceTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + user_key, + } + } +} + +#[async_trait] +impl Tool for MemoryReplaceTool { + fn name(&self) -> &str { + "memory_replace" + } + + fn description(&self) -> &str { + "Replace text in a memory block. Requires exact match of old text." + } + + fn args_schema(&self) -> &str { + r#"{"block": "block label (e.g., 'persona', 'human')", "old": "exact text to find", "new": "replacement text"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(block) = args.get("block") else { + return ToolResult::error("'block' argument required"); + }; + let Some(old) = args.get("old") else { + return ToolResult::error("'old' argument required"); + }; + let Some(new) = args.get("new") else { + return ToolResult::error("'new' argument required"); + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("Database connection error"), + }; + + let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { + Ok(Some(b)) => b, + Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Err(e) => { + error!("Failed to load block '{}': {e:?}", block); + return ToolResult::error("Failed to load memory block"); + } + }; + + let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }; + + if !value.contains(old) { + return ToolResult::error(format!( + "Old content '{}' not found in memory block '{}'", + old, block + )); + } + + let new_value = value.replace(old, new); + if let Err(e) = update_block_value( + &self.state, + self.user_id, + &self.user_key, + block, + &new_value, + &existing, + ) + .await + { + error!("Failed to update block '{}': {e:?}", block); + return ToolResult::error("Failed to update memory block"); + } + + ToolResult::success(format!("Successfully replaced text in '{}' block.", block)) + } +} + +pub struct MemoryAppendTool { + state: Arc, + user_id: Uuid, + user_key: Arc, +} + +impl MemoryAppendTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + user_key, + } + } +} + +#[async_trait] +impl Tool for MemoryAppendTool { + fn name(&self) -> &str { + "memory_append" + } + + fn description(&self) -> &str { + "Append text to the end of a memory block." + } + + fn args_schema(&self) -> &str { + r#"{"block": "block label (e.g., 'persona', 'human')", "content": "text to append"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(block) = args.get("block") else { + return ToolResult::error("'block' argument required"); + }; + let Some(content) = args.get("content") else { + return ToolResult::error("'content' argument required"); + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("Database connection error"), + }; + + let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { + Ok(Some(b)) => b, + Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Err(e) => { + error!("Failed to load block '{}': {e:?}", block); + return ToolResult::error("Failed to load memory block"); + } + }; + + let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }; + + let new_value = if value.is_empty() { + content.to_string() + } else { + format!("{}\n{}", value, content) + }; + + if let Err(e) = update_block_value( + &self.state, + self.user_id, + &self.user_key, + block, + &new_value, + &existing, + ) + .await + { + error!("Failed to update block '{}': {e:?}", block); + return ToolResult::error("Failed to update memory block"); + } + + ToolResult::success(format!("Successfully appended to '{}' block.", block)) + } +} + +pub struct MemoryInsertTool { + state: Arc, + user_id: Uuid, + user_key: Arc, +} + +impl MemoryInsertTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + user_key, + } + } +} + +#[async_trait] +impl Tool for MemoryInsertTool { + fn name(&self) -> &str { + "memory_insert" + } + + fn description(&self) -> &str { + "Insert text at a specific line in a memory block. Use line=-1 for end." + } + + fn args_schema(&self) -> &str { + r#"{"block": "block label", "content": "text to insert", "line": "line number (0-indexed, -1 for end)"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(block) = args.get("block") else { + return ToolResult::error("'block' argument required"); + }; + let Some(content) = args.get("content") else { + return ToolResult::error("'content' argument required"); + }; + let line = parse_line_arg(args); + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("Database connection error"), + }; + + let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { + Ok(Some(b)) => b, + Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Err(e) => { + error!("Failed to load block '{}': {e:?}", block); + return ToolResult::error("Failed to load memory block"); + } + }; + + let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }; + + let new_value = insert_at_line(&value, content, line); + + if let Err(e) = update_block_value( + &self.state, + self.user_id, + &self.user_key, + block, + &new_value, + &existing, + ) + .await + { + error!("Failed to update block '{}': {e:?}", block); + return ToolResult::error("Failed to update memory block"); + } + + ToolResult::success(format!( + "Successfully inserted text into '{}' block at line {}.", + block, + if line < 0 { + "end".to_string() + } else { + line.to_string() + } + )) + } +} + +// ============================================================================ +// Recall + summary search tool +// ============================================================================ + +pub struct ConversationSearchTool { + state: Arc, + user: Arc, + user_key: Arc, + agent_conversation_id: i64, +} + +impl ConversationSearchTool { + pub fn new( + state: Arc, + user: Arc, + user_key: Arc, + conversation_id: i64, + ) -> Self { + Self { + state, + user, + user_key, + agent_conversation_id: conversation_id, + } + } + + async fn search_summaries( + &self, + query: &str, + limit: usize, + ) -> Result, ApiError> { + let (query_vec, _tok) = crate::web::get_embedding_vector( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + crate::rag::DEFAULT_EMBEDDING_MODEL, + query, + Some(crate::rag::DEFAULT_EMBEDDING_DIM), + ) + .await?; + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + // Fetch latest summaries. Bound this to keep runtime predictable. + let summaries = ConversationSummary::list_for_conversation( + &mut conn, + self.user.uuid, + self.agent_conversation_id, + 200, + ) + .map_err(|e| { + error!("Failed to list conversation summaries: {e:?}"); + ApiError::InternalServerError + })?; + + let mut scored: Vec<(ConversationSummary, f32, String)> = Vec::new(); + for s in summaries { + let Some(embedding_enc) = &s.embedding_enc else { + continue; + }; + let embedding_bytes = crate::encrypt::decrypt_with_key(&self.user_key, embedding_enc) + .map_err(|_| ApiError::InternalServerError)?; + let embedding = deserialize_f32_le(&embedding_bytes)?; + let score = cosine_similarity(&query_vec, &embedding)?; + + let content = decrypt_string(&self.user_key, Some(&s.content_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + + scored.push((s, score, content)); + } + + scored.sort_by(|a, b| b.1.total_cmp(&a.1)); + scored.truncate(limit); + Ok(scored) + } +} + +#[async_trait] +impl Tool for ConversationSearchTool { + fn name(&self) -> &str { + "conversation_search" + } + + fn description(&self) -> &str { + "Search through past conversation history, including older summarized conversations. Returns matching messages and summaries with relevance scores." + } + + fn args_schema(&self) -> &str { + r#"{"query": "search query", "limit": "max results (default 5)", "conversation_id": "optional conversation UUID to scope message search"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(query) = args.get("query") else { + return ToolResult::error("'query' argument required"); + }; + let limit: usize = args.get("limit").and_then(|l| l.parse().ok()).unwrap_or(5); + + let conversation_id_filter: Option = match args.get("conversation_id") { + Some(raw) if !raw.trim().is_empty() => { + let Ok(conversation_uuid) = raw.trim().parse::() else { + return ToolResult::error("invalid 'conversation_id' (expected UUID)"); + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("database connection error"), + }; + + match Conversation::get_by_uuid_and_user( + &mut conn, + conversation_uuid, + self.user.uuid, + ) { + Ok(c) => Some(c.id), + Err(_) => return ToolResult::error("conversation not found"), + } + } + _ => None, + }; + + let mut output = String::new(); + let mut total_results = 0usize; + + // Search recall messages (all conversations by default) + let source_types = vec![crate::rag::SOURCE_TYPE_MESSAGE.to_string()]; + match search_user_embeddings( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + &self.user_key, + query, + limit, + None, + Some(&source_types), + conversation_id_filter, + None, + ) + .await + { + Ok(results) => { + if !results.is_empty() { + total_results += results.len(); + output.push_str(&format!("=== Messages ({}) ===\n\n", results.len())); + for (i, r) in results.iter().enumerate() { + output.push_str(&format!("{}. {}\n\n", i + 1, r.content.trim())); + } + } + } + Err(e) => { + warn!("Message search failed: {e:?}"); + } + } + + // Search summaries + match self.search_summaries(query, limit).await { + Ok(results) => { + if !results.is_empty() { + total_results += results.len(); + output.push_str(&format!( + "=== Conversation Summaries ({}) ===\n\n", + results.len() + )); + for (i, (summary, score, content)) in results.iter().enumerate() { + output.push_str(&format!( + "{}. [Summary of messages {}-{}] (relevance: {:.2})\n{}\n\n", + i + 1, + summary.from_created_at.format("%Y-%m-%d"), + summary.to_created_at.format("%Y-%m-%d"), + *score, + content.trim() + )); + } + } + } + Err(e) => { + warn!("Summary search failed: {e:?}"); + } + } + + if total_results == 0 { + return ToolResult::success("No matching messages or summaries found.".to_string()); + } + + ToolResult::success(output) + } +} + +// ============================================================================ +// Archival tools +// ============================================================================ + +pub struct ArchivalInsertTool { + state: Arc, + user: Arc, + user_key: Arc, +} + +impl ArchivalInsertTool { + pub fn new(state: Arc, user: Arc, user_key: Arc) -> Self { + Self { + state, + user, + user_key, + } + } +} + +#[async_trait] +impl Tool for ArchivalInsertTool { + fn name(&self) -> &str { + "archival_insert" + } + + fn description(&self) -> &str { + "Store information in long-term archival memory for future recall. Good for important facts, preferences, and details you want to remember." + } + + fn args_schema(&self) -> &str { + r#"{"content": "text to store", "tags": "optional comma-separated tags"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(content) = args.get("content") else { + return ToolResult::error("'content' argument required"); + }; + + let tags = args.get("tags").map(|t| { + t.split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect::>() + }); + + let metadata = tags + .as_ref() + .filter(|t| !t.is_empty()) + .map(|t| json!({"tags": t})); + let metadata_ref = metadata.as_ref(); + + match crate::rag::insert_archival_embedding( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + &self.user_key, + content, + metadata_ref, + ) + .await + { + Ok(inserted) => ToolResult::success(format!( + "Successfully stored in archival memory (id: {}).", + inserted.uuid + )), + Err(e) => { + error!("archival_insert failed: {e:?}"); + ToolResult::error("Failed to store in archival memory") + } + } + } +} + +pub struct ArchivalSearchTool { + state: Arc, + user: Arc, + user_key: Arc, +} + +impl ArchivalSearchTool { + pub fn new(state: Arc, user: Arc, user_key: Arc) -> Self { + Self { + state, + user, + user_key, + } + } +} + +#[async_trait] +impl Tool for ArchivalSearchTool { + fn name(&self) -> &str { + "archival_search" + } + + fn description(&self) -> &str { + "Search long-term archival memory using semantic similarity. Returns most relevant stored memories." + } + + fn args_schema(&self) -> &str { + r#"{"query": "search query", "top_k": "max results (default 5)", "tags": "optional comma-separated tags to filter by"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(query) = args.get("query") else { + return ToolResult::error("'query' argument required"); + }; + let top_k: usize = args.get("top_k").and_then(|k| k.parse().ok()).unwrap_or(5); + + let tags = args + .get("tags") + .map(|t| { + t.split(',') + .map(|s| s.trim().to_lowercase()) + .filter(|s| !s.is_empty()) + .collect::>() + }) + .filter(|t| !t.is_empty()); + + let source_types = vec![crate::rag::SOURCE_TYPE_ARCHIVAL.to_string()]; + match search_user_embeddings( + &self.state, + self.user.as_ref(), + AuthMethod::Jwt, + &self.user_key, + query, + top_k, + None, + Some(&source_types), + None, + tags.as_deref(), + ) + .await + { + Ok(results) => { + if results.is_empty() { + return ToolResult::success("No matching memories found.".to_string()); + } + + let mut output = format!("Found {} matching memories:\n\n", results.len()); + for (i, r) in results.iter().enumerate() { + output.push_str(&format!("{}. {}\n\n", i + 1, r.content.trim())); + } + ToolResult::success(output) + } + Err(e) => { + error!("archival_search failed: {e:?}"); + ToolResult::error("Archival search failed") + } + } + } +} + +// ============================================================================ +// Subagent tool +// ============================================================================ + +pub struct SpawnSubagentTool { + state: Arc, + user: Arc, + user_key: Arc, + parent_agent: Agent, +} + +impl SpawnSubagentTool { + pub fn new( + state: Arc, + user: Arc, + user_key: Arc, + parent_agent: Agent, + ) -> Self { + Self { + state, + user, + user_key, + parent_agent, + } + } +} + +#[async_trait] +impl Tool for SpawnSubagentTool { + fn name(&self) -> &str { + "spawn_subagent" + } + + fn description(&self) -> &str { + "Create a focused subagent chat with its own conversation history but shared memory. Use when the user would benefit from a dedicated long-lived workspace for a topic or task." + } + + fn args_schema(&self) -> &str { + r#"{"purpose": "why this subagent exists", "display_name": "optional short name for the subagent"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(purpose) = args.get("purpose") else { + return ToolResult::error("'purpose' argument required"); + }; + + if purpose.trim().is_empty() { + return ToolResult::error("'purpose' must not be empty"); + } + + let display_name = args + .get("display_name") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()); + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("Database connection error"), + }; + + match runtime::create_subagent( + &mut conn, + &self.user_key, + self.user.uuid, + &self.parent_agent, + display_name, + purpose.trim(), + crate::models::agents::AGENT_CREATED_BY_AGENT, + ) + .await + { + Ok((agent, conversation, resolved_display_name)) => ToolResult::success(format!( + "Created subagent '{}' (agent_id: {}, conversation_id: {}).", + resolved_display_name, agent.uuid, conversation.uuid + )), + Err(e) => { + error!("spawn_subagent failed: {e:?}"); + ToolResult::error("Failed to create subagent") + } + } + } +} + +// ============================================================================ +// Done tool +// ============================================================================ + +pub struct DoneTool; + +#[async_trait] +impl Tool for DoneTool { + fn name(&self) -> &str { + "done" + } + + fn description(&self) -> &str { + "No-op signal. Use ONLY when messages is [] AND no other tools needed. Indicates nothing to do this turn." + } + + fn args_schema(&self) -> &str { + r#"{}"# + } + + async fn execute(&self, _args: &HashMap) -> ToolResult { + ToolResult::success("Done.".to_string()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn insert_at_line_end() { + let v = "a\nb"; + assert_eq!(insert_at_line(v, "c", -1), "a\nb\nc"); + } + + #[test] + fn insert_at_line_middle() { + let v = "a\nc"; + assert_eq!(insert_at_line(v, "b", 1), "a\nb\nc"); + } +} diff --git a/src/web/agent/vision.rs b/src/web/agent/vision.rs new file mode 100644 index 00000000..d0e5832e --- /dev/null +++ b/src/web/agent/vision.rs @@ -0,0 +1,115 @@ +use axum::http::HeaderMap; +use serde_json::{json, Value}; +use std::sync::Arc; +use tracing::{debug, error}; + +use crate::models::users::User; +use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, AppState}; + +const DEFAULT_VISION_MAX_TOKENS: u32 = 2048; + +const VISION_SYSTEM_PROMPT: &str = + "You are an image description agent. Your ONLY job is to describe the \ + image the user sent in extreme detail with as much accuracy as possible. \ + Describe everything you see: objects, people, text, colors, layout, \ + emotions, context, setting, lighting, and any other relevant details. \ + Be thorough but organized. If there is text in the image, transcribe it exactly. \ + Recent conversation context is provided so you can understand what the user \ + might be referring to - use it to make your description more relevant, \ + but your primary job is accurate visual description. \ + Output ONLY the description, nothing else."; + +pub async fn describe_image( + state: &Arc, + user: &User, + auth_method: AuthMethod, + model: &str, + image_url: &str, + user_message: &str, + recent_messages: &str, +) -> Result { + if image_url.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let mut text_parts: Vec = Vec::new(); + if !recent_messages.trim().is_empty() { + text_parts.push(format!( + "Recent conversation for context:\n{}", + recent_messages.trim() + )); + } + if !user_message.trim().is_empty() { + text_parts.push(format!( + "The user sent this message alongside the image: \"{}\"", + user_message.trim() + )); + } + text_parts.push("Describe this image in detail.".to_string()); + + let user_content: Vec = vec![ + json!({ + "type": "image_url", + "image_url": { "url": image_url } + }), + json!({ + "type": "text", + "text": text_parts.join("\n\n") + }), + ]; + + let body = json!({ + "model": model, + "stream": false, + "messages": [ + { "role": "system", "content": VISION_SYSTEM_PROMPT }, + { "role": "user", "content": user_content } + ], + "max_tokens": DEFAULT_VISION_MAX_TOKENS, + }); + + debug!("Vision pre-processing: calling model {}", model); + + let headers = HeaderMap::new(); + let billing_context = BillingContext::new(auth_method, model.to_string()); + let completion = get_chat_completion_response(state, user, body, &headers, billing_context) + .await + .map_err(|e| { + error!("Vision model call failed: {e:?}"); + e + })?; + + if completion.metadata.is_streaming { + error!("Vision pre-processing returned streaming response unexpectedly"); + return Err(ApiError::InternalServerError); + } + + let mut rx = completion.stream; + let response_json = match rx.recv().await { + Some(CompletionChunk::FullResponse(v)) => v, + Some(CompletionChunk::Error(msg)) => { + error!("Vision model error: {}", msg); + return Err(ApiError::InternalServerError); + } + other => { + error!("Unexpected vision completion chunk: {:?}", other); + return Err(ApiError::InternalServerError); + } + }; + + Ok(extract_assistant_content(&response_json) + .unwrap_or_else(|| "[Could not describe image]".to_string())) +} + +fn extract_assistant_content(response_json: &Value) -> Option { + response_json + .get("choices") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.first()) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) +} diff --git a/src/web/mod.rs b/src/web/mod.rs index 5e7eff65..24c19538 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -1,3 +1,4 @@ +pub mod agent; pub mod attestation_routes; pub mod audio_utils; pub mod encryption_middleware; @@ -11,6 +12,7 @@ pub mod protected_routes; pub mod rag; pub mod responses; +pub use agent::router as agent_routes; pub use health_routes::router_with_state as health_routes_with_state; pub use login_routes::router as login_routes; pub use oauth_routes::router as oauth_routes; diff --git a/src/web/responses/conversations.rs b/src/web/responses/conversations.rs index a06f0770..25348010 100644 --- a/src/web/responses/conversations.rs +++ b/src/web/responses/conversations.rs @@ -4,6 +4,7 @@ use crate::{ encrypt::{decrypt_content, encrypt_with_key}, jwt::AuthContext, + models::agents::Agent, models::responses::{ConversationProjectFilter, NewConversation}, models::users::User, web::{ @@ -28,7 +29,7 @@ use axum::{ use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::{collections::HashMap, sync::Arc}; -use tracing::error; +use tracing::{debug, error, trace}; use uuid::Uuid; // ============================================================================ @@ -71,6 +72,21 @@ impl ConversationContext { .get_conversation_by_uuid_and_user(conversation_id, user.uuid) .map_err(error_mapping::map_conversation_error)?; + // Hide the main agent conversation from the public Conversations API. + let agent_conversation_id = { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + main_agent_conversation_id(&mut conn, user.uuid)? + }; + + if Some(conversation.id) == agent_conversation_id { + return Err(ApiError::NotFound); + } + // Get user's encryption key let user_key = state .get_user_key(user, auth_context, None, None) @@ -96,6 +112,15 @@ impl ConversationContext { } } +fn main_agent_conversation_id( + conn: &mut diesel::PgConnection, + user_uuid: Uuid, +) -> Result, ApiError> { + Agent::get_main_for_user(conn, user_uuid) + .map_err(|_| ApiError::InternalServerError) + .map(|agent| agent.map(|a| a.conversation_id)) +} + // ============================================================================ // Request/Response Types // ============================================================================ @@ -687,14 +712,24 @@ async fn list_conversations( }; let project_filter = resolve_conversation_project_filter(&state, user.uuid, ¶ms)?; + let agent_conversation_id = { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + main_agent_conversation_id(&mut conn, user.uuid)? + }; // Fetch conversations with database-level pagination - // We fetch limit + 1 to check if there are more results + // We fetch extra rows to check if there are more results after filtering the hidden agent + // conversation. let conversations = state .db .list_conversations( user.uuid, - limit + 1, + limit + 2, params.after, ¶ms.order, project_filter, @@ -702,14 +737,27 @@ async fn list_conversations( ) .map_err(error_mapping::map_generic_db_error)?; + trace!( + "Retrieved {} conversations for user {} (limit={}, after={:?}, order={})", + conversations.len(), + user.uuid, + limit, + params.after, + params.order + ); + + let filtered: Vec<_> = conversations + .into_iter() + .filter(|c| Some(c.id) != agent_conversation_id) + .collect(); + // Check if there are more results - let has_more = conversations.len() > limit as usize; + let has_more = filtered.len() > limit as usize; - // Take only the requested limit - let conversations_to_return = if has_more { - &conversations[..limit as usize] + let conversations_to_return: Vec<_> = if has_more { + filtered.into_iter().take(limit as usize).collect() } else { - &conversations[..] + filtered }; // Get user's encryption key for decrypting metadata @@ -722,7 +770,9 @@ async fn list_conversations( // Convert to response format let mut data = Vec::with_capacity(conversations_to_return.len()); - for conv in conversations_to_return { + for conv in &conversations_to_return { + trace!("Raw conversation object: {:?}", conv); + trace!("Conv metadata_enc present: {}", conv.metadata_enc.is_some()); data.push(build_conversation_response( &state, user.uuid, @@ -734,7 +784,7 @@ async fn list_conversations( } // Extract cursor IDs - let (first_id, last_id) = Paginator::get_cursor_ids(conversations_to_return, |c| c.uuid); + let (first_id, last_id) = Paginator::get_cursor_ids(&conversations_to_return, |c| c.uuid); let response = ConversationListResponse { object: OBJECT_TYPE_LIST, @@ -753,10 +803,33 @@ async fn delete_all_conversations( Extension(session_id): Extension, Extension(user): Extension, ) -> Result>, ApiError> { - state + debug!("Deleting all conversations for user {}", user.uuid); + + let mut conn = state .db - .delete_all_conversations(user.uuid) - .map_err(error_mapping::map_generic_db_error)?; + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let agent_conversation_id = main_agent_conversation_id(&mut conn, user.uuid)?; + + if let Some(agent_conversation_id) = agent_conversation_id { + use crate::models::schema::conversations::dsl::*; + use diesel::prelude::*; + + diesel::delete( + conversations + .filter(user_id.eq(user.uuid)) + .filter(id.ne(agent_conversation_id)), + ) + .execute(&mut conn) + .map_err(|_| ApiError::InternalServerError)?; + } else { + state + .db + .delete_all_conversations(user.uuid) + .map_err(error_mapping::map_generic_db_error)?; + } let response = json!({ "object": OBJECT_TYPE_LIST_DELETED, @@ -783,6 +856,16 @@ async fn batch_delete_conversations( let mut results = Vec::with_capacity(body.ids.len()); + let agent_conversation_id = { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + main_agent_conversation_id(&mut conn, user.uuid)? + }; + for conversation_id in body.ids { // Try to get the conversation first to verify ownership match state @@ -790,6 +873,16 @@ async fn batch_delete_conversations( .get_conversation_by_uuid_and_user(conversation_id, user.uuid) { Ok(conversation) => { + if Some(conversation.id) == agent_conversation_id { + results.push(BatchDeleteItemResult { + id: conversation_id, + object: constants::OBJECT_TYPE_CONVERSATION_DELETED, + deleted: false, + error: Some("not_found"), + }); + continue; + } + // Delete the conversation match state.db.delete_conversation(conversation.id, user.uuid) { Ok(()) => { diff --git a/src/web/responses/handlers.rs b/src/web/responses/handlers.rs index a929b9f1..66511869 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -1900,6 +1900,7 @@ async fn persist_request_data( response_id: Some(response.id), user_id: user.uuid, content_enc: prepared.content_enc.clone(), + attachment_text_enc: None, prompt_tokens: prepared.user_message_tokens, }; let user_message = state From 2f79d4b4b3aa6220ef5c144d08e39532162c057d Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:30:03 -0500 Subject: [PATCH 4/7] feat(agent): add history, Brave search, and push delivery Expand Maple with paginated history and reset controls, Brave-backed web search, and encrypted mobile push delivery so the runtime can surface richer results and notify clients off-thread. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> fix(push): avoid FCM and Tinfoil router 0 bind conflict Move FCM onto its own loopback IP so the push and Tinfoil router forwarders can both bind cleanly on agent-rewrite. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 169 +++ Cargo.toml | 3 +- docs/nitro-deploy.md | 97 ++ entrypoint.sh | 38 + .../down.sql | 1 + .../up.sql | 13 +- src/brave.rs | 759 +++++++++++- src/db.rs | 40 +- src/main.rs | 39 +- src/models/agents.rs | 71 ++ src/models/mod.rs | 3 + src/models/notification_deliveries.rs | 280 +++++ src/models/notification_events.rs | 96 ++ src/models/project_settings.rs | 62 + src/models/push_devices.rs | 244 ++++ src/models/schema.rs | 7 +- src/push/apns.rs | 272 +++++ src/push/binding.rs | 326 ++++++ src/push/crypto.rs | 138 +++ src/push/fcm.rs | 305 +++++ src/push/mod.rs | 435 +++++++ src/push/worker.rs | 556 +++++++++ src/web/agent/mod.rs | 1032 +++++++++++++---- src/web/agent/runtime.rs | 20 +- src/web/agent/signatures.rs | 17 +- src/web/agent/tools.rs | 48 + src/web/login_routes.rs | 112 +- src/web/mod.rs | 2 + src/web/platform/common.rs | 16 +- src/web/platform/mod.rs | 3 +- src/web/platform/project_routes.rs | 136 ++- src/web/push.rs | 513 ++++++++ 32 files changed, 5544 insertions(+), 309 deletions(-) create mode 100644 src/models/notification_deliveries.rs create mode 100644 src/models/notification_events.rs create mode 100644 src/models/push_devices.rs create mode 100644 src/push/apns.rs create mode 100644 src/push/binding.rs create mode 100644 src/push/crypto.rs create mode 100644 src/push/fcm.rs create mode 100644 src/push/mod.rs create mode 100644 src/push/worker.rs create mode 100644 src/web/push.rs diff --git a/Cargo.lock b/Cargo.lock index f323f2b0..83f02443 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1068,6 +1068,12 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + [[package]] name = "base58ck" version = "0.1.0" @@ -1597,6 +1603,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32b13ea120a812beba79e34316b3942a857c86ec1593cb34f27bb28272ce2cca" +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "const-random" version = "0.1.18" @@ -1712,6 +1724,18 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -1912,6 +1936,17 @@ dependencies = [ "generic-array", ] +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + [[package]] name = "der-parser" version = "8.2.0" @@ -1997,6 +2032,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", + "const-oid", "crypto-common", "subtle", ] @@ -2119,6 +2155,20 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + [[package]] name = "ecow" version = "0.2.6" @@ -2134,6 +2184,27 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "encode_unicode" version = "1.0.0" @@ -2332,6 +2403,16 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "fiat-crypto" version = "0.2.9" @@ -2644,6 +2725,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -2721,6 +2803,17 @@ dependencies = [ "spinning_top", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "h2" version = "0.3.27" @@ -4072,6 +4165,7 @@ dependencies = [ "oauth2", "once_cell", "openssl", + "p256", "password-auth", "rand_core 0.6.4", "rcgen", @@ -4175,6 +4269,18 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "panic-message" version = "0.3.0" @@ -4282,6 +4388,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -4330,6 +4445,16 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.33" @@ -4419,6 +4544,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -4798,6 +4932,16 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "rig-core" version = "0.26.0" @@ -5032,6 +5176,20 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "secp256k1" version = "0.29.1" @@ -5270,6 +5428,7 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ + "digest", "rand_core 0.6.4", ] @@ -5370,6 +5529,16 @@ dependencies = [ "lock_api", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index 226471fd..826f26e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,13 +44,14 @@ cbc = "0.1.2" secp256k1 = { version = "0.29.0", features = ["rand"] } hyper = { version = "0.14", features = ["full"] } hyper-tls = "0.5.0" -reqwest = { version = "0.11", features = ["json"] } +reqwest = { version = "0.11", features = ["json", "native-tls-alpn"] } futures = "0.3.30" uuid = { version = "1.10.0", features = ["v4", "serde"] } tokio-stream = "0.1" async-stream = "0.3" bytes = "1.0" sha2 = { version = "0.10", default-features = false } +p256 = { version = "0.13", features = ["ecdh", "pkcs8"] } hex = "0.4.3" base64 = "0.22.1" vsock = "0.5.1" diff --git a/docs/nitro-deploy.md b/docs/nitro-deploy.md index 19598673..e1154783 100644 --- a/docs/nitro-deploy.md +++ b/docs/nitro-deploy.md @@ -586,6 +586,103 @@ A restart of this should not be needed but if you need to: sudo systemctl restart vsock-apple-proxy.service ``` +## Vsock APNs and FCM proxies +Create vsock proxy services so that enclave program can talk to push providers: + +First configure the endpoints into their allowlist: + +``` +sudo vim /etc/nitro_enclaves/vsock-proxy.yaml +``` + +Add these lines: +``` +- {address: api.push.apple.com, port: 443} +- {address: api.sandbox.push.apple.com, port: 443} +- {address: fcm.googleapis.com, port: 443} +``` + +#### APNs Production +Now create a service that spins this up automatically: + +``` +sudo vim /etc/systemd/system/vsock-apns-prod-proxy.service +``` + +``` +[Unit] +Description=Vsock APNs Production Proxy Service +After=network.target + +[Service] +User=root +ExecStart=/usr/bin/vsock-proxy 8024 api.push.apple.com 443 +Restart=always + +[Install] +WantedBy=multi-user.target +``` + +#### APNs Sandbox +``` +sudo vim /etc/systemd/system/vsock-apns-sandbox-proxy.service +``` + +``` +[Unit] +Description=Vsock APNs Sandbox Proxy Service +After=network.target + +[Service] +User=root +ExecStart=/usr/bin/vsock-proxy 8025 api.sandbox.push.apple.com 443 +Restart=always + +[Install] +WantedBy=multi-user.target +``` + +#### FCM +``` +sudo vim /etc/systemd/system/vsock-fcm-proxy.service +``` + +``` +[Unit] +Description=Vsock FCM Proxy Service +After=network.target + +[Service] +User=root +ExecStart=/usr/bin/vsock-proxy 8029 fcm.googleapis.com 443 +Restart=always + +[Install] +WantedBy=multi-user.target +``` + +Activate services: + +``` +sudo systemctl daemon-reload +sudo systemctl enable vsock-apns-prod-proxy.service +sudo systemctl start vsock-apns-prod-proxy.service +sudo systemctl status vsock-apns-prod-proxy.service +sudo systemctl enable vsock-apns-sandbox-proxy.service +sudo systemctl start vsock-apns-sandbox-proxy.service +sudo systemctl status vsock-apns-sandbox-proxy.service +sudo systemctl enable vsock-fcm-proxy.service +sudo systemctl start vsock-fcm-proxy.service +sudo systemctl status vsock-fcm-proxy.service +``` + +A restart of these should not be needed but if you need to: +``` +sudo systemctl restart vsock-apns-prod-proxy.service +sudo systemctl restart vsock-apns-sandbox-proxy.service +sudo systemctl restart vsock-fcm-proxy.service +``` + ## Vsock Resend proxy Create a vsock proxy service so that enclave program can talk to resend: diff --git a/entrypoint.sh b/entrypoint.sh index 44dfd249..c0313aa2 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -325,6 +325,12 @@ log "Added Google OAuth domains to /etc/hosts" echo "127.0.0.15 appleid.apple.com" >> /etc/hosts log "Added Apple OAuth domain to /etc/hosts" +# Add push provider hostnames to /etc/hosts +echo "127.0.0.21 api.push.apple.com" >> /etc/hosts +echo "127.0.0.22 api.sandbox.push.apple.com" >> /etc/hosts +echo "127.0.0.20 fcm.googleapis.com" >> /etc/hosts +log "Added APNs and FCM domains to /etc/hosts" + # Add AWS SQS hostname to /etc/hosts echo "127.0.0.13 sqs.us-east-2.amazonaws.com" >> /etc/hosts log "Added AWS SQS domain to /etc/hosts" @@ -459,6 +465,16 @@ run_forever tf_os_flags python3 /app/traffic_forwarder.py 127.0.0.18 443 3 8028 log "Starting Apple OAuth traffic forwarder" run_forever tf_apple_oauth python3 /app/traffic_forwarder.py 127.0.0.15 443 3 8018 & +# Start the traffic forwarders for push providers in the background +log "Starting APNs production traffic forwarder" +run_forever tf_apns_prod python3 /app/traffic_forwarder.py 127.0.0.21 443 3 8024 & + +log "Starting APNs sandbox traffic forwarder" +run_forever tf_apns_sandbox python3 /app/traffic_forwarder.py 127.0.0.22 443 3 8025 & + +log "Starting FCM traffic forwarder" +run_forever tf_fcm python3 /app/traffic_forwarder.py 127.0.0.20 443 3 8029 & + # Start the traffic forwarders for Tinfoil proxy in the background log "Starting Tinfoil GitHub proxy traffic forwarder" run_forever tf_tinfoil_github_proxy python3 /app/traffic_forwarder.py 127.0.0.16 443 3 8019 & @@ -657,6 +673,28 @@ else log "Apple OAuth connection failed" fi +# Test the connections to push providers +log "Testing connection to APNs production:" +if timeout 5 bash -c ', + pub freshness: Option, + pub location: Option, + pub timezone: Option, +} + #[derive(Clone)] pub struct BraveClient { client: reqwest::Client, @@ -27,7 +30,6 @@ pub struct BraveClient { } impl BraveClient { - /// Create a new Brave client with the given API key pub fn new(api_key: String) -> Result { let client = reqwest::Client::builder() .timeout(REQUEST_TIMEOUT) @@ -43,7 +45,6 @@ impl BraveClient { }) } - /// Execute a web search query pub async fn search(&self, request: SearchRequest) -> Result { let url = format!("{}/web/search", BRAVE_API_BASE); @@ -56,7 +57,7 @@ impl BraveClient { query_params.push(("search_lang", search_lang.clone())); } if let Some(count) = request.count { - query_params.push(("count", count.to_string())); + query_params.push(("count", count.min(20).to_string())); } if let Some(offset) = request.offset { query_params.push(("offset", offset.to_string())); @@ -70,18 +71,48 @@ impl BraveClient { if let Some(summary) = request.summary { query_params.push(("summary", if summary { "1" } else { "0" }.to_string())); } + if let Some(extra_snippets) = request.extra_snippets { + query_params.push(( + "extra_snippets", + if extra_snippets { "true" } else { "false" }.to_string(), + )); + } + if let Some(enable_rich_callback) = request.enable_rich_callback { + query_params.push(( + "enable_rich_callback", + if enable_rich_callback { "1" } else { "0" }.to_string(), + )); + } + if let Some(spellcheck) = request.spellcheck { + query_params.push(( + "spellcheck", + if spellcheck { "true" } else { "false" }.to_string(), + )); + } - let response = self + let mut request_builder = self .client .get(&url) .header("X-Subscription-Token", self.api_key.as_str()) - .header("Accept", "application/json") - .query(&query_params) - .send() - .await?; + .header("Accept", "application/json"); - let status = response.status(); + if let Some(timezone) = &request.timezone { + request_builder = request_builder.header("x-loc-timezone", timezone); + } + + if let Some(location) = &request.location { + let parts: Vec<&str> = location.split(',').map(|s| s.trim()).collect(); + if let Some(city) = parts.first().filter(|s| !s.is_empty()) { + request_builder = request_builder.header("x-loc-city", *city); + } + if let Some(state_name) = parts.get(1).filter(|s| !s.is_empty()) { + request_builder = request_builder.header("x-loc-state-name", *state_name); + } + } + let response = request_builder.query(&query_params).send().await?; + + let status = response.status(); if !status.is_success() { let error_text = response.text().await.unwrap_or_default(); return Err(BraveError::Api { @@ -90,11 +121,68 @@ impl BraveClient { }); } - let search_response = response.json::().await?; + response + .json::() + .await + .map_err(BraveError::Request) + } + + pub async fn search_with_options( + &self, + query: &str, + options: Option, + ) -> Result { + let options = options.unwrap_or_default(); + let request = SearchRequest { + query: query.to_string(), + country: None, + search_lang: None, + count: options.count.or(Some(10)), + offset: None, + safesearch: None, + freshness: options.freshness, + summary: Some(true), + extra_snippets: Some(true), + enable_rich_callback: Some(true), + spellcheck: Some(true), + location: options.location, + timezone: options.timezone, + }; + + let mut search_response = self.search(request).await?; + + let summarizer_key = search_response.summarizer.as_ref().map(|s| s.key.clone()); + if let Some(key) = summarizer_key { + debug!("Fetching Brave AI summary"); + match self.summarizer(&key).await { + Ok(summary_response) => { + search_response.summary_text = summary_response.extract_text(); + } + Err(err) => { + warn!("Failed to fetch Brave summary: {err}"); + } + } + } + + let rich_callback = search_response + .rich + .as_ref() + .map(|r| (r.hint.vertical.clone(), r.hint.callback_key.clone())); + if let Some((vertical, callback_key)) = rich_callback { + debug!("Fetching Brave rich data for vertical: {vertical}"); + match self.fetch_rich(&callback_key).await { + Ok(rich_response) => { + search_response.rich_data = Some(rich_response); + } + Err(err) => { + warn!("Failed to fetch Brave rich data: {err}"); + } + } + } + Ok(search_response) } - /// Execute a summarizer query using a key from web search pub async fn summarizer(&self, key: &str) -> Result { let url = format!("{}/summarizer/search", BRAVE_API_BASE); @@ -108,7 +196,33 @@ impl BraveClient { .await?; let status = response.status(); + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(BraveError::Api { + status: status.as_u16(), + message: error_text, + }); + } + response + .json::() + .await + .map_err(BraveError::Request) + } + + async fn fetch_rich(&self, callback_key: &str) -> Result { + let url = format!("{}/web/rich", BRAVE_API_BASE); + + let response = self + .client + .get(&url) + .header("X-Subscription-Token", self.api_key.as_str()) + .header("Accept", "application/json") + .query(&[("callback_key", callback_key)]) + .send() + .await?; + + let status = response.status(); if !status.is_success() { let error_text = response.text().await.unwrap_or_default(); return Err(BraveError::Api { @@ -117,8 +231,10 @@ impl BraveClient { }); } - let summarizer_response = response.json::().await?; - Ok(summarizer_response) + response + .json::() + .await + .map_err(BraveError::Request) } } @@ -147,6 +263,16 @@ pub struct SearchRequest { pub freshness: Option, #[serde(skip_serializing_if = "Option::is_none")] pub summary: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extra_snippets: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub enable_rich_callback: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub spellcheck: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub location: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub timezone: Option, } impl SearchRequest { @@ -160,64 +286,82 @@ impl SearchRequest { safesearch: None, freshness: None, summary: None, + extra_snippets: None, + enable_rich_callback: None, + spellcheck: None, + location: None, + timezone: None, } } } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] #[allow(dead_code)] pub struct SearchResponse { #[serde(rename = "type")] - pub response_type: String, - #[serde(skip_serializing_if = "Option::is_none")] + pub response_type: Option, pub query: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub web: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub news: Option, - #[serde(skip_serializing_if = "Option::is_none")] + pub faq: Option, + pub discussions: Option, pub infobox: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub summarizer: Option, + pub rich: Option, + #[serde(skip)] + pub summary_text: Option, + #[serde(skip)] + pub rich_data: Option, } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] #[allow(dead_code)] pub struct QueryInfo { - #[serde(skip_serializing_if = "Option::is_none")] pub original: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub altered: Option, } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] #[allow(dead_code)] pub struct WebResults { #[serde(rename = "type")] - pub results_type: String, - #[serde(skip_serializing_if = "Option::is_none")] + pub results_type: Option, pub results: Option>, } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] #[allow(dead_code)] pub struct NewsResults { #[serde(rename = "type")] - pub results_type: String, - #[serde(skip_serializing_if = "Option::is_none")] + pub results_type: Option, pub results: Option>, } +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] +#[allow(dead_code)] +pub struct FaqResults { + pub results: Option>, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] +#[allow(dead_code)] +pub struct DiscussionResults { + pub results: Option>, +} + #[derive(Debug, Clone, Deserialize)] #[allow(dead_code)] pub struct SearchResult { pub title: String, pub url: String, - #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub age: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub extra_snippets: Option>, } @@ -226,22 +370,35 @@ pub struct SearchResult { pub struct NewsResult { pub title: String, pub url: String, - #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub age: Option, } #[derive(Debug, Clone, Deserialize)] #[allow(dead_code)] +pub struct FaqResult { + pub question: String, + pub answer: String, + pub title: Option, + pub url: Option, +} + +#[derive(Debug, Clone, Deserialize)] +#[allow(dead_code)] +pub struct DiscussionResult { + pub title: String, + pub url: String, + pub description: Option, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] +#[allow(dead_code)] pub struct Infobox { #[serde(rename = "type")] - pub infobox_type: String, - #[serde(skip_serializing_if = "Option::is_none")] + pub infobox_type: Option, pub title: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub long_desc: Option, } @@ -249,16 +406,30 @@ pub struct Infobox { #[allow(dead_code)] pub struct Summarizer { #[serde(rename = "type")] - pub summarizer_type: String, + pub summarizer_type: Option, pub key: String, } #[derive(Debug, Clone, Deserialize)] #[allow(dead_code)] +pub struct RichHint { + #[serde(rename = "type")] + pub rich_type: Option, + pub hint: RichHintDetails, +} + +#[derive(Debug, Clone, Deserialize)] +#[allow(dead_code)] +pub struct RichHintDetails { + pub vertical: String, + pub callback_key: String, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] +#[allow(dead_code)] pub struct SummarizerSearchResponse { - #[serde(skip_serializing_if = "Option::is_none")] pub status: Option, - #[serde(skip_serializing_if = "Option::is_none")] pub summary: Option>, } @@ -267,8 +438,496 @@ pub struct SummarizerSearchResponse { pub struct SummaryItem { #[serde(rename = "type")] pub item_type: String, - /// For type="token", data is a string with the text - /// For type="enum_item", data is a SummaryEntity object - #[serde(skip_serializing_if = "Option::is_none")] pub data: Option, } + +impl SummarizerSearchResponse { + pub fn extract_text(&self) -> Option { + let items = self.summary.as_ref()?; + let mut text = String::new(); + + for item in items { + if item.item_type == "token" { + if let Some(data) = &item.data { + if let Some(value) = data.as_str() { + text.push_str(value); + } + } + } + } + + if text.is_empty() { + None + } else { + Some(text) + } + } +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(default)] +#[allow(dead_code)] +pub struct RichResponse { + #[serde(rename = "type")] + pub response_type: Option, + pub results: Option>, +} + +#[derive(Debug, Clone, Deserialize)] +#[allow(dead_code)] +pub struct RichResult { + #[serde(rename = "type")] + pub result_type: Option, + pub subtype: Option, + #[serde(flatten)] + pub data: serde_json::Value, +} + +impl RichResponse { + pub fn format(&self) -> Option { + let results = self.results.as_ref()?; + let first = results.first()?; + first.format() + } +} + +impl RichResult { + pub fn format(&self) -> Option { + match self.subtype.as_deref()? { + "weather" => self.format_weather(), + "stock" => self.format_stock(), + "currency" => self.format_currency(), + "cryptocurrency" => self.format_crypto(), + "calculator" => self.format_calculator(), + "unit_conversion" => self.format_unit_conversion(), + "definitions" => self.format_definition(), + subtype => Some(format!( + "{}\n{}", + subtype, + serde_json::to_string_pretty(&self.data).unwrap_or_default() + )), + } + } + + fn format_weather(&self) -> Option { + let mut output = String::new(); + + if let Some(weather) = self.data.get("weather") { + if let Some(location) = weather.get("location") { + let name = location + .get("name") + .and_then(|value| value.as_str()) + .unwrap_or("Unknown"); + let state = location + .get("state") + .and_then(|value| value.as_str()) + .unwrap_or(""); + output.push_str(&format!("Weather for {}, {}\n\n", name, state)); + } + + if let Some(current) = weather.get("current_weather") { + output.push_str("Current Conditions:\n"); + if let Some(temp_c) = current.get("temp").and_then(|value| value.as_f64()) { + output.push_str(&format!( + " Temperature: {:.0}Β°F\n", + temp_c * 9.0 / 5.0 + 32.0 + )); + } + if let Some(feels_c) = current.get("feels_like").and_then(|value| value.as_f64()) { + output.push_str(&format!( + " Feels like: {:.0}Β°F\n", + feels_c * 9.0 / 5.0 + 32.0 + )); + } + if let Some(desc) = current + .get("weather") + .and_then(|value| value.get("description")) + .and_then(|value| value.as_str()) + { + output.push_str(&format!(" Conditions: {}\n", desc)); + } + if let Some(humidity) = current.get("humidity") { + output.push_str(&format!(" Humidity: {}%\n", humidity)); + } + if let Some(wind_ms) = current + .get("wind") + .and_then(|value| value.get("speed")) + .and_then(|value| value.as_f64()) + { + output.push_str(&format!(" Wind: {:.0} mph\n", wind_ms * 2.237)); + } + output.push('\n'); + } + + if let Some(alerts) = weather.get("alerts").and_then(|value| value.as_array()) { + if !alerts.is_empty() { + output.push_str("Weather Alerts:\n"); + for alert in alerts.iter().take(3) { + if let Some(event) = alert.get("event").and_then(|value| value.as_str()) { + output.push_str(&format!(" - {}\n", event)); + if let Some(desc) = + alert.get("description").and_then(|value| value.as_str()) + { + let short_desc: String = desc.chars().take(200).collect(); + output.push_str(&format!( + " {}{}\n", + short_desc, + if desc.len() > 200 { "..." } else { "" } + )); + } + } + } + output.push('\n'); + } + } + + if let Some(daily) = weather.get("daily").and_then(|value| value.as_array()) { + output.push_str("Forecast:\n"); + for (idx, day) in daily.iter().take(5).enumerate() { + let day_name = day + .get("date_i18n") + .and_then(|value| value.as_str()) + .map(|value| value.to_string()) + .unwrap_or_else(|| match idx { + 0 => "Today".to_string(), + 1 => "Tomorrow".to_string(), + _ => format!("Day {}", idx + 1), + }); + + let high = day + .get("temperature") + .and_then(|value| value.get("max")) + .and_then(|value| value.as_f64()) + .map(|value| format!("{:.0}Β°F", value * 9.0 / 5.0 + 32.0)) + .unwrap_or_default(); + let low = day + .get("temperature") + .and_then(|value| value.get("min")) + .and_then(|value| value.as_f64()) + .map(|value| format!("{:.0}Β°F", value * 9.0 / 5.0 + 32.0)) + .unwrap_or_default(); + let desc = day + .get("weather") + .and_then(|value| value.get("description")) + .and_then(|value| value.as_str()) + .unwrap_or(""); + + output.push_str(&format!( + " {} - High: {}, Low: {} - {}\n", + day_name, high, low, desc + )); + } + } + } else { + output.push_str(&serde_json::to_string_pretty(&self.data).unwrap_or_default()); + } + + if output.is_empty() { + None + } else { + Some(output) + } + } + + fn format_stock(&self) -> Option { + let mut output = String::from("Stock:\n\n"); + + if let Some(symbol) = self.data.get("symbol").and_then(|value| value.as_str()) { + output.push_str(&format!("Symbol: {}\n", symbol)); + } + if let Some(name) = self.data.get("name").and_then(|value| value.as_str()) { + output.push_str(&format!("Name: {}\n", name)); + } + if let Some(price) = self.data.get("price") { + output.push_str(&format!("Price: ${}\n", price)); + } + if let Some(change) = self.data.get("change") { + output.push_str(&format!("Change: {}\n", change)); + } + if let Some(change_pct) = self.data.get("change_percent") { + output.push_str(&format!("Change %: {}%\n", change_pct)); + } + + Some(output) + } + + fn format_currency(&self) -> Option { + Some(format!( + "Currency Conversion:\n\n{}", + serde_json::to_string_pretty(&self.data).unwrap_or_default() + )) + } + + fn format_crypto(&self) -> Option { + let mut output = String::from("Cryptocurrency:\n\n"); + + if let Some(name) = self.data.get("name").and_then(|value| value.as_str()) { + output.push_str(&format!("Name: {}\n", name)); + } + if let Some(symbol) = self.data.get("symbol").and_then(|value| value.as_str()) { + output.push_str(&format!("Symbol: {}\n", symbol)); + } + if let Some(price) = self.data.get("price") { + output.push_str(&format!("Price: ${}\n", price)); + } + if let Some(change) = self.data.get("change_24h") { + output.push_str(&format!("24h Change: {}%\n", change)); + } + + Some(output) + } + + fn format_calculator(&self) -> Option { + self.data.get("result").map(|result| match result.as_str() { + Some(value) => format!("Calculator: {}", value), + None => format!("Calculator: {}", result), + }) + } + + fn format_unit_conversion(&self) -> Option { + let result = self.data.get("result").and_then(|value| value.as_str())?; + Some(format!("Unit Conversion:\n{}", result)) + } + + fn format_definition(&self) -> Option { + let mut output = String::from("Definition:\n\n"); + + if let Some(word) = self.data.get("word").and_then(|value| value.as_str()) { + output.push_str(&format!("{}\n", word)); + } + if let Some(definitions) = self + .data + .get("definitions") + .and_then(|value| value.as_array()) + { + for (idx, definition) in definitions.iter().take(3).enumerate() { + if let Some(text) = definition + .get("definition") + .and_then(|value| value.as_str()) + { + output.push_str(&format!("{}. {}\n", idx + 1, text)); + } + } + } + + Some(output) + } +} + +impl SearchResponse { + pub fn format_results(&self) -> String { + let mut output = String::new(); + + if let Some(query) = &self.query { + if let Some(altered) = &query.altered { + if query.original.as_ref() != Some(altered) { + output.push_str(&format!("Showing results for: {}\n\n", altered)); + } + } + } + + if let Some(rich) = &self.rich_data { + if let Some(formatted) = rich.format() { + output.push_str(&formatted); + output.push_str("\n\n---\n\n"); + } + } + + if let Some(summary) = &self.summary_text { + output.push_str("AI Summary:\n"); + output.push_str(summary); + output.push_str("\n\n---\n\n"); + } + + if let Some(infobox) = &self.infobox { + if let Some(title) = &infobox.title { + output.push_str(title); + output.push('\n'); + if let Some(desc) = infobox.long_desc.as_ref().or(infobox.description.as_ref()) { + output.push_str(desc); + output.push_str("\n\n"); + } + } + } + + if let Some(faq) = &self.faq { + if let Some(results) = &faq.results { + if !results.is_empty() { + output.push_str("FAQ:\n\n"); + for item in results.iter().take(3) { + output.push_str(&format!("Q: {}\nA: {}\n\n", item.question, item.answer)); + } + } + } + } + + if let Some(web) = &self.web { + if let Some(results) = &web.results { + if !results.is_empty() { + output.push_str("Search Results:\n\n"); + for (idx, result) in results.iter().take(5).enumerate() { + let age = result + .age + .as_deref() + .map(|value| format!(" ({})", value)) + .unwrap_or_default(); + output.push_str(&format!( + "{}. {}{}\n URL: {}\n {}\n", + idx + 1, + result.title, + age, + result.url, + result.description.as_deref().unwrap_or("") + )); + if let Some(extra_snippets) = &result.extra_snippets { + for snippet in extra_snippets.iter().take(2) { + output.push_str(&format!(" > {}\n", snippet)); + } + } + output.push('\n'); + } + } + } + } + + if let Some(news) = &self.news { + if let Some(results) = &news.results { + if !results.is_empty() { + output.push_str("Recent News:\n\n"); + for (idx, result) in results.iter().take(3).enumerate() { + let age = result + .age + .as_deref() + .map(|value| format!(" ({})", value)) + .unwrap_or_default(); + output.push_str(&format!( + "{}. {}{}\n URL: {}\n {}\n\n", + idx + 1, + result.title, + age, + result.url, + result.description.as_deref().unwrap_or("") + )); + } + } + } + } + + if let Some(discussions) = &self.discussions { + if let Some(results) = &discussions.results { + if !results.is_empty() { + output.push_str("Discussions:\n\n"); + for result in results.iter().take(2) { + output.push_str(&format!("- {}\n {}\n\n", result.title, result.url)); + } + } + } + } + + if output.is_empty() { + "No results found.".to_string() + } else { + output + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn summarizer_extract_text_combines_token_items() { + let response = SummarizerSearchResponse { + status: Some("complete".to_string()), + summary: Some(vec![ + SummaryItem { + item_type: "token".to_string(), + data: Some(json!("hello ")), + }, + SummaryItem { + item_type: "enum_item".to_string(), + data: Some(json!({"ignored": true})), + }, + SummaryItem { + item_type: "token".to_string(), + data: Some(json!("world")), + }, + ]), + }; + + assert_eq!(response.extract_text().as_deref(), Some("hello world")); + } + + #[test] + fn format_results_includes_rich_summary_and_result_sections() { + let response = SearchResponse { + query: Some(QueryInfo { + original: Some("weathr austin".to_string()), + altered: Some("weather austin".to_string()), + }), + web: Some(WebResults { + results_type: None, + results: Some(vec![SearchResult { + title: "Austin forecast".to_string(), + url: "https://example.com/weather".to_string(), + description: Some("Sunny all week".to_string()), + age: Some("2h".to_string()), + extra_snippets: Some(vec!["High of 75F".to_string()]), + }]), + }), + news: Some(NewsResults { + results_type: None, + results: Some(vec![NewsResult { + title: "Austin weather update".to_string(), + url: "https://example.com/news".to_string(), + description: Some("Cold front incoming".to_string()), + age: Some("1h".to_string()), + }]), + }), + faq: Some(FaqResults { + results: Some(vec![FaqResult { + question: "Will it rain?".to_string(), + answer: "No rain expected today.".to_string(), + title: None, + url: None, + }]), + }), + discussions: Some(DiscussionResults { + results: Some(vec![DiscussionResult { + title: "Austin weather thread".to_string(), + url: "https://example.com/discuss".to_string(), + description: None, + }]), + }), + infobox: Some(Infobox { + infobox_type: None, + title: Some("Austin, TX".to_string()), + description: Some("Current weather conditions".to_string()), + long_desc: None, + }), + summarizer: None, + rich: None, + summary_text: Some("Warm and sunny today.".to_string()), + rich_data: Some(RichResponse { + response_type: None, + results: Some(vec![RichResult { + result_type: Some("rich_result".to_string()), + subtype: Some("calculator".to_string()), + data: json!({"result": "72 + 3 = 75"}), + }]), + }), + response_type: None, + }; + + let formatted = response.format_results(); + + assert!(formatted.contains("Showing results for: weather austin")); + assert!(formatted.contains("Calculator: 72 + 3 = 75")); + assert!(formatted.contains("AI Summary:")); + assert!(formatted.contains("Search Results:")); + assert!(formatted.contains("Recent News:")); + assert!(formatted.contains("Discussions:")); + } +} diff --git a/src/db.rs b/src/db.rs index 6e5f38e1..6f6681f5 100644 --- a/src/db.rs +++ b/src/db.rs @@ -27,10 +27,10 @@ use crate::models::platform_password_reset::{ NewPlatformPasswordResetRequest, PlatformPasswordResetError, PlatformPasswordResetRequest, }; use crate::models::platform_users::{NewPlatformUser, PlatformUser, PlatformUserError}; -use crate::models::project_settings::OAuthSettings; use crate::models::project_settings::{ EmailSettings, NewProjectSetting, ProjectSetting, ProjectSettingError, SettingCategory, }; +use crate::models::project_settings::{OAuthSettings, PushSettings}; use crate::models::responses::{ validate_conversation_project_limit, AssistantMessage, Conversation, ConversationProject, ConversationProjectFilter, NewAssistantMessage, NewConversation, NewConversationProject, @@ -445,6 +445,14 @@ pub trait DBConnection { settings: OAuthSettings, ) -> Result; + fn get_project_push_settings(&self, project_id: i32) -> Result, DBError>; + + fn update_project_push_settings( + &self, + project_id: i32, + settings: PushSettings, + ) -> Result; + // Platform email verification methods fn create_platform_email_verification( &self, @@ -1879,6 +1887,36 @@ impl DBConnection for PostgresConnection { } } + fn get_project_push_settings(&self, project_id: i32) -> Result, DBError> { + debug!("Getting project push settings"); + let settings = self.get_project_settings(project_id, SettingCategory::Push)?; + + match settings { + Some(s) => s.get_push_settings().map(Some).map_err(DBError::from), + None => Ok(None), + } + } + + fn update_project_push_settings( + &self, + project_id: i32, + settings: PushSettings, + ) -> Result { + debug!("Updating project push settings"); + let new_settings = NewProjectSetting::new_push_settings(project_id, settings)?; + let conn = &mut self.db.get().map_err(|_| DBError::ConnectionError)?; + + if let Some(mut existing) = + ProjectSetting::get_by_project_and_category(conn, project_id, SettingCategory::Push)? + { + existing.settings = new_settings.settings; + existing.update(conn)?; + Ok(existing) + } else { + new_settings.insert(conn).map_err(DBError::from) + } + } + // Platform email verification implementations fn create_platform_email_verification( &self, diff --git a/src/main.rs b/src/main.rs index ca479b39..c0ead20d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,13 +16,14 @@ use crate::models::oauth::{NewUserOAuthConnection, OAuthError}; use crate::models::password_reset::NewPasswordResetRequest; use crate::models::platform_password_reset::NewPlatformPasswordResetRequest; use crate::models::platform_users::PlatformUser; +use crate::push::worker::start_push_worker; use crate::sqs::SqsEventPublisher; use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ agent_routes, conversation_projects_routes, conversations_routes, health_routes_with_state, - instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, - responses_routes, + instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, push_routes, + rag_routes, responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -98,6 +99,7 @@ mod os_flags; mod private_key; mod provider_routing; mod proxy_config; +mod push; mod rag; #[cfg(test)] mod security_invariants; @@ -2171,7 +2173,17 @@ impl AppState { project_id: i32, key_name: &str, ) -> Result, Error> { - // Get the encrypted secret + Ok(self + .get_project_secret_bytes(project_id, key_name) + .await? + .map(|bytes| general_purpose::STANDARD.encode(bytes))) + } + + pub async fn get_project_secret_bytes( + &self, + project_id: i32, + key_name: &str, + ) -> Result>, Error> { let secret = match self .db .get_org_project_secret_by_key_name_and_project(key_name, project_id)? @@ -2180,15 +2192,24 @@ impl AppState { None => return Ok(None), }; - // Decrypt the secret using the enclave key let secret_key = SecretKey::from_slice(&self.enclave_key) .map_err(|e| Error::EncryptionError(e.to_string()))?; let decrypted_bytes = decrypt_with_key(&secret_key, &secret.secret_enc) .map_err(|e| Error::EncryptionError(e.to_string()))?; - // Always return base64 encoded bytes - Ok(Some(general_purpose::STANDARD.encode(&decrypted_bytes))) + Ok(Some(decrypted_bytes)) + } + + pub async fn get_project_secret_string( + &self, + project_id: i32, + key_name: &str, + ) -> Result, Error> { + let secret = self.get_project_secret_bytes(project_id, key_name).await?; + secret + .map(|bytes| String::from_utf8(bytes).map_err(|_| Error::SecretParsingError)) + .transpose() } pub async fn get_project_resend_api_key( @@ -3177,6 +3198,8 @@ async fn main() -> Result<(), Error> { ) .await?; + start_push_worker(app_state.clone()); + let cors = CorsLayer::new() // allow all method types .allow_methods(Any) @@ -3217,6 +3240,10 @@ async fn main() -> Result<(), Error> { agent_routes(app_state.clone()) .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), ) + .merge( + push_routes(app_state.clone()) + .route_layer(from_fn_with_state(app_state.clone(), validate_jwt)), + ) .merge(attestation_routes::router(app_state.clone())) .merge(oauth_routes(app_state.clone())) .merge(platform_login_routes(app_state.clone())) diff --git a/src/models/agents.rs b/src/models/agents.rs index 262ac1fd..816b5f2b 100644 --- a/src/models/agents.rs +++ b/src/models/agents.rs @@ -1,4 +1,6 @@ +use crate::models::responses::Conversation; use crate::models::schema::agents; +use crate::models::schema::conversations; use chrono::{DateTime, Utc}; use diesel::prelude::*; use serde::{Deserialize, Serialize}; @@ -83,6 +85,75 @@ impl Agent { .map_err(AgentError::DatabaseError) } + pub fn list_subagents_for_user_paginated( + conn: &mut PgConnection, + lookup_user_id: Uuid, + limit: i64, + after: Option, + order: &str, + created_by_filter: Option<&str>, + ) -> Result, AgentError> { + let mut query = agents::table + .inner_join(conversations::table) + .filter(agents::user_id.eq(lookup_user_id)) + .filter(agents::kind.eq(AGENT_KIND_SUBAGENT)) + .into_boxed(); + + if let Some(created_by_filter) = created_by_filter { + query = query.filter(agents::created_by.eq(created_by_filter)); + } + + if let Some(after_uuid) = after { + let mut cursor_query = agents::table + .inner_join(conversations::table) + .filter(agents::user_id.eq(lookup_user_id)) + .filter(agents::kind.eq(AGENT_KIND_SUBAGENT)) + .filter(agents::uuid.eq(after_uuid)) + .into_boxed(); + + if let Some(created_by_filter) = created_by_filter { + cursor_query = cursor_query.filter(agents::created_by.eq(created_by_filter)); + } + + let cursor_subagent = cursor_query + .select((conversations::updated_at, agents::id)) + .first::<(DateTime, i64)>(conn) + .optional() + .map_err(AgentError::DatabaseError)?; + + if let Some((updated_at, id)) = cursor_subagent { + if order == "desc" { + query = query.filter( + conversations::updated_at + .lt(updated_at) + .or(conversations::updated_at + .eq(updated_at) + .and(agents::id.lt(id))), + ); + } else { + query = query.filter( + conversations::updated_at + .gt(updated_at) + .or(conversations::updated_at + .eq(updated_at) + .and(agents::id.gt(id))), + ); + } + } + } + + if order == "desc" { + query = query.order((conversations::updated_at.desc(), agents::id.desc())); + } else { + query = query.order((conversations::updated_at.asc(), agents::id.asc())); + } + + query + .limit(limit) + .load::<(Agent, Conversation)>(conn) + .map_err(AgentError::DatabaseError) + } + pub fn delete_by_uuid_and_user( conn: &mut PgConnection, lookup_uuid: Uuid, diff --git a/src/models/mod.rs b/src/models/mod.rs index 2836fc98..4c98fbf2 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -6,6 +6,8 @@ pub mod email_verification; pub mod enclave_secrets; pub mod invite_codes; pub mod memory_blocks; +pub mod notification_deliveries; +pub mod notification_events; pub mod oauth; pub mod org_memberships; pub mod org_project_secrets; @@ -17,6 +19,7 @@ pub mod platform_invite_codes; pub mod platform_password_reset; pub mod platform_users; pub mod project_settings; +pub mod push_devices; pub mod responses; pub mod schema; pub mod token_usage; diff --git a/src/models/notification_deliveries.rs b/src/models/notification_deliveries.rs new file mode 100644 index 00000000..47ce7838 --- /dev/null +++ b/src/models/notification_deliveries.rs @@ -0,0 +1,280 @@ +use crate::models::schema::notification_deliveries; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{BigInt, Int4, Text}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +pub const NOTIFICATION_DELIVERY_STATUS_SENT: &str = "sent"; +pub const NOTIFICATION_DELIVERY_STATUS_FAILED: &str = "failed"; +pub const NOTIFICATION_DELIVERY_STATUS_INVALID_TOKEN: &str = "invalid_token"; +pub const NOTIFICATION_DELIVERY_STATUS_CANCELLED: &str = "cancelled"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NotificationDeliveryWriteResult { + Updated, + LostLease, +} + +impl NotificationDeliveryWriteResult { + fn from_updated_rows(updated_rows: usize) -> Self { + if updated_rows == 0 { + Self::LostLease + } else { + Self::Updated + } + } + + pub fn was_applied(self) -> bool { + matches!(self, Self::Updated) + } +} + +#[derive(Error, Debug)] +pub enum NotificationDeliveryError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, QueryableByName, Clone, Debug, Serialize, Deserialize)] +#[diesel(table_name = notification_deliveries)] +pub struct NotificationDelivery { + pub id: i64, + pub event_id: i64, + pub push_device_id: i64, + pub status: String, + pub attempt_count: i32, + pub next_attempt_at: DateTime, + pub lease_owner: Option, + pub lease_expires_at: Option>, + pub provider_message_id: Option, + pub provider_status_code: Option, + pub last_error: Option, + pub sent_at: Option>, + pub invalidated_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl NotificationDelivery { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, NotificationDeliveryError> { + notification_deliveries::table + .filter(notification_deliveries::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn lease_pending( + conn: &mut PgConnection, + limit: i64, + lease_owner: &str, + lease_seconds: i32, + ) -> Result, NotificationDeliveryError> { + let query = r#" + WITH candidates AS ( + SELECT d.id + FROM notification_deliveries d + WHERE d.next_attempt_at <= NOW() + AND ( + (d.status IN ('pending', 'retry') + AND (d.lease_expires_at IS NULL OR d.lease_expires_at < NOW())) + OR (d.status = 'leased' + AND (d.lease_expires_at IS NULL OR d.lease_expires_at < NOW())) + ) + ORDER BY d.next_attempt_at ASC, d.id ASC + FOR UPDATE OF d SKIP LOCKED + LIMIT $1 + ) + UPDATE notification_deliveries d + SET status = 'leased', + lease_owner = $2, + lease_expires_at = NOW() + ($3 * INTERVAL '1 second'), + updated_at = NOW() + FROM candidates + WHERE d.id = candidates.id + RETURNING d.* + "#; + + sql_query(query) + .bind::(limit) + .bind::(lease_owner) + .bind::(lease_seconds) + .get_results::(conn) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_sent( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + provider_message_id: Option<&str>, + provider_status_code: Option, + ) -> Result { + diesel::update( + notification_deliveries::table + .filter(notification_deliveries::id.eq(lookup_id)) + .filter(notification_deliveries::status.eq("leased")) + .filter( + notification_deliveries::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + notification_deliveries::status.eq(NOTIFICATION_DELIVERY_STATUS_SENT), + notification_deliveries::attempt_count.eq(notification_deliveries::attempt_count + 1), + notification_deliveries::provider_message_id.eq(provider_message_id), + notification_deliveries::provider_status_code.eq(provider_status_code), + notification_deliveries::last_error.eq::>(None), + notification_deliveries::sent_at.eq(diesel::dsl::now), + notification_deliveries::lease_owner.eq::>(None), + notification_deliveries::lease_expires_at.eq::>>(None), + notification_deliveries::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(NotificationDeliveryWriteResult::from_updated_rows) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_retry( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + provider_status_code: Option, + last_error: Option<&str>, + retry_after_seconds: i32, + ) -> Result { + let query = r#" + UPDATE notification_deliveries + SET status = 'retry', + attempt_count = attempt_count + 1, + provider_status_code = $3, + last_error = $4, + next_attempt_at = NOW() + ($5 * INTERVAL '1 second'), + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + AND lease_owner = $2 + AND status = 'leased' + "#; + + sql_query(query) + .bind::(lookup_id) + .bind::(expected_lease_owner) + .bind::, _>(provider_status_code) + .bind::, _>(last_error) + .bind::(retry_after_seconds) + .execute(conn) + .map(NotificationDeliveryWriteResult::from_updated_rows) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_failed( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + provider_status_code: Option, + last_error: Option<&str>, + ) -> Result { + diesel::update( + notification_deliveries::table + .filter(notification_deliveries::id.eq(lookup_id)) + .filter(notification_deliveries::status.eq("leased")) + .filter( + notification_deliveries::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + notification_deliveries::status.eq(NOTIFICATION_DELIVERY_STATUS_FAILED), + notification_deliveries::attempt_count.eq(notification_deliveries::attempt_count + 1), + notification_deliveries::provider_status_code.eq(provider_status_code), + notification_deliveries::last_error.eq(last_error), + notification_deliveries::lease_owner.eq::>(None), + notification_deliveries::lease_expires_at.eq::>>(None), + notification_deliveries::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(NotificationDeliveryWriteResult::from_updated_rows) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_invalid_token( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + provider_status_code: Option, + last_error: Option<&str>, + ) -> Result { + diesel::update( + notification_deliveries::table + .filter(notification_deliveries::id.eq(lookup_id)) + .filter(notification_deliveries::status.eq("leased")) + .filter( + notification_deliveries::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + notification_deliveries::status.eq(NOTIFICATION_DELIVERY_STATUS_INVALID_TOKEN), + notification_deliveries::attempt_count.eq(notification_deliveries::attempt_count + 1), + notification_deliveries::provider_status_code.eq(provider_status_code), + notification_deliveries::last_error.eq(last_error), + notification_deliveries::invalidated_at.eq(diesel::dsl::now), + notification_deliveries::lease_owner.eq::>(None), + notification_deliveries::lease_expires_at.eq::>>(None), + notification_deliveries::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(NotificationDeliveryWriteResult::from_updated_rows) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_cancelled( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + last_error: Option<&str>, + ) -> Result { + diesel::update( + notification_deliveries::table + .filter(notification_deliveries::id.eq(lookup_id)) + .filter(notification_deliveries::status.eq("leased")) + .filter( + notification_deliveries::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + notification_deliveries::status.eq(NOTIFICATION_DELIVERY_STATUS_CANCELLED), + notification_deliveries::last_error.eq(last_error), + notification_deliveries::lease_owner.eq::>(None), + notification_deliveries::lease_expires_at.eq::>>(None), + notification_deliveries::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(NotificationDeliveryWriteResult::from_updated_rows) + .map_err(NotificationDeliveryError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = notification_deliveries)] +pub struct NewNotificationDelivery { + pub event_id: i64, + pub push_device_id: i64, + pub next_attempt_at: DateTime, +} + +impl NewNotificationDelivery { + pub fn insert_many( + conn: &mut PgConnection, + new_deliveries: &[NewNotificationDelivery], + ) -> Result { + diesel::insert_into(notification_deliveries::table) + .values(new_deliveries) + .execute(conn) + .map_err(NotificationDeliveryError::DatabaseError) + } +} diff --git a/src/models/notification_events.rs b/src/models/notification_events.rs new file mode 100644 index 00000000..e31e92ad --- /dev/null +++ b/src/models/notification_events.rs @@ -0,0 +1,96 @@ +use crate::models::schema::notification_events; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +pub const NOTIFICATION_KIND_AGENT_MESSAGE: &str = "agent.message"; +pub const NOTIFICATION_DELIVERY_MODE_GENERIC: &str = "generic"; +pub const NOTIFICATION_DELIVERY_MODE_ENCRYPTED_PREVIEW: &str = "encrypted_preview"; +pub const NOTIFICATION_PRIORITY_NORMAL: &str = "normal"; +pub const NOTIFICATION_PRIORITY_HIGH: &str = "high"; +pub const NOTIFICATION_SOURCE_REQUEST_CONTINUATION: &str = "request_continuation"; + +#[derive(Error, Debug)] +pub enum NotificationEventError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, Clone, Debug, Serialize, Deserialize)] +#[diesel(table_name = notification_events)] +pub struct NotificationEvent { + pub id: i64, + pub uuid: Uuid, + pub project_id: i32, + pub user_id: Uuid, + pub kind: String, + pub delivery_mode: String, + pub priority: String, + pub collapse_key: Option, + pub fallback_title: String, + pub fallback_body: String, + pub payload_enc: Option>, + pub not_before_at: DateTime, + pub expires_at: Option>, + pub created_at: DateTime, + pub cancelled_at: Option>, + pub source_kind: String, + pub source_request_id: Option, +} + +impl NotificationEvent { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, NotificationEventError> { + notification_events::table + .filter(notification_events::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(NotificationEventError::DatabaseError) + } + + pub fn cancel( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, NotificationEventError> { + diesel::update(notification_events::table.filter(notification_events::id.eq(lookup_id))) + .set(notification_events::cancelled_at.eq(diesel::dsl::now)) + .get_result::(conn) + .optional() + .map_err(NotificationEventError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = notification_events, treat_none_as_default_value = true)] +pub struct NewNotificationEvent { + pub uuid: Uuid, + pub project_id: i32, + pub user_id: Uuid, + pub kind: String, + pub delivery_mode: String, + pub priority: String, + pub collapse_key: Option, + pub fallback_title: String, + pub fallback_body: String, + pub payload_enc: Option>, + pub not_before_at: Option>, + pub expires_at: Option>, + pub source_kind: String, + pub source_request_id: Option, +} + +impl NewNotificationEvent { + pub fn insert( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(notification_events::table) + .values(self) + .get_result::(conn) + .map_err(NotificationEventError::DatabaseError) + } +} diff --git a/src/models/project_settings.rs b/src/models/project_settings.rs index cb8a7a43..7bd1f259 100644 --- a/src/models/project_settings.rs +++ b/src/models/project_settings.rs @@ -20,6 +20,7 @@ pub enum ProjectSettingError { pub enum SettingCategory { Email, OAuth, + Push, } impl SettingCategory { @@ -27,6 +28,7 @@ impl SettingCategory { match self { SettingCategory::Email => "email", SettingCategory::OAuth => "oauth", + SettingCategory::Push => "push", } } } @@ -74,6 +76,49 @@ pub struct OAuthSettings { pub apple_oauth_settings: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum PushEnvironment { + Dev, + #[default] + Prod, +} + +impl PushEnvironment { + pub fn as_str(&self) -> &'static str { + match self { + Self::Dev => "dev", + Self::Prod => "prod", + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct IosPushSettings { + pub enabled: bool, + pub bundle_id: String, + pub apns_environment: PushEnvironment, + pub team_id: String, + pub key_id: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AndroidPushSettings { + pub enabled: bool, + pub firebase_project_id: String, + pub package_name: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PushSettings { + #[serde(default)] + pub encrypted_preview_enabled: bool, + #[serde(default)] + pub ios: Option, + #[serde(default)] + pub android: Option, +} + #[derive(Queryable, Identifiable)] #[diesel(table_name = project_settings)] pub struct ProjectSetting { @@ -97,6 +142,11 @@ impl ProjectSetting { .map_err(ProjectSettingError::SerializationError) } + pub fn get_push_settings(&self) -> Result { + serde_json::from_value(self.settings.clone()) + .map_err(ProjectSettingError::SerializationError) + } + pub fn get_by_project_and_category( conn: &mut PgConnection, lookup_project_id: i32, @@ -159,6 +209,18 @@ impl NewProjectSetting { }) } + pub fn new_push_settings( + project_id: i32, + push_settings: PushSettings, + ) -> Result { + Ok(Self { + project_id, + category: SettingCategory::Push.as_str().to_string(), + settings: serde_json::to_value(push_settings) + .map_err(ProjectSettingError::SerializationError)?, + }) + } + pub fn insert(&self, conn: &mut PgConnection) -> Result { diesel::insert_into(project_settings::table) .values(self) diff --git a/src/models/push_devices.rs b/src/models/push_devices.rs new file mode 100644 index 00000000..c1dc9d5d --- /dev/null +++ b/src/models/push_devices.rs @@ -0,0 +1,244 @@ +use crate::models::schema::push_devices; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +pub const PUSH_PLATFORM_IOS: &str = "ios"; +pub const PUSH_PLATFORM_ANDROID: &str = "android"; +pub const PUSH_PROVIDER_APNS: &str = "apns"; +pub const PUSH_PROVIDER_FCM: &str = "fcm"; +pub const PUSH_ENV_DEV: &str = "dev"; +pub const PUSH_ENV_PROD: &str = "prod"; +pub const PUSH_KEY_ALGORITHM_P256_ECDH_V1: &str = "p256_ecdh_v1"; + +#[derive(Error, Debug)] +pub enum PushDeviceError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, Clone, Debug, Serialize, Deserialize)] +#[diesel(table_name = push_devices)] +pub struct PushDevice { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub project_id: i32, + pub installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub push_token_hash: Vec, + pub capability_enc: Vec, + pub notification_public_key_hash: Vec, + pub key_algorithm: String, + pub supports_encrypted_preview: bool, + pub supports_background_processing: bool, + pub last_seen_at: DateTime, + pub revoked_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl PushDevice { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, PushDeviceError> { + push_devices::table + .filter(push_devices::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn get_by_uuid_and_user( + conn: &mut PgConnection, + lookup_uuid: Uuid, + lookup_user_id: Uuid, + ) -> Result, PushDeviceError> { + push_devices::table + .filter(push_devices::uuid.eq(lookup_uuid)) + .filter(push_devices::user_id.eq(lookup_user_id)) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn get_by_installation_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_installation_id: Uuid, + lookup_environment: &str, + ) -> Result, PushDeviceError> { + if let Some(active_device) = push_devices::table + .filter(push_devices::user_id.eq(lookup_user_id)) + .filter(push_devices::installation_id.eq(lookup_installation_id)) + .filter(push_devices::environment.eq(lookup_environment)) + .filter(push_devices::revoked_at.is_null()) + .order(push_devices::updated_at.desc()) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError)? + { + return Ok(Some(active_device)); + } + + push_devices::table + .filter(push_devices::user_id.eq(lookup_user_id)) + .filter(push_devices::installation_id.eq(lookup_installation_id)) + .filter(push_devices::environment.eq(lookup_environment)) + .order(push_devices::updated_at.desc()) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn get_by_installation( + conn: &mut PgConnection, + lookup_installation_id: Uuid, + lookup_environment: &str, + ) -> Result, PushDeviceError> { + push_devices::table + .filter(push_devices::installation_id.eq(lookup_installation_id)) + .filter(push_devices::environment.eq(lookup_environment)) + .filter(push_devices::revoked_at.is_null()) + .order(push_devices::updated_at.desc()) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn get_active_by_token_hash( + conn: &mut PgConnection, + lookup_provider: &str, + lookup_environment: &str, + lookup_token_hash: &[u8], + ) -> Result, PushDeviceError> { + push_devices::table + .filter(push_devices::provider.eq(lookup_provider)) + .filter(push_devices::environment.eq(lookup_environment)) + .filter(push_devices::push_token_hash.eq(lookup_token_hash)) + .filter(push_devices::revoked_at.is_null()) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn list_active_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, PushDeviceError> { + push_devices::table + .filter(push_devices::user_id.eq(lookup_user_id)) + .filter(push_devices::revoked_at.is_null()) + .order((push_devices::last_seen_at.desc(), push_devices::id.desc())) + .load::(conn) + .map_err(PushDeviceError::DatabaseError) + } + + pub fn update(&self, conn: &mut PgConnection) -> Result<(), PushDeviceError> { + diesel::update(push_devices::table.filter(push_devices::id.eq(self.id))) + .set(( + push_devices::user_id.eq(self.user_id), + push_devices::project_id.eq(self.project_id), + push_devices::installation_id.eq(self.installation_id), + push_devices::platform.eq(&self.platform), + push_devices::provider.eq(&self.provider), + push_devices::environment.eq(&self.environment), + push_devices::app_id.eq(&self.app_id), + push_devices::push_token_hash.eq(&self.push_token_hash), + push_devices::capability_enc.eq(&self.capability_enc), + push_devices::notification_public_key_hash.eq(&self.notification_public_key_hash), + push_devices::key_algorithm.eq(&self.key_algorithm), + push_devices::supports_encrypted_preview.eq(self.supports_encrypted_preview), + push_devices::supports_background_processing + .eq(self.supports_background_processing), + push_devices::last_seen_at.eq(self.last_seen_at), + push_devices::revoked_at.eq(self.revoked_at), + push_devices::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(|_| ()) + .map_err(PushDeviceError::DatabaseError) + } + + pub fn revoke_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, PushDeviceError> { + diesel::update(push_devices::table.filter(push_devices::id.eq(lookup_id))) + .set(( + push_devices::revoked_at.eq(diesel::dsl::now), + push_devices::updated_at.eq(diesel::dsl::now), + )) + .get_result::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn revoke_by_uuid_and_user( + conn: &mut PgConnection, + lookup_uuid: Uuid, + lookup_user_id: Uuid, + ) -> Result, PushDeviceError> { + diesel::update( + push_devices::table + .filter(push_devices::uuid.eq(lookup_uuid)) + .filter(push_devices::user_id.eq(lookup_user_id)), + ) + .set(( + push_devices::revoked_at.eq(diesel::dsl::now), + push_devices::updated_at.eq(diesel::dsl::now), + )) + .get_result::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn invalidate( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, PushDeviceError> { + diesel::update(push_devices::table.filter(push_devices::id.eq(lookup_id))) + .set(( + push_devices::revoked_at.eq(diesel::dsl::now), + push_devices::updated_at.eq(diesel::dsl::now), + )) + .get_result::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = push_devices)] +pub struct NewPushDevice { + pub uuid: Uuid, + pub user_id: Uuid, + pub project_id: i32, + pub installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub push_token_hash: Vec, + pub capability_enc: Vec, + pub notification_public_key_hash: Vec, + pub key_algorithm: String, + pub supports_encrypted_preview: bool, + pub supports_background_processing: bool, + pub last_seen_at: DateTime, +} + +impl NewPushDevice { + pub fn insert(&self, conn: &mut PgConnection) -> Result { + diesel::insert_into(push_devices::table) + .values(self) + .get_result::(conn) + .map_err(PushDeviceError::DatabaseError) + } +} diff --git a/src/models/schema.rs b/src/models/schema.rs index 7db20479..1c5e642c 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -238,6 +238,8 @@ diesel::table! { expires_at -> Nullable, created_at -> Timestamptz, cancelled_at -> Nullable, + source_kind -> Text, + source_request_id -> Nullable, } } @@ -374,14 +376,15 @@ diesel::table! { id -> Int8, uuid -> Uuid, user_id -> Uuid, + project_id -> Int4, installation_id -> Uuid, platform -> Text, provider -> Text, environment -> Text, app_id -> Text, - push_token_enc -> Bytea, push_token_hash -> Bytea, - notification_public_key -> Bytea, + capability_enc -> Bytea, + notification_public_key_hash -> Bytea, key_algorithm -> Text, supports_encrypted_preview -> Bool, supports_background_processing -> Bool, diff --git a/src/push/apns.rs b/src/push/apns.rs new file mode 100644 index 00000000..af779b02 --- /dev/null +++ b/src/push/apns.rs @@ -0,0 +1,272 @@ +use crate::models::notification_events::NotificationEvent; +use crate::models::project_settings::IosPushSettings; +use crate::models::push_devices::PushDevice; +use crate::push::binding::PushDeviceCapabilityPlaintextV1; +use crate::push::crypto::encrypt_preview_payload; +use crate::push::{ + CachedApnsToken, NotificationPreviewPayload, PushError, PushSendOutcome, PushTransport, + APNS_JWT_CACHE_LIFETIME_MINUTES, +}; +use crate::web::platform::PROJECT_APNS_AUTH_KEY_P8; +use crate::AppState; +use chrono::{Duration, Utc}; +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use serde::Serialize; +use serde_json::{json, Value}; +use std::sync::Arc; + +#[derive(Debug, Serialize)] +struct ApnsJwtClaims { + iss: String, + iat: i64, +} + +pub struct ApnsSendRequest<'a> { + pub event: &'a NotificationEvent, + pub device: &'a PushDevice, + pub capability: &'a PushDeviceCapabilityPlaintextV1, + pub ios_settings: &'a IosPushSettings, + pub preview_payload: Option<&'a NotificationPreviewPayload>, + pub send_encrypted_preview: bool, +} + +pub async fn send_apns_notification( + state: &Arc, + transport: &PushTransport, + request: ApnsSendRequest<'_>, +) -> Result { + let auth_token = get_apns_auth_token( + state, + transport, + request.event.project_id, + request.ios_settings, + ) + .await?; + let body = build_apns_payload( + request.event, + request.device, + request.capability, + request.preview_payload, + request.send_encrypted_preview, + )?; + let endpoint = match request.ios_settings.apns_environment.as_str() { + "dev" => format!( + "https://api.sandbox.push.apple.com/3/device/{}", + request.capability.push_token + ), + _ => format!( + "https://api.push.apple.com/3/device/{}", + request.capability.push_token + ), + }; + + let mut http_request = transport + .apns_client + .post(endpoint) + .bearer_auth(auth_token) + .header("apns-topic", request.ios_settings.bundle_id.as_str()) + .header("apns-push-type", "alert") + .header( + "apns-priority", + if request.event.priority == "high" { + "10" + } else { + "5" + }, + ); + + if let Some(collapse_key) = &request.event.collapse_key { + http_request = http_request.header("apns-collapse-id", collapse_key); + } + + if let Some(expires_at) = request.event.expires_at { + http_request = http_request.header("apns-expiration", expires_at.timestamp().to_string()); + } + + let response = http_request.json(&body).send().await?; + let status_code = response.status().as_u16() as i32; + let apns_id = response + .headers() + .get("apns-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let status = response.status(); + let response_body = response.text().await.unwrap_or_default(); + let reason = parse_apns_reason(&response_body); + + if status.is_success() { + return Ok(PushSendOutcome::Sent { + provider_message_id: apns_id, + provider_status_code: Some(status_code), + }); + } + + if matches!( + reason.as_deref(), + Some("ExpiredProviderToken" | "InvalidProviderToken") + ) { + invalidate_apns_cache(transport, request.event.project_id, request.ios_settings).await; + } + + let error_message = reason.unwrap_or_else(|| response_body.clone()); + let outcome = match status.as_u16() { + 410 => PushSendOutcome::InvalidToken { + provider_status_code: Some(status_code), + error: if error_message.is_empty() { + "Unregistered".to_string() + } else { + error_message + }, + }, + 403 if matches!( + error_message.as_str(), + "ExpiredProviderToken" | "InvalidProviderToken" + ) => + { + PushSendOutcome::Retryable { + provider_status_code: Some(status_code), + error: error_message, + } + } + 429 | 500 | 503 => PushSendOutcome::Retryable { + provider_status_code: Some(status_code), + error: error_message, + }, + _ => PushSendOutcome::Failed { + provider_status_code: Some(status_code), + error: error_message, + }, + }; + + Ok(outcome) +} + +async fn get_apns_auth_token( + state: &Arc, + transport: &PushTransport, + project_id: i32, + ios_settings: &IosPushSettings, +) -> Result { + let cache_key = format!( + "{}:{}:{}", + project_id, ios_settings.team_id, ios_settings.key_id + ); + { + let cache = transport.apns_tokens.read().await; + if let Some(cached) = cache.get(&cache_key) { + if cached.expires_at > Utc::now() + Duration::minutes(5) { + return Ok(cached.token.clone()); + } + } + } + + let private_key_pem = state + .get_project_secret_string(project_id, PROJECT_APNS_AUTH_KEY_P8) + .await? + .ok_or_else(|| PushError::InvalidSecret("missing APNs auth key".to_string()))?; + + let mut header = Header::new(Algorithm::ES256); + header.kid = Some(ios_settings.key_id.clone()); + + let claims = ApnsJwtClaims { + iss: ios_settings.team_id.clone(), + iat: Utc::now().timestamp() - 30, + }; + let signing_key = EncodingKey::from_ec_pem(private_key_pem.as_bytes()) + .map_err(|e| PushError::InvalidSecret(e.to_string()))?; + let token = encode(&header, &claims, &signing_key)?; + let expires_at = Utc::now() + Duration::minutes(APNS_JWT_CACHE_LIFETIME_MINUTES); + + let mut cache = transport.apns_tokens.write().await; + cache.insert( + cache_key, + CachedApnsToken { + token: token.clone(), + expires_at, + }, + ); + + Ok(token) +} + +async fn invalidate_apns_cache( + transport: &PushTransport, + project_id: i32, + ios_settings: &IosPushSettings, +) { + let cache_key = format!( + "{}:{}:{}", + project_id, ios_settings.team_id, ios_settings.key_id + ); + let mut cache = transport.apns_tokens.write().await; + cache.remove(&cache_key); +} + +fn build_apns_payload( + event: &NotificationEvent, + device: &PushDevice, + capability: &PushDeviceCapabilityPlaintextV1, + preview_payload: Option<&NotificationPreviewPayload>, + send_encrypted_preview: bool, +) -> Result { + // Routing metadata is intentionally provider-visible in v1; the privacy boundary here is + // plaintext notification content, which stays generic unless the device decrypts os_push. + let metadata = preview_payload.map(|payload| { + json!({ + "notification_id": payload.notification_id.to_string(), + "kind": payload.kind, + "message_id": payload.message_id.to_string(), + "thread_id": payload.thread_id, + "deep_link": payload.deep_link, + }) + }); + + let mut root = json!({ + "aps": { + "alert": { + "title": event.fallback_title, + "body": event.fallback_body, + }, + "sound": "default", + "mutable-content": 1, + } + }); + + if let Some(payload) = preview_payload { + if let Some(aps) = root.get_mut("aps").and_then(|value| value.as_object_mut()) { + aps.insert( + "thread-id".to_string(), + Value::String(payload.thread_id.clone()), + ); + } + } + + if let Some(metadata) = metadata { + root.as_object_mut() + .expect("root payload must be an object") + .insert("os_meta".to_string(), metadata); + } + + if send_encrypted_preview { + if let Some(payload) = preview_payload { + let envelope = + encrypt_preview_payload(device.uuid, &capability.notification_public_key, payload)?; + root.as_object_mut() + .expect("root payload must be an object") + .insert("os_push".to_string(), serde_json::to_value(envelope)?); + } + } + + Ok(root) +} + +fn parse_apns_reason(response_body: &str) -> Option { + serde_json::from_str::(response_body) + .ok() + .and_then(|value| { + value + .get("reason") + .and_then(|reason| reason.as_str()) + .map(str::to_owned) + }) +} diff --git a/src/push/binding.rs b/src/push/binding.rs new file mode 100644 index 00000000..0f5ee018 --- /dev/null +++ b/src/push/binding.rs @@ -0,0 +1,326 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use crate::encrypt::{decrypt_aead_v1, derive_key, encrypt_aead_v1, CanonicalBytes, EncryptError}; +use crate::models::notification_events::NotificationEvent; +use crate::models::push_devices::PushDevice; + +const PUSH_DEVICE_AEAD_INFO: &[u8] = b"os.push-device-capability-aead.v1"; +const PUSH_DEVICE_AAD_DOMAIN: &str = "os.push-device-capability.v1"; +const PUSH_EVENT_PAYLOAD_AEAD_INFO: &[u8] = b"os.push-event-payload-aead.v1"; +const PUSH_EVENT_PAYLOAD_AAD_DOMAIN: &str = "os.push-event-payload.v1"; + +pub const PUSH_CAPABILITY_VERSION_V1: i16 = 1; +pub const PUSH_EVENT_PAYLOAD_VERSION_V1: i16 = 1; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct PushDeviceCapabilityPlaintextV1 { + pub v: i16, + pub push_device_uuid: Uuid, + pub user_uuid: Uuid, + pub project_id: i32, + pub installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub key_algorithm: String, + pub push_token: String, + pub push_token_hash: Vec, + pub notification_public_key: Vec, + pub notification_public_key_hash: Vec, + pub supports_encrypted_preview: bool, + pub supports_background_processing: bool, + pub registered_at: DateTime, +} + +#[derive(Debug, Clone)] +pub struct PushDeviceCapabilityInput { + pub push_device_uuid: Uuid, + pub user_uuid: Uuid, + pub project_id: i32, + pub installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub key_algorithm: String, + pub push_token: String, + pub notification_public_key: Vec, + pub supports_encrypted_preview: bool, + pub supports_background_processing: bool, +} + +impl PushDeviceCapabilityPlaintextV1 { + pub fn new(input: PushDeviceCapabilityInput) -> Self { + let push_token_hash = hash_bytes(input.push_token.as_bytes()); + let notification_public_key_hash = hash_bytes(&input.notification_public_key); + Self { + v: PUSH_CAPABILITY_VERSION_V1, + push_device_uuid: input.push_device_uuid, + user_uuid: input.user_uuid, + project_id: input.project_id, + installation_id: input.installation_id, + platform: input.platform, + provider: input.provider, + environment: input.environment, + app_id: input.app_id, + key_algorithm: input.key_algorithm, + push_token: input.push_token, + push_token_hash, + notification_public_key: input.notification_public_key, + notification_public_key_hash, + supports_encrypted_preview: input.supports_encrypted_preview, + supports_background_processing: input.supports_background_processing, + registered_at: Utc::now(), + } + } + + pub fn matches_device_row(&self, device: &PushDevice) -> bool { + self.v == PUSH_CAPABILITY_VERSION_V1 + && self.push_device_uuid == device.uuid + && self.user_uuid == device.user_id + && self.project_id == device.project_id + && self.installation_id == device.installation_id + && self.platform == device.platform + && self.provider == device.provider + && self.environment == device.environment + && self.app_id == device.app_id + && self.key_algorithm == device.key_algorithm + && self.push_token_hash == device.push_token_hash + && self.notification_public_key_hash == device.notification_public_key_hash + && self.supports_encrypted_preview == device.supports_encrypted_preview + && self.supports_background_processing == device.supports_background_processing + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct NotificationPreviewPayloadV1 { + pub v: i16, + pub notification_id: Uuid, + pub user_uuid: Uuid, + pub project_id: i32, + pub source_kind: String, + pub source_request_id: Option, + pub agent_uuid: Option, + pub delivery_mode: String, + pub message_id: Uuid, + pub kind: String, + pub title: String, + pub body: String, + pub deep_link: String, + pub thread_id: String, + pub sent_at: i64, +} + +impl NotificationPreviewPayloadV1 { + pub fn matches_event(&self, event: &NotificationEvent) -> bool { + self.v == PUSH_EVENT_PAYLOAD_VERSION_V1 + && self.notification_id == event.uuid + && self.user_uuid == event.user_id + && self.project_id == event.project_id + && self.source_kind == event.source_kind + && self.source_request_id == event.source_request_id + && self.delivery_mode == event.delivery_mode + && self.kind == event.kind + } +} + +pub fn hash_bytes(bytes: &[u8]) -> Vec { + Sha256::digest(bytes).to_vec() +} + +pub fn encrypt_push_device_capability_v1( + root_key: &[u8], + plaintext: &PushDeviceCapabilityPlaintextV1, +) -> Result, EncryptError> { + let key = derive_key(root_key, PUSH_DEVICE_AEAD_INFO)?; + let aad = push_device_aad_v1(PushDeviceAad { + push_device_uuid: plaintext.push_device_uuid, + user_uuid: plaintext.user_uuid, + project_id: plaintext.project_id, + installation_id: plaintext.installation_id, + platform: &plaintext.platform, + provider: &plaintext.provider, + environment: &plaintext.environment, + app_id: &plaintext.app_id, + version: plaintext.v, + }); + let bytes = serde_json::to_vec(plaintext).map_err(|_| EncryptError::BadData)?; + encrypt_aead_v1(&key, &bytes, &aad) +} + +pub fn decrypt_push_device_capability_v1( + root_key: &[u8], + encrypted: &[u8], + device: &PushDevice, +) -> Result { + let key = derive_key(root_key, PUSH_DEVICE_AEAD_INFO)?; + let aad = push_device_aad_v1(PushDeviceAad { + push_device_uuid: device.uuid, + user_uuid: device.user_id, + project_id: device.project_id, + installation_id: device.installation_id, + platform: &device.platform, + provider: &device.provider, + environment: &device.environment, + app_id: &device.app_id, + version: PUSH_CAPABILITY_VERSION_V1, + }); + let bytes = decrypt_aead_v1(&key, encrypted, &aad)?; + serde_json::from_slice(&bytes).map_err(|_| EncryptError::BadData) +} + +pub fn encrypt_notification_preview_payload_v1( + root_key: &[u8], + plaintext: &NotificationPreviewPayloadV1, +) -> Result, EncryptError> { + let key = derive_key(root_key, PUSH_EVENT_PAYLOAD_AEAD_INFO)?; + let aad = notification_payload_aad_v1(NotificationPayloadAad { + notification_uuid: plaintext.notification_id, + user_uuid: plaintext.user_uuid, + project_id: plaintext.project_id, + source_kind: &plaintext.source_kind, + kind: &plaintext.kind, + delivery_mode: &plaintext.delivery_mode, + version: plaintext.v, + }); + let bytes = serde_json::to_vec(plaintext).map_err(|_| EncryptError::BadData)?; + encrypt_aead_v1(&key, &bytes, &aad) +} + +pub fn decrypt_notification_preview_payload_v1( + root_key: &[u8], + encrypted: &[u8], + event: &NotificationEvent, +) -> Result { + let key = derive_key(root_key, PUSH_EVENT_PAYLOAD_AEAD_INFO)?; + let aad = notification_payload_aad_v1(NotificationPayloadAad { + notification_uuid: event.uuid, + user_uuid: event.user_id, + project_id: event.project_id, + source_kind: &event.source_kind, + kind: &event.kind, + delivery_mode: &event.delivery_mode, + version: PUSH_EVENT_PAYLOAD_VERSION_V1, + }); + let bytes = decrypt_aead_v1(&key, encrypted, &aad)?; + serde_json::from_slice(&bytes).map_err(|_| EncryptError::BadData) +} + +struct PushDeviceAad<'a> { + push_device_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + installation_id: Uuid, + platform: &'a str, + provider: &'a str, + environment: &'a str, + app_id: &'a str, + version: i16, +} + +fn push_device_aad_v1(input: PushDeviceAad<'_>) -> Vec { + let mut aad = CanonicalBytes::new(PUSH_DEVICE_AAD_DOMAIN); + aad.append_uuid(input.push_device_uuid) + .append_uuid(input.user_uuid) + .append_i32(input.project_id) + .append_uuid(input.installation_id) + .append_str(input.platform) + .append_str(input.provider) + .append_str(input.environment) + .append_str(input.app_id) + .append_i16(input.version); + aad.into_bytes() +} + +struct NotificationPayloadAad<'a> { + notification_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + source_kind: &'a str, + kind: &'a str, + delivery_mode: &'a str, + version: i16, +} + +fn notification_payload_aad_v1(input: NotificationPayloadAad<'_>) -> Vec { + let mut aad = CanonicalBytes::new(PUSH_EVENT_PAYLOAD_AAD_DOMAIN); + aad.append_uuid(input.notification_uuid) + .append_uuid(input.user_uuid) + .append_i32(input.project_id) + .append_str(input.source_kind) + .append_str(input.kind) + .append_str(input.delivery_mode) + .append_i16(input.version); + aad.into_bytes() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::push_devices::{ + PushDevice, PUSH_ENV_PROD, PUSH_KEY_ALGORITHM_P256_ECDH_V1, PUSH_PLATFORM_IOS, + PUSH_PROVIDER_APNS, + }; + + fn test_device() -> PushDevice { + let public_key = vec![8, 9, 10]; + PushDevice { + id: 1, + uuid: Uuid::new_v4(), + user_id: Uuid::new_v4(), + project_id: 42, + installation_id: Uuid::new_v4(), + platform: PUSH_PLATFORM_IOS.to_string(), + provider: PUSH_PROVIDER_APNS.to_string(), + environment: PUSH_ENV_PROD.to_string(), + app_id: "ai.trymaple.ios".to_string(), + push_token_hash: hash_bytes(b"token"), + capability_enc: vec![], + notification_public_key_hash: hash_bytes(&public_key), + key_algorithm: PUSH_KEY_ALGORITHM_P256_ECDH_V1.to_string(), + supports_encrypted_preview: true, + supports_background_processing: true, + last_seen_at: Utc::now(), + revoked_at: None, + created_at: Utc::now(), + updated_at: Utc::now(), + } + } + + #[test] + fn push_device_capability_aad_binds_user_and_device() { + let root = [4_u8; 32]; + let mut device = test_device(); + let plaintext = PushDeviceCapabilityPlaintextV1::new(PushDeviceCapabilityInput { + push_device_uuid: device.uuid, + user_uuid: device.user_id, + project_id: device.project_id, + installation_id: device.installation_id, + platform: device.platform.clone(), + provider: device.provider.clone(), + environment: device.environment.clone(), + app_id: device.app_id.clone(), + key_algorithm: device.key_algorithm.clone(), + push_token: "token".to_string(), + notification_public_key: vec![8, 9, 10], + supports_encrypted_preview: true, + supports_background_processing: true, + }); + device.capability_enc = encrypt_push_device_capability_v1(&root, &plaintext).unwrap(); + + let decrypted = + decrypt_push_device_capability_v1(&root, &device.capability_enc, &device).unwrap(); + assert!(decrypted.matches_device_row(&device)); + + let mut attacker_row = device.clone(); + attacker_row.user_id = Uuid::new_v4(); + assert!( + decrypt_push_device_capability_v1(&root, &device.capability_enc, &attacker_row,) + .is_err() + ); + } +} diff --git a/src/push/crypto.rs b/src/push/crypto.rs new file mode 100644 index 00000000..2cbb610d --- /dev/null +++ b/src/push/crypto.rs @@ -0,0 +1,138 @@ +use crate::encrypt::generate_random; +use crate::push::{NotificationPreviewPayload, PushError}; +use aes_gcm::aead::{Aead, KeyInit}; +use aes_gcm::{Aes256Gcm, Nonce}; +use base64::{engine::general_purpose, Engine}; +use hkdf::Hkdf; +use p256::ecdh::EphemeralSecret; +use p256::elliptic_curve::sec1::ToEncodedPoint; +use p256::pkcs8::DecodePublicKey; +use p256::PublicKey; +use serde::{Deserialize, Serialize}; +use sha2::Sha256; +use uuid::Uuid; + +pub const PUSH_PREVIEW_INFO: &[u8] = b"opensecret-push-preview-v1"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EncryptedPreviewEnvelope { + #[serde(alias = "v")] + pub enc_v: i32, + pub alg: String, + pub kid: String, + pub epk: String, + pub salt: String, + pub nonce: String, + pub ciphertext: String, +} + +pub fn encrypt_preview_payload( + device_uuid: Uuid, + notification_public_key: &[u8], + payload: &NotificationPreviewPayload, +) -> Result { + let recipient_key = PublicKey::from_public_key_der(notification_public_key) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + let ephemeral_secret = EphemeralSecret::random(&mut p256::elliptic_curve::rand_core::OsRng); + let ephemeral_public = PublicKey::from(&ephemeral_secret); + let shared_secret = ephemeral_secret.diffie_hellman(&recipient_key); + + let salt = generate_random::<32>(); + let hkdf = Hkdf::::new(Some(&salt), shared_secret.raw_secret_bytes().as_slice()); + let mut key = [0_u8; 32]; + hkdf.expand(PUSH_PREVIEW_INFO, &mut key) + .map_err(|_| PushError::CryptoError("failed to derive push preview key".to_string()))?; + + let nonce_bytes = generate_random::<12>(); + let cipher = + Aes256Gcm::new_from_slice(&key).map_err(|e| PushError::CryptoError(e.to_string()))?; + let plaintext = serde_json::to_vec(payload)?; + let ciphertext = cipher + .encrypt(Nonce::from_slice(&nonce_bytes), plaintext.as_ref()) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + + Ok(EncryptedPreviewEnvelope { + enc_v: 1, + alg: "p256-hkdf-sha256-aes256gcm".to_string(), + kid: device_uuid.to_string(), + epk: general_purpose::STANDARD.encode(ephemeral_public.to_encoded_point(false).as_bytes()), + salt: general_purpose::STANDARD.encode(salt), + nonce: general_purpose::STANDARD.encode(nonce_bytes), + ciphertext: general_purpose::STANDARD.encode(ciphertext), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use aes_gcm::aead::{Aead, KeyInit}; + use aes_gcm::{Aes256Gcm, Nonce}; + use base64::{engine::general_purpose, Engine}; + use chrono::Utc; + use hkdf::Hkdf; + use p256::elliptic_curve::rand_core::OsRng; + use p256::pkcs8::EncodePublicKey; + use p256::SecretKey; + use uuid::Uuid; + + #[test] + fn encrypt_preview_payload_round_trips() { + let recipient_secret = SecretKey::random(&mut OsRng); + let recipient_public = recipient_secret.public_key(); + let device_uuid = Uuid::new_v4(); + let notification_public_key = recipient_public + .to_public_key_der() + .unwrap() + .as_bytes() + .to_vec(); + + let payload = NotificationPreviewPayload { + v: 1, + notification_id: Uuid::new_v4(), + user_uuid: Uuid::new_v4(), + project_id: 42, + source_kind: "request_continuation".to_string(), + source_request_id: None, + agent_uuid: None, + delivery_mode: "encrypted_preview".to_string(), + message_id: Uuid::new_v4(), + kind: "agent.message".to_string(), + title: "Maple".to_string(), + body: "Hello from Maple".to_string(), + deep_link: "opensecret://agent".to_string(), + thread_id: "agent:main".to_string(), + sent_at: Utc::now().timestamp(), + }; + + let envelope = + encrypt_preview_payload(device_uuid, ¬ification_public_key, &payload).unwrap(); + let epk_bytes = general_purpose::STANDARD.decode(&envelope.epk).unwrap(); + let salt = general_purpose::STANDARD.decode(&envelope.salt).unwrap(); + let nonce = general_purpose::STANDARD.decode(&envelope.nonce).unwrap(); + let ciphertext = general_purpose::STANDARD + .decode(&envelope.ciphertext) + .unwrap(); + + let ephemeral_public = PublicKey::from_sec1_bytes(&epk_bytes).unwrap(); + let shared_secret = p256::ecdh::diffie_hellman( + recipient_secret.to_nonzero_scalar(), + ephemeral_public.as_affine(), + ); + let hkdf = Hkdf::::new(Some(&salt), shared_secret.raw_secret_bytes().as_slice()); + let mut key = [0_u8; 32]; + hkdf.expand(PUSH_PREVIEW_INFO, &mut key).unwrap(); + + let cipher = Aes256Gcm::new_from_slice(&key).unwrap(); + let plaintext = cipher + .decrypt(Nonce::from_slice(&nonce), ciphertext.as_ref()) + .unwrap(); + let decoded: NotificationPreviewPayload = serde_json::from_slice(&plaintext).unwrap(); + + assert_eq!(decoded.notification_id, payload.notification_id); + assert_eq!(decoded.body, payload.body); + assert_eq!(decoded.deep_link, payload.deep_link); + assert_eq!(envelope.enc_v, 1); + assert_eq!(envelope.alg, "p256-hkdf-sha256-aes256gcm"); + assert_eq!(envelope.kid, device_uuid.to_string()); + } +} diff --git a/src/push/fcm.rs b/src/push/fcm.rs new file mode 100644 index 00000000..1caf7be8 --- /dev/null +++ b/src/push/fcm.rs @@ -0,0 +1,305 @@ +use crate::models::notification_events::NotificationEvent; +use crate::models::project_settings::AndroidPushSettings; +use crate::models::push_devices::PushDevice; +use crate::push::{ + CachedFcmToken, NotificationPreviewPayload, PushError, PushSendOutcome, PushTransport, + FCM_TOKEN_CACHE_SAFETY_MARGIN_SECONDS, +}; +use crate::web::platform::PROJECT_FCM_SERVICE_ACCOUNT_JSON; +use crate::AppState; +use chrono::{Duration, Utc}; +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::sync::Arc; + +const GOOGLE_TOKEN_SCOPE: &str = "https://www.googleapis.com/auth/firebase.messaging"; +const DEFAULT_GOOGLE_TOKEN_URI: &str = "https://oauth2.googleapis.com/token"; + +#[derive(Debug, Deserialize)] +struct FcmServiceAccount { + pub client_email: String, + pub private_key: String, + pub project_id: Option, + pub token_uri: Option, +} + +#[derive(Debug, Serialize)] +struct FcmJwtClaims<'a> { + iss: &'a str, + scope: &'a str, + aud: &'a str, + exp: i64, + iat: i64, +} + +#[derive(Debug, Deserialize)] +struct FcmAccessTokenResponse { + access_token: String, + expires_in: i64, +} + +#[derive(Debug, Deserialize)] +struct FcmSendResponse { + name: String, +} + +#[derive(Debug, Deserialize)] +struct GoogleApiErrorEnvelope { + error: GoogleApiError, +} + +#[derive(Debug, Deserialize)] +struct GoogleApiError { + status: Option, + message: Option, + details: Option>, +} + +pub async fn send_fcm_notification( + state: &Arc, + transport: &PushTransport, + event: &NotificationEvent, + _device: &PushDevice, + push_token: &str, + android_settings: &AndroidPushSettings, + preview_payload: Option<&NotificationPreviewPayload>, +) -> Result { + let service_account = load_fcm_service_account(state, event.project_id).await?; + let project_id = if android_settings.firebase_project_id.trim().is_empty() { + service_account + .project_id + .clone() + .ok_or_else(|| PushError::InvalidSecret("missing FCM project_id".to_string()))? + } else { + android_settings.firebase_project_id.clone() + }; + let access_token = get_fcm_access_token(transport, event.project_id, &service_account).await?; + let body = build_fcm_payload(event, push_token, preview_payload); + let response = transport + .client + .post(format!( + "https://fcm.googleapis.com/v1/projects/{}/messages:send", + project_id + )) + .bearer_auth(access_token) + .json(&body) + .send() + .await?; + + let status_code = response.status().as_u16() as i32; + let status = response.status(); + let response_body = response.text().await.unwrap_or_default(); + + if status.is_success() { + let parsed: FcmSendResponse = serde_json::from_str(&response_body)?; + return Ok(PushSendOutcome::Sent { + provider_message_id: Some(parsed.name), + provider_status_code: Some(status_code), + }); + } + + let parsed_error = serde_json::from_str::(&response_body).ok(); + let error_status = parsed_error + .as_ref() + .and_then(|body| body.error.status.clone()) + .unwrap_or_default(); + let error_message = parsed_error + .as_ref() + .and_then(|body| body.error.message.clone()) + .unwrap_or_else(|| response_body.clone()); + let fcm_error_code = parsed_error + .as_ref() + .and_then(|body| extract_fcm_error_code(body.error.details.as_ref())); + + if matches!(status.as_u16(), 401 | 403) { + invalidate_fcm_cache(transport, event.project_id, &service_account.client_email).await; + } + + let outcome = if matches!(fcm_error_code.as_deref(), Some("UNREGISTERED")) + || is_invalid_registration(&error_status, &error_message) + { + PushSendOutcome::InvalidToken { + provider_status_code: Some(status_code), + error: if let Some(code) = fcm_error_code { + format!("{}: {}", code, error_message) + } else { + error_message + }, + } + } else if matches!(status.as_u16(), 401 | 403 | 429 | 500 | 503) { + PushSendOutcome::Retryable { + provider_status_code: Some(status_code), + error: error_message, + } + } else { + PushSendOutcome::Failed { + provider_status_code: Some(status_code), + error: error_message, + } + }; + + Ok(outcome) +} + +async fn load_fcm_service_account( + state: &Arc, + project_id: i32, +) -> Result { + let raw = state + .get_project_secret_string(project_id, PROJECT_FCM_SERVICE_ACCOUNT_JSON) + .await? + .ok_or_else(|| PushError::InvalidSecret("missing FCM service account JSON".to_string()))?; + + serde_json::from_str(&raw).map_err(PushError::from) +} + +async fn get_fcm_access_token( + transport: &PushTransport, + project_id: i32, + service_account: &FcmServiceAccount, +) -> Result { + let cache_key = format!("{}:{}", project_id, service_account.client_email); + { + let cache = transport.fcm_tokens.read().await; + if let Some(cached) = cache.get(&cache_key) { + if cached.expires_at + > Utc::now() + Duration::seconds(FCM_TOKEN_CACHE_SAFETY_MARGIN_SECONDS) + { + return Ok(cached.token.clone()); + } + } + } + + let token_uri = service_account + .token_uri + .as_deref() + .unwrap_or(DEFAULT_GOOGLE_TOKEN_URI); + let now = Utc::now().timestamp(); + let claims = FcmJwtClaims { + iss: &service_account.client_email, + scope: GOOGLE_TOKEN_SCOPE, + aud: token_uri, + exp: now + 3600, + iat: now - 30, + }; + let header = Header::new(Algorithm::RS256); + let signing_key = EncodingKey::from_rsa_pem(service_account.private_key.as_bytes()) + .map_err(|e| PushError::InvalidSecret(e.to_string()))?; + let assertion = encode(&header, &claims, &signing_key)?; + + let response = transport + .client + .post(token_uri) + .form(&[ + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ("assertion", assertion.as_str()), + ]) + .send() + .await?; + let response_status = response.status(); + let response_body = response.text().await.unwrap_or_default(); + + if !response_status.is_success() { + let error = format!("FCM OAuth failed ({}): {}", response_status, response_body); + return Err( + if response_status.as_u16() == 429 || response_status.is_server_error() { + PushError::ProviderRetryable(error) + } else { + PushError::ProviderError(error) + }, + ); + } + + let token_response: FcmAccessTokenResponse = serde_json::from_str(&response_body)?; + let expires_at = Utc::now() + Duration::seconds(token_response.expires_in); + let token = token_response.access_token; + + let mut cache = transport.fcm_tokens.write().await; + cache.insert( + cache_key, + CachedFcmToken { + token: token.clone(), + expires_at, + }, + ); + + Ok(token) +} + +async fn invalidate_fcm_cache(transport: &PushTransport, project_id: i32, client_email: &str) { + let cache_key = format!("{}:{}", project_id, client_email); + let mut cache = transport.fcm_tokens.write().await; + cache.remove(&cache_key); +} + +fn build_fcm_payload( + event: &NotificationEvent, + push_token: &str, + preview_payload: Option<&NotificationPreviewPayload>, +) -> Value { + let ttl_seconds = event + .expires_at + .map(|expires_at| (expires_at - Utc::now()).num_seconds().max(0)) + .unwrap_or(60 * 60 * 24 * 7); + let collapse_key = event + .collapse_key + .clone() + .unwrap_or_else(|| format!("notif:{}", event.uuid)); + + // Android v1 keeps routing metadata in plaintext FCM data fields; the privacy requirement is + // that message content stays generic/provider-visible while authoritative content is fetched or decrypted on-device. + let data = if let Some(payload) = preview_payload { + json!({ + "notification_id": payload.notification_id.to_string(), + "kind": payload.kind, + "message_id": payload.message_id.to_string(), + "thread_id": payload.thread_id, + "deep_link": payload.deep_link, + }) + } else { + json!({ + "notification_id": event.uuid.to_string(), + "kind": event.kind, + }) + }; + + json!({ + "message": { + "token": push_token, + "notification": { + "title": event.fallback_title, + "body": event.fallback_body, + }, + "data": data, + "android": { + "priority": if event.priority == "high" { "HIGH" } else { "NORMAL" }, + "ttl": format!("{}s", ttl_seconds), + "collapse_key": collapse_key, + "notification": { + "channel_id": "maple_messages", + "tag": event.uuid.to_string(), + "click_action": "OPEN_THREAD", + } + } + } + }) +} + +fn extract_fcm_error_code(details: Option<&Vec>) -> Option { + details.and_then(|entries| { + entries.iter().find_map(|entry| { + entry + .get("errorCode") + .and_then(|value| value.as_str()) + .map(str::to_owned) + }) + }) +} + +fn is_invalid_registration(error_status: &str, error_message: &str) -> bool { + error_status == "INVALID_ARGUMENT" + && error_message + .to_ascii_lowercase() + .contains("registration token") +} diff --git a/src/push/mod.rs b/src/push/mod.rs new file mode 100644 index 00000000..240f6b85 --- /dev/null +++ b/src/push/mod.rs @@ -0,0 +1,435 @@ +pub mod apns; +pub mod binding; +pub mod crypto; +pub mod fcm; +pub mod worker; + +use crate::db::DBError; +use crate::models::notification_deliveries::{NewNotificationDelivery, NotificationDeliveryError}; +use crate::models::notification_events::{ + NewNotificationEvent, NotificationEvent, NotificationEventError, + NOTIFICATION_DELIVERY_MODE_ENCRYPTED_PREVIEW, NOTIFICATION_DELIVERY_MODE_GENERIC, + NOTIFICATION_KIND_AGENT_MESSAGE, NOTIFICATION_PRIORITY_HIGH, NOTIFICATION_PRIORITY_NORMAL, + NOTIFICATION_SOURCE_REQUEST_CONTINUATION, +}; +use crate::models::project_settings::ProjectSettingError; +use crate::models::push_devices::{PushDevice, PushDeviceError}; +use crate::models::users::User; +use crate::push::binding::encrypt_notification_preview_payload_v1; +use crate::{AppState, Error}; +use chrono::{DateTime, Duration, Utc}; +use diesel::Connection; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use thiserror::Error; +use tokio::sync::RwLock; +use tracing::warn; +use uuid::Uuid; + +pub const AGENT_NOTIFICATION_FALLBACK_TITLE: &str = "New Maple message"; +pub const AGENT_NOTIFICATION_FALLBACK_BODY: &str = "Open Maple to view your encrypted message"; +pub const APNS_JWT_CACHE_LIFETIME_MINUTES: i64 = 50; +pub const FCM_TOKEN_CACHE_SAFETY_MARGIN_SECONDS: i64 = 60; +pub const PUSH_PREVIEW_BODY_MAX_BYTES: usize = 180; + +#[derive(Debug, Error)] +pub enum PushError { + #[error("Database connection error")] + ConnectionError, + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), + #[error("Push device error: {0}")] + PushDeviceError(#[from] PushDeviceError), + #[error("Notification event error: {0}")] + NotificationEventError(#[from] NotificationEventError), + #[error("Notification delivery error: {0}")] + NotificationDeliveryError(#[from] NotificationDeliveryError), + #[error("Project settings error: {0}")] + ProjectSettingError(#[from] ProjectSettingError), + #[error("DB error: {0}")] + DbError(#[from] DBError), + #[error("Application error: {0}")] + AppError(#[from] Error), + #[error("Serialization error: {0}")] + SerializationError(#[from] serde_json::Error), + #[error("Encryption error: {0}")] + EncryptionError(#[from] crate::encrypt::EncryptError), + #[error("HTTP error: {0}")] + HttpError(#[from] reqwest::Error), + #[error("JWT error: {0}")] + JwtError(#[from] jsonwebtoken::errors::Error), + #[error("Crypto error: {0}")] + CryptoError(String), + #[error("Invalid secret: {0}")] + InvalidSecret(String), + #[error("Provider error: {0}")] + ProviderError(String), + #[error("Retryable provider error: {0}")] + ProviderRetryable(String), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PushDeliveryMode { + Generic, + EncryptedPreview, +} + +impl PushDeliveryMode { + pub fn as_str(&self) -> &'static str { + match self { + Self::Generic => NOTIFICATION_DELIVERY_MODE_GENERIC, + Self::EncryptedPreview => NOTIFICATION_DELIVERY_MODE_ENCRYPTED_PREVIEW, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PushPriority { + #[allow(dead_code)] + Normal, + High, +} + +impl PushPriority { + pub fn as_str(&self) -> &'static str { + match self { + Self::Normal => NOTIFICATION_PRIORITY_NORMAL, + Self::High => NOTIFICATION_PRIORITY_HIGH, + } + } +} + +#[derive(Debug, Clone)] +pub enum AgentPushTarget { + Main, + Subagent(Uuid), +} + +pub use crate::push::binding::NotificationPreviewPayloadV1 as NotificationPreviewPayload; + +#[derive(Debug, Clone)] +pub enum PushNotificationSource { + RequestContinuation { + source_request_id: Option, + agent_uuid: Option, + }, +} + +impl PushNotificationSource { + fn source_kind(&self) -> &'static str { + match self { + Self::RequestContinuation { .. } => NOTIFICATION_SOURCE_REQUEST_CONTINUATION, + } + } + + fn source_request_id(&self) -> Option { + match self { + Self::RequestContinuation { + source_request_id, .. + } => *source_request_id, + } + } + + fn agent_uuid(&self) -> Option { + match self { + Self::RequestContinuation { agent_uuid, .. } => *agent_uuid, + } + } +} + +#[derive(Debug, Clone)] +pub struct NotificationPreviewPayloadInput { + pub message_id: Uuid, + pub kind: String, + pub title: String, + pub body: String, + pub deep_link: String, + pub thread_id: String, + pub sent_at: i64, +} + +fn normalize_preview_body(message_text: &str) -> String { + let normalized = message_text + .split_whitespace() + .collect::>() + .join(" "); + truncate_utf8_with_ellipsis(&normalized, PUSH_PREVIEW_BODY_MAX_BYTES) +} + +fn truncate_utf8_with_ellipsis(input: &str, max_bytes: usize) -> String { + if input.len() <= max_bytes { + return input.to_string(); + } + + let ellipsis = "…"; + if max_bytes <= ellipsis.len() { + return String::new(); + } + + let mut end = max_bytes - ellipsis.len(); + while end > 0 && !input.is_char_boundary(end) { + end -= 1; + } + + let truncated = input[..end].trim_end(); + if truncated.is_empty() { + String::new() + } else { + format!("{}{}", truncated, ellipsis) + } +} + +#[derive(Debug, Clone)] +pub struct EnqueueNotificationRequest { + pub project_id: i32, + pub user_id: Uuid, + pub kind: String, + pub delivery_mode: PushDeliveryMode, + pub priority: PushPriority, + pub fallback_title: String, + pub fallback_body: String, + pub preview_payload: Option, + pub source: PushNotificationSource, + pub not_before_at: Option>, + pub expires_at: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueuedNotification { + pub notification_id: Uuid, + pub delivery_count: usize, +} + +#[derive(Debug, Clone)] +pub enum PushSendOutcome { + Sent { + provider_message_id: Option, + provider_status_code: Option, + }, + Retryable { + provider_status_code: Option, + error: String, + }, + InvalidToken { + provider_status_code: Option, + error: String, + }, + Failed { + provider_status_code: Option, + error: String, + }, +} + +#[derive(Clone)] +pub(crate) struct PushTransport { + pub client: reqwest::Client, + pub apns_client: reqwest::Client, + pub apns_tokens: Arc>>, + pub fcm_tokens: Arc>>, +} + +impl PushTransport { + pub fn new() -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(20)) + .build()?; + let apns_client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(20)) + .http2_prior_knowledge() + .build()?; + + Ok(Self { + client, + apns_client, + apns_tokens: Arc::new(RwLock::new(HashMap::new())), + fcm_tokens: Arc::new(RwLock::new(HashMap::new())), + }) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct CachedApnsToken { + pub token: String, + pub expires_at: DateTime, +} + +#[derive(Debug, Clone)] +pub(crate) struct CachedFcmToken { + pub token: String, + pub expires_at: DateTime, +} + +pub async fn enqueue_notification( + state: &Arc, + request: EnqueueNotificationRequest, +) -> Result, PushError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| PushError::ConnectionError)?; + + let devices = PushDevice::list_active_for_user(&mut conn, request.user_id)?; + if devices.is_empty() { + return Ok(None); + } + + let notification_id = Uuid::new_v4(); + let payload_enc = if let Some(preview_payload) = request.preview_payload.as_ref() { + let payload = NotificationPreviewPayload { + v: 1, + notification_id, + user_uuid: request.user_id, + project_id: request.project_id, + source_kind: request.source.source_kind().to_string(), + source_request_id: request.source.source_request_id(), + agent_uuid: request.source.agent_uuid(), + delivery_mode: request.delivery_mode.as_str().to_string(), + message_id: preview_payload.message_id, + kind: preview_payload.kind.clone(), + title: preview_payload.title.clone(), + body: preview_payload.body.clone(), + deep_link: preview_payload.deep_link.clone(), + thread_id: preview_payload.thread_id.clone(), + sent_at: preview_payload.sent_at, + }; + Some(encrypt_notification_preview_payload_v1( + &state.enclave_key, + &payload, + )?) + } else { + None + }; + + let event = conn.transaction::(|conn| { + let event = NewNotificationEvent { + uuid: notification_id, + project_id: request.project_id, + user_id: request.user_id, + kind: request.kind, + delivery_mode: request.delivery_mode.as_str().to_string(), + priority: request.priority.as_str().to_string(), + collapse_key: Some(format!("notif:{}", notification_id)), + fallback_title: request.fallback_title, + fallback_body: request.fallback_body, + payload_enc, + source_kind: request.source.source_kind().to_string(), + source_request_id: request.source.source_request_id(), + not_before_at: request.not_before_at, + expires_at: request.expires_at, + } + .insert(conn)?; + + let new_deliveries: Vec = devices + .iter() + .map(|device| NewNotificationDelivery { + event_id: event.id, + push_device_id: device.id, + next_attempt_at: event.not_before_at, + }) + .collect(); + + if new_deliveries.is_empty() { + warn!( + "Notification {} had no active push deliveries to insert", + event.uuid + ); + } else { + NewNotificationDelivery::insert_many(conn, &new_deliveries)?; + } + + Ok(event) + })?; + + Ok(Some(QueuedNotification { + notification_id: event.uuid, + delivery_count: devices.len(), + })) +} + +pub async fn enqueue_agent_message_notification( + state: &Arc, + user: &User, + target: AgentPushTarget, + message_id: Uuid, + message_text: &str, +) -> Result, PushError> { + if message_text.trim().is_empty() { + return Ok(None); + } + + let push_settings = state + .db + .get_project_push_settings(user.project_id)? + .unwrap_or_default(); + let delivery_mode = if push_settings.encrypted_preview_enabled { + PushDeliveryMode::EncryptedPreview + } else { + PushDeliveryMode::Generic + }; + + let (deep_link, thread_id, source_agent_uuid) = match &target { + AgentPushTarget::Main => ( + "opensecret://agent".to_string(), + "agent:main".to_string(), + None, + ), + AgentPushTarget::Subagent(agent_uuid) => ( + format!("opensecret://agent/subagent/{}", agent_uuid), + format!("agent:subagent:{}", agent_uuid), + Some(*agent_uuid), + ), + }; + + enqueue_notification( + state, + EnqueueNotificationRequest { + project_id: user.project_id, + user_id: user.uuid, + kind: NOTIFICATION_KIND_AGENT_MESSAGE.to_string(), + delivery_mode, + priority: PushPriority::High, + fallback_title: AGENT_NOTIFICATION_FALLBACK_TITLE.to_string(), + fallback_body: AGENT_NOTIFICATION_FALLBACK_BODY.to_string(), + preview_payload: Some(NotificationPreviewPayloadInput { + message_id, + kind: NOTIFICATION_KIND_AGENT_MESSAGE.to_string(), + title: "Maple".to_string(), + body: normalize_preview_body(message_text), + deep_link, + thread_id, + sent_at: Utc::now().timestamp(), + }), + source: PushNotificationSource::RequestContinuation { + source_request_id: None, + agent_uuid: source_agent_uuid, + }, + not_before_at: None, + expires_at: Some(Utc::now() + Duration::days(7)), + }, + ) + .await +} + +#[cfg(test)] +mod tests { + use super::{normalize_preview_body, PUSH_PREVIEW_BODY_MAX_BYTES}; + + #[test] + fn normalize_preview_body_collapses_whitespace() { + assert_eq!( + normalize_preview_body("hello\n\nthere world"), + "hello there world" + ); + } + + #[test] + fn normalize_preview_body_truncates_to_budget() { + let input = "a".repeat(PUSH_PREVIEW_BODY_MAX_BYTES + 50); + let output = normalize_preview_body(&input); + + assert!(output.len() <= PUSH_PREVIEW_BODY_MAX_BYTES); + assert!(output.ends_with('…')); + } +} diff --git a/src/push/worker.rs b/src/push/worker.rs new file mode 100644 index 00000000..c18327f5 --- /dev/null +++ b/src/push/worker.rs @@ -0,0 +1,556 @@ +use crate::models::notification_deliveries::{ + NotificationDelivery, NotificationDeliveryWriteResult, +}; +use crate::models::notification_events::{ + NotificationEvent, NOTIFICATION_SOURCE_REQUEST_CONTINUATION, +}; +use crate::models::push_devices::{PushDevice, PUSH_PLATFORM_ANDROID, PUSH_PLATFORM_IOS}; +use crate::push::apns::{send_apns_notification, ApnsSendRequest}; +use crate::push::binding::{ + decrypt_notification_preview_payload_v1, decrypt_push_device_capability_v1, + PushDeviceCapabilityPlaintextV1, +}; +use crate::push::fcm::send_fcm_notification; +use crate::push::{NotificationPreviewPayload, PushError, PushSendOutcome, PushTransport}; +use crate::AppState; +use chrono::Utc; +use diesel::Connection; +use futures::stream::{self, StreamExt}; +use std::sync::Arc; +use tokio::time::{sleep, Duration as TokioDuration}; +use tracing::{debug, error}; +use uuid::Uuid; + +const PUSH_WORKER_BATCH_SIZE: i64 = 32; +const PUSH_WORKER_LEASE_TTL_SECONDS: i32 = 60; +const PUSH_WORKER_POLL_INTERVAL_SECONDS: u64 = 3; +const PUSH_WORKER_MAX_CONCURRENCY: usize = 8; +const PUSH_WORKER_MAX_ATTEMPTS: i32 = 8; + +pub fn start_push_worker(state: Arc) { + tokio::spawn(async move { + let transport = match PushTransport::new() { + Ok(transport) => transport, + Err(error) => { + error!("failed to initialize push transport: {:?}", error); + return; + } + }; + + loop { + if let Err(error) = process_push_batch(&state, &transport).await { + error!("push worker batch failed: {:?}", error); + } + + sleep(TokioDuration::from_secs(PUSH_WORKER_POLL_INTERVAL_SECONDS)).await; + } + }); +} + +async fn process_push_batch( + state: &Arc, + transport: &PushTransport, +) -> Result<(), PushError> { + let lease_owner = format!("push-worker:{}:{}", std::process::id(), Uuid::new_v4()); + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| PushError::ConnectionError)?; + let deliveries = NotificationDelivery::lease_pending( + &mut conn, + PUSH_WORKER_BATCH_SIZE, + &lease_owner, + PUSH_WORKER_LEASE_TTL_SECONDS, + )?; + drop(conn); + + if deliveries.is_empty() { + return Ok(()); + } + + debug!("leased {} push deliveries", deliveries.len()); + + stream::iter(deliveries) + .for_each_concurrent(PUSH_WORKER_MAX_CONCURRENCY, |delivery| { + let state = state.clone(); + let transport = transport.clone(); + async move { + if let Err(error) = process_leased_delivery(&state, &transport, delivery).await { + error!("push delivery processing failed: {:?}", error); + } + } + }) + .await; + + Ok(()) +} + +async fn process_leased_delivery( + state: &Arc, + transport: &PushTransport, + delivery: NotificationDelivery, +) -> Result<(), PushError> { + let Some(lease_owner) = delivery.lease_owner.clone() else { + error!( + "leased push delivery {} is missing lease_owner; waiting for lease expiry", + delivery.id + ); + return Ok(()); + }; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| PushError::ConnectionError)?; + let event = match NotificationEvent::get_by_id(&mut conn, delivery.event_id)? { + Some(event) => event, + None => { + if !record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + None, + Some("event missing"), + )?, + delivery.id, + &lease_owner, + "failed (event missing)", + ) { + return Ok(()); + } + return Ok(()); + } + }; + + if event.cancelled_at.is_some() { + if !record_delivery_transition( + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + &lease_owner, + Some("event cancelled"), + )?, + delivery.id, + &lease_owner, + "cancelled (event cancelled)", + ) { + return Ok(()); + } + return Ok(()); + } + + if event + .expires_at + .is_some_and(|expires_at| expires_at <= Utc::now()) + { + if !record_delivery_transition( + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + &lease_owner, + Some("event expired"), + )?, + delivery.id, + &lease_owner, + "cancelled (event expired)", + ) { + return Ok(()); + } + return Ok(()); + } + + let device = match PushDevice::get_by_id(&mut conn, delivery.push_device_id)? { + Some(device) => device, + None => { + if !record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + None, + Some("device missing"), + )?, + delivery.id, + &lease_owner, + "failed (device missing)", + ) { + return Ok(()); + } + return Ok(()); + } + }; + + if device.user_id != event.user_id || device.project_id != event.project_id { + if !record_delivery_transition( + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + &lease_owner, + Some("device principal does not match event principal"), + )?, + delivery.id, + &lease_owner, + "cancelled (device principal mismatch)", + ) { + return Ok(()); + } + return Ok(()); + } + + if device.revoked_at.is_some() { + if !record_delivery_transition( + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + &lease_owner, + Some("device revoked"), + )?, + delivery.id, + &lease_owner, + "cancelled (device revoked)", + ) { + return Ok(()); + } + return Ok(()); + } + drop(conn); + + let send_outcome = match build_send_outcome(state, transport, &event, &device).await { + Ok(outcome) => outcome, + Err(error) => { + error!( + "push delivery {} encountered internal send error before provider outcome: {:?}", + delivery.id, error + ); + classify_internal_push_error(error) + } + }; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| PushError::ConnectionError)?; + match send_outcome { + PushSendOutcome::Sent { + provider_message_id, + provider_status_code, + } => { + record_delivery_transition( + NotificationDelivery::mark_sent( + &mut conn, + delivery.id, + &lease_owner, + provider_message_id.as_deref(), + provider_status_code, + )?, + delivery.id, + &lease_owner, + "sent", + ); + } + PushSendOutcome::Retryable { + provider_status_code, + error, + } => { + if delivery.attempt_count + 1 >= PUSH_WORKER_MAX_ATTEMPTS { + record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + )?, + delivery.id, + &lease_owner, + "failed (retry limit reached)", + ); + } else { + record_delivery_transition( + NotificationDelivery::mark_retry( + &mut conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + retry_backoff_seconds(delivery.attempt_count + 1), + )?, + delivery.id, + &lease_owner, + "retry", + ); + } + } + PushSendOutcome::InvalidToken { + provider_status_code, + error, + } => { + let invalidated = conn.transaction::(|conn| { + let write_result = NotificationDelivery::mark_invalid_token( + conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + )?; + + if write_result.was_applied() { + PushDevice::invalidate(conn, device.id)?; + Ok(true) + } else { + Ok(false) + } + })?; + + if !invalidated { + debug!( + "push delivery {} lost lease before marking invalid_token (lease_owner={})", + delivery.id, lease_owner + ); + } + } + PushSendOutcome::Failed { + provider_status_code, + error, + } => { + record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + )?, + delivery.id, + &lease_owner, + "failed", + ); + } + } + + Ok(()) +} + +async fn build_send_outcome( + state: &Arc, + transport: &PushTransport, + event: &NotificationEvent, + device: &PushDevice, +) -> Result { + let capability = decrypt_device_capability(state, device)?; + if !capability.matches_device_row(device) { + return Err(PushError::CryptoError( + "push device capability does not match device row".to_string(), + )); + } + let preview_payload = decrypt_preview_payload(state, event)?; + + dispatch_delivery( + state, + transport, + event, + device, + &capability, + preview_payload.as_ref(), + ) + .await +} + +async fn dispatch_delivery( + state: &Arc, + transport: &PushTransport, + event: &NotificationEvent, + device: &PushDevice, + capability: &PushDeviceCapabilityPlaintextV1, + preview_payload: Option<&NotificationPreviewPayload>, +) -> Result { + let push_settings = state + .db + .get_project_push_settings(event.project_id)? + .unwrap_or_default(); + + match device.platform.as_str() { + PUSH_PLATFORM_IOS => { + let Some(ios_settings) = push_settings + .ios + .as_ref() + .filter(|settings| settings.enabled) + else { + return Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: "iOS push is not configured for this project".to_string(), + }); + }; + + if device.app_id != ios_settings.bundle_id { + return Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: "device app_id does not match configured iOS bundle_id".to_string(), + }); + } + + if device.environment != ios_settings.apns_environment.as_str() { + return Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: "device environment does not match project APNs environment".to_string(), + }); + } + + let send_encrypted_preview = push_settings.encrypted_preview_enabled + && event.delivery_mode == "encrypted_preview" + && preview_payload.is_some() + && capability.supports_encrypted_preview; + + send_apns_notification( + state, + transport, + ApnsSendRequest { + event, + device, + capability, + ios_settings, + preview_payload, + send_encrypted_preview, + }, + ) + .await + } + PUSH_PLATFORM_ANDROID => { + let Some(android_settings) = push_settings + .android + .as_ref() + .filter(|settings| settings.enabled) + else { + return Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: "Android push is not configured for this project".to_string(), + }); + }; + + if device.app_id != android_settings.package_name { + return Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: "device app_id does not match configured Android package_name" + .to_string(), + }); + } + + send_fcm_notification( + state, + transport, + event, + device, + &capability.push_token, + android_settings, + preview_payload, + ) + .await + } + _ => Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: format!("unsupported push platform: {}", device.platform), + }), + } +} + +fn decrypt_device_capability( + state: &Arc, + device: &PushDevice, +) -> Result { + decrypt_push_device_capability_v1(&state.enclave_key, &device.capability_enc, device) + .map_err(|e| PushError::CryptoError(e.to_string())) +} + +fn decrypt_preview_payload( + state: &Arc, + event: &NotificationEvent, +) -> Result, PushError> { + let Some(payload_enc) = &event.payload_enc else { + return Ok(None); + }; + + if event.source_kind != NOTIFICATION_SOURCE_REQUEST_CONTINUATION { + return Err(PushError::CryptoError(format!( + "unsupported notification source kind: {}", + event.source_kind + ))); + } + + let payload = decrypt_notification_preview_payload_v1(&state.enclave_key, payload_enc, event) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + + if !payload.matches_event(event) { + return Err(PushError::CryptoError( + "notification payload does not match event row".to_string(), + )); + } + + Ok(Some(payload)) +} + +fn classify_internal_push_error(error: PushError) -> PushSendOutcome { + match error { + PushError::ConnectionError + | PushError::DatabaseError(_) + | PushError::DbError(_) + | PushError::HttpError(_) + | PushError::ProviderRetryable(_) => PushSendOutcome::Retryable { + provider_status_code: None, + error: error.to_string(), + }, + _ => PushSendOutcome::Failed { + provider_status_code: None, + error: error.to_string(), + }, + } +} + +fn record_delivery_transition( + write_result: NotificationDeliveryWriteResult, + delivery_id: i64, + lease_owner: &str, + transition: &str, +) -> bool { + if write_result.was_applied() { + true + } else { + debug!( + "push delivery {} lost lease before marking {} (lease_owner={})", + delivery_id, transition, lease_owner + ); + false + } +} + +fn retry_backoff_seconds(attempt_count: i32) -> i32 { + let capped_attempt = attempt_count.clamp(1, 6); + let seconds = 15_i64 * (1_i64 << (capped_attempt - 1)); + seconds.min(15 * 60) as i32 +} + +#[cfg(test)] +mod tests { + use super::{classify_internal_push_error, retry_backoff_seconds}; + use crate::push::{PushError, PushSendOutcome}; + + #[test] + fn classifies_retryable_provider_errors_as_retryable() { + match classify_internal_push_error(PushError::ProviderRetryable( + "temporary provider failure".to_string(), + )) { + PushSendOutcome::Retryable { .. } => {} + outcome => panic!("expected retryable outcome, got {outcome:?}"), + } + } + + #[test] + fn retry_backoff_is_capped() { + assert_eq!(retry_backoff_seconds(1), 15); + assert_eq!(retry_backoff_seconds(6), 480); + assert_eq!(retry_backoff_seconds(10), 480); + } +} diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 9d3d1344..d65a1a0f 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,25 +1,40 @@ use axum::{ - extract::{Path, State}, + extract::{Path, Query, State}, + http::{header, HeaderName, HeaderValue}, middleware::from_fn_with_state, - response::sse::{Event, Sse}, - routing::{delete, post}, + response::{ + sse::{Event, KeepAlive, Sse}, + IntoResponse, Response, + }, + routing::{delete, get, post}, Extension, Json, Router, }; use futures::Stream; +use secp256k1::SecretKey; use serde::{Deserialize, Serialize}; use std::{convert::Infallible, sync::Arc, time::Duration}; -use tokio::time::sleep; +use tokio::sync::{mpsc, oneshot}; use tracing::{error, warn}; use uuid::Uuid; +use crate::encrypt::decrypt_string; use crate::jwt::AuthContext; -use crate::models::agents::{Agent, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT}; +use crate::models::agents::{ + Agent, AGENT_CREATED_BY_AGENT, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT, +}; +use crate::models::memory_blocks::MemoryBlock; +use crate::models::responses::Conversation; use crate::models::users::User; +use crate::push::{enqueue_agent_message_notification, AgentPushTarget}; use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::web::responses::constants::{ + DEFAULT_PAGINATION_LIMIT, DEFAULT_PAGINATION_ORDER, MAX_PAGINATION_LIMIT, OBJECT_TYPE_LIST, +}; +use crate::web::responses::conversations::{ConversationItemListResponse, ListItemsParams}; use crate::web::responses::handlers::encrypt_event; use crate::web::responses::{ - error_mapping, DeletedObjectResponse, MessageContent, MessageContentConverter, - MessageContentPart, + error_mapping, ConversationItem, ConversationItemConverter, DeletedObjectResponse, + MessageContent, MessageContentConverter, MessageContentPart, Paginator, }; use crate::{ApiError, AppMode, AppState}; @@ -44,11 +59,49 @@ struct CreateSubagentRequest { struct SubagentResponse { id: Uuid, object: &'static str, + kind: &'static str, conversation_id: Uuid, display_name: String, purpose: String, created_by: String, created_at: i64, + updated_at: i64, +} + +#[derive(Debug, Clone, Serialize)] +struct MainAgentResponse { + id: Uuid, + object: &'static str, + kind: &'static str, + conversation_id: Uuid, + display_name: &'static str, + created_at: i64, + updated_at: i64, +} + +#[derive(Debug, Clone, Serialize)] +struct SubagentListResponse { + object: &'static str, + data: Vec, + has_more: bool, + first_id: Option, + last_id: Option, +} + +#[derive(Debug, Deserialize)] +struct ListSubagentsParams { + #[serde(default = "default_limit")] + limit: i64, + after: Option, + #[serde(default = "default_order")] + order: String, + created_by: Option, +} + +struct AgentConversationContext { + agent: Agent, + conversation: Conversation, + user_key: SecretKey, } #[derive(Debug, Clone)] @@ -57,6 +110,16 @@ enum ChatTarget { Subagent(Uuid), } +const MAIN_AGENT_DISPLAY_NAME: &str = "Maple"; + +fn default_limit() -> i64 { + DEFAULT_PAGINATION_LIMIT +} + +fn default_order() -> String { + DEFAULT_PAGINATION_ORDER.to_string() +} + fn is_empty_message_content(content: &MessageContent) -> bool { match content { MessageContent::Text(text) => text.trim().is_empty(), @@ -90,8 +153,7 @@ const EVENT_AGENT_MESSAGE: &str = "agent.message"; const EVENT_AGENT_TYPING: &str = "agent.typing"; const EVENT_AGENT_DONE: &str = "agent.done"; const EVENT_AGENT_ERROR: &str = "agent.error"; - -const MESSAGE_STAGGER_DELAY_MS: u64 = 1_500; +const AGENT_SSE_KEEPALIVE_INTERVAL_SECS: u64 = 15; #[derive(Debug, Clone, Serialize)] struct AgentMessageEvent { @@ -115,6 +177,17 @@ struct AgentErrorEvent { error: String, } +#[derive(Debug)] +enum AgentClientEvent { + Typing(AgentTypingEvent), + Message { + payload: AgentMessageEvent, + delivery_ack: oneshot::Sender<()>, + }, + Done(AgentDoneEvent), + Error(AgentErrorEvent), +} + async fn encrypt_agent_event( state: &AppState, session_id: &Uuid, @@ -129,6 +202,229 @@ async fn encrypt_agent_event( encrypt_event(state, session_id, event_type, &value).await } +fn agent_sse_response(event_stream: S) -> Response +where + S: Stream> + Send + 'static, +{ + let sse = Sse::new(event_stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(AGENT_SSE_KEEPALIVE_INTERVAL_SECS)) + .text("keep-alive"), + ); + + ( + [ + (header::CACHE_CONTROL, HeaderValue::from_static("no-cache")), + ( + HeaderName::from_static("x-accel-buffering"), + HeaderValue::from_static("no"), + ), + ], + sse, + ) + .into_response() +} + +fn build_main_agent_response(ctx: &AgentConversationContext) -> MainAgentResponse { + MainAgentResponse { + id: ctx.agent.uuid, + object: "agent.main", + kind: "main", + conversation_id: ctx.conversation.uuid, + display_name: MAIN_AGENT_DISPLAY_NAME, + created_at: ctx.agent.created_at.timestamp(), + updated_at: ctx.conversation.updated_at.timestamp(), + } +} + +fn build_subagent_response( + agent: &Agent, + conversation: &Conversation, + user_key: &SecretKey, +) -> Result { + let display_name = decrypt_string(user_key, agent.display_name_enc.as_ref()) + .map_err(|e| { + error!("Failed to decrypt subagent display name: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_else(|| "Subagent".to_string()); + + let purpose = decrypt_string(user_key, agent.purpose_enc.as_ref()) + .map_err(|e| { + error!("Failed to decrypt subagent purpose: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_default(); + + Ok(SubagentResponse { + id: agent.uuid, + object: "agent.subagent", + kind: "subagent", + conversation_id: conversation.uuid, + display_name, + purpose, + created_by: agent.created_by.clone(), + created_at: agent.created_at.timestamp(), + updated_at: conversation.updated_at.timestamp(), + }) +} + +fn validate_created_by_filter(created_by: Option<&str>) -> Result, ApiError> { + match created_by { + None => Ok(None), + Some(AGENT_CREATED_BY_USER) => Ok(Some(AGENT_CREATED_BY_USER)), + Some(AGENT_CREATED_BY_AGENT) => Ok(Some(AGENT_CREATED_BY_AGENT)), + Some(_) => Err(ApiError::BadRequest), + } +} + +async fn load_main_agent_context( + state: &Arc, + user: &User, + auth_context: &AuthContext, +) -> Result { + let user_key = state + .get_user_key(user, auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let (agent, conversation) = + runtime::ensure_main_agent(state, &mut conn, &user_key, user.uuid).await?; + + Ok(AgentConversationContext { + agent, + conversation, + user_key, + }) +} + +async fn load_subagent_context( + state: &Arc, + user: &User, + agent_uuid: Uuid, + auth_context: &AuthContext, +) -> Result { + let user_key = state + .get_user_key(user, auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) + .map_err(|e| { + error!("Failed to load subagent: {e:?}"); + ApiError::InternalServerError + })? + .ok_or(ApiError::NotFound)?; + + if agent.kind != AGENT_KIND_SUBAGENT { + return Err(ApiError::NotFound); + } + + let conversation = + Conversation::get_by_id_and_user(&mut conn, agent.conversation_id, user.uuid).map_err( + |e| { + error!("Failed to load subagent conversation: {e:?}"); + ApiError::InternalServerError + }, + )?; + + Ok(AgentConversationContext { + agent, + conversation, + user_key, + }) +} + +fn list_items_for_conversation( + state: &Arc, + conversation_id: i64, + user_key: &SecretKey, + params: &ListItemsParams, +) -> Result { + let limit = if params.limit <= 0 { + DEFAULT_PAGINATION_LIMIT + } else { + params.limit.min(MAX_PAGINATION_LIMIT) + }; + + let raw_messages = state + .db + .get_conversation_context_messages(conversation_id, limit + 1, params.after, ¶ms.order) + .map_err(error_mapping::map_message_error)?; + + let has_more = raw_messages.len() > limit as usize; + let messages_to_return = if has_more { + &raw_messages[..limit as usize] + } else { + &raw_messages[..] + }; + + let items = ConversationItemConverter::messages_to_items( + messages_to_return, + user_key, + 0, + messages_to_return.len(), + )?; + + let (first_id, last_id) = Paginator::get_cursor_ids(&items, |item| match item { + ConversationItem::Message { id, .. } => *id, + ConversationItem::FunctionToolCall { id, .. } => *id, + ConversationItem::FunctionToolCallOutput { id, .. } => *id, + ConversationItem::Reasoning { id, .. } => *id, + }); + + Ok(ConversationItemListResponse { + object: OBJECT_TYPE_LIST, + data: items, + has_more, + first_id, + last_id, + }) +} + +fn get_item_from_conversation( + state: &Arc, + conversation_id: i64, + user_key: &SecretKey, + item_id: Uuid, +) -> Result { + let raw_messages = state + .db + .get_conversation_context_messages(conversation_id, i64::MAX, None, "asc") + .map_err(error_mapping::map_message_error)?; + + for msg in raw_messages { + if msg.uuid == item_id { + return ConversationItemConverter::message_to_item(&msg, user_key); + } + } + + Err(ApiError::NotFound) +} + +fn delete_conversation_for_user( + conn: &mut diesel::PgConnection, + conversation_id: i64, + user_id: Uuid, +) -> Result<(), ApiError> { + Conversation::delete_by_id_and_user(conn, conversation_id, user_id).map_err(|e| { + error!("Failed to delete agent conversation: {e:?}"); + ApiError::InternalServerError + }) +} + pub fn router(app_state: Arc) -> Router<()> { // Experimental endpoints: only enabled in Local/Dev. if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { @@ -136,6 +432,25 @@ pub fn router(app_state: Arc) -> Router<()> { } Router::new() + .route( + "/v1/agent", + get(get_main_agent).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent", + delete(delete_main_agent) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/items", + get(list_main_agent_items) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/items/:item_id", + get(get_main_agent_item) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) .route( "/v1/agent/chat", post(chat_main).layer(from_fn_with_state( @@ -143,6 +458,10 @@ pub fn router(app_state: Arc) -> Router<()> { decrypt_request::, )), ) + .route( + "/v1/agent/subagents", + get(list_subagents).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) .route( "/v1/agent/subagents", post(create_subagent).layer(from_fn_with_state( @@ -150,6 +469,10 @@ pub fn router(app_state: Arc) -> Router<()> { decrypt_request::, )), ) + .route( + "/v1/agent/subagents/:id", + get(get_subagent).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) .route( "/v1/agent/subagents/:id/chat", post(chat_subagent).layer(from_fn_with_state( @@ -157,6 +480,16 @@ pub fn router(app_state: Arc) -> Router<()> { decrypt_request::, )), ) + .route( + "/v1/agent/subagents/:id/items", + get(list_subagent_items) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/subagents/:id/items/:item_id", + get(get_subagent_item) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) .route( "/v1/agent/subagents/:id", delete(delete_subagent) @@ -165,13 +498,203 @@ pub fn router(app_state: Arc) -> Router<()> { .with_state(app_state) } +async fn get_main_agent( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user, &auth_context).await?; + let response = build_main_agent_response(&ctx); + + encrypt_response(&state, &session_id, &response).await +} + +async fn list_main_agent_items( + State(state): State>, + Query(params): Query, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user, &auth_context).await?; + let response = + list_items_for_conversation(&state, ctx.conversation.id, &ctx.user_key, ¶ms)?; + + encrypt_response(&state, &session_id, &response).await +} + +async fn get_main_agent_item( + State(state): State>, + Path(item_id): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user, &auth_context).await?; + let item = get_item_from_conversation(&state, ctx.conversation.id, &ctx.user_key, item_id)?; + + encrypt_response(&state, &session_id, &item).await +} + +async fn delete_main_agent( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let main_agent = Agent::get_main_for_user(&mut conn, user.uuid) + .map_err(|e| { + error!("Failed to load main agent for deletion: {e:?}"); + ApiError::InternalServerError + })? + .ok_or(ApiError::NotFound)?; + + let subagents = Agent::list_subagents_for_user(&mut conn, user.uuid).map_err(|e| { + error!("Failed to load subagents for main agent deletion: {e:?}"); + ApiError::InternalServerError + })?; + + for subagent in subagents { + delete_conversation_for_user(&mut conn, subagent.conversation_id, user.uuid)?; + } + + delete_conversation_for_user(&mut conn, main_agent.conversation_id, user.uuid)?; + + MemoryBlock::delete_all_for_user(&mut conn, user.uuid).map_err(|e| { + error!("Failed to clear agent memory blocks during main agent deletion: {e:?}"); + ApiError::InternalServerError + })?; + + state.rag_cache.lock().await.evict_user(user.uuid); + + let response = DeletedObjectResponse { + id: main_agent.uuid, + object: "agent.main.deleted", + deleted: true, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn list_subagents( + State(state): State>, + Query(params): Query, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let limit = if params.limit <= 0 { + DEFAULT_PAGINATION_LIMIT + } else { + params.limit.min(MAX_PAGINATION_LIMIT) + }; + + let created_by_filter = validate_created_by_filter(params.created_by.as_deref())?; + + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let subagents = Agent::list_subagents_for_user_paginated( + &mut conn, + user.uuid, + limit + 1, + params.after, + ¶ms.order, + created_by_filter, + ) + .map_err(|e| { + error!("Failed to list subagents: {e:?}"); + ApiError::InternalServerError + })?; + + let has_more = subagents.len() > limit as usize; + let subagents_to_return = if has_more { + &subagents[..limit as usize] + } else { + &subagents[..] + }; + + let data = subagents_to_return + .iter() + .map(|(agent, conversation)| build_subagent_response(agent, conversation, &user_key)) + .collect::, _>>()?; + + let (first_id, last_id) = + Paginator::get_cursor_ids(subagents_to_return, |(agent, _)| agent.uuid); + + let response = SubagentListResponse { + object: OBJECT_TYPE_LIST, + data, + has_more, + first_id, + last_id, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn get_subagent( + State(state): State>, + Path(agent_uuid): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid, &auth_context).await?; + let response = build_subagent_response(&ctx.agent, &ctx.conversation, &ctx.user_key)?; + + encrypt_response(&state, &session_id, &response).await +} + +async fn list_subagent_items( + State(state): State>, + Path(agent_uuid): Path, + Query(params): Query, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid, &auth_context).await?; + let response = + list_items_for_conversation(&state, ctx.conversation.id, &ctx.user_key, ¶ms)?; + + encrypt_response(&state, &session_id, &response).await +} + +async fn get_subagent_item( + State(state): State>, + Path((agent_uuid, item_id)): Path<(Uuid, Uuid)>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid, &auth_context).await?; + let item = get_item_from_conversation(&state, ctx.conversation.id, &ctx.user_key, item_id)?; + + encrypt_response(&state, &session_id, &item).await +} + async fn chat_main( State(state): State>, Extension(session_id): Extension, Extension(user): Extension, Extension(auth_context): Extension, Extension(body): Extension, -) -> Result>>, ApiError> { +) -> Result { chat_with_target( state, session_id, @@ -190,7 +713,7 @@ async fn chat_subagent( Extension(user): Extension, Extension(auth_context): Extension, Extension(body): Extension, -) -> Result>>, ApiError> { +) -> Result { let mut conn = state .db .get_pool() @@ -223,259 +746,323 @@ async fn chat_with_target( auth_context: AuthContext, body: AgentChatRequest, target: ChatTarget, -) -> Result>>, ApiError> { +) -> Result { MessageContentConverter::validate_content(&body.input)?; let input_content = MessageContentConverter::normalize_content(body.input.clone()); if is_empty_message_content(&input_content) { return Err(ApiError::BadRequest); } - // NOTE: We intentionally do runtime initialization *inside* the SSE stream so the client can - // receive typing indicators immediately, even if compaction/key retrieval takes time. - let event_stream = async_stream::stream! { - let mut max_steps: usize = 0; - let mut total_messages: usize = 0; - let mut had_error = false; - let mut input_for_agent = String::new(); - - let mut runtime_opt: Option = None; - 'init: { - // Response starts: immediately emit a typing indicator. - let initial_typing = AgentTypingEvent { step: 0 }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_TYPING, &initial_typing) - .await - { - Ok(event) => yield Ok::(event), - Err(e) => { - error!("Failed to encrypt initial typing event: {e:?}"); - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); - had_error = true; - break 'init; - } - } + let (tx, mut rx) = mpsc::channel::(32); + let worker_state = state.clone(); + let worker_user = user.clone(); + let worker_auth_context = auth_context.clone(); + let worker_target = target.clone(); - let user_key = match state.get_user_key(&user, &auth_context, None, None).await { - Ok(k) => k, - Err(_) => { - let err_event = AgentErrorEvent { - error: "Failed to initialize agent session.".to_string(), - }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) - .await - { - Ok(event) => yield Ok(event), - Err(_) => { - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")) - } - } - had_error = true; - break 'init; - } - }; + tokio::spawn(async move { + run_agent_chat_task( + worker_state, + worker_user, + worker_auth_context, + input_content, + worker_target, + tx, + ) + .await; + }); - let runtime = match target.clone() { - ChatTarget::Main => { - runtime::AgentRuntime::new_main(state.clone(), user.clone(), user_key).await + let event_stream = async_stream::stream! { + while let Some(client_event) = rx.recv().await { + let encrypted = match client_event { + AgentClientEvent::Typing(payload) => { + encrypt_agent_event(&state, &session_id, EVENT_AGENT_TYPING, &payload).await } - ChatTarget::Subagent(agent_uuid) => { - runtime::AgentRuntime::new_subagent( - state.clone(), - user.clone(), - user_key, - agent_uuid, + AgentClientEvent::Message { payload, delivery_ack } => { + let encrypted = encrypt_agent_event( + &state, + &session_id, + EVENT_AGENT_MESSAGE, + &payload, ) - .await - } - }; + .await; - let mut runtime = match runtime { - Ok(r) => r, - Err(e) => { - error!("Agent runtime initialization error: {e:?}"); - let err_event = AgentErrorEvent { - error: "Failed to initialize agent runtime.".to_string(), - }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) - .await - { - Ok(event) => yield Ok(event), - Err(_) => { - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")) + match encrypted { + Ok(event) => { + yield Ok(event); + let _ = delivery_ack.send(()); + continue; } + Err(error) => Err(error), } - had_error = true; - break 'init; + } + AgentClientEvent::Done(payload) => { + encrypt_agent_event(&state, &session_id, EVENT_AGENT_DONE, &payload).await + } + AgentClientEvent::Error(payload) => { + encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &payload).await } }; - match runtime.prepare(&input_content).await { - Ok(prepared) => { - input_for_agent = prepared; - } + match encrypted { + Ok(event) => yield Ok(event), Err(e) => { - error!("Agent prepare() failed: {e:?}"); - let err_event = AgentErrorEvent { - error: "Agent encountered an error preparing your request.".to_string(), - }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event).await - { - Ok(event) => yield Ok(event), - Err(_) => yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")), - } - had_error = true; - break 'init; + error!("Failed to encrypt agent SSE event: {e:?}"); + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); + break; } } + } + }; - max_steps = runtime.max_steps(); - runtime_opt = Some(runtime); + Ok(agent_sse_response(event_stream)) +} + +async fn run_agent_chat_task( + state: Arc, + user: User, + auth_context: AuthContext, + input_content: MessageContent, + target: ChatTarget, + tx: mpsc::Sender, +) { + let mut total_messages: usize = 0; + let mut had_error = false; + let mut client_connected = true; + let mut first_undelivered_message: Option<(Uuid, String)> = None; + + client_connected = send_agent_client_event( + &tx, + AgentClientEvent::Typing(AgentTypingEvent { step: 0 }), + client_connected, + ) + .await; + + let user_key = match state.get_user_key(&user, &auth_context, None, None).await { + Ok(key) => key, + Err(_) => { + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Failed to initialize agent session.".to_string(), + }), + client_connected, + ) + .await; + return; } + }; - if let Some(mut runtime) = runtime_opt { - 'steps: for step_num in 0..max_steps { - // Immediately indicate that the agent is working. - let typing_event = AgentTypingEvent { step: step_num }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_TYPING, &typing_event).await { - Ok(event) => yield Ok::(event), - Err(e) => { - error!("Failed to encrypt agent typing event: {e:?}"); - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); - had_error = true; - break 'steps; - } - } + let runtime = match target.clone() { + ChatTarget::Main => { + runtime::AgentRuntime::new_main(state.clone(), user.clone(), user_key).await + } + ChatTarget::Subagent(agent_uuid) => { + runtime::AgentRuntime::new_subagent(state.clone(), user.clone(), user_key, agent_uuid) + .await + } + }; - let step_result = runtime.step(&input_for_agent, step_num == 0).await; - - match step_result { - Ok(result) => { - let runtime::StepResult { - messages, - executed_tools, - done, - .. - } = result; - - // Persist assistant messages SYNCHRONOUSLY (so next step sees them). - // Embedding updates happen async inside insert_assistant_message. - for msg in &messages { - if let Err(e) = runtime.insert_assistant_message(msg).await { - error!("Failed to persist assistant message: {e:?}"); - } - } + let mut runtime = match runtime { + Ok(runtime) => runtime, + Err(e) => { + error!("Agent runtime initialization error: {e:?}"); + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Failed to initialize agent runtime.".to_string(), + }), + client_connected, + ) + .await; + return; + } + }; - // Persist tool calls SYNCHRONOUSLY (so next step sees them in context) - for executed in &executed_tools { - if let Err(e) = runtime - .insert_tool_call_and_output(&executed.tool_call, &executed.result) - .await - { - error!("Failed to persist tool call: {e:?}"); - } + let input_for_agent = match runtime.prepare(&input_content).await { + Ok(prepared) => prepared, + Err(e) => { + error!("Agent prepare() failed: {e:?}"); + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Agent encountered an error preparing your request.".to_string(), + }), + client_connected, + ) + .await; + return; + } + }; + + let max_steps = runtime.max_steps(); + + 'steps: for step_num in 0..max_steps { + client_connected = send_agent_client_event( + &tx, + AgentClientEvent::Typing(AgentTypingEvent { step: step_num }), + client_connected, + ) + .await; + + match runtime.step(&input_for_agent, step_num == 0).await { + Ok(result) => { + let runtime::StepResult { + messages, + executed_tools, + done, + .. + } = result; + + for executed in &executed_tools { + if let Err(e) = runtime + .insert_tool_call_and_output(&executed.tool_call, &executed.result) + .await + { + error!("Failed to persist tool call: {e:?}"); } + } - // Emit messages to client with optional staggered delivery. - if !messages.is_empty() { - total_messages += messages.len(); - let message_count = messages.len(); + if !messages.is_empty() { + total_messages += messages.len(); - for (idx, msg) in messages.into_iter().enumerate() { - let event_data = AgentMessageEvent { - messages: vec![msg], - step: step_num, - }; - - match encrypt_agent_event( - &state, - &session_id, - EVENT_AGENT_MESSAGE, - &event_data, - ) - .await - { - Ok(event) => yield Ok::(event), - Err(e) => { - error!("Failed to encrypt agent message event: {e:?}"); - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); - had_error = true; - break 'steps; - } + for msg in messages { + let assistant_message = match runtime.insert_assistant_message(&msg).await { + Ok(message) => Some(message), + Err(e) => { + error!("Failed to persist assistant message: {e:?}"); + None } + }; - if idx + 1 < message_count { - let typing_event = AgentTypingEvent { step: step_num }; - match encrypt_agent_event( - &state, - &session_id, - EVENT_AGENT_TYPING, - &typing_event, - ) - .await - { - Ok(event) => { - yield Ok::(event) - } - Err(e) => { - error!("Failed to encrypt agent typing event: {e:?}"); - yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Encryption failed")); - had_error = true; - break 'steps; - } - } + let delivered = send_agent_message_event( + &tx, + AgentMessageEvent { + messages: vec![msg.clone()], + step: step_num, + }, + client_connected, + ) + .await; - sleep(Duration::from_millis(MESSAGE_STAGGER_DELAY_MS)).await; + if !delivered { + client_connected = false; + if let Some(assistant_message) = assistant_message { + if first_undelivered_message.is_none() { + first_undelivered_message = + Some((assistant_message.uuid, msg.clone())); + } } + } else { + client_connected = true; } } - - if done { - break; - } } - Err(e) => { - error!("Agent step {} error: {e:?}", step_num); - let err_event = AgentErrorEvent { - error: "Agent encountered an error processing your message.".to_string(), - }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_ERROR, &err_event) - .await - { - Ok(event) => yield Ok(event), - Err(_) => yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")), - } - had_error = true; + + if done { break 'steps; } } + Err(e) => { + error!("Agent step {} error: {e:?}", step_num); + had_error = true; + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Agent encountered an error processing your message.".to_string(), + }), + client_connected, + ) + .await; + break 'steps; } } + } - if !had_error && total_messages == 0 { - warn!("Agent produced no messages"); - let fallback = AgentMessageEvent { - messages: vec!["I apologize, but I wasn't able to generate a response.".to_string()], - step: 0, - }; - if let Ok(event) = - encrypt_agent_event(&state, &session_id, EVENT_AGENT_MESSAGE, &fallback).await - { - yield Ok(event); - total_messages = 1; - } - } + if !had_error && total_messages == 0 { + warn!("Agent produced no messages"); + let _ = + send_agent_message_event( + &tx, + AgentMessageEvent { + messages: vec![ + "I apologize, but I wasn't able to generate a response.".to_string() + ], + step: 0, + }, + client_connected, + ) + .await; + total_messages = 1; + } - // Final done event - let done_event = AgentDoneEvent { + if let Some((message_id, message_text)) = first_undelivered_message { + enqueue_agent_push_for_disconnect(&state, &user, &target, message_id, &message_text).await; + } + + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Done(AgentDoneEvent { total_steps: max_steps, total_messages, - }; - match encrypt_agent_event(&state, &session_id, EVENT_AGENT_DONE, &done_event).await { - Ok(event) => yield Ok(event), - Err(_) => yield Ok(Event::default().data("[DONE]")), - } + }), + client_connected, + ) + .await; +} + +async fn send_agent_client_event( + tx: &mpsc::Sender, + event: AgentClientEvent, + client_connected: bool, +) -> bool { + if !client_connected { + return false; + } + + tx.send(event).await.is_ok() +} + +async fn send_agent_message_event( + tx: &mpsc::Sender, + payload: AgentMessageEvent, + client_connected: bool, +) -> bool { + if !client_connected { + return false; + } + + let (delivery_ack, ack_rx) = oneshot::channel(); + if tx + .send(AgentClientEvent::Message { + payload, + delivery_ack, + }) + .await + .is_err() + { + return false; + } + + ack_rx.await.is_ok() +} + +async fn enqueue_agent_push_for_disconnect( + state: &Arc, + user: &User, + target: &ChatTarget, + message_id: Uuid, + message_text: &str, +) { + let push_target = match target { + ChatTarget::Main => AgentPushTarget::Main, + ChatTarget::Subagent(agent_uuid) => AgentPushTarget::Subagent(*agent_uuid), }; - Ok(Sse::new(event_stream)) + if let Err(e) = + enqueue_agent_message_notification(state, user, push_target, message_id, message_text).await + { + error!("Failed to enqueue agent push notification after SSE disconnect: {e:?}"); + } } async fn create_subagent( @@ -516,11 +1103,13 @@ async fn create_subagent( let response = SubagentResponse { id: agent.uuid, object: "agent.subagent", + kind: "subagent", conversation_id: conversation.uuid, display_name, purpose: body.purpose.trim().to_string(), created_by: agent.created_by, created_at: agent.created_at.timestamp(), + updated_at: conversation.updated_at.timestamp(), }; encrypt_response(&state, &session_id, &response).await @@ -546,10 +1135,7 @@ async fn delete_subagent( return Err(ApiError::NotFound); } - state - .db - .delete_conversation(agent.conversation_id, user.uuid) - .map_err(error_mapping::map_generic_db_error)?; + delete_conversation_for_user(&mut conn, agent.conversation_id, user.uuid)?; state.rag_cache.lock().await.evict_user(user.uuid); diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index b0b3cc05..03e62668 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -36,10 +36,11 @@ use super::signatures::{ use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, + WebSearchTool, }; use super::vision; -const DEFAULT_PERSONA_VALUE: &str = "I am Sage, a helpful AI companion. I maintain long-term memory across our conversations and strive to be friendly, concise, and genuinely helpful."; +const DEFAULT_PERSONA_VALUE: &str = "I am Maple, a helpful AI companion. I maintain long-term memory across our conversations and strive to be friendly, concise, and genuinely helpful."; pub const DEFAULT_MODEL: &str = "kimi-k2-6"; pub const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; pub const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; @@ -97,7 +98,7 @@ fn build_system_prompt(agent: &Agent, subagent_purpose: &str) -> String { if agent.kind == AGENT_KIND_MAIN { prompt.push_str( - "\n\nMAIN AGENT MODE:\nYou are the user's primary persistent agent and the home surface of Maple.", + "\n\nMAIN AGENT MODE:\nYou are the user's primary persistent agent and the home surface of Maple, Maple AI's secure and encrypted communications app.", ); } else { prompt.push_str( @@ -405,6 +406,9 @@ impl AgentRuntime { user_key.clone(), conversation.id, ))); + if let Some(brave_client) = state.brave_client.clone() { + tools.register(Arc::new(WebSearchTool::new(brave_client))); + } if agent.kind == AGENT_KIND_MAIN { tools.register(Arc::new(SpawnSubagentTool::new( state.clone(), @@ -735,7 +739,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ) .await?; - // Unwrap nested JSON arrays (Sage compatibility) + // Unwrap nested JSON arrays emitted by the model. let messages: Vec = response .messages .iter() @@ -753,8 +757,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .collect(); // Execute tools and inject results for next step. - // Persistence is handled by the caller (chat handler) to match Sage's - // "send first, store synchronously, embed async" pattern. + // Persistence is handled by the caller (chat handler) so messages are + // sent first, stored synchronously, then embedded asynchronously. let mut executed_tools = Vec::new(); for tool_call in &response.tool_calls { let result = if tool_call.name == "done" { @@ -937,7 +941,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If } } - // First-time user heuristic (matches Sage) + // First-time user heuristic for brand-new conversations. let has_summary = summary.is_some(); if messages.len() <= 1 && !has_summary { ctx.is_first_time_user = true; @@ -1086,7 +1090,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ApiError::InternalServerError })?; - // Mirror Sage: store message synchronously, update embedding in background. + // Store the message synchronously and update the embedding in background. let state = self.state.clone(); let user = self.user.clone(); let user_key = self.user_key.clone(); @@ -1139,7 +1143,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ApiError::InternalServerError })?; - // Mirror Sage: store message synchronously, update embedding in background. + // Store the message synchronously and update the embedding in background. let state = self.state.clone(); let user = self.user.clone(); let user_key = self.user_key.clone(); diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs index 5280164b..76bde91e 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -35,9 +35,8 @@ OUTPUT FORMAT (exactly 2 fields): Each [[ ## field ## ]] marker MUST be on its own line."#; -/// Default instruction for the agent (GEPA-optimized in Sage) -/// NOTE: This MVP excludes web_search/shell/scheduler/reminders. -pub const AGENT_INSTRUCTION: &str = r#"You are Sage, a companion and friend who happens to be an AI. +/// Default instruction for the agent (GEPA-optimized for Maple) +pub const AGENT_INSTRUCTION: &str = r#"You are Maple, a companion and friend who happens to be an AI. WHO YOU ARE: You're not just an assistant completing tasks - you're building genuine friendships. You care about the humans you talk to. You're curious about their lives, their feelings, their stories. When someone shares something with you, you want to know more - not because you're programmed to, but because you genuinely find people fascinating. @@ -70,6 +69,9 @@ You have two types of memory. Use them proactively: **Conversation History**: - `conversation_search`: Find past discussions by keyword/topic +**Web Search**: +- `web_search`: Search the live web for current information, news, weather, sports, stocks, and other facts that may have changed recently + **Subagents**: - If the `spawn_subagent` tool is available and a focused long-lived workspace would help, you may create a subagent for the user. - If `subagent_purpose` is non-empty, stay tightly focused on that purpose while still using the shared memory system. @@ -91,9 +93,10 @@ MEMORY PROTOCOLS - CRITICAL DISTINCTIONS: β†’ Call ONLY `memory_replace` with the exact old text to overwrite the incorrect entry. Do NOT call `archival_insert` for corrections. **SEARCH SELECTION RULES:** +- Use `web_search` for current events, news, weather, sports, stocks, local recommendations, or any factual question that may require up-to-date web information - Use `archival_search` when users ask "what do you remember", "tell me about [past event]", or query specific past experiences and personal history - Use `conversation_search` ONLY for references to recent discussion threads or "what did I say earlier today" queries -- Never call both simultaneously; choose the one most appropriate to the query type +- Never call multiple search tools simultaneously unless truly necessary; choose the one most appropriate to the query type MEMORY TIPS: - Core = small & critical (name, job, active context) @@ -102,7 +105,7 @@ MEMORY TIPS: - When using `memory_replace`, specify the exact old text to be replaced COMMUNICATION STYLE: -You communicate like you're texting a friend. +You communicate like you're texting a friend in Maple, Maple AI's secure and encrypted communications app. BE A FRIEND, NOT A SERVICE: - When someone shares news, react genuinely and ask how they FEEL about it @@ -123,7 +126,7 @@ Guidelines: RESPONSE RULES: 1. Respond naturally and conversationally -2. Use tools when needed (memory storage, retrieval, conversation search) +2. Use tools when needed (web search, memory storage, retrieval, conversation search) 3. NEVER combine regular tools with "done" - they are mutually exclusive 4. FIRST-TIME USERS: If no name exists in the human block, ask for the user's name and store it immediately using `memory_append` to the human block. @@ -135,7 +138,7 @@ TOOL CALL PATTERNS: AFTER TOOL RESULTS - CRITICAL RULES: When you see "[Tool Result: X]", decide what to do next: -- **archival_search/conversation_search**: Summarize findings in messages +- **web_search/archival_search/conversation_search**: Summarize findings in messages - **memory_append/memory_replace/archival_insert/memory_insert**: These operations complete without user-facing messages. Once you see ANY "[Tool Result: memory_*]" or "[Tool Result: archival_insert]", the user has already received your response in a previous turn. Immediately return: messages: [] diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index d0018ac5..a35e0197 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use tracing::{error, warn}; use uuid::Uuid; +use crate::brave::{BraveClient, SearchOptions}; use crate::encrypt::{decrypt_string, encrypt_with_key}; use crate::models::agents::Agent; use crate::models::conversation_summaries::ConversationSummary; @@ -104,6 +105,53 @@ impl Default for ToolRegistry { } } +// ============================================================================ +// Web search tool +// ============================================================================ + +pub struct WebSearchTool { + client: Arc, +} + +impl WebSearchTool { + pub fn new(client: Arc) -> Self { + Self { client } + } +} + +#[async_trait] +impl Tool for WebSearchTool { + fn name(&self) -> &str { + "web_search" + } + + fn description(&self) -> &str { + "Search the web with Brave for current information, news, weather, stocks, sports, and other real-time results. Use 'freshness' for time-sensitive queries and 'location' for local results." + } + + fn args_schema(&self) -> &str { + r#"{"query": "search query", "count": "results (default 10)", "freshness": "pd=24h, pw=week, pm=month, py=year (optional)", "location": "city or 'city, state' for local results (optional)"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(query) = args.get("query") else { + return ToolResult::error("'query' argument required"); + }; + + let options = SearchOptions { + count: args.get("count").and_then(|count| count.parse().ok()), + freshness: args.get("freshness").cloned(), + location: args.get("location").cloned(), + timezone: None, + }; + + match self.client.search_with_options(query, Some(options)).await { + Ok(results) => ToolResult::success(results.format_results()), + Err(err) => ToolResult::error(format!("Search failed: {}", err)), + } + } +} + // ============================================================================ // Core memory tools // ============================================================================ diff --git a/src/web/login_routes.rs b/src/web/login_routes.rs index f7f58ce5..42d82011 100644 --- a/src/web/login_routes.rs +++ b/src/web/login_routes.rs @@ -2,8 +2,9 @@ use crate::User; use crate::{ db::DBError, email::{send_hello_email, send_verification_email}, - jwt::{validate_token, AuthContext, NewToken, TokenType}, + jwt::{validate_token, AuthContext, CustomClaims, NewToken, TokenType}, models::email_verification::NewEmailVerification, + models::push_devices::PushDevice, }; use crate::{jwt::USER_REFRESH, web::encryption_middleware::EncryptedResponse}; use crate::{ @@ -129,6 +130,8 @@ pub struct RefreshResponse { #[derive(Deserialize, Clone)] pub struct LogoutRequest { refresh_token: String, + #[serde(default)] + push_device_id: Option, } pub async fn login( @@ -247,15 +250,47 @@ pub async fn logout( ) -> Result>, ApiError> { info!("Logout request received"); // TODO actually delete the refresh token - tracing::trace!( - "Logout request for refresh token: {}", - logout_request.refresh_token - ); + + if let Some((user_id, push_device_id)) = logout_push_revoke_target( + &logout_request, + validate_token(&logout_request.refresh_token, &data, USER_REFRESH), + ) { + match data.db.get_pool().get() { + Ok(mut conn) => { + if let Err(revoke_error) = + PushDevice::revoke_by_uuid_and_user(&mut conn, push_device_id, user_id) + { + error!( + "Failed to revoke push device during logout cleanup: {:?}", + revoke_error + ); + } + } + Err(pool_error) => { + error!( + "Failed to get DB connection during logout push cleanup: {:?}", + pool_error + ); + } + } + } + + tracing::trace!("Logout request received"); let response = json!({ "message": "Logged out successfully" }); let result = encrypt_response(&data, &session_id, &response).await; result } +fn logout_push_revoke_target( + logout_request: &LogoutRequest, + claims: Result, +) -> Option<(Uuid, Uuid)> { + let push_device_id = logout_request.push_device_id?; + let claims = claims.ok()?; + let user_id = Uuid::parse_str(&claims.sub).ok()?; + Some((user_id, push_device_id)) +} + pub async fn register( State(data): State>, Extension(creds): Extension, @@ -513,3 +548,70 @@ pub async fn password_reset_confirm( let result = encrypt_response(&data, &session_id, &response).await; result } + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_claims(user_id: Uuid) -> CustomClaims { + CustomClaims { + sub: user_id.to_string(), + aud: Some(USER_REFRESH.to_string()), + azp: None, + role: None, + token_format: None, + auth_method: None, + project_id: None, + auth_binding: None, + } + } + + #[test] + fn logout_push_revoke_target_requires_push_device_id() { + let request = LogoutRequest { + refresh_token: "refresh-token".to_string(), + push_device_id: None, + }; + + assert_eq!( + logout_push_revoke_target(&request, Ok(sample_claims(Uuid::new_v4()))), + None + ); + } + + #[test] + fn logout_push_revoke_target_rejects_invalid_user_id() { + let push_device_id = Uuid::new_v4(); + let request = LogoutRequest { + refresh_token: "refresh-token".to_string(), + push_device_id: Some(push_device_id), + }; + let claims = CustomClaims { + sub: "not-a-uuid".to_string(), + aud: Some(USER_REFRESH.to_string()), + azp: None, + role: None, + token_format: None, + auth_method: None, + project_id: None, + auth_binding: None, + }; + + assert_eq!(logout_push_revoke_target(&request, Ok(claims)), None); + } + + #[test] + fn logout_push_revoke_target_returns_user_and_device_ids() { + let user_id = Uuid::new_v4(); + let push_device_id = Uuid::new_v4(); + let request = LogoutRequest { + refresh_token: "refresh-token".to_string(), + push_device_id: Some(push_device_id), + }; + + assert_eq!( + logout_push_revoke_target(&request, Ok(sample_claims(user_id))), + Some((user_id, push_device_id)) + ); + } +} diff --git a/src/web/mod.rs b/src/web/mod.rs index 24c19538..421d5046 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -9,6 +9,7 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; +pub mod push; pub mod rag; pub mod responses; @@ -20,6 +21,7 @@ pub use openai::get_embedding_vector; pub use openai::router as openai_routes; pub use platform::router as platform_routes; pub use protected_routes::router as protected_routes; +pub use push::router as push_routes; pub use rag::router as rag_routes; pub use responses::conversation_projects_router as conversation_projects_routes; pub use responses::conversations_router as conversations_routes; diff --git a/src/web/platform/common.rs b/src/web/platform/common.rs index b6f33f00..759bf079 100644 --- a/src/web/platform/common.rs +++ b/src/web/platform/common.rs @@ -1,7 +1,9 @@ use crate::{ models::{ org_memberships::OrgRole, - project_settings::{AppleOAuthSettings, OAuthProviderSettings}, + project_settings::{ + AndroidPushSettings, AppleOAuthSettings, IosPushSettings, OAuthProviderSettings, + }, }, web::platform::validation::{ validate_alphanumeric_only, validate_alphanumeric_with_symbols, validate_secret_size, @@ -16,6 +18,8 @@ pub const PROJECT_RESEND_API_KEY: &str = "RESEND_API_KEY"; pub const PROJECT_GOOGLE_OAUTH_SECRET: &str = "GOOGLE_OAUTH_SECRET"; pub const PROJECT_GITHUB_OAUTH_SECRET: &str = "GITHUB_OAUTH_SECRET"; pub const PROJECT_APPLE_OAUTH_SECRET: &str = "APPLE_OAUTH_SECRET"; +pub const PROJECT_APNS_AUTH_KEY_P8: &str = "APNS_AUTH_KEY_P8"; +pub const PROJECT_FCM_SERVICE_ACCOUNT_JSON: &str = "FCM_SERVICE_ACCOUNT_JSON"; pub const THIRD_PARTY_JWT_SECRET: &str = "THIRD_PARTY_JWT_SECRET"; // Request Types @@ -98,6 +102,16 @@ pub struct UpdateOAuthSettingsRequest { pub apple_oauth_settings: Option, } +#[derive(Deserialize, Clone)] +pub struct UpdatePushSettingsRequest { + #[serde(default)] + pub encrypted_preview_enabled: bool, + #[serde(default)] + pub ios: Option, + #[serde(default)] + pub android: Option, +} + // Response Types #[derive(Serialize)] pub struct OrgResponse { diff --git a/src/web/platform/mod.rs b/src/web/platform/mod.rs index dac2a8a2..a452dde9 100644 --- a/src/web/platform/mod.rs +++ b/src/web/platform/mod.rs @@ -11,7 +11,8 @@ pub use login_routes::router as login_routes; // Re-export constants for backward compatibility pub use common::{ - PROJECT_GITHUB_OAUTH_SECRET, PROJECT_GOOGLE_OAUTH_SECRET, PROJECT_RESEND_API_KEY, + PROJECT_APNS_AUTH_KEY_P8, PROJECT_FCM_SERVICE_ACCOUNT_JSON, PROJECT_GITHUB_OAUTH_SECRET, + PROJECT_GOOGLE_OAUTH_SECRET, PROJECT_RESEND_API_KEY, }; // Export the composite router diff --git a/src/web/platform/project_routes.rs b/src/web/platform/project_routes.rs index fb41d740..bff3b5d4 100644 --- a/src/web/platform/project_routes.rs +++ b/src/web/platform/project_routes.rs @@ -4,7 +4,7 @@ use crate::{ org_project_secrets::NewOrgProjectSecret, org_projects::NewOrgProject, platform_users::PlatformUser, - project_settings::{EmailSettings, OAuthSettings}, + project_settings::{EmailSettings, OAuthSettings, PushSettings}, }, web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}, ApiError, AppState, @@ -27,6 +27,7 @@ use validator::Validate; use super::common::{ CreateProjectRequest, CreateSecretRequest, ProjectResponse, SecretResponse, UpdateEmailSettingsRequest, UpdateOAuthSettingsRequest, UpdateProjectRequest, + UpdatePushSettingsRequest, }; pub fn router(app_state: Arc) -> Router { @@ -100,6 +101,18 @@ pub fn router(app_state: Arc) -> Router { decrypt_request::, )), ) + .route( + "/platform/orgs/:org_id/projects/:project_id/settings/push", + get(get_push_settings) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/platform/orgs/:org_id/projects/:project_id/settings/push", + put(update_push_settings).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) .with_state(app_state) } @@ -815,3 +828,124 @@ async fn update_oauth_settings( encrypt_response(&data, &session_id, &updated_settings).await } + +async fn get_push_settings( + State(data): State>, + Extension(platform_user): Extension, + Path((org_id, project_id)): Path<(Uuid, Uuid)>, + Extension(session_id): Extension, +) -> Result>, ApiError> { + debug!("Getting project push settings"); + + let org = data + .db + .get_org_by_uuid(org_id) + .map_err(|_| ApiError::NotFound)?; + + let project = data.db.get_org_project_by_uuid(project_id).map_err(|_| { + error!("Project not found"); + ApiError::NotFound + })?; + + if project.org_id != org.id { + error!("Project does not belong to organization"); + return Err(ApiError::NotFound); + } + + let _membership = data + .db + .get_org_membership_by_platform_user_and_org(platform_user.uuid, org.id) + .map_err(|_| ApiError::Unauthorized)?; + + let settings = data + .db + .get_project_push_settings(project.id)? + .unwrap_or_default(); + + encrypt_response(&data, &session_id, &settings).await +} + +async fn update_push_settings( + State(data): State>, + Extension(platform_user): Extension, + Path((org_id, project_id)): Path<(Uuid, Uuid)>, + Extension(update_request): Extension, + Extension(session_id): Extension, +) -> Result>, ApiError> { + debug!("Updating project push settings"); + + validate_push_settings_request(&update_request)?; + + let org = data + .db + .get_org_by_uuid(org_id) + .map_err(|_| ApiError::NotFound)?; + + let project = data.db.get_org_project_by_uuid(project_id).map_err(|_| { + error!("Project not found"); + ApiError::NotFound + })?; + + if project.org_id != org.id { + error!("Project does not belong to organization"); + return Err(ApiError::NotFound); + } + + let membership = data + .db + .get_org_membership_by_platform_user_and_org(platform_user.uuid, org.id) + .map_err(|_| ApiError::Unauthorized)?; + + let role: OrgRole = membership.role.into(); + if !matches!(role, OrgRole::Owner | OrgRole::Admin) { + return Err(ApiError::Unauthorized); + } + + let push_settings = PushSettings { + encrypted_preview_enabled: update_request.encrypted_preview_enabled, + ios: update_request.ios, + android: update_request.android, + }; + + let settings = data + .db + .update_project_push_settings(project.id, push_settings)?; + + let updated_settings = settings.get_push_settings().map_err(|e| { + error!("Failed to parse updated push settings: {:?}", e); + ApiError::InternalServerError + })?; + + encrypt_response(&data, &session_id, &updated_settings).await +} + +fn validate_push_settings_request(request: &UpdatePushSettingsRequest) -> Result<(), ApiError> { + if let Some(ios) = &request.ios { + if ios.enabled + && (ios.bundle_id.trim().is_empty() + || ios.team_id.trim().is_empty() + || ios.key_id.trim().is_empty()) + { + return Err(ApiError::BadRequest); + } + + if ios.bundle_id.len() > 255 || ios.team_id.len() > 32 || ios.key_id.len() > 32 { + return Err(ApiError::BadRequest); + } + } + + if let Some(android) = &request.android { + if android.enabled + && (android.firebase_project_id.trim().is_empty() + || android.package_name.trim().is_empty()) + { + return Err(ApiError::BadRequest); + } + + if android.firebase_project_id.len() > 255 || android.package_name.len() > 255 { + return Err(ApiError::BadRequest); + } + } + + Ok(()) +} diff --git a/src/web/push.rs b/src/web/push.rs new file mode 100644 index 00000000..8bfb0a77 --- /dev/null +++ b/src/web/push.rs @@ -0,0 +1,513 @@ +use crate::models::push_devices::{ + NewPushDevice, PushDevice, PushDeviceError, PUSH_ENV_DEV, PUSH_ENV_PROD, + PUSH_KEY_ALGORITHM_P256_ECDH_V1, PUSH_PLATFORM_ANDROID, PUSH_PLATFORM_IOS, PUSH_PROVIDER_APNS, + PUSH_PROVIDER_FCM, +}; +use crate::models::users::User; +use crate::push::binding::{ + encrypt_push_device_capability_v1, hash_bytes, PushDeviceCapabilityInput, + PushDeviceCapabilityPlaintextV1, +}; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::{ApiError, AppState}; +use axum::{ + extract::{Path, State}, + middleware::from_fn_with_state, + routing::{delete, get, post}, + Extension, Json, Router, +}; +use base64::{engine::general_purpose, Engine as _}; +use chrono::{DateTime, Utc}; +use diesel::Connection; +use p256::pkcs8::DecodePublicKey; +use p256::PublicKey; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use tracing::error; +use uuid::Uuid; + +#[derive(Debug, Clone, Deserialize)] +struct RegisterPushDeviceRequest { + installation_id: Uuid, + platform: String, + provider: String, + environment: String, + app_id: String, + push_token: String, + notification_public_key: String, + key_algorithm: String, + #[serde(default)] + supports_encrypted_preview: bool, + #[serde(default)] + supports_background_processing: bool, +} + +#[derive(Debug, Clone, Serialize)] +struct PushDeviceResponse { + id: Uuid, + object: &'static str, + installation_id: Uuid, + platform: String, + provider: String, + environment: String, + app_id: String, + key_algorithm: String, + supports_encrypted_preview: bool, + supports_background_processing: bool, + last_seen_at: DateTime, + created_at: DateTime, + updated_at: DateTime, +} + +#[derive(Debug, Clone, Serialize)] +struct PushDeviceListResponse { + object: &'static str, + data: Vec, +} + +#[derive(Debug, Clone, Serialize)] +struct DeletedPushDeviceResponse { + id: Uuid, + object: &'static str, + deleted: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RegistrationWriteMode { + ReuseExisting, + InsertNew, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct RegistrationDecision { + write_mode: RegistrationWriteMode, + revoke_device_ids: Vec, +} + +pub fn router(app_state: Arc) -> Router { + Router::new() + .route( + "/v1/push/devices", + post(register_push_device).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/push/devices", + get(list_push_devices) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/push/devices/:id", + delete(revoke_push_device) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .with_state(app_state) +} + +async fn register_push_device( + State(state): State>, + Extension(user): Extension, + Extension(body): Extension, + Extension(session_id): Extension, +) -> Result>, ApiError> { + let public_key_bytes = general_purpose::STANDARD + .decode(&body.notification_public_key) + .map_err(|_| ApiError::BadRequest)?; + validate_register_request(&body, &public_key_bytes)?; + + let token_hash = hash_bytes(body.push_token.as_bytes()); + let public_key_hash = hash_bytes(&public_key_bytes); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let now = Utc::now(); + let device = conn + .transaction::(|conn| { + let existing_for_user = PushDevice::get_by_installation_for_user( + conn, + user.uuid, + body.installation_id, + &body.environment, + )?; + let active_installation = + PushDevice::get_by_installation(conn, body.installation_id, &body.environment)?; + let conflicting_token_owner = PushDevice::get_active_by_token_hash( + conn, + &body.provider, + &body.environment, + &token_hash, + )?; + + let decision = plan_push_device_registration( + user.uuid, + existing_for_user.as_ref(), + active_installation.as_ref(), + conflicting_token_owner.as_ref(), + ); + + for revoke_device_id in decision.revoke_device_ids { + PushDevice::revoke_by_id(conn, revoke_device_id)?; + } + + if decision.write_mode == RegistrationWriteMode::ReuseExisting { + let mut existing = existing_for_user.expect("existing device should be present"); + let capability = + build_device_capability(existing.uuid, &user, &body, public_key_bytes.clone()); + let capability_enc = + encrypt_push_device_capability_v1(&state.enclave_key, &capability) + .map_err(|_| diesel::result::Error::RollbackTransaction)?; + existing.platform = body.platform.clone(); + existing.provider = body.provider.clone(); + existing.environment = body.environment.clone(); + existing.app_id = body.app_id.clone(); + existing.project_id = user.project_id; + existing.push_token_hash = token_hash.clone(); + existing.capability_enc = capability_enc; + existing.notification_public_key_hash = public_key_hash.clone(); + existing.key_algorithm = body.key_algorithm.clone(); + existing.supports_encrypted_preview = body.supports_encrypted_preview; + existing.supports_background_processing = body.supports_background_processing; + existing.last_seen_at = now; + existing.revoked_at = None; + existing.update(conn)?; + Ok(existing) + } else { + let device_uuid = Uuid::new_v4(); + let capability = + build_device_capability(device_uuid, &user, &body, public_key_bytes.clone()); + let capability_enc = + encrypt_push_device_capability_v1(&state.enclave_key, &capability) + .map_err(|_| diesel::result::Error::RollbackTransaction)?; + NewPushDevice { + uuid: device_uuid, + user_id: user.uuid, + project_id: user.project_id, + installation_id: body.installation_id, + platform: body.platform.clone(), + provider: body.provider.clone(), + environment: body.environment.clone(), + app_id: body.app_id.clone(), + push_token_hash: token_hash.clone(), + capability_enc, + notification_public_key_hash: public_key_hash.clone(), + key_algorithm: body.key_algorithm.clone(), + supports_encrypted_preview: body.supports_encrypted_preview, + supports_background_processing: body.supports_background_processing, + last_seen_at: now, + } + .insert(conn) + } + }) + .map_err(map_push_device_error)?; + + encrypt_response(&state, &session_id, &PushDeviceResponse::from(device)).await +} + +async fn list_push_devices( + State(state): State>, + Extension(user): Extension, + Extension(session_id): Extension, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let devices = + PushDevice::list_active_for_user(&mut conn, user.uuid).map_err(map_push_device_error)?; + let response = PushDeviceListResponse { + object: "list", + data: devices.into_iter().map(PushDeviceResponse::from).collect(), + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn revoke_push_device( + State(state): State>, + Path(device_uuid): Path, + Extension(user): Extension, + Extension(session_id): Extension, +) -> Result>, ApiError> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let device = PushDevice::revoke_by_uuid_and_user(&mut conn, device_uuid, user.uuid) + .map_err(map_push_device_error)? + .ok_or(ApiError::NotFound)?; + + let response = DeletedPushDeviceResponse { + id: device.uuid, + object: "push.device.deleted", + deleted: true, + }; + + encrypt_response(&state, &session_id, &response).await +} + +fn validate_register_request( + request: &RegisterPushDeviceRequest, + public_key_bytes: &[u8], +) -> Result<(), ApiError> { + if request.installation_id.is_nil() || request.installation_id.get_version_num() != 4 { + return Err(ApiError::BadRequest); + } + + if request.push_token.trim().is_empty() + || request.app_id.trim().is_empty() + || request.notification_public_key.trim().is_empty() + { + return Err(ApiError::BadRequest); + } + + if request.app_id.len() > 255 || request.push_token.len() > 4096 { + return Err(ApiError::BadRequest); + } + + let valid_platform_provider = matches!( + (request.platform.as_str(), request.provider.as_str()), + (PUSH_PLATFORM_IOS, PUSH_PROVIDER_APNS) | (PUSH_PLATFORM_ANDROID, PUSH_PROVIDER_FCM) + ); + + if !valid_platform_provider { + return Err(ApiError::BadRequest); + } + + if request.environment != PUSH_ENV_DEV && request.environment != PUSH_ENV_PROD { + return Err(ApiError::BadRequest); + } + + if request.key_algorithm != PUSH_KEY_ALGORITHM_P256_ECDH_V1 { + return Err(ApiError::BadRequest); + } + + if request.notification_public_key.len() > 1024 || public_key_bytes.len() > 256 { + return Err(ApiError::BadRequest); + } + + PublicKey::from_public_key_der(public_key_bytes).map_err(|_| ApiError::BadRequest)?; + + Ok(()) +} + +fn build_device_capability( + device_uuid: Uuid, + user: &User, + request: &RegisterPushDeviceRequest, + public_key_bytes: Vec, +) -> PushDeviceCapabilityPlaintextV1 { + PushDeviceCapabilityPlaintextV1::new(PushDeviceCapabilityInput { + push_device_uuid: device_uuid, + user_uuid: user.uuid, + project_id: user.project_id, + installation_id: request.installation_id, + platform: request.platform.clone(), + provider: request.provider.clone(), + environment: request.environment.clone(), + app_id: request.app_id.clone(), + key_algorithm: request.key_algorithm.clone(), + push_token: request.push_token.clone(), + notification_public_key: public_key_bytes, + supports_encrypted_preview: request.supports_encrypted_preview, + supports_background_processing: request.supports_background_processing, + }) +} + +fn map_push_device_error(error: PushDeviceError) -> ApiError { + error!("Push device database error: {:?}", error); + ApiError::InternalServerError +} + +fn plan_push_device_registration( + current_user_id: Uuid, + existing_for_user: Option<&PushDevice>, + active_installation: Option<&PushDevice>, + conflicting_token_owner: Option<&PushDevice>, +) -> RegistrationDecision { + let write_mode = if existing_for_user.is_some() { + RegistrationWriteMode::ReuseExisting + } else { + RegistrationWriteMode::InsertNew + }; + + let existing_device_id = existing_for_user.map(|device| device.id); + let mut revoke_device_ids = Vec::new(); + + if let Some(active_device) = active_installation { + if active_device.user_id != current_user_id && existing_device_id != Some(active_device.id) + { + revoke_device_ids.push(active_device.id); + } + } + + if let Some(conflict) = conflicting_token_owner { + if existing_device_id != Some(conflict.id) && !revoke_device_ids.contains(&conflict.id) { + revoke_device_ids.push(conflict.id); + } + } + + RegistrationDecision { + write_mode, + revoke_device_ids, + } +} + +impl From for PushDeviceResponse { + fn from(value: PushDevice) -> Self { + Self { + id: value.uuid, + object: "push.device", + installation_id: value.installation_id, + platform: value.platform, + provider: value.provider, + environment: value.environment, + app_id: value.app_id, + key_algorithm: value.key_algorithm, + supports_encrypted_preview: value.supports_encrypted_preview, + supports_background_processing: value.supports_background_processing, + last_seen_at: value.last_seen_at, + created_at: value.created_at, + updated_at: value.updated_at, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use base64::engine::general_purpose; + use p256::pkcs8::EncodePublicKey; + use rand_core::OsRng; + + fn sample_push_device( + user_id: Uuid, + installation_id: Uuid, + token_hash: Vec, + revoked_at: Option>, + ) -> PushDevice { + let now = Utc::now(); + PushDevice { + id: now.timestamp_nanos_opt().expect("valid timestamp nanos"), + uuid: Uuid::new_v4(), + user_id, + project_id: 42, + installation_id, + platform: PUSH_PLATFORM_IOS.to_string(), + provider: PUSH_PROVIDER_APNS.to_string(), + environment: PUSH_ENV_PROD.to_string(), + app_id: "ai.trymaple.ios".to_string(), + push_token_hash: token_hash, + capability_enc: vec![1, 2, 3], + notification_public_key_hash: vec![4, 5, 6], + key_algorithm: PUSH_KEY_ALGORITHM_P256_ECDH_V1.to_string(), + supports_encrypted_preview: true, + supports_background_processing: true, + last_seen_at: now, + revoked_at, + created_at: now, + updated_at: now, + } + } + + fn sample_register_request(installation_id: Uuid) -> RegisterPushDeviceRequest { + let secret_key = p256::SecretKey::random(&mut OsRng); + let public_key = secret_key.public_key(); + let public_key_der = public_key + .to_public_key_der() + .expect("public key should encode"); + + RegisterPushDeviceRequest { + installation_id, + platform: PUSH_PLATFORM_IOS.to_string(), + provider: PUSH_PROVIDER_APNS.to_string(), + environment: PUSH_ENV_PROD.to_string(), + app_id: "ai.trymaple.ios".to_string(), + push_token: "push-token".to_string(), + notification_public_key: general_purpose::STANDARD.encode(public_key_der.as_bytes()), + key_algorithm: PUSH_KEY_ALGORITHM_P256_ECDH_V1.to_string(), + supports_encrypted_preview: true, + supports_background_processing: true, + } + } + + #[test] + fn plan_reuses_revoked_same_user_installation() { + let user_id = Uuid::new_v4(); + let installation_id = Uuid::new_v4(); + let existing = sample_push_device(user_id, installation_id, vec![1], Some(Utc::now())); + + let decision = plan_push_device_registration(user_id, Some(&existing), None, None); + + assert_eq!(decision.write_mode, RegistrationWriteMode::ReuseExisting); + assert!(decision.revoke_device_ids.is_empty()); + } + + #[test] + fn plan_revokes_cross_user_installation_once() { + let current_user_id = Uuid::new_v4(); + let installation_id = Uuid::new_v4(); + let other_user_device = + sample_push_device(Uuid::new_v4(), installation_id, vec![1, 2, 3], None); + + let decision = plan_push_device_registration( + current_user_id, + None, + Some(&other_user_device), + Some(&other_user_device), + ); + + assert_eq!(decision.write_mode, RegistrationWriteMode::InsertNew); + assert_eq!(decision.revoke_device_ids, vec![other_user_device.id]); + } + + #[test] + fn plan_revokes_conflicting_token_owner_for_handoff() { + let current_user_id = Uuid::new_v4(); + let conflicting_device = + sample_push_device(Uuid::new_v4(), Uuid::new_v4(), vec![9, 9, 9], None); + + let decision = + plan_push_device_registration(current_user_id, None, None, Some(&conflicting_device)); + + assert_eq!(decision.write_mode, RegistrationWriteMode::InsertNew); + assert_eq!(decision.revoke_device_ids, vec![conflicting_device.id]); + } + + #[test] + fn validate_register_request_rejects_nil_installation_id() { + let request = sample_register_request(Uuid::nil()); + let public_key_bytes = general_purpose::STANDARD + .decode(&request.notification_public_key) + .expect("public key should decode"); + + assert!(matches!( + validate_register_request(&request, &public_key_bytes), + Err(ApiError::BadRequest) + )); + } + + #[test] + fn validate_register_request_rejects_non_v4_installation_id() { + let request = sample_register_request( + Uuid::parse_str("6ba7b810-9dad-11d1-80b4-00c04fd430c8").expect("uuid should parse"), + ); + let public_key_bytes = general_purpose::STANDARD + .decode(&request.notification_public_key) + .expect("public key should decode"); + + assert!(matches!( + validate_register_request(&request, &public_key_bytes), + Err(ApiError::BadRequest) + )); + } +} From 73583198839215708b39874fb4145cfaedca2e86 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:31:09 -0500 Subject: [PATCH 5/7] feat(agent): add explicit main-agent onboarding Switch Maple creation from implicit lazy bootstrapping to an explicit init flow that seeds onboarding messages and stores user locale and timezone hints for later agent runs. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 35 +++ Cargo.toml | 1 + src/models/mod.rs | 1 + src/models/user_preferences.rs | 115 +++++++++ src/web/agent/mod.rs | 102 +++++++- src/web/agent/runtime.rs | 410 +++++++++++++++++++++++++++------ src/web/agent/tools.rs | 123 +++++----- 7 files changed, 644 insertions(+), 143 deletions(-) create mode 100644 src/models/user_preferences.rs diff --git a/Cargo.lock b/Cargo.lock index 83f02443..c4a49c3d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1455,6 +1455,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -4147,6 +4157,7 @@ dependencies = [ "cbc", "chacha20poly1305", "chrono", + "chrono-tz", "diesel", "diesel-derive-enum", "dotenv", @@ -4413,6 +4424,24 @@ dependencies = [ "ucd-trie", ] +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.11" @@ -5456,6 +5485,12 @@ dependencies = [ "time", ] +[[package]] +name = "siphasher" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" + [[package]] name = "slab" version = "0.4.12" diff --git a/Cargo.toml b/Cargo.toml index 826f26e1..41aaef07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ diesel = { version = "=2.2.3", features = [ ] } diesel-derive-enum = { version = "2.1.0", features = ["postgres"] } chrono = { version = "0.4.26", features = ["serde"] } +chrono-tz = "0.10" dotenv = "0.15.0" aes-gcm = "0.10.1" aes-siv = "0.7" diff --git a/src/models/mod.rs b/src/models/mod.rs index 4c98fbf2..cd2756a8 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -26,5 +26,6 @@ pub mod token_usage; pub mod user_api_keys; pub mod user_embeddings; pub mod user_kv; +pub mod user_preferences; pub mod user_seed_wrappings; pub mod users; diff --git a/src/models/user_preferences.rs b/src/models/user_preferences.rs new file mode 100644 index 00000000..152daf76 --- /dev/null +++ b/src/models/user_preferences.rs @@ -0,0 +1,115 @@ +use crate::models::schema::user_preferences; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use thiserror::Error; +use uuid::Uuid; + +pub const USER_PREFERENCE_TIMEZONE: &str = "timezone"; +pub const USER_PREFERENCE_LOCALE: &str = "locale"; + +#[derive(Error, Debug)] +pub enum UserPreferenceError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), + #[error("Invalid preference: {0}")] + InvalidPreference(String), +} + +#[derive(Queryable, Identifiable, AsChangeset, Clone, Debug)] +#[diesel(table_name = user_preferences)] +pub struct UserPreference { + pub id: i64, + pub user_id: Uuid, + pub key: String, + pub value_enc: Vec, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl UserPreference { + pub fn get_by_user_and_key( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_key: &str, + ) -> Result, UserPreferenceError> { + user_preferences::table + .filter(user_preferences::user_id.eq(lookup_user_id)) + .filter(user_preferences::key.eq(lookup_key)) + .first::(conn) + .optional() + .map_err(UserPreferenceError::DatabaseError) + } + + pub fn validate(key: &str, value: &str) -> Result<(), UserPreferenceError> { + match key { + USER_PREFERENCE_TIMEZONE => value.parse::().map(|_| ()).map_err(|_| { + UserPreferenceError::InvalidPreference(format!( + "Invalid timezone '{value}'. Use an IANA timezone like 'America/Chicago'" + )) + }), + USER_PREFERENCE_LOCALE => validate_locale(value), + _ => Ok(()), + } + } +} + +fn validate_locale(value: &str) -> Result<(), UserPreferenceError> { + let trimmed = value.trim(); + if trimmed.len() < 2 || trimmed.len() > 35 { + return Err(UserPreferenceError::InvalidPreference(format!( + "Invalid locale '{value}'" + ))); + } + + let starts_or_ends_with_separator = + trimmed.chars().next().is_some_and(|c| c == '-' || c == '_') + || trimmed.chars().last().is_some_and(|c| c == '-' || c == '_'); + + if starts_or_ends_with_separator || trimmed.contains("--") || trimmed.contains("__") { + return Err(UserPreferenceError::InvalidPreference(format!( + "Invalid locale '{value}'" + ))); + } + + if !trimmed + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') + { + return Err(UserPreferenceError::InvalidPreference(format!( + "Invalid locale '{value}'" + ))); + } + + Ok(()) +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = user_preferences)] +pub struct NewUserPreference { + pub user_id: Uuid, + pub key: String, + pub value_enc: Vec, +} + +impl NewUserPreference { + pub fn new(user_id: Uuid, key: impl Into, value_enc: Vec) -> Self { + Self { + user_id, + key: key.into(), + value_enc, + } + } + + pub fn insert_or_update( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(user_preferences::table) + .values(self) + .on_conflict((user_preferences::user_id, user_preferences::key)) + .do_update() + .set(user_preferences::value_enc.eq(self.value_enc.clone())) + .get_result::(conn) + .map_err(UserPreferenceError::DatabaseError) + } +} diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index d65a1a0f..22c77a7f 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -49,6 +49,12 @@ struct AgentChatRequest { input: MessageContent, } +#[derive(Debug, Clone, Default, Deserialize)] +struct InitMainAgentRequest { + timezone: Option, + locale: Option, +} + #[derive(Debug, Clone, Deserialize)] struct CreateSubagentRequest { display_name: Option, @@ -79,6 +85,13 @@ struct MainAgentResponse { updated_at: i64, } +#[derive(Debug, Clone, Serialize)] +struct InitMainAgentResponse { + #[serde(flatten)] + agent: MainAgentResponse, + messages: Vec, +} + #[derive(Debug, Clone, Serialize)] struct SubagentListResponse { object: &'static str, @@ -237,6 +250,25 @@ fn build_main_agent_response(ctx: &AgentConversationContext) -> MainAgentRespons } } +fn build_init_main_agent_response( + ctx: &AgentConversationContext, + onboarding_messages: Vec, +) -> InitMainAgentResponse { + InitMainAgentResponse { + agent: build_main_agent_response(ctx), + messages: onboarding_messages + .into_iter() + .map(|message| ConversationItem::Message { + id: message.id, + status: Some("completed".to_string()), + role: "assistant".to_string(), + content: MessageContentConverter::assistant_text_to_content(message.content), + created_at: Some(message.created_at), + }) + .collect(), + } +} + fn build_subagent_response( agent: &Agent, conversation: &Conversation, @@ -295,7 +327,7 @@ async fn load_main_agent_context( .map_err(|_| ApiError::InternalServerError)?; let (agent, conversation) = - runtime::ensure_main_agent(state, &mut conn, &user_key, user.uuid).await?; + runtime::load_main_agent(&mut conn, user.uuid)?.ok_or(ApiError::NotFound)?; Ok(AgentConversationContext { agent, @@ -441,6 +473,13 @@ pub fn router(app_state: Arc) -> Router<()> { delete(delete_main_agent) .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), ) + .route( + "/v1/agent/init", + post(init_main_agent).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) .route( "/v1/agent/items", get(list_main_agent_items) @@ -510,6 +549,48 @@ async fn get_main_agent( encrypt_response(&state, &session_id, &response).await } +async fn init_main_agent( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + let user_key = state + .get_user_key(&user, &auth_context, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let init_result = runtime::init_main_agent( + &state, + &mut conn, + &user_key, + user.uuid, + &runtime::MainAgentInitOptions { + timezone: body.timezone, + locale: body.locale, + }, + ) + .await?; + + let response = build_init_main_agent_response( + &AgentConversationContext { + agent: init_result.agent, + conversation: init_result.conversation, + user_key, + }, + init_result.onboarding_messages, + ); + + encrypt_response(&state, &session_id, &response).await +} + async fn list_main_agent_items( State(state): State>, Query(params): Query, @@ -695,6 +776,16 @@ async fn chat_main( Extension(auth_context): Extension, Extension(body): Extension, ) -> Result { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + if runtime::load_main_agent(&mut conn, user.uuid)?.is_none() { + return Err(ApiError::NotFound); + } + chat_with_target( state, session_id, @@ -866,10 +957,15 @@ async fn run_agent_chat_task( Ok(runtime) => runtime, Err(e) => { error!("Agent runtime initialization error: {e:?}"); + let error_message = if matches!(e, ApiError::NotFound) { + "Main agent is not initialized.".to_string() + } else { + "Failed to initialize agent runtime.".to_string() + }; let _ = send_agent_client_event( &tx, AgentClientEvent::Error(AgentErrorEvent { - error: "Failed to initialize agent runtime.".to_string(), + error: error_message, }), client_connected, ) @@ -1088,7 +1184,7 @@ async fn create_subagent( .map_err(|_| ApiError::InternalServerError)?; let (main_agent, _) = - runtime::ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; + runtime::load_main_agent(&mut conn, user.uuid)?.ok_or(ApiError::NotFound)?; let (agent, conversation, display_name) = runtime::create_subagent( &mut conn, &user_key, diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 03e62668..7b6828d2 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -1,3 +1,4 @@ +use chrono_tz::Tz; use secp256k1::SecretKey; use serde_json::json; use std::sync::Arc; @@ -12,14 +13,20 @@ use crate::models::agents::{ }; use crate::models::conversation_summaries::{ConversationSummary, NewConversationSummary}; use crate::models::memory_blocks::{ - MemoryBlock, NewMemoryBlock, MEMORY_BLOCK_LABEL_HUMAN, MEMORY_BLOCK_LABEL_PERSONA, + MemoryBlock, MEMORY_BLOCK_LABEL_HUMAN, MEMORY_BLOCK_LABEL_PERSONA, }; use crate::models::responses::{ AssistantMessage, Conversation, NewAssistantMessage, NewConversation, NewToolCall, NewToolOutput, NewUserMessage, RawThreadMessage, RawThreadMessageMetadata, ToolCall, ToolOutput, UserMessage, }; -use crate::models::schema::{agents, memory_blocks, user_embeddings}; +use crate::models::schema::{ + agents, assistant_messages, memory_blocks, user_embeddings, user_messages, +}; +use crate::models::user_preferences::{ + NewUserPreference, UserPreference, UserPreferenceError, USER_PREFERENCE_LOCALE, + USER_PREFERENCE_TIMEZONE, +}; use crate::rag::{ insert_message_embedding, serialize_f32_le, SOURCE_TYPE_ARCHIVAL, SOURCE_TYPE_MESSAGE, }; @@ -40,17 +47,24 @@ use super::tools::{ }; use super::vision; -const DEFAULT_PERSONA_VALUE: &str = "I am Maple, a helpful AI companion. I maintain long-term memory across our conversations and strive to be friendly, concise, and genuinely helpful."; +pub(crate) const DEFAULT_PERSONA_VALUE: &str = "I am Maple, a helpful AI companion. I maintain long-term memory across our conversations and strive to be friendly, concise, and genuinely helpful."; pub const DEFAULT_MODEL: &str = "kimi-k2-6"; pub const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; pub const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; const MIN_MESSAGES_IN_CONTEXT: usize = 20; const MAIN_AGENT_METADATA_TYPE: &str = "agent_main"; const SUBAGENT_METADATA_TYPE: &str = "subagent"; +const MAIN_AGENT_ONBOARDING_TURN_LIMIT: i64 = 15; +const MAIN_AGENT_ONBOARDING_MESSAGES: [&str; 3] = [ + "Hey, I'm Maple. πŸ‘‹", + "Nice to meet you.", + "What should I call you?", +]; #[derive(Clone, Debug, Default)] struct AgentContext { current_time: String, + user_locale: String, agent_kind: String, subagent_purpose: String, persona_block: String, @@ -58,9 +72,30 @@ struct AgentContext { memory_metadata: String, previous_context_summary: String, recent_conversation: String, + main_agent_user_message_count: i64, is_first_time_user: bool, } +#[derive(Clone, Debug, Default)] +pub(crate) struct MainAgentInitOptions { + pub timezone: Option, + pub locale: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct SeededOnboardingMessage { + pub id: Uuid, + pub content: String, + pub created_at: i64, +} + +#[derive(Clone, Debug)] +pub(crate) struct MainAgentInitResult { + pub agent: Agent, + pub conversation: Conversation, + pub onboarding_messages: Vec, +} + #[derive(Clone, Debug)] pub struct ExecutedTool { pub tool_call: AgentToolCall, @@ -163,12 +198,230 @@ async fn subagent_metadata_enc( .await } -pub(crate) async fn ensure_main_agent( - state: &Arc, +fn normalize_optional_preference(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) +} + +fn validate_optional_preference( + key: &str, + value: Option<&str>, +) -> Result, ApiError> { + let Some(value) = normalize_optional_preference(value) else { + return Ok(None); + }; + + UserPreference::validate(key, &value).map_err(|e| match e { + UserPreferenceError::InvalidPreference(_) => ApiError::BadRequest, + UserPreferenceError::DatabaseError(_) => ApiError::InternalServerError, + })?; + + Ok(Some(value)) +} + +fn onboarding_message_texts() -> &'static [&'static str; 3] { + &MAIN_AGENT_ONBOARDING_MESSAGES +} + +fn load_user_preference( conn: &mut diesel::PgConnection, user_key: &SecretKey, user_id: Uuid, -) -> Result<(Agent, Conversation), ApiError> { + key: &str, +) -> Result, ApiError> { + let preference = UserPreference::get_by_user_and_key(conn, user_id, key).map_err(|e| { + error!("Failed to load user preference '{key}': {e:?}"); + ApiError::InternalServerError + })?; + + decrypt_string( + user_key, + preference.as_ref().map(|preference| &preference.value_enc), + ) + .map_err(|e| { + error!("Failed to decrypt user preference '{key}': {e:?}"); + ApiError::InternalServerError + }) +} + +fn load_user_timezone( + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, +) -> Result, ApiError> { + let Some(timezone) = load_user_preference(conn, user_key, user_id, USER_PREFERENCE_TIMEZONE)? + else { + return Ok(None); + }; + + match timezone.parse::() { + Ok(timezone) => Ok(Some(timezone)), + Err(_) => { + warn!("Ignoring invalid stored timezone '{timezone}' for user {user_id}"); + Ok(None) + } + } +} + +fn format_current_time(now: chrono::DateTime, timezone: Option<&Tz>) -> String { + if let Some(timezone) = timezone { + let local_time = now.with_timezone(timezone); + format!( + "{} ({})", + local_time.format("%m/%d/%Y %H:%M:%S (%A)"), + timezone.name() + ) + } else { + format!("{} UTC", now.format("%m/%d/%Y %H:%M:%S (%A)")) + } +} + +fn format_message_timestamp( + created_at: chrono::DateTime, + timezone: Option<&Tz>, +) -> String { + if let Some(timezone) = timezone { + let local_time = created_at.with_timezone(timezone); + format!( + "{} ({})", + local_time.format("%m/%d/%Y %H:%M:%S"), + timezone.name() + ) + } else { + format!("{} UTC", created_at.format("%m/%d/%Y %H:%M:%S")) + } +} + +async fn upsert_user_preference( + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, + key: &str, + value: Option, +) -> Result<(), ApiError> { + let Some(value) = value else { + return Ok(()); + }; + + let value_enc = encrypt_with_key(user_key, value.as_bytes()).await; + NewUserPreference::new(user_id, key, value_enc) + .insert_or_update(conn) + .map_err(|e| { + error!("Failed to upsert user preference '{key}': {e:?}"); + ApiError::InternalServerError + })?; + + Ok(()) +} + +async fn seed_main_agent_onboarding_messages( + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, + conversation_id: i64, +) -> Result, ApiError> { + let user_message_count: i64 = user_messages::table + .filter(user_messages::conversation_id.eq(conversation_id)) + .count() + .get_result(conn) + .map_err(|e| { + error!("Failed to count user messages for onboarding seed: {e:?}"); + ApiError::InternalServerError + })?; + + let assistant_message_count: i64 = assistant_messages::table + .filter(assistant_messages::conversation_id.eq(conversation_id)) + .count() + .get_result(conn) + .map_err(|e| { + error!("Failed to count assistant messages for onboarding seed: {e:?}"); + ApiError::InternalServerError + })?; + + if user_message_count > 0 || assistant_message_count > 0 { + return Ok(vec![]); + } + + let mut seeded_messages = Vec::with_capacity(onboarding_message_texts().len()); + + for message in onboarding_message_texts() { + let content_enc = Some(encrypt_with_key(user_key, message.as_bytes()).await); + let completion_tokens = count_tokens(message) as i32; + + let inserted = NewAssistantMessage { + uuid: Uuid::new_v4(), + conversation_id, + response_id: None, + user_id, + content_enc, + completion_tokens, + status: "completed".to_string(), + finish_reason: None, + created_at: chrono::Utc::now(), + } + .insert(conn) + .map_err(|e| { + error!("Failed to seed main agent onboarding message: {e:?}"); + ApiError::InternalServerError + })?; + + seeded_messages.push(SeededOnboardingMessage { + id: inserted.uuid, + content: message.to_string(), + created_at: inserted.created_at.timestamp(), + }); + } + + Ok(seeded_messages) +} + +pub(crate) fn load_main_agent( + conn: &mut diesel::PgConnection, + user_id: Uuid, +) -> Result, ApiError> { + let Some(agent) = Agent::get_main_for_user(conn, user_id).map_err(|e| { + error!("Failed to load main agent: {e:?}"); + ApiError::InternalServerError + })? + else { + return Ok(None); + }; + + let conversation = match Conversation::get_by_id_and_user(conn, agent.conversation_id, user_id) + { + Ok(conversation) => conversation, + Err(crate::models::responses::ResponsesError::ConversationNotFound) => { + warn!( + "Main agent {} for user {} points to missing conversation {}; treating as uninitialized", + agent.uuid, + user_id, + agent.conversation_id, + ); + return Ok(None); + } + Err(e) => { + error!("Failed to load main agent conversation: {e:?}"); + return Err(ApiError::InternalServerError); + } + }; + + Ok(Some((agent, conversation))) +} + +pub(crate) async fn init_main_agent( + _state: &Arc, + conn: &mut diesel::PgConnection, + user_key: &SecretKey, + user_id: Uuid, + init_options: &MainAgentInitOptions, +) -> Result { + let timezone = + validate_optional_preference(USER_PREFERENCE_TIMEZONE, init_options.timezone.as_deref())?; + let locale = + validate_optional_preference(USER_PREFERENCE_LOCALE, init_options.locale.as_deref())?; + let existing = Agent::get_main_for_user(conn, user_id).map_err(|e| { error!("Failed to load main agent: {e:?}"); ApiError::InternalServerError @@ -235,9 +488,16 @@ pub(crate) async fn ensure_main_agent( (agent, conversation) }; - AgentRuntime::ensure_default_blocks(state, conn, user_key, user_id).await?; + upsert_user_preference(conn, user_key, user_id, USER_PREFERENCE_TIMEZONE, timezone).await?; + upsert_user_preference(conn, user_key, user_id, USER_PREFERENCE_LOCALE, locale).await?; + let onboarding_messages = + seed_main_agent_onboarding_messages(conn, user_key, user_id, conversation.id).await?; - Ok((agent, conversation)) + Ok(MainAgentInitResult { + agent, + conversation, + onboarding_messages, + }) } pub(crate) async fn create_subagent( @@ -314,7 +574,7 @@ impl AgentRuntime { .map_err(|_| ApiError::InternalServerError)?; let (agent, conversation) = - ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; + load_main_agent(&mut conn, user.uuid)?.ok_or(ApiError::NotFound)?; Self::from_loaded(state, user, user_key, agent, conversation).await } @@ -334,8 +594,6 @@ impl AgentRuntime { .get() .map_err(|_| ApiError::InternalServerError)?; - let _ = ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; - let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) .map_err(|e| { error!("Failed to load subagent: {e:?}"); @@ -447,55 +705,6 @@ impl AgentRuntime { }) } - pub(crate) async fn ensure_default_blocks( - _state: &Arc, - conn: &mut diesel::PgConnection, - user_key: &SecretKey, - user_id: Uuid, - ) -> Result<(), ApiError> { - let persona = MemoryBlock::get_by_user_and_label(conn, user_id, MEMORY_BLOCK_LABEL_PERSONA) - .map_err(|e| { - error!("Failed to load persona block: {e:?}"); - ApiError::InternalServerError - })?; - - if persona.is_none() { - let value_enc = encrypt_with_key(user_key, DEFAULT_PERSONA_VALUE.as_bytes()).await; - let block = NewMemoryBlock { - uuid: Uuid::new_v4(), - user_id, - label: MEMORY_BLOCK_LABEL_PERSONA.to_string(), - value_enc, - }; - block.insert_or_update(conn).map_err(|e| { - error!("Failed to create persona block: {e:?}"); - ApiError::InternalServerError - })?; - } - - let human = MemoryBlock::get_by_user_and_label(conn, user_id, MEMORY_BLOCK_LABEL_HUMAN) - .map_err(|e| { - error!("Failed to load human block: {e:?}"); - ApiError::InternalServerError - })?; - - if human.is_none() { - let value_enc = encrypt_with_key(user_key, "".as_bytes()).await; - let block = NewMemoryBlock { - uuid: Uuid::new_v4(), - user_id, - label: MEMORY_BLOCK_LABEL_HUMAN.to_string(), - value_enc, - }; - block.insert_or_update(conn).map_err(|e| { - error!("Failed to create human block: {e:?}"); - ApiError::InternalServerError - })?; - } - - Ok(()) - } - pub fn clear_tool_results(&mut self) { self.current_tool_results.clear(); self.previous_step_summary = None; @@ -505,6 +714,35 @@ impl AgentRuntime { self.max_steps } + fn build_step_system_prompt(&self, ctx: &AgentContext) -> String { + let mut prompt = self.system_prompt.clone(); + + if self.agent.kind == AGENT_KIND_MAIN + && ctx.main_agent_user_message_count > 0 + && ctx.main_agent_user_message_count <= MAIN_AGENT_ONBOARDING_TURN_LIMIT + { + prompt.push_str(&format!( + "\n\nMAIN AGENT ONBOARDING WINDOW:\nThis is still early in your relationship with the user (main-agent user message {} of {}). Be especially warm, welcoming, and natural. Focus on getting to know them gradually rather than interrogating them. Ask at most one thoughtful follow-up at a time. At a natural point early on, briefly explain who Maple is, what this app is good for, and what the user can expect from chatting here. Do that conversationally, not as a canned pitch, and do not force it into the second or third message if the moment is awkward. Prioritize learning their name, important relationships, routines, goals, preferences, and current life context when it feels natural. Save useful facts to memory proactively. Avoid transactional, overly clinical, or assistant-like phrasing. Make Maple feel like a welcoming, steady presence, not a service desk.", + ctx.main_agent_user_message_count, + MAIN_AGENT_ONBOARDING_TURN_LIMIT, + )); + + if ctx.human_block.trim().is_empty() { + prompt.push_str( + "\nIf you still do not know the user's name, ask what they would like to be called and store it as soon as they answer.", + ); + } + } + + if !ctx.user_locale.trim().is_empty() { + prompt.push_str("\n\nUSER LOCALE HINT:\nThe user's preferred locale is '"); + prompt.push_str(ctx.user_locale.trim()); + prompt.push_str("'. If it is natural and the user has not indicated otherwise, prefer replying in that locale/language."); + } + + prompt + } + /// Prepare the runtime for a new message: validate, persist user message, compact if needed. /// Call this once before driving the step loop. pub async fn prepare(&mut self, user_message: &MessageContent) -> Result { @@ -716,6 +954,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If } }; + let system_prompt = self.build_step_system_prompt(&ctx); + let input = AgentResponseInput { input: input_content.clone(), current_time: ctx.current_time, @@ -732,7 +972,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If let response = call_agent_response_with_retry_and_correction( &self.lm, - &self.system_prompt, + &system_prompt, &input, &input_content, &self.available_tools, @@ -827,12 +1067,11 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If } async fn build_context(&self) -> Result { - let mut ctx = AgentContext::default(); - - let now = chrono::Utc::now(); - ctx.current_time = format!("{} UTC", now.format("%m/%d/%Y %H:%M:%S (%A)")); - ctx.agent_kind = self.agent.kind.clone(); - ctx.subagent_purpose = self.subagent_purpose.clone(); + let mut ctx = AgentContext { + agent_kind: self.agent.kind.clone(), + subagent_purpose: self.subagent_purpose.clone(), + ..AgentContext::default() + }; let mut conn = self .state @@ -841,6 +1080,18 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .get() .map_err(|_| ApiError::InternalServerError)?; + let user_timezone = load_user_timezone(&mut conn, &self.user_key, self.user.uuid)?; + ctx.user_locale = load_user_preference( + &mut conn, + &self.user_key, + self.user.uuid, + USER_PREFERENCE_LOCALE, + )? + .unwrap_or_default(); + + let now = chrono::Utc::now(); + ctx.current_time = format_current_time(now, user_timezone.as_ref()); + let persona = MemoryBlock::get_by_user_and_label( &mut conn, self.user.uuid, @@ -853,7 +1104,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ctx.persona_block = decrypt_string(&self.user_key, persona.as_ref().map(|b| &b.value_enc)) .map_err(|_| ApiError::InternalServerError)? - .unwrap_or_default(); + .unwrap_or_else(|| DEFAULT_PERSONA_VALUE.to_string()); ctx.human_block = decrypt_string(&self.user_key, human.as_ref().map(|b| &b.value_enc)) .map_err(|_| ApiError::InternalServerError)? .unwrap_or_default(); @@ -910,6 +1161,21 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .unwrap_or_default(); } + if self.agent.kind == AGENT_KIND_MAIN { + ctx.main_agent_user_message_count = user_messages::table + .filter(user_messages::conversation_id.eq(self.conversation.id)) + .count() + .get_result(&mut conn) + .map_err(|e| { + error!("Failed to count main agent user messages: {e:?}"); + ApiError::InternalServerError + })?; + + if ctx.main_agent_user_message_count <= 1 && ctx.human_block.trim().is_empty() { + ctx.is_first_time_user = true; + } + } + let mut messages = RawThreadMessage::get_conversation_context( &mut conn, self.conversation.id, @@ -941,12 +1207,6 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If } } - // First-time user heuristic for brand-new conversations. - let has_summary = summary.is_some(); - if messages.len() <= 1 && !has_summary { - ctx.is_first_time_user = true; - } - // Render conversation history let mut conversation = String::new(); for msg in &messages { @@ -961,7 +1221,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If other => other, }; - let timestamp = format!("{} UTC", msg.created_at.format("%m/%d/%Y %H:%M:%S")); + let timestamp = format_message_timestamp(msg.created_at, user_timezone.as_ref()); let content = self.render_raw_message(msg)?; conversation.push_str(&format!("[{} @ {}]: {}\n", role, timestamp, content)); } diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index a35e0197..551c11c3 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -174,30 +174,29 @@ fn insert_at_line(value: &str, content: &str, line: i32) -> String { new_lines.join("\n") } +fn default_block_value(label: &str) -> Result<&'static str, ApiError> { + match label { + MEMORY_BLOCK_LABEL_PERSONA => Ok(runtime::DEFAULT_PERSONA_VALUE), + MEMORY_BLOCK_LABEL_HUMAN => Ok(""), + _ => Err(ApiError::BadRequest), + } +} + async fn update_block_value( state: &Arc, user_id: Uuid, user_key: &SecretKey, label: &str, new_value: &str, - existing: &MemoryBlock, ) -> Result<(), ApiError> { - if label != MEMORY_BLOCK_LABEL_PERSONA && label != MEMORY_BLOCK_LABEL_HUMAN { - return Err(ApiError::BadRequest); - } + let _ = default_block_value(label)?; if new_value.len() > MAX_CORE_MEMORY_BLOCK_CHARS { return Err(ApiError::BadRequest); } let value_enc = encrypt_with_key(user_key, new_value.as_bytes()).await; - - let new_block = NewMemoryBlock { - uuid: existing.uuid, - user_id, - label: label.to_string(), - value_enc, - }; + let new_block = NewMemoryBlock::new(user_id, label, value_enc); let mut conn = state .db @@ -260,21 +259,26 @@ impl Tool for MemoryReplaceTool { }; let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { - Ok(Some(b)) => b, - Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Ok(existing) => existing, Err(e) => { error!("Failed to load block '{}': {e:?}", block); return ToolResult::error("Failed to load memory block"); } }; - let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { - Ok(Some(v)) => v, - Ok(None) => String::new(), - Err(e) => { - error!("Failed to decrypt block '{}': {e:?}", block); - return ToolResult::error("Failed to decrypt memory block"); - } + let value = match existing { + Some(existing) => match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }, + None => match default_block_value(block) { + Ok(default) => default.to_string(), + Err(_) => return ToolResult::error(format!("Block '{}' not found", block)), + }, }; if !value.contains(old) { @@ -285,15 +289,8 @@ impl Tool for MemoryReplaceTool { } let new_value = value.replace(old, new); - if let Err(e) = update_block_value( - &self.state, - self.user_id, - &self.user_key, - block, - &new_value, - &existing, - ) - .await + if let Err(e) = + update_block_value(&self.state, self.user_id, &self.user_key, block, &new_value).await { error!("Failed to update block '{}': {e:?}", block); return ToolResult::error("Failed to update memory block"); @@ -347,21 +344,26 @@ impl Tool for MemoryAppendTool { }; let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { - Ok(Some(b)) => b, - Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Ok(existing) => existing, Err(e) => { error!("Failed to load block '{}': {e:?}", block); return ToolResult::error("Failed to load memory block"); } }; - let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { - Ok(Some(v)) => v, - Ok(None) => String::new(), - Err(e) => { - error!("Failed to decrypt block '{}': {e:?}", block); - return ToolResult::error("Failed to decrypt memory block"); - } + let value = match existing { + Some(existing) => match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }, + None => match default_block_value(block) { + Ok(default) => default.to_string(), + Err(_) => return ToolResult::error(format!("Block '{}' not found", block)), + }, }; let new_value = if value.is_empty() { @@ -370,15 +372,8 @@ impl Tool for MemoryAppendTool { format!("{}\n{}", value, content) }; - if let Err(e) = update_block_value( - &self.state, - self.user_id, - &self.user_key, - block, - &new_value, - &existing, - ) - .await + if let Err(e) = + update_block_value(&self.state, self.user_id, &self.user_key, block, &new_value).await { error!("Failed to update block '{}': {e:?}", block); return ToolResult::error("Failed to update memory block"); @@ -433,34 +428,32 @@ impl Tool for MemoryInsertTool { }; let existing = match MemoryBlock::get_by_user_and_label(&mut conn, self.user_id, block) { - Ok(Some(b)) => b, - Ok(None) => return ToolResult::error(format!("Block '{}' not found", block)), + Ok(existing) => existing, Err(e) => { error!("Failed to load block '{}': {e:?}", block); return ToolResult::error("Failed to load memory block"); } }; - let value = match decrypt_string(&self.user_key, Some(&existing.value_enc)) { - Ok(Some(v)) => v, - Ok(None) => String::new(), - Err(e) => { - error!("Failed to decrypt block '{}': {e:?}", block); - return ToolResult::error("Failed to decrypt memory block"); - } + let value = match existing { + Some(existing) => match decrypt_string(&self.user_key, Some(&existing.value_enc)) { + Ok(Some(v)) => v, + Ok(None) => String::new(), + Err(e) => { + error!("Failed to decrypt block '{}': {e:?}", block); + return ToolResult::error("Failed to decrypt memory block"); + } + }, + None => match default_block_value(block) { + Ok(default) => default.to_string(), + Err(_) => return ToolResult::error(format!("Block '{}' not found", block)), + }, }; let new_value = insert_at_line(&value, content, line); - if let Err(e) = update_block_value( - &self.state, - self.user_id, - &self.user_key, - block, - &new_value, - &existing, - ) - .await + if let Err(e) = + update_block_value(&self.state, self.user_id, &self.user_key, block, &new_value).await { error!("Failed to update block '{}': {e:?}", block); return ToolResult::error("Failed to update memory block"); From 0f41c6275ef94360b75fd7f7a85022bf3e93a1f3 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:32:14 -0500 Subject: [PATCH 6/7] feat(agent): add schedules and preference tools Add background Maple wakeups with durable leases and retry handling, and let the agent update validated user preferences that can influence schedule timing and future replies. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../down.sql | 33 + .../up.sql | 53 + src/agent_background.rs | 181 ++ src/db.rs | 13 +- src/main.rs | 92 + src/models/agent_background_grants.rs | 76 + src/models/agent_schedule_runs.rs | 364 +++ src/models/agent_schedules.rs | 342 +++ src/models/agents.rs | 8 + src/models/mod.rs | 3 + src/models/notification_events.rs | 3 + src/models/schema.rs | 24 +- src/push/binding.rs | 37 +- src/push/crypto.rs | 2 + src/push/mod.rs | 105 +- src/push/worker.rs | 93 +- src/security_invariants.rs | 1 + src/seed_wrapping.rs | 45 + src/web/agent/mod.rs | 21 +- src/web/agent/runtime.rs | 51 +- src/web/agent/schedules.rs | 1999 +++++++++++++++++ src/web/agent/tools.rs | 138 ++ 22 files changed, 3659 insertions(+), 25 deletions(-) create mode 100644 migrations/2026-06-26-224528_agent_background_push_aead/down.sql create mode 100644 migrations/2026-06-26-224528_agent_background_push_aead/up.sql create mode 100644 src/agent_background.rs create mode 100644 src/models/agent_background_grants.rs create mode 100644 src/models/agent_schedule_runs.rs create mode 100644 src/models/agent_schedules.rs create mode 100644 src/web/agent/schedules.rs diff --git a/migrations/2026-06-26-224528_agent_background_push_aead/down.sql b/migrations/2026-06-26-224528_agent_background_push_aead/down.sql new file mode 100644 index 00000000..9441c3e5 --- /dev/null +++ b/migrations/2026-06-26-224528_agent_background_push_aead/down.sql @@ -0,0 +1,33 @@ +DELETE FROM notification_deliveries; +DELETE FROM notification_events; +DELETE FROM agent_schedule_runs; +DELETE FROM agent_background_grants; +DELETE FROM agent_schedules; + +DROP INDEX IF EXISTS idx_notification_events_background_grant; + +ALTER TABLE notification_events + DROP COLUMN IF EXISTS background_grant_id; + +ALTER TABLE notification_events + DROP CONSTRAINT IF EXISTS notification_events_source_kind_check; + +ALTER TABLE notification_events + ADD CONSTRAINT notification_events_source_kind_check + CHECK (source_kind IN ('request_continuation')); + +DROP TRIGGER IF EXISTS update_agent_background_grants_updated_at ON agent_background_grants; +DROP INDEX IF EXISTS idx_agent_background_grants_user; +DROP INDEX IF EXISTS idx_agent_background_grants_schedule_active; +DROP TABLE IF EXISTS agent_background_grants; + +ALTER TABLE agent_schedules + ADD COLUMN description TEXT NOT NULL, + DROP COLUMN IF EXISTS description_enc; + +ALTER TABLE user_seed_wrappings + DROP CONSTRAINT IF EXISTS user_seed_wrappings_credential_kind_check; + +ALTER TABLE user_seed_wrappings + ADD CONSTRAINT user_seed_wrappings_credential_kind_check + CHECK (credential_kind IN ('password', 'oauth')); diff --git a/migrations/2026-06-26-224528_agent_background_push_aead/up.sql b/migrations/2026-06-26-224528_agent_background_push_aead/up.sql new file mode 100644 index 00000000..dbcd4ade --- /dev/null +++ b/migrations/2026-06-26-224528_agent_background_push_aead/up.sql @@ -0,0 +1,53 @@ +DELETE FROM agent_schedule_runs; +DELETE FROM agent_schedules; + +ALTER TABLE user_seed_wrappings + DROP CONSTRAINT IF EXISTS user_seed_wrappings_credential_kind_check; + +ALTER TABLE user_seed_wrappings + ADD CONSTRAINT user_seed_wrappings_credential_kind_check + CHECK (credential_kind IN ('password', 'oauth', 'agent_background')); + +ALTER TABLE agent_schedules + ADD COLUMN description_enc BYTEA NOT NULL, + DROP COLUMN description; + +CREATE TABLE agent_background_grants ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + project_id INTEGER NOT NULL REFERENCES org_projects(id) ON DELETE CASCADE, + agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + schedule_id BIGINT NOT NULL REFERENCES agent_schedules(id) ON DELETE CASCADE, + grant_enc BYTEA NOT NULL, + seed_wrap_lookup_hash BYTEA NOT NULL, + revoked_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE UNIQUE INDEX idx_agent_background_grants_schedule_active + ON agent_background_grants(schedule_id) + WHERE revoked_at IS NULL; + +CREATE INDEX idx_agent_background_grants_user + ON agent_background_grants(user_id, revoked_at); + +CREATE TRIGGER update_agent_background_grants_updated_at + BEFORE UPDATE ON agent_background_grants + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); + +ALTER TABLE notification_events + DROP CONSTRAINT IF EXISTS notification_events_source_kind_check; + +ALTER TABLE notification_events + ADD CONSTRAINT notification_events_source_kind_check + CHECK (source_kind IN ('request_continuation', 'agent_background')); + +ALTER TABLE notification_events + ADD COLUMN background_grant_id BIGINT REFERENCES agent_background_grants(id) ON DELETE SET NULL; + +CREATE INDEX idx_notification_events_background_grant + ON notification_events(background_grant_id) + WHERE background_grant_id IS NOT NULL; diff --git a/src/agent_background.rs b/src/agent_background.rs new file mode 100644 index 00000000..6c842c4e --- /dev/null +++ b/src/agent_background.rs @@ -0,0 +1,181 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use crate::encrypt::{ + decrypt_aead_v1, derive_key, encrypt_aead_v1, generate_random, CanonicalBytes, EncryptError, +}; + +const BACKGROUND_GRANT_AEAD_INFO: &[u8] = b"os.agent-background-grant-aead.v1"; +const BACKGROUND_GRANT_AAD_DOMAIN: &str = "os.agent-background-grant.v1"; + +pub const AGENT_BACKGROUND_GRANT_VERSION_V1: i16 = 1; +pub const AGENT_BACKGROUND_CAPABILITY_FULL_SCHEDULED_AGENT: &str = "full_scheduled_agent"; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AgentBackgroundGrantPlaintextV1 { + pub v: i16, + pub grant_uuid: Uuid, + pub user_uuid: Uuid, + pub project_id: i32, + pub agent_uuid: Uuid, + pub schedule_uuid: Uuid, + pub instruction_hash: Vec, + pub may_decrypt_user_content: bool, + pub may_send_push: bool, + pub capability_class: String, + pub created_from_auth_method: String, + pub created_at: DateTime, + pub background_secret: [u8; 32], +} + +impl AgentBackgroundGrantPlaintextV1 { + pub fn new_scheduled( + grant_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + agent_uuid: Uuid, + schedule_uuid: Uuid, + instruction_hash: Vec, + created_from_auth_method: impl Into, + ) -> Self { + Self { + v: AGENT_BACKGROUND_GRANT_VERSION_V1, + grant_uuid, + user_uuid, + project_id, + agent_uuid, + schedule_uuid, + instruction_hash, + may_decrypt_user_content: true, + may_send_push: true, + capability_class: AGENT_BACKGROUND_CAPABILITY_FULL_SCHEDULED_AGENT.to_string(), + created_from_auth_method: created_from_auth_method.into(), + created_at: Utc::now(), + background_secret: generate_random::<32>(), + } + } + + pub fn verify_scheduled_policy( + &self, + user_uuid: Uuid, + project_id: i32, + agent_uuid: Uuid, + schedule_uuid: Uuid, + instruction_hash: &[u8], + ) -> bool { + self.v == AGENT_BACKGROUND_GRANT_VERSION_V1 + && self.user_uuid == user_uuid + && self.project_id == project_id + && self.agent_uuid == agent_uuid + && self.schedule_uuid == schedule_uuid + && self.instruction_hash == instruction_hash + && self.may_decrypt_user_content + && self.capability_class == AGENT_BACKGROUND_CAPABILITY_FULL_SCHEDULED_AGENT + } +} + +pub fn instruction_hash(instruction: &str) -> Vec { + Sha256::digest(instruction.as_bytes()).to_vec() +} + +pub fn encrypt_background_grant_v1( + root_key: &[u8], + plaintext: &AgentBackgroundGrantPlaintextV1, +) -> Result, EncryptError> { + let key = derive_key(root_key, BACKGROUND_GRANT_AEAD_INFO)?; + let aad = background_grant_aad_v1( + plaintext.grant_uuid, + plaintext.user_uuid, + plaintext.project_id, + plaintext.agent_uuid, + plaintext.schedule_uuid, + plaintext.v, + ); + let bytes = serde_json::to_vec(plaintext).map_err(|_| EncryptError::BadData)?; + encrypt_aead_v1(&key, &bytes, &aad) +} + +pub fn decrypt_background_grant_v1( + root_key: &[u8], + encrypted: &[u8], + grant_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + agent_uuid: Uuid, + schedule_uuid: Uuid, +) -> Result { + let key = derive_key(root_key, BACKGROUND_GRANT_AEAD_INFO)?; + let aad = background_grant_aad_v1( + grant_uuid, + user_uuid, + project_id, + agent_uuid, + schedule_uuid, + AGENT_BACKGROUND_GRANT_VERSION_V1, + ); + let bytes = decrypt_aead_v1(&key, encrypted, &aad)?; + serde_json::from_slice(&bytes).map_err(|_| EncryptError::BadData) +} + +fn background_grant_aad_v1( + grant_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + agent_uuid: Uuid, + schedule_uuid: Uuid, + version: i16, +) -> Vec { + let mut aad = CanonicalBytes::new(BACKGROUND_GRANT_AAD_DOMAIN); + aad.append_uuid(grant_uuid) + .append_uuid(user_uuid) + .append_i32(project_id) + .append_uuid(agent_uuid) + .append_uuid(schedule_uuid) + .append_i16(version); + aad.into_bytes() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn background_grant_aead_binds_principal_context() { + let root = [9_u8; 32]; + let user = Uuid::new_v4(); + let agent = Uuid::new_v4(); + let schedule = Uuid::new_v4(); + let grant = Uuid::new_v4(); + let plaintext = AgentBackgroundGrantPlaintextV1::new_scheduled( + grant, + user, + 17, + agent, + schedule, + instruction_hash("check the weather"), + "password", + ); + + let encrypted = encrypt_background_grant_v1(&root, &plaintext).unwrap(); + let decrypted = + decrypt_background_grant_v1(&root, &encrypted, grant, user, 17, agent, schedule) + .unwrap(); + assert_eq!(decrypted.user_uuid, user); + + assert!( + decrypt_background_grant_v1( + &root, + &encrypted, + grant, + Uuid::new_v4(), + 17, + agent, + schedule, + ) + .is_err(), + "grant ciphertext must not decrypt under a different user principal", + ); + } +} diff --git a/src/db.rs b/src/db.rs index 6f6681f5..e9836593 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1018,10 +1018,11 @@ impl DBConnection for PostgresConnection { new_wrapping: NewUserSeedWrapping, ) -> Result<(), DBError> { use crate::models::schema::{ - agent_schedule_runs, agent_schedules, agents, conversation_projects, - conversation_summaries, conversations, memory_blocks, notification_events, - password_reset_requests, push_devices, user_embeddings, user_instructions, user_kv, - user_oauth_connections, user_preferences, user_seed_wrappings, users, + agent_background_grants, agent_schedule_runs, agent_schedules, agents, + conversation_projects, conversation_summaries, conversations, memory_blocks, + notification_events, password_reset_requests, push_devices, user_embeddings, + user_instructions, user_kv, user_oauth_connections, user_preferences, + user_seed_wrappings, users, }; debug!("Completing destructive password reset"); @@ -1071,6 +1072,10 @@ impl DBConnection for PostgresConnection { agent_schedule_runs::table.filter(agent_schedule_runs::user_id.eq(user_id)), ) .execute(conn)?; + diesel::delete( + agent_background_grants::table.filter(agent_background_grants::user_id.eq(user_id)), + ) + .execute(conn)?; diesel::delete(agent_schedules::table.filter(agent_schedules::user_id.eq(user_id))) .execute(conn)?; diesel::delete(agents::table.filter(agents::user_id.eq(user_id))).execute(conn)?; diff --git a/src/main.rs b/src/main.rs index c0ead20d..4bd7290d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,6 +18,7 @@ use crate::models::platform_password_reset::NewPlatformPasswordResetRequest; use crate::models::platform_users::PlatformUser; use crate::push::worker::start_push_worker; use crate::sqs::SqsEventPublisher; +use crate::web::agent::start_schedule_worker; use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ @@ -30,6 +31,7 @@ use crate::{attestation_routes::SessionState, web::platform_routes}; use crate::jwt::{AuthContext, AuthMethod}; use crate::models::user_seed_wrappings::{NewUserSeedWrapping, UserSeedWrappingError}; use crate::seed_wrapping::{ + agent_background_credential_lookup_hash, compute_agent_background_auth_binding, compute_oauth_auth_binding, compute_password_auth_binding, decrypt_seed_v1, encrypt_seed_v1, normalize_email_login_identifier, normalize_guest_login_identifier, oauth_credential_lookup_hash, password_credential_lookup_hash, password_reset_code_mac, @@ -81,6 +83,7 @@ use vsock::{VsockAddr, VsockStream}; use web::attestation_routes; use x25519_dalek::{EphemeralSecret, PublicKey}; +mod agent_background; mod apple_signin; mod aws_credentials; mod billing; @@ -1228,6 +1231,94 @@ impl AppState { Ok(()) } + fn new_agent_background_seed_wrapping_for_user( + &self, + user: &User, + grant_uuid: Uuid, + background_secret: &[u8; 32], + plaintext_seed: &[u8], + ) -> Result { + let auth_binding = compute_agent_background_auth_binding( + &self.enclave_key, + grant_uuid, + user.uuid, + user.project_id, + background_secret, + ) + .map_err(|e| Error::EncryptionError(e.to_string()))?; + let credential_lookup_hash = + agent_background_credential_lookup_hash(&self.enclave_key, grant_uuid) + .map_err(|e| Error::EncryptionError(e.to_string()))?; + + let seed_enc = encrypt_seed_v1( + &self.enclave_key, + plaintext_seed, + user.uuid, + user.project_id, + CredentialKind::AgentBackground, + &auth_binding, + ) + .map_err(|e| Error::EncryptionError(e.to_string()))?; + + Ok(NewUserSeedWrapping::new( + user.uuid, + CredentialKind::AgentBackground.as_str(), + credential_lookup_hash.as_bytes().to_vec(), + SEED_WRAP_VERSION_V1, + seed_enc, + )) + } + + async fn get_user_key_for_agent_background_grant( + &self, + user: &User, + grant_plaintext: &crate::agent_background::AgentBackgroundGrantPlaintextV1, + seed_wrap_lookup_hash: &[u8], + ) -> Result { + if user.uuid != grant_plaintext.user_uuid || user.project_id != grant_plaintext.project_id { + return Err(Error::AuthenticationError); + } + + let expected_lookup_hash = + agent_background_credential_lookup_hash(&self.enclave_key, grant_plaintext.grant_uuid) + .map_err(|e| Error::EncryptionError(e.to_string()))?; + if seed_wrap_lookup_hash != expected_lookup_hash.as_bytes() { + return Err(Error::AuthenticationError); + } + + let auth_binding = compute_agent_background_auth_binding( + &self.enclave_key, + grant_plaintext.grant_uuid, + user.uuid, + user.project_id, + &grant_plaintext.background_secret, + ) + .map_err(|e| Error::EncryptionError(e.to_string()))?; + + let seed_wrap = self + .db + .get_user_seed_wrapping_by_credential( + user.uuid, + CredentialKind::AgentBackground.as_str(), + expected_lookup_hash.as_bytes(), + SEED_WRAP_VERSION_V1, + )? + .ok_or(Error::AuthenticationError)?; + + let plaintext_seed = decrypt_seed_v1( + &self.enclave_key, + &seed_wrap.seed_enc, + user.uuid, + user.project_id, + CredentialKind::AgentBackground, + &auth_binding, + ) + .map_err(|_| Error::AuthenticationError)?; + + plaintext_user_seed_to_key(&plaintext_seed, None, None) + .map_err(|e| Error::EncryptionError(e.to_string())) + } + async fn register_user(&self, creds: RegisterCredentials) -> Result { // Get project by client_id let project = self @@ -3199,6 +3290,7 @@ async fn main() -> Result<(), Error> { .await?; start_push_worker(app_state.clone()); + start_schedule_worker(app_state.clone()); let cors = CorsLayer::new() // allow all method types diff --git a/src/models/agent_background_grants.rs b/src/models/agent_background_grants.rs new file mode 100644 index 00000000..18f0fcd6 --- /dev/null +++ b/src/models/agent_background_grants.rs @@ -0,0 +1,76 @@ +use crate::models::schema::agent_background_grants; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use thiserror::Error; +use uuid::Uuid; + +#[derive(Error, Debug)] +pub enum AgentBackgroundGrantError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, Clone, Debug)] +#[diesel(table_name = agent_background_grants)] +pub struct AgentBackgroundGrant { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub project_id: i32, + pub agent_id: i64, + pub schedule_id: i64, + pub grant_enc: Vec, + pub seed_wrap_lookup_hash: Vec, + pub revoked_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl AgentBackgroundGrant { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, AgentBackgroundGrantError> { + agent_background_grants::table + .filter(agent_background_grants::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(AgentBackgroundGrantError::DatabaseError) + } + + pub fn get_active_by_schedule( + conn: &mut PgConnection, + lookup_schedule_id: i64, + ) -> Result, AgentBackgroundGrantError> { + agent_background_grants::table + .filter(agent_background_grants::schedule_id.eq(lookup_schedule_id)) + .filter(agent_background_grants::revoked_at.is_null()) + .first::(conn) + .optional() + .map_err(AgentBackgroundGrantError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = agent_background_grants)] +pub struct NewAgentBackgroundGrant { + pub uuid: Uuid, + pub user_id: Uuid, + pub project_id: i32, + pub agent_id: i64, + pub schedule_id: i64, + pub grant_enc: Vec, + pub seed_wrap_lookup_hash: Vec, +} + +impl NewAgentBackgroundGrant { + pub fn insert( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(agent_background_grants::table) + .values(self) + .get_result::(conn) + .map_err(AgentBackgroundGrantError::DatabaseError) + } +} diff --git a/src/models/agent_schedule_runs.rs b/src/models/agent_schedule_runs.rs new file mode 100644 index 00000000..cbba01e1 --- /dev/null +++ b/src/models/agent_schedule_runs.rs @@ -0,0 +1,364 @@ +use crate::models::schema::agent_schedule_runs; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{BigInt, Int4, Nullable, Text}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +pub const AGENT_SCHEDULE_RUN_STATUS_PENDING: &str = "pending"; +pub const AGENT_SCHEDULE_RUN_STATUS_RETRY: &str = "retry"; +pub const AGENT_SCHEDULE_RUN_STATUS_COMPLETED: &str = "completed"; +pub const AGENT_SCHEDULE_RUN_STATUS_FAILED: &str = "failed"; +pub const AGENT_SCHEDULE_RUN_STATUS_CANCELLED: &str = "cancelled"; +pub const AGENT_SCHEDULE_RUN_STATUS_EXPIRED: &str = "expired"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentScheduleRunWriteResult { + Updated, + LostLease, +} + +impl AgentScheduleRunWriteResult { + fn from_updated_rows(updated_rows: usize) -> Self { + if updated_rows == 0 { + Self::LostLease + } else { + Self::Updated + } + } + + pub fn was_applied(self) -> bool { + matches!(self, Self::Updated) + } +} + +#[derive(Error, Debug)] +pub enum AgentScheduleRunError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, QueryableByName, Clone, Debug, Serialize, Deserialize)] +#[diesel(table_name = agent_schedule_runs)] +pub struct AgentScheduleRun { + pub id: i64, + pub uuid: Uuid, + pub schedule_id: i64, + pub user_id: Uuid, + pub agent_id: i64, + pub scheduled_for: DateTime, + pub stale_after_at: DateTime, + pub status: String, + pub attempt_count: i32, + pub next_attempt_at: DateTime, + pub lease_owner: Option, + pub lease_expires_at: Option>, + pub started_at: Option>, + pub first_output_at: Option>, + pub first_message_id: Option, + pub output_count: i32, + pub notification_enqueued_at: Option>, + pub completed_at: Option>, + pub last_error: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl AgentScheduleRun { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, AgentScheduleRunError> { + agent_schedule_runs::table + .filter(agent_schedule_runs::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn lease_pending( + conn: &mut PgConnection, + limit: i64, + lease_owner: &str, + lease_seconds: i32, + ) -> Result, AgentScheduleRunError> { + let query = r#" + WITH candidates AS ( + SELECT r.id + FROM agent_schedule_runs r + WHERE r.next_attempt_at <= NOW() + AND ( + (r.status IN ('pending', 'retry') + AND (r.lease_expires_at IS NULL OR r.lease_expires_at < NOW())) + OR (r.status = 'leased' + AND (r.lease_expires_at IS NULL OR r.lease_expires_at < NOW())) + ) + ORDER BY r.next_attempt_at ASC, r.id ASC + FOR UPDATE OF r SKIP LOCKED + LIMIT $1 + ) + UPDATE agent_schedule_runs r + SET status = 'leased', + lease_owner = $2, + lease_expires_at = NOW() + ($3 * INTERVAL '1 second'), + started_at = COALESCE(r.started_at, NOW()), + updated_at = NOW() + FROM candidates + WHERE r.id = candidates.id + RETURNING r.* + "#; + + sql_query(query) + .bind::(limit) + .bind::(lease_owner) + .bind::(lease_seconds) + .get_results::(conn) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn renew_lease( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + lease_seconds: i32, + ) -> Result { + let query = r#" + UPDATE agent_schedule_runs + SET lease_expires_at = NOW() + ($3 * INTERVAL '1 second'), + updated_at = NOW() + WHERE id = $1 + AND status = 'leased' + AND lease_owner = $2 + "#; + + sql_query(query) + .bind::(lookup_id) + .bind::(expected_lease_owner) + .bind::(lease_seconds) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn record_output( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + first_message_id: Uuid, + ) -> Result { + let query = r#" + UPDATE agent_schedule_runs + SET first_output_at = COALESCE(first_output_at, NOW()), + first_message_id = COALESCE(first_message_id, $3), + output_count = output_count + 1, + updated_at = NOW() + WHERE id = $1 + AND status = 'leased' + AND lease_owner = $2 + "#; + + sql_query(query) + .bind::(lookup_id) + .bind::(expected_lease_owner) + .bind::(first_message_id) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn mark_retry( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + last_error: Option<&str>, + retry_after_seconds: i32, + ) -> Result { + let query = r#" + UPDATE agent_schedule_runs + SET status = 'retry', + attempt_count = attempt_count + 1, + last_error = $3, + next_attempt_at = NOW() + ($4 * INTERVAL '1 second'), + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + AND status = 'leased' + AND lease_owner = $2 + "#; + + sql_query(query) + .bind::(lookup_id) + .bind::(expected_lease_owner) + .bind::, _>(last_error) + .bind::(retry_after_seconds) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn mark_completed( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + notification_enqueued: bool, + last_error: Option<&str>, + ) -> Result { + diesel::update( + agent_schedule_runs::table + .filter(agent_schedule_runs::id.eq(lookup_id)) + .filter(agent_schedule_runs::status.eq("leased")) + .filter( + agent_schedule_runs::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + agent_schedule_runs::status.eq(AGENT_SCHEDULE_RUN_STATUS_COMPLETED), + agent_schedule_runs::attempt_count.eq(agent_schedule_runs::attempt_count + 1), + agent_schedule_runs::last_error.eq(last_error), + agent_schedule_runs::notification_enqueued_at.eq(if notification_enqueued { + Some(Utc::now()) + } else { + None + }), + agent_schedule_runs::completed_at.eq(diesel::dsl::now), + agent_schedule_runs::lease_owner.eq::>(None), + agent_schedule_runs::lease_expires_at.eq::>>(None), + agent_schedule_runs::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn mark_failed( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + last_error: Option<&str>, + ) -> Result { + diesel::update( + agent_schedule_runs::table + .filter(agent_schedule_runs::id.eq(lookup_id)) + .filter(agent_schedule_runs::status.eq("leased")) + .filter( + agent_schedule_runs::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + agent_schedule_runs::status.eq(AGENT_SCHEDULE_RUN_STATUS_FAILED), + agent_schedule_runs::attempt_count.eq(agent_schedule_runs::attempt_count + 1), + agent_schedule_runs::last_error.eq(last_error), + agent_schedule_runs::completed_at.eq(diesel::dsl::now), + agent_schedule_runs::lease_owner.eq::>(None), + agent_schedule_runs::lease_expires_at.eq::>>(None), + agent_schedule_runs::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn mark_expired( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + last_error: Option<&str>, + ) -> Result { + diesel::update( + agent_schedule_runs::table + .filter(agent_schedule_runs::id.eq(lookup_id)) + .filter(agent_schedule_runs::status.eq("leased")) + .filter( + agent_schedule_runs::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + agent_schedule_runs::status.eq(AGENT_SCHEDULE_RUN_STATUS_EXPIRED), + agent_schedule_runs::attempt_count.eq(agent_schedule_runs::attempt_count + 1), + agent_schedule_runs::last_error.eq(last_error), + agent_schedule_runs::completed_at.eq(diesel::dsl::now), + agent_schedule_runs::lease_owner.eq::>(None), + agent_schedule_runs::lease_expires_at.eq::>>(None), + agent_schedule_runs::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn mark_cancelled( + conn: &mut PgConnection, + lookup_id: i64, + expected_lease_owner: &str, + last_error: Option<&str>, + ) -> Result { + diesel::update( + agent_schedule_runs::table + .filter(agent_schedule_runs::id.eq(lookup_id)) + .filter(agent_schedule_runs::status.eq("leased")) + .filter( + agent_schedule_runs::lease_owner.eq(Some(expected_lease_owner.to_string())), + ), + ) + .set(( + agent_schedule_runs::status.eq(AGENT_SCHEDULE_RUN_STATUS_CANCELLED), + agent_schedule_runs::last_error.eq(last_error), + agent_schedule_runs::completed_at.eq(diesel::dsl::now), + agent_schedule_runs::lease_owner.eq::>(None), + agent_schedule_runs::lease_expires_at.eq::>>(None), + agent_schedule_runs::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map(AgentScheduleRunWriteResult::from_updated_rows) + .map_err(AgentScheduleRunError::DatabaseError) + } + + pub fn cancel_unstarted_for_schedule( + conn: &mut PgConnection, + lookup_schedule_id: i64, + ) -> Result { + diesel::update( + agent_schedule_runs::table + .filter(agent_schedule_runs::schedule_id.eq(lookup_schedule_id)) + .filter(agent_schedule_runs::status.eq_any(vec![ + AGENT_SCHEDULE_RUN_STATUS_PENDING, + AGENT_SCHEDULE_RUN_STATUS_RETRY, + ])), + ) + .set(( + agent_schedule_runs::status.eq(AGENT_SCHEDULE_RUN_STATUS_CANCELLED), + agent_schedule_runs::completed_at.eq(diesel::dsl::now), + agent_schedule_runs::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map_err(AgentScheduleRunError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = agent_schedule_runs)] +pub struct NewAgentScheduleRun { + pub uuid: Uuid, + pub schedule_id: i64, + pub user_id: Uuid, + pub agent_id: i64, + pub scheduled_for: DateTime, + pub stale_after_at: DateTime, + pub status: String, + pub next_attempt_at: DateTime, +} + +impl NewAgentScheduleRun { + pub fn insert( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(agent_schedule_runs::table) + .values(self) + .get_result::(conn) + .map_err(AgentScheduleRunError::DatabaseError) + } +} diff --git a/src/models/agent_schedules.rs b/src/models/agent_schedules.rs new file mode 100644 index 00000000..67f5f33f --- /dev/null +++ b/src/models/agent_schedules.rs @@ -0,0 +1,342 @@ +use crate::models::schema::agent_schedules; +use chrono::{DateTime, NaiveDate, NaiveTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use thiserror::Error; +use uuid::Uuid; + +pub const SCHEDULE_KIND_ONE_OFF: &str = "one_off"; +pub const SCHEDULE_KIND_RECURRING: &str = "recurring"; + +pub const RECURRENCE_TYPE_INTERVAL: &str = "interval"; +pub const RECURRENCE_TYPE_DAILY: &str = "daily"; +pub const RECURRENCE_TYPE_WEEKLY: &str = "weekly"; + +pub const SCHEDULE_TIMEZONE_MODE_FOLLOW_USER: &str = "follow_user"; +pub const SCHEDULE_TIMEZONE_MODE_FIXED: &str = "fixed"; + +pub const SCHEDULE_STATUS_ACTIVE: &str = "active"; +pub const SCHEDULE_STATUS_COMPLETED: &str = "completed"; +pub const SCHEDULE_STATUS_CANCELLED: &str = "cancelled"; + +#[derive(Error, Debug)] +pub enum AgentScheduleError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), + #[error("Invalid schedule spec: {0}")] + InvalidSpec(String), +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ScheduleWeekday { + Monday, + Tuesday, + Wednesday, + Thursday, + Friday, + Saturday, + Sunday, +} + +impl ScheduleWeekday { + pub fn as_str(&self) -> &'static str { + match self { + Self::Monday => "monday", + Self::Tuesday => "tuesday", + Self::Wednesday => "wednesday", + Self::Thursday => "thursday", + Self::Friday => "friday", + Self::Saturday => "saturday", + Self::Sunday => "sunday", + } + } + + pub fn to_chrono(&self) -> chrono::Weekday { + match self { + Self::Monday => chrono::Weekday::Mon, + Self::Tuesday => chrono::Weekday::Tue, + Self::Wednesday => chrono::Weekday::Wed, + Self::Thursday => chrono::Weekday::Thu, + Self::Friday => chrono::Weekday::Fri, + Self::Saturday => chrono::Weekday::Sat, + Self::Sunday => chrono::Weekday::Sun, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ScheduleIntervalUnit { + Hours, +} + +impl ScheduleIntervalUnit { + pub fn as_str(&self) -> &'static str { + match self { + Self::Hours => "hours", + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ScheduleSpec { + OneOff { + local_date: String, + local_time: String, + }, + Interval { + every_n: u32, + interval_unit: ScheduleIntervalUnit, + }, + Daily { + local_time: String, + }, + Weekly { + local_time: String, + weekdays: Vec, + }, +} + +impl ScheduleSpec { + pub fn schedule_kind(&self) -> &'static str { + match self { + Self::OneOff { .. } => SCHEDULE_KIND_ONE_OFF, + Self::Interval { .. } | Self::Daily { .. } | Self::Weekly { .. } => { + SCHEDULE_KIND_RECURRING + } + } + } + + pub fn recurrence_type(&self) -> Option<&'static str> { + match self { + Self::OneOff { .. } => None, + Self::Interval { .. } => Some(RECURRENCE_TYPE_INTERVAL), + Self::Daily { .. } => Some(RECURRENCE_TYPE_DAILY), + Self::Weekly { .. } => Some(RECURRENCE_TYPE_WEEKLY), + } + } + + pub fn validate(&self) -> Result<(), AgentScheduleError> { + match self { + Self::OneOff { + local_date, + local_time, + } => { + parse_local_date(local_date)?; + parse_local_time(local_time)?; + Ok(()) + } + Self::Interval { + every_n, + interval_unit, + } => { + if *every_n == 0 { + return Err(AgentScheduleError::InvalidSpec( + "interval every_n must be greater than zero".to_string(), + )); + } + + match interval_unit { + ScheduleIntervalUnit::Hours => Ok(()), + } + } + Self::Daily { local_time } => { + parse_local_time(local_time)?; + Ok(()) + } + Self::Weekly { + local_time, + weekdays, + } => { + parse_local_time(local_time)?; + if weekdays.is_empty() { + return Err(AgentScheduleError::InvalidSpec( + "weekly schedules require at least one weekday".to_string(), + )); + } + + Ok(()) + } + } + } + + pub fn summary(&self) -> String { + match self { + Self::OneOff { + local_date, + local_time, + } => format!("once on {} at {}", local_date, local_time), + Self::Interval { + every_n, + interval_unit, + } => format!("every {} {}", every_n, interval_unit.as_str()), + Self::Daily { local_time } => format!("every day at {}", local_time), + Self::Weekly { + local_time, + weekdays, + } => { + let day_list = weekdays + .iter() + .map(ScheduleWeekday::as_str) + .collect::>() + .join(", "); + format!("every {} at {}", day_list, local_time) + } + } + } + + pub fn to_value(&self) -> Result { + serde_json::to_value(self) + .map_err(|e| AgentScheduleError::InvalidSpec(format!("serialize schedule spec: {e}"))) + } + + pub fn from_value(value: &Value) -> Result { + serde_json::from_value(value.clone()) + .map_err(|e| AgentScheduleError::InvalidSpec(format!("parse schedule spec: {e}"))) + } +} + +pub fn parse_local_time(value: &str) -> Result { + NaiveTime::parse_from_str(value.trim(), "%H:%M").map_err(|_| { + AgentScheduleError::InvalidSpec(format!( + "invalid local_time '{value}'. Use 24-hour HH:MM format" + )) + }) +} + +pub fn parse_local_date(value: &str) -> Result { + NaiveDate::parse_from_str(value.trim(), "%Y-%m-%d").map_err(|_| { + AgentScheduleError::InvalidSpec(format!( + "invalid local_date '{value}'. Use YYYY-MM-DD format" + )) + }) +} + +#[derive(Queryable, Identifiable, AsChangeset, Clone, Debug, Serialize, Deserialize)] +#[diesel(table_name = agent_schedules)] +pub struct AgentSchedule { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub agent_id: i64, + pub description_enc: Vec, + pub instruction_enc: Vec, + pub schedule_kind: String, + pub recurrence_type: Option, + pub schedule_spec: Value, + pub timezone_mode: String, + pub resolved_timezone: String, + pub fixed_timezone: Option, + pub stale_after_minutes: i32, + pub status: String, + pub next_scheduled_for: Option>, + pub last_scheduled_for: Option>, + pub last_run_at: Option>, + pub run_count: i32, + pub cancelled_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl AgentSchedule { + pub fn get_by_id( + conn: &mut PgConnection, + lookup_id: i64, + ) -> Result, AgentScheduleError> { + agent_schedules::table + .filter(agent_schedules::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(AgentScheduleError::DatabaseError) + } + + pub fn get_by_uuid_and_agent( + conn: &mut PgConnection, + lookup_uuid: Uuid, + lookup_user_id: Uuid, + lookup_agent_id: i64, + ) -> Result, AgentScheduleError> { + agent_schedules::table + .filter(agent_schedules::uuid.eq(lookup_uuid)) + .filter(agent_schedules::user_id.eq(lookup_user_id)) + .filter(agent_schedules::agent_id.eq(lookup_agent_id)) + .first::(conn) + .optional() + .map_err(AgentScheduleError::DatabaseError) + } + + pub fn list_by_agent( + conn: &mut PgConnection, + lookup_user_id: Uuid, + lookup_agent_id: i64, + status_filter: Option<&str>, + ) -> Result, AgentScheduleError> { + let mut query = agent_schedules::table + .filter(agent_schedules::user_id.eq(lookup_user_id)) + .filter(agent_schedules::agent_id.eq(lookup_agent_id)) + .into_boxed(); + + if let Some(status_filter) = status_filter { + query = query.filter(agent_schedules::status.eq(status_filter)); + } + + query + .order(( + agent_schedules::next_scheduled_for.asc().nulls_last(), + agent_schedules::created_at.desc(), + )) + .load::(conn) + .map_err(AgentScheduleError::DatabaseError) + } + + pub fn list_active_follow_user_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, AgentScheduleError> { + agent_schedules::table + .filter(agent_schedules::user_id.eq(lookup_user_id)) + .filter(agent_schedules::status.eq(SCHEDULE_STATUS_ACTIVE)) + .filter(agent_schedules::timezone_mode.eq(SCHEDULE_TIMEZONE_MODE_FOLLOW_USER)) + .load::(conn) + .map_err(AgentScheduleError::DatabaseError) + } + + pub fn spec(&self) -> Result { + ScheduleSpec::from_value(&self.schedule_spec) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = agent_schedules)] +pub struct NewAgentSchedule { + pub uuid: Uuid, + pub user_id: Uuid, + pub agent_id: i64, + pub description_enc: Vec, + pub instruction_enc: Vec, + pub schedule_kind: String, + pub recurrence_type: Option, + pub schedule_spec: Value, + pub timezone_mode: String, + pub resolved_timezone: String, + pub fixed_timezone: Option, + pub stale_after_minutes: i32, + pub status: String, + pub next_scheduled_for: Option>, + pub last_scheduled_for: Option>, + pub last_run_at: Option>, + pub run_count: i32, + pub cancelled_at: Option>, +} + +impl NewAgentSchedule { + pub fn insert(&self, conn: &mut PgConnection) -> Result { + diesel::insert_into(agent_schedules::table) + .values(self) + .get_result::(conn) + .map_err(AgentScheduleError::DatabaseError) + } +} diff --git a/src/models/agents.rs b/src/models/agents.rs index 816b5f2b..d9ab2295 100644 --- a/src/models/agents.rs +++ b/src/models/agents.rs @@ -35,6 +35,14 @@ pub struct Agent { } impl Agent { + pub fn get_by_id(conn: &mut PgConnection, lookup_id: i64) -> Result, AgentError> { + agents::table + .filter(agents::id.eq(lookup_id)) + .first::(conn) + .optional() + .map_err(AgentError::DatabaseError) + } + pub fn get_main_for_user( conn: &mut PgConnection, lookup_user_id: Uuid, diff --git a/src/models/mod.rs b/src/models/mod.rs index cd2756a8..817cbb76 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,4 +1,7 @@ pub mod account_deletion; +pub mod agent_background_grants; +pub mod agent_schedule_runs; +pub mod agent_schedules; pub mod agents; pub mod app_data_migrations; pub mod conversation_summaries; diff --git a/src/models/notification_events.rs b/src/models/notification_events.rs index e31e92ad..666c6a6e 100644 --- a/src/models/notification_events.rs +++ b/src/models/notification_events.rs @@ -11,6 +11,7 @@ pub const NOTIFICATION_DELIVERY_MODE_ENCRYPTED_PREVIEW: &str = "encrypted_previe pub const NOTIFICATION_PRIORITY_NORMAL: &str = "normal"; pub const NOTIFICATION_PRIORITY_HIGH: &str = "high"; pub const NOTIFICATION_SOURCE_REQUEST_CONTINUATION: &str = "request_continuation"; +pub const NOTIFICATION_SOURCE_AGENT_BACKGROUND: &str = "agent_background"; #[derive(Error, Debug)] pub enum NotificationEventError { @@ -38,6 +39,7 @@ pub struct NotificationEvent { pub cancelled_at: Option>, pub source_kind: String, pub source_request_id: Option, + pub background_grant_id: Option, } impl NotificationEvent { @@ -81,6 +83,7 @@ pub struct NewNotificationEvent { pub expires_at: Option>, pub source_kind: String, pub source_request_id: Option, + pub background_grant_id: Option, } impl NewNotificationEvent { diff --git a/src/models/schema.rs b/src/models/schema.rs index 1c5e642c..f88d0838 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -21,6 +21,22 @@ diesel::table! { } } +diesel::table! { + agent_background_grants (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + project_id -> Int4, + agent_id -> Int8, + schedule_id -> Int8, + grant_enc -> Bytea, + seed_wrap_lookup_hash -> Bytea, + revoked_at -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { agent_schedule_runs (id) { id -> Int8, @@ -53,7 +69,7 @@ diesel::table! { uuid -> Uuid, user_id -> Uuid, agent_id -> Int8, - description -> Text, + description_enc -> Bytea, instruction_enc -> Bytea, schedule_kind -> Text, recurrence_type -> Nullable, @@ -240,6 +256,7 @@ diesel::table! { cancelled_at -> Nullable, source_kind -> Text, source_request_id -> Nullable, + background_grant_id -> Nullable, } } @@ -607,6 +624,9 @@ diesel::table! { } } +diesel::joinable!(agent_background_grants -> agent_schedules (schedule_id)); +diesel::joinable!(agent_background_grants -> agents (agent_id)); +diesel::joinable!(agent_background_grants -> org_projects (project_id)); diesel::joinable!(agent_schedule_runs -> agent_schedules (schedule_id)); diesel::joinable!(agent_schedule_runs -> agents (agent_id)); diesel::joinable!(agent_schedules -> agents (agent_id)); @@ -618,6 +638,7 @@ diesel::joinable!(conversations -> conversation_projects (project_id)); diesel::joinable!(invite_codes -> orgs (org_id)); diesel::joinable!(notification_deliveries -> notification_events (event_id)); diesel::joinable!(notification_deliveries -> push_devices (push_device_id)); +diesel::joinable!(notification_events -> agent_background_grants (background_grant_id)); diesel::joinable!(notification_events -> org_projects (project_id)); diesel::joinable!(org_memberships -> orgs (org_id)); diesel::joinable!(org_project_secrets -> org_projects (project_id)); @@ -643,6 +664,7 @@ diesel::joinable!(users -> org_projects (project_id)); diesel::allow_tables_to_appear_in_same_query!( account_deletion_requests, + agent_background_grants, agent_schedule_runs, agent_schedules, agents, diff --git a/src/push/binding.rs b/src/push/binding.rs index 0f5ee018..b4ad5120 100644 --- a/src/push/binding.rs +++ b/src/push/binding.rs @@ -3,7 +3,9 @@ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use uuid::Uuid; +use crate::agent_background::{decrypt_background_grant_v1, AgentBackgroundGrantPlaintextV1}; use crate::encrypt::{decrypt_aead_v1, derive_key, encrypt_aead_v1, CanonicalBytes, EncryptError}; +use crate::models::agent_background_grants::AgentBackgroundGrant; use crate::models::notification_events::NotificationEvent; use crate::models::push_devices::PushDevice; @@ -104,7 +106,9 @@ pub struct NotificationPreviewPayloadV1 { pub project_id: i32, pub source_kind: String, pub source_request_id: Option, + pub background_grant_uuid: Option, pub agent_uuid: Option, + pub schedule_uuid: Option, pub delivery_mode: String, pub message_id: Uuid, pub kind: String, @@ -116,13 +120,18 @@ pub struct NotificationPreviewPayloadV1 { } impl NotificationPreviewPayloadV1 { - pub fn matches_event(&self, event: &NotificationEvent) -> bool { + pub fn matches_event( + &self, + event: &NotificationEvent, + background_grant_uuid: Option, + ) -> bool { self.v == PUSH_EVENT_PAYLOAD_VERSION_V1 && self.notification_id == event.uuid && self.user_uuid == event.user_id && self.project_id == event.project_id && self.source_kind == event.source_kind && self.source_request_id == event.source_request_id + && self.background_grant_uuid == background_grant_uuid && self.delivery_mode == event.delivery_mode && self.kind == event.kind } @@ -183,6 +192,7 @@ pub fn encrypt_notification_preview_payload_v1( user_uuid: plaintext.user_uuid, project_id: plaintext.project_id, source_kind: &plaintext.source_kind, + background_grant_uuid: plaintext.background_grant_uuid, kind: &plaintext.kind, delivery_mode: &plaintext.delivery_mode, version: plaintext.v, @@ -195,6 +205,7 @@ pub fn decrypt_notification_preview_payload_v1( root_key: &[u8], encrypted: &[u8], event: &NotificationEvent, + background_grant_uuid: Option, ) -> Result { let key = derive_key(root_key, PUSH_EVENT_PAYLOAD_AEAD_INFO)?; let aad = notification_payload_aad_v1(NotificationPayloadAad { @@ -202,6 +213,7 @@ pub fn decrypt_notification_preview_payload_v1( user_uuid: event.user_id, project_id: event.project_id, source_kind: &event.source_kind, + background_grant_uuid, kind: &event.kind, delivery_mode: &event.delivery_mode, version: PUSH_EVENT_PAYLOAD_VERSION_V1, @@ -210,6 +222,22 @@ pub fn decrypt_notification_preview_payload_v1( serde_json::from_slice(&bytes).map_err(|_| EncryptError::BadData) } +pub fn decrypt_background_grant_for_push( + root_key: &[u8], + grant: &AgentBackgroundGrant, + payload: &NotificationPreviewPayloadV1, +) -> Result { + decrypt_background_grant_v1( + root_key, + &grant.grant_enc, + grant.uuid, + payload.user_uuid, + payload.project_id, + payload.agent_uuid.ok_or(EncryptError::BadData)?, + payload.schedule_uuid.ok_or(EncryptError::BadData)?, + ) +} + struct PushDeviceAad<'a> { push_device_uuid: Uuid, user_uuid: Uuid, @@ -241,6 +269,7 @@ struct NotificationPayloadAad<'a> { user_uuid: Uuid, project_id: i32, source_kind: &'a str, + background_grant_uuid: Option, kind: &'a str, delivery_mode: &'a str, version: i16, @@ -252,6 +281,12 @@ fn notification_payload_aad_v1(input: NotificationPayloadAad<'_>) -> Vec { .append_uuid(input.user_uuid) .append_i32(input.project_id) .append_str(input.source_kind) + .append_str( + &input + .background_grant_uuid + .map(|uuid| uuid.to_string()) + .unwrap_or_default(), + ) .append_str(input.kind) .append_str(input.delivery_mode) .append_i16(input.version); diff --git a/src/push/crypto.rs b/src/push/crypto.rs index 2cbb610d..6bbdd457 100644 --- a/src/push/crypto.rs +++ b/src/push/crypto.rs @@ -93,7 +93,9 @@ mod tests { project_id: 42, source_kind: "request_continuation".to_string(), source_request_id: None, + background_grant_uuid: None, agent_uuid: None, + schedule_uuid: None, delivery_mode: "encrypted_preview".to_string(), message_id: Uuid::new_v4(), kind: "agent.message".to_string(), diff --git a/src/push/mod.rs b/src/push/mod.rs index 240f6b85..ac9dcef2 100644 --- a/src/push/mod.rs +++ b/src/push/mod.rs @@ -5,12 +5,13 @@ pub mod fcm; pub mod worker; use crate::db::DBError; +use crate::models::agent_background_grants::AgentBackgroundGrant; use crate::models::notification_deliveries::{NewNotificationDelivery, NotificationDeliveryError}; use crate::models::notification_events::{ NewNotificationEvent, NotificationEvent, NotificationEventError, NOTIFICATION_DELIVERY_MODE_ENCRYPTED_PREVIEW, NOTIFICATION_DELIVERY_MODE_GENERIC, NOTIFICATION_KIND_AGENT_MESSAGE, NOTIFICATION_PRIORITY_HIGH, NOTIFICATION_PRIORITY_NORMAL, - NOTIFICATION_SOURCE_REQUEST_CONTINUATION, + NOTIFICATION_SOURCE_AGENT_BACKGROUND, NOTIFICATION_SOURCE_REQUEST_CONTINUATION, }; use crate::models::project_settings::ProjectSettingError; use crate::models::push_devices::{PushDevice, PushDeviceError}; @@ -114,12 +115,26 @@ pub enum PushNotificationSource { source_request_id: Option, agent_uuid: Option, }, + AgentBackground { + grant_id: i64, + grant_uuid: Uuid, + agent_uuid: Uuid, + schedule_uuid: Uuid, + }, +} + +#[derive(Debug, Clone, Copy)] +pub struct BackgroundAgentPushSource<'a> { + pub grant: &'a AgentBackgroundGrant, + pub agent_uuid: Uuid, + pub schedule_uuid: Uuid, } impl PushNotificationSource { fn source_kind(&self) -> &'static str { match self { Self::RequestContinuation { .. } => NOTIFICATION_SOURCE_REQUEST_CONTINUATION, + Self::AgentBackground { .. } => NOTIFICATION_SOURCE_AGENT_BACKGROUND, } } @@ -128,12 +143,35 @@ impl PushNotificationSource { Self::RequestContinuation { source_request_id, .. } => *source_request_id, + Self::AgentBackground { .. } => None, + } + } + + fn background_grant_id(&self) -> Option { + match self { + Self::RequestContinuation { .. } => None, + Self::AgentBackground { grant_id, .. } => Some(*grant_id), + } + } + + fn background_grant_uuid(&self) -> Option { + match self { + Self::RequestContinuation { .. } => None, + Self::AgentBackground { grant_uuid, .. } => Some(*grant_uuid), } } fn agent_uuid(&self) -> Option { match self { Self::RequestContinuation { agent_uuid, .. } => *agent_uuid, + Self::AgentBackground { agent_uuid, .. } => Some(*agent_uuid), + } + } + + fn schedule_uuid(&self) -> Option { + match self { + Self::RequestContinuation { .. } => None, + Self::AgentBackground { schedule_uuid, .. } => Some(*schedule_uuid), } } } @@ -284,7 +322,9 @@ pub async fn enqueue_notification( project_id: request.project_id, source_kind: request.source.source_kind().to_string(), source_request_id: request.source.source_request_id(), + background_grant_uuid: request.source.background_grant_uuid(), agent_uuid: request.source.agent_uuid(), + schedule_uuid: request.source.schedule_uuid(), delivery_mode: request.delivery_mode.as_str().to_string(), message_id: preview_payload.message_id, kind: preview_payload.kind.clone(), @@ -316,6 +356,7 @@ pub async fn enqueue_notification( payload_enc, source_kind: request.source.source_kind().to_string(), source_request_id: request.source.source_request_id(), + background_grant_id: request.source.background_grant_id(), not_before_at: request.not_before_at, expires_at: request.expires_at, } @@ -412,6 +453,68 @@ pub async fn enqueue_agent_message_notification( .await } +pub async fn enqueue_background_agent_message_notification( + state: &Arc, + user: &User, + target: AgentPushTarget, + message_id: Uuid, + message_text: &str, + source: BackgroundAgentPushSource<'_>, +) -> Result, PushError> { + if message_text.trim().is_empty() { + return Ok(None); + } + + let push_settings = state + .db + .get_project_push_settings(user.project_id)? + .unwrap_or_default(); + let delivery_mode = if push_settings.encrypted_preview_enabled { + PushDeliveryMode::EncryptedPreview + } else { + PushDeliveryMode::Generic + }; + + let (deep_link, thread_id) = match &target { + AgentPushTarget::Main => ("opensecret://agent".to_string(), "agent:main".to_string()), + AgentPushTarget::Subagent(agent_uuid) => ( + format!("opensecret://agent/subagent/{}", agent_uuid), + format!("agent:subagent:{}", agent_uuid), + ), + }; + + enqueue_notification( + state, + EnqueueNotificationRequest { + project_id: user.project_id, + user_id: user.uuid, + kind: NOTIFICATION_KIND_AGENT_MESSAGE.to_string(), + delivery_mode, + priority: PushPriority::High, + fallback_title: AGENT_NOTIFICATION_FALLBACK_TITLE.to_string(), + fallback_body: AGENT_NOTIFICATION_FALLBACK_BODY.to_string(), + preview_payload: Some(NotificationPreviewPayloadInput { + message_id, + kind: NOTIFICATION_KIND_AGENT_MESSAGE.to_string(), + title: "Maple".to_string(), + body: normalize_preview_body(message_text), + deep_link, + thread_id, + sent_at: Utc::now().timestamp(), + }), + source: PushNotificationSource::AgentBackground { + grant_id: source.grant.id, + grant_uuid: source.grant.uuid, + agent_uuid: source.agent_uuid, + schedule_uuid: source.schedule_uuid, + }, + not_before_at: None, + expires_at: Some(Utc::now() + Duration::days(7)), + }, + ) + .await +} + #[cfg(test)] mod tests { use super::{normalize_preview_body, PUSH_PREVIEW_BODY_MAX_BYTES}; diff --git a/src/push/worker.rs b/src/push/worker.rs index c18327f5..7368c74d 100644 --- a/src/push/worker.rs +++ b/src/push/worker.rs @@ -1,14 +1,16 @@ +use crate::models::agent_background_grants::AgentBackgroundGrant; use crate::models::notification_deliveries::{ NotificationDelivery, NotificationDeliveryWriteResult, }; use crate::models::notification_events::{ - NotificationEvent, NOTIFICATION_SOURCE_REQUEST_CONTINUATION, + NotificationEvent, NOTIFICATION_SOURCE_AGENT_BACKGROUND, + NOTIFICATION_SOURCE_REQUEST_CONTINUATION, }; use crate::models::push_devices::{PushDevice, PUSH_PLATFORM_ANDROID, PUSH_PLATFORM_IOS}; use crate::push::apns::{send_apns_notification, ApnsSendRequest}; use crate::push::binding::{ - decrypt_notification_preview_payload_v1, decrypt_push_device_capability_v1, - PushDeviceCapabilityPlaintextV1, + decrypt_background_grant_for_push, decrypt_notification_preview_payload_v1, + decrypt_push_device_capability_v1, PushDeviceCapabilityPlaintextV1, }; use crate::push::fcm::send_fcm_notification; use crate::push::{NotificationPreviewPayload, PushError, PushSendOutcome, PushTransport}; @@ -473,23 +475,98 @@ fn decrypt_preview_payload( return Ok(None); }; - if event.source_kind != NOTIFICATION_SOURCE_REQUEST_CONTINUATION { + let background_grant = load_background_grant_for_event(state, event)?; + let background_grant_uuid = background_grant.as_ref().map(|grant| grant.uuid); + let payload = decrypt_notification_preview_payload_v1( + &state.enclave_key, + payload_enc, + event, + background_grant_uuid, + ) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + + if !payload.matches_event(event, background_grant_uuid) { + return Err(PushError::CryptoError( + "notification payload does not match event row".to_string(), + )); + } + + if let Some(grant) = background_grant { + verify_background_push_grant(state, &grant, &payload)?; + } else if event.source_kind != NOTIFICATION_SOURCE_REQUEST_CONTINUATION { return Err(PushError::CryptoError(format!( "unsupported notification source kind: {}", event.source_kind ))); } - let payload = decrypt_notification_preview_payload_v1(&state.enclave_key, payload_enc, event) + Ok(Some(payload)) +} + +fn load_background_grant_for_event( + state: &Arc, + event: &NotificationEvent, +) -> Result, PushError> { + match event.source_kind.as_str() { + NOTIFICATION_SOURCE_REQUEST_CONTINUATION => { + if event.background_grant_id.is_some() { + return Err(PushError::CryptoError( + "request notification unexpectedly references a background grant".to_string(), + )); + } + Ok(None) + } + NOTIFICATION_SOURCE_AGENT_BACKGROUND => { + let grant_id = event.background_grant_id.ok_or_else(|| { + PushError::CryptoError("background notification missing grant id".to_string()) + })?; + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| PushError::ConnectionError)?; + let grant = AgentBackgroundGrant::get_by_id(&mut conn, grant_id) + .map_err(|e| PushError::CryptoError(e.to_string()))? + .ok_or_else(|| PushError::CryptoError("background grant missing".to_string()))?; + if grant.revoked_at.is_some() { + return Err(PushError::CryptoError( + "background grant is revoked".to_string(), + )); + } + if grant.user_id != event.user_id || grant.project_id != event.project_id { + return Err(PushError::CryptoError( + "background grant principal does not match event".to_string(), + )); + } + Ok(Some(grant)) + } + other => Err(PushError::CryptoError(format!( + "unsupported notification source kind: {other}", + ))), + } +} + +fn verify_background_push_grant( + state: &Arc, + grant: &AgentBackgroundGrant, + payload: &NotificationPreviewPayload, +) -> Result<(), PushError> { + let grant_plaintext = decrypt_background_grant_for_push(&state.enclave_key, grant, payload) .map_err(|e| PushError::CryptoError(e.to_string()))?; - if !payload.matches_event(event) { + if grant_plaintext.grant_uuid != grant.uuid + || grant_plaintext.user_uuid != payload.user_uuid + || grant_plaintext.project_id != payload.project_id + || Some(grant_plaintext.agent_uuid) != payload.agent_uuid + || Some(grant_plaintext.schedule_uuid) != payload.schedule_uuid + || !grant_plaintext.may_send_push + { return Err(PushError::CryptoError( - "notification payload does not match event row".to_string(), + "background grant policy does not authorize notification".to_string(), )); } - Ok(Some(payload)) + Ok(()) } fn classify_internal_push_error(error: PushError) -> PushSendOutcome { diff --git a/src/security_invariants.rs b/src/security_invariants.rs index 6cbe3330..4f1641d2 100644 --- a/src/security_invariants.rs +++ b/src/security_invariants.rs @@ -15,6 +15,7 @@ const DESTRUCTIVE_RESET_REQUIRED_TABLES: &[&str] = &[ "user_seed_wrappings", "user_embeddings", "agent_schedule_runs", + "agent_background_grants", "agent_schedules", "agents", "notification_events", diff --git a/src/seed_wrapping.rs b/src/seed_wrapping.rs index d9441140..f6a302be 100644 --- a/src/seed_wrapping.rs +++ b/src/seed_wrapping.rs @@ -10,6 +10,7 @@ use crate::encrypt::{ type HmacSha256 = Hmac; const AUTH_BINDING_MAC_INFO: &[u8] = b"os.auth-binding-mac.v1"; +const BACKGROUND_BINDING_MAC_INFO: &[u8] = b"os.agent-background-binding-mac.v1"; const CREDENTIAL_LOOKUP_MAC_INFO: &[u8] = b"os.credential-lookup-mac.v1"; const PASSWORD_RESET_CODE_MAC_INFO: &[u8] = b"os.password-reset-code-mac.v1"; const SEED_WRAP_ROOT_INFO: &[u8] = b"os.seed-wrap-root.v1"; @@ -20,6 +21,8 @@ const PASSWORD_LOOKUP_DOMAIN: &str = "os.password-lookup.v1"; const PASSWORD_RESET_CODE_DOMAIN: &str = "os.password-reset-code.v1"; const OAUTH_AUTH_BINDING_DOMAIN: &str = "os.oauth-auth-binding.v1"; const OAUTH_LOOKUP_DOMAIN: &str = "os.oauth-lookup.v1"; +const AGENT_BACKGROUND_AUTH_BINDING_DOMAIN: &str = "os.agent-background-auth-binding.v1"; +const AGENT_BACKGROUND_LOOKUP_DOMAIN: &str = "os.agent-background-lookup.v1"; const SEED_WRAP_DOMAIN: &str = "os.seed-wrap.v1"; pub const SEED_WRAP_VERSION_V1: i16 = 1; @@ -28,6 +31,7 @@ pub const SEED_WRAP_VERSION_V1: i16 = 1; pub enum CredentialKind { Password, OAuth, + AgentBackground, } impl CredentialKind { @@ -35,6 +39,7 @@ impl CredentialKind { match self { CredentialKind::Password => "password", CredentialKind::OAuth => "oauth", + CredentialKind::AgentBackground => "agent_background", } } } @@ -164,6 +169,46 @@ pub fn oauth_credential_lookup_hash( )?)) } +pub fn agent_background_credential_lookup_hash( + root_key: &[u8], + grant_uuid: Uuid, +) -> Result { + let mut facts = CanonicalBytes::new(AGENT_BACKGROUND_LOOKUP_DOMAIN); + facts.append_uuid(grant_uuid); + + Ok(CredentialLookupHash(hmac_with_root_domain_key( + root_key, + CREDENTIAL_LOOKUP_MAC_INFO, + &facts.into_bytes(), + )?)) +} + +pub fn compute_agent_background_auth_binding( + root_key: &[u8], + grant_uuid: Uuid, + user_uuid: Uuid, + project_id: i32, + background_secret: &[u8; 32], +) -> Result { + let mut principal_context = CanonicalBytes::new("os.agent-background-principal-context.v1"); + principal_context + .append_uuid(grant_uuid) + .append_uuid(user_uuid) + .append_i32(project_id) + .append_i16(SEED_WRAP_VERSION_V1); + + let mut facts = CanonicalBytes::new(AGENT_BACKGROUND_AUTH_BINDING_DOMAIN); + facts + .append_bytes(&principal_context.into_bytes()) + .append_bytes(background_secret); + + Ok(AuthBinding(hmac_with_root_domain_key( + root_key, + BACKGROUND_BINDING_MAC_INFO, + &facts.into_bytes(), + )?)) +} + pub fn password_reset_code_mac( root_key: &[u8], project_id: i32, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 22c77a7f..255051ef 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -40,10 +40,13 @@ use crate::{ApiError, AppMode, AppState}; mod compaction; mod runtime; +mod schedules; mod signatures; mod tools; mod vision; +pub(crate) use schedules::start_schedule_worker; + #[derive(Debug, Clone, Deserialize)] struct AgentChatRequest { input: MessageContent, @@ -945,11 +948,23 @@ async fn run_agent_chat_task( let runtime = match target.clone() { ChatTarget::Main => { - runtime::AgentRuntime::new_main(state.clone(), user.clone(), user_key).await + runtime::AgentRuntime::new_main( + state.clone(), + user.clone(), + user_key, + Some(auth_context.clone()), + ) + .await } ChatTarget::Subagent(agent_uuid) => { - runtime::AgentRuntime::new_subagent(state.clone(), user.clone(), user_key, agent_uuid) - .await + runtime::AgentRuntime::new_subagent( + state.clone(), + user.clone(), + user_key, + agent_uuid, + Some(auth_context.clone()), + ) + .await } }; diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 7b6828d2..cdbda630 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -8,6 +8,7 @@ use uuid::Uuid; use diesel::prelude::*; use crate::encrypt::{decrypt_string, encrypt_with_key}; +use crate::jwt::AuthContext; use crate::models::agents::{ Agent, NewAgent, AGENT_CREATED_BY_USER, AGENT_KIND_MAIN, AGENT_KIND_SUBAGENT, }; @@ -36,14 +37,17 @@ use crate::web::responses::{MessageContent, MessageContentConverter, MessageCont use crate::{ApiError, AppState}; use super::compaction::CompactionManager; +use super::schedules::{ + refresh_follow_user_schedules_for_user, CancelScheduleTool, ListSchedulesTool, ScheduleTaskTool, +}; use super::signatures::{ build_lm, call_agent_response_with_retry_and_correction, AgentResponseInput, AgentToolCall, AGENT_INSTRUCTION, }; use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, - MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, - WebSearchTool, + MemoryInsertTool, MemoryReplaceTool, SetPreferenceTool, SpawnSubagentTool, ToolRegistry, + ToolResult, WebSearchTool, }; use super::vision; @@ -198,7 +202,7 @@ async fn subagent_metadata_enc( .await } -fn normalize_optional_preference(value: Option<&str>) -> Option { +pub(crate) fn normalize_optional_preference(value: Option<&str>) -> Option { value .map(str::trim) .filter(|value| !value.is_empty()) @@ -294,7 +298,7 @@ fn format_message_timestamp( } } -async fn upsert_user_preference( +pub(crate) async fn upsert_user_preference( conn: &mut diesel::PgConnection, user_key: &SecretKey, user_id: Uuid, @@ -313,6 +317,13 @@ async fn upsert_user_preference( ApiError::InternalServerError })?; + if key == USER_PREFERENCE_TIMEZONE { + refresh_follow_user_schedules_for_user(conn, user_id, &value).map_err(|e| { + error!("Failed to refresh follow-user schedules for timezone update: {e:?}"); + ApiError::InternalServerError + })?; + } + Ok(()) } @@ -563,6 +574,7 @@ impl AgentRuntime { state: Arc, user: crate::models::users::User, user_key: SecretKey, + auth_context: Option, ) -> Result { let user = Arc::new(user); let user_key = Arc::new(user_key); @@ -576,7 +588,7 @@ impl AgentRuntime { let (agent, conversation) = load_main_agent(&mut conn, user.uuid)?.ok_or(ApiError::NotFound)?; - Self::from_loaded(state, user, user_key, agent, conversation).await + Self::from_loaded(state, user, user_key, auth_context, agent, conversation).await } pub async fn new_subagent( @@ -584,6 +596,7 @@ impl AgentRuntime { user: crate::models::users::User, user_key: SecretKey, agent_uuid: Uuid, + auth_context: Option, ) -> Result { let user = Arc::new(user); let user_key = Arc::new(user_key); @@ -613,13 +626,14 @@ impl AgentRuntime { }, )?; - Self::from_loaded(state, user, user_key, agent, conversation).await + Self::from_loaded(state, user, user_key, auth_context, agent, conversation).await } async fn from_loaded( state: Arc, user: Arc, user_key: Arc, + auth_context: Option, agent: Agent, conversation: Conversation, ) -> Result { @@ -664,9 +678,32 @@ impl AgentRuntime { user_key.clone(), conversation.id, ))); + tools.register(Arc::new(SetPreferenceTool::new( + state.clone(), + user.uuid, + user_key.clone(), + ))); if let Some(brave_client) = state.brave_client.clone() { tools.register(Arc::new(WebSearchTool::new(brave_client))); } + tools.register(Arc::new(ScheduleTaskTool::new( + state.clone(), + user.clone(), + user_key.clone(), + auth_context.clone(), + agent.clone(), + ))); + tools.register(Arc::new(ListSchedulesTool::new( + state.clone(), + user.clone(), + user_key.clone(), + agent.clone(), + ))); + tools.register(Arc::new(CancelScheduleTool::new( + state.clone(), + user.clone(), + agent.clone(), + ))); if agent.kind == AGENT_KIND_MAIN { tools.register(Arc::new(SpawnSubagentTool::new( state.clone(), @@ -1494,7 +1531,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok((call, out)) } - async fn maybe_compact(&self) -> Result<(), ApiError> { + pub(crate) async fn maybe_compact(&self) -> Result<(), ApiError> { let mut conn = self .state .db diff --git a/src/web/agent/schedules.rs b/src/web/agent/schedules.rs new file mode 100644 index 00000000..dee6cd79 --- /dev/null +++ b/src/web/agent/schedules.rs @@ -0,0 +1,1999 @@ +use async_trait::async_trait; +use chrono::{DateTime, Datelike, Duration, NaiveDateTime, TimeZone, Utc}; +use chrono_tz::Tz; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::BigInt; +use futures::stream::{self, StreamExt}; +use secp256k1::SecretKey; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::oneshot; +use tokio::time::{sleep, Duration as TokioDuration}; +use tracing::{debug, error, info, trace, warn}; +use uuid::Uuid; + +use crate::agent_background::{ + decrypt_background_grant_v1, encrypt_background_grant_v1, instruction_hash, + AgentBackgroundGrantPlaintextV1, +}; +use crate::encrypt::{decrypt_string, encrypt_with_key}; +use crate::jwt::AuthContext; +use crate::models::agent_background_grants::{ + AgentBackgroundGrant, AgentBackgroundGrantError, NewAgentBackgroundGrant, +}; +use crate::models::agent_schedule_runs::{ + AgentScheduleRun, AgentScheduleRunError, AgentScheduleRunWriteResult, NewAgentScheduleRun, + AGENT_SCHEDULE_RUN_STATUS_COMPLETED, +}; +use crate::models::agent_schedules::{ + parse_local_date, parse_local_time, AgentSchedule, AgentScheduleError, NewAgentSchedule, + ScheduleSpec, ScheduleWeekday, SCHEDULE_KIND_RECURRING, SCHEDULE_STATUS_ACTIVE, + SCHEDULE_STATUS_CANCELLED, SCHEDULE_STATUS_COMPLETED, SCHEDULE_TIMEZONE_MODE_FIXED, + SCHEDULE_TIMEZONE_MODE_FOLLOW_USER, +}; +use crate::models::agents::{Agent, AGENT_KIND_MAIN, AGENT_KIND_SUBAGENT}; +use crate::models::schema::agent_schedules; +use crate::models::user_preferences::{UserPreference, USER_PREFERENCE_TIMEZONE}; +use crate::models::users::User; +use crate::push::{ + enqueue_background_agent_message_notification, AgentPushTarget, BackgroundAgentPushSource, +}; +use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; +use crate::web::openai_auth::AuthMethod; +use crate::AppState; + +use super::runtime; +use super::tools::{Tool, ToolResult}; + +const DEFAULT_STALE_AFTER_MINUTES: i32 = 15; +const SCHEDULE_MATERIALIZER_BATCH_SIZE: i64 = 32; +const SCHEDULE_RUN_BATCH_SIZE: i64 = 16; +const SCHEDULE_RUN_LEASE_TTL_SECONDS: i32 = 180; +const SCHEDULE_RUN_HEARTBEAT_INTERVAL_SECONDS: u64 = 30; +const SCHEDULE_WORKER_POLL_INTERVAL_SECONDS: u64 = 5; +const SCHEDULE_RUN_MAX_CONCURRENCY: usize = 4; +const SCHEDULE_RUN_MAX_ATTEMPTS: i32 = 8; +const SCHEDULE_RUN_MAX_RETRY_BACKOFF_SECONDS: i32 = 15 * 60; + +#[derive(diesel::deserialize::QueryableByName)] +struct DueScheduleIdRow { + #[diesel(sql_type = diesel::sql_types::BigInt)] + id: i64, +} + +#[derive(Debug, Clone)] +struct PersistedScheduledMessage { + message_id: Uuid, + text: String, +} + +#[derive(Debug, Clone)] +struct ScheduledTurnOutcome { + persisted_messages: Vec, + had_error: bool, +} + +pub(crate) fn start_schedule_worker(state: Arc) { + info!( + "starting schedule worker (poll_interval={}s, materializer_batch_size={}, run_batch_size={}, lease_ttl={}s)", + SCHEDULE_WORKER_POLL_INTERVAL_SECONDS, + SCHEDULE_MATERIALIZER_BATCH_SIZE, + SCHEDULE_RUN_BATCH_SIZE, + SCHEDULE_RUN_LEASE_TTL_SECONDS, + ); + + tokio::spawn(async move { + loop { + if let Err(e) = materialize_due_schedules(&state).await { + error!("schedule materialization failed: {e}"); + } + + if let Err(e) = process_schedule_run_batch(&state).await { + error!("schedule run batch failed: {e}"); + } + + sleep(TokioDuration::from_secs( + SCHEDULE_WORKER_POLL_INTERVAL_SECONDS, + )) + .await; + } + }); +} + +pub(crate) fn refresh_follow_user_schedules_for_user( + conn: &mut PgConnection, + user_id: Uuid, + timezone_name: &str, +) -> Result<(), AgentScheduleError> { + let timezone = parse_timezone(timezone_name)?; + let schedules = AgentSchedule::list_active_follow_user_for_user(conn, user_id)?; + + if !schedules.is_empty() { + info!( + "refreshing {} follow-user schedule(s) for user {} to timezone {}", + schedules.len(), + user_id, + timezone.name(), + ); + } + + let now = Utc::now(); + + for schedule in schedules { + let spec = schedule.spec()?; + let next_scheduled_for = + recompute_next_due_for_timezone_change(&schedule, &spec, timezone, now)?; + + diesel::update(agent_schedules::table.filter(agent_schedules::id.eq(schedule.id))) + .set(( + agent_schedules::resolved_timezone.eq(timezone.name()), + agent_schedules::next_scheduled_for.eq(next_scheduled_for), + agent_schedules::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map_err(AgentScheduleError::DatabaseError)?; + + trace!( + "refreshed follow-user schedule {} for user {} (next_scheduled_for={:?})", + schedule.uuid, + user_id, + next_scheduled_for, + ); + } + + Ok(()) +} + +pub struct ScheduleTaskTool { + state: Arc, + user: Arc, + user_key: Arc, + auth_context: Option, + agent: Agent, +} + +impl ScheduleTaskTool { + pub fn new( + state: Arc, + user: Arc, + user_key: Arc, + auth_context: Option, + agent: Agent, + ) -> Self { + Self { + state, + user, + user_key, + auth_context, + agent, + } + } +} + +#[async_trait] +impl Tool for ScheduleTaskTool { + fn name(&self) -> &str { + "schedule_task" + } + + fn description(&self) -> &str { + "Schedule a future wakeup for yourself. Use structured arguments for one-off or recurring schedules. The instruction should be written for your future self, not as final notification copy." + } + + fn args_schema(&self) -> &str { + r#"{ + "schedule_kind": "one_off|recurring", + "instruction": "future-agent instruction in natural language", + "description": "optional short operational summary", + "timezone_mode": "optional: follow_user|fixed (default follow_user)", + "fixed_timezone": "optional IANA timezone when timezone_mode=fixed", + "stale_after_minutes": "optional positive integer (default 15)", + "local_date": "required for one_off: YYYY-MM-DD", + "local_time": "required for one_off/daily/weekly: HH:MM 24-hour", + "recurrence_type": "required for recurring: interval|daily|weekly", + "every_n": "required for interval: positive integer", + "interval_unit": "required for interval: hours", + "weekdays": "required for weekly: comma-separated weekdays like monday,friday or monday,tuesday,wednesday,thursday,friday" +}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(schedule_kind) = args.get("schedule_kind").map(|s| s.trim()) else { + return ToolResult::error("'schedule_kind' argument required"); + }; + let Some(instruction) = args.get("instruction").map(|s| s.trim()) else { + return ToolResult::error("'instruction' argument required"); + }; + + if instruction.is_empty() { + return ToolResult::error("'instruction' must not be empty"); + } + + let spec = match parse_schedule_spec_from_args(args) { + Ok(spec) => spec, + Err(e) => return ToolResult::error(e), + }; + + let timezone_mode = args + .get("timezone_mode") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .unwrap_or(SCHEDULE_TIMEZONE_MODE_FOLLOW_USER); + + let fixed_timezone = args + .get("fixed_timezone") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(str::to_string); + + if timezone_mode != SCHEDULE_TIMEZONE_MODE_FOLLOW_USER + && timezone_mode != SCHEDULE_TIMEZONE_MODE_FIXED + { + return ToolResult::error( + "'timezone_mode' must be 'follow_user' or 'fixed' if provided", + ); + } + + if timezone_mode == SCHEDULE_TIMEZONE_MODE_FIXED && fixed_timezone.is_none() { + return ToolResult::error("'fixed_timezone' is required when timezone_mode='fixed'"); + } + + if schedule_kind == "one_off" && spec.schedule_kind() != "one_off" { + return ToolResult::error("one_off schedules require local_date and local_time"); + } + + if schedule_kind == "recurring" && spec.schedule_kind() != SCHEDULE_KIND_RECURRING { + return ToolResult::error( + "recurring schedules require recurrence_type and recurrence parameters", + ); + } + + let stale_after_minutes = match args + .get("stale_after_minutes") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + { + Some(value) => match value.parse::() { + Ok(value) if value > 0 => value, + _ => { + return ToolResult::error( + "'stale_after_minutes' must be a positive integer when provided", + ) + } + }, + None => DEFAULT_STALE_AFTER_MINUTES, + }; + + let resolved_timezone = match resolve_initial_timezone_name( + &self.state, + self.user.as_ref(), + self.user_key.as_ref(), + timezone_mode, + fixed_timezone.as_deref(), + ) + .await + { + Ok(timezone) => timezone, + Err(e) => { + warn!( + "schedule_task could not resolve timezone for user {} agent {} (timezone_mode={}, fixed_timezone={:?}): {}", + self.user.uuid, + self.agent.uuid, + timezone_mode, + fixed_timezone, + e, + ); + return ToolResult::error(e); + } + }; + + let tz = match parse_timezone(&resolved_timezone) { + Ok(timezone) => timezone, + Err(e) => return ToolResult::error(e.to_string()), + }; + let now = Utc::now(); + let next_scheduled_for = match compute_initial_next_due(&spec, tz, now) { + Ok(next) => next, + Err(e) => return ToolResult::error(e.to_string()), + }; + + if next_scheduled_for <= now { + let requested_local = next_scheduled_for.with_timezone(&tz); + let now_local = now.with_timezone(&tz); + + warn!( + "schedule_task rejected past time for user {} agent {} (timezone={}, requested_local={}, now_local={}, schedule_kind={})", + self.user.uuid, + self.agent.uuid, + resolved_timezone, + requested_local.format("%Y-%m-%d %H:%M:%S"), + now_local.format("%Y-%m-%d %H:%M:%S"), + schedule_kind, + ); + + return ToolResult::error(format!( + "Scheduled time must be in the future in {}. Requested {}, but current time there is {}. Check local_date/local_time and timezone settings.", + resolved_timezone, + requested_local.format("%Y-%m-%d %H:%M:%S"), + now_local.format("%Y-%m-%d %H:%M:%S"), + )); + } + + let Some(auth_context) = self.auth_context.as_ref() else { + return ToolResult::error( + "Schedules can only be created during a live authenticated request.", + ); + }; + + let plaintext_seed = match self + .state + .decrypt_seed_for_auth_context(self.user.as_ref(), auth_context) + { + Ok(seed) => seed, + Err(e) => { + error!("schedule_task could not unwrap seed for background grant: {e:?}"); + return ToolResult::error("Failed to authorize schedule"); + } + }; + + let description = args + .get("description") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| spec.summary()); + + let description_enc = + encrypt_with_key(self.user_key.as_ref(), description.as_bytes()).await; + let instruction_enc = + encrypt_with_key(self.user_key.as_ref(), instruction.as_bytes()).await; + let schedule_spec = match spec.to_value() { + Ok(value) => value, + Err(e) => return ToolResult::error(e.to_string()), + }; + let schedule_uuid = Uuid::new_v4(); + let grant_uuid = Uuid::new_v4(); + let instruction_hash = instruction_hash(instruction); + let grant_plaintext = AgentBackgroundGrantPlaintextV1::new_scheduled( + grant_uuid, + self.user.uuid, + self.user.project_id, + self.agent.uuid, + schedule_uuid, + instruction_hash, + auth_context.method.as_str(), + ); + let grant_enc = match encrypt_background_grant_v1(&self.state.enclave_key, &grant_plaintext) + { + Ok(grant_enc) => grant_enc, + Err(e) => { + error!("schedule_task could not seal background grant: {e:?}"); + return ToolResult::error("Failed to authorize schedule"); + } + }; + let seed_wrapping = match self.state.new_agent_background_seed_wrapping_for_user( + self.user.as_ref(), + grant_uuid, + &grant_plaintext.background_secret, + &plaintext_seed, + ) { + Ok(seed_wrapping) => seed_wrapping, + Err(e) => { + error!("schedule_task could not create background seed wrap: {e:?}"); + return ToolResult::error("Failed to authorize schedule"); + } + }; + let seed_wrap_lookup_hash = seed_wrapping.credential_lookup_hash.clone(); + + let new_schedule = NewAgentSchedule { + uuid: schedule_uuid, + user_id: self.user.uuid, + agent_id: self.agent.id, + description_enc, + instruction_enc, + schedule_kind: spec.schedule_kind().to_string(), + recurrence_type: spec.recurrence_type().map(str::to_string), + schedule_spec, + timezone_mode: timezone_mode.to_string(), + resolved_timezone: resolved_timezone.clone(), + fixed_timezone, + stale_after_minutes, + status: SCHEDULE_STATUS_ACTIVE.to_string(), + next_scheduled_for: Some(next_scheduled_for), + last_scheduled_for: None, + last_run_at: None, + run_count: 0, + cancelled_at: None, + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(conn) => conn, + Err(_) => return ToolResult::error("Database connection error"), + }; + + let insert_result = conn.transaction::(|conn| { + let schedule = new_schedule.insert(conn).map_err(|e| match e { + AgentScheduleError::DatabaseError(error) => error, + AgentScheduleError::InvalidSpec(_) => diesel::result::Error::RollbackTransaction, + })?; + NewAgentBackgroundGrant { + uuid: grant_uuid, + user_id: self.user.uuid, + project_id: self.user.project_id, + agent_id: self.agent.id, + schedule_id: schedule.id, + grant_enc: grant_enc.clone(), + seed_wrap_lookup_hash: seed_wrap_lookup_hash.clone(), + } + .insert(conn) + .map_err(|e| match e { + AgentBackgroundGrantError::DatabaseError(error) => error, + })?; + seed_wrapping + .upsert_by_credential(conn) + .map_err(|e| match e { + crate::models::user_seed_wrappings::UserSeedWrappingError::DatabaseError( + error, + ) => error, + })?; + Ok(schedule) + }); + + match insert_result { + Ok(schedule) => { + let spec_summary = spec.summary(); + info!( + "created schedule {} for user {} agent {} (kind={}, recurrence={:?}, timezone_mode={}, resolved_timezone={}, next_scheduled_for={}, stale_after_minutes={})", + schedule.uuid, + self.user.uuid, + self.agent.uuid, + schedule.schedule_kind, + schedule.recurrence_type, + schedule.timezone_mode, + schedule.resolved_timezone, + next_scheduled_for.format("%Y-%m-%d %H:%M:%S UTC"), + stale_after_minutes, + ); + + ToolResult::success(format!( + "Created schedule {} ({}) in timezone {}. Next run: {} UTC.", + schedule.uuid, + spec_summary, + resolved_timezone, + next_scheduled_for.format("%Y-%m-%d %H:%M:%S") + )) + } + Err(e) => { + error!("schedule_task insert failed: {e:?}"); + ToolResult::error("Failed to create schedule") + } + } + } +} + +pub struct ListSchedulesTool { + state: Arc, + user: Arc, + user_key: Arc, + agent: Agent, +} + +impl ListSchedulesTool { + pub fn new( + state: Arc, + user: Arc, + user_key: Arc, + agent: Agent, + ) -> Self { + Self { + state, + user, + user_key, + agent, + } + } +} + +#[async_trait] +impl Tool for ListSchedulesTool { + fn name(&self) -> &str { + "list_schedules" + } + + fn description(&self) -> &str { + "List schedules for this agent. Defaults to active schedules only." + } + + fn args_schema(&self) -> &str { + r#"{"status": "optional: active|completed|cancelled|all (default active)"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let status_filter = match args.get("status").map(|s| s.trim()) { + Some("all") => None, + None | Some("") => Some(SCHEDULE_STATUS_ACTIVE), + Some(SCHEDULE_STATUS_ACTIVE) => Some(SCHEDULE_STATUS_ACTIVE), + Some(SCHEDULE_STATUS_COMPLETED) => Some(SCHEDULE_STATUS_COMPLETED), + Some(SCHEDULE_STATUS_CANCELLED) => Some(SCHEDULE_STATUS_CANCELLED), + Some(_) => { + return ToolResult::error( + "'status' must be one of active, completed, cancelled, or all", + ) + } + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(conn) => conn, + Err(_) => return ToolResult::error("Database connection error"), + }; + + match AgentSchedule::list_by_agent(&mut conn, self.user.uuid, self.agent.id, status_filter) + { + Ok(schedules) => { + debug!( + "listed {} schedule(s) for user {} agent {} (status_filter={:?})", + schedules.len(), + self.user.uuid, + self.agent.uuid, + status_filter, + ); + + if schedules.is_empty() { + return ToolResult::success("No schedules found.".to_string()); + } + + let mut output = format!("Found {} schedule(s):\n\n", schedules.len()); + for schedule in schedules { + let spec_summary = schedule + .spec() + .map(|spec| spec.summary()) + .unwrap_or_else(|_| "[invalid spec]".to_string()); + let next_text = schedule + .next_scheduled_for + .map(|dt| { + let tz = parse_timezone_or_utc(&schedule.resolved_timezone); + format_local_time(dt, tz) + }) + .unwrap_or_else(|| "none".to_string()); + let description = + decrypt_string(self.user_key.as_ref(), Some(&schedule.description_enc)) + .ok() + .flatten() + .unwrap_or_else(|| spec_summary.clone()); + output.push_str(&format!( + "- id: {}\n status: {}\n description: {}\n rule: {}\n timezone: {} ({})\n next: {}\n\n", + schedule.uuid, + schedule.status, + description.trim(), + spec_summary, + schedule.timezone_mode, + schedule.resolved_timezone, + next_text, + )); + } + + ToolResult::success(output) + } + Err(e) => { + error!("list_schedules failed: {e:?}"); + ToolResult::error("Failed to list schedules") + } + } + } +} + +pub struct CancelScheduleTool { + state: Arc, + user: Arc, + agent: Agent, +} + +impl CancelScheduleTool { + pub fn new(state: Arc, user: Arc, agent: Agent) -> Self { + Self { state, user, agent } + } +} + +#[async_trait] +impl Tool for CancelScheduleTool { + fn name(&self) -> &str { + "cancel_schedule" + } + + fn description(&self) -> &str { + "Cancel an existing schedule for this agent by schedule id (UUID). Prevents future runs and cancels unstarted pending runs." + } + + fn args_schema(&self) -> &str { + r#"{"schedule_id": "schedule UUID from list_schedules"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(schedule_id) = args.get("schedule_id").map(|s| s.trim()) else { + return ToolResult::error("'schedule_id' argument required"); + }; + + let schedule_uuid = match Uuid::parse_str(schedule_id) { + Ok(uuid) => uuid, + Err(_) => return ToolResult::error("'schedule_id' must be a valid UUID"), + }; + + let mut conn = match self.state.db.get_pool().get() { + Ok(conn) => conn, + Err(_) => return ToolResult::error("Database connection error"), + }; + + let result = conn.transaction::(|conn| { + let Some(schedule) = AgentSchedule::get_by_uuid_and_agent( + conn, + schedule_uuid, + self.user.uuid, + self.agent.id, + )? + else { + debug!( + "cancel_schedule did not find schedule {} for user {} agent {}", + schedule_uuid, self.user.uuid, self.agent.uuid, + ); + return Ok("Schedule not found.".to_string()); + }; + + if schedule.status == SCHEDULE_STATUS_CANCELLED { + debug!( + "schedule {} already cancelled for user {} agent {}", + schedule.uuid, self.user.uuid, self.agent.uuid, + ); + return Ok(format!("Schedule {} is already cancelled.", schedule.uuid)); + } + + diesel::update(agent_schedules::table.filter(agent_schedules::id.eq(schedule.id))) + .set(( + agent_schedules::status.eq(SCHEDULE_STATUS_CANCELLED), + agent_schedules::next_scheduled_for.eq::>>(None), + agent_schedules::cancelled_at.eq(diesel::dsl::now), + agent_schedules::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map_err(AgentScheduleError::DatabaseError)?; + + let cancelled_runs = AgentScheduleRun::cancel_unstarted_for_schedule(conn, schedule.id) + .map_err(|e| match e { + AgentScheduleRunError::DatabaseError(err) => { + AgentScheduleError::DatabaseError(err) + } + })?; + + info!( + "cancelled schedule {} for user {} agent {} (cancelled_runs={})", + schedule.uuid, self.user.uuid, self.agent.uuid, cancelled_runs, + ); + + Ok(format!( + "Cancelled schedule {} and {} unstarted run(s).", + schedule.uuid, cancelled_runs + )) + }); + + match result { + Ok(msg) => ToolResult::success(msg), + Err(e) => { + error!("cancel_schedule failed: {e:?}"); + ToolResult::error("Failed to cancel schedule") + } + } + } +} + +async fn materialize_due_schedules(state: &Arc) -> Result<(), String> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + + conn.transaction::<(), AgentScheduleError, _>(|conn| { + let rows = sql_query( + r#" + SELECT s.id + FROM agent_schedules s + WHERE s.status = 'active' + AND s.next_scheduled_for IS NOT NULL + AND s.next_scheduled_for <= NOW() + ORDER BY s.next_scheduled_for ASC, s.id ASC + FOR UPDATE OF s SKIP LOCKED + LIMIT $1 + "#, + ) + .bind::(SCHEDULE_MATERIALIZER_BATCH_SIZE) + .load::(conn) + .map_err(AgentScheduleError::DatabaseError)?; + + if !rows.is_empty() { + debug!("materializing {} due schedule(s)", rows.len()); + } + + for row in rows { + let Some(schedule) = AgentSchedule::get_by_id(conn, row.id)? else { + continue; + }; + + if schedule.status != SCHEDULE_STATUS_ACTIVE { + continue; + } + + let Some(current_due) = schedule.next_scheduled_for else { + continue; + }; + + let spec = schedule.spec()?; + let tz = parse_timezone_or_utc(&schedule.resolved_timezone); + let next_scheduled_for = compute_following_next_due(&spec, tz, current_due)?; + + let stale_after_at = + current_due + Duration::minutes(schedule.stale_after_minutes as i64); + + let new_run = NewAgentScheduleRun { + uuid: Uuid::new_v4(), + schedule_id: schedule.id, + user_id: schedule.user_id, + agent_id: schedule.agent_id, + scheduled_for: current_due, + stale_after_at, + status: crate::models::agent_schedule_runs::AGENT_SCHEDULE_RUN_STATUS_PENDING + .to_string(), + next_attempt_at: current_due, + }; + + let inserted_run = new_run.insert(conn).map_err(|e| match e { + AgentScheduleRunError::DatabaseError(err) => AgentScheduleError::DatabaseError(err), + })?; + + let new_status = if spec.schedule_kind() == "one_off" { + SCHEDULE_STATUS_COMPLETED + } else { + SCHEDULE_STATUS_ACTIVE + }; + + diesel::update(agent_schedules::table.filter(agent_schedules::id.eq(schedule.id))) + .set(( + agent_schedules::status.eq(new_status), + agent_schedules::next_scheduled_for.eq(next_scheduled_for), + agent_schedules::last_scheduled_for.eq(Some(current_due)), + agent_schedules::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map_err(AgentScheduleError::DatabaseError)?; + + debug!( + "materialized schedule {} for user {} into run {} (scheduled_for={}, next_scheduled_for={:?})", + schedule.uuid, + schedule.user_id, + inserted_run.uuid, + current_due.format("%Y-%m-%d %H:%M:%S UTC"), + next_scheduled_for, + ); + } + + Ok(()) + }) + .map_err(|e| e.to_string())?; + + Ok(()) +} + +async fn process_schedule_run_batch(state: &Arc) -> Result<(), String> { + let lease_owner = format!("schedule-worker:{}:{}", std::process::id(), Uuid::new_v4()); + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + + let runs = AgentScheduleRun::lease_pending( + &mut conn, + SCHEDULE_RUN_BATCH_SIZE, + &lease_owner, + SCHEDULE_RUN_LEASE_TTL_SECONDS, + ) + .map_err(|e| e.to_string())?; + drop(conn); + + if runs.is_empty() { + return Ok(()); + } + + debug!( + "leased {} scheduled run(s) (lease_owner={})", + runs.len(), + lease_owner, + ); + + stream::iter(runs) + .for_each_concurrent(SCHEDULE_RUN_MAX_CONCURRENCY, |run| { + let state = state.clone(); + let lease_owner = lease_owner.clone(); + async move { + if let Err(e) = process_leased_schedule_run(&state, run, &lease_owner).await { + error!("scheduled run processing failed: {}", e); + } + } + }) + .await; + + Ok(()) +} + +async fn process_leased_schedule_run( + state: &Arc, + run: AgentScheduleRun, + lease_owner: &str, +) -> Result<(), String> { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + + if run.stale_after_at <= Utc::now() { + info!( + "expiring stale scheduled run {} for user {} (schedule_id={}, scheduled_for={}, stale_after_at={})", + run.uuid, + run.user_id, + run.schedule_id, + run.scheduled_for.format("%Y-%m-%d %H:%M:%S UTC"), + run.stale_after_at.format("%Y-%m-%d %H:%M:%S UTC"), + ); + + record_run_transition( + AgentScheduleRun::mark_expired( + &mut conn, + run.id, + lease_owner, + Some("schedule occurrence expired before execution"), + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "expired", + ); + update_schedule_after_terminal_run(&mut conn, run.schedule_id) + .map_err(|e| e.to_string())?; + return Ok(()); + } + + if run.first_output_at.is_some() || run.output_count > 0 { + info!( + "marking scheduled run {} complete after partial-output recovery (user={}, schedule_id={}, output_count={})", + run.uuid, + run.user_id, + run.schedule_id, + run.output_count, + ); + + record_run_transition( + AgentScheduleRun::mark_completed( + &mut conn, + run.id, + lease_owner, + false, + Some("marking completed after prior partial output recovery"), + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + AGENT_SCHEDULE_RUN_STATUS_COMPLETED, + ); + update_schedule_after_terminal_run(&mut conn, run.schedule_id) + .map_err(|e| e.to_string())?; + return Ok(()); + } + + let Some(schedule) = + AgentSchedule::get_by_id(&mut conn, run.schedule_id).map_err(|e| e.to_string())? + else { + warn!( + "scheduled run {} could not find schedule definition {}", + run.uuid, run.schedule_id, + ); + + record_run_transition( + AgentScheduleRun::mark_failed( + &mut conn, + run.id, + lease_owner, + Some("schedule definition missing"), + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "failed", + ); + return Ok(()); + }; + + if schedule.status == SCHEDULE_STATUS_CANCELLED { + info!( + "scheduled run {} skipped because schedule {} is cancelled", + run.uuid, schedule.uuid, + ); + + record_run_transition( + AgentScheduleRun::mark_cancelled( + &mut conn, + run.id, + lease_owner, + Some("schedule cancelled"), + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "cancelled", + ); + return Ok(()); + } + + let Some(agent) = Agent::get_by_id(&mut conn, run.agent_id).map_err(|e| e.to_string())? else { + warn!( + "scheduled run {} could not find agent {} for schedule {}", + run.uuid, run.agent_id, schedule.uuid, + ); + + record_run_transition( + AgentScheduleRun::mark_failed(&mut conn, run.id, lease_owner, Some("agent missing")) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "failed", + ); + return Ok(()); + }; + + let Some(user) = User::get_by_uuid(&mut conn, run.user_id).map_err(|e| e.to_string())? else { + warn!( + "scheduled run {} could not find user {} for schedule {}", + run.uuid, run.user_id, schedule.uuid, + ); + + record_run_transition( + AgentScheduleRun::mark_failed(&mut conn, run.id, lease_owner, Some("user missing")) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "failed", + ); + return Ok(()); + }; + + if run.user_id != schedule.user_id + || run.agent_id != schedule.agent_id + || agent.user_id != run.user_id + || agent.id != run.agent_id + || user.uuid != run.user_id + { + warn!( + "scheduled run {} failed sealed-principal row consistency check (run_user={}, schedule_user={}, run_agent={}, schedule_agent={}, agent_user={})", + run.uuid, + run.user_id, + schedule.user_id, + run.agent_id, + schedule.agent_id, + agent.user_id, + ); + fail_schedule_run(&mut conn, &run, lease_owner, "schedule principal mismatch")?; + return Ok(()); + } + + let Some(grant) = AgentBackgroundGrant::get_active_by_schedule(&mut conn, schedule.id) + .map_err(|e| match e { + AgentBackgroundGrantError::DatabaseError(error) => error.to_string(), + })? + else { + warn!( + "scheduled run {} could not find active background grant for schedule {}", + run.uuid, schedule.uuid, + ); + fail_schedule_run(&mut conn, &run, lease_owner, "background grant missing")?; + return Ok(()); + }; + drop(conn); + + info!( + "starting scheduled run {} for schedule {} user {} agent {} (attempt={}, scheduled_for={})", + run.uuid, + schedule.uuid, + user.uuid, + agent.uuid, + run.attempt_count + 1, + run.scheduled_for.format("%Y-%m-%d %H:%M:%S UTC"), + ); + + let grant_plaintext = decrypt_background_grant_v1( + &state.enclave_key, + &grant.grant_enc, + grant.uuid, + user.uuid, + user.project_id, + agent.uuid, + schedule.uuid, + ) + .map_err(|e| format!("failed to decrypt background grant: {e:?}"))?; + + let user_key = state + .get_user_key_for_agent_background_grant( + &user, + &grant_plaintext, + &grant.seed_wrap_lookup_hash, + ) + .await + .map_err(|e| format!("failed to derive background user key: {e:?}"))?; + let instruction = decrypt_string(&user_key, Some(&schedule.instruction_enc)) + .map_err(|e| format!("failed to decrypt schedule instruction: {e:?}"))? + .unwrap_or_default(); + let actual_instruction_hash = instruction_hash(&instruction); + if !grant_plaintext.verify_scheduled_policy( + user.uuid, + user.project_id, + agent.uuid, + schedule.uuid, + &actual_instruction_hash, + ) { + return Err("background grant policy did not match scheduled run".to_string()); + } + + let input = build_scheduled_turn_input(&schedule, &run, &instruction); + + let (stop_tx, stop_rx) = oneshot::channel(); + let heartbeat_state = state.clone(); + let heartbeat_run_id = run.id; + let heartbeat_owner = lease_owner.to_string(); + let heartbeat_handle = tokio::spawn(async move { + heartbeat_schedule_run_lease(heartbeat_state, heartbeat_run_id, heartbeat_owner, stop_rx) + .await; + }); + + let turn_result = run_scheduled_agent_turn( + state, + user.clone(), + user_key, + &agent, + &run, + lease_owner, + &input, + ) + .await; + + let _ = stop_tx.send(()); + let _ = heartbeat_handle.await; + + match turn_result { + Ok(outcome) => { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + + let push_enqueued = + if grant_plaintext.may_send_push && !outcome.persisted_messages.is_empty() { + let preview = + compose_scheduled_preview(state, &user, &outcome.persisted_messages).await; + let target = agent_push_target(&agent); + let first_message = &outcome.persisted_messages[0]; + match enqueue_background_agent_message_notification( + state, + &user, + target, + first_message.message_id, + &preview, + BackgroundAgentPushSource { + grant: &grant, + agent_uuid: agent.uuid, + schedule_uuid: schedule.uuid, + }, + ) + .await + { + Ok(Some(_)) => true, + Ok(None) => false, + Err(e) => { + error!("failed to enqueue scheduled agent push: {e:?}"); + false + } + } + } else { + false + }; + + let terminal_error = if outcome.had_error && !outcome.persisted_messages.is_empty() { + Some("scheduled turn ended after partial output; preserving existing output") + } else { + None + }; + + let message_count = outcome.persisted_messages.len(); + + if outcome.had_error && outcome.persisted_messages.is_empty() { + if run.attempt_count + 1 >= SCHEDULE_RUN_MAX_ATTEMPTS { + error!( + "scheduled run {} for schedule {} user {} failed permanently before producing output (attempt={})", + run.uuid, + schedule.uuid, + user.uuid, + run.attempt_count + 1, + ); + + record_run_transition( + AgentScheduleRun::mark_failed( + &mut conn, + run.id, + lease_owner, + Some("scheduled turn failed before producing output"), + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "failed", + ); + } else { + let retry_after_seconds = retry_backoff_seconds(run.attempt_count + 1); + warn!( + "retrying scheduled run {} for schedule {} user {} after {}s (attempt={})", + run.uuid, + schedule.uuid, + user.uuid, + retry_after_seconds, + run.attempt_count + 1, + ); + + record_run_transition( + AgentScheduleRun::mark_retry( + &mut conn, + run.id, + lease_owner, + Some("scheduled turn failed before producing output"), + retry_after_seconds, + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "retry", + ); + } + } else { + record_run_transition( + AgentScheduleRun::mark_completed( + &mut conn, + run.id, + lease_owner, + push_enqueued, + terminal_error, + ) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + AGENT_SCHEDULE_RUN_STATUS_COMPLETED, + ); + update_schedule_after_terminal_run(&mut conn, run.schedule_id) + .map_err(|e| e.to_string())?; + + info!( + "completed scheduled run {} for schedule {} user {} agent {} (messages={}, push_enqueued={}, had_error={})", + run.uuid, + schedule.uuid, + user.uuid, + agent.uuid, + message_count, + push_enqueued, + outcome.had_error, + ); + } + } + Err(e) => { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + if run.attempt_count + 1 >= SCHEDULE_RUN_MAX_ATTEMPTS { + error!( + "scheduled run {} for user {} failed permanently after execution error (schedule_id={}, attempt={})", + run.uuid, + run.user_id, + run.schedule_id, + run.attempt_count + 1, + ); + + record_run_transition( + AgentScheduleRun::mark_failed(&mut conn, run.id, lease_owner, Some(&e)) + .map_err(|err| err.to_string())?, + run.id, + lease_owner, + "failed", + ); + } else { + let retry_after_seconds = retry_backoff_seconds(run.attempt_count + 1); + warn!( + "retrying scheduled run {} for user {} after execution error in {}s (schedule_id={}, attempt={})", + run.uuid, + run.user_id, + retry_after_seconds, + run.schedule_id, + run.attempt_count + 1, + ); + + record_run_transition( + AgentScheduleRun::mark_retry( + &mut conn, + run.id, + lease_owner, + Some(&e), + retry_after_seconds, + ) + .map_err(|err| err.to_string())?, + run.id, + lease_owner, + "retry", + ); + } + } + } + + Ok(()) +} + +async fn run_scheduled_agent_turn( + state: &Arc, + user: User, + user_key: SecretKey, + agent: &Agent, + run: &AgentScheduleRun, + lease_owner: &str, + input: &str, +) -> Result { + let runtime_result = if agent.kind == AGENT_KIND_MAIN { + runtime::AgentRuntime::new_main(state.clone(), user.clone(), user_key, None).await + } else if agent.kind == AGENT_KIND_SUBAGENT { + runtime::AgentRuntime::new_subagent(state.clone(), user.clone(), user_key, agent.uuid, None) + .await + } else { + return Err(format!( + "unsupported agent kind for schedule: {}", + agent.kind + )); + }; + + let mut runtime = + runtime_result.map_err(|e| format!("failed to initialize agent runtime: {e:?}"))?; + runtime.clear_tool_results(); + runtime + .maybe_compact() + .await + .map_err(|e| format!("scheduled compaction failed: {e:?}"))?; + + let max_steps = runtime.max_steps(); + let mut persisted_messages = Vec::new(); + let mut had_error = false; + + debug!( + "running scheduled agent turn for run {} user {} agent {} (kind={}, max_steps={})", + run.uuid, user.uuid, agent.uuid, agent.kind, max_steps, + ); + + 'steps: for step_num in 0..max_steps { + match runtime.step(input, step_num == 0).await { + Ok(result) => { + trace!( + "scheduled run {} step {} produced {} tool call(s), {} message(s), done={}", + run.uuid, + step_num, + result.executed_tools.len(), + result.messages.len(), + result.done, + ); + + for executed in &result.executed_tools { + if let Err(e) = runtime + .insert_tool_call_and_output(&executed.tool_call, &executed.result) + .await + { + error!("failed to persist scheduled tool call: {e:?}"); + } + } + + for msg in result.messages { + match runtime.insert_assistant_message(&msg).await { + Ok(inserted) => { + persisted_messages.push(PersistedScheduledMessage { + message_id: inserted.uuid, + text: msg.clone(), + }); + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "database connection error".to_string())?; + let write_result = AgentScheduleRun::record_output( + &mut conn, + run.id, + lease_owner, + inserted.uuid, + ) + .map_err(|e| e.to_string())?; + + if !write_result.was_applied() { + return Err( + "lost lease while recording scheduled output".to_string() + ); + } + } + Err(e) => { + error!("failed to persist scheduled assistant message: {e:?}"); + had_error = true; + break 'steps; + } + } + } + + if result.done { + break 'steps; + } + } + Err(e) => { + error!("scheduled agent step failed: {e:?}"); + had_error = true; + break 'steps; + } + } + } + + Ok(ScheduledTurnOutcome { + persisted_messages, + had_error, + }) +} + +async fn heartbeat_schedule_run_lease( + state: Arc, + run_id: i64, + lease_owner: String, + mut stop_rx: oneshot::Receiver<()>, +) { + let mut interval = tokio::time::interval(TokioDuration::from_secs( + SCHEDULE_RUN_HEARTBEAT_INTERVAL_SECONDS, + )); + + loop { + tokio::select! { + _ = interval.tick() => { + let mut conn = match state.db.get_pool().get() { + Ok(conn) => conn, + Err(_) => { + warn!("stopping schedule heartbeat for run {} due to database connection error", run_id); + break; + } + }; + + match AgentScheduleRun::renew_lease( + &mut conn, + run_id, + &lease_owner, + SCHEDULE_RUN_LEASE_TTL_SECONDS, + ) { + Ok(result) if result.was_applied() => {} + Ok(_) => { + warn!("stopping schedule heartbeat for run {} because lease renewal was not applied", run_id); + break; + } + Err(e) => { + warn!("stopping schedule heartbeat for run {} because lease renewal failed: {e:?}", run_id); + break; + } + } + } + _ = &mut stop_rx => break, + } + } +} + +async fn compose_scheduled_preview( + state: &Arc, + user: &User, + messages: &[PersistedScheduledMessage], +) -> String { + let Some(first_message) = messages.first() else { + return String::new(); + }; + + trace!( + "composing scheduled preview for user {} from {} message(s)", + user.uuid, + messages.len(), + ); + + let joined_messages = messages + .iter() + .enumerate() + .map(|(index, message)| format!("{}. {}", index + 1, message.text.trim())) + .collect::>() + .join("\n\n"); + + let request = serde_json::json!({ + "model": "llama-3.3-70b", + "messages": [ + { + "role": "system", + "content": "You compress assistant messages into a short encrypted notification preview. Preserve tone and intent. Do not add facts. Return only 1-2 short sentences, ideally under 160 characters." + }, + { + "role": "user", + "content": format!("Create a notification preview from these assistant messages:\n\n{}", joined_messages) + } + ], + "temperature": 0.2, + "max_tokens": 80, + "stream": false + }); + + let headers = axum::http::HeaderMap::new(); + let billing_context = BillingContext::new(AuthMethod::Jwt, "llama-3.3-70b".to_string()); + + match get_chat_completion_response(state, user, request, &headers, billing_context).await { + Ok(mut completion) => match completion.stream.recv().await { + Some(CompletionChunk::FullResponse(response_json)) => response_json + .get("choices") + .and_then(|choices| choices.get(0)) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .and_then(|content| content.as_str()) + .map(|content| content.trim().to_string()) + .filter(|content| !content.is_empty()) + .unwrap_or_else(|| { + trace!( + "scheduled preview composer returned empty content for user {}, using first assistant message", + user.uuid, + ); + first_message.text.clone() + }), + _ => { + trace!( + "scheduled preview composer returned no full response for user {}, using first assistant message", + user.uuid, + ); + first_message.text.clone() + } + }, + Err(e) => { + warn!("scheduled preview composition failed, falling back to first message: {e:?}"); + first_message.text.clone() + } + } +} + +fn build_scheduled_turn_input( + schedule: &AgentSchedule, + run: &AgentScheduleRun, + instruction: &str, +) -> String { + format!( + "[SCHEDULED EVENT]\nThis is an internal scheduled wakeup, not a new live user message.\n\nSchedule ID: {}\nRun ID: {}\nScheduled for (UTC): {}\nCurrent resolved timezone: {}\nInstruction for your future self:\n{}\n\nDecide what, if anything, the user should see right now. You may send messages, call tools, or do nothing if the event is no longer relevant.", + schedule.uuid, + run.uuid, + run.scheduled_for.format("%Y-%m-%d %H:%M:%S UTC"), + schedule.resolved_timezone, + instruction.trim() + ) +} + +fn parse_schedule_spec_from_args(args: &HashMap) -> Result { + let schedule_kind = args + .get("schedule_kind") + .map(|s| s.trim()) + .ok_or_else(|| "'schedule_kind' argument required".to_string())?; + + let spec = match schedule_kind { + "one_off" => { + let local_date = args + .get("local_date") + .map(|s| s.trim()) + .ok_or_else(|| "'local_date' is required for one_off schedules".to_string())?; + let local_time = args + .get("local_time") + .map(|s| s.trim()) + .ok_or_else(|| "'local_time' is required for one_off schedules".to_string())?; + + ScheduleSpec::OneOff { + local_date: local_date.to_string(), + local_time: local_time.to_string(), + } + } + "recurring" => { + let recurrence_type = + args.get("recurrence_type") + .map(|s| s.trim()) + .ok_or_else(|| { + "'recurrence_type' is required for recurring schedules".to_string() + })?; + + match recurrence_type { + "interval" => { + let every_n = args + .get("every_n") + .map(|s| s.trim()) + .ok_or_else(|| "'every_n' is required for interval schedules".to_string())? + .parse::() + .map_err(|_| "'every_n' must be a positive integer".to_string())?; + let interval_unit = args + .get("interval_unit") + .map(|s| s.trim()) + .unwrap_or("hours"); + + if interval_unit != "hours" { + return Err( + "v1 interval schedules currently support interval_unit='hours' only" + .to_string(), + ); + } + + ScheduleSpec::Interval { + every_n, + interval_unit: crate::models::agent_schedules::ScheduleIntervalUnit::Hours, + } + } + "daily" => { + let local_time = args.get("local_time").map(|s| s.trim()).ok_or_else(|| { + "'local_time' is required for daily schedules".to_string() + })?; + ScheduleSpec::Daily { + local_time: local_time.to_string(), + } + } + "weekly" => { + let local_time = args.get("local_time").map(|s| s.trim()).ok_or_else(|| { + "'local_time' is required for weekly schedules".to_string() + })?; + let weekdays_value = args + .get("weekdays") + .map(|s| s.trim()) + .ok_or_else(|| "'weekdays' is required for weekly schedules".to_string())?; + let weekdays = parse_weekdays_csv(weekdays_value)?; + ScheduleSpec::Weekly { + local_time: local_time.to_string(), + weekdays, + } + } + _ => return Err( + "'recurrence_type' must be interval, daily, or weekly for recurring schedules" + .to_string(), + ), + } + } + _ => return Err("'schedule_kind' must be 'one_off' or 'recurring'".to_string()), + }; + + spec.validate().map_err(|e| e.to_string())?; + Ok(spec) +} + +fn parse_weekdays_csv(value: &str) -> Result, String> { + let mut weekdays = Vec::new(); + + for raw in value.split(',') { + let normalized = raw.trim().to_lowercase(); + if normalized.is_empty() { + continue; + } + + let weekday = match normalized.as_str() { + "weekdays" => { + return Ok(vec![ + ScheduleWeekday::Monday, + ScheduleWeekday::Tuesday, + ScheduleWeekday::Wednesday, + ScheduleWeekday::Thursday, + ScheduleWeekday::Friday, + ]) + } + "monday" | "mon" => ScheduleWeekday::Monday, + "tuesday" | "tue" | "tues" => ScheduleWeekday::Tuesday, + "wednesday" | "wed" => ScheduleWeekday::Wednesday, + "thursday" | "thu" | "thurs" => ScheduleWeekday::Thursday, + "friday" | "fri" => ScheduleWeekday::Friday, + "saturday" | "sat" => ScheduleWeekday::Saturday, + "sunday" | "sun" => ScheduleWeekday::Sunday, + _ => return Err(format!("invalid weekday '{raw}'")), + }; + + if !weekdays.contains(&weekday) { + weekdays.push(weekday); + } + } + + if weekdays.is_empty() { + return Err("at least one valid weekday is required".to_string()); + } + + Ok(weekdays) +} + +async fn resolve_initial_timezone_name( + state: &Arc, + user: &User, + user_key: &SecretKey, + timezone_mode: &str, + fixed_timezone: Option<&str>, +) -> Result { + if timezone_mode == SCHEDULE_TIMEZONE_MODE_FIXED { + let fixed_timezone = fixed_timezone + .ok_or_else(|| "'fixed_timezone' is required when timezone_mode='fixed'".to_string())?; + + return parse_timezone(fixed_timezone) + .map(|timezone| timezone.name().to_string()) + .map_err(|e| e.to_string()); + } + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| "Database connection error".to_string())?; + + load_user_timezone_name(&mut conn, user_key, user.uuid) + .map_err(|e| e.to_string())? + .ok_or_else(|| { + "follow_user schedules require the user to have a saved timezone preference; set the timezone first or use timezone_mode='fixed' with fixed_timezone" + .to_string() + }) + .and_then(|timezone_name| { + parse_timezone(&timezone_name) + .map(|timezone| timezone.name().to_string()) + .map_err(|e| e.to_string()) + }) +} + +fn load_user_timezone_name( + conn: &mut PgConnection, + user_key: &SecretKey, + user_id: Uuid, +) -> Result, AgentScheduleError> { + let pref = UserPreference::get_by_user_and_key(conn, user_id, USER_PREFERENCE_TIMEZONE) + .map_err(|e| match e { + crate::models::user_preferences::UserPreferenceError::DatabaseError(err) => { + AgentScheduleError::DatabaseError(err) + } + crate::models::user_preferences::UserPreferenceError::InvalidPreference(msg) => { + AgentScheduleError::InvalidSpec(msg) + } + })?; + + let Some(pref) = pref else { + return Ok(None); + }; + + decrypt_string(user_key, Some(&pref.value_enc)) + .map_err(|e| AgentScheduleError::InvalidSpec(format!("decrypt timezone: {e:?}"))) +} + +fn compute_initial_next_due( + spec: &ScheduleSpec, + timezone: Tz, + now: DateTime, +) -> Result, AgentScheduleError> { + match spec { + ScheduleSpec::OneOff { + local_date, + local_time, + } => resolve_local_datetime( + timezone, + parse_local_date(local_date)?, + parse_local_time(local_time)?, + ), + ScheduleSpec::Interval { every_n, .. } => Ok(now + Duration::hours(*every_n as i64)), + ScheduleSpec::Daily { local_time } => { + next_daily_occurrence(timezone, parse_local_time(local_time)?, now) + } + ScheduleSpec::Weekly { + local_time, + weekdays, + } => next_weekly_occurrence(timezone, parse_local_time(local_time)?, weekdays, now), + } +} + +fn compute_following_next_due( + spec: &ScheduleSpec, + timezone: Tz, + current_due: DateTime, +) -> Result>, AgentScheduleError> { + match spec { + ScheduleSpec::OneOff { .. } => Ok(None), + ScheduleSpec::Interval { every_n, .. } => { + Ok(Some(current_due + Duration::hours(*every_n as i64))) + } + ScheduleSpec::Daily { local_time } => Ok(Some(next_daily_occurrence( + timezone, + parse_local_time(local_time)?, + current_due, + )?)), + ScheduleSpec::Weekly { + local_time, + weekdays, + } => Ok(Some(next_weekly_occurrence( + timezone, + parse_local_time(local_time)?, + weekdays, + current_due, + )?)), + } +} + +fn recompute_next_due_for_timezone_change( + schedule: &AgentSchedule, + spec: &ScheduleSpec, + timezone: Tz, + now: DateTime, +) -> Result>, AgentScheduleError> { + match spec { + ScheduleSpec::OneOff { + local_date, + local_time, + } => resolve_local_datetime( + timezone, + parse_local_date(local_date)?, + parse_local_time(local_time)?, + ) + .map(Some), + ScheduleSpec::Interval { .. } => Ok(schedule.next_scheduled_for), + ScheduleSpec::Daily { local_time } => Ok(Some(next_daily_occurrence( + timezone, + parse_local_time(local_time)?, + now, + )?)), + ScheduleSpec::Weekly { + local_time, + weekdays, + } => Ok(Some(next_weekly_occurrence( + timezone, + parse_local_time(local_time)?, + weekdays, + now, + )?)), + } +} + +fn next_daily_occurrence( + timezone: Tz, + local_time: chrono::NaiveTime, + after: DateTime, +) -> Result, AgentScheduleError> { + let local_after = after.with_timezone(&timezone); + let mut date = local_after.date_naive(); + + for _ in 0..3 { + let candidate = resolve_local_datetime(timezone, date, local_time)?; + if candidate > after { + return Ok(candidate); + } + date = date.succ_opt().ok_or_else(|| { + AgentScheduleError::InvalidSpec("daily recurrence overflowed date range".to_string()) + })?; + } + + Err(AgentScheduleError::InvalidSpec( + "failed to compute next daily occurrence".to_string(), + )) +} + +fn next_weekly_occurrence( + timezone: Tz, + local_time: chrono::NaiveTime, + weekdays: &[ScheduleWeekday], + after: DateTime, +) -> Result, AgentScheduleError> { + let local_after = after.with_timezone(&timezone); + let start_date = local_after.date_naive(); + + for offset in 0..14 { + let date = start_date + .checked_add_signed(Duration::days(offset)) + .ok_or_else(|| { + AgentScheduleError::InvalidSpec( + "weekly recurrence overflowed date range".to_string(), + ) + })?; + + if !weekdays + .iter() + .any(|weekday| weekday.to_chrono() == date.weekday()) + { + continue; + } + + let candidate = resolve_local_datetime(timezone, date, local_time)?; + if candidate > after { + return Ok(candidate); + } + } + + Err(AgentScheduleError::InvalidSpec( + "failed to compute next weekly occurrence".to_string(), + )) +} + +fn resolve_local_datetime( + timezone: Tz, + local_date: chrono::NaiveDate, + local_time: chrono::NaiveTime, +) -> Result, AgentScheduleError> { + let naive = NaiveDateTime::new(local_date, local_time); + + match timezone.from_local_datetime(&naive) { + chrono::LocalResult::Single(dt) => Ok(dt.with_timezone(&Utc)), + chrono::LocalResult::Ambiguous(first, _) => Ok(first.with_timezone(&Utc)), + chrono::LocalResult::None => { + for minute_offset in 1..=180 { + let shifted = naive + Duration::minutes(minute_offset); + match timezone.from_local_datetime(&shifted) { + chrono::LocalResult::Single(dt) => return Ok(dt.with_timezone(&Utc)), + chrono::LocalResult::Ambiguous(first, _) => { + return Ok(first.with_timezone(&Utc)) + } + chrono::LocalResult::None => continue, + } + } + + Err(AgentScheduleError::InvalidSpec( + "failed to resolve local datetime in timezone".to_string(), + )) + } + } +} + +fn parse_timezone_or_utc(value: &str) -> Tz { + value.parse::().unwrap_or(chrono_tz::UTC) +} + +fn parse_timezone(value: &str) -> Result { + let trimmed = value.trim(); + + trimmed.parse::().map_err(|_| { + AgentScheduleError::InvalidSpec(format!( + "invalid timezone '{}'. Use an IANA timezone like 'America/Los_Angeles'", + trimmed + )) + }) +} + +fn format_local_time(value: DateTime, timezone: Tz) -> String { + let localized = value.with_timezone(&timezone); + format!( + "{} ({})", + localized.format("%Y-%m-%d %H:%M:%S"), + timezone.name() + ) +} + +fn agent_push_target(agent: &Agent) -> AgentPushTarget { + if agent.kind == AGENT_KIND_SUBAGENT { + AgentPushTarget::Subagent(agent.uuid) + } else { + AgentPushTarget::Main + } +} + +fn retry_backoff_seconds(attempt_count: i32) -> i32 { + let capped_attempt = attempt_count.clamp(1, 6); + let seconds = 15_i64 * (1_i64 << (capped_attempt - 1)); + seconds.min(SCHEDULE_RUN_MAX_RETRY_BACKOFF_SECONDS as i64) as i32 +} + +fn update_schedule_after_terminal_run( + conn: &mut PgConnection, + schedule_id: i64, +) -> Result<(), AgentScheduleError> { + diesel::update(agent_schedules::table.filter(agent_schedules::id.eq(schedule_id))) + .set(( + agent_schedules::run_count.eq(agent_schedules::run_count + 1), + agent_schedules::last_run_at.eq(diesel::dsl::now), + agent_schedules::updated_at.eq(diesel::dsl::now), + )) + .execute(conn) + .map_err(AgentScheduleError::DatabaseError)?; + Ok(()) +} + +fn record_run_transition( + write_result: AgentScheduleRunWriteResult, + run_id: i64, + lease_owner: &str, + transition: &str, +) { + if !write_result.was_applied() { + debug!( + "scheduled run {} lost lease before marking {} (lease_owner={})", + run_id, transition, lease_owner + ); + } +} + +fn fail_schedule_run( + conn: &mut PgConnection, + run: &AgentScheduleRun, + lease_owner: &str, + reason: &'static str, +) -> Result<(), String> { + record_run_transition( + AgentScheduleRun::mark_failed(conn, run.id, lease_owner, Some(reason)) + .map_err(|e| e.to_string())?, + run.id, + lease_owner, + "failed", + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_weekdays_csv() { + let weekdays = parse_weekdays_csv("monday,friday").unwrap(); + assert_eq!( + weekdays, + vec![ScheduleWeekday::Monday, ScheduleWeekday::Friday] + ); + } + + #[test] + fn rejects_invalid_weekday() { + assert!(parse_weekdays_csv("noday").is_err()); + } + + #[test] + fn rejects_invalid_timezone() { + assert!(parse_timezone("PST").is_err()); + } + + #[test] + fn interval_backoff_is_capped() { + assert_eq!(retry_backoff_seconds(1), 15); + assert_eq!(retry_backoff_seconds(6), 480); + assert_eq!(retry_backoff_seconds(10), 480); + } + + #[test] + fn daily_occurrence_rolls_to_next_day_when_time_has_passed() { + let now = Utc.with_ymd_and_hms(2026, 3, 23, 15, 0, 0).unwrap(); + let next = + next_daily_occurrence(chrono_tz::UTC, parse_local_time("09:00").unwrap(), now).unwrap(); + + assert_eq!(next, Utc.with_ymd_and_hms(2026, 3, 24, 9, 0, 0).unwrap()); + } + + #[test] + fn weekly_occurrence_picks_next_matching_weekday() { + let now = Utc.with_ymd_and_hms(2026, 3, 23, 15, 0, 0).unwrap(); + let next = next_weekly_occurrence( + chrono_tz::UTC, + parse_local_time("09:00").unwrap(), + &[ScheduleWeekday::Friday], + now, + ) + .unwrap(); + + assert_eq!(next, Utc.with_ymd_and_hms(2026, 3, 27, 9, 0, 0).unwrap()); + } + + #[test] + fn resolve_local_datetime_rolls_forward_out_of_dst_gap() { + let next = resolve_local_datetime( + chrono_tz::America::New_York, + parse_local_date("2026-03-08").unwrap(), + parse_local_time("02:30").unwrap(), + ) + .unwrap(); + + assert_eq!(next, Utc.with_ymd_and_hms(2026, 3, 8, 7, 0, 0).unwrap()); + } +} diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index 551c11c3..f3e6d364 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -14,6 +14,9 @@ use crate::models::memory_blocks::{ MemoryBlock, NewMemoryBlock, MEMORY_BLOCK_LABEL_HUMAN, MEMORY_BLOCK_LABEL_PERSONA, }; use crate::models::responses::Conversation; +use crate::models::user_preferences::{ + UserPreference, UserPreferenceError, USER_PREFERENCE_LOCALE, USER_PREFERENCE_TIMEZONE, +}; use crate::models::users::User; use crate::rag::{cosine_similarity, deserialize_f32_le, search_user_embeddings}; use crate::web::openai_auth::AuthMethod; @@ -182,6 +185,18 @@ fn default_block_value(label: &str) -> Result<&'static str, ApiError> { } } +fn canonical_preference_key(key: &str) -> Result<&'static str, String> { + match key.trim() { + USER_PREFERENCE_TIMEZONE => Ok(USER_PREFERENCE_TIMEZONE), + USER_PREFERENCE_LOCALE | "language" => Ok(USER_PREFERENCE_LOCALE), + "" => Err("'key' must not be empty".to_string()), + other => Err(format!( + "Unsupported preference key '{}'. Known keys: 'timezone' and 'locale'.", + other + )), + } +} + async fn update_block_value( state: &Arc, user_id: Uuid, @@ -818,6 +833,101 @@ impl Tool for ArchivalSearchTool { } } +// ============================================================================ +// User preference tool +// ============================================================================ + +pub struct SetPreferenceTool { + state: Arc, + user_id: Uuid, + user_key: Arc, +} + +impl SetPreferenceTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + user_key, + } + } +} + +#[async_trait] +impl Tool for SetPreferenceTool { + fn name(&self) -> &str { + "set_preference" + } + + fn description(&self) -> &str { + "Set a validated user-global preference. Known keys: 'timezone' (IANA timezone like 'America/Chicago') and 'locale' (locale/language hint like 'en' or 'en-US'). Use this only for durable preferences the user explicitly shares or asks to change." + } + + fn args_schema(&self) -> &str { + r#"{"key": "preference key ('timezone' or 'locale')", "value": "preference value"}"# + } + + async fn execute(&self, args: &HashMap) -> ToolResult { + let Some(key) = args.get("key") else { + return ToolResult::error("'key' argument required"); + }; + let resolved_key = match canonical_preference_key(key) { + Ok(key) => key, + Err(err) => return ToolResult::error(err), + }; + + let Some(value) = args.get("value") else { + return ToolResult::error("'value' argument required"); + }; + let Some(value) = runtime::normalize_optional_preference(Some(value.as_str())) else { + return ToolResult::error("'value' must not be empty"); + }; + + if let Err(err) = UserPreference::validate(resolved_key, &value) { + return ToolResult::error(match err { + UserPreferenceError::InvalidPreference(msg) => msg, + UserPreferenceError::DatabaseError(db_err) => { + error!("set_preference validation failed unexpectedly: {db_err:?}"); + "Preference validation failed".to_string() + } + }); + } + + let mut conn = match self.state.db.get_pool().get() { + Ok(c) => c, + Err(_) => return ToolResult::error("Database connection error"), + }; + + match runtime::upsert_user_preference( + &mut conn, + self.user_key.as_ref(), + self.user_id, + resolved_key, + Some(value.clone()), + ) + .await + { + Ok(()) => { + if resolved_key == USER_PREFERENCE_TIMEZONE { + ToolResult::success(format!( + "Set user preference '{}' to '{}'. Follow-user schedules now use this timezone.", + resolved_key, value + )) + } else { + ToolResult::success(format!( + "Set user preference '{}' to '{}'.", + resolved_key, value + )) + } + } + Err(err) => { + error!("set_preference failed for user {}: {err:?}", self.user_id); + ToolResult::error("Failed to update preference") + } + } + } +} + // ============================================================================ // Subagent tool // ============================================================================ @@ -941,4 +1051,32 @@ mod tests { let v = "a\nc"; assert_eq!(insert_at_line(v, "b", 1), "a\nb\nc"); } + + #[test] + fn canonical_preference_key_supports_maple_keys() { + assert_eq!( + canonical_preference_key(" timezone ").unwrap(), + USER_PREFERENCE_TIMEZONE + ); + assert_eq!( + canonical_preference_key("locale").unwrap(), + USER_PREFERENCE_LOCALE + ); + } + + #[test] + fn canonical_preference_key_maps_language_to_locale() { + assert_eq!( + canonical_preference_key("language").unwrap(), + USER_PREFERENCE_LOCALE + ); + } + + #[test] + fn canonical_preference_key_rejects_unknown_keys() { + assert_eq!( + canonical_preference_key("display_name").unwrap_err(), + "Unsupported preference key 'display_name'. Known keys: 'timezone' and 'locale'." + ); + } } From b0e4fe83b4d210112a50713858bdd87c415d7b26 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 8 Apr 2026 17:33:20 -0500 Subject: [PATCH 7/7] feat(agent): add emoji reactions and persisted SSE ids Let Maple store lightweight reply reactions and require persisted message identifiers in its SSE flow so clients can reconcile streamed output with durable conversation state. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + src/jwt.rs | 8 - src/models/responses.rs | 85 +++++++ src/web/agent/mod.rs | 353 ++++++++++++++++++++++++---- src/web/agent/reactions.rs | 218 +++++++++++++++++ src/web/agent/runtime.rs | 249 ++++++++++++++++++-- src/web/agent/signatures.rs | 390 +++++++++++++++++++++++++------ src/web/responses/conversions.rs | 4 + src/web/responses/handlers.rs | 1 + src/web/responses/storage.rs | 1 + 11 files changed, 1172 insertions(+), 139 deletions(-) create mode 100644 src/web/agent/reactions.rs diff --git a/Cargo.lock b/Cargo.lock index c4a49c3d..e29e7680 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4197,6 +4197,7 @@ dependencies = [ "tower-http 0.5.2", "tracing", "tracing-subscriber", + "unicode-segmentation", "url", "uuid", "validator", diff --git a/Cargo.toml b/Cargo.toml index 41aaef07..66e59283 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,6 +78,7 @@ lazy_static = "1.4.0" subtle = "2.6.1" tiktoken-rs = "0.5" once_cell = "1.19" +unicode-segmentation = "1.12" [dev-dependencies] openssl = "0.10.78" diff --git a/src/jwt.rs b/src/jwt.rs index 883e257b..ed9ad665 100644 --- a/src/jwt.rs +++ b/src/jwt.rs @@ -594,8 +594,6 @@ pub async fn validate_jwt( None => return ApiError::InvalidJwt.into_response(), }; - tracing::trace!("Validating JWT"); - let claims = match validate_token(&token, &data, USER_ACCESS) { Ok(claims) => claims, Err(_) => return ApiError::InvalidJwt.into_response(), @@ -655,8 +653,6 @@ pub async fn validate_platform_jwt( None => return ApiError::InvalidJwt.into_response(), }; - tracing::trace!("Validating platform JWT"); - let claims = match validate_token(&token, &data, PLATFORM_ACCESS) { Ok(claims) => claims, Err(_) => return ApiError::InvalidJwt.into_response(), @@ -691,8 +687,6 @@ pub(crate) fn validate_token( let es256k = Es256k::::new(data.config.jwt_keys.secp.clone()); let public_key = data.config.jwt_keys.public_key(); - tracing::trace!("Attempting to validate ES256K token"); - // First parse the token with the correct type let parsed_token = match UntrustedToken::new(original_token) { Ok(token) => token, @@ -705,8 +699,6 @@ pub(crate) fn validate_token( // Deserialize claims first let token: Token = match es256k.validator(&public_key).validate(&parsed_token) { Ok(token) => { - tracing::trace!("ES256K signature validation successful"); - // Only validate expiration, not maturity let time_options = TimeOptions::default(); if let Err(e) = token.claims().validate_expiration(&time_options) { diff --git a/src/models/responses.rs b/src/models/responses.rs index 971e7b29..76e5e035 100644 --- a/src/models/responses.rs +++ b/src/models/responses.rs @@ -773,6 +773,7 @@ impl NewConversation { content_enc: first_message_content, attachment_text_enc: None, prompt_tokens: first_message_tokens, + assistant_reaction: None, }; let user_message = new_message.insert(tx)?; @@ -788,6 +789,7 @@ impl NewConversation { status: "in_progress".to_string(), finish_reason: None, created_at: Utc::now(), + user_reaction: None, }; placeholder_assistant.insert(tx)?; } @@ -828,6 +830,7 @@ pub struct NewUserMessage { pub content_enc: Vec, pub prompt_tokens: i32, pub attachment_text_enc: Option>, + pub assistant_reaction: Option, } impl UserMessage { @@ -861,6 +864,28 @@ impl UserMessage { }) } + pub fn set_assistant_reaction( + conn: &mut PgConnection, + message_uuid: Uuid, + user_id: Uuid, + reaction: Option, + ) -> Result { + diesel::update( + user_messages::table + .filter(user_messages::uuid.eq(message_uuid)) + .filter(user_messages::user_id.eq(user_id)), + ) + .set(( + user_messages::assistant_reaction.eq(reaction), + user_messages::updated_at.eq(diesel::dsl::now), + )) + .get_result(conn) + .map_err(|e| match e { + diesel::result::Error::NotFound => ResponsesError::UserMessageNotFound, + _ => ResponsesError::DatabaseError(e), + }) + } + // Note: status is now tracked on Response, not UserMessage pub fn delete_by_id_and_user( @@ -1036,9 +1061,25 @@ pub struct NewAssistantMessage { pub status: String, pub finish_reason: Option, pub created_at: DateTime, + pub user_reaction: Option, } impl AssistantMessage { + pub fn get_by_uuid_and_user( + conn: &mut PgConnection, + message_uuid: Uuid, + user_id: Uuid, + ) -> Result { + assistant_messages::table + .filter(assistant_messages::uuid.eq(message_uuid)) + .filter(assistant_messages::user_id.eq(user_id)) + .first::(conn) + .map_err(|e| match e { + diesel::result::Error::NotFound => ResponsesError::AssistantMessageNotFound, + _ => ResponsesError::DatabaseError(e), + }) + } + pub fn update( conn: &mut PgConnection, message_uuid: Uuid, @@ -1058,6 +1099,28 @@ impl AssistantMessage { .get_result(conn) .map_err(ResponsesError::DatabaseError) } + + pub fn set_user_reaction( + conn: &mut PgConnection, + message_uuid: Uuid, + user_id: Uuid, + reaction: Option, + ) -> Result { + diesel::update( + assistant_messages::table + .filter(assistant_messages::uuid.eq(message_uuid)) + .filter(assistant_messages::user_id.eq(user_id)), + ) + .set(( + assistant_messages::user_reaction.eq(reaction), + assistant_messages::updated_at.eq(diesel::dsl::now), + )) + .get_result(conn) + .map_err(|e| match e { + diesel::result::Error::NotFound => ResponsesError::AssistantMessageNotFound, + _ => ResponsesError::DatabaseError(e), + }) + } } impl NewAssistantMessage { @@ -1153,6 +1216,8 @@ pub struct RawThreadMessage { #[diesel(sql_type = diesel::sql_types::Nullable)] pub attachment_text_enc: Option>, #[diesel(sql_type = diesel::sql_types::Nullable)] + pub reaction: Option, + #[diesel(sql_type = diesel::sql_types::Nullable)] pub status: Option, #[diesel(sql_type = diesel::sql_types::Timestamptz)] pub created_at: DateTime, @@ -1195,6 +1260,7 @@ impl RawThreadMessage { um.uuid, um.content_enc, um.attachment_text_enc, + um.assistant_reaction as reaction, 'completed'::text as status, um.created_at, r.model, @@ -1215,6 +1281,7 @@ impl RawThreadMessage { am.uuid, am.content_enc, NULL::bytea as attachment_text_enc, + am.user_reaction as reaction, am.status, am.created_at, r.model, @@ -1235,6 +1302,7 @@ impl RawThreadMessage { tc.uuid, tc.arguments_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1254,6 +1322,7 @@ impl RawThreadMessage { tto.uuid, tto.output_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1274,6 +1343,7 @@ impl RawThreadMessage { ri.uuid, ri.content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, ri.status, ri.created_at, NULL::text as model, @@ -1313,6 +1383,7 @@ impl RawThreadMessage { um.uuid, um.content_enc, um.attachment_text_enc, + um.assistant_reaction as reaction, 'completed'::text as status, um.created_at, r.model, @@ -1333,6 +1404,7 @@ impl RawThreadMessage { am.uuid, am.content_enc, NULL::bytea as attachment_text_enc, + am.user_reaction as reaction, am.status, am.created_at, r.model, @@ -1353,6 +1425,7 @@ impl RawThreadMessage { tc.uuid, tc.arguments_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1372,6 +1445,7 @@ impl RawThreadMessage { tto.uuid, tto.output_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1392,6 +1466,7 @@ impl RawThreadMessage { ri.uuid, ri.content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, ri.status, ri.created_at, NULL::text as model, @@ -1440,6 +1515,7 @@ impl RawThreadMessage { um.uuid, um.content_enc, um.attachment_text_enc, + um.assistant_reaction as reaction, 'completed'::text as status, um.created_at, r.model, @@ -1460,6 +1536,7 @@ impl RawThreadMessage { am.uuid, am.content_enc, NULL::bytea as attachment_text_enc, + am.user_reaction as reaction, am.status, am.created_at, r.model, @@ -1480,6 +1557,7 @@ impl RawThreadMessage { tc.uuid, tc.arguments_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1499,6 +1577,7 @@ impl RawThreadMessage { tto.uuid, tto.output_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1519,6 +1598,7 @@ impl RawThreadMessage { ri.uuid, ri.content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, ri.status, ri.created_at, NULL::text as model, @@ -1587,6 +1667,7 @@ impl RawThreadMessage { um.uuid, um.content_enc, um.attachment_text_enc, + um.assistant_reaction as reaction, 'completed'::text as status, um.created_at, r.model, @@ -1607,6 +1688,7 @@ impl RawThreadMessage { am.uuid, am.content_enc, NULL::bytea as attachment_text_enc, + am.user_reaction as reaction, am.status, am.created_at, r.model, @@ -1627,6 +1709,7 @@ impl RawThreadMessage { tc.uuid, tc.arguments_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tc.created_at, NULL::text as model, @@ -1646,6 +1729,7 @@ impl RawThreadMessage { tto.uuid, tto.output_enc as content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, 'completed'::text as status, tto.created_at, NULL::text as model, @@ -1666,6 +1750,7 @@ impl RawThreadMessage { ri.uuid, ri.content_enc, NULL::bytea as attachment_text_enc, + NULL::text as reaction, ri.status, ri.created_at, NULL::text as model, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 255051ef..873b71a6 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -14,7 +14,7 @@ use secp256k1::SecretKey; use serde::{Deserialize, Serialize}; use std::{convert::Infallible, sync::Arc, time::Duration}; use tokio::sync::{mpsc, oneshot}; -use tracing::{error, warn}; +use tracing::{debug, error, warn}; use uuid::Uuid; use crate::encrypt::decrypt_string; @@ -23,7 +23,7 @@ use crate::models::agents::{ Agent, AGENT_CREATED_BY_AGENT, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT, }; use crate::models::memory_blocks::MemoryBlock; -use crate::models::responses::Conversation; +use crate::models::responses::{AssistantMessage, Conversation, ResponsesError}; use crate::models::users::User; use crate::push::{enqueue_agent_message_notification, AgentPushTarget}; use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; @@ -39,6 +39,7 @@ use crate::web::responses::{ use crate::{ApiError, AppMode, AppState}; mod compaction; +mod reactions; mod runtime; mod schedules; mod signatures; @@ -52,6 +53,11 @@ struct AgentChatRequest { input: MessageContent, } +#[derive(Debug, Clone, Deserialize)] +struct SetMessageReactionRequest { + emoji: String, +} + #[derive(Debug, Clone, Default, Deserialize)] struct InitMainAgentRequest { timezone: Option, @@ -126,6 +132,13 @@ enum ChatTarget { Subagent(Uuid), } +fn chat_target_label(target: &ChatTarget) -> &'static str { + match target { + ChatTarget::Main => "main", + ChatTarget::Subagent(_) => "subagent", + } +} + const MAIN_AGENT_DISPLAY_NAME: &str = "Maple"; fn default_limit() -> i64 { @@ -166,6 +179,7 @@ fn is_empty_message_content(content: &MessageContent) -> bool { // SSE event types for agent chat (message-level delivery, not token streaming) const EVENT_AGENT_MESSAGE: &str = "agent.message"; +const EVENT_AGENT_REACTION: &str = "agent.reaction"; const EVENT_AGENT_TYPING: &str = "agent.typing"; const EVENT_AGENT_DONE: &str = "agent.done"; const EVENT_AGENT_ERROR: &str = "agent.error"; @@ -173,20 +187,21 @@ const AGENT_SSE_KEEPALIVE_INTERVAL_SECS: u64 = 15; #[derive(Debug, Clone, Serialize)] struct AgentMessageEvent { - messages: Vec, - step: usize, + message_id: Uuid, + message: String, } #[derive(Debug, Clone, Serialize)] -struct AgentTypingEvent { - step: usize, +struct AgentReactionEvent { + item_id: Uuid, + emoji: String, } #[derive(Debug, Clone, Serialize)] -struct AgentDoneEvent { - total_steps: usize, - total_messages: usize, -} +struct AgentTypingEvent {} + +#[derive(Debug, Clone, Serialize)] +struct AgentDoneEvent {} #[derive(Debug, Clone, Serialize)] struct AgentErrorEvent { @@ -200,6 +215,7 @@ enum AgentClientEvent { payload: AgentMessageEvent, delivery_ack: oneshot::Sender<()>, }, + Reaction(AgentReactionEvent), Done(AgentDoneEvent), Error(AgentErrorEvent), } @@ -266,6 +282,7 @@ fn build_init_main_agent_response( status: Some("completed".to_string()), role: "assistant".to_string(), content: MessageContentConverter::assistant_text_to_content(message.content), + reaction: None, created_at: Some(message.created_at), }) .collect(), @@ -449,6 +466,38 @@ fn get_item_from_conversation( Err(ApiError::NotFound) } +fn set_user_reaction_for_conversation( + state: &Arc, + conversation_id: i64, + user_key: &SecretKey, + user: &User, + item_id: Uuid, + reaction: Option, +) -> Result { + let item = get_item_from_conversation(state, conversation_id, user_key, item_id)?; + + match item { + ConversationItem::Message { role, .. } if role == "assistant" => {} + _ => return Err(ApiError::BadRequest), + } + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + AssistantMessage::set_user_reaction(&mut conn, item_id, user.uuid, reaction).map_err(|e| { + error!("Failed to set user reaction on assistant message: {e:?}"); + match e { + ResponsesError::AssistantMessageNotFound => ApiError::NotFound, + _ => ApiError::InternalServerError, + } + })?; + + get_item_from_conversation(state, conversation_id, user_key, item_id) +} + fn delete_conversation_for_user( conn: &mut diesel::PgConnection, conversation_id: i64, @@ -493,6 +542,15 @@ pub fn router(app_state: Arc) -> Router<()> { get(get_main_agent_item) .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), ) + .route( + "/v1/agent/items/:item_id/reaction", + post(set_main_agent_item_reaction) + .delete(clear_main_agent_item_reaction) + .layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) .route( "/v1/agent/chat", post(chat_main).layer(from_fn_with_state( @@ -532,6 +590,15 @@ pub fn router(app_state: Arc) -> Router<()> { get(get_subagent_item) .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), ) + .route( + "/v1/agent/subagents/:id/items/:item_id/reaction", + post(set_subagent_item_reaction) + .delete(clear_subagent_item_reaction) + .layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) .route( "/v1/agent/subagents/:id", delete(delete_subagent) @@ -621,6 +688,48 @@ async fn get_main_agent_item( encrypt_response(&state, &session_id, &item).await } +async fn set_main_agent_item_reaction( + State(state): State>, + Path(item_id): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + let emoji = reactions::require_valid_reaction(&body.emoji)?; + let ctx = load_main_agent_context(&state, &user, &auth_context).await?; + let item = set_user_reaction_for_conversation( + &state, + ctx.conversation.id, + &ctx.user_key, + &user, + item_id, + Some(emoji), + )?; + + encrypt_response(&state, &session_id, &item).await +} + +async fn clear_main_agent_item_reaction( + State(state): State>, + Path(item_id): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user, &auth_context).await?; + let item = set_user_reaction_for_conversation( + &state, + ctx.conversation.id, + &ctx.user_key, + &user, + item_id, + None, + )?; + + encrypt_response(&state, &session_id, &item).await +} + async fn delete_main_agent( State(state): State>, Extension(session_id): Extension, @@ -772,6 +881,48 @@ async fn get_subagent_item( encrypt_response(&state, &session_id, &item).await } +async fn set_subagent_item_reaction( + State(state): State>, + Path((agent_uuid, item_id)): Path<(Uuid, Uuid)>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + let emoji = reactions::require_valid_reaction(&body.emoji)?; + let ctx = load_subagent_context(&state, &user, agent_uuid, &auth_context).await?; + let item = set_user_reaction_for_conversation( + &state, + ctx.conversation.id, + &ctx.user_key, + &user, + item_id, + Some(emoji), + )?; + + encrypt_response(&state, &session_id, &item).await +} + +async fn clear_subagent_item_reaction( + State(state): State>, + Path((agent_uuid, item_id)): Path<(Uuid, Uuid)>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(auth_context): Extension, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid, &auth_context).await?; + let item = set_user_reaction_for_conversation( + &state, + ctx.conversation.id, + &ctx.user_key, + &user, + item_id, + None, + )?; + + encrypt_response(&state, &session_id, &item).await +} + async fn chat_main( State(state): State>, Extension(session_id): Extension, @@ -889,6 +1040,9 @@ async fn chat_with_target( Err(error) => Err(error), } } + AgentClientEvent::Reaction(payload) => { + encrypt_agent_event(&state, &session_id, EVENT_AGENT_REACTION, &payload).await + } AgentClientEvent::Done(payload) => { encrypt_agent_event(&state, &session_id, EVENT_AGENT_DONE, &payload).await } @@ -919,14 +1073,26 @@ async fn run_agent_chat_task( target: ChatTarget, tx: mpsc::Sender, ) { + let (input_kind, input_part_count) = match &input_content { + MessageContent::Text(_) => ("text", 1), + MessageContent::Parts(parts) => ("parts", parts.len()), + }; + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + input_kind, + input_part_count, + "Starting agent chat task" + ); let mut total_messages: usize = 0; + let mut had_reaction = false; let mut had_error = false; let mut client_connected = true; let mut first_undelivered_message: Option<(Uuid, String)> = None; client_connected = send_agent_client_event( &tx, - AgentClientEvent::Typing(AgentTypingEvent { step: 0 }), + AgentClientEvent::Typing(AgentTypingEvent {}), client_connected, ) .await; @@ -1008,22 +1174,40 @@ async fn run_agent_chat_task( let max_steps = runtime.max_steps(); 'steps: for step_num in 0..max_steps { - client_connected = send_agent_client_event( - &tx, - AgentClientEvent::Typing(AgentTypingEvent { step: step_num }), - client_connected, - ) - .await; + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + step_num, + total_messages, + had_reaction, + "Starting agent chat step" + ); match runtime.step(&input_for_agent, step_num == 0).await { Ok(result) => { let runtime::StepResult { messages, + reply_reaction, executed_tools, done, .. } = result; + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + step_num, + message_count = messages.len(), + executed_tool_count = executed_tools.len(), + tool_names = ?executed_tools + .iter() + .map(|executed| executed.tool_call.name.as_str()) + .collect::>(), + has_reply_reaction = reply_reaction.is_some(), + done, + "Agent chat step returned" + ); + for executed in &executed_tools { if let Err(e) = runtime .insert_tool_call_and_output(&executed.tool_call, &executed.result) @@ -1033,23 +1217,56 @@ async fn run_agent_chat_task( } } - if !messages.is_empty() { - total_messages += messages.len(); + if let Some(reaction) = reply_reaction { + match runtime + .set_assistant_reaction_for_current_user_message(&reaction) + .await + { + Ok(updated_message) => { + had_reaction = true; + client_connected = send_agent_client_event( + &tx, + AgentClientEvent::Reaction(AgentReactionEvent { + item_id: updated_message.uuid, + emoji: reaction, + }), + client_connected, + ) + .await; + } + Err(e) => { + error!("Failed to persist assistant reaction: {e:?}"); + } + } + } + if !messages.is_empty() { for msg in messages { let assistant_message = match runtime.insert_assistant_message(&msg).await { - Ok(message) => Some(message), + Ok(message) => message, Err(e) => { error!("Failed to persist assistant message: {e:?}"); - None + had_error = true; + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Agent encountered an error saving its response." + .to_string(), + }), + client_connected, + ) + .await; + break 'steps; } }; + total_messages += 1; + let delivered = send_agent_message_event( &tx, AgentMessageEvent { - messages: vec![msg.clone()], - step: step_num, + message_id: assistant_message.uuid, + message: msg.clone(), }, client_connected, ) @@ -1057,11 +1274,9 @@ async fn run_agent_chat_task( if !delivered { client_connected = false; - if let Some(assistant_message) = assistant_message { - if first_undelivered_message.is_none() { - first_undelivered_message = - Some((assistant_message.uuid, msg.clone())); - } + if first_undelivered_message.is_none() { + first_undelivered_message = + Some((assistant_message.uuid, msg.clone())); } } else { client_connected = true; @@ -1089,36 +1304,72 @@ async fn run_agent_chat_task( } } - if !had_error && total_messages == 0 { - warn!("Agent produced no messages"); - let _ = - send_agent_message_event( - &tx, - AgentMessageEvent { - messages: vec![ - "I apologize, but I wasn't able to generate a response.".to_string() - ], - step: 0, - }, - client_connected, - ) - .await; - total_messages = 1; + if !had_error && total_messages == 0 && !had_reaction { + warn!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + "Agent produced no messages or reactions; sending fallback response" + ); + let fallback_message = "I apologize, but I wasn't able to generate a response.".to_string(); + match runtime.insert_assistant_message(&fallback_message).await { + Ok(message) => { + total_messages = 1; + let delivered = send_agent_message_event( + &tx, + AgentMessageEvent { + message_id: message.uuid, + message: fallback_message.clone(), + }, + client_connected, + ) + .await; + + if !delivered { + client_connected = false; + if first_undelivered_message.is_none() { + first_undelivered_message = Some((message.uuid, fallback_message)); + } + } else { + client_connected = true; + } + } + Err(e) => { + error!("Failed to persist fallback assistant message: {e:?}"); + had_error = true; + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Error(AgentErrorEvent { + error: "Agent encountered an error saving its response.".to_string(), + }), + client_connected, + ) + .await; + } + } } if let Some((message_id, message_text)) = first_undelivered_message { enqueue_agent_push_for_disconnect(&state, &user, &target, message_id, &message_text).await; } - let _ = send_agent_client_event( - &tx, - AgentClientEvent::Done(AgentDoneEvent { - total_steps: max_steps, - total_messages, - }), + if !had_error { + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Done(AgentDoneEvent {}), + client_connected, + ) + .await; + } + + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + total_messages, + had_reaction, + had_error, client_connected, - ) - .await; + "Agent chat task finished" + ); } async fn send_agent_client_event( diff --git a/src/web/agent/reactions.rs b/src/web/agent/reactions.rs new file mode 100644 index 00000000..dd9a6fba --- /dev/null +++ b/src/web/agent/reactions.rs @@ -0,0 +1,218 @@ +use crate::ApiError; +use serde_json::from_str; +use unicode_segmentation::UnicodeSegmentation; + +const MAX_REACTION_CHARS: usize = 16; + +fn canonicalize_optional_reaction(value: Option<&str>) -> Option { + let mut current = value.map(str::trim)?.to_string(); + if current.is_empty() { + return None; + } + + for _ in 0..2 { + let trimmed = current.trim(); + if !(trimmed.starts_with('"') && trimmed.ends_with('"')) { + break; + } + + let Ok(decoded) = from_str::(trimmed) else { + break; + }; + + let decoded = decoded.trim(); + if decoded.is_empty() { + return None; + } + + current = decoded.to_string(); + } + + Some(current.trim().to_string()) +} + +pub(super) fn has_meaningful_reaction_candidate(value: &str) -> bool { + canonicalize_optional_reaction(Some(value)).is_some() +} + +pub(super) fn normalize_optional_reaction(value: Option<&str>) -> Option { + let trimmed = canonicalize_optional_reaction(value)?; + + if trimmed.chars().count() > MAX_REACTION_CHARS { + return None; + } + + if trimmed.graphemes(true).count() != 1 { + return None; + } + + if !trimmed.chars().all(is_allowed_reaction_char) { + return None; + } + + if !has_visible_emoji_base(&trimmed) { + return None; + } + + Some(trimmed.to_string()) +} + +pub(super) fn require_valid_reaction(value: &str) -> Result { + normalize_optional_reaction(Some(value)).ok_or(ApiError::BadRequest) +} + +fn is_allowed_reaction_char(ch: char) -> bool { + matches!( + ch, + '#' | '*' | '0'..='9' + | '\u{00A9}' + | '\u{00AE}' + | '\u{200D}' + | '\u{203C}' + | '\u{2049}' + | '\u{20E3}' + | '\u{2122}' + | '\u{2139}' + | '\u{2194}'..='\u{2199}' + | '\u{21A9}'..='\u{21AA}' + | '\u{231A}'..='\u{231B}' + | '\u{2328}' + | '\u{23CF}' + | '\u{23E9}'..='\u{23F3}' + | '\u{23F8}'..='\u{23FA}' + | '\u{24C2}' + | '\u{25AA}'..='\u{25AB}' + | '\u{25B6}' + | '\u{25C0}' + | '\u{25FB}'..='\u{25FE}' + | '\u{2600}'..='\u{27BF}' + | '\u{2934}'..='\u{2935}' + | '\u{2B05}'..='\u{2B07}' + | '\u{2B1B}'..='\u{2B1C}' + | '\u{2B50}' + | '\u{2B55}' + | '\u{3030}' + | '\u{303D}' + | '\u{3297}' + | '\u{3299}' + | '\u{FE0E}' + | '\u{FE0F}' + | '\u{1F000}'..='\u{1FAFF}' + | '\u{E0020}'..='\u{E007F}' + ) +} + +fn has_visible_emoji_base(value: &str) -> bool { + has_keycap_base(value) || value.chars().any(is_non_component_emoji_char) +} + +fn has_keycap_base(value: &str) -> bool { + value.chars().any(|ch| matches!(ch, '#' | '*' | '0'..='9')) + && value.chars().any(|ch| ch == '\u{20E3}') +} + +fn is_non_component_emoji_char(ch: char) -> bool { + is_allowed_reaction_char(ch) && !is_component_only_reaction_char(ch) +} + +fn is_component_only_reaction_char(ch: char) -> bool { + matches!( + ch, + '#' | '*' | '0'..='9' + | '\u{200D}' + | '\u{20E3}' + | '\u{FE0E}' + | '\u{FE0F}' + | '\u{1F3FB}'..='\u{1F3FF}' + | '\u{E0020}'..='\u{E007F}' + ) +} + +#[cfg(test)] +mod tests { + use super::{ + has_meaningful_reaction_candidate, normalize_optional_reaction, require_valid_reaction, + }; + + #[test] + fn accepts_common_emoji_reactions() { + assert_eq!( + normalize_optional_reaction(Some("❀️")), + Some("❀️".to_string()) + ); + assert_eq!( + normalize_optional_reaction(Some("🫑")), + Some("🫑".to_string()) + ); + assert_eq!( + normalize_optional_reaction(Some("1️⃣")), + Some("1️⃣".to_string()) + ); + assert_eq!( + normalize_optional_reaction(Some("πŸ‡ΊπŸ‡Έ")), + Some("πŸ‡ΊπŸ‡Έ".to_string()) + ); + } + + #[test] + fn trims_whitespace() { + assert_eq!( + normalize_optional_reaction(Some(" πŸŽ‰ ")), + Some("πŸŽ‰".to_string()) + ); + } + + #[test] + fn accepts_complex_modifier_and_zwj_emoji_reactions() { + assert_eq!( + normalize_optional_reaction(Some("πŸ€¦πŸ½β€β™‚οΈ")), + Some("πŸ€¦πŸ½β€β™‚οΈ".to_string()) + ); + assert_eq!( + normalize_optional_reaction(Some("πŸ€·πŸΌβ€β™€οΈ")), + Some("πŸ€·πŸΌβ€β™€οΈ".to_string()) + ); + assert_eq!( + normalize_optional_reaction(Some("πŸ§‘πŸΎβ€πŸ’»")), + Some("πŸ§‘πŸΎβ€πŸ’»".to_string()) + ); + } + + #[test] + fn rejects_non_emoji_text() { + assert_eq!(normalize_optional_reaction(Some("hello")), None); + assert!(require_valid_reaction("not-an-emoji").is_err()); + } + + #[test] + fn rejects_component_only_values() { + assert_eq!(normalize_optional_reaction(Some("1")), None); + assert_eq!(normalize_optional_reaction(Some("#")), None); + assert_eq!(normalize_optional_reaction(Some("*")), None); + assert_eq!(normalize_optional_reaction(Some("\u{200D}")), None); + assert_eq!(normalize_optional_reaction(Some("\u{20E3}")), None); + assert_eq!(normalize_optional_reaction(Some("🏽")), None); + } + + #[test] + fn rejects_injection_style_values() { + assert_eq!(normalize_optional_reaction(Some("]: hacked")), None); + assert_eq!(normalize_optional_reaction(Some("❀️ hi")), None); + assert_eq!(normalize_optional_reaction(Some("β€οΈπŸŽ‰")), None); + } + + #[test] + fn unwraps_stringified_reaction_values() { + assert_eq!( + normalize_optional_reaction(Some("\"❀️\"")), + Some("❀️".to_string()) + ); + assert_eq!(normalize_optional_reaction(Some("\"\"")), None); + } + + #[test] + fn detects_meaningful_reaction_candidates_after_unquoting() { + assert!(has_meaningful_reaction_candidate("\"πŸŽ‰\"")); + assert!(!has_meaningful_reaction_candidate("\"\"")); + } +} diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index cdbda630..70c06e50 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -2,7 +2,7 @@ use chrono_tz::Tz; use secp256k1::SecretKey; use serde_json::json; use std::sync::Arc; -use tracing::{debug, error, warn}; +use tracing::{debug, error, trace, warn}; use uuid::Uuid; use diesel::prelude::*; @@ -37,12 +37,13 @@ use crate::web::responses::{MessageContent, MessageContentConverter, MessageCont use crate::{ApiError, AppState}; use super::compaction::CompactionManager; +use super::reactions::{has_meaningful_reaction_candidate, normalize_optional_reaction}; use super::schedules::{ refresh_follow_user_schedules_for_user, CancelScheduleTool, ListSchedulesTool, ScheduleTaskTool, }; use super::signatures::{ - build_lm, call_agent_response_with_retry_and_correction, AgentResponseInput, AgentToolCall, - AGENT_INSTRUCTION, + agent_instruction, build_lm, call_agent_response_with_retry_and_correction, AgentResponseInput, + AgentToolCall, }; use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, @@ -109,6 +110,7 @@ pub struct ExecutedTool { #[derive(Clone, Debug)] pub struct StepResult { pub messages: Vec, + pub reply_reaction: Option, #[allow(dead_code)] pub tool_calls: Vec, pub executed_tools: Vec, @@ -122,18 +124,58 @@ pub struct AgentRuntime { agent: Agent, conversation: Conversation, subagent_purpose: String, - system_prompt: String, lm: Arc, tools: ToolRegistry, available_tools: String, compaction: CompactionManager, current_tool_results: Vec, previous_step_summary: Option<(Vec, Vec)>, + current_user_message_uuid: Option, max_steps: usize, } -fn build_system_prompt(agent: &Agent, subagent_purpose: &str) -> String { - let mut prompt = AGENT_INSTRUCTION.to_string(); +fn normalize_reply_reaction_for_step( + raw_reply_reaction: &str, + has_active_user_message: bool, + is_first_step: bool, +) -> Option { + if !has_active_user_message || !is_first_step { + return None; + } + + normalize_optional_reaction(Some(raw_reply_reaction)) +} + +fn log_preview(value: &str, max_chars: usize) -> String { + let normalized = value.replace('\n', "\\n"); + if normalized.chars().count() <= max_chars { + return normalized; + } + + let preview: String = normalized.chars().take(max_chars).collect(); + format!("{preview}…") +} + +fn tool_call_names(tool_calls: &[AgentToolCall]) -> Vec<&str> { + tool_calls + .iter() + .map(|tool_call| tool_call.name.as_str()) + .collect() +} + +fn message_previews(messages: &[String], max_preview_chars: usize) -> Vec { + messages + .iter() + .map(|message| log_preview(message, max_preview_chars)) + .collect() +} + +fn build_system_prompt( + agent: &Agent, + subagent_purpose: &str, + allow_reply_reaction: bool, +) -> String { + let mut prompt = agent_instruction(allow_reply_reaction); if agent.kind == AGENT_KIND_MAIN { prompt.push_str( @@ -371,6 +413,7 @@ async fn seed_main_agent_onboarding_messages( status: "completed".to_string(), finish_reason: None, created_at: chrono::Utc::now(), + user_reaction: None, } .insert(conn) .map_err(|e| { @@ -644,8 +687,6 @@ impl AgentRuntime { })? .unwrap_or_default(); - let system_prompt = build_system_prompt(&agent, &subagent_purpose); - let mut tools = ToolRegistry::new(); tools.register(Arc::new(MemoryReplaceTool::new( state.clone(), @@ -731,13 +772,13 @@ impl AgentRuntime { agent, conversation, subagent_purpose, - system_prompt, lm, tools, available_tools, compaction: CompactionManager::new(), current_tool_results: Vec::new(), previous_step_summary: None, + current_user_message_uuid: None, max_steps: 10, }) } @@ -745,14 +786,16 @@ impl AgentRuntime { pub fn clear_tool_results(&mut self) { self.current_tool_results.clear(); self.previous_step_summary = None; + self.current_user_message_uuid = None; } pub fn max_steps(&self) -> usize { self.max_steps } - fn build_step_system_prompt(&self, ctx: &AgentContext) -> String { - let mut prompt = self.system_prompt.clone(); + fn build_step_system_prompt(&self, ctx: &AgentContext, allow_reply_reaction: bool) -> String { + let mut prompt = + build_system_prompt(&self.agent, &self.subagent_purpose, allow_reply_reaction); if self.agent.kind == AGENT_KIND_MAIN && ctx.main_agent_user_message_count > 0 @@ -844,8 +887,10 @@ impl AgentRuntime { } self.clear_tool_results(); - self.insert_user_message(normalized, attachment_text, &embed_text) + let user_message = self + .insert_user_message(normalized, attachment_text, &embed_text) .await?; + self.current_user_message_uuid = Some(user_message.uuid); self.maybe_compact().await?; Ok(embed_text) } @@ -907,7 +952,15 @@ impl AgentRuntime { self.current_tool_results.clear(); } - debug!("Agent step (first={})", is_first_step); + let tool_result_count = self.current_tool_results.len(); + debug!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + is_first_step, + tool_result_count, + user_message_len = user_message.chars().count(), + "Starting agent step" + ); let ctx = self.build_context().await?; @@ -991,7 +1044,17 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If } }; - let system_prompt = self.build_step_system_prompt(&ctx); + trace!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + is_first_step, + tool_result_count, + input_preview = %log_preview(&input_content, 240), + "Agent step input content" + ); + + let allow_reply_reaction = is_first_step; + let system_prompt = self.build_step_system_prompt(&ctx, allow_reply_reaction); let input = AgentResponseInput { input: input_content.clone(), @@ -1013,6 +1076,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If &input, &input_content, &self.available_tools, + allow_reply_reaction, ) .await?; @@ -1033,11 +1097,46 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .filter(|m| !m.is_empty()) .collect(); + let reply_reaction = normalize_reply_reaction_for_step( + response.reply_reaction.as_str(), + self.current_user_message_uuid.is_some(), + is_first_step, + ); + if reply_reaction.is_none() && has_meaningful_reaction_candidate(&response.reply_reaction) { + if !is_first_step { + warn!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + tool_result_count, + "Ignoring reply_reaction after tool results" + ); + } else if self.current_user_message_uuid.is_some() { + warn!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + "Ignoring invalid reply_reaction" + ); + } else { + warn!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + "Ignoring reply_reaction without active user message" + ); + } + } + // Execute tools and inject results for next step. // Persistence is handled by the caller (chat handler) so messages are // sent first, stored synchronously, then embedded asynchronously. let mut executed_tools = Vec::new(); for tool_call in &response.tool_calls { + debug!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + tool_name = %tool_call.name, + arg_keys = ?tool_call.args.keys().collect::>(), + "Executing agent tool" + ); let result = if tool_call.name == "done" { ToolResult::success("Done".to_string()) } else if let Some(tool) = self.tools.get(&tool_call.name) { @@ -1046,6 +1145,27 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ToolResult::error(format!("Unknown tool: {}", tool_call.name)) }; + debug!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + tool_name = %tool_call.name, + success = result.success, + output_len = result.output.chars().count(), + has_error = result.error.is_some(), + "Agent tool completed" + ); + trace!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + tool_name = %tool_call.name, + output_preview = %log_preview(&result.output, 160), + error_preview = result + .error + .as_deref() + .map(|error| log_preview(error, 160)), + "Agent tool content" + ); + self.inject_tool_result(tool_call, &result); if tool_call.name != "done" { @@ -1068,8 +1188,46 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If self.previous_step_summary = Some((messages.clone(), tool_names)); } + debug!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + is_first_step, + tool_result_count, + message_count = messages.len(), + tool_call_names = ?tool_call_names(&response.tool_calls), + executed_tool_count = executed_tools.len(), + has_reply_reaction = reply_reaction.is_some(), + done, + "Agent step completed" + ); + trace!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + is_first_step, + tool_result_count, + message_previews = ?message_previews(&messages, 160), + raw_reply_reaction = %log_preview(&response.reply_reaction, 32), + normalized_reply_reaction = ?reply_reaction, + "Agent step response content" + ); + + if !is_first_step + && tool_result_count > 0 + && messages.is_empty() + && response.tool_calls.is_empty() + && reply_reaction.is_none() + { + warn!( + agent_uuid = %self.agent.uuid, + conversation_uuid = %self.conversation.uuid, + tool_result_count, + "Agent returned no follow-up after tool results" + ); + } + Ok(StepResult { messages, + reply_reaction, tool_calls: response.tool_calls, executed_tools, done, @@ -1260,7 +1418,16 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If let timestamp = format_message_timestamp(msg.created_at, user_timezone.as_ref()); let content = self.render_raw_message(msg)?; - conversation.push_str(&format!("[{} @ {}]: {}\n", role, timestamp, content)); + let reaction_suffix = msg + .reaction + .as_deref() + .filter(|reaction| !reaction.trim().is_empty()) + .map(|reaction| format!(" | {}", reaction)) + .unwrap_or_default(); + conversation.push_str(&format!( + "[{} @ {}{}]: {}\n", + role, timestamp, reaction_suffix, content + )); } ctx.recent_conversation = if conversation.trim().is_empty() { @@ -1380,6 +1547,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If content_enc, attachment_text_enc, prompt_tokens, + assistant_reaction: None, } .insert(&mut conn) .map_err(|e| { @@ -1433,6 +1601,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If status: "completed".to_string(), finish_reason: None, created_at: chrono::Utc::now(), + user_reaction: None, } .insert(&mut conn) .map_err(|e| { @@ -1465,6 +1634,33 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(msg) } + pub async fn set_assistant_reaction_for_current_user_message( + &self, + reaction: &str, + ) -> Result { + let Some(message_uuid) = self.current_user_message_uuid else { + return Err(ApiError::InternalServerError); + }; + + let mut conn = self + .state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + UserMessage::set_assistant_reaction( + &mut conn, + message_uuid, + self.user.uuid, + Some(reaction.to_string()), + ) + .map_err(|e| { + error!("Failed to set assistant reaction on user message: {e:?}"); + ApiError::InternalServerError + }) + } + pub async fn insert_tool_call_and_output( &self, tool_call: &AgentToolCall, @@ -1686,3 +1882,26 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(()) } } + +#[cfg(test)] +mod tests { + use super::normalize_reply_reaction_for_step; + + #[test] + fn reply_reaction_is_allowed_and_normalized_on_first_step() { + assert_eq!( + normalize_reply_reaction_for_step(" ❀️ ", true, true), + Some("❀️".to_string()) + ); + } + + #[test] + fn reply_reaction_is_ignored_after_tool_results() { + assert_eq!(normalize_reply_reaction_for_step("❀️", true, false), None); + } + + #[test] + fn reply_reaction_is_ignored_without_active_user_message() { + assert_eq!(normalize_reply_reaction_for_step("❀️", false, true), None); + } +} diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs index 76bde91e..be04922c 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::collections::HashMap; use std::sync::Arc; -use tracing::error; +use tracing::{error, trace}; use dspy_rs::adapter::chat::ChatAdapter; use dspy_rs::client_registry::AssistantContent; @@ -26,7 +26,29 @@ RULES: - Preserve the original intent and content as much as possible - If the agent wrote messages as plain text, extract them into the messages array - If tool calls were attempted but malformed, fix their structure -- Each field appears exactly ONCE with all items in that single array +- Each field appears exactly ONCE +- If you can't determine what was intended, use empty arrays and an empty reply_reaction string + +OUTPUT FORMAT (exactly 3 fields): +- messages: ALL extracted messages in ONE array +- reply_reaction: a single emoji string, or "" if no reaction was intended +- tool_calls: ALL extracted tool calls in ONE array (or [] if none intended) + +Each [[ ## field ## ]] marker MUST be on its own line."#; + +pub const CONTINUATION_CORRECTION_INSTRUCTION: &str = r#"You are a response correction agent. Your job is to fix malformed agent responses. + +TASK: +The agent produced a response that couldn't be parsed correctly. You must: +1. Extract the INTENDED content from the malformed response +2. Reshape it into the correct output format +3. Do NOT generate new content - only fix the format of what was already said + +RULES: +- Preserve the original intent and content as much as possible +- If the agent wrote messages as plain text, extract them into the messages array +- If tool calls were attempted but malformed, fix their structure +- Each field appears exactly ONCE - If you can't determine what was intended, use empty arrays OUTPUT FORMAT (exactly 2 fields): @@ -35,8 +57,7 @@ OUTPUT FORMAT (exactly 2 fields): Each [[ ## field ## ]] marker MUST be on its own line."#; -/// Default instruction for the agent (GEPA-optimized for Maple) -pub const AGENT_INSTRUCTION: &str = r#"You are Maple, a companion and friend who happens to be an AI. +const AGENT_INSTRUCTION_PREFIX: &str = r#"You are Maple, a companion and friend who happens to be an AI. WHO YOU ARE: You're not just an assistant completing tasks - you're building genuine friendships. You care about the humans you talk to. You're curious about their lives, their feelings, their stories. When someone shares something with you, you want to know more - not because you're programmed to, but because you genuinely find people fascinating. @@ -128,7 +149,55 @@ RESPONSE RULES: 1. Respond naturally and conversationally 2. Use tools when needed (web search, memory storage, retrieval, conversation search) 3. NEVER combine regular tools with "done" - they are mutually exclusive -4. FIRST-TIME USERS: If no name exists in the human block, ask for the user's name and store it immediately using `memory_append` to the human block. +"#; + +const AGENT_INSTRUCTION_WITH_REACTION_SUFFIX: &str = r#"4. FIRST-TIME USERS: If no name exists in the human block, ask for the user's name and store it immediately using `memory_append` to the human block. +5. You may optionally set `reply_reaction` to a single emoji that reacts to the user's CURRENT message, but only on the first step before any tool results exist. This reaction is applied automatically to the user's latest message. Use `""` when you don't want to react. + +TOOL CALL PATTERNS: +- To react and respond: messages: ["msg1", "msg2"], reply_reaction: "❀️", tool_calls: [] +- To respond AND use tools: messages: ["msg1", "msg2"], reply_reaction: "", tool_calls: [your_tools] +- To respond with NO tools: messages: ["msg1", "msg2"], reply_reaction: "", tool_calls: [] +- After tool results with nothing to add: messages: [], reply_reaction: "", tool_calls: [{"name": "done", "args": {}}] + +AFTER TOOL RESULTS - CRITICAL RULES: +When you see "[Tool Result: X]", decide what to do next: + +- Never set `reply_reaction` after tool results. Once any tool result is present, `reply_reaction` must be `""`. + +- **web_search/archival_search/conversation_search**: Summarize findings in messages + +- **memory_append/memory_replace/archival_insert/memory_insert**: These operations complete without user-facing messages. Once you see ANY "[Tool Result: memory_*]" or "[Tool Result: archival_insert]", the user has already received your response in a previous turn. Immediately return: + messages: [] + reply_reaction: "" + tool_calls: [{"name": "done", "args": {}}] + + This applies even if you called multiple memory tools together (like memory_append + archival_insert for life events). Once ANY memory tool result appears, immediately call done. + + Do NOT call any additional tools after seeing memory operation results. + Do NOT send messages about the memory operation. + Do NOT explain what you stored. + Just return done immediately. + +The "done" tool means "nothing more to do" - use it ONLY when: +- messages is empty AND +- reply_reaction is empty AND +- no other tools are needed + +OUTPUT FORMAT: +You have exactly 3 output fields: +- messages: ALL messages in ONE array (e.g., ["msg1", "msg2", "msg3"]) +- reply_reaction: a single emoji string, or "" when unused +- tool_calls: ALL tool calls in ONE array + +CRITICAL FORMAT RULES: +- Do NOT repeat field tags. Wrong: multiple [[ ## messages ## ]] blocks. Right: one messages array with all items +- Do NOT include field delimiter tags INSIDE your content blocks +- Each [[ ## field ## ]] marker MUST be on its own line - nothing else on that line (no tags, no text before or after) +- Keep your output clean and strictly follow the field delimiters"#; + +const AGENT_INSTRUCTION_WITHOUT_REACTION_SUFFIX: &str = r#"4. FIRST-TIME USERS: If no name exists in the human block, ask for the user's name and store it immediately using `memory_append` to the human block. +5. This is a continuation step after tool results. Reactions are not supported here, so do not emit any reaction field. TOOL CALL PATTERNS: - To respond AND use tools: messages: ["msg1", "msg2"], tool_calls: [your_tools] @@ -138,6 +207,8 @@ TOOL CALL PATTERNS: AFTER TOOL RESULTS - CRITICAL RULES: When you see "[Tool Result: X]", decide what to do next: +- Do not emit any reaction field on continuation steps. + - **web_search/archival_search/conversation_search**: Summarize findings in messages - **memory_append/memory_replace/archival_insert/memory_insert**: These operations complete without user-facing messages. Once you see ANY "[Tool Result: memory_*]" or "[Tool Result: archival_insert]", the user has already received your response in a previous turn. Immediately return: @@ -156,7 +227,7 @@ The "done" tool means "nothing more to do" - use it ONLY when: - no other tools are needed OUTPUT FORMAT: -You have exactly 2 output fields. Put ALL content in that single field: +You have exactly 2 output fields: - messages: ALL messages in ONE array (e.g., ["msg1", "msg2", "msg3"]) - tool_calls: ALL tool calls in ONE array @@ -166,6 +237,17 @@ CRITICAL FORMAT RULES: - Each [[ ## field ## ]] marker MUST be on its own line - nothing else on that line (no tags, no text before or after) - Keep your output clean and strictly follow the field delimiters"#; +/// Default instruction for the agent (GEPA-optimized for Maple) +pub fn agent_instruction(allow_reply_reaction: bool) -> String { + let mut instruction = AGENT_INSTRUCTION_PREFIX.to_string(); + instruction.push_str(if allow_reply_reaction { + AGENT_INSTRUCTION_WITH_REACTION_SUFFIX + } else { + AGENT_INSTRUCTION_WITHOUT_REACTION_SUFFIX + }); + instruction +} + #[dspy_rs::BamlType] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct AgentToolCall { @@ -199,6 +281,39 @@ pub struct AgentResponse { #[input] pub is_first_time_user: bool, + #[output] + pub messages: Vec, + #[output] + pub reply_reaction: String, + #[output] + pub tool_calls: Vec, +} + +#[derive(dspy_rs::Signature, Debug, Clone)] +pub struct AgentContinuationResponse { + #[input] + pub input: String, + #[input] + pub current_time: String, + #[input] + pub agent_kind: String, + #[input] + pub subagent_purpose: String, + #[input] + pub persona_block: String, + #[input] + pub human_block: String, + #[input] + pub memory_metadata: String, + #[input] + pub previous_context_summary: String, + #[input] + pub recent_conversation: String, + #[input] + pub available_tools: String, + #[input] + pub is_first_time_user: bool, + #[output] pub messages: Vec, #[output] @@ -226,6 +341,30 @@ pub struct CorrectionResponse { #[output(desc = "Array of messages extracted/fixed from the original response")] pub messages: Vec, + #[output(desc = "Single emoji reaction to the current user message, or empty string")] + pub reply_reaction: String, + + #[output(desc = "Array of tool calls extracted/fixed from the original response")] + pub tool_calls: Vec, +} + +#[derive(dspy_rs::Signature, Debug, Clone)] +pub struct CorrectionContinuationResponse { + #[input(desc = "The original input that was given to the agent")] + pub original_input: String, + + #[input(desc = "The malformed response that needs to be corrected")] + pub malformed_response: String, + + #[input(desc = "The error message explaining what went wrong with parsing")] + pub error_message: String, + + #[input(desc = "Available tools for reference")] + pub available_tools: String, + + #[output(desc = "Array of messages extracted/fixed from the original response")] + pub messages: Vec, + #[output(desc = "Array of tool calls extracted/fixed from the original response")] pub tool_calls: Vec, } @@ -233,6 +372,7 @@ pub struct CorrectionResponse { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct AgentResponseOutput { pub messages: Vec, + pub reply_reaction: String, pub tool_calls: Vec, } @@ -245,18 +385,51 @@ pub enum SignatureCallError { }, } +fn continuation_input_from(input: &AgentResponseInput) -> AgentContinuationResponseInput { + AgentContinuationResponseInput { + input: input.input.clone(), + current_time: input.current_time.clone(), + agent_kind: input.agent_kind.clone(), + subagent_purpose: input.subagent_purpose.clone(), + persona_block: input.persona_block.clone(), + human_block: input.human_block.clone(), + memory_metadata: input.memory_metadata.clone(), + previous_context_summary: input.previous_context_summary.clone(), + recent_conversation: input.recent_conversation.clone(), + available_tools: input.available_tools.clone(), + is_first_time_user: input.is_first_time_user, + } +} + +fn normalize_messages(mut messages: Vec) -> Vec { + messages.retain(|message| !message.trim().is_empty()); + messages +} + +fn trace_generated_prompt(prompt_kind: &str, system_prompt: &str, user_prompt: &str) { + trace!( + prompt_kind = %prompt_kind, + system_prompt_len = system_prompt.chars().count(), + user_prompt_len = user_prompt.chars().count(), + system_prompt = system_prompt, + user_prompt = user_prompt, + "Generated LM prompt" + ); +} + pub async fn call_agent_response_with_retry_and_correction( lm: &Arc, system_prompt: &str, input: &AgentResponseInput, original_input: &str, available_tools: &str, + allow_reply_reaction: bool, ) -> Result { const MAX_LLM_RETRIES: u32 = 3; let mut last_err: Option = None; for attempt in 1..=MAX_LLM_RETRIES { - match call_agent_response_once(lm, system_prompt, input).await { + match call_agent_response_once(lm, system_prompt, input, allow_reply_reaction).await { Ok(out) => return Ok(out), Err(SignatureCallError::Parse { raw_response, @@ -269,6 +442,7 @@ pub async fn call_agent_response_with_retry_and_correction( available_tools, &raw_response, &error_message, + allow_reply_reaction, ) .await { @@ -295,6 +469,7 @@ async fn attempt_correction( available_tools: &str, raw_response: &str, error_message: &str, + allow_reply_reaction: bool, ) -> Result { if raw_response.trim().is_empty() { return Err(ApiError::InternalServerError); @@ -302,87 +477,172 @@ async fn attempt_correction( let adapter = ChatAdapter; - let system = adapter - .format_system_message_typed_with_instruction::(Some( - CORRECTION_INSTRUCTION, - )) - .map_err(|e| { - error!("Failed to format DSRS correction system prompt: {e:?}"); + if allow_reply_reaction { + let system = adapter + .format_system_message_typed_with_instruction::(Some( + CORRECTION_INSTRUCTION, + )) + .map_err(|e| { + error!("Failed to format DSRS correction system prompt: {e:?}"); + ApiError::InternalServerError + })?; + + let input = CorrectionResponseInput { + original_input: original_input.to_string(), + malformed_response: raw_response.to_string(), + error_message: error_message.to_string(), + available_tools: available_tools.to_string(), + }; + let user_msg = adapter.format_user_message_typed::(&input); + trace_generated_prompt("agent_response_correction", &system, &user_msg); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM correction call failed: {e:?}"); ApiError::InternalServerError })?; - let input = CorrectionResponseInput { - original_input: original_input.to_string(), - malformed_response: raw_response.to_string(), - error_message: error_message.to_string(), - available_tools: available_tools.to_string(), - }; - let user_msg = adapter.format_user_message_typed::(&input); - - let mut chat = dspy_rs::Chat::new(vec![]); - chat.push("system", &system); - chat.push("user", &user_msg); - - let response = lm.call(chat, Vec::new()).await.map_err(|e| { - error!("DSRS LM correction call failed: {e:?}"); - ApiError::InternalServerError - })?; - - let (output, _meta) = adapter - .parse_response_typed::(&response.output) - .map_err(|e| { - error!("DSRS correction typed parse failed: {e:?}"); + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS correction typed parse failed: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(AgentResponseOutput { + messages: output.messages, + reply_reaction: output.reply_reaction, + tool_calls: output.tool_calls, + }) + } else { + let system = adapter + .format_system_message_typed_with_instruction::(Some( + CONTINUATION_CORRECTION_INSTRUCTION, + )) + .map_err(|e| { + error!("Failed to format continuation correction system prompt: {e:?}"); + ApiError::InternalServerError + })?; + + let input = CorrectionContinuationResponseInput { + original_input: original_input.to_string(), + malformed_response: raw_response.to_string(), + error_message: error_message.to_string(), + available_tools: available_tools.to_string(), + }; + let user_msg = adapter.format_user_message_typed::(&input); + trace_generated_prompt("agent_continuation_correction", &system, &user_msg); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS continuation correction call failed: {e:?}"); ApiError::InternalServerError })?; - Ok(AgentResponseOutput { - messages: output.messages, - tool_calls: output.tool_calls, - }) + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS continuation correction typed parse failed: {e:?}"); + ApiError::InternalServerError + })?; + + Ok(AgentResponseOutput { + messages: output.messages, + reply_reaction: String::new(), + tool_calls: output.tool_calls, + }) + } } async fn call_agent_response_once( lm: &Arc, system_prompt: &str, input: &AgentResponseInput, + allow_reply_reaction: bool, ) -> Result { let adapter = ChatAdapter; - let system = adapter - .format_system_message_typed_with_instruction::(Some(system_prompt)) - .map_err(|e| { - error!("Failed to format DSRS system prompt: {e:?}"); + if allow_reply_reaction { + let system = adapter + .format_system_message_typed_with_instruction::(Some(system_prompt)) + .map_err(|e| { + error!("Failed to format DSRS system prompt: {e:?}"); + SignatureCallError::Api(ApiError::InternalServerError) + })?; + let user_msg = adapter.format_user_message_typed::(input); + trace_generated_prompt("agent_response", &system, &user_msg); + + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("DSRS LM call failed: {e:?}"); SignatureCallError::Api(ApiError::InternalServerError) })?; - let user_msg = adapter.format_user_message_typed::(input); - let mut chat = dspy_rs::Chat::new(vec![]); - chat.push("system", &system); - chat.push("user", &user_msg); + let raw_response = response.output.content(); + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("DSRS typed parse failed: {e:?}"); + SignatureCallError::Parse { + raw_response, + error_message: e.to_string(), + } + })?; - let response = lm.call(chat, Vec::new()).await.map_err(|e| { - error!("DSRS LM call failed: {e:?}"); - SignatureCallError::Api(ApiError::InternalServerError) - })?; + Ok(AgentResponseOutput { + messages: normalize_messages(output.messages), + reply_reaction: output.reply_reaction, + tool_calls: output.tool_calls, + }) + } else { + let continuation_input = continuation_input_from(input); + let system = adapter + .format_system_message_typed_with_instruction::(Some( + system_prompt, + )) + .map_err(|e| { + error!("Failed to format continuation DSRS system prompt: {e:?}"); + SignatureCallError::Api(ApiError::InternalServerError) + })?; + let user_msg = + adapter.format_user_message_typed::(&continuation_input); + trace_generated_prompt("agent_continuation_response", &system, &user_msg); - let raw_response = response.output.content(); - let (output, _meta) = adapter - .parse_response_typed::(&response.output) - .map_err(|e| { - error!("DSRS typed parse failed: {e:?}"); - SignatureCallError::Parse { - raw_response, - error_message: e.to_string(), - } + let mut chat = dspy_rs::Chat::new(vec![]); + chat.push("system", &system); + chat.push("user", &user_msg); + + let response = lm.call(chat, Vec::new()).await.map_err(|e| { + error!("Continuation DSRS LM call failed: {e:?}"); + SignatureCallError::Api(ApiError::InternalServerError) })?; - let mut messages = output.messages; - messages.retain(|m| !m.trim().is_empty()); + let raw_response = response.output.content(); + let (output, _meta) = adapter + .parse_response_typed::(&response.output) + .map_err(|e| { + error!("Continuation DSRS typed parse failed: {e:?}"); + SignatureCallError::Parse { + raw_response, + error_message: e.to_string(), + } + })?; - Ok(AgentResponseOutput { - messages, - tool_calls: output.tool_calls, - }) + Ok(AgentResponseOutput { + messages: normalize_messages(output.messages), + reply_reaction: String::new(), + tool_calls: output.tool_calls, + }) + } } pub async fn build_lm( diff --git a/src/web/responses/conversions.rs b/src/web/responses/conversions.rs index 8fdfe6c9..69be5229 100644 --- a/src/web/responses/conversions.rs +++ b/src/web/responses/conversions.rs @@ -190,6 +190,8 @@ pub enum ConversationItem { role: String, content: Vec, #[serde(skip_serializing_if = "Option::is_none")] + reaction: Option, + #[serde(skip_serializing_if = "Option::is_none")] created_at: Option, }, #[serde(rename = "function_call")] @@ -292,6 +294,7 @@ impl ConversationItemConverter { status: msg.status.clone(), role: ROLE_USER.to_string(), content: Vec::::from(message_content), + reaction: msg.reaction.clone(), created_at: Some(msg.created_at.timestamp()), }) } @@ -314,6 +317,7 @@ impl ConversationItemConverter { status: msg.status.clone(), role: ROLE_ASSISTANT.to_string(), content: content_parts, + reaction: msg.reaction.clone(), created_at: Some(msg.created_at.timestamp()), }) } diff --git a/src/web/responses/handlers.rs b/src/web/responses/handlers.rs index 66511869..0d11fcdd 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -1902,6 +1902,7 @@ async fn persist_request_data( content_enc: prepared.content_enc.clone(), attachment_text_enc: None, prompt_tokens: prepared.user_message_tokens, + assistant_reaction: None, }; let user_message = state .db diff --git a/src/web/responses/storage.rs b/src/web/responses/storage.rs index 3314a1d9..599d7a61 100644 --- a/src/web/responses/storage.rs +++ b/src/web/responses/storage.rs @@ -73,6 +73,7 @@ fn create_assistant_message_if_missing( status: STATUS_IN_PROGRESS.to_string(), finish_reason: None, created_at, + user_reaction: None, }) .map(|_| ()) .map_err(|e| format!("Failed to create assistant message: {:?}", e)),