DRIFT: Dynamic Rule-Based Defense with Injection Isolation for Securing LLM Agents

Hao Li, Xiaogeng Liu, Hung-Chun Chiu, Dianqi Li, Ning Zhang, Chaowei Xiao.

The official implementation of NeurIPS 2025 paper "DRIFT: Dynamic Rule-Based Defense with Injection Isolation for Securing LLM Agents".

How to Start

We provide the evaluation of DRIFT on GPT-4o-mini and GPT-4o, full code (including other models) will be released later, you can reproduce the results following:

Construct Your Environment

pip install "agentdojo==0.1.26"
pip install -r requirements.txt

Set Your OPENAI API KEY

export OPENAI_API_KEY=your_key

run task with no attack

python pipeline_main.py \
--model gpt-4o-mini-2024-07-18 \
--build_constraints --injection_isolation --dynamic_validation

run task under attack

python pipeline_main.py \
--model gpt-4o-mini-2024-07-18 --do_attack \
--attack_type important_instructions \
--build_constraints --injection_isolation --dynamic_validation

If you want to evaluate under adaptive attack, add configure of --adaptive_attack.

References

If you find this work useful in your research or applications, we appreciate that if you can kindly cite:

@articles{DRIFT,
  title={DRIFT: Dynamic Rule-Based Defense with Injection Isolation for Securing LLM Agents},
  author={Hao Li and Xiaogeng Liu and Hung-Chun Chiu and Dianqi Li and Ning Zhang and Chaowei Xiao},
  journal = {NeurIPS},
  year={2025}
}