Threat-model an AI application from data flows

A list of AI risks becomes actionable only when it is tied to the application’s assets, identities, boundaries and operations. Draw the real data flows before selecting mitigations.

Map each trust crossing

The LLMOps and security lab models four fixture flows from user prompt through retrieval, generation and an approved billing write.

from llmops_cases import threat_model_case

result = threat_model_case()
assert len(result["flows"]) == 4
assert result["all_crossings_mapped"] is True
assert result["threats"]["documents"] == "indirect injection"
assert result["controls"]["arguments"] == "schema and authorization"
assert result["residual_risk_review"] == "pending"

The mapping is a teaching record. “Pending” matters: attaching a control name does not prove that the control is effective or accepted.

Draw identities and stores

Include the browser or calling service, application, policy layer, retriever, model provider, tool gateway, external systems, caches, memory, logs and evaluation store. Label data class, direction, protocol, encryption boundary, tenant and retention. Distinguish user identity, application identity and delegated authority.

Mark assets such as personal data, system instructions, credentials, proprietary documents, approval records and write authority. Then ask what happens if each input is hostile, stale or belongs to another tenant. AI-specific paths include direct and indirect prompt injection, poisoned retrieval, sensitive output, improper output handling, excessive agency and unbounded consumption. Conventional threats such as broken access control, SSRF and dependency compromise still apply.

Turn threats into tests

For every material threat, assign prevention, detection, response and an owner. “Use a guardrail” is too vague. “Reject a tool call when tenant ID differs from authenticated context, before data access” produces a test. Link each test to evidence, release gates and a residual-risk decision.

Update the model when tools, providers, data sources, permissions or deployment boundaries change. A diagram without versions and review date becomes historical decoration.

The Generative & Agentic AI course uses threat models to connect prompts with application security and operations.

Reproducible exercise

Diagram one read-only and one write path. Add five attack paths, write a deterministic test for each boundary, and document one residual risk that needs an accountable human decision.

Continue learning

This article is part of the LLMOps, security and operational evaluation sequence. Use the neighbouring tasks when you need the prerequisite or the next application.

• Review the prerequisite or neighbouring task in Build a redaction check for generated output.
• Continue with Test an AI system when its upstream service is unavailable.

References: OWASP Top 10 for LLM and GenAI Applications and NIST AI RMF Generative AI Profile.