Skip to main content

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

Project description

ToolShield: A Package to Guard Your Agent

PyPI Paper Homepage License HuggingFace

Supports:  Claude Code Codex Cursor OpenHands OpenClaw


Quickstart | Pre-Generated Safety Experiences | Generate Your Own Safety Experiences |
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 6 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 (gmail example)
toolshield import --exp-file gmail-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 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/SOUL.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 Gmail
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 Safety Experiences

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 — both legacy SSE and modern Streamable HTTP transports:

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 × 6 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.2.tar.gz (9.8 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.2-py3-none-any.whl (121.8 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: toolshield-0.1.2.tar.gz
  • Upload date:
  • Size: 9.8 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.2.tar.gz
Algorithm Hash digest
SHA256 94fcc5d0e5842b957e8ed3cff9002533bd137c2e50f4d669f6608d71e9b850d9
MD5 b132a15ffaa8716c6cc305f6dd09b9ce
BLAKE2b-256 b7e25b8840d8ffdf3cc3f2d89fb70fb239bfd2ce7278d7d9d8791ab6eea7d953

See more details on using hashes here.

File details

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

File metadata

  • Download URL: toolshield-0.1.2-py3-none-any.whl
  • Upload date:
  • Size: 121.8 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.2-py3-none-any.whl
Algorithm Hash digest
SHA256 9c772472a71701d0ff9af9962a7bf90ea8b9ee257b1d6d00070988907c775892
MD5 b488927b9974a9112fbbe46b98072923
BLAKE2b-256 314a58cc524d38f99a738962c8eba97649783df428684ae633b6d81481f480e4

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