From 9d989a4e590715d80dc120f9e0129d9e962b4eea Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 10 Feb 2026 09:49:47 -0600 Subject: [PATCH 01/34] feat: add phase 1 brute-force RAG infrastructure Enables encrypted user embedding storage + cosine search with experimental /v1/rag endpoints for local/dev testing. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../down.sql | 11 + .../2026-02-10-020612_user_embeddings/up.sql | 83 ++ src/main.rs | 9 +- src/models/mod.rs | 1 + src/models/schema.rs | 24 + src/models/user_embeddings.rs | 83 ++ src/rag.rs | 880 ++++++++++++++++++ src/web/mod.rs | 3 + src/web/openai.rs | 175 +++- src/web/rag.rs | 209 +++++ 10 files changed, 1426 insertions(+), 52 deletions(-) create mode 100644 migrations/2026-02-10-020612_user_embeddings/down.sql create mode 100644 migrations/2026-02-10-020612_user_embeddings/up.sql create mode 100644 src/models/user_embeddings.rs create mode 100644 src/rag.rs create mode 100644 src/web/rag.rs diff --git a/migrations/2026-02-10-020612_user_embeddings/down.sql b/migrations/2026-02-10-020612_user_embeddings/down.sql new file mode 100644 index 00000000..167f3308 --- /dev/null +++ b/migrations/2026-02-10-020612_user_embeddings/down.sql @@ -0,0 +1,11 @@ +DROP TRIGGER IF EXISTS update_user_embeddings_updated_at ON user_embeddings; + +DROP INDEX IF EXISTS idx_user_embeddings_user_id; +DROP INDEX IF EXISTS idx_user_embeddings_user_created; +DROP INDEX IF EXISTS idx_user_embeddings_user_source; +DROP INDEX IF EXISTS idx_user_embeddings_user_conversation; +DROP INDEX IF EXISTS idx_user_embeddings_user_message_id; +DROP INDEX IF EXISTS idx_user_embeddings_assistant_message_id; +DROP INDEX IF EXISTS idx_user_embeddings_model; + +DROP TABLE IF EXISTS user_embeddings; diff --git a/migrations/2026-02-10-020612_user_embeddings/up.sql b/migrations/2026-02-10-020612_user_embeddings/up.sql new file mode 100644 index 00000000..54071675 --- /dev/null +++ b/migrations/2026-02-10-020612_user_embeddings/up.sql @@ -0,0 +1,83 @@ +-- =========================================================== +-- User Embeddings (Phase 1 RAG brute-force infrastructure) +-- =========================================================== + +CREATE TABLE user_embeddings ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + -- Source tracking + -- v1 uses: 'message', 'archival'. Keep as open TEXT for future types. + source_type TEXT NOT NULL DEFAULT 'message', + + -- Message provenance (ONLY for source_type='message') + user_message_id BIGINT REFERENCES user_messages(id) ON DELETE CASCADE, + assistant_message_id BIGINT REFERENCES assistant_messages(id) ON DELETE CASCADE, + conversation_id BIGINT REFERENCES conversations(id) ON DELETE CASCADE, + + -- Embedding vector + vector_enc BYTEA NOT NULL, -- AES-256-GCM encrypted float32 array + embedding_model TEXT NOT NULL, -- e.g. "nomic-embed-text" + vector_dim INTEGER NOT NULL DEFAULT 768, + + -- Content that was embedded + content_enc BYTEA NOT NULL, + metadata_enc BYTEA, + + -- Plaintext metadata + token_count INTEGER NOT NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- Enforce invariants for the v1 source types without blocking future types. + CONSTRAINT user_embeddings_message_source_check CHECK ( + (source_type <> 'message') + OR ( + -- exactly one of the message FKs must be set + (user_message_id IS NOT NULL) <> (assistant_message_id IS NOT NULL) + ) + ), + CONSTRAINT user_embeddings_message_conversation_check CHECK ( + (source_type <> 'message') OR (conversation_id IS NOT NULL) + ), + CONSTRAINT user_embeddings_archival_source_check CHECK ( + (source_type <> 'archival') + OR ( + user_message_id IS NULL + AND assistant_message_id IS NULL + AND conversation_id IS NULL + ) + ) +); + +-- Primary query path: load all vectors for a user (brute-force scan) +CREATE INDEX idx_user_embeddings_user_id ON user_embeddings(user_id); + +-- For time-filtered searches (recency bias) +CREATE INDEX idx_user_embeddings_user_created ON user_embeddings(user_id, created_at DESC); + +-- For source-type filtered searches (message vs archival) +CREATE INDEX idx_user_embeddings_user_source ON user_embeddings(user_id, source_type); + +-- For conversation-scoped searches (message recall) +CREATE INDEX idx_user_embeddings_user_conversation ON user_embeddings(user_id, conversation_id); + +-- Idempotency/deduplication: prevent double-indexing the same message +CREATE UNIQUE INDEX idx_user_embeddings_user_message_id + ON user_embeddings(user_message_id) + WHERE user_message_id IS NOT NULL; + +CREATE UNIQUE INDEX idx_user_embeddings_assistant_message_id + ON user_embeddings(assistant_message_id) + WHERE assistant_message_id IS NOT NULL; + +-- For staleness queries: find vectors that need re-embedding after model change +CREATE INDEX idx_user_embeddings_model ON user_embeddings(user_id, embedding_model); + +-- Trigger for updated_at +-- Note: update_updated_at_column() function already exists from previous migrations +CREATE TRIGGER update_user_embeddings_updated_at +BEFORE UPDATE ON user_embeddings +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/main.rs b/src/main.rs index 3063fa53..c69a9e3b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,7 @@ use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ conversations_routes, document_routes, health_routes_with_state, instructions_routes, - login_routes, oauth_routes, openai_routes, protected_routes, responses_routes, + login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -85,6 +85,7 @@ mod oauth; mod os_flags; mod private_key; mod proxy_config; +mod rag; mod sqs; mod tokens; mod web; @@ -408,6 +409,7 @@ pub struct AppState { aws_credential_manager: Arc>>, enclave_key: Vec, proxy_router: Arc, + rag_cache: Arc>, resend_api_key: Option, ephemeral_keys: Arc>>, session_states: Arc>>, @@ -726,6 +728,7 @@ impl AppStateBuilder { aws_credential_manager, enclave_key, proxy_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())), @@ -2786,6 +2789,10 @@ async fn main() -> Result<(), Error> { document_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 39416e19..bb76de90 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -17,5 +17,6 @@ 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 users; diff --git a/src/models/schema.rs b/src/models/schema.rs index fffdd9a7..626df6cb 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -307,6 +307,26 @@ diesel::table! { } } +diesel::table! { + user_embeddings (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + source_type -> Text, + user_message_id -> Nullable, + assistant_message_id -> Nullable, + conversation_id -> Nullable, + vector_enc -> Bytea, + embedding_model -> Text, + vector_dim -> Int4, + content_enc -> Bytea, + metadata_enc -> Nullable, + token_count -> Int4, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { user_instructions (id) { id -> Int8, @@ -391,6 +411,9 @@ diesel::joinable!(tool_calls -> responses (response_id)); diesel::joinable!(tool_outputs -> conversations (conversation_id)); diesel::joinable!(tool_outputs -> responses (response_id)); diesel::joinable!(tool_outputs -> tool_calls (tool_call_fk)); +diesel::joinable!(user_embeddings -> assistant_messages (assistant_message_id)); +diesel::joinable!(user_embeddings -> conversations (conversation_id)); +diesel::joinable!(user_embeddings -> user_messages (user_message_id)); diesel::joinable!(user_messages -> conversations (conversation_id)); diesel::joinable!(user_messages -> responses (response_id)); diesel::joinable!(user_oauth_connections -> oauth_providers (provider_id)); @@ -420,6 +443,7 @@ diesel::allow_tables_to_appear_in_same_query!( tool_calls, tool_outputs, user_api_keys, + user_embeddings, user_instructions, user_kv, user_messages, diff --git a/src/models/user_embeddings.rs b/src/models/user_embeddings.rs new file mode 100644 index 00000000..fb3bbac3 --- /dev/null +++ b/src/models/user_embeddings.rs @@ -0,0 +1,83 @@ +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 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 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..d8b4d09a --- /dev/null +++ b/src/rag.rs @@ -0,0 +1,880 @@ +use crate::encrypt::{decrypt_with_key, encrypt_with_key}; +use crate::models::schema::user_embeddings; +use crate::models::user_embeddings::NewUserEmbedding; +use crate::{ApiError, AppState}; +use diesel::prelude::*; +use secp256k1::SecretKey; +use serde::Serialize; +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, 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: &AppState, text: &str) -> Result<(Vec, i32), ApiError> { + crate::web::get_embedding_vector( + state, + DEFAULT_EMBEDDING_MODEL, + text, + Some(DEFAULT_EMBEDDING_DIM), + ) + .await +} + +pub async fn insert_archival_embedding( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + text: &str, + metadata: Option<&serde_json::Value>, +) -> Result { + let (vector, token_count) = embed_text_via_tinfoil(state, 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 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, + 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) +} + +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)) +} + +#[allow(clippy::too_many_arguments)] +pub async fn search_user_embeddings( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + query: &str, + top_k: usize, + max_tokens: Option, + source_types: Option<&[String]>, + conversation_id: Option, +) -> Result, ApiError> { + let top_k = top_k.clamp(1, 20); + + let (query_vec, _query_tokens) = embed_text_via_tinfoil(state, query).await?; + + let cached = { + let mut cache = state.rag_cache.lock().await; + cache.get(user_id) + }; + + let embeddings = 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 af0d75d5..8dfc6657 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -9,15 +9,18 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; +pub mod rag; pub mod responses; pub use documents::router as document_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; +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::conversations_router as conversations_routes; pub use responses::instructions_router as instructions_routes; pub use responses::responses_router as responses_routes; diff --git a/src/web/openai.rs b/src/web/openai.rs index 89c1e18a..7d275923 100644 --- a/src/web/openai.rs +++ b/src/web/openai.rs @@ -1558,55 +1558,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> { - debug!("Entering proxy_embeddings function"); - - // 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> { // Get the model route configuration let route = match state.proxy_router.get_model_route(&embedding_request.model) { Some(r) => r, @@ -1626,8 +1581,15 @@ async fn proxy_embeddings( .pool_max_idle_per_host(10) .build::<_, HyperBody>(https); + // Translate model name for the provider (if needed) + let provider_model_name = state + .proxy_router + .get_model_name_for_provider(&embedding_request.model, &route.primary.provider_name); + let mut provider_request = embedding_request.clone(); + provider_request.model = provider_model_name; + // 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 })?; @@ -1692,6 +1654,117 @@ async fn proxy_embeddings( ApiError::InternalServerError })?; + Ok((response_json, route.primary.provider_name)) +} + +pub async fn get_embedding_vector( + state: &AppState, + 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) = request_embeddings_from_primary(state, &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); + } + } + + let prompt_tokens = match response_json + .get("usage") + .and_then(|u| u.get("prompt_tokens")) + .and_then(|v| v.as_i64()) + { + Some(v) => v as i32, + None => { + error!("Embeddings response missing usage.prompt_tokens"); + 0 + } + }; + + 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"); + + // 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); + } + + let (response_json, provider_name) = + request_embeddings_from_primary(&state, &embedding_request).await?; + // Handle billing - embeddings only have prompt_tokens (no completion_tokens) if let Some(usage) = response_json.get("usage") { let prompt_tokens = usage @@ -1711,7 +1784,7 @@ async fn proxy_embeddings( &user, &billing_context, embedding_usage, - &route.primary.provider_name, + &provider_name, ) .await; } diff --git a/src/web/rag.rs b/src/web/rag.rs new file mode 100644 index 00000000..bf3949d5 --- /dev/null +++ b/src/web/rag.rs @@ -0,0 +1,209 @@ +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::models::users::User; +use crate::rag; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +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(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.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let inserted = rag::insert_archival_embedding( + &state, + user.uuid, + &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(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.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let results = rag::search_user_embeddings( + &state, + user.uuid, + &user_key, + &body.query, + top_k, + body.max_tokens, + body.source_types.as_deref(), + conversation_internal_id, + ) + .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 d95d27d676aa18e5bb608d7f9521df9e428d3c64 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 10 Feb 2026 16:02:21 -0600 Subject: [PATCH 02/34] feat: add sage mvp agent storage tables Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../down.sql | 10 ++ .../up.sql | 78 +++++++++++ src/models/agent_config.rs | 88 ++++++++++++ src/models/conversation_summaries.rs | 107 +++++++++++++++ src/models/memory_blocks.rs | 129 ++++++++++++++++++ src/models/mod.rs | 3 + src/models/schema.rs | 55 ++++++++ 7 files changed, 470 insertions(+) create mode 100644 migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql create mode 100644 migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql create mode 100644 src/models/agent_config.rs create mode 100644 src/models/conversation_summaries.rs create mode 100644 src/models/memory_blocks.rs diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql new file mode 100644 index 00000000..db111084 --- /dev/null +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql @@ -0,0 +1,10 @@ +DROP TRIGGER IF EXISTS update_agent_config_updated_at ON agent_config; +DROP TABLE IF EXISTS agent_config; + +DROP INDEX IF EXISTS idx_conversation_summaries_chain; +DROP INDEX IF EXISTS idx_conversation_summaries_user_conv; +DROP TABLE IF EXISTS conversation_summaries; + +DROP TRIGGER IF EXISTS update_memory_blocks_updated_at ON memory_blocks; +DROP INDEX IF EXISTS idx_memory_blocks_user_id; +DROP TABLE IF EXISTS memory_blocks; diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql new file mode 100644 index 00000000..a39836b7 --- /dev/null +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql @@ -0,0 +1,78 @@ +-- Sage MVP agent storage tables +-- +-- Note: update_updated_at_column() already exists from previous migrations. + +CREATE TABLE memory_blocks ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + label TEXT NOT NULL, + description TEXT, + value_enc BYTEA NOT NULL, + char_limit INTEGER NOT NULL DEFAULT 5000, + read_only BOOLEAN NOT NULL DEFAULT FALSE, + version INTEGER NOT NULL DEFAULT 1, + + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + + UNIQUE(user_id, label) +); + +CREATE INDEX idx_memory_blocks_user_id ON memory_blocks(user_id); + +CREATE TRIGGER update_memory_blocks_updated_at +BEFORE UPDATE ON memory_blocks +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); + +CREATE TABLE conversation_summaries ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + + from_created_at TIMESTAMPTZ NOT NULL, + to_created_at TIMESTAMPTZ NOT NULL, + message_count INTEGER NOT NULL, + + content_enc BYTEA NOT NULL, + content_tokens INTEGER NOT NULL, + + embedding_enc BYTEA, + + previous_summary_id BIGINT REFERENCES conversation_summaries(id) ON DELETE SET NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT valid_time_range CHECK (from_created_at <= to_created_at) +); + +CREATE INDEX idx_conversation_summaries_user_conv + ON conversation_summaries(user_id, conversation_id, created_at DESC); + +CREATE INDEX idx_conversation_summaries_chain + ON conversation_summaries(previous_summary_id); + +CREATE TABLE agent_config ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE UNIQUE, + + conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, + + enabled BOOLEAN NOT NULL DEFAULT FALSE, + model TEXT NOT NULL DEFAULT 'deepseek-r1-0528', + max_context_tokens INTEGER NOT NULL DEFAULT 100000, + compaction_threshold REAL NOT NULL DEFAULT 0.80, + + system_prompt_enc BYTEA, + preferences_enc BYTEA, + + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TRIGGER update_agent_config_updated_at +BEFORE UPDATE ON agent_config +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/models/agent_config.rs b/src/models/agent_config.rs new file mode 100644 index 00000000..633beb10 --- /dev/null +++ b/src/models/agent_config.rs @@ -0,0 +1,88 @@ +use crate::models::schema::agent_config; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +#[derive(Error, Debug)] +pub enum AgentConfigError { + #[error("Database error: {0}")] + DatabaseError(#[from] diesel::result::Error), +} + +#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] +#[diesel(table_name = agent_config)] +pub struct AgentConfig { + pub id: i64, + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: Option, + pub enabled: bool, + pub model: String, + pub max_context_tokens: i32, + pub compaction_threshold: f32, + pub system_prompt_enc: Option>, + pub preferences_enc: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl AgentConfig { + pub fn get_by_user_id( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result, AgentConfigError> { + agent_config::table + .filter(agent_config::user_id.eq(lookup_user_id)) + .first::(conn) + .optional() + .map_err(AgentConfigError::DatabaseError) + } + + pub fn delete_for_user( + conn: &mut PgConnection, + lookup_user_id: Uuid, + ) -> Result { + diesel::delete(agent_config::table.filter(agent_config::user_id.eq(lookup_user_id))) + .execute(conn) + .map_err(AgentConfigError::DatabaseError) + } +} + +#[derive(Insertable, Debug, Clone)] +#[diesel(table_name = agent_config)] +pub struct NewAgentConfig { + pub uuid: Uuid, + pub user_id: Uuid, + pub conversation_id: Option, + pub enabled: bool, + pub model: String, + pub max_context_tokens: i32, + pub compaction_threshold: f32, + pub system_prompt_enc: Option>, + pub preferences_enc: Option>, +} + +impl NewAgentConfig { + pub fn insert_or_update( + &self, + conn: &mut PgConnection, + ) -> Result { + diesel::insert_into(agent_config::table) + .values(self) + .on_conflict(agent_config::user_id) + .do_update() + .set(( + agent_config::conversation_id.eq(self.conversation_id), + agent_config::enabled.eq(self.enabled), + agent_config::model.eq(self.model.clone()), + agent_config::max_context_tokens.eq(self.max_context_tokens), + agent_config::compaction_threshold.eq(self.compaction_threshold), + agent_config::system_prompt_enc.eq(self.system_prompt_enc.clone()), + agent_config::preferences_enc.eq(self.preferences_enc.clone()), + )) + .get_result::(conn) + .map_err(AgentConfigError::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..e9128520 --- /dev/null +++ b/src/models/memory_blocks.rs @@ -0,0 +1,129 @@ +use crate::models::schema::memory_blocks; +use chrono::{DateTime, Utc}; +use diesel::prelude::*; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +#[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 description: Option, + pub value_enc: Vec, + pub char_limit: i32, + pub read_only: bool, + pub version: i32, + 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 description: Option, + pub value_enc: Vec, + pub char_limit: i32, + pub read_only: bool, + pub version: i32, +} + +impl NewMemoryBlock { + pub fn new( + user_id: Uuid, + label: impl Into, + description: Option, + value_enc: Vec, + ) -> Self { + NewMemoryBlock { + uuid: Uuid::new_v4(), + user_id, + label: label.into(), + description, + value_enc, + char_limit: 5000, + read_only: false, + version: 1, + } + } + + 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::description.eq(self.description.clone()), + memory_blocks::value_enc.eq(self.value_enc.clone()), + memory_blocks::char_limit.eq(self.char_limit), + memory_blocks::read_only.eq(self.read_only), + memory_blocks::version.eq(memory_blocks::version + 1), + )) + .get_result::(conn) + .map_err(MemoryBlockError::DatabaseError) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index bb76de90..5127021c 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,7 +1,10 @@ pub mod account_deletion; +pub mod agent_config; +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/schema.rs b/src/models/schema.rs index 626df6cb..2df1656f 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -21,6 +21,23 @@ diesel::table! { } } +diesel::table! { + agent_config (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + conversation_id -> Nullable, + enabled -> Bool, + model -> Text, + max_context_tokens -> Int4, + compaction_threshold -> Float4, + system_prompt_enc -> Nullable, + preferences_enc -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { assistant_messages (id) { id -> Int8, @@ -37,6 +54,23 @@ diesel::table! { } } +diesel::table! { + conversation_summaries (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + conversation_id -> Int8, + from_created_at -> Timestamptz, + to_created_at -> Timestamptz, + message_count -> Int4, + content_enc -> Bytea, + content_tokens -> Int4, + embedding_enc -> Nullable, + previous_summary_id -> Nullable, + created_at -> Timestamptz, + } +} + diesel::table! { conversations (id) { id -> Int8, @@ -82,6 +116,22 @@ diesel::table! { } } +diesel::table! { + memory_blocks (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + label -> Text, + description -> Nullable, + value_enc -> Bytea, + char_limit -> Int4, + read_only -> Bool, + version -> Int4, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { oauth_providers (id) { id -> Int4, @@ -395,8 +445,10 @@ diesel::table! { } } +diesel::joinable!(agent_config -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> responses (response_id)); +diesel::joinable!(conversation_summaries -> conversations (conversation_id)); diesel::joinable!(invite_codes -> orgs (org_id)); diesel::joinable!(org_memberships -> orgs (org_id)); diesel::joinable!(org_project_secrets -> org_projects (project_id)); @@ -421,11 +473,14 @@ diesel::joinable!(users -> org_projects (project_id)); diesel::allow_tables_to_appear_in_same_query!( account_deletion_requests, + agent_config, assistant_messages, + conversation_summaries, conversations, email_verifications, enclave_secrets, invite_codes, + memory_blocks, oauth_providers, org_memberships, org_project_secrets, From 57d73252c5ebc8166c7d8d8fd05fa62a06c02a4a Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 10 Feb 2026 21:17:38 -0600 Subject: [PATCH 03/34] feat: add DSRS-backed agent chat endpoint Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 3076 ++++++++++++++++++++++++++++++++--- Cargo.toml | 1 + src/main.rs | 9 +- src/web/agent/mod.rs | 131 ++ src/web/agent/signatures.rs | 319 ++++ src/web/mod.rs | 2 + 6 files changed, 3329 insertions(+), 209 deletions(-) create mode 100644 src/web/agent/mod.rs create mode 100644 src/web/agent/signatures.rs diff --git a/Cargo.lock b/Cargo.lock index c8e7775b..56fa1f80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[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" @@ -68,6 +74,20 @@ 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 0.8.39", +] + [[package]] name = "aho-corasick" version = "1.1.3" @@ -78,10 +98,25 @@ dependencies = [ ] [[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" @@ -92,12 +127,71 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "0.6.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" + +[[package]] +name = "anstyle-parse" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" +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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +[[package]] +name = "arc-swap" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ded5f9a03ac8f24d1b8a25101ee812cd32cdc8c50a4c50237de2c4915850e73" +dependencies = [ + "rustversion", +] + [[package]] name = "argon2" version = "0.5.3" @@ -110,12 +204,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.4.1", + "hashbrown 0.16.0", + "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.4.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.4.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.4.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.4.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.4.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" @@ -128,7 +442,7 @@ dependencies = [ "nom", "num-traits", "rusticata-macros", - "thiserror", + "thiserror 1.0.63", "time", ] @@ -155,6 +469,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" @@ -177,6 +503,12 @@ dependencies = [ "syn 2.0.106", ] +[[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" @@ -188,6 +520,15 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -196,9 +537,9 @@ 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" @@ -385,7 +726,7 @@ dependencies = [ "hex", "hmac", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "once_cell", "percent-encoding", "sha2", @@ -455,17 +796,17 @@ dependencies = [ "aws-smithy-types", "bytes", "fastrand", - "h2", + "h2 0.3.26", "http 0.2.12", "http-body 0.4.6", "http-body 1.0.1", "httparse", "hyper 0.14.30", - "hyper-rustls", + "hyper-rustls 0.24.2", "once_cell", "pin-project-lite", "pin-utils", - "rustls", + "rustls 0.21.12", "tokio", "tracing", ] @@ -480,7 +821,7 @@ dependencies = [ "aws-smithy-types", "bytes", "http 0.2.12", - "http 1.1.0", + "http 1.4.0", "pin-project-lite", "tokio", "tracing", @@ -498,7 +839,7 @@ dependencies = [ "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", @@ -547,7 +888,7 @@ dependencies = [ "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", @@ -581,7 +922,7 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "http-body-util", "mime", @@ -612,10 +953,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b62ddb9cb1ec0a098ad4bbf9344d0713fa193ae1a80af55febcff2627b6a00c1" dependencies = [ "futures-core", - "getrandom", + "getrandom 0.2.15", "instant", "pin-project-lite", - "rand", + "rand 0.8.5", "tokio", ] @@ -629,19 +970,78 @@ dependencies = [ "cc", "cfg-if", "libc", - "miniz_oxide", + "miniz_oxide 0.7.4", "object", "rustc-demangle", ] +[[package]] +name = "baml-ids" +version = "0.0.1" +dependencies = [ + "anyhow", + "getrandom 0.2.15", + "serde", + "time", + "type-safe-id", + "uuid", +] + +[[package]] +name = "baml-types" +version = "0.1.0" +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" +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" +dependencies = [ + "convert_case", + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.106", +] + [[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]] @@ -684,7 +1084,7 @@ version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "36915bbaca237c626689b5bd14d02f2ba7a5a359d30a2a08be697392e3718079" dependencies = [ - "thiserror", + "thiserror 1.0.63", ] [[package]] @@ -707,13 +1107,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", ] @@ -754,22 +1163,16 @@ checksum = "ce6bc65742dea50536e35ad42492b234c27904a27f0abdcbce605015cb4ea026" 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" @@ -791,20 +1194,10 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5285c8bcaa25876d07f37e3d30c303f2609179716e11d688f51e8f1fe70063e2" 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" @@ -812,7 +1205,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb18c03d0db0247e147a21a6faafd5a7eb851c743db062de72018b6b7e8e4d16" dependencies = [ "bitcoin-io", - "hex-conservative 0.2.1", + "hex-conservative", "serde", ] @@ -824,9 +1217,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.6.0" +version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" [[package]] name = "blake2" @@ -856,39 +1249,95 @@ dependencies = [ ] [[package]] -name = "bstr" -version = "1.12.0" +name = "bon" +version = "3.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "234113d19d0d7d613b40e86fb654acf958910802bcceab913a4f9e7cda03b1a4" +checksum = "234655ec178edd82b891e262ea7cf71f6584bcd09eff94db786be23f1821825c" dependencies = [ - "memchr", - "regex-automata 0.4.11", - "serde", + "bon-macros", + "rustversion", ] [[package]] -name = "bumpalo" -version = "3.16.0" +name = "bon-macros" +version = "3.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" +checksum = "89ec27229c38ed0eb3c0feee3d2c1d6a4379ae44f418a29a658890e062d8f365" +dependencies = [ + "darling 0.20.10", + "ident_case", + "prettyplease", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.106", +] [[package]] -name = "byteorder" -version = "1.5.0" +name = "brotli" +version = "8.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] [[package]] -name = "bytes" -version = "1.10.1" +name = "brotli-decompressor" +version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] [[package]] -name = "bytes-utils" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" +name = "bstd" +version = "0.1.0" +dependencies = [ + "anyhow", + "num", + "rand 0.8.5", + "regex", +] + +[[package]] +name = "bstr" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234113d19d0d7d613b40e86fb654acf958910802bcceab913a4f9e7cda03b1a4" +dependencies = [ + "memchr", + "regex-automata", + "serde", +] + +[[package]] +name = "bumpalo" +version = "3.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" + +[[package]] +name = "bytes-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" dependencies = [ "bytes", "either", @@ -908,6 +1357,10 @@ name = "cc" version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e9e8aabfac534be767c909e0690571677d49f41bd8465ae876fe043d52ba5292" +dependencies = [ + "jobserver", + "libc", +] [[package]] name = "cfg-if" @@ -947,17 +1400,16 @@ dependencies = [ [[package]] name = "chrono" -version = "0.4.38" +version = "0.4.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118" dependencies = [ - "android-tzdata", "iana-time-zone", "js-sys", "num-traits", "serde", "wasm-bindgen", - "windows-targets 0.52.6", + "windows-link 0.2.1", ] [[package]] @@ -998,6 +1450,46 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clap" +version = "4.5.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6899ea499e3fb9305a65d5ebf6e3d2248c5fab291f300ad0a704fbe142eae31a" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b12c8b680195a62a8364d16b8447b01b6c2c8f9aaf68bee653be34d4245e238" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim 0.11.1", +] + +[[package]] +name = "clap_derive" +version = "4.5.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a92793da1a46a5f2a02a6f4c46c6496b28c43638adea8306fcb0caa1634f24e5" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "syn 2.0.106", +] + +[[package]] +name = "clap_lex" +version = "0.7.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e64b0cc0439b12df2fa678eae89a1c56a529fd067a9115f7827f1fffd22b32" + [[package]] name = "cmac" version = "0.7.2" @@ -1009,6 +1501,85 @@ dependencies = [ "digest", ] +[[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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" + +[[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.48.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.15", + "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" @@ -1025,6 +1596,17 @@ 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" @@ -1034,6 +1616,34 @@ 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" @@ -1053,10 +1663,31 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" 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" @@ -1092,14 +1723,38 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "darling" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b750cb3417fd1b327431a470f388520309479ab0bf5e323505daf0290cd3850" +dependencies = [ + "darling_core 0.14.4", + "darling_macro 0.14.4", +] + [[package]] name = "darling" version = "0.20.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f63b86c8a8826a49b8c21f08a2d07338eec8d900540f8630dc76284be802989" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.20.10", + "darling_macro 0.20.10", +] + +[[package]] +name = "darling_core" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "109c1ca6e6b7f82cc233a97004ea8ed7ca123a9af07a8230878fcfda9b158bf0" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim 0.10.0", + "syn 1.0.109", ] [[package]] @@ -1112,17 +1767,28 @@ dependencies = [ "ident_case", "proc-macro2", "quote", - "strsim", + "strsim 0.11.1", "syn 2.0.106", ] +[[package]] +name = "darling_macro" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4aab4dbc9f7611d8b55048a3a16d2d010c2c8334e46304b40ac1cc14bf3b48e" +dependencies = [ + "darling_core 0.14.4", + "quote", + "syn 1.0.109", +] + [[package]] name = "darling_macro" version = "0.20.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ - "darling_core", + "darling_core 0.20.10", "quote", "syn 2.0.106", ] @@ -1185,7 +1851,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf97ee7261bb708fa3402fa9c17a54b70e90e3cb98afb3dc8999d5512cb03f94" dependencies = [ "bigdecimal", - "bitflags 2.6.0", + "bitflags 2.10.0", "byteorder", "chrono", "diesel_derives", @@ -1244,6 +1910,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.59.0", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -1261,13 +1948,19 @@ 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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c5d9abe6314103864cc2d8901b7ae224e0ab1a103a0a416661b4097b0779b607" dependencies = [ - "darling", + "darling 0.20.10", "either", "heck 0.5.0", "proc-macro2", @@ -1275,6 +1968,58 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "dspy-rs" +version = "0.7.3" +dependencies = [ + "anyhow", + "arrow", + "async-trait", + "bamltype", + "bon", + "csv", + "dsrs_macros", + "enum_dispatch", + "facet", + "foyer", + "futures", + "hf-hub", + "indexmap", + "kdam", + "parquet", + "rand 0.8.5", + "rayon", + "regex", + "reqwest 0.12.23", + "rig-core", + "rstest", + "schemars", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "dsrs_macros" +version = "0.7.2" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.106", +] + +[[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" @@ -1290,6 +2035,12 @@ version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "encoding_rs" version = "0.8.34" @@ -1299,11 +2050,44 @@ dependencies = [ "cfg-if", ] +[[package]] +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.106", +] + +[[package]] +name = "env_filter" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bf3c259d255ca70051b30e2e95b5446cdb8949ac4cd22c0d7fd634d89f568e2" +dependencies = [ + "log", +] + +[[package]] +name = "env_logger" +version = "0.11.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c863f0904021b108aa8b2f55046443e6b1ebde8fd4a15c399893aae4fa069f" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "log", +] + [[package]] name = "equivalent" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" @@ -1315,6 +2099,114 @@ dependencies = [ "windows-sys 0.59.0", ] +[[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" @@ -1325,6 +2217,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.1.1" @@ -1337,12 +2239,57 @@ version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags 2.10.0", + "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 0.8.9", + "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" @@ -1367,6 +2314,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.10.0", + "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.2", + "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 = "futures" version = "0.3.31" @@ -1485,6 +2549,31 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", + "wasip3", +] + [[package]] name = "ghash" version = "0.5.1" @@ -1501,6 +2590,12 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd" +[[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" @@ -1516,7 +2611,7 @@ dependencies = [ "parking_lot", "portable-atomic", "quanta", - "rand", + "rand 0.8.5", "smallvec", "spinning_top", ] @@ -1540,6 +2635,25 @@ dependencies = [ "tracing", ] +[[package]] +name = "h2" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http 1.4.0", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "half" version = "1.8.3" @@ -1554,6 +2668,7 @@ checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" dependencies = [ "cfg-if", "crunchy", + "num-traits", ] [[package]] @@ -1562,11 +2677,25 @@ version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +[[package]] +name = "hashbrown" +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.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" +dependencies = [ + "allocator-api2", +] [[package]] name = "heck" @@ -1582,9 +2711,15 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.3.9" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" + +[[package]] +name = "hermit-abi" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" [[package]] name = "hex" @@ -1592,19 +2727,13 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" 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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5313b072ce3c597065a808dbf612c4c8e8590bdbf8b579508bf7a762c5eae6cd" dependencies = [ - "arrayvec", + "arrayvec 0.7.6", ] [[package]] @@ -1613,6 +2742,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.2", + "reqwest 0.12.23", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "ureq", + "windows-sys 0.60.2", +] + [[package]] name = "hmac" version = "0.12.1" @@ -1635,12 +2788,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", ] @@ -1662,7 +2814,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http 1.1.0", + "http 1.4.0", ] [[package]] @@ -1673,7 +2825,7 @@ checksum = "793429d76616a256bcb62c2a2ec2bed781c8307e797e2598c50010f2bee2544f" dependencies = [ "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "pin-project-lite", ] @@ -1700,7 +2852,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2", + "h2 0.3.26", "http 0.2.12", "http-body 0.4.6", "httparse", @@ -1724,7 +2876,8 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "http 1.1.0", + "h2 0.4.13", + "http 1.4.0", "http-body 1.0.1", "httparse", "httpdate", @@ -1746,10 +2899,26 @@ dependencies = [ "http 0.2.12", "hyper 0.14.30", "log", - "rustls", + "rustls 0.21.12", "rustls-native-certs", "tokio", - "tokio-rustls", + "tokio-rustls 0.24.1", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" +dependencies = [ + "http 1.4.0", + "hyper 1.7.0", + "hyper-util", + "rustls 0.23.14", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.0", + "tower-service", ] [[package]] @@ -1792,7 +2961,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "hyper 1.7.0", "ipnet", @@ -1800,9 +2969,11 @@ dependencies = [ "percent-encoding", "pin-project-lite", "socket2 0.6.0", + "system-configuration 0.6.1", "tokio", "tower-service", "tracing", + "windows-registry", ] [[package]] @@ -1946,6 +3117,25 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +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.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b215e67ed1d1a4b1702acd787c487d16e4c977c5dcbcc4587bdb5ea26b6ce06" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "hashbrown 0.16.0", + "rustc-hash 2.1.1", +] + [[package]] name = "ident_case" version = "1.0.1" @@ -1983,6 +3173,12 @@ dependencies = [ "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" @@ -1991,6 +3187,21 @@ checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" dependencies = [ "equivalent", "hashbrown 0.16.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]] @@ -2012,13 +3223,44 @@ dependencies = [ "cfg-if", ] +[[package]] +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" +dependencies = [ + "anyhow", + "colored", + "pest", + "serde", + "strsim 0.11.1", +] + +[[package]] +name = "internal-baml-jinja" +version = "0.1.0" +dependencies = [ + "anyhow", + "baml-types", + "indexmap", + "minijinja", + "serde_json", + "serde_yaml", + "strum", + "toon", +] + [[package]] name = "io-uring" version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "046fa2d4d00aea763528b4950358d0ead425372445dc8ff86312b3c69ff7727b" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "cfg-if", "libc", ] @@ -2039,12 +3281,37 @@ dependencies = [ "serde", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + [[package]] name = "js-sys" version = "0.3.81" @@ -2055,6 +3322,28 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonish" +version = "0.1.0" +dependencies = [ + "anyhow", + "baml-types", + "bstd", + "getrandom 0.2.15", + "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 = "9.3.0" @@ -2082,7 +3371,7 @@ dependencies = [ "ciborium", "hmac", "lazy_static", - "rand_core", + "rand_core 0.6.4", "secp256k1", "serde", "serde_json", @@ -2092,6 +3381,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.5" @@ -2107,6 +3406,69 @@ 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" @@ -2119,11 +3481,21 @@ version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" +[[package]] +name = "libredox" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616" +dependencies = [ + "bitflags 2.10.0", + "libc", +] + [[package]] name = "linux-raw-sys" -version = "0.4.14" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" [[package]] name = "litemap" @@ -2147,13 +3519,96 @@ version = "0.4.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +[[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.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08ab2867e3eeeca90e844d1940eab391c9dc5228783db2ed999acbc0a9ed375a" +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.5", + "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 = "7d3eb2acc57c82d21d699119b859e2df70a91dbdb84734885a1e72be83bdecb5" +dependencies = [ + "madsim", + "spin", + "tokio", +] + [[package]] name = "matchers" -version = "0.1.0" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" dependencies = [ - "regex-automata 0.1.10", + "regex-automata", ] [[package]] @@ -2203,6 +3658,29 @@ 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.11.0" +source = "git+https://github.com/boundaryml/minijinja.git?branch=main#ee25c021c8bb78d7ca11fa99972918eea4585330" +dependencies = [ + "aho-corasick", + "indexmap", + "serde", + "serde_json", + "unicase", + "unicode-ident", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -2218,18 +3696,38 @@ dependencies = [ "adler", ] +[[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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "80e04d1dcff3aae0704555fe5fee3bcfaf3d1fdf8a7e521d5b9d2b42acb52cec" dependencies = [ - "hermit-abi", + "hermit-abi 0.3.9", "libc", "wasi", "windows-sys 0.52.0", ] +[[package]] +name = "mixtrics" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb252c728b9d77c6ef9103f0c81524fa0a3d3b161d0a936295d7fbeff6e04c11" +dependencies = [ + "itertools", + "parking_lot", +] + [[package]] name = "multer" version = "3.1.0" @@ -2239,7 +3737,7 @@ dependencies = [ "bytes", "encoding_rs", "futures-util", - "http 1.1.0", + "http 1.4.0", "httparse", "memchr", "mime", @@ -2247,6 +3745,27 @@ 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.15", +] + [[package]] name = "native-tls" version = "0.2.12" @@ -2283,7 +3802,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "cfg-if", "cfg_aliases", "libc", @@ -2314,12 +3833,25 @@ checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" [[package]] name = "nu-ansi-term" -version = "0.46.0" +version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "overload", - "winapi", + "windows-sys 0.59.0", +] + +[[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]] @@ -2332,6 +3864,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.1.0" @@ -2347,6 +3888,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" @@ -2354,8 +3917,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 0.5.2", + "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" @@ -2364,15 +3944,15 @@ checksum = "c38841cdd844847e3e7c8d29cef9dcfed8877f8f56f9071f77843ecf3baf937f" dependencies = [ "base64 0.13.1", "chrono", - "getrandom", + "getrandom 0.2.15", "http 0.2.12", - "rand", + "rand 0.8.5", "reqwest 0.11.27", "serde", "serde_json", "serde_path_to_error", "sha2", - "thiserror", + "thiserror 1.0.63", "url", ] @@ -2400,6 +3980,12 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "opaque-debug" version = "0.3.1" @@ -2434,9 +4020,10 @@ dependencies = [ "diesel", "diesel-derive-enum", "dotenv", + "dspy-rs", "futures", "generic-array", - "getrandom", + "getrandom 0.2.15", "hex", "hmac", "hyper 0.14.30", @@ -2447,7 +4034,7 @@ dependencies = [ "oauth2", "once_cell", "password-auth", - "rand_core", + "rand_core 0.6.4", "rcgen", "regex", "reqwest 0.11.27", @@ -2459,7 +4046,7 @@ dependencies = [ "serde_json", "sha2", "subtle", - "thiserror", + "thiserror 1.0.63", "tiktoken-rs", "tokio", "tokio-stream", @@ -2481,7 +4068,7 @@ version = "0.10.66" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9529f4786b70a3e8c61e11179af17ab6188ad8d0ded78c5529441ed39d4bd9c1" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "cfg-if", "foreign-types", "libc", @@ -2519,6 +4106,30 @@ 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.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f4779c6901a562440c3786d08192c6fbda7c1c2060edd10006b05ee35d10f2d" +dependencies = [ + "num-traits", +] + [[package]] name = "outref" version = "0.5.1" @@ -2526,10 +4137,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a" [[package]] -name = "overload" -version = "0.1.1" +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 = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" [[package]] name = "parking_lot" @@ -2554,6 +4171,39 @@ dependencies = [ "windows-targets 0.52.6", ] +[[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.4.1", + "hashbrown 0.16.0", + "lz4_flex", + "num", + "num-bigint", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "twox-hash", + "zstd", +] + [[package]] name = "password-auth" version = "1.0.0" @@ -2561,9 +4211,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a2a4764cc1f8d961d802af27193c6f4f0124bd0e76e8393cf818e18880f0524" dependencies = [ "argon2", - "getrandom", + "getrandom 0.2.15", "password-hash", - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -2573,10 +4223,16 @@ 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" @@ -2593,6 +4249,16 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[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" @@ -2615,9 +4281,9 @@ dependencies = [ [[package]] name = "pin-project-lite" -version = "0.2.14" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" [[package]] name = "pin-utils" @@ -2672,7 +4338,7 @@ version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" dependencies = [ - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -2684,6 +4350,36 @@ 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.106", +] + +[[package]] +name = "proc-macro-crate" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro-error-attr2" version = "2.0.0" @@ -2739,6 +4435,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r2d2" version = "0.8.10" @@ -2757,8 +4459,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -2768,7 +4480,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]] @@ -2777,7 +4499,25 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.15", +] + +[[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]] @@ -2786,7 +4526,27 @@ version = "11.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb9ee317cfe3fbd54b36a511efc1edd42e216903c9cd575e686dd68a2ba90d8d" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", +] + +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", ] [[package]] @@ -2808,28 +4568,50 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a908a6e00f1fdd0dfd9c0eb08ce85126f6d8bbda50017e74bc4a4b7d4a926a4" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", ] [[package]] -name = "regex" -version = "1.11.3" +name = "redox_users" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b5288124840bee7b386bc413c487869b360b2b4ec421ea56425128692f2a82c" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ - "aho-corasick", - "memchr", - "regex-automata 0.4.11", - "regex-syntax 0.8.6", + "getrandom 0.2.15", + "libredox", + "thiserror 2.0.18", ] [[package]] -name = "regex-automata" -version = "0.1.10" +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 = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.106", +] + +[[package]] +name = "regex" +version = "1.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +checksum = "8b5288124840bee7b386bc413c487869b360b2b4ec421ea56425128692f2a82c" dependencies = [ - "regex-syntax 0.6.29", + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", ] [[package]] @@ -2840,7 +4622,7 @@ checksum = "833eb9ce86d40ef33cb1306d8accf7bc8ec2bfea4355cbdebb3df68b40925cad" dependencies = [ "aho-corasick", "memchr", - "regex-syntax 0.8.6", + "regex-syntax", ] [[package]] @@ -2851,15 +4633,15 @@ checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a" [[package]] name = "regex-syntax" -version = "0.6.29" +version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" +checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" [[package]] -name = "regex-syntax" -version = "0.8.6" +name = "relative-path" +version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" [[package]] name = "reqwest" @@ -2872,11 +4654,11 @@ dependencies = [ "encoding_rs", "futures-core", "futures-util", - "h2", + "h2 0.3.26", "http 0.2.12", "http-body 0.4.6", "hyper 0.14.30", - "hyper-rustls", + "hyper-rustls 0.24.2", "hyper-tls 0.5.0", "ipnet", "js-sys", @@ -2886,22 +4668,22 @@ dependencies = [ "once_cell", "percent-encoding", "pin-project-lite", - "rustls", + "rustls 0.21.12", "rustls-pemfile", "serde", "serde_json", "serde_urlencoded", "sync_wrapper 0.1.2", - "system-configuration", + "system-configuration 0.5.1", "tokio", "tokio-native-tls", - "tokio-rustls", + "tokio-rustls 0.24.1", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "webpki-roots", + "webpki-roots 0.25.4", "winreg", ] @@ -2913,15 +4695,22 @@ checksum = "d429f34c8092b2d42c7c93cec323bb4adeb7c67698f70839adec842ec10c7ceb" 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-rustls 0.27.7", "hyper-tls 0.6.0", "hyper-util", "js-sys", "log", + "mime", + "mime_guess", "native-tls", "percent-encoding", "pin-project-lite", @@ -2932,12 +4721,14 @@ dependencies = [ "sync_wrapper 1.0.1", "tokio", "tokio-native-tls", + "tokio-util", "tower 0.5.2", "tower-http 0.6.6", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", ] @@ -2950,10 +4741,38 @@ dependencies = [ "ecow", "governor", "maybe-async", - "rand", + "rand 0.8.5", "reqwest 0.12.23", "serde", - "thiserror", + "thiserror 1.0.63", +] + +[[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.1.0", + "pin-project-lite", + "reqwest 0.12.23", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", + "tracing-futures", + "url", ] [[package]] @@ -2964,13 +4783,43 @@ checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" dependencies = [ "cc", "cfg-if", - "getrandom", + "getrandom 0.2.15", "libc", "spin", "untrusted", "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.106", + "unicode-ident", +] + [[package]] name = "rustc-demangle" version = "0.1.24" @@ -2983,6 +4832,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustc_version" version = "0.4.1" @@ -3003,15 +4858,15 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.34" +version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70dc5ec042f7a43c4a73241207cecc9873a06d45debb38b329f8541d85c2730f" +checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3022,10 +4877,25 @@ checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" dependencies = [ "log", "ring", - "rustls-webpki", + "rustls-webpki 0.101.7", "sct", ] +[[package]] +name = "rustls" +version = "0.23.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "415d9944693cb90382053259f89fbb077ea730ad7273047ec63b19bc9b160ba8" +dependencies = [ + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki 0.102.8", + "subtle", + "zeroize", +] + [[package]] name = "rustls-native-certs" version = "0.6.3" @@ -3063,6 +4933,17 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustls-webpki" +version = "0.102.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -3090,7 +4971,32 @@ version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" dependencies = [ - "parking_lot", + "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.106", ] [[package]] @@ -3115,8 +5021,8 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e0cc0f1cf93f4969faf3ea1c7d8a9faed25918d96affa959720823dfe86d4f3" dependencies = [ - "bitcoin_hashes 0.14.0", - "rand", + "bitcoin_hashes", + "rand 0.8.5", "secp256k1-sys", "serde", ] @@ -3130,13 +5036,23 @@ dependencies = [ "cc", ] +[[package]] +name = "secrecy" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a" +dependencies = [ + "serde", + "zeroize", +] + [[package]] name = "security-framework" version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "core-foundation", "core-foundation-sys", "libc", @@ -3159,6 +5075,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b" +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -3208,12 +5130,24 @@ dependencies = [ "syn 2.0.106", ] +[[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.106", +] + [[package]] name = "serde_json" version = "1.0.145" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" dependencies = [ + "indexmap", "itoa", "memchr", "ryu", @@ -3231,6 +5165,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8bbf91e5a4d6315eee45e704372590b30e260ee83af6639d64557f51b067776" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -3243,11 +5186,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", @@ -3282,6 +5238,18 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + +[[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" @@ -3290,7 +5258,7 @@ checksum = "adc4e5204eb1910f40f9cfa375f6f05b68c3abac4b6fd879c8ff5e7ae8a0a085" dependencies = [ "num-bigint", "num-traits", - "thiserror", + "thiserror 1.0.63", "time", ] @@ -3303,12 +5271,24 @@ dependencies = [ "autocfg", ] +[[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.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +[[package]] +name = "snap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" + [[package]] name = "socket2" version = "0.5.7" @@ -3329,11 +5309,25 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + [[package]] 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" @@ -3350,12 +5344,40 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "strsim" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" + [[package]] name = "strsim" 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.106", +] + [[package]] name = "subtle" version = "2.6.1" @@ -3430,7 +5452,18 @@ checksum = "ba3a3adc5c275d719af8cb4272ea1c4a6d668a777f37e115f6d11ddbc1c8e0e7" dependencies = [ "bitflags 1.3.2", "core-foundation", - "system-configuration-sys", + "system-configuration-sys 0.5.0", +] + +[[package]] +name = "system-configuration" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c879d448e9d986b661742763247d3693ed13609438cf3d006f51f5368a5ba6b" +dependencies = [ + "bitflags 2.10.0", + "core-foundation", + "system-configuration-sys 0.6.0", ] [[package]] @@ -3443,26 +5476,77 @@ 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.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbcdd0c794ebb0d4cf35e88edd2f7d2c4c3e9a5a6dab322839b321c6a87a64" +checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" dependencies = [ - "cfg-if", "fastrand", + "getrandom 0.4.1", "once_cell", "rustix", "windows-sys 0.59.0", ] +[[package]] +name = "terminal_size" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b8cb979cb11c32ce1603f8137b22262a9d131aaa5c37b5678025f22b8becd0" +dependencies = [ + "rustix", + "windows-sys 0.60.2", +] + +[[package]] +name = "test-log" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37d53ac171c92a39e4769491c4b4dde7022c60042254b5fc044ae409d34a24d4" +dependencies = [ + "env_logger", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-macros" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be35209fd0781c5401458ab66e4f98accf63553e8fae7425503e92fdd319783b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.106", +] + [[package]] name = "thiserror" version = "1.0.63" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.63", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", ] [[package]] @@ -3476,6 +5560,17 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.106", +] + [[package]] name = "thread_local" version = "1.1.8" @@ -3486,6 +5581,17 @@ dependencies = [ "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]] name = "tiktoken-rs" version = "0.5.9" @@ -3498,7 +5604,7 @@ dependencies = [ "fancy-regex", "lazy_static", "parking_lot", - "rustc-hash", + "rustc-hash 1.1.0", ] [[package]] @@ -3532,6 +5638,15 @@ 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" @@ -3604,7 +5719,18 @@ version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" dependencies = [ - "rustls", + "rustls 0.21.12", + "tokio", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4" +dependencies = [ + "rustls 0.23.14", + "rustls-pki-types", "tokio", ] @@ -3632,6 +5758,67 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + +[[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_edit" +version = "0.23.10+spec-1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.0.7+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "247eaa3197818b831697600aadf81514e577e0cba5eab10f7e064e78ae154df1" +dependencies = [ + "winnow", +] + +[[package]] +name = "toml_writer" +version = "1.0.6+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" + +[[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.4.13" @@ -3669,9 +5856,9 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "bytes", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "http-body-util", "pin-project-lite", @@ -3685,10 +5872,10 @@ version = "0.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" dependencies = [ - "bitflags 2.6.0", + "bitflags 2.10.0", "bytes", "futures-util", - "http 1.1.0", + "http 1.4.0", "http-body 1.0.1", "iri-string", "pin-project-lite", @@ -3711,9 +5898,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", @@ -3723,9 +5910,9 @@ 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", @@ -3734,14 +5921,26 @@ dependencies = [ [[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" @@ -3755,14 +5954,14 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.18" +version = "0.3.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e" dependencies = [ "matchers", "nu-ansi-term", "once_cell", - "regex", + "regex-automata", "sharded-slab", "smallvec", "thread_local", @@ -3777,12 +5976,51 @@ 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.2", +] + +[[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.5", + "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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +[[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 = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + [[package]] name = "unicode-bidi" version = "0.3.15" @@ -3797,27 +6035,56 @@ checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" [[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.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[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" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" [[package]] -name = "universal-hash" -version = "0.5.1" +name = "unsynn" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +checksum = "501a7adf1a4bd9951501e5c66621e972ef8874d787628b7f90e64f936ef7ec0a" dependencies = [ - "crypto-common", - "subtle", + "mutants", + "proc-macro2", + "rustc-hash 2.1.1", ] [[package]] @@ -3826,6 +6093,26 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74fc6b57825be3373f7054754755f03ac3a8f5d70015ccad699ba2029956f4a" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "native-tls", + "once_cell", + "rustls 0.23.14", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + [[package]] name = "url" version = "2.5.2" @@ -3856,14 +6143,21 @@ 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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" dependencies = [ - "getrandom", + "getrandom 0.2.15", "serde", + "wasm-bindgen", ] [[package]] @@ -3888,7 +6182,7 @@ version = "0.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7df16e474ef958526d1205f6dda359fdfab79d9aa6d54bafcb92dcd07673dca" dependencies = [ - "darling", + "darling 0.20.10", "once_cell", "proc-macro-error2", "proc-macro2", @@ -3945,6 +6239,24 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wasip2" +version = "1.0.2+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +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 = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.104" @@ -4017,6 +6329,53 @@ 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.10.0", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + [[package]] name = "web-sys" version = "0.3.81" @@ -4027,12 +6386,40 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-roots" version = "0.25.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1" +[[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.6", +] + +[[package]] +name = "webpki-roots" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -4064,6 +6451,47 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-registry" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b8a9ed28765efc97bbc954883f4e6796c33a06546ebafacbabee9696967499e" +dependencies = [ + "windows-link 0.1.3", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", +] + +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows-sys" version = "0.48.0" @@ -4091,6 +6519,24 @@ 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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link 0.2.1", +] + [[package]] name = "windows-targets" version = "0.48.5" @@ -4115,13 +6561,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 0.2.1", + "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" @@ -4134,6 +6597,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" @@ -4146,6 +6615,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" @@ -4158,12 +6633,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" @@ -4176,6 +6663,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" @@ -4188,6 +6681,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" @@ -4200,6 +6699,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" @@ -4212,6 +6717,21 @@ 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.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" +dependencies = [ + "memchr", +] + [[package]] name = "winreg" version = "0.50.0" @@ -4222,6 +6742,94 @@ 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-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.106", + "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.106", + "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.10.0", + "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 = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + [[package]] name = "write16" version = "1.0.0" @@ -4241,7 +6849,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", ] @@ -4259,7 +6867,7 @@ dependencies = [ "nom", "oid-registry", "rusticata-macros", - "thiserror", + "thiserror 1.0.63", "time", ] @@ -4309,7 +6917,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" dependencies = [ "byteorder", - "zerocopy-derive", + "zerocopy-derive 0.7.35", +] + +[[package]] +name = "zerocopy" +version = "0.8.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db6d35d663eadb6c932438e763b262fe1a70987f9ae936e60158176d710cae4a" +dependencies = [ + "zerocopy-derive 0.8.39", ] [[package]] @@ -4323,6 +6940,17 @@ dependencies = [ "syn 2.0.106", ] +[[package]] +name = "zerocopy-derive" +version = "0.8.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.106", +] + [[package]] name = "zerofrom" version = "0.1.5" @@ -4385,3 +7013,37 @@ dependencies = [ "quote", "syn 2.0.106", ] + +[[package]] +name = "zlib-rs" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7948af682ccbc3342b6e9420e8c51c1fe5d7bf7756002b4a3c6cabfe96a7e3c" + +[[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 714f347a..ea111016 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ thiserror = "1.0.63" async-trait = "0.1.81" jsonwebtoken = "9.3.0" jwt-compact = { version = "0.9.0-beta.1", features = ["es256k"] } +dspy-rs = { path = "../../ThirdParties/DSRs/crates/dspy-rs" } diesel = { version = "=2.2.2", features = [ "postgres", "postgres_backend", diff --git a/src/main.rs b/src/main.rs index c69a9e3b..4c7e2473 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,8 +18,9 @@ use crate::sqs::SqsEventPublisher; use crate::web::openai_auth::validate_openai_auth; use crate::web::platform_login_routes; use crate::web::{ - conversations_routes, document_routes, health_routes_with_state, instructions_routes, - login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, responses_routes, + agent_routes, conversations_routes, document_routes, health_routes_with_state, + instructions_routes, login_routes, oauth_routes, openai_routes, protected_routes, rag_routes, + responses_routes, }; use crate::{attestation_routes::SessionState, web::platform_routes}; @@ -2793,6 +2794,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/web/agent/mod.rs b/src/web/agent/mod.rs new file mode 100644 index 00000000..c513819b --- /dev/null +++ b/src/web/agent/mod.rs @@ -0,0 +1,131 @@ +use axum::{ + extract::State, middleware::from_fn_with_state, routing::post, Extension, Json, Router, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use tracing::error; +use uuid::Uuid; + +use crate::encrypt::decrypt_string; +use crate::models::agent_config::AgentConfig; +use crate::models::memory_blocks::MemoryBlock; +use crate::models::users::User; +use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; +use crate::web::responses::error_mapping; +use crate::{ApiError, AppMode, AppState}; + +mod signatures; + +#[derive(Debug, Clone, Deserialize)] +struct AgentChatRequest { + input: String, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentChatResponse { + messages: Vec, + tool_calls: 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/agent/chat", + post(chat).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .with_state(app_state) +} + +async fn chat( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if body.input.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let user_key = state + .get_user_key(user.uuid, 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_config = AgentConfig::get_by_user_id(&mut conn, user.uuid).map_err(|e| { + error!("Failed to load agent config: {e:?}"); + ApiError::InternalServerError + })?; + + let persona = + MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, "persona").map_err(|e| { + error!("Failed to load persona memory block: {e:?}"); + ApiError::InternalServerError + })?; + let human = MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, "human").map_err(|e| { + error!("Failed to load human memory block: {e:?}"); + ApiError::InternalServerError + })?; + + let persona_text = decrypt_string(&user_key, persona.as_ref().map(|b| &b.value_enc)) + .map_err(|e| { + error!("Failed to decrypt persona memory block: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_default(); + + let human_text = decrypt_string(&user_key, human.as_ref().map(|b| &b.value_enc)) + .map_err(|e| { + error!("Failed to decrypt human memory block: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_default(); + + let is_first_time_user = persona.is_none() && human.is_none(); + + let input = signatures::AgentResponseInput { + input: body.input, + current_time: Utc::now().to_rfc3339(), + persona_block: persona_text, + human_block: human_text, + memory_metadata: String::new(), + previous_context_summary: String::new(), + recent_conversation: String::new(), + available_tools: "- done: call when you're finished".to_string(), + is_first_time_user, + }; + + let model = agent_config + .as_ref() + .map(|c| c.model.as_str()) + .unwrap_or("deepseek-r1-0528"); + let output = signatures::run_agent_response_signature( + &state, + &user, + model, + signatures::DEFAULT_AGENT_SYSTEM_PROMPT, + &input, + ) + .await?; + + let response = AgentChatResponse { + messages: output.messages, + tool_calls: output.tool_calls, + }; + + encrypt_response(&state, &session_id, &response).await +} diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs new file mode 100644 index 00000000..73b46b6c --- /dev/null +++ b/src/web/agent/signatures.rs @@ -0,0 +1,319 @@ +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}; + +pub const DEFAULT_AGENT_SYSTEM_PROMPT: &str = r#"You are Maple, a helpful AI assistant. + +Return `tool_calls` as an empty JSON array when no tools are needed."#; + +#[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 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, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AgentResponseOutput { + pub messages: Vec, + pub tool_calls: Vec, +} + +pub async fn run_agent_response_signature( + state: &Arc, + user: &crate::models::users::User, + model: &str, + system_prompt: &str, + input: &AgentResponseInput, +) -> Result { + let lm = build_lm(state.clone(), Arc::new(user.clone()), model.to_string()).await?; + 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:?}"); + 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 + })?; + + let mut messages = output.messages; + messages.retain(|m| !m.trim().is_empty()); + + Ok(AgentResponseOutput { + messages, + tool_calls: output.tool_calls, + }) +} + +async fn build_lm( + state: Arc, + user: Arc, + model: String, +) -> 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(0.7) + .max_tokens(32768) + .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/mod.rs b/src/web/mod.rs index 8dfc6657..feb5e90b 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 documents; @@ -12,6 +13,7 @@ pub mod protected_routes; pub mod rag; pub mod responses; +pub use agent::router as agent_routes; pub use documents::router as document_routes; pub use health_routes::router_with_state as health_routes_with_state; pub use login_routes::router as login_routes; From 2d825cff26e91b67c0824912848f9297a2221d36 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 09:46:20 -0600 Subject: [PATCH 04/34] feat: implement sage-style agent runtime and api surface Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/rag.rs | 51 ++ src/web/agent/compaction.rs | 247 +++++++ src/web/agent/mod.rs | 859 +++++++++++++++++++++-- src/web/agent/runtime.rs | 1048 ++++++++++++++++++++++++++++ src/web/agent/signatures.rs | 303 +++++++- src/web/agent/tools.rs | 824 ++++++++++++++++++++++ src/web/responses/conversations.rs | 100 ++- 7 files changed, 3362 insertions(+), 70 deletions(-) create mode 100644 src/web/agent/compaction.rs create mode 100644 src/web/agent/runtime.rs create mode 100644 src/web/agent/tools.rs diff --git a/src/rag.rs b/src/rag.rs index d8b4d09a..118c2287 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -345,6 +345,57 @@ pub async fn insert_archival_embedding( Ok(inserted) } +#[allow(clippy::too_many_arguments)] +pub async fn insert_message_embedding( + state: &AppState, + user_id: Uuid, + user_key: &SecretKey, + text: &str, + conversation_id: i64, + user_message_id: Option, + assistant_message_id: Option, +) -> Result { + let text = text.trim(); + if text.is_empty() { + return Err(ApiError::BadRequest); + } + + let (vector, token_count) = embed_text_via_tinfoil(state, 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, + 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, 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 index c513819b..acfe2c85 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,21 +1,38 @@ use axum::{ - extract::State, middleware::from_fn_with_state, routing::post, Extension, Json, Router, + extract::{Path, Query, State}, + middleware::from_fn_with_state, + routing::{delete, get, post, put}, + Extension, Json, Router, }; -use chrono::Utc; +use diesel::prelude::*; use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; use std::sync::Arc; use tracing::error; use uuid::Uuid; -use crate::encrypt::decrypt_string; -use crate::models::agent_config::AgentConfig; -use crate::models::memory_blocks::MemoryBlock; +use crate::encrypt::{decrypt_content, decrypt_string, encrypt_with_key}; +use crate::models::agent_config::{AgentConfig, NewAgentConfig}; +use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; +use crate::models::responses::{Conversation, NewConversation}; +use crate::models::schema::user_embeddings; use crate::models::users::User; +use crate::rag; use crate::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; -use crate::web::responses::error_mapping; +use crate::web::responses::constants::{DEFAULT_PAGINATION_LIMIT, MAX_PAGINATION_LIMIT}; +use crate::web::responses::conversations::{ + ConversationItemListResponse, ConversationListResponse, ConversationResponse, ListItemsParams, +}; +use crate::web::responses::{ + error_mapping, ConversationBuilder, ConversationItem, ConversationItemConverter, + DeletedObjectResponse, Paginator, +}; use crate::{ApiError, AppMode, AppState}; +mod compaction; +mod runtime; mod signatures; +mod tools; #[derive(Debug, Clone, Deserialize)] struct AgentChatRequest { @@ -28,6 +45,71 @@ struct AgentChatResponse { tool_calls: Vec, } +#[derive(Debug, Clone, Serialize)] +struct AgentConfigResponse { + enabled: bool, + model: String, + max_context_tokens: i32, + compaction_threshold: f32, + system_prompt: Option, + conversation_id: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct UpdateAgentConfigRequest { + enabled: Option, + model: Option, + max_context_tokens: Option, + compaction_threshold: Option, + system_prompt: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct MemoryBlockResponse { + label: String, + description: Option, + value: String, + char_limit: i32, + read_only: bool, + version: i32, +} + +#[derive(Debug, Clone, Deserialize)] +struct UpdateMemoryBlockRequest { + description: Option, + value: Option, + char_limit: Option, + read_only: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct InsertArchivalRequest { + text: String, + metadata: Option, +} + +#[derive(Debug, Clone, Serialize)] +struct InsertArchivalResponse { + id: Uuid, + source_type: String, + embedding_model: String, + token_count: i32, + created_at: chrono::DateTime, +} + +#[derive(Debug, Clone, Deserialize)] +struct MemorySearchRequest { + query: String, + top_k: Option, + max_tokens: Option, + source_types: Option>, +} + +#[derive(Debug, Clone, Serialize)] +struct MemorySearchResponse { + 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) { @@ -42,6 +124,68 @@ pub fn router(app_state: Arc) -> Router<()> { decrypt_request::, )), ) + .route( + "/v1/agent/config", + get(get_config).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/config", + put(update_config).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/memory/blocks", + get(list_memory_blocks) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/memory/blocks/:label", + get(get_memory_block) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/memory/blocks/:label", + put(update_memory_block).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/memory/archival", + post(insert_archival).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/memory/archival/:id", + delete(delete_archival) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/memory/search", + post(memory_search).layer(from_fn_with_state( + app_state.clone(), + decrypt_request::, + )), + ) + .route( + "/v1/agent/conversations", + get(list_agent_conversations) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/conversations/:id/items", + get(list_agent_conversation_items) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) + .route( + "/v1/agent/conversations/:id", + delete(delete_agent_conversation) + .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), + ) .with_state(app_state) } @@ -55,6 +199,27 @@ async fn chat( return Err(ApiError::BadRequest); } + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut runtime = runtime::AgentRuntime::new(state.clone(), user.clone(), user_key).await?; + let (messages, tool_calls) = runtime.process_message(&body.input).await?; + + let response = AgentChatResponse { + messages, + tool_calls, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn get_config( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { let user_key = state .get_user_key(user.uuid, None, None) .await @@ -66,66 +231,668 @@ async fn chat( .get() .map_err(|_| ApiError::InternalServerError)?; - let agent_config = AgentConfig::get_by_user_id(&mut conn, user.uuid).map_err(|e| { - error!("Failed to load agent config: {e:?}"); + let cfg = get_or_create_agent_config(&mut conn, user.uuid).await?; + + let system_prompt = decrypt_string(&user_key, cfg.system_prompt_enc.as_ref()).map_err(|e| { + error!("Failed to decrypt agent system prompt: {e:?}"); ApiError::InternalServerError })?; - let persona = - MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, "persona").map_err(|e| { - error!("Failed to load persona memory block: {e:?}"); + let conversation_uuid = if let Some(conversation_id) = cfg.conversation_id { + match Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) { + Ok(c) => Some(c.uuid), + Err(_) => None, + } + } else { + None + }; + + let response = AgentConfigResponse { + enabled: cfg.enabled, + model: cfg.model, + max_context_tokens: cfg.max_context_tokens, + compaction_threshold: cfg.compaction_threshold, + system_prompt, + conversation_id: conversation_uuid, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn update_config( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let mut cfg = get_or_create_agent_config(&mut conn, user.uuid).await?; + + if let Some(enabled) = body.enabled { + cfg.enabled = enabled; + } + if let Some(model) = body.model { + if model.trim().is_empty() { + return Err(ApiError::BadRequest); + } + cfg.model = model; + } + if let Some(max_context_tokens) = body.max_context_tokens { + if max_context_tokens <= 0 { + return Err(ApiError::BadRequest); + } + cfg.max_context_tokens = max_context_tokens; + } + if let Some(threshold) = body.compaction_threshold { + cfg.compaction_threshold = threshold.clamp(0.5, 0.95); + } + + let system_prompt_enc = match body.system_prompt { + Some(p) if !p.trim().is_empty() => Some(encrypt_with_key(&user_key, p.as_bytes()).await), + Some(_) => None, + None => cfg.system_prompt_enc.clone(), + }; + + // If enabling and conversation is missing, initialize agent thread + blocks. + let conversation_id = if cfg.enabled && cfg.conversation_id.is_none() { + let metadata_enc = Some( + encrypt_with_key( + &user_key, + json!({"type":"agent_main"}).to_string().as_bytes(), + ) + .await, + ); + + let conversation = NewConversation { + uuid: Uuid::new_v4(), + user_id: user.uuid, + metadata_enc, + } + .insert(&mut conn) + .map_err(|e| { + error!("Failed to create agent conversation: {e:?}"); ApiError::InternalServerError })?; - let human = MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, "human").map_err(|e| { - error!("Failed to load human memory block: {e:?}"); + + runtime::AgentRuntime::ensure_default_blocks(&state, &mut conn, &user_key, user.uuid) + .await?; + + Some(conversation.id) + } else { + cfg.conversation_id + }; + + let updated = NewAgentConfig { + uuid: cfg.uuid, + user_id: cfg.user_id, + conversation_id, + enabled: cfg.enabled, + model: cfg.model.clone(), + max_context_tokens: cfg.max_context_tokens, + compaction_threshold: cfg.compaction_threshold, + system_prompt_enc, + preferences_enc: cfg.preferences_enc.clone(), + } + .insert_or_update(&mut conn) + .map_err(|e| { + error!("Failed to update agent config: {e:?}"); ApiError::InternalServerError })?; - let persona_text = decrypt_string(&user_key, persona.as_ref().map(|b| &b.value_enc)) - .map_err(|e| { - error!("Failed to decrypt persona memory block: {e:?}"); + let system_prompt = + decrypt_string(&user_key, updated.system_prompt_enc.as_ref()).map_err(|e| { + error!("Failed to decrypt agent system prompt: {e:?}"); + ApiError::InternalServerError + })?; + + let conversation_uuid = if let Some(conversation_id) = updated.conversation_id { + match Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) { + Ok(c) => Some(c.uuid), + Err(_) => None, + } + } else { + None + }; + + let response = AgentConfigResponse { + enabled: updated.enabled, + model: updated.model, + max_context_tokens: updated.max_context_tokens, + compaction_threshold: updated.compaction_threshold, + system_prompt, + conversation_id: conversation_uuid, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn list_memory_blocks( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>>, ApiError> { + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let blocks = MemoryBlock::get_all_for_user(&mut conn, user.uuid).map_err(|e| { + error!("Failed to list memory blocks: {e:?}"); + ApiError::InternalServerError + })?; + + let mut out: Vec = Vec::with_capacity(blocks.len()); + for b in blocks { + let value = decrypt_string(&user_key, Some(&b.value_enc)) + .map_err(|e| { + error!("Failed to decrypt memory block: {e:?}"); + ApiError::InternalServerError + })? + .unwrap_or_default(); + + out.push(MemoryBlockResponse { + label: b.label, + description: b.description, + value, + char_limit: b.char_limit, + read_only: b.read_only, + version: b.version, + }); + } + + encrypt_response(&state, &session_id, &out).await +} + +async fn get_memory_block( + State(state): State>, + Path(label): Path, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + if label.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let Some(block) = + MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, &label).map_err(|e| { + error!("Failed to load memory block: {e:?}"); ApiError::InternalServerError })? + else { + return Err(ApiError::NotFound); + }; + + let value = decrypt_string(&user_key, Some(&block.value_enc)) + .map_err(|_| ApiError::InternalServerError)? .unwrap_or_default(); - let human_text = decrypt_string(&user_key, human.as_ref().map(|b| &b.value_enc)) - .map_err(|e| { - error!("Failed to decrypt human memory block: {e:?}"); + let response = MemoryBlockResponse { + label: block.label, + description: block.description, + value, + char_limit: block.char_limit, + read_only: block.read_only, + version: block.version, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn update_memory_block( + State(state): State>, + Path(label): Path, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if label.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let existing = + MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, &label).map_err(|e| { + error!("Failed to load memory block: {e:?}"); ApiError::InternalServerError - })? + })?; + + if let Some(b) = &existing { + if b.read_only { + return Err(ApiError::Unauthorized); + } + } + + let value_enc = match body.value { + Some(value) => encrypt_with_key(&user_key, value.as_bytes()).await, + None => match &existing { + Some(b) => b.value_enc.clone(), + None => encrypt_with_key(&user_key, b"").await, + }, + }; + + let new_block = NewMemoryBlock { + uuid: existing + .as_ref() + .map(|b| b.uuid) + .unwrap_or_else(Uuid::new_v4), + user_id: user.uuid, + label: label.clone(), + description: body + .description + .or_else(|| existing.as_ref().and_then(|b| b.description.clone())), + value_enc, + char_limit: body + .char_limit + .or_else(|| existing.as_ref().map(|b| b.char_limit)) + .unwrap_or(5000), + read_only: body + .read_only + .or_else(|| existing.as_ref().map(|b| b.read_only)) + .unwrap_or(false), + version: 1, + }; + + let updated = new_block.insert_or_update(&mut conn).map_err(|e| { + error!("Failed to upsert memory block: {e:?}"); + ApiError::InternalServerError + })?; + + let decrypted_value = decrypt_string(&user_key, Some(&updated.value_enc)) + .map_err(|_| ApiError::InternalServerError)? .unwrap_or_default(); - let is_first_time_user = persona.is_none() && human.is_none(); - - let input = signatures::AgentResponseInput { - input: body.input, - current_time: Utc::now().to_rfc3339(), - persona_block: persona_text, - human_block: human_text, - memory_metadata: String::new(), - previous_context_summary: String::new(), - recent_conversation: String::new(), - available_tools: "- done: call when you're finished".to_string(), - is_first_time_user, - }; - - let model = agent_config - .as_ref() - .map(|c| c.model.as_str()) - .unwrap_or("deepseek-r1-0528"); - let output = signatures::run_agent_response_signature( + let response = MemoryBlockResponse { + label: updated.label, + description: updated.description, + value: decrypted_value, + char_limit: updated.char_limit, + read_only: updated.read_only, + version: updated.version, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn insert_archival( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + 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.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let inserted = rag::insert_archival_embedding( &state, - &user, - model, - signatures::DEFAULT_AGENT_SYSTEM_PROMPT, - &input, + user.uuid, + &user_key, + &body.text, + body.metadata.as_ref(), ) .await?; - let response = AgentChatResponse { - messages: output.messages, - tool_calls: output.tool_calls, + let response = InsertArchivalResponse { + id: inserted.uuid, + source_type: inserted.source_type, + embedding_model: inserted.embedding_model, + token_count: inserted.token_count, + created_at: inserted.created_at, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn delete_archival( + State(state): State>, + Path(embedding_uuid): Path, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + use crate::models::schema::user_embeddings::dsl::*; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let source: Option = user_embeddings + .filter(user_id.eq(user.uuid)) + .filter(uuid.eq(embedding_uuid)) + .select(source_type) + .first::(&mut conn) + .optional() + .map_err(|_| ApiError::InternalServerError)?; + + match source.as_deref() { + Some(crate::rag::SOURCE_TYPE_ARCHIVAL) => {} + Some(_) => return Err(ApiError::NotFound), + None => return Err(ApiError::NotFound), + } + + rag::delete_user_embedding_by_uuid(&state, user.uuid, embedding_uuid).await?; + + let response = DeletedObjectResponse { + id: embedding_uuid, + object: "archival.deleted", + deleted: true, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn memory_search( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, + Extension(body): Extension, +) -> Result>, ApiError> { + if body.query.trim().is_empty() { + return Err(ApiError::BadRequest); + } + + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let source_types = body.source_types.unwrap_or_else(|| { + vec![ + crate::rag::SOURCE_TYPE_MESSAGE.to_string(), + crate::rag::SOURCE_TYPE_ARCHIVAL.to_string(), + ] + }); + + let results = rag::search_user_embeddings( + &state, + user.uuid, + &user_key, + &body.query, + body.top_k.unwrap_or(5), + body.max_tokens, + Some(&source_types), + None, + ) + .await?; + + let response = MemorySearchResponse { results }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn list_agent_conversations( + State(state): State>, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let cfg = AgentConfig::get_by_user_id(&mut conn, user.uuid) + .map_err(|_| ApiError::InternalServerError)?; + + let mut data: Vec = Vec::new(); + + if let Some(cfg) = cfg { + if let Some(conversation_id) = cfg.conversation_id { + if let Ok(conv) = + Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) + { + let metadata = decrypt_content(&user_key, conv.metadata_enc.as_ref()) + .map_err(|_| ApiError::InternalServerError)?; + + data.push( + ConversationBuilder::from_conversation(&conv) + .metadata(metadata) + .build(), + ); + } + } + } + + let (first_id, last_id) = Paginator::get_cursor_ids(&data, |c| c.id); + let response = ConversationListResponse { + object: "list", + data, + has_more: false, + first_id, + last_id, }; encrypt_response(&state, &session_id, &response).await } + +async fn list_agent_conversation_items( + State(state): State>, + Path(conversation_uuid): Path, + Query(params): Query, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let ctx = load_agent_conversation_context(&state, user.uuid, conversation_uuid).await?; + + 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( + ctx.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, + &ctx.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, + }); + + let response = ConversationItemListResponse { + object: "list", + data: items, + has_more, + first_id, + last_id, + }; + + encrypt_response(&state, &session_id, &response).await +} + +async fn delete_agent_conversation( + State(state): State>, + Path(conversation_uuid): Path, + Extension(session_id): Extension, + Extension(user): Extension, +) -> Result>, ApiError> { + let ctx = load_agent_conversation_context(&state, user.uuid, conversation_uuid).await?; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + conn.transaction(|tx| -> Result<(), diesel::result::Error> { + use crate::models::schema::{agent_config, conversations}; + + diesel::delete( + user_embeddings::table + .filter(user_embeddings::user_id.eq(user.uuid)) + .filter(user_embeddings::conversation_id.eq(Some(ctx.conversation.id))) + .filter(user_embeddings::source_type.eq(crate::rag::SOURCE_TYPE_MESSAGE)), + ) + .execute(tx)?; + + // Delete the conversation (cascades). We keep the filter on user_id for safety. + let deleted = diesel::delete( + conversations::table + .filter(conversations::id.eq(ctx.conversation.id)) + .filter(conversations::user_id.eq(user.uuid)), + ) + .execute(tx)?; + + if deleted == 0 { + return Err(diesel::result::Error::NotFound); + } + + diesel::update(agent_config::table.filter(agent_config::user_id.eq(user.uuid))) + .set(agent_config::conversation_id.eq::>(None)) + .execute(tx)?; + + Ok(()) + }) + .map_err(|_| ApiError::InternalServerError)?; + + state.rag_cache.lock().await.evict_user(user.uuid); + + let response = DeletedObjectResponse::conversation(conversation_uuid); + + encrypt_response(&state, &session_id, &response).await +} + +struct AgentConversationContext { + conversation: Conversation, + user_key: secp256k1::SecretKey, +} + +async fn load_agent_conversation_context( + state: &AppState, + user_id: Uuid, + conversation_uuid: Uuid, +) -> Result { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let cfg = AgentConfig::get_by_user_id(&mut conn, user_id) + .map_err(|_| ApiError::InternalServerError)? + .ok_or(ApiError::NotFound)?; + + let Some(agent_conversation_id) = cfg.conversation_id else { + return Err(ApiError::NotFound); + }; + + let conversation = Conversation::get_by_uuid_and_user(&mut conn, conversation_uuid, user_id) + .map_err(|_| ApiError::NotFound)?; + + if conversation.id != agent_conversation_id { + return Err(ApiError::NotFound); + } + + let user_key = state + .get_user_key(user_id, None, None) + .await + .map_err(|_| error_mapping::map_key_retrieval_error())?; + + Ok(AgentConversationContext { + conversation, + user_key, + }) +} + +async fn get_or_create_agent_config( + conn: &mut PgConnection, + user_id: Uuid, +) -> Result { + let Some(cfg) = + AgentConfig::get_by_user_id(conn, user_id).map_err(|_| ApiError::InternalServerError)? + else { + let new_cfg = NewAgentConfig { + uuid: Uuid::new_v4(), + user_id, + conversation_id: None, + enabled: true, + model: "deepseek-r1-0528".to_string(), + max_context_tokens: 100_000, + compaction_threshold: 0.80, + system_prompt_enc: None, + preferences_enc: None, + }; + + return new_cfg + .insert_or_update(conn) + .map_err(|_| ApiError::InternalServerError); + }; + + Ok(cfg) +} diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs new file mode 100644 index 00000000..de7286c8 --- /dev/null +++ b/src/web/agent/runtime.rs @@ -0,0 +1,1048 @@ +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::agent_config::{AgentConfig, NewAgentConfig}; +use crate::models::conversation_summaries::{ConversationSummary, NewConversationSummary}; +use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; +use crate::models::responses::{ + AssistantMessage, Conversation, NewAssistantMessage, NewConversation, NewToolCall, + NewToolOutput, NewUserMessage, RawThreadMessage, RawThreadMessageMetadata, ToolCall, + ToolOutput, UserMessage, +}; +use crate::models::schema::{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::responses::{MessageContent, MessageContentConverter}; +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, ToolRegistry, ToolResult, +}; + +// Mirrors Sage defaults +const DEFAULT_PERSONA_DESCRIPTION: &str = "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions."; +const DEFAULT_HUMAN_DESCRIPTION: &str = "The human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation."; +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_CONTEXT_WINDOW: i32 = 100_000; +const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; +const MIN_MESSAGES_IN_CONTEXT: usize = 20; + +#[derive(Clone, Debug, Default)] +struct AgentContext { + current_time: 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)] +struct StepResult { + messages: Vec, + tool_calls: Vec, + done: bool, +} + +pub struct AgentRuntime { + state: Arc, + user: Arc, + user_key: Arc, + agent_config: AgentConfig, + conversation: Conversation, + 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, +} + +impl AgentRuntime { + pub async fn new( + 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)?; + + // Load or create config + let agent_config = match AgentConfig::get_by_user_id(&mut conn, user.uuid).map_err(|e| { + error!("Failed to load agent config: {e:?}"); + ApiError::InternalServerError + })? { + Some(cfg) => cfg, + None => { + let new_cfg = NewAgentConfig { + uuid: Uuid::new_v4(), + user_id: user.uuid, + conversation_id: None, + enabled: true, + model: "deepseek-r1-0528".to_string(), + max_context_tokens: DEFAULT_CONTEXT_WINDOW, + compaction_threshold: DEFAULT_COMPACTION_THRESHOLD, + system_prompt_enc: None, + preferences_enc: None, + }; + + new_cfg.insert_or_update(&mut conn).map_err(|e| { + error!("Failed to create agent config: {e:?}"); + ApiError::InternalServerError + })? + } + }; + + if !agent_config.enabled { + return Err(ApiError::Unauthorized); + } + + // Ensure conversation exists + let conversation = if let Some(conversation_id) = agent_config.conversation_id { + Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid).map_err( + |e| { + error!("Failed to load agent conversation: {e:?}"); + ApiError::InternalServerError + }, + )? + } else { + let metadata = json!({"type":"agent_main"}); + let metadata_enc = encrypt_with_key(&user_key, metadata.to_string().as_bytes()).await; + + let new_conversation = NewConversation { + uuid: Uuid::new_v4(), + user_id: user.uuid, + metadata_enc: Some(metadata_enc), + }; + let conversation = new_conversation.insert(&mut conn).map_err(|e| { + error!("Failed to create agent conversation: {e:?}"); + ApiError::InternalServerError + })?; + + let updated_cfg = NewAgentConfig { + uuid: agent_config.uuid, + user_id: agent_config.user_id, + conversation_id: Some(conversation.id), + enabled: agent_config.enabled, + model: agent_config.model.clone(), + max_context_tokens: agent_config.max_context_tokens, + compaction_threshold: agent_config.compaction_threshold, + system_prompt_enc: agent_config.system_prompt_enc.clone(), + preferences_enc: agent_config.preferences_enc.clone(), + }; + let _ = updated_cfg.insert_or_update(&mut conn).map_err(|e| { + error!("Failed to update agent config with conversation_id: {e:?}"); + ApiError::InternalServerError + })?; + + conversation + }; + + // Ensure default memory blocks exist + Self::ensure_default_blocks(&state, &mut conn, &user_key, user.uuid).await?; + + // System prompt override + let system_prompt = match decrypt_string(&user_key, agent_config.system_prompt_enc.as_ref()) + .map_err(|e| { + error!("Failed to decrypt system_prompt_enc: {e:?}"); + ApiError::InternalServerError + })? { + Some(s) if !s.trim().is_empty() => s, + _ => AGENT_INSTRUCTION.to_string(), + }; + + // Tool registry + 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.uuid, + user_key.clone(), + ))); + tools.register(Arc::new(ArchivalSearchTool::new( + state.clone(), + user.uuid, + user_key.clone(), + ))); + tools.register(Arc::new(ConversationSearchTool::new( + state.clone(), + user.uuid, + user_key.clone(), + conversation.id, + ))); + tools.register(Arc::new(DoneTool)); + + let available_tools = tools.generate_description(); + + let lm = build_lm( + state.clone(), + user.clone(), + agent_config.model.clone(), + 0.7, + 32768, + ) + .await?; + + Ok(Self { + state, + user, + user_key, + agent_config, + conversation, + 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, "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: "persona".to_string(), + description: Some(DEFAULT_PERSONA_DESCRIPTION.to_string()), + value_enc, + char_limit: 2000, + read_only: false, + version: 1, + }; + 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, "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: "human".to_string(), + description: Some(DEFAULT_HUMAN_DESCRIPTION.to_string()), + value_enc, + char_limit: 2000, + read_only: false, + version: 1, + }; + 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 async fn process_message( + &mut self, + user_message: &str, + ) -> Result<(Vec, Vec), ApiError> { + let trimmed = user_message.trim(); + if trimmed.is_empty() { + return Err(ApiError::BadRequest); + } + + // Clear at start of request + self.clear_tool_results(); + + // Persist user message (and recall embedding) + self.insert_user_message(trimmed).await?; + + // Compaction (summary memory) if needed + self.maybe_compact().await?; + + let mut all_messages: Vec = Vec::new(); + let mut last_tool_calls: Vec = Vec::new(); + + for step_num in 0..self.max_steps { + let result = self.step(trimmed, step_num == 0).await?; + + all_messages.extend(result.messages); + last_tool_calls = result.tool_calls; + + if result.done { + break; + } + } + + if all_messages.is_empty() { + warn!("Agent produced no messages"); + all_messages.push("I apologize, but I wasn't able to generate a response.".to_string()); + } + + Ok((all_messages, last_tool_calls)) + } + + 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, + 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(); + + // Persist assistant messages (before tool execution to preserve ordering) + for msg in &messages { + self.insert_assistant_message(msg).await?; + } + + 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)) + }; + + // Inject tool result for next step + self.inject_tool_result(tool_call, &result); + + // Persist tool call + output (skip done) + if tool_call.name != "done" { + self.insert_tool_call_and_output(tool_call, &result).await?; + } + } + + 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, + 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)")); + + 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, "persona") + .map_err(|_| ApiError::InternalServerError)?; + let human = MemoryBlock::get_by_user_and_label(&mut conn, self.user.uuid, "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)) + .filter(user_embeddings::conversation_id.eq(Some(self.conversation.id))) + .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 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 == "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) + } + + async fn insert_user_message(&self, text: &str) -> Result { + let content = super::tools::normalize_user_message_content(text); + let prompt_tokens = super::tools::count_user_message_tokens(&content); + let content_enc = encrypt_with_key(&self.user_key, content.as_bytes()).await; + + 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, + 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_id = self.user.uuid; + let user_key = self.user_key.clone(); + let text = 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_id, + user_key.as_ref(), + &text, + conversation_id, + user_message_id, + None, + ) + .await; + }); + + Ok(msg) + } + + 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, + } + .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_id = self.user.uuid; + 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_id, + user_key.as_ref(), + &text, + conversation_id, + None, + assistant_message_id, + ) + .await; + }); + + Ok(msg) + } + + 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 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(), + } + .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(), + } + .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, + self.agent_config.max_context_tokens as usize, + self.agent_config.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, + 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 index 73b46b6c..594525e1 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -13,9 +13,151 @@ use dspy_rs::client_registry::AssistantContent; use dspy_rs::{CompletionError, CompletionRequest, CompletionResponse}; use dspy_rs::{CustomCompletionModel, LMClient, OneOrMany, LM}; -pub const DEFAULT_AGENT_SYSTEM_PROMPT: &str = r#"You are Maple, a helpful AI assistant. - -Return `tool_calls` as an empty JSON array when no tools are needed."#; +/// 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 main 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 + +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)] @@ -52,27 +194,154 @@ pub struct AgentResponse { 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, } -pub async fn run_agent_response_signature( - state: &Arc, - user: &crate::models::users::User, - model: &str, +#[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 { - let lm = build_lm(state.clone(), Arc::new(user.clone()), model.to_string()).await?; + 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:?}"); - ApiError::InternalServerError + SignatureCallError::Api(ApiError::InternalServerError) })?; let user_msg = adapter.format_user_message_typed::(input); @@ -82,14 +351,18 @@ pub async fn run_agent_response_signature( let response = lm.call(chat, Vec::new()).await.map_err(|e| { error!("DSRS LM call failed: {e:?}"); - ApiError::InternalServerError + 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:?}"); - ApiError::InternalServerError + SignatureCallError::Parse { + raw_response, + error_message: e.to_string(), + } })?; let mut messages = output.messages; @@ -101,10 +374,12 @@ pub async fn run_agent_response_signature( }) } -async fn build_lm( +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| { @@ -159,8 +434,8 @@ async fn build_lm( let lm = LM::builder() .base_url("http://localhost".to_string()) .model(model) - .temperature(0.7) - .max_tokens(32768) + .temperature(temperature) + .max_tokens(max_tokens) .cache(false) .build() .await diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs new file mode 100644 index 00000000..dd783577 --- /dev/null +++ b/src/web/agent/tools.rs @@ -0,0 +1,824 @@ +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::conversation_summaries::ConversationSummary; +use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; +use crate::models::responses::Conversation; +use crate::rag::{cosine_similarity, deserialize_f32_le, search_user_embeddings}; +use crate::tokens::count_tokens; +use crate::web::responses::MessageContentConverter; +use crate::web::responses::{MessageContent, MessageContentPart}; +use crate::{ApiError, AppState}; + +// ============================================================================ +// 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 existing.read_only { + return Err(ApiError::BadRequest); + } + + if new_value.len() > existing.char_limit.max(0) as usize { + 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(), + description: existing.description.clone(), + value_enc, + char_limit: existing.char_limit, + read_only: existing.read_only, + version: existing.version, + }; + + 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_id: Uuid, + user_key: Arc, + agent_conversation_id: i64, +} + +impl ConversationSearchTool { + pub fn new( + state: Arc, + user_id: Uuid, + user_key: Arc, + conversation_id: i64, + ) -> Self { + Self { + state, + user_id, + 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, + 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_id, + 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_id) + { + 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_id, + &self.user_key, + query, + limit, + None, + Some(&source_types), + conversation_id_filter, + ) + .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_id: Uuid, + user_key: Arc, +} + +impl ArchivalInsertTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + 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_id, + &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_id: Uuid, + user_key: Arc, +} + +impl ArchivalSearchTool { + pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + Self { + state, + user_id, + 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); + + // NOTE: Tag filtering isn't indexable because metadata is encrypted. + // We accept the argument for schema compatibility but don't apply it. + + let source_types = vec![crate::rag::SOURCE_TYPE_ARCHIVAL.to_string()]; + match search_user_embeddings( + &self.state, + self.user_id, + &self.user_key, + query, + top_k, + None, + Some(&source_types), + None, + ) + .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") + } + } + } +} + +// ============================================================================ +// 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()) + } +} + +// ============================================================================ +// Formatting helpers +// ============================================================================ + +pub fn normalize_user_message_content(text: &str) -> String { + // Store user messages in the DB as MessageContent JSON, like Responses API. + // Use input_text part for compatibility with Conversations API. + let content = MessageContent::Parts(vec![MessageContentPart::InputText { + text: text.to_string(), + }]); + serde_json::to_string(&content).unwrap_or_else(|_| format!("\"{}\"", text)) +} + +pub fn count_user_message_tokens(content_json: &str) -> i32 { + let parsed: Result = serde_json::from_str(content_json); + match parsed { + Ok(content) => count_tokens(&MessageContentConverter::extract_text_for_token_counting( + &content, + )) as i32, + Err(_) => count_tokens(content_json) as i32, + } +} + +#[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/responses/conversations.rs b/src/web/responses/conversations.rs index d3be4843..4e063184 100644 --- a/src/web/responses/conversations.rs +++ b/src/web/responses/conversations.rs @@ -3,6 +3,7 @@ use crate::{ encrypt::{decrypt_content, encrypt_with_key}, + models::agent_config::AgentConfig, models::responses::NewConversation, models::users::User, web::{ @@ -69,6 +70,23 @@ 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)?; + + AgentConfig::get_by_user_id(&mut conn, user_uuid) + .map_err(|_| ApiError::InternalServerError)? + .and_then(|cfg| cfg.conversation_id) + }; + + if Some(conversation.id) == agent_conversation_id { + return Err(ApiError::NotFound); + } + // Get user's encryption key let user_key = state .get_user_key(user_uuid, None, None) @@ -545,11 +563,24 @@ async fn list_conversations( params.limit.min(MAX_PAGINATION_LIMIT) }; + let agent_conversation_id = { + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + AgentConfig::get_by_user_id(&mut conn, user.uuid) + .map_err(|_| ApiError::InternalServerError)? + .and_then(|cfg| cfg.conversation_id) + }; + // Fetch conversations with database-level pagination // We fetch limit + 1 to check if there are more results let conversations = state .db - .list_conversations(user.uuid, limit + 1, params.after, ¶ms.order) + // Fetch an extra conversation since we may filter out the hidden agent thread. + .list_conversations(user.uuid, limit + 2, params.after, ¶ms.order) .map_err(error_mapping::map_generic_db_error)?; trace!( @@ -561,14 +592,18 @@ async fn list_conversations( 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 @@ -595,7 +630,7 @@ async fn list_conversations( .collect::, _>>()?; // 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, @@ -618,10 +653,33 @@ async fn delete_all_conversations( ) -> Result>, ApiError> { debug!("Deleting all conversations for user {}", user.uuid); - state + 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 = AgentConfig::get_by_user_id(&mut conn, user.uuid) + .map_err(|_| ApiError::InternalServerError)? + .and_then(|cfg| cfg.conversation_id); + + 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, @@ -654,6 +712,18 @@ 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)?; + + AgentConfig::get_by_user_id(&mut conn, user.uuid) + .map_err(|_| ApiError::InternalServerError)? + .and_then(|cfg| cfg.conversation_id) + }; + for conversation_id in body.ids { // Try to get the conversation first to verify ownership match state @@ -661,6 +731,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(()) => { From 4121f82bf8b4b7b250261be4460e0f879f138e5e Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 10:26:31 -0600 Subject: [PATCH 05/34] fix: enforce billing for internal embeddings Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/rag.rs | 37 +++++++--- src/web/agent/mod.rs | 7 +- src/web/agent/runtime.rs | 19 ++++-- src/web/agent/tools.rs | 40 ++++++----- src/web/openai.rs | 141 ++++++++++++++++++++------------------- src/web/rag.rs | 7 +- 6 files changed, 145 insertions(+), 106 deletions(-) diff --git a/src/rag.rs b/src/rag.rs index 118c2287..822c31e9 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -1,6 +1,8 @@ use crate::encrypt::{decrypt_with_key, 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 diesel::prelude::*; use secp256k1::SecretKey; @@ -285,9 +287,16 @@ fn apply_token_budget(results: Vec, budget: i32) -> Vec Result<(Vec, i32), ApiError> { +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), @@ -296,13 +305,15 @@ async fn embed_text_via_tinfoil(state: &AppState, text: &str) -> Result<(Vec, + user: &User, + auth_method: AuthMethod, user_key: &SecretKey, text: &str, metadata: Option<&serde_json::Value>, ) -> Result { - let (vector, token_count) = embed_text_via_tinfoil(state, text).await?; + 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; @@ -347,20 +358,22 @@ pub async fn insert_archival_embedding( #[allow(clippy::too_many_arguments)] pub async fn insert_message_embedding( - state: &AppState, - user_id: Uuid, + 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, text).await?; + 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; @@ -491,8 +504,9 @@ async fn load_all_user_embeddings( #[allow(clippy::too_many_arguments)] pub async fn search_user_embeddings( - state: &AppState, - user_id: Uuid, + state: &Arc, + user: &User, + auth_method: AuthMethod, user_key: &SecretKey, query: &str, top_k: usize, @@ -502,7 +516,10 @@ pub async fn search_user_embeddings( ) -> Result, ApiError> { let top_k = top_k.clamp(1, 20); - let (query_vec, _query_tokens) = embed_text_via_tinfoil(state, query).await?; + let user_id = user.uuid; + + let (query_vec, _query_tokens) = + embed_text_via_tinfoil(state, user, auth_method, query).await?; let cached = { let mut cache = state.rag_cache.lock().await; diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index acfe2c85..81181730 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -19,6 +19,7 @@ use crate::models::schema::user_embeddings; 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::constants::{DEFAULT_PAGINATION_LIMIT, MAX_PAGINATION_LIMIT}; use crate::web::responses::conversations::{ ConversationItemListResponse, ConversationListResponse, ConversationResponse, ListItemsParams, @@ -572,7 +573,8 @@ async fn insert_archival( let inserted = rag::insert_archival_embedding( &state, - user.uuid, + &user, + AuthMethod::Jwt, &user_key, &body.text, body.metadata.as_ref(), @@ -653,7 +655,8 @@ async fn memory_search( let results = rag::search_user_embeddings( &state, - user.uuid, + &user, + AuthMethod::Jwt, &user_key, &body.query, body.top_k.unwrap_or(5), diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index de7286c8..2badb93b 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -20,6 +20,7 @@ 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}; use crate::{ApiError, AppState}; @@ -193,17 +194,17 @@ impl AgentRuntime { ))); tools.register(Arc::new(ArchivalInsertTool::new( state.clone(), - user.uuid, + user.clone(), user_key.clone(), ))); tools.register(Arc::new(ArchivalSearchTool::new( state.clone(), - user.uuid, + user.clone(), user_key.clone(), ))); tools.register(Arc::new(ConversationSearchTool::new( state.clone(), - user.uuid, + user.clone(), user_key.clone(), conversation.id, ))); @@ -757,7 +758,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If // Mirror Sage: store message synchronously, update embedding in background. let state = self.state.clone(); - let user_id = self.user.uuid; + let user = self.user.clone(); let user_key = self.user_key.clone(); let text = text.to_string(); let conversation_id = self.conversation.id; @@ -766,7 +767,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If tokio::spawn(async move { let _ = insert_message_embedding( &state, - user_id, + user.as_ref(), + AuthMethod::Jwt, user_key.as_ref(), &text, conversation_id, @@ -808,7 +810,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If // Mirror Sage: store message synchronously, update embedding in background. let state = self.state.clone(); - let user_id = self.user.uuid; + let user = self.user.clone(); let user_key = self.user_key.clone(); let text = text.to_string(); let conversation_id = self.conversation.id; @@ -817,7 +819,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If tokio::spawn(async move { let _ = insert_message_embedding( &state, - user_id, + user.as_ref(), + AuthMethod::Jwt, user_key.as_ref(), &text, conversation_id, @@ -1004,6 +1007,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If // 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), diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index dd783577..df0e4e88 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -10,8 +10,10 @@ use crate::encrypt::{decrypt_string, encrypt_with_key}; use crate::models::conversation_summaries::ConversationSummary; use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; use crate::models::responses::Conversation; +use crate::models::users::User; use crate::rag::{cosine_similarity, deserialize_f32_le, search_user_embeddings}; use crate::tokens::count_tokens; +use crate::web::openai_auth::AuthMethod; use crate::web::responses::MessageContentConverter; use crate::web::responses::{MessageContent, MessageContentPart}; use crate::{ApiError, AppState}; @@ -434,7 +436,7 @@ impl Tool for MemoryInsertTool { pub struct ConversationSearchTool { state: Arc, - user_id: Uuid, + user: Arc, user_key: Arc, agent_conversation_id: i64, } @@ -442,13 +444,13 @@ pub struct ConversationSearchTool { impl ConversationSearchTool { pub fn new( state: Arc, - user_id: Uuid, + user: Arc, user_key: Arc, conversation_id: i64, ) -> Self { Self { state, - user_id, + user, user_key, agent_conversation_id: conversation_id, } @@ -461,6 +463,8 @@ impl ConversationSearchTool { ) -> 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), @@ -477,7 +481,7 @@ impl ConversationSearchTool { // Fetch latest summaries. Bound this to keep runtime predictable. let summaries = ConversationSummary::list_for_conversation( &mut conn, - self.user_id, + self.user.uuid, self.agent_conversation_id, 200, ) @@ -540,8 +544,11 @@ impl Tool for ConversationSearchTool { Err(_) => return ToolResult::error("database connection error"), }; - match Conversation::get_by_uuid_and_user(&mut conn, conversation_uuid, self.user_id) - { + 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"), } @@ -556,7 +563,8 @@ impl Tool for ConversationSearchTool { let source_types = vec![crate::rag::SOURCE_TYPE_MESSAGE.to_string()]; match search_user_embeddings( &self.state, - self.user_id, + self.user.as_ref(), + AuthMethod::Jwt, &self.user_key, query, limit, @@ -620,15 +628,15 @@ impl Tool for ConversationSearchTool { pub struct ArchivalInsertTool { state: Arc, - user_id: Uuid, + user: Arc, user_key: Arc, } impl ArchivalInsertTool { - pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + pub fn new(state: Arc, user: Arc, user_key: Arc) -> Self { Self { state, - user_id, + user, user_key, } } @@ -668,7 +676,8 @@ impl Tool for ArchivalInsertTool { match crate::rag::insert_archival_embedding( &self.state, - self.user_id, + self.user.as_ref(), + AuthMethod::Jwt, &self.user_key, content, metadata_ref, @@ -689,15 +698,15 @@ impl Tool for ArchivalInsertTool { pub struct ArchivalSearchTool { state: Arc, - user_id: Uuid, + user: Arc, user_key: Arc, } impl ArchivalSearchTool { - pub fn new(state: Arc, user_id: Uuid, user_key: Arc) -> Self { + pub fn new(state: Arc, user: Arc, user_key: Arc) -> Self { Self { state, - user_id, + user, user_key, } } @@ -729,7 +738,8 @@ impl Tool for ArchivalSearchTool { let source_types = vec![crate::rag::SOURCE_TYPE_ARCHIVAL.to_string()]; match search_user_embeddings( &self.state, - self.user_id, + self.user.as_ref(), + AuthMethod::Jwt, &self.user_key, query, top_k, diff --git a/src/web/openai.rs b/src/web/openai.rs index 7d275923..5b084874 100644 --- a/src/web/openai.rs +++ b/src/web/openai.rs @@ -1657,8 +1657,73 @@ async fn request_embeddings_from_primary( Ok((response_json, route.primary.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 prompt_tokens > 0 { + let billing_context = BillingContext::new(auth_method, embedding_request.model.clone()); + let embedding_usage = CompletionUsage { + prompt_tokens, + completion_tokens: 0, + }; + 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: &AppState, + state: &Arc, + user: &User, + auth_method: AuthMethod, model: &str, input: &str, dimensions: Option, @@ -1671,7 +1736,8 @@ pub async fn get_embedding_vector( user: None, }; - let (response_json, _provider) = request_embeddings_from_primary(state, &request).await?; + let (response_json, _provider, prompt_tokens) = + request_embeddings_with_billing(state, user, auth_method, &request).await?; let embedding = response_json .get("data") @@ -1698,18 +1764,6 @@ pub async fn get_embedding_vector( } } - let prompt_tokens = match response_json - .get("usage") - .and_then(|u| u.get("prompt_tokens")) - .and_then(|v| v.as_i64()) - { - Some(v) => v as i32, - None => { - error!("Embeddings response missing usage.prompt_tokens"); - 0 - } - }; - Ok((vector, prompt_tokens)) } @@ -1718,39 +1772,11 @@ async fn proxy_embeddings( _headers: HeaderMap, axum::Extension(session_id): axum::Extension, axum::Extension(user): axum::Extension, - axum::Extension(_auth_method): axum::Extension, + axum::Extension(auth_method): axum::Extension, axum::Extension(embedding_request): axum::Extension, ) -> Result>, ApiError> { debug!("Entering proxy_embeddings function"); - // 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(), @@ -1762,33 +1788,8 @@ async fn proxy_embeddings( return Err(ApiError::BadRequest); } - let (response_json, provider_name) = - request_embeddings_from_primary(&state, &embedding_request).await?; - - // 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 - }; - publish_usage_event_internal( - &state, - &user, - &billing_context, - embedding_usage, - &provider_name, - ) - .await; - } - } + let (response_json, _provider_name, _prompt_tokens) = + request_embeddings_with_billing(&state, &user, auth_method, &embedding_request).await?; debug!("Exiting proxy_embeddings function"); diff --git a/src/web/rag.rs b/src/web/rag.rs index bf3949d5..c0e5621f 100644 --- a/src/web/rag.rs +++ b/src/web/rag.rs @@ -13,6 +13,7 @@ use uuid::Uuid; 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}; @@ -105,7 +106,8 @@ async fn insert_archival_embedding( let inserted = rag::insert_archival_embedding( &state, - user.uuid, + &user, + AuthMethod::Jwt, &user_key, &body.text, body.metadata.as_ref(), @@ -168,7 +170,8 @@ async fn search( let results = rag::search_user_embeddings( &state, - user.uuid, + &user, + AuthMethod::Jwt, &user_key, &body.query, top_k, From 5d31e44c581637a7b24981fc50a9cce39f64ddbc Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 11:03:42 -0600 Subject: [PATCH 06/34] fix: enforce agent memory block char limits Use Sage's 20k default char_limit and reject oversized updates via /v1/agent/memory/blocks/:label. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/models/memory_blocks.rs | 5 ++++- src/web/agent/mod.rs | 32 +++++++++++++++++++++++++++----- src/web/agent/runtime.rs | 6 +++--- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/models/memory_blocks.rs b/src/models/memory_blocks.rs index e9128520..ff399a29 100644 --- a/src/models/memory_blocks.rs +++ b/src/models/memory_blocks.rs @@ -5,6 +5,9 @@ use serde::{Deserialize, Serialize}; use thiserror::Error; use uuid::Uuid; +/// Default character limit per core memory block (mirrors Sage/Letta). +pub const DEFAULT_BLOCK_CHAR_LIMIT: i32 = 20_000; + #[derive(Error, Debug)] pub enum MemoryBlockError { #[error("Database error: {0}")] @@ -102,7 +105,7 @@ impl NewMemoryBlock { label: label.into(), description, value_enc, - char_limit: 5000, + char_limit: DEFAULT_BLOCK_CHAR_LIMIT, read_only: false, version: 1, } diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 81181730..2e374069 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -13,7 +13,7 @@ use uuid::Uuid; use crate::encrypt::{decrypt_content, decrypt_string, encrypt_with_key}; use crate::models::agent_config::{AgentConfig, NewAgentConfig}; -use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; +use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock, DEFAULT_BLOCK_CHAR_LIMIT}; use crate::models::responses::{Conversation, NewConversation}; use crate::models::schema::user_embeddings; use crate::models::users::User; @@ -500,6 +500,31 @@ async fn update_memory_block( } } + let char_limit = body + .char_limit + .or_else(|| existing.as_ref().map(|b| b.char_limit)) + .unwrap_or(DEFAULT_BLOCK_CHAR_LIMIT); + + if char_limit <= 0 { + return Err(ApiError::BadRequest); + } + + if let Some(value) = body.value.as_ref() { + if value.len() > char_limit as usize { + return Err(ApiError::BadRequest); + } + } else if body.char_limit.is_some() { + if let Some(b) = &existing { + let existing_value = decrypt_string(&user_key, Some(&b.value_enc)) + .map_err(|_| ApiError::InternalServerError)? + .unwrap_or_default(); + + if existing_value.len() > char_limit as usize { + return Err(ApiError::BadRequest); + } + } + } + let value_enc = match body.value { Some(value) => encrypt_with_key(&user_key, value.as_bytes()).await, None => match &existing { @@ -519,10 +544,7 @@ async fn update_memory_block( .description .or_else(|| existing.as_ref().and_then(|b| b.description.clone())), value_enc, - char_limit: body - .char_limit - .or_else(|| existing.as_ref().map(|b| b.char_limit)) - .unwrap_or(5000), + char_limit, read_only: body .read_only .or_else(|| existing.as_ref().map(|b| b.read_only)) diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 2badb93b..db74908e 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -9,7 +9,7 @@ use diesel::prelude::*; use crate::encrypt::{decrypt_string, encrypt_with_key}; use crate::models::agent_config::{AgentConfig, NewAgentConfig}; use crate::models::conversation_summaries::{ConversationSummary, NewConversationSummary}; -use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; +use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock, DEFAULT_BLOCK_CHAR_LIMIT}; use crate::models::responses::{ AssistantMessage, Conversation, NewAssistantMessage, NewConversation, NewToolCall, NewToolOutput, NewUserMessage, RawThreadMessage, RawThreadMessageMetadata, ToolCall, @@ -258,7 +258,7 @@ impl AgentRuntime { label: "persona".to_string(), description: Some(DEFAULT_PERSONA_DESCRIPTION.to_string()), value_enc, - char_limit: 2000, + char_limit: DEFAULT_BLOCK_CHAR_LIMIT, read_only: false, version: 1, }; @@ -281,7 +281,7 @@ impl AgentRuntime { label: "human".to_string(), description: Some(DEFAULT_HUMAN_DESCRIPTION.to_string()), value_enc, - char_limit: 2000, + char_limit: DEFAULT_BLOCK_CHAR_LIMIT, read_only: false, version: 1, }; From 8a478415240572f10cee225f85d48b4856e4ba2f Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 12:03:24 -0600 Subject: [PATCH 07/34] feat: add encrypted tag filtering for archival memory Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../down.sql | 1 + .../2026-02-10-020612_user_embeddings/up.sql | 8 + src/models/schema.rs | 1 + src/models/user_embeddings.rs | 2 + src/rag.rs | 199 +++++++++++++++++- src/web/agent/mod.rs | 1 + src/web/agent/tools.rs | 13 +- src/web/rag.rs | 1 + 8 files changed, 213 insertions(+), 13 deletions(-) diff --git a/migrations/2026-02-10-020612_user_embeddings/down.sql b/migrations/2026-02-10-020612_user_embeddings/down.sql index 167f3308..2be284ca 100644 --- a/migrations/2026-02-10-020612_user_embeddings/down.sql +++ b/migrations/2026-02-10-020612_user_embeddings/down.sql @@ -3,6 +3,7 @@ DROP TRIGGER IF EXISTS update_user_embeddings_updated_at ON user_embeddings; DROP INDEX IF EXISTS idx_user_embeddings_user_id; DROP INDEX IF EXISTS idx_user_embeddings_user_created; DROP INDEX IF EXISTS idx_user_embeddings_user_source; +DROP INDEX IF EXISTS idx_user_embeddings_archival_tags_enc; DROP INDEX IF EXISTS idx_user_embeddings_user_conversation; DROP INDEX IF EXISTS idx_user_embeddings_user_message_id; DROP INDEX IF EXISTS idx_user_embeddings_assistant_message_id; diff --git a/migrations/2026-02-10-020612_user_embeddings/up.sql b/migrations/2026-02-10-020612_user_embeddings/up.sql index 54071675..76621fb6 100644 --- a/migrations/2026-02-10-020612_user_embeddings/up.sql +++ b/migrations/2026-02-10-020612_user_embeddings/up.sql @@ -25,6 +25,9 @@ CREATE TABLE user_embeddings ( content_enc BYTEA NOT NULL, metadata_enc BYTEA, + -- Deterministically-encrypted tags (base64) for SQL-level filtering + tags_enc TEXT[] NOT NULL DEFAULT '{}', + -- Plaintext metadata token_count INTEGER NOT NULL, @@ -61,6 +64,11 @@ CREATE INDEX idx_user_embeddings_user_created ON user_embeddings(user_id, create -- For source-type filtered searches (message vs archival) CREATE INDEX idx_user_embeddings_user_source ON user_embeddings(user_id, source_type); +-- For tag-filtered archival searches (encrypted, per-user tokens) +CREATE INDEX idx_user_embeddings_archival_tags_enc + ON user_embeddings USING GIN(tags_enc) + WHERE source_type = 'archival'; + -- For conversation-scoped searches (message recall) CREATE INDEX idx_user_embeddings_user_conversation ON user_embeddings(user_id, conversation_id); diff --git a/src/models/schema.rs b/src/models/schema.rs index 2df1656f..3844eac8 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -371,6 +371,7 @@ diesel::table! { vector_dim -> Int4, content_enc -> Bytea, metadata_enc -> Nullable, + tags_enc -> Array>, token_count -> Int4, created_at -> Timestamptz, updated_at -> Timestamptz, diff --git a/src/models/user_embeddings.rs b/src/models/user_embeddings.rs index fb3bbac3..804cb470 100644 --- a/src/models/user_embeddings.rs +++ b/src/models/user_embeddings.rs @@ -26,6 +26,7 @@ pub struct UserEmbedding { 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, @@ -45,6 +46,7 @@ pub struct NewUserEmbedding { pub vector_dim: i32, pub content_enc: Vec, pub metadata_enc: Option>, + pub tags_enc: Vec>, pub token_count: i32, } diff --git a/src/rag.rs b/src/rag.rs index 822c31e9..21380403 100644 --- a/src/rag.rs +++ b/src/rag.rs @@ -1,14 +1,15 @@ -use crate::encrypt::{decrypt_with_key, encrypt_with_key}; +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, VecDeque}; +use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; use std::sync::Arc; use std::time::{Duration, Instant}; use tracing::{debug, error}; @@ -304,6 +305,49 @@ async fn embed_text_via_tinfoil( .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, @@ -326,6 +370,9 @@ pub async fn insert_archival_embedding( None }; + let tags = extract_tags_from_metadata(metadata); + let tags_enc = encrypt_tags_b64(user_key, &tags); + let mut conn = state .db .get_pool() @@ -344,6 +391,7 @@ pub async fn insert_archival_embedding( vector_dim: DEFAULT_EMBEDDING_DIM, content_enc, metadata_enc, + tags_enc, token_count, } .insert(&mut conn) @@ -397,6 +445,7 @@ pub async fn insert_message_embedding( vector_dim: DEFAULT_EMBEDDING_DIM, content_enc, metadata_enc: None, + tags_enc: Vec::new(), token_count, } .insert(&mut conn) @@ -502,6 +551,116 @@ async fn load_all_user_embeddings( 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, @@ -513,6 +672,7 @@ pub async fn search_user_embeddings( max_tokens: Option, source_types: Option<&[String]>, conversation_id: Option, + tags: Option<&[String]>, ) -> Result, ApiError> { let top_k = top_k.clamp(1, 20); @@ -521,17 +681,34 @@ pub async fn search_user_embeddings( let (query_vec, _query_tokens) = embed_text_via_tinfoil(state, user, auth_method, query).await?; - let cached = { - let mut cache = state.rag_cache.lock().await; - cache.get(user_id) - }; + 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(hit) = cached { - hit + 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 loaded = load_all_user_embeddings(state, user_id, user_key).await?; - state.rag_cache.lock().await.put(user_id, loaded.clone()); - loaded + 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( diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 2e374069..9a969f69 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -685,6 +685,7 @@ async fn memory_search( body.max_tokens, Some(&source_types), None, + None, ) .await?; diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index df0e4e88..119fa1b6 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -571,6 +571,7 @@ impl Tool for ConversationSearchTool { None, Some(&source_types), conversation_id_filter, + None, ) .await { @@ -732,8 +733,15 @@ impl Tool for ArchivalSearchTool { }; let top_k: usize = args.get("top_k").and_then(|k| k.parse().ok()).unwrap_or(5); - // NOTE: Tag filtering isn't indexable because metadata is encrypted. - // We accept the argument for schema compatibility but don't apply it. + 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( @@ -746,6 +754,7 @@ impl Tool for ArchivalSearchTool { None, Some(&source_types), None, + tags.as_deref(), ) .await { diff --git a/src/web/rag.rs b/src/web/rag.rs index c0e5621f..7efe32d3 100644 --- a/src/web/rag.rs +++ b/src/web/rag.rs @@ -178,6 +178,7 @@ async fn search( body.max_tokens, body.source_types.as_deref(), conversation_internal_id, + None, ) .await?; From cf96827fd7994aace5005b31a2075735b3bc717b Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 12:28:32 -0600 Subject: [PATCH 08/34] feat: SSE message delivery for agent chat (Sage-style step loop) Restructure agent chat to emit messages per-step via SSE instead of batching all messages into a single JSON response. This mirrors Sage's Signal delivery pattern: messages are sent immediately as produced, persisted synchronously (so the next step sees them in context), and embeddings update async. - Split process_message into prepare() + public step() so the handler drives the loop - Chat handler emits agent.message, agent.done, agent.error SSE events - Persistence moved from runtime to handler (send first, store sync, embed async) - Added ExecutedTool to carry tool results for handler-side persistence - Updated sage-in-maple doc to reflect conversations API isolation done Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- internal_docs/sage-in-maple-architecture.md | 630 ++++++++++++++++++++ src/web/agent/mod.rs | 119 +++- src/web/agent/runtime.rs | 81 ++- 3 files changed, 775 insertions(+), 55 deletions(-) create mode 100644 internal_docs/sage-in-maple-architecture.md diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md new file mode 100644 index 00000000..1ade2186 --- /dev/null +++ b/internal_docs/sage-in-maple-architecture.md @@ -0,0 +1,630 @@ +# Sage-in-Maple: Agent Memory Architecture + +## Bringing Persistent Agent Memory into OpenSecret + +**Date:** February 2026 +**Status:** MVP implemented (Local/Dev only, non-streaming; Responses API auto-embedding + production feature flags pending) +**Related Docs:** +- `potential-rag-integration-brute-force.md` -- RAG/vector storage layer (prerequisite) +- `architecture-for-rag-integration.md` -- OpenSecret encryption and data model reference +- Sage V2 Design Doc (`~/Dev/Personal/sage/docs/SAGE_V2_DESIGN.md`) -- proven prototype +- Sage V2 Codebase (`~/Dev/Personal/sage/crates/sage-core`) -- DSRs signatures + BAML parsing + multi-step tool loop prototype + +--- + +## 1. Goal + +Bring Sage's proven 4-tier memory architecture (core, recall, archival, summary) into Maple as a first-class feature, without disrupting the existing Responses API that third-party developers build against. + +The end state: Maple users can have a persistent AI agent that remembers across conversations, has editable personality/user memory blocks, can store and search long-term memories, and auto-compacts conversation history when context windows fill up -- all with the same per-user encryption guarantees that exist today. + +### 1.1 The User Model + +**One user : one main agent : one persistent conversation.** + +The agent lives in a single long-running conversation thread. Unlike Responses API threads (which are isolated, stateless, and disposable), the agent conversation is the user's ongoing relationship with Maple. Memory blocks, archival memory, and conversation summaries all attach to this single thread. + +**Implementation note:** this mapping must be explicit in the database. We should not infer "the agent thread" from encrypted `conversations.metadata_enc`. Instead, store the thread pointer in `agent_config.conversation_id` (BIGINT FK to `conversations.id`) when the agent is initialized. + +The Responses API continues to exist alongside the agent. Users can still have one-off stateless threads for quick questions. But the agent is the primary interface for users who want memory and persistence. + +### 1.2 Cross-Thread Memory Visibility + +**The agent sees everything.** When the agent's `conversation_search` tool fires, it searches across all of the user's embedded messages -- both the agent's own conversation and any Responses API threads. The `user_embeddings` table indexes messages from all conversations by default. + +Rationale: if a user had a Responses API thread about a topic, then asks the agent about that topic, the agent should know about it. The user opted into the agent system; that's the trust boundary. Individual threads are not a privacy boundary within the same user account. + +The agent can still scope searches to its own thread via `conversation_id` filtering, but the default is broad. See the RAG proposal's "Data Visibility and Isolation" section for the full filtering model. + +**Future: incognito threads.** A `private` flag on conversations will let users create Responses API threads that are excluded from embedding entirely. The agent cannot see them. This is deferred to post-v1. + +Additionally, requests made with `store=false` should not be embedded/indexed into `user_embeddings` (per-request opt-out), even if the Responses API continues to persist messages today. + +### 1.3 Future Vision: Agent-Backed Responses API + +Eventually, the Responses API threads themselves could be backed by the agent's memory system. Instead of stateless threads with no memory, a Responses API thread would: +- Start with a fresh conversation history (no prior messages in context) +- But still have access to memory blocks, archival search, and conversation search +- Effectively: the agent "knows" the user (from core + archival memory), but the thread is a clean slate + +This would give users the best of both: the lightweight feel of a new thread, with the accumulated knowledge of the persistent agent. This is a post-v1 evolution that requires proving out the agent system first -- particularly that memory blocks and search produce consistently good results without the agent's own conversation context. + +### 1.4 Future: Subagents (Multi-Agent per User) + +Subagents are out of scope for the MVP, but we should keep the storage design compatible with them. + +**Key idea:** each agent is just: +1) an **agent identity/config** record, and +2) a **single long-running conversation thread** (`conversations.id`) that the agent reads/writes. + +The existing message tables remain unchanged; different agents simply write to different `conversation_id`s. + +We also want subagents to support **configurable memory scoping**: +- **Shared/inherited memory:** subagents read/write the user's shared memory blocks and archival memory. +- **Isolated memory:** subagents have their own blocks/archival memory (no visibility into the user's shared memory unless explicitly allowed). + +The cleanest way to support this is to introduce an `agents` table and add an `agent_id` FK to *agent-specific* tables (Section 4.5). Shared vs isolated memory becomes a query policy (and optionally a config field) rather than a schema rewrite. + +--- + +## 2. Key Architectural Decision: Separate API, Shared Storage + +### What stays the same + +The existing Responses API (`/v1/responses/*`, `/v1/conversations/*`, `/v1/instructions/*`) is **untouched**. It continues to be a stateless OpenAI-compatible chat API. Third-party developers who use it see no changes. + +The existing database tables for conversations and messages are **reused** by the agent system. The agent reads from and writes to the same `conversations`, `user_messages`, `assistant_messages`, `tool_calls`, `tool_outputs`, and `reasoning_items` tables. The encrypted content format is identical. + +### What's new + +A new `/v1/agent/*` API surface with its own handler module (`src/web/agent/`), its own step loop, and its own context assembly logic. This API implements the Sage-style regenerated-context pattern instead of the Responses API's middle-truncation pattern. + +New database tables for agent-specific concerns (detailed in Section 4): +- `memory_blocks` -- core memory (persona, human, custom blocks) +- `user_embeddings` -- archival memory + chat embeddings (from the RAG proposal) +- `conversation_summaries` -- compaction artifacts +- `agent_config` -- per-user agent settings + +### 2.1 Feature Flags and Rollout Safety (MVP requirement) + +Sage is a large feature that touches storage, tool execution, and LLM calls. We need a lightweight, production-friendly feature flag system to safely ship in enclaves. + +At minimum, ship with **two levels of gating**: +- **Global kill switch** (env/config): hard-disable all `/v1/agent/*` routes and background jobs (auto-embedding, compaction, reminders). +- **Per-user opt-in** (`agent_config.enabled`): enables the persistent agent for that user. + +Flags should also gate sub-features independently (even if `agent_config.enabled=true`): +- auto-embedding into `user_embeddings` (recall memory) +- agent tool execution (memory tools, web search) +- compaction/summarization +- GEPA runs (optimizer execution) vs simply *consuming* an optimized instruction + +### 2.2 Conversations API Isolation (Hide the Main Agent Thread) + +Even though the agent's persistent thread is stored in the shared `conversations`/message tables, the **main agent conversation should be treated as an internal implementation detail** and **excluded from the public Conversations API** (`/v1/conversations/*`). + +Rationale: +- Prevents confusion for developers using the stateless Responses/Conversations APIs. +- Avoids muddying thread lists with an always-on, long-running agent thread. +- Keeps `/v1/agent/*` as the single supported interface for agent state. + +**Status: Implemented (2026-02-11).** The agent thread is hidden from all public Conversations API endpoints: +- `agent_config.conversation_id` stores the agent thread pointer. +- `ConversationContext::load()` returns `404` when the requested conversation is the agent thread (guards `GET/POST/DELETE /v1/conversations/:id` and `/items`). +- `GET /v1/conversations` filters out the agent thread before pagination. +- `DELETE /v1/conversations` (delete-all) excludes the agent conversation via `id.ne(agent_conversation_id)`. +- `POST /v1/conversations/batch-delete` returns `not_found` for the agent thread instead of deleting it. + +--- + +## 3. Why Reuse the Message Tables + +This was the hardest decision. We did a column-by-column comparison between Sage's `messages` table and Maple's split message tables. The conclusion: + +**Maple's schema is a strict superset of Sage's.** Sage stores `(id, role, content, agent_id, sequence_id, embedding, tool_calls, tool_results, created_at)` in a single table. Maple stores the same information across typed tables with more metadata: token counts per message, response status tracking, response_id linking, conversation_id scoping, and UUIDs for external references. + +**What Sage does with messages at runtime:** + +| Operation | Sage approach | Maple equivalent | +|---|---|---| +| Read recent messages in order | `SELECT * FROM messages WHERE agent_id = ? ORDER BY sequence_id DESC LIMIT N` | The existing `RawThreadMessage::get_conversation_context` UNION ALL query across all message tables | +| Track which messages are "in context" | `agents.message_ids` UUID array | Compute dynamically from conversation + summaries (or add field to `agent_config`) | +| Store new messages | `INSERT INTO messages` with role | `INSERT INTO user_messages` / `assistant_messages` / etc. (existing flow) | +| Search messages semantically | pgvector on messages.embedding | `user_embeddings` table with in-process brute-force search (from RAG proposal) | +| Mark messages as summarized | Doesn't mark messages; records sequence ranges in summaries table | Same approach: `conversation_summaries` stores ranges, messages are untouched | + +**Nothing in Sage's message access patterns requires modifying the existing Maple message table columns.** The agent needs to *read* from those tables differently (regenerated context with memory blocks instead of middle-truncation), but the storage format is identical. + +**Potential future addition:** An `is_in_context` boolean or a `summary_id` FK on messages would be an optimization for fast filtering. But it's not required -- the summaries table's ranges can determine this. If needed, it would be an additive nullable column with a migration default, not a structural change. + +--- + +## 4. New Database Tables + +### 4.1 `memory_blocks` -- Core Memory + +The always-in-context memory blocks that define agent personality and user information. Directly modeled on Sage's `blocks` table, adapted for Maple's encryption model. + +```sql +CREATE TABLE memory_blocks ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + label TEXT NOT NULL, -- plaintext: "persona", "human", custom labels + description TEXT, -- plaintext: how this block should be used + value_enc BYTEA NOT NULL, -- AES-256-GCM encrypted block content + char_limit INTEGER NOT NULL DEFAULT 5000, + read_only BOOLEAN NOT NULL DEFAULT FALSE, + version INTEGER NOT NULL DEFAULT 1, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(user_id, label) +); + +CREATE INDEX idx_memory_blocks_user_id ON memory_blocks(user_id); + +CREATE TRIGGER update_memory_blocks_updated_at +BEFORE UPDATE ON memory_blocks +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**Key differences from Sage's `blocks` table:** +- `value_enc` (BYTEA) instead of `value` (TEXT) -- content encrypted with user's per-user key +- `user_id` instead of `agent_id` -- MVP keeps blocks per-user. For multi-agent/subagents, add a nullable `agent_id` so blocks can be shared (`agent_id IS NULL`) or agent-scoped (`agent_id = `). +- `label` and `description` are plaintext -- these are structural identifiers ("persona", "human"), not user content. They don't need encryption. This allows querying by label without decryption. + +**Default blocks created on agent initialization:** +- `persona`: "I am a helpful AI assistant." (editable by agent) +- `human`: "" (populated by agent as it learns about the user) + +### 4.2 `user_embeddings` -- General-Purpose Embedding Store + +Defined in the RAG proposal. A single table that serves all embedding use cases through a `source_type` discriminator: + +- `source_type = 'message'`: auto-indexed chat history (recall memory) +- `source_type = 'archival'`: agent-inserted long-term memories (archival memory) +- `source_type = 'document'`: user-uploaded document chunks (document RAG) + +- Future source types added without migration + +Key columns for the agent system: `embedding_model` (tracks which model produced the vector, enables re-embedding on model upgrade), `content_enc` (always present -- the embedded text itself, encrypted), and `tags_enc` (deterministically-encrypted, base64 tags enabling SQL-indexable filtering for archival memories; extracted from `metadata.tags` when present). + +**Future:** add `chunk_index` (ordering within multi-chunk sources like documents). + +**Future (multi-agent/subagents):** consider adding an optional `agent_id` column for `source_type='archival'` (and possibly `source_type='document'`) rows so archival/document memory can be shared or isolated per agent. Recall/message embeddings remain scoped by `conversation_id` and can still default to cross-thread search. + +See `potential-rag-integration-brute-force.md` for full schema and API design. + +### 4.3 `conversation_summaries` -- Compaction Artifacts + +Stores rolling summaries when conversation context exceeds limits. Modeled on Sage's `summaries` table, adapted for Maple's conversation model. + +```sql +CREATE TABLE conversation_summaries ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + + -- Range of messages this summary covers + -- Uses message created_at timestamps as boundaries + from_created_at TIMESTAMPTZ NOT NULL, + to_created_at TIMESTAMPTZ NOT NULL, + message_count INTEGER NOT NULL, -- how many messages were summarized + + -- Encrypted summary content + content_enc BYTEA NOT NULL, -- AES-256-GCM encrypted summary text + content_tokens INTEGER NOT NULL, -- plaintext token count of summary + + -- Encrypted embedding for semantic search over summaries + embedding_enc BYTEA, -- AES-256-GCM encrypted float32 array + + -- Chain: previous summary that was absorbed into this one + previous_summary_id BIGINT REFERENCES conversation_summaries(id) ON DELETE SET NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + CONSTRAINT valid_time_range CHECK (from_created_at <= to_created_at) +); + +CREATE INDEX idx_conversation_summaries_user_conv + ON conversation_summaries(user_id, conversation_id, created_at DESC); + +CREATE INDEX idx_conversation_summaries_chain + ON conversation_summaries(previous_summary_id); +``` + +**Key differences from Sage:** +- Uses `from_created_at` / `to_created_at` timestamp ranges instead of Sage's `from_sequence_id` / `to_sequence_id`. Maple's message tables don't have a monotonic sequence_id across the UNION -- they have per-table auto-increment IDs. Timestamps are the reliable ordering key that works across all message types. +- `embedding_enc` instead of pgvector `embedding` column. Follows the same encrypted pattern as everything else. +- `content_tokens` plaintext for context budgeting. + +### 4.4 `agent_config` -- Per-User Agent Settings + +```sql +CREATE TABLE agent_config ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE UNIQUE, + + -- Which conversation thread is the user's "main agent" thread + -- Nullable until agent initialization creates/assigns a thread. + conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, + + -- Agent behavior settings (plaintext, not sensitive) + enabled BOOLEAN NOT NULL DEFAULT FALSE, + model TEXT NOT NULL DEFAULT 'deepseek-r1-0528', + max_context_tokens INTEGER NOT NULL DEFAULT 100000, + compaction_threshold REAL NOT NULL DEFAULT 0.80, + + -- Agent system prompt (encrypted, user-customizable) + system_prompt_enc BYTEA, + + -- User preferences (encrypted JSON: timezone, response style, etc.) + preferences_enc BYTEA, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE TRIGGER update_agent_config_updated_at +BEFORE UPDATE ON agent_config +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**Notes:** +- One agent config per user (MVP). The `UNIQUE(user_id)` constraint enforces this. +- **Future (multi-agent/subagents):** prefer introducing an `agents` table and making `agent_config` keyed by `agent_id` (or merging `agent_config` into `agents`). This makes agent identity explicit and enables agent-scoped memory via `agent_id` FKs (Section 4.5). The simpler alternative (remove `UNIQUE(user_id)` + add a `name` column) is viable but makes shared vs isolated memory harder to represent cleanly. +- `conversation_id` is the canonical pointer to the user's persistent agent thread. This avoids relying on encrypted conversation metadata to locate the agent conversation. +- `enabled` allows opt-in activation. Users who don't want agent features continue using the stateless Responses API. +- `system_prompt_enc` is the base agent instruction. Equivalent to Sage's `agents.system_prompt`, but encrypted. If NULL, a default system prompt is used. +- `preferences_enc` absorbs what Sage stores in the separate `user_preferences` table. A single encrypted JSON blob is simpler than a KV table for a small number of preferences (timezone, response style, language). + +### 4.5 Future: `agents` Table + `agent_id` FKs (Multi-Agent/Subagents) + +If we want multiple agents per user (subagents), the most extensible model is to make agent identity explicit. + +**Proposed `agents` table (future):** + +```sql +CREATE TABLE agents ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + -- User-visible identifier (plaintext) + name TEXT NOT NULL DEFAULT 'main', + + -- Each agent has exactly one long-running conversation thread + conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + + enabled BOOLEAN NOT NULL DEFAULT FALSE, + + -- Optional: policy for shared vs isolated memory + -- Examples: 'shared', 'isolated', 'overlay' + memory_mode TEXT NOT NULL DEFAULT 'shared', + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(user_id, name), + UNIQUE(conversation_id) +); + +CREATE TRIGGER update_agents_updated_at +BEFORE UPDATE ON agents +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**Agent-scoped foreign keys (future):** add `agent_id` (nullable or required) to agent-specific tables: +- `memory_blocks.agent_id` (nullable) + - `agent_id IS NULL` => shared blocks for the user + - `agent_id = X` => blocks scoped to agent X +- `conversation_summaries.agent_id` (optional; can be derived from `conversation_id`) +- `user_embeddings.agent_id` (optional; most useful for `source_type='archival'` / `'document'`) +- `scheduled_tasks.agent_id` (for reminders) +- `agent_config.agent_id` (if/when `agent_config` becomes per-agent) + +**Uniqueness note (Postgres):** if we use `memory_blocks.agent_id IS NULL` to represent shared blocks, we likely want **partial unique indexes** to enforce: +- exactly one shared block per `(user_id, label)` and +- exactly one agent-scoped block per `(user_id, agent_id, label)`. + +For MVP, we can start with the simpler per-user `agent_config` model, and only introduce `agents` once subagents are in scope. + +--- + +## 5. Agent Context Assembly + +This is the fundamental difference between the Responses API and the agent system. Instead of the Responses API's approach (load all messages, middle-truncate to fit context window), the agent regenerates a structured context on every turn. + +### 5.1 Sage V2 Prototype Pattern (DSRs Signatures + BAML Parsing) + +The current Sage V2 prototype uses DSRs signatures for typed input/output and BAML parsing for structured extraction, rather than provider-native tool calling. + +Conceptually, the agent is called with a set of **separated context fields** (so GEPA can optimize them independently), and returns **two typed outputs**: `messages[]` and `tool_calls[]`. + +``` +AgentResponse (inputs) + - input (user message OR tool-result continuation) + - current_time + - persona_block + - human_block + - memory_metadata + - previous_context_summary + - recent_conversation + - available_tools + - is_first_time_user + +AgentResponse (outputs) + - messages: string[] + - tool_calls: {name: string, args: map}[] +``` + +``` +[System Prompt] + - Base instruction (GEPA-optimized or user-customized) + - + ... + ... + - + memory_replace, memory_append, archival_insert, archival_search, + conversation_search, web_search, ... + - + archival_count, recall_count, last_compaction, ... + +[Previous Context Summary] (if conversation was compacted) + +[Recent Messages] (last N messages that fit in token budget) + +[Current User Message] +``` + +### 5.2 Maple Agent Context Builder + +A new `src/web/agent/context_builder.rs` that replaces the Responses API's `context_builder.rs` for agent requests: + +1. **Load agent config** -- model, max_context_tokens, compaction_threshold +2. **Build system prompt:** + - Start with base instruction from `agent_config.system_prompt_enc` (or default) + - Decrypt and inject memory blocks from `memory_blocks` table + - Inject tool descriptions (from agent tool registry) + - Inject memory metadata (counts from user_embeddings, memory_blocks, conversation_summaries) +3. **Load summary** -- most recent `conversation_summaries` entry for this conversation, if any +4. **Load recent messages** -- from the existing message tables (same UNION ALL query), but: + - Start from after the summary's `to_created_at` timestamp + - Load messages until token budget is consumed + - No middle-truncation -- if budget exceeded, trigger compaction instead +5. **Check compaction threshold** -- if total tokens > threshold, compact oldest messages into a new summary before proceeding +6. **Assemble final prompt** -- system + summary + recent messages + current input + +In practice, if we follow the Sage V2 prototype, the "assembled prompt" is expressed as the DSRs signature inputs above (plus an instruction string), and BAML parses the structured outputs. + +### 5.3 Compaction vs Truncation + +The Responses API truncates: it drops middle messages and inserts `[Previous messages truncated due to context limits]`. Information is lost. + +The agent system compacts: it summarizes old messages into a `conversation_summaries` entry using an LLM call, then removes those messages from the in-context window (but they remain in the database and are searchable via recall memory). Information is preserved in compressed form. + +Compaction uses a DSRs-style signature (following Sage's `SummarizeConversation` pattern): + +``` +Input: + - previous_summary (if chaining summaries) + - messages_to_summarize +Output: + - summary (100 word limit) +``` + +The LLM call for summarization uses a fast/cheap model (e.g., `gpt-oss-120b`) to minimize latency and cost. + +--- + +## 6. Agent Step Loop + +The agent processes messages through a multi-step loop, following Sage's proven pattern: + +``` +User sends message + -> Step 1: Build context, call LLM + -> LLM returns: messages to user + tool calls + -> Execute tool calls (memory_replace, archival_insert, web_search, etc.) + -> If tools were called: inject results, go to Step 2 + -> Step 2: Build context (with tool results), call LLM + -> LLM returns: messages to user + tool calls + -> ... repeat until LLM returns messages with no tool calls, or max steps reached + -> Send final messages to user +``` + +**Max steps:** 10 (same as Sage). Prevents infinite loops. + +**Tool execution:** Each tool call is persisted to the existing `tool_calls` and `tool_outputs` tables. Memory tools (`memory_replace`, `archival_insert`, etc.) also modify the memory tables directly. + +### 6.2 Structured Output + Parse Correction (Recommended) + +To maximize reliability across providers/models (and to avoid native tool calling issues), align with the Sage V2 prototype: + +- The agent returns structured `messages[]` and `tool_calls[]` via DSRs signatures + BAML parsing. +- If the LLM returns malformed output that fails to parse, run a **correction signature** that reshapes the raw text into the expected structure (preserving intent, not generating new content). +- Include a `done` tool as a no-op stop signal for tool-result continuation steps. + +This pattern dramatically reduces sensitivity to provider-specific tool calling bugs, while still supporting multi-step tool loops. + +### 6.3 DSRs Integration in OpenSecret (Nitro-safe LLM routing) + +DSRs must not call providers directly. All agent LLM calls (main steps + summarization + correction) must route through OpenSecret's existing LLM pipeline (billing, auth, retries, provider routing): `web::openai::get_chat_completion_response()`. + +Implementation options: +- **Preferred:** implement a custom DSRs LM backend/hook that calls `get_chat_completion_response()` and returns the model text + usage. +- **Fallback:** vendor/fork DSRs at a pinned commit to add the hook (so we control upgrades and avoid billing bypass regressions). + +**Current status (implemented 2026-02-10):** OpenSecret now uses a locally-patched `dspy-rs` (DSRs) with a `LMClient::Custom(CustomCompletionModel)` hook that routes completions through `get_chat_completion_response()` (see `src/web/agent/signatures.rs`). Default `temperature=0.7` and `max_tokens=32768` match the Sage prototype. + +**Billing note:** embedding calls used by the agent system (archival insert/search, auto-embedding of agent messages, summary embeddings) must also route through OpenSecret's billing-aware embedding pipeline (`web::get_embedding_vector()`), and must never call providers directly. + +### 6.1 Memory Tools + +These are the agent's interface to the memory system. Following Sage's design: + +| Tool | Action | Storage | +|---|---|---| +| `memory_replace` | Replace text in a core memory block | UPDATE `memory_blocks` | +| `memory_append` | Append text to a core memory block | UPDATE `memory_blocks` | +| `memory_insert` | Insert text at a specific line in a block | UPDATE `memory_blocks` | +| `archival_insert` | Store information in long-term memory | INSERT into `user_embeddings` (source_type='archival') | +| `archival_search` | Search long-term memory semantically (optional tag filter) | Query `user_embeddings` (source_type='archival') with brute-force cosine similarity | +| `conversation_search` | Search conversation history semantically | Query `user_embeddings` (source_type='message') -- **all conversations by default**, not just the agent's own thread | + +All memory tool inputs/outputs are encrypted before storage. The tool execution happens in-enclave where the user key is available. + +**`conversation_search` visibility note:** By default, this tool searches across all of the user's embedded messages, including Responses API threads. This is intentional -- the agent should surface relevant context regardless of which API surface generated it. The tool can accept an optional `conversation_id` parameter to scope to a specific thread if needed, but the default is broad. See the RAG proposal's "Data Visibility and Isolation" section. + +**Archival tags (implemented 2026-02-11):** `archival_insert` supports `metadata.tags` (string or string[]). Tags are normalized (trim + lowercase), deterministically encrypted per-tag (base64) and stored in `user_embeddings.tags_enc`. `archival_search` supports a `tags` argument and applies an ANY-match SQL filter (`tags_enc && `) backed by a partial GIN index. + +--- + +## 7. API Surface + +### 7.1 New Routes (`/v1/agent/*`) + +**Current status (implemented 2026-02-11):** the MVP `/v1/agent/*` surface below is implemented (Local/Dev only, non-streaming). + +``` +POST /v1/agent/chat -- Send a message to the agent (step loop, non-streaming JSON) +GET /v1/agent/config -- Get agent settings +PUT /v1/agent/config -- Update agent settings (model, system prompt, etc.) + +GET /v1/agent/memory/blocks -- List all memory blocks +GET /v1/agent/memory/blocks/:label -- Get a specific block +PUT /v1/agent/memory/blocks/:label -- Manually edit a block + +POST /v1/agent/memory/archival -- Manually insert archival memory +POST /v1/agent/memory/search -- Search archival + recall memory +DELETE /v1/agent/memory/archival/:id -- Delete specific archival entry + +GET /v1/agent/conversations -- List agent conversations (reuses conversations table) +GET /v1/agent/conversations/:id/items -- Get conversation items (reuses existing item format) +DELETE /v1/agent/conversations/:id -- Delete conversation + associated summaries +``` + +### 7.2 Chat Endpoint Detail + +**Current implementation (2026-02-11):** `POST /v1/agent/chat` accepts `{ "input": "..." }` (encrypted request/response like other endpoints) and returns `{ messages: string[], tool_calls: ToolCall[] }`. This endpoint is **non-streaming** in the MVP (no SSE). + +``` +POST /v1/agent/chat +{ + "input": "What did we discuss about deployment strategies last week?" +} +``` + +Response: + +```json +{ + "messages": ["Based on our previous discussion, you mentioned wanting to use blue-green deployments."], + "tool_calls": [] +} +``` + +### 7.3 Relationship to Responses API + +The two API surfaces are independent but share storage: + +``` +/v1/responses/* -- Stateless chat API (existing, unchanged) + reads/writes: conversations, user_messages, assistant_messages, + tool_calls, tool_outputs, reasoning_items, user_instructions + triggers: (TODO) async embedding into user_embeddings (store=true and non-private threads only) + +/v1/agent/* -- Persistent agent API (new) + reads/writes: conversations, user_messages, assistant_messages, + tool_calls, tool_outputs, reasoning_items + also reads/writes: memory_blocks, user_embeddings, + conversation_summaries, agent_config + triggers: async embedding into user_embeddings for agent messages (implemented) + searches: user_embeddings across ALL user conversations (broad default) +``` + +A user can use both APIs simultaneously. The key data flow: + +- **Responses API -> Agent visibility (pending):** Responses API threads will be auto-embedded into `user_embeddings` when eligible (stored + not-private). Today, only the agent's own messages are auto-embedded. +- **Agent conversation visibility:** The agent conversation is stored in the shared tables, but the **main agent thread should be hidden from `/v1/conversations/*`** (Section 2.2). Clients should use `/v1/agent/conversations/*` to access agent threads. +- **Isolation boundary:** `store=false` requests and future incognito/private threads on the Responses API will NOT be embedded. The agent cannot see them. This is the opt-out mechanism. + +--- + +## 8. What We're Not Deciding Yet + +These are intentionally deferred: + +- **Native tool calling in the agent loop.** The Responses API will continue using native tool calling via the Tinfoil proxy. For the agent loop, prefer DSRs signatures + BAML parsing (Section 6.2) for reliability; native tool calling can be revisited later. + +- **GEPA operationalization in production.** MVP should be GEPA-*ready*: consume a GEPA-optimized instruction stored in `agent_config.system_prompt_enc`. Running GEPA itself can be feature-flagged/offline until we have safe evaluation + trace capture. + +- **Multi-agent per user (subagents).** MVP assumes one agent per user. When subagents are in scope, prefer an `agents` table + `agent_id` FKs on agent-specific tables (Section 4.5) so memory can be shared or isolated per agent. The simpler approach (remove `UNIQUE(user_id)` + add `name` on `agent_config`) is viable but less expressive for memory scoping. + +- **Group/shared memory.** All memory is per-user. Shared memory blocks between users in the same project would require a new sharing model. Deferred. + +- **Voice/image input handling.** Sage has a vision pipeline for Signal image attachments. Maple's Responses API already handles multimodal input. The agent system inherits this capability from the shared message tables. + +- **Reminders / scheduling.** Post-MVP. Sage V2 has a working scheduler + tools (`schedule_task`, `list_schedules`, `cancel_schedule`) backed by a `scheduled_tasks` table. We can port this once the core agent loop is stable. + +- **Code sandbox execution.** Long-lead, deferred. Requires a secure execution substrate (likely isolated runtime) and careful Nitro threat modeling. + +- **Incognito / private threads.** A `private` flag on Responses API conversations that excludes their messages from embedding (and thus from agent visibility). In addition, `store=false` should suppress embedding as a per-request opt-out. The conversation-level `private` flag is deferred to post-v1. A UI toggle to set incognito as the default for new threads is further out. + +- **Agent-backed Responses API threads.** The long-term vision where Responses API threads inherit agent memory (memory blocks + search) but start with a clean conversation history. Requires proving out the agent system first. See Section 1.3. + +--- + +## 9. Implementation Ordering + +A suggested sequence, building on the RAG layer as the foundation: + +### Phase 1: RAG Foundation +- [x] Implement `user_embeddings` table + migration +- [x] Implement brute-force search + LRU cache +- [x] Implement `/v1/rag/*` API endpoints +- [ ] Wire up async embedding generation after message creation in Responses API (respect `store=false` and future `private` threads) +- **Milestone:** Recall + archival memory storage/search exists (even before the agent loop ships) + +### Phase 2: Agent Storage + Feature Flags +- [x] Implement `memory_blocks`, `conversation_summaries`, `agent_config` tables + migrations +- [x] Implement Diesel models for new tables +- [x] Add `agent_config.conversation_id` pointer to the user's persistent agent thread +- [x] Implement memory block CRUD operations +- [ ] Implement global + per-user feature flagging for `/v1/agent/*` and background jobs +- **Milestone:** Agent storage layer exists and is tested + +### Phase 3: Agent LLM Runtime (DSRs + BAML) +- [x] Integrate DSRs signatures + BAML parsing (AgentResponse) +- [x] Integrate correction signatures for parse repair +- [x] Implement Nitro-safe DSRs LM routing via `get_chat_completion_response()` +- [x] Implement summarization signatures for compaction +- **Milestone:** Agent can reliably produce `messages[]` + `tool_calls[]` without native tool calling + +### Phase 4: Agent Context Builder + Compaction +- [x] Implement regenerated agent context builder (currently in `src/web/agent/runtime.rs`) +- [x] Implement compaction (LLM summarization of old messages) into `conversation_summaries` +- **Milestone:** Agent can sustain long-running conversations without truncation + +### Phase 5: Memory Tools + Multi-Step Agent Loop +- [x] Implement agent tool registry (core memory + archival + conversation search + web search) +- [x] Implement multi-step execution loop (max steps, tool-result continuation) +- **Milestone:** DSRs-based agent can have multi-turn conversations with memory tool use + +### Phase 6: Agent API Surface +- [x] Implement `/v1/agent/chat` (non-streaming) +- [x] Implement remaining `/v1/agent/*` endpoints (config, memory management, conversations) +- [x] Wire up memory block initialization + agent thread creation on opt-in +- [x] Wire up async embedding generation after agent message creation +- **Milestone:** Full agent API surface available + +### Phase 7: Post-MVP +- Reminders/scheduling (`scheduled_tasks` + scheduler loop + tools) +- Monitoring/metrics + tuning +- (Optional) agent-backed Responses API threads (Section 1.3) +- Code sandboxes (separate major project) diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 9a969f69..2515eb42 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,14 +1,16 @@ use axum::{ extract::{Path, Query, State}, middleware::from_fn_with_state, + response::sse::{Event, Sse}, routing::{delete, get, post, put}, Extension, Json, Router, }; use diesel::prelude::*; +use futures::Stream; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::sync::Arc; -use tracing::error; +use tracing::{error, warn}; use uuid::Uuid; use crate::encrypt::{decrypt_content, decrypt_string, encrypt_with_key}; @@ -24,6 +26,7 @@ use crate::web::responses::constants::{DEFAULT_PAGINATION_LIMIT, MAX_PAGINATION_ use crate::web::responses::conversations::{ ConversationItemListResponse, ConversationListResponse, ConversationResponse, ListItemsParams, }; +use crate::web::responses::handlers::encrypt_event; use crate::web::responses::{ error_mapping, ConversationBuilder, ConversationItem, ConversationItemConverter, DeletedObjectResponse, Paginator, @@ -40,10 +43,26 @@ struct AgentChatRequest { input: String, } +// SSE event types for agent chat (message-level delivery, not token streaming) +const EVENT_AGENT_MESSAGE: &str = "agent.message"; +const EVENT_AGENT_DONE: &str = "agent.done"; +const EVENT_AGENT_ERROR: &str = "agent.error"; + #[derive(Debug, Clone, Serialize)] -struct AgentChatResponse { +struct AgentMessageEvent { messages: Vec, - tool_calls: Vec, + step: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentDoneEvent { + total_steps: usize, + total_messages: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct AgentErrorEvent { + error: String, } #[derive(Debug, Clone, Serialize)] @@ -195,7 +214,7 @@ async fn chat( Extension(session_id): Extension, Extension(user): Extension, Extension(body): Extension, -) -> Result>, ApiError> { +) -> Result>>, ApiError> { if body.input.trim().is_empty() { return Err(ApiError::BadRequest); } @@ -206,14 +225,96 @@ async fn chat( .map_err(|_| error_mapping::map_key_retrieval_error())?; let mut runtime = runtime::AgentRuntime::new(state.clone(), user.clone(), user_key).await?; - let (messages, tool_calls) = runtime.process_message(&body.input).await?; + runtime.prepare(&body.input).await?; + + let input = body.input.clone(); + let max_steps = runtime.max_steps(); + + let event_stream = async_stream::stream! { + let mut total_messages: usize = 0; + let mut had_error = false; + + for step_num in 0..max_steps { + let step_result = runtime.step(&input, step_num == 0).await; + + match step_result { + Ok(result) => { + // Persist assistant messages SYNCHRONOUSLY (so next step sees them). + // Embedding updates happen async inside insert_assistant_message. + for msg in &result.messages { + if let Err(e) = runtime.insert_assistant_message(msg).await { + error!("Failed to persist assistant message: {e:?}"); + } + } + + // Emit messages to client immediately (like Sage sending via Signal) + if !result.messages.is_empty() { + total_messages += result.messages.len(); + let event_data = AgentMessageEvent { + messages: result.messages, + step: step_num, + }; + match encrypt_event(&state, &session_id, EVENT_AGENT_MESSAGE, &serde_json::to_value(&event_data).unwrap_or_default()).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; + } + } + } + + // Persist tool calls SYNCHRONOUSLY (so next step sees them in context) + 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 tool call: {e:?}"); + } + } + + if result.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_event(&state, &session_id, EVENT_AGENT_ERROR, &serde_json::to_value(&err_event).unwrap_or_default()).await { + Ok(event) => yield Ok(event), + Err(_) => yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")), + } + had_error = true; + break; + } + } + } + + 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_event(&state, &session_id, EVENT_AGENT_MESSAGE, &serde_json::to_value(&fallback).unwrap_or_default()).await { + yield Ok(event); + total_messages = 1; + } + } - let response = AgentChatResponse { - messages, - tool_calls, + // Final done event + let done_event = AgentDoneEvent { + total_steps: max_steps, + total_messages, + }; + match encrypt_event(&state, &session_id, EVENT_AGENT_DONE, &serde_json::to_value(&done_event).unwrap_or_default()).await { + Ok(event) => yield Ok(event), + Err(_) => yield Ok(Event::default().data("[DONE]")), + } }; - encrypt_response(&state, &session_id, &response).await + Ok(Sse::new(event_stream)) } async fn get_config( diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index db74908e..65a6ab75 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -54,10 +54,18 @@ struct AgentContext { } #[derive(Clone, Debug)] -struct StepResult { - messages: Vec, - tool_calls: Vec, - done: bool, +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 { @@ -299,47 +307,27 @@ impl AgentRuntime { self.previous_step_summary = None; } - pub async fn process_message( - &mut self, - user_message: &str, - ) -> Result<(Vec, Vec), ApiError> { + 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: &str) -> Result<(), ApiError> { let trimmed = user_message.trim(); if trimmed.is_empty() { return Err(ApiError::BadRequest); } - // Clear at start of request self.clear_tool_results(); - - // Persist user message (and recall embedding) self.insert_user_message(trimmed).await?; - - // Compaction (summary memory) if needed self.maybe_compact().await?; - - let mut all_messages: Vec = Vec::new(); - let mut last_tool_calls: Vec = Vec::new(); - - for step_num in 0..self.max_steps { - let result = self.step(trimmed, step_num == 0).await?; - - all_messages.extend(result.messages); - last_tool_calls = result.tool_calls; - - if result.done { - break; - } - } - - if all_messages.is_empty() { - warn!("Agent produced no messages"); - all_messages.push("I apologize, but I wasn't able to generate a response.".to_string()); - } - - Ok((all_messages, last_tool_calls)) + Ok(()) } - async fn step( + /// 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, @@ -470,11 +458,10 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .filter(|m| !m.is_empty()) .collect(); - // Persist assistant messages (before tool execution to preserve ordering) - for msg in &messages { - self.insert_assistant_message(msg).await?; - } - + // 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()) @@ -484,12 +471,13 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If ToolResult::error(format!("Unknown tool: {}", tool_call.name)) }; - // Inject tool result for next step self.inject_tool_result(tool_call, &result); - // Persist tool call + output (skip done) if tool_call.name != "done" { - self.insert_tool_call_and_output(tool_call, &result).await?; + executed_tools.push(ExecutedTool { + tool_call: tool_call.clone(), + result, + }); } } @@ -508,6 +496,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(StepResult { messages, tool_calls: response.tool_calls, + executed_tools, done, }) } @@ -730,7 +719,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(content) } - async fn insert_user_message(&self, text: &str) -> Result { + pub async fn insert_user_message(&self, text: &str) -> Result { let content = super::tools::normalize_user_message_content(text); let prompt_tokens = super::tools::count_user_message_tokens(&content); let content_enc = encrypt_with_key(&self.user_key, content.as_bytes()).await; @@ -781,7 +770,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(msg) } - async fn insert_assistant_message(&self, text: &str) -> Result { + 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); @@ -833,7 +822,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(msg) } - async fn insert_tool_call_and_output( + pub async fn insert_tool_call_and_output( &self, tool_call: &AgentToolCall, result: &ToolResult, From 01a60e41646ba7ffc1cc49aa2e3bb4e5bfe95b22 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 12:57:49 -0600 Subject: [PATCH 09/34] fix: align agent defaults with Sage (kimi-k2, 256k context) Sage uses kimi-k2 as the default model with a 256k context window. The agent config was incorrectly defaulting to deepseek-r1-0528 with 100k context. Fixed in both the runtime code and the migration SQL. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql | 4 ++-- src/web/agent/runtime.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql index a39836b7..26e35b69 100644 --- a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql @@ -62,8 +62,8 @@ CREATE TABLE agent_config ( conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, enabled BOOLEAN NOT NULL DEFAULT FALSE, - model TEXT NOT NULL DEFAULT 'deepseek-r1-0528', - max_context_tokens INTEGER NOT NULL DEFAULT 100000, + model TEXT NOT NULL DEFAULT 'kimi-k2', + max_context_tokens INTEGER NOT NULL DEFAULT 256000, compaction_threshold REAL NOT NULL DEFAULT 0.80, system_prompt_enc BYTEA, diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 65a6ab75..4459b5bb 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -38,7 +38,7 @@ use super::tools::{ const DEFAULT_PERSONA_DESCRIPTION: &str = "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions."; const DEFAULT_HUMAN_DESCRIPTION: &str = "The human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation."; 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_CONTEXT_WINDOW: i32 = 100_000; +const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; const MIN_MESSAGES_IN_CONTEXT: usize = 20; @@ -111,7 +111,7 @@ impl AgentRuntime { user_id: user.uuid, conversation_id: None, enabled: true, - model: "deepseek-r1-0528".to_string(), + model: "kimi-k2".to_string(), max_context_tokens: DEFAULT_CONTEXT_WINDOW, compaction_threshold: DEFAULT_COMPACTION_THRESHOLD, system_prompt_enc: None, From 51eff603b737492851c35d5dec18487023e88045 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 11 Feb 2026 13:02:51 -0600 Subject: [PATCH 10/34] fix: deduplicate agent config defaults into shared constants The get_or_create_agent_config helper in mod.rs still had the old deepseek-r1-0528 / 100k defaults. Extracted DEFAULT_MODEL, DEFAULT_CONTEXT_WINDOW, DEFAULT_COMPACTION_THRESHOLD as pub constants in runtime.rs and use them from both creation sites. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql | 2 +- src/web/agent/mod.rs | 6 +++--- src/web/agent/runtime.rs | 7 ++++--- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql index 26e35b69..bd1980ed 100644 --- a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql @@ -62,7 +62,7 @@ CREATE TABLE agent_config ( conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, enabled BOOLEAN NOT NULL DEFAULT FALSE, - model TEXT NOT NULL DEFAULT 'kimi-k2', + model TEXT NOT NULL DEFAULT 'kimi-k2-5', max_context_tokens INTEGER NOT NULL DEFAULT 256000, compaction_threshold REAL NOT NULL DEFAULT 0.80, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 2515eb42..fcba3858 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1009,9 +1009,9 @@ async fn get_or_create_agent_config( user_id, conversation_id: None, enabled: true, - model: "deepseek-r1-0528".to_string(), - max_context_tokens: 100_000, - compaction_threshold: 0.80, + model: runtime::DEFAULT_MODEL.to_string(), + max_context_tokens: runtime::DEFAULT_CONTEXT_WINDOW, + compaction_threshold: runtime::DEFAULT_COMPACTION_THRESHOLD, system_prompt_enc: None, preferences_enc: None, }; diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 4459b5bb..69aad6e4 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -38,8 +38,9 @@ use super::tools::{ const DEFAULT_PERSONA_DESCRIPTION: &str = "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions."; const DEFAULT_HUMAN_DESCRIPTION: &str = "The human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation."; 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_CONTEXT_WINDOW: i32 = 256_000; -const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; +pub const DEFAULT_MODEL: &str = "kimi-k2-5"; +pub const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; +pub const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; const MIN_MESSAGES_IN_CONTEXT: usize = 20; #[derive(Clone, Debug, Default)] @@ -111,7 +112,7 @@ impl AgentRuntime { user_id: user.uuid, conversation_id: None, enabled: true, - model: "kimi-k2".to_string(), + model: DEFAULT_MODEL.to_string(), max_context_tokens: DEFAULT_CONTEXT_WINDOW, compaction_threshold: DEFAULT_COMPACTION_THRESHOLD, system_prompt_enc: None, From 00426bd5a581c322a01d0af200304a053fab9937 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 23 Feb 2026 14:21:41 -0600 Subject: [PATCH 11/34] feat: stream agent typing indicators between messages Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/web/agent/mod.rs | 232 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 196 insertions(+), 36 deletions(-) diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index fcba3858..3c247cbc 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -9,7 +9,8 @@ use diesel::prelude::*; use futures::Stream; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; +use tokio::time::sleep; use tracing::{error, warn}; use uuid::Uuid; @@ -45,15 +46,23 @@ struct AgentChatRequest { // 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, @@ -65,6 +74,20 @@ 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 +} + #[derive(Debug, Clone, Serialize)] struct AgentConfigResponse { enabled: bool, @@ -219,60 +242,192 @@ async fn chat( return Err(ApiError::BadRequest); } - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut runtime = runtime::AgentRuntime::new(state.clone(), user.clone(), user_key).await?; - runtime.prepare(&body.input).await?; - + // 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 input = body.input.clone(); - let max_steps = runtime.max_steps(); let event_stream = async_stream::stream! { + let mut max_steps: usize = 0; let mut total_messages: usize = 0; let mut had_error = false; - for step_num in 0..max_steps { + 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.uuid, 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 mut runtime = + match runtime::AgentRuntime::new(state.clone(), user.clone(), user_key).await { + 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; + } + }; + + if let Err(e) = runtime.prepare(&input).await { + 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, 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 &result.messages { + for msg in &messages { if let Err(e) = runtime.insert_assistant_message(msg).await { error!("Failed to persist assistant message: {e:?}"); } } - // Emit messages to client immediately (like Sage sending via Signal) - if !result.messages.is_empty() { - total_messages += result.messages.len(); - let event_data = AgentMessageEvent { - messages: result.messages, - step: step_num, - }; - match encrypt_event(&state, &session_id, EVENT_AGENT_MESSAGE, &serde_json::to_value(&event_data).unwrap_or_default()).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; - } + // 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:?}"); } } - // Persist tool calls SYNCHRONOUSLY (so next step sees them in context) - 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 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 result.done { + if done { break; } } @@ -281,14 +436,17 @@ async fn chat( let err_event = AgentErrorEvent { error: "Agent encountered an error processing your message.".to_string(), }; - match encrypt_event(&state, &session_id, EVENT_AGENT_ERROR, &serde_json::to_value(&err_event).unwrap_or_default()).await { + 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; + break 'steps; } } + } } if !had_error && total_messages == 0 { @@ -297,7 +455,9 @@ async fn chat( messages: vec!["I apologize, but I wasn't able to generate a response.".to_string()], step: 0, }; - if let Ok(event) = encrypt_event(&state, &session_id, EVENT_AGENT_MESSAGE, &serde_json::to_value(&fallback).unwrap_or_default()).await { + if let Ok(event) = + encrypt_agent_event(&state, &session_id, EVENT_AGENT_MESSAGE, &fallback).await + { yield Ok(event); total_messages = 1; } @@ -308,7 +468,7 @@ async fn chat( total_steps: max_steps, total_messages, }; - match encrypt_event(&state, &session_id, EVENT_AGENT_DONE, &serde_json::to_value(&done_event).unwrap_or_default()).await { + 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]")), } From b259ad7755a830241cd98763c03407333d567ab7 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 24 Feb 2026 09:38:39 -0600 Subject: [PATCH 12/34] Updated docs --- internal_docs/sage-in-maple-architecture.md | 55 +++++++++++++++------ 1 file changed, 40 insertions(+), 15 deletions(-) diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md index 1ade2186..10c30f82 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/sage-in-maple-architecture.md @@ -3,7 +3,7 @@ ## Bringing Persistent Agent Memory into OpenSecret **Date:** February 2026 -**Status:** MVP implemented (Local/Dev only, non-streaming; Responses API auto-embedding + production feature flags pending) +**Status:** MVP implemented (Local/Dev only; `/v1/agent/chat` streams message-level SSE; Responses API auto-embedding + production feature flags pending) **Related Docs:** - `potential-rag-integration-brute-force.md` -- RAG/vector storage layer (prerequisite) - `architecture-for-rag-integration.md` -- OpenSecret encryption and data model reference @@ -489,10 +489,10 @@ All memory tool inputs/outputs are encrypted before storage. The tool execution ### 7.1 New Routes (`/v1/agent/*`) -**Current status (implemented 2026-02-11):** the MVP `/v1/agent/*` surface below is implemented (Local/Dev only, non-streaming). +**Current status (implemented 2026-02-11):** the MVP `/v1/agent/*` surface below is implemented (Local/Dev only). `POST /v1/agent/chat` streams message-level SSE (not token streaming); other endpoints are encrypted JSON. ``` -POST /v1/agent/chat -- Send a message to the agent (step loop, non-streaming JSON) +POST /v1/agent/chat -- Send a message to the agent (step loop, request-scoped SSE) GET /v1/agent/config -- Get agent settings PUT /v1/agent/config -- Update agent settings (model, system prompt, etc.) @@ -507,11 +507,22 @@ DELETE /v1/agent/memory/archival/:id -- Delete specific archival entry GET /v1/agent/conversations -- List agent conversations (reuses conversations table) GET /v1/agent/conversations/:id/items -- Get conversation items (reuses existing item format) DELETE /v1/agent/conversations/:id -- Delete conversation + associated summaries + +GET /v1/agent/events -- Long-lived SSE channel for proactive agent delivery + fan-out (post-MVP) ``` ### 7.2 Chat Endpoint Detail -**Current implementation (2026-02-11):** `POST /v1/agent/chat` accepts `{ "input": "..." }` (encrypted request/response like other endpoints) and returns `{ messages: string[], tool_calls: ToolCall[] }`. This endpoint is **non-streaming** in the MVP (no SSE). +**Current implementation (2026-02-11):** `POST /v1/agent/chat` accepts `{ "input": "..." }` (encrypted request body like other endpoints) and returns request-scoped SSE (message-level events, not token streaming). + +Event types (payloads are JSON, encrypted per-session and base64-encoded in the `data:` field): + +| Event type | Meaning | +|---|---| +| `agent.typing` | Agent is working on step `step` | +| `agent.message` | One or more user-visible messages for step `step` | +| `agent.done` | Turn completed (`total_steps`, `total_messages`) | +| `agent.error` | Turn failed (`error`) | ``` POST /v1/agent/chat @@ -520,16 +531,29 @@ POST /v1/agent/chat } ``` -Response: +### 7.3 Proactive Agent Delivery: Long-Lived SSE Channel (Post-MVP) -```json -{ - "messages": ["Based on our previous discussion, you mentioned wanting to use blue-green deployments."], - "tool_calls": [] -} -``` +`POST /v1/agent/chat` uses request-scoped SSE: the stream opens on request, delivers step-loop events (`agent.typing`, `agent.message`, `agent.done`), and closes. This is correct for request-response interactions but insufficient for proactive agent behavior (reminders, scheduled task results, agent-initiated messages). + +**Decision: long-lived SSE over WebSockets.** Three reasons specific to our architecture: + +1. **Encryption model fit.** The per-session encryption middleware operates on HTTP request/response cycles. SSE is still HTTP and slots in naturally. WebSockets would require rethinking per-message encryption (WS frames aren't HTTP responses, so `encrypt_event` / session key lookup needs a different path). +2. **Proxy chain simplicity.** tinfoil-proxy and continuum-proxy sit in front of the enclave via vsock. HTTP/SSE flows through standard reverse proxies. WebSocket upgrade handling through that chain adds operational complexity -- sticky sessions, connection draining, timeout tuning at every hop. +3. **Unidirectionality is sufficient.** Proactive agent messages are server-to-client. The client already has `POST /v1/agent/chat` for the other direction. + +**Proposed endpoint:** `GET /v1/agent/events` + +The client opens this connection once and keeps it open. The server pushes **the same `agent.*` SSE events used by `/v1/agent/chat`** (`agent.typing`, `agent.message`, `agent.done`, `agent.error`) whenever the agent produces output outside of an active chat request. + +**Resilience requirements:** +- **At-least-once delivery.** Clients MUST dedupe by SSE event `id`. +- **SSE `id` support + `Last-Event-ID` resumption.** Every non-heartbeat event MUST include an SSE `id:` field. On reconnect, clients send `Last-Event-ID` to resume from the next event. +- **DB-backed event log (fan-out).** Events are persisted so clients that were offline can catch up on reconnect. Because this is fan-out, events are retained by TTL (e.g., 24h) rather than deleted-on-delivery. +- **Heartbeat keepalive.** Emit SSE comment frames (e.g., `: ping\n\n`) every ~30s to prevent proxy idle timeouts. + +**Relationship to `POST /v1/agent/chat`:** The two channels are independent. Chat continues to use request-scoped SSE for immediate step-loop delivery. The long-lived channel handles async/proactive events only. A client that only uses chat (no proactive features) never needs to open `/v1/agent/events`. -### 7.3 Relationship to Responses API +### 7.4 Relationship to Responses API The two API surfaces are independent but share storage: @@ -570,7 +594,7 @@ These are intentionally deferred: - **Voice/image input handling.** Sage has a vision pipeline for Signal image attachments. Maple's Responses API already handles multimodal input. The agent system inherits this capability from the shared message tables. -- **Reminders / scheduling.** Post-MVP. Sage V2 has a working scheduler + tools (`schedule_task`, `list_schedules`, `cancel_schedule`) backed by a `scheduled_tasks` table. We can port this once the core agent loop is stable. +- **Reminders / scheduling.** Post-MVP. Sage V2 has a working scheduler + tools (`schedule_task`, `list_schedules`, `cancel_schedule`) backed by a `scheduled_tasks` table. We can port this once the core agent loop is stable. Delivery of fired reminders will use the long-lived SSE channel (`GET /v1/agent/events`, Section 7.3). - **Code sandbox execution.** Long-lead, deferred. Requires a secure execution substrate (likely isolated runtime) and careful Nitro threat modeling. @@ -617,14 +641,15 @@ A suggested sequence, building on the RAG layer as the foundation: - **Milestone:** DSRs-based agent can have multi-turn conversations with memory tool use ### Phase 6: Agent API Surface -- [x] Implement `/v1/agent/chat` (non-streaming) +- [x] Implement `/v1/agent/chat` (request-scoped SSE) - [x] Implement remaining `/v1/agent/*` endpoints (config, memory management, conversations) - [x] Wire up memory block initialization + agent thread creation on opt-in - [x] Wire up async embedding generation after agent message creation - **Milestone:** Full agent API surface available ### Phase 7: Post-MVP -- Reminders/scheduling (`scheduled_tasks` + scheduler loop + tools) +- Long-lived SSE event channel (`GET /v1/agent/events`) -- proactive notification delivery (Section 7.3) +- Reminders/scheduling (`scheduled_tasks` + scheduler loop + tools), delivered via the event channel - Monitoring/metrics + tuning - (Optional) agent-backed Responses API threads (Section 1.3) - Code sandboxes (separate major project) From e441405f4c683ce64ac61debbd312fe4665b7a74 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 24 Feb 2026 11:36:08 -0600 Subject: [PATCH 13/34] docs: annotate sage-in-maple architecture implementation status Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- internal_docs/sage-in-maple-architecture.md | 80 +++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md index 10c30f82..af283082 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/sage-in-maple-architecture.md @@ -10,6 +10,8 @@ - Sage V2 Design Doc (`~/Dev/Personal/sage/docs/SAGE_V2_DESIGN.md`) -- proven prototype - Sage V2 Codebase (`~/Dev/Personal/sage/crates/sage-core`) -- DSRs signatures + BAML parsing + multi-step tool loop prototype +**Overall implementation status:** Partially implemented (MVP shipped behind Local/Dev gating; Responses API auto-embedding + production feature flags not implemented yet). + --- ## 1. Goal @@ -28,6 +30,8 @@ The agent lives in a single long-running conversation thread. Unlike Responses A The Responses API continues to exist alongside the agent. Users can still have one-off stateless threads for quick questions. But the agent is the primary interface for users who want memory and persistence. +**Implementation status:** Complete (agent_config.conversation_id is the canonical pointer to the main agent thread). + ### 1.2 Cross-Thread Memory Visibility **The agent sees everything.** When the agent's `conversation_search` tool fires, it searches across all of the user's embedded messages -- both the agent's own conversation and any Responses API threads. The `user_embeddings` table indexes messages from all conversations by default. @@ -40,6 +44,8 @@ The agent can still scope searches to its own thread via `conversation_id` filte Additionally, requests made with `store=false` should not be embedded/indexed into `user_embeddings` (per-request opt-out), even if the Responses API continues to persist messages today. +**Implementation status:** Partially implemented (cross-thread search works where embeddings exist; Responses API auto-embedding + private/store=false exclusions not implemented yet). + ### 1.3 Future Vision: Agent-Backed Responses API Eventually, the Responses API threads themselves could be backed by the agent's memory system. Instead of stateless threads with no memory, a Responses API thread would: @@ -49,6 +55,8 @@ Eventually, the Responses API threads themselves could be backed by the agent's This would give users the best of both: the lightweight feel of a new thread, with the accumulated knowledge of the persistent agent. This is a post-v1 evolution that requires proving out the agent system first -- particularly that memory blocks and search produce consistently good results without the agent's own conversation context. +**Implementation status:** Not implemented yet. + ### 1.4 Future: Subagents (Multi-Agent per User) Subagents are out of scope for the MVP, but we should keep the storage design compatible with them. @@ -65,6 +73,10 @@ We also want subagents to support **configurable memory scoping**: The cleanest way to support this is to introduce an `agents` table and add an `agent_id` FK to *agent-specific* tables (Section 4.5). Shared vs isolated memory becomes a query policy (and optionally a config field) rather than a schema rewrite. +**Implementation status:** Not implemented yet. + +**Overall implementation status:** Partially implemented (MVP main-agent model shipped; agent-backed Responses API + subagents not implemented yet). + --- ## 2. Key Architectural Decision: Separate API, Shared Storage @@ -75,6 +87,8 @@ The existing Responses API (`/v1/responses/*`, `/v1/conversations/*`, `/v1/instr The existing database tables for conversations and messages are **reused** by the agent system. The agent reads from and writes to the same `conversations`, `user_messages`, `assistant_messages`, `tool_calls`, `tool_outputs`, and `reasoning_items` tables. The encrypted content format is identical. +**Implementation status:** Complete (agent reuses existing message tables and encryption format). + ### What's new A new `/v1/agent/*` API surface with its own handler module (`src/web/agent/`), its own step loop, and its own context assembly logic. This API implements the Sage-style regenerated-context pattern instead of the Responses API's middle-truncation pattern. @@ -85,6 +99,8 @@ New database tables for agent-specific concerns (detailed in Section 4): - `conversation_summaries` -- compaction artifacts - `agent_config` -- per-user agent settings +**Implementation status:** Complete (MVP /v1/agent surface + new tables + RAG foundation implemented; Local/Dev gated). + ### 2.1 Feature Flags and Rollout Safety (MVP requirement) Sage is a large feature that touches storage, tool execution, and LLM calls. We need a lightweight, production-friendly feature flag system to safely ship in enclaves. @@ -99,6 +115,8 @@ Flags should also gate sub-features independently (even if `agent_config.enabled - compaction/summarization - GEPA runs (optimizer execution) vs simply *consuming* an optimized instruction +**Implementation status:** Not implemented yet (only Local/Dev AppMode gating today; os_flags client exists but is unused for agent rollout). + ### 2.2 Conversations API Isolation (Hide the Main Agent Thread) Even though the agent's persistent thread is stored in the shared `conversations`/message tables, the **main agent conversation should be treated as an internal implementation detail** and **excluded from the public Conversations API** (`/v1/conversations/*`). @@ -115,6 +133,10 @@ Rationale: - `DELETE /v1/conversations` (delete-all) excludes the agent conversation via `id.ne(agent_conversation_id)`. - `POST /v1/conversations/batch-delete` returns `not_found` for the agent thread instead of deleting it. +**Implementation status:** Complete. + +**Overall implementation status:** Partially implemented (API split + conversation isolation complete; feature flags not implemented yet). + --- ## 3. Why Reuse the Message Tables @@ -137,6 +159,8 @@ This was the hardest decision. We did a column-by-column comparison between Sage **Potential future addition:** An `is_in_context` boolean or a `summary_id` FK on messages would be an optimization for fast filtering. But it's not required -- the summaries table's ranges can determine this. If needed, it would be an additive nullable column with a migration default, not a structural change. +**Overall implementation status:** Complete (agent uses existing message tables via RawThreadMessage UNION queries; no message-table schema changes required). + --- ## 4. New Database Tables @@ -180,6 +204,8 @@ FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); - `persona`: "I am a helpful AI assistant." (editable by agent) - `human`: "" (populated by agent as it learns about the user) +**Implementation status:** Complete (table + models + CRUD endpoints + tools; code default char_limit=20000). + ### 4.2 `user_embeddings` -- General-Purpose Embedding Store Defined in the RAG proposal. A single table that serves all embedding use cases through a `source_type` discriminator: @@ -198,6 +224,8 @@ Key columns for the agent system: `embedding_model` (tracks which model produced See `potential-rag-integration-brute-force.md` for full schema and API design. +**Implementation status:** Partially implemented (table + in-process search/cache + archival tags implemented; Responses API auto-embedding not implemented yet). + ### 4.3 `conversation_summaries` -- Compaction Artifacts Stores rolling summaries when conversation context exceeds limits. Modeled on Sage's `summaries` table, adapted for Maple's conversation model. @@ -242,6 +270,8 @@ CREATE INDEX idx_conversation_summaries_chain - `embedding_enc` instead of pgvector `embedding` column. Follows the same encrypted pattern as everything else. - `content_tokens` plaintext for context budgeting. +**Implementation status:** Complete (compaction writes summaries + embeds them for search). + ### 4.4 `agent_config` -- Per-User Agent Settings ```sql @@ -283,6 +313,8 @@ FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); - `system_prompt_enc` is the base agent instruction. Equivalent to Sage's `agents.system_prompt`, but encrypted. If NULL, a default system prompt is used. - `preferences_enc` absorbs what Sage stores in the separate `user_preferences` table. A single encrypted JSON blob is simpler than a KV table for a small number of preferences (timezone, response style, language). +**Implementation status:** Partially implemented (config + conversation_id + system_prompt implemented; preferences + robust opt-in/flags not implemented yet). + ### 4.5 Future: `agents` Table + `agent_id` FKs (Multi-Agent/Subagents) If we want multiple agents per user (subagents), the most extensible model is to make agent identity explicit. @@ -334,6 +366,10 @@ FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); For MVP, we can start with the simpler per-user `agent_config` model, and only introduce `agents` once subagents are in scope. +**Implementation status:** Not implemented yet. + +**Overall implementation status:** Partially implemented (MVP agent tables shipped; multi-agent/subagent schema not implemented yet). + --- ## 5. Agent Context Assembly @@ -382,6 +418,8 @@ AgentResponse (outputs) [Current User Message] ``` +**Implementation status:** Complete (DSRs-style signature inputs/outputs implemented via dspy-rs + structured parsing). + ### 5.2 Maple Agent Context Builder A new `src/web/agent/context_builder.rs` that replaces the Responses API's `context_builder.rs` for agent requests: @@ -402,6 +440,8 @@ A new `src/web/agent/context_builder.rs` that replaces the Responses API's `cont In practice, if we follow the Sage V2 prototype, the "assembled prompt" is expressed as the DSRs signature inputs above (plus an instruction string), and BAML parses the structured outputs. +**Implementation status:** Partially implemented (context assembly implemented in `src/web/agent/runtime.rs`; does not strictly pack messages to a token budget). + ### 5.3 Compaction vs Truncation The Responses API truncates: it drops middle messages and inserts `[Previous messages truncated due to context limits]`. Information is lost. @@ -420,6 +460,10 @@ Output: The LLM call for summarization uses a fast/cheap model (e.g., `gpt-oss-120b`) to minimize latency and cost. +**Implementation status:** Partially implemented (compaction + summaries implemented; summarization currently uses the agent model, not a dedicated cheap model). + +**Overall implementation status:** Partially implemented (regenerated context + compaction shipped; strict token-budget packing + cheap summarizer model not implemented yet). + --- ## 6. Agent Step Loop @@ -452,6 +496,8 @@ To maximize reliability across providers/models (and to avoid native tool callin This pattern dramatically reduces sensitivity to provider-specific tool calling bugs, while still supporting multi-step tool loops. +**Implementation status:** Complete (AgentResponse + correction signatures + done tool implemented). + ### 6.3 DSRs Integration in OpenSecret (Nitro-safe LLM routing) DSRs must not call providers directly. All agent LLM calls (main steps + summarization + correction) must route through OpenSecret's existing LLM pipeline (billing, auth, retries, provider routing): `web::openai::get_chat_completion_response()`. @@ -464,6 +510,8 @@ Implementation options: **Billing note:** embedding calls used by the agent system (archival insert/search, auto-embedding of agent messages, summary embeddings) must also route through OpenSecret's billing-aware embedding pipeline (`web::get_embedding_vector()`), and must never call providers directly. +**Implementation status:** Complete (DSRs LM hook routes via `get_chat_completion_response()`; embeddings route via billing-aware `get_embedding_vector()`). + ### 6.1 Memory Tools These are the agent's interface to the memory system. Following Sage's design: @@ -483,6 +531,10 @@ All memory tool inputs/outputs are encrypted before storage. The tool execution **Archival tags (implemented 2026-02-11):** `archival_insert` supports `metadata.tags` (string or string[]). Tags are normalized (trim + lowercase), deterministically encrypted per-tag (base64) and stored in `user_embeddings.tags_enc`. `archival_search` supports a `tags` argument and applies an ANY-match SQL filter (`tags_enc && `) backed by a partial GIN index. +**Implementation status:** Partially implemented (core memory + archival + conversation_search implemented; web_search tool not implemented yet). + +**Overall implementation status:** Partially implemented (step loop + tool persistence implemented; web_search + other post-MVP tools pending). + --- ## 7. API Surface @@ -511,6 +563,8 @@ DELETE /v1/agent/conversations/:id -- Delete conversation + associated summar GET /v1/agent/events -- Long-lived SSE channel for proactive agent delivery + fan-out (post-MVP) ``` +**Implementation status:** Partially implemented (all MVP routes except `GET /v1/agent/events`; Local/Dev only). + ### 7.2 Chat Endpoint Detail **Current implementation (2026-02-11):** `POST /v1/agent/chat` accepts `{ "input": "..." }` (encrypted request body like other endpoints) and returns request-scoped SSE (message-level events, not token streaming). @@ -531,6 +585,8 @@ POST /v1/agent/chat } ``` +**Implementation status:** Complete (request-scoped, message-level SSE with typing/message/done/error events). + ### 7.3 Proactive Agent Delivery: Long-Lived SSE Channel (Post-MVP) `POST /v1/agent/chat` uses request-scoped SSE: the stream opens on request, delivers step-loop events (`agent.typing`, `agent.message`, `agent.done`), and closes. This is correct for request-response interactions but insufficient for proactive agent behavior (reminders, scheduled task results, agent-initiated messages). @@ -553,6 +609,8 @@ The client opens this connection once and keeps it open. The server pushes **the **Relationship to `POST /v1/agent/chat`:** The two channels are independent. Chat continues to use request-scoped SSE for immediate step-loop delivery. The long-lived channel handles async/proactive events only. A client that only uses chat (no proactive features) never needs to open `/v1/agent/events`. +**Implementation status:** Not implemented yet. + ### 7.4 Relationship to Responses API The two API surfaces are independent but share storage: @@ -578,6 +636,10 @@ A user can use both APIs simultaneously. The key data flow: - **Agent conversation visibility:** The agent conversation is stored in the shared tables, but the **main agent thread should be hidden from `/v1/conversations/*`** (Section 2.2). Clients should use `/v1/agent/conversations/*` to access agent threads. - **Isolation boundary:** `store=false` requests and future incognito/private threads on the Responses API will NOT be embedded. The agent cannot see them. This is the opt-out mechanism. +**Implementation status:** Partially implemented (agent messages auto-embedded; Responses API auto-embedding + private/store=false exclusions not implemented yet). + +**Overall implementation status:** Partially implemented (MVP /v1/agent surface shipped Local/Dev; proactive events channel pending). + --- ## 8. What We're Not Deciding Yet @@ -602,6 +664,8 @@ These are intentionally deferred: - **Agent-backed Responses API threads.** The long-term vision where Responses API threads inherit agent memory (memory blocks + search) but start with a clean conversation history. Requires proving out the agent system first. See Section 1.3. +**Overall implementation status:** Not implemented yet (intentionally deferred). + --- ## 9. Implementation Ordering @@ -615,6 +679,8 @@ A suggested sequence, building on the RAG layer as the foundation: - [ ] Wire up async embedding generation after message creation in Responses API (respect `store=false` and future `private` threads) - **Milestone:** Recall + archival memory storage/search exists (even before the agent loop ships) +**Implementation status:** Partially implemented (RAG + /v1/rag done; Responses API auto-embedding not implemented yet). + ### Phase 2: Agent Storage + Feature Flags - [x] Implement `memory_blocks`, `conversation_summaries`, `agent_config` tables + migrations - [x] Implement Diesel models for new tables @@ -623,6 +689,8 @@ A suggested sequence, building on the RAG layer as the foundation: - [ ] Implement global + per-user feature flagging for `/v1/agent/*` and background jobs - **Milestone:** Agent storage layer exists and is tested +**Implementation status:** Partially implemented (storage tables/models done; feature flags not implemented yet). + ### Phase 3: Agent LLM Runtime (DSRs + BAML) - [x] Integrate DSRs signatures + BAML parsing (AgentResponse) - [x] Integrate correction signatures for parse repair @@ -630,16 +698,22 @@ A suggested sequence, building on the RAG layer as the foundation: - [x] Implement summarization signatures for compaction - **Milestone:** Agent can reliably produce `messages[]` + `tool_calls[]` without native tool calling +**Implementation status:** Complete. + ### Phase 4: Agent Context Builder + Compaction - [x] Implement regenerated agent context builder (currently in `src/web/agent/runtime.rs`) - [x] Implement compaction (LLM summarization of old messages) into `conversation_summaries` - **Milestone:** Agent can sustain long-running conversations without truncation +**Implementation status:** Complete. + ### Phase 5: Memory Tools + Multi-Step Agent Loop - [x] Implement agent tool registry (core memory + archival + conversation search + web search) - [x] Implement multi-step execution loop (max steps, tool-result continuation) - **Milestone:** DSRs-based agent can have multi-turn conversations with memory tool use +**Implementation status:** Partially implemented (memory tools + step loop done; web_search tool not implemented yet). + ### Phase 6: Agent API Surface - [x] Implement `/v1/agent/chat` (request-scoped SSE) - [x] Implement remaining `/v1/agent/*` endpoints (config, memory management, conversations) @@ -647,9 +721,15 @@ A suggested sequence, building on the RAG layer as the foundation: - [x] Wire up async embedding generation after agent message creation - **Milestone:** Full agent API surface available +**Implementation status:** Complete (Local/Dev only). + ### Phase 7: Post-MVP - Long-lived SSE event channel (`GET /v1/agent/events`) -- proactive notification delivery (Section 7.3) - Reminders/scheduling (`scheduled_tasks` + scheduler loop + tools), delivered via the event channel - Monitoring/metrics + tuning - (Optional) agent-backed Responses API threads (Section 1.3) - Code sandboxes (separate major project) + +**Implementation status:** Not implemented yet. + +**Overall implementation status:** Partially implemented (Phases 1-6 mostly complete; Phase 7 items pending). From f04a15717f8083706dfaa19ca23dd2d83bffa93a Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 24 Feb 2026 17:45:28 -0600 Subject: [PATCH 14/34] feat: sage-style vision pre-processing for agent images Adds encrypted attachment_text storage and a vision pre-step for input_image so agents and embeddings see [Uploaded Image: ...] descriptions. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- internal_docs/sage-in-maple-architecture.md | 31 +++- .../down.sql | 2 + .../up.sql | 7 + src/models/responses.rs | 25 +++ src/models/schema.rs | 1 + src/web/agent/mod.rs | 69 +++++-- src/web/agent/runtime.rs | 168 ++++++++++++++++-- src/web/agent/tools.rs | 26 --- src/web/agent/vision.rs | 115 ++++++++++++ src/web/responses/handlers.rs | 1 + 10 files changed, 389 insertions(+), 56 deletions(-) create mode 100644 migrations/2026-02-24-150000_user_messages_attachment_text/down.sql create mode 100644 migrations/2026-02-24-150000_user_messages_attachment_text/up.sql create mode 100644 src/web/agent/vision.rs diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md index af283082..4da9aa2d 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/sage-in-maple-architecture.md @@ -464,6 +464,35 @@ The LLM call for summarization uses a fast/cheap model (e.g., `gpt-oss-120b`) to **Overall implementation status:** Partially implemented (regenerated context + compaction shipped; strict token-budget packing + cheap summarizer model not implemented yet). +### 5.4 Vision / Image Handling (How Sage Does It) + +Sage V2 does **not** pass images through the DSRs signature/tool loop. Instead, it runs a **vision pre-processing** step that converts an image attachment into a **text description**, then injects that text into the conversation so the rest of the system stays text-only. + +**End-to-end flow (Sage codebase):** + +1. **Signal ingestion** parses `attachments[]` from `signal-cli` JSON-RPC (`crates/sage-core/src/signal.rs::parse_incoming_message`). Each attachment includes a MIME type (`contentType`) and an ID/filename (`id`). +2. **Pre-processing trigger** in the main receive loop (`crates/sage-core/src/main.rs`) checks for the first supported image attachment (`crates/sage-core/src/vision.rs::is_supported_image`, currently `image/jpeg`, `image/png`, `image/webp`, `image/gif`). +3. **Load image bytes** from the `signal-cli` attachments directory mounted into the Sage container (default path used by Sage: `/signal-cli-data/.local/share/signal-cli/attachments/`). +4. **Build a small β€œrecent context” string** for the vision model (Sage calls `SageAgent::get_recent_messages_for_vision(6)` to format the last few user/assistant turns as `[role]: …` lines). +5. **Call a vision-capable model** via an OpenAI-compatible endpoint (`crates/sage-core/src/vision.rs::describe_image`), using `MAPLE_VISION_MODEL` (defaults to `MAPLE_MODEL`): + - Reads the image from disk, base64-encodes it, and wraps it as `data:;base64,<...>`. + - Sends `POST {MAPLE_API_URL}/chat/completions` with `messages` containing: + - a dedicated **system prompt** (β€œYou are an image description agent… transcribe text exactly…”) and + - a **user content array** with an `image_url` part plus a `text` part that includes recent conversation context + the user’s accompanying text. + - Uses `max_tokens=2048` and returns the model’s text description. +6. **Inject description into the user message** as a bracketed suffix/payload (`[Uploaded Image: …]`) and pass the resulting **text-only** `user_message` into the DSRs agent step loop (`agent_guard.step(&user_message, ...)`). +7. **Persist both the original text and the derived description**: + - Sage stores the original user text in `messages.content` and stores the derived description in `messages.attachment_text` (migration: `crates/sage-core/migrations/2026-02-05-000000_add_attachment_text`). + - It updates the message embedding asynchronously using the *combined* text (original message + injected description) so recall search can match on image content. +8. **Render attachment descriptions in agent context**: when building `recent_conversation` for the DSRs signature, Sage appends `attachment_text` alongside the user message (`crates/sage-core/src/sage_agent.rs`, β€œRender attachment_text alongside user messages”). + +**Key implication for Maple/OpenSecret agent:** DSRs stays **text-in/text-out**; β€œimage understanding” is an explicit **pre-step** that produces storable text (and in OpenSecret, that derived text should be encrypted at rest) that can participate in compaction, embeddings, and memory tools. + +**Notable Sage quirks (worth copying or fixing intentionally):** +- Only the **first** supported image attachment is processed per inbound message. +- On vision failure, Sage injects a placeholder description (so the agent still sees β€œsomething happened”). +- `conversation_search` results are formatted from `messages.content` and (today) do **not** include `attachment_text`, so image-only messages can look blank in tool output even though the agent’s in-context transcript includes the description. + --- ## 6. Agent Step Loop @@ -654,7 +683,7 @@ These are intentionally deferred: - **Group/shared memory.** All memory is per-user. Shared memory blocks between users in the same project would require a new sharing model. Deferred. -- **Voice/image input handling.** Sage has a vision pipeline for Signal image attachments. Maple's Responses API already handles multimodal input. The agent system inherits this capability from the shared message tables. +- **Voice/image input handling.** Sage converts image attachments into injected text via a separate vision pre-processing call (Section 5.4). Maple’s Responses API supports multimodal input directly, but the agent runtime likely needs its own normalization path if it wants Sage-style text-only context + embeddings. - **Reminders / scheduling.** Post-MVP. Sage V2 has a working scheduler + tools (`schedule_task`, `list_schedules`, `cancel_schedule`) backed by a `scheduled_tasks` table. We can port this once the core agent loop is stable. Delivery of fired reminders will use the long-lived SSE channel (`GET /v1/agent/events`, Section 7.3). diff --git a/migrations/2026-02-24-150000_user_messages_attachment_text/down.sql b/migrations/2026-02-24-150000_user_messages_attachment_text/down.sql new file mode 100644 index 00000000..4687cd06 --- /dev/null +++ b/migrations/2026-02-24-150000_user_messages_attachment_text/down.sql @@ -0,0 +1,2 @@ +ALTER TABLE user_messages + DROP COLUMN attachment_text_enc; diff --git a/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql b/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql new file mode 100644 index 00000000..1ad13d0f --- /dev/null +++ b/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql @@ -0,0 +1,7 @@ +-- Add optional derived image description storage for user messages +-- +-- This stores Sage-style vision pre-processing results (encrypted in app layer) +-- without mutating the original MessageContent JSON. + +ALTER TABLE user_messages + ADD COLUMN attachment_text_enc BYTEA; diff --git a/src/models/responses.rs b/src/models/responses.rs index 5ed58b3f..abe10b15 100644 --- a/src/models/responses.rs +++ b/src/models/responses.rs @@ -420,6 +420,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)?; @@ -461,6 +462,7 @@ pub struct UserMessage { pub prompt_tokens: i32, pub created_at: DateTime, pub updated_at: DateTime, + pub attachment_text_enc: Option>, } #[derive(Insertable, Debug)] @@ -472,6 +474,7 @@ pub struct NewUserMessage { pub user_id: Uuid, pub content_enc: Vec, pub prompt_tokens: i32, + pub attachment_text_enc: Option>, } impl UserMessage { @@ -790,6 +793,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)] @@ -832,6 +837,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -851,6 +857,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -870,6 +877,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, @@ -888,6 +896,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, @@ -907,6 +916,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, @@ -945,6 +955,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -964,6 +975,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -983,6 +995,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, @@ -1001,6 +1014,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, @@ -1020,6 +1034,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, @@ -1067,6 +1082,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1086,6 +1102,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1105,6 +1122,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, @@ -1123,6 +1141,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, @@ -1142,6 +1161,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, @@ -1209,6 +1229,7 @@ impl RawThreadMessage { um.id, um.uuid, um.content_enc, + um.attachment_text_enc, 'completed'::text as status, um.created_at, r.model, @@ -1228,6 +1249,7 @@ impl RawThreadMessage { am.id, am.uuid, am.content_enc, + NULL::bytea as attachment_text_enc, am.status, am.created_at, r.model, @@ -1247,6 +1269,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, @@ -1265,6 +1288,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, @@ -1284,6 +1308,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/models/schema.rs b/src/models/schema.rs index 3844eac8..83ea7ab4 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -414,6 +414,7 @@ diesel::table! { prompt_tokens -> Int4, created_at -> Timestamptz, updated_at -> Timestamptz, + attachment_text_enc -> Nullable, } } diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 3c247cbc..21bd8e8a 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -30,7 +30,7 @@ use crate::web::responses::conversations::{ use crate::web::responses::handlers::encrypt_event; use crate::web::responses::{ error_mapping, ConversationBuilder, ConversationItem, ConversationItemConverter, - DeletedObjectResponse, Paginator, + DeletedObjectResponse, MessageContent, MessageContentConverter, MessageContentPart, Paginator, }; use crate::{ApiError, AppMode, AppState}; @@ -38,10 +38,39 @@ mod compaction; mod runtime; mod signatures; mod tools; +mod vision; #[derive(Debug, Clone, Deserialize)] struct AgentChatRequest { - input: String, + input: MessageContent, +} + +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) @@ -238,18 +267,19 @@ async fn chat( Extension(user): Extension, Extension(body): Extension, ) -> Result>>, ApiError> { - if body.input.trim().is_empty() { + 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 input = body.input.clone(); - 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: { @@ -312,18 +342,23 @@ async fn chat( } }; - if let Err(e) = runtime.prepare(&input).await { - 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")), + 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; } - had_error = true; - break 'init; } max_steps = runtime.max_steps(); @@ -344,7 +379,7 @@ async fn chat( } } - let step_result = runtime.step(&input, step_num == 0).await; + let step_result = runtime.step(&input_for_agent, step_num == 0).await; match step_result { Ok(result) => { diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 69aad6e4..97770efe 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -21,7 +21,7 @@ use crate::rag::{ }; use crate::tokens::count_tokens; use crate::web::openai_auth::AuthMethod; -use crate::web::responses::{MessageContent, MessageContentConverter}; +use crate::web::responses::{MessageContent, MessageContentConverter, MessageContentPart}; use crate::{ApiError, AppState}; use super::compaction::CompactionManager; @@ -33,6 +33,7 @@ use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, MemoryInsertTool, MemoryReplaceTool, ToolRegistry, ToolResult, }; +use super::vision; // Mirrors Sage defaults const DEFAULT_PERSONA_DESCRIPTION: &str = "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions."; @@ -314,16 +315,118 @@ impl AgentRuntime { /// 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: &str) -> Result<(), ApiError> { - let trimmed = user_message.trim(); - if trimmed.is_empty() { + 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(trimmed).await?; + self.insert_user_message(normalized, attachment_text, &embed_text) + .await?; self.maybe_compact().await?; - Ok(()) + 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. @@ -681,7 +784,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .map_err(|_| ApiError::InternalServerError)? .unwrap_or_default(); - let content = if msg.message_type == "user" { + 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), @@ -707,6 +810,21 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If 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 { @@ -720,10 +838,35 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If Ok(content) } - pub async fn insert_user_message(&self, text: &str) -> Result { - let content = super::tools::normalize_user_message_content(text); - let prompt_tokens = super::tools::count_user_message_tokens(&content); - let content_enc = encrypt_with_key(&self.user_key, content.as_bytes()).await; + 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 @@ -738,6 +881,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If response_id: None, user_id: self.user.uuid, content_enc, + attachment_text_enc, prompt_tokens, } .insert(&mut conn) @@ -750,7 +894,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If let state = self.state.clone(); let user = self.user.clone(); let user_key = self.user_key.clone(); - let text = text.to_string(); + let text = embed_text.to_string(); let conversation_id = self.conversation.id; let user_message_id = Some(msg.id); diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index 119fa1b6..9199996f 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -12,10 +12,7 @@ use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock}; use crate::models::responses::Conversation; use crate::models::users::User; use crate::rag::{cosine_similarity, deserialize_f32_le, search_user_embeddings}; -use crate::tokens::count_tokens; use crate::web::openai_auth::AuthMethod; -use crate::web::responses::MessageContentConverter; -use crate::web::responses::{MessageContent, MessageContentPart}; use crate::{ApiError, AppState}; // ============================================================================ @@ -802,29 +799,6 @@ impl Tool for DoneTool { } } -// ============================================================================ -// Formatting helpers -// ============================================================================ - -pub fn normalize_user_message_content(text: &str) -> String { - // Store user messages in the DB as MessageContent JSON, like Responses API. - // Use input_text part for compatibility with Conversations API. - let content = MessageContent::Parts(vec![MessageContentPart::InputText { - text: text.to_string(), - }]); - serde_json::to_string(&content).unwrap_or_else(|_| format!("\"{}\"", text)) -} - -pub fn count_user_message_tokens(content_json: &str) -> i32 { - let parsed: Result = serde_json::from_str(content_json); - match parsed { - Ok(content) => count_tokens(&MessageContentConverter::extract_text_for_token_counting( - &content, - )) as i32, - Err(_) => count_tokens(content_json) as i32, - } -} - #[cfg(test)] mod tests { use super::*; 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/responses/handlers.rs b/src/web/responses/handlers.rs index eba3b4d2..231d3c7c 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -1259,6 +1259,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 98479d0a3ff00656944333a678b6cef8c8152c7d Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 2 Mar 2026 07:31:57 -0600 Subject: [PATCH 15/34] Update dsrs reference --- Cargo.lock | 10 ++++++++++ Cargo.toml | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 56fa1f80..4c84aa67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -978,6 +978,7 @@ dependencies = [ [[package]] name = "baml-ids" version = "0.0.1" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "getrandom 0.2.15", @@ -990,6 +991,7 @@ dependencies = [ [[package]] name = "baml-types" version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "baml-ids", @@ -1009,6 +1011,7 @@ dependencies = [ [[package]] name = "bamltype" version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "baml-types", @@ -1026,6 +1029,7 @@ dependencies = [ [[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", @@ -1297,6 +1301,7 @@ dependencies = [ [[package]] name = "bstd" version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "num", @@ -1971,6 +1976,7 @@ dependencies = [ [[package]] name = "dspy-rs" version = "0.7.3" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "arrow", @@ -2006,6 +2012,7 @@ dependencies = [ [[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", @@ -3232,6 +3239,7 @@ 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", @@ -3243,6 +3251,7 @@ dependencies = [ [[package]] name = "internal-baml-jinja" version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "baml-types", @@ -3325,6 +3334,7 @@ dependencies = [ [[package]] name = "jsonish" version = "0.1.0" +source = "git+https://github.com/OpenSecretCloud/DSRs.git?branch=main#be22bef0008bd3631ca4ebd2b873d47c0dbc5916" dependencies = [ "anyhow", "baml-types", diff --git a/Cargo.toml b/Cargo.toml index ea111016..869312d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,7 @@ thiserror = "1.0.63" async-trait = "0.1.81" jsonwebtoken = "9.3.0" jwt-compact = { version = "0.9.0-beta.1", features = ["es256k"] } -dspy-rs = { path = "../../ThirdParties/DSRs/crates/dspy-rs" } +dspy-rs = { git = "https://github.com/OpenSecretCloud/DSRs.git", branch = "main" } diesel = { version = "=2.2.2", features = [ "postgres", "postgres_backend", From 7475e53bedcd1d83507632120ee53ae56f69d68c Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Thu, 5 Mar 2026 12:40:10 -0600 Subject: [PATCH 16/34] fix: use rust 1.90 in nix build and add git dep output hashes Use custom rustPlatform from rust-overlay to match rust-toolchain.toml. Add outputHashes for baml-ids, minijinja, and rig-core git dependencies. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- flake.nix | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/flake.nix b/flake.nix index 4972bd80..76fd27cb 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 @@ -225,7 +229,7 @@ }; # Build the main Rust package - opensecret = pkgs.rustPlatform.buildRustPackage { + opensecret = rustPlatform.buildRustPackage { pname = "opensecret"; version = "0.1.0"; src = pkgs.lib.cleanSourceWith { @@ -248,6 +252,11 @@ }; cargoLock = { lockFile = ./Cargo.lock; + outputHashes = { + "baml-ids-0.0.1" = "sha256-wjgwOcZvZIWEAjpq0gZwtZccQ+FupudBzpxR4+udKEQ="; + "minijinja-2.11.0" = "sha256-i3rXVnwRBWr7ewdFteezZWOrvzW89/VaJ3Yh6Bb40ZY="; + "rig-core-0.26.0" = "sha256-/1C2/UTuJrre4IYNW7i1pYaOGy0dxzhLyogR+5NL1nk="; + }; }; nativeBuildInputs = [ pkgs.pkg-config From 77df6017cbcb37b1f621df4985f5872b3522447e Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Thu, 5 Mar 2026 19:08:39 -0600 Subject: [PATCH 17/34] feat: model Sage as a main agent with subagents Hide the primary Maple relationship behind a main agent thread while giving users focused subagent chats that share memory and code-owned runtime defaults. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- internal_docs/sage-in-maple-architecture.md | 959 ++++++++--------- .../down.sql | 7 +- .../up.sql | 53 +- src/models/agent_config.rs | 88 -- src/models/agents.rs | 121 +++ src/models/memory_blocks.rs | 31 +- src/models/mod.rs | 2 +- src/models/schema.rs | 23 +- src/web/agent/mod.rs | 965 +++--------------- src/web/agent/runtime.rs | 406 ++++++-- src/web/agent/signatures.rs | 10 +- src/web/agent/tools.rs | 100 +- src/web/responses/conversations.rs | 27 +- 13 files changed, 1160 insertions(+), 1632 deletions(-) delete mode 100644 src/models/agent_config.rs create mode 100644 src/models/agents.rs diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md index 4da9aa2d..dadda1da 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/sage-in-maple-architecture.md @@ -1,173 +1,269 @@ # Sage-in-Maple: Agent Memory Architecture -## Bringing Persistent Agent Memory into OpenSecret +## Main Agent + Shared-Memory Subagents -**Date:** February 2026 -**Status:** MVP implemented (Local/Dev only; `/v1/agent/chat` streams message-level SSE; Responses API auto-embedding + production feature flags pending) +**Date:** March 2026 +**Status:** Target architecture (local-only reset; no backwards-compatibility constraints) **Related Docs:** -- `potential-rag-integration-brute-force.md` -- RAG/vector storage layer (prerequisite) +- `potential-rag-integration-brute-force.md` -- RAG/vector storage layer - `architecture-for-rag-integration.md` -- OpenSecret encryption and data model reference - Sage V2 Design Doc (`~/Dev/Personal/sage/docs/SAGE_V2_DESIGN.md`) -- proven prototype - Sage V2 Codebase (`~/Dev/Personal/sage/crates/sage-core`) -- DSRs signatures + BAML parsing + multi-step tool loop prototype -**Overall implementation status:** Partially implemented (MVP shipped behind Local/Dev gating; Responses API auto-embedding + production feature flags not implemented yet). +**Implementation note:** the current local codebase still reflects an older single-agent MVP in places. This document is now the source of truth for the rewrite. --- ## 1. Goal -Bring Sage's proven 4-tier memory architecture (core, recall, archival, summary) into Maple as a first-class feature, without disrupting the existing Responses API that third-party developers build against. +Bring Sage's proven 4-tier memory architecture (core, recall, archival, summary) into Maple as a first-class product experience, without breaking the existing Responses API that third-party developers already understand. -The end state: Maple users can have a persistent AI agent that remembers across conversations, has editable personality/user memory blocks, can store and search long-term memories, and auto-compacts conversation history when context windows fill up -- all with the same per-user encryption guarantees that exist today. +The target product shape is: +- **one main persistent agent** as the app's home surface +- **unlimited subagents** as topic-specific chats/workspaces +- **shared memory** across the main agent and all subagents +- **separate conversation history** per agent thread +- **no end-user configuration burden** for prompts, models, memory block definitions, or tuning knobs + +The end state is not β€œusers manually configure an agent.” The end state is β€œusers talk to Maple, Maple manages the agent system for them.” ### 1.1 The User Model -**One user : one main agent : one persistent conversation.** +**One user : one main agent : unlimited subagents.** + +Each agent owns exactly one conversation thread (`conversations.id`): +- The **main agent** is the user's long-running relationship with Maple. +- A **subagent** is a topic-specific extension of the main agent with its **own conversation history**, but access to the **same shared memory** and **same recall surface**. +- A subagent can be created either by the **user** or by the **main agent**. +- A subagent can remain long-lived forever if the user keeps using it. + +This means the durable model is: +1. **Agent identity** +2. **One conversation thread per agent** +3. **User-shared memory across all of that user's agents** -The agent lives in a single long-running conversation thread. Unlike Responses API threads (which are isolated, stateless, and disposable), the agent conversation is the user's ongoing relationship with Maple. Memory blocks, archival memory, and conversation summaries all attach to this single thread. +**Implementation note:** do not infer agent identity from encrypted `conversations.metadata_enc`. Agent identity must be explicit in a dedicated `agents` table. -**Implementation note:** this mapping must be explicit in the database. We should not infer "the agent thread" from encrypted `conversations.metadata_enc`. Instead, store the thread pointer in `agent_config.conversation_id` (BIGINT FK to `conversations.id`) when the agent is initialized. +### 1.2 Shared Memory, Isolated Transcripts -The Responses API continues to exist alongside the agent. Users can still have one-off stateless threads for quick questions. But the agent is the primary interface for users who want memory and persistence. +All agents belonging to a user share: +- `memory_blocks` (core memory) +- archival memory in `user_embeddings` +- recall search over embedded messages in `user_embeddings` -**Implementation status:** Complete (agent_config.conversation_id is the canonical pointer to the main agent thread). +Each individual agent keeps its own: +- recent conversation history +- `conversation_summaries` chain +- active thread-local context window -### 1.2 Cross-Thread Memory Visibility +A subagent does **not** inherit the main agent's recent transcript. If it needs older context, it should use `conversation_search` just like the main agent would. -**The agent sees everything.** When the agent's `conversation_search` tool fires, it searches across all of the user's embedded messages -- both the agent's own conversation and any Responses API threads. The `user_embeddings` table indexes messages from all conversations by default. +### 1.3 Conversation List Model -Rationale: if a user had a Responses API thread about a topic, then asks the agent about that topic, the agent should know about it. The user opted into the agent system; that's the trust boundary. Individual threads are not a privacy boundary within the same user account. +The **main agent** is the app's first-class home surface, not just another conversation row. -The agent can still scope searches to its own thread via `conversation_id` filtering, but the default is broad. See the RAG proposal's "Data Visibility and Isolation" section for the full filtering model. +The **conversation list** should show: +- legacy Responses API threads +- subagent chats -**Future: incognito threads.** A `private` flag on conversations will let users create Responses API threads that are excluded from embedding entirely. The agent cannot see them. This is deferred to post-v1. +The conversation list should **not** show: +- the main agent thread -Additionally, requests made with `store=false` should not be embedded/indexed into `user_embeddings` (per-request opt-out), even if the Responses API continues to persist messages today. +So the product model becomes: +- **Main agent:** always-available home agent surface +- **Subagents:** the new first-class β€œchat/workspace” objects +- **Responses threads:** legacy/stateless threads that still exist alongside the new system -**Implementation status:** Partially implemented (cross-thread search works where embeddings exist; Responses API auto-embedding + private/store=false exclusions not implemented yet). +### 1.4 Cross-Thread Memory Visibility -### 1.3 Future Vision: Agent-Backed Responses API +**All agents for a user see the same memory layer.** -Eventually, the Responses API threads themselves could be backed by the agent's memory system. Instead of stateless threads with no memory, a Responses API thread would: -- Start with a fresh conversation history (no prior messages in context) -- But still have access to memory blocks, archival search, and conversation search -- Effectively: the agent "knows" the user (from core + archival memory), but the thread is a clean slate +When `conversation_search` runs, it should search across all eligible embedded messages for that user by default: +- main agent thread messages +- subagent thread messages +- Responses API thread messages -This would give users the best of both: the lightweight feel of a new thread, with the accumulated knowledge of the persistent agent. This is a post-v1 evolution that requires proving out the agent system first -- particularly that memory blocks and search produce consistently good results without the agent's own conversation context. +Rationale: subagents are extensions of the same agent system, not separate privacy domains. If the user or the main agent creates a subagent for taxes, writing, or deployment planning, it should be able to benefit from the same user-level memory and searchable history. -**Implementation status:** Not implemented yet. +The tool can still accept an optional `conversation_id` to scope search to one thread, but the default should be broad. -### 1.4 Future: Subagents (Multi-Agent per User) +**Future:** +- `private` / incognito conversations excluded from embedding +- `store=false` suppressing embedding/indexing -Subagents are out of scope for the MVP, but we should keep the storage design compatible with them. +Those remain opt-out boundaries for recall visibility. -**Key idea:** each agent is just: -1) an **agent identity/config** record, and -2) a **single long-running conversation thread** (`conversations.id`) that the agent reads/writes. +### 1.5 Relationship to the Responses API -The existing message tables remain unchanged; different agents simply write to different `conversation_id`s. +The Responses API remains a separate surface: +- good for stateless or developer-facing usage +- preserved for backwards compatibility at the product/API level +- still eligible for auto-embedding into `user_embeddings` when policy allows -We also want subagents to support **configurable memory scoping**: -- **Shared/inherited memory:** subagents read/write the user's shared memory blocks and archival memory. -- **Isolated memory:** subagents have their own blocks/archival memory (no visibility into the user's shared memory unless explicitly allowed). +Responses threads are **not** subagents. They remain distinct objects. But they can still contribute to the user's shared recall memory. -The cleanest way to support this is to introduce an `agents` table and add an `agent_id` FK to *agent-specific* tables (Section 4.5). Shared vs isolated memory becomes a query policy (and optionally a config field) rather than a schema rewrite. +### 1.6 Configuration Philosophy -**Implementation status:** Not implemented yet. +The product should be **agent-managed**, not **user-configured**. -**Overall implementation status:** Partially implemented (MVP main-agent model shipped; agent-backed Responses API + subagents not implemented yet). +That means the following should be owned by code / rollout config, not by per-user mutable database rows: +- default model selection +- prompt templates +- context windows +- compaction thresholds +- fixed memory block definitions and limits + +The database should persist **durable state**, not **tuning policy**. --- -## 2. Key Architectural Decision: Separate API, Shared Storage +## 2. Core Architectural Decisions -### What stays the same +### 2.1 Separate Runtime, Shared Storage -The existing Responses API (`/v1/responses/*`, `/v1/conversations/*`, `/v1/instructions/*`) is **untouched**. It continues to be a stateless OpenAI-compatible chat API. Third-party developers who use it see no changes. +The agent runtime remains separate from the Responses API runtime. -The existing database tables for conversations and messages are **reused** by the agent system. The agent reads from and writes to the same `conversations`, `user_messages`, `assistant_messages`, `tool_calls`, `tool_outputs`, and `reasoning_items` tables. The encrypted content format is identical. +What stays shared: +- `conversations` +- `user_messages` +- `assistant_messages` +- `tool_calls` +- `tool_outputs` +- `reasoning_items` -**Implementation status:** Complete (agent reuses existing message tables and encryption format). +What becomes agent-specific: +- `agents` +- `memory_blocks` +- `conversation_summaries` +- agent-only tools / orchestration logic -### What's new +This keeps the OpenAI-compatible Responses API stable while allowing Sage-style regenerated context and memory behavior for the main agent and subagents. -A new `/v1/agent/*` API surface with its own handler module (`src/web/agent/`), its own step loop, and its own context assembly logic. This API implements the Sage-style regenerated-context pattern instead of the Responses API's middle-truncation pattern. +### 2.2 Introduce Explicit Agent Identity Now -New database tables for agent-specific concerns (detailed in Section 4): -- `memory_blocks` -- core memory (persona, human, custom blocks) -- `user_embeddings` -- archival memory + chat embeddings (from the RAG proposal) -- `conversation_summaries` -- compaction artifacts -- `agent_config` -- per-user agent settings +The old MVP's `agent_config` shape bakes in β€œone agent per user.” That is no longer the right abstraction. -**Implementation status:** Complete (MVP /v1/agent surface + new tables + RAG foundation implemented; Local/Dev gated). +The architecture should move directly to an explicit `agents` table now. -### 2.1 Feature Flags and Rollout Safety (MVP requirement) +Why: +- one user must have exactly one **main** agent +- one user can have many **subagents** +- each agent maps to exactly one conversation thread +- subagents must be distinguishable from Responses API threads in list views and routing +- the main agent must be hideable from `/v1/conversations/*` while subagents remain visible -Sage is a large feature that touches storage, tool execution, and LLM calls. We need a lightweight, production-friendly feature flag system to safely ship in enclaves. +### 2.3 Keep Shared Memory Per User -At minimum, ship with **two levels of gating**: -- **Global kill switch** (env/config): hard-disable all `/v1/agent/*` routes and background jobs (auto-embedding, compaction, reminders). -- **Per-user opt-in** (`agent_config.enabled`): enables the persistent agent for that user. +For the subagent model described above, the correct v1 choice is: +- **shared per-user core memory** +- **shared per-user archival memory** +- **shared per-user recall search** +- **per-agent thread-local conversation history** -Flags should also gate sub-features independently (even if `agent_config.enabled=true`): -- auto-embedding into `user_embeddings` (recall memory) -- agent tool execution (memory tools, web search) -- compaction/summarization -- GEPA runs (optimizer execution) vs simply *consuming* an optimized instruction +This means we should **not** add `agent_id` to `memory_blocks` or `user_embeddings` in v1. Doing so would complicate the schema for a capability we explicitly do not want yet. -**Implementation status:** Not implemented yet (only Local/Dev AppMode gating today; os_flags client exists but is unused for agent rollout). +### 2.4 Treat the Main Agent as Special, but Not All Agents as Hidden -### 2.2 Conversations API Isolation (Hide the Main Agent Thread) +Only the **main agent** should be hidden from the generic conversation list. -Even though the agent's persistent thread is stored in the shared `conversations`/message tables, the **main agent conversation should be treated as an internal implementation detail** and **excluded from the public Conversations API** (`/v1/conversations/*`). +Subagents should appear in `/v1/conversations/*` because they are the new product-level chat objects. They are exactly what users should browse, reopen, delete, and continue. -Rationale: -- Prevents confusion for developers using the stateless Responses/Conversations APIs. -- Avoids muddying thread lists with an always-on, long-running agent thread. -- Keeps `/v1/agent/*` as the single supported interface for agent state. +So the rule is: +- **hide the main agent thread** +- **show subagent threads** +- **show Responses API threads** -**Status: Implemented (2026-02-11).** The agent thread is hidden from all public Conversations API endpoints: -- `agent_config.conversation_id` stores the agent thread pointer. -- `ConversationContext::load()` returns `404` when the requested conversation is the agent thread (guards `GET/POST/DELETE /v1/conversations/:id` and `/items`). -- `GET /v1/conversations` filters out the agent thread before pagination. -- `DELETE /v1/conversations` (delete-all) excludes the agent conversation via `id.ne(agent_conversation_id)`. -- `POST /v1/conversations/batch-delete` returns `not_found` for the agent thread instead of deleting it. +### 2.5 No Backwards Compatibility in Local-Only Schema -**Implementation status:** Complete. +This work has not been deployed anywhere yet. We do **not** need compatibility-preserving additive migrations for the current Sage work. -**Overall implementation status:** Partially implemented (API split + conversation isolation complete; feature flags not implemented yet). +So for the rewrite: +- edit the existing local-only Sage migrations directly +- replace `agent_config` with `agents` +- simplify `memory_blocks` now instead of carrying legacy fields forward +- move code-owned defaults out of SQL and into Rust/config --- ## 3. Why Reuse the Message Tables -This was the hardest decision. We did a column-by-column comparison between Sage's `messages` table and Maple's split message tables. The conclusion: - -**Maple's schema is a strict superset of Sage's.** Sage stores `(id, role, content, agent_id, sequence_id, embedding, tool_calls, tool_results, created_at)` in a single table. Maple stores the same information across typed tables with more metadata: token counts per message, response status tracking, response_id linking, conversation_id scoping, and UUIDs for external references. +This decision does not change. -**What Sage does with messages at runtime:** +**Maple's split message tables are still the right storage layer** for Sage-style agents. The main difference is no longer β€œone persistent agent thread per user”; it is now β€œone persistent thread per agent identity.” -| Operation | Sage approach | Maple equivalent | +| Operation | Agent requirement | Maple storage choice | |---|---|---| -| Read recent messages in order | `SELECT * FROM messages WHERE agent_id = ? ORDER BY sequence_id DESC LIMIT N` | The existing `RawThreadMessage::get_conversation_context` UNION ALL query across all message tables | -| Track which messages are "in context" | `agents.message_ids` UUID array | Compute dynamically from conversation + summaries (or add field to `agent_config`) | -| Store new messages | `INSERT INTO messages` with role | `INSERT INTO user_messages` / `assistant_messages` / etc. (existing flow) | -| Search messages semantically | pgvector on messages.embedding | `user_embeddings` table with in-process brute-force search (from RAG proposal) | -| Mark messages as summarized | Doesn't mark messages; records sequence ranges in summaries table | Same approach: `conversation_summaries` stores ranges, messages are untouched | +| Read recent messages in order | Load messages for the active agent's thread | Existing `RawThreadMessage` UNION query filtered by that agent's `conversation_id` | +| Store new messages | Persist by role and message type | Existing `user_messages` / `assistant_messages` / etc. | +| Search semantically | Search across all eligible user messages | `user_embeddings` keyed by `user_id` with optional `conversation_id` filter | +| Track compacted history | Keep summaries per thread | `conversation_summaries` keyed by `conversation_id` | -**Nothing in Sage's message access patterns requires modifying the existing Maple message table columns.** The agent needs to *read* from those tables differently (regenerated context with memory blocks instead of middle-truncation), but the storage format is identical. +Nothing about the subagent model requires changing the core message tables themselves. The agent system still needs to **read them differently**, not **store them differently**. -**Potential future addition:** An `is_in_context` boolean or a `summary_id` FK on messages would be an optimization for fast filtering. But it's not required -- the summaries table's ranges can determine this. If needed, it would be an additive nullable column with a migration default, not a structural change. +--- -**Overall implementation status:** Complete (agent uses existing message tables via RawThreadMessage UNION queries; no message-table schema changes required). +## 4. Data Model ---- +### 4.1 `agents` -- Explicit Agent Identity + +This replaces the old `agent_config`-as-identity model. + +```sql +CREATE TABLE agents ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + -- One conversation thread per agent + conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + + -- 'main' or 'subagent' + kind TEXT NOT NULL, + + -- Main agent has NULL parent. Subagents should point at the main agent. + parent_agent_id BIGINT REFERENCES agents(id) ON DELETE SET NULL, + + -- User-visible metadata; encrypted because names/purposes may be sensitive. + display_name_enc BYTEA, + purpose_enc BYTEA, + + -- Provenance only: 'user' or 'agent' + created_by TEXT NOT NULL DEFAULT 'user', + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), -## 4. New Database Tables + CONSTRAINT agents_kind_check CHECK (kind IN ('main', 'subagent')), + CONSTRAINT agents_created_by_check CHECK (created_by IN ('user', 'agent')), + UNIQUE(conversation_id) +); + +CREATE UNIQUE INDEX idx_agents_one_main_per_user + ON agents(user_id) + WHERE kind = 'main'; -### 4.1 `memory_blocks` -- Core Memory +CREATE INDEX idx_agents_user_kind_created + ON agents(user_id, kind, created_at DESC); -The always-in-context memory blocks that define agent personality and user information. Directly modeled on Sage's `blocks` table, adapted for Maple's encryption model. +CREATE INDEX idx_agents_parent_agent_id + ON agents(parent_agent_id); + +CREATE TRIGGER update_agents_updated_at +BEFORE UPDATE ON agents +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**Notes:** +- There is exactly **one** `kind='main'` agent per user. +- There may be **many** `kind='subagent'` rows per user. +- `conversation_id` is the canonical pointer to the agent's thread. +- `display_name_enc` and `purpose_enc` are encrypted because they may contain user-sensitive content. +- `parent_agent_id` exists so a subagent is explicitly modeled as an extension of the main agent. +- There is intentionally **no per-agent model/prompt/tuning config** here. + +### 4.2 `memory_blocks` -- Shared Core Memory + +Core memory remains per-user and shared across the main agent plus all subagents. ```sql CREATE TABLE memory_blocks ( @@ -175,12 +271,9 @@ CREATE TABLE memory_blocks ( uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, - label TEXT NOT NULL, -- plaintext: "persona", "human", custom labels - description TEXT, -- plaintext: how this block should be used - value_enc BYTEA NOT NULL, -- AES-256-GCM encrypted block content - char_limit INTEGER NOT NULL DEFAULT 5000, - read_only BOOLEAN NOT NULL DEFAULT FALSE, - version INTEGER NOT NULL DEFAULT 1, + -- Fixed built-in labels for now: 'persona', 'human' + label TEXT NOT NULL, + value_enc BYTEA NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), @@ -195,40 +288,42 @@ BEFORE UPDATE ON memory_blocks FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); ``` -**Key differences from Sage's `blocks` table:** -- `value_enc` (BYTEA) instead of `value` (TEXT) -- content encrypted with user's per-user key -- `user_id` instead of `agent_id` -- MVP keeps blocks per-user. For multi-agent/subagents, add a nullable `agent_id` so blocks can be shared (`agent_id IS NULL`) or agent-scoped (`agent_id = `). -- `label` and `description` are plaintext -- these are structural identifiers ("persona", "human"), not user content. They don't need encryption. This allows querying by label without decryption. - -**Default blocks created on agent initialization:** -- `persona`: "I am a helpful AI assistant." (editable by agent) -- `human`: "" (populated by agent as it learns about the user) +**V1 choices:** +- Shared per user, **not** per agent +- No `description` +- No persisted `char_limit` +- No `read_only` +- No `version` +- No public manual CRUD requirement -**Implementation status:** Complete (table + models + CRUD endpoints + tools; code default char_limit=20000). +**Built-in blocks for v1:** +- `persona` +- `human` -### 4.2 `user_embeddings` -- General-Purpose Embedding Store +Any limits, formatting rules, or block semantics should live in code. If we later add more built-in blocks, we can do that by code first without needing a new table shape. -Defined in the RAG proposal. A single table that serves all embedding use cases through a `source_type` discriminator: +### 4.3 `user_embeddings` -- Shared Recall + Archival Store -- `source_type = 'message'`: auto-indexed chat history (recall memory) -- `source_type = 'archival'`: agent-inserted long-term memories (archival memory) -- `source_type = 'document'`: user-uploaded document chunks (document RAG) +This remains the shared embedding layer for the user. -- Future source types added without migration +- `source_type = 'message'`: embedded chat history +- `source_type = 'archival'`: long-term memory passages inserted by the agent +- `source_type = 'document'`: document chunks -Key columns for the agent system: `embedding_model` (tracks which model produced the vector, enables re-embedding on model upgrade), `content_enc` (always present -- the embedded text itself, encrypted), and `tags_enc` (deterministically-encrypted, base64 tags enabling SQL-indexable filtering for archival memories; extracted from `metadata.tags` when present). +**V1 rule:** do **not** add `agent_id` here. -**Future:** add `chunk_index` (ordering within multi-chunk sources like documents). +Why: +- recall memory is intentionally shared across the main agent and all subagents +- archival memory is also intentionally shared in v1 +- `conversation_id` already provides the necessary filter for message-scope search when needed -**Future (multi-agent/subagents):** consider adding an optional `agent_id` column for `source_type='archival'` (and possibly `source_type='document'`) rows so archival/document memory can be shared or isolated per agent. Recall/message embeddings remain scoped by `conversation_id` and can still default to cross-thread search. +If we ever decide to support isolated subagent memory later, we can add explicit scope then. We should not pre-complicate the current design for a behavior we do not want yet. -See `potential-rag-integration-brute-force.md` for full schema and API design. +See `potential-rag-integration-brute-force.md` for the full schema and search behavior. -**Implementation status:** Partially implemented (table + in-process search/cache + archival tags implemented; Responses API auto-embedding not implemented yet). +### 4.4 `conversation_summaries` -- Per-Thread Compaction Artifacts -### 4.3 `conversation_summaries` -- Compaction Artifacts - -Stores rolling summaries when conversation context exceeds limits. Modeled on Sage's `summaries` table, adapted for Maple's conversation model. +This table stays conceptually the same. ```sql CREATE TABLE conversation_summaries ( @@ -237,22 +332,15 @@ CREATE TABLE conversation_summaries ( user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, - -- Range of messages this summary covers - -- Uses message created_at timestamps as boundaries from_created_at TIMESTAMPTZ NOT NULL, to_created_at TIMESTAMPTZ NOT NULL, - message_count INTEGER NOT NULL, -- how many messages were summarized - - -- Encrypted summary content - content_enc BYTEA NOT NULL, -- AES-256-GCM encrypted summary text - content_tokens INTEGER NOT NULL, -- plaintext token count of summary + message_count INTEGER NOT NULL, - -- Encrypted embedding for semantic search over summaries - embedding_enc BYTEA, -- AES-256-GCM encrypted float32 array + content_enc BYTEA NOT NULL, + content_tokens INTEGER NOT NULL, + embedding_enc BYTEA, - -- Chain: previous summary that was absorbed into this one previous_summary_id BIGINT REFERENCES conversation_summaries(id) ON DELETE SET NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), CONSTRAINT valid_time_range CHECK (from_created_at <= to_created_at) @@ -265,127 +353,84 @@ CREATE INDEX idx_conversation_summaries_chain ON conversation_summaries(previous_summary_id); ``` -**Key differences from Sage:** -- Uses `from_created_at` / `to_created_at` timestamp ranges instead of Sage's `from_sequence_id` / `to_sequence_id`. Maple's message tables don't have a monotonic sequence_id across the UNION -- they have per-table auto-increment IDs. Timestamps are the reliable ordering key that works across all message types. -- `embedding_enc` instead of pgvector `embedding` column. Follows the same encrypted pattern as everything else. -- `content_tokens` plaintext for context budgeting. +**Key point:** summaries are per conversation thread, which naturally means per agent. No `agent_id` is needed because `conversation_id` already maps back to `agents`. -**Implementation status:** Complete (compaction writes summaries + embeds them for search). +### 4.5 Remove `agent_config` -### 4.4 `agent_config` -- Per-User Agent Settings +`agent_config` should not exist in the target architecture. -```sql -CREATE TABLE agent_config ( - id BIGSERIAL PRIMARY KEY, - uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, - user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE UNIQUE, +Why it should go away: +- it mixes durable identity with tuning policy +- it assumes one agent per user +- it stores values that should now be code-owned (`model`, `system_prompt`, `max_context_tokens`, `compaction_threshold`) +- it encourages user-facing configuration we do not want as a product - -- Which conversation thread is the user's "main agent" thread - -- Nullable until agent initialization creates/assigns a thread. - conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, +The durable pieces of state we actually need are: +- **agent identity** -> `agents` +- **shared core memory** -> `memory_blocks` +- **shared recall/archival storage** -> `user_embeddings` +- **thread-local compaction state** -> `conversation_summaries` - -- Agent behavior settings (plaintext, not sensitive) - enabled BOOLEAN NOT NULL DEFAULT FALSE, - model TEXT NOT NULL DEFAULT 'deepseek-r1-0528', - max_context_tokens INTEGER NOT NULL DEFAULT 100000, - compaction_threshold REAL NOT NULL DEFAULT 0.80, - - -- Agent system prompt (encrypted, user-customizable) - system_prompt_enc BYTEA, - - -- User preferences (encrypted JSON: timezone, response style, etc.) - preferences_enc BYTEA, - - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() -); - -CREATE TRIGGER update_agent_config_updated_at -BEFORE UPDATE ON agent_config -FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); -``` - -**Notes:** -- One agent config per user (MVP). The `UNIQUE(user_id)` constraint enforces this. -- **Future (multi-agent/subagents):** prefer introducing an `agents` table and making `agent_config` keyed by `agent_id` (or merging `agent_config` into `agents`). This makes agent identity explicit and enables agent-scoped memory via `agent_id` FKs (Section 4.5). The simpler alternative (remove `UNIQUE(user_id)` + add a `name` column) is viable but makes shared vs isolated memory harder to represent cleanly. -- `conversation_id` is the canonical pointer to the user's persistent agent thread. This avoids relying on encrypted conversation metadata to locate the agent conversation. -- `enabled` allows opt-in activation. Users who don't want agent features continue using the stateless Responses API. -- `system_prompt_enc` is the base agent instruction. Equivalent to Sage's `agents.system_prompt`, but encrypted. If NULL, a default system prompt is used. -- `preferences_enc` absorbs what Sage stores in the separate `user_preferences` table. A single encrypted JSON blob is simpler than a KV table for a small number of preferences (timezone, response style, language). - -**Implementation status:** Partially implemented (config + conversation_id + system_prompt implemented; preferences + robust opt-in/flags not implemented yet). - -### 4.5 Future: `agents` Table + `agent_id` FKs (Multi-Agent/Subagents) - -If we want multiple agents per user (subagents), the most extensible model is to make agent identity explicit. - -**Proposed `agents` table (future):** - -```sql -CREATE TABLE agents ( - id BIGSERIAL PRIMARY KEY, - uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, - user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, - - -- User-visible identifier (plaintext) - name TEXT NOT NULL DEFAULT 'main', - - -- Each agent has exactly one long-running conversation thread - conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, - - enabled BOOLEAN NOT NULL DEFAULT FALSE, - - -- Optional: policy for shared vs isolated memory - -- Examples: 'shared', 'isolated', 'overlay' - memory_mode TEXT NOT NULL DEFAULT 'shared', +--- - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), +## 5. Agent Context Assembly - UNIQUE(user_id, name), - UNIQUE(conversation_id) -); +The main agent and every subagent use the same high-level regenerated-context pattern. -CREATE TRIGGER update_agents_updated_at -BEFORE UPDATE ON agents -FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); -``` +### 5.1 Code-Owned Prompting, Not User-Owned Prompting -**Agent-scoped foreign keys (future):** add `agent_id` (nullable or required) to agent-specific tables: -- `memory_blocks.agent_id` (nullable) - - `agent_id IS NULL` => shared blocks for the user - - `agent_id = X` => blocks scoped to agent X -- `conversation_summaries.agent_id` (optional; can be derived from `conversation_id`) -- `user_embeddings.agent_id` (optional; most useful for `source_type='archival'` / `'document'`) -- `scheduled_tasks.agent_id` (for reminders) -- `agent_config.agent_id` (if/when `agent_config` becomes per-agent) +The active agent runtime should load: +1. the `agents` row for the active agent +2. code-owned prompt template(s) for that agent kind +3. shared memory blocks for the user +4. the latest summary for the active conversation +5. recent messages for the active conversation -**Uniqueness note (Postgres):** if we use `memory_blocks.agent_id IS NULL` to represent shared blocks, we likely want **partial unique indexes** to enforce: -- exactly one shared block per `(user_id, label)` and -- exactly one agent-scoped block per `(user_id, agent_id, label)`. +**Prompt/model selection should be code-owned.** -For MVP, we can start with the simpler per-user `agent_config` model, and only introduce `agents` once subagents are in scope. +That means: +- no `system_prompt_enc` per user +- no per-user model selection row +- no per-user compaction threshold row +- no per-user max context row -**Implementation status:** Not implemented yet. +Instead, code chooses: +- default main-agent instruction +- default subagent instruction +- default model(s) +- context window / packing rules +- compaction thresholds -**Overall implementation status:** Partially implemented (MVP agent tables shipped; multi-agent/subagent schema not implemented yet). +### 5.2 Main Agent vs Subagent Context ---- +The main agent and subagents share memory, but they do **not** share recent transcript. -## 5. Agent Context Assembly +The context builder should behave like this: -This is the fundamental difference between the Responses API and the agent system. Instead of the Responses API's approach (load all messages, middle-truncate to fit context window), the agent regenerates a structured context on every turn. +1. **Load active agent** from `agents` +2. **Build system prompt** + - start with code-owned instruction for `kind='main'` or `kind='subagent'` + - if subagent, inject decrypted `purpose_enc` + - inject shared `persona` and `human` blocks + - inject available tools + - inject memory metadata +3. **Load latest summary** for the active `conversation_id` +4. **Load recent messages** for the active `conversation_id` +5. **Pack within token budget** +6. **Compact old messages** in that same conversation when thresholds are reached -### 5.1 Sage V2 Prototype Pattern (DSRs Signatures + BAML Parsing) +A subagent should never silently get the main agent's recent transcript. If it needs information from another thread, it should use recall tools intentionally. -The current Sage V2 prototype uses DSRs signatures for typed input/output and BAML parsing for structured extraction, rather than provider-native tool calling. +### 5.3 DSRs Signature Shape -Conceptually, the agent is called with a set of **separated context fields** (so GEPA can optimize them independently), and returns **two typed outputs**: `messages[]` and `tool_calls[]`. +The DSRs pattern still fits well. Conceptually the active agent call becomes: -``` +```text AgentResponse (inputs) - - input (user message OR tool-result continuation) + - input - current_time + - agent_kind + - subagent_purpose - persona_block - human_block - memory_metadata @@ -399,366 +444,216 @@ AgentResponse (outputs) - tool_calls: {name: string, args: map}[] ``` -``` -[System Prompt] - - Base instruction (GEPA-optimized or user-customized) - - - ... - ... - - - memory_replace, memory_append, archival_insert, archival_search, - conversation_search, web_search, ... - - - archival_count, recall_count, last_compaction, ... - -[Previous Context Summary] (if conversation was compacted) - -[Recent Messages] (last N messages that fit in token budget) - -[Current User Message] -``` - -**Implementation status:** Complete (DSRs-style signature inputs/outputs implemented via dspy-rs + structured parsing). - -### 5.2 Maple Agent Context Builder - -A new `src/web/agent/context_builder.rs` that replaces the Responses API's `context_builder.rs` for agent requests: - -1. **Load agent config** -- model, max_context_tokens, compaction_threshold -2. **Build system prompt:** - - Start with base instruction from `agent_config.system_prompt_enc` (or default) - - Decrypt and inject memory blocks from `memory_blocks` table - - Inject tool descriptions (from agent tool registry) - - Inject memory metadata (counts from user_embeddings, memory_blocks, conversation_summaries) -3. **Load summary** -- most recent `conversation_summaries` entry for this conversation, if any -4. **Load recent messages** -- from the existing message tables (same UNION ALL query), but: - - Start from after the summary's `to_created_at` timestamp - - Load messages until token budget is consumed - - No middle-truncation -- if budget exceeded, trigger compaction instead -5. **Check compaction threshold** -- if total tokens > threshold, compact oldest messages into a new summary before proceeding -6. **Assemble final prompt** -- system + summary + recent messages + current input - -In practice, if we follow the Sage V2 prototype, the "assembled prompt" is expressed as the DSRs signature inputs above (plus an instruction string), and BAML parses the structured outputs. +### 5.4 Compaction vs Truncation -**Implementation status:** Partially implemented (context assembly implemented in `src/web/agent/runtime.rs`; does not strictly pack messages to a token budget). +This also stays the same conceptually: +- Responses API truncates +- agents compact -### 5.3 Compaction vs Truncation +Each agent thread owns its own summary chain. Main agent compaction and subagent compaction are independent because they are keyed by different `conversation_id`s. -The Responses API truncates: it drops middle messages and inserts `[Previous messages truncated due to context limits]`. Information is lost. +### 5.5 Vision / Image Handling -The agent system compacts: it summarizes old messages into a `conversation_summaries` entry using an LLM call, then removes those messages from the in-context window (but they remain in the database and are searchable via recall memory). Information is preserved in compressed form. - -Compaction uses a DSRs-style signature (following Sage's `SummarizeConversation` pattern): - -``` -Input: - - previous_summary (if chaining summaries) - - messages_to_summarize -Output: - - summary (100 word limit) -``` +No architectural change is needed here. -The LLM call for summarization uses a fast/cheap model (e.g., `gpt-oss-120b`) to minimize latency and cost. +The Sage-style approach still fits: +- preprocess image input into text +- inject the derived text into the conversation flow +- persist the derived text so it can be embedded and recalled later -**Implementation status:** Partially implemented (compaction + summaries implemented; summarization currently uses the agent model, not a dedicated cheap model). - -**Overall implementation status:** Partially implemented (regenerated context + compaction shipped; strict token-budget packing + cheap summarizer model not implemented yet). - -### 5.4 Vision / Image Handling (How Sage Does It) - -Sage V2 does **not** pass images through the DSRs signature/tool loop. Instead, it runs a **vision pre-processing** step that converts an image attachment into a **text description**, then injects that text into the conversation so the rest of the system stays text-only. - -**End-to-end flow (Sage codebase):** - -1. **Signal ingestion** parses `attachments[]` from `signal-cli` JSON-RPC (`crates/sage-core/src/signal.rs::parse_incoming_message`). Each attachment includes a MIME type (`contentType`) and an ID/filename (`id`). -2. **Pre-processing trigger** in the main receive loop (`crates/sage-core/src/main.rs`) checks for the first supported image attachment (`crates/sage-core/src/vision.rs::is_supported_image`, currently `image/jpeg`, `image/png`, `image/webp`, `image/gif`). -3. **Load image bytes** from the `signal-cli` attachments directory mounted into the Sage container (default path used by Sage: `/signal-cli-data/.local/share/signal-cli/attachments/`). -4. **Build a small β€œrecent context” string** for the vision model (Sage calls `SageAgent::get_recent_messages_for_vision(6)` to format the last few user/assistant turns as `[role]: …` lines). -5. **Call a vision-capable model** via an OpenAI-compatible endpoint (`crates/sage-core/src/vision.rs::describe_image`), using `MAPLE_VISION_MODEL` (defaults to `MAPLE_MODEL`): - - Reads the image from disk, base64-encodes it, and wraps it as `data:;base64,<...>`. - - Sends `POST {MAPLE_API_URL}/chat/completions` with `messages` containing: - - a dedicated **system prompt** (β€œYou are an image description agent… transcribe text exactly…”) and - - a **user content array** with an `image_url` part plus a `text` part that includes recent conversation context + the user’s accompanying text. - - Uses `max_tokens=2048` and returns the model’s text description. -6. **Inject description into the user message** as a bracketed suffix/payload (`[Uploaded Image: …]`) and pass the resulting **text-only** `user_message` into the DSRs agent step loop (`agent_guard.step(&user_message, ...)`). -7. **Persist both the original text and the derived description**: - - Sage stores the original user text in `messages.content` and stores the derived description in `messages.attachment_text` (migration: `crates/sage-core/migrations/2026-02-05-000000_add_attachment_text`). - - It updates the message embedding asynchronously using the *combined* text (original message + injected description) so recall search can match on image content. -8. **Render attachment descriptions in agent context**: when building `recent_conversation` for the DSRs signature, Sage appends `attachment_text` alongside the user message (`crates/sage-core/src/sage_agent.rs`, β€œRender attachment_text alongside user messages”). - -**Key implication for Maple/OpenSecret agent:** DSRs stays **text-in/text-out**; β€œimage understanding” is an explicit **pre-step** that produces storable text (and in OpenSecret, that derived text should be encrypted at rest) that can participate in compaction, embeddings, and memory tools. - -**Notable Sage quirks (worth copying or fixing intentionally):** -- Only the **first** supported image attachment is processed per inbound message. -- On vision failure, Sage injects a placeholder description (so the agent still sees β€œsomething happened”). -- `conversation_search` results are formatted from `messages.content` and (today) do **not** include `attachment_text`, so image-only messages can look blank in tool output even though the agent’s in-context transcript includes the description. +This should work the same way for the main agent and for subagents. --- -## 6. Agent Step Loop - -The agent processes messages through a multi-step loop, following Sage's proven pattern: +## 6. Tools -``` -User sends message - -> Step 1: Build context, call LLM - -> LLM returns: messages to user + tool calls - -> Execute tool calls (memory_replace, archival_insert, web_search, etc.) - -> If tools were called: inject results, go to Step 2 - -> Step 2: Build context (with tool results), call LLM - -> LLM returns: messages to user + tool calls - -> ... repeat until LLM returns messages with no tool calls, or max steps reached - -> Send final messages to user -``` - -**Max steps:** 10 (same as Sage). Prevents infinite loops. - -**Tool execution:** Each tool call is persisted to the existing `tool_calls` and `tool_outputs` tables. Memory tools (`memory_replace`, `archival_insert`, etc.) also modify the memory tables directly. - -### 6.2 Structured Output + Parse Correction (Recommended) - -To maximize reliability across providers/models (and to avoid native tool calling issues), align with the Sage V2 prototype: - -- The agent returns structured `messages[]` and `tool_calls[]` via DSRs signatures + BAML parsing. -- If the LLM returns malformed output that fails to parse, run a **correction signature** that reshapes the raw text into the expected structure (preserving intent, not generating new content). -- Include a `done` tool as a no-op stop signal for tool-result continuation steps. - -This pattern dramatically reduces sensitivity to provider-specific tool calling bugs, while still supporting multi-step tool loops. - -**Implementation status:** Complete (AgentResponse + correction signatures + done tool implemented). +These remain the core tools for the shared-memory model: -### 6.3 DSRs Integration in OpenSecret (Nitro-safe LLM routing) - -DSRs must not call providers directly. All agent LLM calls (main steps + summarization + correction) must route through OpenSecret's existing LLM pipeline (billing, auth, retries, provider routing): `web::openai::get_chat_completion_response()`. - -Implementation options: -- **Preferred:** implement a custom DSRs LM backend/hook that calls `get_chat_completion_response()` and returns the model text + usage. -- **Fallback:** vendor/fork DSRs at a pinned commit to add the hook (so we control upgrades and avoid billing bypass regressions). +| Tool | Action | Storage | +|---|---|---| +| `memory_replace` | Replace text in a core memory block | `memory_blocks` | +| `memory_append` | Append text to a core memory block | `memory_blocks` | +| `memory_insert` | Insert text into a core memory block | `memory_blocks` | +| `archival_insert` | Store long-term memory | `user_embeddings` (`source_type='archival'`) | +| `archival_search` | Search long-term memory | `user_embeddings` (`source_type='archival'`) | +| `conversation_search` | Search embedded message history | `user_embeddings` (`source_type='message'`) | +| `spawn_subagent` | Create a topic-specific subagent chat | `conversations` + `agents` | +| `done` | Stop signal after tool-result continuation | no-op | -**Current status (implemented 2026-02-10):** OpenSecret now uses a locally-patched `dspy-rs` (DSRs) with a `LMClient::Custom(CustomCompletionModel)` hook that routes completions through `get_chat_completion_response()` (see `src/web/agent/signatures.rs`). Default `temperature=0.7` and `max_tokens=32768` match the Sage prototype. +### 6.1 Memory Tool Semantics -**Billing note:** embedding calls used by the agent system (archival insert/search, auto-embedding of agent messages, summary embeddings) must also route through OpenSecret's billing-aware embedding pipeline (`web::get_embedding_vector()`), and must never call providers directly. +For v1, core memory tools should operate on the built-in shared blocks only. -**Implementation status:** Complete (DSRs LM hook routes via `get_chat_completion_response()`; embeddings route via billing-aware `get_embedding_vector()`). +That means: +- `persona` +- `human` -### 6.1 Memory Tools +This is still enough to get the product benefits of shared always-in-context memory without exposing a general-purpose user-editable block system. -These are the agent's interface to the memory system. Following Sage's design: +### 6.2 `conversation_search` Visibility -| Tool | Action | Storage | -|---|---|---| -| `memory_replace` | Replace text in a core memory block | UPDATE `memory_blocks` | -| `memory_append` | Append text to a core memory block | UPDATE `memory_blocks` | -| `memory_insert` | Insert text at a specific line in a block | UPDATE `memory_blocks` | -| `archival_insert` | Store information in long-term memory | INSERT into `user_embeddings` (source_type='archival') | -| `archival_search` | Search long-term memory semantically (optional tag filter) | Query `user_embeddings` (source_type='archival') with brute-force cosine similarity | -| `conversation_search` | Search conversation history semantically | Query `user_embeddings` (source_type='message') -- **all conversations by default**, not just the agent's own thread | +By default, `conversation_search` should search across all eligible message embeddings for the user: +- main agent thread +- all subagent threads +- Responses API threads -All memory tool inputs/outputs are encrypted before storage. The tool execution happens in-enclave where the user key is available. +The tool may optionally accept a `conversation_id` filter to scope to one thread. -**`conversation_search` visibility note:** By default, this tool searches across all of the user's embedded messages, including Responses API threads. This is intentional -- the agent should surface relevant context regardless of which API surface generated it. The tool can accept an optional `conversation_id` parameter to scope to a specific thread if needed, but the default is broad. See the RAG proposal's "Data Visibility and Isolation" section. +### 6.3 `spawn_subagent` -**Archival tags (implemented 2026-02-11):** `archival_insert` supports `metadata.tags` (string or string[]). Tags are normalized (trim + lowercase), deterministically encrypted per-tag (base64) and stored in `user_embeddings.tags_enc`. `archival_search` supports a `tags` argument and applies an ANY-match SQL filter (`tags_enc && `) backed by a partial GIN index. +Subagent creation should be a first-class operation. -**Implementation status:** Partially implemented (core memory + archival + conversation_search implemented; web_search tool not implemented yet). +A `spawn_subagent` tool should: +- create a new conversation row +- create a new `agents` row with `kind='subagent'` +- link it to the main agent via `parent_agent_id` +- store a display name and/or purpose +- return identifiers the client can use to hand off the user into that chat -**Overall implementation status:** Partially implemented (step loop + tool persistence implemented; web_search + other post-MVP tools pending). +This is how the main agent can proactively create a focused workspace for the user. --- ## 7. API Surface -### 7.1 New Routes (`/v1/agent/*`) +### 7.1 Public Agent-Facing Routes -**Current status (implemented 2026-02-11):** the MVP `/v1/agent/*` surface below is implemented (Local/Dev only). `POST /v1/agent/chat` streams message-level SSE (not token streaming); other endpoints are encrypted JSON. +The public product surface should focus on the main agent and subagent lifecycle, not user configuration. +```text +POST /v1/agent/chat -- chat with the main agent (request-scoped SSE) +POST /v1/agent/subagents -- create a new subagent +POST /v1/agent/subagents/:id/chat -- chat with a subagent (request-scoped SSE) +DELETE /v1/agent/subagents/:id -- delete a subagent +GET /v1/agent/events -- long-lived SSE for proactive delivery (post-MVP) ``` -POST /v1/agent/chat -- Send a message to the agent (step loop, request-scoped SSE) -GET /v1/agent/config -- Get agent settings -PUT /v1/agent/config -- Update agent settings (model, system prompt, etc.) -GET /v1/agent/memory/blocks -- List all memory blocks -GET /v1/agent/memory/blocks/:label -- Get a specific block -PUT /v1/agent/memory/blocks/:label -- Manually edit a block +### 7.2 Routes That Should Not Be Public Product API -POST /v1/agent/memory/archival -- Manually insert archival memory -POST /v1/agent/memory/search -- Search archival + recall memory -DELETE /v1/agent/memory/archival/:id -- Delete specific archival entry +These may exist temporarily as local debugging surfaces, but they should not be part of the target user-facing architecture: +- manual agent config update endpoints +- manual memory block CRUD endpoints +- manual archival insert/delete endpoints for normal product usage -GET /v1/agent/conversations -- List agent conversations (reuses conversations table) -GET /v1/agent/conversations/:id/items -- Get conversation items (reuses existing item format) -DELETE /v1/agent/conversations/:id -- Delete conversation + associated summaries +The product intent is **agent-managed memory**, not β€œsettings panels for the user to tune Sage.” -GET /v1/agent/events -- Long-lived SSE channel for proactive agent delivery + fan-out (post-MVP) -``` +### 7.3 Conversation List Integration -**Implementation status:** Partially implemented (all MVP routes except `GET /v1/agent/events`; Local/Dev only). +`/v1/conversations/*` should become the list/read/delete surface for: +- Responses API threads +- subagent chats -### 7.2 Chat Endpoint Detail +It should **exclude**: +- the main agent thread -**Current implementation (2026-02-11):** `POST /v1/agent/chat` accepts `{ "input": "..." }` (encrypted request body like other endpoints) and returns request-scoped SSE (message-level events, not token streaming). +The API should expose enough metadata to distinguish conversation kinds in the UI, e.g.: +- `response` +- `subagent` -Event types (payloads are JSON, encrypted per-session and base64-encoded in the `data:` field): +That distinction should come from the `agents` join / derived server response, not from making encrypted conversation metadata the canonical source of truth. -| Event type | Meaning | -|---|---| -| `agent.typing` | Agent is working on step `step` | -| `agent.message` | One or more user-visible messages for step `step` | -| `agent.done` | Turn completed (`total_steps`, `total_messages`) | -| `agent.error` | Turn failed (`error`) | +### 7.4 Chat Event Model -``` -POST /v1/agent/chat -{ - "input": "What did we discuss about deployment strategies last week?" -} -``` - -**Implementation status:** Complete (request-scoped, message-level SSE with typing/message/done/error events). - -### 7.3 Proactive Agent Delivery: Long-Lived SSE Channel (Post-MVP) - -`POST /v1/agent/chat` uses request-scoped SSE: the stream opens on request, delivers step-loop events (`agent.typing`, `agent.message`, `agent.done`), and closes. This is correct for request-response interactions but insufficient for proactive agent behavior (reminders, scheduled task results, agent-initiated messages). - -**Decision: long-lived SSE over WebSockets.** Three reasons specific to our architecture: +The SSE event model can stay the same for both main agent and subagents: +- `agent.typing` +- `agent.message` +- `agent.done` +- `agent.error` -1. **Encryption model fit.** The per-session encryption middleware operates on HTTP request/response cycles. SSE is still HTTP and slots in naturally. WebSockets would require rethinking per-message encryption (WS frames aren't HTTP responses, so `encrypt_event` / session key lookup needs a different path). -2. **Proxy chain simplicity.** tinfoil-proxy and continuum-proxy sit in front of the enclave via vsock. HTTP/SSE flows through standard reverse proxies. WebSocket upgrade handling through that chain adds operational complexity -- sticky sessions, connection draining, timeout tuning at every hop. -3. **Unidirectionality is sufficient.** Proactive agent messages are server-to-client. The client already has `POST /v1/agent/chat` for the other direction. +The difference is only which agent identity / conversation thread is being driven. -**Proposed endpoint:** `GET /v1/agent/events` - -The client opens this connection once and keeps it open. The server pushes **the same `agent.*` SSE events used by `/v1/agent/chat`** (`agent.typing`, `agent.message`, `agent.done`, `agent.error`) whenever the agent produces output outside of an active chat request. - -**Resilience requirements:** -- **At-least-once delivery.** Clients MUST dedupe by SSE event `id`. -- **SSE `id` support + `Last-Event-ID` resumption.** Every non-heartbeat event MUST include an SSE `id:` field. On reconnect, clients send `Last-Event-ID` to resume from the next event. -- **DB-backed event log (fan-out).** Events are persisted so clients that were offline can catch up on reconnect. Because this is fan-out, events are retained by TTL (e.g., 24h) rather than deleted-on-delivery. -- **Heartbeat keepalive.** Emit SSE comment frames (e.g., `: ping\n\n`) every ~30s to prevent proxy idle timeouts. +--- -**Relationship to `POST /v1/agent/chat`:** The two channels are independent. Chat continues to use request-scoped SSE for immediate step-loop delivery. The long-lived channel handles async/proactive events only. A client that only uses chat (no proactive features) never needs to open `/v1/agent/events`. +## 8. Relationship to Other API Surfaces -**Implementation status:** Not implemented yet. +```text +/v1/responses/* + Stateless chat surface + Uses conversations + message tables + Threads remain visible in /v1/conversations + Eligible for shared recall embedding when policy allows -### 7.4 Relationship to Responses API +/v1/agent/chat + Main agent home surface + Uses agents + shared memory + thread-local summaries + Main thread is hidden from /v1/conversations -The two API surfaces are independent but share storage: +/v1/agent/subagents/* + Subagent lifecycle + chat surface + Uses agents + shared memory + thread-local summaries + Subagent threads are visible in /v1/conversations -``` -/v1/responses/* -- Stateless chat API (existing, unchanged) - reads/writes: conversations, user_messages, assistant_messages, - tool_calls, tool_outputs, reasoning_items, user_instructions - triggers: (TODO) async embedding into user_embeddings (store=true and non-private threads only) - -/v1/agent/* -- Persistent agent API (new) - reads/writes: conversations, user_messages, assistant_messages, - tool_calls, tool_outputs, reasoning_items - also reads/writes: memory_blocks, user_embeddings, - conversation_summaries, agent_config - triggers: async embedding into user_embeddings for agent messages (implemented) - searches: user_embeddings across ALL user conversations (broad default) +/v1/conversations/* + Unified list/read/delete surface for response threads + subagent chats + Does not expose the main agent thread ``` -A user can use both APIs simultaneously. The key data flow: - -- **Responses API -> Agent visibility (pending):** Responses API threads will be auto-embedded into `user_embeddings` when eligible (stored + not-private). Today, only the agent's own messages are auto-embedded. -- **Agent conversation visibility:** The agent conversation is stored in the shared tables, but the **main agent thread should be hidden from `/v1/conversations/*`** (Section 2.2). Clients should use `/v1/agent/conversations/*` to access agent threads. -- **Isolation boundary:** `store=false` requests and future incognito/private threads on the Responses API will NOT be embedded. The agent cannot see them. This is the opt-out mechanism. - -**Implementation status:** Partially implemented (agent messages auto-embedded; Responses API auto-embedding + private/store=false exclusions not implemented yet). - -**Overall implementation status:** Partially implemented (MVP /v1/agent surface shipped Local/Dev; proactive events channel pending). +The key data-flow rule is: +- all eligible stored messages from all three surfaces feed the same user-level recall memory +- all agents for that user can search that recall memory +- only thread-local history and summaries stay isolated per agent --- -## 8. What We're Not Deciding Yet - -These are intentionally deferred: - -- **Native tool calling in the agent loop.** The Responses API will continue using native tool calling via the Tinfoil proxy. For the agent loop, prefer DSRs signatures + BAML parsing (Section 6.2) for reliability; native tool calling can be revisited later. - -- **GEPA operationalization in production.** MVP should be GEPA-*ready*: consume a GEPA-optimized instruction stored in `agent_config.system_prompt_enc`. Running GEPA itself can be feature-flagged/offline until we have safe evaluation + trace capture. - -- **Multi-agent per user (subagents).** MVP assumes one agent per user. When subagents are in scope, prefer an `agents` table + `agent_id` FKs on agent-specific tables (Section 4.5) so memory can be shared or isolated per agent. The simpler approach (remove `UNIQUE(user_id)` + add `name` on `agent_config`) is viable but less expressive for memory scoping. - -- **Group/shared memory.** All memory is per-user. Shared memory blocks between users in the same project would require a new sharing model. Deferred. +## 9. What We Are Not Deciding Yet -- **Voice/image input handling.** Sage converts image attachments into injected text via a separate vision pre-processing call (Section 5.4). Maple’s Responses API supports multimodal input directly, but the agent runtime likely needs its own normalization path if it wants Sage-style text-only context + embeddings. +These remain intentionally deferred: +- isolated subagent memory +- nested subagent trees beyond the β€œmain agent -> subagent” model +- per-agent custom model selection +- per-agent custom prompt editing by users +- reminders / scheduling / background task execution +- code sandbox execution +- private/incognito threads excluding embedding +- agent-backed Responses API threads -- **Reminders / scheduling.** Post-MVP. Sage V2 has a working scheduler + tools (`schedule_task`, `list_schedules`, `cancel_schedule`) backed by a `scheduled_tasks` table. We can port this once the core agent loop is stable. Delivery of fired reminders will use the long-lived SSE channel (`GET /v1/agent/events`, Section 7.3). - -- **Code sandbox execution.** Long-lead, deferred. Requires a secure execution substrate (likely isolated runtime) and careful Nitro threat modeling. - -- **Incognito / private threads.** A `private` flag on Responses API conversations that excludes their messages from embedding (and thus from agent visibility). In addition, `store=false` should suppress embedding as a per-request opt-out. The conversation-level `private` flag is deferred to post-v1. A UI toggle to set incognito as the default for new threads is further out. - -- **Agent-backed Responses API threads.** The long-term vision where Responses API threads inherit agent memory (memory blocks + search) but start with a clean conversation history. Requires proving out the agent system first. See Section 1.3. - -**Overall implementation status:** Not implemented yet (intentionally deferred). +The important v1 decision is already made: **subagents share memory and search, but not transcript history.** --- -## 9. Implementation Ordering +## 10. Implementation Ordering -A suggested sequence, building on the RAG layer as the foundation: +Because this is still local-only, we should update the existing Sage docs and schema directly instead of preserving the old MVP shape. ### Phase 1: RAG Foundation -- [x] Implement `user_embeddings` table + migration -- [x] Implement brute-force search + LRU cache -- [x] Implement `/v1/rag/*` API endpoints -- [ ] Wire up async embedding generation after message creation in Responses API (respect `store=false` and future `private` threads) -- **Milestone:** Recall + archival memory storage/search exists (even before the agent loop ships) - -**Implementation status:** Partially implemented (RAG + /v1/rag done; Responses API auto-embedding not implemented yet). - -### Phase 2: Agent Storage + Feature Flags -- [x] Implement `memory_blocks`, `conversation_summaries`, `agent_config` tables + migrations -- [x] Implement Diesel models for new tables -- [x] Add `agent_config.conversation_id` pointer to the user's persistent agent thread -- [x] Implement memory block CRUD operations -- [ ] Implement global + per-user feature flagging for `/v1/agent/*` and background jobs -- **Milestone:** Agent storage layer exists and is tested - -**Implementation status:** Partially implemented (storage tables/models done; feature flags not implemented yet). - -### Phase 3: Agent LLM Runtime (DSRs + BAML) -- [x] Integrate DSRs signatures + BAML parsing (AgentResponse) -- [x] Integrate correction signatures for parse repair -- [x] Implement Nitro-safe DSRs LM routing via `get_chat_completion_response()` -- [x] Implement summarization signatures for compaction -- **Milestone:** Agent can reliably produce `messages[]` + `tool_calls[]` without native tool calling - -**Implementation status:** Complete. - -### Phase 4: Agent Context Builder + Compaction -- [x] Implement regenerated agent context builder (currently in `src/web/agent/runtime.rs`) -- [x] Implement compaction (LLM summarization of old messages) into `conversation_summaries` -- **Milestone:** Agent can sustain long-running conversations without truncation - -**Implementation status:** Complete. - -### Phase 5: Memory Tools + Multi-Step Agent Loop -- [x] Implement agent tool registry (core memory + archival + conversation search + web search) -- [x] Implement multi-step execution loop (max steps, tool-result continuation) -- **Milestone:** DSRs-based agent can have multi-turn conversations with memory tool use - -**Implementation status:** Partially implemented (memory tools + step loop done; web_search tool not implemented yet). - -### Phase 6: Agent API Surface -- [x] Implement `/v1/agent/chat` (request-scoped SSE) -- [x] Implement remaining `/v1/agent/*` endpoints (config, memory management, conversations) -- [x] Wire up memory block initialization + agent thread creation on opt-in -- [x] Wire up async embedding generation after agent message creation -- **Milestone:** Full agent API surface available - -**Implementation status:** Complete (Local/Dev only). +- Keep `user_embeddings` +- Keep brute-force search + cache +- Keep experimental `/v1/rag/*` endpoints for validation if useful +- Add Responses API auto-embedding when the agent rewrite lands + +### Phase 2: Schema Reset (edit local-only migrations in place) +- Replace `agent_config` with `agents` +- Simplify `memory_blocks` +- Keep `conversation_summaries` +- Keep `user_embeddings` shared per user +- Do **not** add per-agent memory scoping columns + +### Phase 3: Main Agent Runtime Rewrite +- Load the main agent from `agents` +- Move model/prompt/tuning defaults into code +- Remove user-facing config assumptions +- Keep regenerated context + compaction + +### Phase 4: Subagent Lifecycle +- Create subagents through API and tool flow +- Give each subagent its own conversation + summary chain +- Inject subagent purpose into prompt assembly + +### Phase 5: Conversation List Integration +- Hide the main agent thread from `/v1/conversations/*` +- Show subagent threads there +- Continue showing Responses API threads +- Add derived conversation kind metadata for UI rendering + +### Phase 6: Remove Non-Product Surfaces +- Remove or internalize manual config endpoints +- Remove or internalize manual memory block endpoints +- Keep only the public surfaces that match the product model ### Phase 7: Post-MVP -- Long-lived SSE event channel (`GET /v1/agent/events`) -- proactive notification delivery (Section 7.3) -- Reminders/scheduling (`scheduled_tasks` + scheduler loop + tools), delivered via the event channel -- Monitoring/metrics + tuning -- (Optional) agent-backed Responses API threads (Section 1.3) -- Code sandboxes (separate major project) - -**Implementation status:** Not implemented yet. +- Long-lived SSE event channel +- reminders / scheduling +- background delegation patterns +- optional future scoped memory if product requirements change -**Overall implementation status:** Partially implemented (Phases 1-6 mostly complete; Phase 7 items pending). +**Overall implementation status:** target architecture documented; code rewrite still required. diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql index db111084..21a3e325 100644 --- a/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql @@ -1,5 +1,8 @@ -DROP TRIGGER IF EXISTS update_agent_config_updated_at ON agent_config; -DROP TABLE IF EXISTS agent_config; +DROP TRIGGER IF EXISTS update_agents_updated_at ON agents; +DROP INDEX IF EXISTS idx_agents_parent_agent_id; +DROP INDEX IF EXISTS idx_agents_user_kind_created; +DROP INDEX IF EXISTS idx_agents_one_main_per_user; +DROP TABLE IF EXISTS agents; DROP INDEX IF EXISTS idx_conversation_summaries_chain; DROP INDEX IF EXISTS idx_conversation_summaries_user_conv; diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql index bd1980ed..2753e6bd 100644 --- a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql +++ b/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql @@ -8,11 +8,7 @@ CREATE TABLE memory_blocks ( user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, label TEXT NOT NULL, - description TEXT, value_enc BYTEA NOT NULL, - char_limit INTEGER NOT NULL DEFAULT 5000, - read_only BOOLEAN NOT NULL DEFAULT FALSE, - version INTEGER NOT NULL DEFAULT 1, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -54,25 +50,40 @@ CREATE INDEX idx_conversation_summaries_user_conv CREATE INDEX idx_conversation_summaries_chain ON conversation_summaries(previous_summary_id); -CREATE TABLE agent_config ( - id BIGSERIAL PRIMARY KEY, - uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, - user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE UNIQUE, - - conversation_id BIGINT REFERENCES conversations(id) ON DELETE SET NULL, +CREATE TABLE agents ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + conversation_id BIGINT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + kind TEXT NOT NULL, + parent_agent_id BIGINT REFERENCES agents(id) ON DELETE SET NULL, + display_name_enc BYTEA, + purpose_enc BYTEA, + created_by TEXT NOT NULL DEFAULT 'user', + + created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT agents_kind_check CHECK (kind IN ('main', 'subagent')), + CONSTRAINT agents_created_by_check CHECK (created_by IN ('user', 'agent')), + CONSTRAINT agents_parent_check CHECK ( + (kind = 'main' AND parent_agent_id IS NULL) + OR (kind = 'subagent' AND parent_agent_id IS NOT NULL) + ), + UNIQUE(conversation_id) +); - enabled BOOLEAN NOT NULL DEFAULT FALSE, - model TEXT NOT NULL DEFAULT 'kimi-k2-5', - max_context_tokens INTEGER NOT NULL DEFAULT 256000, - compaction_threshold REAL NOT NULL DEFAULT 0.80, +CREATE UNIQUE INDEX idx_agents_one_main_per_user + ON agents(user_id) + WHERE kind = 'main'; - system_prompt_enc BYTEA, - preferences_enc BYTEA, +CREATE INDEX idx_agents_user_kind_created + ON agents(user_id, kind, created_at DESC); - created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP -); +CREATE INDEX idx_agents_parent_agent_id + ON agents(parent_agent_id); -CREATE TRIGGER update_agent_config_updated_at -BEFORE UPDATE ON agent_config +CREATE TRIGGER update_agents_updated_at +BEFORE UPDATE ON agents FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/models/agent_config.rs b/src/models/agent_config.rs deleted file mode 100644 index 633beb10..00000000 --- a/src/models/agent_config.rs +++ /dev/null @@ -1,88 +0,0 @@ -use crate::models::schema::agent_config; -use chrono::{DateTime, Utc}; -use diesel::prelude::*; -use serde::{Deserialize, Serialize}; -use thiserror::Error; -use uuid::Uuid; - -#[derive(Error, Debug)] -pub enum AgentConfigError { - #[error("Database error: {0}")] - DatabaseError(#[from] diesel::result::Error), -} - -#[derive(Queryable, Identifiable, AsChangeset, Serialize, Deserialize, Clone, Debug)] -#[diesel(table_name = agent_config)] -pub struct AgentConfig { - pub id: i64, - pub uuid: Uuid, - pub user_id: Uuid, - pub conversation_id: Option, - pub enabled: bool, - pub model: String, - pub max_context_tokens: i32, - pub compaction_threshold: f32, - pub system_prompt_enc: Option>, - pub preferences_enc: Option>, - pub created_at: DateTime, - pub updated_at: DateTime, -} - -impl AgentConfig { - pub fn get_by_user_id( - conn: &mut PgConnection, - lookup_user_id: Uuid, - ) -> Result, AgentConfigError> { - agent_config::table - .filter(agent_config::user_id.eq(lookup_user_id)) - .first::(conn) - .optional() - .map_err(AgentConfigError::DatabaseError) - } - - pub fn delete_for_user( - conn: &mut PgConnection, - lookup_user_id: Uuid, - ) -> Result { - diesel::delete(agent_config::table.filter(agent_config::user_id.eq(lookup_user_id))) - .execute(conn) - .map_err(AgentConfigError::DatabaseError) - } -} - -#[derive(Insertable, Debug, Clone)] -#[diesel(table_name = agent_config)] -pub struct NewAgentConfig { - pub uuid: Uuid, - pub user_id: Uuid, - pub conversation_id: Option, - pub enabled: bool, - pub model: String, - pub max_context_tokens: i32, - pub compaction_threshold: f32, - pub system_prompt_enc: Option>, - pub preferences_enc: Option>, -} - -impl NewAgentConfig { - pub fn insert_or_update( - &self, - conn: &mut PgConnection, - ) -> Result { - diesel::insert_into(agent_config::table) - .values(self) - .on_conflict(agent_config::user_id) - .do_update() - .set(( - agent_config::conversation_id.eq(self.conversation_id), - agent_config::enabled.eq(self.enabled), - agent_config::model.eq(self.model.clone()), - agent_config::max_context_tokens.eq(self.max_context_tokens), - agent_config::compaction_threshold.eq(self.compaction_threshold), - agent_config::system_prompt_enc.eq(self.system_prompt_enc.clone()), - agent_config::preferences_enc.eq(self.preferences_enc.clone()), - )) - .get_result::(conn) - .map_err(AgentConfigError::DatabaseError) - } -} 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/memory_blocks.rs b/src/models/memory_blocks.rs index ff399a29..744d46f2 100644 --- a/src/models/memory_blocks.rs +++ b/src/models/memory_blocks.rs @@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize}; use thiserror::Error; use uuid::Uuid; -/// Default character limit per core memory block (mirrors Sage/Letta). -pub const DEFAULT_BLOCK_CHAR_LIMIT: i32 = 20_000; +pub const MEMORY_BLOCK_LABEL_PERSONA: &str = "persona"; +pub const MEMORY_BLOCK_LABEL_HUMAN: &str = "human"; #[derive(Error, Debug)] pub enum MemoryBlockError { @@ -21,11 +21,7 @@ pub struct MemoryBlock { pub uuid: Uuid, pub user_id: Uuid, pub label: String, - pub description: Option, pub value_enc: Vec, - pub char_limit: i32, - pub read_only: bool, - pub version: i32, pub created_at: DateTime, pub updated_at: DateTime, } @@ -85,29 +81,16 @@ pub struct NewMemoryBlock { pub uuid: Uuid, pub user_id: Uuid, pub label: String, - pub description: Option, pub value_enc: Vec, - pub char_limit: i32, - pub read_only: bool, - pub version: i32, } impl NewMemoryBlock { - pub fn new( - user_id: Uuid, - label: impl Into, - description: Option, - value_enc: Vec, - ) -> Self { + pub fn new(user_id: Uuid, label: impl Into, value_enc: Vec) -> Self { NewMemoryBlock { uuid: Uuid::new_v4(), user_id, label: label.into(), - description, value_enc, - char_limit: DEFAULT_BLOCK_CHAR_LIMIT, - read_only: false, - version: 1, } } @@ -119,13 +102,7 @@ impl NewMemoryBlock { .values(self) .on_conflict((memory_blocks::user_id, memory_blocks::label)) .do_update() - .set(( - memory_blocks::description.eq(self.description.clone()), - memory_blocks::value_enc.eq(self.value_enc.clone()), - memory_blocks::char_limit.eq(self.char_limit), - memory_blocks::read_only.eq(self.read_only), - memory_blocks::version.eq(memory_blocks::version + 1), - )) + .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 5127021c..79a3578f 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,5 +1,5 @@ pub mod account_deletion; -pub mod agent_config; +pub mod agents; pub mod conversation_summaries; pub mod email_verification; pub mod enclave_secrets; diff --git a/src/models/schema.rs b/src/models/schema.rs index 83ea7ab4..760ffb23 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -22,17 +22,16 @@ diesel::table! { } diesel::table! { - agent_config (id) { + agents (id) { id -> Int8, uuid -> Uuid, user_id -> Uuid, - conversation_id -> Nullable, - enabled -> Bool, - model -> Text, - max_context_tokens -> Int4, - compaction_threshold -> Float4, - system_prompt_enc -> Nullable, - preferences_enc -> Nullable, + conversation_id -> Int8, + kind -> Text, + parent_agent_id -> Nullable, + display_name_enc -> Nullable, + purpose_enc -> Nullable, + created_by -> Text, created_at -> Timestamptz, updated_at -> Timestamptz, } @@ -122,11 +121,7 @@ diesel::table! { uuid -> Uuid, user_id -> Uuid, label -> Text, - description -> Nullable, value_enc -> Bytea, - char_limit -> Int4, - read_only -> Bool, - version -> Int4, created_at -> Timestamptz, updated_at -> Timestamptz, } @@ -447,7 +442,7 @@ diesel::table! { } } -diesel::joinable!(agent_config -> conversations (conversation_id)); +diesel::joinable!(agents -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> responses (response_id)); diesel::joinable!(conversation_summaries -> conversations (conversation_id)); @@ -475,7 +470,7 @@ diesel::joinable!(users -> org_projects (project_id)); diesel::allow_tables_to_appear_in_same_query!( account_deletion_requests, - agent_config, + agents, assistant_messages, conversation_summaries, conversations, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 21bd8e8a..183761a0 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,36 +1,24 @@ use axum::{ - extract::{Path, Query, State}, + extract::{Path, State}, middleware::from_fn_with_state, response::sse::{Event, Sse}, - routing::{delete, get, post, put}, + routing::{delete, post}, Extension, Json, Router, }; -use diesel::prelude::*; use futures::Stream; use serde::{Deserialize, Serialize}; -use serde_json::{json, Value}; -use std::{sync::Arc, time::Duration}; +use std::{convert::Infallible, sync::Arc, time::Duration}; use tokio::time::sleep; use tracing::{error, warn}; use uuid::Uuid; -use crate::encrypt::{decrypt_content, decrypt_string, encrypt_with_key}; -use crate::models::agent_config::{AgentConfig, NewAgentConfig}; -use crate::models::memory_blocks::{MemoryBlock, NewMemoryBlock, DEFAULT_BLOCK_CHAR_LIMIT}; -use crate::models::responses::{Conversation, NewConversation}; -use crate::models::schema::user_embeddings; +use crate::models::agents::{Agent, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT}; 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::constants::{DEFAULT_PAGINATION_LIMIT, MAX_PAGINATION_LIMIT}; -use crate::web::responses::conversations::{ - ConversationItemListResponse, ConversationListResponse, ConversationResponse, ListItemsParams, -}; use crate::web::responses::handlers::encrypt_event; use crate::web::responses::{ - error_mapping, ConversationBuilder, ConversationItem, ConversationItemConverter, - DeletedObjectResponse, MessageContent, MessageContentConverter, MessageContentPart, Paginator, + error_mapping, DeletedObjectResponse, MessageContent, MessageContentConverter, + MessageContentPart, }; use crate::{ApiError, AppMode, AppState}; @@ -45,6 +33,29 @@ 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(), @@ -117,71 +128,6 @@ async fn encrypt_agent_event( encrypt_event(state, session_id, event_type, &value).await } -#[derive(Debug, Clone, Serialize)] -struct AgentConfigResponse { - enabled: bool, - model: String, - max_context_tokens: i32, - compaction_threshold: f32, - system_prompt: Option, - conversation_id: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct UpdateAgentConfigRequest { - enabled: Option, - model: Option, - max_context_tokens: Option, - compaction_threshold: Option, - system_prompt: Option, -} - -#[derive(Debug, Clone, Serialize)] -struct MemoryBlockResponse { - label: String, - description: Option, - value: String, - char_limit: i32, - read_only: bool, - version: i32, -} - -#[derive(Debug, Clone, Deserialize)] -struct UpdateMemoryBlockRequest { - description: Option, - value: Option, - char_limit: Option, - read_only: Option, -} - -#[derive(Debug, Clone, Deserialize)] -struct InsertArchivalRequest { - text: String, - metadata: Option, -} - -#[derive(Debug, Clone, Serialize)] -struct InsertArchivalResponse { - id: Uuid, - source_type: String, - embedding_model: String, - token_count: i32, - created_at: chrono::DateTime, -} - -#[derive(Debug, Clone, Deserialize)] -struct MemorySearchRequest { - query: String, - top_k: Option, - max_tokens: Option, - source_types: Option>, -} - -#[derive(Debug, Clone, Serialize)] -struct MemorySearchResponse { - 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) { @@ -191,82 +137,80 @@ pub fn router(app_state: Arc) -> Router<()> { Router::new() .route( "/v1/agent/chat", - post(chat).layer(from_fn_with_state( + post(chat_main).layer(from_fn_with_state( app_state.clone(), decrypt_request::, )), ) .route( - "/v1/agent/config", - get(get_config).layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/config", - put(update_config).layer(from_fn_with_state( - app_state.clone(), - decrypt_request::, - )), - ) - .route( - "/v1/agent/memory/blocks", - get(list_memory_blocks) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/memory/blocks/:label", - get(get_memory_block) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/memory/blocks/:label", - put(update_memory_block).layer(from_fn_with_state( - app_state.clone(), - decrypt_request::, - )), - ) - .route( - "/v1/agent/memory/archival", - post(insert_archival).layer(from_fn_with_state( + "/v1/agent/subagents", + post(create_subagent).layer(from_fn_with_state( app_state.clone(), - decrypt_request::, + decrypt_request::, )), ) .route( - "/v1/agent/memory/archival/:id", - delete(delete_archival) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/memory/search", - post(memory_search).layer(from_fn_with_state( + "/v1/agent/subagents/:id/chat", + post(chat_subagent).layer(from_fn_with_state( app_state.clone(), - decrypt_request::, + decrypt_request::, )), ) .route( - "/v1/agent/conversations", - get(list_agent_conversations) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/conversations/:id/items", - get(list_agent_conversation_items) - .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), - ) - .route( - "/v1/agent/conversations/:id", - delete(delete_agent_conversation) + "/v1/agent/subagents/:id", + delete(delete_subagent) .layer(from_fn_with_state(app_state.clone(), decrypt_request::<()>)), ) .with_state(app_state) } -async fn chat( +async fn chat_main( State(state): State>, Extension(session_id): Extension, Extension(user): Extension, Extension(body): Extension, -) -> Result>>, ApiError> { +) -> Result>>, ApiError> { + chat_with_target(state, session_id, user, body, ChatTarget::Main).await +} + +async fn chat_subagent( + State(state): State>, + Path(agent_uuid): Path, + Extension(session_id): Extension, + Extension(user): 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, + body, + ChatTarget::Subagent(agent_uuid), + ) + .await +} + +async fn chat_with_target( + state: Arc, + session_id: Uuid, + user: User, + 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) { @@ -316,31 +260,40 @@ async fn chat( } }; - let mut runtime = - match runtime::AgentRuntime::new(state.clone(), user.clone(), user_key).await { - 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, - ) + 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")) - } + { + Ok(event) => yield Ok(event), + Err(_) => { + yield Ok(Event::default().event(EVENT_AGENT_ERROR).data("Error")) } - had_error = true; - break 'init; } - }; + had_error = true; + break 'init; + } + }; match runtime.prepare(&input_content).await { Ok(prepared) => { @@ -512,217 +465,13 @@ async fn chat( Ok(Sse::new(event_stream)) } -async fn get_config( +async fn create_subagent( State(state): State>, Extension(session_id): Extension, Extension(user): Extension, -) -> Result>, ApiError> { - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let cfg = get_or_create_agent_config(&mut conn, user.uuid).await?; - - let system_prompt = decrypt_string(&user_key, cfg.system_prompt_enc.as_ref()).map_err(|e| { - error!("Failed to decrypt agent system prompt: {e:?}"); - ApiError::InternalServerError - })?; - - let conversation_uuid = if let Some(conversation_id) = cfg.conversation_id { - match Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) { - Ok(c) => Some(c.uuid), - Err(_) => None, - } - } else { - None - }; - - let response = AgentConfigResponse { - enabled: cfg.enabled, - model: cfg.model, - max_context_tokens: cfg.max_context_tokens, - compaction_threshold: cfg.compaction_threshold, - system_prompt, - conversation_id: conversation_uuid, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn update_config( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(body): Extension, -) -> Result>, ApiError> { - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let mut cfg = get_or_create_agent_config(&mut conn, user.uuid).await?; - - if let Some(enabled) = body.enabled { - cfg.enabled = enabled; - } - if let Some(model) = body.model { - if model.trim().is_empty() { - return Err(ApiError::BadRequest); - } - cfg.model = model; - } - if let Some(max_context_tokens) = body.max_context_tokens { - if max_context_tokens <= 0 { - return Err(ApiError::BadRequest); - } - cfg.max_context_tokens = max_context_tokens; - } - if let Some(threshold) = body.compaction_threshold { - cfg.compaction_threshold = threshold.clamp(0.5, 0.95); - } - - let system_prompt_enc = match body.system_prompt { - Some(p) if !p.trim().is_empty() => Some(encrypt_with_key(&user_key, p.as_bytes()).await), - Some(_) => None, - None => cfg.system_prompt_enc.clone(), - }; - - // If enabling and conversation is missing, initialize agent thread + blocks. - let conversation_id = if cfg.enabled && cfg.conversation_id.is_none() { - let metadata_enc = Some( - encrypt_with_key( - &user_key, - json!({"type":"agent_main"}).to_string().as_bytes(), - ) - .await, - ); - - let conversation = NewConversation { - uuid: Uuid::new_v4(), - user_id: user.uuid, - metadata_enc, - } - .insert(&mut conn) - .map_err(|e| { - error!("Failed to create agent conversation: {e:?}"); - ApiError::InternalServerError - })?; - - runtime::AgentRuntime::ensure_default_blocks(&state, &mut conn, &user_key, user.uuid) - .await?; - - Some(conversation.id) - } else { - cfg.conversation_id - }; - - let updated = NewAgentConfig { - uuid: cfg.uuid, - user_id: cfg.user_id, - conversation_id, - enabled: cfg.enabled, - model: cfg.model.clone(), - max_context_tokens: cfg.max_context_tokens, - compaction_threshold: cfg.compaction_threshold, - system_prompt_enc, - preferences_enc: cfg.preferences_enc.clone(), - } - .insert_or_update(&mut conn) - .map_err(|e| { - error!("Failed to update agent config: {e:?}"); - ApiError::InternalServerError - })?; - - let system_prompt = - decrypt_string(&user_key, updated.system_prompt_enc.as_ref()).map_err(|e| { - error!("Failed to decrypt agent system prompt: {e:?}"); - ApiError::InternalServerError - })?; - - let conversation_uuid = if let Some(conversation_id) = updated.conversation_id { - match Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) { - Ok(c) => Some(c.uuid), - Err(_) => None, - } - } else { - None - }; - - let response = AgentConfigResponse { - enabled: updated.enabled, - model: updated.model, - max_context_tokens: updated.max_context_tokens, - compaction_threshold: updated.compaction_threshold, - system_prompt, - conversation_id: conversation_uuid, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn list_memory_blocks( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>>, ApiError> { - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let blocks = MemoryBlock::get_all_for_user(&mut conn, user.uuid).map_err(|e| { - error!("Failed to list memory blocks: {e:?}"); - ApiError::InternalServerError - })?; - - let mut out: Vec = Vec::with_capacity(blocks.len()); - for b in blocks { - let value = decrypt_string(&user_key, Some(&b.value_enc)) - .map_err(|e| { - error!("Failed to decrypt memory block: {e:?}"); - ApiError::InternalServerError - })? - .unwrap_or_default(); - - out.push(MemoryBlockResponse { - label: b.label, - description: b.description, - value, - char_limit: b.char_limit, - read_only: b.read_only, - version: b.version, - }); - } - - encrypt_response(&state, &session_id, &out).await -} - -async fn get_memory_block( - State(state): State>, - Path(label): Path, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>, ApiError> { - if label.trim().is_empty() { + Extension(body): Extension, +) -> Result>, ApiError> { + if body.purpose.trim().is_empty() { return Err(ApiError::BadRequest); } @@ -737,484 +486,64 @@ async fn get_memory_block( .get() .map_err(|_| ApiError::InternalServerError)?; - let Some(block) = - MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, &label).map_err(|e| { - error!("Failed to load memory block: {e:?}"); - ApiError::InternalServerError - })? - else { - return Err(ApiError::NotFound); - }; - - let value = decrypt_string(&user_key, Some(&block.value_enc)) - .map_err(|_| ApiError::InternalServerError)? - .unwrap_or_default(); - - let response = MemoryBlockResponse { - label: block.label, - description: block.description, - value, - char_limit: block.char_limit, - read_only: block.read_only, - version: block.version, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn update_memory_block( - State(state): State>, - Path(label): Path, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(body): Extension, -) -> Result>, ApiError> { - if label.trim().is_empty() { - return Err(ApiError::BadRequest); - } - - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let existing = - MemoryBlock::get_by_user_and_label(&mut conn, user.uuid, &label).map_err(|e| { - error!("Failed to load memory block: {e:?}"); - ApiError::InternalServerError - })?; - - if let Some(b) = &existing { - if b.read_only { - return Err(ApiError::Unauthorized); - } - } - - let char_limit = body - .char_limit - .or_else(|| existing.as_ref().map(|b| b.char_limit)) - .unwrap_or(DEFAULT_BLOCK_CHAR_LIMIT); - - if char_limit <= 0 { - return Err(ApiError::BadRequest); - } - - if let Some(value) = body.value.as_ref() { - if value.len() > char_limit as usize { - return Err(ApiError::BadRequest); - } - } else if body.char_limit.is_some() { - if let Some(b) = &existing { - let existing_value = decrypt_string(&user_key, Some(&b.value_enc)) - .map_err(|_| ApiError::InternalServerError)? - .unwrap_or_default(); - - if existing_value.len() > char_limit as usize { - return Err(ApiError::BadRequest); - } - } - } - - let value_enc = match body.value { - Some(value) => encrypt_with_key(&user_key, value.as_bytes()).await, - None => match &existing { - Some(b) => b.value_enc.clone(), - None => encrypt_with_key(&user_key, b"").await, - }, - }; - - let new_block = NewMemoryBlock { - uuid: existing - .as_ref() - .map(|b| b.uuid) - .unwrap_or_else(Uuid::new_v4), - user_id: user.uuid, - label: label.clone(), - description: body - .description - .or_else(|| existing.as_ref().and_then(|b| b.description.clone())), - value_enc, - char_limit, - read_only: body - .read_only - .or_else(|| existing.as_ref().map(|b| b.read_only)) - .unwrap_or(false), - version: 1, - }; - - let updated = new_block.insert_or_update(&mut conn).map_err(|e| { - error!("Failed to upsert memory block: {e:?}"); - ApiError::InternalServerError - })?; - - let decrypted_value = decrypt_string(&user_key, Some(&updated.value_enc)) - .map_err(|_| ApiError::InternalServerError)? - .unwrap_or_default(); - - let response = MemoryBlockResponse { - label: updated.label, - description: updated.description, - value: decrypted_value, - char_limit: updated.char_limit, - read_only: updated.read_only, - version: updated.version, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn insert_archival( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(body): Extension, -) -> Result>, ApiError> { - 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.uuid, 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 = InsertArchivalResponse { - id: inserted.uuid, - source_type: inserted.source_type, - embedding_model: inserted.embedding_model, - token_count: inserted.token_count, - created_at: inserted.created_at, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn delete_archival( - State(state): State>, - Path(embedding_uuid): Path, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>, ApiError> { - use crate::models::schema::user_embeddings::dsl::*; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let source: Option = user_embeddings - .filter(user_id.eq(user.uuid)) - .filter(uuid.eq(embedding_uuid)) - .select(source_type) - .first::(&mut conn) - .optional() - .map_err(|_| ApiError::InternalServerError)?; - - match source.as_deref() { - Some(crate::rag::SOURCE_TYPE_ARCHIVAL) => {} - Some(_) => return Err(ApiError::NotFound), - None => return Err(ApiError::NotFound), - } - - rag::delete_user_embedding_by_uuid(&state, user.uuid, embedding_uuid).await?; - - let response = DeletedObjectResponse { - id: embedding_uuid, - object: "archival.deleted", - deleted: true, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn memory_search( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, - Extension(body): Extension, -) -> Result>, ApiError> { - if body.query.trim().is_empty() { - return Err(ApiError::BadRequest); - } - - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let source_types = body.source_types.unwrap_or_else(|| { - vec![ - crate::rag::SOURCE_TYPE_MESSAGE.to_string(), - crate::rag::SOURCE_TYPE_ARCHIVAL.to_string(), - ] - }); - - let results = rag::search_user_embeddings( - &state, - &user, - AuthMethod::Jwt, + 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, - &body.query, - body.top_k.unwrap_or(5), - body.max_tokens, - Some(&source_types), - None, - None, + user.uuid, + &main_agent, + body.display_name.as_deref(), + body.purpose.trim(), + AGENT_CREATED_BY_USER, ) .await?; - let response = MemorySearchResponse { results }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn list_agent_conversations( - State(state): State>, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>, ApiError> { - let user_key = state - .get_user_key(user.uuid, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - let cfg = AgentConfig::get_by_user_id(&mut conn, user.uuid) - .map_err(|_| ApiError::InternalServerError)?; - - let mut data: Vec = Vec::new(); - - if let Some(cfg) = cfg { - if let Some(conversation_id) = cfg.conversation_id { - if let Ok(conv) = - Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid) - { - let metadata = decrypt_content(&user_key, conv.metadata_enc.as_ref()) - .map_err(|_| ApiError::InternalServerError)?; - - data.push( - ConversationBuilder::from_conversation(&conv) - .metadata(metadata) - .build(), - ); - } - } - } - - let (first_id, last_id) = Paginator::get_cursor_ids(&data, |c| c.id); - let response = ConversationListResponse { - object: "list", - data, - has_more: false, - first_id, - last_id, + 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 list_agent_conversation_items( +async fn delete_subagent( State(state): State>, - Path(conversation_uuid): Path, - Query(params): Query, - Extension(session_id): Extension, - Extension(user): Extension, -) -> Result>, ApiError> { - let ctx = load_agent_conversation_context(&state, user.uuid, conversation_uuid).await?; - - 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( - ctx.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, - &ctx.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, - }); - - let response = ConversationItemListResponse { - object: "list", - data: items, - has_more, - first_id, - last_id, - }; - - encrypt_response(&state, &session_id, &response).await -} - -async fn delete_agent_conversation( - State(state): State>, - Path(conversation_uuid): Path, + Path(agent_uuid): Path, Extension(session_id): Extension, Extension(user): Extension, ) -> Result>, ApiError> { - let ctx = load_agent_conversation_context(&state, user.uuid, conversation_uuid).await?; - - let mut conn = state - .db - .get_pool() - .get() - .map_err(|_| ApiError::InternalServerError)?; - - conn.transaction(|tx| -> Result<(), diesel::result::Error> { - use crate::models::schema::{agent_config, conversations}; - - diesel::delete( - user_embeddings::table - .filter(user_embeddings::user_id.eq(user.uuid)) - .filter(user_embeddings::conversation_id.eq(Some(ctx.conversation.id))) - .filter(user_embeddings::source_type.eq(crate::rag::SOURCE_TYPE_MESSAGE)), - ) - .execute(tx)?; - - // Delete the conversation (cascades). We keep the filter on user_id for safety. - let deleted = diesel::delete( - conversations::table - .filter(conversations::id.eq(ctx.conversation.id)) - .filter(conversations::user_id.eq(user.uuid)), - ) - .execute(tx)?; - - if deleted == 0 { - return Err(diesel::result::Error::NotFound); - } - - diesel::update(agent_config::table.filter(agent_config::user_id.eq(user.uuid))) - .set(agent_config::conversation_id.eq::>(None)) - .execute(tx)?; - - Ok(()) - }) - .map_err(|_| ApiError::InternalServerError)?; - - state.rag_cache.lock().await.evict_user(user.uuid); - - let response = DeletedObjectResponse::conversation(conversation_uuid); - - encrypt_response(&state, &session_id, &response).await -} - -struct AgentConversationContext { - conversation: Conversation, - user_key: secp256k1::SecretKey, -} - -async fn load_agent_conversation_context( - state: &AppState, - user_id: Uuid, - conversation_uuid: Uuid, -) -> Result { let mut conn = state .db .get_pool() .get() .map_err(|_| ApiError::InternalServerError)?; - let cfg = AgentConfig::get_by_user_id(&mut conn, user_id) + let agent = Agent::get_by_uuid_and_user(&mut conn, agent_uuid, user.uuid) .map_err(|_| ApiError::InternalServerError)? .ok_or(ApiError::NotFound)?; - let Some(agent_conversation_id) = cfg.conversation_id else { - return Err(ApiError::NotFound); - }; - - let conversation = Conversation::get_by_uuid_and_user(&mut conn, conversation_uuid, user_id) - .map_err(|_| ApiError::NotFound)?; - - if conversation.id != agent_conversation_id { + if agent.kind != AGENT_KIND_SUBAGENT { return Err(ApiError::NotFound); } - let user_key = state - .get_user_key(user_id, None, None) - .await - .map_err(|_| error_mapping::map_key_retrieval_error())?; - - Ok(AgentConversationContext { - conversation, - user_key, - }) -} + state + .db + .delete_conversation(agent.conversation_id, user.uuid) + .map_err(error_mapping::map_generic_db_error)?; -async fn get_or_create_agent_config( - conn: &mut PgConnection, - user_id: Uuid, -) -> Result { - let Some(cfg) = - AgentConfig::get_by_user_id(conn, user_id).map_err(|_| ApiError::InternalServerError)? - else { - let new_cfg = NewAgentConfig { - uuid: Uuid::new_v4(), - user_id, - conversation_id: None, - enabled: true, - model: runtime::DEFAULT_MODEL.to_string(), - max_context_tokens: runtime::DEFAULT_CONTEXT_WINDOW, - compaction_threshold: runtime::DEFAULT_COMPACTION_THRESHOLD, - system_prompt_enc: None, - preferences_enc: None, - }; + state.rag_cache.lock().await.evict_user(user.uuid); - return new_cfg - .insert_or_update(conn) - .map_err(|_| ApiError::InternalServerError); + let response = DeletedObjectResponse { + id: agent.uuid, + object: "agent.subagent.deleted", + deleted: true, }; - Ok(cfg) + encrypt_response(&state, &session_id, &response).await } diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 97770efe..34504b06 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -7,15 +7,19 @@ use uuid::Uuid; use diesel::prelude::*; use crate::encrypt::{decrypt_string, encrypt_with_key}; -use crate::models::agent_config::{AgentConfig, NewAgentConfig}; +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, DEFAULT_BLOCK_CHAR_LIMIT}; +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::{memory_blocks, user_embeddings}; +use crate::models::schema::{agents, memory_blocks, user_embeddings}; use crate::rag::{ insert_message_embedding, serialize_f32_le, SOURCE_TYPE_ARCHIVAL, SOURCE_TYPE_MESSAGE, }; @@ -31,22 +35,23 @@ use super::signatures::{ }; use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, - MemoryInsertTool, MemoryReplaceTool, ToolRegistry, ToolResult, + MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, }; use super::vision; -// Mirrors Sage defaults -const DEFAULT_PERSONA_DESCRIPTION: &str = "The persona block: Stores details about your current persona, guiding how you behave and respond. This helps you to maintain consistency and personality in your interactions."; -const DEFAULT_HUMAN_DESCRIPTION: &str = "The human block: Stores key details about the person you are conversing with, allowing for more personalized and friend-like conversation."; 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-5"; 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, @@ -74,8 +79,9 @@ pub struct AgentRuntime { state: Arc, user: Arc, user_key: Arc, - agent_config: AgentConfig, + agent: Agent, conversation: Conversation, + subagent_purpose: String, system_prompt: String, lm: Arc, tools: ToolRegistry, @@ -86,8 +92,207 @@ pub struct AgentRuntime { 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, + 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, + 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, + 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( + pub async fn new_main( state: Arc, user: crate::models::users::User, user_key: SecretKey, @@ -101,91 +306,67 @@ impl AgentRuntime { .get() .map_err(|_| ApiError::InternalServerError)?; - // Load or create config - let agent_config = match AgentConfig::get_by_user_id(&mut conn, user.uuid).map_err(|e| { - error!("Failed to load agent config: {e:?}"); - ApiError::InternalServerError - })? { - Some(cfg) => cfg, - None => { - let new_cfg = NewAgentConfig { - uuid: Uuid::new_v4(), - user_id: user.uuid, - conversation_id: None, - enabled: true, - model: DEFAULT_MODEL.to_string(), - max_context_tokens: DEFAULT_CONTEXT_WINDOW, - compaction_threshold: DEFAULT_COMPACTION_THRESHOLD, - system_prompt_enc: None, - preferences_enc: None, - }; + let (agent, conversation) = + ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; - new_cfg.insert_or_update(&mut conn).map_err(|e| { - error!("Failed to create agent config: {e:?}"); - ApiError::InternalServerError - })? - } - }; + Self::from_loaded(state, user, user_key, agent, conversation).await + } - if !agent_config.enabled { - return Err(ApiError::Unauthorized); - } + 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); - // Ensure conversation exists - let conversation = if let Some(conversation_id) = agent_config.conversation_id { - Conversation::get_by_id_and_user(&mut conn, conversation_id, user.uuid).map_err( - |e| { - error!("Failed to load agent conversation: {e:?}"); - ApiError::InternalServerError - }, - )? - } else { - let metadata = json!({"type":"agent_main"}); - let metadata_enc = encrypt_with_key(&user_key, metadata.to_string().as_bytes()).await; + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; - let new_conversation = NewConversation { - uuid: Uuid::new_v4(), - user_id: user.uuid, - metadata_enc: Some(metadata_enc), - }; - let conversation = new_conversation.insert(&mut conn).map_err(|e| { - error!("Failed to create agent conversation: {e:?}"); - ApiError::InternalServerError - })?; + let _ = ensure_main_agent(&state, &mut conn, &user_key, user.uuid).await?; - let updated_cfg = NewAgentConfig { - uuid: agent_config.uuid, - user_id: agent_config.user_id, - conversation_id: Some(conversation.id), - enabled: agent_config.enabled, - model: agent_config.model.clone(), - max_context_tokens: agent_config.max_context_tokens, - compaction_threshold: agent_config.compaction_threshold, - system_prompt_enc: agent_config.system_prompt_enc.clone(), - preferences_enc: agent_config.preferences_enc.clone(), - }; - let _ = updated_cfg.insert_or_update(&mut conn).map_err(|e| { - error!("Failed to update agent config with conversation_id: {e:?}"); + 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)?; - conversation - }; + 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 + }, + )?; - // Ensure default memory blocks exist - Self::ensure_default_blocks(&state, &mut conn, &user_key, user.uuid).await?; + Self::from_loaded(state, user, user_key, agent, conversation).await + } - // System prompt override - let system_prompt = match decrypt_string(&user_key, agent_config.system_prompt_enc.as_ref()) + 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 system_prompt_enc: {e:?}"); + error!("Failed to decrypt subagent purpose: {e:?}"); ApiError::InternalServerError - })? { - Some(s) if !s.trim().is_empty() => s, - _ => AGENT_INSTRUCTION.to_string(), - }; + })? + .unwrap_or_default(); + + let system_prompt = build_system_prompt(&agent, &subagent_purpose); - // Tool registry let mut tools = ToolRegistry::new(); tools.register(Arc::new(MemoryReplaceTool::new( state.clone(), @@ -218,14 +399,21 @@ impl AgentRuntime { 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(), - agent_config.model.clone(), + DEFAULT_MODEL.to_string(), 0.7, 32768, ) @@ -235,8 +423,9 @@ impl AgentRuntime { state, user, user_key, - agent_config, + agent, conversation, + subagent_purpose, system_prompt, lm, tools, @@ -254,8 +443,8 @@ impl AgentRuntime { user_key: &SecretKey, user_id: Uuid, ) -> Result<(), ApiError> { - let persona = - MemoryBlock::get_by_user_and_label(conn, user_id, "persona").map_err(|e| { + 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 })?; @@ -265,12 +454,8 @@ impl AgentRuntime { let block = NewMemoryBlock { uuid: Uuid::new_v4(), user_id, - label: "persona".to_string(), - description: Some(DEFAULT_PERSONA_DESCRIPTION.to_string()), + label: MEMORY_BLOCK_LABEL_PERSONA.to_string(), value_enc, - char_limit: DEFAULT_BLOCK_CHAR_LIMIT, - read_only: false, - version: 1, }; block.insert_or_update(conn).map_err(|e| { error!("Failed to create persona block: {e:?}"); @@ -278,22 +463,19 @@ impl AgentRuntime { })?; } - let human = MemoryBlock::get_by_user_and_label(conn, user_id, "human").map_err(|e| { - error!("Failed to load human 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: "human".to_string(), - description: Some(DEFAULT_HUMAN_DESCRIPTION.to_string()), + label: MEMORY_BLOCK_LABEL_HUMAN.to_string(), value_enc, - char_limit: DEFAULT_BLOCK_CHAR_LIMIT, - read_only: false, - version: 1, }; block.insert_or_update(conn).map_err(|e| { error!("Failed to create human block: {e:?}"); @@ -527,6 +709,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If 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, @@ -637,6 +821,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If 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 @@ -645,10 +831,15 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If .get() .map_err(|_| ApiError::InternalServerError)?; - let persona = MemoryBlock::get_by_user_and_label(&mut conn, self.user.uuid, "persona") - .map_err(|_| ApiError::InternalServerError)?; - let human = MemoryBlock::get_by_user_and_label(&mut conn, self.user.uuid, "human") - .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)? @@ -667,7 +858,6 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If 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)) - .filter(user_embeddings::conversation_id.eq(Some(self.conversation.id))) .count() .get_result(&mut conn) .map_err(|_| ApiError::InternalServerError)?; @@ -1080,8 +1270,8 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If if !self.compaction.should_compact( current_tokens as usize, - self.agent_config.max_context_tokens as usize, - self.agent_config.compaction_threshold, + DEFAULT_CONTEXT_WINDOW as usize, + DEFAULT_COMPACTION_THRESHOLD, ) { return Ok(()); } diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs index 594525e1..498c81ec 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -17,7 +17,7 @@ use dspy_rs::{CustomCompletionModel, LMClient, OneOrMany, LM}; pub const CORRECTION_INSTRUCTION: &str = r#"You are a response correction agent. Your job is to fix malformed agent responses. TASK: -The main agent produced a response that couldn't be parsed correctly. You must: +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 @@ -70,6 +70,10 @@ You have two types of memory. Use them proactively: **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:** @@ -174,6 +178,10 @@ pub struct AgentResponse { #[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, diff --git a/src/web/agent/tools.rs b/src/web/agent/tools.rs index 9199996f..d0018ac5 100644 --- a/src/web/agent/tools.rs +++ b/src/web/agent/tools.rs @@ -7,14 +7,21 @@ 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}; +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 // ============================================================================ @@ -127,11 +134,11 @@ async fn update_block_value( new_value: &str, existing: &MemoryBlock, ) -> Result<(), ApiError> { - if existing.read_only { + if label != MEMORY_BLOCK_LABEL_PERSONA && label != MEMORY_BLOCK_LABEL_HUMAN { return Err(ApiError::BadRequest); } - if new_value.len() > existing.char_limit.max(0) as usize { + if new_value.len() > MAX_CORE_MEMORY_BLOCK_CHARS { return Err(ApiError::BadRequest); } @@ -141,11 +148,7 @@ async fn update_block_value( uuid: existing.uuid, user_id, label: label.to_string(), - description: existing.description.clone(), value_enc, - char_limit: existing.char_limit, - read_only: existing.read_only, - version: existing.version, }; let mut conn = state @@ -774,6 +777,89 @@ impl Tool for ArchivalSearchTool { } } +// ============================================================================ +// 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 // ============================================================================ diff --git a/src/web/responses/conversations.rs b/src/web/responses/conversations.rs index 4e063184..1b297220 100644 --- a/src/web/responses/conversations.rs +++ b/src/web/responses/conversations.rs @@ -3,7 +3,7 @@ use crate::{ encrypt::{decrypt_content, encrypt_with_key}, - models::agent_config::AgentConfig, + models::agents::Agent, models::responses::NewConversation, models::users::User, web::{ @@ -78,9 +78,7 @@ impl ConversationContext { .get() .map_err(|_| ApiError::InternalServerError)?; - AgentConfig::get_by_user_id(&mut conn, user_uuid) - .map_err(|_| ApiError::InternalServerError)? - .and_then(|cfg| cfg.conversation_id) + main_agent_conversation_id(&mut conn, user_uuid)? }; if Some(conversation.id) == agent_conversation_id { @@ -112,6 +110,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 // ============================================================================ @@ -570,9 +577,7 @@ async fn list_conversations( .get() .map_err(|_| ApiError::InternalServerError)?; - AgentConfig::get_by_user_id(&mut conn, user.uuid) - .map_err(|_| ApiError::InternalServerError)? - .and_then(|cfg| cfg.conversation_id) + main_agent_conversation_id(&mut conn, user.uuid)? }; // Fetch conversations with database-level pagination @@ -659,9 +664,7 @@ async fn delete_all_conversations( .get() .map_err(|_| ApiError::InternalServerError)?; - let agent_conversation_id = AgentConfig::get_by_user_id(&mut conn, user.uuid) - .map_err(|_| ApiError::InternalServerError)? - .and_then(|cfg| cfg.conversation_id); + 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::*; @@ -719,9 +722,7 @@ async fn batch_delete_conversations( .get() .map_err(|_| ApiError::InternalServerError)?; - AgentConfig::get_by_user_id(&mut conn, user.uuid) - .map_err(|_| ApiError::InternalServerError)? - .and_then(|cfg| cfg.conversation_id) + main_agent_conversation_id(&mut conn, user.uuid)? }; for conversation_id in body.ids { From 8da00a0e2518f6842c52132e64b10f83b80025cd Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Thu, 5 Mar 2026 20:02:19 -0600 Subject: [PATCH 18/34] feat: expose paginated agent history endpoints Mirror the conversations API so clients can list agents and browse main-agent and subagent items with consistent cursor semantics. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/models/agents.rs | 71 ++++++++ src/web/agent/mod.rs | 426 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 492 insertions(+), 5 deletions(-) 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/web/agent/mod.rs b/src/web/agent/mod.rs index 183761a0..17ced6e8 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,24 +1,33 @@ use axum::{ - extract::{Path, State}, + extract::{Path, Query, State}, middleware::from_fn_with_state, response::sse::{Event, Sse}, - routing::{delete, post}, + 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 tracing::{error, warn}; use uuid::Uuid; -use crate::models::agents::{Agent, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT}; +use crate::encrypt::decrypt_string; +use crate::models::agents::{ + Agent, AGENT_CREATED_BY_AGENT, AGENT_CREATED_BY_USER, AGENT_KIND_SUBAGENT, +}; +use crate::models::responses::Conversation; use crate::models::users::User; 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}; @@ -43,11 +52,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)] @@ -56,6 +103,16 @@ enum ChatTarget { Subagent(Uuid), } +const MAIN_AGENT_DISPLAY_NAME: &str = "Sage"; + +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(), @@ -128,6 +185,193 @@ async fn encrypt_agent_event( encrypt_event(state, session_id, event_type, &value).await } +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, +) -> Result { + let user_key = state + .get_user_key(user.uuid, 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, +) -> Result { + let user_key = state + .get_user_key(user.uuid, 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) +} + pub fn router(app_state: Arc) -> Router<()> { // Experimental endpoints: only enabled in Local/Dev. if !matches!(app_state.app_mode, AppMode::Local | AppMode::Dev) { @@ -135,6 +379,20 @@ 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/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( @@ -142,6 +400,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( @@ -149,6 +411,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( @@ -156,6 +422,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) @@ -164,6 +440,144 @@ 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, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user).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, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user).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, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user).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 list_subagents( + State(state): State>, + Query(params): Query, + Extension(session_id): Extension, + Extension(user): 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.uuid, 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, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid).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, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid).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, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid).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, @@ -502,11 +916,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 From e892e9d18c9eed4b1a41cc25bb95feb4cf6b4ea7 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Fri, 6 Mar 2026 14:50:47 -0600 Subject: [PATCH 19/34] feat: add a dev-only main agent reset endpoint Allow Sage development flows to fully reset the main agent, its subagents, and shared memory without changing production behavior. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/web/agent/mod.rs | 67 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 63 insertions(+), 4 deletions(-) diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 17ced6e8..61ed882a 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -17,6 +17,7 @@ use crate::encrypt::decrypt_string; 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::web::encryption_middleware::{decrypt_request, encrypt_response, EncryptedResponse}; @@ -372,6 +373,17 @@ fn get_item_from_conversation( 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) { @@ -383,6 +395,11 @@ pub fn router(app_state: Arc) -> Router<()> { "/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) @@ -476,6 +493,51 @@ async fn get_main_agent_item( 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, @@ -948,10 +1010,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); From 486b61dd6ffb80b26b5bf052fd22161c626e2d58 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Sat, 7 Mar 2026 21:20:50 -0600 Subject: [PATCH 20/34] feat: add encrypted push delivery for Sage agents Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 200 ++- Cargo.toml | 2 + docs/encrypted-mobile-push-notifications.md | 1296 +++++++++++++++++ docs/nitro-deploy.md | 102 ++ entrypoint.sh | 16 + .../down.sql | 10 + .../up.sql | 86 ++ src/db.rs | 40 +- src/main.rs | 39 +- src/models/mod.rs | 3 + src/models/notification_deliveries.rs | 226 +++ src/models/notification_events.rs | 91 ++ src/models/project_settings.rs | 62 + src/models/push_devices.rs | 237 +++ src/models/schema.rs | 69 + src/push/apns.rs | 263 ++++ src/push/crypto.rs | 151 ++ src/push/fcm.rs | 299 ++++ src/push/mod.rs | 388 +++++ src/push/worker.rs | 376 +++++ src/web/agent/mod.rs | 489 ++++--- 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 | 286 ++++ 26 files changed, 4668 insertions(+), 220 deletions(-) create mode 100644 docs/encrypted-mobile-push-notifications.md create mode 100644 migrations/2026-03-07-120000_push_notifications_v1/down.sql create mode 100644 migrations/2026-03-07-120000_push_notifications_v1/up.sql 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/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 4c84aa67..3289b421 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1038,6 +1038,12 @@ dependencies = [ "syn 2.0.106", ] +[[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" @@ -1525,7 +1531,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -1556,6 +1562,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" @@ -1661,6 +1673,18 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +[[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.6" @@ -1826,6 +1850,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" @@ -1911,6 +1946,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", + "const-oid", "crypto-common", "subtle", ] @@ -1933,7 +1969,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2027,6 +2063,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.2" @@ -2042,6 +2092,27 @@ version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +[[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" @@ -2103,7 +2174,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2240,6 +2311,16 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6" +[[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" @@ -2541,6 +2622,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -2623,6 +2705,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.26" @@ -2773,6 +2866,15 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + [[package]] name = "hmac" version = "0.12.1" @@ -3847,7 +3949,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4035,6 +4137,7 @@ dependencies = [ "generic-array", "getrandom 0.2.15", "hex", + "hkdf", "hmac", "hyper 0.14.30", "hyper-tls 0.5.0", @@ -4043,6 +4146,7 @@ dependencies = [ "lazy_static", "oauth2", "once_cell", + "p256", "password-auth", "rand_core 0.6.4", "rcgen", @@ -4146,6 +4250,18 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a" +[[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" @@ -4253,6 +4369,15 @@ dependencies = [ "serde", ] +[[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.1" @@ -4301,6 +4426,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.30" @@ -4381,6 +4516,15 @@ dependencies = [ "syn 2.0.106", ] +[[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.4.0" @@ -4757,6 +4901,16 @@ dependencies = [ "thiserror 1.0.63", ] +[[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" @@ -4876,7 +5030,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5025,6 +5179,20 @@ dependencies = [ "untrusted", ] +[[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.0" @@ -5248,6 +5416,16 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "digest", + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.8" @@ -5348,6 +5526,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.0" @@ -5506,7 +5694,7 @@ dependencies = [ "getrandom 0.4.1", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 869312d6..df49f79e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,8 @@ 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"] } +hkdf = "0.12" hex = "0.4.3" base64 = "0.22.1" vsock = "0.5.1" diff --git a/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md new file mode 100644 index 00000000..61c10eb2 --- /dev/null +++ b/docs/encrypted-mobile-push-notifications.md @@ -0,0 +1,1296 @@ +# Encrypted Mobile Push Notifications for Maple / OpenSecret + +## Implementation-Ready Backend, Enclave, and Mobile Design + +**Date:** March 2026 +**Status:** Implementation target / pre-build spec +**Related Docs:** +- `sage-in-maple-architecture.md` +- `architecture-for-rag-integration.md` +- `potential-rag-integration-brute-force.md` + +--- + +## 1. Executive Summary + +This document defines the concrete v1 design for encrypted mobile push notifications in OpenSecret / Maple. + +The final architectural decisions are: + +1. **Use direct APNs for iOS and direct FCM HTTP v1 for Android.** +2. **Run the push sender inside the OpenSecret enclave app process.** +3. **Store provider configuration per project using `project_settings` + `org_project_secrets`.** +4. **Register one notification keypair per device, not one shared key per user.** +5. **Use a durable Postgres outbox + delivery worker, not synchronous request-path sends.** +6. **Treat successful live Sage SSE delivery as the acknowledgement for request/response flows and skip push entirely when streaming succeeds.** +7. **Prefer standard visible-notification behavior on Android in v1; keep Android data-only encrypted preview as an optional later enhancement.** + +This gives us Signal-like transport privacy for push content without pretending we are building a full Signal protocol. + +--- + +## 2. What Problem This Solves + +We want Maple / Sage to send notifications for events like: + +- reminder due +- long-running agent task complete +- async follow-up from Sage +- security or account alerts + +while ensuring that: + +- Apple and Google do not need plaintext notification content +- raw push tokens are not stored plaintext in Postgres +- losing one device does not compromise every other device +- the design fits the current OpenSecret enclave + project-scoped secret model + +--- + +## 3. Non-Goals + +This design does **not** attempt to provide: + +- a double-ratchet messaging protocol +- server blindness to the notification plaintext at creation time +- guaranteed hidden lock-screen content after the device itself decrypts it +- a single cross-device shared notification private key + +This is **encrypted push delivery**, not a general secure messaging redesign. + +--- + +## 4. Important Clarifications + +### 4.1 Apple and Google do not run our decryption logic + +They only deliver payloads. The only code we control at receipt time runs: + +- in an iOS **Notification Service Extension** +- in an Android **FirebaseMessagingService** / `WorkManager` flow + +So decryption must happen **on the device**. + +### 4.2 There is no special vsock registration with Apple or Google + +We do **not** register AWS vsock or Nitro plumbing with APNs or FCM. + +We only need: + +- outbound HTTPS access from the enclave to APNs / FCM endpoints +- parent-instance `vsock-proxy` allowlist entries +- enclave-local `/etc/hosts` entries plus `traffic_forwarder.py` mappings + +### 4.3 Android and iOS do not behave the same + +iOS supports the strongest version of the β€œmessaging app” pattern: + +- send a generic visible alert +- include encrypted preview payload +- let NSE replace the visible content before display +- if the NSE fails or times out, iOS shows the original generic alert +- if the app is already foregrounded, the app can suppress the visible notification via `willPresent` + +Android does **not** offer the same β€œgeneric visible fallback + encrypted replacement” behavior for encrypted previews. The standard Android messaging-app approach is usually: + +- send a standard visible FCM notification with non-sensitive text +- include extra `data` fields for routing / grouping / dedup +- if the app is already foregrounded, handle it in-app and avoid showing another local notification + +Android **can** do encrypted preview with **high-priority data-only** FCM, but that is less standard and depends more heavily on background execution. + +So the v1 system should explicitly model Android delivery as: + +1. **Standard visible notification**: default v1 path, most reliable, non-sensitive visible text, app-open fetch for authoritative content +2. **Encrypted preview best-effort**: optional later path, richer UX, but depends on app background execution + +This is the most important platform caveat in the whole design. We should follow the standard platform patterns rather than force parity where the OSes behave differently. + +--- + +## 5. Final Architectural Decisions + +### 5.1 Per-device notification keypairs + +Each device generates its own notification keypair locally and registers only the public key. + +Why: + +- push routing is already device-specific +- revocation is per-device +- one compromised device does not force full-account rotation +- mobile secure storage works best with locally generated keys + +### 5.2 Separate push crypto from existing user-key crypto + +Push crypto must **not** reuse the current server-derived secp256k1 user-key flow. + +Use a dedicated **P-256 ECDH** notification keypair instead. + +Why: + +- better native support on iOS and Android +- easier hardware-backed storage +- avoids importing a shared account root into device secure storage + +### 5.3 Push sender runs inside the enclave app + +The sender should live inside the main OpenSecret server process, behind the current enclave boundary. + +Why: + +- push payload plaintext can be sensitive +- APNs / FCM credentials are sensitive +- push tokens are sensitive enough to encrypt at rest +- this keeps the privacy story aligned with the rest of OpenSecret + +### 5.4 Successful live Sage SSE delivery suppresses push in v1 + +For the common β€œuser sent a message to Sage and is waiting on an SSE stream” flow: + +- if the final Sage response is delivered successfully over the open SSE stream, do **not** enqueue push for any device +- if the SSE stream closes or errors before the final response is delivered, enqueue **exactly one** notification event for that interrupted turn +- if the interrupted turn produced multiple assistant messages, use the **first missed assistant message** as the preview source for that one notification +- this SSE success signal acts as the acknowledgement; no extra per-thread presence service or explicit ACK protocol is required in v1 + +This is intentionally simple and matches the normal messaging-app intuition: if the user is clearly active on one device and the live response succeeded, don’t notify every other device. + +### 5.5 Outbox + worker, not inline send + +Notification generation and transport delivery must be separate steps. + +Why: + +- APNs / FCM reject, throttle, and invalidate tokens +- we need retries and backoff +- future reminder scheduling requires durable queue semantics +- request handlers should not block on provider APIs + +### 5.6 Stable notification IDs + client dedup are required + +Every logical notification must have a stable `notification_id`. + +Use `notification_events.uuid` as that canonical ID. + +Rules: + +- retries must reuse the same `notification_id` +- clients should keep a small cache of recently seen `notification_id` values and no-op duplicates +- APNs `apns-collapse-id` and FCM `collapse_key` should only be reused for retries of the same logical notification, not for every message in a thread by default + +--- + +## 6. Existing Codebase Integration Map + +This section is the implementation map for a follow-on coding agent. + +### 6.1 Existing files to modify + +#### Database and models + +- `migrations//up.sql` + - Add `push_devices`, `notification_events`, and `notification_deliveries`. + - Add indexes and `updated_at` triggers matching the existing schema style. + +- `migrations//down.sql` + - Drop the new push tables and indexes. + +- `src/models/mod.rs` + - Export the new push model modules. + +- `src/models/schema.rs` + - Regenerate via Diesel after the migration is added. + +- `src/models/project_settings.rs` + - Extend `SettingCategory` with `Push`. + - Add `PushSettings`, `IosPushSettings`, `AndroidPushSettings`, and related helpers. + +- `src/db.rs` + - Add trait methods and `PostgresConnection` implementations for push device CRUD, event enqueue, delivery leasing, retry, and invalidation. + - Mirror the existing settings / secrets patterns already used for email and OAuth. + +#### Web API routing + +- `src/web/mod.rs` + - Export a new `push` router module. + +- `src/main.rs` + - Merge the new push router under `validate_jwt`. + - Start the background push worker after app state initialization and migrations. + +- `src/web/responses/handlers.rs` + - Gate push enqueue on whether a live SSE response finished successfully; this is the v1 suppression signal for the normal Sage request/response flow. + +- `src/web/agent/mod.rs` + - If mobile uses the direct agent SSE endpoints, apply the same β€œsuccessful SSE delivery => no push” rule there. + +- `src/web/platform/common.rs` + - Add new secret constants and request / response types for project push settings. + +- `src/web/platform/project_routes.rs` + - Add `GET` / `PUT /platform/orgs/:org_id/projects/:project_id/settings/push`. + - Reuse the existing project secret endpoints for actual APNs / FCM credentials. + +#### Deployment / runtime + +- `entrypoint.sh` + - Add `/etc/hosts` entries and `traffic_forwarder.py` processes for APNs production, APNs sandbox, and FCM. + +- `docs/nitro-deploy.md` + - Add parent-instance `vsock-proxy` instructions for the new outbound hosts. + - This is a deployment follow-up, not required for compiling the backend. + +### 6.2 New files to add + +- `src/models/push_devices.rs` + - Diesel model and helper methods for device registrations. + +- `src/models/notification_events.rs` + - Durable logical notification rows. + +- `src/models/notification_deliveries.rs` + - Per-device delivery state rows. + +- `src/web/push.rs` + - App-user authenticated routes like `POST /v1/push/devices`. + +- `src/push/mod.rs` + - Shared types, enqueue helpers, and internal interfaces. + +- `src/push/crypto.rs` + - P-256 ECDH + HKDF + AES-GCM envelope logic. + +- `src/push/apns.rs` + - APNs request building, JWT auth, and response handling. + +- `src/push/fcm.rs` + - FCM OAuth token exchange, request building, and response handling. + +- `src/push/worker.rs` + - Polling / leasing loop that sends pending deliveries. + +### 6.3 Codebase gotchas to preserve + +1. `src/main.rs::get_project_secret()` currently returns **base64-encoded decrypted bytes**. + That is fine for some existing flows, but APNs `.p8` material and FCM service account JSON are easier to consume as raw bytes or raw UTF-8. Add a raw helper rather than forcing each caller to decode a second time. + +2. `src/encrypt.rs` is useful for **at-rest encryption** but not for device ECDH push envelopes. + Its current primitives take `secp256k1::SecretKey` and do symmetric encryption only. + +3. `src/main.rs::create_ephemeral_key()` uses **x25519** for session establishment. + Do not overload that codepath for push notifications. + +4. New user-facing push APIs should live under **`/v1/push/*`**, not under the older `/protected/*` namespace. + That matches the newer `agent` and `responses` route style. + +--- + +## 7. Project Configuration Model + +Push config should follow the same split already used elsewhere in the repo. + +### 7.1 Non-secret settings in `project_settings` + +Add a new `SettingCategory::Push` and store a JSON payload like: + +```rust +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PushSettings { + pub encrypted_preview_enabled: bool, + pub ios: Option, + pub android: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +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)] +pub struct AndroidPushSettings { + pub enabled: bool, + pub firebase_project_id: String, + pub package_name: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PushEnvironment { + Dev, + Prod, +} +``` + +`encrypted_preview_enabled` controls whether the server attaches the encrypted preview envelope for supported devices. If disabled, the system still uses a generic provider-visible notification shell; it must **not** fall back to plaintext message content. + +Notes: + +- no new SQL table is required for settings; `project_settings` already supports new categories via JSONB +- `team_id` and `key_id` are identifiers, not secrets +- APNs environment stays explicit because sandbox vs production matters operationally + +### 7.2 Secrets in `org_project_secrets` + +Add constants in `src/web/platform/common.rs`: + +```rust +pub const PROJECT_APNS_AUTH_KEY_P8: &str = "APNS_AUTH_KEY_P8"; +pub const PROJECT_FCM_SERVICE_ACCOUNT_JSON: &str = "FCM_SERVICE_ACCOUNT_JSON"; +``` + +Use the existing platform secret endpoints to store them. + +Important: + +- the existing platform secret API expects **base64-encoded raw bytes** +- APNs `.p8` content should be uploaded as raw PEM bytes, then base64-encoded for transport to the API +- FCM service account JSON should be uploaded as raw JSON bytes, then base64-encoded for transport to the API + +--- + +## 8. Data Model + +### 8.1 `push_devices` + +One row per device installation. + +```sql +CREATE TABLE push_devices ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + installation_id UUID NOT NULL, + platform TEXT NOT NULL, + provider TEXT NOT NULL, + environment TEXT NOT NULL, + app_id TEXT NOT NULL, + + push_token_enc BYTEA NOT NULL, + push_token_hash BYTEA NOT NULL, + + notification_public_key BYTEA NOT NULL, + key_algorithm TEXT NOT NULL, + + supports_encrypted_preview BOOLEAN NOT NULL DEFAULT FALSE, + supports_background_processing BOOLEAN NOT NULL DEFAULT FALSE, + + last_seen_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + revoked_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE (user_id, installation_id, environment), + UNIQUE (provider, environment, push_token_hash), + CHECK ( + (platform = 'ios' AND provider = 'apns') OR + (platform = 'android' AND provider = 'fcm') + ), + CHECK (environment IN ('dev', 'prod')), + CHECK (key_algorithm = 'p256_ecdh_v1') +); + +CREATE INDEX idx_push_devices_user_active + ON push_devices(user_id, revoked_at); +``` + +Implementation notes: + +- `push_token_enc` should be encrypted at rest with the enclave key using the same style as project secrets +- `push_token_hash` should be `SHA-256(push_token)` raw bytes for dedupe and lookup +- `notification_public_key` should store raw DER-encoded SPKI bytes, not a base64 string + +### 8.2 `notification_events` + +Logical notification rows created by product code. + +```sql +CREATE TABLE notification_events ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + project_id INTEGER NOT NULL REFERENCES org_projects(id) ON DELETE CASCADE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + kind TEXT NOT NULL, + delivery_mode TEXT NOT NULL, + priority TEXT NOT NULL DEFAULT 'normal', + collapse_key TEXT, + + fallback_title TEXT NOT NULL, + fallback_body TEXT NOT NULL, + payload_enc BYTEA, + + not_before_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + cancelled_at TIMESTAMPTZ, + + CHECK (delivery_mode IN ('generic', 'encrypted_preview')), + CHECK (priority IN ('normal', 'high')) +); + +CREATE INDEX idx_notification_events_due + ON notification_events(user_id, not_before_at, cancelled_at); +``` + +Implementation notes: + +- `payload_enc` should contain the encrypted preview payload or deep-link metadata encrypted at rest with the enclave key +- `fallback_title` and `fallback_body` are always explicit so the worker never has to invent fallback text at send time +- `notification_events.uuid` is the canonical `notification_id` exposed to clients and reused across retries + +### 8.3 `notification_deliveries` + +Per-device delivery records. + +```sql +CREATE TABLE notification_deliveries ( + id BIGSERIAL PRIMARY KEY, + event_id BIGINT NOT NULL REFERENCES notification_events(id) ON DELETE CASCADE, + push_device_id BIGINT NOT NULL REFERENCES push_devices(id) ON DELETE CASCADE, + + status TEXT NOT NULL DEFAULT 'pending', + attempt_count INTEGER NOT NULL DEFAULT 0, + next_attempt_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + lease_owner TEXT, + lease_expires_at TIMESTAMPTZ, + + provider_message_id TEXT, + provider_status_code INTEGER, + last_error TEXT, + sent_at TIMESTAMPTZ, + invalidated_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE (event_id, push_device_id), + CHECK ( + status IN ( + 'pending', + 'leased', + 'sent', + 'retry', + 'failed', + 'invalid_token', + 'cancelled' + ) + ) +); + +CREATE INDEX idx_notification_deliveries_pending + ON notification_deliveries(status, next_attempt_at); +``` + +Implementation notes: + +- deliveries should be materialized when the event is enqueued, not lazily at send time +- the worker should still re-check `push_devices.revoked_at` before sending + +--- + +## 9. API Surface + +### 9.1 App-user push device API + +Create a new router in `src/web/push.rs` and mount it in `main.rs` with `validate_jwt` plus the usual encrypted request / response middleware. + +#### Register or rotate device + +```http +POST /v1/push/devices +``` + +Request payload: + +```json +{ + "installation_id": "uuid", + "platform": "ios", + "provider": "apns", + "environment": "prod", + "app_id": "ai.trymaple.ios", + "push_token": "opaque platform token string", + "notification_public_key": "base64-encoded SPKI DER", + "key_algorithm": "p256_ecdh_v1", + "supports_encrypted_preview": true, + "supports_background_processing": true +} +``` + +Behavior: + +- upsert by `(user_id, installation_id, environment)` +- rotate token if it changed +- rotate public key if it changed +- clear `revoked_at` if this is a re-registration +- refresh `last_seen_at` + +#### List current user's devices + +```http +GET /v1/push/devices +``` + +Return metadata only, never the raw push token. + +#### Revoke one device + +```http +DELETE /v1/push/devices/:id +``` + +Behavior: + +- user-scoped revoke only +- set `revoked_at` +- optionally return the standard deleted-object response shape already used in newer APIs + +### 9.2 Platform project settings API + +Add a new project settings endpoint in `src/web/platform/project_routes.rs`: + +```http +GET /platform/orgs/:org_id/projects/:project_id/settings/push +PUT /platform/orgs/:org_id/projects/:project_id/settings/push +``` + +Use existing platform auth, org membership checks, and the same validation style already used for email / OAuth settings. + +Do **not** add a dedicated secret endpoint. Reuse the existing generic project secret endpoints with: + +- `APNS_AUTH_KEY_P8` +- `FCM_SERVICE_ACCOUNT_JSON` + +--- + +## 10. DB Trait and Service Methods + +Add the following categories of methods to `src/db.rs`. + +### 10.1 Push device methods + +- `upsert_push_device(...)` +- `get_push_device_by_uuid_for_user(...)` +- `list_active_push_devices_for_user(...)` +- `revoke_push_device_for_user(...)` +- `invalidate_push_device(...)` + +### 10.2 Push settings methods + +- `get_project_push_settings(project_id)` +- `update_project_push_settings(project_id, settings)` + +This mirrors the existing email / OAuth settings pattern. + +### 10.3 Event and delivery methods + +- `create_notification_event(...)` +- `create_notification_deliveries_for_event(...)` +- `lease_pending_notification_deliveries(limit, lease_owner, lease_ttl)` +- `mark_notification_delivery_sent(...)` +- `mark_notification_delivery_retry(...)` +- `mark_notification_delivery_failed(...)` +- `mark_notification_delivery_invalid_token(...)` + +The lease query should use a transaction plus `FOR UPDATE SKIP LOCKED` semantics so the design works across multiple workers / enclaves from day one. +Use Postgres time (`now()`) for lease and retry timestamps instead of local enclave clock time. + +--- + +## 11. Worker Architecture + +### 11.1 Where the worker should start + +Initialize the worker from `src/main.rs` after: + +- app state build +- runtime migrations +- provider setup that already happens during boot + +Do **not** initialize it from request handlers. + +### 11.2 Worker structure + +Create `src/push/worker.rs` with a loop like: + +1. sleep / interval tick +2. lease due deliveries from Postgres +3. for each leased delivery: + - load event + device + project settings + - skip if device revoked or event cancelled / expired + - send via APNs or FCM + - update delivery status based on response + +Recommended v1 behavior: + +- poll every 2-5 seconds +- cap concurrent sends with a small semaphore +- exponential backoff on provider retryable failures + +### 11.3 Multi-enclave operation + +Run the worker in **every** OpenSecret backend / enclave instance. + +Shared Postgres is the coordination plane. + +Rules: + +- workers claim batches using row-level leases plus `FOR UPDATE SKIP LOCKED` +- no leader election is required +- any enclave can resume work after another enclave crashes +- `lease_owner`, `lease_expires_at`, `next_attempt_at`, and `status` must all live in Postgres + +### 11.4 Delivery semantics + +This design is **at-least-once**, not exactly-once. + +If a worker sends successfully to APNs / FCM but crashes before marking the row `sent`, another worker may retry after the lease expires. + +That is acceptable if we keep these protections in place: + +- stable `notification_id` +- one delivery row per `(event_id, push_device_id)` +- client-side recent-ID dedup cache +- provider collapse identifiers reused only for retries of the same logical notification + +### 11.5 Why Postgres outbox instead of SQS for v1 + +The repo already has some SQS plumbing, but v1 push delivery should still use Postgres as the source of truth. + +Why: + +- fewer moving parts +- easier local/dev operation +- delivery state stays queryable in one place +- no need to solve dual-write consistency on day one + +SQS can be added later as a wake-up optimization if volume requires it. + +--- + +## 12. Cryptographic Design + +### 12.1 What can be reused from the current repo + +Reusable: + +- `src/encrypt.rs::encrypt_with_key` for encrypting raw bytes at rest with the enclave key +- `src/encrypt.rs::decrypt_with_key` for decrypting those bytes +- existing `sha2` dependency for token hashing or HKDF backing hash + +Not reusable as-is: + +- current secp256k1 user-key flow +- current x25519 session setup in `AppState` +- any API that assumes a `secp256k1::SecretKey` is the transport key + +### 12.2 New dependencies to add + +Recommended new crates: + +- `p256` for P-256 ECDH and public-key parsing +- `hkdf` for key derivation + +Existing crates already cover the rest: + +- `aes-gcm` +- `sha2` +- `jsonwebtoken` +- `reqwest` + +### 12.3 Registered public-key format + +The client should register the public key as: + +- **base64-encoded DER SubjectPublicKeyInfo bytes** over the API +- stored as raw bytes in Postgres + +This is the best interoperability point because: + +- Android naturally exports X.509 / SPKI bytes +- iOS can export DER representation +- Rust `p256` can parse SPKI cleanly + +### 12.4 Envelope algorithm + +Use: + +- `P-256 ECDH` +- `HKDF-SHA256` +- `AES-256-GCM` + +Recommended envelope: + +```json +{ + "enc_v": 1, + "alg": "p256-hkdf-sha256-aes256gcm", + "kid": "push-device-uuid", + "epk": "base64-encoded ephemeral SEC1 public key", + "salt": "base64-encoded 32-byte salt", + "nonce": "base64-encoded 12-byte nonce", + "ciphertext": "base64-encoded AES-GCM ciphertext" +} +``` + +### 12.5 Plaintext preview payload + +```json +{ + "v": 1, + "notification_id": "uuid", + "message_id": "uuid", + "kind": "sage.reminder", + "title": "Reminder", + "body": "Follow up on the deployment thread", + "deep_link": "opensecret://agent/subagent/uuid", + "thread_id": "agent:uuid", + "sent_at": 1772800000 +} +``` + +`notification_id` should be identical to `notification_events.uuid` and remain stable across retries. + +Keep this intentionally small. + +For the initial implementation, normalize whitespace and truncate the preview `body` to a conservative byte budget before encryption so the final APNs payload stays comfortably under provider limits. + +### 12.6 At-rest encryption rules + +- encrypt `push_token_enc` with the enclave key +- encrypt `notification_events.payload_enc` with the enclave key +- do **not** store preview plaintext or raw push tokens in logs + +--- + +## 13. APNs Integration + +### 13.1 Endpoint and auth + +Send requests to: + +- production: `https://api.push.apple.com/3/device/` +- development: `https://api.sandbox.push.apple.com/3/device/` + +Use APNs token auth: + +- sign JWT with the stored `.p8` key +- header: `alg=ES256`, `kid=` +- claims: `iss=`, `iat=` +- refresh before one hour elapses; cache for about 50 minutes + +### 13.2 Required headers + +Use: + +- `authorization: bearer ` +- `apns-topic: ` +- `apns-push-type: alert` +- `apns-priority: 10` +- `apns-collapse-id: notif:` when retrying the same logical notification + +Do **not** reuse one collapse ID for every message in a thread by default. +If desired, use `aps.thread-id` for notification-center grouping separately from collapse behavior. + +### 13.3 Payload shape for encrypted preview on iOS + +```json +{ + "aps": { + "alert": { + "title": "New Maple update", + "body": "Open Maple to view it" + }, + "mutable-content": 1, + "sound": "default" + }, + "os_meta": { + "notification_id": "...", + "kind": "agent.message", + "message_id": "...", + "thread_id": "agent:subagent:...", + "deep_link": "opensecret://agent/subagent/..." + }, + "os_push": { + "enc_v": 1, + "alg": "p256-hkdf-sha256-aes256gcm", + "kid": "...", + "epk": "...", + "salt": "...", + "nonce": "...", + "ciphertext": "..." + } +} +``` + +Important: + +- the fallback `alert.title` and `alert.body` must be non-empty +- the fallback alert must remain **generic**; never place plaintext assistant content in the provider-visible payload +- `mutable-content: 1` is required for the NSE to run +- if the NSE times out, iOS shows the original generic alert +- `os_meta` is allowed to carry minimal routing metadata (`notification_id`, `message_id`, `thread_id`, `deep_link`, `kind`), but not plaintext notification content + +### 13.4 APNs error handling + +Treat these as permanent invalid-token outcomes: + +- `410 Unregistered` + +Treat these as permanent non-retryable failures, but **do not** auto-revoke the device row based on them alone: + +- `400 BadDeviceToken` +- `400 DeviceTokenNotForTopic` + +Treat these as retryable: + +- `429` +- `500` +- `503` + +Persist: + +- HTTP status +- APNs `apns-id` +- failure reason string if returned + +Note: APNs token invalidation signals can be delayed. Use them for cleanup, but do not infer uninstall with certainty. + +--- + +## 14. FCM Integration + +### 14.1 Endpoint and auth + +Send requests to: + +```text +https://fcm.googleapis.com/v1/projects//messages:send +``` + +Use OAuth 2 service-account auth: + +1. load decrypted service account JSON +2. create a JWT assertion signed with the service-account RSA private key +3. exchange it at `https://oauth2.googleapis.com/token` (or the service account `token_uri`) +4. cache the returned access token until shortly before expiry + +Required scope: + +```text +https://www.googleapis.com/auth/firebase.messaging +``` + +### 14.2 Recommended Android v1 behavior + +For v1, prefer the standard messaging-app approach: + +- send a normal visible FCM notification with non-sensitive text +- include `data` fields for routing, grouping, and duplicate protection +- let the app fetch or render authoritative content after open + +This is the most standard and reliable path. +It avoids relying on background execution for every Sage message. + +### 14.3 Standard FCM payload for v1 + +```json +{ + "message": { + "token": "", + "notification": { + "title": "New Sage message", + "body": "Open Maple to view it" + }, + "data": { + "notification_id": "", + "thread_id": "agent:uuid", + "message_id": "", + "deep_link": "opensecret://agent/subagent/uuid" + }, + "android": { + "priority": "HIGH", + "collapse_key": "sage:notif:", + "ttl": "300s" + } + } +} +``` + +Notes: + +- visible text should remain generic / non-sensitive +- FCM `data` values must be strings +- FCM `data` is provider-visible routing metadata, so it may include fields like `notification_id`, `message_id`, `thread_id`, `deep_link`, and `kind`, but never plaintext message content +- reuse the same `collapse_key` only for retries of the same logical notification +- `notification_id` is for exact dedup; `thread_id` is for grouping / clearing when the user opens the conversation + +### 14.4 Optional Android encrypted-preview path (deferred) + +If later we want closer parity with iOS: + +- send a **high-priority data-only** message +- `FirebaseMessagingService.onMessageReceived()` decrypts it +- app posts a local notification after decryption +- if the process does not run, the encrypted preview may not be shown + +That makes Android encrypted preview **best-effort**, not a required part of the v1 design. + +### 14.5 FCM error handling + +Treat these as permanent invalid-token outcomes: + +- `404 UNREGISTERED` +- `400 INVALID_ARGUMENT` when the error clearly identifies token invalidity + +Treat these as retryable: + +- `429` +- `500` +- `503` + +Persist: + +- HTTP status +- returned FCM message name on success +- error code / status string on failure + +--- + +## 15. Mobile Client Requirements + +### 15.1 iOS + +#### Registration flow + +1. request notification permission +2. receive APNs device token +3. generate `P256.KeyAgreement.PrivateKey()` locally +4. store private key in keychain / Secure Enclave-backed flow when available +5. expose the public key in DER / SPKI form +6. register with `POST /v1/push/devices` + +#### Storage requirements + +The key must be accessible from both: + +- main app target +- Notification Service Extension target + +Recommended keychain accessibility: + +- `AfterFirstUnlockThisDeviceOnly`-style semantics + +This allows NSE access after first unlock while keeping the key device-bound. + +#### NSE behavior + +The Notification Service Extension should: + +1. read `userInfo["os_push"]` +2. load the private key from shared keychain storage +3. decrypt the envelope +4. rewrite `title`, `body`, and routing metadata +5. call the completion handler promptly + +If the key is unavailable or decryption fails, do nothing and let iOS show the generic fallback alert. + +#### Foreground / open-thread behavior + +If the app is already foregrounded and showing the target Sage conversation or active response screen: + +- implement `userNotificationCenter(_:willPresent:withCompletionHandler:)` +- return `[]` to suppress banner / sound / badge presentation for that in-app event +- clear previously delivered notifications for that `thread_id` when the user opens the conversation + +The NSE should not be treated as the primary β€œsuppress if user is already looking” mechanism. For v1, successful SSE delivery is the primary suppression signal. + +### 15.2 Android + +#### Registration flow + +1. request notification permission on Android 13+ +2. obtain FCM registration token +3. generate an EC keypair in Android Keystore +4. export the public key as SPKI bytes +5. register with `POST /v1/push/devices` + +#### Runtime behavior + +- use `FirebaseMessagingService.onNewToken()` to rotate tokens +- for v1, assume the default incoming path is a standard visible FCM notification plus `data` +- if the app is already foregrounded and showing the target Sage conversation, update UI and do not show an extra local notification +- when the user opens a conversation, clear notifications associated with that `thread_id` +- if we later enable Android data-only encrypted preview, handle it in `onMessageReceived()` and only post a local notification after successful decrypt +- use `WorkManager` for longer processing if needed + +#### Direct Boot caveat + +Before first unlock after reboot, Android secure key material may be unavailable. + +So: + +- encrypted preview mode should degrade to app-open fetch or no preview +- if a notification must be visible pre-unlock, use generic mode instead + +### 15.3 Cross-platform client guidance + +#### Recent-ID dedup cache + +Each client should persist a small cache of recently seen `notification_id` values for at least 1-7 days. + +Rule: + +- if a notification arrives with an already-seen `notification_id`, treat it as a no-op + +This is the client-side protection against at-least-once retries or server bugs. + +#### Thread grouping and cleanup + +Each payload should include a `thread_id` when applicable. + +Use it to: + +- group notifications in app UI if desired +- clear thread notifications when the user opens that thread +- avoid leaving stale notifications around once the user has seen the conversation + +#### Foreground suppression is defense-in-depth + +If a visible notification arrives while the app is already open, suppress or no-op locally where the platform allows it. + +But the primary v1 suppression rule is still server-side: + +- successful live SSE delivery => no push enqueue at all + +--- + +## 16. Nitro / Enclave Networking Changes + +### 16.1 Parent-instance allowlist additions + +Add these hosts to `/etc/nitro_enclaves/vsock-proxy.yaml`: + +```yaml +- {address: api.push.apple.com, port: 443} +- {address: api.sandbox.push.apple.com, port: 443} +- {address: fcm.googleapis.com, port: 443} +``` + +`oauth2.googleapis.com` is already in the repo's current setup and can be reused for the FCM access-token exchange. + +### 16.2 Parent-instance proxy services + +Following the existing patterns in `docs/nitro-deploy.md`, add new parent services such as: + +- `vsock-apns-prod-proxy.service` +- `vsock-apns-sandbox-proxy.service` +- `vsock-fcm-proxy.service` + +Example mappings: + +- APNs prod on parent port `8024` +- APNs sandbox on parent port `8025` +- FCM on parent port `8029` + +Any unused ports are acceptable; they just need to stay consistent with enclave `traffic_forwarder.py` mappings. + +### 16.3 Enclave `entrypoint.sh` changes + +Add new `/etc/hosts` entries and forwarders, for example: + +```sh +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.34 fcm.googleapis.com" >> /etc/hosts + +run_forever tf_apns_prod python3 /app/traffic_forwarder.py 127.0.0.21 443 3 8024 & +run_forever tf_apns_sandbox python3 /app/traffic_forwarder.py 127.0.0.22 443 3 8025 & +run_forever tf_fcm python3 /app/traffic_forwarder.py 127.0.0.34 443 3 8029 & +``` + +The exact IPs and ports can be changed, but they should use currently unused slots. + +--- + +## 17. Product Event Model + +### 17.1 Live Sage request/response flows + +For the common flow where the user sends a message to Sage and the client keeps an SSE stream open waiting for responses: + +1. generate and stream the Sage response as normal +2. if the final response is successfully delivered over that SSE stream, do **not** enqueue push +3. if the SSE stream closes or errors before final response delivery completes, enqueue exactly one notification event for that interrupted turn +4. if multiple assistant messages were generated after the disconnect point, use the first missed assistant message as the preview source for that one notification +5. that push fan-out then goes to all active devices + +This deliberately treats a successful live SSE stream as β€œthe user likely saw it” and avoids more complex per-thread presence logic in v1. + +### 17.2 Internal enqueue API + +Do not create a public "send push" API in v1. + +Instead add an internal enqueue helper in `src/push/mod.rs` shaped roughly like: + +```rust +pub struct EnqueueNotificationRequest { + pub project_id: i32, + pub user_id: Uuid, + pub kind: String, + pub delivery_mode: PushDeliveryMode, + pub priority: PushPriority, + pub collapse_key: Option, + pub fallback_title: String, + pub fallback_body: String, + pub payload: Option, + pub not_before_at: Option>, + pub expires_at: Option>, +} +``` + +### 17.3 Initial event sources + +Likely first integrations: + +- `src/web/agent/runtime.rs` + - subagent follow-up or background completion signals + +- `src/web/responses/handlers.rs` + - async response completion hooks + - live Sage SSE completion / failure handling + +Reminder scheduling can come later; the outbox design already supports `not_before_at`. + +### 17.4 Notification identity and duplicate protection + +- `notification_events.uuid` is the canonical `notification_id` +- include `thread_id` and `message_id` in payloads whenever applicable +- retries must reuse the same `notification_id` +- APNs `apns-collapse-id` and FCM `collapse_key` should key off that `notification_id` +- clients should keep a recent-ID cache and no-op duplicates + +--- + +## 18. Validation and Test Plan + +### 18.1 Unit tests + +Add focused tests for: + +- P-256 envelope roundtrip in `src/push/crypto.rs` +- APNs JWT generation and refresh behavior +- FCM service-account assertion generation and token parsing +- push-device upsert / revoke / invalidate DB behavior +- delivery leasing and retry transitions + +### 18.2 Integration tests + +Add route tests for: + +- `POST /v1/push/devices` +- `GET /v1/push/devices` +- `DELETE /v1/push/devices/:id` +- `GET` / `PUT` platform push settings routes + +### 18.3 Provider client testability + +Do not hardcode transport base URLs so deeply that they cannot be mocked. + +The APNs and FCM sender code should allow test-only override of base URLs, making it possible to verify: + +- request headers +- request body shapes +- retry classification +- invalid-token handling + +--- + +## 19. Implementation Order + +### Phase 1: Schema and project settings + +- add migration for push tables +- add push model files and Diesel schema +- add `SettingCategory::Push` +- add project push settings routes +- add new project secret constants + +### Phase 2: Device registry API + +- add `src/web/push.rs` +- add push-device DB methods +- implement register / list / revoke routes +- add user-scoped tests + +### Phase 3: Outbox and generic delivery + +- add event / delivery DB methods +- add worker leasing loop +- wire live Sage SSE paths so successful stream completion skips push enqueue +- implement APNs generic delivery +- implement standard FCM visible delivery + +### Phase 4: Client dedup + foreground suppression + iOS encrypted preview + +- add stable `notification_id` handling end-to-end +- add client recent-ID cache guidance and integration +- add thread-based notification cleanup when a conversation opens +- add `p256` + `hkdf` +- implement per-device envelope encryption +- add iOS encrypted preview payload / NSE rewrite + +### Phase 5: Android optional encrypted-preview path + +- keep standard visible FCM delivery as the default +- optionally add Android data-only encrypted preview if product later decides it is worth the extra complexity +- if enabled, implement it in `FirebaseMessagingService` with local-notification posting only after successful decrypt + +### Phase 6: Product integrations and hardening + +- connect Maple / Sage event sources +- add collapse keys, TTL, and backoff policy tuning +- add metrics and invalid token cleanup dashboards +- optionally add stronger device attestation later + +--- + +## 20. Rejected Alternatives + +### 20.1 Shared per-user notification key + +Rejected because it makes revocation and secure local storage worse. + +### 20.2 Sender outside the enclave + +Rejected for v1 because it weakens the trust boundary around payload plaintext and provider secrets. + +### 20.3 Silent-only push as the only strategy + +Rejected because silent / background behavior is too unreliable for user-facing alerts. + +### 20.4 Reusing existing user-key encryption APIs + +Rejected because those APIs are built around server-derived secp256k1 secrets and do not map cleanly to per-device mobile decryption. + +--- + +## 21. Final Recommendation + +The v1 implementation should be: + +- **direct APNs + direct FCM HTTP v1** +- **push sender inside the enclave app** +- **project-scoped provider config in `project_settings` + `org_project_secrets`** +- **per-device P-256 notification keypairs** +- **Postgres outbox + multi-enclave delivery worker coordinated via Postgres row leases** +- **successful live Sage SSE delivery suppresses push entirely** +- **iOS encrypted preview with generic fallback** +- **Android v1 uses standard visible FCM notifications plus data, with app-open fetch and local foreground suppression** +- **stable `notification_id` plus client-side dedup / thread cleanup** + +This is the simplest design that still matches the current OpenSecret architecture and avoids painting us into a corner later. diff --git a/docs/nitro-deploy.md b/docs/nitro-deploy.md index 8426a8ba..58d39b94 100644 --- a/docs/nitro-deploy.md +++ b/docs/nitro-deploy.md @@ -530,6 +530,108 @@ A restart of this should not be needed but if you need to: sudo systemctl restart vsock-apple-proxy.service ``` +## Vsock mobile push proxies +Create vsock proxy services so that enclave program can talk to APNs and FCM: + +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} +``` + +Note: FCM OAuth token exchange reuses the existing Google OAuth proxy for `oauth2.googleapis.com` on port `8014`, so no additional oauth token proxy is needed here. + +These values should match the enclave-side mappings in `entrypoint.sh`: + +- `127.0.0.21 api.push.apple.com` -> parent port `8024` +- `127.0.0.22 api.sandbox.push.apple.com` -> parent port `8025` +- `127.0.0.34 fcm.googleapis.com` -> parent port `8029` + +Now create services that spin these 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 +``` + +``` +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 +``` + +``` +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 the 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 f4e56aaa..9db8f993 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -323,6 +323,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.34 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" @@ -442,6 +448,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.34 443 3 8029 & + # Start the traffic forwarders for Tinfoil proxy in the background log "Starting Tinfoil API GitHub proxy traffic forwarder" run_forever tf_tinfoil_api_github_proxy python3 /app/traffic_forwarder.py 127.0.0.16 443 3 8019 & diff --git a/migrations/2026-03-07-120000_push_notifications_v1/down.sql b/migrations/2026-03-07-120000_push_notifications_v1/down.sql new file mode 100644 index 00000000..e578d4bd --- /dev/null +++ b/migrations/2026-03-07-120000_push_notifications_v1/down.sql @@ -0,0 +1,10 @@ +DROP TRIGGER IF EXISTS update_notification_deliveries_updated_at ON notification_deliveries; +DROP TRIGGER IF EXISTS update_push_devices_updated_at ON push_devices; + +DROP INDEX IF EXISTS idx_notification_deliveries_pending; +DROP INDEX IF EXISTS idx_notification_events_user_due; +DROP INDEX IF EXISTS idx_push_devices_user_active; + +DROP TABLE IF EXISTS notification_deliveries; +DROP TABLE IF EXISTS notification_events; +DROP TABLE IF EXISTS push_devices; diff --git a/migrations/2026-03-07-120000_push_notifications_v1/up.sql b/migrations/2026-03-07-120000_push_notifications_v1/up.sql new file mode 100644 index 00000000..766cfd16 --- /dev/null +++ b/migrations/2026-03-07-120000_push_notifications_v1/up.sql @@ -0,0 +1,86 @@ +CREATE TABLE push_devices ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT uuid_generate_v4() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + installation_id UUID NOT NULL, + platform TEXT NOT NULL CHECK (platform IN ('ios', 'android')), + provider TEXT NOT NULL CHECK (provider IN ('apns', 'fcm')), + environment TEXT NOT NULL CHECK (environment IN ('dev', 'prod')), + app_id TEXT NOT NULL, + push_token_enc BYTEA NOT NULL, + push_token_hash BYTEA NOT NULL, + notification_public_key BYTEA NOT NULL, + key_algorithm TEXT NOT NULL CHECK (key_algorithm IN ('p256_ecdh_v1')), + supports_encrypted_preview BOOLEAN NOT NULL DEFAULT false, + supports_background_processing BOOLEAN NOT NULL DEFAULT false, + last_seen_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + 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 TABLE notification_events ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT uuid_generate_v4() UNIQUE, + project_id INTEGER NOT NULL REFERENCES org_projects(id) ON DELETE CASCADE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + kind TEXT NOT NULL, + delivery_mode TEXT NOT NULL CHECK (delivery_mode IN ('generic', 'encrypted_preview')), + priority TEXT NOT NULL DEFAULT 'normal' CHECK (priority IN ('normal', 'high')), + collapse_key TEXT, + fallback_title TEXT NOT NULL, + fallback_body TEXT NOT NULL, + payload_enc BYTEA, + not_before_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + expires_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + cancelled_at TIMESTAMP WITH TIME ZONE +); + +CREATE TABLE notification_deliveries ( + id BIGSERIAL PRIMARY KEY, + event_id BIGINT NOT NULL REFERENCES notification_events(id) ON DELETE CASCADE, + push_device_id BIGINT NOT NULL REFERENCES push_devices(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending' CHECK ( + status IN ('pending', 'leased', 'sent', 'retry', 'failed', 'invalid_token', 'cancelled') + ), + attempt_count INTEGER NOT NULL DEFAULT 0, + next_attempt_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + lease_owner TEXT, + lease_expires_at TIMESTAMP WITH TIME ZONE, + provider_message_id TEXT, + provider_status_code INTEGER, + last_error TEXT, + sent_at TIMESTAMP WITH TIME ZONE, + invalidated_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, + UNIQUE (event_id, push_device_id) +); + +CREATE INDEX idx_push_devices_user_active + ON push_devices(user_id, revoked_at); + +CREATE UNIQUE INDEX idx_push_devices_installation_active + ON push_devices(installation_id, environment) + WHERE revoked_at IS NULL; + +CREATE UNIQUE INDEX idx_push_devices_token_active + ON push_devices(provider, environment, push_token_hash) + WHERE revoked_at IS NULL; + +CREATE INDEX idx_notification_events_user_due + ON notification_events(user_id, not_before_at, cancelled_at); + +CREATE INDEX idx_notification_deliveries_pending + ON notification_deliveries(status, next_attempt_at); + +CREATE TRIGGER update_push_devices_updated_at + BEFORE UPDATE ON push_devices + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); + +CREATE TRIGGER update_notification_deliveries_updated_at + BEFORE UPDATE ON notification_deliveries + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/db.rs b/src/db.rs index 18429983..f7709569 100644 --- a/src/db.rs +++ b/src/db.rs @@ -24,10 +24,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::{ AssistantMessage, Conversation, NewAssistantMessage, NewConversation, NewReasoningItem, NewResponse, NewToolCall, NewToolOutput, NewUserInstruction, NewUserMessage, RawThreadMessage, @@ -390,6 +390,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, @@ -1634,6 +1642,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 4c7e2473..a2faccdc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,13 +14,14 @@ use crate::models::account_deletion::{AccountDeletionError, NewAccountDeletionRe 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, conversations_routes, document_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}; @@ -86,6 +87,7 @@ mod oauth; mod os_flags; mod private_key; mod proxy_config; +mod push; mod rag; mod sqs; mod tokens; @@ -1702,7 +1704,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)? @@ -1711,15 +1723,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( @@ -2758,6 +2779,8 @@ async fn main() -> Result<(), Error> { ) .await?; + start_push_worker(app_state.clone()); + let cors = CorsLayer::new() // allow all method types .allow_methods(Any) @@ -2798,6 +2821,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/mod.rs b/src/models/mod.rs index 79a3578f..b85ba507 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -5,6 +5,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; @@ -16,6 +18,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..46e6015e --- /dev/null +++ b/src/models/notification_deliveries.rs @@ -0,0 +1,226 @@ +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(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 id + FROM notification_deliveries + WHERE status IN ('pending', 'retry') + AND next_attempt_at <= NOW() + AND (lease_expires_at IS NULL OR lease_expires_at < NOW()) + ORDER BY next_attempt_at ASC, id ASC + FOR UPDATE SKIP LOCKED + LIMIT $1 + ) + UPDATE notification_deliveries + SET status = 'leased', + lease_owner = $2, + lease_expires_at = NOW() + ($3 * INTERVAL '1 second'), + updated_at = NOW() + WHERE id IN (SELECT id FROM candidates) + RETURNING * + "#; + + 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, + provider_message_id: Option<&str>, + provider_status_code: Option, + ) -> Result<(), NotificationDeliveryError> { + diesel::update( + notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + ) + .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(|_| ()) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_retry( + conn: &mut PgConnection, + lookup_id: i64, + provider_status_code: Option, + last_error: Option<&str>, + retry_after_seconds: i32, + ) -> Result<(), NotificationDeliveryError> { + let query = r#" + UPDATE notification_deliveries + SET status = 'retry', + attempt_count = attempt_count + 1, + provider_status_code = $2, + last_error = $3, + next_attempt_at = NOW() + ($4 * INTERVAL '1 second'), + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + "#; + + sql_query(query) + .bind::(lookup_id) + .bind::, _>(provider_status_code) + .bind::, _>(last_error) + .bind::(retry_after_seconds) + .execute(conn) + .map(|_| ()) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_failed( + conn: &mut PgConnection, + lookup_id: i64, + provider_status_code: Option, + last_error: Option<&str>, + ) -> Result<(), NotificationDeliveryError> { + diesel::update( + notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + ) + .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(|_| ()) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_invalid_token( + conn: &mut PgConnection, + lookup_id: i64, + provider_status_code: Option, + last_error: Option<&str>, + ) -> Result<(), NotificationDeliveryError> { + diesel::update( + notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + ) + .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(|_| ()) + .map_err(NotificationDeliveryError::DatabaseError) + } + + pub fn mark_cancelled( + conn: &mut PgConnection, + lookup_id: i64, + last_error: Option<&str>, + ) -> Result<(), NotificationDeliveryError> { + diesel::update( + notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + ) + .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(|_| ()) + .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, +} + +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..7b4492bd --- /dev/null +++ b/src/models/notification_events.rs @@ -0,0 +1,91 @@ +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"; + +#[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>, +} + +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>, +} + +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..ee63f862 --- /dev/null +++ b/src/models/push_devices.rs @@ -0,0 +1,237 @@ +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 installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub push_token_enc: Vec, + pub push_token_hash: Vec, + pub notification_public_key: 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> { + 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)) + .first::(conn) + .optional() + .map_err(PushDeviceError::DatabaseError) + } + + pub fn get_by_installation( + conn: &mut PgConnection, + lookup_installation_id: Uuid, + lookup_environment: &str, + ) -> Result, PushDeviceError> { + if let Some(active_device) = 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)? + { + return Ok(Some(active_device)); + } + + push_devices::table + .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_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::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_enc.eq(&self.push_token_enc), + push_devices::push_token_hash.eq(&self.push_token_hash), + push_devices::notification_public_key.eq(&self.notification_public_key), + 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 user_id: Uuid, + pub installation_id: Uuid, + pub platform: String, + pub provider: String, + pub environment: String, + pub app_id: String, + pub push_token_enc: Vec, + pub push_token_hash: Vec, + pub notification_public_key: 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 760ffb23..06b2b37d 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -127,6 +127,46 @@ diesel::table! { } } +diesel::table! { + notification_deliveries (id) { + id -> Int8, + event_id -> Int8, + push_device_id -> Int8, + status -> Text, + attempt_count -> Int4, + next_attempt_at -> Timestamptz, + lease_owner -> Nullable, + lease_expires_at -> Nullable, + provider_message_id -> Nullable, + provider_status_code -> Nullable, + last_error -> Nullable, + sent_at -> Nullable, + invalidated_at -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + +diesel::table! { + notification_events (id) { + id -> Int8, + uuid -> Uuid, + project_id -> Int4, + user_id -> Uuid, + kind -> Text, + delivery_mode -> Text, + priority -> Text, + collapse_key -> Nullable, + fallback_title -> Text, + fallback_body -> Text, + payload_enc -> Nullable, + not_before_at -> Timestamptz, + expires_at -> Nullable, + created_at -> Timestamptz, + cancelled_at -> Nullable, + } +} + diesel::table! { oauth_providers (id) { id -> Int4, @@ -255,6 +295,29 @@ diesel::table! { } } +diesel::table! { + push_devices (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + installation_id -> Uuid, + platform -> Text, + provider -> Text, + environment -> Text, + app_id -> Text, + push_token_enc -> Bytea, + push_token_hash -> Bytea, + notification_public_key -> Bytea, + key_algorithm -> Text, + supports_encrypted_preview -> Bool, + supports_background_processing -> Bool, + last_seen_at -> Timestamptz, + revoked_at -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { reasoning_items (id) { id -> Int8, @@ -447,6 +510,9 @@ diesel::joinable!(assistant_messages -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> responses (response_id)); diesel::joinable!(conversation_summaries -> conversations (conversation_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 -> org_projects (project_id)); diesel::joinable!(org_memberships -> orgs (org_id)); diesel::joinable!(org_project_secrets -> org_projects (project_id)); diesel::joinable!(org_projects -> orgs (org_id)); @@ -478,6 +544,8 @@ diesel::allow_tables_to_appear_in_same_query!( enclave_secrets, invite_codes, memory_blocks, + notification_deliveries, + notification_events, oauth_providers, org_memberships, org_project_secrets, @@ -489,6 +557,7 @@ diesel::allow_tables_to_appear_in_same_query!( platform_password_reset_requests, platform_users, project_settings, + push_devices, reasoning_items, responses, token_usage, diff --git a/src/push/apns.rs b/src/push/apns.rs new file mode 100644 index 00000000..c624f39a --- /dev/null +++ b/src/push/apns.rs @@ -0,0 +1,263 @@ +use crate::models::notification_events::NotificationEvent; +use crate::models::project_settings::IosPushSettings; +use crate::models::push_devices::PushDevice; +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 push_token: &'a str, + 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.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.push_token + ), + _ => format!("https://api.push.apple.com/3/device/{}", request.push_token), + }; + + let mut http_request = transport + .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, + preview_payload: Option<&NotificationPreviewPayload>, + send_encrypted_preview: bool, +) -> Result { + 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, 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/crypto.rs b/src/push/crypto.rs new file mode 100644 index 00000000..e1ce2a57 --- /dev/null +++ b/src/push/crypto.rs @@ -0,0 +1,151 @@ +use crate::encrypt::generate_random; +use crate::models::push_devices::PushDevice; +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; + +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: &PushDevice, + payload: &NotificationPreviewPayload, +) -> Result { + let recipient_key = PublicKey::from_public_key_der(&device.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 crate::models::push_devices::{ + PushDevice, PUSH_ENV_PROD, PUSH_KEY_ALGORITHM_P256_ECDH_V1, PUSH_PLATFORM_IOS, + PUSH_PROVIDER_APNS, + }; + 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 = PushDevice { + id: 1, + uuid: Uuid::new_v4(), + user_id: Uuid::new_v4(), + 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_enc: vec![1, 2, 3], + push_token_hash: vec![4, 5, 6], + notification_public_key: recipient_public + .to_public_key_der() + .unwrap() + .as_bytes() + .to_vec(), + 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(), + }; + + let payload = NotificationPreviewPayload { + v: 1, + notification_id: Uuid::new_v4(), + message_id: Uuid::new_v4(), + kind: "agent.message".to_string(), + title: "Sage".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, &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"); + } +} diff --git a/src/push/fcm.rs b/src/push/fcm.rs new file mode 100644 index 00000000..db547385 --- /dev/null +++ b/src/push/fcm.rs @@ -0,0 +1,299 @@ +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() { + return Err(PushError::ProviderError(format!( + "FCM OAuth failed ({}): {}", + response_status, response_body + ))); + } + + 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)); + + 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": "sage_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..924cb138 --- /dev/null +++ b/src/push/mod.rs @@ -0,0 +1,388 @@ +pub mod apns; +pub mod crypto; +pub mod fcm; +pub mod worker; + +use crate::db::DBError; +use crate::encrypt::encrypt_with_key; +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, +}; +use crate::models::project_settings::ProjectSettingError; +use crate::models::push_devices::{PushDevice, PushDeviceError}; +use crate::models::users::User; +use crate::{AppState, Error}; +use chrono::{DateTime, Duration, Utc}; +use diesel::Connection; +use secp256k1::SecretKey; +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 Sage message"; +pub const AGENT_NOTIFICATION_FALLBACK_BODY: &str = "Open Maple to view it"; +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), +} + +#[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), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NotificationPreviewPayload { + pub v: i32, + pub notification_id: Uuid, + 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, +} + +#[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 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_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()?; + + Ok(Self { + 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 { + let payload = NotificationPreviewPayload { + v: 1, + notification_id, + message_id: preview_payload.message_id, + kind: preview_payload.kind, + title: preview_payload.title, + body: preview_payload.body, + deep_link: preview_payload.deep_link, + thread_id: preview_payload.thread_id, + sent_at: preview_payload.sent_at, + }; + let enclave_key = SecretKey::from_slice(&state.enclave_key) + .map_err(|e| PushError::InvalidSecret(e.to_string()))?; + let payload_bytes = serde_json::to_vec(&payload)?; + Some(encrypt_with_key(&enclave_key, &payload_bytes).await) + } 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, + 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, + }) + .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) = 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: "Sage".to_string(), + body: normalize_preview_body(message_text), + deep_link, + thread_id, + sent_at: Utc::now().timestamp(), + }), + 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..efe59842 --- /dev/null +++ b/src/push/worker.rs @@ -0,0 +1,376 @@ +use crate::encrypt::decrypt_with_key; +use crate::models::notification_deliveries::NotificationDelivery; +use crate::models::notification_events::NotificationEvent; +use crate::models::push_devices::{PushDevice, PUSH_PLATFORM_ANDROID, PUSH_PLATFORM_IOS}; +use crate::push::apns::{send_apns_notification, ApnsSendRequest}; +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 secp256k1::SecretKey; +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 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 => { + NotificationDelivery::mark_failed(&mut conn, delivery.id, None, Some("event missing"))?; + return Ok(()); + } + }; + + if event.cancelled_at.is_some() { + NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("event cancelled"))?; + return Ok(()); + } + + if event + .expires_at + .is_some_and(|expires_at| expires_at <= Utc::now()) + { + NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("event expired"))?; + return Ok(()); + } + + let device = match PushDevice::get_by_id(&mut conn, delivery.push_device_id)? { + Some(device) => device, + None => { + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + None, + Some("device missing"), + )?; + return Ok(()); + } + }; + + if device.revoked_at.is_some() { + NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("device revoked"))?; + 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, + } => { + NotificationDelivery::mark_sent( + &mut conn, + delivery.id, + provider_message_id.as_deref(), + provider_status_code, + )?; + } + PushSendOutcome::Retryable { + provider_status_code, + error, + } => { + if delivery.attempt_count + 1 >= PUSH_WORKER_MAX_ATTEMPTS { + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + provider_status_code, + Some(&error), + )?; + } else { + NotificationDelivery::mark_retry( + &mut conn, + delivery.id, + provider_status_code, + Some(&error), + retry_backoff_seconds(delivery.attempt_count + 1), + )?; + } + } + PushSendOutcome::InvalidToken { + provider_status_code, + error, + } => { + conn.transaction::<(), PushError, _>(|conn| { + NotificationDelivery::mark_invalid_token( + conn, + delivery.id, + provider_status_code, + Some(&error), + )?; + PushDevice::invalidate(conn, device.id)?; + Ok(()) + })?; + } + PushSendOutcome::Failed { + provider_status_code, + error, + } => { + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + provider_status_code, + Some(&error), + )?; + } + } + + Ok(()) +} + +async fn build_send_outcome( + state: &Arc, + transport: &PushTransport, + event: &NotificationEvent, + device: &PushDevice, +) -> Result { + let preview_payload = decrypt_preview_payload(state, event)?; + let push_token = decrypt_push_token(state, device)?; + + dispatch_delivery( + state, + transport, + event, + device, + &push_token, + preview_payload.as_ref(), + ) + .await +} + +async fn dispatch_delivery( + state: &Arc, + transport: &PushTransport, + event: &NotificationEvent, + device: &PushDevice, + push_token: &str, + 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() + && device.supports_encrypted_preview; + + send_apns_notification( + state, + transport, + ApnsSendRequest { + event, + device, + push_token, + 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, + push_token, + android_settings, + preview_payload, + ) + .await + } + _ => Ok(PushSendOutcome::Failed { + provider_status_code: None, + error: format!("unsupported push platform: {}", device.platform), + }), + } +} + +fn decrypt_preview_payload( + state: &Arc, + event: &NotificationEvent, +) -> Result, PushError> { + let Some(payload_enc) = &event.payload_enc else { + return Ok(None); + }; + + let secret_key = SecretKey::from_slice(&state.enclave_key) + .map_err(|e| PushError::InvalidSecret(e.to_string()))?; + let plaintext = decrypt_with_key(&secret_key, payload_enc) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + let payload = serde_json::from_slice::(&plaintext)?; + + Ok(Some(payload)) +} + +fn decrypt_push_token(state: &Arc, device: &PushDevice) -> Result { + let secret_key = SecretKey::from_slice(&state.enclave_key) + .map_err(|e| PushError::InvalidSecret(e.to_string()))?; + let plaintext = decrypt_with_key(&secret_key, &device.push_token_enc) + .map_err(|e| PushError::CryptoError(e.to_string()))?; + + String::from_utf8(plaintext).map_err(|e| PushError::InvalidSecret(e.to_string())) +} + +fn classify_internal_push_error(error: PushError) -> PushSendOutcome { + match error { + PushError::ConnectionError + | PushError::DatabaseError(_) + | PushError::DbError(_) + | PushError::HttpError(_) => PushSendOutcome::Retryable { + provider_status_code: None, + error: error.to_string(), + }, + _ => PushSendOutcome::Failed { + provider_status_code: None, + error: error.to_string(), + }, + } +} + +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 +} diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 61ed882a..4a99f552 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -9,6 +9,7 @@ use futures::Stream; use secp256k1::SecretKey; use serde::{Deserialize, Serialize}; use std::{convert::Infallible, sync::Arc, time::Duration}; +use tokio::sync::{mpsc, oneshot}; use tokio::time::sleep; use tracing::{error, warn}; use uuid::Uuid; @@ -20,6 +21,7 @@ use crate::models::agents::{ 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, @@ -172,6 +174,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, @@ -693,252 +706,320 @@ async fn chat_with_target( 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_target = target.clone(); - let user_key = match state.get_user_key(user.uuid, 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, 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(Sse::new(event_stream)) +} + +async fn run_agent_chat_task( + state: Arc, + user: User, + 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.uuid, 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(); + let message_count = messages.len(); - for (idx, msg) in messages.into_iter().enumerate() { - let event_data = AgentMessageEvent { - messages: vec![msg], - step: step_num, - }; + for (idx, msg) in messages.into_iter().enumerate() { + let assistant_message = match runtime.insert_assistant_message(&msg).await { + Ok(message) => Some(message), + Err(e) => { + error!("Failed to persist assistant message: {e:?}"); + None + } + }; - 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; + let delivered = send_agent_message_event( + &tx, + AgentMessageEvent { + messages: vec![msg.clone()], + step: step_num, + }, + client_connected, + ) + .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 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; - } - } + if idx + 1 < message_count && client_connected { + client_connected = send_agent_client_event( + &tx, + AgentClientEvent::Typing(AgentTypingEvent { step: step_num }), + client_connected, + ) + .await; + if client_connected { 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; + + 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( diff --git a/src/web/mod.rs b/src/web/mod.rs index feb5e90b..f7c96e9d 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -10,6 +10,7 @@ mod openai; pub mod openai_auth; pub mod platform; pub mod protected_routes; +pub mod push; pub mod rag; pub mod responses; @@ -22,6 +23,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::conversations_router as conversations_routes; pub use responses::instructions_router as instructions_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..3b4d8359 --- /dev/null +++ b/src/web/push.rs @@ -0,0 +1,286 @@ +use crate::encrypt::encrypt_with_key; +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::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 secp256k1::SecretKey; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +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, +} + +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 = Sha256::digest(body.push_token.as_bytes()).to_vec(); + let enclave_key = + SecretKey::from_slice(&state.enclave_key).map_err(|_| ApiError::InternalServerError)?; + let push_token_enc = encrypt_with_key(&enclave_key, body.push_token.as_bytes()).await; + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + let now = Utc::now(); + let device = conn + .transaction::(|conn| { + let existing = + 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, + )?; + + if let Some(conflict) = conflicting_token_owner.as_ref() { + if existing.as_ref().map(|device| device.id) != Some(conflict.id) { + PushDevice::revoke_by_id(conn, conflict.id)?; + } + } + + if let Some(mut existing) = existing { + existing.user_id = user.uuid; + existing.platform = body.platform.clone(); + existing.provider = body.provider.clone(); + existing.environment = body.environment.clone(); + existing.app_id = body.app_id.clone(); + existing.push_token_enc = push_token_enc.clone(); + existing.push_token_hash = token_hash.clone(); + existing.notification_public_key = public_key_bytes.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 { + NewPushDevice { + user_id: user.uuid, + 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_enc: push_token_enc.clone(), + push_token_hash: token_hash.clone(), + notification_public_key: public_key_bytes.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.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 map_push_device_error(error: PushDeviceError) -> ApiError { + error!("Push device database error: {:?}", error); + ApiError::InternalServerError +} + +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, + } + } +} From 206837375e6eb410a816ee9f8c9573598c15a7b2 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 9 Mar 2026 20:36:39 -0500 Subject: [PATCH 21/34] feat: align persisted agent branding with Maple Reflect the Maple AI secure messaging identity across agent prompts, push copy, and migration naming. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- .../down.sql | 0 .../up.sql | 2 +- .../up.sql | 2 +- src/push/crypto.rs | 2 +- src/push/fcm.rs | 2 +- src/push/mod.rs | 6 +++--- src/web/agent/mod.rs | 2 +- src/web/agent/runtime.rs | 16 ++++++++-------- src/web/agent/signatures.rs | 6 +++--- 9 files changed, 19 insertions(+), 19 deletions(-) rename migrations/{2026-02-10-214235_sage_agent_mvp_storage => 2026-02-10-214235_maple_agent_mvp_storage}/down.sql (100%) rename migrations/{2026-02-10-214235_sage_agent_mvp_storage => 2026-02-10-214235_maple_agent_mvp_storage}/up.sql (98%) diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql b/migrations/2026-02-10-214235_maple_agent_mvp_storage/down.sql similarity index 100% rename from migrations/2026-02-10-214235_sage_agent_mvp_storage/down.sql rename to migrations/2026-02-10-214235_maple_agent_mvp_storage/down.sql diff --git a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql b/migrations/2026-02-10-214235_maple_agent_mvp_storage/up.sql similarity index 98% rename from migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql rename to migrations/2026-02-10-214235_maple_agent_mvp_storage/up.sql index 2753e6bd..8feb4e4c 100644 --- a/migrations/2026-02-10-214235_sage_agent_mvp_storage/up.sql +++ b/migrations/2026-02-10-214235_maple_agent_mvp_storage/up.sql @@ -1,4 +1,4 @@ --- Sage MVP agent storage tables +-- Maple MVP agent storage tables -- -- Note: update_updated_at_column() already exists from previous migrations. diff --git a/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql b/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql index 1ad13d0f..6a49f909 100644 --- a/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql +++ b/migrations/2026-02-24-150000_user_messages_attachment_text/up.sql @@ -1,6 +1,6 @@ -- Add optional derived image description storage for user messages -- --- This stores Sage-style vision pre-processing results (encrypted in app layer) +-- This stores Maple-style vision pre-processing results (encrypted in app layer) -- without mutating the original MessageContent JSON. ALTER TABLE user_messages diff --git a/src/push/crypto.rs b/src/push/crypto.rs index e1ce2a57..59af8b8d 100644 --- a/src/push/crypto.rs +++ b/src/push/crypto.rs @@ -112,7 +112,7 @@ mod tests { notification_id: Uuid::new_v4(), message_id: Uuid::new_v4(), kind: "agent.message".to_string(), - title: "Sage".to_string(), + title: "Maple".to_string(), body: "Hello from Maple".to_string(), deep_link: "opensecret://agent".to_string(), thread_id: "agent:main".to_string(), diff --git a/src/push/fcm.rs b/src/push/fcm.rs index db547385..b0591b4c 100644 --- a/src/push/fcm.rs +++ b/src/push/fcm.rs @@ -271,7 +271,7 @@ fn build_fcm_payload( "ttl": format!("{}s", ttl_seconds), "collapse_key": collapse_key, "notification": { - "channel_id": "sage_messages", + "channel_id": "maple_messages", "tag": event.uuid.to_string(), "click_action": "OPEN_THREAD", } diff --git a/src/push/mod.rs b/src/push/mod.rs index 924cb138..52582381 100644 --- a/src/push/mod.rs +++ b/src/push/mod.rs @@ -26,8 +26,8 @@ use tokio::sync::RwLock; use tracing::warn; use uuid::Uuid; -pub const AGENT_NOTIFICATION_FALLBACK_TITLE: &str = "New Sage message"; -pub const AGENT_NOTIFICATION_FALLBACK_BODY: &str = "Open Maple to view it"; +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; @@ -352,7 +352,7 @@ pub async fn enqueue_agent_message_notification( preview_payload: Some(NotificationPreviewPayloadInput { message_id, kind: NOTIFICATION_KIND_AGENT_MESSAGE.to_string(), - title: "Sage".to_string(), + title: "Maple".to_string(), body: normalize_preview_body(message_text), deep_link, thread_id, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 4a99f552..a1210060 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -106,7 +106,7 @@ enum ChatTarget { Subagent(Uuid), } -const MAIN_AGENT_DISPLAY_NAME: &str = "Sage"; +const MAIN_AGENT_DISPLAY_NAME: &str = "Maple"; fn default_limit() -> i64 { DEFAULT_PAGINATION_LIMIT diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 34504b06..d7545250 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -39,7 +39,7 @@ use super::tools::{ }; 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-5"; pub const DEFAULT_CONTEXT_WINDOW: i32 = 256_000; pub const DEFAULT_COMPACTION_THRESHOLD: f32 = 0.80; @@ -97,7 +97,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( @@ -729,7 +729,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() @@ -747,8 +747,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" { @@ -931,7 +931,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; @@ -1080,7 +1080,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(); @@ -1132,7 +1132,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 498c81ec..07d42355 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -35,9 +35,9 @@ OUTPUT FORMAT (exactly 2 fields): Each [[ ## field ## ]] marker MUST be on its own line."#; -/// Default instruction for the agent (GEPA-optimized in Sage) +/// Default instruction for the agent (GEPA-optimized for Maple) /// 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. +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. @@ -102,7 +102,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 From dd634aa9b9751c518edf952719e4f237dcfbd99b Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 10 Mar 2026 13:41:38 -0500 Subject: [PATCH 22/34] feat: add Brave web search to Maple agent Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/brave.rs | 759 +++++++++++++++++++++++++++++++++--- src/web/agent/runtime.rs | 4 + src/web/agent/signatures.rs | 11 +- src/web/agent/tools.rs | 48 +++ 4 files changed, 768 insertions(+), 54 deletions(-) diff --git a/src/brave.rs b/src/brave.rs index 7f34e345..d345b794 100644 --- a/src/brave.rs +++ b/src/brave.rs @@ -1,11 +1,7 @@ -//! Minimal Brave Search API client -//! -//! This module provides a lightweight client for the Brave Search API. -//! Only includes what we actually use - no bloat from auto-generated code. - use serde::{Deserialize, Serialize}; use std::sync::Arc; use std::time::Duration; +use tracing::{debug, warn}; const BRAVE_API_BASE: &str = "https://api.search.brave.com/res/v1"; const REQUEST_TIMEOUT: Duration = Duration::from_secs(20); @@ -19,7 +15,14 @@ pub enum BraveError { Api { status: u16, message: String }, } -/// Brave API client with reusable HTTP client and stored API key +#[derive(Debug, Clone, Default)] +pub struct SearchOptions { + pub count: Option, + 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/web/agent/runtime.rs b/src/web/agent/runtime.rs index d7545250..5b2904bd 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -36,6 +36,7 @@ use super::signatures::{ use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, + WebSearchTool, }; use super::vision; @@ -399,6 +400,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(), diff --git a/src/web/agent/signatures.rs b/src/web/agent/signatures.rs index 07d42355..e3d0d682 100644 --- a/src/web/agent/signatures.rs +++ b/src/web/agent/signatures.rs @@ -36,7 +36,6 @@ OUTPUT FORMAT (exactly 2 fields): Each [[ ## field ## ]] marker MUST be on its own line."#; /// Default instruction for the agent (GEPA-optimized for Maple) -/// NOTE: This MVP excludes web_search/shell/scheduler/reminders. pub const AGENT_INSTRUCTION: &str = r#"You are Maple, a companion and friend who happens to be an AI. WHO YOU ARE: @@ -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) @@ -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 // ============================================================================ From 2d7c95010144e8fb2d8cacfedeb65f89945673b3 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 10 Mar 2026 20:21:00 -0500 Subject: [PATCH 23/34] Move agent msg timing to client --- src/web/agent/mod.rs | 57 ++++++++++++++++++++++++++------------------ 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index a1210060..2ed4b356 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -1,7 +1,11 @@ use axum::{ extract::{Path, Query, State}, + http::{header, HeaderName, HeaderValue}, middleware::from_fn_with_state, - response::sse::{Event, Sse}, + response::{ + sse::{Event, KeepAlive, Sse}, + IntoResponse, Response, + }, routing::{delete, get, post}, Extension, Json, Router, }; @@ -10,7 +14,6 @@ use secp256k1::SecretKey; use serde::{Deserialize, Serialize}; use std::{convert::Infallible, sync::Arc, time::Duration}; use tokio::sync::{mpsc, oneshot}; -use tokio::time::sleep; use tracing::{error, warn}; use uuid::Uuid; @@ -149,8 +152,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 { @@ -199,6 +201,29 @@ 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, @@ -658,7 +683,7 @@ async fn chat_main( Extension(session_id): Extension, Extension(user): Extension, Extension(body): Extension, -) -> Result>>, ApiError> { +) -> Result { chat_with_target(state, session_id, user, body, ChatTarget::Main).await } @@ -668,7 +693,7 @@ async fn chat_subagent( Extension(session_id): Extension, Extension(user): Extension, Extension(body): Extension, -) -> Result>>, ApiError> { +) -> Result { let mut conn = state .db .get_pool() @@ -699,7 +724,7 @@ async fn chat_with_target( user: User, 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) { @@ -758,7 +783,7 @@ async fn chat_with_target( } }; - Ok(Sse::new(event_stream)) + Ok(agent_sse_response(event_stream)) } async fn run_agent_chat_task( @@ -867,9 +892,8 @@ async fn run_agent_chat_task( if !messages.is_empty() { total_messages += messages.len(); - let message_count = messages.len(); - for (idx, msg) in messages.into_iter().enumerate() { + for msg in messages { let assistant_message = match runtime.insert_assistant_message(&msg).await { Ok(message) => Some(message), Err(e) => { @@ -899,19 +923,6 @@ async fn run_agent_chat_task( } else { client_connected = true; } - - if idx + 1 < message_count && client_connected { - client_connected = send_agent_client_event( - &tx, - AgentClientEvent::Typing(AgentTypingEvent { step: step_num }), - client_connected, - ) - .await; - - if client_connected { - sleep(Duration::from_millis(MESSAGE_STAGGER_DELAY_MS)).await; - } - } } } From 094fdc2e71a125815b40a378c20846f0f7e85e8d Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Fri, 13 Mar 2026 00:19:41 -0500 Subject: [PATCH 24/34] fix: harden push delivery isolation and scheduling Prevent cross-user device reuse from leaking queued notifications, reclaim expired leases, respect deferred sends, and clarify that v1 protects plaintext notification content rather than all provider-visible metadata. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- docs/encrypted-mobile-push-notifications.md | 4 +++ src/models/notification_deliveries.rs | 26 +++++++++------ src/models/push_devices.rs | 27 +++++++++------- src/push/apns.rs | 2 ++ src/push/fcm.rs | 14 ++++++--- src/push/mod.rs | 3 ++ src/push/worker.rs | 35 ++++++++++++++++++++- src/web/push.rs | 27 +++++++++++++--- 8 files changed, 107 insertions(+), 31 deletions(-) diff --git a/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md index 61c10eb2..d1154cf9 100644 --- a/docs/encrypted-mobile-push-notifications.md +++ b/docs/encrypted-mobile-push-notifications.md @@ -41,6 +41,7 @@ We want Maple / Sage to send notifications for events like: while ensuring that: - Apple and Google do not need plaintext notification content +- minimal routing metadata may remain provider-visible in APNs / FCM payloads, and that is acceptable for v1 - raw push tokens are not stored plaintext in Postgres - losing one device does not compromise every other device - the design fits the current OpenSecret enclave + project-scoped secret model @@ -53,6 +54,7 @@ This design does **not** attempt to provide: - a double-ratchet messaging protocol - server blindness to the notification plaintext at creation time +- secrecy for all push metadata, timing information, routing identifiers, or provider-visible fields - guaranteed hidden lock-screen content after the device itself decrypts it - a single cross-device shared notification private key @@ -834,6 +836,7 @@ Important: - `mutable-content: 1` is required for the NSE to run - if the NSE times out, iOS shows the original generic alert - `os_meta` is allowed to carry minimal routing metadata (`notification_id`, `message_id`, `thread_id`, `deep_link`, `kind`), but not plaintext notification content +- this is intentional: the privacy requirement is protecting notification **content**, not hiding all routing metadata from the push provider ### 13.4 APNs error handling @@ -926,6 +929,7 @@ Notes: - visible text should remain generic / non-sensitive - FCM `data` values must be strings - FCM `data` is provider-visible routing metadata, so it may include fields like `notification_id`, `message_id`, `thread_id`, `deep_link`, and `kind`, but never plaintext message content +- this is intentional in v1: the goal is to keep message content encrypted from the push provider, not to hide all metadata needed for client routing - reuse the same `collapse_key` only for retries of the same logical notification - `notification_id` is for exact dedup; `thread_id` is for grouping / clearing when the user opens the conversation diff --git a/src/models/notification_deliveries.rs b/src/models/notification_deliveries.rs index 46e6015e..572d7225 100644 --- a/src/models/notification_deliveries.rs +++ b/src/models/notification_deliveries.rs @@ -57,22 +57,27 @@ impl NotificationDelivery { ) -> Result, NotificationDeliveryError> { let query = r#" WITH candidates AS ( - SELECT id - FROM notification_deliveries - WHERE status IN ('pending', 'retry') - AND next_attempt_at <= NOW() - AND (lease_expires_at IS NULL OR lease_expires_at < NOW()) - ORDER BY next_attempt_at ASC, id ASC - FOR UPDATE SKIP LOCKED + 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 + UPDATE notification_deliveries d SET status = 'leased', lease_owner = $2, lease_expires_at = NOW() + ($3 * INTERVAL '1 second'), updated_at = NOW() - WHERE id IN (SELECT id FROM candidates) - RETURNING * + FROM candidates + WHERE d.id = candidates.id + RETURNING d.* "#; sql_query(query) @@ -211,6 +216,7 @@ impl NotificationDelivery { pub struct NewNotificationDelivery { pub event_id: i64, pub push_device_id: i64, + pub next_attempt_at: DateTime, } impl NewNotificationDelivery { diff --git a/src/models/push_devices.rs b/src/models/push_devices.rs index ee63f862..1a7a57a0 100644 --- a/src/models/push_devices.rs +++ b/src/models/push_devices.rs @@ -73,10 +73,24 @@ impl PushDevice { 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) @@ -87,21 +101,10 @@ impl PushDevice { lookup_installation_id: Uuid, lookup_environment: &str, ) -> Result, PushDeviceError> { - if let Some(active_device) = 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)? - { - return Ok(Some(active_device)); - } - 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() diff --git a/src/push/apns.rs b/src/push/apns.rs index c624f39a..d087d3cb 100644 --- a/src/push/apns.rs +++ b/src/push/apns.rs @@ -203,6 +203,8 @@ fn build_apns_payload( 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(), diff --git a/src/push/fcm.rs b/src/push/fcm.rs index b0591b4c..1caf7be8 100644 --- a/src/push/fcm.rs +++ b/src/push/fcm.rs @@ -201,10 +201,14 @@ async fn get_fcm_access_token( let response_body = response.text().await.unwrap_or_default(); if !response_status.is_success() { - return Err(PushError::ProviderError(format!( - "FCM OAuth failed ({}): {}", - response_status, response_body - ))); + 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)?; @@ -243,6 +247,8 @@ fn build_fcm_payload( .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(), diff --git a/src/push/mod.rs b/src/push/mod.rs index 52582381..aec9896c 100644 --- a/src/push/mod.rs +++ b/src/push/mod.rs @@ -64,6 +64,8 @@ pub enum PushError { InvalidSecret(String), #[error("Provider error: {0}")] ProviderError(String), + #[error("Retryable provider error: {0}")] + ProviderRetryable(String), } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -289,6 +291,7 @@ pub async fn enqueue_notification( .map(|device| NewNotificationDelivery { event_id: event.id, push_device_id: device.id, + next_attempt_at: event.not_before_at, }) .collect(); diff --git a/src/push/worker.rs b/src/push/worker.rs index efe59842..b8b44822 100644 --- a/src/push/worker.rs +++ b/src/push/worker.rs @@ -124,6 +124,15 @@ async fn process_leased_delivery( } }; + if device.user_id != event.user_id { + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + Some("device user does not match event user"), + )?; + return Ok(()); + } + if device.revoked_at.is_some() { NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("device revoked"))?; return Ok(()); @@ -358,7 +367,8 @@ fn classify_internal_push_error(error: PushError) -> PushSendOutcome { PushError::ConnectionError | PushError::DatabaseError(_) | PushError::DbError(_) - | PushError::HttpError(_) => PushSendOutcome::Retryable { + | PushError::HttpError(_) + | PushError::ProviderRetryable(_) => PushSendOutcome::Retryable { provider_status_code: None, error: error.to_string(), }, @@ -374,3 +384,26 @@ fn retry_backoff_seconds(attempt_count: i32) -> i32 { 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/push.rs b/src/web/push.rs index 3b4d8359..640a133a 100644 --- a/src/web/push.rs +++ b/src/web/push.rs @@ -118,7 +118,13 @@ async fn register_push_device( let now = Utc::now(); let device = conn .transaction::(|conn| { - let existing = + 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, @@ -127,14 +133,27 @@ async fn register_push_device( &token_hash, )?; + let already_revoked_installation_id = if let Some(active_device) = active_installation + .as_ref() + .filter(|device| device.user_id != user.uuid) + { + PushDevice::revoke_by_id(conn, active_device.id)?; + Some(active_device.id) + } else { + None + }; + + let mut existing = existing_for_user.filter(|device| device.revoked_at.is_none()); + if let Some(conflict) = conflicting_token_owner.as_ref() { - if existing.as_ref().map(|device| device.id) != Some(conflict.id) { + if existing.as_ref().map(|device| device.id) != Some(conflict.id) + && already_revoked_installation_id != Some(conflict.id) + { PushDevice::revoke_by_id(conn, conflict.id)?; } } - if let Some(mut existing) = existing { - existing.user_id = user.uuid; + if let Some(mut existing) = existing.take() { existing.platform = body.platform.clone(); existing.provider = body.provider.clone(); existing.environment = body.environment.clone(); From bc599e0b4e50a2847c9ecb620e2913afa2a88a20 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Fri, 13 Mar 2026 22:27:19 -0500 Subject: [PATCH 25/34] fix: standardize push device registration handoff Reuse same-install device bindings across re-registration, add logout cleanup revoke behavior, and document the expected push lifecycle so account switches and reinstalls behave predictably. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- docs/encrypted-mobile-push-notifications.md | 60 ++++-- src/web/login_routes.rs | 104 +++++++++- src/web/push.rs | 207 ++++++++++++++++++-- 3 files changed, 333 insertions(+), 38 deletions(-) diff --git a/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md index d1154cf9..3f80f076 100644 --- a/docs/encrypted-mobile-push-notifications.md +++ b/docs/encrypted-mobile-push-notifications.md @@ -22,7 +22,7 @@ The final architectural decisions are: 3. **Store provider configuration per project using `project_settings` + `org_project_secrets`.** 4. **Register one notification keypair per device, not one shared key per user.** 5. **Use a durable Postgres outbox + delivery worker, not synchronous request-path sends.** -6. **Treat successful live Sage SSE delivery as the acknowledgement for request/response flows and skip push entirely when streaming succeeds.** +6. **Treat successful live Sage SSE delivery as the acknowledgement for agent chat flows and skip push entirely when streaming succeeds.** 7. **Prefer standard visible-notification behavior on Android in v1; keep Android data-only encrypted preview as an optional later enhancement.** This gives us Signal-like transport privacy for push content without pretending we are building a full Signal protocol. @@ -220,11 +220,8 @@ This section is the implementation map for a follow-on coding agent. - Merge the new push router under `validate_jwt`. - Start the background push worker after app state initialization and migrations. -- `src/web/responses/handlers.rs` - - Gate push enqueue on whether a live SSE response finished successfully; this is the v1 suppression signal for the normal Sage request/response flow. - - `src/web/agent/mod.rs` - - If mobile uses the direct agent SSE endpoints, apply the same β€œsuccessful SSE delivery => no push” rule there. + - Gate push enqueue on whether a live SSE response finished successfully; this is the v1 suppression signal for the normal Sage agent flow. - `src/web/platform/common.rs` - Add new secret constants and request / response types for project push settings. @@ -358,7 +355,7 @@ Important: ### 8.1 `push_devices` -One row per device installation. +One row per active installation binding, with revoked rows retained for re-registration history and delivery auditing. ```sql CREATE TABLE push_devices ( @@ -386,8 +383,6 @@ CREATE TABLE push_devices ( created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - UNIQUE (user_id, installation_id, environment), - UNIQUE (provider, environment, push_token_hash), CHECK ( (platform = 'ios' AND provider = 'apns') OR (platform = 'android' AND provider = 'fcm') @@ -398,13 +393,24 @@ CREATE TABLE push_devices ( CREATE INDEX idx_push_devices_user_active ON push_devices(user_id, revoked_at); + +CREATE UNIQUE INDEX idx_push_devices_installation_active + ON push_devices(installation_id, environment) + WHERE revoked_at IS NULL; + +CREATE UNIQUE INDEX idx_push_devices_token_active + ON push_devices(provider, environment, push_token_hash) + WHERE revoked_at IS NULL; ``` Implementation notes: +- `installation_id` is an app-generated UUIDv4 for one installed copy of the app; it is not an APNs or FCM identifier +- store `installation_id` in normal app-scoped local storage so uninstall / reinstall creates a fresh installation identity - `push_token_enc` should be encrypted at rest with the enclave key using the same style as project secrets - `push_token_hash` should be `SHA-256(push_token)` raw bytes for dedupe and lookup - `notification_public_key` should store raw DER-encoded SPKI bytes, not a base64 string +- keep at most one active binding per `(installation_id, environment)` and per `(provider, environment, push_token_hash)`; older rows are revoked rather than overwritten ### 8.2 `notification_events` @@ -525,10 +531,10 @@ Request payload: Behavior: -- upsert by `(user_id, installation_id, environment)` -- rotate token if it changed -- rotate public key if it changed -- clear `revoked_at` if this is a re-registration +- `installation_id` must be a non-nil UUIDv4 generated and stored by the app +- if the authenticated user already has a row for that installation, update it in place and clear `revoked_at` +- if another active row currently owns the same `installation_id` or the same push token, revoke that older binding before activating the current authenticated user +- rotate token / public key / capabilities when they change - refresh `last_seen_at` #### List current user's devices @@ -551,6 +557,31 @@ Behavior: - set `revoked_at` - optionally return the standard deleted-object response shape already used in newer APIs +#### Client lifecycle expectations + +- first install: generate a fresh UUIDv4 `installation_id`, generate the notification keypair, fetch the APNs / FCM token, then call `POST /v1/push/devices` +- app start, login, or token rotation: call `POST /v1/push/devices` again; the route is idempotent for the current authenticated user and installation +- explicit logout: while the access token is still valid, call `DELETE /v1/push/devices/:id` for the current device before dropping auth +- account switch on the same installed app: revoke the old device binding on logout, then register again after the new user logs in; if the old binding is still active, the server treats the matching installation or token as a handoff and revokes the older binding +- uninstall / reinstall: create a new `installation_id` and notification keypair, then register as a new installation on next login; stale old rows are cleaned up by explicit revoke or later invalid-token handling + +#### Logout cleanup fallback + +The existing `/logout` request may also include an optional `push_device_id` field so the backend can best-effort revoke the current push binding during auth teardown: + +```json +{ + "refresh_token": "jwt", + "push_device_id": "uuid" +} +``` + +Preferred client flow: + +1. call `DELETE /v1/push/devices/:id` while the access token is still valid +2. call `/logout` +3. include `push_device_id` in `/logout` as a cleanup fallback in case the device revoke races with local auth teardown + ### 9.2 Platform project settings API Add a new project settings endpoint in `src/web/platform/project_routes.rs`: @@ -1123,7 +1154,7 @@ The exact IPs and ports can be changed, but they should use currently unused slo ## 17. Product Event Model -### 17.1 Live Sage request/response flows +### 17.1 Live Sage agent SSE flows For the common flow where the user sends a message to Sage and the client keeps an SSE stream open waiting for responses: @@ -1164,8 +1195,7 @@ Likely first integrations: - `src/web/agent/runtime.rs` - subagent follow-up or background completion signals -- `src/web/responses/handlers.rs` - - async response completion hooks +- `src/web/agent/mod.rs` - live Sage SSE completion / failure handling Reminder scheduling can come later; the outbox design already supports `not_before_at`. diff --git a/src/web/login_routes.rs b/src/web/login_routes.rs index 7b14d2be..a624fabb 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, NewToken, TokenType}, + jwt::{validate_token, 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( @@ -240,16 +243,48 @@ pub async fn logout( debug!("Entering logout function"); 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; debug!("Exiting logout function"); 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,62 @@ pub async fn password_reset_confirm( debug!("Exiting password_reset_confirm function"); 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, + } + } + + #[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, + }; + + 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/push.rs b/src/web/push.rs index 640a133a..99396ad6 100644 --- a/src/web/push.rs +++ b/src/web/push.rs @@ -71,6 +71,18 @@ struct DeletedPushDeviceResponse { 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( @@ -133,27 +145,19 @@ async fn register_push_device( &token_hash, )?; - let already_revoked_installation_id = if let Some(active_device) = active_installation - .as_ref() - .filter(|device| device.user_id != user.uuid) - { - PushDevice::revoke_by_id(conn, active_device.id)?; - Some(active_device.id) - } else { - None - }; - - let mut existing = existing_for_user.filter(|device| device.revoked_at.is_none()); + let decision = plan_push_device_registration( + user.uuid, + existing_for_user.as_ref(), + active_installation.as_ref(), + conflicting_token_owner.as_ref(), + ); - if let Some(conflict) = conflicting_token_owner.as_ref() { - if existing.as_ref().map(|device| device.id) != Some(conflict.id) - && already_revoked_installation_id != Some(conflict.id) - { - PushDevice::revoke_by_id(conn, conflict.id)?; - } + for revoke_device_id in decision.revoke_device_ids { + PushDevice::revoke_by_id(conn, revoke_device_id)?; } - if let Some(mut existing) = existing.take() { + if decision.write_mode == RegistrationWriteMode::ReuseExisting { + let mut existing = existing_for_user.expect("existing device should be present"); existing.platform = body.platform.clone(); existing.provider = body.provider.clone(); existing.environment = body.environment.clone(); @@ -242,6 +246,10 @@ 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() @@ -284,6 +292,40 @@ fn map_push_device_error(error: PushDeviceError) -> ApiError { 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 { @@ -303,3 +345,132 @@ impl From for PushDeviceResponse { } } } + +#[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, + 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_enc: vec![1, 2, 3], + push_token_hash: token_hash, + notification_public_key: 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 94e06d6ad600836757e53bbaf7bb3f91e06af8bc Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 17 Mar 2026 23:23:43 -0500 Subject: [PATCH 26/34] fix: harden push delivery lease ownership Prevent stale workers from overwriting notification delivery results, fold the platform/provider integrity check into the base push migration, and align the push design doc with the shipped backend behavior. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- docs/encrypted-mobile-push-notifications.md | 344 ++++++++++-------- .../up.sql | 6 +- src/models/notification_deliveries.rs | 82 ++++- src/push/worker.rs | 207 +++++++++-- 4 files changed, 428 insertions(+), 211 deletions(-) diff --git a/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md index 3f80f076..2ad14839 100644 --- a/docs/encrypted-mobile-push-notifications.md +++ b/docs/encrypted-mobile-push-notifications.md @@ -1,9 +1,9 @@ # Encrypted Mobile Push Notifications for Maple / OpenSecret -## Implementation-Ready Backend, Enclave, and Mobile Design +## Backend, Enclave, and Mobile Design Reference **Date:** March 2026 -**Status:** Implementation target / pre-build spec +**Status:** Backend v1 is implemented for the Sage agent disconnect flow; broader event sources, mobile-client integrations, and some hardening remain open **Related Docs:** - `sage-in-maple-architecture.md` - `architecture-for-rag-integration.md` @@ -27,17 +27,38 @@ The final architectural decisions are: This gives us Signal-like transport privacy for push content without pretending we are building a full Signal protocol. +### 1.1 Current backend scope in this repo + +As of this revision, the backend implementation in this repo currently includes: + +- project-scoped push settings and secrets +- device register / list / revoke routes plus logout cleanup fallback +- encrypted token storage, durable notification events, and per-device delivery rows +- direct APNs + direct FCM sender implementations +- background delivery worker with Postgres leasing +- successful Sage agent SSE delivery suppression plus disconnect-triggered `agent.message` enqueue + +The currently shipped backend scope is narrower than the full product design described later in this document: + +- the implemented event source in this repo is currently the Sage agent SSE disconnect path in `src/web/agent/mod.rs` +- reminder / task-complete / account-alert sources remain follow-up work +- iOS NSE behavior, Android client behavior, recent-ID caches, and thread cleanup live in mobile codebases rather than this backend repo + --- ## 2. What Problem This Solves -We want Maple / Sage to send notifications for events like: +Longer-term Maple / Sage product goals include notifications for events like: - reminder due - long-running agent task complete - async follow-up from Sage - security or account alerts +Current backend implementation note: + +- this repo currently enqueues only `agent.message` notifications for interrupted Sage agent SSE turns + while ensuring that: - Apple and Google do not need plaintext notification content @@ -182,96 +203,77 @@ Rules: --- -## 6. Existing Codebase Integration Map +## 6. Current Codebase Integration Map -This section is the implementation map for a follow-on coding agent. +This section reflects the current backend code layout rather than the original pre-build implementation plan. -### 6.1 Existing files to modify +### 6.1 Database and model ownership -#### Database and models +- `migrations/2026-03-07-120000_push_notifications_v1/{up,down}.sql` + - Defines `push_devices`, `notification_events`, and `notification_deliveries` plus active indexes and triggers. -- `migrations//up.sql` - - Add `push_devices`, `notification_events`, and `notification_deliveries`. - - Add indexes and `updated_at` triggers matching the existing schema style. - -- `migrations//down.sql` - - Drop the new push tables and indexes. +- `src/models/push_devices.rs` + - Owns device registration lookups, same-user reuse, revoke, invalidate, and active-device listing. -- `src/models/mod.rs` - - Export the new push model modules. +- `src/models/notification_events.rs` + - Owns durable logical notification rows. -- `src/models/schema.rs` - - Regenerate via Diesel after the migration is added. +- `src/models/notification_deliveries.rs` + - Owns delivery insertion, leasing, retry scheduling, and final state transitions. - `src/models/project_settings.rs` - - Extend `SettingCategory` with `Push`. - - Add `PushSettings`, `IosPushSettings`, `AndroidPushSettings`, and related helpers. + - Owns `PushSettings`, `IosPushSettings`, and `AndroidPushSettings`. - `src/db.rs` - - Add trait methods and `PostgresConnection` implementations for push device CRUD, event enqueue, delivery leasing, retry, and invalidation. - - Mirror the existing settings / secrets patterns already used for email and OAuth. + - Currently handles project push settings read / write helpers. + - The main device / event / delivery CRUD surface lives in the model modules above, not in `src/db.rs`. -#### Web API routing +### 6.2 Web, runtime, and routing integration -- `src/web/mod.rs` - - Export a new `push` router module. - -- `src/main.rs` - - Merge the new push router under `validate_jwt`. - - Start the background push worker after app state initialization and migrations. +- `src/web/push.rs` + - App-user authenticated device register / list / revoke routes. - `src/web/agent/mod.rs` - - Gate push enqueue on whether a live SSE response finished successfully; this is the v1 suppression signal for the normal Sage agent flow. + - Sage agent SSE acknowledgement tracking plus the current disconnect-triggered `agent.message` enqueue path. - `src/web/platform/common.rs` - - Add new secret constants and request / response types for project push settings. + - Push secret constants and project-settings request / response types. - `src/web/platform/project_routes.rs` - - Add `GET` / `PUT /platform/orgs/:org_id/projects/:project_id/settings/push`. - - Reuse the existing project secret endpoints for actual APNs / FCM credentials. - -#### Deployment / runtime - -- `entrypoint.sh` - - Add `/etc/hosts` entries and `traffic_forwarder.py` processes for APNs production, APNs sandbox, and FCM. + - `GET` / `PUT /platform/orgs/:org_id/projects/:project_id/settings/push`. -- `docs/nitro-deploy.md` - - Add parent-instance `vsock-proxy` instructions for the new outbound hosts. - - This is a deployment follow-up, not required for compiling the backend. - -### 6.2 New files to add - -- `src/models/push_devices.rs` - - Diesel model and helper methods for device registrations. - -- `src/models/notification_events.rs` - - Durable logical notification rows. +- `src/main.rs` + - Raw project secret helpers, route mounting, and background push-worker startup. -- `src/models/notification_deliveries.rs` - - Per-device delivery state rows. - -- `src/web/push.rs` - - App-user authenticated routes like `POST /v1/push/devices`. +### 6.3 Push delivery modules - `src/push/mod.rs` - - Shared types, enqueue helpers, and internal interfaces. + - Shared types, enqueue helpers, generated `notification_id`, automatic `notif:` collapse keys, and current `agent.message` defaults. - `src/push/crypto.rs` - P-256 ECDH + HKDF + AES-GCM envelope logic. - `src/push/apns.rs` - - APNs request building, JWT auth, and response handling. + - APNs request building, JWT auth, cache invalidation, and response handling. - `src/push/fcm.rs` - - FCM OAuth token exchange, request building, and response handling. + - FCM OAuth token exchange, payload building, cached-token invalidation, and retry classification. - `src/push/worker.rs` - Polling / leasing loop that sends pending deliveries. -### 6.3 Codebase gotchas to preserve +### 6.4 Deployment / runtime integration + +- `entrypoint.sh` + - Adds `/etc/hosts` entries and `traffic_forwarder.py` processes for APNs production, APNs sandbox, and FCM. + +- `docs/nitro-deploy.md` + - Documents the parent-instance `vsock-proxy` instructions for those outbound hosts. + +### 6.5 Codebase gotchas to preserve -1. `src/main.rs::get_project_secret()` currently returns **base64-encoded decrypted bytes**. - That is fine for some existing flows, but APNs `.p8` material and FCM service account JSON are easier to consume as raw bytes or raw UTF-8. Add a raw helper rather than forcing each caller to decode a second time. +1. `src/main.rs` now exposes raw project-secret helpers for PEM / JSON consumption. + Keep using raw bytes / raw UTF-8 for APNs `.p8` material and FCM service account JSON rather than double-decoding base64. 2. `src/encrypt.rs` is useful for **at-rest encryption** but not for device ECDH push envelopes. Its current primitives take `secp256k1::SecretKey` and do symmetric encryption only. @@ -279,18 +281,18 @@ This section is the implementation map for a follow-on coding agent. 3. `src/main.rs::create_ephemeral_key()` uses **x25519** for session establishment. Do not overload that codepath for push notifications. -4. New user-facing push APIs should live under **`/v1/push/*`**, not under the older `/protected/*` namespace. - That matches the newer `agent` and `responses` route style. +4. User-facing push APIs live under **`/v1/push/*`**, not under the older `/protected/*` namespace. + That matches the newer route style already used elsewhere in the repo. --- ## 7. Project Configuration Model -Push config should follow the same split already used elsewhere in the repo. +Current push config follows the same split already used elsewhere in the repo. ### 7.1 Non-secret settings in `project_settings` -Add a new `SettingCategory::Push` and store a JSON payload like: +The current implementation uses `SettingCategory::Push` and stores a JSON payload like: ```rust #[derive(Debug, Clone, Serialize, Deserialize, Default)] @@ -334,7 +336,7 @@ Notes: ### 7.2 Secrets in `org_project_secrets` -Add constants in `src/web/platform/common.rs`: +The current implementation uses constants in `src/web/platform/common.rs`: ```rust pub const PROJECT_APNS_AUTH_KEY_P8: &str = "APNS_AUTH_KEY_P8"; @@ -360,20 +362,20 @@ One row per active installation binding, with revoked rows retained for re-regis ```sql CREATE TABLE push_devices ( id BIGSERIAL PRIMARY KEY, - uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + uuid UUID NOT NULL DEFAULT uuid_generate_v4() UNIQUE, user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, installation_id UUID NOT NULL, - platform TEXT NOT NULL, - provider TEXT NOT NULL, - environment TEXT NOT NULL, + platform TEXT NOT NULL CHECK (platform IN ('ios', 'android')), + provider TEXT NOT NULL CHECK (provider IN ('apns', 'fcm')), + environment TEXT NOT NULL CHECK (environment IN ('dev', 'prod')), app_id TEXT NOT NULL, push_token_enc BYTEA NOT NULL, push_token_hash BYTEA NOT NULL, notification_public_key BYTEA NOT NULL, - key_algorithm TEXT NOT NULL, + key_algorithm TEXT NOT NULL CHECK (key_algorithm IN ('p256_ecdh_v1')), supports_encrypted_preview BOOLEAN NOT NULL DEFAULT FALSE, supports_background_processing BOOLEAN NOT NULL DEFAULT FALSE, @@ -386,9 +388,7 @@ CREATE TABLE push_devices ( CHECK ( (platform = 'ios' AND provider = 'apns') OR (platform = 'android' AND provider = 'fcm') - ), - CHECK (environment IN ('dev', 'prod')), - CHECK (key_algorithm = 'p256_ecdh_v1') + ) ); CREATE INDEX idx_push_devices_user_active @@ -410,6 +410,7 @@ Implementation notes: - `push_token_enc` should be encrypted at rest with the enclave key using the same style as project secrets - `push_token_hash` should be `SHA-256(push_token)` raw bytes for dedupe and lookup - `notification_public_key` should store raw DER-encoded SPKI bytes, not a base64 string +- route-level validation and the DB constraint both enforce valid `(platform, provider)` pairs for defense in depth - keep at most one active binding per `(installation_id, environment)` and per `(provider, environment, push_token_hash)`; older rows are revoked rather than overwritten ### 8.2 `notification_events` @@ -419,7 +420,7 @@ Logical notification rows created by product code. ```sql CREATE TABLE notification_events ( id BIGSERIAL PRIMARY KEY, - uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + uuid UUID NOT NULL DEFAULT uuid_generate_v4() UNIQUE, project_id INTEGER NOT NULL REFERENCES org_projects(id) ON DELETE CASCADE, user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, @@ -441,7 +442,7 @@ CREATE TABLE notification_events ( CHECK (priority IN ('normal', 'high')) ); -CREATE INDEX idx_notification_events_due +CREATE INDEX idx_notification_events_user_due ON notification_events(user_id, not_before_at, cancelled_at); ``` @@ -504,7 +505,7 @@ Implementation notes: ### 9.1 App-user push device API -Create a new router in `src/web/push.rs` and mount it in `main.rs` with `validate_jwt` plus the usual encrypted request / response middleware. +The current router lives in `src/web/push.rs` and is mounted in `main.rs` with `validate_jwt` plus the usual encrypted request / response middleware. #### Register or rotate device @@ -584,7 +585,7 @@ Preferred client flow: ### 9.2 Platform project settings API -Add a new project settings endpoint in `src/web/platform/project_routes.rs`: +The current project settings endpoints live in `src/web/platform/project_routes.rs`: ```http GET /platform/orgs/:org_id/projects/:project_id/settings/push @@ -600,37 +601,39 @@ Do **not** add a dedicated secret endpoint. Reuse the existing generic project s --- -## 10. DB Trait and Service Methods - -Add the following categories of methods to `src/db.rs`. +## 10. Current Backend Module Split -### 10.1 Push device methods +### 10.1 `src/db.rs` and push settings -- `upsert_push_device(...)` -- `get_push_device_by_uuid_for_user(...)` -- `list_active_push_devices_for_user(...)` -- `revoke_push_device_for_user(...)` -- `invalidate_push_device(...)` - -### 10.2 Push settings methods +`src/db.rs` currently owns the project push settings helpers: - `get_project_push_settings(project_id)` - `update_project_push_settings(project_id, settings)` This mirrors the existing email / OAuth settings pattern. -### 10.3 Event and delivery methods +### 10.2 Device / event / delivery DB logic + +Most push DB behavior currently lives in model modules rather than `src/db.rs`: -- `create_notification_event(...)` -- `create_notification_deliveries_for_event(...)` -- `lease_pending_notification_deliveries(limit, lease_owner, lease_ttl)` -- `mark_notification_delivery_sent(...)` -- `mark_notification_delivery_retry(...)` -- `mark_notification_delivery_failed(...)` -- `mark_notification_delivery_invalid_token(...)` +- `src/models/push_devices.rs` + - installation lookups, active-token lookup, revoke, invalidate, and active device listing +- `src/models/notification_events.rs` + - event insertion and loading helpers +- `src/models/notification_deliveries.rs` + - bulk delivery insertion, leasing, retry scheduling, and `mark_*` transitions -The lease query should use a transaction plus `FOR UPDATE SKIP LOCKED` semantics so the design works across multiple workers / enclaves from day one. -Use Postgres time (`now()`) for lease and retry timestamps instead of local enclave clock time. +The delivery lease query uses Postgres row locking / leasing semantics, and retry scheduling uses Postgres time rather than enclave-local clock math. + +### 10.3 `src/push/mod.rs` enqueue behavior + +`src/push/mod.rs` currently owns: + +- `EnqueueNotificationRequest` +- generated canonical `notification_id` +- automatic `collapse_key = notif:` generation +- durable event creation plus per-device delivery materialization +- the current `agent.message` enqueue policy, including the 7-day expiry currently used by that path --- @@ -676,6 +679,7 @@ Rules: - no leader election is required - any enclave can resume work after another enclave crashes - `lease_owner`, `lease_expires_at`, `next_attempt_at`, and `status` must all live in Postgres +- final `mark_*` delivery transitions should only succeed when the current row still belongs to that worker's `lease_owner`; stale workers should become lost-lease no-ops instead of overwriting newer results ### 11.4 Delivery semantics @@ -824,8 +828,9 @@ Use: - `apns-topic: ` - `apns-push-type: alert` - `apns-priority: 10` -- `apns-collapse-id: notif:` when retrying the same logical notification +- `apns-collapse-id: notif:` for the current logical notification event +Retries reuse that same value because `notification_id` stays stable across retries. Do **not** reuse one collapse ID for every message in a thread by default. If desired, use `aps.thread-id` for notification-center grouping separately from collapse behavior. @@ -942,14 +947,20 @@ It avoids relying on background execution for every Sage message. }, "data": { "notification_id": "", + "kind": "agent.message", "thread_id": "agent:uuid", "message_id": "", "deep_link": "opensecret://agent/subagent/uuid" }, "android": { "priority": "HIGH", - "collapse_key": "sage:notif:", - "ttl": "300s" + "collapse_key": "notif:", + "ttl": "604800s", + "notification": { + "channel_id": "maple_messages", + "tag": "", + "click_action": "OPEN_THREAD" + } } } } @@ -959,9 +970,12 @@ Notes: - visible text should remain generic / non-sensitive - FCM `data` values must be strings +- current backend derives Android TTL from `expires_at`; when an event has no explicit expiry, the sender defaults to 7 days +- the current `agent.message` enqueue path sets a 7-day expiry, so the common Sage-agent payload today is effectively close to `604800s` - FCM `data` is provider-visible routing metadata, so it may include fields like `notification_id`, `message_id`, `thread_id`, `deep_link`, and `kind`, but never plaintext message content - this is intentional in v1: the goal is to keep message content encrypted from the push provider, not to hide all metadata needed for client routing -- reuse the same `collapse_key` only for retries of the same logical notification +- current implementation auto-keys `collapse_key` to `notification_id` without the older `sage:` prefix shown in earlier drafts +- reuse the same `collapse_key` only for the same logical notification - `notification_id` is for exact dedup; `thread_id` is for grouping / clearing when the user opens the conversation ### 14.4 Optional Android encrypted-preview path (deferred) @@ -984,10 +998,14 @@ Treat these as permanent invalid-token outcomes: Treat these as retryable: +- `401` +- `403` - `429` - `500` - `503` +For `401` / `403`, the current implementation also invalidates the cached OAuth token first so the retry can fetch a fresh credential. + Persist: - HTTP status @@ -1179,24 +1197,29 @@ pub struct EnqueueNotificationRequest { pub kind: String, pub delivery_mode: PushDeliveryMode, pub priority: PushPriority, - pub collapse_key: Option, pub fallback_title: String, pub fallback_body: String, - pub payload: Option, + pub preview_payload: Option, pub not_before_at: Option>, pub expires_at: Option>, } ``` +The current helper generates `notification_id` internally and derives `collapse_key` as `notif:` when the event row is created. + ### 17.3 Initial event sources -Likely first integrations: +Current implemented backend source: + +- `src/web/agent/mod.rs` + - live Sage agent SSE completion / failure handling and one `agent.message` notification for an interrupted turn + +Future integrations: - `src/web/agent/runtime.rs` - subagent follow-up or background completion signals -- `src/web/agent/mod.rs` - - live Sage SSE completion / failure handling +- reminder scheduling, task-complete hooks, and account / security alerts Reminder scheduling can come later; the outbox design already supports `not_before_at`. @@ -1210,19 +1233,18 @@ Reminder scheduling can come later; the outbox design already supports `not_befo --- -## 18. Validation and Test Plan +## 18. Validation Coverage and Remaining Test Work -### 18.1 Unit tests +### 18.1 Current unit coverage -Add focused tests for: +Current repo coverage includes focused tests for: - P-256 envelope roundtrip in `src/push/crypto.rs` -- APNs JWT generation and refresh behavior -- FCM service-account assertion generation and token parsing -- push-device upsert / revoke / invalidate DB behavior -- delivery leasing and retry transitions +- preview normalization and enqueue-side helpers in `src/push/mod.rs` +- retry classification and backoff behavior in `src/push/worker.rs` +- push handoff planning and logout cleanup parsing in `src/web/push.rs` and `src/web/login_routes.rs` -### 18.2 Integration tests +### 18.2 Remaining integration tests Add route tests for: @@ -1230,12 +1252,15 @@ Add route tests for: - `GET /v1/push/devices` - `DELETE /v1/push/devices/:id` - `GET` / `PUT` platform push settings routes +- `/logout` push cleanup fallback +- same-device account handoff behavior across explicit revoke + re-register flows ### 18.3 Provider client testability -Do not hardcode transport base URLs so deeply that they cannot be mocked. +APNs and FCM sender endpoints are currently hardcoded. +For stronger provider-client tests, add test-only overrideable base URLs. -The APNs and FCM sender code should allow test-only override of base URLs, making it possible to verify: +That should make it possible to verify: - request headers - request body shapes @@ -1244,52 +1269,29 @@ The APNs and FCM sender code should allow test-only override of base URLs, makin --- -## 19. Implementation Order - -### Phase 1: Schema and project settings - -- add migration for push tables -- add push model files and Diesel schema -- add `SettingCategory::Push` -- add project push settings routes -- add new project secret constants - -### Phase 2: Device registry API +## 19. Suggested Remaining Implementation Order -- add `src/web/push.rs` -- add push-device DB methods -- implement register / list / revoke routes -- add user-scoped tests +### Phase 1: Delivery hardening -### Phase 3: Outbox and generic delivery +- tighten permanent failure invalidation / quarantine behavior +- add targeted lost-lease transition tests around the current `mark_*` behavior +- add metrics for lost-lease and invalid-token outcomes -- add event / delivery DB methods -- add worker leasing loop -- wire live Sage SSE paths so successful stream completion skips push enqueue -- implement APNs generic delivery -- implement standard FCM visible delivery +### Phase 2: Provider testability and integration tests -### Phase 4: Client dedup + foreground suppression + iOS encrypted preview +- add test-only APNs / FCM base URL overrides +- add route / lease / provider-client integration coverage -- add stable `notification_id` handling end-to-end -- add client recent-ID cache guidance and integration -- add thread-based notification cleanup when a conversation opens -- add `p256` + `hkdf` -- implement per-device envelope encryption -- add iOS encrypted preview payload / NSE rewrite +### Phase 3: Product event source expansion -### Phase 5: Android optional encrypted-preview path +- add reminder, task-complete, and account / security event hooks +- decide which event types should ship in the current backend scope versus later phases -- keep standard visible FCM delivery as the default -- optionally add Android data-only encrypted preview if product later decides it is worth the extra complexity -- if enabled, implement it in `FirebaseMessagingService` with local-notification posting only after successful decrypt +### Phase 4: Mobile client alignment and tuning -### Phase 6: Product integrations and hardening - -- connect Maple / Sage event sources -- add collapse keys, TTL, and backoff policy tuning -- add metrics and invalid token cleanup dashboards -- optionally add stronger device attestation later +- align mobile repos on recent-ID dedup, thread cleanup, and foreground suppression +- tune TTL, collapse-key policy, and other delivery defaults once real client behavior is exercised +- optionally add stronger device attestation later if the product needs it --- @@ -1328,3 +1330,37 @@ The v1 implementation should be: - **stable `notification_id` plus client-side dedup / thread cleanup** This is the simplest design that still matches the current OpenSecret architecture and avoids painting us into a corner later. + +Current implementation note: the backend in this repo already follows this design for the Sage agent disconnect path, while broader event sources and some hardening items remain follow-up work. + +--- + +## 22. Remaining Implementation Details + +These items are the main remaining hardening work after the current push handoff, delivery-isolation, and doc-alignment changes. + +### 22.1 Permanent failure hygiene + +- tighten device / config permanent-failure handling so obviously dead or ineligible targets are invalidated or quarantined more aggressively +- keep the worker from repeatedly retrying rows that are no longer realistically deliverable +- continue improving provider-specific invalid-token classification, especially for APNs / FCM permanent device failures + +### 22.2 Integration-level test coverage + +- add route-level tests for push registration, revoke, logout cleanup, and same-device account handoff +- add leasing tests that exercise stale-worker and lost-lease behavior against the actual DB transition code +- add provider-client tests with injectable APNs / FCM base URLs so delivery classification can be verified without real provider traffic + +### 22.3 Migration cleanup and consistency + +- clean up migration / down-migration consistency around index definitions where helpful + +### 22.4 Delivery policy flexibility + +- decide whether future callers need explicit internal enqueue control over `collapse_key`, TTL, or similar delivery-policy knobs +- document any chosen defaults centrally so backend and client expectations stay aligned + +### 22.5 Scope expansion and client alignment + +- extend backend event sources beyond the current `agent.message` Sage disconnect path when product priorities require it +- keep the mobile-client behavior docs aligned with what is actually implemented in the iOS / Android repos diff --git a/migrations/2026-03-07-120000_push_notifications_v1/up.sql b/migrations/2026-03-07-120000_push_notifications_v1/up.sql index 766cfd16..61d1a7ef 100644 --- a/migrations/2026-03-07-120000_push_notifications_v1/up.sql +++ b/migrations/2026-03-07-120000_push_notifications_v1/up.sql @@ -16,7 +16,11 @@ CREATE TABLE push_devices ( last_seen_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, 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 + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + CHECK ( + (platform = 'ios' AND provider = 'apns') OR + (platform = 'android' AND provider = 'fcm') + ) ); CREATE TABLE notification_events ( diff --git a/src/models/notification_deliveries.rs b/src/models/notification_deliveries.rs index 572d7225..47ce7838 100644 --- a/src/models/notification_deliveries.rs +++ b/src/models/notification_deliveries.rs @@ -11,6 +11,26 @@ 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}")] @@ -91,11 +111,17 @@ impl NotificationDelivery { pub fn mark_sent( conn: &mut PgConnection, lookup_id: i64, + expected_lease_owner: &str, provider_message_id: Option<&str>, provider_status_code: Option, - ) -> Result<(), NotificationDeliveryError> { + ) -> Result { diesel::update( - notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + 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), @@ -109,48 +135,58 @@ impl NotificationDelivery { notification_deliveries::updated_at.eq(diesel::dsl::now), )) .execute(conn) - .map(|_| ()) + .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<(), NotificationDeliveryError> { + ) -> Result { let query = r#" UPDATE notification_deliveries SET status = 'retry', attempt_count = attempt_count + 1, - provider_status_code = $2, - last_error = $3, - next_attempt_at = NOW() + ($4 * INTERVAL '1 second'), + 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(|_| ()) + .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<(), NotificationDeliveryError> { + ) -> Result { diesel::update( - notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + 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), @@ -162,18 +198,24 @@ impl NotificationDelivery { notification_deliveries::updated_at.eq(diesel::dsl::now), )) .execute(conn) - .map(|_| ()) + .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<(), NotificationDeliveryError> { + ) -> Result { diesel::update( - notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + 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), @@ -186,17 +228,23 @@ impl NotificationDelivery { notification_deliveries::updated_at.eq(diesel::dsl::now), )) .execute(conn) - .map(|_| ()) + .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<(), NotificationDeliveryError> { + ) -> Result { diesel::update( - notification_deliveries::table.filter(notification_deliveries::id.eq(lookup_id)), + 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), @@ -206,7 +254,7 @@ impl NotificationDelivery { notification_deliveries::updated_at.eq(diesel::dsl::now), )) .execute(conn) - .map(|_| ()) + .map(NotificationDeliveryWriteResult::from_updated_rows) .map_err(NotificationDeliveryError::DatabaseError) } } diff --git a/src/push/worker.rs b/src/push/worker.rs index b8b44822..e354e4a1 100644 --- a/src/push/worker.rs +++ b/src/push/worker.rs @@ -1,5 +1,7 @@ use crate::encrypt::decrypt_with_key; -use crate::models::notification_deliveries::NotificationDelivery; +use crate::models::notification_deliveries::{ + NotificationDelivery, NotificationDeliveryWriteResult, +}; use crate::models::notification_events::NotificationEvent; use crate::models::push_devices::{PushDevice, PUSH_PLATFORM_ANDROID, PUSH_PLATFORM_IOS}; use crate::push::apns::{send_apns_notification, ApnsSendRequest}; @@ -85,6 +87,14 @@ async fn process_leased_delivery( 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() @@ -93,13 +103,38 @@ async fn process_leased_delivery( let event = match NotificationEvent::get_by_id(&mut conn, delivery.event_id)? { Some(event) => event, None => { - NotificationDelivery::mark_failed(&mut conn, delivery.id, None, Some("event missing"))?; + 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() { - NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("event cancelled"))?; + 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(()); } @@ -107,34 +142,74 @@ async fn process_leased_delivery( .expires_at .is_some_and(|expires_at| expires_at <= Utc::now()) { - NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("event expired"))?; + 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 => { - NotificationDelivery::mark_failed( - &mut conn, + if !record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + None, + Some("device missing"), + )?, delivery.id, - None, - Some("device missing"), - )?; + &lease_owner, + "failed (device missing)", + ) { + return Ok(()); + } return Ok(()); } }; if device.user_id != event.user_id { - NotificationDelivery::mark_cancelled( - &mut conn, + if !record_delivery_transition( + NotificationDelivery::mark_cancelled( + &mut conn, + delivery.id, + &lease_owner, + Some("device user does not match event user"), + )?, delivery.id, - Some("device user does not match event user"), - )?; + &lease_owner, + "cancelled (device user mismatch)", + ) { + return Ok(()); + } return Ok(()); } if device.revoked_at.is_some() { - NotificationDelivery::mark_cancelled(&mut conn, delivery.id, Some("device revoked"))?; + 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); @@ -160,59 +235,96 @@ async fn process_leased_delivery( provider_message_id, provider_status_code, } => { - NotificationDelivery::mark_sent( - &mut conn, + record_delivery_transition( + NotificationDelivery::mark_sent( + &mut conn, + delivery.id, + &lease_owner, + provider_message_id.as_deref(), + provider_status_code, + )?, delivery.id, - provider_message_id.as_deref(), - provider_status_code, - )?; + &lease_owner, + "sent", + ); } PushSendOutcome::Retryable { provider_status_code, error, } => { if delivery.attempt_count + 1 >= PUSH_WORKER_MAX_ATTEMPTS { - NotificationDelivery::mark_failed( - &mut conn, + record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + )?, delivery.id, - provider_status_code, - Some(&error), - )?; + &lease_owner, + "failed (retry limit reached)", + ); } else { - NotificationDelivery::mark_retry( - &mut conn, + 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, - provider_status_code, - Some(&error), - retry_backoff_seconds(delivery.attempt_count + 1), - )?; + &lease_owner, + "retry", + ); } } PushSendOutcome::InvalidToken { provider_status_code, error, } => { - conn.transaction::<(), PushError, _>(|conn| { - NotificationDelivery::mark_invalid_token( + let invalidated = conn.transaction::(|conn| { + let write_result = NotificationDelivery::mark_invalid_token( conn, delivery.id, + &lease_owner, provider_status_code, Some(&error), )?; - PushDevice::invalidate(conn, device.id)?; - Ok(()) + + 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, } => { - NotificationDelivery::mark_failed( - &mut conn, + record_delivery_transition( + NotificationDelivery::mark_failed( + &mut conn, + delivery.id, + &lease_owner, + provider_status_code, + Some(&error), + )?, delivery.id, - provider_status_code, - Some(&error), - )?; + &lease_owner, + "failed", + ); } } @@ -379,6 +491,23 @@ fn classify_internal_push_error(error: PushError) -> PushSendOutcome { } } +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)); From 269ae4ac3734ae15db6f4d4c91865419da6eafe1 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Thu, 19 Mar 2026 23:04:41 -0500 Subject: [PATCH 27/34] fix(push): force APNs onto HTTP/2 in Nitro APNs requires ALPN-negotiated HTTP/2, so the Nitro push path now uses an HTTP/2-only client and exposes the same startup TCP checks as the other outbound push endpoints. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.toml | 2 +- entrypoint.sh | 22 ++++++++++++++++++++++ src/push/apns.rs | 2 +- src/push/mod.rs | 6 ++++++ 4 files changed, 30 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index df49f79e..b239edef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,7 +44,7 @@ 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" diff --git a/entrypoint.sh b/entrypoint.sh index 9db8f993..fa969369 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -626,6 +626,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 fcm_tokens: Arc>>, } @@ -212,9 +213,14 @@ impl PushTransport { 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())), }) From c4220f5367808b3c59b9f7366450c5fd5caf0939 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Sat, 21 Mar 2026 16:08:12 -0500 Subject: [PATCH 28/34] docs(push): sync v1 status with shipped clients Record the current iOS, macOS, and Android push behavior so the backend design doc reflects the shipped disconnect-triggered v1 scope and the remaining hardening gaps. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- docs/encrypted-mobile-push-notifications.md | 45 ++++++++++++++++----- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md index 2ad14839..2addcd1b 100644 --- a/docs/encrypted-mobile-push-notifications.md +++ b/docs/encrypted-mobile-push-notifications.md @@ -3,7 +3,7 @@ ## Backend, Enclave, and Mobile Design Reference **Date:** March 2026 -**Status:** Backend v1 is implemented for the Sage agent disconnect flow; broader event sources, mobile-client integrations, and some hardening remain open +**Status:** Backend v1 is implemented for the Sage agent disconnect flow; Maple iOS encrypted preview and Android release delivery are implemented, macOS delivery works with generic fallback while encrypted-preview parity remains follow-up work, and broader event sources plus some client hardening remain open **Related Docs:** - `sage-in-maple-architecture.md` - `architecture-for-rag-integration.md` @@ -44,6 +44,13 @@ The currently shipped backend scope is narrower than the full product design des - reminder / task-complete / account-alert sources remain follow-up work - iOS NSE behavior, Android client behavior, recent-ID caches, and thread cleanup live in mobile codebases rather than this backend repo +### 1.2 Current Maple client status + +- **iOS release / TestFlight**: APNs registration, shared-keychain private-key access, notification toggle / revoke handling, and Notification Service Extension decryption are implemented. +- **macOS release / TestFlight**: APNs registration and delivery are implemented with a macOS notification service target, but the currently observed user-facing behavior still falls back to the generic notification text, so encrypted-preview parity remains follow-up work. +- **Android release**: FCM HTTP v1 registration / revoke flows, permission handling, token rotation, foreground refresh behavior, and generic visible notifications are implemented for the real release package `ai.trymaple.assistant`. +- **Android debug (`.dev`)**: not the intended end-to-end push validation lane. + --- ## 2. What Problem This Solves @@ -1022,9 +1029,9 @@ Persist: 1. request notification permission 2. receive APNs device token -3. generate `P256.KeyAgreement.PrivateKey()` locally -4. store private key in keychain / Secure Enclave-backed flow when available -5. expose the public key in DER / SPKI form +3. generate a dedicated P-256 notification keypair locally +4. store the private key in shared keychain storage so both the app and Notification Service Extension can read it +5. derive the public key in DER / SPKI form from that stored private key 6. register with `POST /v1/push/devices` #### Storage requirements @@ -1052,6 +1059,12 @@ The Notification Service Extension should: If the key is unavailable or decryption fails, do nothing and let iOS show the generic fallback alert. +Current Maple status: + +- implemented in the release / TestFlight lane +- shared key access is already wired for the main app and NSE targets +- notification toggle / logout cleanup is implemented + #### Foreground / open-thread behavior If the app is already foregrounded and showing the target Sage conversation or active response screen: @@ -1068,16 +1081,22 @@ The NSE should not be treated as the primary β€œsuppress if user is already look 1. request notification permission on Android 13+ 2. obtain FCM registration token -3. generate an EC keypair in Android Keystore -4. export the public key as SPKI bytes +3. generate a dedicated P-256 notification keypair locally +4. export / derive the public key as SPKI bytes 5. register with `POST /v1/push/devices` +Current Maple status: + +- the current Android client uses Rust helper functions to generate the keypair and stores the PKCS#8 private key in `EncryptedSharedPreferences` +- moving that private key into Android Keystore remains future hardening, but it does not change the backend contract +- end-to-end push validation is currently intended for the release package `ai.trymaple.assistant` + #### Runtime behavior - use `FirebaseMessagingService.onNewToken()` to rotate tokens - for v1, assume the default incoming path is a standard visible FCM notification plus `data` -- if the app is already foregrounded and showing the target Sage conversation, update UI and do not show an extra local notification -- when the user opens a conversation, clear notifications associated with that `thread_id` +- if the app is already foregrounded, refresh the in-app chat state and do not show an extra local notification +- current notification taps route back into the main Maple chat surface and trigger a refresh; thread-specific notification clearing remains follow-up work - if we later enable Android data-only encrypted preview, handle it in `onMessageReceived()` and only post a local notification after successful decrypt - use `WorkManager` for longer processing if needed @@ -1102,6 +1121,8 @@ Rule: This is the client-side protection against at-least-once retries or server bugs. +Current Maple status: this cache is still recommended but is not yet implemented in the mobile clients, so provider collapse keys and server-side SSE suppression remain the main duplicate controls for the current v1 scope. + #### Thread grouping and cleanup Each payload should include a `thread_id` when applicable. @@ -1112,6 +1133,8 @@ Use it to: - clear thread notifications when the user opens that thread - avoid leaving stale notifications around once the user has seen the conversation +Current Maple status: thread-specific cleanup is not yet implemented consistently across Maple clients. That is follow-up hardening work rather than a blocker for the current disconnect-triggered v1 scope. + #### Foreground suppression is defense-in-depth If a visible notification arrives while the app is already open, suppress or no-op locally where the platform allows it. @@ -1360,7 +1383,9 @@ These items are the main remaining hardening work after the current push handoff - decide whether future callers need explicit internal enqueue control over `collapse_key`, TTL, or similar delivery-policy knobs - document any chosen defaults centrally so backend and client expectations stay aligned -### 22.5 Scope expansion and client alignment +### 22.5 Client parity and scope expansion +- finish macOS encrypted-preview parity so the macOS notification service extension rewrites the generic fallback text like iOS +- add client-side `notification_id` dedup caches and thread cleanup where they materially improve UX - extend backend event sources beyond the current `agent.message` Sage disconnect path when product priorities require it -- keep the mobile-client behavior docs aligned with what is actually implemented in the iOS / Android repos +- keep the mobile-client behavior docs aligned with what is actually implemented in the Maple repos From 64ab90f7bd6805c35b63d31aab9b355555ef7760 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Sun, 22 Mar 2026 00:59:40 -0500 Subject: [PATCH 29/34] feat(agent): add explicit main-agent init flow Seed onboarding messages during init, return them directly to clients, and persist user locale/timezone so Maple's first-run experience is durable and predictable. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 35 ++ Cargo.toml | 1 + internal_docs/sage-in-maple-architecture.md | 164 ++++++- .../down.sql | 3 + .../up.sql | 16 + src/models/mod.rs | 1 + src/models/schema.rs | 12 + src/models/user_preferences.rs | 115 +++++ src/web/agent/mod.rs | 101 ++++- src/web/agent/runtime.rs | 409 ++++++++++++++---- src/web/agent/tools.rs | 123 +++--- 11 files changed, 827 insertions(+), 153 deletions(-) create mode 100644 migrations/2026-03-22-033535_add_user_preferences/down.sql create mode 100644 migrations/2026-03-22-033535_add_user_preferences/up.sql create mode 100644 src/models/user_preferences.rs diff --git a/Cargo.lock b/Cargo.lock index 3289b421..861b396c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1423,6 +1423,16 @@ dependencies = [ "windows-link 0.2.1", ] +[[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" @@ -4129,6 +4139,7 @@ dependencies = [ "cbc", "chacha20poly1305", "chrono", + "chrono-tz", "diesel", "diesel-derive-enum", "dotenv", @@ -4394,6 +4405,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.5" @@ -5450,6 +5479,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.9" diff --git a/Cargo.toml b/Cargo.toml index b239edef..f373d730 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ diesel = { version = "=2.2.2", 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/internal_docs/sage-in-maple-architecture.md b/internal_docs/sage-in-maple-architecture.md index dadda1da..523fa81f 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/sage-in-maple-architecture.md @@ -23,6 +23,7 @@ The target product shape is: - **unlimited subagents** as topic-specific chats/workspaces - **shared memory** across the main agent and all subagents - **separate conversation history** per agent thread +- **lightweight user-global preferences** for durable context like timezone/locale - **no end-user configuration burden** for prompts, models, memory block definitions, or tuning knobs The end state is not β€œusers manually configure an agent.” The end state is β€œusers talk to Maple, Maple manages the agent system for them.” @@ -115,6 +116,18 @@ That means the following should be owned by code / rollout config, not by per-us The database should persist **durable state**, not **tuning policy**. +Valid durable state still includes product-owned per-user context such as: +- agent identity +- shared memory +- user-global preferences like `timezone` and `locale` +- persisted onboarding messages in the main agent's conversation history + +What should **not** live in per-user rows: +- model selection +- prompt editing +- compaction thresholds +- context window tuning + --- ## 2. Core Architectural Decisions @@ -183,6 +196,25 @@ So for the rewrite: - simplify `memory_blocks` now instead of carrying legacy fields forward - move code-owned defaults out of SQL and into Rust/config +### 2.6 Make Main-Agent Initialization Explicit + +The backend should **not** lazily create the main agent as a side effect of read or chat routes. + +The target product flow should be: +- client tries to load the main agent surface (`GET /v1/agent` and/or initial main-agent history load) +- server returns the main agent if it already exists, otherwise `404` +- client calls `POST /v1/agent/init` when the main agent does not exist + +`POST /v1/agent/init` should be the single creation path for the user's main agent. It should: +- create the main conversation row +- create the main `agents` row +- persist any init-time user-global preferences such as `timezone` and `locale` +- seed the first assistant onboarding messages into conversation history + +The init response should return the main-agent metadata **plus** the seeded onboarding `messages` so clients can render and stagger them immediately without an extra history fetch. + +Those initial onboarding messages should be **server-authored and persisted as real assistant messages**. They should not be frontend-only placeholders, and they should not be model-generated by default. + --- ## 3. Why Reuse the Message Tables @@ -302,7 +334,45 @@ FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); Any limits, formatting rules, or block semantics should live in code. If we later add more built-in blocks, we can do that by code first without needing a new table shape. -### 4.3 `user_embeddings` -- Shared Recall + Archival Store +### 4.3 `user_preferences` -- User-Global Durable Preferences + +Some durable state is legitimately user-global and should not be inferred from the client every time. + +This is distinct from prompt/model/tuning config. It is product context shared across the user's main agent and all subagents. + +```sql +CREATE TABLE user_preferences ( + id BIGSERIAL PRIMARY KEY, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + -- Code-owned preference key, e.g. 'timezone', 'locale' + key TEXT NOT NULL, + + -- Encrypted because values may contain user-sensitive context + value_enc BYTEA NOT NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(user_id, key) +); + +CREATE INDEX idx_user_preferences_user_id ON user_preferences(user_id); + +CREATE TRIGGER update_user_preferences_updated_at +BEFORE UPDATE ON user_preferences +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**V1 choices:** +- Shared per user, **not** per agent +- Plaintext `key`, encrypted `value_enc` +- Preference keys are code-owned and validated in Rust; they do **not** need SQL enum enforcement in v1 +- Initial keys should stay small and product-owned: `timezone`, `locale` +- This table is for durable user-global context only, **not** prompt/model/tuning config +- Onboarding state does **not** need to live here in v1; onboarding can be derived from user-message counts at runtime + +### 4.4 `user_embeddings` -- Shared Recall + Archival Store This remains the shared embedding layer for the user. @@ -321,7 +391,7 @@ If we ever decide to support isolated subagent memory later, we can add explicit See `potential-rag-integration-brute-force.md` for the full schema and search behavior. -### 4.4 `conversation_summaries` -- Per-Thread Compaction Artifacts +### 4.5 `conversation_summaries` -- Per-Thread Compaction Artifacts This table stays conceptually the same. @@ -355,7 +425,7 @@ CREATE INDEX idx_conversation_summaries_chain **Key point:** summaries are per conversation thread, which naturally means per agent. No `agent_id` is needed because `conversation_id` already maps back to `agents`. -### 4.5 Remove `agent_config` +### 4.6 Remove `agent_config` `agent_config` should not exist in the target architecture. @@ -383,8 +453,9 @@ The active agent runtime should load: 1. the `agents` row for the active agent 2. code-owned prompt template(s) for that agent kind 3. shared memory blocks for the user -4. the latest summary for the active conversation -5. recent messages for the active conversation +4. user-global preferences for the user (at least `timezone` and `locale`) +5. the latest summary for the active conversation +6. recent messages for the active conversation **Prompt/model selection should be code-owned.** @@ -401,6 +472,10 @@ Instead, code chooses: - context window / packing rules - compaction thresholds +User-global preferences still matter here: +- `timezone` should localize `current_time` and rendered conversation timestamps +- `locale` can inform prompt hints and future language behavior + ### 5.2 Main Agent vs Subagent Context The main agent and subagents share memory, but they do **not** share recent transcript. @@ -412,6 +487,7 @@ The context builder should behave like this: - start with code-owned instruction for `kind='main'` or `kind='subagent'` - if subagent, inject decrypted `purpose_enc` - inject shared `persona` and `human` blocks + - inject any relevant user-global preferences/context derived from `user_preferences` - inject available tools - inject memory metadata 3. **Load latest summary** for the active `conversation_id` @@ -421,7 +497,42 @@ The context builder should behave like this: A subagent should never silently get the main agent's recent transcript. If it needs information from another thread, it should use recall tools intentionally. -### 5.3 DSRs Signature Shape +### 5.3 Main-Agent Onboarding + +The main agent should feel warm and present from the first screen. A brand-new user should not land in a blank conversation waiting to type first. + +During `POST /v1/agent/init`, the backend should seed the first three assistant messages into the main-agent conversation history. + +That same init response should include those seeded messages directly so clients can replay them with normal typing / stagger behavior without making an immediate follow-up history call. + +Those first messages should be: +- server-authored +- persisted as real `assistant_messages` +- visible in normal conversation history APIs +- not frontend-hardcoded placeholders +- not model-generated by default + +Suggested default English copy for v1: + +1. `Hey, I'm Maple. πŸ‘‹` +2. `Nice to meet you.` +3. `What should I call you?` + +After init, the main agent should continue using a stronger onboarding prompt overlay for approximately the first **10-15 user messages** in the main-agent thread. + +That onboarding window should be derived at prompt-build time from the count of **user messages** in the main-agent conversation. No separate persisted onboarding-state column is required in v1. + +During this onboarding window, Maple should: +- be especially warm, inviting, and emotionally open +- get to know the user gradually rather than interrogating them +- ask at most one thoughtful follow-up at a time +- proactively save key facts, routines, preferences, goals, and relationships to memory +- avoid transactional assistant language +- prioritize relationship-building when the user is simply arriving and getting settled + +Only actual user turns should count toward this window. Seeded assistant messages, tool calls, and tool outputs should not. + +### 5.4 DSRs Signature Shape The DSRs pattern still fits well. Conceptually the active agent call becomes: @@ -429,6 +540,7 @@ The DSRs pattern still fits well. Conceptually the active agent call becomes: AgentResponse (inputs) - input - current_time + - user_locale - agent_kind - subagent_purpose - persona_block @@ -437,14 +549,16 @@ AgentResponse (inputs) - previous_context_summary - recent_conversation - available_tools - - is_first_time_user + - main_agent_user_message_count AgentResponse (outputs) - messages: string[] - tool_calls: {name: string, args: map}[] ``` -### 5.4 Compaction vs Truncation +`main_agent_user_message_count` lets code-owned prompt logic keep onboarding active for the first ~10-15 main-agent user turns without persisting separate onboarding state. + +### 5.5 Compaction vs Truncation This also stays the same conceptually: - Responses API truncates @@ -452,7 +566,7 @@ This also stays the same conceptually: Each agent thread owns its own summary chain. Main agent compaction and subagent compaction are independent because they are keyed by different `conversation_id`s. -### 5.5 Vision / Image Handling +### 5.6 Vision / Image Handling No architectural change is needed here. @@ -477,6 +591,7 @@ These remain the core tools for the shared-memory model: | `archival_insert` | Store long-term memory | `user_embeddings` (`source_type='archival'`) | | `archival_search` | Search long-term memory | `user_embeddings` (`source_type='archival'`) | | `conversation_search` | Search embedded message history | `user_embeddings` (`source_type='message'`) | +| `set_preference` | Set validated user-global preferences such as timezone or locale | `user_preferences` | | `spawn_subagent` | Create a topic-specific subagent chat | `conversations` + `agents` | | `done` | Stop signal after tool-result continuation | no-op | @@ -512,6 +627,16 @@ A `spawn_subagent` tool should: This is how the main agent can proactively create a focused workspace for the user. +### 6.4 `set_preference` + +`set_preference` should exist for agent-managed user-global preferences. + +V1 known keys should be small and validated in code, for example: +- `timezone` -- IANA timezone like `America/Chicago` +- `locale` -- locale/language hint like `en` or `en-US` + +This tool is **not** a general-purpose config surface for prompt/model tuning. It is a narrow way for Maple to persist durable user-global context when the user explicitly shares it or when the client provides it during init. + --- ## 7. API Surface @@ -520,8 +645,13 @@ This is how the main agent can proactively create a focused workspace for the us The public product surface should focus on the main agent and subagent lifecycle, not user configuration. +The main agent should be created explicitly. Read/chat routes must not lazily initialize it as a side effect. + ```text -POST /v1/agent/chat -- chat with the main agent (request-scoped SSE) +GET /v1/agent -- load the main agent or return 404 if not initialized +POST /v1/agent/init -- initialize the main agent, store init-time preferences, seed onboarding messages +GET /v1/agent/items -- read the main agent transcript/history +POST /v1/agent/chat -- chat with the main agent (request-scoped SSE; must not implicitly init) POST /v1/agent/subagents -- create a new subagent POST /v1/agent/subagents/:id/chat -- chat with a subagent (request-scoped SSE) DELETE /v1/agent/subagents/:id -- delete a subagent @@ -576,8 +706,16 @@ The difference is only which agent identity / conversation thread is being drive /v1/agent/chat Main agent home surface Uses agents + shared memory + thread-local summaries + Requires prior explicit init Main thread is hidden from /v1/conversations +/v1/agent/init + Explicit main-agent bootstrap + Creates/repairs the main agent + conversation + Stores user-global preferences provided at init + Seeds the first assistant onboarding messages into history + Returns seeded onboarding messages directly in the init response + /v1/agent/subagents/* Subagent lifecycle + chat surface Uses agents + shared memory + thread-local summaries @@ -624,14 +762,20 @@ Because this is still local-only, we should update the existing Sage docs and sc ### Phase 2: Schema Reset (edit local-only migrations in place) - Replace `agent_config` with `agents` - Simplify `memory_blocks` +- Add `user_preferences` shared per user - Keep `conversation_summaries` - Keep `user_embeddings` shared per user - Do **not** add per-agent memory scoping columns ### Phase 3: Main Agent Runtime Rewrite +- Add explicit `POST /v1/agent/init` - Load the main agent from `agents` +- Remove lazy main-agent creation from read/chat paths - Move model/prompt/tuning defaults into code - Remove user-facing config assumptions +- Localize `current_time` and transcript timestamps from `user_preferences.timezone` +- Seed the first three onboarding messages during init +- Keep richer onboarding guidance active for the first ~10-15 main-agent user turns - Keep regenerated context + compaction ### Phase 4: Subagent Lifecycle diff --git a/migrations/2026-03-22-033535_add_user_preferences/down.sql b/migrations/2026-03-22-033535_add_user_preferences/down.sql new file mode 100644 index 00000000..b476af4e --- /dev/null +++ b/migrations/2026-03-22-033535_add_user_preferences/down.sql @@ -0,0 +1,3 @@ +DROP TRIGGER IF EXISTS update_user_preferences_updated_at ON user_preferences; +DROP INDEX IF EXISTS idx_user_preferences_user_id; +DROP TABLE IF EXISTS user_preferences; diff --git a/migrations/2026-03-22-033535_add_user_preferences/up.sql b/migrations/2026-03-22-033535_add_user_preferences/up.sql new file mode 100644 index 00000000..bf39303f --- /dev/null +++ b/migrations/2026-03-22-033535_add_user_preferences/up.sql @@ -0,0 +1,16 @@ +CREATE TABLE user_preferences ( + id BIGSERIAL PRIMARY KEY, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + key TEXT NOT NULL, + value_enc BYTEA NOT NULL, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE (user_id, key) +); + +CREATE INDEX idx_user_preferences_user_id ON user_preferences(user_id); + +CREATE TRIGGER update_user_preferences_updated_at +BEFORE UPDATE ON user_preferences +FOR EACH ROW +EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/models/mod.rs b/src/models/mod.rs index b85ba507..2adbf1b4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -25,4 +25,5 @@ pub mod token_usage; pub mod user_api_keys; pub mod user_embeddings; pub mod user_kv; +pub mod user_preferences; pub mod users; diff --git a/src/models/schema.rs b/src/models/schema.rs index 06b2b37d..f492eba1 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -491,6 +491,17 @@ diesel::table! { } } +diesel::table! { + user_preferences (id) { + id -> Int8, + user_id -> Uuid, + key -> Text, + value_enc -> Bytea, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { users (id) { id -> Int4, @@ -569,5 +580,6 @@ diesel::allow_tables_to_appear_in_same_query!( user_kv, user_messages, user_oauth_connections, + user_preferences, 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 2ed4b356..59b4b688 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -48,6 +48,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, @@ -78,6 +84,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, @@ -236,6 +249,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, @@ -293,7 +325,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, @@ -438,6 +470,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) @@ -506,6 +545,47 @@ 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(body): Extension, +) -> Result>, ApiError> { + let user_key = state + .get_user_key(user.uuid, 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, @@ -684,6 +764,16 @@ async fn chat_main( Extension(user): 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, user, body, ChatTarget::Main).await } @@ -834,10 +924,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, ) @@ -1055,7 +1150,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 5b2904bd..eb920260 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-5"; 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,229 @@ 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, + 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, + } + .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, -) -> Result<(Agent, Conversation), ApiError> { + 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 @@ -231,9 +483,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( @@ -308,7 +567,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 } @@ -328,8 +587,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:?}"); @@ -441,55 +698,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; @@ -499,6 +707,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 { @@ -710,6 +947,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, @@ -726,7 +965,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, @@ -821,12 +1060,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 @@ -835,6 +1073,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, @@ -847,7 +1097,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(); @@ -904,6 +1154,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, @@ -935,12 +1200,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 { @@ -955,7 +1214,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 cc7e512c24ac87fc93563c5a40676147f10d7905 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 23 Mar 2026 17:54:05 -0500 Subject: [PATCH 30/34] feat(agent): add scheduled wakeups Let Maple schedule future agent turns with recurring rules, timezone-aware execution, and normal agent.message push delivery. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- docs/encrypted-mobile-push-notifications.md | 48 +- ....md => maple-agent-memory-architecture.md} | 22 +- .../potential-rag-integration-brute-force.md | 647 ++++++ ...notifications-implementation-scratchpad.md | 110 + internal_docs/scheduled-agent-wakeups.md | 1725 ++++++++++++++++ .../down.sql | 2 + .../up.sql | 89 + src/main.rs | 2 + src/models/agent_schedule_runs.rs | 364 ++++ src/models/agent_schedules.rs | 342 ++++ src/models/agents.rs | 8 + src/models/mod.rs | 2 + src/models/schema.rs | 57 + src/web/agent/mod.rs | 3 + src/web/agent/runtime.rs | 28 +- src/web/agent/schedules.rs | 1769 +++++++++++++++++ 16 files changed, 5182 insertions(+), 36 deletions(-) rename internal_docs/{sage-in-maple-architecture.md => maple-agent-memory-architecture.md} (96%) create mode 100644 internal_docs/potential-rag-integration-brute-force.md create mode 100644 internal_docs/push-notifications-implementation-scratchpad.md create mode 100644 internal_docs/scheduled-agent-wakeups.md create mode 100644 migrations/2026-03-23-180000_agent_schedules_v1/down.sql create mode 100644 migrations/2026-03-23-180000_agent_schedules_v1/up.sql 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/docs/encrypted-mobile-push-notifications.md b/docs/encrypted-mobile-push-notifications.md index 2addcd1b..b4b97110 100644 --- a/docs/encrypted-mobile-push-notifications.md +++ b/docs/encrypted-mobile-push-notifications.md @@ -3,9 +3,9 @@ ## Backend, Enclave, and Mobile Design Reference **Date:** March 2026 -**Status:** Backend v1 is implemented for the Sage agent disconnect flow; Maple iOS encrypted preview and Android release delivery are implemented, macOS delivery works with generic fallback while encrypted-preview parity remains follow-up work, and broader event sources plus some client hardening remain open +**Status:** Backend v1 is implemented for the Maple agent disconnect flow; Maple iOS encrypted preview and Android release delivery are implemented, macOS delivery works with generic fallback while encrypted-preview parity remains follow-up work, and broader event sources plus some client hardening remain open **Related Docs:** -- `sage-in-maple-architecture.md` +- `maple-agent-memory-architecture.md` - `architecture-for-rag-integration.md` - `potential-rag-integration-brute-force.md` @@ -22,7 +22,7 @@ The final architectural decisions are: 3. **Store provider configuration per project using `project_settings` + `org_project_secrets`.** 4. **Register one notification keypair per device, not one shared key per user.** 5. **Use a durable Postgres outbox + delivery worker, not synchronous request-path sends.** -6. **Treat successful live Sage SSE delivery as the acknowledgement for agent chat flows and skip push entirely when streaming succeeds.** +6. **Treat successful live Maple SSE delivery as the acknowledgement for agent chat flows and skip push entirely when streaming succeeds.** 7. **Prefer standard visible-notification behavior on Android in v1; keep Android data-only encrypted preview as an optional later enhancement.** This gives us Signal-like transport privacy for push content without pretending we are building a full Signal protocol. @@ -36,11 +36,11 @@ As of this revision, the backend implementation in this repo currently includes: - encrypted token storage, durable notification events, and per-device delivery rows - direct APNs + direct FCM sender implementations - background delivery worker with Postgres leasing -- successful Sage agent SSE delivery suppression plus disconnect-triggered `agent.message` enqueue +- successful Maple agent SSE delivery suppression plus disconnect-triggered `agent.message` enqueue The currently shipped backend scope is narrower than the full product design described later in this document: -- the implemented event source in this repo is currently the Sage agent SSE disconnect path in `src/web/agent/mod.rs` +- the implemented event source in this repo is currently the Maple agent SSE disconnect path in `src/web/agent/mod.rs` - reminder / task-complete / account-alert sources remain follow-up work - iOS NSE behavior, Android client behavior, recent-ID caches, and thread cleanup live in mobile codebases rather than this backend repo @@ -55,16 +55,16 @@ The currently shipped backend scope is narrower than the full product design des ## 2. What Problem This Solves -Longer-term Maple / Sage product goals include notifications for events like: +Longer-term Maple product goals include notifications for events like: - reminder due - long-running agent task complete -- async follow-up from Sage +- async follow-up from Maple - security or account alerts Current backend implementation note: -- this repo currently enqueues only `agent.message` notifications for interrupted Sage agent SSE turns +- this repo currently enqueues only `agent.message` notifications for interrupted Maple agent SSE turns while ensuring that: @@ -174,11 +174,11 @@ Why: - push tokens are sensitive enough to encrypt at rest - this keeps the privacy story aligned with the rest of OpenSecret -### 5.4 Successful live Sage SSE delivery suppresses push in v1 +### 5.4 Successful live Maple SSE delivery suppresses push in v1 -For the common β€œuser sent a message to Sage and is waiting on an SSE stream” flow: +For the common β€œuser sent a message to Maple and is waiting on an SSE stream” flow: -- if the final Sage response is delivered successfully over the open SSE stream, do **not** enqueue push for any device +- if the final Maple response is delivered successfully over the open SSE stream, do **not** enqueue push for any device - if the SSE stream closes or errors before the final response is delivered, enqueue **exactly one** notification event for that interrupted turn - if the interrupted turn produced multiple assistant messages, use the **first missed assistant message** as the preview source for that one notification - this SSE success signal acts as the acknowledgement; no extra per-thread presence service or explicit ACK protocol is required in v1 @@ -241,7 +241,7 @@ This section reflects the current backend code layout rather than the original p - App-user authenticated device register / list / revoke routes. - `src/web/agent/mod.rs` - - Sage agent SSE acknowledgement tracking plus the current disconnect-triggered `agent.message` enqueue path. + - Maple agent SSE acknowledgement tracking plus the current disconnect-triggered `agent.message` enqueue path. - `src/web/platform/common.rs` - Push secret constants and project-settings request / response types. @@ -788,7 +788,7 @@ Recommended envelope: "v": 1, "notification_id": "uuid", "message_id": "uuid", - "kind": "sage.reminder", + "kind": "maple.reminder", "title": "Reminder", "body": "Follow up on the deployment thread", "deep_link": "opensecret://agent/subagent/uuid", @@ -940,7 +940,7 @@ For v1, prefer the standard messaging-app approach: - let the app fetch or render authoritative content after open This is the most standard and reliable path. -It avoids relying on background execution for every Sage message. +It avoids relying on background execution for every Maple message. ### 14.3 Standard FCM payload for v1 @@ -949,7 +949,7 @@ It avoids relying on background execution for every Sage message. "message": { "token": "", "notification": { - "title": "New Sage message", + "title": "New Maple message", "body": "Open Maple to view it" }, "data": { @@ -978,7 +978,7 @@ Notes: - visible text should remain generic / non-sensitive - FCM `data` values must be strings - current backend derives Android TTL from `expires_at`; when an event has no explicit expiry, the sender defaults to 7 days -- the current `agent.message` enqueue path sets a 7-day expiry, so the common Sage-agent payload today is effectively close to `604800s` +- the current `agent.message` enqueue path sets a 7-day expiry, so the common Maple-agent payload today is effectively close to `604800s` - FCM `data` is provider-visible routing metadata, so it may include fields like `notification_id`, `message_id`, `thread_id`, `deep_link`, and `kind`, but never plaintext message content - this is intentional in v1: the goal is to keep message content encrypted from the push provider, not to hide all metadata needed for client routing - current implementation auto-keys `collapse_key` to `notification_id` without the older `sage:` prefix shown in earlier drafts @@ -1067,7 +1067,7 @@ Current Maple status: #### Foreground / open-thread behavior -If the app is already foregrounded and showing the target Sage conversation or active response screen: +If the app is already foregrounded and showing the target Maple conversation or active response screen: - implement `userNotificationCenter(_:willPresent:withCompletionHandler:)` - return `[]` to suppress banner / sound / badge presentation for that in-app event @@ -1195,11 +1195,11 @@ The exact IPs and ports can be changed, but they should use currently unused slo ## 17. Product Event Model -### 17.1 Live Sage agent SSE flows +### 17.1 Live Maple agent SSE flows -For the common flow where the user sends a message to Sage and the client keeps an SSE stream open waiting for responses: +For the common flow where the user sends a message to Maple and the client keeps an SSE stream open waiting for responses: -1. generate and stream the Sage response as normal +1. generate and stream the Maple response as normal 2. if the final response is successfully delivered over that SSE stream, do **not** enqueue push 3. if the SSE stream closes or errors before final response delivery completes, enqueue exactly one notification event for that interrupted turn 4. if multiple assistant messages were generated after the disconnect point, use the first missed assistant message as the preview source for that one notification @@ -1235,7 +1235,7 @@ The current helper generates `notification_id` internally and derives `collapse_ Current implemented backend source: - `src/web/agent/mod.rs` - - live Sage agent SSE completion / failure handling and one `agent.message` notification for an interrupted turn + - live Maple agent SSE completion / failure handling and one `agent.message` notification for an interrupted turn Future integrations: @@ -1347,14 +1347,14 @@ The v1 implementation should be: - **project-scoped provider config in `project_settings` + `org_project_secrets`** - **per-device P-256 notification keypairs** - **Postgres outbox + multi-enclave delivery worker coordinated via Postgres row leases** -- **successful live Sage SSE delivery suppresses push entirely** +- **successful live Maple SSE delivery suppresses push entirely** - **iOS encrypted preview with generic fallback** - **Android v1 uses standard visible FCM notifications plus data, with app-open fetch and local foreground suppression** - **stable `notification_id` plus client-side dedup / thread cleanup** This is the simplest design that still matches the current OpenSecret architecture and avoids painting us into a corner later. -Current implementation note: the backend in this repo already follows this design for the Sage agent disconnect path, while broader event sources and some hardening items remain follow-up work. +Current implementation note: the backend in this repo already follows this design for the Maple agent disconnect path, while broader event sources and some hardening items remain follow-up work. --- @@ -1387,5 +1387,5 @@ These items are the main remaining hardening work after the current push handoff - finish macOS encrypted-preview parity so the macOS notification service extension rewrites the generic fallback text like iOS - add client-side `notification_id` dedup caches and thread cleanup where they materially improve UX -- extend backend event sources beyond the current `agent.message` Sage disconnect path when product priorities require it +- extend backend event sources beyond the current `agent.message` Maple disconnect path when product priorities require it - keep the mobile-client behavior docs aligned with what is actually implemented in the Maple repos diff --git a/internal_docs/sage-in-maple-architecture.md b/internal_docs/maple-agent-memory-architecture.md similarity index 96% rename from internal_docs/sage-in-maple-architecture.md rename to internal_docs/maple-agent-memory-architecture.md index 523fa81f..6d695474 100644 --- a/internal_docs/sage-in-maple-architecture.md +++ b/internal_docs/maple-agent-memory-architecture.md @@ -1,4 +1,4 @@ -# Sage-in-Maple: Agent Memory Architecture +# Maple Agent Memory Architecture ## Main Agent + Shared-Memory Subagents @@ -7,8 +7,8 @@ **Related Docs:** - `potential-rag-integration-brute-force.md` -- RAG/vector storage layer - `architecture-for-rag-integration.md` -- OpenSecret encryption and data model reference -- Sage V2 Design Doc (`~/Dev/Personal/sage/docs/SAGE_V2_DESIGN.md`) -- proven prototype -- Sage V2 Codebase (`~/Dev/Personal/sage/crates/sage-core`) -- DSRs signatures + BAML parsing + multi-step tool loop prototype +- Prototype design doc -- proven 4-tier memory architecture reference +- Prototype agent codebase -- DSR signatures, structured parsing, and multi-step tool loop reference **Implementation note:** the current local codebase still reflects an older single-agent MVP in places. This document is now the source of truth for the rewrite. @@ -16,7 +16,7 @@ ## 1. Goal -Bring Sage's proven 4-tier memory architecture (core, recall, archival, summary) into Maple as a first-class product experience, without breaking the existing Responses API that third-party developers already understand. +Bring the proven 4-tier memory architecture (core, recall, archival, summary) into Maple as a first-class product experience, without breaking the existing Responses API that third-party developers already understand. The target product shape is: - **one main persistent agent** as the app's home surface @@ -150,7 +150,7 @@ What becomes agent-specific: - `conversation_summaries` - agent-only tools / orchestration logic -This keeps the OpenAI-compatible Responses API stable while allowing Sage-style regenerated context and memory behavior for the main agent and subagents. +This keeps the OpenAI-compatible Responses API stable while allowing Maple-style regenerated context and memory behavior for the main agent and subagents. ### 2.2 Introduce Explicit Agent Identity Now @@ -188,10 +188,10 @@ So the rule is: ### 2.5 No Backwards Compatibility in Local-Only Schema -This work has not been deployed anywhere yet. We do **not** need compatibility-preserving additive migrations for the current Sage work. +This work has not been deployed anywhere yet. We do **not** need compatibility-preserving additive migrations for the current Maple work. So for the rewrite: -- edit the existing local-only Sage migrations directly +- edit the existing local-only Maple migrations directly - replace `agent_config` with `agents` - simplify `memory_blocks` now instead of carrying legacy fields forward - move code-owned defaults out of SQL and into Rust/config @@ -221,7 +221,7 @@ Those initial onboarding messages should be **server-authored and persisted as r This decision does not change. -**Maple's split message tables are still the right storage layer** for Sage-style agents. The main difference is no longer β€œone persistent agent thread per user”; it is now β€œone persistent thread per agent identity.” +**Maple's split message tables are still the right storage layer** for Maple-style agents. The main difference is no longer β€œone persistent agent thread per user”; it is now β€œone persistent thread per agent identity.” | Operation | Agent requirement | Maple storage choice | |---|---|---| @@ -570,7 +570,7 @@ Each agent thread owns its own summary chain. Main agent compaction and subagent No architectural change is needed here. -The Sage-style approach still fits: +The Maple-style approach still fits: - preprocess image input into text - inject the derived text into the conversation flow - persist the derived text so it can be embedded and recalled later @@ -665,7 +665,7 @@ These may exist temporarily as local debugging surfaces, but they should not be - manual memory block CRUD endpoints - manual archival insert/delete endpoints for normal product usage -The product intent is **agent-managed memory**, not β€œsettings panels for the user to tune Sage.” +The product intent is **agent-managed memory**, not β€œsettings panels for the user to tune Maple.” ### 7.3 Conversation List Integration @@ -751,7 +751,7 @@ The important v1 decision is already made: **subagents share memory and search, ## 10. Implementation Ordering -Because this is still local-only, we should update the existing Sage docs and schema directly instead of preserving the old MVP shape. +Because this is still local-only, we should update the existing Maple docs and schema directly instead of preserving the old MVP shape. ### Phase 1: RAG Foundation - Keep `user_embeddings` diff --git a/internal_docs/potential-rag-integration-brute-force.md b/internal_docs/potential-rag-integration-brute-force.md new file mode 100644 index 00000000..77cb9ff4 --- /dev/null +++ b/internal_docs/potential-rag-integration-brute-force.md @@ -0,0 +1,647 @@ +# OpenSecret RAG Architecture Proposal + +## Chat Memory via Encrypted Vector Search + +**Author:** Claude (with T), updated with Droid +**Date:** February 2026 +**Status:** Approved for Implementation + +--- + +## Executive Summary + +This proposal describes a RAG (Retrieval-Augmented Generation) system for OpenSecret that enables semantic search over encrypted user chat history and agent archival memory. The system stores per-message embeddings generated by the existing Tinfoil proxy, encrypted at rest with per-user keys, and performs brute-force cosine similarity search in-process inside the Nitro Enclave after decryption. + +The design prioritizes simplicity, correctness, and architectural consistency with OpenSecret's existing encryption model. It is divided into two parts: what to build now, and what to add as the system scales. + +### Revision Notes + +This document was revised after an architectural analysis comparing the OpenSecret backend with the Maple agent prototype (a privacy-first personal AI agent with a Letta-inspired 4-tier memory system). Key changes from the original proposal: + +- **Single database, not separate PlanetScale instance.** The vector workload doesn't justify a second database given brute-force in-process search. Eliminates cross-DB consistency concerns and operational complexity. +- **Narrowed v1 scope: chat + archival only.** Document ingestion/chunking and batch indexing are explicitly out of scope for the first implementation. +- **`user_embeddings` references messages via explicit BIGINT foreign keys.** Message-derived embeddings reference `user_messages.id` / `assistant_messages.id` with `ON DELETE CASCADE`, ensuring unambiguous provenance (user vs assistant) and strong deletion consistency. +- **Content duplication is retained (for now).** Even with explicit message FKs, storing the embedded text alongside the vector keeps search results and future re-embedding self-contained. A later optimization can drop duplicated message content and fetch it from the message tables only for the final top-k. +- **Explicit hooks for future agent integration.** The schema and API are designed so a Maple-style agent can use `archival_insert` / `archival_search` / `conversation_search` tools against this same storage layer. + +### Implementation Rollout Strategy + +The implementation is split into two phases: + +**Phase 1 (Now): Storage and search infrastructure.** +Build the `user_embeddings` table, Diesel models, encryption/serialization helpers, brute-force search engine, LRU cache, and a set of **experimental REST endpoints** (`/v1/rag/*`) for manual testing. These endpoints exist solely to validate the infrastructure before Maple lands and are **not intended for production client use**. No existing code paths (Responses API, etc.) are modified to auto-embed messages. + +**Phase 2 (When Maple lands): Agent integration and auto-embedding.** +- The Responses API message creation flow gains `tokio::spawn` hooks to auto-embed user and assistant messages after persistence. +- Maple's agent loop calls the storage/search layer directly via internal Rust functions (NOT via HTTP) using `archival_insert`, `archival_search`, and `conversation_search` tool implementations. +- The experimental REST endpoints are removed. All embedding reads/writes go through the agent's tool-use interface or internal auto-embedding hooks. No embeddings functionality is exposed over HTTP in production. +- `store=false` rejection is added to the Responses API handler at this point (when auto-embedding is active and the invariant matters). +- Billing hooks (`publish_usage_event_internal`) for auto-embedding are wired up at this point. + +--- + +## Part 1: What To Build Now + +### 1.1 Architecture Decision Summary + +| Decision | Choice | Rationale | +|---|---|---| +| Where does search run? | In-process, inside the existing Nitro Enclave | User keys are only available inside the enclave. No external vector DB can search encrypted vectors. | +| Search algorithm | Brute-force cosine similarity | At target per-user scale (<=40K vectors), brute force completes in ~4ms with 100% recall. No approximation errors, no index to maintain. | +| Vector storage | `user_embeddings` table in the **main PostgreSQL database** | Same encryption patterns, same connection pool, same Diesel migrations. No cross-DB consistency issues. Workload isolation can be achieved later via read replicas if needed. | +| Embedding model | `nomic-embed-text` via existing Tinfoil proxy (768 dimensions) | Already deployed at `localhost:8093`. Embeddings never leave the TEE chain. | +| Chunking strategy | One embedding per message (user and assistant separately) | Chat messages are short enough that per-message embedding captures intent well. Preserves granularity for retrieval. | +| Embedding generation | Server-side (user sends text, server calls Tinfoil) | Keeps the embedding model consistent. Removes client-side dependency on a specific model. | +| Content duplication | Store encrypted content copy in embeddings table | Keeps search results and future re-embedding self-contained (no extra DB round-trips to fetch/decrypt message content for the final top-k). Storage cost is negligible vs the vector itself. | + +### 1.2 Data Flow + +**Indexing (on message creation):** + +``` +User sends message + -> OpenSecret processes chat as normal + -> After storing encrypted message in main DB (existing flow) + -> Call Tinfoil proxy: POST localhost:8093/v1/embeddings + -> Receive 768-dim float vector + -> Encrypt vector with user's per-user key (AES-256-GCM) + -> Encrypt message text with user's per-user key (AES-256-GCM) + -> Store in user_embeddings: + (user_id, source_type='message', conversation_id, + user_message_id OR assistant_message_id, + vector_enc, content_enc, metadata_enc, token_count) +``` + +**Indexing (agent archival insert -- future):** + +``` +Agent decides to store a memory + -> Agent calls archival_insert tool with content + optional tags + -> Server calls Tinfoil proxy for embedding + -> Encrypt vector + content with user's per-user key + -> Store in user_embeddings with source_type = 'archival' +``` + +**Querying (on RAG-augmented chat or agent archival_search):** + +``` +User sends message (or agent triggers search) + -> Derive user key from enclave_key + seed_enc + -> Embed the query text via Tinfoil proxy + -> Load encrypted vectors from user_embeddings (streaming in batches of 1K) + -> For each batch: + -> Decrypt vectors in-process + -> Compute cosine similarity against query embedding + -> Maintain running top-k heap + -> Return top-k texts (decrypted) for context injection +``` + +### 1.3 Database Schema + +A single new table in the main PostgreSQL database, following existing Diesel migration conventions. + +#### `user_embeddings` table (v1 scope: chat history + archival memory) + +This v1 implementation intentionally narrows scope to two embedding sources: + +- `message`: auto-indexed chat history (user and assistant messages) +- `archival`: agent-inserted long-term memories (this row is the source of truth) + +Document ingestion/chunking and batch indexing are explicitly out of scope for v1. + +**Why explicit BIGINT message foreign keys (not `source_uuid`)** + +In PostgreSQL, a single column cannot be a foreign key to *either* `user_messages` *or* `assistant_messages`. A polymorphic `source_uuid` would be ambiguous (user vs assistant), hard to enforce, and would not give us `ON DELETE CASCADE` semantics for data consistency. + +Instead, message-derived embeddings reference the internal auto-incrementing IDs directly: + +- `user_message_id BIGINT REFERENCES user_messages(id) ON DELETE CASCADE` +- `assistant_message_id BIGINT REFERENCES assistant_messages(id) ON DELETE CASCADE` + +This aligns with OpenSecret conventions (BIGINT `id` for internal references; UUIDs for external API objects) and guarantees deletion consistency: + +- Deleting a `response` cascades to messages β†’ cascades to embeddings. +- Deleting a `conversation` cascades to messages β†’ cascades to embeddings. + +**Source type matrix (v1):** + +| `source_type` | `content_enc` | Message reference | `conversation_id` | Use case | +|---|---|---|---|---| +| `message` | Copy of message text | Exactly one of `user_message_id` / `assistant_message_id` | `conversations.id` (BIGINT FK) | Auto-indexed chat turns | +| `archival` | The passage itself (source of truth) | NULL | NULL | Agent stores "The user has a dog named Bob" | + +```sql +CREATE TABLE user_embeddings ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + + -- Source tracking + -- v1 uses: 'message', 'archival'. We intentionally keep this as open TEXT. + source_type TEXT NOT NULL DEFAULT 'message', + + -- Message provenance (ONLY for source_type='message') + user_message_id BIGINT REFERENCES user_messages(id) ON DELETE CASCADE, + assistant_message_id BIGINT REFERENCES assistant_messages(id) ON DELETE CASCADE, + conversation_id BIGINT REFERENCES conversations(id) ON DELETE CASCADE, + + -- Embedding vector + vector_enc BYTEA NOT NULL, -- AES-256-GCM encrypted float32 array + embedding_model TEXT NOT NULL, -- e.g. "nomic-embed-text", for staleness detection + vector_dim INTEGER NOT NULL DEFAULT 768, + + -- Content that was embedded + content_enc BYTEA NOT NULL, -- AES-256-GCM encrypted text that was embedded + metadata_enc BYTEA, -- AES-256-GCM encrypted JSON (tags, optional extras) + + -- Plaintext metadata for querying/budgeting + token_count INTEGER NOT NULL, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + -- Enforce invariants for the v1 source types without blocking future types. + CONSTRAINT user_embeddings_message_source_check CHECK ( + (source_type <> 'message') + OR ( + -- exactly one of the message FKs must be set + (user_message_id IS NOT NULL) <> (assistant_message_id IS NOT NULL) + ) + ), + CONSTRAINT user_embeddings_message_conversation_check CHECK ( + (source_type <> 'message') OR (conversation_id IS NOT NULL) + ), + CONSTRAINT user_embeddings_archival_source_check CHECK ( + (source_type <> 'archival') + OR ( + user_message_id IS NULL + AND assistant_message_id IS NULL + AND conversation_id IS NULL + ) + ) +); + +-- Primary query path: load all vectors for a user (brute-force scan) +CREATE INDEX idx_user_embeddings_user_id ON user_embeddings(user_id); + +-- For time-filtered searches (recency bias) +CREATE INDEX idx_user_embeddings_user_created ON user_embeddings(user_id, created_at DESC); + +-- For source-type filtered searches (message vs archival) +CREATE INDEX idx_user_embeddings_user_source ON user_embeddings(user_id, source_type); + +-- For conversation-scoped searches (message recall) +CREATE INDEX idx_user_embeddings_user_conversation ON user_embeddings(user_id, conversation_id); + +-- Idempotency/deduplication: prevent double-indexing the same message +CREATE UNIQUE INDEX idx_user_embeddings_user_message_id + ON user_embeddings(user_message_id) + WHERE user_message_id IS NOT NULL; + +CREATE UNIQUE INDEX idx_user_embeddings_assistant_message_id + ON user_embeddings(assistant_message_id) + WHERE assistant_message_id IS NOT NULL; + +-- For staleness queries: find vectors that need re-embedding after model change +CREATE INDEX idx_user_embeddings_model ON user_embeddings(user_id, embedding_model); + +-- Trigger for updated_at +CREATE TRIGGER update_user_embeddings_updated_at +BEFORE UPDATE ON user_embeddings +FOR EACH ROW EXECUTE FUNCTION update_updated_at_column(); +``` + +**Column notes:** + +- `vector_enc`: Float32 array serialized to bytes using little-endian byte order (`f32::to_le_bytes()` / `f32::from_le_bytes()`), then AES-256-GCM encrypted with a random 12-byte nonce prepended. For 768 dims: 768Γ—4 = 3,072 bytes, plus nonce overhead. +- `embedding_model`: Plaintext model identifier (e.g. `"nomic-embed-text"`) for staleness detection and re-embedding migrations. +- `vector_dim`: Plaintext dimension count for validation and potential future model upgrades. +- `user_message_id` / `assistant_message_id`: Internal message references (BIGINT) with hard `ON DELETE CASCADE`. Exactly one is set for `source_type='message'`. +- `conversation_id`: Internal conversation reference (BIGINT). Set for message embeddings to support efficient conversation-scoped search without joins. +- `content_enc`: The text that was embedded, encrypted with the user's per-user key. For `message`, this duplicates message text (keeps retrieval + future re-embedding self-contained). For `archival`, this is the source of truth. +- `metadata_enc`: Encrypted JSON. Recommended shapes: + - `archival`: `{tags: ["preferences"], description: "..."}` + - `message`: optional and generally unnecessary (provenance is determined by which FK column is set) +- `token_count`: Plaintext token count for budget management without decryption. + +#### Why a single table instead of per-source-type tables? + +Considered alternatives: +- **Separate tables per source type** (`message_embeddings`, `archival_embeddings`). Rejected: the data shape is almost identical and the search path is the same. Separate tables means duplicate Diesel models and duplicate search code. One table with a `source_type` discriminator keeps implementation simpler. +- **Adding `embedding_enc BYTEA` to existing message tables.** Rejected because: (1) not all messages should be embedded (tool calls, reasoning items, incomplete messages); (2) bloats the context builder's UNION ALL query with ~3 KB per row even when RAG isn't used; (3) couples embedding lifecycle to message lifecycle (re-embedding after model upgrade would require updating message rows); (4) archival passages have no corresponding message row. + +#### Why `source_type` is not a CHECK constraint + +The column uses open `TEXT` instead of `CHECK (source_type IN (...))`. This is intentional: adding a new source type later can start as an application-level change without immediately requiring a DB migration. The v1 CHECK constraints above only enforce invariants when `source_type` is one of the known v1 types (`message`, `archival`). + +### 1.4 Search Implementation + +Brute-force cosine similarity in Rust. No external dependencies beyond standard float math. + +```rust +fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); + let norm_a: f32 = a.iter().map(|x| x * x).sum::().sqrt(); + let norm_b: f32 = b.iter().map(|x| x * x).sum::().sqrt(); + dot / (norm_a * norm_b) +} +``` + +**Streaming batch approach:** + +Rather than loading all of a user's vectors in one query, use paginated batches: + +1. Query (keyset pagination): `SELECT id, vector_enc, content_enc, metadata_enc, token_count FROM user_embeddings WHERE user_id = $1 AND id > $2 ORDER BY id LIMIT 1000` +2. Decrypt the batch of vectors in-process +3. Compute cosine similarity for each against the query vector +4. Maintain a min-heap of size k (top-k results across all batches) +5. Repeat until no more rows +6. Return the top-k results with decrypted content + +Optional filters can be applied at the SQL level to reduce the scan set before any decryption happens: +- `AND source_type = 'archival'` -- search only archival memory +- `AND source_type = 'message'` -- search only chat history +- `AND source_type IN ('message', 'archival')` -- search multiple types +- `AND conversation_id = $3` -- search within a specific conversation (internal `conversations.id`) +- `AND created_at > $5` -- recency filter +- `AND embedding_model = $6` -- only search vectors from a specific model (useful during migration) + +For users with <=1,000 vectors (the vast majority), this completes in a single batch. + +### 1.5 In-Memory Cache (Simple LRU) + +A basic LRU cache of decrypted vector sets, keyed by `user_id`. Purpose: avoid re-loading and re-decrypting vectors on repeated queries within the same session. + +**Initial implementation:** + +- `HashMap` with a `VecDeque` for LRU ordering +- `CachedVectorSet`: the decrypted vectors + metadata, plus a timestamp +- **Max entries:** Start with 100 users cached simultaneously +- **Eviction:** LRU eviction when cache is full. Also evict entries older than 5 minutes (TTL). +- **Invalidation:** Evict the user's cache entry any time embeddings change (insert/update/delete). In practice: evict on successful inserts/updates; also evict on delete paths (delete-all embeddings and conversation deletion) so deleted content doesn't appear in search results until TTL expiry. +- **Memory budget:** 100 users x ~280 KB average (1,000 vectors x 280 bytes per decrypted entry) = ~28 MB. Conservative and safe for enclave memory. + +The cache is populated as a side effect of the streaming query: as batches are decrypted for search, the results are also written into the cache. Subsequent queries for the same user hit the cache directly and skip the DB entirely. + +This is deliberately simple. No distributed cache, no cross-instance coordination. Each enclave instance has its own independent LRU. With horizontal scaling, a user routed to a different instance gets a cache miss and loads fresh -- this is fine. + +### 1.6 API Design + +The production API surface for RAG is **internal only** -- no HTTP endpoints are exposed to clients. All embedding reads and writes happen through: + +1. **Auto-embedding hooks** in the Responses API message persistence flow (Phase 2, when Maple lands) +2. **Agent memory tools** -- `archival_insert`, `archival_search`, `conversation_search` -- called internally by the Maple agent loop as Rust function calls, not HTTP requests + +#### Experimental REST Endpoints (Phase 1 Only -- Will Be Removed) + +During Phase 1, a set of `/v1/rag/*` HTTP endpoints are provided **solely for development and testing** of the storage/search infrastructure. These endpoints: + +- Are **not documented** for SDK consumers +- Are **not called by any existing code path** (no auto-embedding, no Responses API integration) +- Will be **deleted before Maple ships to production** -- replaced by internal Rust function calls from the agent loop +- Serve as a convenient way to manually insert embeddings, run searches, and validate the system end-to-end during development + +#### Authentication and Encryption Middleware + +All `/v1/rag/*` experimental routes use the existing authentication and session encryption middleware, identical to the KV store and Responses API routes. This means: +- Session-based ECDH encryption (request/response bodies encrypted via `x-session-id`) +- User authentication via JWT (provides `Extension`, `Extension` session ID) +- Per-user key derivation via `get_user_key(user_uuid, None, None)` for encrypting/decrypting embeddings +- No new auth mechanisms or middleware are introduced. + +#### Experimental Endpoints (`/v1/rag/*`) + +**Auto-index chat messages (Phase 2 -- not built in Phase 1):** + +When Maple lands, message embeddings will be created automatically as a background `tokio::spawn` task after a `user_message` / `assistant_message` is persisted. **No auto-embedding hooks are added to existing code in Phase 1.** The auto-embedding flow is documented here for completeness: + +- There is **no client-facing endpoint** to index messages (not even experimentally). +- Each message embedding row uses: + - `source_type = 'message'` + - `conversation_id = conversations.id` + - exactly one of `user_message_id` / `assistant_message_id` set + +This gives us unambiguous provenance (user vs assistant) and hard deletion consistency via `ON DELETE CASCADE`. + +**Conversation UUID resolution (UUID -> BIGINT):** Externally, experimental APIs accept the public `conversations.uuid` for consistency with existing endpoints. For conversation-scoped search, resolve UUID -> internal `conversations.id` via `conversations::table.filter(conversations::uuid.eq(conv_uuid)).select(conversations::id).first::(conn)` (the `uuid` column has a unique index). The resolved internal ID is then used to filter on `user_embeddings.conversation_id`. + +**Insert archival memory (experimental -- tests the insert + embed + encrypt pipeline):** + +In production, archival inserts will be performed internally by the Maple agent via `archival_insert` tool calls. This endpoint exists for Phase 1 testing only. + +``` +POST /v1/rag/embeddings +{ + "text": "User prefers dark mode and concise responses", + "metadata": { "tags": ["preferences"] } // optional, encrypted as metadata_enc +} +-> 201 Created +{ + "id": "embedding-uuid", + "source_type": "archival", + "embedding_model": "nomic-embed-text", + "token_count": 8, + "created_at": "2026-02-09T..." +} +``` + +**Search (experimental -- tests the decrypt + brute-force + top-k pipeline):** + +In production, searches will be performed internally by the Maple agent via `conversation_search` and `archival_search` tool calls. + +``` +POST /v1/rag/search +{ + "query": "what did we discuss about deployment strategies?", + "top_k": 5, // default 5, max 20 + "max_tokens": 2000, // optional token budget for results + "source_types": ["message"], // optional: filter by source type(s), null for all + "conversation_id": "uuid" // optional: scope to specific conversation (public conversation UUID; server resolves to conversations.id) +} +-> 200 OK +{ + "results": [ + { + "content": "decrypted text...", + "score": 0.87, + "token_count": 145 + } + ] +} +``` + +**Search response shape (intentionally minimal):** The only required consumer is the agent/context-builder, which only needs decrypted text + ranking score (and `token_count` for packing within a context budget). Filtering happens in the query, not the response. By not returning any provenance fields (conversation/message IDs, source types, embedding IDs), the implementation avoids any BIGINT-to-UUID reverse-resolution and can simply decrypt the top-k `content_enc` values and return them. + +**Delete all embeddings for a user (account deletion / data reset):** + +``` +DELETE /v1/rag/embeddings +-> 204 No Content +``` + +Deletes all embeddings for the authenticated user. Also evicts the LRU cache entry. Note: this endpoint is likely to survive into production (useful for account deletion flows), unlike the insert/search endpoints. + +**Delete a specific embedding by ID:** + +``` +DELETE /v1/rag/embeddings/:id +-> 204 No Content +``` + +**Re-embedding support (model migration):** + +``` +GET /v1/rag/embeddings/status +-> 200 OK +{ + "total_embeddings": 4521, + "by_model": { + "nomic-embed-text": 4521 + }, + "stale_count": 0 // embeddings using a deprecated model +} +``` + +When a model is deprecated, a background job can query `WHERE embedding_model != 'current-model'`, decrypt the `content_enc`, re-embed via Tinfoil, and update `vector_enc` + `embedding_model`. The `content_enc` being always present is what makes this possible without needing to join back to source tables. + +### 1.7 Data Visibility and Isolation + +A key design decision: **which messages get embedded, and who can search what?** + +#### The Model + +One user has three ways to interact with Maple: + +1. **Responses API threads** (`/v1/responses/*`) -- stateless, one-off conversations. A user can have many of these. +2. **Main agent** (`/v1/agent/chat`) -- a single persistent home agent with one long-running thread. This thread is hidden from the generic conversation list. +3. **Subagent chats** -- topic-specific agent conversations. A user can have many of these. Each subagent has its own long-running thread and appears in `/v1/conversations/*` like a first-class chat. + +The main agent and all subagents share the same user-level memory layer. What differs between them is only thread-local history and summaries. + +#### Default: Broad Visibility + +**All stored messages from all three surfaces are embedded into `user_embeddings` when eligible.** An agent's `conversation_search` searches across all of the user's messages by default -- main-agent messages, subagent messages, and Responses API thread messages. + +Rationale: the entire agent system is the user's persistent memory layer. If a user had a detailed Responses API conversation about Kubernetes deployment strategies, or already discussed the topic in another subagent, then asks the main agent or a new subagent about deployment, that agent should surface that knowledge. Users who opt into the agent system expect cross-conversation memory. That's the whole point. + +The `conversation_id` column on each embedding enables filtering when needed: +- **Broad search (default):** `WHERE user_id = $1` -- agent sees everything +- **Main-thread-only search:** `WHERE user_id = $1 AND conversation_id = $main_agent_conversation_id` +- **Specific subagent or Responses thread search:** `WHERE user_id = $1 AND conversation_id = $specific_thread_id` + +Since search is brute-force in-process, these filters are just WHERE clauses applied before decryption. Zero architectural cost to support any combination. + +#### `store=false` Rejection (Phase 2) + +OpenSecret does not support `store=false`. The `store` field exists on `ResponsesCreateRequest` (defaults to `true`) and is persisted on the `responses` row, but no code path currently changes behavior based on it. **When auto-embedding is activated in Phase 2**, requests with `store=false` will be rejected with a 400 Bad Request error, added early in the Responses API handler (in `handlers.rs`), before any DB writes or streaming begins. This is deferred to Phase 2 because the rejection only matters when auto-embedding is active and the "all stored messages get embedded" invariant needs enforcement. + +A future `private` / incognito flag on conversations may allow users to have Responses API threads that are excluded from embedding entirely. This is out of scope for v1. + +### 1.8 Integration Points + +#### Phase 1: No Integration With Existing Code + +**In Phase 1, no existing code paths are modified.** The Responses API, conversations system, and all other handlers remain untouched. The RAG infrastructure (table, models, search engine, cache) is built as a standalone module with experimental HTTP endpoints for testing. This keeps Phase 1 low-risk and independently deployable. + +#### Phase 2: Responses API Auto-Embedding (When Maple Lands) + +The Responses API message creation flow will gain background embedding hooks: + +``` +ResponsesCreateRequest (store=false is rejected with 400 before reaching this point) + -> persist user_message (existing) + -> persist assistant_message (existing, after streaming) + -> [async, tokio::spawn] embed user_message and store in user_embeddings + -> [async, tokio::spawn] embed assistant_message and store in user_embeddings +``` + +Embedding failures are non-fatal. Messages are stored regardless. Embeddings are a best-effort enrichment. + +**Async embedding via `tokio::spawn` (no formal queue):** + +Embedding is fired as a `tokio::spawn` task after the message is persisted and the response stream completes. This is a simple fire-and-forget async task -- no background queue, no worker pool, no separate infrastructure. At current scale, concurrent Tinfoil proxy calls from simultaneous users are not a concern. If this becomes a bottleneck (e.g., thundering herd during peak usage), a bounded semaphore or formal queue can be added later. Embedding failures are logged but non-fatal -- the message is already stored. The unique indexes on `user_message_id` / `assistant_message_id` ensure one embedding per message, making retries or future backfill safe and idempotent. + +#### Phase 2: Billing (User-Billed, Activated With Auto-Embedding) + +Auto-indexing embedding calls will be billed per-user, identical to the existing `/v1/embeddings` endpoint. The existing `proxy_embeddings` handler already calls `publish_usage_event_internal` with `prompt_tokens` after every Tinfoil proxy call. The RAG auto-embedding path will reuse the same `publish_usage_event_internal` function with the same model name (`nomic-embed-text`) and prompt token accounting. No new billing mechanism is needed. **In Phase 1, no billing hooks are added** since there is no auto-embedding; any manual testing via experimental endpoints does not generate billing events. + +#### Phase 2: Agent System Integration (When Maple Lands) + +The agent will call the storage/search layer directly via internal Rust functions -- **not via HTTP**: + +- `archival_insert` -> calls the embed + encrypt + store function internally with `source_type: "archival"` +- `archival_search` -> calls the search function internally with `source_types: ["archival"]` +- `conversation_search` -> calls the search function internally with `source_types: ["message"]` (default: all conversations, not just the agent's own thread) + +The agent's context builder injects search results into the system prompt as memory blocks, rather than the Responses API's current middle-truncation approach. The experimental HTTP endpoints are removed at this point. + +### 1.9 What This Gets You + +For a typical user with a year of chat history (~2,000-4,000 messages -> 4,000-8,000 embeddings): + +- **Storage:** ~28-56 MB per user in the embeddings table +- **Search latency:** Single-digit milliseconds (brute force on cached vectors) to ~50ms (cold start, streaming from DB) +- **Accuracy:** 100% recall -- true nearest neighbors, not approximate +- **Privacy:** Identical to existing encryption model. Vectors encrypted at rest with per-user keys. Decrypted only in-enclave, in-process. No external vector DB has access to plaintext. + +--- + +## Part 2: Scaling Roadmap + +This section describes what to watch, when to act, and what to build as the system grows. These are ordered by when you're likely to need them. + +### 2.0 Operational note: Postgres maintenance (vacuum / bloat) + +Keeping `user_embeddings` in the primary database is the simplest operational model, but it does give up the vacuum isolation benefit of a separate DB. Large re-embedding jobs (mass `UPDATE`) and high churn (`DELETE` from user resets / conversation deletions) will generate dead tuples that autovacuum must clean up, potentially competing with hot-path tables. If this becomes a real issue, mitigate with: small-batch re-embedding, per-table autovacuum tuning, table partitioning, and/or moving scan-heavy reads to a replica (or splitting databases later). + +### 2.1 Monitoring: What to Instrument From Day One + +You need visibility into the vector workload to know when scaling pressure is real vs. theoretical. Add metrics for: + +| Metric | What it tells you | Watch threshold | +|---|---|---| +| `rag.search.latency_p99` | End-to-end search time including DB load + decrypt + search | >100ms | +| `rag.search.db_load_ms` | Time spent loading encrypted vectors from DB | >50ms | +| `rag.search.decrypt_ms` | Time spent decrypting vectors | >20ms | +| `rag.search.compare_ms` | Time spent on cosine similarity | >10ms | +| `rag.cache.hit_rate` | LRU cache effectiveness | <60% | +| `rag.cache.memory_bytes` | Total memory consumed by cached vector sets | Approaching enclave memory limit | +| `rag.user.vector_count` | Distribution of per-user vector counts | p99 approaching 50K | +| `rag.index.insert_latency` | Time to embed + encrypt + store a new vector | >200ms | + +### 2.2 Scaling Tier 1: Cache Optimization (Likely Needed at ~5K-10K DAU) + +**Trigger:** Cache hit rate drops below 70%, or enclave memory pressure increases. + +**What to do:** + +- **Tiered cache:** Store only the vectors (not the content text) in the hot cache. Content is only decrypted for the final top-k results, not the full scan set. This cuts cache memory per user by ~40%. +- **Adaptive TTL:** Power users who query frequently get longer TTLs. Users with a single query per session get evicted faster. +- **Cache warming on session start:** When a user authenticates, pre-load their vectors into the cache asynchronously before any RAG query is triggered. +- **Compressed vectors:** Store vectors as float16 in the cache (not in the DB -- the DB stores full precision encrypted vectors). Half-precision is sufficient for cosine similarity ranking and halves cache memory. + +### 2.3 Scaling Tier 2: Query Optimization (Likely Needed at ~10K-25K DAU) + +**Trigger:** `rag.search.db_load_ms` consistently above 50ms for power users, or `rag.search.latency_p99` above 200ms. + +**What to do:** + +- **Pre-filtering by time window:** Add a `created_at` filter to the vector query. Most RAG use cases for chat memory benefit from recency bias -- searching the last 6 months instead of all history cuts the scan set significantly. The `idx_user_embeddings_user_created` index supports this efficiently. +- **Source-type filtering:** When the agent only needs archival memory or only chat history, filter at the SQL level using `source_type`. +- **Token budget pre-filtering:** If the request includes `max_tokens`, use the plaintext `token_count` column to estimate how many results you need and stop scanning early once you have enough high-scoring candidates to fill the budget. +- **Parallel batch decryption:** Decrypt vector batches on a thread pool (Rayon or Tokio spawn_blocking) while the next batch is being fetched from the DB. Overlaps CPU work with I/O wait. + +### 2.4 Scaling Tier 3: Index Structures (Likely Needed When p99 Users Hit ~50K+ Vectors) + +**Trigger:** `rag.search.compare_ms` consistently above 30ms for power users, and time-window pre-filtering isn't sufficient because users need deep history search. + +You are unlikely to need this for ~12-18 months; brute-force scan + caching should carry the initial product. Treat ANN as an emergency lever once power users' vector sets get large enough to matter. + +**What to do:** + +This is where you introduce an in-process approximate nearest neighbor (ANN) index. The recommended path is: + +**Option A: `instant-distance` (pure Rust HNSW)** + +- Pure Rust, no C++ dependencies, no BLAS/LAPACK +- Builds cleanly inside a Nitro Enclave without cross-compilation pain +- Hierarchical Navigable Small World (HNSW) graph -- 95-99% recall at orders of magnitude faster than brute force for large vector sets +- Build the index from decrypted vectors on cache warm-up, query the index instead of doing a linear scan +- Index building is the slow part (~100ms for 50K vectors), but it's a one-time cost per cache entry + +**Option B: FAISS (via `faiss-rs` bindings)** + +- More mature, more index types (IVF, PQ, HNSW, flat) +- Significantly harder to build inside a Nitro Enclave (C++ library, BLAS dependency) +- Only worth the build complexity if `instant-distance` becomes insufficient + +**In both cases, the architecture doesn't change.** Vectors are still encrypted at rest in PostgreSQL, decrypted in-process, and the index is built ephemerally from decrypted vectors. The DB remains the source of truth. + +### 2.5 Scaling Tier 4: Storage Optimization (Likely Needed at ~100K+ Users) + +**Trigger:** Vector table exceeds 500 GB, or storage costs become a concern. + +**What to do:** + +- **Dimensional reduction:** If `nomic-embed-text` supports Matryoshka representations (v1.5 does), truncate embeddings from 768 to 256 or 384 dimensions at indexing time. Cuts storage by 50-67% with modest recall loss. Requires re-embedding existing vectors. +- **Quantized storage:** Store vectors as int8 instead of float32 in the DB (4x storage reduction). Decrypt and upcast to float32 for search. +- **Archival tiering:** Move embeddings older than N months to cold storage (S3 as encrypted blobs). Only search hot vectors by default. +- **Content deduplication:** For `message` source type, drop `content_enc` and join back to the main message tables for the final top-k results only. The cross-table join only happens for k results (5-20), not the full scan set. + +--- + +## Appendix A: Size Estimates + +### Per-User Storage + +| User type | Messages/year | Embeddings/year | Storage/year | +|---|---|---|---| +| Light (casual chat) | 500 | 1,000 | 7 MB | +| Medium (regular use) | 2,000 | 4,000 | 28 MB | +| Heavy (power user) | 10,000 | 20,000 | 140 MB | +| Max projected (2yr heavy) | 20,000 | 40,000 | 280 MB | + +### Fleet-Wide Storage (Projected 1 Year) + +| Scenario | Total users | Avg embeddings/user | DB size | +|---|---|---|---| +| Current scale | 2,500 MAU | 2,000 | ~35 GB | +| 10x growth | 25,000 MAU | 2,000 | ~350 GB | +| 10x growth + power users | 25,000 MAU + 500 power | mixed | ~430 GB | + +### Search Latency (Brute Force, 768 Dimensions) + +| Vectors per user | Cosine similarity time | Total w/ DB load (cold) | Total w/ cache (warm) | +|---|---|---|---| +| 1,000 | <1 ms | ~20 ms | <2 ms | +| 10,000 | ~1 ms | ~50 ms | ~3 ms | +| 40,000 | ~4 ms | ~200 ms | ~6 ms | +| 100,000 | ~10 ms | ~500 ms | ~15 ms | + +## Appendix B: Encryption Consistency + +The vector storage follows the exact same encryption pattern as existing OpenSecret data: + +| Existing pattern | Vector DB equivalent | +|---|---| +| `user_messages.content_enc` (AES-GCM) | `user_embeddings.content_enc` (AES-GCM) | +| `user_messages.prompt_tokens` (plaintext int) | `user_embeddings.token_count` (plaintext int) | +| `conversations.metadata_enc` (AES-GCM JSON) | `user_embeddings.metadata_enc` (AES-GCM JSON) | +| `user_kv.value_enc` (AES-GCM) | `user_embeddings.vector_enc` (AES-GCM) | +| Key: `get_user_key(user_uuid, None, None)` | Key: `get_user_key(user_uuid, None, None)` | + +No new encryption primitives, key derivation paths, or trust boundaries are introduced. + +## Appendix C: What We're Not Building (Yet) + +For clarity, these are explicitly out of scope. Items are grouped by phase. + +**Out of scope for Phase 1 (infrastructure only):** + +- **Auto-embedding hooks in the Responses API.** No existing code paths are modified. No `tokio::spawn` embedding tasks. This is Phase 2 work. +- **Agent memory tools.** The `archival_insert` / `archival_search` / `conversation_search` tools are designed for but not implemented until Maple lands (Phase 2). +- **`store=false` rejection.** Deferred to Phase 2 when auto-embedding makes the invariant meaningful. +- **Billing hooks for auto-embedding.** No `publish_usage_event_internal` calls from RAG code in Phase 1. +- **Production HTTP API for embeddings.** The Phase 1 experimental endpoints are for testing only and will be removed when Maple ships. In production, all embedding operations go through internal Rust function calls from the agent, not HTTP. + +**Out of scope for both phases (future work):** + +- **Document ingestion / chunking.** v1 only supports chat-history embeddings (`source_type='message'`) and agent archival memory (`source_type='archival'`). Document indexing can be added later with a separate design + migration. +- **Cross-user search** -- architecturally impossible by design and not desired. +- **Automatic RAG injection into Responses API** -- the Responses API stays stateless; RAG injection happens through the agent system. +- **Re-embedding background job.** The schema tracks `embedding_model` per vector and the status endpoint reports staleness. But the actual background re-embedding worker is not part of v1. Re-embedding can be triggered manually or built as a follow-up. +- **Hybrid search (keyword + semantic)** -- brute-force semantic search only for now. +- **External vector database (Pinecone, Weaviate, etc.)** -- incompatible with per-user encryption model. All search must happen inside the enclave. + +## Appendix D: Test Strategy + +**Phase 1 tests** validate the infrastructure in isolation: + +- **Unit tests:** Cosine similarity correctness (known vectors with known scores), vector serialization/encryption round-trips (f32 array -> LE bytes -> AES-GCM encrypt -> decrypt -> LE bytes -> f32 array), and top-k heap behavior (verify correct ranking and tie-breaking). +- **Integration tests:** Cover the experimental `/v1/rag/*` endpoints (archival insert, search, delete, status) using the existing test harness patterns from the Responses API and KV store tests. These tests exercise the full pipeline: HTTP request -> embed via Tinfoil -> encrypt -> store -> search -> decrypt -> response. + +**Phase 2 tests** (when Maple lands) will additionally cover: +- Auto-embedding hooks: verify that creating a message via the Responses API produces a corresponding `user_embeddings` row. +- Agent tool integration: verify `archival_insert`, `archival_search`, `conversation_search` tool calls produce correct results through the internal Rust API. +- Deletion cascades: verify that deleting a conversation/response cascades to embeddings. +- Billing: verify that auto-embedding generates `publish_usage_event_internal` calls. diff --git a/internal_docs/push-notifications-implementation-scratchpad.md b/internal_docs/push-notifications-implementation-scratchpad.md new file mode 100644 index 00000000..bd3c94b4 --- /dev/null +++ b/internal_docs/push-notifications-implementation-scratchpad.md @@ -0,0 +1,110 @@ +# Push Notifications Implementation Scratchpad + +## 2026-03-07 β€” Kickoff + +- Started backend implementation for the mobile push notification system. +- Moved the main design doc from `internal_docs/encrypted-mobile-push-notifications.md` to `docs/encrypted-mobile-push-notifications.md` so it can be shared with app developers. +- Re-read the design doc before starting implementation. +- Current implementation constraints being followed: + - successful live Maple SSE delivery suppresses push in v1 + - Postgres is the durable coordination plane across multiple enclaves + - Android v1 uses standard visible FCM notifications plus `data` + - iOS keeps encrypted preview support with generic fallback + - stable `notification_id` is required for retry-safe client dedup +- Initial audit focus: + - agent SSE flow in `src/web/agent/mod.rs` + - project settings + secrets patterns already used in platform routes and DB helpers + +## 2026-03-07 β€” Scope correction + +- User clarified that implementation scope is **Maple agent APIs only** for now. +- Do **not** modify or depend on the Responses API implementation for this push work. +- Current architecture question discovered during audit: + - agent chat generation currently runs inside the SSE stream in `src/web/agent/mod.rs` + - if the client disconnects, generation likely stops with the stream + - to implement β€œSSE success suppresses push; disconnect triggers push,” we may need to refactor agent generation to continue in a background task after disconnect + +## 2026-03-07 β€” Section 1 complete: schema and core model scaffolding + +- Re-read the push design doc sections covering data model, worker leasing, and notification identity before writing schema changes. +- Added a new migration: `2026-03-07-120000_push_notifications_v1`. +- Added schema for: + - `push_devices` + - `notification_events` + - `notification_deliveries` +- Added model files: + - `src/models/push_devices.rs` + - `src/models/notification_events.rs` + - `src/models/notification_deliveries.rs` +- Updated `src/models/mod.rs` and `src/models/schema.rs` to register the new tables. +- Extended `src/models/project_settings.rs` with typed push settings (`PushSettings`, `IosPushSettings`, `AndroidPushSettings`, `PushEnvironment`) and `SettingCategory::Push`. +- Added typed DB helpers in `src/db.rs` for project push settings. +- Current implementation choice: + - use the existing project-settings DB pattern for push config + - keep push-device / event / delivery transactional flows on direct Diesel connections so enqueue and delivery leasing can stay atomic + +## 2026-03-07 β€” Section 2 complete: routes, config plumbing, and worker skeleton + +- Added platform push settings API under: + - `GET /platform/orgs/:org_id/projects/:project_id/settings/push` + - `PUT /platform/orgs/:org_id/projects/:project_id/settings/push` +- Added app-user push device API under: + - `POST /v1/push/devices` + - `GET /v1/push/devices` + - `DELETE /v1/push/devices/:id` +- Added raw project secret helpers in `AppState` so push senders can consume PEM / JSON secrets without the old base64 wrapper. +- Added push module structure: + - `src/push/mod.rs` + - `src/push/crypto.rs` + - `src/push/apns.rs` + - `src/push/fcm.rs` + - `src/push/worker.rs` +- Added: + - per-device P-256 ECDH encrypted preview envelope generation + - APNs token-auth caching + send path + - FCM OAuth token caching + HTTP v1 send path + - leased delivery worker loop with retry / invalid-token handling +- Wired the push worker into server startup. + +## 2026-03-07 β€” Section 3 complete: Maple SSE disconnect behavior + +- Refactored `src/web/agent/mod.rs` so agent generation now runs in a background task behind an internal channel instead of inside the SSE stream body. +- New behavior: + - if the SSE channel is still alive, agent message events are streamed and push is suppressed + - if the SSE channel is dropped, generation continues and the backend enqueues exactly one notification for that interrupted turn + - if multiple assistant messages were generated after disconnect, the first missed assistant message is used as the preview source for that one notification +- Stream success detection was tightened beyond simple in-process queueing: + - the background task now waits for a stream-side delivery acknowledgement before treating an assistant message event as delivered +- Push enqueue helper currently wires agent-originated messages into `notification_events` + `notification_deliveries` with a stable `notification_id` and generic provider-visible fallback text. + +## 2026-03-08 β€” Review fixups in progress + +- Updated encrypted preview envelope serialization to match the main spec (`enc_v`, `p256-hkdf-sha256-aes256gcm`). +- Added conservative preview-body truncation so encrypted iOS previews stay within a small payload budget. +- Tightened worker failure handling so pre-send errors now transition deliveries to `retry` or `failed` instead of leaving rows stuck in leased retry loops. +- Tightened push registration: + - validate registered public keys as P-256 SPKI + - allow account-switch / stale-token recovery by moving active ownership to the latest installation registration + - stop revoked rows from permanently occupying the global token uniqueness slot +- Tightened APNs invalidation behavior so topic/config mistakes do not auto-revoke valid devices. + +## 2026-03-07 β€” Section 4 complete: enclave networking glue in-repo + +- Updated `entrypoint.sh` with local host mappings for: + - `api.push.apple.com` + - `api.sandbox.push.apple.com` + - `fcm.googleapis.com` +- Added enclave traffic forwarders for: + - APNs prod (`8024`) + - APNs sandbox (`8025`) + - FCM (`8029`) +- Important remaining deployment note: + - the parent-instance `vsock-proxy` allowlist and systemd unit changes described in the design doc are not represented as editable runtime config files in this repo, so those still need to be applied in deployment infrastructure outside this codebase. + +## 2026-03-07 β€” Validation + +- Ran `cargo fmt --all` +- Ran `cargo clippy --all-targets --all-features -- -D warnings` +- Ran `cargo test` +- Result: all validators passed (`104` tests) +- Re-ran `cargo clippy --all-targets --all-features -- -D warnings && cargo test` after the enclave networking script updates; still green. diff --git a/internal_docs/scheduled-agent-wakeups.md b/internal_docs/scheduled-agent-wakeups.md new file mode 100644 index 00000000..4d7df989 --- /dev/null +++ b/internal_docs/scheduled-agent-wakeups.md @@ -0,0 +1,1725 @@ +# Scheduled Agent Wakeups and Background Turns + +## Status + +- **Date:** March 2026 +- **Status:** v1 design reviewed; initial implementation landed +- **Audience:** backend / agent / mobile product design +- **Purpose:** capture the agreed architecture for scheduled agent wakeups and record the initial implementation status + +This document started as a forward-looking design draft. It now also records the initial implementation that landed on 2026-03-23. + +Historical research notes are still preserved below. Any "open questions" inside the section 17 research log should be read as point-in-time notes from before implementation unless they are restated later as still-open items. + +## Implementation update (2026-03-23) + +The agreed v1 scheduler shape is now wired into the codebase. + +Implemented pieces: + +- Postgres-backed schedule definitions plus per-occurrence run records via `agent_schedules` and `agent_schedule_runs` +- Structured recurrence for `one_off`, `interval`, `daily`, and `weekly` schedules +- Timezone mode support for `follow_user` versus `fixed`, including recomputation when the user's timezone preference changes +- Agent tools for `schedule_task`, `list_schedules`, and `cancel_schedule` +- Background schedule worker startup in `main.rs` +- Lease-based execution with retries, heartbeat renewal, stale-expiry handling, and partial-output completion behavior +- Scheduled turns reusing the normal agent runtime for tool calls and assistant-message persistence, without inserting fake user messages +- Scheduled reminder push enqueueing with one encrypted preview per turn, using `llama-3.3-70b` summarization plus first-message fallback +- Coverage for recurrence parsing / calculation edge cases including DST-gap handling + +Still intentionally out of scope for v1: + +- direct client schedule-management APIs beyond the agent tools +- strict per-agent turn serialization beyond the accepted v1 parallel-turn tradeoff +- advanced calendar rules like third Monday / last business day + +--- + +## 1. Core Product Idea + +The key design decision is: + +**scheduled tasks should wake the agent, not directly send notifications.** + +That means a scheduled event is not fundamentally a push job. It is a future agent-triggered turn. + +The LLM remains the brain: + +- the agent decides whether to message the user at all +- the agent decides what tone and wording to use +- the agent may use normal tools before replying +- the agent may decide the event is no longer relevant and do nothing + +Push delivery is a downstream consequence of that background turn, not the primary product abstraction. + +This is the main behavioral difference from a traditional reminder system. We are not storing β€œnotification text to send later.” We are storing β€œinstructions for the future agent to consider when that time arrives.” + +--- + +## 2. Example Product Behavior + +### 2.1 Wake-up example + +User says: + +> I want a wake up notification at 8am tomorrow. + +Agent says: + +> Okay. + +The agent schedules a one-off event for the next day at 8am local time. The scheduled payload is **not** the notification text. It is a future-facing instruction for the agent, such as: + +> At 8am tomorrow, wake the user up. Check whether they have already been active in the conversation this morning before sending anything. If there is relevant recent context about why they need to wake up, use it. + +At 8am, the event fires and wakes the agent in the background. + +The agent then has normal context available: + +- recent thread history +- memory blocks +- archival memory +- conversation search +- normal tools +- current local time / timezone context + +At that point the agent might: + +- send nothing because the user already woke up and has been chatting since 7am +- send a personalized wake-up message based on newer context +- check the weather first, then send a more useful message +- send several messages if that is what the normal turn logic decides + +### 2.2 Weather-aware wake-up example + +User says: + +> Wake me up at 8 tomorrow and let me know if I need a coat. + +At 8am, the agent wakes, sees the scheduled instruction, calls `web_search`, gets the weather, and then sends: + +> Hey, it’s 30 degrees out. Wake up and put on a coat. + +This is exactly why the scheduled event should enter the normal agent runtime rather than bypassing it with prewritten push text. + +--- + +## 3. Primary v1 Decisions + +### 3.1 Best-effort agent wakeup + +The system is **best-effort**, not alarm-clock-grade guaranteed delivery. + +If the model stack is down or a dependency fails, the system should retry according to durable queue semantics. But if the full turn never succeeds, we accept that no user-visible notification may be delivered. + +This is consistent with the product philosophy: + +- the LLM is the decision-maker +- the scheduled event exists to wake the agent +- there is no separate deterministic fallback reminder engine in v1 + +### 3.2 Internal event, not fake user or tool-result persistence + +When a scheduled event fires, it should enter the runtime as a **new internal event type**. + +It should **not** be persisted as: + +- a fake `user_message` +- a fake `assistant_message` +- a fake `tool_output` +- a synthetic row in existing message tables pretending a tool call stayed open for hours + +The internal event should still reuse almost all of the normal runtime behavior, but it is not itself normal chat history. + +### 3.3 Reuse the normal agent runtime as much as possible + +The scheduled-event turn should reuse the same major machinery as a normal user-driven agent turn: + +- thread loading +- memory loading +- prompt assembly +- tool registry +- tool execution chain +- assistant message persistence +- post-turn delivery / push planning + +In other words, the internal trigger is new, but nearly all of the downstream execution model should stay aligned with the existing agent flow. + +### 3.4 Normal outputs are persisted normally + +The **results** of the scheduled turn should be treated as ordinary assistant behavior. + +If the agent sends messages, those messages should be stored as real assistant messages. + +If the agent calls tools, those tool calls and tool outputs should follow the normal persistence rules used by the existing runtime. + +If the agent sends three messages, then calls tools, then sends three more messages, that is allowed. The scheduled trigger should not artificially force a single-message reminder format. + +### 3.5 No explicit push tool + +There should be no β€œsend push notification” tool in this design. + +The agent should keep doing normal things: + +- send messages +- search +- consult memory +- reason over recent context + +If the scheduled background turn produces assistant messages and there is no live SSE consumer for that turn, then the delivery system should treat those outputs as eligible for push, subject to the user’s push settings. + +Push remains delivery plumbing, not agent-facing product logic. + +### 3.6 One push per scheduled turn by default + +For v1, scheduled background turns should usually produce **at most one push notification event**. + +Even if the agent emits multiple assistant messages during the turn, we should derive a single encrypted preview summary for notification delivery. + +This keeps scheduled events from producing noisy notification bursts while still preserving the full assistant output in conversation history. + +### 3.7 No separate notifier agent + +We should **not** create a second autonomous β€œnotifier” agent. + +Instead, after the main scheduled turn completes, we should run a lightweight **notification-preview composer** over the turn’s assistant outputs. + +That post-processing step is allowed to use a cheap model, but it is not another reasoning agent with independent product behavior. + +--- + +## 4. Scheduling Model + +### 4.1 Agent-scoped future wakeups + +Scheduled events should belong to an agent identity, not to a generic global notification system. + +That means the event targets: + +- the main agent, or +- a specific subagent + +When the event fires, that exact agent/thread is what wakes up. + +### 4.2 Store instructions for the future agent + +The stored payload should be agent-oriented instructions, not notification copy. + +Good examples: + +- β€œAt 8am tomorrow, wake the user up for yoga class. If they have already been active this morning, avoid a redundant wake-up message. If weather seems relevant, check it first.” +- β€œTomorrow at 6pm, remind the user to leave for dinner. If recent conversation changed the plan, adapt to the latest context.” + +Bad examples: + +- β€œWake up! Time for yoga!” +- β€œIt’s cold outside, bring a coat!” +- β€œReminder: leave now.” + +Those bad examples are candidate notification text, not the right stored instruction shape. + +### 4.3 One-off first, with limited recurring support in v1 + +The primary implementation milestone should still be **one-off scheduled events first**. + +However, based on current product discussion, the broader v1 design should also be able to support a **limited structured recurrence model** for common cases like: + +- every hour +- every morning at 8am +- every Friday at 5pm +- weekdays at 7:30am + +What remains out of scope is not recurrence altogether, but rather: + +- raw free-form cron syntax as the main product surface +- advanced calendar-style recurrences like β€œthird Monday” or β€œlast business day” + +### 4.4 Schedule time is authoritative; message wording is not + +The scheduler is responsible for firing the event at the right time. + +The agent is responsible for deciding whether and how to turn that event into user-visible output. + +This division of responsibility is important: + +- time semantics are deterministic infrastructure +- language semantics stay inside the agent runtime + +--- + +## 5. Tool Contract for Scheduling + +### 5.1 Agent-facing scheduling tool goal + +The scheduling tool should explicitly teach the model that it is writing instructions for its **future self**. + +The tool prompt / schema guidance should say that the agent must: + +- write in normal human language +- describe the goal and relevant context for the future turn +- mention conditional behavior when appropriate +- avoid drafting final notification copy +- avoid assuming the future state will match the current one exactly + +### 5.2 Recommended instruction-writing guidance + +The tool contract should push the agent toward guidance like this: + +> You are scheduling a future wakeup for yourself. Write instructions for the future agent turn, not the final notification text. Be specific about the goal, timing, relevant context, and any useful conditional behavior. Use normal natural language. Do not optimize for brevity if detail would help the future agent make a better decision. + +### 5.3 Suggested base scheduling arguments + +The exact schema can be finalized later, but the base shape is roughly: + +- `run_at`: absolute scheduled time +- `timezone`: optional user-local timezone override if needed +- `instruction`: detailed future-agent instruction in natural language +- `description`: short operational description for listing/debugging +- `expires_at`: optional staleness cutoff for cases where old reminders stop being useful + +The important field is `instruction`, and it should be modeled as the sensitive core payload. + +Recurring-specific structured fields are described later in section `19.15`. + +### 5.4 Example of a good scheduled instruction + +> At 8am local time tomorrow, wake the user up. They want this as a real wake-up nudge, not just a passive reminder. Before sending anything, consider whether they have already been active in chat this morning. If there is recent context about what they need to wake up for, use that context. If helpful, you may check weather or other relevant information before replying. + +This is the right shape because it tells the future agent what job it is trying to do without locking it into exact wording. + +--- + +## 6. Runtime Trigger Model + +### 6.1 Scheduled event wakeup path + +When a scheduled event becomes due: + +1. a worker leases the event row +2. the worker claims execution responsibility for that target agent/thread on that server +3. the runtime starts a background agent turn using a scheduled-event trigger +4. the runtime injects hidden event context +5. the agent executes its normal loop +6. normal outputs are persisted +7. post-turn delivery planning decides whether to enqueue push +8. the scheduled event row is completed or retried + +### 6.2 Hidden event context + +The runtime should inject hidden context that includes things like: + +- scheduled event id +- scheduled-for time +- actual fired-at time +- local timezone / local wall-clock representation +- the stored instruction +- optional source message reference or creation context + +This hidden context should behave similarly to other runtime-only inputs, but it should not create user-visible history rows by itself. + +### 6.3 Prompt shape + +The scheduled event should enter prompt assembly in a form that clearly communicates: + +- this is an internal scheduled wakeup +- the user did not just send a new live message right now +- the agent should consider the latest context before acting +- it may decide to no-op + +The system should make it easy for the agent to understand that the stored instruction came from an earlier planning moment and may need reinterpretation based on newer context. + +--- + +## 7. Delivery and Push Behavior + +### 7.1 Push is downstream of the background turn + +The scheduled event itself should not enqueue push directly. + +Instead: + +- the agent finishes its background turn +- the system inspects the assistant outputs from that turn +- if there are user-visible assistant messages and push is enabled, a notification event may be enqueued + +### 7.2 Notification preview generation + +For the encrypted notification preview text: + +- use the existing available **Llama** model as a compression/summarization step +- its job is to compress the turn’s assistant outputs into a short preview while preserving tone and intent +- it must not invent facts not present in the assistant outputs + +The summarizer should receive the assistant outputs from that scheduled turn and produce a compact 1-2 sentence preview suitable for encrypted notification display. + +### 7.3 Preview fallback behavior + +If the Llama summarizer fails, fall back to: + +- the **first assistant message** produced in that scheduled turn + +This fallback does not need to inspect which substep or tool boundary produced it. If it is the first normal assistant message from the turn, it is the fallback preview source. + +### 7.4 Preserve full assistant output in chat history + +The push preview is only a delivery summary. + +The full assistant outputs of the scheduled turn should still be persisted normally. The notification preview should not replace or reshape the underlying conversation history. + +### 7.5 One encrypted preview per scheduled turn + +If the agent produced multiple messages, the summarizer should compress across the whole turn and yield one preview. + +That preview becomes the encrypted notification text. The provider-visible shell can remain generic, consistent with the current encrypted push model. + +--- + +## 8. Retry, Leasing, and Recovery + +### 8.1 Durable queue semantics + +Scheduled events need real durable worker semantics, not an in-memory timer. + +At minimum this requires: + +- due-time polling +- row leasing +- lease expiry recovery +- explicit retry states +- backoff +- attempt counters +- terminal failure handling + +### 8.2 Use both lease expiry and explicit retry + +We want both failure mechanisms: + +1. **lease expiry** for crash recovery +2. **retry status with next attempt time** for known transient failures + +That means: + +- if a worker crashes mid-turn, another worker can recover the event after lease expiry +- if the runtime returns a transient failure, the current worker should explicitly schedule a retry with backoff + +### 8.3 Extend or heartbeat the lease while the turn runs + +Scheduled turns may take longer than a small lease interval, especially if they involve tool calls. + +The worker should therefore renew the lease while the turn is active so another worker does not mistakenly steal the same event during a long-running but healthy execution. + +### 8.4 Per-agent turn serialization + +The long-term safest model is still: + +- every agent has one serialized turn lane +- scheduled turns join that same lane +- user turns still take precedence where appropriate, but concurrency is controlled centrally + +This avoids context corruption, duplicate outputs, and confusing interleaving. + +However, the current v1 decision is more relaxed: + +- a scheduled wakeup may run with its own lifetime even if a live turn for that same agent is already in progress + +So serialization remains the preferred end-state, but it is **not** a blocker for the first implementation pass. + +### 8.5 Staleness and expiry + +Many scheduled events are time-sensitive. + +Examples: + +- a wake-up reminder is useful at 8:00am, much less useful at 2:00pm +- a β€œleave now for dinner” reminder may be stale shortly after it misses its target + +So scheduled events should support an explicit expiry or stale-after threshold. Once stale, the system should stop retrying and mark the event terminal. + +### 8.6 Best-effort failure model + +Because this feature is explicitly best-effort, there is no requirement in v1 to synthesize a system-generated reminder if the LLM never completes successfully. + +If all retries fail or the event expires, the event simply fails. + +--- + +## 9. Data Model Intent + +### 9.1 Separate scheduled-events table + +This feature should use a dedicated scheduled-events table rather than overloading `notification_events`. + +`notification_events` should remain the push outbox. + +The new scheduled-events table is the source of truth for: + +- when to wake the agent +- which agent to wake +- what instruction to inject +- current lease / retry / completion state + +### 9.2 Plaintext fields + +We should keep operational fields plaintext so PostgreSQL can index, filter, and lease them efficiently. + +That likely includes: + +- row id / uuid +- user id +- agent id +- conversation id or agent-thread reference if needed +- status +- run time +- next attempt time +- lease owner / lease expiry +- attempt count +- created at / updated at +- optional expiry time +- short operational description +- source message id references where useful + +These are infrastructure fields, not the sensitive semantic payload. + +### 9.3 Encrypted fields + +The main sensitive field is likely: + +- `instruction_enc` + +That is the actual future-agent instruction and is the most obvious content that should be encrypted at rest. + +We may also want an optional encrypted metadata blob for future-sensitive context if we later decide we need more than just the instruction text. + +### 9.4 Plaintext minimum principle + +We should keep only the minimum necessary fields encrypted. + +The goal is: + +- preserve queryability, scheduling efficiency, and operational simplicity in Postgres +- encrypt only the sensitive content-bearing parts of the event + +At the current intent level, that likely means **the instruction is the primary encrypted field**, while scheduling and operational metadata stay plaintext. + +### 9.5 Error storage should be sanitized + +Retry / failure diagnostics should avoid storing highly sensitive prompt or model output content in plaintext error columns. + +Plaintext `last_error` fields should stay operational and sanitized. + +--- + +## 10. Execution Semantics + +### 10.1 The scheduled event itself is not chat history + +The scheduled trigger should not appear in user-visible history as if the user just sent a new message. + +The user should see only the outputs the agent actually chose to send. + +### 10.2 The scheduled event can still lead to rich behavior + +Even though the trigger itself is hidden, the turn can still produce: + +- multiple assistant messages +- tool calls +- tool outputs +- follow-up reasoning using the normal runtime + +This is intentional. Scheduled turns should not be artificially crippled compared with ordinary turns. + +### 10.3 No-op is valid behavior + +The agent should be allowed to conclude that no user-visible action is needed. + +Examples: + +- the user is already awake and active +- the plan changed in later conversation +- the event is now irrelevant given newer context + +In those cases the event should complete successfully with no assistant message and no push. + +--- + +## 11. Notification Preview Composer + +### 11.1 Purpose + +The notification-preview composer exists only to create a compact encrypted push preview after a scheduled turn. + +It is not a new product-level agent. + +### 11.2 Inputs + +The composer should consume: + +- the assistant messages produced by the scheduled turn +- optional turn metadata like thread target or message ids + +It should not independently search, reason, or reinterpret the user’s world beyond compressing the already-produced assistant output. + +### 11.3 Model choice + +Use the existing available **Llama** model because it is already present and is a strong fit for compression / summarization. + +### 11.4 Behavior requirements + +The preview composer should: + +- preserve tone and intent +- stay faithful to the assistant outputs +- avoid adding new facts +- keep the result short +- target notification-preview length rather than full chat prose + +### 11.5 Deterministic fallback + +If the composer fails for any reason, the fallback is deterministic: + +- use the first assistant message from the turn as the encrypted preview text + +This keeps the delivery path robust without adding another complex branch. + +--- + +## 12. Security and Privacy Boundaries + +### 12.1 Sensitive data should remain encrypted + +The instruction payload is sensitive because it contains semantic user intent and planning context. + +That content should be encrypted at rest in the scheduled-events table. + +### 12.2 Operational metadata may remain plaintext + +We do not need to encrypt every field just because it belongs to a scheduled event. + +In particular, fields required for: + +- polling +- leasing +- retry management +- indexing +- expiry checks +- list views + +should remain plaintext unless they themselves carry sensitive user content. + +### 12.3 No plaintext instruction logging + +The system should avoid logging: + +- decrypted scheduled instructions +- raw notification preview plaintext +- full assistant turn content in operational logs + +Logs should stay operational rather than content-rich. + +### 12.4 Push encryption remains unchanged in principle + +Scheduled turns should continue to use the same encrypted-push architecture already adopted for mobile push delivery. + +The only new part is how the preview text is sourced: + +- from the assistant outputs of the scheduled turn, +- compressed by Llama, +- with first-message fallback. + +--- + +## 13. Product Scope for the First Implementation Pass + +### In scope + +- one-off scheduled agent wakeups +- limited structured recurring schedules for common cases (interval / daily / weekly) +- agent-scoped scheduled events for main agent and subagents +- encrypted-at-rest scheduled instructions +- durable leasing and retries +- hidden internal event trigger into the normal agent runtime +- normal persistence of assistant outputs +- one encrypted push preview derived after turn completion +- Llama preview compression with first-message fallback +- no-op outcomes when the agent decides nothing should be sent + +### Out of scope for the first pass + +- raw cron as the primary product/API surface +- advanced calendar-style recurrences (e.g. third Monday, last business day) +- a separate notifier agent +- a user-facing generic push tool +- persisting the internal scheduled trigger into chat history +- guaranteed alarm-clock delivery semantics +- a deterministic system-generated reminder fallback when the LLM never succeeds + +--- + +## 14. Main Rejected Alternatives + +### 14.1 Scheduled events directly enqueue notifications + +Rejected because it prevents the agent from adapting to newer context and reduces reminders to prewritten notification copy. + +### 14.2 Persist the fired event as a fake user message + +Rejected because it pollutes chat history and misrepresents what actually happened. + +### 14.3 Persist the fired event as a delayed tool result + +Rejected because it is structurally misleading and creates awkward long-lived pseudo-tool state. + +### 14.4 Second autonomous notifier agent + +Rejected because it adds a second reasoning layer when a lightweight preview composer is enough. + +### 14.5 Explicit push-notification tool + +Rejected because push should remain a delivery concern, not a first-class agent product surface. + +--- + +## 15. Open Questions for the Next Pass + +These are intentionally left for a later code-aware design pass: + +- exact table schema and naming +- exact worker / lease heartbeat implementation +- exact agent-turn locking strategy relative to live user turns +- whether `conversation_id` should be stored on the scheduled row or derived from `agent_id` +- how user-visible schedule listing / cancellation UX should work +- how much scheduling metadata should be exposed to clients +- whether multiple due scheduled events for one agent should coalesce before execution +- whether future versions should support richer recurrence families or additional event classes beyond the initial interval/daily/weekly set + +Note: the default stale / expiry policy for v1 is now set to **15 minutes**; what remains open is whether some schedule families should override that by default. + +--- + +## 16. Bottom Line + +The intended architecture is: + +**schedule now -> wake the agent later -> let the agent decide what to do -> optionally push the result** + +That preserves the core Maple product idea: + +- the LLM is the brain +- reminders are context-aware agent behaviors, not canned notifications +- push is a downstream transport layer +- scheduled instructions are written for the future agent, not for the lock screen + +This should be the starting point for the implementation design pass. + +--- + +## 17. Code Research Log + +This section is intentionally incremental. It captures code-informed findings as they are discovered, rather than waiting for one final summary. + +### 17.1 Research pass 1 β€” current live agent turn flow and push hook + +Relevant files reviewed in this pass: + +- `src/web/agent/mod.rs` +- `src/web/agent/runtime.rs` +- `src/push/mod.rs` +- `src/models/notification_events.rs` +- `src/models/notification_deliveries.rs` +- `src/push/worker.rs` + +#### What the code does today + +The current `agent.message` push path is tightly coupled to the live SSE chat flow. + +In `src/web/agent/mod.rs`, `run_agent_chat_task(...)` currently does this: + +1. initialize `AgentRuntime` +2. call `runtime.prepare(&input_content)` +3. run `runtime.step(...)` in a loop +4. persist any tool call + tool output rows +5. persist assistant messages with `runtime.insert_assistant_message(...)` +6. attempt to stream each assistant message over SSE +7. if SSE delivery fails, remember only the **first persisted-but-undelivered** assistant message +8. after the turn completes, enqueue **one** push notification from that first missed assistant message + +This means the current implementation already has the important product behavior of: + +- persisting assistant output first +- treating push as downstream delivery +- sending at most one push per interrupted turn + +That is directionally aligned with the scheduled-wakeup design. + +#### Important code-level constraints discovered + +##### `prepare()` currently persists a user message + +`AgentRuntime::prepare(...)` in `src/web/agent/runtime.rs` currently inserts a real `user_messages` row and schedules embedding work. + +That is correct for live user chat, but it is **not** correct for scheduled wakeups if the internal scheduled trigger must stay hidden. + +This strongly suggests the scheduled-turn implementation will need either: + +- a parallel `prepare_internal_event(...)` path, or +- a refactor that separates β€œbuild step input” from β€œpersist user message.” + +##### Assistant/tool persistence is already reusable + +`AgentRuntime` already exposes reusable persistence helpers for the outputs we do want: + +- `insert_assistant_message(...)` +- `insert_tool_call_and_output(...)` + +Those helpers also already spawn background embedding work for assistant messages, which is good: scheduled turns should benefit from the same memory/indexing behavior as ordinary turns. + +##### Current push enqueue expects one `message_id` and one text body + +`enqueue_agent_push_for_disconnect(...)` calls `enqueue_agent_message_notification(...)`, which takes: + +- a single `message_id` +- a single `message_text` + +That is sufficient for the current disconnect flow, but it is narrower than the scheduled-turn design, where one push preview may need to summarize **multiple** assistant messages from one background turn. + +##### The push outbox is already durable and leased + +The push system already has a strong distributed-worker pattern: + +- durable `notification_events` +- per-device `notification_deliveries` +- leasing via `FOR UPDATE SKIP LOCKED` +- explicit retry transitions +- lease ownership checks on writeback +- expiry/cancellation handling + +This is a strong conceptual template for scheduled-event execution, even though the scheduled-event table and worker do not exist yet. + +#### Early implications for the feature + +- scheduled wakeups should likely reuse the existing assistant/tool persistence paths +- scheduled wakeups should **not** reuse the live `prepare(...)` path unchanged +- push for scheduled turns should likely happen **after the whole background turn finishes**, not on the first produced assistant message +- the current push helper shape probably needs either a lower-level enqueue path or a new scheduled-turn-specific wrapper + +#### Open questions discovered in this pass + +- What is the cleanest hidden-trigger entrypoint into `AgentRuntime`? +- Should scheduled turns share most of `run_agent_chat_task(...)` through an extracted non-SSE turn runner? +- What should `message_id` mean when one push preview summarizes several assistant messages from one turn? +- How do we make one-push-per-turn idempotent across scheduled-worker retries? +- I did **not** find an obvious existing per-agent turn lock in the reviewed code, so where should scheduled-turn serialization live relative to live user turns? + +### 17.2 Research pass 2 β€” context assembly, model hooks, and concurrency gaps + +Relevant files reviewed in this pass: + +- `src/web/agent/runtime.rs` +- `src/web/agent/signatures.rs` +- `src/web/agent/compaction.rs` +- `src/proxy_config.rs` +- `src/web/openai.rs` +- `src/web/responses/handlers.rs` +- `src/main.rs` + +#### What the runtime already gives us for a scheduled turn + +`AgentRuntime::build_context(...)` already assembles most of the context we want a scheduled wakeup to inherit. + +It currently loads: + +- the user’s timezone from `user_preferences` +- the user’s locale from `user_preferences` +- a formatted `current_time` string localized to that timezone +- decrypted persona + human memory blocks +- memory metadata counts +- the latest conversation summary +- recent conversation history +- main-agent onboarding state +- subagent purpose + +This is important because it means a scheduled wakeup does **not** need a special prompt stack to feel β€œnormal.” Once it reaches the runtime correctly, it can already inherit most of the same state as a live user turn. + +#### Tool-result continuation behavior is already encoded in the runtime + +`AgentRuntime::step(...)` already has a continuation mode for follow-up steps after tools run. + +It tracks: + +- `current_tool_results` +- `previous_step_summary` + +and then injects tool-result instructions like: + +- this is a continuation of the previous turn +- previous messages are already visible to the user +- silence is the default unless there is genuinely new information + +This does **not** directly implement scheduled events, but it shows the runtime already understands β€œnon-fresh-user-message continuation” as a first-class pattern. + +That is encouraging for the scheduled-event design, which also wants a hidden internal trigger rather than a fake new user message. + +#### `prepare()` currently does more than validation + +This pass confirmed that `prepare(...)` currently does all of the following together: + +- validates / normalizes the user message +- runs vision pre-processing when needed +- clears tool-result state +- inserts the user message +- triggers background embedding +- runs compaction checks + +This means a hidden scheduled-event entrypoint will need a deliberate decision about which of those behaviors to keep. + +Most likely: + +- keep tool-result reset +- probably keep compaction behavior +- skip user-message insertion +- skip any semantics that assume a new live user-authored message exists + +#### LLM helper infrastructure is already flexible enough for a preview composer + +There are two useful existing LLM helper patterns in the codebase. + +##### Pattern A: direct small non-streaming helper calls + +`src/web/responses/handlers.rs::spawn_title_generation_task(...)` already runs a background helper call using: + +- model: `llama-3.3-70b` +- `get_chat_completion_response(...)` +- non-streaming completion extraction +- best-effort logging-only failure behavior + +This is a strong template for a post-turn preview-composition step. + +##### Pattern B: typed DSRS summarization with correction + +`src/web/agent/compaction.rs` already implements typed summarization with correction retries. + +That path uses: + +- `build_lm(...)` from `src/web/agent/signatures.rs` +- typed signatures +- malformed-output correction + +This means we already have a reusable pattern if we decide the notification-preview composer should be typed rather than raw string extraction. + +#### Llama availability is real, but conditional + +`src/proxy_config.rs` confirms that `llama-3.3-70b` is a routed canonical model in this repo. + +However, it is only available when `TINFOIL_API_BASE` is configured. Without Tinfoil, the router falls back to Continuum-only models and Llama is not included. + +That means the current product decision of β€œuse Llama for preview composition” is viable, but there is an implementation question about what to do in environments where Tinfoil is absent. + +#### Current agent model defaults + +The main agent runtime currently builds its LM with: + +- request model: `kimi-k2-5-agent` +- billing model: `kimi-k2-5` + +So the preview composer would be a genuine secondary helper model path, not just a reuse of the main-agent LM. + +#### I still did not find an agent turn lock + +I searched the agent runtime and broader app state for agent/conversation turn serialization and did **not** find an obvious existing lock, semaphore, or in-flight turn registry for agent chat. + +This remains one of the most important implementation gaps discovered so far. + +#### Additional structural finding for subagents + +`AgentRuntime::new_subagent(...)` already reconstructs the runtime from the subagent UUID by loading: + +- the subagent row +- its conversation +- its decrypted purpose + +That means scheduled events targeting subagents can likely anchor on subagent identity cleanly without needing to persist every prompt-level field on the scheduled row. + +#### Open questions discovered in this pass + +- Should the scheduled-event row store `agent_id`, `agent_uuid`, `conversation_id`, or some combination? +- Should the scheduled-event entrypoint run compaction before the turn, the same way live `prepare(...)` does? +- Should the preview composer use the lightweight raw-completion pattern like title generation, or a typed DSRS summarizer with correction? +- What should happen in environments where `llama-3.3-70b` is unavailable because Tinfoil is not configured? +- Where should agent-turn serialization live if the current runtime has no obvious lock/queue for it? + +### 17.3 Research pass 3 β€” agent signature shape and queue-schema implications + +Relevant files reviewed in this pass: + +- `src/web/agent/signatures.rs` +- `src/web/agent/runtime.rs` +- `migrations/2026-03-07-120000_push_notifications_v1/up.sql` + +#### Current agent signature shape is input-centric + +The main agent signature in `src/web/agent/signatures.rs` currently takes these top-level inputs: + +- `input` +- `current_time` +- `agent_kind` +- `subagent_purpose` +- `persona_block` +- `human_block` +- `memory_metadata` +- `previous_context_summary` +- `recent_conversation` +- `available_tools` +- `is_first_time_user` + +There is currently **no explicit field for trigger kind** such as: + +- live user message +- tool-result continuation +- scheduled internal event + +So if we want scheduled events to be a first-class trigger type at the prompt layer, we will likely need one of two approaches: + +1. encode the scheduled-event semantics entirely inside the `input` string, or +2. extend the agent signature with an explicit field such as `turn_source`, `event_kind`, or equivalent + +This is an important design choice because it affects how visible and stable the scheduled-event concept becomes in prompt assembly. + +#### The runtime already supports first-step raw string input + +`AgentRuntime::step(...)` accepts a plain `user_message: &str` for the first step after `prepare(...)` returns the normalized text. + +That means a scheduled internal event could probably enter the runtime as a first-step string without needing to fake a `MessageContent` object. + +So there is already a plausible implementation path where the scheduled-event worker: + +- skips `prepare(...)` +- builds a hidden event input string +- feeds that directly into `step(..., is_first_step = true)` + +Whether that is the final shape depends on how much explicit structure we want in the signature. + +#### Retry and correction behavior already exists for agent outputs + +`call_agent_response_with_retry_and_correction(...)` in `src/web/agent/signatures.rs` already provides: + +- up to 3 LM attempts +- correction on parse failures +- a fallback correction prompt that reshapes malformed outputs rather than generating new content + +This means scheduled turns would inherit the same output-hardening behavior as live turns if they reuse the same signature path. + +That is useful for the best-effort design: the scheduled worker should not need a separate output-repair mechanism if it is truly using the same runtime. + +#### Push schema confirms a separation we should preserve + +The push migration confirms that `notification_events` is structurally a delivery/outbox table, not a scheduler table. + +Notable fields: + +- `payload_enc` +- `not_before_at` +- `expires_at` +- `cancelled_at` +- one delivery row per device with lease/retry state + +This reinforces the design decision already captured earlier: + +- `notification_events` should remain the push outbox +- scheduled-agent wakeups should have their own source-of-truth table + +Even though `notification_events.not_before_at` exists, it is the wrong abstraction for β€œwake the agent and let it decide what to do.” + +#### New open questions from this pass + +- Is it better to add an explicit `event_kind` / `turn_source` field to the agent signature now, or keep scheduled-event semantics hidden in the first-step input text? +- If scheduled turns bypass `prepare(...)`, what is the cleanest place to run the subset of prepare-time behavior we still want, like compaction and tool-state reset? +- Should the scheduled-event worker persist a turn-level execution record for idempotency / observability, even if the trigger itself stays out of chat history? + +### 17.4 Research pass 4 β€” preview payload and delivery-shape details + +Relevant files reviewed in this pass: + +- `src/push/mod.rs` +- `src/web/agent/mod.rs` + +#### The current encrypted preview payload is message-centric + +`NotificationPreviewPayload` currently contains: + +- `notification_id` +- `message_id` +- `kind` +- `title` +- `body` +- `deep_link` +- `thread_id` +- `sent_at` + +This is important for scheduled turns because the current payload model assumes a notification is anchored to **one message id**. + +That works naturally for the current disconnect flow, where the preview is derived from one first-undelivered assistant message. + +It is a bit more awkward for scheduled turns, where the preview may summarize an entire background turn. + +#### Current preview text budget is very small + +`src/push/mod.rs` sets: + +- `PUSH_PREVIEW_BODY_MAX_BYTES = 180` + +and `normalize_preview_body(...)` currently: + +- collapses whitespace +- truncates to that byte budget +- adds an ellipsis when needed + +This means the Llama preview composer should target a genuinely short result. Even if we ask for 1-2 sentences, the implementation still needs to respect this small body budget. + +#### Deep-linking already distinguishes main vs subagent + +The current push helper already derives: + +- main agent deep link: `opensecret://agent` +- subagent deep link: `opensecret://agent/subagent/` + +and matching `thread_id` values. + +That is useful because a scheduled-turn push can likely reuse the same targeting rules without inventing a new notification routing system. + +#### Current scheduled-turn preview implication + +For scheduled turns, if we keep the existing preview payload shape, then we probably need to choose a canonical message id for the turn, most likely: + +- the first assistant message of the scheduled turn, or +- the last assistant message of the scheduled turn + +The first-message option aligns better with the already-planned fallback behavior and with the current disconnect path, which also anchors on the first user-visible missed message. + +#### Additional open questions from this pass + +- Should scheduled-turn notifications keep the existing message-centric payload shape and simply use the first assistant message as the canonical anchor? +- Do we want a future turn-level identifier in preview payloads, or is that unnecessary if one notification maps to one scheduled turn? +- Should the preview composer enforce a tighter output limit than β€œ1-2 sentences” so it naturally fits the 180-byte cap? + +### 17.5 Research pass 5 β€” concrete worker numbers and missing lease-heartbeat support + +Relevant files reviewed in this pass: + +- `src/push/worker.rs` +- `src/models/notification_deliveries.rs` + +#### Current push worker defaults + +The existing push worker uses these concrete values: + +- batch size: `32` +- lease TTL: `60` seconds +- poll interval: `3` seconds +- max concurrency: `8` +- max attempts: `8` + +Retry backoff is exponential and currently caps at `480` seconds: + +- attempt 1 -> `15s` +- attempt 2 -> `30s` +- attempt 3 -> `60s` +- attempt 4 -> `120s` +- attempt 5 -> `240s` +- attempt 6+ -> `480s` + +These values are useful reference points for scheduled-event execution, but scheduled agent turns will probably need longer-running lease behavior than push deliveries do. + +#### Important limitation: the existing worker does not heartbeat or extend leases + +The push worker leases a delivery row once and then processes it. There is no lease-renewal path because provider send attempts are expected to be short. + +That is likely **not** enough for scheduled agent turns, which may: + +- call tools +- wait on network operations +- run multiple LLM steps +- take meaningfully longer than a push send + +So the scheduled-event worker probably needs one of these patterns: + +- an explicit lease heartbeat / renewal update, or +- a substantially longer lease window plus separate stuck-turn handling + +The first option still looks cleaner. + +#### No advisory-lock pattern found in current code + +I searched for advisory-lock / turn-lock patterns and did **not** find existing Postgres advisory-lock usage or an equivalent in-memory agent turn registry in the current agent code. + +This strengthens the earlier conclusion that scheduled turns will likely need a brand-new serialization mechanism rather than plugging into an existing one. + +#### Additional open questions from this pass + +- Should scheduled-event execution use the same general lease-state machine shape as push deliveries, but with a lease-renewal API added? +- What should the initial lease TTL be for a scheduled turn, given that tool-augmented agent turns can be much longer than provider sends? +- Should turn serialization be implemented at the database level, in app memory, or both? + +### 17.6 Research pass 6 β€” current agent route surface + +Relevant files reviewed in this pass: + +- `src/web/agent/mod.rs` + +#### Current agent API surface + +The current agent router exposes: + +- `/v1/agent` +- `/v1/agent/init` +- `/v1/agent/items` +- `/v1/agent/chat` +- `/v1/agent/subagents` +- `/v1/agent/subagents/:id/chat` +- related item and delete routes + +There is currently **no** schedule-specific API surface yet for: + +- create scheduled task +- list scheduled tasks +- cancel scheduled task + +That is expected, but it is now confirmed from code. + +#### Current routing is gated to local/dev + +The current agent router is only enabled when `app_mode` is `Local` or `Dev`. + +That means any initial schedule API work in this area will naturally follow the same experimental surface unless we intentionally broaden that exposure later. + +#### Implication for the feature + +If we introduce schedule endpoints as part of the agent API, the most natural home appears to be under the existing agent surface, likely alongside main-agent and subagent routes rather than under push routes. + +That fits the product model already captured earlier: scheduled events are agent behavior, not push-device plumbing. + +#### Open questions from this pass + +- Should scheduling routes live under `/v1/agent/...` and `/v1/agent/subagents/:id/...`, or should creation happen only via agent tool use at first? +- Do we want a direct client-management surface for list/cancel in v1, or is the first pass tool-driven only? +- If these remain under the current experimental agent router, is that sufficient for the intended rollout path? + +--- + +## 18. Decisions Captured During Review + +These are user-confirmed decisions made after the research pass. + +### 18.1 Concurrency with live turns in v1 + +For v1, a scheduled wakeup may run with its **own lifetime** even if that same agent is already mid-turn on a live conversation. + +That means: + +- we are **not** blocking v1 on solving agent-turn merging +- we are **not** blocking v1 on building a strong per-agent serialization mechanism first +- parallel execution for this edge case is acceptable in v1 + +This should be treated as an explicit temporary product tradeoff, not an accidental bug. + +### 18.2 Partial-failure policy after user-visible output + +Normal completion-level retries inside the agent runtime can stay as they are. + +However, if a scheduled turn has already produced and persisted user-visible assistant output, and then later fails after exhausting its normal retries, v1 should: + +- keep the already-produced output +- treat the scheduled event as effectively **completed** +- **not** retry the entire scheduled event from the top + +Rationale: + +- retrying after user-visible output risks duplicate or confusing messages +- preserving partial useful output is better than replaying the whole scheduled wakeup + +This implies an important execution rule for v1: + +- scheduled-event retries are appropriate when the turn remains effectively silent +- once the user has already received meaningful output from that scheduled turn, the system should prefer completion over replay + +### 18.3 Default expiry policy + +The default stale/expiry cutoff for scheduled wakeups in v1 should be: + +- **15 minutes** + +This is especially appropriate for wake-up-style events where usefulness drops quickly after the intended time. + +### 18.4 Client/API surface for v1 + +The first pass should be: + +- **tool-driven only** + +That means we do **not** need a broad direct schedule-management API surface in the initial implementation just to validate the scheduled-agent-wakeup model. + +--- + +## 19. Recurring Schedule Model + +Recurring schedules are importantly different from one-off wakeups. + +For a one-off event, the system stores one future wakeup and runs it once. + +For a recurring event like: + +> Every morning at 8am, look up the weather and tell me. + +the system should **not** be modeled as an infinite stream of pre-created future tasks. Instead, it should be modeled as a durable recurring schedule definition that produces concrete occurrences over time. + +### 19.1 Do not use literal OS cron jobs + +We should not depend on host-level cron or per-user OS cron entries. + +This should stay an application-level scheduler with Postgres as the durable source of truth, just like the broader scheduled-agent-wakeup design. + +The recurrence engine may use cron-like logic internally, but the product abstraction should remain: + +- store schedule definition +- compute next due occurrence +- wake the agent for that occurrence +- update the schedule to the next occurrence + +### 19.2 Two-layer model: schedule definition plus occurrences + +The clean recurring model is: + +1. **Recurring schedule definition** + - which agent it belongs to + - recurrence rule + - timezone + - encrypted future-agent instruction template + - default expiry / stale window + - active / cancelled state + +2. **Concrete occurrence / run** + - this specific 8am firing + - scheduled-for timestamp + - lease / retry state + - execution result / terminal status + +This matters for several reasons: + +- recurring schedules need stable long-lived identity +- each occurrence needs its own retry/lease behavior +- occurrence-level state should not mutate the canonical recurrence definition +- observability and idempotency are easier when each firing is concrete + +### 19.3 Recurrence should be stored in local-time semantics + +For human schedules like: + +- every day at 8am +- weekdays at 7:30am +- every Monday at 9am + +the important semantic is **wall-clock local time**, not a fixed UTC interval. + +So the durable schedule definition should preserve: + +- the user’s IANA timezone +- the local recurrence rule + +Example: + +- timezone: `America/Chicago` +- recurrence: daily at `08:00` + +The resulting UTC timestamp will naturally change across DST transitions, but the user still experiences β€œ8am every morning,” which is what we want. + +### 19.4 Prefer structured recurrence over raw cron as the primary product shape + +For v1, I would strongly prefer a structured recurrence model over exposing raw cron expressions to the agent or user. + +For example, represent recurring schedules as fields like: + +- frequency: daily / weekly +- local_time: `08:00` +- weekdays: optional set +- timezone: IANA timezone + +rather than asking the model to write raw cron strings. + +Reasons: + +- easier for the agent to author correctly +- easier to validate +- clearer DST behavior +- fewer parsing edge cases +- better product ergonomics for later client surfaces + +If needed, we can still compile that structured rule into cron-like evaluation logic internally. + +### 19.5 Compute the next occurrence from the schedule, not from β€œnow” + +The next firing time should be computed from the schedule definition and the last scheduled occurrence, not by repeatedly adding a fixed duration to the current time. + +This avoids drift. + +For example, β€œevery morning at 8am” should not gradually wander because a worker happened to process yesterday’s run at 8:02am. + +The schedule should still target the next proper local 8:00am occurrence. + +### 19.6 Materialize occurrences one at a time + +We should not eagerly create months of future runs. + +Instead, the scheduler should typically: + +- maintain the next due occurrence for the recurring schedule +- when that occurrence is claimed, create or mark a concrete run +- after terminal completion / skip, compute and persist the next due occurrence + +This keeps state smaller and makes edits/cancellation much simpler. + +### 19.7 Recurring schedules should still wake the agent, not send canned reminders + +The recurring model does **not** change the core product philosophy. + +Each occurrence should still wake the agent as a normal background turn. + +So for: + +> Every morning at 8am, look up the weather and tell me. + +the recurring schedule should store a future-agent instruction template like: + +> Each morning at 8am local time, check the weather for the user and send a concise useful morning weather update. If there is newer context that changes what would be helpful, adapt to it. + +At each occurrence, the runtime should inject occurrence-specific metadata like: + +- this occurrence was scheduled for 2026-03-24 08:00 America/Chicago +- this is part of a recurring schedule +- current local time is ... + +The agent then decides what to do with current context just like a one-off scheduled wakeup. + +### 19.8 Backfill policy matters a lot for recurring schedules + +Recurring schedules raise an important question: + +If the system is down or delayed, should it run missed occurrences later? + +For the kind of recurring schedules we are discussing, the default v1 answer should be: + +- **no large backfills by default** + +Instead: + +- if the occurrence is still inside its stale window, it can run late +- if it is past the stale window, mark that occurrence skipped/expired and move on to the next recurrence + +This prevents nonsense behavior like sending several stale morning weather notifications after the fact. + +### 19.9 DST and ambiguous-time policy must be explicit + +Recurring schedules have timezone edge cases that one-off schedules mostly avoid. + +At minimum, the implementation needs explicit behavior for: + +- daylight saving time spring-forward gaps +- daylight saving time fall-back repeated hours + +For common cases like 8am daily reminders, the behavior is straightforward: fire at local 8am and let the UTC time shift with DST. + +For rarer ambiguous times, we should document a policy rather than leave it accidental. + +### 19.10 Cancellation semantics + +Cancelling a recurring schedule should mean: + +- stop generating future occurrences +- cancel any pending unstarted occurrence rows if they already exist +- do not affect already completed historical runs + +This is another reason the schedule-definition / occurrence split is helpful. + +### 19.11 Relationship to v1 tool-driven scope + +Even if v1 is tool-driven only, recurring schedules still need enough internal structure for the agent to create them safely. + +That means the schedule tool should be able to distinguish: + +- one-off wakeup +- recurring wakeup + +without forcing the agent to encode every recurrence as a raw cron string. + +### 19.12 Recommended v1 stance on recurrence + +Given the current product discussion, the safest v1 recurrence shape is: + +- daily / weekly recurring schedules only +- timezone-aware local-time execution +- structured recurrence fields, not raw cron strings +- one concrete occurrence at a time +- stale occurrences skipped instead of heavily backfilled +- same background-agent execution model as one-off events + +That would cover common product cases like: + +- every morning at 8am, send weather +- weekdays at 7am, remind me to wake up +- every Sunday evening, remind me to plan the week + +without taking on the full complexity of arbitrary cron semantics immediately. + +### 19.13 Distributed scheduling across multiple OpenSecret servers + +This system must assume that multiple OpenSecret servers may be running the same code concurrently against the same Postgres database. + +That means recurring scheduling must be **database-coordinated**, not process-local. + +Important implications: + +- no server-local in-memory timers as the source of truth +- no assumption that one server β€œowns” a user’s schedules +- no schedule state that only exists in process memory + +Postgres must remain the durable coordination plane. + +#### Recommended distributed model + +For both one-off and recurring schedules: + +- workers on any server may poll for due rows +- claiming work must happen through row leasing / transactional state changes +- lease ownership must be checked again when writing terminal results + +For recurring schedules specifically, there are two related coordination problems: + +1. **Who claims the due occurrence for execution?** +2. **Who advances the recurring schedule to its next occurrence?** + +Both need to be safe under concurrent workers on different servers. + +#### Suggested safety pattern for recurring schedules + +The safest shape is: + +- keep the recurring schedule definition row in Postgres +- when an occurrence is due, claim/update it transactionally +- generate a concrete occurrence/run row with a unique occurrence key +- advance the schedule definition to the next due local-time occurrence in the same transaction, or under the same row lock + +To avoid duplicates, recurring occurrences should have a uniqueness guarantee like: + +- unique `(schedule_id, scheduled_for_at)` + +or another stable occurrence identity derived from the recurrence. + +That way, even if two servers race, the database prevents duplicate occurrence creation. + +#### Reuse the existing push-worker philosophy + +The current push system already uses the right distributed philosophy: + +- `FOR UPDATE SKIP LOCKED` +- lease owner fields +- lease expiry +- compare-and-set style writeback + +Scheduled-agent wakeups should follow the same philosophy, even if the exact schema differs. + +#### Important consequence for future turn serialization + +Because multiple OpenSecret servers may execute turns, any future β€œone active turn per agent” policy cannot rely only on in-memory locks. + +If we later tighten serialization beyond the v1 parallel-turn tradeoff, that coordination will need to be database-backed or otherwise distributed-safe. + +### 19.14 Timezone-following versus fixed-timezone schedules + +User timezone updates matter a lot for human schedules like: + +- wake me up at 8am +- every morning at 8am, tell me the weather +- remind me at 6pm to leave + +These often mean: + +- β€œ8am where I am” + +not: + +- β€œ8am forever in the timezone I happened to be in when I created the schedule” + +So the scheduler should support two distinct timezone behaviors: + +1. **Floating / follow-user-timezone schedules** + - the schedule follows the user’s current timezone preference + - best default for personal routine reminders and wake-up-style schedules + +2. **Fixed-timezone schedules** + - the schedule stays anchored to a specific IANA timezone + - appropriate when the user explicitly anchors it, e.g. β€œ9am Eastern” + +This distinction is standard and important. Calendaring systems often make the same separation between floating local-time events and explicitly zoned events. + +#### Recommended product default + +For the kind of personal schedules we are discussing, the default should usually be: + +- **follow the user’s timezone preference** + +unless the user or agent explicitly anchors the schedule to a fixed timezone. + +#### What the timezone update tool should do + +If the user’s timezone preference is updated, the system should proactively recompute future schedules that are marked as: + +- active, and +- follow-user-timezone + +It should **not** touch: + +- historical completed occurrences +- cancelled schedules +- schedules explicitly pinned to a fixed timezone + +Because the design already prefers one concrete occurrence at a time, this recalculation should stay relatively cheap: + +- update the schedule definition’s `next_due_at` +- update or replace any unclaimed pending future occurrence for that schedule if one already exists + +There is no need to rewrite a large history of future rows if we keep materialization bounded. + +#### Important schema implication + +For follow-user-timezone schedules, we cannot store only an absolute UTC timestamp and expect timezone changes to work correctly. + +We also need to preserve the **local-time intent**, for example: + +- local recurrence rule +- local clock time +- timezone mode (`follow_user` vs `fixed`) +- fixed timezone value when applicable + +Then `next_due_at` becomes a computed operational field derived from that intent. + +#### One-off schedules + +This distinction also applies to one-off schedules. + +If the product wants a one-off β€œ8am tomorrow” wakeup to follow the user when they travel, then that one-off event also needs to preserve local-time intent rather than only a resolved UTC timestamp. + +If instead a one-off is meant to stay pinned to the original timezone, it should be marked fixed. + +So the schedule tool should ideally decide this explicitly when creating the event, rather than leaving the system to guess later. + +### 19.15 Agent-friendly v1 recurrence parameters + +The scheduling interface should be simple enough that the agent can reliably map common user intent into correct parameters without being forced to invent raw cron syntax. + +The right v1 model is not β€œfree-form cron strings.” It is a small structured recurrence vocabulary that covers the common cases well. + +#### Recommended v1 recurrence families + +Support these recurring families in v1: + +1. **Interval** + - example: every hour + - fields like: + - `every_n` + - `unit` = `hours` + +2. **Daily wall-clock** + - examples: every morning at 8am, every day at 6pm + - fields like: + - `local_time` + - timezone behavior + +3. **Weekly wall-clock** + - examples: every Friday at 9am, weekdays at 7:30am + - fields like: + - `weekdays` + - `local_time` + - timezone behavior + +This already covers most of the common recurring schedules we care about. + +#### Example parameter shape + +Conceptually, a tool/API shape like this is enough for v1: + +- `schedule_kind`: `one_off` | `recurring` +- `instruction`: future-agent instruction text +- `description`: short operational summary +- `timezone_mode`: `follow_user` | `fixed` +- `fixed_timezone`: optional IANA timezone when pinned +- `stale_after_minutes`: optional override + +For one-off: + +- `local_date` +- `local_time` + +For recurring interval: + +- `recurrence_type`: `interval` +- `every_n` +- `interval_unit`: `hours` + +For recurring daily: + +- `recurrence_type`: `daily` +- `local_time` + +For recurring weekly: + +- `recurrence_type`: `weekly` +- `weekdays` +- `local_time` + +That is much easier for the agent to use correctly than arbitrary cron. + +#### How common examples map + +- β€œevery morning at 8am” + - recurring + - daily + - `local_time = 08:00` + - `timezone_mode = follow_user` + +- β€œevery Friday at 5pm” + - recurring + - weekly + - `weekdays = [friday]` + - `local_time = 17:00` + +- β€œevery hour” + - recurring + - interval + - `every_n = 1` + - `interval_unit = hours` + +- β€œevery weekday at 7:30am” + - recurring + - weekly + - `weekdays = [monday, tuesday, wednesday, thursday, friday]` + - `local_time = 07:30` + +#### Important distinction: interval versus wall-clock recurrence + +β€œEvery hour” is fundamentally different from β€œevery day at 8am.” + +- interval schedules are duration-based +- daily/weekly schedules are wall-clock based + +That distinction should be explicit in the schedule definition, because it affects DST handling, drift, and how the next occurrence is computed. + +#### Is this standard practice? + +Yes. + +It is very normal to: + +- own the schedule schema and distributed execution model +- use a limited structured recurrence model for product ergonomics +- defer complex calendar recurrences until later +- rely on a recurrence library internally rather than exposing raw recurrence syntax directly + +This is a sound and scalable engineering approach, especially for a multi-server application using Postgres as the coordination plane. diff --git a/migrations/2026-03-23-180000_agent_schedules_v1/down.sql b/migrations/2026-03-23-180000_agent_schedules_v1/down.sql new file mode 100644 index 00000000..213f4b73 --- /dev/null +++ b/migrations/2026-03-23-180000_agent_schedules_v1/down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS agent_schedule_runs; +DROP TABLE IF EXISTS agent_schedules; diff --git a/migrations/2026-03-23-180000_agent_schedules_v1/up.sql b/migrations/2026-03-23-180000_agent_schedules_v1/up.sql new file mode 100644 index 00000000..9b2b9958 --- /dev/null +++ b/migrations/2026-03-23-180000_agent_schedules_v1/up.sql @@ -0,0 +1,89 @@ +CREATE TABLE agent_schedules ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + description TEXT NOT NULL, + instruction_enc BYTEA NOT NULL, + schedule_kind TEXT NOT NULL CHECK (schedule_kind IN ('one_off', 'recurring')), + recurrence_type TEXT CHECK ( + recurrence_type IS NULL OR recurrence_type IN ('interval', 'daily', 'weekly') + ), + schedule_spec JSONB NOT NULL, + timezone_mode TEXT NOT NULL CHECK (timezone_mode IN ('follow_user', 'fixed')), + resolved_timezone TEXT NOT NULL, + fixed_timezone TEXT, + stale_after_minutes INTEGER NOT NULL DEFAULT 15 CHECK (stale_after_minutes > 0), + status TEXT NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'completed', 'cancelled')), + next_scheduled_for TIMESTAMP WITH TIME ZONE, + last_scheduled_for TIMESTAMP WITH TIME ZONE, + last_run_at TIMESTAMP WITH TIME ZONE, + run_count INTEGER NOT NULL DEFAULT 0 CHECK (run_count >= 0), + cancelled_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, + CONSTRAINT agent_schedules_fixed_timezone_check CHECK ( + (timezone_mode = 'follow_user' AND fixed_timezone IS NULL) + OR (timezone_mode = 'fixed' AND fixed_timezone IS NOT NULL) + ), + CONSTRAINT agent_schedules_kind_recurrence_check CHECK ( + (schedule_kind = 'one_off' AND recurrence_type IS NULL) + OR (schedule_kind = 'recurring' AND recurrence_type IS NOT NULL) + ) +); + +CREATE TABLE agent_schedule_runs ( + id BIGSERIAL PRIMARY KEY, + uuid UUID NOT NULL DEFAULT gen_random_uuid() UNIQUE, + schedule_id BIGINT NOT NULL REFERENCES agent_schedules(id) ON DELETE CASCADE, + user_id UUID NOT NULL REFERENCES users(uuid) ON DELETE CASCADE, + agent_id BIGINT NOT NULL REFERENCES agents(id) ON DELETE CASCADE, + scheduled_for TIMESTAMP WITH TIME ZONE NOT NULL, + stale_after_at TIMESTAMP WITH TIME ZONE NOT NULL, + status TEXT NOT NULL DEFAULT 'pending' CHECK ( + status IN ('pending', 'leased', 'retry', 'completed', 'failed', 'cancelled', 'expired') + ), + attempt_count INTEGER NOT NULL DEFAULT 0 CHECK (attempt_count >= 0), + next_attempt_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + lease_owner TEXT, + lease_expires_at TIMESTAMP WITH TIME ZONE, + started_at TIMESTAMP WITH TIME ZONE, + first_output_at TIMESTAMP WITH TIME ZONE, + first_message_id UUID, + output_count INTEGER NOT NULL DEFAULT 0 CHECK (output_count >= 0), + notification_enqueued_at TIMESTAMP WITH TIME ZONE, + completed_at TIMESTAMP WITH TIME ZONE, + last_error TEXT, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE (schedule_id, scheduled_for) +); + +CREATE INDEX idx_agent_schedules_active_due + ON agent_schedules(status, next_scheduled_for, id) + WHERE status = 'active'; + +CREATE INDEX idx_agent_schedules_agent_status + ON agent_schedules(agent_id, status, next_scheduled_for); + +CREATE INDEX idx_agent_schedules_user_status + ON agent_schedules(user_id, status, created_at DESC); + +CREATE INDEX idx_agent_schedule_runs_pending + ON agent_schedule_runs(status, next_attempt_at, id); + +CREATE INDEX idx_agent_schedule_runs_schedule + ON agent_schedule_runs(schedule_id, scheduled_for DESC); + +CREATE INDEX idx_agent_schedule_runs_agent_status + ON agent_schedule_runs(agent_id, status, scheduled_for DESC); + +CREATE TRIGGER update_agent_schedules_updated_at + BEFORE UPDATE ON agent_schedules + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); + +CREATE TRIGGER update_agent_schedule_runs_updated_at + BEFORE UPDATE ON agent_schedule_runs + FOR EACH ROW + EXECUTE FUNCTION update_updated_at_column(); diff --git a/src/main.rs b/src/main.rs index a2faccdc..c5511e81 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,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::{ @@ -2780,6 +2781,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_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..360a7295 --- /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: String, + 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: String, + 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 2adbf1b4..2907e6a6 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,4 +1,6 @@ pub mod account_deletion; +pub mod agent_schedule_runs; +pub mod agent_schedules; pub mod agents; pub mod conversation_summaries; pub mod email_verification; diff --git a/src/models/schema.rs b/src/models/schema.rs index f492eba1..4f890a51 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -21,6 +21,58 @@ diesel::table! { } } +diesel::table! { + agent_schedule_runs (id) { + id -> Int8, + uuid -> Uuid, + schedule_id -> Int8, + user_id -> Uuid, + agent_id -> Int8, + scheduled_for -> Timestamptz, + stale_after_at -> Timestamptz, + status -> Text, + attempt_count -> Int4, + next_attempt_at -> Timestamptz, + lease_owner -> Nullable, + lease_expires_at -> Nullable, + started_at -> Nullable, + first_output_at -> Nullable, + first_message_id -> Nullable, + output_count -> Int4, + notification_enqueued_at -> Nullable, + completed_at -> Nullable, + last_error -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + +diesel::table! { + agent_schedules (id) { + id -> Int8, + uuid -> Uuid, + user_id -> Uuid, + agent_id -> Int8, + description -> Text, + instruction_enc -> Bytea, + schedule_kind -> Text, + recurrence_type -> Nullable, + schedule_spec -> Jsonb, + timezone_mode -> Text, + resolved_timezone -> Text, + fixed_timezone -> Nullable, + stale_after_minutes -> Int4, + status -> Text, + next_scheduled_for -> Nullable, + last_scheduled_for -> Nullable, + last_run_at -> Nullable, + run_count -> Int4, + cancelled_at -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + diesel::table! { agents (id) { id -> Int8, @@ -516,6 +568,9 @@ diesel::table! { } } +diesel::joinable!(agent_schedule_runs -> agent_schedules (schedule_id)); +diesel::joinable!(agent_schedule_runs -> agents (agent_id)); +diesel::joinable!(agent_schedules -> agents (agent_id)); diesel::joinable!(agents -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> conversations (conversation_id)); diesel::joinable!(assistant_messages -> responses (response_id)); @@ -547,6 +602,8 @@ diesel::joinable!(users -> org_projects (project_id)); diesel::allow_tables_to_appear_in_same_query!( account_deletion_requests, + agent_schedule_runs, + agent_schedules, agents, assistant_messages, conversation_summaries, diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 59b4b688..3c9419a6 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -39,10 +39,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, diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index eb920260..42f525af 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -36,6 +36,9 @@ 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, @@ -313,6 +316,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(()) } @@ -660,6 +670,22 @@ impl AgentRuntime { 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(), + agent.clone(), + ))); + tools.register(Arc::new(ListSchedulesTool::new( + state.clone(), + user.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(), @@ -1483,7 +1509,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..d3ca693a --- /dev/null +++ b/src/web/agent/schedules.rs @@ -0,0 +1,1769 @@ +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::encrypt::decrypt_string; +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_agent_message_notification, AgentPushTarget}; +use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; +use crate::web::openai_auth::AuthMethod; +use crate::{ApiError, 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_or_utc(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, + agent: Agent, +} + +impl ScheduleTaskTool { + pub fn new( + state: Arc, + user: Arc, + user_key: Arc, + agent: Agent, + ) -> Self { + Self { + state, + user, + user_key, + 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) => { + error!("schedule_task failed to resolve timezone: {e:?}"); + return ToolResult::error("Failed to resolve timezone for schedule"); + } + }; + + let tz = parse_timezone_or_utc(&resolved_timezone); + 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 { + return ToolResult::error("Scheduled time must be in the future"); + } + + let description = args + .get("description") + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| truncate_for_description(instruction)); + + let instruction_enc = + crate::encrypt::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 new_schedule = NewAgentSchedule { + uuid: Uuid::new_v4(), + user_id: self.user.uuid, + agent_id: self.agent.id, + description, + 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"), + }; + + match new_schedule.insert(&mut conn) { + 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, + agent: Agent, +} + +impl ListSchedulesTool { + pub fn new(state: Arc, user: Arc, agent: Agent) -> Self { + Self { state, user, 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()); + output.push_str(&format!( + "- id: {}\n status: {}\n description: {}\n rule: {}\n timezone: {} ({})\n next: {}\n\n", + schedule.uuid, + schedule.status, + schedule.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(()); + }; + 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 user_key = state + .get_user_key(user.uuid, None, None) + .await + .map_err(|e| format!("failed to derive 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 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 !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_agent_message_notification( + state, + &user, + target, + first_message.message_id, + &preview, + ) + .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).await + } else if agent.kind == AGENT_KIND_SUBAGENT { + runtime::AgentRuntime::new_subagent(state.clone(), user.clone(), user_key, agent.uuid).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: {}\nDescription: {}\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, + schedule.description.trim(), + 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 { + return Ok(parse_timezone_or_utc(fixed_timezone.unwrap_or("UTC")) + .name() + .to_string()); + } + + let mut conn = state + .db + .get_pool() + .get() + .map_err(|_| ApiError::InternalServerError)?; + + load_user_timezone_name(&mut conn, user_key, user.uuid) + .map(|timezone| timezone.unwrap_or_else(|| "UTC".to_string())) + .map_err(|_| ApiError::InternalServerError) +} + +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 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 truncate_for_description(instruction: &str) -> String { + let trimmed = instruction.trim(); + if trimmed.len() <= 80 { + return trimmed.to_string(); + } + + let mut end = 80; + while end > 0 && !trimmed.is_char_boundary(end) { + end -= 1; + } + format!("{}…", trimmed[..end].trim_end()) +} + +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 + ); + } +} + +#[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 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()); + } +} From 003cde8c2db80a1d6f569d6fb4ceaed7a19c0dc5 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 23 Mar 2026 18:37:24 -0500 Subject: [PATCH 31/34] fix(agent): validate schedule timezones strictly Avoid silently treating follow-user schedules as UTC so future local reminder times are not rejected as already in the past. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/web/agent/schedules.rs | 81 ++++++++++++++++++++++++++++++++------ 1 file changed, 68 insertions(+), 13 deletions(-) diff --git a/src/web/agent/schedules.rs b/src/web/agent/schedules.rs index d3ca693a..154792c3 100644 --- a/src/web/agent/schedules.rs +++ b/src/web/agent/schedules.rs @@ -31,7 +31,7 @@ use crate::models::users::User; use crate::push::{enqueue_agent_message_notification, AgentPushTarget}; use crate::web::openai::{get_chat_completion_response, BillingContext, CompletionChunk}; use crate::web::openai_auth::AuthMethod; -use crate::{ApiError, AppState}; +use crate::AppState; use super::runtime; use super::tools::{Tool, ToolResult}; @@ -96,7 +96,7 @@ pub(crate) fn refresh_follow_user_schedules_for_user( user_id: Uuid, timezone_name: &str, ) -> Result<(), AgentScheduleError> { - let timezone = parse_timezone_or_utc(timezone_name); + let timezone = parse_timezone(timezone_name)?; let schedules = AgentSchedule::list_active_follow_user_for_user(conn, user_id)?; if !schedules.is_empty() { @@ -263,12 +263,22 @@ impl Tool for ScheduleTaskTool { { Ok(timezone) => timezone, Err(e) => { - error!("schedule_task failed to resolve timezone: {e:?}"); - return ToolResult::error("Failed to resolve timezone for schedule"); + 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 = parse_timezone_or_utc(&resolved_timezone); + 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, @@ -276,7 +286,25 @@ impl Tool for ScheduleTaskTool { }; if next_scheduled_for <= now { - return ToolResult::error("Scheduled time must be in the future"); + 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 description = args @@ -1423,22 +1451,33 @@ async fn resolve_initial_timezone_name( user_key: &SecretKey, timezone_mode: &str, fixed_timezone: Option<&str>, -) -> Result { +) -> Result { if timezone_mode == SCHEDULE_TIMEZONE_MODE_FIXED { - return Ok(parse_timezone_or_utc(fixed_timezone.unwrap_or("UTC")) - .name() - .to_string()); + 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(|_| ApiError::InternalServerError)?; + .map_err(|_| "Database connection error".to_string())?; load_user_timezone_name(&mut conn, user_key, user.uuid) - .map(|timezone| timezone.unwrap_or_else(|| "UTC".to_string())) - .map_err(|_| ApiError::InternalServerError) + .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( @@ -1642,6 +1681,17 @@ 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!( @@ -1725,6 +1775,11 @@ mod tests { 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); From 2d590b004aaa944d4a9a6d5f650eaa188e6886c1 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Mon, 23 Mar 2026 19:43:39 -0500 Subject: [PATCH 32/34] feat(agent): let Maple update user preferences via tools Expose a set_preference tool so Maple can persist timezone and locale during conversations and keep follow-user schedules aligned with timezone changes. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/web/agent/runtime.rs | 13 ++-- src/web/agent/tools.rs | 138 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 4 deletions(-) diff --git a/src/web/agent/runtime.rs b/src/web/agent/runtime.rs index 42f525af..bd29e0b6 100644 --- a/src/web/agent/runtime.rs +++ b/src/web/agent/runtime.rs @@ -45,8 +45,8 @@ use super::signatures::{ }; use super::tools::{ ArchivalInsertTool, ArchivalSearchTool, ConversationSearchTool, DoneTool, MemoryAppendTool, - MemoryInsertTool, MemoryReplaceTool, SpawnSubagentTool, ToolRegistry, ToolResult, - WebSearchTool, + MemoryInsertTool, MemoryReplaceTool, SetPreferenceTool, SpawnSubagentTool, ToolRegistry, + ToolResult, WebSearchTool, }; use super::vision; @@ -201,7 +201,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()) @@ -297,7 +297,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, @@ -667,6 +667,11 @@ 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))); } 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 4f69377b680eadc35e35b44966f9adda658b1133 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 24 Mar 2026 18:11:22 -0500 Subject: [PATCH 33/34] feat(agent): add first-class Maple emoji reactions Persist reactions as message metadata and surface them in Maple SSE and prompts while keeping continuation turns and observability predictable. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + internal_docs/maple-emoji-reactions-plan.md | 375 +++++++++++++++++ .../down.sql | 5 + .../up.sql | 5 + src/db.rs | 2 - src/encrypt.rs | 5 - src/jwt.rs | 12 - src/main.rs | 10 +- src/models/responses.rs | 87 ++++ src/models/schema.rs | 2 + src/private_key.rs | 19 +- src/web/agent/mod.rs | 240 ++++++++++- src/web/agent/reactions.rs | 218 ++++++++++ src/web/agent/runtime.rs | 249 ++++++++++- src/web/agent/signatures.rs | 390 +++++++++++++++--- src/web/encryption_middleware.rs | 5 - src/web/responses/conversions.rs | 4 + src/web/responses/handlers.rs | 2 + 19 files changed, 1499 insertions(+), 133 deletions(-) create mode 100644 internal_docs/maple-emoji-reactions-plan.md create mode 100644 migrations/2026-03-24-220000_agent_message_reactions/down.sql create mode 100644 migrations/2026-03-24-220000_agent_message_reactions/up.sql create mode 100644 src/web/agent/reactions.rs diff --git a/Cargo.lock b/Cargo.lock index 861b396c..6f2cbedb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4178,6 +4178,7 @@ dependencies = [ "tower-http 0.5.2", "tracing", "tracing-subscriber", + "unicode-segmentation", "url", "uuid", "validator", diff --git a/Cargo.toml b/Cargo.toml index f373d730..6725556a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,3 +79,4 @@ lazy_static = "1.4.0" subtle = "2.6.1" tiktoken-rs = "0.5" once_cell = "1.19" +unicode-segmentation = "1.12" diff --git a/internal_docs/maple-emoji-reactions-plan.md b/internal_docs/maple-emoji-reactions-plan.md new file mode 100644 index 00000000..a9ed122a --- /dev/null +++ b/internal_docs/maple-emoji-reactions-plan.md @@ -0,0 +1,375 @@ +# Maple Emoji Reactions + +**Date:** March 2026 +**Status:** Implemented +**Related Docs:** +- `internal_docs/maple-agent-memory-architecture.md` +- `docs/encrypted-mobile-push-notifications.md` + +## Implementation Notes + +- `reply_reaction` is only available on the first assistant step of a turn. +- Continuation turns after tool results omit the `reply_reaction` field from the model output template entirely. +- User reaction mutation routes rely on the HTTP response; separate SSE fan-out for those mutations is out of scope in v1. + +--- + +## 1. Goal + +Add first-class emoji reactions to Maple conversations for both directions: + +- users reacting to assistant messages +- assistants reacting to the user message they are replying to + +This should be implemented as real message metadata in storage, API objects, and SSE, while also being rendered into the Maple prompt transcript so the model can see reaction context. + +--- + +## 2. Product Shape + +### 2.1 Supported v1 Behavior + +Each message can have at most one reaction from the opposite side: + +- a `user_messages` row may hold one assistant reaction +- an `assistant_messages` row may hold one user reaction + +This matches the current Maple chat model and keeps the first implementation small. + +### 2.2 Prompt Rendering + +Reactions should appear in the generated prompt transcript header, not as separate text in the message body. + +Example: + +```text +[user @ 10:42 AM | ❀️]: i got the job +[assistant @ 10:43 AM | πŸŽ‰]: NO WAY +``` + +The reaction glyph by itself is enough. We do not need extra wording like `assistant reacted`. + +--- + +## 3. Why Reactions Should Be First-Class + +Reactions should not be hidden inside prompt-only formatting or conversation-level metadata. + +They should be first-class because we need them to work consistently across: + +- database persistence +- normal item fetch/list APIs +- Maple SSE live delivery +- prompt construction for future turns + +The prompt rendering is a consumer of reaction data, not the source of truth. + +--- + +## 4. Storage Plan + +Use direct nullable columns on the existing shared message tables: + +```sql +ALTER TABLE user_messages + ADD COLUMN assistant_reaction TEXT; + +ALTER TABLE assistant_messages + ADD COLUMN user_reaction TEXT; +``` + +### 4.1 Why Direct Columns + +This is the smallest change that fits the current product requirements: + +- one reaction per message +- only the opposite side may react +- no reaction history required +- no multi-user aggregation required + +A separate `message_reactions` table can be introduced later if Maple needs multiple reactions, audit history, or multi-actor support. + +### 4.2 Encoding + +Use PostgreSQL `TEXT`. + +- PostgreSQL already supports UTF-8 text well +- emoji do not require a special column type +- we do not need encryption for reactions in v1 + +The main complexity is validation and normalization of multi-codepoint emoji, not storage type. + +--- + +## 5. API Model Plan + +Extend `ConversationItem::Message` with a first-class reaction field. + +Conceptually: + +```rust +Message { + id, + status, + role, + content, + reaction: Option, + created_at, +} +``` + +Interpretation: + +- for `role = "user"`, `reaction` means the assistant's reaction +- for `role = "assistant"`, `reaction` means the user's reaction + +This keeps the API simple for clients and avoids separate reaction wrapper objects in v1. + +### 5.1 Surfaces Covered + +Once threaded through shared message conversion, reactions will flow through: + +- `/v1/agent/items` +- `/v1/agent/items/:item_id` +- `/v1/agent/subagents/:id/items` +- `/v1/agent/subagents/:id/items/:item_id` +- shared conversation item APIs where those conversation rows are visible + +--- + +## 6. Assistant Reaction Generation + +Assistant reactions should be a first-class Maple agent output, not a tool call. + +Current Maple output is: + +```rust +messages: Vec +tool_calls: Vec +``` + +Implemented first-step Maple output: + +```rust +messages: Vec +reply_reaction: String +tool_calls: Vec +``` + +The model uses `""` when no first-step reaction is intended. + +Continuation turns after tool results use a reduced output shape with no `reply_reaction` field. + +### 6.1 Why Not a Tool + +Normal Maple tools currently imply all of the following: + +- execute a tool implementation +- inject a `[Tool Result: ...]` block into the next model turn +- persist `tool_call` and `tool_output` rows into transcript history + +That behavior is correct for real tools, but it is the wrong shape for emoji reactions. + +An assistant emoji reaction is closer to a first-class chat output, like a message, than to a tool execution. + +### 6.2 Targeting Semantics + +`reply_reaction` should implicitly target the current user message for the turn. + +No explicit message id needs to be exposed to the model in v1. + +--- + +## 7. Runtime Plan + +### 7.1 User -> Assistant Flow + +When a user sends a Maple message: + +1. persist the user message as normal +2. retain the inserted user message id/uuid in runtime state for the current turn +3. run the Maple agent loop +4. persist assistant `messages` as normal assistant message rows +5. on the first assistant step only, if `reply_reaction` is present, update the just-inserted user message row with `assistant_reaction` + +### 7.2 Why This Works Well + +- the model only needs to decide whether to react and which emoji to use +- runtime already knows which user message is being replied to +- no transcript noise is created from fake tool rows + +--- + +## 8. User Reaction Mutation Plan + +Add Maple-specific mutation endpoints for reacting to assistant messages. + +Suggested shape: + +- `POST /v1/agent/items/:item_id/reaction` +- `DELETE /v1/agent/items/:item_id/reaction` +- `POST /v1/agent/subagents/:id/items/:item_id/reaction` +- `DELETE /v1/agent/subagents/:id/items/:item_id/reaction` + +Expected behavior: + +- only assistant message items can receive a user reaction through these routes +- `POST` sets or replaces the reaction +- `DELETE` clears it + +Main-agent mutation routes should remain under `/v1/agent/...` because the public conversations API intentionally hides the main-agent thread. + +--- + +## 9. SSE Plan + +Maple SSE needs to become more structured so reactions can update live. + +### 9.1 Current Limitation + +Current `agent.message` events only send: + +```json +{ "messages": [...], "step": N } +``` + +That is too thin for live reactions because clients do not receive message item ids. + +### 9.2 Proposed Changes + +#### `agent.message` + +Include the persisted assistant message id with each delivered assistant message. + +Conceptually: + +```json +{ + "message_id": "uuid", + "messages": ["..."], + "step": 0 +} +``` + +If Maple continues to emit one assistant message per event, this shape is sufficient. + +#### `agent.reaction` + +Add a dedicated reaction event: + +```json +{ + "item_id": "uuid", + "emoji": "🫑" +} +``` + +We do not need an explicit `actor` field. In Maple, the target message role already implies who reacted: + +- reaction on a user message means assistant reacted +- reaction on an assistant message means user reacted + +In the implemented v1 scope, live `agent.reaction` SSE is used for assistant-generated reactions during chat turns. User `POST/DELETE .../reaction` mutations return the updated item over HTTP; separate multi-client fan-out is intentionally out of scope. + +--- + +## 10. Prompt Construction Plan + +Prompt building should include reaction metadata in the transcript header. + +This applies when rendering `recent_conversation` from stored thread items. + +Example rendering rules: + +- no reaction: + ```text + [user @ 10:42 AM]: hey + ``` +- with reaction: + ```text + [user @ 10:42 AM | ❀️]: hey + ``` + +This keeps the signal lightweight while still letting the model perceive social context. + +--- + +## 11. Validation / Normalization + +V1 accepts arbitrary emoji input without a fixed palette, but the server still applies lightweight normalization: + +- trim surrounding whitespace +- unwrap quoted JSON-string values +- reject empty strings +- enforce a 16-character maximum +- require a single grapheme cluster +- reject component-only values such as bare `1`, `#`, `*`, joiners, or variation selectors without an emoji base + +We should treat the stored value as a single reaction value, even if the emoji is composed of multiple Unicode codepoints. + +We do not need a fixed reaction palette in v1. + +--- + +## 12. File-Level Impact + +Expected implementation areas: + +- migrations for `user_messages` and `assistant_messages` +- `src/models/schema.rs` +- `src/models/responses.rs` +- `src/db.rs` +- `src/web/responses/conversions.rs` +- `src/web/responses/conversations.rs` +- `src/web/agent/runtime.rs` +- `src/web/agent/signatures.rs` +- `src/web/agent/mod.rs` + +Potentially also: + +- client-facing request/response structs for Maple reaction mutation routes +- Responses API item serializers if reaction support should also appear there + +--- + +## 13. Implementation Order + +Recommended order: + +1. add DB columns and model/schema wiring +2. extend raw conversation message queries to carry reaction fields +3. extend `ConversationItem::Message` with `reaction` +4. wire reactions through Maple item list/get routes +5. add Maple user-reaction mutation endpoints +6. add `reply_reaction` to Maple DSR output +7. update runtime to persist assistant reactions onto the current user message +8. extend Maple SSE with `message_id` and `agent.reaction` +9. render reactions into prompt transcript headers + +--- + +## 14. Non-Goals for v1 + +Not included in the first implementation: + +- multiple reactions per message +- reaction history / audit log +- arbitrary target-item reactions by the assistant +- cross-user aggregated reactions +- reaction counts +- custom reaction metadata objects beyond a single emoji string + +--- + +## 15. Summary + +The v1 Maple reaction design should be: + +- direct nullable reaction columns on message rows +- first-class `reaction` field on message items +- first-class first-step agent output `reply_reaction` +- continuation-step templates with no `reply_reaction` field +- live `agent.reaction` SSE events for assistant-generated reactions during chat +- prompt rendering using compact header metadata like `| ❀️` + +This keeps the feature native to Maple chat semantics without abusing the tool-call system. diff --git a/migrations/2026-03-24-220000_agent_message_reactions/down.sql b/migrations/2026-03-24-220000_agent_message_reactions/down.sql new file mode 100644 index 00000000..b9b00b64 --- /dev/null +++ b/migrations/2026-03-24-220000_agent_message_reactions/down.sql @@ -0,0 +1,5 @@ +ALTER TABLE assistant_messages + DROP COLUMN IF EXISTS user_reaction; + +ALTER TABLE user_messages + DROP COLUMN IF EXISTS assistant_reaction; diff --git a/migrations/2026-03-24-220000_agent_message_reactions/up.sql b/migrations/2026-03-24-220000_agent_message_reactions/up.sql new file mode 100644 index 00000000..bf4b9361 --- /dev/null +++ b/migrations/2026-03-24-220000_agent_message_reactions/up.sql @@ -0,0 +1,5 @@ +ALTER TABLE user_messages + ADD COLUMN assistant_reaction TEXT; + +ALTER TABLE assistant_messages + ADD COLUMN user_reaction TEXT; diff --git a/src/db.rs b/src/db.rs index f7709569..ddbe68b8 100644 --- a/src/db.rs +++ b/src/db.rs @@ -643,7 +643,6 @@ impl DBConnection for PostgresConnection { } fn get_user_by_uuid(&self, uuid: Uuid) -> Result { - debug!("Getting user by UUID"); let conn = &mut self.db.get().map_err(|_| DBError::ConnectionError)?; let result = User::get_by_uuid(conn, uuid)?.ok_or(DBError::UserNotFound); if let Err(ref e) = result { @@ -2456,7 +2455,6 @@ impl DBConnection for PostgresConnection { after: Option, order: &str, ) -> Result, DBError> { - debug!("Getting conversation context messages"); let conn = &mut self.db.get().map_err(|_| DBError::ConnectionError)?; RawThreadMessage::get_conversation_context(conn, conversation_id, limit, after, order) .map_err(DBError::from) diff --git a/src/encrypt.rs b/src/encrypt.rs index 8f2cd0e9..30f8039b 100644 --- a/src/encrypt.rs +++ b/src/encrypt.rs @@ -32,7 +32,6 @@ pub enum EncryptError { } pub async fn encrypt_with_key(encryption_key: &SecretKey, bytes: &[u8]) -> Vec { - tracing::debug!("Entering encrypt_with_key"); let cipher = Aes256Gcm::new_from_slice(&encryption_key.secret_bytes()).expect("should convert"); // Generate a random 96-bit nonce @@ -46,13 +45,10 @@ pub async fn encrypt_with_key(encryption_key: &SecretKey, bytes: &[u8]) -> Vec Result, EncryptError> { - tracing::trace!("Entering decrypt_with_key"); - if bytes.len() < 12 { tracing::error!( "Decrypt failed: Input too short (length {}), minimum 12 bytes required", @@ -81,7 +77,6 @@ pub fn decrypt_with_key(encryption_key: &SecretKey, bytes: &[u8]) -> Result, next: Next, ) -> impl IntoResponse { - tracing::debug!("Entering validate_jwt"); let token = match req .headers() .get(header::AUTHORIZATION) @@ -430,8 +429,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(), @@ -454,7 +451,6 @@ pub async fn validate_jwt( }; req.extensions_mut().insert(user); - tracing::debug!("Exiting validate_jwt"); next.run(req).await } @@ -463,7 +459,6 @@ pub async fn validate_platform_jwt( mut req: Request, next: Next, ) -> impl IntoResponse { - tracing::debug!("Entering validate_platform_jwt"); let token = match req .headers() .get(header::AUTHORIZATION) @@ -474,8 +469,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(), @@ -498,7 +491,6 @@ pub async fn validate_platform_jwt( }; req.extensions_mut().insert(platform_user); - tracing::debug!("Exiting validate_platform_jwt"); next.run(req).await } @@ -511,8 +503,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, @@ -525,8 +515,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/main.rs b/src/main.rs index c5511e81..27849ece 100644 --- a/src/main.rs +++ b/src/main.rs @@ -907,8 +907,6 @@ impl AppState { derivation_path: Option<&str>, seed_phrase_derivation_path: Option<&str>, ) -> Result { - info!("Getting user key for UUID: {}", user_uuid); - if let Some(path) = derivation_path { debug!("Using BIP-32 derivation path: {}", path); } @@ -918,13 +916,9 @@ impl AppState { } let user = self.get_user(user_uuid).await?; - debug!("User retrieved successfully"); let encrypted_seed = match user.get_seed_encrypted().await { - Some(es) => { - debug!("Found existing encrypted seed for user"); - es - } + Some(es) => es, None => { // create seed if not already exists info!( @@ -939,7 +933,6 @@ impl AppState { } }; - debug!("Decrypting user seed and deriving key"); let user_secret_key = decrypt_user_seed_to_key( self.enclave_key.clone(), encrypted_seed, @@ -947,7 +940,6 @@ impl AppState { seed_phrase_derivation_path, )?; - info!("Successfully derived key for user: {}", user_uuid); Ok(user_secret_key) } diff --git a/src/models/responses.rs b/src/models/responses.rs index abe10b15..0870082c 100644 --- a/src/models/responses.rs +++ b/src/models/responses.rs @@ -422,6 +422,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)?; @@ -436,6 +437,7 @@ impl NewConversation { completion_tokens: 0, status: "in_progress".to_string(), finish_reason: None, + user_reaction: None, }; placeholder_assistant.insert(tx)?; } @@ -463,6 +465,7 @@ pub struct UserMessage { pub created_at: DateTime, pub updated_at: DateTime, pub attachment_text_enc: Option>, + pub assistant_reaction: Option, } #[derive(Insertable, Debug)] @@ -475,6 +478,7 @@ pub struct NewUserMessage { pub content_enc: Vec, pub prompt_tokens: i32, pub attachment_text_enc: Option>, + pub assistant_reaction: Option, } impl UserMessage { @@ -508,6 +512,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( @@ -666,6 +692,7 @@ pub struct AssistantMessage { pub finish_reason: Option, pub created_at: DateTime, pub updated_at: DateTime, + pub user_reaction: Option, } #[derive(Insertable, Debug)] @@ -679,9 +706,25 @@ pub struct NewAssistantMessage { pub completion_tokens: i32, pub status: String, pub finish_reason: Option, + 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, @@ -701,6 +744,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 { @@ -796,6 +861,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, @@ -838,6 +905,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, @@ -858,6 +926,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, @@ -878,6 +947,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, @@ -897,6 +967,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, @@ -917,6 +988,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, @@ -956,6 +1028,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, @@ -976,6 +1049,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, @@ -996,6 +1070,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, @@ -1015,6 +1090,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, @@ -1035,6 +1111,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, @@ -1083,6 +1160,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, @@ -1103,6 +1181,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, @@ -1123,6 +1202,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, @@ -1142,6 +1222,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, @@ -1162,6 +1243,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, @@ -1230,6 +1312,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, @@ -1250,6 +1333,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, @@ -1270,6 +1354,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, @@ -1289,6 +1374,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, @@ -1309,6 +1395,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/models/schema.rs b/src/models/schema.rs index 4f890a51..dda9d0cd 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -102,6 +102,7 @@ diesel::table! { finish_reason -> Nullable, created_at -> Timestamptz, updated_at -> Timestamptz, + user_reaction -> Nullable, } } @@ -525,6 +526,7 @@ diesel::table! { created_at -> Timestamptz, updated_at -> Timestamptz, attachment_text_enc -> Nullable, + assistant_reaction -> Nullable, } } diff --git a/src/private_key.rs b/src/private_key.rs index 9d45aae8..1c20460b 100644 --- a/src/private_key.rs +++ b/src/private_key.rs @@ -37,29 +37,17 @@ pub fn decrypt_user_seed_to_key( seed_phrase_derivation_path: Option<&str>, ) -> Result { // If seed_phrase_derivation_path is provided, derive a child mnemonic using BIP-85 - let (source_mnemonic, uses_bip85) = if let Some(bip85_path) = seed_phrase_derivation_path { + let source_mnemonic = if let Some(bip85_path) = seed_phrase_derivation_path { info!("Using BIP-85 derivation with path: {}", bip85_path); - ( - decrypt_and_derive_bip85_mnemonic(enclave_key, encrypted_seed, bip85_path)?, - true, - ) + decrypt_and_derive_bip85_mnemonic(enclave_key, encrypted_seed, bip85_path)? } else { - debug!("Using root mnemonic (no BIP-85 derivation)"); - ( - decrypt_user_seed_to_mnemonic(enclave_key, encrypted_seed)?, - false, - ) + decrypt_user_seed_to_mnemonic(enclave_key, encrypted_seed)? }; // Generate seed from the appropriate mnemonic (either the root or the BIP-85 derived one) - info!( - "Generating seed from {} mnemonic", - if uses_bip85 { "BIP-85 derived" } else { "root" } - ); let seed = source_mnemonic.to_seed(""); // Create extended private key from the seed - debug!("Creating extended private key from seed"); let xprivkey = Xpriv::new_master(Network::Bitcoin, &seed) .map_err(|e| Error::EncryptionError(e.to_string()))?; @@ -74,7 +62,6 @@ pub fn decrypt_user_seed_to_key( debug!("Successfully derived child key using BIP-32 path"); Ok(derived_key.private_key) } else { - debug!("Using master key (no BIP-32 derivation)"); Ok(xprivkey.private_key) } } diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index 3c9419a6..ac58f59d 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; @@ -22,7 +22,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}; @@ -38,6 +38,7 @@ use crate::web::responses::{ use crate::{ApiError, AppMode, AppState}; mod compaction; +mod reactions; mod runtime; mod schedules; mod signatures; @@ -51,6 +52,11 @@ struct AgentChatRequest { input: MessageContent, } +#[derive(Debug, Clone, Deserialize)] +struct SetMessageReactionRequest { + emoji: String, +} + #[derive(Debug, Clone, Default, Deserialize)] struct InitMainAgentRequest { timezone: Option, @@ -125,6 +131,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 { @@ -165,6 +178,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"; @@ -172,10 +186,18 @@ const AGENT_SSE_KEEPALIVE_INTERVAL_SECS: u64 = 15; #[derive(Debug, Clone, Serialize)] struct AgentMessageEvent { + #[serde(skip_serializing_if = "Option::is_none")] + message_id: Option, messages: Vec, step: usize, } +#[derive(Debug, Clone, Serialize)] +struct AgentReactionEvent { + item_id: Uuid, + emoji: String, +} + #[derive(Debug, Clone, Serialize)] struct AgentTypingEvent { step: usize, @@ -199,6 +221,7 @@ enum AgentClientEvent { payload: AgentMessageEvent, delivery_ack: oneshot::Sender<()>, }, + Reaction(AgentReactionEvent), Done(AgentDoneEvent), Error(AgentErrorEvent), } @@ -265,6 +288,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(), @@ -446,6 +470,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, @@ -490,6 +546,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( @@ -529,6 +594,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) @@ -614,6 +688,46 @@ 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(body): Extension, +) -> Result>, ApiError> { + let emoji = reactions::require_valid_reaction(&body.emoji)?; + let ctx = load_main_agent_context(&state, &user).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, +) -> Result>, ApiError> { + let ctx = load_main_agent_context(&state, &user).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, @@ -761,6 +875,46 @@ 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(body): Extension, +) -> Result>, ApiError> { + let emoji = reactions::require_valid_reaction(&body.emoji)?; + let ctx = load_subagent_context(&state, &user, agent_uuid).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, +) -> Result>, ApiError> { + let ctx = load_subagent_context(&state, &user, agent_uuid).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, @@ -857,6 +1011,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 } @@ -886,7 +1043,19 @@ 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; @@ -963,6 +1132,14 @@ async fn run_agent_chat_task( let max_steps = runtime.max_steps(); 'steps: for step_num in 0..max_steps { + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + step_num, + total_messages, + had_reaction, + "Starting agent chat step" + ); client_connected = send_agent_client_event( &tx, AgentClientEvent::Typing(AgentTypingEvent { step: step_num }), @@ -974,11 +1151,27 @@ async fn run_agent_chat_task( 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) @@ -988,6 +1181,29 @@ async fn run_agent_chat_task( } } + 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() { total_messages += messages.len(); @@ -1003,6 +1219,7 @@ async fn run_agent_chat_task( let delivered = send_agent_message_event( &tx, AgentMessageEvent { + message_id: assistant_message.as_ref().map(|message| message.uuid), messages: vec![msg.clone()], step: step_num, }, @@ -1044,12 +1261,17 @@ async fn run_agent_chat_task( } } - if !had_error && total_messages == 0 { - warn!("Agent produced no messages"); + 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 _ = send_agent_message_event( &tx, AgentMessageEvent { + message_id: None, messages: vec![ "I apologize, but I wasn't able to generate a response.".to_string() ], @@ -1074,6 +1296,16 @@ async fn run_agent_chat_task( client_connected, ) .await; + + debug!( + user_uuid = %user.uuid, + target = chat_target_label(&target), + total_messages, + had_reaction, + had_error, + client_connected, + "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 bd29e0b6..be301819 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::*; @@ -36,12 +36,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, @@ -108,6 +109,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, @@ -121,18 +123,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( @@ -369,6 +411,7 @@ async fn seed_main_agent_onboarding_messages( completion_tokens, status: "completed".to_string(), finish_reason: None, + user_reaction: None, } .insert(conn) .map_err(|e| { @@ -633,8 +676,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(), @@ -718,13 +759,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, }) } @@ -732,14 +773,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 @@ -831,8 +874,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) } @@ -894,7 +939,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?; @@ -978,7 +1031,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(), @@ -1000,6 +1063,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?; @@ -1020,11 +1084,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) { @@ -1033,6 +1132,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" { @@ -1055,8 +1175,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, @@ -1247,7 +1405,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() { @@ -1367,6 +1534,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| { @@ -1419,6 +1587,7 @@ SELF-CHECK: Before ANY message, ask: "Is this new info the user hasn't seen?" If completion_tokens, status: "completed".to_string(), finish_reason: None, + user_reaction: None, } .insert(&mut conn) .map_err(|e| { @@ -1451,6 +1620,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, @@ -1669,3 +1865,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 e3d0d682..66886ff5 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/encryption_middleware.rs b/src/web/encryption_middleware.rs index 37e21d95..62941bf2 100644 --- a/src/web/encryption_middleware.rs +++ b/src/web/encryption_middleware.rs @@ -46,7 +46,6 @@ pub async fn decrypt_request( where T: DeserializeOwned + Send + Sync + Clone + 'static, { - tracing::debug!("Entering decrypt_request"); let session_id = headers .get("x-session-id") .and_then(|v| v.to_str().ok()) @@ -85,8 +84,6 @@ where request.extensions_mut().insert(decrypted); request.extensions_mut().insert(session_id); - - tracing::debug!("Exiting decrypt_request"); Ok(next.run(request).await) } @@ -95,12 +92,10 @@ pub async fn encrypt_response( session_id: &Uuid, response: &T, ) -> Result>, ApiError> { - tracing::debug!("Entering encrypt_response"); let response_json = serde_json::to_vec(response).map_err(|_| ApiError::InternalServerError)?; let encrypted_response = state .encrypt_session_data(session_id, &response_json) .await?; - tracing::debug!("Exiting encrypt_response"); Ok(Json(EncryptedResponse::new( base64::engine::general_purpose::STANDARD.encode(encrypted_response), ))) 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 231d3c7c..45a7d72d 100644 --- a/src/web/responses/handlers.rs +++ b/src/web/responses/handlers.rs @@ -1261,6 +1261,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 @@ -1654,6 +1655,7 @@ async fn setup_completion_processor( completion_tokens: 0, status: STATUS_IN_PROGRESS.to_string(), finish_reason: None, + user_reaction: None, }; state .db From 7117817a8e11313df4deed109287ff5acf7e0a6f Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 24 Mar 2026 22:48:01 -0500 Subject: [PATCH 34/34] fix(agent): require persisted message ids in Maple SSE Simplify agent SSE delivery by sending one persisted message per event and using minimal typing and done signals for clients. Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/web/agent/mod.rs | 121 ++++++++++++++++++++++++------------------- 1 file changed, 68 insertions(+), 53 deletions(-) diff --git a/src/web/agent/mod.rs b/src/web/agent/mod.rs index ac58f59d..bf55e916 100644 --- a/src/web/agent/mod.rs +++ b/src/web/agent/mod.rs @@ -186,10 +186,8 @@ const AGENT_SSE_KEEPALIVE_INTERVAL_SECS: u64 = 15; #[derive(Debug, Clone, Serialize)] struct AgentMessageEvent { - #[serde(skip_serializing_if = "Option::is_none")] - message_id: Option, - messages: Vec, - step: usize, + message_id: Uuid, + message: String, } #[derive(Debug, Clone, Serialize)] @@ -199,15 +197,10 @@ struct AgentReactionEvent { } #[derive(Debug, Clone, Serialize)] -struct AgentTypingEvent { - step: usize, -} +struct AgentTypingEvent {} #[derive(Debug, Clone, Serialize)] -struct AgentDoneEvent { - total_steps: usize, - total_messages: usize, -} +struct AgentDoneEvent {} #[derive(Debug, Clone, Serialize)] struct AgentErrorEvent { @@ -1062,7 +1055,7 @@ async fn run_agent_chat_task( client_connected = send_agent_client_event( &tx, - AgentClientEvent::Typing(AgentTypingEvent { step: 0 }), + AgentClientEvent::Typing(AgentTypingEvent {}), client_connected, ) .await; @@ -1140,12 +1133,6 @@ async fn run_agent_chat_task( had_reaction, "Starting agent chat step" ); - 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) => { @@ -1205,23 +1192,32 @@ async fn run_agent_chat_task( } if !messages.is_empty() { - total_messages += messages.len(); - 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 { - message_id: assistant_message.as_ref().map(|message| message.uuid), - messages: vec![msg.clone()], - step: step_num, + message_id: assistant_message.uuid, + message: msg.clone(), }, client_connected, ) @@ -1229,11 +1225,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; @@ -1267,35 +1261,56 @@ async fn run_agent_chat_task( target = chat_target_label(&target), "Agent produced no messages or reactions; sending fallback response" ); - let _ = - send_agent_message_event( - &tx, - AgentMessageEvent { - message_id: None, - messages: vec![ - "I apologize, but I wasn't able to generate a response.".to_string() - ], - step: 0, - }, - client_connected, - ) - .await; - total_messages = 1; + 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, - }), - client_connected, - ) - .await; + if !had_error { + let _ = send_agent_client_event( + &tx, + AgentClientEvent::Done(AgentDoneEvent {}), + client_connected, + ) + .await; + } debug!( user_uuid = %user.uuid,