Replace the policy rebuild flag with shared generation counters

backend/erato/src/models/assistant_hub.rs

@@ -442,24 +442,25 @@ pub async fn build_submission_diff(
 pub async fn submit_version(
     conn: &DatabaseConnection,
     policy: &PolicyEngine,
-    config: &AssistantHubConfig,
+    config: &crate::config::AppConfig,
     subject: &Subject,
     source_assistant_id: Uuid,
     profile: HubSubmissionProfile,
     audience_grants: Vec<HubAudienceGrantInput>,
 ) -> Result<HubVersionRecord, Report> {
-    ensure_enabled(config)?;
+    let hub_config = &config.assistant_hub;
+    ensure_enabled(hub_config)?;
     let mut profile = profile;
     profile.version_number = profile.version_number.trim().to_string();
-    validate_profile(config, &profile)?;
+    validate_profile(hub_config, &profile)?;
 
     let hub_assistant = get_or_create_hub_assistant(conn, subject, source_assistant_id).await?;
     ensure_unique_version_number(conn, hub_assistant.id, &profile.version_number).await?;
     let cloned = clone_source_assistant(conn, source_assistant_id).await?;
     policy.invalidate_data().await;
 
     let diff_summary =
-        build_submission_diff(conn, config, subject, source_assistant_id, &profile).await?;
+        build_submission_diff(conn, hub_config, subject, source_assistant_id, &profile).await?;
     let now = Utc::now().into();
     let version = assistant_hub_assistant_versions::ActiveModel {
         id: Set(Uuid::new_v4()),
@@ -497,6 +498,7 @@ pub async fn submit_version(
             conn,
             policy,
             subject,
+            config,
             "assistant".to_string(),
             cloned.id.to_string(),
             grant.subject_type,

backend/erato/src/models/file_upload.rs

@@ -405,10 +405,6 @@ pub async fn get_file_upload_by_id(
         .await?
         .wrap_err("File upload not found")?;
 
-    policy
-        .rebuild_data_if_needed(conn, &crate::config::AppConfig::default())
-        .await?;
-
     authorize!(
         policy,
         subject,

backend/erato/src/models/share_grant.rs

@@ -46,17 +46,18 @@ pub async fn create_share_grant(
     conn: &DatabaseConnection,
     policy: &PolicyEngine,
     subject: &Subject,
+    config: &crate::config::AppConfig,
     resource_type: String,
     resource_id: String,
     subject_type: String,
     subject_id_type: String,
     subject_id_value: String,
     role: String,
 ) -> Result<share_grants::Model, Report> {
-    // Rebuild policy data if needed
-    policy
-        .rebuild_data_if_needed(conn, &crate::config::AppConfig::default())
-        .await?;
+    // The resource may have been created earlier in the same request (e.g. the
+    // assistant cloned during a hub submission); the authorize below needs it
+    // in the policy data.
+    policy.rebuild_data_if_needed(conn, config).await?;
 
     // Get the user ID from subject
     let user_id_str = subject.user_id();
@@ -198,12 +199,12 @@ pub async fn delete_share_grant(
     conn: &DatabaseConnection,
     policy: &PolicyEngine,
     subject: &Subject,
+    config: &crate::config::AppConfig,
     grant_id: Uuid,
 ) -> Result<(), Report> {
-    // Rebuild policy data if needed
-    policy
-        .rebuild_data_if_needed(conn, &crate::config::AppConfig::default())
-        .await?;
+    // The grant may have been created since this request's engine was handed
+    // out; the authorize below needs it in the policy data.
+    policy.rebuild_data_if_needed(conn, config).await?;
 
     // Get the user ID from subject
     let user_id_str = subject.user_id();

backend/erato/src/policy/engine.rs

@@ -16,6 +16,7 @@ use sea_orm::{DatabaseConnection, EntityTrait, FromQueryResult, QuerySelect};
 use serde_json::{Value as JsonValue, json};
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use tokio::sync::{Mutex, RwLock};
 use tracing::instrument;
 
@@ -336,13 +337,19 @@ macro_rules! authorize {
 }
 pub(crate) use authorize;
 
+/// All clones (including request-scoped ones) share the generation counters and
+/// the rebuild lock, so an invalidation from any task is visible process-wide and
+/// a rebuild from any task satisfies it process-wide.
 #[derive(Debug, Clone)]
 pub struct PolicyEngine {
     engine: Arc<RwLock<Engine>>,
-    data_needs_rebuild: Arc<RwLock<bool>>,
-    /// Serializes rebuilds process-wide (shared by request-scoped clones):
-    /// concurrent stale observers queue here and re-check staleness after
-    /// acquiring, so they coalesce into a single rebuild.
+    /// Generation the policy data must reach to be considered fresh.
+    target_generation: Arc<AtomicU64>,
+    /// Generation the engine data was built at. Data is stale while this is
+    /// behind `target_generation`; `0` means the data has never been built.
+    built_generation: Arc<AtomicU64>,
+    /// Serializes rebuilds: concurrent stale observers queue here and re-check
+    /// freshness after acquiring, so they coalesce into a single rebuild.
     rebuild_lock: Arc<Mutex<()>>,
 }
 
@@ -363,36 +370,35 @@ impl PolicyEngine {
             .unwrap();
         Self {
             engine: Arc::new(RwLock::new(engine)),
-            data_needs_rebuild: Arc::new(RwLock::new(true)),
+            target_generation: Arc::new(AtomicU64::new(1)),
+            built_generation: Arc::new(AtomicU64::new(0)),
             rebuild_lock: Arc::new(Mutex::new(())),
         }
     }
 
-    /// Clone the engine for use in a request handler.
-    /// Unlike regular Clone, this creates an independent `data_needs_rebuild` state
-    /// set to `false`, so that invalidating the global engine doesn't affect
-    /// cloned request-scoped engines.
-    pub fn clone_for_request(&self) -> Self {
-        Self {
-            engine: self.engine.clone(),
-            data_needs_rebuild: Arc::new(RwLock::new(false)),
-            rebuild_lock: self.rebuild_lock.clone(),
-        }
+    fn is_stale(&self) -> bool {
+        self.built_generation.load(Ordering::SeqCst) < self.target_generation.load(Ordering::SeqCst)
     }
 
+    #[cfg(test)]
     async fn set_data(&self, data: JsonValue) -> Result<(), Report> {
+        let generation = self.target_generation.load(Ordering::SeqCst);
+        self.install_data(data, generation).await
+    }
+
+    async fn install_data(&self, data: JsonValue, generation: u64) -> Result<(), Report> {
         let mut guard = self.engine.write().await;
         guard.clear_data();
         guard
             .add_data_json(&data.to_string())
             .map_err(|e| eyre!(e))?;
-        *self.data_needs_rebuild.write().await = false;
+        self.built_generation
+            .fetch_max(generation, Ordering::SeqCst);
         Ok(())
     }
 
     pub async fn invalidate_data(&self) {
-        *self.data_needs_rebuild.write().await = true;
-        // info!("Invalidated policy data");
+        self.target_generation.fetch_add(1, Ordering::SeqCst);
     }
 
     #[instrument(skip_all)]
@@ -410,6 +416,10 @@ impl PolicyEngine {
         db: &DatabaseConnection,
         config: &AppConfig,
     ) -> Result<(), Report> {
+        // Capture the target before reading from the database: an invalidation
+        // that lands while the reads are in flight may not be reflected in what
+        // was read, so it must remain pending after this rebuild completes.
+        let generation = self.target_generation.load(Ordering::SeqCst);
         // Fetch policy data for each resource type
         let chat_data = fetch_chat_policy_data(db).await?;
         let assistant_data = fetch_assistant_policy_data(db).await?;
@@ -452,8 +462,7 @@ impl PolicyEngine {
             "config_permissions": build_config_permissions_policy_data(config),
         });
 
-        self.set_data(policy_data).await?;
-        // info!("Finished policy data rebuild");
+        self.install_data(policy_data, generation).await?;
         Ok(())
     }
 
@@ -463,13 +472,13 @@ impl PolicyEngine {
         db: &DatabaseConnection,
         config: &AppConfig,
     ) -> Result<(), Report> {
-        if !*self.data_needs_rebuild.read().await {
+        if !self.is_stale() {
             return Ok(());
         }
         let _rebuild_guard = self.rebuild_lock.lock().await;
         // Re-check after acquiring: a rebuild that finished while we waited on
-        // the lock already covers this invalidation.
-        if !*self.data_needs_rebuild.read().await {
+        // the lock already covers our pending invalidation.
+        if !self.is_stale() {
             return Ok(());
         }
         self.rebuild_data_locked(db, config).await
@@ -523,10 +532,12 @@ impl PolicyEngine {
         organization_group_ids: &[String],
         groups: &[String],
     ) -> Result<(), Report> {
-        // info!("Authorizing");
-        if *self.data_needs_rebuild.read().await {
+        // Freshness is enforced when the engine is handed out (middleware and
+        // explicit rebuild_data_if_needed calls); here we only refuse to
+        // authorize against data that has never been built at all.
+        if self.built_generation.load(Ordering::SeqCst) == 0 {
             return Err(eyre!(
-                "Policy data is stale and needs to be rebuilt before authorization"
+                "Policy data has not been built yet; rebuild before authorization"
             ));
         }
         // First validate the resource_kind-action combination as an assertion
@@ -808,6 +819,77 @@ mod tests {
         .unwrap();
     }
 
+    #[tokio::test]
+    async fn test_authorize_fails_before_first_data_build() {
+        let engine = PolicyEngine::new();
+        let result = authorize!(
+            engine,
+            &Subject::User("user_1".to_string()),
+            &Resource::Chat("chat_1".to_string()),
+            Action::Read
+        );
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_invalidation_marks_stale_but_authorize_keeps_working() {
+        let engine = PolicyEngine::new();
+        engine
+            .set_data(json!({
+                "resource_attributes": {
+                    "chat": {
+                        "chat_1": { "id": "chat_1", "owner_id": "user_1" }
+                    }
+                }
+            }))
+            .await
+            .unwrap();
+        assert!(!engine.is_stale());
+
+        engine.invalidate_data().await;
+        assert!(engine.is_stale());
+
+        // Authorization proceeds on the already-built data; freshness is
+        // enforced where engines are handed out, not per authorize call.
+        let result = authorize!(
+            engine,
+            &Subject::User("user_1".to_string()),
+            &Resource::Chat("chat_1".to_string()),
+            Action::Read
+        );
+        assert!(result.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_invalidation_is_shared_across_clones() {
+        let engine = PolicyEngine::new();
+        engine.set_data(json!({})).await.unwrap();
+
+        let request_clone = engine.clone();
+        engine.invalidate_data().await;
+        assert!(request_clone.is_stale());
+
+        // A rebuild through any clone satisfies the pending invalidation
+        // process-wide (set_data stands in for a full rebuild here).
+        request_clone.set_data(json!({})).await.unwrap();
+        assert!(!engine.is_stale());
+        assert!(!request_clone.is_stale());
+    }
+
+    #[tokio::test]
+    async fn test_install_data_ignores_generation_regression() {
+        let engine = PolicyEngine::new();
+        engine.set_data(json!({})).await.unwrap();
+        engine.invalidate_data().await;
+        engine.set_data(json!({})).await.unwrap();
+        assert!(!engine.is_stale());
+
+        // A rebuild that captured its target before an already-applied newer
+        // one must not move built_generation backwards.
+        engine.install_data(json!({}), 1).await.unwrap();
+        assert!(!engine.is_stale());
+    }
+
     #[tokio::test]
     async fn test_authorize_short_form() {
         let subject = Subject::User("user_1".to_string());

backend/erato/src/server/api/v1beta/assistant_hub.rs

@@ -626,7 +626,7 @@ pub async fn submit_assistant_hub_version(
     let record = assistant_hub::submit_version(
         &app_state.db,
         &policy,
-        &app_state.config.assistant_hub,
+        &app_state.config,
         &me_user.to_subject(),
         source_assistant_id,
         profile,

backend/erato/src/server/api/v1beta/assistants.rs

@@ -495,6 +495,7 @@ pub async fn create_assistant(
                 &app_state.db,
                 &policy,
                 &me_user.to_subject(),
+                &app_state.config,
                 "assistant".to_string(),
                 created_assistant.id.to_string(),
                 grant_input.subject_type,

backend/erato/src/server/api/v1beta/message_streaming.rs

@@ -7200,10 +7200,7 @@ async fn run_message_submit_task(
             .map_err(Report::msg)?;
 
         tracing::info!("Rebuilding policy data for newly created chat");
-        // Via the global engine: `policy` is a request-scoped clone whose
-        // severed staleness flag cannot see the handler's invalidation.
-        app_state
-            .global_policy_engine
+        policy
             .rebuild_data_if_needed(&app_state.db, &app_state.config)
             .await
             .wrap_err("Failed to rebuild policy data after chat creation")?;

backend/erato/src/server/api/v1beta/share_grants.rs

@@ -293,6 +293,7 @@ pub async fn create_share_grant(
         &app_state.db,
         &policy,
         &me_user.to_subject(),
+        &app_state.config,
         request.resource_type,
         request.resource_id,
         request.subject_type,
@@ -492,22 +493,28 @@ pub async fn delete_share_grant(
     let grant_id = Uuid::parse_str(&grant_id).map_err(|_| StatusCode::BAD_REQUEST)?;
 
     // Delete the share grant
-    share_grant::delete_share_grant(&app_state.db, &policy, &me_user.to_subject(), grant_id)
-        .await
-        .map_err(|e| {
-            if e.to_string().contains("Access denied") || e.to_string().contains("does not own") {
-                tracing::warn!(
-                    "User {} attempted to delete a share grant for a resource they don't own: {}",
-                    me_user.id,
-                    e
-                );
-                StatusCode::FORBIDDEN
-            } else if e.to_string().contains("not found") {
-                StatusCode::NOT_FOUND
-            } else {
-                log_internal_server_error(e)
-            }
-        })?;
+    share_grant::delete_share_grant(
+        &app_state.db,
+        &policy,
+        &me_user.to_subject(),
+        &app_state.config,
+        grant_id,
+    )
+    .await
+    .map_err(|e| {
+        if e.to_string().contains("Access denied") || e.to_string().contains("does not own") {
+            tracing::warn!(
+                "User {} attempted to delete a share grant for a resource they don't own: {}",
+                me_user.id,
+                e
+            );
+            StatusCode::FORBIDDEN
+        } else if e.to_string().contains("not found") {
+            StatusCode::NOT_FOUND
+        } else {
+            log_internal_server_error(e)
+        }
+    })?;
 
     tracing::info!("User {} deleted share grant {}", me_user.id, grant_id);