Spaces:
Sleeping
Sleeping
| """ | |
| Agent Request Manager | |
| Handles agent requests for user input, decisions, and permissions | |
| with full governance tracking and audit trail. | |
| Features: | |
| - Permission requests | |
| - Decision requests | |
| - Input requests | |
| - Confirmation requests | |
| - Request expiration and revocation | |
| """ | |
| import asyncio | |
| from datetime import datetime, timedelta | |
| import logging | |
| from typing import Any, Dict, List, Optional | |
| import uuid | |
| from sqlalchemy.orm import Session | |
| from core.agent_governance_service import AgentGovernanceService | |
| from core.models import AgentRegistry, AgentRequestLog, CanvasAudit | |
| from core.websockets import manager as ws_manager | |
| logger = logging.getLogger(__name__) | |
| # Feature flags | |
| import os | |
| AGENT_REQUESTS_ENABLED = os.getenv("AGENT_REQUESTS_ENABLED", "true").lower() == "true" | |
| class AgentRequestManager: | |
| """ | |
| Manages agent requests for user input/decisions. | |
| Provides a structured way for agents to request permissions, decisions, | |
| or input from users with full audit trail and governance. | |
| """ | |
| # Request timeouts by urgency | |
| REQUEST_TIMEOUTS = { | |
| "low": 3600, # 1 hour | |
| "medium": 600, # 10 minutes | |
| "high": 60, # 1 minute | |
| "blocking": 30 # 30 seconds | |
| } | |
| def __init__(self, db: Session): | |
| self.db = db | |
| self.governance = AgentGovernanceService(db) | |
| self._pending_requests: Dict[str, asyncio.Event] = {} | |
| async def create_permission_request( | |
| self, | |
| user_id: str, | |
| agent_id: str, | |
| title: str, | |
| permission: str, | |
| context: Dict[str, Any], | |
| urgency: str = "medium", | |
| expires_in: Optional[int] = None | |
| ) -> str: | |
| """ | |
| Create a permission request from agent to user. | |
| Args: | |
| user_id: User ID | |
| agent_id: Agent ID requesting permission | |
| title: Request title | |
| permission: Permission being requested | |
| context: Context dict with operation, impact, alternatives | |
| urgency: Urgency level (low, medium, high, blocking) | |
| expires_in: Optional custom expiration time in seconds | |
| Returns: | |
| request_id: Unique request ID | |
| """ | |
| if not AGENT_REQUESTS_ENABLED: | |
| return str(uuid.uuid4()) | |
| try: | |
| # Generate request ID | |
| request_id = str(uuid.uuid4()) | |
| # Get agent | |
| agent = self.db.query(AgentRegistry).filter( | |
| AgentRegistry.id == agent_id | |
| ).first() | |
| agent_name = agent.name if agent else "Agent" | |
| # Calculate expiration | |
| timeout = expires_in or self.REQUEST_TIMEOUTS.get(urgency, 600) | |
| expires_at = datetime.utcnow() + timedelta(seconds=timeout) | |
| # Create options | |
| options = [ | |
| { | |
| "label": "Approve", | |
| "description": f"Allow {agent_name} to use this permission", | |
| "consequences": f"{agent_name} will be able to perform this action now and in the future", | |
| "action": "approve" | |
| }, | |
| { | |
| "label": "Approve Once", | |
| "description": "Allow this single action", | |
| "consequences": f"{agent_name} will perform this action once, then ask again", | |
| "action": "approve_once" | |
| }, | |
| { | |
| "label": "Deny", | |
| "description": "Don't allow this action", | |
| "consequences": f"{agent_name} will not be able to perform this action", | |
| "action": "deny" | |
| } | |
| ] | |
| # Create request data | |
| request_data = { | |
| "request_id": request_id, | |
| "agent_id": agent_id, | |
| "agent_name": agent_name, | |
| "request_type": "permission", | |
| "urgency": urgency, | |
| "title": title, | |
| "explanation": f"I need permission to: {permission}", | |
| "permission": permission, | |
| "context": context, | |
| "options": options, | |
| "suggested_option": 1, # Approve Once - safe default | |
| "governance": { | |
| "requires_signature": urgency == "blocking", | |
| "audit_log_required": True, | |
| "revocable": True | |
| } | |
| } | |
| # Create log entry | |
| request_log = AgentRequestLog( | |
| id=str(uuid.uuid4()), | |
| agent_id=agent_id, | |
| user_id=user_id, | |
| request_id=request_id, | |
| request_type="permission", | |
| request_data=request_data, | |
| expires_at=expires_at | |
| ) | |
| self.db.add(request_log) | |
| self.db.commit() | |
| # Create event for response | |
| self._pending_requests[request_id] = asyncio.Event() | |
| # Broadcast request | |
| await ws_manager.broadcast( | |
| f"user:{user_id}", | |
| { | |
| "type": "agent:request", | |
| "data": request_data | |
| } | |
| ) | |
| # Create audit | |
| await self._create_audit( | |
| agent_id=agent_id, | |
| user_id=user_id, | |
| request_id=request_id, | |
| action="create_permission_request" | |
| ) | |
| logger.info( | |
| f"Created permission request {request_id} from agent {agent_id}, " | |
| f"user {user_id}" | |
| ) | |
| return request_id | |
| except Exception as e: | |
| logger.error(f"Failed to create permission request: {e}") | |
| return str(uuid.uuid4()) | |
| async def create_decision_request( | |
| self, | |
| user_id: str, | |
| agent_id: str, | |
| title: str, | |
| explanation: str, | |
| options: List[Dict[str, Any]], | |
| context: Dict[str, Any], | |
| urgency: str = "low", | |
| suggested_option: int = 0, | |
| expires_in: Optional[int] = None | |
| ) -> str: | |
| """ | |
| Create a decision request from agent to user. | |
| Args: | |
| user_id: User ID | |
| agent_id: Agent ID requesting decision | |
| title: Request title | |
| explanation: Why agent needs this decision | |
| options: List of decision options | |
| context: Context dict | |
| urgency: Urgency level | |
| suggested_option: Index of suggested option | |
| expires_in: Optional expiration time | |
| Returns: | |
| request_id: Unique request ID | |
| """ | |
| if not AGENT_REQUESTS_ENABLED: | |
| return str(uuid.uuid4()) | |
| try: | |
| request_id = str(uuid.uuid4()) | |
| # Get agent | |
| agent = self.db.query(AgentRegistry).filter( | |
| AgentRegistry.id == agent_id | |
| ).first() | |
| agent_name = agent.name if agent else "Agent" | |
| # Calculate expiration | |
| timeout = expires_in or self.REQUEST_TIMEOUTS.get(urgency, 3600) | |
| expires_at = datetime.utcnow() + timedelta(seconds=timeout) | |
| # Create request data | |
| request_data = { | |
| "request_id": request_id, | |
| "agent_id": agent_id, | |
| "agent_name": agent_name, | |
| "request_type": "decision", | |
| "urgency": urgency, | |
| "title": title, | |
| "explanation": explanation, | |
| "context": context, | |
| "options": options, | |
| "suggested_option": suggested_option, | |
| "governance": { | |
| "requires_signature": False, | |
| "audit_log_required": True, | |
| "revocable": True | |
| } | |
| } | |
| # Create log entry | |
| request_log = AgentRequestLog( | |
| id=str(uuid.uuid4()), | |
| agent_id=agent_id, | |
| user_id=user_id, | |
| request_id=request_id, | |
| request_type="decision", | |
| request_data=request_data, | |
| expires_at=expires_at | |
| ) | |
| self.db.add(request_log) | |
| self.db.commit() | |
| # Create event for response | |
| self._pending_requests[request_id] = asyncio.Event() | |
| # Broadcast request | |
| await ws_manager.broadcast( | |
| f"user:{user_id}", | |
| { | |
| "type": "agent:request", | |
| "data": request_data | |
| } | |
| ) | |
| # Create audit | |
| await self._create_audit( | |
| agent_id=agent_id, | |
| user_id=user_id, | |
| request_id=request_id, | |
| action="create_decision_request" | |
| ) | |
| logger.info( | |
| f"Created decision request {request_id} from agent {agent_id}" | |
| ) | |
| return request_id | |
| except Exception as e: | |
| logger.error(f"Failed to create decision request: {e}") | |
| return str(uuid.uuid4()) | |
| async def wait_for_response( | |
| self, | |
| request_id: str, | |
| timeout: Optional[int] = None | |
| ) -> Optional[Dict[str, Any]]: | |
| """ | |
| Wait for user response to a request. | |
| Args: | |
| request_id: Request ID to wait for | |
| timeout: Optional timeout in seconds (default: use request urgency) | |
| Returns: | |
| User response dict or None if timeout | |
| """ | |
| if request_id not in self._pending_requests: | |
| logger.warning(f"Request {request_id} not found") | |
| return None | |
| try: | |
| # Get timeout from request if not provided | |
| if timeout is None: | |
| request_log = self.db.query(AgentRequestLog).filter( | |
| AgentRequestLog.request_id == request_id | |
| ).first() | |
| if request_log and request_log.expires_at: | |
| timeout = int((request_log.expires_at - datetime.utcnow()).total_seconds()) | |
| else: | |
| timeout = 600 # Default 10 minutes | |
| # Wait for response | |
| event = self._pending_requests[request_id] | |
| try: | |
| await asyncio.wait_for(event.wait(), timeout=timeout) | |
| except asyncio.TimeoutError: | |
| logger.warning(f"Request {request_id} timed out") | |
| # Mark as expired | |
| request_log = self.db.query(AgentRequestLog).filter( | |
| AgentRequestLog.request_id == request_id | |
| ).first() | |
| if request_log: | |
| request_log.revoked = True | |
| self.db.commit() | |
| return None | |
| # Get response | |
| request_log = self.db.query(AgentRequestLog).filter( | |
| AgentRequestLog.request_id == request_id | |
| ).first() | |
| if request_log: | |
| return request_log.user_response | |
| return None | |
| except Exception as e: | |
| logger.error(f"Failed to wait for response: {e}") | |
| return None | |
| finally: | |
| # Clean up | |
| self._pending_requests.pop(request_id, None) | |
| async def handle_response( | |
| self, | |
| user_id: str, | |
| request_id: str, | |
| response: Dict[str, Any] | |
| ): | |
| """ | |
| Handle user response to a request. | |
| Args: | |
| user_id: User ID | |
| request_id: Request ID | |
| response: User's response | |
| """ | |
| if not AGENT_REQUESTS_ENABLED: | |
| return | |
| try: | |
| # Get request log | |
| request_log = self.db.query(AgentRequestLog).filter( | |
| AgentRequestLog.request_id == request_id, | |
| AgentRequestLog.user_id == user_id | |
| ).first() | |
| if not request_log: | |
| logger.warning(f"Request {request_id} not found for user {user_id}") | |
| return | |
| # Check if expired | |
| if request_log.expires_at and datetime.utcnow() > request_log.expires_at: | |
| logger.warning(f"Request {request_id} has expired") | |
| return | |
| # Update log | |
| request_log.user_response = response | |
| request_log.responded_at = datetime.utcnow() | |
| request_log.response_time_seconds = ( | |
| datetime.utcnow() - request_log.created_at | |
| ).total_seconds() | |
| self.db.commit() | |
| # Trigger event | |
| if request_id in self._pending_requests: | |
| self._pending_requests[request_id].set() | |
| # Create audit | |
| await self._create_audit( | |
| agent_id=request_log.agent_id, | |
| user_id=user_id, | |
| request_id=request_id, | |
| action="handle_response", | |
| metadata={"response": response} | |
| ) | |
| logger.info(f"Handled response for request {request_id}") | |
| except Exception as e: | |
| logger.error(f"Failed to handle response: {e}") | |
| async def revoke_request( | |
| self, | |
| request_id: str | |
| ): | |
| """ | |
| Revoke a pending request. | |
| Args: | |
| request_id: Request ID to revoke | |
| """ | |
| try: | |
| # Get request log | |
| request_log = self.db.query(AgentRequestLog).filter( | |
| AgentRequestLog.request_id == request_id | |
| ).first() | |
| if request_log: | |
| request_log.revoked = True | |
| self.db.commit() | |
| # Trigger event with None response | |
| if request_id in self._pending_requests: | |
| self._pending_requests[request_id].set() | |
| logger.info(f"Revoked request {request_id}") | |
| except Exception as e: | |
| logger.error(f"Failed to revoke request: {e}") | |
| async def _create_audit( | |
| self, | |
| agent_id: str, | |
| user_id: str, | |
| request_id: str, | |
| action: str, | |
| metadata: Optional[Dict[str, Any]] = None | |
| ): | |
| """Create canvas audit entry.""" | |
| try: | |
| audit = CanvasAudit( | |
| id=str(uuid.uuid4()), | |
| workspace_id="default", | |
| agent_id=agent_id, | |
| agent_execution_id=None, | |
| user_id=user_id, | |
| canvas_id=None, | |
| session_id=None, | |
| component_type="agent_request_prompt", | |
| component_name="agent_request_manager", | |
| action=action, | |
| audit_metadata={ | |
| "request_id": request_id, | |
| **(metadata or {}) | |
| }, | |
| governance_check_passed=True | |
| ) | |
| self.db.add(audit) | |
| self.db.commit() | |
| except Exception as e: | |
| logger.error(f"Failed to create audit: {e}") | |
| # Singleton instance helper | |
| def get_agent_request_manager(db: Session) -> AgentRequestManager: | |
| """Get or create agent request manager instance.""" | |
| return AgentRequestManager(db) | |