Este é o repositório de código do nosso artigo: Your Agentic LLMs Secretly Encode Latent Signals of Indirect Prompt-Injection Exposure.

Versão ArXiv e link do artigo: https://arxiv.org/abs/2608.02657

Este repositório implementa o pipeline de sondagem para sinais latentes de exposição a IPI: coleta de rastros do AgentDojo, rotulagem de risco de IPI, extração de características de estados ocultos e treinamento/avaliação de sondas de exposição a IPI.

Citação

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

Início Rápido

Comece pelas verificações compatíveis com CPU:

1. Clonar e Configurar

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. Executar Testes

uv run pytest

A suíte de testes é projetada para ser executada sem recursos de GPU. Ela verifica a serialização de mensagens, o registro de ataques de longo horizonte do AgentDojo, a lógica de rotulagem, a construção do conjunto de dados de sondas, as métricas de treinamento, o controle de tokens e o empacotamento do patch do vLLM.

3. Configurar um Endpoint de Modelo

Para a coleta, aponte o coletor para um servidor de modelo compatível com OpenAI:

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

Para extração de estados ocultos e treinamento de sondas, instale as dependências do modelo e a build do PyTorch que corresponda ao seu ambiente CUDA:

uv sync --extra dev --extra models

Visão Geral do Pipeline

O fluxo de trabalho principal é exposto por meio de ipi-signal-probe:

Exemplo:

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.

Os protocolos de rotulagem suportados são risk_faced, risk_visible e risk_actual.

Exemplo Reprodutível com Qwen3.5-0.8B

O diretório examples/qwen3_5_0_8b/ contém uma reprodução executável em modelo pequeno e um teste de usabilidade de código aberto para Qwen/Qwen3.5-0.8B.

Inicie oito servidores vLLM com uma GPU cada:

examples/qwen3_5_0_8b/serve_qwen3_5_0_8b_8x.sh start

Execute um teste de fumaça (smoke test):

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

Execute o exemplo completo:

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

O exemplo completo usa 8 servidores vLLM independentes com TP=1, 8 processos de coleta, 256 workers por processo/servidor, 8 partições de extração de características e partições de treinamento de sonda com uma GPU. Consulte examples/qwen3_5_0_8b/README.md para todas as variáveis de ambiente, modos de retomada, configurações de cache e comandos de status/parada.

Captura Direta via vLLM

O backend padrão de extração de características é transformers_hook. O backend de captura direta de estados ocultos via vLLM requer o overlay de vLLM 0.19.0 incluso no repositório:

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

Siga docs/vllm_patches_v0.19.0.md antes de usar:

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

Ataques de Longo Horizonte do AgentDojo

Este repositório inclui uma cópia (checkout) do AgentDojo corrigida em vendor/agentdojo. O AgentDojo upstream não expõe nativamente os ataques de longo horizonte exigidos pelo artigo; portanto, os comandos de coleta devem ser executados a partir deste checkout de origem, ou com vendor/agentdojo/src colocado à frente de qualquer AgentDojo upstream instalado no PYTHONPATH.

Estrutura de Diretórios

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

Configuração de Ambiente

Licença

Este projeto é licenciado sob Apache-2.0. Veja LICENSE.