From 7e58eac0a03f7cd4fda7f2466795a79c8237ea49 Mon Sep 17 00:00:00 2001 From: Chirag Date: Mon, 24 Aug 2026 17:05:23 +0530 Subject: [PATCH 1/2] fix(output): redact findings after output scan --- sdk/src/unplug/pipelines/output.py | 28 ++++++++++----- sdk/tests/integration/test_guard_v2.py | 50 ++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/sdk/src/unplug/pipelines/output.py b/sdk/src/unplug/pipelines/output.py index b2000b7..4030593 100644 --- a/sdk/src/unplug/pipelines/output.py +++ b/sdk/src/unplug/pipelines/output.py @@ -10,6 +10,7 @@ from unplug.core.agent.boundaries import strip_boundary_markers from unplug.core.context import ExecutionContext from unplug.core.privacy.secrets import SecretsSanitizer +from unplug.core.redaction import apply_span_redactions from unplug.core.runtime.stats import MetricsCollector from unplug.core.taint import TaintedText, TrustLevel from unplug.models import Finding, ScanResult @@ -94,14 +95,25 @@ def _redact( *, policy: ScanPolicy | None = None, ) -> str | None: - _ = policy + resolved_policy = policy or self._config.policy text = self._extract_text(input_data) if text is None or not findings: return None - if self._sanitizer: - return self._sanitizer.sanitize(text).clean_text - return super()._redact( - input_data, - findings, - policy=policy or self._config.policy, - ) + if self._sanitizer is None: + return apply_span_redactions(text, findings, resolved_policy) + # Finding spans refer to the original text, so they must be applied + # before sanitization (whose replacements change string lengths). + # Spans already covered by a detected secret are left to the sanitizer, + # which keeps its named [REDACTED:] placeholder. + secret_spans = [ + (m.span_start, m.span_end) for m in self._sanitizer.sanitize(text).secrets_found + ] + span_findings = [ + f + for f in findings + if not any(f.span_start < end and start < f.span_end for start, end in secret_spans) + ] + span_redacted = apply_span_redactions(text, span_findings, resolved_policy) + if span_redacted is None: + span_redacted = text + return self._sanitizer.sanitize(span_redacted).clean_text diff --git a/sdk/tests/integration/test_guard_v2.py b/sdk/tests/integration/test_guard_v2.py index cddb00c..5d0569a 100644 --- a/sdk/tests/integration/test_guard_v2.py +++ b/sdk/tests/integration/test_guard_v2.py @@ -1,6 +1,8 @@ """Tests for Guard v2: backward compatibility + new methods.""" from unplug import Guard, ScanResult, TaintedText, TrustLevel +from unplug.config.guard import GuardConfig +from unplug.config.policy import RedactionMode, ScanPolicy from unplug.core.privacy.secrets import SecretsRegistry from unplug.models import Source @@ -81,6 +83,54 @@ def test_accepts_tainted_text(self): result = guard.scan_output(text) assert result.safe is True + def test_redacts_ip_literal_url(self): + guard = Guard() + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert any( + f.category == "urls" and f.subcategory == "ip_literal_url" for f in result.findings + ) + assert result.redacted_text is not None + assert "http://198.51.100.7/collect" not in result.redacted_text + assert "[BLOCKED:urls]" in result.redacted_text + + def test_redacts_url_after_registered_secret(self): + registry = SecretsRegistry() + registry.register("TESTKEY", "sk-test-abc123secretvalue99") + guard = Guard(secrets_registry=registry) + result = guard.scan_output( + "Key: sk-test-abc123secretvalue99 Link: http://198.51.100.7/collect" + ) + assert result.redacted_text is not None + assert "sk-test-abc123secretvalue99" not in result.redacted_text + assert "[REDACTED:TESTKEY]" in result.redacted_text + assert "198.51.100.7" not in result.redacted_text + + def test_redacts_leakage_finding(self): + guard = Guard() + result = guard.scan_output("see this email: user@example.com now") + assert any(f.category == "leakage" for f in result.findings) + assert result.redacted_text is not None + assert "user@example.com" not in result.redacted_text + + def test_url_redaction_mode_strip(self): + guard = Guard(config=GuardConfig(policy=ScanPolicy(redaction_mode=RedactionMode.STRIP))) + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert result.redacted_text == "Link: /collect" + + def test_url_redaction_mode_redacted_tags(self): + guard = Guard( + config=GuardConfig(policy=ScanPolicy(redaction_mode=RedactionMode.REDACTED_TAGS)) + ) + result = guard.scan_output("Link: http://198.51.100.7/collect") + # the url finding supplies its own replacement text + assert result.redacted_text == "Link: [BLOCKED:url]/collect" + + def test_redaction_mode_none_keeps_redacted_text_absent(self): + guard = Guard(config=GuardConfig(policy=ScanPolicy(redaction_mode=RedactionMode.NONE))) + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert any(f.category == "urls" for f in result.findings) + assert result.redacted_text is None + class TestCheckToolCall: def test_detects_destructive_tool(self): From db207854ab1f0207f752c6273b58dd7394a4e82e Mon Sep 17 00:00:00 2001 From: Chirag Date: Mon, 24 Aug 2026 17:25:40 +0530 Subject: [PATCH 2/2] fix(output): preserve redaction around secret overlaps --- sdk/src/unplug/pipelines/output.py | 48 ++++++++++++++++++---- sdk/tests/integration/test_guard_v2.py | 30 ++++++++++++++ sdk/tests/unit/pipelines/test_pipelines.py | 41 ++++++++++++++++++ 3 files changed, 111 insertions(+), 8 deletions(-) diff --git a/sdk/src/unplug/pipelines/output.py b/sdk/src/unplug/pipelines/output.py index 4030593..55b2ce1 100644 --- a/sdk/src/unplug/pipelines/output.py +++ b/sdk/src/unplug/pipelines/output.py @@ -18,6 +18,30 @@ from unplug.scanners.base import BaseScanner +def _subtract_spans( + finding: Finding, + secret_spans: list[tuple[int, int]], +) -> list[Finding]: + """Split a finding into copies covering the parts outside secret spans. + + Secret spans must be sorted and merged. Spans are half-open, matching + ``apply_span_redactions()``. + """ + residuals: list[tuple[int, int]] = [] + cursor = finding.span_start + for start, end in secret_spans: + if end <= start or start >= finding.span_end or end <= cursor: + continue + if start > cursor: + residuals.append((cursor, start)) + cursor = end + if cursor >= finding.span_end: + break + if cursor < finding.span_end: + residuals.append((cursor, finding.span_end)) + return [finding.model_copy(update={"span_start": s, "span_end": e}) for s, e in residuals] + + class OutputPipeline(BasePipeline): name = "output" @@ -103,15 +127,23 @@ def _redact( return apply_span_redactions(text, findings, resolved_policy) # Finding spans refer to the original text, so they must be applied # before sanitization (whose replacements change string lengths). - # Spans already covered by a detected secret are left to the sanitizer, - # which keeps its named [REDACTED:] placeholder. - secret_spans = [ - (m.span_start, m.span_end) for m in self._sanitizer.sanitize(text).secrets_found - ] + # Secret spans are subtracted from finding spans so policy redaction + # covers the residual parts while the sanitizer keeps its named + # [REDACTED:] placeholder for the secret itself. + secret_spans = sorted( + {(m.span_start, m.span_end) for m in self._sanitizer.sanitize(text).secrets_found} + ) + merged_secret_spans: list[tuple[int, int]] = [] + for start, end in secret_spans: + if end <= start: + continue + if merged_secret_spans and start <= merged_secret_spans[-1][1]: + prev = merged_secret_spans[-1] + merged_secret_spans[-1] = (prev[0], max(prev[1], end)) + else: + merged_secret_spans.append((start, end)) span_findings = [ - f - for f in findings - if not any(f.span_start < end and start < f.span_end for start, end in secret_spans) + residual for f in findings for residual in _subtract_spans(f, merged_secret_spans) ] span_redacted = apply_span_redactions(text, span_findings, resolved_policy) if span_redacted is None: diff --git a/sdk/tests/integration/test_guard_v2.py b/sdk/tests/integration/test_guard_v2.py index 5d0569a..fdafb3a 100644 --- a/sdk/tests/integration/test_guard_v2.py +++ b/sdk/tests/integration/test_guard_v2.py @@ -131,6 +131,36 @@ def test_redaction_mode_none_keeps_redacted_text_absent(self): assert any(f.category == "urls" for f in result.findings) assert result.redacted_text is None + def test_secret_inside_url_finding_redacts_both(self): + # The registered value sits inside the ip_literal_url finding span + # (http://198.51.100.7), so this covers both residual sides of the + # split finding: the prefix and the suffix around the secret. + registry = SecretsRegistry() + registry.register("HOSTID", "198.51.100") + guard = Guard(secrets_registry=registry) + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert result.redacted_text == "Link: [BLOCKED:urls][REDACTED:HOSTID][BLOCKED:urls]/collect" + + def test_secret_inside_url_strip_mode(self): + registry = SecretsRegistry() + registry.register("HOSTID", "198.51.100") + guard = Guard( + secrets_registry=registry, + config=GuardConfig(policy=ScanPolicy(redaction_mode=RedactionMode.STRIP)), + ) + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert result.redacted_text == "Link: [REDACTED:HOSTID]/collect" + + def test_secret_inside_url_redacted_tags_mode(self): + registry = SecretsRegistry() + registry.register("HOSTID", "198.51.100") + guard = Guard( + secrets_registry=registry, + config=GuardConfig(policy=ScanPolicy(redaction_mode=RedactionMode.REDACTED_TAGS)), + ) + result = guard.scan_output("Link: http://198.51.100.7/collect") + assert result.redacted_text == "Link: [BLOCKED:url][REDACTED:HOSTID][BLOCKED:url]/collect" + class TestCheckToolCall: def test_detects_destructive_tool(self): diff --git a/sdk/tests/unit/pipelines/test_pipelines.py b/sdk/tests/unit/pipelines/test_pipelines.py index 5341e9a..eb48ee8 100644 --- a/sdk/tests/unit/pipelines/test_pipelines.py +++ b/sdk/tests/unit/pipelines/test_pipelines.py @@ -161,6 +161,47 @@ def test_low_confidence_finding_allows(self): action = pipeline._decide(0.1, [low], text_len=20, policy=pipeline.config.policy) assert action == Action.ALLOW + def test_subtract_spans_splits_finding_around_secrets(self): + from unplug.models import Finding + from unplug.pipelines.output import _subtract_spans + + finding = Finding( + category="urls", + subcategory="ip_literal_url", + stage="regex", + span_start=10, + span_end=40, + score=0.7, + evidence="two secrets inside", + replacement="[BLOCKED:url]", + ) + # two disjoint secret regions inside the finding span + residuals = _subtract_spans(finding, [(12, 18), (25, 30)]) + assert [(r.span_start, r.span_end) for r in residuals] == [(10, 12), (18, 25), (30, 40)] + for residual in residuals: + assert residual.category == finding.category + assert residual.score == finding.score + assert residual.replacement == finding.replacement + + def test_subtract_spans_fully_covered_finding(self): + from unplug.models import Finding + from unplug.pipelines.output import _subtract_spans + + finding = Finding( + category="urls", + subcategory="ip_literal_url", + stage="regex", + span_start=5, + span_end=15, + score=0.7, + evidence="exact match", + ) + assert _subtract_spans(finding, [(5, 15)]) == [] + assert _subtract_spans(finding, [(0, 20)]) == [] + assert _subtract_spans(finding, []) == [finding] + # invalid and zero-length secret spans are ignored + assert _subtract_spans(finding, [(7, 7), (9, 4)]) == [finding] + class TestToolCallPipeline: def test_detects_destructive(self):