Fixing LlamaFirewall: Nested Event Loop Errors

Hey folks! Ever run into a snag where your LlamaFirewall just won't play nice? I recently wrestled with an issue, and I'm here to break down what happened and how we can potentially fix it. The core problem revolves around nested asyncio event loops, and it's a bit of a head-scratcher, but let's dive in and unravel this together.

The Core Issue: Nested Event Loops

So, what's the deal with these nested event loops? Well, in the context of our LlamaFirewall and the mcp-context-protector, the problem arises because the asyncio.run() function is being called from within an existing event loop. Think of it like trying to start a movie inside another movie – it just doesn't work that way. Python's asyncio library is designed to have a single event loop per thread, and calling asyncio.run() when one is already running throws a RuntimeError. This is precisely the error we see in the traceback:

ERROR:llama_firewall_provider:Error in LlamaFirewallProvider.check_tool_response
Traceback (most recent call last):
  File "xxx/mcp-context-protector/src/contextprotector/guardrail_providers/llama_firewall.py", line 124, in check_tool_response
    result = lf.scan(message)
             ^^^^^^^^^^^^^^^^
  File "xxx/mcp-context-protector/.venv/lib/python3.11/site-packages/llamafirewall/llamafirewall.py", line 122, in scan
    scanner_result = asyncio.run(scanner_instance.scan(input, trace))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/asyncio/runners.py", line 186, in run
    raise RuntimeError(
RuntimeError: asyncio.run() cannot be called from a running event loop
WARNING:mcp_wrapper:Guardrail alert triggered for tool 'get_alerts': Error checking tool response: asyncio.run()
                    cannot be called from a running event loop

The LlamaFirewall internally calls asyncio.run() within its scan method. However, the mcp-context-protector (where our guardrail provider lives) is already running within an asyncio event loop. This clash leads to the error.

Where the Problem Lies: Code Analysis

Let's take a closer look at the code snippets to understand where this issue surfaces. First, in the llamafirewall/llamafirewall.py file:

# external llamafirewall/llamafirewall.py
for scanner_type in scanners:
    scanner_instance = create_scanner(scanner_type)
    LOG.debug(
        f"[LlamaFirewall] Scanning with {scanner_instance.name}, for the input {str(input.content)[:20]}"
    )
    scanner_result = asyncio.run(scanner_instance.scan(input, trace)) # <-- this call fails

This section of code iterates through various scanners and calls the scan method using asyncio.run(). This is problematic because it blocks the current event loop and tries to start a new one.

Now, let's examine the src/contextprotector/__main__.py file:

# src/contextprotector/__main__.py
def main() -> None:
    """Launch async main function."""
    asyncio.run(main_async()) # <-- entire mcp-context-protector already inside an event loop

Here, the entire mcp-context-protector is already within an asyncio event loop. This means that when the LlamaFirewall attempts to start another one, it runs into the RuntimeError.

Solution: Embrace async and await

The fix, in essence, involves making the LlamaFirewall more async-aware. The library provides an scan_async function which is designed for this use case. However, the current code needs some refactoring to leverage this properly. Many functions need to be marked as async, and the code must use await where appropriate to avoid blocking the event loop. The key is to avoid using asyncio.run() within the mcp-context-protector's event loop. Instead, the scan_async function should be used and awaited.

Refactoring Strategy

1. Identify Asynchronous Operations: First, identify all functions within the LlamaFirewall that perform asynchronous operations (e.g., network calls, file I/O).
2. Mark Functions as async: Decorate these functions with the async keyword.
3. Use await: Within these async functions, use the await keyword before any call to an asynchronous function or method.
4. Replace asyncio.run(): Replace any calls to asyncio.run() with appropriate calls to the scan_async method, and await them. This is the crucial step to ensure that the code works correctly within an existing event loop.
5. Test Thoroughly: After refactoring, thoroughly test the code to ensure that the LlamaFirewall integrates seamlessly with the mcp-context-protector without raising the nested event loop error.

Reproducing the Issue

To see this issue in action, you can use the official demo repository. Here's how:

1. Get the Demo Repo: Clone the demo repo from the provided GitHub link.
2. Run the Inspector: Execute the command provided in the original issue to run the demo using the inspector. This should trigger the LlamaFirewall and expose the error.

npx @modelcontextprotocol/inspector bash $PWD/mcp-context-protector.sh \
	--guardrail-provider LlamaFirewall --command-args $(which uv) run \
	--directory xxx/mcp-quickstart-resources/weather-server-python/ weather.py

Conclusion: A Path Forward

The LlamaFirewallProvider issue stems from nested asyncio event loops. To resolve this, the code needs refactoring to fully embrace asynchronous programming. This involves marking functions as async, using await for asynchronous operations, and avoiding asyncio.run() within the context of an already running event loop. By using the scan_async method and making the necessary adjustments, we can integrate the LlamaFirewall seamlessly into the mcp-context-protector, thus avoiding the RuntimeError. This refactoring should ensure a more robust and efficient implementation. So, if you're experiencing this, get ready to dive into some async and await magic, and your LlamaFirewall will be back in action in no time!