"""
Agent Lifecycle and Communication Contracts
This module defines comprehensive contracts for agent lifecycle management and communication
within the Agent Orchestration Platform, ensuring secure and reliable agent operations.
Architecture Integration:
- Design Patterns: State Machine for agent lifecycle, Observer for state changes
- Security Model: Agent isolation with secure communication channels and resource boundaries
- Performance Profile: O(1) state validation with cached agent status checks
Technical Decisions:
- Immutable State: Agent state transitions validated through immutable state objects
- Process Isolation: Each agent runs in separate process with secure communication
- Resource Monitoring: Continuous monitoring of agent resource usage and health
- Contract Enforcement: All agent operations validated through precondition/postcondition contracts
Dependencies & Integration:
- External: psutil for process monitoring and resource tracking
- Internal: security contracts, audit logging, filesystem boundaries, input validation
Quality Assurance:
- Test Coverage: Property-based testing for state transitions and resource limits
- Error Handling: Graceful degradation with automatic recovery and restart capabilities
Author: Adder_4 | Created: 2025-06-26 | Last Modified: 2025-06-26
"""
import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum, auto
from pathlib import Path
from typing import Any, Dict, List, Optional, Set
import psutil
from src.boundaries.audit import AuditCategory, AuditLevel, get_audit_logger
from src.boundaries.filesystem import FilesystemBoundary, ResourceLimits
from src.contracts.security import (
SecurityContractContext,
agent_creation_contract,
agent_deletion_contract,
get_security_context,
message_sending_contract,
status_read_contract,
)
from src.models.agent import (
AgentSpecialization,
AgentState,
AgentStatus,
ClaudeConfig,
ResourceMetrics,
)
from src.models.ids import AgentId, SessionId, validate_agent_name
from src.models.security import Permission, SecurityContext
from src.models.session import SecurityLevel, SessionState
from src.utils.contracts_shim import contract, ensure, require
from src.validators.input import InputType, validate_agent_name, validate_message_content
class AgentLifecycleError(Exception):
"""Base exception for agent lifecycle operations."""
pass
class AgentStateTransitionError(AgentLifecycleError):
"""Invalid agent state transition attempted."""
pass
class AgentResourceError(AgentLifecycleError):
"""Agent resource limit exceeded or unavailable."""
pass
class AgentCommunicationError(AgentLifecycleError):
"""Agent communication failure."""
pass
@dataclass(frozen=True)
class AgentCreationRequest:
"""Immutable request for agent creation with validation."""
session_id: SessionId
agent_name: str
specialization: AgentSpecialization = AgentSpecialization.GENERAL
system_prompt_suffix: str = ""
claude_config: Optional[ClaudeConfig] = None
resource_limits: Optional[ResourceLimits] = None
def __post_init__(self):
"""Validate creation request."""
validate_agent_name(self.agent_name)
if len(self.system_prompt_suffix) > 5000:
raise ValueError("System prompt suffix too long")
@dataclass(frozen=True)
class AgentDeletionRequest:
"""Immutable request for agent deletion with validation."""
agent_id: AgentId
force: bool = False
preserve_logs: bool = True
cleanup_resources: bool = True
@dataclass(frozen=True)
class MessageRequest:
"""Immutable request for sending message to agent."""
agent_id: AgentId
content: str
prepend_adder: bool = True
wait_for_response: bool = False
timeout_seconds: int = 30
def __post_init__(self):
"""Validate message request."""
if not self.content.strip():
raise ValueError("Message content cannot be empty")
if len(self.content) > 100000: # 100KB limit
raise ValueError("Message content too large")
class AgentStateValidator:
"""
Validator for agent state transitions and resource constraints.
Implements comprehensive validation for all agent lifecycle operations
with security-focused checks and resource monitoring.
"""
# Valid state transitions
VALID_TRANSITIONS = {
AgentStatus.CREATED: {
AgentStatus.STARTING,
AgentStatus.ERROR,
AgentStatus.TERMINATED,
},
AgentStatus.STARTING: {
AgentStatus.ACTIVE,
AgentStatus.ERROR,
AgentStatus.TERMINATED,
},
AgentStatus.ACTIVE: {
AgentStatus.IDLE,
AgentStatus.WORKING,
AgentStatus.ERROR,
AgentStatus.TERMINATED,
},
AgentStatus.IDLE: {
AgentStatus.WORKING,
AgentStatus.ACTIVE,
AgentStatus.ERROR,
AgentStatus.TERMINATED,
},
AgentStatus.WORKING: {
AgentStatus.IDLE,
AgentStatus.ACTIVE,
AgentStatus.ERROR,
AgentStatus.TERMINATED,
},
AgentStatus.ERROR: {AgentStatus.STARTING, AgentStatus.TERMINATED},
AgentStatus.TERMINATED: set(), # Terminal state
}
@classmethod
@contract(current_status="$AgentStatus", target_status="$AgentStatus")
def validate_state_transition(
cls, current_status: AgentStatus, target_status: AgentStatus
) -> bool:
"""
Validate agent state transition.
Contract: State transitions must follow valid state machine rules.
"""
return target_status in cls.VALID_TRANSITIONS.get(current_status, set())
@classmethod
@contract(agent_state="$AgentState")
def validate_agent_health(cls, agent_state: AgentState) -> List[str]:
"""
Validate agent health and return list of issues.
Contract: Returns empty list for healthy agents, issues list for unhealthy agents.
"""
issues = []
# Check resource limits
if not agent_state.resource_metrics.is_within_limits(agent_state.claude_config):
issues.append("Resource limits exceeded")
# Check heartbeat freshness
time_since_heartbeat = datetime.utcnow() - agent_state.last_heartbeat
if time_since_heartbeat > timedelta(minutes=5):
issues.append("Agent heartbeat stale")
# Check error count
if agent_state.error_count > 10:
issues.append("Too many errors")
# Check restart count
if agent_state.restart_count > 5:
issues.append("Too many restarts")
return issues
@classmethod
@contract(session_id="$SessionId", agent_count="int,>=0")
def validate_session_capacity(cls, session_id: SessionId, agent_count: int) -> bool:
"""
Validate session has capacity for additional agents.
Contract: Session must not exceed maximum agent limit (8).
"""
MAX_AGENTS_PER_SESSION = 8
return agent_count < MAX_AGENTS_PER_SESSION
class AgentLifecycleManager:
"""
Comprehensive agent lifecycle management with contract enforcement.
Manages agent creation, deletion, state transitions, and communication
with full security contract validation and audit logging.
"""
def __init__(self):
"""Initialize lifecycle manager."""
self.agents: Dict[AgentId, AgentState] = {}
self.agent_processes: Dict[AgentId, psutil.Process] = {}
self.state_validator = AgentStateValidator()
@agent_creation_contract
@require(lambda request: request.agent_name.startswith("Agent_"))
@require(lambda request: len(request.system_prompt_suffix) <= 5000)
@ensure(lambda result: result is not None)
async def create_agent(self, request: AgentCreationRequest) -> AgentState:
"""
Create new agent with comprehensive validation and security checks.
Contracts:
Preconditions:
- Agent name must follow Agent_# format
- System prompt suffix must be reasonable length
- Security context must be valid
- Session must exist and have capacity
Postconditions:
- Agent created with valid state
- Process spawned and monitored
- Agent registered in system
- Creation event audited
Invariants:
- Agent IDs are unique system-wide
- Resource limits enforced
- Security boundaries established
"""
# Validate input
name_validation = validate_agent_name(request.agent_name)
if not name_validation.is_valid:
raise AgentLifecycleError(
f"Invalid agent name: {name_validation.violations}"
)
# Check session capacity
session_agent_count = sum(
1
for agent in self.agents.values()
if agent.session_id == request.session_id
)
if not self.state_validator.validate_session_capacity(
request.session_id, session_agent_count
):
raise AgentLifecycleError("Session at maximum agent capacity")
# Generate unique agent ID
from src.models.ids import create_agent_id
agent_id = create_agent_id()
# Create agent state
agent_state = AgentState(
agent_id=agent_id,
session_id=request.session_id,
name=request.agent_name,
process_id=None, # Will be set when process starts
iterm_tab_id=None, # Will be set by iTerm integration
status=AgentStatus.CREATED,
specialization=request.specialization,
system_prompt_suffix=request.system_prompt_suffix,
claude_config=request.claude_config or ClaudeConfig(),
created_at=datetime.utcnow(),
last_heartbeat=datetime.utcnow(),
resource_metrics=ResourceMetrics(
cpu_percent=0.0,
memory_mb=0,
file_descriptors=0,
uptime_seconds=0,
last_activity=datetime.utcnow(),
),
)
# Register agent
self.agents[agent_id] = agent_state
# Log creation
audit_logger = get_audit_logger()
await audit_logger.log_agent_created(
agent_id=str(agent_id),
session_id=str(request.session_id),
user_id=get_security_context().user_id,
metadata={
"agent_name": request.agent_name,
"specialization": request.specialization.value,
"resource_limits": str(request.resource_limits),
},
)
return agent_state
@agent_deletion_contract
@require(lambda agent_id: agent_id is not None)
@ensure(lambda result, agent_id: agent_id not in self.agents)
async def delete_agent(self, request: AgentDeletionRequest) -> bool:
"""
Delete agent with cleanup and resource reclamation.
Contracts:
Preconditions:
- Agent must exist in system
- Security context must allow deletion
Postconditions:
- Agent removed from system
- Process terminated gracefully
- Resources cleaned up
- Deletion event audited
Invariants:
- No orphaned processes remain
- All resources properly reclaimed
- Audit trail maintained
"""
agent_id = request.agent_id
# Check if agent exists
if agent_id not in self.agents:
raise AgentLifecycleError(f"Agent not found: {agent_id}")
agent_state = self.agents[agent_id]
try:
# Terminate process if running
if agent_id in self.agent_processes:
process = self.agent_processes[agent_id]
if not request.force:
# Graceful shutdown
process.terminate()
try:
process.wait(timeout=10)
except psutil.TimeoutExpired:
process.kill()
else:
# Force kill
process.kill()
del self.agent_processes[agent_id]
# Remove from registry
del self.agents[agent_id]
# Log deletion
audit_logger = get_audit_logger()
await audit_logger.log_event(
level=AuditLevel.INFO,
category=AuditCategory.AGENT_LIFECYCLE,
operation="delete_agent",
resource_type="agent",
resource_id=str(agent_id),
success=True,
user_id=get_security_context().user_id,
metadata={
"agent_name": agent_state.name,
"force_delete": request.force,
"preserve_logs": request.preserve_logs,
},
)
return True
except Exception as e:
# Log failure
audit_logger = get_audit_logger()
await audit_logger.log_event(
level=AuditLevel.ERROR,
category=AuditCategory.AGENT_LIFECYCLE,
operation="delete_agent",
resource_type="agent",
resource_id=str(agent_id),
success=False,
error_message=str(e),
user_id=get_security_context().user_id,
)
raise AgentLifecycleError(f"Failed to delete agent: {e}")
@message_sending_contract
@require(lambda request: len(request.content.strip()) > 0)
@require(lambda request: len(request.content) <= 100000)
@ensure(lambda result: result is not None)
async def send_message(self, request: MessageRequest) -> Dict[str, Any]:
"""
Send message to agent with validation and audit logging.
Contracts:
Preconditions:
- Message content must not be empty
- Message must be within size limits
- Agent must exist and be active
- Security context must allow messaging
Postconditions:
- Message delivered to agent
- ADDER+ prompt prepended if requested
- Communication event audited
- Response tracked if requested
Invariants:
- Message content properly sanitized
- Agent conversation history updated
- No unauthorized information leakage
"""
agent_id = request.agent_id
# Check if agent exists
if agent_id not in self.agents:
raise AgentCommunicationError(f"Agent not found: {agent_id}")
agent_state = self.agents[agent_id]
# Validate agent is in communicable state
if agent_state.status not in {AgentStatus.ACTIVE, AgentStatus.IDLE}:
raise AgentCommunicationError(
f"Agent not ready for communication: {agent_state.status}"
)
# Validate and sanitize message content
content_validation = validate_message_content(
request.content, context="general"
)
if not content_validation.is_valid:
raise AgentCommunicationError(
f"Invalid message content: {content_validation.violations}"
)
sanitized_content = content_validation.sanitized_input or request.content
try:
# Prepare message with ADDER+ prompt if requested
final_message = sanitized_content
if request.prepend_adder:
final_message = f"You are {agent_state.name}\n\n{sanitized_content}"
# Update agent state to working
updated_state = agent_state.with_status(AgentStatus.WORKING)
self.agents[agent_id] = updated_state
# Simulate message delivery (actual implementation would use iTerm2 API)
delivery_result = {
"agent_id": str(agent_id),
"message_length": len(final_message),
"delivered_at": datetime.utcnow().isoformat(),
"response_expected": request.wait_for_response,
}
# Log communication
audit_logger = get_audit_logger()
await audit_logger.log_event(
level=AuditLevel.INFO,
category=AuditCategory.COMMUNICATION,
operation="send_message",
resource_type="agent",
resource_id=str(agent_id),
success=True,
user_id=get_security_context().user_id,
agent_id=str(agent_id),
metadata={
"message_length": len(sanitized_content),
"prepend_adder": request.prepend_adder,
"wait_for_response": request.wait_for_response,
},
)
return delivery_result
except Exception as e:
# Log failure
audit_logger = get_audit_logger()
await audit_logger.log_event(
level=AuditLevel.ERROR,
category=AuditCategory.COMMUNICATION,
operation="send_message",
resource_type="agent",
resource_id=str(agent_id),
success=False,
error_message=str(e),
user_id=get_security_context().user_id,
)
raise AgentCommunicationError(f"Failed to send message: {e}")
@status_read_contract
@require(lambda agent_id: agent_id is not None)
@ensure(lambda result: result is not None)
async def get_agent_status(self, agent_id: AgentId) -> Dict[str, Any]:
"""
Get comprehensive agent status with health validation.
Contracts:
Preconditions:
- Agent ID must be valid
- Security context must allow status reading
Postconditions:
- Returns complete agent status
- Health validation performed
- Status read event audited
Invariants:
- Status data is read-only
- No sensitive information exposed
- Resource metrics are current
"""
if agent_id not in self.agents:
raise AgentLifecycleError(f"Agent not found: {agent_id}")
agent_state = self.agents[agent_id]
# Validate agent health
health_issues = self.state_validator.validate_agent_health(agent_state)
# Get process information if available
process_info = None
if agent_id in self.agent_processes:
try:
process = self.agent_processes[agent_id]
process_info = {
"pid": process.pid,
"cpu_percent": process.cpu_percent(),
"memory_mb": process.memory_info().rss / 1024 / 1024,
"status": process.status(),
"create_time": datetime.fromtimestamp(
process.create_time()
).isoformat(),
}
except (psutil.NoSuchProcess, psutil.AccessDenied):
process_info = {"error": "Process not accessible"}
status_data = {
"agent_id": str(agent_id),
"name": agent_state.name,
"status": agent_state.status.value,
"specialization": agent_state.specialization.value,
"session_id": str(agent_state.session_id),
"created_at": agent_state.created_at.isoformat(),
"last_heartbeat": agent_state.last_heartbeat.isoformat(),
"is_healthy": len(health_issues) == 0,
"health_issues": health_issues,
"error_count": agent_state.error_count,
"restart_count": agent_state.restart_count,
"resource_metrics": {
"cpu_percent": agent_state.resource_metrics.cpu_percent,
"memory_mb": agent_state.resource_metrics.memory_mb,
"file_descriptors": agent_state.resource_metrics.file_descriptors,
"uptime_seconds": agent_state.resource_metrics.uptime_seconds,
"last_activity": agent_state.resource_metrics.last_activity.isoformat(),
},
"process_info": process_info,
}
return status_data
@status_read_contract
@ensure(lambda result: isinstance(result, dict))
async def get_all_agents_status(
self, session_id: Optional[SessionId] = None
) -> Dict[str, Any]:
"""
Get status of all agents, optionally filtered by session.
Contracts:
Preconditions:
- Security context must allow status reading
Postconditions:
- Returns comprehensive status for all agents
- Optional session filtering applied
- Status read event audited
Invariants:
- Only authorized agents included
- Status data is consistent
- Resource summaries are accurate
"""
agents_to_include = self.agents.values()
# Filter by session if specified
if session_id:
agents_to_include = [
agent for agent in agents_to_include if agent.session_id == session_id
]
# Get status for each agent
agent_statuses = []
total_cpu = 0.0
total_memory = 0
healthy_count = 0
for agent_state in agents_to_include:
status = await self.get_agent_status(agent_state.agent_id)
agent_statuses.append(status)
total_cpu += agent_state.resource_metrics.cpu_percent
total_memory += agent_state.resource_metrics.memory_mb
if status["is_healthy"]:
healthy_count += 1
return {
"total_agents": len(agent_statuses),
"healthy_agents": healthy_count,
"unhealthy_agents": len(agent_statuses) - healthy_count,
"total_cpu_percent": total_cpu,
"total_memory_mb": total_memory,
"session_id": str(session_id) if session_id else None,
"agents": agent_statuses,
"timestamp": datetime.utcnow().isoformat(),
}
# Global lifecycle manager instance
_lifecycle_manager_instance: Optional[AgentLifecycleManager] = None
def get_lifecycle_manager() -> AgentLifecycleManager:
"""Get global agent lifecycle manager instance."""
global _lifecycle_manager_instance
if _lifecycle_manager_instance is None:
_lifecycle_manager_instance = AgentLifecycleManager()
return _lifecycle_manager_instance
# Convenience functions for agent lifecycle operations
async def create_agent_with_contracts(
session_id: SessionId,
agent_name: str,
specialization: AgentSpecialization = AgentSpecialization.GENERAL,
system_prompt_suffix: str = "",
claude_config: Optional[ClaudeConfig] = None,
) -> AgentState:
"""Create agent with comprehensive contract validation."""
manager = get_lifecycle_manager()
request = AgentCreationRequest(
session_id=session_id,
agent_name=agent_name,
specialization=specialization,
system_prompt_suffix=system_prompt_suffix,
claude_config=claude_config,
)
return await manager.create_agent(request)
async def delete_agent_with_contracts(
agent_id: AgentId, force: bool = False, preserve_logs: bool = True
) -> bool:
"""Delete agent with comprehensive contract validation."""
manager = get_lifecycle_manager()
request = AgentDeletionRequest(
agent_id=agent_id, force=force, preserve_logs=preserve_logs
)
return await manager.delete_agent(request)
async def send_message_with_contracts(
agent_id: AgentId,
content: str,
prepend_adder: bool = True,
wait_for_response: bool = False,
) -> Dict[str, Any]:
"""Send message to agent with comprehensive contract validation."""
manager = get_lifecycle_manager()
request = MessageRequest(
agent_id=agent_id,
content=content,
prepend_adder=prepend_adder,
wait_for_response=wait_for_response,
)
return await manager.send_message(request)
