diff --git a/api/v1/mcpgatewayextension_types.go b/api/v1/mcpgatewayextension_types.go index da3c0ddde..0cd911a31 100644 --- a/api/v1/mcpgatewayextension_types.go +++ b/api/v1/mcpgatewayextension_types.go @@ -149,6 +149,7 @@ type MCPGatewayExtensionSpec struct { // Applies to request/response prefix stripping and guardrails checks. // +optional // +default=1048576 + // +kubebuilder:validation:Minimum=1 MaxBodyBytes *int32 `json:"maxBodyBytes,omitempty"` } diff --git a/bundle/manifests/mcp-gateway.clusterserviceversion.yaml b/bundle/manifests/mcp-gateway.clusterserviceversion.yaml index 8c94b31aa..b922c1f05 100644 --- a/bundle/manifests/mcp-gateway.clusterserviceversion.yaml +++ b/bundle/manifests/mcp-gateway.clusterserviceversion.yaml @@ -6,7 +6,7 @@ metadata: capabilities: Basic Install categories: Integration & Delivery containerImage: ghcr.io/kuadrant/mcp-controller:latest - createdAt: "2026-08-14T15:16:44Z" + createdAt: "2026-08-20T18:43:06Z" description: An Envoy-based gateway for Model Context Protocol (MCP) servers operators.operatorframework.io/builder: operator-sdk-v1.38.0 operators.operatorframework.io/project_layout: go.kubebuilder.io/v4 diff --git a/bundle/manifests/mcp.kuadrant.io_mcpgatewayextensions.yaml b/bundle/manifests/mcp.kuadrant.io_mcpgatewayextensions.yaml index 96fe4ed6a..623606545 100644 --- a/bundle/manifests/mcp.kuadrant.io_mcpgatewayextensions.yaml +++ b/bundle/manifests/mcp.kuadrant.io_mcpgatewayextensions.yaml @@ -117,6 +117,7 @@ spec: maxBodyBytes caps the size of any body the router buffers, in bytes. Applies to request/response prefix stripping and guardrails checks. format: int32 + minimum: 1 type: integer oauthProtectedResource: description: |- diff --git a/charts/mcp-gateway/crds/mcp.kuadrant.io_mcpgatewayextensions.yaml b/charts/mcp-gateway/crds/mcp.kuadrant.io_mcpgatewayextensions.yaml index ec28ec380..dec559aea 100644 --- a/charts/mcp-gateway/crds/mcp.kuadrant.io_mcpgatewayextensions.yaml +++ b/charts/mcp-gateway/crds/mcp.kuadrant.io_mcpgatewayextensions.yaml @@ -117,6 +117,7 @@ spec: maxBodyBytes caps the size of any body the router buffers, in bytes. Applies to request/response prefix stripping and guardrails checks. format: int32 + minimum: 1 type: integer oauthProtectedResource: description: |- diff --git a/config/crd/mcp.kuadrant.io_mcpgatewayextensions.yaml b/config/crd/mcp.kuadrant.io_mcpgatewayextensions.yaml index ec28ec380..dec559aea 100644 --- a/config/crd/mcp.kuadrant.io_mcpgatewayextensions.yaml +++ b/config/crd/mcp.kuadrant.io_mcpgatewayextensions.yaml @@ -117,6 +117,7 @@ spec: maxBodyBytes caps the size of any body the router buffers, in bytes. Applies to request/response prefix stripping and guardrails checks. format: int32 + minimum: 1 type: integer oauthProtectedResource: description: |- diff --git a/config/mcp-gateway/components/controller/deployment-controller.yaml b/config/mcp-gateway/components/controller/deployment-controller.yaml index f00c7cde0..89450969a 100644 --- a/config/mcp-gateway/components/controller/deployment-controller.yaml +++ b/config/mcp-gateway/components/controller/deployment-controller.yaml @@ -20,14 +20,14 @@ spec: serviceAccountName: mcp-controller containers: - name: mcp-controller - image: ghcr.io/kuadrant/mcp-controller:v0.9.0 + image: ghcr.io/kuadrant/mcp-controller:latest imagePullPolicy: IfNotPresent command: - ./mcp_controller - --log-level=0 # info level env: - name: RELATED_IMAGE_ROUTER_BROKER - value: ghcr.io/kuadrant/mcp-gateway:v0.9.0 + value: ghcr.io/kuadrant/mcp-gateway:latest ports: - name: health containerPort: 8081 diff --git a/internal/controller/session_store.go b/internal/controller/session_store.go index df246250c..0c0cc2727 100644 --- a/internal/controller/session_store.go +++ b/internal/controller/session_store.go @@ -50,7 +50,8 @@ func (r *MCPGatewayExtensionReconciler) validateSessionStore(ctx context.Context } // enqueueMCPGatewayExtForSecret maps a secret change to MCPGatewayExtension reconcile requests. -// It enqueues extensions that reference the secret via trustedHeadersKey or sessionStore. +// It enqueues extensions that reference the secret via sessionStore, trustedHeadersKey, +// caCertBundleRef, the session signing key, or the guardrails-ref annotation. func (r *MCPGatewayExtensionReconciler) enqueueMCPGatewayExtForSecret(ctx context.Context, obj client.Object) []reconcile.Request { secret := obj.(*corev1.Secret) @@ -83,6 +84,13 @@ func (r *MCPGatewayExtensionReconciler) enqueueMCPGatewayExtForSecret(ctx contex requests = append(requests, reconcile.Request{ NamespacedName: types.NamespacedName{Name: ext.Name, Namespace: ext.Namespace}, }) + continue + } + if ref := ext.Annotations[labelGuardrailsReference]; ref != "" && ref == secret.Name { + requests = append(requests, reconcile.Request{ + NamespacedName: types.NamespacedName{Name: ext.Name, Namespace: ext.Namespace}, + }) + continue } } return requests diff --git a/internal/guardrails/checker.go b/internal/guardrails/checker.go new file mode 100644 index 000000000..c035613a0 --- /dev/null +++ b/internal/guardrails/checker.go @@ -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 { + *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) { + 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) + } + 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 +} diff --git a/internal/guardrails/checker_test.go b/internal/guardrails/checker_test.go new file mode 100644 index 000000000..2aa1ec821 --- /dev/null +++ b/internal/guardrails/checker_test.go @@ -0,0 +1,211 @@ +package guardrails + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/Kuadrant/mcp-gateway/internal/config" +) + +func newTestChecker(t *testing.T, handler http.HandlerFunc, failMode string) Checker { + t.Helper() + return newTestCheckerWithMaxBodyBytes(t, handler, failMode, 0) +} + +func newTestCheckerWithMaxBodyBytes(t *testing.T, handler http.HandlerFunc, failMode string, maxBodyBytes int64) Checker { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + return NewChecker(&config.GuardrailsConfig{ + URL: server.URL, + Model: "meta/llama-3.1-8b-instruct", + ConfigIDs: []string{"global-1"}, + FailMode: failMode, + }, nil, 0, maxBodyBytes) +} + +func TestNeMoChecker_CheckRequest(t *testing.T) { + t.Run("success verdict is allowed", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + require.Equal(t, "/v1/guardrail/checks", r.URL.Path) + + messages := body["messages"].([]any) + msg := messages[0].(map[string]any) + require.Equal(t, "user", msg["role"]) + require.Equal(t, "execute_sql", msg["name"]) + + guardrails := body["guardrails"].(map[string]any) + require.Equal(t, []any{"global-1", "server-1"}, guardrails["config_ids"]) + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"success","content":"ok"}`)) + }, FailModeDeny) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{"query":"SELECT 1"}`), []string{"server-1"}) + require.NoError(t, err) + require.Equal(t, &Decision{Status: StatusAllowed, Content: "ok"}, decision) + }) + + t.Run("blocked verdict carries the triggering rail as reason", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"blocked","content":"denied","rail":"tool-safety-v1"}`)) + }, FailModeDeny) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, &Decision{Status: StatusBlocked, Content: "denied", Reason: "tool-safety-v1"}, decision) + }) + + t.Run("translation failure is a hard error regardless of failMode", func(t *testing.T) { + checker := newTestChecker(t, func(_ http.ResponseWriter, _ *http.Request) { + t.Fatal("guardrails server should not be called on translation failure") + }, FailModeAllow) + + decision, err := checker.CheckRequest(context.Background(), "", json.RawMessage(`{}`), nil) + require.Error(t, err) + require.Nil(t, decision) + }) + + t.Run("non-2xx applies failMode: deny blocks", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }, FailModeDeny) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusBlocked, decision.Status) + require.Error(t, decision.Err, "a failMode fallback must be distinguishable from a real blocked verdict") + }) + + t.Run("non-2xx applies failMode: allow releases", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + }, FailModeAllow) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusAllowed, decision.Status) + require.Error(t, decision.Err) + }) + + t.Run("malformed guardrails response applies failMode rather than erroring", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`not json`)) + }, FailModeDeny) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusBlocked, decision.Status) + require.Error(t, decision.Err) + }) + + t.Run("oversized guardrails response applies failMode rather than erroring", func(t *testing.T) { + checker := newTestCheckerWithMaxBodyBytes(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"success","content":"` + strings.Repeat("a", 32) + `"}`)) + }, FailModeDeny, 8) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusBlocked, decision.Status) + require.Error(t, decision.Err, "a failMode fallback must be distinguishable from a real blocked verdict") + }) + + t.Run("unreachable guardrails server applies failMode", func(t *testing.T) { + checker := NewChecker(&config.GuardrailsConfig{ + URL: "http://127.0.0.1:1", + Model: "meta/llama-3.1-8b-instruct", + FailMode: FailModeAllow, + }, nil, 0, 0) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusAllowed, decision.Status) + require.Error(t, decision.Err) + }) + + t.Run("malformed URL applies failMode rather than erroring", func(t *testing.T) { + for _, tc := range []struct { + failMode string + status Status + reason string + }{ + {FailModeAllow, StatusAllowed, ""}, + {FailModeDeny, StatusBlocked, "guardrails check failed"}, + } { + t.Run(tc.failMode, func(t *testing.T) { + checker := NewChecker(&config.GuardrailsConfig{ + URL: "http://[", + Model: "meta/llama-3.1-8b-instruct", + FailMode: tc.failMode, + }, nil, 0, 0) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, tc.status, decision.Status) + require.Equal(t, tc.reason, decision.Reason) + require.Error(t, decision.Err) + require.Contains(t, decision.Err.Error(), "failed to build check request") + }) + } + }) + + t.Run("a real blocked verdict carries no Err, unlike a failMode fallback", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"blocked","content":"denied","rail":"tool-safety-v1"}`)) + }, FailModeDeny) + + decision, err := checker.CheckRequest(context.Background(), "execute_sql", json.RawMessage(`{}`), nil) + require.NoError(t, err) + require.Equal(t, StatusBlocked, decision.Status) + require.NoError(t, decision.Err) + }) +} + +func TestNeMoChecker_CheckResponse(t *testing.T) { + t.Run("modified verdict returns substituted content", func(t *testing.T) { + checker := newTestChecker(t, func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + messages := body["messages"].([]any) + msg := messages[0].(map[string]any) + require.Equal(t, "assistant", msg["role"]) + _, hasConfig := msg["config"] + require.False(t, hasConfig, "response checks must not set the tool config tag") + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"modified","content":"[REDACTED]","rail":"pii-detection"}`)) + }, FailModeDeny) + + decision, err := checker.CheckResponse(context.Background(), "execute_sql", []byte("alice@example.com"), nil) + require.NoError(t, err) + require.Equal(t, &Decision{Status: StatusModified, Content: "[REDACTED]", Reason: "pii-detection"}, decision) + }) +} + +func TestMergeConfigIDs(t *testing.T) { + require.Equal(t, []string{"a", "b"}, mergeConfigIDs([]string{"a"}, []string{"b"})) + require.Equal(t, []string{"a"}, mergeConfigIDs([]string{"a"}, nil)) + require.Equal(t, []string{"b"}, mergeConfigIDs(nil, []string{"b"})) + require.Nil(t, mergeConfigIDs(nil, nil)) + + t.Run("deduplicates overlapping global and per-server IDs, keeping global order first", func(t *testing.T) { + require.Equal(t, []string{"a", "b", "c"}, mergeConfigIDs([]string{"a", "b"}, []string{"b", "c"})) + }) + + t.Run("deduplicates within a single side", func(t *testing.T) { + require.Equal(t, []string{"a", "b"}, mergeConfigIDs([]string{"a", "a"}, []string{"b", "b"})) + }) +} diff --git a/internal/guardrails/external/nemo/transformer.go b/internal/guardrails/external/nemo/transformer.go new file mode 100644 index 000000000..f926db768 --- /dev/null +++ b/internal/guardrails/external/nemo/transformer.go @@ -0,0 +1,175 @@ +// Package nemo implements the guardrails Transformer for NeMo Guardrails, +// translating between MCP and NeMo's /v1/guardrail/checks schema. +package nemo + +import ( + "encoding/json" + "fmt" +) + +// Status values in CheckResponse.Status. +const ( + StatusSuccess = "success" + StatusModified = "modified" + StatusBlocked = "blocked" +) + +// CheckRequest is the request body for NeMo's /v1/guardrail/checks endpoint. +type CheckRequest struct { + Model string `json:"model"` + Messages []Message `json:"messages"` + Guardrails GuardrailsConfig `json:"guardrails"` +} + +// Message is a single message in a CheckRequest. +type Message struct { + Role string `json:"role"` + Name string `json:"name,omitempty"` + Content json.RawMessage `json:"content"` + Config string `json:"config,omitempty"` +} + +// GuardrailsConfig stores the merged global+per-server config IDs. +type GuardrailsConfig struct { + ConfigIDs []string `json:"config_ids"` +} + +// CheckResponse is the response body from NeMo's /v1/guardrail/checks endpoint. +type CheckResponse struct { + Status string `json:"status"` + Content string `json:"content"` + Rail string `json:"rail,omitempty"` +} + +// Transformer translates between MCP guardrails checks and NeMo's +// /v1/guardrail/checks request/response schema. +type Transformer struct { + Model string +} + +var emptyJSONObject = json.RawMessage("{}") + +// NewTransformer returns a transformer bound to the model identifier +// resolved from the guardrails Secret. +func NewTransformer(model string) *Transformer { + return &Transformer{Model: model} +} + +// TransformRequest translates a tools/call request into a NeMo +// /v1/guardrail/checks request body. Maps params.name to messages[0].name, +// params.arguments (JSON-encoded) to messages[0].content, and role to +// "user". +func (t *Transformer) TransformRequest(toolName string, arguments json.RawMessage, configIDs []string) ([]byte, error) { + if toolName == "" { + return nil, fmt.Errorf("nemo: tool name is required") + } + + if len(arguments) == 0 { + arguments = emptyJSONObject + } + + quoted, err := quoteJSONString(arguments) + if err != nil { + return nil, fmt.Errorf("nemo: failed to quote arguments: %w", err) + } + + req := CheckRequest{ + Model: t.Model, + Messages: []Message{ + { + Role: "user", + Name: toolName, + Content: quoted, + Config: "tool", + }, + }, + Guardrails: GuardrailsConfig{ + ConfigIDs: nonNilConfigIDs(configIDs), + }, + } + + body, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("nemo: failed to marshal check request: %w", err) + } + return body, nil +} + +// TransformResponse translates a tools/call result's text content into a NeMo +// /v1/guardrail/checks request body. Maps the tool name (from request context) +// to messages[0].name, the text content to messages[0].content, and role to +// "assistant". +func (t *Transformer) TransformResponse(toolName string, content []byte, configIDs []string) ([]byte, error) { + if toolName == "" { + return nil, fmt.Errorf("nemo: tool name is required") + } + + quoted, err := quoteJSONString(content) + if err != nil { + return nil, fmt.Errorf("nemo: failed to quote content: %w", err) + } + + req := CheckRequest{ + Model: t.Model, + Messages: []Message{ + { + Role: "assistant", + Name: toolName, + Content: quoted, + }, + }, + Guardrails: GuardrailsConfig{ + ConfigIDs: nonNilConfigIDs(configIDs), + }, + } + + body, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("nemo: failed to marshal check request: %w", err) + } + return body, nil +} + +// ParseCheckResponse unmarshals a raw /v1/guardrail/checks HTTP response body. +// An unrecognized Status is a translation failure. +func (t *Transformer) ParseCheckResponse(body []byte) (*CheckResponse, error) { + var resp CheckResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("nemo: failed to unmarshal check response: %w", err) + } + + switch resp.Status { + case StatusSuccess, StatusModified, StatusBlocked: + default: + return nil, fmt.Errorf("nemo: unrecognized status %q", resp.Status) + } + + return &resp, nil +} + +// quoteJSONString JSON-quotes b once into a RawMessage so the outer marshal +// can splice it without re-encoding the text as a Go string. +func quoteJSONString(b []byte) (json.RawMessage, error) { + s := rawJSONString(b) + quoted, err := json.Marshal(&s) + if err != nil { + return nil, err + } + return quoted, nil +} + +// rawJSONString quotes a []byte as a JSON string. encoding/json's +// TextMarshaler path quotes the bytes directly, without string(b). +type rawJSONString []byte + +func (s *rawJSONString) MarshalText() ([]byte, error) { + return *s, nil +} + +// nonNilConfigIDs avoids marshaling a nil slice as JSON null. +func nonNilConfigIDs(configIDs []string) []string { + if configIDs == nil { + return []string{} + } + return configIDs +} diff --git a/internal/guardrails/external/nemo/transformer_test.go b/internal/guardrails/external/nemo/transformer_test.go new file mode 100644 index 000000000..b3b26dfac --- /dev/null +++ b/internal/guardrails/external/nemo/transformer_test.go @@ -0,0 +1,131 @@ +package nemo + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNeMoTransformer_TransformRequest(t *testing.T) { + t.Run("maps tool name, arguments, and config IDs", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + body, err := transformer.TransformRequest( + "execute_sql", + json.RawMessage(`{"query": "DROP TABLE users"}`), + []string{"tool-safety-v1", "input_checking"}, + ) + require.NoError(t, err) + require.JSONEq(t, `{ + "model": "meta/llama-3.1-8b-instruct", + "messages": [{ + "role": "user", + "name": "execute_sql", + "content": "{\"query\": \"DROP TABLE users\"}", + "config": "tool" + }], + "guardrails": {"config_ids": ["tool-safety-v1", "input_checking"]} + }`, string(body)) + }) + + t.Run("defaults arguments to an empty object when absent", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + body, err := transformer.TransformRequest("no_args_tool", nil, nil) + require.NoError(t, err) + + var got CheckRequest + require.NoError(t, json.Unmarshal(body, &got)) + require.Equal(t, json.RawMessage(`"{}"`), got.Messages[0].Content) + require.Equal(t, []string{}, got.Guardrails.ConfigIDs) + }) + + t.Run("errors without a tool name", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + _, err := transformer.TransformRequest("", json.RawMessage(`{}`), nil) + require.Error(t, err) + }) +} + +func TestNeMoTransformer_TransformResponse(t *testing.T) { + t.Run("maps tool name, text content, and config IDs with assistant role", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + body, err := transformer.TransformResponse( + "execute_sql", + []byte("Query returned 42 rows. Customer emails: alice@example.com, bob@example.com"), + []string{"tool-safety-v1", "pii-detection"}, + ) + require.NoError(t, err) + require.JSONEq(t, `{ + "model": "meta/llama-3.1-8b-instruct", + "messages": [{ + "role": "assistant", + "name": "execute_sql", + "content": "Query returned 42 rows. Customer emails: alice@example.com, bob@example.com" + }], + "guardrails": {"config_ids": ["tool-safety-v1", "pii-detection"]} + }`, string(body)) + }) + + t.Run("JSON-quotes response text with special characters", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + body, err := transformer.TransformResponse("echo", []byte("say \"hello\"\n"), nil) + require.NoError(t, err) + require.JSONEq(t, `{ + "model": "meta/llama-3.1-8b-instruct", + "messages": [{"role": "assistant", "name": "echo", "content": "say \"hello\"\n"}], + "guardrails": {"config_ids": []} + }`, string(body)) + }) + + t.Run("does not set the tool config tag on response checks", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + body, err := transformer.TransformResponse("execute_sql", []byte("ok"), nil) + require.NoError(t, err) + require.NotContains(t, string(body), `"config"`) + }) + + t.Run("errors without a tool name", func(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + _, err := transformer.TransformResponse("", []byte("ok"), nil) + require.Error(t, err) + }) +} + +func TestNeMoTransformer_ParseCheckResponse(t *testing.T) { + transformer := NewTransformer("meta/llama-3.1-8b-instruct") + + t.Run("parses a success verdict", func(t *testing.T) { + resp, err := transformer.ParseCheckResponse([]byte(`{"status":"success","content":"ok","rail":null}`)) + require.NoError(t, err) + require.Equal(t, &CheckResponse{Status: StatusSuccess, Content: "ok"}, resp) + }) + + t.Run("parses a modified verdict with the substituted content and triggering rail", func(t *testing.T) { + resp, err := transformer.ParseCheckResponse([]byte(`{"status":"modified","content":"[REDACTED]","rail":"pii-detection"}`)) + require.NoError(t, err) + require.Equal(t, &CheckResponse{Status: StatusModified, Content: "[REDACTED]", Rail: "pii-detection"}, resp) + }) + + t.Run("parses a blocked verdict", func(t *testing.T) { + resp, err := transformer.ParseCheckResponse([]byte(`{"status":"blocked","content":"I can't share that.","rail":"pii-detection"}`)) + require.NoError(t, err) + require.Equal(t, StatusBlocked, resp.Status) + }) + + t.Run("errors on an unrecognized status", func(t *testing.T) { + _, err := transformer.ParseCheckResponse([]byte(`{"status":"unknown"}`)) + require.Error(t, err) + }) + + t.Run("errors on malformed JSON", func(t *testing.T) { + _, err := transformer.ParseCheckResponse([]byte(`not json`)) + require.Error(t, err) + }) +} diff --git a/internal/guardrails/secret.go b/internal/guardrails/secret.go index 8442ee25c..df3260c1f 100644 --- a/internal/guardrails/secret.go +++ b/internal/guardrails/secret.go @@ -51,6 +51,9 @@ func EnsureNeMoConfigData(secretType corev1.SecretType, data map[string][]byte) if err != nil || parsed.Scheme == "" || parsed.Host == "" { return nil, fmt.Errorf("%s: url %q is not a valid absolute URL", configDataKey, cfg.URL) } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return nil, fmt.Errorf("%s: url scheme must be http or https, got %q", configDataKey, parsed.Scheme) + } if cfg.Model == "" { return nil, fmt.Errorf("%s: model is required", configDataKey) diff --git a/internal/guardrails/secret_test.go b/internal/guardrails/secret_test.go index 69d1a9444..d9dcf54ea 100644 --- a/internal/guardrails/secret_test.go +++ b/internal/guardrails/secret_test.go @@ -67,6 +67,28 @@ model: meta/llama-3.1-8b-instruct require.Error(t, err) }) + t.Run("errors when url scheme is not http or https", func(t *testing.T) { + _, err := EnsureNeMoConfigData(SecretTypeNeMo, map[string][]byte{ + configDataKey: []byte(` +url: ftp://nemo-guardrails.internal:8080 +model: meta/llama-3.1-8b-instruct +`), + }) + require.Error(t, err) + require.Contains(t, err.Error(), "http or https") + }) + + t.Run("accepts an http url", func(t *testing.T) { + cfg, err := EnsureNeMoConfigData(SecretTypeNeMo, map[string][]byte{ + configDataKey: []byte(` +url: http://nemo-guardrails.internal:8080 +model: meta/llama-3.1-8b-instruct +`), + }) + require.NoError(t, err) + require.Equal(t, "http://nemo-guardrails.internal:8080", cfg.URL) + }) + t.Run("errors when model is missing", func(t *testing.T) { _, err := EnsureNeMoConfigData(SecretTypeNeMo, map[string][]byte{ configDataKey: []byte(`url: https://nemo-guardrails.internal:8080`),