fix(contextual_access): redact post-hook content alongside output [PLT-3790]

.dockerignore

@@ -11,7 +11,7 @@ LICENSE
 
 # Compiled binaries
 advanced_server
-examples/advanced_server/advanced_server
+examples/contextual_access/advanced_server/advanced_server
 
 # IDE / editor
 .vscode

README.md

@@ -46,6 +46,8 @@ These servers implement webhook endpoints that integrate with an engine's hook s
 | `POST /pre` | Validate/modify tool inputs before execution |
 | `POST /post` | Validate/modify tool outputs after execution |
 
+Post-hook requests from remote MCP servers also carry `content`, the content blocks the server returned alongside `output`. The examples that redact or filter `output` apply the same change to `content` and return it as `override.content`. MCP gateways render only text blocks, so pii_redactor and content_filter handle just those; advanced_server handles every block.
+
 ## Architecture
 
 ```

examples/contextual_access/ab_testing/Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=$BUILDPLATFORM golang:1.25-alpine AS builder
+FROM --platform=$BUILDPLATFORM golang:1.26-alpine AS builder
 
 ARG TARGETOS
 ARG TARGETARCH
@@ -11,15 +11,15 @@ RUN go mod download
 
 # Copy shared package and example source
 COPY pkg/ pkg/
-COPY examples/ab_testing/ examples/ab_testing/
+COPY examples/contextual_access/ab_testing/ examples/contextual_access/ab_testing/
 
 RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \
-    go build -ldflags="-s -w" -trimpath -o /bin/server ./examples/ab_testing
+    go build -ldflags="-s -w" -trimpath -o /bin/server ./examples/contextual_access/ab_testing
 
 FROM gcr.io/distroless/static-debian12
 
 COPY --from=builder /bin/server /bin/server
 
-EXPOSE 8080
+EXPOSE 8888
 
 ENTRYPOINT ["/bin/server"]

examples/contextual_access/ab_testing/README.md

@@ -14,7 +14,7 @@ A minimal hook server that demonstrates how to **A/B test and canary-deploy tool
 
 ```bash
 # Run with experiment config
-go run ./examples/ab_testing -config experiments.yaml
+go run ./examples/contextual_access/ab_testing -config ./examples/contextual_access/ab_testing/example-config.yaml
 ```
 
 ## Config File Format
@@ -79,7 +79,7 @@ experiments:
 
 ```bash
 # Start with example config
-go run ./examples/ab_testing -config experiments.yaml &
+go run ./examples/contextual_access/ab_testing -config ./examples/contextual_access/ab_testing/example-config.yaml &
 
 # Send pre-hook requests for different users
 for i in $(seq 1 20); do

examples/contextual_access/ab_testing/main.go

@@ -8,7 +8,7 @@
 //
 // Usage:
 //
-//	go run ./examples/ab_testing -port 8888 -config experiments.yaml
+//	go run ./examples/contextual_access/ab_testing -port 8888 -config ./examples/contextual_access/ab_testing/example-config.yaml
 package main
 
 import (

examples/contextual_access/advanced_server/Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=$BUILDPLATFORM golang:1.25-alpine AS builder
+FROM --platform=$BUILDPLATFORM golang:1.26-alpine AS builder
 
 ARG TARGETOS
 ARG TARGETARCH
@@ -11,15 +11,15 @@ RUN go mod download
 
 # Copy shared package and example source
 COPY pkg/ pkg/
-COPY examples/advanced_server/ examples/advanced_server/
+COPY examples/contextual_access/advanced_server/ examples/contextual_access/advanced_server/
 
 RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \
-    go build -ldflags="-s -w" -trimpath -o /bin/server ./examples/advanced_server
+    go build -ldflags="-s -w" -trimpath -o /bin/server ./examples/contextual_access/advanced_server
 
 FROM gcr.io/distroless/static-debian12
 
 COPY --from=builder /bin/server /bin/server
 
-EXPOSE 8080
+EXPOSE 8888
 
 ENTRYPOINT ["/bin/server"]

examples/contextual_access/advanced_server/README.md

@@ -7,7 +7,7 @@ A comprehensive hook server with a web dashboard for managing access rules, PII
 ### 1. Basic Rules (Access, Pre, Post)
 - **Access control**: Block users, toolkits, or specific tools from being visible
 - **Pre-execution rules**: Block or modify tool requests before execution
-- **Post-execution rules**: Block or modify tool responses after execution
+- **Post-execution rules**: Block or modify tool responses after execution. A rule that overrides the output also clears the server's `content` blocks, so clients get the new output instead of the original text.
 - **Pattern matching**: Exact, glob (`*`), and regex (`~pattern`) patterns
 - **Input/output matching**: Filter based on request content
 
@@ -35,16 +35,16 @@ A comprehensive hook server with a web dashboard for managing access rules, PII
 
 ```bash
 # Run with defaults (port 8888, no auth)
-go run ./examples/advanced_server
+go run ./examples/contextual_access/advanced_server
 
 # Run with a configuration file
-go run ./examples/advanced_server -config ./examples/advanced_server/example-config.yaml
+go run ./examples/contextual_access/advanced_server -config ./examples/contextual_access/advanced_server/example-config.yaml
 
 # Run with authentication
-go run ./examples/advanced_server -token "my-secret-token"
+go run ./examples/contextual_access/advanced_server -token "my-secret-token"
 
 # Run with TLS
-go run ./examples/advanced_server -tls -cert server.crt -key server.key
+go run ./examples/contextual_access/advanced_server -tls -cert server.crt -key server.key
 ```
 
 Then open `http://localhost:8888/` in your browser to access the dashboard.
@@ -104,7 +104,7 @@ See [example-config.yaml](example-config.yaml) for a full example with all optio
 
 ## PII Redaction Details
 
-The PII redactor scans all string values in tool response outputs. When PII is detected:
+The PII redactor scans all string values in tool response outputs, and the string fields of any `content` blocks (sent by remote MCP servers), returning the redacted blocks as `override.content`. Base64 payloads (`data`, `blob`) are left as-is. When PII is detected:
 
 - **Redact mode**: Replaces PII with labeled placeholders (e.g., `[EMAIL REDACTED]`)
 - **Block mode**: Returns an error response instead of the tool output

examples/contextual_access/advanced_server/ab_testing.go

@@ -17,18 +17,18 @@ import (
 // ABTestManager manages experiment state and variant assignment.
 type ABTestManager struct {
 	mu          sync.RWMutex
-	assignments map[string]string          // "user:experiment" -> variant name
+	assignments map[string]string           // "user:experiment" -> variant name
 	stats       map[string]*ExperimentStats // experiment name -> stats
 }
 
 // ExperimentStats tracks usage statistics for an experiment.
 type ExperimentStats struct {
-	Name            string                  `json:"name"`
-	TotalRequests   int                     `json:"total_requests"`
-	VariantCounts   map[string]int          `json:"variant_counts"`
-	UniqueUsers     map[string]map[string]bool `json:"-"` // variant -> set of user IDs (not serialised)
-	VariantUsers    map[string]int          `json:"variant_users"` // variant -> unique user count
-	LastRequestTime *time.Time              `json:"last_request_time,omitempty"`
+	Name            string                     `json:"name"`
+	TotalRequests   int                        `json:"total_requests"`
+	VariantCounts   map[string]int             `json:"variant_counts"`
+	UniqueUsers     map[string]map[string]bool `json:"-"`             // variant -> set of user IDs (not serialised)
+	VariantUsers    map[string]int             `json:"variant_users"` // variant -> unique user count
+	LastRequestTime *time.Time                 `json:"last_request_time,omitempty"`
 }
 
 // NewABTestManager creates a new A/B test manager.
@@ -208,11 +208,11 @@ type RegistryResponse struct {
 
 // arcadeToolResponse represents a single tool from the Arcade engine API.
 type arcadeToolResponse struct {
-	Name               string                 `json:"name"`
-	Description        string                 `json:"description"`
-	FullyQualifiedName string                 `json:"fully_qualified_name"`
-	QualifiedName      string                 `json:"qualified_name"`
-	Toolkit            arcadeToolkitResponse  `json:"toolkit"`
+	Name               string                `json:"name"`
+	Description        string                `json:"description"`
+	FullyQualifiedName string                `json:"fully_qualified_name"`
+	QualifiedName      string                `json:"qualified_name"`
+	Toolkit            arcadeToolkitResponse `json:"toolkit"`
 }
 
 // arcadeToolkitResponse represents toolkit info nested in a tool response.

examples/contextual_access/advanced_server/example-config.yaml

@@ -3,7 +3,7 @@
 # This file demonstrates all available configuration options.
 # The server hot-reloads this file when it changes.
 #
-# Usage: go run ./examples/advanced_server -config example-config.yaml
+# Usage: go run ./examples/contextual_access/advanced_server -config ./examples/contextual_access/advanced_server/example-config.yaml
 
 # Health endpoint configuration
 health:

examples/contextual_access/advanced_server/main.go

@@ -8,8 +8,8 @@
 //
 // Usage:
 //
-//	go run ./examples/advanced_server -port 8888 -config config.yaml
-//	go run ./examples/advanced_server -port 8888 -token secret123
+//	go run ./examples/contextual_access/advanced_server -port 8888 -config ./examples/contextual_access/advanced_server/example-config.yaml
+//	go run ./examples/contextual_access/advanced_server -port 8888 -token secret123
 package main
 
 import (
@@ -440,7 +440,7 @@ func (s *HookServer) evaluatePostRules(req server.PostHookRequest) (*server.Post
 	ruleMatch := ""
 	for i, rule := range postCfg.Rules {
 		if s.matchPostRule(rule, userID, req) {
-			result = s.applyPostRule(rule)
+			result = s.applyPostRule(rule, req)
 			ruleMatch = fmt.Sprintf("post.rules[%d]", i)
 			break
 		}
@@ -455,30 +455,36 @@ func (s *HookServer) evaluatePostRules(req server.PostHookRequest) (*server.Post
 
 	// Always apply PII redaction on top of whatever result we have.
 	// PII is a security/compliance feature and should never be bypassed by rules.
-	// Scan both inputs and output for PII — inputs may contain sensitive data
-	// that the tool could echo back, and output may not always be populated.
+	// Scan inputs, output, and content for PII — inputs may contain sensitive
+	// data that the tool could echo back, output may not always be populated,
+	// and content (the blocks a remote server sent alongside the output)
+	// reaches the client too.
 	piiFound := false
 	if cfg.PII != nil && cfg.PII.Enabled {
-		hasContent := req.Output != nil || (req.Inputs != nil && len(*req.Inputs) > 0)
-		if hasContent {
+		hasData := req.Output != nil || req.Content != nil || (req.Inputs != nil && len(*req.Inputs) > 0)
+		if hasData {
 			detector := NewPIIDetector(cfg.PII)
 
-			// Scan both output and inputs for PII
-			var outputScan, inputScan PIIScanResult
+			// Scan output, content, and inputs for PII
+			var outputScan, contentScan, inputScan PIIScanResult
 			if req.Output != nil {
 				outputScan = detector.ScanAndSummarizeAny(req.Output)
 			}
+			if req.Content != nil {
+				contentScan = detector.ScanAndSummarizeContent(*req.Content)
+			}
 			if req.Inputs != nil {
 				inputScan = detector.ScanAndSummarizeAny(*req.Inputs)
 			}
+			outputPII := outputScan.ContainsPII || contentScan.ContainsPII
 
-			if outputScan.ContainsPII || inputScan.ContainsPII {
+			if outputPII || inputScan.ContainsPII {
 				piiFound = true
 
 				if cfg.PII.Action == "block" {
 					// Block the response entirely, regardless of rule result
 					errMsg := "Response blocked: PII detected"
-					if outputScan.ContainsPII && inputScan.ContainsPII {
+					if outputPII && inputScan.ContainsPII {
 						errMsg = "Response blocked: PII detected in inputs and output"
 					} else if inputScan.ContainsPII {
 						errMsg = "Response blocked: PII detected in inputs"
@@ -510,6 +516,20 @@ func (s *HookServer) evaluatePostRules(req server.PostHookRequest) (*server.Post
 					}
 					result.Override.Output = redacted
 				}
+
+				// Redact content too, or the server's original text reaches the
+				// client. As with output, a rule's override takes precedence.
+				contentToRedact := req.Content
+				if result.Override != nil && result.Override.Content != nil {
+					contentToRedact = result.Override.Content
+				}
+				if contentToRedact != nil {
+					redacted := detector.RedactContent(*contentToRedact)
+					if result.Override == nil {
+						result.Override = &server.PostHookOverride{}
+					}
+					result.Override.Content = &redacted
+				}
 				return result, joinRuleMatch(ruleMatch, "pii:redact"), true
 			}
 		}
@@ -544,7 +564,7 @@ func (s *HookServer) matchPostRule(rule PostRule, userID string, req server.Post
 	return true
 }
 
-func (s *HookServer) applyPostRule(rule PostRule) *server.PostHookResult {
+func (s *HookServer) applyPostRule(rule PostRule, req server.PostHookRequest) *server.PostHookResult {
 	result := &server.PostHookResult{
 		Code: actionToCode(rule.Action),
 	}
@@ -559,6 +579,11 @@ func (s *HookServer) applyPostRule(rule PostRule) *server.PostHookResult {
 			result.Override = &server.PostHookOverride{
 				Output: output,
 			}
+			// Clear the server's content blocks so clients get the
+			// replacement output (as text) instead of the original result.
+			if req.Content != nil {
+				result.Override.Content = &[]server.ContentBlock{}
+			}
 		}
 	}
 

examples/contextual_access/advanced_server/pii.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"regexp"
 	"strings"
+
+	"github.com/ArcadeAI/logical-extensions-examples/pkg/server"
 )
 
 // =============================================================================
@@ -40,7 +42,7 @@ func NewPIIDetector(cfg *PIIConfig) *PIIDetector {
 		d.labels["ssn"] = "[SSN REDACTED]"
 	}
 	if cfg.Types.Phone {
-		d.patterns["phone"] = regexp.MustCompile(`\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b`)
+		d.patterns["phone"] = regexp.MustCompile(`(?:\+?\b1[-.\s]?\(?|\(|\b)\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b`)
 		d.labels["phone"] = "[PHONE REDACTED]"
 	}
 	if cfg.Types.CreditCard {
@@ -193,6 +195,21 @@ func (d *PIIDetector) ScanAndSummarize(data map[string]interface{}) PIIScanResul
 func (d *PIIDetector) ScanAndSummarizeAny(data interface{}) PIIScanResult {
 	var matches []PIIMatch
 	d.scanValue(data, "", &matches)
+	return summarize(matches)
+}
+
+// ScanAndSummarizeContent scans the string fields of post-hook content blocks
+// (text, uri, annotations, _meta, and an embedded resource's text) for PII and
+// returns a summary.
+func (d *PIIDetector) ScanAndSummarizeContent(blocks []server.ContentBlock) PIIScanResult {
+	var matches []PIIMatch
+	for i, b := range blocks {
+		d.scanBlockFields(b.AdditionalProperties, fmt.Sprintf("content[%d]", i), &matches)
+	}
+	return summarize(matches)
+}
+
+func summarize(matches []PIIMatch) PIIScanResult {
 	counts := make(map[string]int)
 	for _, m := range matches {
 		counts[m.Type]++
@@ -208,3 +225,50 @@ func (d *PIIDetector) ScanAndSummarizeAny(data interface{}) PIIScanResult {
 func (d *PIIDetector) RedactAny(data interface{}) interface{} {
 	return d.redactValue(data)
 }
+
+// RedactContent redacts PII from the string fields of post-hook content blocks,
+// keeping each block's type and base64 payloads unchanged.
+func (d *PIIDetector) RedactContent(blocks []server.ContentBlock) []server.ContentBlock {
+	result := make([]server.ContentBlock, len(blocks))
+	for i, b := range blocks {
+		result[i] = server.ContentBlock{Type: b.Type, AdditionalProperties: d.redactBlockFields(b.AdditionalProperties)}
+	}
+	return result
+}
+
+// isBinaryField reports whether a content block field holds a base64 payload
+// (image and audio "data", a resource's "blob"). Regexes could corrupt these,
+// so they are left as-is.
+func isBinaryField(key string) bool {
+	return key == "data" || key == "blob"
+}
+
+func (d *PIIDetector) scanBlockFields(fields map[string]interface{}, path string, matches *[]PIIMatch) {
+	for key, val := range fields {
+		if isBinaryField(key) {
+			continue
+		}
+		newPath := path + "." + key
+		if resource, ok := val.(map[string]interface{}); ok && key == "resource" {
+			d.scanBlockFields(resource, newPath, matches)
+			continue
+		}
+		d.scanValue(val, newPath, matches)
+	}
+}
+
+func (d *PIIDetector) redactBlockFields(fields map[string]interface{}) map[string]interface{} {
+	result := make(map[string]interface{}, len(fields))
+	for key, val := range fields {
+		resource, isMap := val.(map[string]interface{})
+		switch {
+		case isBinaryField(key):
+			result[key] = val
+		case isMap && key == "resource":
+			result[key] = d.redactBlockFields(resource)
+		default:
+			result[key] = d.redactValue(val)
+		}
+	}
+	return result
+}