"""In-memory AgentRegistry and TaskDispatcher implementations. These are the default concrete implementations for Phase 4 modular contracts. They are intentionally simple and testable; production deployments may substitute database-backed or distributed variants. """ from __future__ import annotations import asyncio import logging from collections import defaultdict from typing import Any from cyberbrain26.extension_layer.agent_registry import Agent, AgentRegistry from cyberbrain26.security_plane.actions import ( DEFAULT_EXECUTION_TIMEOUT_SECONDS, ExecutionAuthorization, bound_execution_authorization, execution_payload_digest, normalize_action, ) from cyberbrain26.security_plane.approval import ApprovalManager, ApprovalStatus from cyberbrain26.security_plane.audit import AuditEventType, AuditLog from cyberbrain26.security_plane.permission import PermissionGateway, SecurityError from cyberbrain26.security_plane.policy import PolicyEngine logger = logging.getLogger(__name__) class InMemoryAgentRegistry: """Default AgentRegistry backed by an in-memory dict. Supports register, unregister, lookup by ID, listing by role, and existence checks. Thread-safe via simple dict operations. """ def __init__(self) -> None: self._agents: dict[str, Agent] = {} self._roles: dict[str, list[str]] = {} def register(self, agent_id: str, agent: Agent) -> None: if agent_id in self._agents: logger.warning( "Re-registering agent '%s'; overwriting previous registration", agent_id, ) self._agents[agent_id] = agent role = getattr(agent, "role", None) if role: if role not in self._roles: self._roles[role] = [] if agent_id not in self._roles[role]: self._roles[role].append(agent_id) def unregister(self, agent_id: str) -> bool: if agent_id not in self._agents: return False del self._agents[agent_id] for role_list in self._roles.values(): if agent_id in role_list: role_list.remove(agent_id) return True def get(self, agent_id: str) -> Agent | None: return self._agents.get(agent_id) def list_by_role(self, role: str) -> list[Agent]: return [self._agents[aid] for aid in self._roles.get(role, []) if aid in self._agents] def list_all(self) -> list[tuple[str, Agent]]: return list(self._agents.items()) def exists(self, agent_id: str) -> bool: return agent_id in self._agents class InMemoryTaskDispatcher: """Default TaskDispatcher backed by an AgentRegistry. Routes tasks to agents by ID or by role. Requires agents to be registered in the registry before dispatching. All dispatches pass through Identity → Policy → Permission → Audit before agent.execute() is called. This dispatcher owns the *outer* authorization gate, and an approval it obtains is the one human decision for the step. When the authorized action is one the agent will really perform, the decision is handed to the agent as an ``ExecutionAuthorization`` bound to that single call, so the inner gate (the execution gateway standing in front of the sandbox) re-verifies it instead of asking the same operator the same question again. The grant is minted only when a human actually decided, and it is bound to a digest of the execution payload, so a step whose payload changes afterwards is refused at the inner gate rather than executed under the earlier approval. """ def __init__( self, registry: AgentRegistry | None = None, policy_engine: PolicyEngine | None = None, audit_log: AuditLog | None = None, approval_manager: ApprovalManager | None = None, approval_timeout_seconds: float = 30.0, approval_poll_interval: float = 0.05, ) -> None: if approval_timeout_seconds <= 0: raise ValueError("approval_timeout_seconds must be positive") if approval_poll_interval <= 0: raise ValueError("approval_poll_interval must be positive") self._registry = registry or InMemoryAgentRegistry() self._policy_engine = policy_engine or PolicyEngine() self._audit_log_instance = audit_log self._policy_engine.audit_log = self._audit_log_instance self._permission_gateway = PermissionGateway( self._policy_engine, audit_log=self._audit_log_instance ) self._approval_manager = approval_manager self._approval_timeout_seconds = approval_timeout_seconds self._approval_poll_interval = approval_poll_interval self._role_indices: dict[str, int] = defaultdict(int) def register_agent(self, agent_id: str, agent: Agent) -> None: self._registry.register(agent_id, agent) def _select_agent_by_role(self, role: str) -> Agent: """Select an agent by role using deterministic round-robin.""" agents = self._registry.list_by_role(role) if not agents: raise KeyError(f"no agent with role: {role}") idx = self._role_indices[role] self._role_indices[role] = (idx + 1) % len(agents) return agents[idx] def _record_audit_event( self, task_id: str, agent_id: str, action: str, decision: str, reason: str, *, event_type: AuditEventType = AuditEventType.PERMISSION_CHECK, ) -> None: """Record a dispatch-level security event with task correlation.""" if self._audit_log_instance is not None: self._audit_log_instance.record( event_type, actor=agent_id, agent_id=agent_id, task_id=task_id, action=action, policy_decision=decision, result=decision, error=reason if decision != "allow" else "", ) async def _authorize( self, task_id: str, registered_agent_id: str, agent: Agent, action: str, resource: str | None = None, ) -> None: """Enforce identity, canonical policy, permission, approval, and audit. Returns nothing on purpose: the structural port the agents module uses for bus-delivered work declares this signature, and a bus task that clears the chain is not followed by a dispatched step whose payload a grant could describe. Callers that do own a step use :meth:`_authorize_step`, which also reports the approval that was obtained. """ await self._authorize_step(task_id, registered_agent_id, agent, action, resource) async def _authorize_step( self, task_id: str, registered_agent_id: str, agent: Agent, action: str, resource: str | None = None, ) -> str | None: """Enforce the authorization chain and report the approval behind it. Returns: The id of the approval request a human approved, or ``None`` when the chain allowed the action outright. ``None`` deliberately means "no human decision exists here", so a caller can never turn a policy "allow" into an approval that authorises executable content. """ identity = getattr(agent, "identity", None) if identity is None: reason = "agent identity is required" self._record_audit_event(task_id, registered_agent_id, action, "deny", reason) raise SecurityError(reason) if identity.agent_id != registered_agent_id: reason = ( f"agent identity {identity.agent_id!r} does not match " f"registered id {registered_agent_id!r}" ) self._record_audit_event(task_id, registered_agent_id, action, "deny", reason) raise SecurityError(reason) # Always use the dispatcher's canonical policy. An agent-supplied gateway # must not be able to weaken the system policy. result = self._permission_gateway.evaluate(identity, action, resource) self._record_audit_event( task_id=task_id, agent_id=identity.agent_id, action=action, decision=result.decision.value, reason=result.reason, ) if result.decision.value == "allow": # Allowed outright: no human decided anything, so there is no # approval to hand down. Stated explicitly rather than as a bare # return, because "no approval" is now a meaningful value here. return None if result.decision.value != "require_approval": raise SecurityError(f"Agent {identity.agent_id} denied for action '{action}'") if self._approval_manager is None: reason = "approval manager is not configured" self._record_audit_event( task_id, identity.agent_id, action, "deny", reason, event_type=AuditEventType.APPROVAL_DENIED, ) raise SecurityError(f"Approval unavailable for action '{action}'") request = self._approval_manager.create_request( identity, action, resource, context={"task_id": task_id}, reason=result.reason, ) self._record_audit_event( task_id, identity.agent_id, action, "pending", request.request_id, event_type=AuditEventType.APPROVAL_REQUIRED, ) wait_seconds = min( self._approval_timeout_seconds, self._approval_manager.timeout_seconds, ) deadline = asyncio.get_running_loop().time() + wait_seconds while asyncio.get_running_loop().time() < deadline: current = self._approval_manager.get_request(request.request_id) if current is None: break if current.status == ApprovalStatus.APPROVED: self._record_audit_event( task_id, identity.agent_id, action, "approved", current.reviewer or "approval granted", event_type=AuditEventType.APPROVAL_GRANTED, ) return request.request_id if current.status != ApprovalStatus.PENDING: break await asyncio.sleep(min(self._approval_poll_interval, wait_seconds)) self._approval_manager.check_expired() current = self._approval_manager.get_request(request.request_id) if current is not None and current.status == ApprovalStatus.PENDING: self._approval_manager.cancel(request.request_id) status = current.status.value if current is not None else "missing" self._record_audit_event( task_id, identity.agent_id, action, "deny", f"approval {status}", event_type=AuditEventType.APPROVAL_DENIED, ) raise SecurityError(f"Approval {status} for action '{action}'") def _require_declared_action( self, task_id: str, agent_id: str, payload: dict[str, Any], ) -> str: """Return the canonical action declared by the payload. Security rationale: the dispatcher used to invent ``"execute"`` when the payload carried no action. That name is in neither the scope map nor the sensitive-action set, so the step skipped the scope gate and was allowed by SAFE mode's non-sensitive rule — an action nobody had ever authorized. No default is inferred here: a payload without an explicit action states nothing to authorize and is refused, exactly as the CLI already refuses an action-less roadmap step. """ action = normalize_action(payload.get("action")) if not action: reason = ( "payload declares no 'action', so there is nothing to authorize; " "it was refused and no default action is inferred" ) self._record_audit_event(task_id, agent_id, "", "deny", reason) raise SecurityError(reason) return action async def dispatch(self, task_id: str, agent_id: str, payload: dict[str, Any]) -> Any: agent = self._registry.get(agent_id) if agent is None: raise KeyError(f"agent not found: {agent_id}") action = self._require_declared_action(task_id, agent_id, payload) resource = payload.get("resource") approval_id = await self._authorize_step(task_id, agent_id, agent, action, resource) return await self._execute_authorized( task_id, agent_id, agent, action, payload, approval_id ) async def dispatch_to_role(self, task_id: str, role: str, payload: dict[str, Any]) -> Any: agent = self._select_agent_by_role(role) agent_id = next( (aid for aid, a in self._registry.list_all() if a is agent), "unknown", ) action = self._require_declared_action(task_id, agent_id, payload) resource = payload.get("resource") approval_id = await self._authorize_step(task_id, agent_id, agent, action, resource) return await self._execute_authorized( task_id, agent_id, agent, action, payload, approval_id ) def _execution_grant( self, task_id: str, agent_id: str, action: str, payload: dict[str, Any], approval_id: str | None, ) -> ExecutionAuthorization | None: """Return the grant to bind to this agent call, or ``None`` for none. The digest is taken over the payload exactly as the agent will read it: the same keys, and the same default timeout, that ``CoderAgent`` passes to the execution gateway. The two sides therefore agree on a well-behaved step and disagree on a tampered one, which is the property the inner gate relies on. """ if approval_id is None: # Policy allowed the action with no human involved. That is not an # approval and must never be passed off as one, so the inner gate # keeps asking for the real execution action on its own. return None return ExecutionAuthorization( agent_id=agent_id, task_id=task_id, action=action, resource=payload.get("resource"), payload_digest=execution_payload_digest( code=payload.get("code"), language=payload.get("language"), workspace=payload.get("workspace"), timeout_seconds=payload.get( "timeout_seconds", DEFAULT_EXECUTION_TIMEOUT_SECONDS ), ), approval_request_id=approval_id, ) async def _execute_authorized( self, task_id: str, agent_id: str, agent: Agent, action: str, payload: dict[str, Any], approval_id: str | None = None, ) -> Any: self._record_audit_event( task_id, agent_id, action, "started", "agent execution started", event_type=AuditEventType.AGENT_STARTED, ) try: # The grant is bound here, not stored anywhere: the agent's own gates # read it during this call only, and it is gone before the next # dispatch, so no later step can spend this approval. with bound_execution_authorization( self._execution_grant(task_id, agent_id, action, payload, approval_id) ): result = await agent.execute(payload) except Exception as exc: self._record_audit_event( task_id, agent_id, action, "failed", type(exc).__name__, event_type=AuditEventType.AGENT_FINISHED, ) raise succeeded = bool(getattr(result, "success", True)) self._record_audit_event( task_id, agent_id, action, "succeeded" if succeeded else "failed", "agent execution completed" if succeeded else "agent returned failure", event_type=AuditEventType.AGENT_FINISHED, ) return result __all__ = ["InMemoryAgentRegistry", "InMemoryTaskDispatcher"]