Create a threat model for an AI enabled full stack feature

An AI feature can answer correctly in a demo while remaining unsafe at the boundary where it retrieves private data or calls a tool. The model may encounter malicious instructions inside a document, produce an invalid action argument, or mix content from another tenant. Threat modeling should focus on those realistic paths rather than treating the model as either universally trustworthy or inherently impossible to use.
Consider an illustrative support assistant that retrieves organization documents, summarizes an answer, and can draft a ticket. It cannot send messages or change account permissions without a separate authorized workflow. The threat model should establish what the assistant can observe, what it can do, and how untrusted content is prevented from gaining authority.
Define assets and allowed outcomes
List the assets: organization documents, user identity, provider credentials, ticket data, tool permissions, and cost budget. State the intended user journey and acceptable outputs. An answer may cite retrieved evidence; a draft ticket may be reviewed before submission. These are different authorities and should not share one vague assistant capability.
Record what the system must never do, such as retrieve another tenant's documents or reveal a secret through a response. Keep the requirements concrete enough to test. A practical artifact is an asset and action inventory with owner, sensitivity, allowed reader, allowed writer, and consequence. This prevents the review from focusing only on prompt wording while overlooking the tools and data stores that create actual risk.
Follow inputs across trust boundaries
Trace user text, uploaded files, retrieved passages, conversation history, and provider output. Mark where each source becomes part of a model request and where model output enters application logic. Retrieved documents remain untrusted even when they come from an internal search index; they may contain copied instructions or attacker-controlled text.
Separate system instructions from evidence and label retrieved content with source identity. Do not let a document choose tool permissions, tenant scope, or approval requirements. Enforce those decisions outside the model. The threat model should include a passage that tells the assistant to reveal credentials or change a ticket destination. The expected outcome is rejection or harmless treatment as content, backed by tool constraints rather than faith in the prompt.
Constrain retrieval before generation
Apply authorization and tenant scope when querying the source, not after the model has already seen the results. Include permission dimensions in caches and indexes where required. A correct final answer filter cannot undo private data exposure to the model or provider.
Validate document ownership and derived artifacts such as chunks, embeddings, and search results. Define deletion and permission-change behavior for those copies. If a user loses access to a document, stale retrieval state must follow the product's revocation policy. Link tenant isolation for the broader data-path inventory. Retrieval relevance and retrieval authorization are separate conditions, and both must be satisfied before content enters the prompt.
Review tool authority as application code
Give each tool a narrow operation, validated arguments, and server-side authorization. A create ticket draft tool should not accept arbitrary database queries or file paths. Derive actor and tenant context from trusted session state, not from model-supplied fields. Bound result sizes and redact sensitive diagnostics.
Distinguish drafting from execution. If a consequential action requires review, present the concrete proposed result and preserve its identity through approval and execution. The model must not substitute a different action after approval. Use idempotency or conditional state where repeats can cause duplicate effects. An instruction asking the assistant to be careful is helpful context, but the tool implementation determines the authority it actually has.
Turn threats into abuse cases
For each boundary, write a concrete attack or failure and expected protection. Examples include a malicious retrieved instruction, a cross-tenant document identifier, an oversized prompt, an invented tool argument, and a replayed approval. Include accidental model errors as well as deliberate misuse.
Prioritize by consequence and plausibility. A secret-bearing tool output or unscoped search query deserves attention before cosmetic response wording. State existing mitigations and remaining gaps. Avoid claiming zero risk or complete protection from a short checklist. The abuse cases should guide implementation and verification: which request is rejected, which field is ignored, which content is never retrieved, and which action requires an authorized review.
Verify application controls and model behavior separately
Test deterministic controls directly. Cross-tenant retrieval should fail regardless of the model's response. Invalid tool arguments should be rejected by schema and semantic validation. Unapproved actions should be inaccessible through the tool interface. These tests establish application boundaries.
Then evaluate model behavior with a small synthetic set of malicious or confusing documents and unsupported questions. Check whether answers remain grounded and whether the assistant attempts prohibited actions. Model evaluations are useful but probabilistic; a passing set does not replace enforced permissions. Record model, prompt, tools, test inputs, and observed results. See building a small evaluation set for comparing behavior without relying on one impressive response.
Include privacy, cost, and recovery
Define which content is sent to the provider and what retention or processing policy applies, verifying relevant vendor facts when making commitments. Minimize private context and keep credentials outside prompts. Log safe operation metadata rather than complete conversations by default. Apply access and retention to evaluation artifacts too.
Bound input size, tool execution, retries, and provider spending. A tool loop can exhaust resources even without data theft. Plan recovery for wrong drafts, failed actions, and uncertain provider outcomes. Keep human escalation or a non-AI path for important unsupported tasks. The threat model should describe who investigates a failure and which durable state enables recovery, not merely list prevention controls.
Use AI as a constrained threat-model reviewer
Give an assistant the asset inventory, journey, trust boundaries, tool contracts, and synthetic abuse cases. Ask for missing paths and concrete counterexamples, with evidence references to implementation where available. Require it to distinguish an implemented control from a proposed mitigation. Do not provide real secrets or private documents to make the model more convincing.
Review its suggestions with the team and prioritize the smallest controls that address consequential risks. Repeat the review when tools gain authority, retrieval scope expands, or output starts driving actions. AI can help enumerate threats, but the final model is an engineering artifact tied to permissions, tests, and operational ownership. Its value is a clear map of what can go wrong and what verified boundary prevents or contains it.
Sources
OWASP threat modeling guidance supplies the primary modeling framework.
OWASP authorization guidance supports enforcing tool and retrieval authority outside model decisions. The support-assistant threat model is illustrative.
About the Author

Kamruz
Full Stack Developer
A passionate developer who loves creating modern web applications and sharing knowledge with the community.
More About Me