Agent Security Best Practices

Why This Matters

Every AI agent deployed in production is a potential attack vector. Unlike traditional software, agents process natural language inputs that can be manipulated, execute tool calls with real-world consequences, and have access to sensitive data. A single prompt injection attack can bypass safety guardrails, exfiltrate private data, or cause unauthorized actions. Security must be built into the architecture from day one—not bolted on afterward.

Real-World Analogy

Think of agent security like airport security. Passengers (inputs) go through multiple checkpoints: ID verification (authentication), baggage screening (input filtering), metal detectors (content classification), and boarding pass validation (authorization). No single checkpoint catches everything—the strength is in the layered approach. If a passenger slips past one layer, the next catches them.

Prompt Injection Detection

import re
import logging
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum

logger = logging.getLogger(__name__)


class ThreatLevel(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


@dataclass(frozen=True)
class ThreatResult:
    detected: bool
    threat_level: ThreatLevel
    threat_type: str
    confidence: float
    details: str


class PromptInjectionDetector:
    def __init__(self) -> None:
        self._injection_patterns: list[tuple[str, ThreatLevel]] = [
            (r"ignore\s+(previous|all|above)\s+(instructions|prompts)", ThreatLevel.HIGH),
            (r"you\s+are\s+now\s+(?:a|an)\s+", ThreatLevel.HIGH),
            (r"system\s*:\s*", ThreatLevel.MEDIUM),
            (r"<\|im_start\|>", ThreatLevel.CRITICAL),
            (r"bypass\s+(safety|restrictions|rules)", ThreatLevel.CRITICAL),
            (r"forget\s+(your|all)\s+(rules|instructions)", ThreatLevel.HIGH),
            (r"act\s+as\s+if\s+you\s+have\s+no\s+restrictions", ThreatLevel.CRITICAL),
            (r"developer\s+mode\s+(on|enabled)", ThreatLevel.CRITICAL),
            (r"override\s+(previous|system)", ThreatLevel.CRITICAL),
            (r"new\s+instructions\s*:", ThreatLevel.HIGH),
        ]
        self._encoding_patterns: list[tuple[str, ThreatLevel]] = [
            (r"base64\s*(encode|decode)", ThreatLevel.MEDIUM),
            (r"rot13", ThreatLevel.MEDIUM),
            (r"hex\s*encode", ThreatLevel.MEDIUM),
        ]

    def detect(self, user_input: str) -> ThreatResult:
        threats: list[ThreatLevel] = []
        for pattern, level in self._injection_patterns + self._encoding_patterns:
            if re.search(pattern, user_input, re.IGNORECASE):
                threats.append(level)

        if not threats:
            return ThreatResult(
                detected=False, threat_level=ThreatLevel.LOW,
                threat_type="none", confidence=0.95,
                details="No injection patterns detected",
            )

        max_level = max(threats, key=lambda lvl: list(ThreatLevel).index(lvl))
        logger.warning("Injection detected: level=%s count=%d", max_level.value, len(threats))
        return ThreatResult(
            detected=True, threat_level=max_level,
            threat_type="prompt_injection",
            confidence=min(0.6 + 0.1 * len(threats), 0.99),
            details=f"Detected {len(threats)} suspicious patterns",
        )

    def sanitize(self, user_input: str) -> str:
        sanitized = user_input
        for pattern, _ in self._injection_patterns + self._encoding_patterns:
            sanitized = re.sub(pattern, "[FILTERED]", sanitized, flags=re.IGNORECASE)
        sanitized = re.sub(r"system\s*:", "system_[FILTERED]", sanitized, flags=re.IGNORECASE)
        return sanitized

Data Leakage Prevention

import re
import logging
from dataclasses import dataclass
from typing import Optional
from enum import Enum

logger = logging.getLogger(__name__)


class DataType(Enum):
    EMAIL = "email"
    PHONE = "phone"
    SSN = "ssn"
    CREDIT_CARD = "credit_card"
    API_KEY = "api_key"
    IP_ADDRESS = "ip_address"


@dataclass(frozen=True)
class PIIDetection:
    data_type: DataType
    value: str
    start: int
    end: int
    confidence: float


class DataLeakagePrevention:
    def __init__(self) -> None:
        self._pii_patterns: dict[DataType, str] = {
            DataType.EMAIL: r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
            DataType.PHONE: r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b",
            DataType.SSN: r"\b\d{3}-\d{2}-\d{4}\b",
            DataType.CREDIT_CARD: r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b",
            DataType.IP_ADDRESS: r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b",
        }
        self._keyword_pattern = re.compile(
            r"(?i)(password|secret|token|api_key|private_key)\s*[:=]\s*[\"']?([^\s\"']+)",
        )

    def detect(self, text: str) -> list[PIIDetection]:
        detections: list[PIIDetection] = []
        for dtype, pattern in self._pii_patterns.items():
            for match in re.finditer(pattern, text):
                detections.append(PIIDetection(
                    data_type=dtype, value=match.group(),
                    start=match.start(), end=match.end(), confidence=0.9,
                ))
        for match in self._keyword_pattern.finditer(text):
            detections.append(PIIDetection(
                data_type=DataType.API_KEY, value=match.group(),
                start=match.start(), end=match.end(), confidence=0.95,
            ))
        return sorted(detections, key=lambda d: d.start, reverse=True)

    def mask(self, text: str) -> str:
        detections = self.detect(text)
        masked = text
        for det in detections:
            replacement = self._mask_value(det.data_type, det.value)
            masked = masked[:det.start] + replacement + masked[det.end:]
        logger.info("Masked %d PII detections in text", len(detections))
        return masked

    @staticmethod
    def _mask_value(dtype: DataType, value: str) -> str:
        if dtype == DataType.EMAIL:
            local, domain = value.split("@", 1)
            return f"{local[0]}***@{domain}"
        elif dtype == DataType.PHONE:
            digits = re.sub(r"\D", "", value)
            return f"({digits[:3]}) ***-{digits[-4:]}"
        elif dtype == DataType.SSN:
            return f"***-**-{value[-4:]}"
        elif dtype == DataType.CREDIT_CARD:
            digits = re.sub(r"\D", "", value)
            return f"****-****-****-{digits[-4:]}"
        elif dtype == DataType.API_KEY:
            return f"{value[:4]}...{value[-4:]}" if len(value) > 8 else "****"
        elif dtype == DataType.IP_ADDRESS:
            parts = value.split(".")
            return f"{parts[0]}.{parts[1]}.xxx.xxx"
        return "***"

Sandboxing System

import asyncio
import io
import contextlib
import logging
import resource
from dataclasses import dataclass, field
from typing import Any, Optional

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SandboxConfig:
    max_memory_mb: int = 256
    max_cpu_time_sec: int = 30
    max_output_size: int = 1_048_576
    network_access: bool = False
    filesystem_access: bool = False
    allowed_imports: tuple[str, ...] = ("json", "math", "datetime", "re")


@dataclass(frozen=True)
class SandboxResult:
    success: bool
    output: str
    error: Optional[str]
    execution_time: float
    memory_used: int
    was_violated: bool
    violation_reason: Optional[str]


_DANGEROUS_IMPORTS = frozenset({
    "subprocess", "os", "shutil", "socket", "urllib", "requests",
    "httpx", "ctypes", "pickle", "pathlib",
})


class SandboxManager:
    def __init__(self, config: Optional[SandboxConfig] = None) -> None:
        self._config = config or SandboxConfig()
        self._violations: list[dict[str, Any]] = []

    async def execute(self, code: str, language: str = "python") -> SandboxResult:
        if language != "python":
            return SandboxResult(
                success=False, output="", error=f"Unsupported language: {language}",
                execution_time=0, memory_used=0, was_violated=False, violation_reason=None,
            )

        violations = self._scan_code(code)
        if violations:
            reason = "; ".join(violations)
            logger.warning("Code blocked: %s", reason)
            return SandboxResult(
                success=False, output="", error=f"Sandbox violation: {reason}",
                execution_time=0, memory_used=0, was_violated=True,
                violation_reason=violations[0],
            )

        try:
            return await asyncio.wait_for(
                self._run_python(code), timeout=self._config.max_cpu_time_sec,
            )
        except asyncio.TimeoutError:
            return SandboxResult(
                success=False, output="", error="Execution timed out",
                execution_time=self._config.max_cpu_time_sec, memory_used=0,
                was_violated=True, violation_reason="timeout",
            )
        except Exception as exc:
            return SandboxResult(
                success=False, output="", error=str(exc),
                execution_time=0, memory_used=0,
                was_violated=False, violation_reason=None,
            )

    def _scan_code(self, code: str) -> list[str]:
        violations: list[str] = []
        if any(imp in code for imp in _DANGEROUS_IMPORTS):
            violations.append("dangerous_imports")
        if not self._config.filesystem_access and any(op in code for op in ("open(", "write(", "os.path")):
            violations.append("filesystem_access")
        if not self._config.network_access and any(op in code for op in ("requests.", "httpx.", "socket.")):
            violations.append("network_access")
        return violations

    async def _run_python(self, code: str) -> SandboxResult:
        output_buffer = io.StringIO()
        mem_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        try:
            with contextlib.redirect_stdout(output_buffer):
                exec(code, {"__builtins__": {}})
            mem_used = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - mem_before
            return SandboxResult(
                success=True, output=output_buffer.getvalue(), error=None,
                execution_time=0, memory_used=mem_used,
                was_violated=False, violation_reason=None,
            )
        except Exception as exc:
            return SandboxResult(
                success=False, output=output_buffer.getvalue(), error=str(exc),
                execution_time=0, memory_used=0,
                was_violated=False, violation_reason=None,
            )

Rate Limiting & Access Control

import time
import logging
from dataclasses import dataclass, field
from typing import Optional
from collections import defaultdict
from enum import Enum

logger = logging.getLogger(__name__)


class AccessLevel(Enum):
    PUBLIC = "public"
    AUTHENTICATED = "authenticated"
    PREMIUM = "premium"
    ADMIN = "admin"


@dataclass(frozen=True)
class RateLimitConfig:
    requests_per_minute: int = 60
    requests_per_hour: int = 1000
    cooldown_seconds: int = 60


@dataclass
class UserPermissions:
    user_id: str
    access_level: AccessLevel
    allowed_tools: frozenset[str] = frozenset()
    max_tokens_per_request: int = 4096


class RateLimiter:
    def __init__(self, config: Optional[RateLimitConfig] = None) -> None:
        self._config = config or RateLimitConfig()
        self._history: dict[str, list[float]] = defaultdict(list)
        self._blocked: dict[str, float] = {}

    def is_allowed(self, user_id: str) -> tuple[bool, Optional[str]]:
        now = time.time()
        if user_id in self._blocked:
            if now < self._blocked[user_id]:
                return False, f"Blocked until {time.strftime('%H:%M:%S', time.localtime(self._blocked[user_id]))}"
            del self._blocked[user_id]

        self._history[user_id] = [t for t in self._history[user_id] if now - t < 3600]
        recent = [t for t in self._history[user_id] if now - t < 60]

        if len(recent) >= self._config.requests_per_minute:
            self._blocked[user_id] = now + self._config.cooldown_seconds
            return False, "Rate limit exceeded (per minute)"
        if len(self._history[user_id]) >= self._config.requests_per_hour:
            return False, "Rate limit exceeded (per hour)"

        self._history[user_id].append(now)
        return True, None


class AccessController:
    def __init__(self) -> None:
        self._permissions: dict[str, UserPermissions] = {}
        self._audit_log: list[dict] = []

    def grant(self, permissions: UserPermissions) -> None:
        self._permissions[permissions.user_id] = permissions
        self._log(permissions.user_id, "access_granted", permissions.access_level.value)

    def check(self, user_id: str, tool: str) -> bool:
        perms = self._permissions.get(user_id)
        if not perms:
            return False
        allowed = perms.access_level in (AccessLevel.ADMIN, AccessLevel.PREMIUM) or tool in perms.allowed_tools
        self._log(user_id, "permission_check", tool, allowed)
        return allowed

    def _log(self, user_id: str, action: str, details: str, success: bool = True) -> None:
        self._audit_log.append({
            "timestamp": time.time(), "user_id": user_id,
            "action": action, "details": details, "success": success,
        })

Mathematical Foundations

Security Score (weighted composite):

where represents normalized scores for injection prevention, leakage control, sandbox integrity, and access control.

Threat Detection Metrics:

Risk Score:

Performance Considerations

Component	Latency	Cost	Accuracy
Regex injection detection	<1ms	Negligible	~85% recall
PII masking	<5ms	Low	~90% precision
Sandbox execution	10-100ms	Medium	High (deterministic)
Rate limiting	<1ms	Negligible	Exact
Content classification	50-200ms	Medium	~92% F1

Security Considerations

Defense in depth: Never rely on a single security control. Layer input filtering, sandboxing, output validation, and monitoring.
Fail closed: When in doubt, deny access. A false negative (missing a threat) is far costlier than a false positive (blocking a legitimate request).
Least privilege: Grant agents only the minimum permissions needed. Use role-based access control (RBAC) with granular tool permissions.
Audit everything: Log all inputs, outputs, tool calls, and security events. Immutable audit logs are essential for incident response and compliance.
Secret management: Never hardcode API keys or credentials. Use environment variables, secret managers (Vault, AWS Secrets Manager), or encrypted config files.

Common Pitfalls

Pitfall	Consequence	Solution
Single defense layer	Bypasses compromise entire system	Implement defense in depth
Ignoring indirect injection	Attacks via retrieved documents	Treat all external data as untrusted
Over-blocking legitimate requests	Poor user experience	Use risk-based access control
Hardcoded secrets in prompts	Credential leakage	Use secret management systems
No output filtering	PII leaks to users	Filter both inputs and outputs
Missing rate limiting	Cost overrun, abuse	Per-user and per-endpoint limits
Static security rules	Miss evolving threats	Regular rule updates and ML classifiers
No incident response plan	Slow breach containment	Maintain and drill response procedures

Interview Q&A

1. What is prompt injection and how do you defend against it?

Prompt injection occurs when an attacker crafts input that overrides system instructions or manipulates the agent's behavior. Defense strategies include: (1) input sanitization with regex pattern matching, (2) content classification using ML-based classifiers, (3) role separation between system and user messages using delimiters, (4) output validation to ensure responses follow guidelines, and (5) sandboxing to limit tool capabilities. The key is defense in depth—never rely on a single layer. Use a classifier to assign confidence scores and route suspicious inputs to a restricted handler.

2. How do you prevent data leakage in AI agents?

Data leakage prevention requires a pipeline approach: (1) PII detection using regex and NER models before processing, (2) output filtering to mask sensitive information before returning to users, (3) access control limiting which data sources the agent can query, (4) audit logging of all data interactions for compliance, and (5) encryption of sensitive data at rest and in transit. Implement a DLP (Data Loss Prevention) pipeline that processes both inputs and outputs, and use token-level analysis to catch partial PII exposures.

3. What are the key components of agent sandboxing?

Sandboxing includes: (1) resource limits (CPU time, memory, output size), (2) network isolation with allowlist-based egress control, (3) filesystem ACLs restricting access to approved paths only, (4) syscall filtering to block dangerous operations, (5) containerization with minimal base images and read-only filesystems, and (6) monitoring with real-time alerting on policy violations. The goal is to minimize the attack surface while allowing necessary functionality. Use gVisor or Firecracker for stronger isolation than Docker alone.

4. How do you implement rate limiting for agent systems?

Rate limiting uses token bucket or sliding window algorithms: (1) track request counts per user per time window, (2) set per-user and per-endpoint limits based on access tier, (3) implement burst handling with temporary overrides for legitimate spikes, (4) apply cooldown periods after limit violations, and (5) provide usage statistics dashboards. Consider both request count and token usage limits. Use Redis-backed distributed counters for multi-instance deployments.

5. What is indirect prompt injection and why is it dangerous?

Indirect injection embeds malicious instructions in retrieved documents (RAG sources, web pages, emails). It's dangerous because: (1) the agent trusts retrieved content, (2) attacks can be hidden in legitimate-looking data, (3) traditional input filtering may miss them since the user didn't type the malicious content, and (4) they can affect any RAG-powered system. Defense requires treating all external data as untrusted, sanitizing retrieved content, and using content security policies that restrict what instructions can appear in non-system messages.

6. How do you balance security with agent functionality?

Balance through: (1) risk-based security—more controls for higher-risk operations like database writes, (2) progressive trust—start with least privilege and expand based on observed behavior, (3) user experience—don't over-block legitimate use cases (tune false positive rates), (4) monitoring over blocking—log and alert rather than block everything, and (5) regular review of security rules based on actual threat patterns. Use confidence thresholds: low-confidence detections get logged, medium gets confirmed, high gets blocked.

7. What audit logging should be implemented for agent security?

Essential audit logs include: (1) all user inputs with PII masked, (2) agent reasoning traces and tool calls, (3) access control decisions, (4) security violations and threats detected, (5) data access patterns, and (6) configuration changes. Logs should be immutable (append-only, checksummed), timestamped with UTC, and stored in a separate security-focused system. Implement log retention policies (7 years for compliance) and regular security reviews. Use structured logging format (JSON) for machine analysis.

8. How do you handle security incidents in agent systems?

Incident response follows NIST phases: (1) preparation—maintain runbooks and communication channels, (2) detection and analysis—correlate logs and alerts, (3) containment—disable affected agents, revoke compromised credentials, (4) eradication—patch vulnerabilities, rotate secrets, (5) recovery—restore from clean state, (6) lessons learned—update controls and conduct drills. For AI agents specifically, maintain model rollback capability, have pre-approved prompt templates, and establish a security review board for high-risk agent behaviors.

KnowledgeCheck

What is the primary defense against prompt injection?
- a) Single regex filter
- b) Defense in depth with multiple layers
- c) Blocking all special characters
- d) Using only whitelisted inputs
Which technique masks sensitive data before processing?
- a) Rate limiting
- b) PII detection and masking
- c) Access control
- d) Output encryption
What is the key principle of agent sandboxing?
- a) Maximize resource access
- b) Minimize attack surface
- c) Allow all network calls
- d) Disable logging
How does indirect prompt injection differ from direct injection?
- a) It uses different attack vectors
- b) It occurs via retrieved documents, not user input
- c) It is easier to detect
- d) It only affects specific models
What should audit logs be?
- a) Editable for corrections
- b) Immutable and timestamped
- c) Stored in the agent database
- d) Cleared daily
What is the recommended approach when a security check has low confidence?
- a) Block immediately
- b) Allow silently
- c) Log and route to restricted handler
- d) Ignore completely

Answers: 1-b, 2-b, 3-b, 4-b, 5-b, 6-c

Agent Security Best Practices: Prompt Injection, Data Leakage & Sandboxing