Build a Basic AI Agent From Scratch: Security III

Previous parts of Build a Basic AI Agent From Scratch:

• Basic Agent
• Tools
• Long Task Planning
• Human in the Loop & Security
• Security II

You can find and clone this code in this blog series' Github repo.

In the previous part we started closing the gaps left open by human-in-the-loop: a Docker sandbox to contain runaway commands, prompt-injection defenses so the model stops trusting tool output as instructions, and schema validation so malformed tool calls never reach execution. In this part we finish the job. We will harden the tool policy gate with path scoping, a shell denylist, and an SSRF guard, enforce resource and cost limits so a stuck loop cannot run forever, scrub secrets out of the container environment, and add audit logging and a kill switch so every decision is recorded and any session can be aborted mid-flight.

Tool Policy Hardening

In the Human in the Loop & Security part, check_permission was a single function: read tools and planning tools were always allowed, write tools were allowed inside the working directory in acceptEdits, everything else asked. That was a mode-based decision.

We now add a policy gate that runs before the mode decision. The gate has three layers:

• Path scoping : every path argument is resolved (handling relative paths, symlinks, and .. traversal) and rejected if it escapes the project working directory. This stops the agent from touching ~/.ssh/id_rsa, /etc/passwd, or anything outside the project tree.
• Shell policy : every run_bash command is screened against a regex denylist of dangerous patterns (fork bombs, dd, mkfs, redirects into /etc/, ...) and a token-level denylist of binaries the agent should never invoke (sudo, nc, curl, chmod, docker, ...). This catches destructive and exfiltration commands.
• SSRF guard : Server-Side Request Forgery is an attack where a server-side process is tricked into making requests to internal resources it shouldn't be able to reach. In an agent context, a prompt injection could make webfetch hit internal services or the cloud metadata endpoint. The guard resolves the URL's host, checks the resulting IP against loopback, link-local, multicast, and RFC1918 private ranges, and blocks them.

def check_tool_policy(
    tool_name: str,
    args: dict[str, Any],
    working_dir: Path,
) -> tuple[bool, str | None]:
    """Run all policy layers. Returns (allowed, reason).

    Called by ``agent.check_permission`` BEFORE the mode-based decision.
    A False here is a hard block that no mode can override.
    """
    # Layer 1: path scope.
    ok, reason = check_path_scope(tool_name, args, working_dir)
    if not ok:
        return False, reason

    # Layer 2: shell policy.
    if tool_name == "run_bash":
        ok, reason = check_shell_policy(args.get("command", ""))
        if not ok:
            return False, reason

    # Layer 3: web / SSRF policy.
    if tool_name == "webfetch":
        ok, reason = check_web_policy(args.get("url", ""))
        if not ok:
            return False, reason

    return True, None

Enter fullscreen mode Exit fullscreen mode

Path Scoping

The old check only ran on write tools. Now, the check is generalized to every tool that uses paths: read_file, glob_files, grep, write_file, edit_file. Each tool's path argument is resolved (handling relative paths, symlinks, and .. traversal) and rejected if it escapes the working directory:

PATH_TOOLS: dict[str, str] = {
    "read_file": "path",
    "glob_files": "path",
    "grep": "path",
    "write_file": "path",
    "edit_file": "path",
}

def check_path_scope(
    tool_name: str,
    args: dict[str, Any],
    working_dir: Path,
) -> tuple[bool, str | None]:
    """Reject path-bearing tool calls whose target escapes working_dir.

    Note: the Docker mount already constrains the *container's* view of
    the filesystem; this check is a defense-in-depth layer on the host
    side so a malicious path is rejected before it ever reaches docker.
    """
    if tool_name not in PATH_TOOLS:
        return True, None
    raw = args.get(PATH_TOOLS[tool_name])
    if not raw:
        return True, None # missing arg is a schema problem, not a scope problem
    try:
        target = Path(raw)
        if not target.is_absolute():
            target = working_dir / target
        target.resolve().relative_to(working_dir.resolve())
        return True, None
    except (ValueError, OSError, RuntimeError) as e:
        return False, (
            f"Path '{raw}' is outside the working directory "
            f"({working_dir}). File tools may only touch paths inside "
            f"the project root. ({e})"
        )

Enter fullscreen mode Exit fullscreen mode

This is defense-in-depth on the host side. The Docker mount already constrained the container's view of the file system, but a malicious path is rejected here before it even reaches the sandbox.

Shell Policy

By far, run_bash is the most dangerous tool, so it gets its own screening. Two checks run against every command.

First, a regex denylist of dangerous patterns:

SHELL_DENYLIST_PATTERNS = [
    (re.compile(r"\brm\s+-rf?\s+(/|~|\*|\$HOME|\.\.)", re.I),
     "recursive delete of a broad or root target"),
    (re.compile(r">\s*/etc/", re.I),
     "redirect into /etc/ (system files)"),
    (re.compile(r"\bmkfs\b", re.I), "filesystem format command"),
    (re.compile(r"\bdd\b\s+if=", re.I), "raw disk write via dd"),
    (re.compile(r":\(\)\s*\{", re.I), "fork-bomb pattern"),
    (re.compile(r"\b(eval|exec)\b", re.I),
     "eval/exec in a shell command (injection risk)"),
    (re.compile(r">\s*/dev/sd", re.I), "write to a block device"),
    (re.compile(r"\bhistory\s+-c\b", re.I), "history wipe"),
    (re.compile(r"\bexport\s+PATH=", re.I),
     "PATH override (could shadow binaries)"),
]

Enter fullscreen mode Exit fullscreen mode

Then, the command is lexically parsed and every token is checked against a deny list. The agent should never use any of those binaries in any command:

SHELL_DENYLIST_BINARIES = frozenset({
    "docker", # sandbox escape / host control
    "sudo", "su", # privilege escalation
    "nc", "netcat", "ncat", # reverse shells / exfil
    "curl", "wget", # exfil / SSRF
    "chmod", "chown", # permission tampering
    "mkfs", "dd", # destructive disk ops
    "shutdown", "reboot", "halt", "poweroff",
    "systemctl", "service",
    "crontab", "at",
})

Enter fullscreen mode Exit fullscreen mode

Web Policy (SSRF Guard)

webfetch is screened against server-side request forgery. Cloud metadata endpoints, localhost, and private IP ranges are all blocked:

WEB_DENYLIST_HOSTS = frozenset({
    "169.254.169.254", # AWS / GCP / Azure cloud metadata
    "metadata.google.internal", # GCP metadata
    "metadata.azure.com", # Azure metadata
    "0.0.0.0",
    "::1",
    "localhost",
})

def check_web_policy(url: str) -> tuple[bool, str | None]:
    """Reject URLs that target loopback / link-local / private ranges."""
    from urllib.parse import urlparse
    try:
        parsed = urlparse(url)
    except ValueError as e:
        return False, f"Unparseable URL: {e}"
    if parsed.scheme not in ("http", "https"):
        return False, f"Unsupported scheme '{parsed.scheme}'."
    host = parsed.hostname
    if not host:
        return False, "URL has no host component."
    if host.lower() in WEB_DENYLIST_HOSTS:
        return False, f"Blocked host '{host}' (loopback / metadata)."

    # Resolve and check the IP family.
    try:
        infos = socket.getaddrinfo(host, None)
    except socket.gaierror:
        # Let the actual fetcher surface the DNS error.
        return True, None
    for info in infos:
        ip = info[4][0]
        try:
            addr = ipaddress.ip_address(ip)
        except ValueError:
            continue
        if addr.is_loopback or addr.is_link_local or addr.is_multicast:
            return False, f"Blocked IP '{ip}' (loopback / link-local / multicast)."
        if addr.is_private:
            return False, f"Blocked IP '{ip}' (RFC1918 private range SSRF guard)."
    return True, None

Enter fullscreen mode Exit fullscreen mode

localhost and 127.0.0.1 are where services that aren't exposed to the public network live: databases (Postgres on 5432, Redis on 6379), admin panels, metrics endpoints, internal APIs, and sometimes cloud-metadata proxies. A prompt-injected webfetch("http://localhost:8080/admin/delete-all") could trigger destructive actions on a service that trusts local connections and therefore has no auth. Blocking loopback prevents the agent from being used as a pivot into the host's internal network.

The host is resolved via getaddrinfo and each returned IP is checked with ipaddress. Loopback, link-local, multicast, and RFC1918 private ranges are all blocked. This stops the model from fetching http://169.254.169.254/latest/meta-data/iam/security-credentials/ to steal cloud credentials, or probing internal services on 10.0.0.0/8.

Always-Confirm Patterns

On top of the hard policy gate, some patterns are dangerous enough that they require explicit confirmation regardless of permission mode. These are the irreversible / exfiltration class:

# NOTE: for run_bash, anything already covered by SHELL_DENYLIST_BINARIES
# or SHELL_DENYLIST_PATTERNS (rm -rf, sudo/su, docker, chmod,
# curl/wget/nc/netcat/ncat) is a hard block in layer 1 and never
# reaches this layer, so it is deliberately NOT duplicated here. Only
# patterns that layer 1 does *not* already hard-block belong in this list.
ALWAYS_CONFIRM_ARG_PATTERNS: list[tuple[str, re.Pattern[str], str]] = [
    # force-push to git (git itself is not on the shell denylist)
    ("run_bash", re.compile(r"\bgit\s+push\s+(-f|--force)", re.I),
     "force-push to git (rewrites remote history)"),
    # overwriting a file with empty content = delete
    ("write_file", re.compile(r'"content"\s*:\s*"\s*"'),
     "write_file with empty content on an existing file (delete-via-empty)"),
]

Enter fullscreen mode Exit fullscreen mode

The check_permission function now has a 3-layer structure:

1. Hard policy gate (check_tool_policy). Checks path scope, shell policy, SSRF guard. A False here blocks the call regardless of mode.
2. Always-confirm. If always_confirm_required returns (True, reason):
   • in dangerouslySkipPermissions: the call is refused outright (irreversible actions are never auto-run, even with the user's blanket opt-in);
   • in default / acceptEdits: the user is prompted with a [DESTRUCTIVE] label.
3. Mode decision. The original default / acceptEdits / dangerouslySkipPermissions logic, with the addition that write_file emptying an existing file forces a [DELETE-via-empty] confirmation even in acceptEdits.

def check_permission(
    tool_name: str,
    args: dict,
    mode: PermissionMode,
    working_dir: Path,
) -> tuple[bool, str | None]:
    # Layer 1: hard policy gate (path scope, shell, SSRF).
    ok, reason = tool_policy.check_tool_policy(tool_name, args, working_dir)
    if not ok:
        return False, reason

    # Layer 2: always-confirm & irreversible patterns.
    confirm_required, confirm_reason = tool_policy.always_confirm_required(
        tool_name, args)
    if confirm_required:
        if mode == PermissionMode.DANGEROUSLY_SKIP_PERMISSIONS:
            return False, (
                f"Blocked even in dangerouslySkipPermissions: "
                f"{confirm_reason or 'irreversible action'}. "
                "Run this command manually outside the agent if it is truly intended."
            )
        return _ask_permission(f"{tool_name} [DESTRUCTIVE]", args), None

    # Layer 3: mode-based decision (read/planning free, acceptEdits for in-tree writes, etc.)
    # ...

Enter fullscreen mode Exit fullscreen mode

always_confirm_required returns (bool, reason) rather than a bare bool, so this layer carries its own specific reason instead of falling back on whatever layer 1 happened to return (which, on the success path, is always None). check_permission itself also returns (allowed, reason), so rejections carry a machine-readable reason surfaced to the LLM and the audit log.

Principle of Least Privilege

One more control fits here: a --tools CLI flag that accepts a comma-separated allow list of tool names. The harness filters both the in-process registry and the schemas exposed to the LLM, so the model never even sees tools it can't call:

$ python agent.py --tools read_file,glob_files,grep,run_bash

Enter fullscreen mode Exit fullscreen mode

If we know our agent will never need to write files, we can reduce the danger blast radius enormously by limiting the allowed tools.

Resource & Cost Controls

We will add to the harness some resource and cost controls to limit the agent from going out of hand while we are not looking at it.

Hard Iteration Caps

IterationCaps tracks two counters:

• turns : incremented once per LLM response within a single user turn. Reset on each new user message. The default is 40 turns, but it can be overwritten with --max-turns.
• tool_calls : incremented once per tool dispatch, across the whole session. The default cap is 200, but it can be overwritten with --max-tool-calls.

@dataclass
class IterationCaps:
    max_turns_per_user_msg: int = DEFAULT_MAX_TURNS_PER_USER_MSG
    max_tool_calls_per_session: int = DEFAULT_MAX_TOOL_CALLS_PER_SESSION
    turns: int = 0
    tool_calls: int = 0

    def reset_turn(self) -> None:
        self.turns = 0

    def bump_turn(self) -> str | None:
        self.turns += 1
        if self.turns > self.max_turns_per_user_msg:
            return (
                f"Reached the per-turn iteration cap "
                f"({self.max_turns_per_user_msg} LLM turns). Stop calling "
                f"tools and give the user a concise summary of progress."
            )
        return None

Enter fullscreen mode Exit fullscreen mode

When a cap is reached, the loop injects a "stop and summarize" message into messages and logs an iteration_cap_hit audit event (with scope of per_turn or session). The model is told to stop calling tools and give the user a concise summary.

The defaults are chosen to be generous enough for real coding tasks but short enough that a stuck loop is killed quickly.

Token Budget Enforcement

ContextBudget tracks cumulative tokens (we do a rough estimation of 4 chars per token) and trims the message history before each LLM call:

The 4-chars-per-token estimate is a rough heuristic, but it's good enough for a safety budget. Instead, we could get the actual token count from the specific model API, but to avoid implementing a method to get the actual token count for each different LLM Provider, we will use the rough estimation for now.

@dataclass
class ContextBudget:
    max_tokens: int = DEFAULT_MAX_CONTEXT_TOKENS
    trim_threshold: float = 0.8
    keep_recent: int = DEFAULT_KEEP_RECENT_MESSAGES

    def check_and_trim(self, messages: list) -> tuple[bool, str | None]:
        """Trim *messages* in place if over budget."""
        self.last_estimate = self.estimate_total(messages)
        if self.last_estimate <= int(self.max_tokens * self.trim_threshold):
            return False, None

        # We always keep messages[0] (system prompt) and the last
        # ``keep_recent`` messages. Everything in between is a
        # candidate for trimming.
        if len(messages) <= self.keep_recent + 1:
            return False, None

        cut_start = 1
        cut_end = len(messages) - self.keep_recent
        dropped = messages[cut_start:cut_end]

        summary = self._summarize(dropped)
        messages[cut_start:cut_end] = [{
            "role": "system",
            "content": summary,
        }]
        # ...

Enter fullscreen mode Exit fullscreen mode

If the estimate exceeds max_tokens * trim_threshold (by default it's 80%), it replaces everything between the system prompt and the last n keep_recent messages (by default 8 messages) with a summary message.

The summary is deterministic and LLM-free: it records dropped message count, total chars, tool-call names, a SHA-256 hash of the dropped conversation (first 16 hex), and an instruction to re-read files rather than relying on dropped context:

@staticmethod
def _summarize(dropped: list) -> str:
    tool_calls: list[str] = []
    total_chars = 0
    for m in dropped:
        # ... collect tool-call names and char counts
    digest = hashlib.sha256(blob.encode("utf-8", "replace")).hexdigest()[:16]
    lines = [
        "[context-trim summary]",
        f"Earlier conversation dropped to stay within the token budget.",
        f"Dropped messages: {len(dropped)} ({total_chars} chars).",
        f"Tools called in dropped section: {', '.join(tool_calls) or 'none'}.",
        f"Conversation hash (first 16 hex): {digest}.",
        "Re-read any files you need rather than relying on the dropped context.",
    ]
    return "\n".join(lines)

Enter fullscreen mode Exit fullscreen mode

A companion function, cap_tool_result, caps each tool result at 32 KB before insertion into messages, with a notice telling the model to use read_file offset/limit for more. Without this, a single read_file of a 50k-line file would blow the budget in one call.

By default, the budget is 24,000 tokens, but it can be overwritten with --max-context-tokens.

Timeouts

We already saw the per-tool-call timeout in the previous part. The LLM call itself now carries timeout=llm_timeout (default 120s, via --llm-timeout). A timeout or connection failure is caught and reported to the user rather than crashing the process, with an llm_call_error audit event:

try:
    response = client.chat.completions.create(
        model="gemma4",
        messages=messages,
        tools=active_schemas,
        temperature=0.7,
        timeout=llm_timeout,
    )
except Exception as e:
    audit.log("llm_call_error", error=str(e))
    print(f"\n [llm error] {e}")
    print("Assistant: I could not reach the model. "
          "Please check the endpoint and retry.")
    break

Enter fullscreen mode Exit fullscreen mode

Cost Circuit Breakers

CostTracker accumulates API spend. After each chat.completions.create, record_usage reads response.usage (or None for local backends like Ollama that don't report usage) and accumulates total_tokens_in, total_tokens_out, and total_cost_usd:

@dataclass
class CostTracker:
    max_cost_usd: float = DEFAULT_MAX_COST_USD
    price_in: float = DEFAULT_PRICE_PER_1K_IN
    price_out: float = DEFAULT_PRICE_PER_1K_OUT
    total_tokens_in: int = 0
    total_tokens_out: int = 0
    total_cost_usd: float = 0.0
    calls: int = 0

    def record_usage(self, usage: Any | None) -> None:
        self.calls += 1
        if usage is None:
            return
        pt = getattr(usage, "prompt_tokens", None)
        ct = getattr(usage, "completion_tokens", None)
        # ... handle dict or object form
        self.total_tokens_in += pt or 0
        self.total_tokens_out += ct or 0
        self.total_cost_usd = (
            self.total_tokens_in / 1000.0 * self.price_in
            + self.total_tokens_out / 1000.0 * self.price_out
        )

    def check(self) -> str | None:
        if self.total_cost_usd >= self.max_cost_usd:
            return (
                f"Cost limit reached: ${self.total_cost_usd:.4f} >= "
                f"${self.max_cost_usd:.4f} cap. Stop and report to the user."
            )
        return None

Enter fullscreen mode Exit fullscreen mode

check() returns a reason when spend ≥ max_cost_usd; the loop logs a cost_limit_hit audit event and stops with a user-facing message. CLI flag: --max-cost-usd (default $5). The session-end summary now also prints the full cost breakdown (calls, tokens_in, tokens_out, cost).

Secret & Credential Management

The model can leak anything it can see. If secrets get in the context window of our agent they can easily leak via logs, or tool uses. So we will have to make sure that the model never sees secrets. To enforce this, we will use three layers:

Never Put Secrets in the System Prompt

Two complementary checks run at startup.

scan_environment_for_secrets() lists all host env vars matching the pattern KEY|SECRET|TOKEN|PASSWORD|PASSWD|CREDENTIAL|APIKEY|API_KEY|AUTH (case-insensitive). The count is logged in the audit config event as secret_env_count.

SECRET_ENV_PATTERN = re.compile(
    r"(.*(?:KEY|SECRET|TOKEN|PASSWORD|PASSWD|CREDENTIAL|APIKEY|API_KEY|AUTH).*)",
    re.IGNORECASE,
)

def scan_environment_for_secrets() -> list[tuple[str, str]]:
    """Return a list of (name, source) for env vars that look like secrets."""
    found: list[tuple[str, str]] = []
    for name in sorted(os.environ):
        if SECRET_ENV_PATTERN.match(name):
            found.append((name, "env"))
    return found

Enter fullscreen mode Exit fullscreen mode

audit_system_prompt(prompt) statically checks the system prompt template before it is sent to the model. It flags os.environ / os.getenv interpolation patterns in the template (e.g. f"... {os.environ['API_KEY']} ...") and literal occurrences of host secret env-var names in the prompt.

PROMPT_INTERPOLATION_PATTERN = re.compile(
    r"\{.*(?:os\.environ|os\.getenv|getenv|environ).*\}",
    re.IGNORECASE,
)

def audit_system_prompt(prompt: str) -> list[str]:
    """Statically check *prompt* for patterns that could leak secrets."""
    warnings: list[str] = []

    if PROMPT_INTERPOLATION_PATTERN.search(prompt):
        warnings.append(
            "System prompt appears to interpolate os.environ / os.getenv. "
            "Never put secrets in the system prompt, the model can leak "
            "them in tool calls or responses."
        )

    for name in os.environ:
        if SECRET_ENV_PATTERN.match(name) and name in prompt:
            warnings.append(
                f"System prompt literally contains the env-var name "
                f"'{name}'. Even if this is the name and not the value, "
                f"its presence may prompt the model to look it up."
            )

    return warnings

Enter fullscreen mode Exit fullscreen mode

If any warning is found, the harness refuses to start (RuntimeError). This is a defense-in-depth check: the current prompt is a static string with no env-var interpolation, but if someone later edits it to inject a key, this catches it at startup before the model ever sees the prompt.

Credential Injection at the Harness Level

The sandbox container only inherits an allow list of env vars from the host. Everything else is stripped out before the container starts:

ALLOWED_CONTAINER_ENV = frozenset({
    "PATH",
    "HOME",
    "USER",
    "LANG",
    "LC_ALL",
    "TERM",
    "AGENT_SESSION_ID",
    "AGENT_SESSION_TOKEN"
})

def build_container_env(session_id: str, session_token: str) -> dict[str, str]:
    """Return the minimal environment dict for the sandbox container.

    Only ALLOWED_CONTAINER_ENV variables are inherited from the host;
    everything else (including any secret-looking env vars) is stripped.
    """
    env: dict[str, str] = {}
    for name in ALLOWED_CONTAINER_ENV:
        if name in os.environ:
            env[name] = os.environ[name]
    env["AGENT_SESSION_ID"] = session_id
    env["AGENT_SESSION_TOKEN"] = session_token
    return env

Enter fullscreen mode Exit fullscreen mode

DockerSandbox. __init__ accepts a container_env dict; _start_container passes each entry as a -e NAME=VALUE flag to docker run.

A companion check lists host paths that must never be bind-mounted into the container: ~/.aws, ~/.ssh, ~/.config/gcloud, ~/.docker, ~/.netrc, ~/.kube, ~/.gnupg. check_credential_mounts() reports which exist on the host so the operator can verify the docker run command never mounts them:

CREDENTIAL_MOUNT_PATHS: list[Path] = [
    Path.home() / ".aws",
    Path.home() / ".ssh",
    Path.home() / ".config" / "gcloud",
    Path.home() / ".docker",
    Path.home() / ".netrc",
    Path.home() / ".kube",
    Path.home() / ".gnupg",
]

def check_credential_mounts() -> list[Path]:
    """Return the subset of credential paths that exist on the host."""
    return [p for p in CREDENTIAL_MOUNT_PATHS if p.exists()]

Enter fullscreen mode Exit fullscreen mode

Rotate Credentials per Session

SessionCredentials generates fresh per-session credentials:

class SessionCredentials:
    def __init__ (self) -> None:
        self.session_id = secrets.token_urlsafe(12)
        self.session_token = secrets.token_urlsafe(32)
        self._revoked = False

    def revoke(self) -> None:
        """Mark the session credentials as revoked.

        In the current single-process model this is bookkeeping; in a
        multi-process / server scenario it would also invalidate the
        token in a shared store.
        """
        self._revoked = True
        # Rotate: generate a new token so any stale reference is useless.
        self.session_token = secrets.token_urlsafe(32)

    def __repr__ (self) -> str:
        # Never include the token itself in repr/log output.
        return f"SessionCredentials(id={self.session_id!r}, revoked={self._revoked})"

Enter fullscreen mode Exit fullscreen mode

The token is injected into the container env as AGENT_SESSION_TOKEN by the harness. Any harness-authenticated tool (e.g. a future github_api tool) uses this token rather than a long-lived host credential.

revoke() is called in the finally block of __main__ on session end. Recreating the container rotates the credentials, which DockerSandbox already does once per session.

Observability & Kill Switch

Security is great. But you also have to think about what happens when your security measures break down. For the last measure of security, we will ask ourselves: if something goes wrong anyway, can we see it, and can we stop it?

Append-Only Audit Log

AuditLog writes one JSON object per line to logs/audit-<session>-<timestamp>.jsonl, flushing after every record so a crash still leaves a complete trail. The audit includes event types for: permission checks, LLM requests/responses, tool results, cap hits, and aborts.

class AuditLog:
    """Append-only JSONL audit log of agent activity."""

    def __init__ (self, log_dir: Path, session_id: str | None = None):
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.session_id = session_id or uuid.uuid4().hex[:12]
        ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        self.path = self.log_dir / f"audit-{self.session_id}-{ts}.jsonl"
        self._fh = self.path.open