llamaindex-relayshield

LlamaIndex tools and a mandatory pre-execution gate for RelayShield's agentic-security endpoints — MCP server registry risk and AI-agent-sourced credential breach detection.

Install

pip install llamaindex-relayshield

Tools

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.llms.openai import OpenAI
from llamaindex_relayshield import check_mcp_server_risk_tool, check_prompt_injection_breach_tool

agent = FunctionAgent(
    tools=[check_mcp_server_risk_tool, check_prompt_injection_breach_tool],
    llm=OpenAI(model="gpt-4o-mini"),
)

result = await agent.run(
    "Is it safe to connect to the MCP server at https://mcp.example.com/sse? My RelayShield key is rs_live_..."
)

• check_mcp_server_risk — flags known-malicious IOC matches, typosquat domains, and newly-registered domains hosting an MCP server, before an agent connects to or installs it.
• check_prompt_injection_breach — checks whether an email appears in RelayShield's stolen-session corpus with a suspected-agentic-source marker (a session/token exposure that shows signs of having been captured via a compromised AI agent).

Both tools take api_key as a call argument rather than reading it from the environment implicitly — a shared agent process can act safely on behalf of multiple callers with different RelayShield keys.

Get a key at api.relayshield.net/developers.

Mandatory gate

Most "AI agent security" checks are optional — the agent can call them, but nothing stops it skipping the call and taking the risky action anyway. The RelayShield gate is the other kind: enforced before a protected action (connecting to or installing an MCP server) can happen at all.

LlamaIndex has no dedicated pre-execution hook API like LangChain's wrap_tool_call or the OpenAI Agents SDK's @tool_input_guardrail. What it does have is call_tool() — a workflow step every built-in agent (FunctionAgent, ReActAgent, CodeActAgent) inherits unmodified from BaseWorkflowAgent. This package subclasses it:

from llamaindex_relayshield import RelayShieldGatedFunctionAgent
from llamaindex_relayshield import check_mcp_server_risk_tool

agent = RelayShieldGatedFunctionAgent(
    tools=[connect_mcp_server_tool, check_mcp_server_risk_tool],
    llm=OpenAI(model="gpt-4o-mini"),
    # Names of tools this gate applies to -- everything else runs unmodified.
    protected_tools={"connect_mcp_server"},
)

RelayShieldGatedReActAgent and RelayShieldGatedCodeActAgent are the same pattern for the other two built-in agent types.

Properties, all non-negotiable by design:

• A hook exception defaults to defer (blocked, with an explanatory message), never silently to allow — a gate failure must not become a pass.
• Bounded retry applies only to transient upstream failures (timeout/429/5xx) — auth failures, malformed responses, and payment-required states are terminal after one attempt.
• The gate logs the decision, reason codes, check version, target, and timestamp — never keys, payment proofs, or session material.
• Only tool names listed in protected_tools are gated; everything else passes straight through to normal execution.

A note on how this is implemented, since it's less standard than the other two integrations: overriding call_tool() on a subclass only works because the override re-applies LlamaIndex's @step decorator. Verified directly against the installed package — LlamaIndex's step registry only recognizes methods carrying the _step_config attribute that @step sets at definition time; an override without it would silently fail to register as a step at all, breaking every tool call in the agent, not just skipping the gate.

Same normalized policy as langchain-relayshield's RelayShieldMCPGateMiddleware, openai-agents-relayshield's relayshield_mcp_gate, and the original standalone reference implementation, relayshield-langchain-gate — ported rather than imported, so this package has no dependency on LangChain/LangGraph or the OpenAI Agents SDK.

License

MIT