toolshield 0.1.0

Training-free defense for multi-turn safety risks in tool-using AI agents

ToolShield: Training-Free Defense for Multi-Turn Safety Risks in Tool-Using AI Agents

Supports: 

---

Quickstart | Pre-generated | Generate Your Own | Extend | Benchmark | Citation

ToolShield is a training-free, tool-agnostic defense for AI agents that use MCP tools. It works by letting the agent explore tool functionality in a sandbox, learning from its own executions, and distilling safety guidelines before deployment. Reduces attack success rate by 30% on average — with zero fine-tuning.

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. Inject them in one command:

# For Claude Code
toolshield import \
  --exp-file experiences/claude-sonnet-4.5/experience_list_claude-sonnet-4.5_postgres.json \
  --agent claude_code

# For Codex
toolshield import \
  --exp-file experiences/claude-sonnet-4.5/experience_list_claude-sonnet-4.5_postgres.json \
  --agent codex

# For OpenClaw
toolshield import \
  --exp-file experiences/claude-sonnet-4.5/experience_list_claude-sonnet-4.5_postgres.json \
  --agent openclaw

# For Cursor (writes to global user rules via SQLite)
toolshield import \
  --exp-file experiences/claude-sonnet-4.5/experience_list_claude-sonnet-4.5_postgres.json \
  --agent cursor

# For OpenHands (creates a microagent)
toolshield import \
  --exp-file experiences/claude-sonnet-4.5/experience_list_claude-sonnet-4.5_postgres.json \
  --agent openhands

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 experiences in experiences/:

Model	📁 Filesystem	🐘 PostgreSQL	💻 Terminal	🎭 Playwright	📝 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

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
├── agentrisk/           # evaluation framework (see agentrisk/README.md)
├── experiences/         # pre-generated safety experiences (6 models × 5 tools)
├── 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:

• OpenAgentSafety
• SafeArena
• MCPMark

Citation

License

MIT