Local model guards

Published from the project repository

This page is generated from the project’s documentation. References to source files, tests and unpublished documents appear as plain text, and test names such as suite::test_name identify the evidence behind a claim. The source repository is not public yet, so that evidence is available on request.

The optional local guard service runs a prompt-injection classifier through the gateway’s existing authenticated webhook. It downloads missing pinned artifacts automatically, verifies the cache on each startup, and loads the classifier once. It reports ready only after a real model warmup succeeds and while the worker remains live. Bounded supervision detects worker exits and idle hangs, then verifies and warms each replacement. It does not change the main model provider or replace the built-in CPU secret detector, authorization, approvals or flow rules.

The companion service is Python on POSIX hosts, outside the single Rust binary. The gateway does not install Python, drivers or dependencies. The runtime environment is exact-pinned and transitively hash-locked in scripts/local_guard_service/uv.lock; the gateway’s Rust dependency graph is unchanged. Windows service hosting is not supported in this first version.

Eligibility and failure behavior

Model format	Available runtime/device	Behavior
Safetensors	CUDA or Apple MPS, including successful allocation/kernel probe	Fetch missing artifacts; load on that GPU; warm up before ready
Safetensors	No usable compatible GPU, or GPU runtime absent	Unavailable; no model download and no CPU fallback
ONNX, --onnx-provider cpu (default)	ONNX Runtime CPU provider	Fetch missing artifacts; load and warm up on CPU, whether or not a GPU exists
ONNX, --onnx-provider cuda	CUDA provider registered; compatible Linux x86_64 runtime/driver required	Fetch missing artifacts; load with explicit fallback policy and warm up before ready
ONNX, explicit CUDA unavailable or lost during load	No requested provider	Unavailable; never downgrade to CPU

GPU presence alone is insufficient: model initialization must also succeed. ONNX provider selection is explicit: CPU remains the default. CUDA defaults to --onnx-cpu-fallback forbid, rejecting graph placement on CPU. An operator can choose allow for mixed CPU/CUDA nodes; CUDA must still register successfully. Session execution fallback is disabled in both modes: no failed subject is retried on another provider. A provider registration is not proof of kernel placement; run the opt-in profiler-backed qualification below. No backend substitutes a model. CUDA uses logical device 0 and disables TF32 reduced-precision arithmetic. Only CUDA is added for ONNX; Safetensors retains its existing CUDA/MPS gate.

An authenticated GET /health returns the state and a fixed failure reason. States include disabled, probing, verifying, waiting_for_cache, downloading, loading, ready, recovering, backoff, stopping, unsupported_backend and failed. Health also returns restarts and restart_limit; it contains no PID, path, text or credential. Only ready returns HTTP 200. Inspection while unready returns an error; the gateway applies the rule’s on_error, which defaults to deny. Hardware gating never removes an existing policy requirement.

There is one active inspection and no waiting queue, including during request body reception. Excess work is refused. Oversized byte/token-window inputs are refused without invalidating the loaded model. A runtime crash, timeout or cancellation terminates the worker process group. That request fails without replay; a single supervisor may recover the worker for subsequent requests. The lifetime restart budget defaults to three replacements (four total initialization attempts), with 1, 2 and 4 seconds of backoff, capped by --restart-backoff-max. Readiness does not reset that budget. Set --max-restarts 0 to disable automatic replacement. After exhaustion, health remains 503 with restart_budget_exhausted until an operator restarts the service. No request silently skips inspection because the GPU or model failed; error responses always follow the gateway rule’s on_error.

The default idle heartbeat interval is one second, using the inspection deadline as its timeout. It uses the same single IPC slot, so a concurrent inspection can receive guard_busy during a heartbeat. An in-flight inspection has its own bound; an idle hang is detected within interval + deadline, subject to host scheduling. Readiness checks worker liveness immediately; heartbeat expiry revokes readiness for a living but unresponsive worker. A heartbeat proves IPC responsiveness, not ongoing model quality or a fresh GPU kernel. Only initialization performs warmup.

Unsupported hardware, invalid settings/manifest, repository access failures, missing or corrupt offline artifacts, download failures and incompatible models are terminal. No automatic artifact/download retry is added. EOF, worker crashes, malformed verdicts and worker deadlines consume the shared worker restart budget. A provisioning deadline is terminal; worker probe/load deadlines can recover. Model initialization exceptions (including allocation errors) require an operator to correct compatibility/resources; process crashes during initialization can recover. Every replacement probes hardware before provisioning and re-verifies the whole cache before loading. There is no model substitution or automatic GPU-to-CPU recovery. The explicit CPU-node placement option does not change that recovery rule.

Shutdown revokes admission first, cancels active inspection and initialization, joins supervision, and finishes process-group TERM/KILL cleanup even under repeated cancellation. Process-group cleanup covers descendants after their leader exits. Run under a cgroup supervisor for supervisor SIGKILL/OOM or descendants that create a new session: a bare CLI cannot guarantee cleanup after its own uncatchable death. See the hardened deployment example. Process liveness alone does not establish guard readiness.

Run on the same host

Use Python 3.13 or 3.14 and uv 0.11.6. From the repository root:

uv sync --locked --project scripts/local_guard_service --extra gpu --no-dev
install -d -m 700 "$HOME/.config/mcp-local-guard" "$HOME/.cache/mcp-local-guard"
(umask 077; openssl rand -hex 32 > "$HOME/.config/mcp-local-guard/token")
scripts/local_guard_service/.venv/bin/python scripts/local_guard_service/serve.py \
  --cache "$HOME/.cache/mcp-local-guard" \
  --token-file "$HOME/.config/mcp-local-guard/token"

The default manifest pins Meta Prompt Guard 2 22M. Before its first download, obtain access to the gated model repository and supply HF_TOKEN through your secret manager/service environment. The service cannot grant access or accept the model terms on your behalf. A 401/403 produces model_access_required without printing the token or upstream response. The token is not passed to the inference subprocess. Model files are about 292 MB; GPU packages and device memory are extra.

Configure the gateway with:

MCP_GUARD_WEBHOOK_URL=http://127.0.0.1:8091/inspect
MCP_GUARD_WEBHOOK_TOKEN=file:///absolute/path/to/.config/mcp-local-guard/token
MCP_GUARD_WEBHOOK_TIMEOUT_MS=2000

The URL must reach this service in the gateway’s network namespace. The service binds only to IPv4 loopback and requires the bearer on both routes. A container must share that loopback namespace; a separate container’s 127.0.0.1 does not reach the host. Keep the guard token distinct from the model-repository token and the main provider credential.

Start with label rules and retain the existing built-in secret rules:

content:
  - id: local-prompt-injection
    subjects: [prompt]
    engine: webhook
    action: label
    on_error: deny
  - id: local-completion-injection
    subjects: [completion]
    engine: webhook
    action: label
    mode: buffer
    on_error: deny

Findings use detector local_model, class prompt_injection, and a count of flagged token windows. No raw score, text or tenant is returned. The gateway applies the rule’s own action, so a classifier finding cannot turn label into deny. Define prompt_injection as source: field in the policy’s class catalog before referencing it in detect or flow rules, matching external webhook classes in this build; source: detector is reserved for the compiled-in detector inventory. One webhook URL is supported; this service occupies that existing engine slot.

ONNX models

For CPU-only hosts, omit --extra gpu when syncing and explicitly select the reviewed full-precision Prompt Guard 2 22M ONNX profile:

scripts/local_guard_service/.venv/bin/python scripts/local_guard_service/serve.py \
  --manifest scripts/local_guard_service/models/prompt-guard-2-22m-onnx.json \
  --cache "$HOME/.cache/mcp-local-guard-onnx" \
  --token-file "$HOME/.config/mcp-local-guard/token"

This uses Gravitee’s public export of the intended Meta model under the Llama 4 license. It downloads missing pinned artifacts automatically and runs on CPU. See the provenance and reproducible evaluation procedure and validation record before deployment. The default remains the gated Safetensors profile and stays unavailable without a GPU. There is no silent model selection or automatic conversion.

The manifest has the same fields as the default manifest, with format: onnx, the ONNX filename in model_file, an immutable Hugging Face repository commit, and a size and SHA-256 digest for every artifact. Include tokenizer.json and any external tensor data. Only flat filenames are accepted. The selected export must take int64 input_ids and optionally attention_mask and token_type_ids, with its first output being [1, labels] logits. Configure the positive label index explicitly. The tokenizer’s special tokens and model context must agree. Warmup rejects incompatible shapes/providers before readiness.

Artifacts and manifests are trusted operator inputs, not tenant inputs. Pinning detects changed bytes; it does not make a malicious model safe. Review ONNX external data references and run the service under a dedicated account with restricted filesystem access. ONNX Runtime itself interprets the graph and its external-data paths; the manifest does not establish a filesystem sandbox for native model code.

Provisioning and cache

Downloads occur only during initialization, after the backend probe. Missing or invalid cache entries are fetched automatically. --offline disables fetching and fails if any required file is absent or invalid. No inspection request can choose a model, change its revision or trigger downloads. Restarting with a valid cache makes no artifact network request.

The source is Hugging Face HTTPS at an immutable commit. The service follows at most five requests through an exact host allowlist, rejects non-public DNS addresses, strips repository authorization on CDN redirects, and bounds time and bytes. Model weights use SHA-256; the default model’s small configuration and tokenizer files use the pinned Git blob IDs exposed by Hugging Face metadata. Those Git SHA-1 IDs are weaker than SHA-256; use independently reviewed SHA-256 manifests where required. Updating any revision or checksum is a reviewed change.

A cross-process cache lock serializes provisioning. Each file is streamed to a private partial file, verified, then atomically published; a model is loaded only after the entire bundle verifies. Unexpected extra files are refused. Failed or cancelled downloads remove partial files. Keep the cache under a private parent directory, writable only by the service account. Do not share it with untrusted users. Cache integrity checks do not protect against an attacker controlling the service account or host after verification.

Runtime model loading is local-only, disallows remote model code and requires Safetensors for the PyTorch CUDA/MPS path; ONNX uses the explicitly selected provider. The worker receives an environment allowlist, offline/telemetry-disabled settings, and no repository bearer. For enforced network containment, restrict service egress after initialization or use an offline preprovisioned cache; offline library settings are not a network sandbox.

Inspection bounds and streaming

The gateway sends decoded strings joined with newlines in bytes_b64. Member names, tool schemas and system instructions are included; role and individual string boundaries are not preserved. A classifier cannot reconstruct every trust boundary from that representation. Test false positives on legitimate code, instructions and quoted attacks before enabling deny rules.

The gateway adds x-mcp-guard-complete-subject: true for complete prompts, tool results and buffered/non-streamed completions, and false for streamed lookback windows. This service refuses missing/false markers, so it requires the matching gateway version. Existing external webhooks can ignore the additive header; the JSON contract is unchanged. Completion rules should use buffer mode. A stream rule is refused at inspection through its on_error; it is not rejected at policy load.

Tokenization retains all input tokens, then evaluates bounded overlapping windows with the model’s special tokens. The default has a 512-token context and 64-token overlap. No tail is silently truncated. Limits default to 256 KiB of decoded text, 256 windows, 1.5 seconds per inspection and 600 seconds for startup. Use --max-subject-bytes, --max-windows, --deadline, --startup-timeout, and --threshold deliberately; keep the gateway timeout above the service deadline. Overlap reduces boundary gaps but does not supply unlimited semantic context.

Validation and rollout

The Python tests execute a real synthetic ONNX classifier through the subprocess and HTTP adapters. They cover cache reuse, tampering, concurrency, partial-download cleanup, repository access failure, no-GPU gating, authentication, complete-subject checks, token-window coverage, overload, timeout and cancellation. Scripted HTTP transport tests cover redirect credentials, origin/DNS restrictions and byte bounds. The Rust regression covers the header and default-deny streaming behavior.

uv sync --locked --project scripts/local_guard_service
PYTHONPATH=scripts/local_guard_service \
  scripts/local_guard_service/.venv/bin/python -m unittest discover \
  -s scripts/local_guard_service/tests -v

The trained ONNX profile has an opt-in CPU validation and curated evaluation record. The ordinary suite remains synthetic; real-model and gateway tests require explicit opt-in. Successful Meta gated-model download/load, CUDA/MPS execution, production quality and load performance remain unverified. ONNX CUDA has an offline qualification harness for actual kernel placement, CPU/CUDA probability and decision comparison, and authenticated service performance; no GPU execution is claimed on a CPU-only host. Run representative benign and attack traffic, including multilingual content, long inputs and quoted attacks; record model revision, precision, thresholds, hardware, missed detections, false positives, p50/p95/p99 latency, queue refusals and RAM/VRAM before switching to deny. Detection is probabilistic, and an allow verdict is not proof that content is safe.

CUDA installation and confinement require the separate deployment procedure. Use uv sync --locked --project scripts/local_guard_service --no-default-groups --group onnx-cuda for the CUDA environment, and explicitly select the ONNX manifest plus --onnx-provider cuda. CPU and CUDA wheel groups conflict intentionally. The existing gpu extra still means PyTorch/Transformers for Safetensors, not ONNX.