From b6b3eab4cf46d340c26797c8660407c5ab782447 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 22 May 2026 14:09:41 +0800 Subject: [PATCH] Support explicit branch access for subagent concurrency --- memoria/crates/memoria-api/src/auth.rs | 70 +++++ memoria/crates/memoria-api/src/models.rs | 9 + .../crates/memoria-api/src/routes/admin.rs | 14 +- .../memoria-api/src/routes/governance.rs | 17 +- memoria/crates/memoria-api/src/routes/mcp.rs | 78 ++++-- .../crates/memoria-api/src/routes/memory.rs | 169 +++++++++--- .../memoria-api/src/routes/snapshots.rs | 4 +- memoria/crates/memoria-cli/src/main.rs | 5 +- memoria/crates/memoria-git/src/service.rs | 6 +- memoria/crates/memoria-mcp/src/purge_args.rs | 7 + memoria/crates/memoria-mcp/src/remote.rs | 22 +- memoria/crates/memoria-mcp/src/server.rs | 4 +- memoria/crates/memoria-mcp/src/tools.rs | 76 ++++-- memoria/crates/memoria-service/src/lib.rs | 4 +- memoria/crates/memoria-service/src/service.rs | 254 ++++++++++++++++-- memoria/crates/memoria-storage/src/store.rs | 62 ++++- .../memoria-storage/tests/branch_ops.rs | 39 +++ .../memoria-storage/tests/store_crud.rs | 9 +- 18 files changed, 716 insertions(+), 133 deletions(-) diff --git a/memoria/crates/memoria-api/src/auth.rs b/memoria/crates/memoria-api/src/auth.rs index 09b198e5..5a3db694 100644 --- a/memoria/crates/memoria-api/src/auth.rs +++ b/memoria/crates/memoria-api/src/auth.rs @@ -142,6 +142,7 @@ pub async fn group_main_write_guard( req: axum::http::Request, next: axum::middleware::Next, ) -> axum::response::Response { + use axum::body::{to_bytes, Body}; use axum::response::IntoResponse; // Extract bearer token from Authorization header @@ -176,6 +177,75 @@ pub async fn group_main_write_guard( if group_main_write_allowed_for_solo_owner(&state, gid, &p.user_id).await { return next.run(req).await; } + let path = req.uri().path(); + let explicit_query_branch = req + .uri() + .query() + .and_then(|query| { + serde_urlencoded::from_str::< + std::collections::HashMap, + >(query) + .ok() + }) + .and_then(|query| query.get("branch").cloned()) + .map(|branch| { + let branch = branch.trim(); + !branch.is_empty() && branch != "main" + }) + .unwrap_or(false); + if req.method() == axum::http::Method::DELETE + && path.starts_with("/v1/memories/") + && !path.ends_with("/correct") + && explicit_query_branch + { + return next.run(req).await; + } + let branch_aware_body_route = matches!( + path, + "/v1/memories" + | "/v1/memories/batch" + | "/v1/memories/correct" + | "/v1/memories/purge" + | "/v1/observe" + ) || (path.starts_with("/v1/memories/") + && path.ends_with("/correct")); + if !branch_aware_body_route { + return ( + StatusCode::FORBIDDEN, + "main is read-only in group mode; \ + create or checkout a branch, then use selective apply" + .to_string(), + ) + .into_response(); + } + let (parts, body) = req.into_parts(); + let bytes = match to_bytes(body, 8 * 1024 * 1024).await { + Ok(bytes) => bytes, + Err(e) => { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + format!("request body could not be buffered: {e}"), + ) + .into_response(); + } + }; + let explicit_non_main_branch = + serde_json::from_slice::(&bytes) + .ok() + .and_then(|body| { + body.get("branch") + .and_then(|branch| branch.as_str()) + .map(str::to_string) + }) + .map(|branch| { + let branch = branch.trim(); + !branch.is_empty() && branch != "main" + }) + .unwrap_or(false); + if explicit_non_main_branch { + let req = axum::http::Request::from_parts(parts, Body::from(bytes)); + return next.run(req).await; + } return ( StatusCode::FORBIDDEN, "main is read-only in group mode; \ diff --git a/memoria/crates/memoria-api/src/models.rs b/memoria/crates/memoria-api/src/models.rs index c6b94da8..050a752b 100644 --- a/memoria/crates/memoria-api/src/models.rs +++ b/memoria/crates/memoria-api/src/models.rs @@ -17,6 +17,7 @@ pub struct StoreRequest { pub initial_confidence: Option, pub observed_at: Option, pub source: Option, + pub branch: Option, } fn default_memory_type() -> String { "semantic".to_string() @@ -25,6 +26,7 @@ fn default_memory_type() -> String { #[derive(Deserialize)] pub struct BatchStoreRequest { pub memories: Vec, + pub branch: Option, } #[derive(Deserialize)] @@ -34,6 +36,7 @@ pub struct RetrieveRequest { pub top_k: i64, pub session_id: Option, pub session_scope: Option, + pub branch: Option, /// Explain level: false/"none" = off, true/"basic" = basic, "verbose" = per-candidate scores, "analyze" = full #[serde(default, deserialize_with = "deserialize_explain")] pub explain: String, @@ -79,6 +82,7 @@ pub struct SearchRequest { pub top_k: i64, pub session_id: Option, pub session_scope: Option, + pub branch: Option, #[serde(default, deserialize_with = "deserialize_explain")] pub explain: String, } @@ -122,6 +126,7 @@ fn deserialize_explain<'de, D: serde::Deserializer<'de>>(d: D) -> Result, + pub branch: Option, } #[derive(Deserialize)] @@ -131,6 +136,7 @@ pub struct CorrectByQueryRequest { pub session_id: Option, pub session_scope: Option, pub reason: Option, + pub branch: Option, } impl CorrectByQueryRequest { @@ -156,6 +162,7 @@ pub struct PurgeRequest { pub session_id: Option, pub memory_types: Option>, pub reason: Option, + pub branch: Option, } pub enum PurgeSelector { @@ -448,6 +455,7 @@ mod tests { session_id: Some("sess-1".to_string()), memory_types: Some(vec![]), reason: None, + branch: None, }; match request.selector().unwrap() { @@ -470,6 +478,7 @@ mod tests { session_id: None, memory_types: Some(vec![]), reason: None, + branch: None, }; assert!(matches!(request.selector().unwrap(), PurgeSelector::None)); diff --git a/memoria/crates/memoria-api/src/routes/admin.rs b/memoria/crates/memoria-api/src/routes/admin.rs index bc11732e..b69f77d8 100644 --- a/memoria/crates/memoria-api/src/routes/admin.rs +++ b/memoria/crates/memoria-api/src/routes/admin.rs @@ -856,7 +856,10 @@ pub async fn user_branch_stats( if name == "main" || raw_table_name.is_empty() { continue; } - if !raw_table_name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + if !raw_table_name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + { tracing::warn!( user_id = %user_id, branch = %name, @@ -867,11 +870,10 @@ pub async fn user_branch_stats( continue; } let bt = user_store.t(raw_table_name); - let count_result = sqlx::query_scalar::<_, i64>(&format!( - "SELECT COUNT(*) FROM {bt} WHERE is_active > 0" - )) - .fetch_one(user_store.pool()) - .await; + let count_result = + sqlx::query_scalar::<_, i64>(&format!("SELECT COUNT(*) FROM {bt} WHERE is_active > 0")) + .fetch_one(user_store.pool()) + .await; let count = match count_result { Ok(n) => n, Err(e) => { diff --git a/memoria/crates/memoria-api/src/routes/governance.rs b/memoria/crates/memoria-api/src/routes/governance.rs index 4df63189..e065674b 100644 --- a/memoria/crates/memoria-api/src/routes/governance.rs +++ b/memoria/crates/memoria-api/src/routes/governance.rs @@ -5,13 +5,28 @@ use memoria_service::{ use serde_json::json; use tracing::warn; -use crate::{auth::AuthUser, models::*, routes::memory::api_err, state::AppState}; +use crate::{ + auth::{group_main_write_allowed_for_solo_owner, AuthUser}, + models::*, + routes::memory::api_err, + state::AppState, +}; pub async fn governance( State(state): State, auth: AuthUser, Json(req): Json, ) -> Result, (StatusCode, String)> { + if let Some(group_id) = auth.group_id.as_deref() { + if !group_main_write_allowed_for_solo_owner(&state, group_id, &auth.user_id).await { + return Err(( + StatusCode::FORBIDDEN, + "governance is disabled in group mode because it mutates main memory state" + .to_string(), + )); + } + } + let sql = state .service .user_sql_store(auth.scope_id()) diff --git a/memoria/crates/memoria-api/src/routes/mcp.rs b/memoria/crates/memoria-api/src/routes/mcp.rs index 3ed4df9b..c5190edd 100644 --- a/memoria/crates/memoria-api/src/routes/mcp.rs +++ b/memoria/crates/memoria-api/src/routes/mcp.rs @@ -95,6 +95,26 @@ const GROUP_MAIN_WRITE_BLOCKED: &[&str] = &[ "memory_rollback", ]; +fn mcp_explicit_non_main_branch(params: Option<&serde_json::Value>) -> bool { + params + .and_then(|params| params.get("arguments")) + .and_then(|args| args.get("branch")) + .and_then(|branch| branch.as_str()) + .map(str::trim) + .is_some_and(|branch| !branch.is_empty() && branch != "main") +} + +fn mcp_tool_supports_explicit_branch(tool: &str) -> bool { + matches!( + tool, + "memory_store" | "memory_correct" | "memory_purge" | "memory_observe" + ) +} + +fn mcp_tool_mutates_main_regardless_checkout(tool: &str) -> bool { + matches!(tool, "memory_governance") +} + fn mcp_tool_dirty_mask(tool: &str) -> Option { use crate::metrics_summary::DirtyMask; match tool { @@ -254,38 +274,50 @@ pub async fn mcp_handler( let blocked_tool: Option = if let Some(gid) = &auth.group_id { if let Some(tool) = tracked_tool.as_deref() { if GROUP_MAIN_WRITE_BLOCKED.contains(&tool) { - // ACTOR_USER_ID is already set by the global actor_scope_layer. - let maybe_blocked = state - .service - .user_sql_store(gid) + if mcp_tool_mutates_main_regardless_checkout(tool) { + if crate::auth::group_main_write_allowed_for_solo_owner( + &state, + gid, + &auth.user_id, + ) .await - .ok() - .map(|sql| { + { + None + } else { + Some(tool.to_string()) + } + } else if mcp_tool_supports_explicit_branch(tool) + && mcp_explicit_non_main_branch(params.as_ref()) + { + None + } else { + // ACTOR_USER_ID is already set by the global actor_scope_layer. + let maybe_blocked = state.service.user_sql_store(gid).await.ok().map(|sql| { let gid_owned = gid.clone(); - memoria_storage::ACTOR_USER_ID - .scope(auth.user_id.clone(), async move { - sql.active_branch_name(&gid_owned).await - }) + memoria_storage::ACTOR_USER_ID.scope(auth.user_id.clone(), async move { + sql.active_branch_name(&gid_owned).await + }) }); - if let Some(fut) = maybe_blocked { - if let Ok(branch) = fut.await { - let is_solo_owner = - crate::auth::group_main_write_allowed_for_solo_owner( - &state, - gid, - &auth.user_id, - ) - .await; - if branch == "main" && !is_solo_owner { - Some(tool.to_string()) + if let Some(fut) = maybe_blocked { + if let Ok(branch) = fut.await { + let is_solo_owner = + crate::auth::group_main_write_allowed_for_solo_owner( + &state, + gid, + &auth.user_id, + ) + .await; + if branch == "main" && !is_solo_owner { + Some(tool.to_string()) + } else { + None + } } else { None } } else { None } - } else { - None } } else { None diff --git a/memoria/crates/memoria-api/src/routes/memory.rs b/memoria/crates/memoria-api/src/routes/memory.rs index 3479fec8..ccdd8dfc 100644 --- a/memoria/crates/memoria-api/src/routes/memory.rs +++ b/memoria/crates/memoria-api/src/routes/memory.rs @@ -9,6 +9,7 @@ use sqlx::Row; use crate::{auth::AuthUser, models::*, state::AppState}; use memoria_core::nullable_str_from_row; +use memoria_service::ListActiveOptions; type ApiResult = Result, (StatusCode, String)>; pub fn api_err(e: impl std::fmt::Display) -> (StatusCode, String) { @@ -36,9 +37,13 @@ pub fn api_err_typed(e: memoria_core::MemoriaError) -> (StatusCode, String) { async fn find_memory_any_user( state: &AppState, memory_id: &str, + branch: Option<&str>, ) -> Result, (StatusCode, String)> { let shared = state.service.shared_sql_store().map_err(api_err_typed)?; let Some(router) = shared.db_router() else { + if branch.is_some() { + return Ok(None); + } let memory = state.service.get(memory_id).await.map_err(api_err_typed)?; return Ok(memory.map(|m| (m.user_id.clone(), m))); }; @@ -50,7 +55,11 @@ async fn find_memory_any_user( .user_sql_store(&candidate) .await .map_err(api_err_typed)?; - let table = sql.active_table(&candidate).await.map_err(api_err_typed)?; + let table = match sql.table_for_branch(&candidate, branch).await { + Ok(table) => table, + Err(memoria_core::MemoriaError::NotFound(_)) if branch.is_some() => continue, + Err(err) => return Err(api_err_typed(err)), + }; if let Some(memory) = sql .get_from(&table, memory_id) .await @@ -68,6 +77,7 @@ pub struct ListQuery { pub memory_type: Option, pub session_id: Option, pub trust_tier: Option, + pub branch: Option, #[serde(default = "default_limit")] pub limit: i64, pub cursor: Option, @@ -76,6 +86,21 @@ fn default_limit() -> i64 { 100 } +#[derive(Deserialize, Default)] +pub struct BranchQuery { + pub branch: Option, +} + +fn branch_param(branch: Option<&str>) -> Option<&str> { + branch.map(str::trim).filter(|branch| !branch.is_empty()) +} + +fn normalize_branch(branch: Option) -> Option { + branch + .map(|branch| branch.trim().to_string()) + .filter(|branch| !branch.is_empty()) +} + pub async fn health() -> &'static str { "ok" } @@ -121,13 +146,16 @@ pub async fn list_memories( let fetch_limit = limit + 1; let mut memories = state .service - .list_active_paged( + .list_active_paged_on_branch( auth.scope_id(), - fetch_limit, - q.memory_type.as_deref(), - q.session_id.as_deref(), - q.trust_tier.as_deref(), - cursor, + branch_param(q.branch.as_deref()), + ListActiveOptions { + limit: fetch_limit, + memory_type: q.memory_type.as_deref(), + session_id: q.session_id.as_deref(), + trust_tier: q.trust_tier.as_deref(), + cursor, + }, ) .await .map_err(api_err)?; @@ -180,8 +208,9 @@ pub async fn store_memory( }; let m = state .service - .store_memory( + .store_memory_on_branch( auth.scope_id(), + branch_param(req.branch.as_deref()), &req.content, mt, req.session_id, @@ -219,10 +248,18 @@ pub async fn batch_store( )); } } + let top_branch = normalize_branch(req.branch); let items: Vec<_> = req .memories .into_iter() .map(|r| { + let item_branch = normalize_branch(r.branch); + if item_branch.is_some() && item_branch.as_deref() != top_branch.as_deref() { + return Err(( + StatusCode::UNPROCESSABLE_ENTITY, + "per-memory branch in batch must match the batch branch".to_string(), + )); + } let mt = parse_memory_type(&r.memory_type) .map_err(|e| (StatusCode::UNPROCESSABLE_ENTITY, e)); let tier = r @@ -231,16 +268,16 @@ pub async fn batch_store( .map(parse_trust_tier) .transpose() .map_err(|e| (StatusCode::UNPROCESSABLE_ENTITY, e)); - (r.content, mt, tier, r.session_id) + Ok((r.content, mt, tier, r.session_id, top_branch.clone())) }) - .collect(); + .collect::, _>>()?; // Validate all types upfront let mut validated = Vec::with_capacity(items.len()); - for (content, mt_result, tier_result, session_id) in items { + for (content, mt_result, tier_result, session_id, branch) in items { let mt = mt_result?; let tier = tier_result?; - validated.push((content, mt, session_id, tier)); + validated.push((content, mt, session_id, tier, branch)); } let author = if auth.group_id.is_some() { @@ -248,9 +285,18 @@ pub async fn batch_store( } else { None }; + let batch_items = validated + .into_iter() + .map(|(content, mt, session_id, tier, _branch)| (content, mt, session_id, tier)) + .collect(); let results = state .service - .store_batch(auth.scope_id(), validated, author) + .store_batch_on_branch( + auth.scope_id(), + branch_param(top_branch.as_deref()), + batch_items, + author, + ) .await .map_err(api_err)?; Ok(( @@ -275,8 +321,9 @@ pub async fn retrieve( if level != memoria_service::ExplainLevel::None { let (results, explain) = state .service - .retrieve_explain_level_with_options( + .retrieve_explain_level_with_options_on_branch( auth.scope_id(), + branch_param(req.branch.as_deref()), &req.query, top_k, level, @@ -291,7 +338,13 @@ pub async fn retrieve( } else { let results = state .service - .retrieve_with_options(auth.scope_id(), &req.query, top_k, &retrieve_options) + .retrieve_with_options_on_branch( + auth.scope_id(), + branch_param(req.branch.as_deref()), + &req.query, + top_k, + &retrieve_options, + ) .await .map_err(api_err)?; let items: Vec = results.into_iter().map(Into::into).collect(); @@ -314,8 +367,9 @@ pub async fn search( if level != memoria_service::ExplainLevel::None { let (results, explain) = state .service - .retrieve_explain_level_with_options( + .retrieve_explain_level_with_options_on_branch( auth.scope_id(), + branch_param(req.branch.as_deref()), &req.query, top_k, level, @@ -330,7 +384,13 @@ pub async fn search( } else { let results = state .service - .retrieve_with_options(auth.scope_id(), &req.query, top_k, &retrieve_options) + .retrieve_with_options_on_branch( + auth.scope_id(), + branch_param(req.branch.as_deref()), + &req.query, + top_k, + &retrieve_options, + ) .await .map_err(api_err)?; Ok(Json(serde_json::json!(results @@ -344,12 +404,18 @@ pub async fn get_memory( State(state): State, auth: AuthUser, Path(id): Path, + Query(q): Query, ) -> ApiResult> { - let owned = state + let branch = branch_param(q.branch.as_deref()); + let owned = match state .service - .get_for_user(auth.scope_id(), &id) + .get_for_user_on_branch(auth.scope_id(), branch, &id) .await - .map_err(api_err_typed)?; + { + Ok(memory) => memory, + Err(memoria_core::MemoriaError::NotFound(_)) if auth.is_master && branch.is_some() => None, + Err(err) => return Err(api_err_typed(err)), + }; if let Some(memory) = owned { return Ok(Json(Some(memory.into()))); } @@ -358,7 +424,7 @@ pub async fn get_memory( return Ok(Json(None)); } - if let Some((_, memory)) = find_memory_any_user(&state, &id).await? { + if let Some((_, memory)) = find_memory_any_user(&state, &id, branch).await? { return Ok(Json(Some(memory.into()))); } @@ -372,14 +438,14 @@ pub async fn correct_memory( Json(req): Json, ) -> ApiResult { let effective_user_id = if auth.is_master { - find_memory_any_user(&state, &id) + find_memory_any_user(&state, &id, branch_param(req.branch.as_deref())) .await? .map(|(owner_id, _)| owner_id) .unwrap_or_else(|| auth.scope_id().to_string()) } else { if state .service - .get_for_user(auth.scope_id(), &id) + .get_for_user_on_branch(auth.scope_id(), branch_param(req.branch.as_deref()), &id) .await .map_err(api_err_typed)? .is_none() @@ -390,7 +456,12 @@ pub async fn correct_memory( }; let m = state .service - .correct(&effective_user_id, &id, &req.new_content) + .correct_on_branch( + &effective_user_id, + branch_param(req.branch.as_deref()), + &id, + &req.new_content, + ) .await .map_err(api_err_typed)?; Ok(Json(m.into())) @@ -408,7 +479,13 @@ pub async fn correct_by_query( .map_err(|err| (StatusCode::UNPROCESSABLE_ENTITY, err))?; let results = state .service - .retrieve_with_options(auth.scope_id(), &req.query, 1, &retrieve_options) + .retrieve_with_options_on_branch( + auth.scope_id(), + branch_param(req.branch.as_deref()), + &req.query, + 1, + &retrieve_options, + ) .await .map_err(api_err)?; let found = results.into_iter().next().ok_or_else(|| { @@ -419,7 +496,12 @@ pub async fn correct_by_query( })?; let m = state .service - .correct(auth.scope_id(), &found.memory_id, &req.new_content) + .correct_on_branch( + auth.scope_id(), + branch_param(req.branch.as_deref()), + &found.memory_id, + &req.new_content, + ) .await .map_err(api_err)?; Ok(Json(m.into())) @@ -429,16 +511,17 @@ pub async fn delete_memory( State(state): State, auth: AuthUser, Path(id): Path, + Query(q): Query, ) -> Result { let effective_user_id = if auth.is_master { - find_memory_any_user(&state, &id) + find_memory_any_user(&state, &id, branch_param(q.branch.as_deref())) .await? .map(|(owner_id, _)| owner_id) .unwrap_or_else(|| auth.scope_id().to_string()) } else { if state .service - .get_for_user(auth.scope_id(), &id) + .get_for_user_on_branch(auth.scope_id(), branch_param(q.branch.as_deref()), &id) .await .map_err(api_err_typed)? .is_none() @@ -449,7 +532,7 @@ pub async fn delete_memory( }; let _ = state .service - .purge(&effective_user_id, &id) + .purge_on_branch(&effective_user_id, branch_param(q.branch.as_deref()), &id) .await .map_err(api_err_typed)?; Ok(StatusCode::NO_CONTENT) @@ -463,13 +546,14 @@ pub async fn purge_memories( let selector = req .selector() .map_err(|e| (StatusCode::UNPROCESSABLE_ENTITY, e))?; + let branch = branch_param(req.branch.as_deref()); let result = match selector { PurgeSelector::MemoryIds(ids) => { if !auth.is_master { for id in &ids { let mem = state .service - .get_for_user(auth.scope_id(), id) + .get_for_user_on_branch(auth.scope_id(), branch, id) .await .map_err(api_err)? .ok_or_else(|| { @@ -483,13 +567,13 @@ pub async fn purge_memories( let id_refs: Vec<&str> = ids.iter().map(|s| s.as_str()).collect(); state .service - .purge_batch(auth.scope_id(), &id_refs) + .purge_batch_on_branch(auth.scope_id(), branch, &id_refs) .await .map_err(api_err)? } PurgeSelector::Topic(topic) => state .service - .purge_by_topic(auth.scope_id(), &topic) + .purge_by_topic_on_branch(auth.scope_id(), branch, &topic) .await .map_err(api_err)?, PurgeSelector::Session { @@ -497,7 +581,12 @@ pub async fn purge_memories( memory_types, } => state .service - .purge_by_session_id(auth.scope_id(), &session_id, memory_types.as_deref()) + .purge_by_session_id_on_branch( + auth.scope_id(), + branch, + &session_id, + memory_types.as_deref(), + ) .await .map_err(api_err)?, PurgeSelector::None => memoria_service::PurgeResult { @@ -591,6 +680,7 @@ pub struct ObserveRequest { pub messages: Vec, pub source_event_ids: Option>, pub session_id: Option, + pub branch: Option, } /// Extract and store memories from a conversation turn. @@ -603,7 +693,12 @@ pub async fn observe_turn( ) -> ApiResult { let (memories, has_llm) = state .service - .observe_turn(auth.scope_id(), &req.messages, req.session_id) + .observe_turn_on_branch( + auth.scope_id(), + branch_param(req.branch.as_deref()), + &req.messages, + req.session_id, + ) .await .map_err(api_err)?; @@ -630,6 +725,7 @@ pub async fn get_memory_history( State(state): State, auth: AuthUser, Path(id): Path, + Query(q): Query, ) -> ApiResult { use sqlx::Row; @@ -638,7 +734,10 @@ pub async fn get_memory_history( .user_sql_store(auth.scope_id()) .await .map_err(api_err)?; - let table = sql.active_table(auth.scope_id()).await.map_err(api_err)?; + let table = sql + .table_for_branch(auth.scope_id(), branch_param(q.branch.as_deref())) + .await + .map_err(api_err_typed)?; let mut chain = Vec::new(); let mut visited = std::collections::HashSet::new(); diff --git a/memoria/crates/memoria-api/src/routes/snapshots.rs b/memoria/crates/memoria-api/src/routes/snapshots.rs index c2de692f..fa73144d 100644 --- a/memoria/crates/memoria-api/src/routes/snapshots.rs +++ b/memoria/crates/memoria-api/src/routes/snapshots.rs @@ -860,7 +860,9 @@ pub async fn list_branches( "name": "main", "active": active_branch == "main", })]; - for (name, _table_name, created_at) in sql.list_branches(auth.scope_id()).await.map_err(api_err)? { + for (name, _table_name, created_at) in + sql.list_branches(auth.scope_id()).await.map_err(api_err)? + { let created_at_str = format_snapshot_timestamp(created_at); branches.push(json!({ "name": name, diff --git a/memoria/crates/memoria-cli/src/main.rs b/memoria/crates/memoria-cli/src/main.rs index 6c567b8b..4ca5c9a0 100644 --- a/memoria/crates/memoria-cli/src/main.rs +++ b/memoria/crates/memoria-cli/src/main.rs @@ -647,10 +647,7 @@ fn normalize_tool_name(tool: Option) -> Option { // trim, lowercase, collapse whitespace runs into a single hyphen. // "My Agent" → "my-agent", "cursor" → "cursor". let lower = raw.trim().to_ascii_lowercase(); - let normalized = lower - .split_whitespace() - .collect::>() - .join("-"); + let normalized = lower.split_whitespace().collect::>().join("-"); if normalized.is_empty() { None } else { diff --git a/memoria/crates/memoria-git/src/service.rs b/memoria/crates/memoria-git/src/service.rs index 4b217b52..1d97dcd7 100644 --- a/memoria/crates/memoria-git/src/service.rs +++ b/memoria/crates/memoria-git/src/service.rs @@ -1079,7 +1079,11 @@ impl GitForDataService { // Cannot use INSERT...ON DUPLICATE KEY UPDATE because MatrixOne branch tables are // read-only views in some builds. Instead: DELETE + INSERT. if !restore_adds.is_empty() { - let ph = restore_adds.iter().map(|_| "?").collect::>().join(","); + let ph = restore_adds + .iter() + .map(|_| "?") + .collect::>() + .join(","); let delete_sql = format!( "DELETE FROM {main_table_ref} WHERE user_id = ? AND memory_id IN ({ph})" ); diff --git a/memoria/crates/memoria-mcp/src/purge_args.rs b/memoria/crates/memoria-mcp/src/purge_args.rs index 174099b0..7f72d466 100644 --- a/memoria/crates/memoria-mcp/src/purge_args.rs +++ b/memoria/crates/memoria-mcp/src/purge_args.rs @@ -8,6 +8,7 @@ pub(crate) struct MemoryPurgeArgs { pub(crate) topic: Option, pub(crate) session_id: Option, pub(crate) memory_types: Option>, + pub(crate) branch: Option, } pub(crate) fn parse_memory_purge_args(args: &Value) -> Result { @@ -26,6 +27,11 @@ pub(crate) fn parse_memory_purge_args(args: &Value) -> Result { .map(str::trim) .filter(|s| !s.is_empty()) .map(str::to_string); + let branch = args["branch"] + .as_str() + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_string); let memory_types = match args.get("memory_types") { None | Some(Value::Null) => None, Some(Value::Array(values)) => Some( @@ -60,5 +66,6 @@ pub(crate) fn parse_memory_purge_args(args: &Value) -> Result { topic, session_id, memory_types, + branch, }) } diff --git a/memoria/crates/memoria-mcp/src/remote.rs b/memoria/crates/memoria-mcp/src/remote.rs index ae8447fc..f789c228 100644 --- a/memoria/crates/memoria-mcp/src/remote.rs +++ b/memoria/crates/memoria-mcp/src/remote.rs @@ -130,6 +130,7 @@ impl RemoteClient { "memory_type": args["memory_type"].as_str().unwrap_or("semantic"), "session_id": args["session_id"], "trust_tier": args["trust_tier"], + "branch": args["branch"], })) .send() .await?; @@ -155,6 +156,7 @@ impl RemoteClient { args["top_k"].as_i64().unwrap_or(5) }, "session_id": args["session_id"], + "branch": args["branch"], }); if let Some(session_scope) = args .get("session_scope") @@ -195,7 +197,11 @@ impl RemoteClient { let r = if !memory_id.is_empty() { self.client .put(self.url(&format!("/v1/memories/{memory_id}/correct"))) - .json(&json!({"new_content": new_content, "reason": args["reason"]})) + .json(&json!({ + "new_content": new_content, + "reason": args["reason"], + "branch": args["branch"] + })) .send() .await? } else { @@ -206,7 +212,8 @@ impl RemoteClient { "new_content": new_content, "session_id": args["session_id"], "session_scope": args["session_scope"], - "reason": args["reason"] + "reason": args["reason"], + "branch": args["branch"] })) .send() .await? @@ -236,7 +243,7 @@ impl RemoteClient { let r = self .client .post(self.url("/v1/memories/purge")) - .json(&json!({"memory_ids": ids})) + .json(&json!({"memory_ids": ids, "branch": purge_args.branch})) .send() .await?; let body = Self::parse_response(r).await?; @@ -248,7 +255,7 @@ impl RemoteClient { let r = self .client .post(self.url("/v1/memories/purge")) - .json(&json!({"topic": topic})) + .json(&json!({"topic": topic, "branch": purge_args.branch})) .send() .await?; let body = Self::parse_response(r).await?; @@ -268,6 +275,7 @@ impl RemoteClient { .map(|memory_type| memory_type.to_string()) .collect::>() }), + "branch": purge_args.branch, })) .send() .await?; @@ -307,6 +315,9 @@ impl RemoteClient { if let Some(session_id) = session_id { req = req.query(&[("session_id", session_id)]); } + if let Some(branch) = args.get("branch").and_then(Value::as_str) { + req = req.query(&[("branch", branch)]); + } let body = Self::parse_response(req.send().await?).await?; let items = body["items"].as_array().cloned().unwrap_or_default(); if items.is_empty() { @@ -688,10 +699,11 @@ impl RemoteClient { "memory_observe" => { let r = self .client - .post(self.url("/v1/memories/observe")) + .post(self.url("/v1/observe")) .json(&json!({ "messages": args["messages"], "session_id": args["session_id"], + "branch": args["branch"], })) .send() .await?; diff --git a/memoria/crates/memoria-mcp/src/server.rs b/memoria/crates/memoria-mcp/src/server.rs index e44eeca9..41cdc1ac 100644 --- a/memoria/crates/memoria-mcp/src/server.rs +++ b/memoria/crates/memoria-mcp/src/server.rs @@ -100,9 +100,7 @@ pub async fn dispatch_http( // (`active_branch_name`, `set_active_branch`, `active_table`) see the // same per-user scope they would in the original async task. match actor_user_id { - Some(actor_id) => handle.block_on( - memoria_storage::ACTOR_USER_ID.scope(actor_id, fut) - ), + Some(actor_id) => handle.block_on(memoria_storage::ACTOR_USER_ID.scope(actor_id, fut)), None => handle.block_on(fut), } }) diff --git a/memoria/crates/memoria-mcp/src/tools.rs b/memoria/crates/memoria-mcp/src/tools.rs index b4d06e4d..8a82cbf4 100644 --- a/memoria/crates/memoria-mcp/src/tools.rs +++ b/memoria/crates/memoria-mcp/src/tools.rs @@ -6,7 +6,8 @@ use anyhow::Result; use memoria_core::{MemoryType, TrustTier}; use memoria_git::GitForDataService; use memoria_service::{ - ConsolidationInput, ConsolidationStrategy, DefaultConsolidationStrategy, MemoryService, + ConsolidationInput, ConsolidationStrategy, DefaultConsolidationStrategy, ListActiveOptions, + MemoryService, }; use memoria_storage::SqlMemoryStore; use serde_json::{json, Value}; @@ -54,6 +55,13 @@ fn parse_retrieve_options_arg(args: &Value) -> Result Option<&str> { + args.get("branch") + .and_then(Value::as_str) + .map(str::trim) + .filter(|branch| !branch.is_empty()) +} + enum ToolCallName { MemoryStore, MemoryRetrieve, @@ -116,6 +124,7 @@ pub fn list() -> Value { "content": {"type": "string"}, "memory_type": {"type": "string", "default": "semantic"}, "session_id": {"type": "string"}, + "branch": {"type": "string", "description": "Optional branch to read/write without changing the active checkout"}, "trust_tier": { "type": "string", "enum": ["T1", "T2", "T3", "T4"], @@ -135,6 +144,7 @@ pub fn list() -> Value { "top_k": {"type": "integer", "default": 5}, "session_id": {"type": "string"}, "session_scope": {"type": "string", "enum": ["prefer", "only"], "description": "How to use session_id: prefer=non-strict retrieval using the provided session_id as context; only=limit results to that session plus unscoped memories"}, + "branch": {"type": "string", "description": "Optional branch to read from without changing the active checkout"}, "explain": {"type": ["boolean", "string"], "default": false, "description": "Explain level: false/\"none\"=off, true/\"basic\"=timing+path, \"verbose\"=+per-candidate scores, \"analyze\"=full"} }, "required": ["query"] @@ -150,6 +160,7 @@ pub fn list() -> Value { "top_k": {"type": "integer", "default": 10}, "session_id": {"type": "string"}, "session_scope": {"type": "string", "enum": ["prefer", "only"], "description": "How to use session_id: prefer=non-strict search using the provided session_id as context; only=limit results to that session plus unscoped memories"}, + "branch": {"type": "string", "description": "Optional branch to search without changing the active checkout"}, "explain": {"type": ["boolean", "string"], "default": false, "description": "Explain level: false/\"none\"=off, true/\"basic\"=timing+path, \"verbose\"=+per-candidate scores, \"analyze\"=full"} }, "required": ["query"] @@ -166,6 +177,7 @@ pub fn list() -> Value { "new_content": {"type": "string"}, "session_id": {"type": "string", "description": "Optional session to use when resolving query-based correction"}, "session_scope": {"type": "string", "enum": ["prefer", "only"], "description": "Only used with query-based correction. prefer=non-strict lookup using the provided session_id as context; only=restrict lookup to the given session plus unscoped memories"}, + "branch": {"type": "string", "description": "Optional branch to correct without changing the active checkout"}, "reason": {"type": "string"} }, "required": ["new_content"] @@ -181,6 +193,7 @@ pub fn list() -> Value { "topic": {"type": "string", "description": "Keyword — bulk-delete all matching memories"}, "session_id": {"type": "string", "description": "Exact session identifier — bulk-delete memories from that session"}, "memory_types": {"type": "array", "items": {"type": "string", "enum": ["semantic", "working", "episodic", "profile", "tool_result", "procedural"]}, "description": "Optional memory type filter. Only valid with session_id"}, + "branch": {"type": "string", "description": "Optional branch to purge from without changing the active checkout"}, "reason": {"type": "string"} } } @@ -203,7 +216,8 @@ pub fn list() -> Value { "properties": { "limit": {"type": "integer", "default": 20}, "memory_type": {"type": "string"}, - "session_id": {"type": "string", "description": "Exact session identifier — only list memories from that session"} + "session_id": {"type": "string", "description": "Exact session identifier — only list memories from that session"}, + "branch": {"type": "string", "description": "Optional branch to list without changing the active checkout"} } } }, @@ -294,8 +308,9 @@ pub async fn call( .map_err(|e| anyhow::anyhow!("{e}"))?; let mt = MemoryType::from_str(memory_type).unwrap_or(MemoryType::Semantic); let m = match service - .store_memory( + .store_memory_on_branch( user_id, + branch_arg(&args), &content, mt, session_id.clone(), @@ -383,8 +398,9 @@ pub async fn call( if level != memoria_service::ExplainLevel::None { let (results, stats) = service - .retrieve_explain_level_with_options( + .retrieve_explain_level_with_options_on_branch( user_id, + branch_arg(&args), &query, top_k, level, @@ -408,7 +424,13 @@ pub async fn call( ))) } else { let results = service - .retrieve_with_options(user_id, &query, top_k, &retrieve_options) + .retrieve_with_options_on_branch( + user_id, + branch_arg(&args), + &query, + top_k, + &retrieve_options, + ) .await?; if results.is_empty() { return Ok(mcp_text("No relevant memories found.")); @@ -436,7 +458,13 @@ pub async fn call( } else if !query.is_empty() { let retrieve_options = parse_retrieve_options_arg(&args)?; let results = service - .retrieve_with_options(user_id, query, 1, &retrieve_options) + .retrieve_with_options_on_branch( + user_id, + branch_arg(&args), + query, + 1, + &retrieve_options, + ) .await?; match results.into_iter().next() { Some(found) => found.memory_id, @@ -446,7 +474,9 @@ pub async fn call( return Ok(mcp_text("Provide memory_id or query")); }; - let m = service.correct(user_id, &old_mid, new_content).await?; + let m = service + .correct_on_branch(user_id, branch_arg(&args), &old_mid, new_content) + .await?; Ok(mcp_text(&format!( "Corrected memory {}: {}", @@ -463,21 +493,30 @@ pub async fn call( .map(str::trim) .filter(|s| !s.is_empty()) .collect(); - let result = service.purge_batch(user_id, &ids).await?; + let result = service + .purge_batch_on_branch(user_id, branch_arg(&args), &ids) + .await?; Ok(mcp_text(&format_purge_msg( &format!("Purged {} memory(s)", result.purged), &result, ))) } else if let Some(topic) = purge_args.topic { // Bulk by keyword: exact text match then purge - let result = service.purge_by_topic(user_id, &topic).await?; + let result = service + .purge_by_topic_on_branch(user_id, branch_arg(&args), &topic) + .await?; Ok(mcp_text(&format_purge_msg( &format!("Purged {} memory(s) matching '{topic}'", result.purged), &result, ))) } else if let Some(session_id) = purge_args.session_id { let result = service - .purge_by_session_id(user_id, &session_id, purge_args.memory_types.as_deref()) + .purge_by_session_id_on_branch( + user_id, + branch_arg(&args), + &session_id, + purge_args.memory_types.as_deref(), + ) .await?; Ok(mcp_text(&format_purge_msg( &format!( @@ -511,13 +550,16 @@ pub async fn call( ToolCallName::MemoryList => { let limit = args["limit"].as_i64().unwrap_or(20); let memories = service - .list_active_filtered( + .list_active_filtered_on_branch( user_id, - limit, - args.get("memory_type").and_then(Value::as_str), - args.get("session_id").and_then(Value::as_str), - args.get("trust_tier").and_then(Value::as_str), - None, + branch_arg(&args), + ListActiveOptions { + limit, + memory_type: args.get("memory_type").and_then(Value::as_str), + session_id: args.get("session_id").and_then(Value::as_str), + trust_tier: args.get("trust_tier").and_then(Value::as_str), + cursor: None, + }, ) .await?; if memories.is_empty() { @@ -1026,7 +1068,7 @@ pub async fn call( let session_id = args["session_id"].as_str().map(String::from); let (memories, has_llm) = service - .observe_turn(user_id, &messages, session_id.clone()) + .observe_turn_on_branch(user_id, branch_arg(&args), &messages, session_id.clone()) .await?; // Graph sync (best-effort) for each stored memory diff --git a/memoria/crates/memoria-service/src/lib.rs b/memoria/crates/memoria-service/src/lib.rs index be3b7446..eb3d7e9e 100644 --- a/memoria/crates/memoria-service/src/lib.rs +++ b/memoria/crates/memoria-service/src/lib.rs @@ -45,8 +45,8 @@ pub use scoring::{ DefaultScoringPlugin, FeedbackTotals, ScoringPlugin, ScoringStore, TuningResult, }; pub use service::{ - CandidateScore, ExplainLevel, InMemoryFlusher, MemoryService, PurgeResult, RetrievalExplain, - RetrieveOptions, SessionScope, ENTITY_EXTRACTION_DROPS, + CandidateScore, ExplainLevel, InMemoryFlusher, ListActiveOptions, MemoryService, PurgeResult, + RetrievalExplain, RetrieveOptions, SessionScope, ENTITY_EXTRACTION_DROPS, }; pub use stats_reporter::StatsReporter; diff --git a/memoria/crates/memoria-service/src/service.rs b/memoria/crates/memoria-service/src/service.rs index c48e97fb..6d69b613 100644 --- a/memoria/crates/memoria-service/src/service.rs +++ b/memoria/crates/memoria-service/src/service.rs @@ -120,6 +120,27 @@ impl RetrieveOptions { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ListActiveOptions<'a> { + pub limit: i64, + pub memory_type: Option<&'a str>, + pub session_id: Option<&'a str>, + pub trust_tier: Option<&'a str>, + pub cursor: Option<&'a str>, +} + +impl ListActiveOptions<'_> { + pub fn new(limit: i64) -> Self { + Self { + limit, + memory_type: None, + session_id: None, + trust_tier: None, + cursor: None, + } + } +} + /// Per-candidate scoring breakdown — answers "why is this memory ranked here?" /// Only populated at Verbose/Analyze level. #[derive(Debug, serde::Serialize)] @@ -1127,7 +1148,6 @@ impl MemoryService { } #[allow(clippy::too_many_arguments)] - #[tracing::instrument(skip(self, content), fields(user_id))] pub async fn store_memory( &self, user_id: &str, @@ -1138,6 +1158,34 @@ impl MemoryService { observed_at: Option>, initial_confidence: Option, author_id: Option, + ) -> Result { + self.store_memory_on_branch( + user_id, + None, + content, + memory_type, + session_id, + trust_tier, + observed_at, + initial_confidence, + author_id, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + #[tracing::instrument(skip(self, content), fields(user_id, branch))] + pub async fn store_memory_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + content: &str, + memory_type: MemoryType, + session_id: Option, + trust_tier: Option, + observed_at: Option>, + initial_confidence: Option, + author_id: Option, ) -> Result { let t0 = std::time::Instant::now(); // Sensitivity check — block HIGH tier, redact MEDIUM tier @@ -1178,7 +1226,7 @@ impl MemoryService { // TODO(concurrency): race between dedup check and insert can create duplicates if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; if let Some(ref emb) = memory.embedding { // L2 threshold from cosine similarity 0.95: sqrt(2*(1-0.95)) ≈ 0.3162 // Only supersede near-identical memories, not contradictions. @@ -1389,9 +1437,21 @@ impl MemoryService { query: &str, top_k: i64, options: &RetrieveOptions, + ) -> Result, MemoriaError> { + self.retrieve_with_options_on_branch(user_id, None, query, top_k, options) + .await + } + + pub async fn retrieve_with_options_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + query: &str, + top_k: i64, + options: &RetrieveOptions, ) -> Result, MemoriaError> { let (mems, _) = self - .retrieve_inner(user_id, query, top_k, ExplainLevel::None, options) + .retrieve_inner(user_id, branch, query, top_k, ExplainLevel::None, options) .await?; self.bump_access_counts(&mems); Ok(mems) @@ -1407,6 +1467,7 @@ impl MemoryService { let (mems, explain) = self .retrieve_inner( user_id, + None, query, top_k, ExplainLevel::Basic, @@ -1442,10 +1503,25 @@ impl MemoryService { top_k: i64, level: ExplainLevel, options: &RetrieveOptions, + ) -> Result<(Vec, RetrievalExplain), MemoriaError> { + self.retrieve_explain_level_with_options_on_branch( + user_id, None, query, top_k, level, options, + ) + .await + } + + pub async fn retrieve_explain_level_with_options_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + query: &str, + top_k: i64, + level: ExplainLevel, + options: &RetrieveOptions, ) -> Result<(Vec, RetrievalExplain), MemoriaError> { let start = std::time::Instant::now(); let (mems, explain) = self - .retrieve_inner(user_id, query, top_k, level, options) + .retrieve_inner(user_id, branch, query, top_k, level, options) .await?; self.bump_access_counts(&mems); @@ -1473,6 +1549,7 @@ impl MemoryService { async fn retrieve_inner( &self, user_id: &str, + branch: Option<&str>, query: &str, top_k: i64, level: ExplainLevel, @@ -1486,7 +1563,7 @@ impl MemoryService { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let strict_session_id = options.strict_session_id(); // Load per-user feedback_weight lazily — only when needed for scoring // (avoids extra DB query when fulltext fallback has no feedback to apply) @@ -1497,7 +1574,8 @@ impl MemoryService { explain.embedding_ms = p0_start.elapsed().as_secs_f64() * 1000.0; // Phase 1: graph retrieval (activation-based) - if strict_session_id.is_none() { + let main_table = sql.t("mem_memories"); + if strict_session_id.is_none() && table == main_table { if let Some(ref embedding) = emb { explain.graph_attempted = true; let g_start = std::time::Instant::now(); @@ -1767,6 +1845,7 @@ impl MemoryService { ) -> Result<(Vec, RetrievalExplain), MemoriaError> { self.retrieve_inner( user_id, + None, query, top_k, ExplainLevel::Basic, @@ -1782,8 +1861,15 @@ impl MemoryService { top_k: i64, level: ExplainLevel, ) -> Result<(Vec, RetrievalExplain), MemoriaError> { - self.retrieve_inner(user_id, query, top_k, level, &RetrieveOptions::default()) - .await + self.retrieve_inner( + user_id, + None, + query, + top_k, + level, + &RetrieveOptions::default(), + ) + .await } // TODO(concurrency): concurrent correct on same memory_id can create duplicate @@ -1793,6 +1879,17 @@ impl MemoryService { user_id: &str, memory_id: &str, new_content: &str, + ) -> Result { + self.correct_on_branch(user_id, None, memory_id, new_content) + .await + } + + pub async fn correct_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + memory_id: &str, + new_content: &str, ) -> Result { // Sensitivity check — same as store_memory let sensitivity = check_sensitivity(new_content); @@ -1810,7 +1907,7 @@ impl MemoryService { // Branch-aware: resolve table and fetch old memory from correct table if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let old = sql .get_from(&table, memory_id) .await? @@ -1913,11 +2010,20 @@ impl MemoryService { } pub async fn purge(&self, user_id: &str, memory_id: &str) -> Result { + self.purge_on_branch(user_id, None, memory_id).await + } + + pub async fn purge_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + memory_id: &str, + ) -> Result { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; let (snap, warning) = sql.create_safety_snapshot("purge").await; self.send_edit_log(user_id, "purge", Some(memory_id), None, "", snap.as_deref()); - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let deactivated = sql.soft_delete_from(&table, memory_id).await?; if deactivated > 0 { self.report(StatsEvent::MemoryDeactivated { @@ -1946,11 +2052,20 @@ impl MemoryService { &self, user_id: &str, ids: &[&str], + ) -> Result { + self.purge_batch_on_branch(user_id, None, ids).await + } + + pub async fn purge_batch_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + ids: &[&str], ) -> Result { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; let (snap, warning) = sql.create_safety_snapshot("purge").await; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; for id in ids { let deactivated = sql.soft_delete_from(&table, id).await?; self.send_edit_log(user_id, "purge", Some(id), None, "", snap.as_deref()); @@ -1983,11 +2098,20 @@ impl MemoryService { &self, user_id: &str, topic: &str, + ) -> Result { + self.purge_by_topic_on_branch(user_id, None, topic).await + } + + pub async fn purge_by_topic_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + topic: &str, ) -> Result { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; let (snap, warning) = sql.create_safety_snapshot("purge").await; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let ids = sql.find_ids_by_topic(&table, user_id, topic).await?; let reason = format!("topic:{topic}"); self.purge_sql_ids( @@ -2018,11 +2142,22 @@ impl MemoryService { user_id: &str, session_id: &str, memory_types: Option<&[MemoryType]>, + ) -> Result { + self.purge_by_session_id_on_branch(user_id, None, session_id, memory_types) + .await + } + + pub async fn purge_by_session_id_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + session_id: &str, + memory_types: Option<&[MemoryType]>, ) -> Result { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; let (snap, warning) = sql.create_safety_snapshot("purge").await; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let ids = sql .find_ids_by_session_id(&table, user_id, session_id, memory_types) .await?; @@ -2099,10 +2234,19 @@ impl MemoryService { &self, user_id: &str, memory_id: &str, + ) -> Result, MemoriaError> { + self.get_for_user_on_branch(user_id, None, memory_id).await + } + + pub async fn get_for_user_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + memory_id: &str, ) -> Result, MemoriaError> { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; return sql.get_from(&table, memory_id).await; } self.store.get(memory_id).await @@ -2131,6 +2275,16 @@ impl MemoryService { .await } + pub async fn list_active_paged_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + options: ListActiveOptions<'_>, + ) -> Result, MemoriaError> { + self.list_active_filtered_on_branch(user_id, branch, options) + .await + } + pub async fn list_active_filtered( &self, user_id: &str, @@ -2139,27 +2293,55 @@ impl MemoryService { session_id: Option<&str>, trust_tier: Option<&str>, cursor: Option<&str>, + ) -> Result, MemoriaError> { + self.list_active_filtered_on_branch( + user_id, + None, + ListActiveOptions { + limit, + memory_type, + session_id, + trust_tier, + cursor, + }, + ) + .await + } + + pub async fn list_active_filtered_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + options: ListActiveOptions<'_>, ) -> Result, MemoriaError> { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; return sql - .list_active_lite(&table, user_id, limit, memory_type, session_id, trust_tier, cursor) + .list_active_lite( + &table, + user_id, + options.limit, + options.memory_type, + options.session_id, + options.trust_tier, + options.cursor, + ) .await; } // Fallback: trait path — no SQL store means no server-side filter/cursor. // Production always uses SQL store; this path is for trait-only test doubles. - let mut mems = self.store.list_active(user_id, limit).await?; - if let Some(mt) = memory_type { + let mut mems = self.store.list_active(user_id, options.limit).await?; + if let Some(mt) = options.memory_type { mems.retain(|m| m.memory_type.to_string() == mt); } - if let Some(session_id) = session_id { + if let Some(session_id) = options.session_id { mems.retain(|m| m.session_id.as_deref() == Some(session_id)); } - if let Some(tt) = trust_tier { + if let Some(tt) = options.trust_tier { mems.retain(|m| m.trust_tier.to_string() == tt); } - if let Some(cursor_id) = cursor { + if let Some(cursor_id) = options.cursor { mems.retain(|m| m.memory_id.as_str() < cursor_id); } Ok(mems) @@ -2218,6 +2400,18 @@ impl MemoryService { user_id: &str, items: Vec<(String, MemoryType, Option, Option)>, author_id: Option, + ) -> Result, MemoriaError> { + self.store_batch_on_branch(user_id, None, items, author_id) + .await + } + + /// Batch store with a single embedding API call, targeting an optional branch. + pub async fn store_batch_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + items: Vec<(String, MemoryType, Option, Option)>, + author_id: Option, ) -> Result, MemoriaError> { if items.is_empty() { return Ok(vec![]); @@ -2270,7 +2464,7 @@ impl MemoryService { } if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; let refs: Vec<&Memory> = results.iter().collect(); sql.batch_insert_into(&table, &refs).await?; let payloads: Vec = results @@ -2308,6 +2502,17 @@ impl MemoryService { user_id: &str, messages: &[serde_json::Value], session_id: Option, + ) -> Result<(Vec, bool), MemoriaError> { + self.observe_turn_on_branch(user_id, None, messages, session_id) + .await + } + + pub async fn observe_turn_on_branch( + &self, + user_id: &str, + branch: Option<&str>, + messages: &[serde_json::Value], + session_id: Option, ) -> Result<(Vec, bool), MemoriaError> { let has_llm = self.llm.is_some(); @@ -2330,7 +2535,7 @@ impl MemoryService { let mut stored = Vec::with_capacity(candidates.len()); for mem in candidates { - match self.persist_with_dedup(user_id, mem).await { + match self.persist_with_dedup(user_id, branch, mem).await { Ok(m) => stored.push(m), Err(MemoriaError::Blocked(_)) => continue, Err(e) => return Err(e), @@ -2448,6 +2653,7 @@ impl MemoryService { async fn persist_with_dedup( &self, user_id: &str, + branch: Option<&str>, mut mem: Memory, ) -> Result { let sensitivity = check_sensitivity(&mem.content); @@ -2465,7 +2671,7 @@ impl MemoryService { if self.sql_store.is_some() { let sql = self.user_sql_store(user_id).await?; - let table = sql.active_table(user_id).await?; + let table = sql.table_for_branch(user_id, branch).await?; if let Some(ref emb) = mem.embedding { let l2_threshold = 0.3162; let mtype = mem.memory_type.to_string(); diff --git a/memoria/crates/memoria-storage/src/store.rs b/memoria/crates/memoria-storage/src/store.rs index 2bff7add..cfaf7f10 100644 --- a/memoria/crates/memoria-storage/src/store.rs +++ b/memoria/crates/memoria-storage/src/store.rs @@ -1611,12 +1611,14 @@ impl SqlMemoryStore { Err(e) if is_duplicate_column(e) => tracing::info!( "migration: author_id column already exists in {memories_table}, skipping" ), - Err(e) => tracing::error!( - "migration: failed to add author_id to {memories_table}: {e}" - ), + Err(e) => { + tracing::error!("migration: failed to add author_id to {memories_table}: {e}") + } } } else { - tracing::debug!("migration: author_id column already exists in {memories_table}, skipping"); + tracing::debug!( + "migration: author_id column already exists in {memories_table}, skipping" + ); } // Ensure the index exists even when the column already existed before this migration. // This covers partially-migrated databases where `author_id` is present but @@ -1661,15 +1663,17 @@ impl SqlMemoryStore { for bt_raw in &branch_table_names { // bt_raw is the raw table name (e.g. br_abc123_my_branch) without DB prefix. // Validate against a strict allowlist before interpolating into DDL. - if !bt_raw.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + if !bt_raw + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + { tracing::warn!( "migration: skipping branch table with invalid identifier '{bt_raw}'" ); continue; } let bt_full = self.t(bt_raw); - let has_col = - info_schema_column_exists(pool, schema_name, bt_raw, "author_id").await; + let has_col = info_schema_column_exists(pool, schema_name, bt_raw, "author_id").await; if !has_col { let r = exec_ddl_with_retry( pool, @@ -1692,8 +1696,7 @@ impl SqlMemoryStore { // Always ensure the index exists regardless of whether the column was just // added or was already present (handles "column exists but index is missing" // on older databases). The error is expected when the index already exists. - let has_idx = - info_schema_index_exists(pool, schema_name, bt_raw, "idx_author").await; + let has_idx = info_schema_index_exists(pool, schema_name, bt_raw, "idx_author").await; if !has_idx { let _ = exec_ddl_with_retry( pool, @@ -2509,6 +2512,43 @@ impl SqlMemoryStore { } } + /// Resolve a memory table for an optional explicit branch. + /// + /// `None` preserves checkout-based behavior via `active_table`. A concrete + /// branch name bypasses `mem_user_state`, so concurrent agents can read/write + /// different branches without racing on the active-branch pointer. + pub async fn table_for_branch( + &self, + user_id: &str, + branch: Option<&str>, + ) -> Result { + let Some(branch) = branch.map(str::trim).filter(|branch| !branch.is_empty()) else { + return self.active_table(user_id).await; + }; + + if branch == "main" { + return Ok(self.t("mem_memories")); + } + + let branches_table = self.t("mem_branches"); + let branch_row = sqlx::query(&format!( + "SELECT table_name FROM {branches_table} WHERE user_id = ? AND name = ? AND status = 'active'" + )) + .bind(user_id) + .bind(branch) + .fetch_optional(&self.pool) + .await + .map_err(db_err)?; + + match branch_row { + Some(r) => { + let raw = r.try_get::("table_name").map_err(db_err)?; + Ok(self.t(&raw)) + } + None => Err(MemoriaError::NotFound(format!("Branch '{branch}'"))), + } + } + pub async fn set_active_branch(&self, user_id: &str, branch: &str) -> Result<(), MemoriaError> { let state_user = self.state_user(user_id); let user_state_table = self.t("mem_user_state"); @@ -2617,7 +2657,9 @@ impl SqlMemoryStore { // Lenient: a corrupt/NULL timestamp yields None rather than failing the // entire branch list. name and table_name are strict because they drive // branch operations. - r.try_get::, _>("created_at").ok().flatten(), + r.try_get::, _>("created_at") + .ok() + .flatten(), )) }) .collect() diff --git a/memoria/crates/memoria-storage/tests/branch_ops.rs b/memoria/crates/memoria-storage/tests/branch_ops.rs index 2fb5f896..157300f2 100644 --- a/memoria/crates/memoria-storage/tests/branch_ops.rs +++ b/memoria/crates/memoria-storage/tests/branch_ops.rs @@ -124,6 +124,45 @@ async fn test_set_active_branch_changes_table() { store.set_active_branch(&user, "main").await.ok(); } +#[tokio::test] +async fn test_table_for_branch_bypasses_active_checkout() { + let store = setup().await; + let user = uid(); + let branch_a = format!("br_a_{}", &uid()[5..]); + let branch_b = format!("br_b_{}", &uid()[5..]); + let table_a = format!("mem_br_a_{}", &uid()[5..]); + let table_b = format!("mem_br_b_{}", &uid()[5..]); + + store + .register_branch(&user, &branch_a, &table_a) + .await + .expect("register a"); + store + .register_branch(&user, &branch_b, &table_b) + .await + .expect("register b"); + store + .set_active_branch(&user, &branch_b) + .await + .expect("set active b"); + + let explicit = store + .table_for_branch(&user, Some(&branch_a)) + .await + .expect("explicit branch table"); + let implicit = store + .table_for_branch(&user, None) + .await + .expect("implicit active table"); + + assert_eq!(explicit, table_a); + assert_eq!(implicit, table_b); + + store.deregister_branch(&user, &branch_a).await.ok(); + store.deregister_branch(&user, &branch_b).await.ok(); + store.set_active_branch(&user, "main").await.ok(); +} + // ── 3. active_table falls back to main if branch deleted ───────────────────── #[tokio::test] diff --git a/memoria/crates/memoria-storage/tests/store_crud.rs b/memoria/crates/memoria-storage/tests/store_crud.rs index 7ee236e2..fd2bc962 100644 --- a/memoria/crates/memoria-storage/tests/store_crud.rs +++ b/memoria/crates/memoria-storage/tests/store_crud.rs @@ -267,7 +267,14 @@ async fn test_search_vector_from_filtered_scoped_prefilters_by_session() { store.insert(&global_mem).await.unwrap(); let global_results = store - .search_vector_from_filtered_scoped("mem_memories", &uid, &mk_vec(1.0, 0.0, 0.0), 1, None, None) + .search_vector_from_filtered_scoped( + "mem_memories", + &uid, + &mk_vec(1.0, 0.0, 0.0), + 1, + None, + None, + ) .await .expect("global vector search"); assert_eq!(