Free lesson · GenAI Security Engineering

Deploy output sanitizer as response middleware on GKE

Package sanitization pipeline as FastAPI response middleware. Deploy with Helm on GKE and configure resource limits for sanitizer pods.

Course: AI Security Engineering · Chapter 3 · Output Sanitization Engineering

Free to read — no subscription required.

Introduction

When you deploy an LLM-backed service to GKE, every model response passes through your stack unfiltered unless you explicitly intercept it. Without a sanitization layer in the request-response path, code execution payloads, unauthorized URLs, and system prompt fragments can reach end users. By the end of this lesson, you'll be able to package the output sanitizer as FastAPI response middleware, deploy it with a Helm chart that includes resource limits and a ConfigMap-mounted configuration, and wire up Kubernetes readiness and liveness probes so traffic only reaches fully initialized sanitizer pods.

Key Terminology

  • ASGI middleware — A component that sits in the request-response path by wrapping the ASGI send callable; OutputSanitizationMiddleware uses this pattern to intercept every HTTP response body through a SanitizationResponder before bytes reach the client.
  • Sanitization pipeline — The ordered list of stage objects (CodeScanningStage, URLValidationStage, LeakDetectionStage) assembled by _build_pipeline from SanitizationConfig boolean flags; each stage receives the buffered response body in turn and may redact or reject it.
  • ConfigMap-mounted configuration — A Kubernetes ConfigMap that stores the sanitizer's YAML settings (URL allowlist, code policy, leak-detection flags) and is volume-mounted as a file inside the pod, decoupling configuration updates from container image rebuilds.
  • Liveness probe — A Kubernetes health check wired to /healthz that calls sanitizer.scan("health check probe") to confirm the pipeline is still responsive; a failing liveness check triggers a pod restart to recover from hung states.
  • Readiness probe — A Kubernetes health check wired to /readyz that verifies three conditions — config_loaded, redis_connected, and stages_initialized — before the pod is added to Service endpoints; any failing check removes the pod from the load-balancer pool until all dependencies are healthy.
  • SanitizationConfig — The @dataclass that drives pipeline assembly by toggling enable_code_scanning, enable_url_validation, and enable_leak_detection, and carrying per-stage parameters such as url_allowlist, code_policy, and system_prompt that the Helm chart surfaces through the mounted ConfigMap.

Concepts

Intercepting Model Responses at the ASGI Layer

LLM responses are just HTTP response bodies — and in a standard FastAPI application, nothing prevents a body containing executable code, an off-allowlist URL, or a leaked system prompt from reaching the client. The only place to intercept every response uniformly, regardless of which route handler produced it, is at the ASGI transport layer.

ASGI middleware works by replacing the send callable the application receives. When OutputSanitizationMiddleware.__call__ is invoked, it constructs a SanitizationResponder that wraps the real send, buffers http.response.body chunks as they arrive, runs them through the pipeline stages, and only then forwards the sanitized bytes downstream. Because this wrapping happens at the transport boundary, no individual route handler needs to call the sanitizer — coverage is guaranteed for every HTTP response the app emits. Non-HTTP scopes such as WebSocket connections and lifespan events fall through the scope["type"] != "http" guard and reach the wrapped application untouched, so the middleware does not interfere with those protocols.

Loading diagram...

Two Probes Answer Two Different Questions

Kubernetes uses two probes because pod health has two distinct meanings. The liveness probe asks: is this pod still alive and capable of doing work at all? The readiness probe asks: is this pod ready to serve production traffic right now?

For the output sanitizer, those questions have different answers during startup. A pod may have a functioning event loop (alive) but not yet have loaded its ConfigMap-mounted configuration or confirmed that its Redis connection — used for token budget tracking — is reachable (not yet ready). The /healthz endpoint is intentionally minimal: it calls sanitizer.scan with a probe string to confirm the pipeline objects are instantiated and not deadlocked. The /readyz endpoint is deliberately stricter: it checks config_loaded, redis_connected, and stages_initialized together, returning the full checks dict so an operator can see exactly which dependency blocked the pod (see Code Walkthrough).

This separation prevents a common deployment failure: a pod that has not finished loading configuration begins receiving real traffic, passes requests through an uninitialized pipeline, and silently emits unsanitized responses. By keeping the pod out of the Service endpoints until /readyz returns 200, GKE ensures that a sanitizer pod either serves correctly or serves nothing at all.

ConfigMap as Externalised, Restart-Triggering Configuration

Baking sanitization policy into the container image couples two independent lifecycles: application code changes and policy changes. A ConfigMap breaks that coupling. The Helm chart stores the URL allowlist, code-block policy, leak-detection thresholds, and the enable_* flags as a YAML file in a ConfigMap, then mounts it as a volume path the application reads at startup.

Because Kubernetes does not restart pods automatically when a ConfigMap changes, the Helm chart uses a checksum annotation on the Deployment's pod template — a SHA256 of the ConfigMap's content baked into metadata.annotations. When helm upgrade detects that the annotation value changed (because the ConfigMap content changed), it triggers a rolling restart, cycling pods one at a time through the readiness gate. The result is that a policy update — say, adding a domain to url_allowlist — deploys with the same zero-downtime guarantees as a code change, and pods with the old policy serve traffic until their replacements pass /readyz.

Code Walkthrough

Now that you understand the Helm chart structure, ConfigMap-mounted configuration, and the role of readiness and liveness probes, the following walkthrough shows how these pieces connect in running code.

FastAPI Response Middleware

The output sanitizer is deployed as FastAPI response middleware that intercepts every model response before it reaches the client. The middleware wraps the ASGI send callable with a SanitizationResponder that buffers body chunks, runs the configured pipeline stages, and forwards the sanitized body downstream.

Code snippetpython
1from dataclasses import dataclass, field 2from typing import Callable 3 4@dataclass 5class SanitizationConfig: 6 enable_code_scanning: bool = True 7 enable_url_validation: bool = True 8 enable_leak_detection: bool = False 9 code_policy: dict = field(default_factory=dict) 10 url_allowlist: list[str] = field(default_factory=list) 11 system_prompt: str = "" 12 13class OutputSanitizationMiddleware: 14 def __init__(self, app: Callable, config: SanitizationConfig): 15 self.app = app 16 self.stages = self._build_pipeline(config) 17 18 async def __call__(self, scope, receive, send): 19 if scope["type"] != "http": 20 await self.app(scope, receive, send) 21 return 22 responder = SanitizationResponder(send, self.stages) 23 await self.app(scope, receive, responder) 24 25 def _build_pipeline(self, config: SanitizationConfig) -> list: 26 stages = [] 27 if config.enable_code_scanning: 28 stages.append(CodeScanningStage(config.code_policy)) 29 if config.enable_url_validation: 30 stages.append(URLValidationStage(config.url_allowlist)) 31 if config.enable_leak_detection: 32 stages.append(LeakDetectionStage(config.system_prompt)) 33 return stages

The _build_pipeline method conditionally assembles stages from configuration flags — the same flags that the Helm chart's ConfigMap exposes as a mounted YAML file. Non-HTTP scopes such as WebSocket and lifespan events bypass the pipeline entirely and pass through to the wrapped application unchanged.

Kubernetes Health Endpoints

The Helm chart configures two probes against these endpoints. The liveness probe calls /healthz to detect hung pods; the readiness probe calls /readyz to confirm that configuration is loaded and external dependencies are reachable before the pod receives traffic.

Code snippetpython
1from fastapi import FastAPI, HTTPException, Depends 2from redis.asyncio import Redis 3 4app = FastAPI() 5 6async def get_redis() -> Redis: 7 return Redis(host="redis-service", port=6379, decode_responses=True) 8 9@app.get("/healthz") 10async def liveness(): 11 try: 12 test_result = sanitizer.scan("health check probe") 13 if test_result is not None: 14 return {"status": "ok"} 15 except Exception: 16 pass 17 raise HTTPException(status_code=503, detail="Sanitizer unhealthy") 18 19@app.get("/readyz") 20async def readiness(redis: Redis = Depends(get_redis)): 21 checks = { 22 "config_loaded": sanitizer.config is not None, 23 "redis_connected": await redis.ping(), 24 "stages_initialized": all(s.initialized for s in sanitizer.stages), 25 } 26 if all(checks.values()): 27 return {"status": "ready", "checks": checks} 28 raise HTTPException(status_code=503, detail=checks)

The readiness endpoint performs three dependency checks: whether the sanitization configuration is loaded from the ConfigMap, whether the Redis connection used for token budget tracking responds to a ping, and whether every pipeline stage reports as initialized. A pod that fails any check is removed from the Service endpoints, so requests never reach an instance that has not finished loading its configuration.

Confirm that the deployment is healthy by running kubectl get pods -l app=output-sanitizer and verifying that READY shows 1/1 for every pod, then send a test request through the service and check that the response passes through without a 503 from either probe path.

Do's and Don'ts

Now that you have worked through the implementation, the practices below separate a durable approach from a fragile one.

Do's

  1. Do mount the sanitizer configuration from a Kubernetes ConfigMapSanitizationConfig flags like enable_code_scanning, enable_url_validation, and url_allowlist change independently of your image; ConfigMap mounting lets you update policies with a rollout rather than a rebuild.
  2. Do implement /readyz as a multi-dependency gate that checks config_loaded, redis_connected, and stages_initialized — a pod that returns ready before all three pass will receive traffic with an uninitialized pipeline, silently forwarding unsanitized model responses to clients.
  3. Do wrap only the ASGI send callable in SanitizationResponder and pass non-HTTP scopes through unmodified — intercepting WebSocket or lifespan events in OutputSanitizationMiddleware.__call__ breaks those ASGI channels without providing any sanitization benefit.

Don'ts

  1. Don't use /healthz as the readiness probe — the liveness endpoint only verifies that the sanitizer can process a trivial string; it does not check whether the ConfigMap configuration is loaded or Redis is reachable, so pods will enter the Service endpoint pool before they are actually ready to sanitize live traffic.
  2. Don't enable enable_leak_detection without supplying a non-empty system_prompt to LeakDetectionStage — the stage compares model output against the prompt it receives; an empty string causes every response to trivially pass the check, defeating the entire purpose of the stage.
  3. Don't deploy the Helm chart without explicit CPU and memory resource limits on the sanitizer pod — the SanitizationResponder buffers all body chunks in memory before forwarding; an unbounded pod under heavy streaming traffic can OOM-kill itself or a neighbor, and a hung pod without a liveness probe limit goes undetected by /healthz.

This lesson is free to read. Its 3 hands-on labs — real code, in a cloud IDE — are part of the GenAI Security Engineering subscription.

From · cancel anytime

More free lessons in AI Security Engineering

All free lessons in GenAI Security Engineering