feat(library): add Polygraf PII detection and masking integration

[ali-khoda]

• Fix Polygraf API auth/payload: Authorization: Bearer $POLYGRAF_API_KEY + request flags; robust response parsing.

• Align Polygraf docs/examples with the working cloud endpoint + entity labels (Email, Person, Phone).

• Local validation: pytest -q tests/test_polygraf.py passed.

Summary by CodeRabbit

Release Notes

• New Features

  • Added Polygraf integration for detecting and masking personally identifiable information (PII) across user input, bot output, and retrieval sources.
  • Supports configurable entity type detection with options for cloud and self-hosted deployments.

• Documentation

  • Added comprehensive Polygraf integration guide with configuration examples for PII detection and masking workflows.

[codecov]

Codecov Report

❌ Patch coverage is 89.61039% with 16 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
nemoguardrails/library/polygraf/actions.py	88.34%	12 Missing ⚠️
nemoguardrails/library/polygraf/request.py	90.24%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

Documentation preview

https://nvidia-nemo.github.io/Guardrails/review/pr-1693

[greptile-apps]

Greptile Summary

This PR integrates Polygraf's PII detection and masking API into NeMo Guardrails, adding new actions, Colang flows (both v1 and v2), configuration models, documentation, and tests. It addresses issues from a prior review cycle by fixing auth headers, adding **kwargs to both actions, exposing a shared-session parameter, adding an API key warning, and handling null entities responses.

• New library module (nemoguardrails/library/polygraf/) with polygraf_detect_pii and polygraf_mask_pii actions, a reusable polygraf_request helper, and both Colang v1 and v2 flow definitions.
• Config model adds PolygrafDetection / PolygrafDetectionOptions under RailsRuntimeConfig, with example configs and full documentation.

Confidence Score: 5/5

Safe to merge; all blocking issues from the prior review round have been addressed and no new functional defects were found.

The previously flagged concerns (missing **kwargs, per-request session creation, silent missing API key, entities: null parsing) are all resolved in this revision. The remaining observations are minor test-quality items that do not affect runtime correctness of the guardrail integration itself.

No files require special attention for merge.

Important Files Changed

Filename	Overview
nemoguardrails/library/polygraf/request.py	Polygraf HTTP helper with shared-session support, timeout, and response parsing; contains a broad TypeError catch intended only for test doubles
nemoguardrails/library/polygraf/actions.py	Detection and masking actions with proper **kwargs, API-key warning, and entity filtering; logic is correct
tests/test_polygraf.py	Good async test coverage; _FakeSession lacks a timeout parameter causing most tests to exercise the fallback code path; dead chat >> line in one test
nemoguardrails/rails/llm/config.py	Adds PolygrafDetection and PolygrafDetectionOptions Pydantic models with appropriate defaults; correctly wired into RailsRuntimeConfig
nemoguardrails/library/polygraf/flows.co	Colang v2 flows for detection and masking on input, output, and retrieval; correct
nemoguardrails/library/polygraf/flows.v1.co	Colang v1 subflow definitions mirroring the v2 flows; correct
docs/configure-rails/guardrail-catalog/community/polygraf.md	New Polygraf integration guide with cloud and self-hosted configuration examples

Sequence Diagram

sequenceDiagram
    participant Rail as NeMo Rail (Colang flow)
    participant Action as polygraf_detect_pii / polygraf_mask_pii
    participant Request as polygraf_request
    participant Polygraf as Polygraf API

    Rail->>Action: "source, text, config, **kwargs"
    Action->>Action: resolve PolygrafDetection config
    Action->>Action: _get_polygraf_api_key()
    Action->>Request: text, server_endpoint, api_key, session?
    Request->>Request: build payload
    Request->>Polygraf: POST /pii/text-detect
    Polygraf-->>Request: 200 JSON
    Request-->>Action: List[Dict] entities
    alt detect PII
        Action-->>Rail: True or False
    else mask PII
        Action-->>Rail: masked text
    end

Loading

Prompt To Fix All With AI

Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 3
nemoguardrails/library/polygraf/request.py:84-89
The `TypeError` catch is broader than intended. If `session.post(...)` raises a `TypeError` for any reason other than an unrecognised `timeout` kwarg (e.g., a misbehaving session wrapper), the exception is silently swallowed and the retry without `timeout` is attempted — which likely raises the same error again, but now without the original traceback context. The real fix is to update `_FakeSession` in the tests to accept `timeout` (as `_FakeSessionWithTimeoutKwarg` already does), removing the need for this try/except entirely.

```suggestion
    post_ctx = session.post(server_endpoint, json=payload, headers=headers, timeout=timeout)
```

### Issue 2 of 3
tests/test_polygraf.py:151-164
`_FakeSession.post` does not accept a `timeout` keyword argument, so every test using it exercises the `except TypeError` fallback path in `_send_polygraf_request` rather than the primary code path. If the primary path (with `timeout`) ever regresses, these tests would not catch it. Consider aligning `_FakeSession` with `_FakeSessionWithTimeoutKwarg` so all tests exercise the same production branch.

```suggestion
class _FakeSession:
    def __init__(self, response):
        self.response = response
        self.requests = []

    def post(self, server_endpoint, json, headers, timeout=None):
        self.requests.append(
            {
                "server_endpoint": server_endpoint,
                "json": json,
                "headers": headers,
                "timeout": timeout,
            }
        )
        return _FakePostContextManager(self.response)
```

### Issue 3 of 3
tests/test_polygraf.py:543-544
The `chat >> "Hi!"` line just appends to `chat.history` (no generation is triggered for Colang 1.0 until `chat.bot()` / `chat <<` is called). The subsequent `chat.app.generate(messages=[...])` is an independent call that doesn't use `chat.history`, so `chat >> "Hi!"` is dead code here. Removing it makes the test intent clearer.

```suggestion
    response = chat.app.generate(messages=[{"role": "user", "content": "Hi!"}])
```

_{Reviews (4): Last reviewed commit: "fix(polygraf): harden HTTP timeout and e..." | Re-trigger Greptile}

[DataMonarch]

Updated the PR to address the review feedback and keep the diff focused on the Polygraf integration.

Changes made:

• Added **kwargs support to polygraf_mask_pii.
• Added a warning when POLYGRAF_API_KEY is not set, so cloud auth failures are easier to diagnose.
• Added optional shared aiohttp.ClientSession support to polygraf_request and passed it through from actions.
• Expanded Polygraf tests for request payload/auth behavior, wrapped responses, error handling, missing API-key diagnostics, and extra kwargs/session passthrough.
• Restored unrelated core files to match current develop, so the PR no longer includes embedding/KV/model-engine/dependency changes.

Validation:

• ruff check passed for touched Polygraf files.
• ruff format --check passed for touched Polygraf files.
• pytest tests/test_polygraf.py -v: 13 passed.
• pytest tests/test_gliner.py tests/test_rails_llm_config.py tests/test_config_loading.py -q: 37 passed.
• Additional smoke checks passed for detect/mask behavior, request payload/header behavior, shared session passthrough, and v1/v2 Polygraf import loading.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds Polygraf, a third-party PII detection and masking integration, to NeMo Guardrails. The implementation includes configuration models, HTTP request handling, two parameterized Rails actions, six flow definitions spanning detection and masking for input/output/retrieval sources, documentation with examples, and comprehensive test coverage.

Changes

Polygraf PII Detection and Masking Integration

Layer / File(s)	Summary
Configuration models and Rails integration point nemoguardrails/rails/llm/config.py	PolygrafDetectionOptions and PolygrafDetection models define the Polygraf server endpoint and per-context entity filtering. RailsConfigData.polygraf field wires the configuration into the top-level Rails config.
HTTP request layer nemoguardrails/library/polygraf/request.py	polygraf_request builds request payloads with detection parameters and optional Bearer auth, reusing a caller-provided aiohttp session. _send_polygraf_request handles HTTP POST, status validation, JSON parsing, and response shape normalization (accepts list or {"entities": list}).
Action implementations nemoguardrails/library/polygraf/__init__.py, nemoguardrails/library/polygraf/actions.py	polygraf_detect_pii queries Polygraf and returns a boolean; polygraf_mask_pii queries and replaces detected entity spans with <EntityType> placeholders. Both validate allowed source flows, filter by configured entity types, read POLYGRAF_API_KEY from environment with warning logging, and support optional session reuse.
Flow definitions nemoguardrails/library/polygraf/flows.co, nemoguardrails/library/polygraf/flows.v1.co	Six flows (input, output, retrieval × detection, masking) in both Colang syntax versions. Detection flows abort with "answer unknown" when PII is found; masking flows replace message/chunk variables with masked results without halting.
Documentation and examples docs/about/overview.md, docs/configure-rails/guardrail-catalog/community/polygraf.md, docs/configure-rails/guardrail-catalog/third-party.md, examples/configs/polygraf/pii_detection/config.yml, examples/configs/polygraf/pii_masking/config.yml	Polygraf is added to the PII detection overview. A comprehensive guide covers API key setup, endpoint configuration (cloud vs. self-hosted), entity type filtering, and includes two runnable YAML example configurations for detection and masking.
Test suite tests/test_polygraf.py	Tests validate request wiring (shared session, Bearer auth), response parsing (valid/invalid shapes), action behavior (missing API key, entity filtering), and end-to-end flow integration (detection/masking on input, output, retrieval) using mock Polygraf responses and the TestChat harness.

Sequence Diagram(s)

sequenceDiagram
  participant User as User/Flow
  participant Detect as polygraf_detect_pii
  participant Mask as polygraf_mask_pii
  participant Req as polygraf_request
  participant HTTP as Polygraf API
  
  User->>Detect: text, source, config
  Detect->>Req: text, server_endpoint, api_key
  Req->>HTTP: POST with payload and Bearer auth
  HTTP-->>Req: entity detections
  Req-->>Detect: entities list
  Detect-->>User: bool (PII found)
  
  User->>Mask: text, source, config
  Mask->>Req: text, server_endpoint, api_key
  Req->>HTTP: POST with payload and Bearer auth
  HTTP-->>Req: entity detections
  Req-->>Mask: entities list
  Mask->>Mask: replace spans with placeholder tags
  Mask-->>User: masked text

Loading

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 23.08% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Test Results For Major Changes	✅ Passed	Major changes documented with test results: 13 tests passing, 37 related tests passing, ruff validation, and smoke checks for functional behavior.
Title check	✅ Passed	The title clearly and concisely summarizes the main change: adding Polygraf PII detection and masking integration.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (2)

docs/configure-rails/guardrail-catalog/community/polygraf.md (2)

81-89: 💤 Low value

Consider clarifying entity type naming conventions.

The entity examples show mixed casing: Person, Email, Phone use title case, while CREDIT_CARD_NUMBER, PASSPORT_NUMBER, SSN, LOCATION, ORGANIZATION use uppercase. This inconsistency might confuse users about the correct format.

If Polygraf supports both formats or if different entities require different casing, consider adding a brief note explaining the convention. Otherwise, standardize the examples to one format.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/configure-rails/guardrail-catalog/community/polygraf.md` around lines 81
- 89, The entity list mixes title-case examples (Person, Email, Phone) with
ALL_CAPS constants (CREDIT_CARD_NUMBER, PASSPORT_NUMBER, SSN, LOCATION,
ORGANIZATION); update the docs to either standardize all examples to one
convention or add a short note explaining that Polygraf supports both formats
and when to use each. Locate the entity examples in the polygraf.md content and
either convert Person/Email/Phone to CREDIT_CARD_STYLE or convert the uppercase
ones to Title Case, or add a single-sentence clarification describing the naming
convention and any mapping/aliases the system accepts.

---

91-91: ⚡ Quick win

Add a direct link to Polygraf entity documentation.

Line 91 references "the Polygraf documentation" but doesn't provide a URL. Adding a direct link would improve user experience and reduce friction.

💡 Suggested addition

-For a complete list of supported entity types, please refer to the Polygraf documentation.
+For a complete list of supported entity types, please refer to the [Polygraf documentation](https://docs.polygraf.ai/entity-types) (or appropriate URL).

Note: Replace with the actual Polygraf documentation URL for entity types.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/configure-rails/guardrail-catalog/community/polygraf.md` at line 91,
Replace the plain-text reference "For a complete list of supported entity types,
please refer to the Polygraf documentation." with the same sentence that
includes a direct hyperlink to the Polygraf entity types page (i.e., embed the
actual Polygraf documentation URL for entity types into the sentence), ensuring
the visible link text remains "Polygraf documentation" so readers can click
through directly.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/configure-rails/guardrail-catalog/community/polygraf.md`:
- Line 5: Add a brief "Retrieval Flows" subsection to the Polygraf docs to
document detection/masking for retrieved documents: describe the
retrieval-specific keys under rails.config.polygraf (e.g., server_endpoint and
retrieval.entities) and show how to register a retrieval flow (referencing
retrieval.flows and the flow names like "polygraf detect pii on retrieval" and
"polygraf mask pii on retrieval"); keep it concise and placed after the existing
Configuration section so readers can see input/output examples plus the
analogous retrieval example.

In `@nemoguardrails/library/polygraf/request.py`:
- Around line 76-81: The except block that catches aiohttp.ContentTypeError when
calling resp.json() should preserve the original exception context by using
exception chaining; change the handler to capture the exception (e.g., "except
aiohttp.ContentTypeError as err:") and re-raise the ValueError with "raise
ValueError(... ) from err" so the original ContentTypeError is attached to the
new ValueError (refer to resp.json(), aiohttp.ContentTypeError, and the
ValueError being raised).

In `@tests/test_polygraf.py`:
- Around line 494-495: The assertion in tests/test_polygraf.py currently uses an
OR with "<NAME>" and checks for "John", which can yield false positives and
mismatches the mocked mask token; update the assertion to require that the
original sensitive text is absent AND the exact mock token is present (i.e.,
replace assert "John" not in response["content"] or "<NAME>" in
response["content"] with assert "John" not in response["content"] and "<Person>"
in response["content"]), and apply the same change to the analogous assertion at
the other location (lines referenced in the comment) so all tests check for
absence of the original string and presence of the mocked "<Person>"
placeholder.
- Around line 478-479: Fix the malformed completion string in the output-masking
test in tests/test_polygraf.py: the list element currently has '  "Hi! I am
John.' which is missing the closing double-quote; update that string (the
completion payload entry used in the test) to include the missing closing quote
so the Colang-style content is properly balanced (e.g., change the item that
contains '  "Hi! I am John.' to include the trailing double-quote).

---

Nitpick comments:
In `@docs/configure-rails/guardrail-catalog/community/polygraf.md`:
- Around line 81-89: The entity list mixes title-case examples (Person, Email,
Phone) with ALL_CAPS constants (CREDIT_CARD_NUMBER, PASSPORT_NUMBER, SSN,
LOCATION, ORGANIZATION); update the docs to either standardize all examples to
one convention or add a short note explaining that Polygraf supports both
formats and when to use each. Locate the entity examples in the polygraf.md
content and either convert Person/Email/Phone to CREDIT_CARD_STYLE or convert
the uppercase ones to Title Case, or add a single-sentence clarification
describing the naming convention and any mapping/aliases the system accepts.
- Line 91: Replace the plain-text reference "For a complete list of supported
entity types, please refer to the Polygraf documentation." with the same
sentence that includes a direct hyperlink to the Polygraf entity types page
(i.e., embed the actual Polygraf documentation URL for entity types into the
sentence), ensuring the visible link text remains "Polygraf documentation" so
readers can click through directly.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b51d82d0-c30f-42b9-b2ed-375092ed7c44

📥 Commits

Reviewing files that changed from the base of the PR and between 16be99a and 222e646.

📒 Files selected for processing (12)

• docs/about/overview.md
• docs/configure-rails/guardrail-catalog/community/polygraf.md
• docs/configure-rails/guardrail-catalog/third-party.md
• examples/configs/polygraf/pii_detection/config.yml
• examples/configs/polygraf/pii_masking/config.yml
• nemoguardrails/library/polygraf/__init__.py
• nemoguardrails/library/polygraf/actions.py
• nemoguardrails/library/polygraf/flows.co
• nemoguardrails/library/polygraf/flows.v1.co
• nemoguardrails/library/polygraf/request.py
• nemoguardrails/rails/llm/config.py
• tests/test_polygraf.py

[DataMonarch]

Addressed the latest Greptile/CodeRabbit follow-up comments.

Updates:

• Treat wrapped Polygraf responses with entities: null as an empty result.
• Preserve JSON parsing exception context with exception chaining.
• Tightened masking assertions and fixed the malformed completion string in tests.
• Added a concise retrieval-flow example to the Polygraf docs and clarified entity label guidance.

Validation:

• ruff check passed for touched code files.
• ruff format --check passed for touched code files.
• pytest tests/test_polygraf.py tests/test_gliner.py tests/test_rails_llm_config.py tests/test_config_loading.py -q: 51 passed.

[DataMonarch]

Pushed self-review hardening for two production risks I found.

• Added a default 30s aiohttp ClientTimeout for Polygraf requests. Without this, a slow or unresponsive Polygraf endpoint could hang the rails pipeline indefinitely (sibling community integrations like policyai/jailbreak/ai_defense/crowdstrike all set timeouts).
• Made the detect/mask actions defensive against malformed Polygraf responses. Non-dict entries, entries with missing entity_type, or with non-integer offsets are now skipped with a warning instead of raising KeyError/AttributeError mid-rail.
• Added tests covering timeout propagation and malformed-entity skipping.

Validation:

• ruff check passed.
• ruff format --check passed.
• pytest tests/test_polygraf.py tests/test_gliner.py tests/test_rails_llm_config.py tests/test_config_loading.py -q → 53 passed.

Note: the earlier pr-tests-matrix (Ubuntu, 3.11) failure was an unrelated flaky perf threshold in tests/v2_x/test_state_serialization.py (avg_time < 0.2 failed at 0.211s). It passes locally and on the other matrix entries.

[Pouyanpi]

@DataMonarch would you please rebase to latest develop branch and resolve the conflicts. It seems you've addressed all the review comments from coderabbit and greptile, once the conflicts are resolved I'll start the review. We will definitely have it for 0.23.0 release, sorry for the delays.

[Pouyanpi]

@DataMonarch @ali-khoda

Thanks for the PR, please have a look at following review comments. We are good to merge once these are addressed. Thank you.

[DataMonarch]

@Pouyanpi addressed your review.

Changes in bbceb7e2e:

Fail-closed semantics (actions.py)

• Provider/network failures from polygraf_request (now caught as ValueError) cause polygraf_detect_pii to return True and polygraf_mask_pii to return a new FAILSAFE_MASK_PLACEHOLDER (<REDACTED>). The original text never propagates downstream.
• If Polygraf reports a configured (selected) entity with missing/invalid offsets or type, both actions also fail closed. Unselected malformed entities are still skipped (no false-positive blocks).
• All warning logs were rewritten to use only structural metadata (shape, keys, invalid_fields); no entity text or response body is ever logged.

Request layer (request.py)

• asyncio.TimeoutError and aiohttp.ClientError are normalized to ValueError, so timeouts/DNS/TLS failures honor the documented exception contract.

Config (config.py)

• Both PolygrafDetection and PolygrafDetectionOptions now use ConfigDict(extra="forbid").

Colang 2 flows (flows.co)

• Every flow now declares global $user_message / $bot_message / $relevant_chunks before reading the rail variable. Mask flows also place the global declaration above the action call so the assignment binds to the rail variable.
• Removed the stray #######... separator lines from flows.co and flows.v1.co.

Tests (test_polygraf.py)

• New coverage:
  • test_polygraf_request_normalizes_timeout_as_value_error
  • test_polygraf_request_normalizes_client_error_as_value_error
  • test_polygraf_config_rejects_unknown_keys
  • test_polygraf_mask_pii_fails_closed_on_malformed_selected_entity (asserts the PII value never appears in result or logs)
  • test_polygraf_mask_pii_skips_unselected_malformed_entity
  • test_polygraf_mask_pii_fails_closed_on_provider_error
  • test_polygraf_detect_pii_fails_closed_on_provider_error
  • test_polygraf_detect_pii_fails_closed_on_malformed_selected_entity
• Colang 2 coverage via parser:
  • test_polygraf_v2_flows_parse_successfully
  • test_polygraf_v2_flows_declare_required_globals (catches missing global regressions in flows.co)

Note on the Check Fern docs failure
The failure log shows Authentication required. Please run 'fern login' or set the FERN_TOKEN environment variable. This is FERN_TOKEN not being available to fork PRs (a standard GitHub Actions limitation), not a problem in this PR. Happy to defer to maintainer guidance.

[Pouyanpi]

Thank you @DataMonarch your main fix/changes work, would you please address following issues, these are not fully addressed yet.

1. malformed spans can still fail open at nemoguardrails/library/polygraf/actions.py.

   _classify_entities still treats some malformed spans as safe or unselected. When entity_type is missing under an enabled filter, has_malformed_selected remains false; reversed or
   out-of-range integer offsets are appended to safe_spans. These cases leave raw PII unchanged. Please fail closed when the type is unknown, reject booleans, and validate 0 <= start < end <= len(text).

2. documentation contract is stale at docs/configure-rails/guardrail-catalog/community/polygraf.mdx

   It says the integration raises ValueError, but actions now catch it and either block detection or return .

3. Coverage gap at tests/test_polygraf.py.

   The new Colang 2 tests inspect parsed declarations but do not execute the rails. Actual local runtime verification passed, but a TestChat regression test would protect the behavior that originally produced text: null.

[DataMonarch]

@Pouyanpi thanks for catching the remaining gaps. Pushed e0a5faf4d:

1. Malformed-span fail-open in _classify_entities → fixed

• Missing/empty entity_type now fails closed (we cannot tell whether the entity would have matched the filter).
• bool offsets explicitly rejected (isinstance(True, int) is True in Python, so a plain int check would pass them through).
• Added 0 <= start < end <= len(text) validation; reversed/negative/empty/out-of-range spans no longer reach safe_spans. If such a span matched the filter (or the filter is None), the action fails closed.
• _classify_entities now takes text_length so it can do the range check.

2. Stale ValueError contract in polygraf.mdx → fixed

• Removed the "raises ValueError" note. Added a ## Failure Handling section that documents the actual contract: detect blocks (True), mask returns <REDACTED>, and what triggers each fail-closed path (provider/network/parse failure, malformed entity span, missing/invalid offsets). Structural-only logging is also called out.

3. Coverage gap (parser-only Colang 2 tests) → fixed

• New test_polygraf_v2_input_flow_passes_actual_user_message_to_action is a TestChat-based regression. It loads the shipped flows.co text, wires it into a v2 input rail via flow input rails $input_text, registers a mock polygraf_mask_pii that records the text it received, then sends a user message. It asserts the action received the actual user text rather than None. Verified locally that removing the global $user_message declaration makes the test fail with the exact bug you described.
• Additional malformed-span tests: missing entity_type (mask + detect), and a parametrized test over bool start, bool end, negative start, reversed offsets, end past text length, and empty spans.

Local validation: ruff clean, all 34 Polygraf tests pass.

[Pouyanpi]

Thanks @DataMonarch for making the changes and your flexibility 👍🏻 merging 🚀