Skip to content

feat: add explicit branch access for memory operations #203

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aptend merged 1 commit into from
May 25, 2026
Merged

feat: add explicit branch access for memory operations #203

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 70 additions & 0 deletions memoria/crates/memoria-api/src/auth.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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; \
Expand Down
9 changes: 9 additions & 0 deletions memoria/crates/memoria-api/src/models.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -25,6 +26,7 @@ fn default_memory_type() -> String {
#[derive(Deserialize)]
pub struct BatchStoreRequest {
pub memories: Vec<StoreRequest>,
pub branch: Option<String>,
}

#[derive(Deserialize)]
Expand All @@ -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,
Expand Down Expand Up @@ -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,
}
Expand Down Expand Up @@ -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)]
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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() {
Expand All @@ -470,6 +478,7 @@ mod tests {
session_id: None,
memory_types: Some(vec![]),
reason: None,
branch: None,
};

assert!(matches!(request.selector().unwrap(), PurgeSelector::None));
Expand Down
14 changes: 8 additions & 6 deletions memoria/crates/memoria-api/src/routes/admin.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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) => {
Expand Down
17 changes: 16 additions & 1 deletion memoria/crates/memoria-api/src/routes/governance.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
78 changes: 55 additions & 23 deletions memoria/crates/memoria-api/src/routes/mcp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading