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".

• [2026.4.19] 🛠️ Update the evaluation on AgentDyn.
• [2026.1.30] 🛠️ Support the evaluation on more agents.
• [2026.1.30] 🛠️ Update the evaluation code on ASB.

We provide the evaluation of DRIFT, you can reproduce the results following:

conda create -n drift python=3.11
source activate drift
pip install "agentdojo==0.1.35"
pip install -r requirements.txt

We provide three API providers, including OpenAI, Google, and OpenRouter. Please set up the API Key as you need.

export OPENAI_API_KEY=your_key
export GOOGLE_API_KEY=your_key
export OPENROUTER_API_KEY=your_key

python pipeline_main.py \
--model gpt-4o-mini-2024-07-18 \
--build_constraints --injection_isolation --dynamic_validation
--suites banking,slack,travel,workspace

python pipeline_main.py \
--model gpt-4o-mini-2024-07-18 --do_attack \
--attack_type important_instructions \
--build_constraints --injection_isolation --dynamic_validation
--suites banking,slack,travel,workspace

You can evaluate any model from the supported providers by passing its model identifier (eg., gemini-2.5-pro) to the --model flag. To evaluate under an adaptive attack, include the --adaptive_attack configuration.

To evaluate on AgentDyn, you can directly replace the AgentDojo dependency with the AgentDyn version. First, run:

git clone [email protected]:SaFo-Lab/AgentDyn.git
cd AgentDyn
pip install -e .

Then, AgentDojo dependency has been replaced with the AgentDyn version, which additionally supports the shopping, github, and dailylife suites. You can use the same commands as for evaluating on AgentDojo to evaluate on these three suites, as shown below:

python pipeline_main.py \
--model gpt-4o-mini-2024-07-18 \
--build_constraints --injection_isolation --dynamic_validation
--suites shopping,github,dailylife

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

Please refer to ASB_DRIFT/README.md.

You can find the cached results in runs/.

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}
}