AI agents · Threat modeling · Security

Threat Model an AI Agent System from First Principles

Turn an agent architecture into a data-flow model, identify attacker paths across tools and memory, and convert the highest risks into testable controls.

Published
Updated
Code examples
Illustrative
Reading time
5 minutes

A useful agent threat model begins with assets, identities, data flows, trust boundaries, and irreversible effects—not a generic list of AI risks. Draw what the system actually does, ask how an attacker can cross each boundary, assign an owner and testable control to every material threat, and update the model whenever tools, memory, models, or deployment boundaries change.

OWASP frames threat modeling around four questions: what are we building, what can go wrong, what will we do, and did we do enough. It recommends maintaining the model with the system rather than treating it as a one-time document (OWASP Threat Modeling Cheat Sheet).

Define the system and security objective

Use one concrete system: a support agent reads customer tickets and knowledge articles, queries account status, drafts replies, and sends a reply only after approval.

Write the objective before the diagram:

The agent may read tickets for the authenticated support tenant, retrieve
approved knowledge, and draft responses. It may send only the exact preview
approved by an authorized operator. It must never expose credentials,
cross-tenant data, or instructions embedded in untrusted content.

This statement creates testable invariants. “The agent should be safe” does not.

List assets and the consequence of compromise:

AssetSecurity propertyExample impact
Customer ticket dataConfidentiality, tenant isolationCross-customer disclosure
Reply authorityIntegrity, user intentFraudulent external message
Agent memoryIntegrity, provenancePersistent goal hijacking
Tool credentialsConfidentiality, audience restrictionAccess outside agent policy
Approval recordIntegrity, bindingSwap approved preview before send
Audit trailIntegrity, availabilityIncident cannot be reconstructed

Include business assets such as reputation and spend. Unbounded loops may harm availability and generate unexpected cost even without traditional data compromise. OWASP's agent security guidance identifies tool abuse, memory poisoning, goal hijacking, excessive autonomy, cascading failures, and denial-of-wallet among agent-specific risks (OWASP AI Agent Security Cheat Sheet).

Draw the real data-flow diagram

Operator browser
   │ authenticated request
   ▼
Agent API ─────► model provider
   │                ▲
   │                │ prompts and outputs
   ├──► ticket tool ┤
   ├──► retrieval ──┤
   ├──► memory store│
   ├──► approval store
   └──► send-reply tool ──► customer email system

External/untrusted inputs:
ticket body, attachments, retrieved documents, tool results, URLs,
model output, delegated-agent artifacts

For every arrow record source identity, destination identity, authentication, authorization decision, data classification, protocol, retention, size limit, and whether the data can influence a tool. Mark trust boundaries around the browser, application, model provider, each tool service, each tenant, and every third party.

The model is not trusted merely because its provider is trusted. Its output is probabilistic, may reflect hostile context, and must be validated before it reaches a privileged interface. Retrieved text and tool results are also attacker-controlled whenever an attacker can influence the underlying ticket, page, document, or API.

Inventory identities and authority

Separate these actors:

  • the human user who requested work;
  • the client application;
  • the agent runtime;
  • each tool service account;
  • a delegated remote agent; and
  • administrators and incident responders.

For each tool, answer:

Who authenticated?
Which tenant is active?
Which exact objects are authorized?
Is this read, preview, commit, or delete?
Can it be retried safely?
What approval is required?
Which credential reaches the downstream service?

Never infer the user or tenant from model output. Never let one broad agent credential silently replace the user's narrower authority. NIST's zero-trust guidance emphasizes protecting resources rather than trusting network location and making access decisions around subjects and resources (NIST SP 800-207).

Walk each boundary with two threat lenses

Use STRIDE for conventional security and an agent-specific abuse pass. STRIDE prompts spoofing, tampering, repudiation, information disclosure, denial of service, and elevation of privilege. OWASP presents it as one useful technique, not a mandatory universal standard.

For the ticket-to-model boundary:

LensThreatControl candidate
TamperingTicket contains indirect prompt injectionPreserve provenance; treat body as data
DisclosureModel receives unrelated customer historyQuery by verified tenant and minimum fields
ElevationInjected text induces send_replySeparate draft from commit; approval required
Agent-specificHostile content becomes durable memoryQuarantine writes; attach source and TTL
DenialHuge attachment exhausts context or parserType, size, page, time, and token limits

Repeat for model output to tool calls, tool results to context, memory reads and writes, delegated-agent messages, approval creation and consumption, and log export.

Ask adversarial questions that ordinary DFD reviews miss:

  • Can untrusted data become an instruction?
  • Can an observation forge another agent's identity or tool result?
  • Can a read tool encode data into an outbound URL or search term?
  • Can a model split a prohibited action across allowed tools?
  • Can memory make a one-request injection persistent?
  • Can the object change between approval and execution?
  • Can retries duplicate an external effect?
  • Can a compromised peer pass authority downstream?

Write threats as falsifiable scenarios

Use a stable identifier and an abuse narrative:

id: AGT-007
asset: reply_authority
precondition: attacker can write a customer ticket
action: ticket instructs the model to send data to attacker@example.net
boundary: untrusted ticket -> model -> send tool
impact: unauthorized external disclosure
controls:
  - send tool accepts an immutable approved_preview_id
  - approval binds recipient, subject, body hash, tenant, and expiry
  - tool rejects raw body supplied at commit time
verification:
  - indirect-injection integration test
  - approval-substitution race test
owner: messaging-platform
status: mitigated

Do not use “sanitize the prompt” as a complete mitigation. Specify where enforcement occurs and how a test proves it.

Prioritize without fake precision

Score impact using concrete harms: data class, financial ceiling, affected tenants, reversibility, visibility, and recovery time. Estimate exploitability from attacker access, automation, required knowledge, existing controls, and observability.

A numeric likelihood times impact score can help sort, but uncertain inputs do not become objective because they are multiplied. Record assumptions and confidence. Escalate any plausible path to credentials, cross-tenant access, or irreversible high-impact action even when frequency is unknown.

NIST's Generative AI Profile is voluntary risk-management guidance, not a security certification. It recommends adapting governance, measurement, oversight, and documentation to generative-AI risks (NIST AI 600-1).

Convert mitigations into engineering work

Every accepted mitigation needs:

  • a code or infrastructure owner;
  • an enforcement point outside model judgment;
  • a negative test or observable metric;
  • a failure behavior;
  • a rollout deadline; and
  • residual risk approved by the appropriate owner.

Prefer eliminating unnecessary power. Removing generic shell access is stronger than prompting the model not to misuse it. Replace broad credentials with narrow capabilities, isolate secrets, block arbitrary egress, separate preview and commit, and require transaction-bound approval.

Review triggers and completion checklist

Re-open the model when adding a tool, data source, memory type, model/provider, tenant, execution environment, protocol, or autonomous action. Also update it after an incident or failed security test.

Before sign-off, verify:

  • the deployed system matches the diagram;
  • every identity and tenant transition is explicit;
  • all model inputs have provenance and trust labels;
  • tool authority is smaller than the agent's general goal;
  • memory writes, delegation, and approvals are modeled;
  • high-impact effects have non-model enforcement;
  • each material threat has an owner and negative test; and
  • residual risks and assumptions are documented.