Fix PII result copy regression

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pii_filter"
-version = "0.3.1"
+version = "0.3.2"
 edition.workspace = true
 authors.workspace = true
 license.workspace = true

plugins/rust/python-package/pii_filter/Cargo.toml

@@ -1,6 +1,6 @@
 description: "Rust-backed PII detection and masking for prompt arguments, tool inputs, and tool outputs"
 author: "ContextForge Contributors"
-version: "0.3.1"
+version: "0.3.2"
 kind: "cpex_pii_filter.pii_filter.PIIFilterPlugin"
 available_hooks:
   - "prompt_pre_fetch"

plugins/rust/python-package/pii_filter/cpex_pii_filter/plugin-manifest.yaml

@@ -5,7 +5,7 @@
 
 use log::{debug, warn};
 use pyo3::prelude::*;
-use pyo3::types::{PyAny, PyDict, PyList, PySet, PyString, PyTuple};
+use pyo3::types::{PyAny, PyDict, PyList, PyMapping, PySet, PyString, PyTuple};
 use pyo3_stub_gen::derive::*;
 use std::collections::HashMap;
 
@@ -357,24 +357,28 @@ impl PIIDetectorRust {
             }
         }
 
-        // Handle dictionaries
-        if let Ok(dict) = data.cast::<PyDict>() {
-            let mut entries: Vec<(Py<PyAny>, Py<PyAny>)> = Vec::with_capacity(dict.len());
+        // Handle mappings through the Python protocol. CPEX isolation wraps
+        // dicts in copy-on-write dict subclasses whose visible entries are not
+        // stored in the underlying PyDict table.
+        if let Ok(mapping) = data.cast::<PyMapping>() {
+            let mapping_len = mapping.len()?;
+            let mut entries: Vec<(Py<PyAny>, Py<PyAny>)> = Vec::with_capacity(mapping_len);
             let mut all_detections = HashMap::new();
-            if dict.len() > self.config.max_collection_items {
+            if mapping_len > self.config.max_collection_items {
                 warn!(
                     "Rejected nested mapping at path '{}' because size {} exceeds max {}",
-                    path,
-                    dict.len(),
-                    self.config.max_collection_items
+                    path, mapping_len, self.config.max_collection_items
                 );
                 return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                     "Nested mapping exceeds maximum size of {} items",
                     self.config.max_collection_items
                 )));
             }
 
-            for (key, value) in dict.iter() {
+            for item in mapping.items()?.iter() {
+                let item = item.cast::<PyTuple>()?;
+                let key = item.get_item(0)?;
+                let value = item.get_item(1)?;
                 let key_str = key.str()?.to_string_lossy().into_owned();
                 let new_path = if path.is_empty() {
                     key_str.clone()
@@ -1659,6 +1663,57 @@ class ConfigModel:
         });
     }
 
+    #[test]
+    fn test_process_nested_mapping_allows_collection_limit_boundary() {
+        Python::initialize();
+        Python::attach(|py| {
+            let config = PyDict::new(py);
+            config.set_item("detect_email", true).unwrap();
+            config.set_item("max_collection_items", 1).unwrap();
+
+            let detector = PIIDetectorRust::new(&config.into_any()).unwrap();
+            let data = PyDict::new(py);
+            data.set_item("email", "alice@example.com").unwrap();
+
+            let (modified, new_data, _) =
+                detector.process_nested(py, &data.into_any(), "").unwrap();
+
+            assert!(modified);
+            assert_eq!(
+                new_data
+                    .bind(py)
+                    .cast::<PyDict>()
+                    .unwrap()
+                    .get_item("email")
+                    .unwrap()
+                    .unwrap()
+                    .extract::<String>()
+                    .unwrap(),
+                "[REDACTED]"
+            );
+        });
+    }
+
+    #[test]
+    fn test_process_nested_mapping_rejects_over_collection_limit() {
+        Python::initialize();
+        Python::attach(|py| {
+            let config = PyDict::new(py);
+            config.set_item("detect_email", true).unwrap();
+            config.set_item("max_collection_items", 1).unwrap();
+
+            let detector = PIIDetectorRust::new(&config.into_any()).unwrap();
+            let data = PyDict::new(py);
+            data.set_item("first", "alice@example.com").unwrap();
+            data.set_item("second", "bob@example.com").unwrap();
+
+            let err = detector
+                .process_nested(py, &data.into_any(), "")
+                .unwrap_err();
+            assert!(err.is_instance_of::<pyo3::exceptions::PyValueError>(py));
+        });
+    }
+
     #[test]
     fn test_detects_plus_prefixed_international_phone_number() {
         let config = PIIConfig {

plugins/rust/python-package/pii_filter/src/detector.rs

@@ -1,6 +1,10 @@
 from dataclasses import dataclass
 import logging
+import os
 from pathlib import Path
+import subprocess
+import sys
+import textwrap
 
 import pytest
 
@@ -314,6 +318,95 @@ async def test_tool_post_invoke_returns_copied_payload_for_frozen_models():
     assert result.modified_payload.result["contact"] == "[REDACTED]"
 
 
+@pytest.mark.asyncio
+async def test_tool_post_invoke_returns_new_nested_result_for_mcp_content():
+    plugin = PIIFilterPlugin(_make_config())
+    payload = ToolPostInvokePayload(
+        name="search",
+        result={
+            "content": [
+                {
+                    "type": "text",
+                    "text": "Contact alice@example.com",
+                }
+            ],
+            "isError": False,
+        },
+    )
+
+    result = await plugin.tool_post_invoke(payload, _make_context())
+
+    assert result.modified_payload is not None
+    assert result.modified_payload is not payload
+    assert result.modified_payload.result is not payload.result
+    assert result.modified_payload.result["content"] is not payload.result["content"]
+    assert result.modified_payload.result["content"][0] is not payload.result["content"][0]
+    assert payload.result["content"][0]["text"] == "Contact alice@example.com"
+    assert result.modified_payload.result["content"][0]["text"] == "Contact [REDACTED]"
+
+
+def test_tool_post_invoke_survives_real_cpex_policy_with_isolated_payload():
+    plugin_root = (
+        Path(__file__).resolve().parents[3]
+        / "plugins"
+        / "rust"
+        / "python-package"
+        / "pii_filter"
+    )
+    script = """
+        import asyncio
+        from cpex.framework import PluginConfig, PluginContext
+        from cpex.framework.hooks.policies import HookPayloadPolicy, apply_policy
+        from cpex.framework.hooks.tools import ToolPostInvokePayload
+        from cpex.framework.memory import wrap_payload_for_isolation
+        from cpex.framework.models import GlobalContext
+        from cpex_pii_filter.pii_filter import PIIFilterPlugin
+
+        async def main():
+            plugin = PIIFilterPlugin(PluginConfig(
+                name="pii_filter",
+                kind="cpex_pii_filter.pii_filter.PIIFilterPlugin",
+                config={"detect_email": True, "detect_ssn": True, "block_on_detection": False},
+            ))
+            payload = ToolPostInvokePayload(
+                name="search",
+                result={
+                    "content": [{"type": "text", "text": "Contact alice@example.com"}],
+                    "isError": False,
+                },
+            )
+            plugin_input = wrap_payload_for_isolation(payload)
+            context = PluginContext(global_context=GlobalContext(request_id="req-pii"))
+
+            result = await plugin.tool_post_invoke(plugin_input, context)
+            assert result.modified_payload is not None
+            filtered = apply_policy(
+                plugin_input,
+                result.modified_payload,
+                HookPayloadPolicy(writable_fields=frozenset({"result"})),
+                apply_to=payload,
+            )
+            assert filtered is not None
+            assert payload.result["content"][0]["text"] == "Contact alice@example.com"
+            assert filtered.result["content"][0]["text"] == "Contact [REDACTED]"
+
+        asyncio.run(main())
+        print("ok")
+    """
+    env = os.environ.copy()
+    env.pop("PYTHONPATH", None)
+    result = subprocess.run(
+        [sys.executable, "-c", textwrap.dedent(script)],
+        cwd=plugin_root,
+        env=env,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr
+    assert result.stdout.strip() == "ok"
+
+
 @pytest.mark.asyncio
 async def test_tool_post_invoke_blocks_when_configured():
     plugin = PIIFilterPlugin(_make_config(block_on_detection=True))