IPI Exposure Signal

This is the code repository for our paper: Your Agentic LLMs Secretly Encode Latent Signals of Indirect Prompt-Injection Exposure.

ArXiv version and paper link: https://arxiv.org/abs/2608.02657

This repository implements the probing pipeline for latent IPI-exposure signals: AgentDojo trace collection, IPI risk labeling, hidden-state featurization, and IPI-exposure probe training/evaluation.

Citation

@misc{dong2026agenticllmssecretlyencode,
  title={Your Agentic LLMs Secretly Encode Latent Signals of Indirect Prompt-Injection Exposure},
  author={Jianshuo Dong and Yiming Liu and Maosen Zhang and Nan Deng and Xu Peng and Xiaoping Zhang and Tianwei Zhang and Jie Zhang and Han Qiu},
  year={2026},
  eprint={2608.02657},
  archivePrefix={arXiv},
  primaryClass={cs.CR},
  url={https://arxiv.org/abs/2608.02657},
}

Quick Start

Get started with the CPU-friendly checks first:

1. Clone and Setup

git clone https://github.com/jianshuod/IPI-exposure-signal.git
cd IPI-exposure-signal

# Install uv if needed.
curl -LsSf https://astral.sh/uv/install.sh | sh

# This repository targets Python 3.12.
uv python install 3.12
uv sync --extra dev

2. Run Tests

uv run pytest

The test suite is designed to run without GPU resources. It checks core message serialization, AgentDojo long-horizon attack registration, labeling logic, probe dataset construction, training metrics, token accounting, and vLLM patch packaging.

3. Configure a Model Endpoint

For collection, point the collector at an OpenAI-compatible model server:

export IPI_AWARE_UPSTREAM_BASE_URL=http://127.0.0.1:8000/v1
export IPI_AWARE_UPSTREAM_API_KEY=EMPTY

For hidden-state extraction and probe training, install model dependencies and the PyTorch build that matches your CUDA environment:

uv sync --extra dev --extra models

Pipeline Overview

The main workflow is exposed through ipi-signal-probe:

Example:

uv run ipi-signal-probe collect \
  --suite workspace \
  --attack direct \
  --injected \
  --max-cases 4 \
  --results-dir results/probe_traces/v2

uv run ipi-signal-probe label \
  --root results/probe_traces/v2/<run-name> \
  --labeling-protocol risk_faced

uv run ipi-signal-probe featurize \
  --root results/probe_traces/v2/<run-name> \
  --model Qwen/Qwen3-8B \
  --backend transformers_hook \
  --model-family qwen3 \
  --selected-position -1

uv run ipi-signal-probe partition \
  --root results/probe_traces/v2/<run-name> \
  --dataset broad

# Training consumes a JSON config. The Qwen example below generates the
# train/eval configs used for its full run.

Supported labeling protocols are risk_faced, risk_visible, and risk_actual.

Reproducible Qwen3.5-0.8B Example

The examples/qwen3_5_0_8b/ directory contains a runnable small-model reproduction and open-source usability test for Qwen/Qwen3.5-0.8B.

Start eight one-GPU vLLM servers:

examples/qwen3_5_0_8b/serve_qwen3_5_0_8b_8x.sh start

Run a smoke test:

QWEN3_5_0_8B_MODE=smoke \
examples/qwen3_5_0_8b/run_probe_pipeline.sh start

Run the full example:

QWEN3_5_0_8B_MODE=full \
examples/qwen3_5_0_8b/run_probe_pipeline.sh start

The full example uses 8 independent TP=1 vLLM servers, 8 collector processes, 256 workers per process/server, 8 featurization shards, and one-GPU probe training shards. See examples/qwen3_5_0_8b/README.md for all environment variables, resume modes, cache settings, and status/stop commands.

vLLM Direct Capture

The default featurization backend is transformers_hook. The direct vLLM hidden-state capture backend requires the checked vLLM 0.19.0 overlay:

uv sync --frozen --extra dev --extra vllm
uv run python scripts/apply_vllm_ipi_aware_patch.py --check

Follow docs/vllm_patches_v0.19.0.md before using:

uv run ipi-signal-probe featurize \
  --root results/probe_traces/v2/<run-name> \
  --model Qwen/Qwen3-8B \
  --backend vllm_direct \
  --model-family qwen3 \
  --selected-position -1

Long-Horizon AgentDojo Attacks

This repository vendors a patched AgentDojo checkout under vendor/agentdojo. Upstream AgentDojo does not natively expose the long-horizon attacks required by the paper, so collection commands should be run from this source checkout, or with vendor/agentdojo/src placed ahead of any installed upstream AgentDojo on PYTHONPATH.

Directory Structure

IPI-exposure-signal/
├── ipi_aware/
│   ├── data_collection/      # AgentDojo traces and decision points
│   ├── message_content.py    # Chat-message normalization used by collection/features
│   └── probes/               # labels, features, partitions, training, eval
├── examples/
│   └── qwen3_5_0_8b/         # 8-GPU Qwen/Qwen3.5-0.8B reproduction scripts
├── patches/                  # vLLM 0.19.0 hidden-state capture overlay
├── vendor/agentdojo/         # patched AgentDojo source for long-horizon attacks
├── docs/                     # vLLM patch notes
├── scripts/                  # utility scripts
└── tests/                    # CPU-friendly regression tests

Environment Configuration

License

This project is licensed under Apache-2.0. See LICENSE.

`collect`	Run AgentDojo tasks and save traces plus decision points
`label`	Assign IPI risk labels to collected decision points
`featurize`	Extract model hidden states for probe examples
`partition`	Build train/validation/evaluation splits
`train`	Train lightweight probes over hidden-state features
`eval`	Evaluate trained probes on configured splits
`eval-groups`	Run grouped evaluation for labeling protocols and held-out settings

`IPI_AWARE_UPSTREAM_BASE_URL`	OpenAI-compatible model endpoint used by collection
`IPI_AWARE_UPSTREAM_API_KEY`	API key for the upstream endpoint, or `EMPTY` for local vLLM
`PYTHONPATH`	Put `vendor/agentdojo/src` first when using the patched AgentDojo source outside this checkout
`QWEN3_5_0_8B_MODEL`	Override the default Qwen/Qwen3.5-0.8B checkpoint path in the example scripts
`QWEN3_5_0_8B_CACHE_DIR`	Redirect model and local caches to a data filesystem
`QWEN3_5_0_8B_FEAT_BACKEND`	Set to `vllm_direct` after applying the vLLM overlay

IPI-exposure-signal

Most Popular