Skip to main content

brainwires_agents/
access_control.rs

1//! Unified access control manager for inter-agent coordination
2//!
3//! Provides a single interface for managing file locks, resource locks (build/test),
4//! and read-before-write enforcement.
5
6use anyhow::{Result, anyhow};
7use serde_json::Value;
8use std::collections::{HashMap, HashSet};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use std::time::Duration;
12use tokio::sync::RwLock;
13
14use crate::file_locks::{FileLockManager, LockGuard, LockType};
15use crate::resource_locks::{ResourceLockGuard, ResourceLockManager, ResourceScope, ResourceType};
16
17const DEFAULT_INITIAL_DELAY_MS: u64 = 100;
18const DEFAULT_MAX_RETRIES: u32 = 5;
19const DEFAULT_MAX_DELAY_SECS: u64 = 5;
20const PERSISTENT_LOCK_TIMEOUT_SECS: u64 = 300;
21const FILE_LOCK_BACKOFF_INITIAL_MS: u64 = 50;
22const FILE_LOCK_BACKOFF_MAX_MS: u64 = 500;
23const RESOURCE_LOCK_BACKOFF_INITIAL_MS: u64 = 100;
24const RESOURCE_LOCK_BACKOFF_MAX_SECS: u64 = 1;
25
26/// Trait for persistent lock storage (inter-process coordination)
27///
28/// Implement this trait to enable cross-process lock coordination.
29/// The default implementation (no-op) is used when no persistent store is configured.
30#[async_trait::async_trait]
31pub trait LockPersistence: Send + Sync {
32    /// Try to acquire a persistent lock
33    ///
34    /// Returns `Ok(true)` if the lock was acquired, `Ok(false)` if it's already held.
35    async fn try_acquire(
36        &self,
37        lock_type: &str,
38        resource_path: &str,
39        agent_id: &str,
40        timeout: Option<Duration>,
41    ) -> Result<bool>;
42
43    /// Release a persistent lock
44    async fn release(&self, lock_type: &str, resource_path: &str, agent_id: &str) -> Result<()>;
45
46    /// Release all locks held by an agent
47    async fn release_all_for_agent(&self, agent_id: &str) -> Result<usize>;
48
49    /// Cleanup stale locks (e.g., from crashed processes)
50    async fn cleanup_stale(&self) -> Result<usize>;
51}
52
53/// Strategy for handling lock contention
54#[derive(Debug, Clone)]
55pub enum ContentionStrategy {
56    /// Fail immediately if lock is unavailable
57    FailFast,
58    /// Wait up to the specified duration for the lock
59    WaitWithTimeout(Duration),
60    /// Retry with exponential backoff
61    RetryWithBackoff {
62        /// Initial delay between retries.
63        initial_delay: Duration,
64        /// Maximum number of retries.
65        max_retries: u32,
66        /// Maximum delay cap.
67        max_delay: Duration,
68    },
69}
70
71impl Default for ContentionStrategy {
72    fn default() -> Self {
73        ContentionStrategy::RetryWithBackoff {
74            initial_delay: Duration::from_millis(DEFAULT_INITIAL_DELAY_MS),
75            max_retries: DEFAULT_MAX_RETRIES,
76            max_delay: Duration::from_secs(DEFAULT_MAX_DELAY_SECS),
77        }
78    }
79}
80
81/// Bundle of locks acquired for a single operation
82pub struct LockBundle {
83    /// File lock guard (if applicable)
84    pub file_lock: Option<LockGuard>,
85    /// Resource lock guard (if applicable)
86    pub resource_lock: Option<ResourceLockGuard>,
87}
88
89impl LockBundle {
90    /// Create an empty lock bundle
91    pub fn empty() -> Self {
92        Self {
93            file_lock: None,
94            resource_lock: None,
95        }
96    }
97
98    /// Check if this bundle contains any locks
99    pub fn has_locks(&self) -> bool {
100        self.file_lock.is_some() || self.resource_lock.is_some()
101    }
102}
103
104/// Unified access control manager
105pub struct AccessControlManager {
106    /// File lock manager (in-memory, intra-process)
107    file_locks: Arc<FileLockManager>,
108    /// Resource lock manager (build/test, in-memory, intra-process)
109    resource_locks: Arc<ResourceLockManager>,
110    /// Strategy for handling contention
111    contention_strategy: ContentionStrategy,
112    /// Tracking of files read by each agent (for read-before-write enforcement)
113    read_tracking: RwLock<HashMap<String, HashSet<PathBuf>>>,
114    /// Project root for determining resource scope
115    project_root: PathBuf,
116    /// Optional persistent lock store for inter-process coordination
117    lock_store: Option<Arc<dyn LockPersistence>>,
118}
119
120impl AccessControlManager {
121    /// Create a new access control manager
122    pub fn new(project_root: PathBuf) -> Self {
123        Self {
124            file_locks: Arc::new(FileLockManager::new()),
125            resource_locks: Arc::new(ResourceLockManager::new()),
126            contention_strategy: ContentionStrategy::default(),
127            read_tracking: RwLock::new(HashMap::new()),
128            project_root,
129            lock_store: None,
130        }
131    }
132
133    /// Create with custom managers (for testing or sharing)
134    pub fn with_managers(
135        file_locks: Arc<FileLockManager>,
136        resource_locks: Arc<ResourceLockManager>,
137        project_root: PathBuf,
138    ) -> Self {
139        Self {
140            file_locks,
141            resource_locks,
142            contention_strategy: ContentionStrategy::default(),
143            read_tracking: RwLock::new(HashMap::new()),
144            project_root,
145            lock_store: None,
146        }
147    }
148
149    /// Set the contention strategy
150    pub fn with_strategy(mut self, strategy: ContentionStrategy) -> Self {
151        self.contention_strategy = strategy;
152        self
153    }
154
155    /// Enable inter-process locking with a persistent lock store
156    pub fn with_lock_persistence(mut self, lock_store: Arc<dyn LockPersistence>) -> Self {
157        self.lock_store = Some(lock_store);
158        self
159    }
160
161    /// Get the lock store (if configured)
162    pub fn lock_store(&self) -> Option<&Arc<dyn LockPersistence>> {
163        self.lock_store.as_ref()
164    }
165
166    /// Get a reference to the file lock manager
167    pub fn file_locks(&self) -> &Arc<FileLockManager> {
168        &self.file_locks
169    }
170
171    /// Get a reference to the resource lock manager
172    pub fn resource_locks(&self) -> &Arc<ResourceLockManager> {
173        &self.resource_locks
174    }
175
176    /// Track that an agent has read a file
177    pub async fn track_file_read(&self, agent_id: &str, path: &Path) {
178        let canonical_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
179        let mut tracking = self.read_tracking.write().await;
180        tracking
181            .entry(agent_id.to_string())
182            .or_default()
183            .insert(canonical_path);
184    }
185
186    /// Check if an agent has read a file
187    pub async fn has_read_file(&self, agent_id: &str, path: &Path) -> bool {
188        let canonical_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
189        let tracking = self.read_tracking.read().await;
190        tracking
191            .get(agent_id)
192            .map(|files| files.contains(&canonical_path))
193            .unwrap_or(false)
194    }
195
196    /// Validate that a write operation is allowed (file must have been read first)
197    pub async fn validate_write(&self, agent_id: &str, path: &Path) -> Result<()> {
198        // New files don't need to be read first
199        if !path.exists() {
200            return Ok(());
201        }
202
203        if !self.has_read_file(agent_id, path).await {
204            return Err(anyhow!(
205                "Must read file before writing: {}. Use read_file first.",
206                path.display()
207            ));
208        }
209        Ok(())
210    }
211
212    /// Clear read tracking for an agent (call on agent shutdown)
213    pub async fn clear_tracking_for_agent(&self, agent_id: &str) {
214        let mut tracking = self.read_tracking.write().await;
215        tracking.remove(agent_id);
216    }
217
218    /// Get the lock requirement for a file operation
219    pub fn get_file_lock_requirement(
220        tool_name: &str,
221        input: &Value,
222    ) -> Option<(PathBuf, LockType)> {
223        let path_str = input
224            .get("path")
225            .or_else(|| input.get("file_path"))
226            .and_then(|v| v.as_str())?;
227
228        let path = PathBuf::from(path_str);
229
230        let lock_type = match tool_name {
231            // Read operations - shared lock
232            "read_file" | "list_directory" | "search_files" => LockType::Read,
233            // Write operations - exclusive lock
234            "write_file" | "edit_file" | "patch_file" | "delete_file" | "create_directory" => {
235                LockType::Write
236            }
237            _ => return None,
238        };
239
240        Some((path, lock_type))
241    }
242
243    /// Detect if a bash command is a build command
244    pub fn detect_build_command(command: &str) -> bool {
245        let build_patterns = [
246            "cargo build",
247            "cargo b ",
248            "cargo b\n",
249            "cargo b$",
250            "make ",
251            "make\n",
252            "make$",
253            "cmake",
254            "npm run build",
255            "npm build",
256            "yarn build",
257            "pnpm build",
258            "go build",
259            "mvn compile",
260            "mvn package",
261            "gradle build",
262            "gradle assemble",
263            "msbuild",
264            "dotnet build",
265            "rustc ",
266            "gcc ",
267            "g++ ",
268            "clang ",
269            "clang++ ",
270            "javac ",
271            "tsc ",
272            "webpack",
273            "vite build",
274            "rollup",
275            "esbuild",
276        ];
277
278        let cmd_lower = command.to_lowercase();
279        build_patterns
280            .iter()
281            .any(|p| cmd_lower.contains(&p.to_lowercase()))
282    }
283
284    /// Detect if a bash command is a test command
285    pub fn detect_test_command(command: &str) -> bool {
286        let test_patterns = [
287            "cargo test",
288            "cargo t ",
289            "cargo t\n",
290            "cargo t$",
291            "npm test",
292            "npm run test",
293            "yarn test",
294            "pnpm test",
295            "go test",
296            "pytest",
297            "python -m pytest",
298            "jest",
299            "mocha",
300            "vitest",
301            "mvn test",
302            "gradle test",
303            "dotnet test",
304            "rspec",
305            "bundle exec rspec",
306            "phpunit",
307            "mix test",
308            "elixir.*test",
309        ];
310
311        let cmd_lower = command.to_lowercase();
312        test_patterns
313            .iter()
314            .any(|p| cmd_lower.contains(&p.to_lowercase()))
315    }
316
317    /// Get resource lock requirement for a bash command
318    pub fn get_resource_requirement(&self, command: &str) -> Option<(ResourceType, ResourceScope)> {
319        let is_build = Self::detect_build_command(command);
320        let is_test = Self::detect_test_command(command);
321
322        let resource_type = match (is_build, is_test) {
323            (true, true) => ResourceType::BuildTest,
324            (true, false) => ResourceType::Build,
325            (false, true) => ResourceType::Test,
326            (false, false) => return None,
327        };
328
329        // Use project scope based on project root
330        let scope = ResourceScope::Project(self.project_root.clone());
331
332        Some((resource_type, scope))
333    }
334
335    /// Acquire all necessary locks for a tool operation
336    pub async fn acquire_for_tool(
337        self: &Arc<Self>,
338        agent_id: &str,
339        tool_name: &str,
340        input: &Value,
341    ) -> Result<LockBundle> {
342        let mut bundle = LockBundle::empty();
343
344        // Handle file operations
345        if let Some((path, lock_type)) = Self::get_file_lock_requirement(tool_name, input) {
346            // For write operations, validate read-before-write
347            if lock_type == LockType::Write {
348                self.validate_write(agent_id, &path).await?;
349            }
350
351            let file_lock = self
352                .acquire_file_lock_with_retry(agent_id, &path, lock_type)
353                .await?;
354            bundle.file_lock = Some(file_lock);
355        }
356
357        // Handle bash commands (build/test)
358        if tool_name == "execute_command"
359            && let Some(command) = input.get("command").and_then(|v| v.as_str())
360            && let Some((resource_type, scope)) = self.get_resource_requirement(command)
361        {
362            let resource_lock = self
363                .acquire_resource_lock_with_retry(agent_id, resource_type, scope)
364                .await?;
365            bundle.resource_lock = Some(resource_lock);
366        }
367
368        Ok(bundle)
369    }
370
371    /// Convert LockType to string for persistent storage
372    fn lock_type_to_string(lock_type: LockType) -> &'static str {
373        match lock_type {
374            LockType::Read => "file_read",
375            LockType::Write => "file_write",
376        }
377    }
378
379    /// Convert ResourceType to string for persistent storage
380    fn resource_type_to_string(resource_type: ResourceType) -> &'static str {
381        match resource_type {
382            ResourceType::Build => "build",
383            ResourceType::Test => "test",
384            ResourceType::BuildTest => "build_test",
385            ResourceType::GitIndex => "git_index",
386            ResourceType::GitCommit => "git_commit",
387            ResourceType::GitRemoteWrite => "git_remote_write",
388            ResourceType::GitRemoteMerge => "git_remote_merge",
389            ResourceType::GitBranch => "git_branch",
390            ResourceType::GitDestructive => "git_destructive",
391        }
392    }
393
394    /// Try to acquire persistent lock (if lock_store is configured)
395    async fn try_acquire_persistent_lock(
396        &self,
397        agent_id: &str,
398        lock_type_str: &str,
399        resource_path: &str,
400    ) -> Result<bool> {
401        if let Some(store) = &self.lock_store {
402            store
403                .try_acquire(
404                    lock_type_str,
405                    resource_path,
406                    agent_id,
407                    Some(Duration::from_secs(PERSISTENT_LOCK_TIMEOUT_SECS)),
408                )
409                .await
410        } else {
411            // No persistent store, always succeed
412            Ok(true)
413        }
414    }
415
416    /// Release persistent lock (if lock_store is configured)
417    async fn release_persistent_lock(
418        &self,
419        agent_id: &str,
420        lock_type_str: &str,
421        resource_path: &str,
422    ) -> Result<()> {
423        if let Some(store) = &self.lock_store {
424            store
425                .release(lock_type_str, resource_path, agent_id)
426                .await?;
427        }
428        Ok(())
429    }
430
431    /// Acquire a file lock with retry based on contention strategy
432    async fn acquire_file_lock_with_retry(
433        &self,
434        agent_id: &str,
435        path: &Path,
436        lock_type: LockType,
437    ) -> Result<LockGuard> {
438        let lock_type_str = Self::lock_type_to_string(lock_type);
439        let resource_path = path.to_string_lossy().to_string();
440
441        match &self.contention_strategy {
442            ContentionStrategy::FailFast => {
443                // Try persistent lock first
444                if !self
445                    .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
446                    .await?
447                {
448                    return Err(anyhow!(
449                        "File {} is locked by another process",
450                        path.display()
451                    ));
452                }
453
454                // Then acquire in-memory lock
455                match self
456                    .file_locks
457                    .acquire_lock(agent_id, path, lock_type)
458                    .await
459                {
460                    Ok(guard) => Ok(guard),
461                    Err(e) => {
462                        // Release persistent lock on failure
463                        let _ = self
464                            .release_persistent_lock(agent_id, lock_type_str, &resource_path)
465                            .await;
466                        Err(e)
467                    }
468                }
469            }
470            ContentionStrategy::WaitWithTimeout(timeout) => {
471                let deadline = tokio::time::Instant::now() + *timeout;
472                let mut delay = Duration::from_millis(FILE_LOCK_BACKOFF_INITIAL_MS);
473
474                loop {
475                    // Try persistent lock first
476                    if self
477                        .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
478                        .await?
479                    {
480                        // Then try in-memory lock
481                        match self
482                            .file_locks
483                            .acquire_lock(agent_id, path, lock_type)
484                            .await
485                        {
486                            Ok(guard) => return Ok(guard),
487                            Err(e) => {
488                                // Release persistent lock and retry
489                                let _ = self
490                                    .release_persistent_lock(
491                                        agent_id,
492                                        lock_type_str,
493                                        &resource_path,
494                                    )
495                                    .await;
496                                if tokio::time::Instant::now() >= deadline {
497                                    return Err(anyhow!(
498                                        "Timeout waiting for file lock on {}: {}",
499                                        path.display(),
500                                        e
501                                    ));
502                                }
503                            }
504                        }
505                    } else if tokio::time::Instant::now() >= deadline {
506                        return Err(anyhow!(
507                            "Timeout waiting for file lock on {} (held by another process)",
508                            path.display()
509                        ));
510                    }
511
512                    tokio::time::sleep(delay).await;
513                    delay =
514                        std::cmp::min(delay * 2, Duration::from_millis(FILE_LOCK_BACKOFF_MAX_MS));
515                }
516            }
517            ContentionStrategy::RetryWithBackoff {
518                initial_delay,
519                max_retries,
520                max_delay,
521            } => {
522                let mut delay = *initial_delay;
523                let mut attempts = 0;
524
525                loop {
526                    // Try persistent lock first
527                    if self
528                        .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
529                        .await?
530                    {
531                        // Then try in-memory lock
532                        match self
533                            .file_locks
534                            .acquire_lock(agent_id, path, lock_type)
535                            .await
536                        {
537                            Ok(guard) => return Ok(guard),
538                            Err(e) => {
539                                // Release persistent lock and retry
540                                let _ = self
541                                    .release_persistent_lock(
542                                        agent_id,
543                                        lock_type_str,
544                                        &resource_path,
545                                    )
546                                    .await;
547                                attempts += 1;
548                                if attempts > *max_retries {
549                                    return Err(anyhow!(
550                                        "Failed to acquire file lock on {} after {} attempts: {}",
551                                        path.display(),
552                                        max_retries,
553                                        e
554                                    ));
555                                }
556                                tracing::debug!(
557                                    "Lock contention on {}, attempt {}/{}, waiting {:?}",
558                                    path.display(),
559                                    attempts,
560                                    max_retries,
561                                    delay
562                                );
563                            }
564                        }
565                    } else {
566                        attempts += 1;
567                        if attempts > *max_retries {
568                            return Err(anyhow!(
569                                "Failed to acquire file lock on {} after {} attempts (held by another process)",
570                                path.display(),
571                                max_retries
572                            ));
573                        }
574                        tracing::debug!(
575                            "Lock contention on {} (inter-process), attempt {}/{}, waiting {:?}",
576                            path.display(),
577                            attempts,
578                            max_retries,
579                            delay
580                        );
581                    }
582
583                    tokio::time::sleep(delay).await;
584                    delay = std::cmp::min(delay * 2, *max_delay);
585                }
586            }
587        }
588    }
589
590    /// Get the resource path string for a scope
591    fn scope_to_resource_path(scope: &ResourceScope) -> String {
592        match scope {
593            ResourceScope::Global => "global".to_string(),
594            ResourceScope::Project(path) => path.to_string_lossy().to_string(),
595        }
596    }
597
598    /// Acquire a resource lock with retry based on contention strategy
599    async fn acquire_resource_lock_with_retry(
600        &self,
601        agent_id: &str,
602        resource_type: ResourceType,
603        scope: ResourceScope,
604    ) -> Result<ResourceLockGuard> {
605        let lock_type_str = Self::resource_type_to_string(resource_type);
606        let resource_path = Self::scope_to_resource_path(&scope);
607
608        match &self.contention_strategy {
609            ContentionStrategy::FailFast => {
610                // Try persistent lock first
611                if !self
612                    .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
613                    .await?
614                {
615                    return Err(anyhow!("{} lock is held by another process", resource_type));
616                }
617
618                // Then acquire in-memory lock
619                let description = format!("{} lock", resource_type);
620                match self
621                    .resource_locks
622                    .acquire_resource(agent_id, resource_type, scope, &description)
623                    .await
624                {
625                    Ok(guard) => Ok(guard),
626                    Err(e) => {
627                        // Release persistent lock on failure
628                        let _ = self
629                            .release_persistent_lock(agent_id, lock_type_str, &resource_path)
630                            .await;
631                        Err(e)
632                    }
633                }
634            }
635            ContentionStrategy::WaitWithTimeout(timeout) => {
636                let deadline = tokio::time::Instant::now() + *timeout;
637                let mut delay = Duration::from_millis(RESOURCE_LOCK_BACKOFF_INITIAL_MS);
638                let description = format!("{} lock", resource_type);
639
640                loop {
641                    // Try persistent lock first
642                    if self
643                        .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
644                        .await?
645                    {
646                        // Then try in-memory lock
647                        match self
648                            .resource_locks
649                            .acquire_resource(agent_id, resource_type, scope.clone(), &description)
650                            .await
651                        {
652                            Ok(guard) => return Ok(guard),
653                            Err(e) => {
654                                // Release persistent lock and retry
655                                let _ = self
656                                    .release_persistent_lock(
657                                        agent_id,
658                                        lock_type_str,
659                                        &resource_path,
660                                    )
661                                    .await;
662                                if tokio::time::Instant::now() >= deadline {
663                                    return Err(anyhow!(
664                                        "Timeout waiting for {} lock: {}",
665                                        resource_type,
666                                        e
667                                    ));
668                                }
669                            }
670                        }
671                    } else if tokio::time::Instant::now() >= deadline {
672                        return Err(anyhow!(
673                            "Timeout waiting for {} lock (held by another process)",
674                            resource_type
675                        ));
676                    }
677
678                    tokio::time::sleep(delay).await;
679                    delay = std::cmp::min(
680                        delay * 2,
681                        Duration::from_secs(RESOURCE_LOCK_BACKOFF_MAX_SECS),
682                    );
683                }
684            }
685            ContentionStrategy::RetryWithBackoff {
686                initial_delay,
687                max_retries,
688                max_delay,
689            } => {
690                let mut delay = *initial_delay;
691                let mut attempts = 0;
692                let description = format!("{} lock", resource_type);
693
694                loop {
695                    // Try persistent lock first
696                    if self
697                        .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
698                        .await?
699                    {
700                        // Then try in-memory lock
701                        match self
702                            .resource_locks
703                            .acquire_resource(agent_id, resource_type, scope.clone(), &description)
704                            .await
705                        {
706                            Ok(guard) => return Ok(guard),
707                            Err(e) => {
708                                // Release persistent lock and retry
709                                let _ = self
710                                    .release_persistent_lock(
711                                        agent_id,
712                                        lock_type_str,
713                                        &resource_path,
714                                    )
715                                    .await;
716                                attempts += 1;
717                                if attempts > *max_retries {
718                                    return Err(anyhow!(
719                                        "Failed to acquire {} lock after {} attempts: {}",
720                                        resource_type,
721                                        max_retries,
722                                        e
723                                    ));
724                                }
725                                tracing::debug!(
726                                    "{} lock contention, attempt {}/{}, waiting {:?}",
727                                    resource_type,
728                                    attempts,
729                                    max_retries,
730                                    delay
731                                );
732                            }
733                        }
734                    } else {
735                        attempts += 1;
736                        if attempts > *max_retries {
737                            return Err(anyhow!(
738                                "Failed to acquire {} lock after {} attempts (held by another process)",
739                                resource_type,
740                                max_retries
741                            ));
742                        }
743                        tracing::debug!(
744                            "{} lock contention (inter-process), attempt {}/{}, waiting {:?}",
745                            resource_type,
746                            attempts,
747                            max_retries,
748                            delay
749                        );
750                    }
751
752                    tokio::time::sleep(delay).await;
753                    delay = std::cmp::min(delay * 2, *max_delay);
754                }
755            }
756        }
757    }
758
759    /// Release all locks and tracking for an agent (call on agent shutdown)
760    pub async fn cleanup_agent(&self, agent_id: &str) -> (usize, usize, usize) {
761        let file_locks_released = self.file_locks.release_all_locks(agent_id).await;
762        let resource_locks_released = self.resource_locks.release_all_for_agent(agent_id).await;
763
764        // Also release persistent locks
765        let persistent_locks_released = if let Some(store) = &self.lock_store {
766            store.release_all_for_agent(agent_id).await.unwrap_or(0)
767        } else {
768            0
769        };
770
771        self.clear_tracking_for_agent(agent_id).await;
772        (
773            file_locks_released,
774            resource_locks_released,
775            persistent_locks_released,
776        )
777    }
778
779    /// Cleanup stale persistent locks (call on startup)
780    pub async fn cleanup_stale_locks(&self) -> Result<usize> {
781        if let Some(store) = &self.lock_store {
782            store.cleanup_stale().await
783        } else {
784            Ok(0)
785        }
786    }
787}
788
789#[cfg(test)]
790mod tests {
791    use super::*;
792    use std::path::PathBuf;
793    use tempfile::tempdir;
794
795    fn create_manager() -> Arc<AccessControlManager> {
796        Arc::new(AccessControlManager::new(PathBuf::from("/test/project")))
797    }
798
799    #[tokio::test]
800    async fn test_track_file_read() {
801        let manager = create_manager();
802
803        let path = PathBuf::from("/test/file.txt");
804        assert!(!manager.has_read_file("agent-1", &path).await);
805
806        manager.track_file_read("agent-1", &path).await;
807        assert!(manager.has_read_file("agent-1", &path).await);
808
809        // Different agent hasn't read it
810        assert!(!manager.has_read_file("agent-2", &path).await);
811    }
812
813    #[tokio::test]
814    async fn test_validate_write_requires_read() {
815        let manager = create_manager();
816
817        // Create a temp file
818        let dir = tempdir().unwrap();
819        let file_path = dir.path().join("test.txt");
820        std::fs::write(&file_path, "test content").unwrap();
821
822        // Write without reading should fail
823        let result = manager.validate_write("agent-1", &file_path).await;
824        assert!(result.is_err());
825
826        // After reading, write should succeed
827        manager.track_file_read("agent-1", &file_path).await;
828        let result = manager.validate_write("agent-1", &file_path).await;
829        assert!(result.is_ok());
830    }
831
832    #[tokio::test]
833    async fn test_validate_write_allows_new_files() {
834        let manager = create_manager();
835
836        // New file that doesn't exist should be allowed
837        let path = PathBuf::from("/nonexistent/new_file.txt");
838        let result = manager.validate_write("agent-1", &path).await;
839        assert!(result.is_ok());
840    }
841
842    #[tokio::test]
843    async fn test_get_file_lock_requirement() {
844        // Read operations
845        let input = serde_json::json!({"path": "/test/file.txt"});
846        let req = AccessControlManager::get_file_lock_requirement("read_file", &input);
847        assert!(matches!(req, Some((_, LockType::Read))));
848
849        // Write operations
850        let req = AccessControlManager::get_file_lock_requirement("write_file", &input);
851        assert!(matches!(req, Some((_, LockType::Write))));
852
853        let req = AccessControlManager::get_file_lock_requirement("edit_file", &input);
854        assert!(matches!(req, Some((_, LockType::Write))));
855
856        // Unknown tool
857        let req = AccessControlManager::get_file_lock_requirement("unknown_tool", &input);
858        assert!(req.is_none());
859    }
860
861    #[tokio::test]
862    async fn test_detect_build_command() {
863        assert!(AccessControlManager::detect_build_command("cargo build"));
864        assert!(AccessControlManager::detect_build_command(
865            "cargo build --release"
866        ));
867        assert!(AccessControlManager::detect_build_command("npm run build"));
868        assert!(AccessControlManager::detect_build_command("make all"));
869        assert!(AccessControlManager::detect_build_command(
870            "gcc -o main main.c"
871        ));
872
873        assert!(!AccessControlManager::detect_build_command("ls -la"));
874        assert!(!AccessControlManager::detect_build_command("cargo test"));
875        assert!(!AccessControlManager::detect_build_command("echo hello"));
876    }
877
878    #[tokio::test]
879    async fn test_detect_test_command() {
880        assert!(AccessControlManager::detect_test_command("cargo test"));
881        assert!(AccessControlManager::detect_test_command(
882            "cargo test --release"
883        ));
884        assert!(AccessControlManager::detect_test_command("npm test"));
885        assert!(AccessControlManager::detect_test_command("pytest"));
886        assert!(AccessControlManager::detect_test_command("jest"));
887
888        assert!(!AccessControlManager::detect_test_command("ls -la"));
889        assert!(!AccessControlManager::detect_test_command("cargo build"));
890        assert!(!AccessControlManager::detect_test_command("echo hello"));
891    }
892
893    #[tokio::test]
894    async fn test_get_resource_requirement() {
895        let manager = create_manager();
896
897        // Build command
898        let req = manager.get_resource_requirement("cargo build");
899        assert!(matches!(req, Some((ResourceType::Build, _))));
900
901        // Test command
902        let req = manager.get_resource_requirement("cargo test");
903        assert!(matches!(req, Some((ResourceType::Test, _))));
904
905        // Build + test (cargo test with build)
906        let req = manager.get_resource_requirement("cargo build && cargo test");
907        assert!(matches!(req, Some((ResourceType::BuildTest, _))));
908
909        // Neither
910        let req = manager.get_resource_requirement("ls -la");
911        assert!(req.is_none());
912    }
913
914    #[tokio::test]
915    async fn test_acquire_for_tool_file_operation() {
916        let manager = create_manager();
917
918        // Read operation should succeed without prior read
919        let input = serde_json::json!({"path": "/test/file.txt"});
920        let result = manager
921            .acquire_for_tool("agent-1", "read_file", &input)
922            .await;
923        assert!(result.is_ok());
924        let bundle = result.unwrap();
925        assert!(bundle.file_lock.is_some());
926        assert!(bundle.resource_lock.is_none());
927    }
928
929    #[tokio::test]
930    async fn test_acquire_for_tool_build_command() {
931        let manager = create_manager();
932
933        let input = serde_json::json!({"command": "cargo build"});
934        let result = manager
935            .acquire_for_tool("agent-1", "execute_command", &input)
936            .await;
937        assert!(result.is_ok());
938        let bundle = result.unwrap();
939        assert!(bundle.file_lock.is_none());
940        assert!(bundle.resource_lock.is_some());
941    }
942
943    #[tokio::test]
944    async fn test_cleanup_agent() {
945        let manager = create_manager();
946
947        // Acquire some locks
948        let input = serde_json::json!({"path": "/test/file.txt"});
949        let bundle = manager
950            .acquire_for_tool("agent-1", "read_file", &input)
951            .await
952            .unwrap();
953
954        // Forget the bundle to prevent auto-release
955        std::mem::forget(bundle);
956
957        // Track a file read
958        manager
959            .track_file_read("agent-1", &PathBuf::from("/test/file.txt"))
960            .await;
961
962        // Cleanup
963        let (file_released, _resource_released, _persistent_released) =
964            manager.cleanup_agent("agent-1").await;
965        assert_eq!(file_released, 1);
966
967        // Tracking should be cleared
968        assert!(
969            !manager
970                .has_read_file("agent-1", &PathBuf::from("/test/file.txt"))
971                .await
972        );
973    }
974
975    #[tokio::test]
976    async fn test_clear_tracking_for_agent() {
977        let manager = create_manager();
978
979        manager
980            .track_file_read("agent-1", &PathBuf::from("/test/file1.txt"))
981            .await;
982        manager
983            .track_file_read("agent-1", &PathBuf::from("/test/file2.txt"))
984            .await;
985        manager
986            .track_file_read("agent-2", &PathBuf::from("/test/file1.txt"))
987            .await;
988
989        manager.clear_tracking_for_agent("agent-1").await;
990
991        assert!(
992            !manager
993                .has_read_file("agent-1", &PathBuf::from("/test/file1.txt"))
994                .await
995        );
996        assert!(
997            !manager
998                .has_read_file("agent-1", &PathBuf::from("/test/file2.txt"))
999                .await
1000        );
1001        // agent-2 tracking should remain
1002        assert!(
1003            manager
1004                .has_read_file("agent-2", &PathBuf::from("/test/file1.txt"))
1005                .await
1006        );
1007    }
1008}