Before / After: Hardening the Three Layers of an LLM App

#security #ai #python #llm

One LLM app has three layers where attacks land: the prompt, the tool call, the agent trace. Here is before / after for each, with real integration code from the RESK docs.

1. The prompt

The vulnerable version forwards raw user text to the model. Prompt injection does not need a bug in your code — just a forward(). A pasted instruction and your assistant works for the attacker.

The fix is one middleware line. ReskMiddleware screens every request through the resk-llm detector pipeline (11 detectors: DirectInjection, Bypass, MemoryPoisoning, Exfiltration...) before your LLM is ever called.

from fastapi import FastAPI
from resk2 import SecurityPipeline
from resk2.integrations import ReskMiddleware

app = FastAPI()
pipeline = SecurityPipeline().add(DirectInjectionDetector())
app.add_middleware(ReskMiddleware, pipeline=pipeline)

2. The tool call

An agent with a tool is an agent with privileges. If your handler runs whatever the model decided, a hijacked completion can send the email, delete the record, call the paid API.

The fix: verify every action against a policy bitmask before execution. reskSecure refuses unauthorized tool calls at the code boundary — the model can propose, the policy decides.

from resksecure import verify_tool_action, load_policy

policy = load_policy("policy.yaml")

def handle(action):
    allowed = verify_tool_action(action.name,
                                 user_mask=7,
                                 policy_set=policy)
    if not allowed:
        raise PermissionError("Not authorized")
    return run_tool(action.name, action.args)

3. The trace

If an agent acts and nothing records it, you cannot audit it. Incident reviews fail on the same missing artifact: what did the agent do, with which parameters, how confident was it.

The fix: capture every agent action with confidence, parameters and result. ReskPoints exports to Datadog, Prometheus, OpenTelemetry and webhooks.

from reskpoints import AgentLogger

logger = AgentLogger()
logger.log("agent-1", "tool_call", 0.95,
           {"tool": "search", "query": "RAG papers"}, "3 results")

Takeaway

Screen the input, gate the actions, record everything. All three are drop-in: one middleware, one verify call, one logger line.

resk-llm: pip install resk-llm — https://resk.fr/projects/resksafety.html
reskSecure: pip install resksecure — https://resk.fr/projects/resksecure.html
ReskPoints: pip install reskpoints — https://resk.fr/projects/reskpoints.html

All open source under github.com/Resk-Security.

Bonus: the same idea, one level up - the prompt itself

Code guards the layers. The prompt is where attacks start. Three before / after diffs, prompt only:

Direct override: the pipeline blocks it before the model sees it. Hidden in the data: the
sanitizer strips it and tells you. The leaked secret: the validator catches the email and the
credential in the answer. Same pipeline, prompt-level view.

DEV Community