aster/agents/context/
isolation.rs

1//! Context Isolation
2//!
3//! Provides sandboxed execution environments for agents with
4//! resource limits and tool permission enforcement.
5//!
6//! This module implements:
7//! - Sandbox creation with configurable resource limits
8//! - Tool permission enforcement (allowed/denied lists)
9//! - Sandbox state management (active, suspended, terminated)
10//! - Automatic cleanup of expired sandboxes
11
12use std::collections::{HashMap, HashSet};
13
14use chrono::{DateTime, Duration, Utc};
15use serde::{Deserialize, Serialize};
16
17use super::types::{AgentContext, AgentContextError, AgentContextResult, ContextUpdate};
18
19/// Sandbox state representing the lifecycle of a sandboxed context
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
21#[serde(rename_all = "camelCase")]
22pub enum SandboxState {
23    /// Sandbox is active and can execute operations
24    #[default]
25    Active,
26    /// Sandbox is suspended due to resource limits or manual suspension
27    Suspended,
28    /// Sandbox is terminated and cannot be resumed
29    Terminated,
30}
31
32impl SandboxState {
33    /// Check if the sandbox can transition to the given state
34    pub fn can_transition_to(&self, target: SandboxState) -> bool {
35        match (self, target) {
36            // Active can go to Suspended or Terminated
37            (SandboxState::Active, SandboxState::Suspended) => true,
38            (SandboxState::Active, SandboxState::Terminated) => true,
39            // Suspended can go to Active (resume) or Terminated
40            (SandboxState::Suspended, SandboxState::Active) => true,
41            (SandboxState::Suspended, SandboxState::Terminated) => true,
42            // Terminated is final - cannot transition
43            (SandboxState::Terminated, _) => false,
44            // Same state transitions are allowed (no-op)
45            (s1, s2) if *s1 == s2 => true,
46            _ => false,
47        }
48    }
49}
50
51/// Resource usage tracking for a sandbox
52#[derive(Debug, Clone, Default, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct ResourceUsage {
55    /// Current token count used
56    pub tokens_used: usize,
57    /// Number of files accessed
58    pub files_accessed: usize,
59    /// Number of tool results stored
60    pub tool_results_count: usize,
61    /// Number of tool calls made
62    pub tool_calls_made: usize,
63}
64
65impl ResourceUsage {
66    /// Create new resource usage tracker
67    pub fn new() -> Self {
68        Self::default()
69    }
70
71    /// Add tokens to usage
72    pub fn add_tokens(&mut self, count: usize) {
73        self.tokens_used += count;
74    }
75
76    /// Increment file access count
77    pub fn add_file_access(&mut self) {
78        self.files_accessed += 1;
79    }
80
81    /// Increment tool results count
82    pub fn add_tool_result(&mut self) {
83        self.tool_results_count += 1;
84    }
85
86    /// Increment tool calls count
87    pub fn add_tool_call(&mut self) {
88        self.tool_calls_made += 1;
89    }
90}
91
92/// Resource restrictions for a sandbox
93#[derive(Debug, Clone, Serialize, Deserialize)]
94#[serde(rename_all = "camelCase")]
95pub struct SandboxRestrictions {
96    /// Maximum tokens allowed in the sandbox
97    pub max_tokens: usize,
98    /// Maximum number of files that can be accessed
99    pub max_files: usize,
100    /// Maximum number of tool results that can be stored
101    pub max_tool_results: usize,
102    /// Set of tools that are explicitly allowed (if Some, only these tools are allowed)
103    pub allowed_tools: Option<HashSet<String>>,
104    /// Set of tools that are explicitly denied (checked after allowed_tools)
105    pub denied_tools: Option<HashSet<String>>,
106}
107
108impl Default for SandboxRestrictions {
109    fn default() -> Self {
110        Self {
111            max_tokens: 100_000,
112            max_files: 50,
113            max_tool_results: 100,
114            allowed_tools: None,
115            denied_tools: None,
116        }
117    }
118}
119
120impl SandboxRestrictions {
121    /// Create restrictions with custom limits
122    pub fn new(max_tokens: usize, max_files: usize, max_tool_results: usize) -> Self {
123        Self {
124            max_tokens,
125            max_files,
126            max_tool_results,
127            allowed_tools: None,
128            denied_tools: None,
129        }
130    }
131
132    /// Set allowed tools (whitelist)
133    pub fn with_allowed_tools(
134        mut self,
135        tools: impl IntoIterator<Item = impl Into<String>>,
136    ) -> Self {
137        self.allowed_tools = Some(tools.into_iter().map(|t| t.into()).collect());
138        self
139    }
140
141    /// Set denied tools (blacklist)
142    pub fn with_denied_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
143        self.denied_tools = Some(tools.into_iter().map(|t| t.into()).collect());
144        self
145    }
146
147    /// Check if a tool is allowed based on the restrictions
148    pub fn is_tool_allowed(&self, tool_name: &str) -> bool {
149        // If there's an allowed list, tool must be in it
150        if let Some(allowed) = &self.allowed_tools {
151            if !allowed.contains(tool_name) {
152                return false;
153            }
154        }
155
156        // If there's a denied list, tool must not be in it
157        if let Some(denied) = &self.denied_tools {
158            if denied.contains(tool_name) {
159                return false;
160            }
161        }
162
163        true
164    }
165
166    /// Check if resource usage exceeds any limit
167    pub fn check_limits(&self, usage: &ResourceUsage) -> Option<ResourceLimitViolation> {
168        if usage.tokens_used > self.max_tokens {
169            return Some(ResourceLimitViolation::TokensExceeded {
170                used: usage.tokens_used,
171                limit: self.max_tokens,
172            });
173        }
174        if usage.files_accessed > self.max_files {
175            return Some(ResourceLimitViolation::FilesExceeded {
176                used: usage.files_accessed,
177                limit: self.max_files,
178            });
179        }
180        if usage.tool_results_count > self.max_tool_results {
181            return Some(ResourceLimitViolation::ToolResultsExceeded {
182                used: usage.tool_results_count,
183                limit: self.max_tool_results,
184            });
185        }
186        None
187    }
188}
189
190/// Types of resource limit violations
191#[derive(Debug, Clone, PartialEq, Eq)]
192pub enum ResourceLimitViolation {
193    /// Token limit exceeded
194    TokensExceeded { used: usize, limit: usize },
195    /// File access limit exceeded
196    FilesExceeded { used: usize, limit: usize },
197    /// Tool results limit exceeded
198    ToolResultsExceeded { used: usize, limit: usize },
199}
200
201impl std::fmt::Display for ResourceLimitViolation {
202    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
203        match self {
204            ResourceLimitViolation::TokensExceeded { used, limit } => {
205                write!(f, "Token limit exceeded: {} used, {} allowed", used, limit)
206            }
207            ResourceLimitViolation::FilesExceeded { used, limit } => {
208                write!(
209                    f,
210                    "File limit exceeded: {} accessed, {} allowed",
211                    used, limit
212                )
213            }
214            ResourceLimitViolation::ToolResultsExceeded { used, limit } => {
215                write!(
216                    f,
217                    "Tool results limit exceeded: {} stored, {} allowed",
218                    used, limit
219                )
220            }
221        }
222    }
223}
224
225/// A sandboxed context with resource limits and state management
226#[derive(Debug, Clone, Serialize, Deserialize)]
227#[serde(rename_all = "camelCase")]
228pub struct SandboxedContext {
229    /// Unique sandbox identifier
230    pub sandbox_id: String,
231    /// Associated agent ID
232    pub agent_id: String,
233    /// The isolated context
234    pub context: AgentContext,
235    /// Resource restrictions
236    pub restrictions: SandboxRestrictions,
237    /// Current sandbox state
238    pub state: SandboxState,
239    /// Creation timestamp
240    pub created_at: DateTime<Utc>,
241    /// Expiration timestamp (if set)
242    pub expires_at: Option<DateTime<Utc>>,
243    /// Current resource usage
244    pub resources: ResourceUsage,
245    /// Reason for suspension (if suspended)
246    pub suspension_reason: Option<String>,
247}
248
249impl SandboxedContext {
250    /// Create a new sandboxed context
251    pub fn new(
252        context: AgentContext,
253        agent_id: impl Into<String>,
254        restrictions: Option<SandboxRestrictions>,
255    ) -> Self {
256        Self {
257            sandbox_id: uuid::Uuid::new_v4().to_string(),
258            agent_id: agent_id.into(),
259            context,
260            restrictions: restrictions.unwrap_or_default(),
261            state: SandboxState::Active,
262            created_at: Utc::now(),
263            expires_at: None,
264            resources: ResourceUsage::new(),
265            suspension_reason: None,
266        }
267    }
268
269    /// Set expiration time
270    pub fn with_expiration(mut self, expires_at: DateTime<Utc>) -> Self {
271        self.expires_at = Some(expires_at);
272        self
273    }
274
275    /// Set expiration duration from now
276    pub fn with_ttl(mut self, ttl: Duration) -> Self {
277        self.expires_at = Some(Utc::now() + ttl);
278        self
279    }
280
281    /// Check if the sandbox has expired
282    pub fn is_expired(&self) -> bool {
283        if let Some(expires_at) = self.expires_at {
284            Utc::now() > expires_at
285        } else {
286            false
287        }
288    }
289
290    /// Check if the sandbox is active
291    pub fn is_active(&self) -> bool {
292        self.state == SandboxState::Active && !self.is_expired()
293    }
294
295    /// Check if a tool is allowed in this sandbox
296    pub fn is_tool_allowed(&self, tool_name: &str) -> bool {
297        self.restrictions.is_tool_allowed(tool_name)
298    }
299
300    /// Check resource limits and return violation if any
301    pub fn check_limits(&self) -> Option<ResourceLimitViolation> {
302        self.restrictions.check_limits(&self.resources)
303    }
304
305    /// Record token usage and check limits
306    pub fn record_tokens(&mut self, count: usize) -> AgentContextResult<()> {
307        self.resources.add_tokens(count);
308        self.check_and_suspend_if_exceeded()
309    }
310
311    /// Record file access and check limits
312    pub fn record_file_access(&mut self) -> AgentContextResult<()> {
313        self.resources.add_file_access();
314        self.check_and_suspend_if_exceeded()
315    }
316
317    /// Record tool result and check limits
318    pub fn record_tool_result(&mut self) -> AgentContextResult<()> {
319        self.resources.add_tool_result();
320        self.check_and_suspend_if_exceeded()
321    }
322
323    /// Check limits and suspend if exceeded
324    fn check_and_suspend_if_exceeded(&mut self) -> AgentContextResult<()> {
325        if let Some(violation) = self.check_limits() {
326            self.state = SandboxState::Suspended;
327            self.suspension_reason = Some(violation.to_string());
328            return Err(AgentContextError::ResourceLimitExceeded(
329                violation.to_string(),
330            ));
331        }
332        Ok(())
333    }
334}
335
336/// Context Isolation Manager
337///
338/// Manages sandboxed execution environments for agents with:
339/// - Resource limit enforcement
340/// - Tool permission management
341/// - Sandbox lifecycle management
342/// - Automatic cleanup of expired sandboxes
343#[derive(Debug, Default)]
344pub struct ContextIsolation {
345    /// Map of sandbox ID to sandboxed context
346    sandboxes: HashMap<String, SandboxedContext>,
347    /// Map of agent ID to sandbox ID for quick lookup
348    agent_sandboxes: HashMap<String, String>,
349}
350
351impl ContextIsolation {
352    /// Create a new context isolation manager
353    pub fn new() -> Self {
354        Self {
355            sandboxes: HashMap::new(),
356            agent_sandboxes: HashMap::new(),
357        }
358    }
359
360    /// Create a new sandbox for an agent context
361    pub fn create_sandbox(
362        &mut self,
363        context: AgentContext,
364        agent_id: Option<String>,
365        restrictions: Option<SandboxRestrictions>,
366    ) -> SandboxedContext {
367        let agent_id = agent_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
368        let sandbox = SandboxedContext::new(context, agent_id.clone(), restrictions);
369
370        let sandbox_id = sandbox.sandbox_id.clone();
371        self.sandboxes.insert(sandbox_id.clone(), sandbox.clone());
372        self.agent_sandboxes.insert(agent_id, sandbox_id);
373
374        sandbox
375    }
376
377    /// Create a sandbox with expiration
378    pub fn create_sandbox_with_ttl(
379        &mut self,
380        context: AgentContext,
381        agent_id: Option<String>,
382        restrictions: Option<SandboxRestrictions>,
383        ttl: Duration,
384    ) -> SandboxedContext {
385        let agent_id = agent_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
386        let sandbox = SandboxedContext::new(context, agent_id.clone(), restrictions).with_ttl(ttl);
387
388        let sandbox_id = sandbox.sandbox_id.clone();
389        self.sandboxes.insert(sandbox_id.clone(), sandbox.clone());
390        self.agent_sandboxes.insert(agent_id, sandbox_id);
391
392        sandbox
393    }
394
395    /// Get a sandbox by ID
396    pub fn get_sandbox(&self, sandbox_id: &str) -> Option<&SandboxedContext> {
397        self.sandboxes.get(sandbox_id)
398    }
399
400    /// Get a mutable sandbox by ID
401    pub fn get_sandbox_mut(&mut self, sandbox_id: &str) -> Option<&mut SandboxedContext> {
402        self.sandboxes.get_mut(sandbox_id)
403    }
404
405    /// Get the isolated context for an agent
406    pub fn get_isolated_context(&self, agent_id: &str) -> Option<&AgentContext> {
407        self.agent_sandboxes
408            .get(agent_id)
409            .and_then(|sandbox_id| self.sandboxes.get(sandbox_id))
410            .map(|sandbox| &sandbox.context)
411    }
412
413    /// Get sandbox by agent ID
414    pub fn get_sandbox_by_agent(&self, agent_id: &str) -> Option<&SandboxedContext> {
415        self.agent_sandboxes
416            .get(agent_id)
417            .and_then(|sandbox_id| self.sandboxes.get(sandbox_id))
418    }
419
420    /// Get mutable sandbox by agent ID
421    pub fn get_sandbox_by_agent_mut(&mut self, agent_id: &str) -> Option<&mut SandboxedContext> {
422        if let Some(sandbox_id) = self.agent_sandboxes.get(agent_id).cloned() {
423            self.sandboxes.get_mut(&sandbox_id)
424        } else {
425            None
426        }
427    }
428
429    /// Update a sandbox's context
430    pub fn update_sandbox(
431        &mut self,
432        sandbox_id: &str,
433        updates: ContextUpdate,
434    ) -> AgentContextResult<()> {
435        let sandbox = self
436            .sandboxes
437            .get_mut(sandbox_id)
438            .ok_or_else(|| AgentContextError::NotFound(sandbox_id.to_string()))?;
439
440        // Check if sandbox is active
441        if sandbox.state != SandboxState::Active {
442            return Err(AgentContextError::InvalidStateTransition(format!(
443                "Cannot update sandbox in {:?} state",
444                sandbox.state
445            )));
446        }
447
448        // Check if expired
449        if sandbox.is_expired() {
450            sandbox.state = SandboxState::Terminated;
451            return Err(AgentContextError::InvalidStateTransition(
452                "Sandbox has expired".to_string(),
453            ));
454        }
455
456        // Apply updates to context
457        let context = &mut sandbox.context;
458
459        if let Some(messages) = updates.add_messages {
460            context.conversation_history.extend(messages);
461        }
462
463        if let Some(files) = updates.add_files {
464            for file in files {
465                sandbox.resources.add_file_access();
466                context.file_context.push(file);
467            }
468        }
469
470        if let Some(results) = updates.add_tool_results {
471            for result in results {
472                sandbox.resources.add_tool_result();
473                context.tool_results.push(result);
474            }
475        }
476
477        if let Some(env) = updates.set_environment {
478            context.environment.extend(env);
479        }
480
481        if let Some(prompt) = updates.set_system_prompt {
482            context.system_prompt = Some(prompt);
483        }
484
485        if let Some(dir) = updates.set_working_directory {
486            context.working_directory = dir;
487        }
488
489        if let Some(tags) = updates.add_tags {
490            for tag in tags {
491                context.metadata.add_tag(tag);
492            }
493        }
494
495        if let Some(custom) = updates.set_custom_metadata {
496            for (key, value) in custom {
497                context.metadata.set_custom(key, value);
498            }
499        }
500
501        context.metadata.touch();
502
503        // Check resource limits after update
504        sandbox.check_and_suspend_if_exceeded()?;
505
506        Ok(())
507    }
508
509    /// Check if a tool is allowed in a sandbox
510    pub fn is_tool_allowed(&self, sandbox_id: &str, tool_name: &str) -> bool {
511        self.sandboxes
512            .get(sandbox_id)
513            .map(|s| s.is_tool_allowed(tool_name))
514            .unwrap_or(false)
515    }
516
517    /// Suspend a sandbox
518    pub fn suspend(&mut self, sandbox_id: &str) -> AgentContextResult<()> {
519        let sandbox = self
520            .sandboxes
521            .get_mut(sandbox_id)
522            .ok_or_else(|| AgentContextError::NotFound(sandbox_id.to_string()))?;
523
524        if !sandbox.state.can_transition_to(SandboxState::Suspended) {
525            return Err(AgentContextError::InvalidStateTransition(format!(
526                "Cannot suspend sandbox in {:?} state",
527                sandbox.state
528            )));
529        }
530
531        sandbox.state = SandboxState::Suspended;
532        sandbox.suspension_reason = Some("Manually suspended".to_string());
533        Ok(())
534    }
535
536    /// Resume a suspended sandbox
537    pub fn resume(&mut self, sandbox_id: &str) -> AgentContextResult<()> {
538        let sandbox = self
539            .sandboxes
540            .get_mut(sandbox_id)
541            .ok_or_else(|| AgentContextError::NotFound(sandbox_id.to_string()))?;
542
543        if !sandbox.state.can_transition_to(SandboxState::Active) {
544            return Err(AgentContextError::InvalidStateTransition(format!(
545                "Cannot resume sandbox in {:?} state",
546                sandbox.state
547            )));
548        }
549
550        // Check if expired before resuming
551        if sandbox.is_expired() {
552            sandbox.state = SandboxState::Terminated;
553            return Err(AgentContextError::InvalidStateTransition(
554                "Cannot resume expired sandbox".to_string(),
555            ));
556        }
557
558        sandbox.state = SandboxState::Active;
559        sandbox.suspension_reason = None;
560        Ok(())
561    }
562
563    /// Terminate a sandbox
564    pub fn terminate(&mut self, sandbox_id: &str) -> AgentContextResult<()> {
565        let sandbox = self
566            .sandboxes
567            .get_mut(sandbox_id)
568            .ok_or_else(|| AgentContextError::NotFound(sandbox_id.to_string()))?;
569
570        if !sandbox.state.can_transition_to(SandboxState::Terminated) {
571            return Err(AgentContextError::InvalidStateTransition(format!(
572                "Cannot terminate sandbox in {:?} state",
573                sandbox.state
574            )));
575        }
576
577        sandbox.state = SandboxState::Terminated;
578        Ok(())
579    }
580
581    /// Cleanup a specific sandbox (remove from memory)
582    pub fn cleanup(&mut self, sandbox_id: &str) {
583        if let Some(sandbox) = self.sandboxes.remove(sandbox_id) {
584            self.agent_sandboxes.remove(&sandbox.agent_id);
585        }
586    }
587
588    /// Cleanup all expired sandboxes
589    /// Returns the number of sandboxes cleaned up
590    pub fn cleanup_expired(&mut self) -> usize {
591        let expired_ids: Vec<String> = self
592            .sandboxes
593            .iter()
594            .filter(|(_, sandbox)| sandbox.is_expired())
595            .map(|(id, _)| id.clone())
596            .collect();
597
598        let count = expired_ids.len();
599
600        for sandbox_id in expired_ids {
601            self.cleanup(&sandbox_id);
602        }
603
604        count
605    }
606
607    /// Get all sandbox IDs
608    pub fn list_sandbox_ids(&self) -> Vec<String> {
609        self.sandboxes.keys().cloned().collect()
610    }
611
612    /// Get all sandboxes in a specific state
613    pub fn list_sandboxes_by_state(&self, state: SandboxState) -> Vec<&SandboxedContext> {
614        self.sandboxes
615            .values()
616            .filter(|s| s.state == state)
617            .collect()
618    }
619
620    /// Get sandbox count
621    pub fn sandbox_count(&self) -> usize {
622        self.sandboxes.len()
623    }
624
625    /// Get active sandbox count
626    pub fn active_sandbox_count(&self) -> usize {
627        self.sandboxes.values().filter(|s| s.is_active()).count()
628    }
629}
630
631#[cfg(test)]
632mod tests {
633    use super::*;
634    use crate::agents::context::types::FileContext;
635
636    #[test]
637    fn test_sandbox_state_transitions() {
638        // Active -> Suspended
639        assert!(SandboxState::Active.can_transition_to(SandboxState::Suspended));
640        // Active -> Terminated
641        assert!(SandboxState::Active.can_transition_to(SandboxState::Terminated));
642        // Suspended -> Active
643        assert!(SandboxState::Suspended.can_transition_to(SandboxState::Active));
644        // Suspended -> Terminated
645        assert!(SandboxState::Suspended.can_transition_to(SandboxState::Terminated));
646        // Terminated -> anything is false
647        assert!(!SandboxState::Terminated.can_transition_to(SandboxState::Active));
648        assert!(!SandboxState::Terminated.can_transition_to(SandboxState::Suspended));
649        // Same state is allowed
650        assert!(SandboxState::Active.can_transition_to(SandboxState::Active));
651    }
652
653    #[test]
654    fn test_resource_usage_tracking() {
655        let mut usage = ResourceUsage::new();
656        assert_eq!(usage.tokens_used, 0);
657        assert_eq!(usage.files_accessed, 0);
658
659        usage.add_tokens(100);
660        assert_eq!(usage.tokens_used, 100);
661
662        usage.add_file_access();
663        assert_eq!(usage.files_accessed, 1);
664
665        usage.add_tool_result();
666        assert_eq!(usage.tool_results_count, 1);
667
668        usage.add_tool_call();
669        assert_eq!(usage.tool_calls_made, 1);
670    }
671
672    #[test]
673    fn test_sandbox_restrictions_default() {
674        let restrictions = SandboxRestrictions::default();
675        assert_eq!(restrictions.max_tokens, 100_000);
676        assert_eq!(restrictions.max_files, 50);
677        assert_eq!(restrictions.max_tool_results, 100);
678        assert!(restrictions.allowed_tools.is_none());
679        assert!(restrictions.denied_tools.is_none());
680    }
681
682    #[test]
683    fn test_sandbox_restrictions_tool_allowed() {
684        // No restrictions - all tools allowed
685        let restrictions = SandboxRestrictions::default();
686        assert!(restrictions.is_tool_allowed("bash"));
687        assert!(restrictions.is_tool_allowed("read_file"));
688
689        // With allowed list
690        let restrictions =
691            SandboxRestrictions::default().with_allowed_tools(vec!["bash", "read_file"]);
692        assert!(restrictions.is_tool_allowed("bash"));
693        assert!(restrictions.is_tool_allowed("read_file"));
694        assert!(!restrictions.is_tool_allowed("write_file"));
695
696        // With denied list
697        let restrictions = SandboxRestrictions::default().with_denied_tools(vec!["bash"]);
698        assert!(!restrictions.is_tool_allowed("bash"));
699        assert!(restrictions.is_tool_allowed("read_file"));
700
701        // With both allowed and denied
702        let restrictions = SandboxRestrictions::default()
703            .with_allowed_tools(vec!["bash", "read_file", "write_file"])
704            .with_denied_tools(vec!["write_file"]);
705        assert!(restrictions.is_tool_allowed("bash"));
706        assert!(restrictions.is_tool_allowed("read_file"));
707        assert!(!restrictions.is_tool_allowed("write_file")); // Denied takes precedence
708        assert!(!restrictions.is_tool_allowed("other")); // Not in allowed list
709    }
710
711    #[test]
712    fn test_sandbox_restrictions_check_limits() {
713        let restrictions = SandboxRestrictions::new(100, 5, 10);
714
715        // Within limits
716        let usage = ResourceUsage {
717            tokens_used: 50,
718            files_accessed: 3,
719            tool_results_count: 5,
720            tool_calls_made: 0,
721        };
722        assert!(restrictions.check_limits(&usage).is_none());
723
724        // Token limit exceeded
725        let usage = ResourceUsage {
726            tokens_used: 150,
727            files_accessed: 3,
728            tool_results_count: 5,
729            tool_calls_made: 0,
730        };
731        assert!(matches!(
732            restrictions.check_limits(&usage),
733            Some(ResourceLimitViolation::TokensExceeded { .. })
734        ));
735
736        // File limit exceeded
737        let usage = ResourceUsage {
738            tokens_used: 50,
739            files_accessed: 10,
740            tool_results_count: 5,
741            tool_calls_made: 0,
742        };
743        assert!(matches!(
744            restrictions.check_limits(&usage),
745            Some(ResourceLimitViolation::FilesExceeded { .. })
746        ));
747
748        // Tool results limit exceeded
749        let usage = ResourceUsage {
750            tokens_used: 50,
751            files_accessed: 3,
752            tool_results_count: 15,
753            tool_calls_made: 0,
754        };
755        assert!(matches!(
756            restrictions.check_limits(&usage),
757            Some(ResourceLimitViolation::ToolResultsExceeded { .. })
758        ));
759    }
760
761    #[test]
762    fn test_sandboxed_context_creation() {
763        let context = AgentContext::new();
764        let sandbox = SandboxedContext::new(context, "agent-1", None);
765
766        assert!(!sandbox.sandbox_id.is_empty());
767        assert_eq!(sandbox.agent_id, "agent-1");
768        assert_eq!(sandbox.state, SandboxState::Active);
769        assert!(sandbox.is_active());
770        assert!(!sandbox.is_expired());
771    }
772
773    #[test]
774    fn test_sandboxed_context_with_ttl() {
775        let context = AgentContext::new();
776        let sandbox = SandboxedContext::new(context, "agent-1", None).with_ttl(Duration::hours(1));
777
778        assert!(sandbox.expires_at.is_some());
779        assert!(!sandbox.is_expired());
780
781        // Create an already expired sandbox
782        let context = AgentContext::new();
783        let sandbox =
784            SandboxedContext::new(context, "agent-2", None).with_ttl(Duration::seconds(-1));
785
786        assert!(sandbox.is_expired());
787        assert!(!sandbox.is_active());
788    }
789
790    #[test]
791    fn test_sandboxed_context_record_resources() {
792        let context = AgentContext::new();
793        let restrictions = SandboxRestrictions::new(100, 5, 10);
794        let mut sandbox = SandboxedContext::new(context, "agent-1", Some(restrictions));
795
796        // Record within limits
797        assert!(sandbox.record_tokens(50).is_ok());
798        assert_eq!(sandbox.resources.tokens_used, 50);
799
800        // Record exceeding limits
801        let result = sandbox.record_tokens(100);
802        assert!(result.is_err());
803        assert_eq!(sandbox.state, SandboxState::Suspended);
804    }
805
806    #[test]
807    fn test_context_isolation_create_sandbox() {
808        let mut isolation = ContextIsolation::new();
809        let context = AgentContext::new();
810
811        let sandbox = isolation.create_sandbox(context, Some("agent-1".to_string()), None);
812
813        assert!(!sandbox.sandbox_id.is_empty());
814        assert_eq!(sandbox.agent_id, "agent-1");
815        assert_eq!(isolation.sandbox_count(), 1);
816
817        // Can retrieve by sandbox ID
818        assert!(isolation.get_sandbox(&sandbox.sandbox_id).is_some());
819
820        // Can retrieve by agent ID
821        assert!(isolation.get_isolated_context("agent-1").is_some());
822    }
823
824    #[test]
825    fn test_context_isolation_suspend_resume() {
826        let mut isolation = ContextIsolation::new();
827        let context = AgentContext::new();
828        let sandbox = isolation.create_sandbox(context, Some("agent-1".to_string()), None);
829        let sandbox_id = sandbox.sandbox_id.clone();
830
831        // Suspend
832        assert!(isolation.suspend(&sandbox_id).is_ok());
833        assert_eq!(
834            isolation.get_sandbox(&sandbox_id).unwrap().state,
835            SandboxState::Suspended
836        );
837
838        // Resume
839        assert!(isolation.resume(&sandbox_id).is_ok());
840        assert_eq!(
841            isolation.get_sandbox(&sandbox_id).unwrap().state,
842            SandboxState::Active
843        );
844
845        // Terminate
846        assert!(isolation.terminate(&sandbox_id).is_ok());
847        assert_eq!(
848            isolation.get_sandbox(&sandbox_id).unwrap().state,
849            SandboxState::Terminated
850        );
851
852        // Cannot resume terminated
853        assert!(isolation.resume(&sandbox_id).is_err());
854    }
855
856    #[test]
857    fn test_context_isolation_cleanup() {
858        let mut isolation = ContextIsolation::new();
859
860        // Create some sandboxes
861        let context1 = AgentContext::new();
862        let sandbox1 = isolation.create_sandbox(context1, Some("agent-1".to_string()), None);
863        let sandbox1_id = sandbox1.sandbox_id.clone();
864
865        let context2 = AgentContext::new();
866        let _sandbox2 = isolation.create_sandbox(context2, Some("agent-2".to_string()), None);
867
868        assert_eq!(isolation.sandbox_count(), 2);
869
870        // Cleanup one
871        isolation.cleanup(&sandbox1_id);
872        assert_eq!(isolation.sandbox_count(), 1);
873        assert!(isolation.get_sandbox(&sandbox1_id).is_none());
874        assert!(isolation.get_isolated_context("agent-1").is_none());
875    }
876
877    #[test]
878    fn test_context_isolation_cleanup_expired() {
879        let mut isolation = ContextIsolation::new();
880
881        // Create an expired sandbox
882        let context1 = AgentContext::new();
883        let _sandbox1 = isolation.create_sandbox_with_ttl(
884            context1,
885            Some("agent-1".to_string()),
886            None,
887            Duration::seconds(-1), // Already expired
888        );
889
890        // Create a non-expired sandbox
891        let context2 = AgentContext::new();
892        let _sandbox2 = isolation.create_sandbox_with_ttl(
893            context2,
894            Some("agent-2".to_string()),
895            None,
896            Duration::hours(1),
897        );
898
899        assert_eq!(isolation.sandbox_count(), 2);
900
901        // Cleanup expired
902        let cleaned = isolation.cleanup_expired();
903        assert_eq!(cleaned, 1);
904        assert_eq!(isolation.sandbox_count(), 1);
905        assert!(isolation.get_isolated_context("agent-1").is_none());
906        assert!(isolation.get_isolated_context("agent-2").is_some());
907    }
908
909    #[test]
910    fn test_context_isolation_is_tool_allowed() {
911        let mut isolation = ContextIsolation::new();
912        let context = AgentContext::new();
913        let restrictions =
914            SandboxRestrictions::default().with_allowed_tools(vec!["bash", "read_file"]);
915
916        let sandbox =
917            isolation.create_sandbox(context, Some("agent-1".to_string()), Some(restrictions));
918
919        assert!(isolation.is_tool_allowed(&sandbox.sandbox_id, "bash"));
920        assert!(isolation.is_tool_allowed(&sandbox.sandbox_id, "read_file"));
921        assert!(!isolation.is_tool_allowed(&sandbox.sandbox_id, "write_file"));
922        assert!(!isolation.is_tool_allowed("nonexistent", "bash"));
923    }
924
925    #[test]
926    fn test_context_isolation_update_sandbox() {
927        let mut isolation = ContextIsolation::new();
928        let context = AgentContext::new();
929        let sandbox = isolation.create_sandbox(context, Some("agent-1".to_string()), None);
930        let sandbox_id = sandbox.sandbox_id.clone();
931
932        let updates = ContextUpdate {
933            add_files: Some(vec![FileContext::new("/test.rs", "fn main() {}")]),
934            ..Default::default()
935        };
936
937        assert!(isolation.update_sandbox(&sandbox_id, updates).is_ok());
938
939        let sandbox = isolation.get_sandbox(&sandbox_id).unwrap();
940        assert_eq!(sandbox.context.file_context.len(), 1);
941        assert_eq!(sandbox.resources.files_accessed, 1);
942    }
943
944    #[test]
945    fn test_context_isolation_update_suspended_sandbox_fails() {
946        let mut isolation = ContextIsolation::new();
947        let context = AgentContext::new();
948        let sandbox = isolation.create_sandbox(context, Some("agent-1".to_string()), None);
949        let sandbox_id = sandbox.sandbox_id.clone();
950
951        // Suspend the sandbox
952        isolation.suspend(&sandbox_id).unwrap();
953
954        // Try to update - should fail
955        let updates = ContextUpdate {
956            add_files: Some(vec![FileContext::new("/test.rs", "fn main() {}")]),
957            ..Default::default()
958        };
959
960        assert!(isolation.update_sandbox(&sandbox_id, updates).is_err());
961    }
962
963    #[test]
964    fn test_context_isolation_list_by_state() {
965        let mut isolation = ContextIsolation::new();
966
967        // Create sandboxes in different states
968        let context1 = AgentContext::new();
969        let sandbox1 = isolation.create_sandbox(context1, Some("agent-1".to_string()), None);
970
971        let context2 = AgentContext::new();
972        let sandbox2 = isolation.create_sandbox(context2, Some("agent-2".to_string()), None);
973        isolation.suspend(&sandbox2.sandbox_id).unwrap();
974
975        let context3 = AgentContext::new();
976        let sandbox3 = isolation.create_sandbox(context3, Some("agent-3".to_string()), None);
977        isolation.terminate(&sandbox3.sandbox_id).unwrap();
978
979        // List by state
980        let active = isolation.list_sandboxes_by_state(SandboxState::Active);
981        assert_eq!(active.len(), 1);
982        assert_eq!(active[0].sandbox_id, sandbox1.sandbox_id);
983
984        let suspended = isolation.list_sandboxes_by_state(SandboxState::Suspended);
985        assert_eq!(suspended.len(), 1);
986
987        let terminated = isolation.list_sandboxes_by_state(SandboxState::Terminated);
988        assert_eq!(terminated.len(), 1);
989    }
990}
aster/agents/context/isolation.rs

aster/agents/context/
isolation.rs
