Detect and contain rogue agent behavior

Introduction

In production, adversaries rarely mount a single-turn attack on an agent — instead they craft multi-step campaigns where every individual tool call stays within policy, but the cumulative pattern signals exfiltration or goal hijacking. Static defenses like argument validation and least-privilege scoping cannot catch this because they examine each turn in isolation. By the end of this lesson, you'll be able to build a behavioral profiling system that models a healthy agent's baseline, score live sessions against it using a composite anomaly detector, and trigger a graduated kill switch — with drift state persisted so that session reconnects cannot be used to reset evidence of an ongoing campaign.

Key Terminology

• Behavioral Profiling — the practice of fitting a statistical model of healthy agent sessions so that future sessions can be scored against it; this lesson's profiler captures two dimensions: sequence length and tool-mix entropy via fit_baseline().
• Tool-Mix Entropy — a Shannon entropy measure of how evenly an agent distributes tool calls across available tools in a session; a rogue exfiltration loop collapses entropy toward zero by hammering a single tool, which anomaly_score() detects as mix_drift.
• Composite Anomaly Score — a weighted drift signal returned by anomaly_score() that blends the sequence-length z-score (weight 0.6) with tool-mix drift (weight 0.4); values above roughly 2.5 warrant containment action.
• Drift State Persistence — the practice of accumulating anomaly scores in DriftStore under a Redis key tied to the agent's first-session timestamp, so reconnects inherit prior evidence rather than resetting the counter to zero.
• Kill Switch — a graduated containment response triggered when cumulative drift exceeds a threshold calibrated to an explicit false-positive budget; upon firing, it stops new tool calls, snapshots session state to incident storage, and revokes all capability tokens.
• False-Positive Budget — an explicit, pre-committed allowance for the fraction of benign sessions the detector may flag (e.g., 2 % per week); when the budget is exceeded, the detector is retuned rather than silenced, preventing threshold creep from hiding real attacks.

Concepts

Why Per-Turn Defenses Miss Multi-Step Campaigns

Argument validation and least-privilege scoping are necessary but not sufficient for agentic security. Both techniques examine a tool call in isolation: does this call's arguments conform to policy? Does the agent have the capability token it's invoking? A sophisticated adversary can answer "yes" to both questions on every individual turn while the sequence of turns encodes malicious intent — slowly exfiltrating data across 30 reads that each stay within rate limits, or pivoting toward a higher-value target one plausible step at a time.

Catching this class of attack requires shifting the unit of analysis from the individual call to the session. A behavioral profiler treats the full sequence of tool invocations as the observable, fits a model of what healthy sequences look like, and scores each live session against that model. Anomalies emerge not because any single call broke a rule, but because the pattern diverges from baseline.

Two Dimensions of Normal: Sequence Length and Tool-Mix Entropy

A healthy agent session has two statistical signatures that are both easy to measure and difficult for an attacker to mimic while executing a real campaign. The first is sequence length — how many tool calls a session makes. The second, and more diagnostic, is tool-mix entropy: how evenly the session's calls are distributed across the available tool set.

Entropy is a powerful signal because rogue campaigns are inherently repetitive. An exfiltration loop that reads the same file class over and over, or a goal-hijacking sequence that repeatedly probes a privilege-escalation tool, drives entropy toward zero. A legitimate agent, by contrast, tends to be versatile — reading, writing, querying, and acting across a range of tools in proportion to the task. The baseline captures this diversity as a single scalar via fit_baseline(), and anomaly_score() measures how far the live session's entropy drifts from it (see Code Walkthrough).

The composite score weights sequence-length deviation at 60% and entropy drift at 40%. The weighting reflects a design choice: the z-score is well-normalized (it accounts for baseline variance via seq_std) while raw entropy drift is not normalized, so it gets a smaller weight to keep the composite stable across baselines with different tool-set sizes.

Closing the Reconnect Bypass with Persistent Drift State

Any detector that resets its accumulator on reconnect can be bypassed trivially: the attacker interrupts the session after every few tool calls, reconnects, and the fresh counter never accumulates enough drift to cross the kill-switch threshold. This bypass is well-known and requires no special tooling — it just requires patience.

The correct fix is to anchor drift state to the agent's first session timestamp, not its current connection identifier. DriftStore (see Code Walkthrough) encodes this invariant by building its Redis key from agent_id plus first_session_at.date(). Every reconnect of the same agent resolves to the same key and inherits the prior running total via redis.incrbyfloat, which is also atomic — preventing a race where two concurrent sessions both read a stale score and each increment from it separately.

This design means the kill switch fires on the campaign, not the session. An adversary who reconnects 10 times to stay under the per-session threshold accumulates drift across all 10 bursts and still crosses the threshold when the campaign as a whole reaches sufficient evidence.

Code Walkthrough

Now that you understand the operating discipline around baseline refresh cadence, false-positive budgets, and mandatory kill-switch alerting, here is how those principles translate into working Python.

Fitting the baseline. A profiler summarizes healthy sessions along two key dimensions: sequence length (how many tool calls a normal session makes) and tool-mix entropy (how evenly the agent distributes its calls across available tools). Entropy matters because a rogue campaign often hammers a single exfiltration tool repeatedly, collapsing the diversity score.

Code snippetpython
1from collections import Counter
2from dataclasses import dataclass
3from math import log2
4from statistics import mean, stdev
5
6@dataclass
7class Baseline:
8    seq_mean: float
9    seq_std: float
10    tool_entropy: float
11    sample_size: int
12
13def fit_baseline(sessions: list[list[str]]) -> Baseline:
14    """sessions: each session is a list of tool-name strings."""
15    seq_lens = [len(s) for s in sessions]
16    all_tools = [tool for s in sessions for tool in s]
17    counts = Counter(all_tools)
18    total = sum(counts.values())
19    entropy = -sum((v / total) * log2(v / total) for v in counts.values())
20    return Baseline(
21        seq_mean=mean(seq_lens),
22        seq_std=stdev(seq_lens) if len(seq_lens) > 1 else 0.0,
23        tool_entropy=entropy,
24        sample_size=len(sessions),
25    )
26
27def anomaly_score(session: list[str], baseline: Baseline) -> float:
28    """Returns a composite drift score; values above ~2.5 warrant action."""
29    seq_z = abs(len(session) - baseline.seq_mean) / (baseline.seq_std or 1.0)
30    counts = Counter(session)
31    total = sum(counts.values()) or 1
32    entropy = -sum((v / total) * log2(v / total) for v in counts.values())
33    mix_drift = abs(entropy - baseline.tool_entropy)
34    return 0.6 * seq_z + 0.4 * mix_drift

Persisting drift state across reconnects. A naive implementation resets the counter each time an agent reconnects — an attacker can exploit this by triggering repeated reconnects to wash out accumulated drift. The fix is to key drift state against the agent's first session timestamp, not the current connection, so every reconnect inherits the prior tally.

Code snippetpython
1import redis
2from datetime import datetime
3
4class DriftStore:
5    def __init__(self, client: redis.Redis) -> None:
6        self.redis = client
7
8    def _key(self, agent_id: str, first_session_at: datetime) -> str:
9        return f"drift:{agent_id}:{first_session_at.date().isoformat()}"
10
11    def accumulate(self, agent_id: str, first_session_at: datetime, score: float) -> float:
12        """Add score to the running tally and return the new cumulative total."""
13        return float(self.redis.incrbyfloat(self._key(agent_id, first_session_at), score))
14
15    def read(self, agent_id: str, first_session_at: datetime) -> float:
16        raw = self.redis.get(self._key(agent_id, first_session_at))
17        return float(raw) if raw else 0.0

When accumulate returns a value above the kill-switch threshold — calibrated against your explicit false-positive budget — raise a containment exception, stop issuing new tool calls, snapshot session state to incident-response storage, and revoke all capability tokens for the session.

Confirm that calling accumulate across two separate DriftStore instances pointed at the same Redis server returns the same running cumulative total — this verifies that session reconnects inherit prior drift rather than restarting from zero.

Do's and Don'ts

Now that you have worked through the implementation, the practices below separate a durable approach from a fragile one.

Do's

1. ✓Do key drift state against the agent's first session timestamp — use first_session_at.date().isoformat() as part of the Redis key in DriftStore so that reconnects accumulate onto the same running tally rather than starting fresh; an attacker who repeatedly reconnects to wash out anomaly evidence will instead find the cumulative score climbing.
2. ✓Do score sessions using both seq_z and mix_drift as a composite — the anomaly_score function weights sequence-length Z-score at 0.6 and tool-entropy drift at 0.4 because either dimension alone misses campaigns that stay within normal call counts but hammer a single exfiltration tool, collapsing diversity while keeping session length ordinary.
3. ✓Do fit Baseline.tool_entropy against the full cross-session tool distribution, not per-session averages — fit_baseline aggregates all tool calls before computing entropy so that a rogue session's entropy is compared against the true population mix; per-session averaging would mask legitimate heavy use of one tool in healthy sessions.

Don'ts

1. ✗Don't reset drift counters on reconnect by instantiating a fresh DriftStore without passing the original first_session_at — omitting the first-session anchor means _key generates a new Redis key on each reconnect, silently resetting cumulative drift to zero and letting a multi-session exfiltration campaign stay permanently below the kill-switch threshold.
2. ✗Don't use baseline.seq_std as the Z-score denominator without guarding the zero case — passing a single-session baseline or a training set where every session has the same length produces seq_std=0.0; dividing by it raises ZeroDivisionError, which is why anomaly_score replaces it with 1.0 — drop that guard and any homogeneous baseline silently crashes the detector at runtime.
3. ✗Don't trigger the kill switch solely on a single-turn anomaly_score spike without accumulating via DriftStore — a one-time high score can be a legitimate burst; it is the cumulative drift returned by accumulate exceeding your calibrated false-positive budget that warrants containment, capability-token revocation, and incident-response snapshotting.