feat(pii): restore request-scoped pseudonyms (#11272)

* feat(pii): restore request-scoped pseudonyms

Replace masked request values with unique per-request tokens when response restoration is enabled, then restore them across JSON and SSE write boundaries. Document the opt-in model setting and expose it in config metadata.\n\nAssisted-by: Codex:gpt-5

* fix(pii): wrap reversible redaction tokens

Use configurable token delimiters to avoid restoring ordinary model text that happens to match an internal identifier. Rename the option and document the confidentiality tradeoff.

Assisted-by: Codex:gpt-5

---------

Co-authored-by: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com>

Author: localai-org-maint-bot

core/config/meta/registry.go

@@ -816,6 +816,25 @@ func DefaultRegistry() map[string]FieldMetaOverride {
 			AutocompleteProvider: "models:token_classify",
 			Order:                201,
 		},
+		"pii.reversible_redactions": {
+			Section:     "pii",
+			Label:       "Reversible Redactions",
+			Description: "Replace masked values with wrapped request-scoped tokens and restore them when the model returns those tokens. Supports streaming responses and never persists the substitution map.",
+			Component:   "toggle",
+			Order:       202,
+		},
+		"pii.reversible_token_prefix": {
+			Section:     "pii",
+			Label:       "Reversible Token Prefix",
+			Description: "Prefix for reversible redaction tokens. Defaults to [REDACTED:.",
+			Order:       203,
+		},
+		"pii.reversible_token_suffix": {
+			Section:     "pii",
+			Label:       "Reversible Token Suffix",
+			Description: "Suffix for reversible redaction tokens. Defaults to ].",
+			Order:       204,
+		},
 
 		// --- PII detection policy (on a token_classify detector model) ---
 		"pii_detection.min_score": {

core/config/model_config.go

@@ -440,8 +440,18 @@ type PIIConfig struct {
 	// model just opts in by listing detectors. Multiple detectors union
 	// their hits; overlapping spans resolve to the strongest action.
 	Detectors []string `yaml:"detectors,omitempty" json:"detectors,omitempty"`
+
+	// ReversibleRedactions replaces request PII with stable, request-scoped
+	// tokens and restores those values when the wrapped tokens appear in the response.
+	ReversibleRedactions  bool   `yaml:"reversible_redactions,omitempty" json:"reversible_redactions,omitempty"`
+	ReversibleTokenPrefix string `yaml:"reversible_token_prefix,omitempty" json:"reversible_token_prefix,omitempty"`
+	ReversibleTokenSuffix string `yaml:"reversible_token_suffix,omitempty" json:"reversible_token_suffix,omitempty"`
 }
 
+func (c ModelConfig) PIIReversibleRedactions() bool    { return c.PII.ReversibleRedactions }
+func (c ModelConfig) PIIReversibleTokenPrefix() string { return c.PII.ReversibleTokenPrefix }
+func (c ModelConfig) PIIReversibleTokenSuffix() string { return c.PII.ReversibleTokenSuffix }
+
 // @Description Detection policy for a token-classification (NER) model
 // used as a PII detector. Lives on the detector model's own config so the
 // model is a self-describing policy unit: consuming models reference it by

core/services/routing/pii/middleware.go

@@ -194,6 +194,16 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
 
 			texts := adapter.Scan(parsed)
 			updates := make([]ScannedText, 0, len(texts))
+			prefix, suffix := defaultReversibleTokenPrefix, defaultReversibleTokenSuffix
+			if cfg, ok := rawCfg.(responsePIIConfig); ok {
+				if cfg.PIIReversibleTokenPrefix() != "" {
+					prefix = cfg.PIIReversibleTokenPrefix()
+				}
+				if cfg.PIIReversibleTokenSuffix() != "" {
+					suffix = cfg.PIIReversibleTokenSuffix()
+				}
+			}
+			pseudonyms := newPseudonymizer(prefix, suffix)
 			var blocked bool
 			var firstEventID string
 
@@ -259,7 +269,11 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
 				if res.Blocked {
 					blocked = true
 				}
-				updates = append(updates, ScannedText{Index: st.Index, Text: res.Redacted})
+				redacted := res.Redacted
+				if cfg, ok := rawCfg.(responsePIIConfig); ok && cfg.PIIReversibleRedactions() {
+					redacted = pseudonyms.replace(st.Text, res.Spans)
+				}
+				updates = append(updates, ScannedText{Index: st.Index, Text: redacted})
 			}
 
 			if blocked {
@@ -279,7 +293,16 @@ func RequestMiddleware(redactor *Redactor, store EventStore, adapter Adapter, fa
 			if firstEventID != "" {
 				c.Set(ctxKeyPIIEventID, firstEventID)
 			}
-			return next(c)
+			if len(pseudonyms.original) == 0 {
+				return next(c)
+			}
+			writer := newRestoringWriter(c.Response().Writer, pseudonyms.original)
+			c.Response().Writer = writer
+			err := next(c)
+			if finishErr := writer.Finish(); err == nil {
+				err = finishErr
+			}
+			return err
 		}
 	}
 }

core/services/routing/pii/middleware_test.go

@@ -60,10 +60,16 @@ func setRequestOnContext(req *fakeRequest) echo.MiddlewareFunc {
 type fakeModelPIIConfig struct {
 	enabled   bool
 	detectors []string
+	reverse   bool
+	prefix    string
+	suffix    string
 }
 
-func (f fakeModelPIIConfig) PIIIsEnabled() bool     { return f.enabled }
-func (f fakeModelPIIConfig) PIIDetectors() []string { return f.detectors }
+func (f fakeModelPIIConfig) PIIIsEnabled() bool               { return f.enabled }
+func (f fakeModelPIIConfig) PIIDetectors() []string           { return f.detectors }
+func (f fakeModelPIIConfig) PIIReversibleRedactions() bool    { return f.reverse }
+func (f fakeModelPIIConfig) PIIReversibleTokenPrefix() string { return f.prefix }
+func (f fakeModelPIIConfig) PIIReversibleTokenSuffix() string { return f.suffix }
 
 func withModelConfig(cfg fakeModelPIIConfig) echo.MiddlewareFunc {
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -129,6 +135,58 @@ var _ = Describe("RequestMiddleware (NER)", func() {
 		Expect(events[0].Direction).To(Equal(DirectionIn))
 	})
 
+	It("restores distinct pseudonyms across streaming write boundaries", func() {
+		body := &fakeRequest{Messages: []string{"Email alice@example.com or bob@example.com"}}
+		mw := RequestMiddleware(&Redactor{}, store(), fakeAdapter(), nil,
+			WithNERResolver(resolverFor(map[string]NERConfig{
+				"privacy-filter": nerCfg(ActionMask,
+					NEREntity{Group: "EMAIL", Start: 6, End: 23, Score: 0.95},
+					NEREntity{Group: "EMAIL", Start: 27, End: 42, Score: 0.95}),
+			})))
+		e := echo.New()
+		e.POST("/chat", func(c echo.Context) error {
+			Expect(body.Messages[0]).To(Equal("Email [REDACTED:EMAIL_001] or [REDACTED:EMAIL_002]"))
+			_, err := c.Response().Write([]byte(`data: {"delta":"EMAIL_001 and [REDACTED:EMAIL_0`))
+			Expect(err).ToNot(HaveOccurred())
+			_, err = c.Response().Write([]byte(`01] and [REDACTED:EMAIL_002]"}` + "\n\n"))
+			return err
+		}, setRequestOnContext(body), withModelConfig(fakeModelPIIConfig{
+			enabled: true, detectors: []string{"privacy-filter"}, reverse: true,
+		}), mw)
+
+		req := httptest.NewRequest(http.MethodPost, "/chat", strings.NewReader(`{}`))
+		w := httptest.NewRecorder()
+		e.ServeHTTP(w, req)
+
+		Expect(w.Code).To(Equal(http.StatusOK))
+		Expect(w.Body.String()).To(Equal("data: {\"delta\":\"EMAIL_001 and alice@example.com and bob@example.com\"}\n\n"))
+	})
+
+	It("uses configured reversible redaction token delimiters", func() {
+		body := &fakeRequest{Messages: []string{"Email alice@example.com"}}
+		mw := RequestMiddleware(&Redactor{}, store(), fakeAdapter(), nil,
+			WithNERResolver(resolverFor(map[string]NERConfig{
+				"privacy-filter": nerCfg(ActionMask,
+					NEREntity{Group: "EMAIL", Start: 6, End: 23, Score: 0.95}),
+			})))
+		e := echo.New()
+		e.POST("/chat", func(c echo.Context) error {
+			Expect(body.Messages[0]).To(Equal("Email <PII:EMAIL_001>"))
+			_, err := c.Response().Write([]byte(`{"text":"<PII:EMAIL_001>"}`))
+			return err
+		}, setRequestOnContext(body), withModelConfig(fakeModelPIIConfig{
+			enabled: true, detectors: []string{"privacy-filter"}, reverse: true,
+			prefix: "<PII:", suffix: ">",
+		}), mw)
+
+		req := httptest.NewRequest(http.MethodPost, "/chat", strings.NewReader(`{}`))
+		w := httptest.NewRecorder()
+		e.ServeHTTP(w, req)
+
+		Expect(w.Code).To(Equal(http.StatusOK))
+		Expect(w.Body.String()).To(Equal(`{"text":"alice@example.com"}`))
+	})
+
 	It("blocks (400) when a detected entity's action is block", func() {
 		st := store()
 		body := &fakeRequest{Messages: []string{"my password is hunter2 ok"}}

core/services/routing/pii/pseudonymizer.go

@@ -0,0 +1,142 @@
+package pii
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+	"unicode"
+)
+
+type responsePIIConfig interface {
+	PIIReversibleRedactions() bool
+	PIIReversibleTokenPrefix() string
+	PIIReversibleTokenSuffix() string
+}
+
+const (
+	defaultReversibleTokenPrefix = "[REDACTED:"
+	defaultReversibleTokenSuffix = "]"
+)
+
+type pseudonymizer struct {
+	byValue  map[string]string
+	original map[string]string
+	counts   map[string]int
+	prefix   string
+	suffix   string
+}
+
+func newPseudonymizer(prefix, suffix string) *pseudonymizer {
+	return &pseudonymizer{
+		byValue:  map[string]string{},
+		original: map[string]string{},
+		counts:   map[string]int{},
+		prefix:   prefix,
+		suffix:   suffix,
+	}
+}
+
+func (p *pseudonymizer) replace(text string, spans []Span) string {
+	var b strings.Builder
+	last := 0
+	for _, span := range spans {
+		if span.Action != ActionMask || span.Start < last || span.End > len(text) {
+			continue
+		}
+		b.WriteString(text[last:span.Start])
+		value := text[span.Start:span.End]
+		token, ok := p.byValue[value]
+		if !ok {
+			group := pseudonymGroup(span.Pattern)
+			p.counts[group]++
+			token = fmt.Sprintf("%s%s_%03d%s", p.prefix, group, p.counts[group], p.suffix)
+			p.byValue[value] = token
+			p.original[token] = value
+		}
+		b.WriteString(token)
+		last = span.End
+	}
+	b.WriteString(text[last:])
+	return b.String()
+}
+
+func pseudonymGroup(pattern string) string {
+	if i := strings.LastIndexByte(pattern, ':'); i >= 0 {
+		pattern = pattern[i+1:]
+	}
+	var b strings.Builder
+	for _, r := range strings.ToUpper(pattern) {
+		if unicode.IsLetter(r) || unicode.IsDigit(r) {
+			b.WriteRune(r)
+		} else {
+			b.WriteByte('_')
+		}
+	}
+	if b.Len() == 0 {
+		return "PII"
+	}
+	return b.String()
+}
+
+type restoringWriter struct {
+	http.ResponseWriter
+	pending      string
+	replacements map[string]string
+}
+
+func newRestoringWriter(w http.ResponseWriter, originals map[string]string) *restoringWriter {
+	replacements := make(map[string]string, len(originals))
+	for token, original := range originals {
+		encoded, _ := json.Marshal(original)
+		replacements[token] = string(encoded[1 : len(encoded)-1])
+	}
+	return &restoringWriter{ResponseWriter: w, replacements: replacements}
+}
+
+func (w *restoringWriter) Write(data []byte) (int, error) {
+	w.pending += string(data)
+	ready, pending := w.splitReady(w.replace(w.pending))
+	w.pending = pending
+	if ready != "" {
+		if _, err := w.ResponseWriter.Write([]byte(ready)); err != nil {
+			return 0, err
+		}
+	}
+	return len(data), nil
+}
+
+func (w *restoringWriter) Flush() {
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (w *restoringWriter) Finish() error {
+	if w.pending == "" {
+		return nil
+	}
+	_, err := w.ResponseWriter.Write([]byte(w.replace(w.pending)))
+	w.pending = ""
+	return err
+}
+
+func (w *restoringWriter) replace(s string) string {
+	for token, original := range w.replacements {
+		s = strings.ReplaceAll(s, token, original)
+	}
+	return s
+}
+
+func (w *restoringWriter) splitReady(s string) (string, string) {
+	keep := 0
+	for token := range w.replacements {
+		limit := min(len(token)-1, len(s))
+		for n := 1; n <= limit; n++ {
+			if strings.HasSuffix(s, token[:n]) && n > keep {
+				keep = n
+			}
+		}
+	}
+	return s[:len(s)-keep], s[len(s)-keep:]
+}

docs/content/operations/middleware.md

@@ -165,15 +165,33 @@ pii:
   enabled: true               # default-on for cloud-proxy; explicit for audit
   detectors:
     - privacy-filter-multilingual
+  reversible_redactions: true # restore request PII if the model echoes its wrapped token
+  reversible_token_prefix: "[REDACTED:" # optional; this is the default
+  reversible_token_suffix: "]"          # optional; this is the default
 ```
 
+`reversible_redactions` enables bijective, request-scoped replacement. Each
+masked value is sent to the backend as a stable wrapped token such as
+`[REDACTED:EMAIL_001]` instead of a generic redaction marker. If the model includes that token in
+its response, LocalAI restores the original value before returning JSON or SSE
+to the caller. The substitution map exists only for that request and is never
+logged or persisted. Leave the option unset (the default) for irreversible
+`[REDACTED:...]` masking.
+
+The prefix and suffix reduce collisions with ordinary model output and can be
+customized with `reversible_token_prefix` and `reversible_token_suffix`.
+Reversible redactions provide less confidentiality than irreversible masking:
+any third party that can observe both the redacted request and restored response
+may be able to infer the original values.
+
 Multiple detectors **union** their detections; overlapping spans resolve to
 the strongest action (`block` > `mask` > `allow`). A configured detector
 that can't be loaded **fails the request closed** (HTTP 503,
 `error.type=pii_ner_unavailable`) rather than silently skipping the check.
 The same NER path runs on the [MITM proxy]({{< relref "mitm-proxy.md" >}})
-request body for intercepted hosts. Response/output redaction is out of
-scope for now.
+request body for intercepted hosts. Reversible response restoration currently
+applies to LocalAI API routes; the MITM proxy keeps its own output-redaction
+policy.
 
 ### Instance-wide default detector