Skip to main content

ToolShield: Training-Free Defense for Tool-Using AI Agents

Project description

ToolShield: Just One Command to Guard Your Coding Agent

PyPI Python License Homepage HuggingFace

Supports:  Claude Code Codex Cursor OpenHands OpenClaw


Quickstart | Pre-Generated Safety Experiences | Generate Your Own | Extend to New Tools | Safety Benchmark | Citation

ToolShield is a training-free, tool-agnostic defense for AI agents that use MCP tools. Just pip install toolshield and a single command guards your coding agent with safety experiences — no API keys, no sandbox setup, no fine-tuning. Reduces attack success rate by 30% on average.

Overview

Quickstart

pip install toolshield

Use Pre-generated Experiences

We ship safety experiences for 6 models across 5 tools, with plug-and-play support for 5 coding agents:

Claude Code Codex Cursor OpenHands OpenClaw

Inject them in one command — no need to know where files are installed:

# For Claude Code (filesystem example)
toolshield import --exp-file filesystem-mcp.json --agent claude_code

# For Codex (postgres example)
toolshield import --exp-file postgres-mcp.json --agent codex

# For OpenClaw (terminal example)
toolshield import --exp-file terminal-mcp.json --agent openclaw

# For Cursor (playwright example)
toolshield import --exp-file playwright-mcp.json --agent cursor

# For OpenHands (notion example)
toolshield import --exp-file notion-mcp.json --agent openhands

Use experiences from a different model with --model:

toolshield import --exp-file filesystem-mcp.json --model gpt-5.2 --agent claude_code

Or import all bundled experiences (all 5 tools) in one shot:

toolshield import --all --agent claude_code

You can also import multiple experience files individually:

toolshield import --exp-file filesystem-mcp.json --agent claude_code
toolshield import --exp-file terminal-mcp.json --agent claude_code
toolshield import --exp-file postgres-mcp.json --agent claude_code

See all available bundled experiences:

toolshield list

This appends safety guidelines to your agent's context file (~/.claude/CLAUDE.md, ~/.codex/AGENTS.md, ~/.openclaw/workspace/AGENTS.md, Cursor's global user rules, or ~/.openhands/microagents/toolshield.md). To remove them:

toolshield unload --agent claude_code

Available bundled experiences (run toolshield list to see all):

Model Filesystem PostgreSQL Terminal Chrome Notion
claude-sonnet-4.5
gpt-5.2
deepseek-v3.2
gemini-3-flash-preview
qwen3-coder-plus
seed-1.6

More plug-and-play experiences for additional tools coming soon (including Toolathlon support)! Have a tool you'd like covered? Open an issue.

Generate Your Own

Point ToolShield at any running MCP server to generate custom safety experiences:

export TOOLSHIELD_MODEL_NAME="anthropic/claude-sonnet-4.5"
export OPENROUTER_API_KEY="your-key"

# Full pipeline: inspect → generate safety tree → test → distill → inject
toolshield \
  --mcp_name postgres \
  --mcp_server http://localhost:9091 \
  --output_path output/postgres \
  --agent codex

Or generate without injecting (useful for review):

toolshield generate \
  --mcp_name postgres \
  --mcp_server http://localhost:9091 \
  --output_path output/postgres

Auto-discover Local MCP Servers

Automatically scan localhost for running MCP servers, run the full pipeline for each, and inject the results:

toolshield auto --agent codex

This scans ports 8000-10000 by default (configurable with --start-port / --end-port).

Extend to New Tools

ToolShield works with any MCP server that has an SSE endpoint:

toolshield generate \
  --mcp_name your_custom_tool \
  --mcp_server http://localhost:PORT \
  --output_path output/your_custom_tool

MT-AgentRisk Benchmark

We also release MT-AgentRisk, a benchmark of 365 harmful tasks across 5 MCP tools, transformed into multi-turn attack sequences. See agentrisk/README.md for full evaluation setup.

Quick evaluation:

# 1. Download benchmark tasks
git clone https://huggingface.co/datasets/CHATS-Lab/MT-AgentRisk
cp -r MT-AgentRisk/workspaces/* workspaces/

# 2. Run a single task (requires OpenHands setup — see agentrisk/README.md)
python agentrisk/run_eval.py \
  --task-path workspaces/terminal/multi_turn_tasks/multi-turn_root-remove \
  --agent-llm-config agent \
  --env-llm-config env \
  --outputs-path output/eval \
  --server-hostname localhost

Add --use-experience <path> to evaluate with ToolShield defense.

Repository Layout

ToolShield/
├── toolshield/              # pip-installable defense package
│   └── experiences/         # bundled safety experiences (6 models × 5 tools)
├── agentrisk/               # evaluation framework (see agentrisk/README.md)
├── workspaces/              # MT-AgentRisk task data (from HuggingFace)
├── docker/                  # Dockerfiles and compose
└── scripts/                 # experiment reproduction guides

Acknowledgments

We thank the authors of the following projects for their contributions:

Citation

@misc{li2026unsaferturnsbenchmarkingdefending,
      title={Unsafer in Many Turns: Benchmarking and Defending Multi-Turn Safety Risks in Tool-Using Agents},
      author={Xu Li and Simon Yu and Minzhou Pan and Yiyou Sun and Bo Li and Dawn Song and Xue Lin and Weiyan Shi},
      year={2026},
      eprint={2602.13379},
      archivePrefix={arXiv},
      primaryClass={cs.CR},
      url={https://arxiv.org/abs/2602.13379},
}

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

toolshield-0.1.1.tar.gz (13.2 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

toolshield-0.1.1-py3-none-any.whl (117.6 kB view details)

Uploaded Python 3

File details

Details for the file toolshield-0.1.1.tar.gz.

File metadata

  • Download URL: toolshield-0.1.1.tar.gz
  • Upload date:
  • Size: 13.2 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.12

File hashes

Hashes for toolshield-0.1.1.tar.gz
Algorithm Hash digest
SHA256 f0393dccbcc5f0f104c4fcfe31084a17c2d1b6da41fb30668022d7ac3d6ff576
MD5 49bf7486d7a97d4c5dcbacbf33040787
BLAKE2b-256 8cc09c33ecdd907fdb6178ae6ba081dc5efb9ce52ffc4bff22e5088413139351

See more details on using hashes here.

File details

Details for the file toolshield-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: toolshield-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 117.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.12

File hashes

Hashes for toolshield-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 6a5793a952d2c9309c6e6ea4a5b91678473641d5da425065ba96e07bf6df06eb
MD5 ca94f8821a345cfeb488a7c1faf54cbf
BLAKE2b-256 cd77fc4052a63786db0b369384a333d92923942fe9579770cb40ddf57544d526

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page