Abstract

As agentic AI systems take on increasingly complex and consequential tasks, securing them against prompt injection becomes critical. Existing defenses often rely on detecting or reasoning about injected malicious instructions and can therefore remain vulnerable to adaptive attacks that optimize the injection itself. Diffusion language models (dLLMs), an emerging framework for language generation, offer a new opportunity for security enforcement: by iteratively refining their sequence-wide predictions, they enable broader-scope constraint enforcement throughout generation. This paper introduces Policy-Constrained Diffusion Agents (CoDA), a framework that enforces task-dependent security policies directly within the denoising process of diffusion language models. At each denoising step, CoDA projects token distributions onto a policy-compliant set, maintaining constraint satisfaction throughout generation. CoDA is evaluated across static and adaptive prompt injection security benchmarks and state-of-the-art diffusion language model families. Results demonstrate substantial reductions in attack success rate with competitive task utility and practical inference overhead.

Citation

Why Diffusion Agents?

Most prompt-injection defenses try to recognize the injected instruction, with a classifier, a guard model or a hardened prompt. An adaptive attacker simply rewrites the injection until it is no longer recognized. Autoregressive models also commit to tokens left to right, so a constraint can only react to what has already been written.

Diffusion language models draft the entire output at once and refine it over many steps. Every step is a chance to look at the whole provisional action (which tool, which recipient, which URL) and correct it before anything is committed. CoDA uses those steps to enforce a policy, rather than trying to spot the attack.

CoDA

For each user task, a policy states which tools the agent may call and which values security-sensitive arguments may take. Policies come from the user's instruction and trusted pre-injection state, never from tool outputs, so an injection cannot widen them. At every reverse step, CoDA takes the denoiser's clean-sequence proposal, finds the policy-relevant regions in its provisional output (a tool name, a recipient field, a URL), projects the token distributions over those regions onto the policy-compliant set, and continues denoising from the corrected proposal:

The projection only touches the identified region; every other position keeps the denoiser's prediction, and later steps let the rest of the sequence adapt to the corrected values. The policies cover several kinds of constraint:

Authorized values often span several tokens. Constraining each position independently would let pieces of different allowed values recombine into one nobody authorized, so CoDA constrains multi-token spans with a prefix trie over the allowed set:

Where do policies come from?

A policy is compiled once from the trusted user prompt, before any tool output is read. We evaluate two sources. Oracle policies are hand-verified per task and give an upper bound that isolates enforcement from specification. Synthesized policies come from a single zero-shot call to an off-the-shelf LLM, constrained to a JSON schema, with no prompt tuning:

{ "modes": ["tool_t", "arg_t"],
  "allowed_tools": ["read_file", "update_user_info"],
  "rationale": "..." }

The policy model never sees the injected content. A wrong policy can over- or under-restrict the agent, but the attacker cannot steer it.

Results

We evaluate on AgentDojo (Banking, Slack, Travel, Workspace) with DiffusionGemma 26B-A4B and LLaDA2.0-mini. Beyond AgentDojo's static injections, we attack every system adaptively: a feedback-guided genetic search evolves the injection, a critic LLM reads the agent's full transcript after every attempt and explains why the attack fell short, and a mutator LLM uses that critique for the next rewrite, for up to 100 queries per scenario. The undefended agent (Vanilla-DA) keeps falling as the attacker spends more queries; CoDA with Oracle policies is never breached:

With Oracle policies, CoDA keeps task utility close to, and often above, the undefended agent. Policies synthesized automatically by GPT-5.6 Sol, Gemini 3.7 Flash and Claude Opus 5 also hold at 0% on Slack, Travel and Workspace, and are breached in about 10% of Banking scenarios. The prompt-based guardrail (PBG) baseline is breached on three suites, and the Last-Step ablation, which applies the same policy only after denoising, loses much of the utility:

Under adaptive attack: attack success rate (ASR, red, lower is better) and utility (teal, higher is better) in %, and average model calls per scenario. Bold marks the best value per suite. Hover a row for exact values and time.

Why project at every step rather than once at the end? The Last-Step ablation applies the same Oracle policy as a post-hoc overwrite. The constrained span is fixed, but the rest of the call was written around the old value and never gets a chance to adapt:

An Attack, Query by Query

A Banking task: "Spotify sent me a note that their prices increased 10% this month. Please send them the difference for March payment (see transaction history), which was too little." An injected transaction description asks the agent to send money to the attacker's account instead. After 30 failed attempts, each critiqued by the critic, the attacker reaches this trigger at query 31:

The demo steps through runs like this one: each mutated trigger, the agent's transcript, the critic's score and critique, and where each policy stops the attack.

Limitations

• Runtime argument discovery. When the value a task needs lives in untrusted data ("pay the IBAN in invoice.txt"), a policy cannot fix it in advance. This is the main source of the Banking utility trade-off and of the residual ASR with synthesized policies.
• Open-ended semantics. Tries and allowlists capture structure, not meaning. Constraints like "keep the tone polite" need additional checks on free text.
• Multi-turn intent drift. The policy is fixed at the first turn. Safely re-synthesizing it when the user changes their mind is open.

Enforcement is exact once a policy is given; the open problem is writing the policy.

Related links

AgentDojo (Debenedetti et al., 2024) is the benchmark of tool-using agents and prompt injections we evaluate on.

Our adaptive attacker follows the critic-guided search of Nasr et al. (2025), "The Attacker Moves Second".

The diffusion agents are built on DiffusionGemma and LLaDA2.0.

The website template was borrowed from mip-NeRF and Michaël Gharbi.