Support explicit branch access for subagent concurrency

Author: aptend

memoria/crates/memoria-api/src/auth.rs

@@ -142,6 +142,7 @@ pub async fn group_main_write_guard(
     req: axum::http::Request<axum::body::Body>,
     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<String, String>,
+                                >(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::<serde_json::Value>(&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; \

memoria/crates/memoria-api/src/models.rs

@@ -17,6 +17,7 @@ pub struct StoreRequest {
     pub initial_confidence: Option<f64>,
     pub observed_at: Option<String>,
     pub source: Option<String>,
+    pub branch: Option<String>,
 }
 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<StoreRequest>,
+    pub branch: Option<String>,
 }
 
 #[derive(Deserialize)]
@@ -34,6 +36,7 @@ pub struct RetrieveRequest {
     pub top_k: i64,
     pub session_id: Option<String>,
     pub session_scope: Option<String>,
+    pub branch: Option<String>,
     /// 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<String>,
     pub session_scope: Option<String>,
+    pub branch: Option<String>,
     #[serde(default, deserialize_with = "deserialize_explain")]
     pub explain: String,
 }
@@ -122,6 +126,7 @@ fn deserialize_explain<'de, D: serde::Deserializer<'de>>(d: D) -> Result<String,
 pub struct CorrectRequest {
     pub new_content: String,
     pub reason: Option<String>,
+    pub branch: Option<String>,
 }
 
 #[derive(Deserialize)]
@@ -131,6 +136,7 @@ pub struct CorrectByQueryRequest {
     pub session_id: Option<String>,
     pub session_scope: Option<String>,
     pub reason: Option<String>,
+    pub branch: Option<String>,
 }
 
 impl CorrectByQueryRequest {
@@ -156,6 +162,7 @@ pub struct PurgeRequest {
     pub session_id: Option<String>,
     pub memory_types: Option<Vec<String>>,
     pub reason: Option<String>,
+    pub branch: Option<String>,
 }
 
 pub enum PurgeSelector {

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) => {

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<AppState>,
     auth: AuthUser,
     Json(req): Json<GovernanceRequest>,
 ) -> Result<Json<serde_json::Value>, (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())

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<crate::metrics_summary::DirtyMask> {
     use crate::metrics_summary::DirtyMask;
     match tool {
@@ -254,38 +274,50 @@ pub async fn mcp_handler(
     let blocked_tool: Option<String> = 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