feat(guardrails): add Checker and NeMo translation layer #1409
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,266 @@ | ||
| // Package guardrails checks tools/call requests and responses against an | ||
| // external guardrails server. Checker owns HTTP transport, timeout, TLS, | ||
| // fail mode, config ID merging, and provider translation. | ||
| package guardrails | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "context" | ||
| "crypto/tls" | ||
| "encoding/json" | ||
| "fmt" | ||
| "io" | ||
| "net" | ||
| "net/http" | ||
| "strings" | ||
| "time" | ||
|
|
||
| "github.com/Kuadrant/mcp-gateway/internal/config" | ||
| "github.com/Kuadrant/mcp-gateway/internal/guardrails/external/nemo" | ||
| ) | ||
|
|
||
| // checksPath is the guardrails server endpoint all checks are sent to. | ||
| const checksPath = "/v1/guardrail/checks" | ||
|
|
||
| // checkTimeout bounds a single guardrails HTTP round trip, comfortably | ||
| // inside the 10s ext_proc message_timeout. | ||
| const checkTimeout = 3 * time.Second | ||
|
|
||
| // dialTimeout bounds DNS/TCP connection setup so an unreachable guardrails | ||
| // server fails fast rather than eating the full checkTimeout on dial alone. | ||
| const dialTimeout = 1 * time.Second | ||
|
|
||
| // defaultMaxIdleConnsPerHost is used when the caller doesn't specify a | ||
| // concurrency hint. | ||
| const defaultMaxIdleConnsPerHost = 100 | ||
|
|
||
| // defaultMaxBodyBytes bounds the guardrails server's check response when the | ||
| // caller doesn't specify a limit, matching the MCPGatewayExtension | ||
| // maxBodyBytes default (1 MiB). | ||
| const defaultMaxBodyBytes = 1 << 20 | ||
|
|
||
| // Status is the outcome of a guardrails check. | ||
| type Status string | ||
|
|
||
| // Status values a Decision can carry. | ||
| const ( | ||
| StatusAllowed Status = "allowed" | ||
| StatusBlocked Status = "blocked" | ||
| StatusModified Status = "modified" | ||
| ) | ||
|
|
||
| // Decision is the outcome of a single guardrails check, translated from the | ||
| // NeMo Guardrails server response into a form the router acts on. | ||
| type Decision struct { | ||
| Status Status | ||
| // Content is the text to forward: the original content unless Status | ||
| // is StatusModified, in which case it's the guardrails modified text. | ||
| Content string | ||
| // Reason names the triggering rail. Empty when Status is StatusAllowed. | ||
| Reason string | ||
| // Err is set when Status was resolved by failMode after a transport | ||
| // failure or unparseable response. | ||
| Err error | ||
| } | ||
|
|
||
| // Checker runs guardrails checks against tools/call requests and responses. | ||
| type Checker interface { | ||
| CheckRequest(ctx context.Context, toolName string, arguments json.RawMessage, configIDs []string) (*Decision, error) | ||
| CheckResponse(ctx context.Context, toolName string, content []byte, configIDs []string) (*Decision, error) | ||
| } | ||
|
|
||
| // provider translates between MCP and a guardrails backend's check | ||
| // request/response schema, and classifies a raw verdict into a Status. | ||
| // Secret type determines which implementation is used; nemoProvider is the | ||
| // only one today. | ||
| type provider interface { | ||
| TransformRequest(toolName string, arguments json.RawMessage, configIDs []string) ([]byte, error) | ||
| TransformResponse(toolName string, content []byte, configIDs []string) ([]byte, error) | ||
| ParseCheckResponse(body []byte) (status Status, content, reason string, err error) | ||
| } | ||
|
|
||
| // nemoProvider adapts *nemo.Transformer to the provider interface, | ||
| // translating NeMo's status strings into the transport-agnostic Status. | ||
| type nemoProvider struct { | ||
|
christinaexyou marked this conversation as resolved.
|
||
| *nemo.Transformer | ||
| } | ||
|
|
||
| func (p *nemoProvider) ParseCheckResponse(body []byte) (Status, string, string, error) { | ||
| resp, err := p.Transformer.ParseCheckResponse(body) | ||
| if err != nil { | ||
| return "", "", "", err | ||
| } | ||
| switch resp.Status { | ||
| case nemo.StatusSuccess: | ||
| return StatusAllowed, resp.Content, resp.Rail, nil | ||
| case nemo.StatusModified: | ||
| return StatusModified, resp.Content, resp.Rail, nil | ||
| case nemo.StatusBlocked: | ||
| return StatusBlocked, resp.Content, resp.Rail, nil | ||
| default: | ||
| // unreachable: nemo.Transformer.ParseCheckResponse already rejects | ||
| // unrecognized status values. | ||
| return "", "", "", fmt.Errorf("guardrails: unrecognized status %q", resp.Status) | ||
| } | ||
| } | ||
|
|
||
| // nemoChecker implements Checker against a NeMo Guardrails server. | ||
| type nemoChecker struct { | ||
| httpClient *http.Client | ||
| baseURL string | ||
| globalConfigIDs []string | ||
| failMode string | ||
| maxBodyBytes int64 | ||
| provider provider | ||
| } | ||
|
|
||
| // NewChecker constructs a Checker for the given resolved guardrails config. | ||
| // maxBodyBytes bounds the guardrails server's check response; non-positive | ||
| // values fall back to defaultMaxBodyBytes. | ||
| func NewChecker(cfg *config.GuardrailsConfig, tlsConfig *tls.Config, maxIdleConnsPerHost int, maxBodyBytes int64) Checker { | ||
| if maxIdleConnsPerHost <= 0 { | ||
| maxIdleConnsPerHost = defaultMaxIdleConnsPerHost | ||
| } | ||
| if maxBodyBytes <= 0 { | ||
| maxBodyBytes = defaultMaxBodyBytes | ||
| } | ||
|
|
||
| transport := http.DefaultTransport.(*http.Transport).Clone() | ||
| dialer := &net.Dialer{Timeout: dialTimeout} | ||
| transport.DialContext = dialer.DialContext | ||
| transport.TLSClientConfig = tlsConfig | ||
| transport.MaxIdleConnsPerHost = maxIdleConnsPerHost | ||
|
|
||
| return &nemoChecker{ | ||
| // no Client.Timeout: each call sets its own context deadline instead | ||
| httpClient: &http.Client{Transport: transport}, | ||
| baseURL: strings.TrimSuffix(cfg.URL, "/"), | ||
| globalConfigIDs: cfg.ConfigIDs, | ||
| failMode: normalizeFailMode(cfg.FailMode), | ||
| maxBodyBytes: maxBodyBytes, | ||
| provider: &nemoProvider{Transformer: nemo.NewTransformer(cfg.Model)}, | ||
| } | ||
| } | ||
|
|
||
| // CheckRequest translates and checks a tools/call request. A translation | ||
| // failure is always a hard deny regardless of failMode; a transport failure | ||
| // or an unparseable guardrails response falls back to failMode instead. | ||
| func (c *nemoChecker) CheckRequest(ctx context.Context, toolName string, arguments json.RawMessage, configIDs []string) (*Decision, error) { | ||
| body, err := c.provider.TransformRequest(toolName, arguments, mergeConfigIDs(c.globalConfigIDs, configIDs)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: request translation failed: %w", err) | ||
| } | ||
| return c.check(ctx, body) | ||
| } | ||
|
|
||
| // CheckResponse translates and checks a tools/call response's text content. | ||
| // Same failure semantics as CheckRequest. | ||
| func (c *nemoChecker) CheckResponse(ctx context.Context, toolName string, content []byte, configIDs []string) (*Decision, error) { | ||
| body, err := c.provider.TransformResponse(toolName, content, mergeConfigIDs(c.globalConfigIDs, configIDs)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: response translation failed: %w", err) | ||
| } | ||
| return c.check(ctx, body) | ||
| } | ||
|
|
||
| // check performs the guardrails HTTP round trip and maps the outcome to a | ||
| // Decision. Non-2xx, transport errors, oversized bodies, and unparseable | ||
| // responses all fall back to failMode rather than propagating an error — | ||
| // only a translation failure (handled by the caller) skips failMode | ||
| // entirely. | ||
| func (c *nemoChecker) check(ctx context.Context, body []byte) (*Decision, error) { | ||
|
Contributor
There was a problem hiding this comment. nit , suggestion: we might benefit from avoiding any missed errors if we split this? func (c *nemoChecker) check(ctx context.Context, body []byte) (*Decision, error)
{
decision, err := c.doCheck(ctx, body)
if err != nil {
return c.failModeDecision(err), nil
}
return decision, nil
}
func (c *nemoChecker) doCheck(ctx context.Context, body []byte) (*Decision,
error) {
ctx, cancel := context.WithTimeout(ctx, checkTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
c.baseURL+checksPath, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("guardrails: failed to build check request: %w",
err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("guardrails: request failed: %w", err)
}
Collaborator
Author
There was a problem hiding this comment. done; i split it into |
||
| decision, err := c.checkResponse(ctx, body) | ||
| if err != nil { | ||
| return c.failModeDecision(err), nil | ||
| } | ||
| return decision, nil | ||
| } | ||
|
|
||
| func (c *nemoChecker) checkResponse(ctx context.Context, body []byte) (*Decision, error) { | ||
| ctx, cancel := context.WithTimeout(ctx, checkTimeout) | ||
| defer cancel() | ||
|
|
||
| req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+checksPath, bytes.NewReader(body)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: failed to build check request: %w", err) | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
| } | ||
| req.Header.Set("Content-Type", "application/json") | ||
| req.Header.Set("Accept", "application/json") | ||
|
|
||
| resp, err := c.httpClient.Do(req) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: request failed: %w", err) | ||
| } | ||
| defer resp.Body.Close() //nolint:errcheck // best-effort close, response already consumed | ||
|
|
||
| if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { | ||
| _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, c.maxBodyBytes+1)) | ||
| return nil, fmt.Errorf("guardrails: server returned status %d", resp.StatusCode) | ||
| } | ||
|
|
||
| respBody, err := io.ReadAll(io.LimitReader(resp.Body, c.maxBodyBytes+1)) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: failed to read response: %w", err) | ||
| } | ||
| if int64(len(respBody)) > c.maxBodyBytes { | ||
| return nil, fmt.Errorf("guardrails: response exceeds %d byte limit", c.maxBodyBytes) | ||
| } | ||
|
|
||
| status, content, reason, err := c.provider.ParseCheckResponse(respBody) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("guardrails: malformed response: %w", err) | ||
| } | ||
|
|
||
| return &Decision{Status: status, Content: content, Reason: reason}, nil | ||
| } | ||
|
|
||
| // failModeDecision resolves a transport failure or unparseable response | ||
| // into a Decision per the configured failMode, keeping cause so callers can | ||
| // tell a failMode fallback apart from a real guardrails verdict. | ||
| func (c *nemoChecker) failModeDecision(cause error) *Decision { | ||
| if c.failMode == FailModeAllow { | ||
| return &Decision{Status: StatusAllowed, Err: cause} | ||
| } | ||
| return &Decision{Status: StatusBlocked, Reason: "guardrails check failed", Err: cause} | ||
| } | ||
|
|
||
| // mergeConfigIDs lists global config IDs first, then per-server ones, so | ||
| // gateway-wide policies evaluate before server-specific ones, deduplicating | ||
| // any overlap between the two. | ||
| func mergeConfigIDs(global, perServer []string) []string { | ||
| if len(global) == 0 { | ||
| return dedup(perServer) | ||
| } | ||
| if len(perServer) == 0 { | ||
| return dedup(global) | ||
| } | ||
| merged := make([]string, 0, len(global)+len(perServer)) | ||
| merged = append(merged, global...) | ||
| merged = append(merged, perServer...) | ||
| return dedup(merged) | ||
| } | ||
|
|
||
| // dedup removes duplicates while preserving first-occurrence order. Returns | ||
| // the input unmodified (including nil) when there's nothing to dedup. | ||
| func dedup(ids []string) []string { | ||
| if len(ids) < 2 { | ||
| return ids | ||
| } | ||
| seen := make(map[string]struct{}, len(ids)) | ||
| out := make([]string, 0, len(ids)) | ||
| for _, id := range ids { | ||
| if _, ok := seen[id]; ok { | ||
| continue | ||
| } | ||
| seen[id] = struct{}{} | ||
| out = append(out, id) | ||
| } | ||
| return out | ||
| } | ||
|
|
||
| func normalizeFailMode(failMode string) string { | ||
| if failMode == "" { | ||
| return FailModeDeny | ||
| } | ||
| return failMode | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
Repository: Kuadrant/mcp-gateway
Length of output: 1527
🏁 Script executed:
Repository: Kuadrant/mcp-gateway
Length of output: 14734
🌐 Web query:
Kubernetes documentation imagePullPolicy IfNotPresent latest tag cached image mutable tags💡 Result:
In Kubernetes, the imagePullPolicy field determines when the Kubelet attempts to pull a container image from a registry. The behavior of IfNotPresent and its interaction with the latest tag or mutable tags is governed by how the Kubelet handles local caching [1][2]. imagePullPolicy: IfNotPresent When set to IfNotPresent, the Kubelet will only pull the image if it is not already present on the node [1][2]. If an image with the specified tag is already cached locally, the Kubelet skips the pull process entirely and does not check the registry for updates [1][3]. Consequently, if a mutable tag (such as:latest) is updated in the remote registry, a Pod using IfNotPresent will continue to use the older, cached version of the image already present on the node [1][4]. Defaulting Behavior If the imagePullPolicy field is omitted, Kubernetes automatically sets a default based on the image tag [1][5]: - If the tag is:latest or if no tag is specified, the policy defaults to Always [1][2]. - If any other specific tag is used, the policy defaults to IfNotPresent [1][2]. - Recent updates have clarified that if a digest is provided (even with the:latest tag), the default may behave as IfNotPresent, as the digest is an immutable identifier [6]. imagePullPolicy: Always When set to Always, the Kubelet contacts the container registry every time it launches a container to resolve the image name to an image digest [1][2]. - If the resolved digest matches an image already present locally, the Kubelet uses the cached version [2][7]. - If the resolved digest differs from what is cached locally, or if the digest is not present, the Kubelet downloads the necessary layers [1][2]. Summary for Mutable Tags Because tags like:latest are mutable, they can point to different image digests over time. Using IfNotPresent with a mutable tag is generally discouraged in production environments because it prevents the automatic adoption of updated images, leading to inconsistent Pod versions across a cluster [8][4]. For deployments requiring the most recent version of a mutable tag, imagePullPolicy: Always is required to ensure the Kubelet verifies the latest image digest with the registry [1][3]. To ensure absolute consistency and avoid registry dependency issues, using immutable image digests (e.g., image@sha256:...) is the recommended best practice [1][6]. top_results: [1][2][4][3][5]
Citations:
imagePullPolicy: Alwaysmay pull stale image based off stale image digest resolution kubernetes/kubernetes#116867imagePullPolicy: Alwayskubernetes/kubernetes#90809🏁 Script executed:
Repository: Kuadrant/mcp-gateway
Length of output: 6177
Pin images and align all install artifacts.
IfNotPresentcan reuse cachedlatestimages, so nodes can run different builds. Pin both images to immutable digests or protected immutable release tags. Update the generated bundle and align theconfig/installpath, which currently deploysconfig/mcp-system/deployment-controller.yamlwithv0.9.0.🤖 Prompt for AI Agents