Threat Modeling AI Systems

Why threat-model an AI system

Threat modeling is a structured way of answering one question before you write the code: if someone wanted to make this system misbehave, where would they push, and what would it cost them? You do it at design time, on a whiteboard or in a document, working from the architecture rather than the running system. That timing is the whole point. The cheapest place to close an attack path is the diagram, before it hardens into an integration you cannot easily unwind and a data flow other teams already depend on.

It is a proactive discipline in a field that leans heavily on reactive ones. Most AI assurance work observes the built system: evals sample expected behavior, red-teaming attacks the deployed thing, monitoring watches production. All of those are necessary and none of them tell you, at design time, that giving an agent both a retrieval tool over untrusted documents and an outbound email tool creates an exfiltration path. Threat modeling reasons about the shape of the system, so it catches the class of problem that only exists because of how the pieces were wired together.

It does not replace red-teaming; it aims it. A threat model is a hypothesis about how the system could be attacked, and red-teaming is the experiment that tests whether the hypothesis holds. Skip the modeling and your red team spends its scarce, expensive attention rediscovering the obvious exposures instead of finding the subtle ones. The two are a matched pair: reason about the attack surface first, then go try to break the places the model flagged.

Threat modeling versus red-teaming

Teams conflate these two constantly, usually by assuming that a red-team engagement makes threat modeling redundant. It does not, and the distinction is worth stating precisely, because they operate at different times, on different artifacts, and produce different kinds of confidence. Threat modeling is analytical and happens on the design; red-teaming is empirical and happens on the build.

The blind spots are complementary, which is the argument for running both. Threat modeling misses what you did not think of, because it is bounded by imagination. Red-teaming misses what the testers did not get to, because it is bounded by time and coverage. Run only the model and you have a tidy list of hypothetical risks with no evidence any of them are real. Run only the red team and you get a scatter of confirmed exploits with no map telling you which regions of the system went unexamined.

In a mature program the two feed each other in a tight cycle. The threat model produces a ranked list of what to attack, so the red team spends its effort on the highest-risk paths instead of improvising. The red team then confirms, refutes, or reprioritizes each entry, and those results feed straight back into the next iteration of the model. Threat modeling gives red-teaming direction; red-teaming gives threat modeling ground truth.

Map the system and its boundaries

You cannot enumerate threats against a system you have not drawn. The first move is a data-flow map: the components, the data that moves between them, and the boundaries that data crosses. For an LLM system the components worth naming explicitly are the model, the retrieval layer and its knowledge sources, the tools the system can call, and any persistent memory or conversation state. Draw how a request flows through them and where each piece of data originates.

The step that matters most is marking the trust boundaries, the lines where data crosses from something you control into something you do not, or the reverse. On a classic web app those boundaries are familiar: the network edge, the auth layer, the database. An AI system adds boundaries that appsec never had to reason about, and the single most important analytical move is to mark every place untrusted input can enter the model's context. There are three you must always account for.

            [ USER ]                       (untrusted #1)
               |  prompt
               v
        +--------------+     retrieved docs   +-----------+
        |   MODEL /    |<---------------------|  RETRIEVAL |
        |  ORCHESTRATOR|   (untrusted #2)     |  + sources |
        +--------------+                      +-----------+
           |        ^
      tool |        | tool results
      call |        | (untrusted #3)
           v        |
        +----------------+        === trust boundary ===
        |     TOOLS      |  --> external systems, data,
        | (email, DB,    |      side effects the model
        |  web, code)    |      can trigger
        +----------------+
            

The reason this matters so much is that the model treats all three sources as text in the same context window, with no inherent notion of which tokens are trustworthy. A retrieved document and a tool result carry exactly the same authority as the system prompt once they land in context. So a boundary is any point where content the attacker might influence becomes model input, and in an agentic system that includes web pages it fetches, records it queries, and the output of one tool feeding the input of the next.

Enumerate the threats

With the map drawn, you walk it methodically and ask, at each component and each boundary, what could go wrong. The value of a framework here is that it forces coverage: rather than brainstorming attacks freeform and stopping when you feel done, you apply a fixed lens to every element so the gaps become visible. Classic STRIDE (spoofing, tampering, repudiation, information disclosure, denial of service, elevation of privilege) still applies to the conventional parts of the system, and you should run it there.

For the AI-specific surface, extend the exercise with the threat classes that STRIDE was never designed to capture. Public references such as the OWASP LLM Top 10 and MITRE ATLAS catalog these at length and are useful as shared checklists; treat them as living inventories to map against rather than reproducing their exact contents here. The classes that recur across nearly every enterprise LLM system are worth naming directly.

Enumeration is a walk, not a list you memorize. Take each boundary from the map and ask which of these classes apply to it. The retrieval layer invites poisoning and indirect injection; the tool layer invites insecure use and excessive agency; the model artifact itself invites supply-chain risk. The output is not prose but a table of concrete threats, each tied to a specific component and boundary, which is exactly the input the next step needs.

What classic appsec misses

A seasoned application-security team can threat-model the conventional parts of an AI system on day one: the APIs, the auth, the datastores, the network edge. What trips them up is that three properties of LLM systems break assumptions their entire discipline rests on, and each one opens threats that traditional modeling has no square for.

The practical consequence is that the AI-specific threat classes from the previous section are not optional add-ons for the paranoid. They are the direct result of these three properties, and a threat model that stops at the conventional surface has examined the plumbing while ignoring the part of the system that actually reasons and acts. Extend the framework, or you have modeled the least novel and least dangerous half of the design.

From threats to mitigations

A threat model that ends at a list of threats is a worry log. The output that matters is a prioritized set of mitigations, each wired into the design and each traceable to the threat it addresses. Prioritize by the familiar product of likelihood and impact, then work down the list, because you will not close everything and you should spend your first, best effort on the paths that are both reachable and costly.

The mitigation is not generic hardening; it is chosen to match the threat class. Each of the AI-specific classes has a natural home in the architecture, and the discipline is to route each finding to the layer that actually neutralizes it rather than reaching for one favorite control.

1. Prompt injection. Constrain and label untrusted content, keep it out of privileged instruction slots, and put a mediation layer between the model and any consequential action. The full treatment lives in prompt-injection defense.
2. Insecure tool use and excessive agency. Least privilege for tools: narrow scopes, validated arguments, human approval on high-impact actions, and sandboxed execution. See guardrails and sandboxing and tool calling.
3. Exfiltration. Control outbound channels, filter what leaves, and apply data-minimization so sensitive content is not in reach of a compromised path in the first place.
4. Poisoning and supply-chain. Provenance checks: verify model, adapter, and data sources; pin and scan dependencies; and treat the knowledge base as an asset whose integrity you monitor. This ties into model risk.

Two rules keep this honest. First, every mitigation must land somewhere in the architecture with an owner, not float as a recommendation, because a control no one built is not a control. Second, a residual-risk decision is a legitimate output: some threats you accept, with the acceptance recorded and signed rather than left implicit. What the process must never produce is a threat that was enumerated, never mitigated, and never consciously accepted.

The architect view

The failure mode at the organizational level is treating threat modeling as a document produced once, filed, and never reopened. The system it described stops existing the moment you add a tool, swap a model, or point retrieval at a new source. Threat modeling is not a deliverable; it is a habit that runs early and then iterates, triggered by material changes to the architecture the same way a code change triggers a review.

The single most useful principle to instill in a team is this: treat every untrusted input source as a trust boundary, and re-ask the enumeration questions whenever a new one appears. Most AI security incidents trace back to a boundary that was never marked as one, usually a retrieval source or a tool result that someone assumed was benign because it came from inside the system. The map is only as good as its honesty about where attacker-influenced data can reach the model.

The closing note for a skeptical CTO is that this is the cheapest security work in the entire AI lifecycle. It is a few hours of structured argument over a diagram, and it repays that time by catching the architectural exposures that are nearly impossible to retrofit once the system is built and depended upon. Model the system early, mark every untrusted boundary honestly, feed the result into red-teaming and guardrail design, and iterate every time the architecture moves. Do that, and the assurance your governance program has to attest to becomes something you designed in rather than something you hope survived.