"""AgentGuard firewall — public entry point."""
from __future__ import annotations
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
from agentguard.capability.manifest import CapabilityManifest
from agentguard.firewall_ops import FirewallEngine
from agentguard.firewall_types import (
ACTION_BLOCK,
ACTION_FORWARD,
ACTION_QUARANTINE,
MAX_MESSAGE_BYTES,
FirewallDecision,
)
F = TypeVar("F", bound=Callable[..., Any])
AF = TypeVar("AF", bound=Callable[..., Awaitable[Any]])
__all__ = [
"ACTION_BLOCK",
"ACTION_FORWARD",
"ACTION_QUARANTINE",
"AgentGuard",
"FirewallDecision",
"MAX_MESSAGE_BYTES",
]
[docs]
class AgentGuard:
"""Inter-agent security middleware entry point."""
def __init__(
self,
risk_threshold: float = 0.85,
enable_trust_attestation: bool = True,
enable_capability_enforcement: bool = True,
audit_log_path: str | None = "./audit.jsonl",
mode: str = "enforce",
task_objective: str | None = None,
enable_consistency_check: bool = True,
consistency_threshold: float = 0.10,
consistency_ml_risk_floor: float = 0.75,
model_path: str | None = None,
require_ml_model: bool = False,
enable_otel_export: bool = False,
include_message_preview: bool = True,
max_message_bytes: int = MAX_MESSAGE_BYTES,
) -> None:
if mode not in ("enforce", "monitor"):
raise ValueError(f"mode must be 'enforce' or 'monitor', got {mode!r}")
if max_message_bytes < 1:
raise ValueError("max_message_bytes must be >= 1")
eng = FirewallEngine(
risk_threshold=risk_threshold,
enable_trust_attestation=enable_trust_attestation,
enable_capability_enforcement=enable_capability_enforcement,
enable_consistency_check=enable_consistency_check,
consistency_ml_risk_floor=consistency_ml_risk_floor,
consistency_threshold=consistency_threshold,
mode=mode,
include_message_preview=include_message_preview,
max_message_bytes=max_message_bytes,
audit_log_path=audit_log_path,
task_objective=task_objective,
model_path=model_path,
require_ml_model=require_ml_model,
enable_otel_export=enable_otel_export,
)
self._engine = eng
self.risk_threshold = eng.risk_threshold
self.enable_trust_attestation = eng.enable_trust_attestation
self.enable_capability_enforcement = eng.enable_capability_enforcement
self.enable_consistency_check = eng.enable_consistency_check
self.consistency_ml_risk_floor = eng.consistency_ml_risk_floor
self.mode = eng.mode
self.include_message_preview = eng.include_message_preview
self.max_message_bytes = eng.max_message_bytes
self._logger, self._trust = eng.logger, eng.trust
self._capabilities, self._rules = eng.capabilities, eng.rules
self._ml, self._consistency = eng.ml, eng.consistency
self._mcp_inspector, self._otel = eng.mcp_inspector, eng.otel
@property
def is_ml_model_loaded(self) -> bool:
"""Return True when the DeBERTa ONNX scorer is active."""
return self._ml.is_model_loaded()
[docs]
def register_agent(self, agent_id: str, manifest: CapabilityManifest) -> None:
"""Register agent identity, trust keys, and capability manifest."""
self._engine.register_agent(agent_id, manifest)
[docs]
def register_delegated_agent(
self, parent_agent_id: str, child_manifest: CapabilityManifest
) -> CapabilityManifest:
"""Register a sub-agent with monotonic capability attenuation."""
return self._engine.register_delegated_agent(parent_agent_id, child_manifest)
[docs]
def rotate_keys(self) -> None:
"""Rotate ephemeral trust keys for all registered agents."""
self._trust.rotate_all_keys()
[docs]
def sign_payload(
self,
sender_id: str,
payload_bytes: bytes,
*,
recipient_id: str,
nonce: bytes | None = None,
) -> bytes:
"""Sign an inter-agent envelope (sender, recipient, nonce, payload).
The returned wire signature binds the recipient and a unique nonce so
a previously signed payload cannot be redirected or silently replayed.
Pass the bytes unchanged to :meth:`inspect_message` as ``signature``.
"""
return self._trust.sign_envelope(
sender_id,
recipient_id,
payload_bytes,
nonce=nonce,
)
[docs]
def inspect_message(
self,
sender_id: str,
recipient_id: str,
message: str,
payload_bytes: bytes,
*,
signature: bytes | None = None,
) -> FirewallDecision:
"""Run inspection, trust, and size checks on a message."""
return self._engine.inspect_message(
sender_id, recipient_id, message, payload_bytes, signature=signature
)
[docs]
def inspect_content(
self,
sender_id: str,
recipient_id: str,
message: str,
) -> FirewallDecision:
"""Inspect unsigned content (rules + ML + consistency, no trust check).
Use this at boundaries where messages cannot carry Ed25519
signatures — user input, framework middleware hooks, or third-party
content. The decision is audited like any other inspection.
"""
return self._engine.inspect_content(sender_id, recipient_id, message)
[docs]
def get_manifest(self, agent_id: str) -> CapabilityManifest | None:
"""Return the registered capability manifest for an agent, if any."""
return self._capabilities.get_manifest(agent_id)
[docs]
def check_data_source(self, agent_id: str, data_source: str) -> bool:
"""Validate data-source access; False when blocked in enforce mode."""
return self._engine.check_data_source(agent_id, data_source)
[docs]
def check_output_tokens(self, agent_id: str, token_count: int) -> bool:
"""Validate output size; False when blocked in enforce mode."""
return self._engine.check_output_tokens(agent_id, token_count)
[docs]
def check_endpoint(self, agent_id: str, endpoint: str) -> bool:
"""Validate an external endpoint; False when blocked in enforce mode."""
return self._engine.check_endpoint(agent_id, endpoint)
[docs]
def wrap(self, graph: object) -> object:
"""Wrap a LangGraph compiled graph with AgentGuard interception."""
from agentguard.adapters.langgraph import LangGraphAdapter # noqa: PLC0415
return LangGraphAdapter.wrap(self, graph)