1use 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#[async_trait::async_trait]
31pub trait LockPersistence: Send + Sync {
32 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 async fn release(&self, lock_type: &str, resource_path: &str, agent_id: &str) -> Result<()>;
45
46 async fn release_all_for_agent(&self, agent_id: &str) -> Result<usize>;
48
49 async fn cleanup_stale(&self) -> Result<usize>;
51}
52
53#[derive(Debug, Clone)]
55pub enum ContentionStrategy {
56 FailFast,
58 WaitWithTimeout(Duration),
60 RetryWithBackoff {
62 initial_delay: Duration,
64 max_retries: u32,
66 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
81pub struct LockBundle {
83 pub file_lock: Option<LockGuard>,
85 pub resource_lock: Option<ResourceLockGuard>,
87}
88
89impl LockBundle {
90 pub fn empty() -> Self {
92 Self {
93 file_lock: None,
94 resource_lock: None,
95 }
96 }
97
98 pub fn has_locks(&self) -> bool {
100 self.file_lock.is_some() || self.resource_lock.is_some()
101 }
102}
103
104pub struct AccessControlManager {
106 file_locks: Arc<FileLockManager>,
108 resource_locks: Arc<ResourceLockManager>,
110 contention_strategy: ContentionStrategy,
112 read_tracking: RwLock<HashMap<String, HashSet<PathBuf>>>,
114 project_root: PathBuf,
116 lock_store: Option<Arc<dyn LockPersistence>>,
118}
119
120impl AccessControlManager {
121 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 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 pub fn with_strategy(mut self, strategy: ContentionStrategy) -> Self {
151 self.contention_strategy = strategy;
152 self
153 }
154
155 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 pub fn lock_store(&self) -> Option<&Arc<dyn LockPersistence>> {
163 self.lock_store.as_ref()
164 }
165
166 pub fn file_locks(&self) -> &Arc<FileLockManager> {
168 &self.file_locks
169 }
170
171 pub fn resource_locks(&self) -> &Arc<ResourceLockManager> {
173 &self.resource_locks
174 }
175
176 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 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 pub async fn validate_write(&self, agent_id: &str, path: &Path) -> Result<()> {
198 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 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 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_file" | "list_directory" | "search_files" => LockType::Read,
233 "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 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 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 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 let scope = ResourceScope::Project(self.project_root.clone());
331
332 Some((resource_type, scope))
333 }
334
335 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 if let Some((path, lock_type)) = Self::get_file_lock_requirement(tool_name, input) {
346 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 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 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 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 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 Ok(true)
413 }
414 }
415
416 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 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 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 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 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 if self
477 .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
478 .await?
479 {
480 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 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 if self
528 .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
529 .await?
530 {
531 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 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 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 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 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 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 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 if self
643 .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
644 .await?
645 {
646 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 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 if self
697 .try_acquire_persistent_lock(agent_id, lock_type_str, &resource_path)
698 .await?
699 {
700 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 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 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 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 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 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 let dir = tempdir().unwrap();
819 let file_path = dir.path().join("test.txt");
820 std::fs::write(&file_path, "test content").unwrap();
821
822 let result = manager.validate_write("agent-1", &file_path).await;
824 assert!(result.is_err());
825
826 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 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 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 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 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 let req = manager.get_resource_requirement("cargo build");
899 assert!(matches!(req, Some((ResourceType::Build, _))));
900
901 let req = manager.get_resource_requirement("cargo test");
903 assert!(matches!(req, Some((ResourceType::Test, _))));
904
905 let req = manager.get_resource_requirement("cargo build && cargo test");
907 assert!(matches!(req, Some((ResourceType::BuildTest, _))));
908
909 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 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 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 std::mem::forget(bundle);
956
957 manager
959 .track_file_read("agent-1", &PathBuf::from("/test/file.txt"))
960 .await;
961
962 let (file_released, _resource_released, _persistent_released) =
964 manager.cleanup_agent("agent-1").await;
965 assert_eq!(file_released, 1);
966
967 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 assert!(
1003 manager
1004 .has_read_file("agent-2", &PathBuf::from("/test/file1.txt"))
1005 .await
1006 );
1007 }
1008}