How to Threat Model an AI Agent: Application, RAG, Memory, Tools, and Identity
A step-by-step method for threat modeling production AI agents across application, model, RAG, memory, tools, MCP, identity, approvals, and infrastructure boundaries.
Threat modelling an AI agent begins with a familiar question: what are we trying to protect, and from whom? The challenge is that an agent's behaviour is assembled at runtime from several systems. User input, model instructions, retrieved data, memory, tools, identities, and external services all influence what happens next.
A diagram containing one box labelled “LLM” misses most of the security boundary. A useful threat model follows data and authority from the user interface to the final system effect.
This guide provides a workshop method that can be completed before launch and updated as the agent evolves. It complements the OWASP Top 10 for Agentic Applications but stays grounded in a particular implementation.
Step 1: define the business outcome
Write down what the agent is meant to accomplish and what it must never accomplish. Avoid feature language such as “uses RAG” or “has five tools.” Describe outcomes.
For a support agent, legitimate outcomes may include summarising the current customer's ticket and proposing an approved refund. Unacceptable outcomes may include accessing another customer, issuing a refund above a threshold, changing account identity, or sending data externally.
These statements give the threat model an anchor. A vulnerability matters because it enables an unacceptable outcome, not because a payload looks clever.
Step 2: identify actors and trust levels
List customer users, employees, administrators, developers, service accounts, external content authors, tool operators, model providers, and other agents. Include compromised and malicious versions of otherwise legitimate actors.
Record how identity is established and where it can be lost. Does a user ID travel from the application into retrieval and tool calls? Does an internal agent receive authenticated delegation or only a text message saying who requested the task?
Natural-language claims are not identity. Treat “the CFO approved this” as untrusted unless an authenticated system confirms it.
Step 3: draw the implementation layers
At minimum, include:
- User interface and application
- API and authentication layer
- System/developer instructions and model
- Retrieval pipeline and data sources
- Session state and persistent memory
- Tools, plugins, and MCP servers
- Service identities and permissions
- External systems and recipients
- Logging, evaluation, and human review
Show data flows and control flows separately where possible. A document flowing into context is different from authority flowing through a credential.
Step 4: mark trust boundaries
Trust boundaries occur when content or authority changes domain: browser to application, tenant to shared service, external document to model context, model decision to tool execution, one agent to another, or application to model provider.
For each boundary ask:
- Is the sender authenticated?
- Is the content trusted as data, instruction, or neither?
- Which identity authorises the next step?
- What validation occurs?
- What is logged?
- What maximum consequence is reachable?
The model context is a particularly important boundary because trusted instructions and untrusted content may be represented in the same medium.
Step 5: inventory assets
Assets include more than data. Record customer records, credentials, internal documents, prompts, code, model access, tool authority, transaction limits, audit evidence, memory integrity, service availability, and user trust.
Classify confidentiality, integrity, and availability requirements. An attacker may gain value by changing a remembered payment address even if no secret is disclosed. A loop that consumes the agent budget may be an availability incident.
Step 6: map every external-content source
List chat messages, uploads, websites, emails, tickets, repositories, documents, RAG records, tool results, MCP metadata, and inter-agent messages. Determine who can influence each source and whether provenance survives into context.
Assume some untrusted text will resemble instructions. Threats include direct injection, indirect prompt injection, malicious examples, false authority, poisoned memory, and content crafted to manipulate tool selection.
Step 7: analyse RAG separately
Map ingestion, parsing, chunking, metadata, embeddings, index namespaces, filters, reranking, caches, and citations. Ask whether source permissions are enforced at retrieval time and whether the generator receives more data than the user can access.
Threats include poisoned sources, cross-tenant retrieval, metadata leakage, hidden instruction retrieval, stale permissions, and summaries that remove provenance. Use the RAG leakage test guide to convert these into controlled cases.
Step 8: analyse memory and session state
Separate temporary context from persistent memory. Record who can write, read, update, and delete each memory type. Determine whether values are user-provided, model-inferred, or imported from tools.
Threats include cross-user leakage, session fixation, stale privilege, persistent instruction attacks, false facts, and low-trust agents influencing high-trust workflows. The memory-poisoning guide provides delayed and cross-session test patterns.
Step 9: analyse tools and MCP
For each tool, document name, server, description, parameters, credentials, data reach, side effects, network reach, and approvals. Review capabilities in combination, not only one at a time.
Threats include unsafe arguments, broken object authorisation, command or URL injection, excessive agency, tool poisoning, server impersonation, result injection, rug pulls, and uncontrolled retries. Treat MCP tool metadata as a trust boundary.
Step 10: trace identity and authority
Create an authority chain from the user to every backend action. Record whether each request executes as the user, a delegated identity, or a shared service account. Identify where tenant, role, and task scope are bound.
Ask what happens after logout, role change, credential revocation, agent cloning, or task completion. Verify that an agent cannot retain stale authority through a long session.
Step 11: model human approvals
Place each approval on the diagram and specify what the reviewer sees. Threats include fabricated approval claims, misleading summaries, parameter changes after approval, approval reuse, fatigue, and social engineering.
Sensitive approval should be tied to the exact action and trusted system state. The agent should not control both the proposed action and the only evidence used to approve it.
Step 12: consider failure and abuse conditions
Threat models often assume every dependency responds correctly. Add timeouts, partial results, malformed data, duplicates, conflicting tools, model refusal, provider outage, and budget exhaustion.
An attacker may deliberately trigger retries or ambiguity. Controls include bounded loops, idempotency, circuit breakers, transaction limits, and safe fallback rather than autonomous improvisation.
Step 13: build concrete attack paths
Combine threats into end-to-end hypotheses:
Poisoned document
→ retrieved into agent context
→ false authority accepted
→ over-scoped lookup tool called
→ another tenant's canary returned
→ response or memory exposes it
For each path, note assumptions and controls. Do not label a chain proven until it has been safely replayed. A threat model identifies possibilities; validation determines observable behaviour.
Step 14: prioritise by reachable consequence
Use likelihood and impact, but include control strength and observability. A prompt injection against a public FAQ bot is different from the same weakness in an agent with payment tools.
Prioritise paths that combine untrusted content, sensitive data, broad credentials, autonomous action, weak approval, and poor logging. Also consider blast radius: one user, one tenant, or the entire organisation.
Step 15: turn threats into test cases
Define a harmless success criterion for each high-priority path. Use synthetic tenants, canary records, mock tools, controlled recipients, and reversible actions. Specify scope, maximum turns, rate, stop conditions, and required evidence.
A test case should state expected boundary, setup, adversarial sequence, observable outcome, and remediation retest. Baseline checks can run frequently; deeper multi-turn campaigns should receive appropriate review.
Step 16: connect controls to owners
Assign each mitigation to the layer that can enforce it. Prompt changes belong to AI engineering; object authorisation to the API; tenant filters to retrieval; token scope to identity; recipient restrictions to the messaging service; approval design to product and security.
Avoid the “prompt fix” reflex. If a tool can access every tenant, the durable fix is not a sentence asking the model to choose carefully.
Step 17: preserve evidence and uncertainty
Record transcripts, retrieved sources, tool calls, arguments, identities, policy decisions, results, and configuration versions. Separate confirmed facts, likely findings, inconclusive runs, and threat hypotheses.
Agent behaviour can vary. Reproduce important outcomes and keep analyst review for consequential results. Evidence makes the threat model operational rather than theoretical.
Step 18: update the model continuously
Revisit the threat model when the agent gains a tool, model, memory feature, data source, permission, integration, user group, or deployment environment. Remote MCP metadata and provider behaviour can change without an application release.
Turn confirmed findings into regression tests. Threat modelling should feed security validation, and validation should correct the threat model.
A 90-minute workshop agenda
Use the first 15 minutes to define purpose and unacceptable outcomes. Spend 20 minutes drawing data and authority flows, 20 on RAG/memory/tools, 15 on identities and approvals, 10 creating attack paths, and 10 assigning controls and test owners.
Invite application engineering, AI engineering, security, identity, and the business owner. The most valuable discoveries often occur when two teams realise they made different assumptions about the same boundary.
Oxyne uses complete-system context to guide agentic AI security validation through supported external interfaces. It does not replace architecture review or claim direct visibility into every internal system. The combination is stronger: threat modelling identifies important paths, and controlled testing determines which paths produce observable failures.
A good agent threat model does not attempt to predict every sentence a model might generate. It limits what manipulated reasoning can reach, identifies where authority changes hands, and creates evidence-based tests for the outcomes the business cannot accept.
Example threat register entries
A useful register is specific enough to assign and test. Instead of “prompt injection,” write: “A supplier-controlled webpage can influence the procurement agent to send an internal pricing file to an unapproved domain because the upload tool accepts arbitrary recipients.” Record the affected assets, existing controls, test method, owner, and residual risk.
Instead of “data leakage,” write: “The retrieval service uses a shared index and relies on a model-proposed tenant filter; a manipulated conversation may select another tenant's document metadata.” This wording tells engineering where enforcement should live.
Other strong entries include stale session authority after role change, memory written by a low-trust agent and read by a privileged workflow, tool definition changes outside release control, and repeated side effects after timeout.
Applying STRIDE without forcing the fit
Traditional categories still help. Spoofing applies to users, servers, tools, and agents. Tampering applies to prompts, memory, documents, metadata, and results. Repudiation appears when tool actions lack traceable identity. Information disclosure includes context, retrieval, logs, and model output. Denial of service includes loops and budget exhaustion. Elevation of privilege appears through over-scoped identities and unsafe delegation.
Use STRIDE to broaden coverage, then add agent-specific questions about goal manipulation, instruction/data confusion, tool composition, persistent context, and human trust. Frameworks organise thinking; the architecture determines the real threats.
Reviewing third-party agents
When source and internal architecture are unavailable, build a black-box threat model from documented interfaces, permissions, provider statements, observed network/data flows, and controlled tests. Clearly label assumptions and blind spots.
Ask vendors for data retention, model providers, subprocessors, tool scope, identity binding, logging, incident response, and change notification. Verify observable boundaries with synthetic accounts where the contract permits.
Do not translate lack of visibility into low risk. It is an uncertainty that may require stronger contractual, architectural, or deployment controls.
Threat-model completion criteria
The document is ready for a release decision when owners, unacceptable outcomes, data and authority flows, trust boundaries, high-priority attack paths, controls, test cases, residual risks, and update triggers are recorded. Every critical path should have either evidence of an effective control or an explicit acceptance by the accountable owner.
Keep the diagram small enough to understand and link detailed inventories separately. A beautifully complex diagram that nobody updates is less useful than a clear model used during every material change.