Key Takeaways from This Guide
- Security Belongs in Code, Not Just Prompts: While a good system prompt is important, a robust security model must be enforced at the architectural level using middleware and hardened tools.
- The “Wrap Tool Call” Hook is Your Best Friend: For stateful access control, a class-based middleware that implements the
wrap_tool_callmethod is the correct, modern pattern for intercepting and controlling tool execution. - Leverage Built-in Middleware: Don’t reinvent the wheel. Use LangChain’s pre-built
PIIMiddlewarefor powerful, configurable detection and redaction of sensitive data. - Defense-in-Depth is Crucial: A multi-layered approach is key. Our solution combines a centralized RBAC firewall, PII sanitizers, and a final security check inside the tool itself for maximum security.
- State Management is Not Optional: Stateful agents require a
checkpointerand a clear state schema (like ourFirewallState) to correctly manage context and interruptions.
The era of Agentic AI is here. We’re moving beyond simple chatbots to building autonomous systems that can interact with the world through tools—APIs, databases, file systems, and even command-line shells. While this leap in capability is transformative, it opens a Pandora’s box of security vulnerabilities that can no longer be ignored.
Giving an Agentic AI a tool is like giving a new employee a password. How do you ensure they only access what they’re supposed to? How do you prevent them from being tricked into deleting critical data or leaking sensitive customer information?
This guide provides a comprehensive, production-ready blueprint for building an Agentic AI Firewall. We’ll architect a security layer that sits between your agent and its tools, enforcing permissions, sanitizing data, and providing a crucial line of defense for any serious Agentic AI application. This isn’t just about theory; it’s a practical, hands-on guide with code that you can run today.
The Business Case: When “Helpful” Agentic AI Becomes a Million-Dollar Liability
For managers and founders, the promise of Agentic AI is immense: automated workflows, hyper-personalized customer support, and data analysis at unprecedented scale. But the risks are just as significant and can translate into real, substantial costs.
Consider these scenarios:
- The Over-Privileged Support Agent: You deploy an AI support agent with access to your customer database. A malicious actor uses a clever prompt injection attack, tricking the agent into running a query like
SELECT * FROM users;. The result? A catastrophic data breach that leaks every customer’s PII, leading to millions in regulatory fines (like GDPR), loss of customer trust, and brand damage that takes years to repair. - The Naive DevOps Assistant: You create an Agentic AI to help your engineering team manage infrastructure. An engineer, trying to be clever, asks it to “clean up temporary files.” The agent, lacking proper guardrails, interprets this as a command to run
rm -rf /tmp, but a small mistake in the logic makes it runrm -rf /instead. The result? Your entire production server is wiped, causing hours or days of downtime, lost revenue, and frantic, costly recovery efforts. - The Accidental Data Leak: An internal AI analyst, tasked with creating a financial report, is given access to a tool that reads sales data. The tool’s output inadvertently includes raw customer emails and internal API keys. The agent, focused only on the numbers, includes this sensitive information in its final report, which is then shared widely within the company. The result? Your internal API keys are now exposed, creating a massive security hole, and customer PII has been improperly shared, violating data handling policies.
These aren’t far-fetched science fiction; they are the practical, high-stakes risks of deploying powerful, tool-using agents without a security-first mindset. A simple system prompt saying “be careful” is not a security strategy.
The Solution: A Multi-Layered Firewall with LangChain Middleware
To address these risks, we need to move security out of the prompt and into the architecture. We will build an Agentic AI Firewall using LangChain’s powerful middleware system. Middleware allows us to intercept and control the agent’s execution flow at critical points, creating a robust, multi-layered defense.
Our firewall will implement:
- Strict Role-Based Access Control (RBAC): We will explicitly define which users (or roles) can access which tools. A “support_user” can’t access the
shell_tool, period. - Input & Output Sanitization: We’ll use built-in PII detection to automatically block or redact sensitive information, preventing both malicious inputs and accidental data leaks.
- Hardened Tools: Instead of relying on middleware to catch everything, we’ll build a final layer of defense directly into our most dangerous tools, ensuring that even if other layers fail, a malicious command cannot be executed.
- Audit Logging: Every attempt to use a tool—whether successful or denied—will be logged, creating an invaluable audit trail for security reviews and incident response.
This architecture treats Agentic AI security as a first-class citizen, providing the guardrails needed for production deployment.
The Complete, Production-Ready Code
Here is the working code for our Agentic Firewall. This notebook is designed to be run in Google Colab (deliberately, so you can reproduce the results).
Step 1: Environment Setup
First, we install the necessary libraries. langgraph is a key dependency for the Agentic AI executor.
# This cell installs all the required Python packages for our project.
# - langchain: The core library for building applications with LLMs.
# - langchain_openai: Provides the specific integration for OpenAI models.
# - langchain-core: Contains the essential abstractions and components of the LangChain framework.
# - langgraph: The library that powers the stateful, cyclic graphs of the modern agent executor.
!pip install -q langchain langchain_openai langchain-core langgraph
Step 2: API Key Configuration
We securely access our OpenAI API key using Colab’s built-in secret management. This is far more secure than hardcoding keys or uploading .env files.
# We import the necessary libraries for handling secrets and environment variables.
from google.colab import userdata
import os
# This is the recommended practice for using secrets in Google Colab.
# It fetches the secret named 'OPENAI_API_KEY' from the Colab Secrets Manager
# and sets it as an environment variable, which the OpenAI library automatically uses.
openai_api_key = userdata.get('OPENAI_API_KEY')
os.environ['OPENAI_API_KEY'] = openai_api_key
Step 3: Defining Our “Dangerous” Tools
Here, we define two powerful tools. Critically, our shell_tool contains its own final layer of defense—a regex check to block command chaining. This demonstrates the “defense-in-depth” principle.
import subprocess
import re # We need the regular expression module for our shell tool sanitizer.
from langchain_core.tools import tool
# This tool interacts with the local file system.
@tool
def file_system_tool(file_path: str, content: str = None, operation: str = 'read'):
"""Interacts with the file system. Operations: 'read', 'write', 'delete'."""
# A basic but important security check to prevent path traversal attacks (e.g., trying to access '../../etc/passwd').
if '..' in os.path.normpath(file_path):
return "Error: Path traversal detected. Access denied."
try:
if operation == 'read':
with open(file_path, 'r') as f: return f.read()
elif operation == 'write':
if content is None: return "Error: 'content' is required."
with open(file_path, 'w') as f: f.write(content)
return f"Successfully wrote to {file_path}."
elif operation == 'delete':
os.remove(file_path)
return f"Successfully deleted {file_path}."
else: return "Error: Invalid operation."
except Exception as e: return f"An error occurred: {e}"
# This tool can execute shell commands, making it extremely powerful and dangerous.
@tool
def shell_tool(command: str):
"""Executes a shell command."""
# DEFENSE-IN-DEPTH: This is our final, most important security check.
# Even if other middleware fails, this check inside the tool itself prevents command chaining.
# It looks for common shell metacharacters that could be used for malicious purposes.
if re.search(r"[;&|`$()<>]", command):
return "Error: For security reasons, shell commands with metacharacters like ;, &, |, etc., are not allowed."
try:
# If the command is safe, execute it in a subprocess.
result = subprocess.run(command, shell=True, capture_output=True, text=True, check=True)
return f"STDOUT:\n{result.stdout}"
except subprocess.CalledProcessError as e:
return f"Error executing command: {e.stderr}"
Step 4: Building the Custom Firewall Middleware
This is the heart of our solution. We create a class, RBACFirewall, that inherits from AgentMiddleware. This allows us to create a stateful firewall that holds our permission rules and uses the wrap_tool_call hook to enforce them.
import logging
from typing import Any, Callable, Dict, List
# These are the imports required for building class-based middleware.
from langchain.agents.middleware import AgentMiddleware, AgentState
from langchain.tools.tool_node import ToolCallRequest
from langchain_core.messages import ToolMessage
# Set up a dedicated logger for our firewall to create a clear audit trail.
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("AgenticFirewall")
# 1. Define a Custom State Schema. This tells the agent that its state can (and should) include a 'user_id'.
# This is the correct way to pass contextual information into the agent's execution flow.
class FirewallState(AgentState):
user_id: str
# 2. Create our custom firewall as a Python class. This is the recommended pattern for complex, configurable middleware.
class RBACFirewall(AgentMiddleware):
# The __init__ method allows us to configure the firewall with our permission rules when it's created.
def __init__(self, user_permissions: Dict[str, List[str]]):
self.user_permissions = user_permissions
# The `wrap_tool_call` method is the perfect hook for our firewall. It's a special method that LangChain
# will automatically call right before any tool is executed.
def wrap_tool_call(self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], Any]) -> Any:
"""This hook intercepts tool calls to enforce role-based access control."""
# The `request` object contains everything we need: the agent's current state and the details of the tool call.
# We can safely get the 'user_id' from the state.
user_id = request.state.get("user_id", "default_user")
tool_name = request.tool_call["name"]
# AUDIT LOGGING: We log every attempt, creating a security audit trail.
logger.info(f"Firewall: User '{user_id}' attempting to run tool '{tool_name}' with args {request.tool_call['args']}")
# THE CORE LOGIC: Check if the user's role exists in our permissions and if the tool is in their allowed list.
if user_id not in self.user_permissions or tool_name not in self.user_permissions[user_id]:
error_message = f"PERMISSION DENIED for user '{user_id}' to use tool '{tool_name}'."
logger.warning(error_message) # Log the denial.
# SHORT-CIRCUIT: We do not call the `handler`. Instead, we immediately return a ToolMessage with the error.
# This prevents the tool from ever running and tells the agent exactly why it failed.
return ToolMessage(content=error_message, tool_call_id=request.tool_call["id"])
# If all checks pass, we log the approval and call the `handler`, which proceeds with the actual tool execution.
logger.info(f"Firewall: Permission GRANTED for user '{user_id}' to run tool '{tool_name}'.")
return handler(request)
Step 5: Assembling the Agent
Here, we bring everything together. We instantiate our custom RBACFirewall and combine it with LangChain’s powerful built-in PIIMiddleware. The order in the middleware list is important—we want our permission check to run first.
import re
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware, PIIMiddleware
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage
from langchain_core.globals import set_debug
# Set to True to see the agent's detailed thought process, or False for clean output.
set_debug(False)
# --- Define Permissions and Instantiate Middleware ---
user_permissions = {
"default_user": ["shell_tool"],
"admin_user": ["file_system_tool", "shell_tool"],
}
# 1. Instantiate our custom class-based firewall with the permission rules.
rbac_firewall = RBACFirewall(user_permissions)
# 2. Instantiate built-in middleware for other security tasks.
# Note: I am not useing the HumanInTheLoopMiddleware for this script to run non-interactively,
# but it's a powerful tool for real applications.
output_sanitizer = PIIMiddleware("email", strategy="redact", apply_to_tool_results=True)
api_key_sanitizer = PIIMiddleware("api_key", detector=r"sk_[a-zA-Z0-9]{20,}", strategy="redact", apply_to_tool_results=True)
# --- Create the Agent ---
llm = ChatOpenAI(model="gpt-5.2", temperature=0)
tools = [file_system_tool, shell_tool]
# This system prompt is crucial. It explicitly tells the agent how to behave when it receives an error from a tool.
system_message = (
"You are a helpful and careful assistant. "
"If a tool returns an error, a block message, or a PERMISSION DENIED message, "
"you MUST report that exact message back to the user as your final answer."
)
# `create_agent` is the factory function for building stateful, graph-based agents.
agent = create_agent(
llm,
tools,
system_prompt=system_message,
# A checkpointer is required for any stateful agent to save its progress.
checkpointer=InMemorySaver(),
# The middleware list defines our security layers. They are executed in order.
middleware=[
rbac_firewall, # First, check permissions.
output_sanitizer, # Then, sanitize the output of allowed tools.
api_key_sanitizer,
],
# We must tell the agent about our custom state schema so it knows how to handle 'user_id'.
state_schema=FirewallState,
)
Step 6: Putting the Firewall to the Test
Finally, we run a series of test cases to prove that every layer of our firewall works exactly as intended.
# --- Create a dummy file for testing ---
with open("customer_record.txt", "w") as f:
f.write("Customer: Jane Doe\nEmail: jane.doe@example.com\nKey: sk_12345thisisafakekey67890\n")
# --- Define unique thread IDs for each conversation ---
# The agent is stateful. Each 'thread_id' represents a separate, independent conversation.
thread_1 = {"configurable": {"thread_id": "thread-1"}}
thread_2 = {"configurable": {"thread_id": "thread-2"}}
thread_3 = {"configurable": {"thread_id": "thread-3"}}
thread_4 = {"configurable": {"thread_id": "thread-4"}}
# --- Test Case 1: RBAC Permission Denied ---
print("\n--- Test Case 1: Denied Action (RBAC) ---")
# We invoke as a 'default_user' trying to use a tool they don't have permission for.
response = agent.invoke(
{"messages": [HumanMessage(content="Read the file 'customer_record.txt'")], "user_id": "default_user"},
thread_1
)
# EXPECTED: The agent should immediately return the "PERMISSION DENIED" message from our RBACFirewall.
print(f"✅ Final Output: {response['messages'][-1].content}")
# --- Test Case 2: Input Sanitization ---
print("\n--- Test Case 2: Input Sanitization (Block Dangerous Characters) ---")
# We invoke as an 'admin_user' but with a malicious command string.
response = agent.invoke(
{"messages": [HumanMessage(content="Run the command `ls; whoami`")], "user_id": "admin_user"},
thread_2
)
# EXPECTED: The agent should return the error message from our hardened `shell_tool`.
print(f"✅ Final Output: {response['messages'][-1].content}")
# --- Test Case 3: Output Sanitization ---
print("\n--- Test Case 3: Output Sanitization (Redact PII) ---")
# We invoke as an 'admin_user' to read a file containing sensitive data.
response = agent.invoke(
{"messages": [HumanMessage(content="Read the file 'customer_record.txt'")], "user_id": "admin_user"},
thread_3
)
# EXPECTED: The agent should return the file's content, but with the email and API key automatically redacted.
print(f"✅ Final Output:\n{response['messages'][-1].content}")
# --- Test Case 4: Normal Approved Use ---
print("\n--- Test Case 4: Normal Approved Use ---")
# A simple, safe command to demonstrate a successful tool call.
response = agent.invoke(
{"messages": [HumanMessage(content="Use the shell tool to list files")], "user_id": "admin_user"},
thread_4
)
print(f"✅ Final Output: {response['messages'][-1].content}")
# --- Cleanup ---
os.remove("customer_record.txt")
set_debug(False)
Here is what I got:

Conclusion: Building Trust in Your Agentic AI
The transition to Agentic AI represents a paradigm shift in what’s possible, but it demands an equal shift in our approach to security. Simply trusting a system prompt is no longer sufficient. As we’ve demonstrated, a robust Agentic Firewall, built with a layered, defense-in-depth strategy, is not just a best practice—it is an absolute requirement for deploying these systems in production.
By using architectural patterns like LangChain’s class-based middleware and wrap_tool_call hooks, you can build a security layer that is both powerful and maintainable.
This blueprint—combining Role-Based Access Control, automated PII sanitization, and hardened tools—provides the foundation you need to unlock the immense potential of Agentic AI while protecting your business, your customers, and your data from costly mistakes. The future of AI is agentic, and with this guide, you can ensure that future is also secure.
FAQ: Frequently Asked Questions
What is Agentic AI?
Agentic AI refers to artificial intelligence systems that are autonomous and can take actions to achieve goals. Unlike traditional models that only respond to prompts, agentic systems can use “tools” (like APIs, databases, or file systems) to interact with the world, make decisions, and execute multi-step plans to complete complex tasks.
Why is Secure AI so important for agents?
Security is critical because tool-using agents have a much larger “attack surface.” A compromised agent doesn’t just give a bad answer; it could delete your production database, leak sensitive customer data, or drain your cloud budget by calling expensive APIs in a loop. A Secure AI approach, like the Agentic Firewall, builds guardrails directly into the agent’s architecture to prevent these disastrous outcomes.
What is LangChain Middleware?
LangChain Middleware provides a powerful way to intercept and control the agent’s execution lifecycle. It uses “hooks” like wrap_tool_call to run custom code before or after a tool is executed. This is the ideal place to implement cross-cutting concerns like security, logging, and caching without cluttering your core application logic.
What is the difference between a decorator and a class for middleware?
Decorators (like `@wrap_tool_call` on a standalone function) are excellent for simple, stateless middleware. However, for more complex, stateful logic—like our firewall, which needs to be configured with a set of permissions—a class-based middleware that inherits from AgentMiddleware is the correct and more powerful pattern. It allows you to initialize the middleware with configuration and maintain state across calls.
Can’t I just put security rules in the system prompt?
While you should always instruct your agent to be careful in its prompt, this is not a reliable security mechanism. Prompts can be manipulated by malicious user input (prompt injection), and LLMs can sometimes misunderstand or ignore instructions. A true security model must be enforced in code via middleware, which cannot be bypassed by the user or the LLM.


