fix(middleware): improve pre-tool middleware guarding logic#1824
Open
cparadis-nvidia wants to merge 22 commits intoNVIDIA:developfrom
Open
fix(middleware): improve pre-tool middleware guarding logic#1824cparadis-nvidia wants to merge 22 commits intoNVIDIA:developfrom
cparadis-nvidia wants to merge 22 commits intoNVIDIA:developfrom
Conversation
…l chunks and verifies each one, with _analyze_chunk handling the LLM call per chunk. A violation anywhere triggers the overall result. Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
…ack scenario that was previously exploitable), all three action modes, targeting, threshold, and streaming. Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughPreToolVerifierMiddleware now HTML-escapes user_input, splits long inputs into overlapping sliding windows (50% stride) bounded by new Changes
Sequence DiagramsequenceDiagram
participant Client
participant Middleware as PreToolVerifierMiddleware
participant Analyzer as _analyze_content()
participant Windower as Sliding Window Builder
participant LLM
participant Aggregator as Result Aggregation
Client->>Middleware: Invoke function with input
Middleware->>Analyzer: Request analysis
alt short input (<= max_content_length)
Analyzer->>LLM: Single HTML-escaped prompt
LLM-->>Analyzer: Window result
else long input (> max_content_length)
Analyzer->>Windower: Build overlapping windows (50% stride)
Windower->>Analyzer: Return sampled/sequence of windows (cap by max_chunks)
loop per window (stop on first should_refuse)
Analyzer->>LLM: Analyze window (HTML-escaped)
LLM-->>Analyzer: Window result
end
end
Analyzer->>Aggregator: Consolidate window results
Aggregator->>Aggregator: any() for booleans, max() for confidence, set() for violation_types, concat reasons from violating windows, set sanitized_input=None for multi-window
Aggregator-->>Analyzer: Aggregated assessment
Analyzer-->>Middleware: Final verdict
alt should_refuse
Middleware->>Client: Raise ValueError (refusal)
else should_redirect
Middleware->>Client: Execute with sanitized_input
else partial_compliance
Middleware->>Client: Log warning, execute with original input
else
Middleware->>Client: Execute with original input
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py`:
- Around line 247-249: The chunking logic that slices inputs into disjoint
_MAX_CONTENT_LENGTH pieces (used by the method that calls _analyze_chunk()) is
bypassable at chunk seams; modify the splitter to include overlap (e.g., sliding
window or a carryover segment equal to the max directive length) or perform
explicit seam analysis by concatenating the last N bytes of chunk i with the
first N bytes of chunk i+1 and passing that to _analyze_chunk() so payloads
spanning boundaries are detected; update the caller that currently iterates over
chunks to run this extra seam check (referencing _MAX_CONTENT_LENGTH and
_analyze_chunk()) and add a regression test that constructs a malicious
directive split across a boundary to assert the analyzer now flags it.
- Around line 258-268: The chunking logic in PreToolVerifierMiddleware currently
slices content_str into unlimited _MAX_CONTENT_LENGTH chunks and awaits
_analyze_chunk for each, which can amplify cost/latency; add a hard ceiling
(e.g., MAX_CHUNKS constant) and enforce it when building chunks so you never
call _analyze_chunk more than MAX_CHUNKS, iterate calling _analyze_chunk
sequentially so you can early-exit if a refusing result is observed, and if the
input would require more than MAX_CHUNKS chunks return a deterministic policy
(fail-open or fail-closed) immediately instead of processing extra chunks;
update the variables/chunk loop around content_str, chunks, results and the call
sites to _analyze_chunk to implement the cap and early-exit behavior.
In
`@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py`:
- Around line 19-25: The mocked verifier responses in the tests are returning
Python list/string representations (e.g. vt = "['prompt_injection']") which are
not valid JSON and cause _analyze_chunk() to hit its exception path; update the
tests to serialize the mocked verifier response bodies with json.dumps() (import
json at the top) so fields like vt, reason, and sanitized are valid JSON
strings—ensure the AsyncMock/patch that simulates the verifier returns
json.dumps({...}) for the response body in both the instances shown and the
other block around lines 61-73 so the positive-path chunking behavior is
actually exercised.
🪄 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: Pro
Run ID: 209f2635-1b3d-40b7-8f07-03f3a2848560
📒 Files selected for processing (2)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.pypackages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Show resolved
Hide resolved
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py (1)
277-289:⚠️ Potential issue | 🔴 CriticalDon't make oversized verification probabilistic.
Once this path samples windows at random, an attacker can pad the input until the malicious region is unlikely to be chosen, and the same payload can pass or fail on different requests. The current
range(0, len(content_str), _STRIDE)also emits a redundant trailing suffix window, so some inputs hit this branch earlier than necessary. For a security gate, cap overflow should return a deterministic fail-open/fail-closed result instead of sampling.🔒 Deterministic cap handling
- windows = [content_str[i:i + _MAX_CONTENT_LENGTH] for i in range(0, len(content_str), _STRIDE)] + last_start = len(content_str) - _MAX_CONTENT_LENGTH + window_starts = list(range(0, last_start + 1, _STRIDE)) + if window_starts[-1] != last_start: + window_starts.append(last_start) + windows = [content_str[i:i + _MAX_CONTENT_LENGTH] for i in window_starts] if len(windows) > _MAX_CHUNKS: logger.warning( - "PreToolVerifierMiddleware: Input to %s requires %d windows (cap=%d); " - "randomly sampling %d windows", + "PreToolVerifierMiddleware: Input to %s requires %d windows (cap=%d)", function_name, len(windows), _MAX_CHUNKS, - _MAX_CHUNKS, ) - windows = random.sample(windows, _MAX_CHUNKS) + if self.config.fail_closed: + return PreToolVerificationResult( + violation_detected=True, + confidence=1.0, + reason="Input blocked: security verification input exceeds scanning limit", + violation_types=[], + sanitized_input=None, + should_refuse=True, + error=True, + ) + return PreToolVerificationResult( + violation_detected=False, + confidence=0.0, + reason="Analysis skipped: input exceeds scanning limit", + violation_types=[], + sanitized_input=None, + should_refuse=False, + error=True, + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py` around lines 277 - 289, The code currently builds sliding windows from content_str using range(0, len(content_str), _STRIDE) then randomly samples when len(windows) > _MAX_CHUNKS, which is non-deterministic and also produces a redundant trailing short window; change this to build only full windows using range(0, max(0, len(content_str) - _MAX_CONTENT_LENGTH + 1), _STRIDE) (so windows are length _MAX_CONTENT_LENGTH without trailing suffix) and remove random.sample; instead, if len(windows) > _MAX_CHUNKS in PreToolVerifierMiddleware (the block using windows, _MAX_CHUNKS, function_name and logger), deterministically handle the overflow by failing the verification path (e.g., log a clear error via logger and return/reject/raise the verification failure) or deterministically select the first/last N windows—do not use randomness.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py`:
- Around line 184-188: The verifier prompt currently interpolates untrusted
chunk directly into user_prompt_parts (see variable chunk and list
user_prompt_parts in defense_middleware_pre_tool_verifier.py and the surrounding
logic around function_name), allowing a payload containing </user_input> to
break the boundary; fix by sanitizing or serializing chunk before appending
(e.g., escape XML/HTML entities or base64/JSON-encode the chunk) and update the
appended string to include a clear marker that the content is escaped (so the
verifier decodes or treats it as literal). Ensure the change is applied where
the code builds the "Input to verify:\n<user_input>...</user_input>" string so
all untrusted input is never inserted verbatim.
---
Duplicate comments:
In
`@packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py`:
- Around line 277-289: The code currently builds sliding windows from
content_str using range(0, len(content_str), _STRIDE) then randomly samples when
len(windows) > _MAX_CHUNKS, which is non-deterministic and also produces a
redundant trailing short window; change this to build only full windows using
range(0, max(0, len(content_str) - _MAX_CONTENT_LENGTH + 1), _STRIDE) (so
windows are length _MAX_CONTENT_LENGTH without trailing suffix) and remove
random.sample; instead, if len(windows) > _MAX_CHUNKS in
PreToolVerifierMiddleware (the block using windows, _MAX_CHUNKS, function_name
and logger), deterministically handle the overflow by failing the verification
path (e.g., log a clear error via logger and return/reject/raise the
verification failure) or deterministically select the first/last N windows—do
not use randomness.
🪄 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: Pro
Run ID: 54cba0d7-ad59-434b-bbab-f0cc1196c55b
📒 Files selected for processing (2)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.pypackages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
✅ Files skipped from review due to trivial changes (1)
- packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Show resolved
Hide resolved
…ts tags as literal characters Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py (2)
376-393: Strengthen the over-cap test to prove sampling logic is actually used.Current assertions can still pass if code processes the first
_MAX_CHUNKSwindows deterministically. Consider asserting that the sampling helper is invoked in this path.As per coding guidelines, "
**/tests/**/*.py: Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code."Suggested test hardening
+import random import json from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import patch @@ - result = await middleware._analyze_content(over_cap_content, function_name=middleware_context.name) + with patch("nat.middleware.defense.defense_middleware_pre_tool_verifier.random.sample", + wraps=random.sample) as sample_mock: + result = await middleware._analyze_content(over_cap_content, function_name=middleware_context.name) + + sample_mock.assert_called_once() @@ assert mock_llm.ainvoke.call_count == _MAX_CHUNKS🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py` around lines 376 - 393, The test should verify that sampling is used when content exceeds the window cap by mocking the sampling operation and asserting it was invoked; update test_over_cap_randomly_samples_max_chunks_windows to patch the sampler used by PreToolVerifierMiddleware (e.g., patch random.sample or the middleware method responsible for sampling such as _sample_windows) before calling middleware._analyze_content, assert the sampler was called exactly once and with expected args, and keep the existing assertions that mock_llm.ainvoke.call_count == _MAX_CHUNKS and no violation/ refusal to ensure behavior is unchanged.
242-263: Assert exact aggregation contract (order + delimiter), not just membership/containment.These tests currently verify presence but not the exact output format/order. Tightening assertions will better guard regressions in dedupe order and semicolon-join behavior.
As per coding guidelines, "
**/tests/**/*.py: Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code."Suggested assertion updates
- assert set(result.violation_types) == {"prompt_injection", "jailbreak", "social_engineering"} - assert len(result.violation_types) == 3 + assert result.violation_types == ["prompt_injection", "jailbreak", "social_engineering"] @@ - assert "reason A" in result.reason - assert "reason B" in result.reason + assert result.reason == "reason A; reason B"Also applies to: 289-312
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py` around lines 242 - 263, Update the test assertions to validate the exact aggregation contract (order and delimiter), not just membership: in test_windowed_violation_types_deduplicated replace the set-based checks on result.violation_types with an exact equality assertion that enforces the expected semicolon-joined order (e.g., assert result.violation_types == "prompt_injection;jailbreak;social_engineering" if the API returns a string, or assert result.violation_types == ["prompt_injection","jailbreak","social_engineering"] if it returns a list). Apply the same change to the analogous assertions in the other test block noted (lines 289-312) so both tests assert exact order and semicolon delimiter rather than using set containment.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py`:
- Around line 91-94: Update the class docstring so it accurately describes the
current over-cap behavior: instead of saying inputs over _MAX_CHUNKS “bypass LLM
calls entirely,” state that the input is capped and only sampled/capped analysis
is performed (up to _MAX_CHUNKS LLM calls) when the input exceeds _MAX_CHUNKS;
update any wording around early loop exit and should_refuse propagation to
reflect that a violation found mid-scan can stop further window checks but that
the LLM call count itself is capped at _MAX_CHUNKS (refer to the class/docstring
near _MAX_CHUNKS and the loop handling should_refuse in the pre-tool verifier
tests).
---
Nitpick comments:
In
`@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py`:
- Around line 376-393: The test should verify that sampling is used when content
exceeds the window cap by mocking the sampling operation and asserting it was
invoked; update test_over_cap_randomly_samples_max_chunks_windows to patch the
sampler used by PreToolVerifierMiddleware (e.g., patch random.sample or the
middleware method responsible for sampling such as _sample_windows) before
calling middleware._analyze_content, assert the sampler was called exactly once
and with expected args, and keep the existing assertions that
mock_llm.ainvoke.call_count == _MAX_CHUNKS and no violation/ refusal to ensure
behavior is unchanged.
- Around line 242-263: Update the test assertions to validate the exact
aggregation contract (order and delimiter), not just membership: in
test_windowed_violation_types_deduplicated replace the set-based checks on
result.violation_types with an exact equality assertion that enforces the
expected semicolon-joined order (e.g., assert result.violation_types ==
"prompt_injection;jailbreak;social_engineering" if the API returns a string, or
assert result.violation_types ==
["prompt_injection","jailbreak","social_engineering"] if it returns a list).
Apply the same change to the analogous assertions in the other test block noted
(lines 289-312) so both tests assert exact order and semicolon delimiter rather
than using set containment.
🪄 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: Pro
Run ID: 38dc8360-221f-49b4-abc8-c9a6c65e9a6b
📒 Files selected for processing (2)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.pypackages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
…HUNKS calls) Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py`:
- Around line 116-118: The current assertion that "</user_input>" is not in
user_message_content is invalid because _analyze_chunk() appends a literal
wrapper; instead restrict the check to the injected payload: extract the
injected user content (e.g., locate the markers used by _analyze_chunk() or
split around the wrapper) and assert that the raw closing tag does not appear
inside that injected portion while the escaped form "</user_input>" does;
alternatively assert that the literal "</user_input>" appears exactly once (the
wrapper added by _analyze_chunk()) and not inside the injected content, and keep
the existing check that "</user_input>" is present in the injected
payload.
- Around line 386-392: The test currently only asserts the LLM call cap but not
that the randomized sampling branch was used; patch or mock random.sample inside
the test for middleware._analyze_content to verify it was called with the full
list of windows (length _MAX_CHUNKS * _STRIDE + 1 or the specific 17-window list
generated) and k=_MAX_CHUNKS, then assert mock_llm.ainvoke.call_count ==
_MAX_CHUNKS and that random.sample was invoked with those parameters so the
over-cap randomized path is exercised; locate this behavior around the test for
middleware._analyze_content and the mock names random.sample and
mock_llm.ainvoke to add the additional assertion.
- Around line 339-348: The test uses homogeneous "a" content so overlapping
windows are indistinguishable; update the test in
test_defense_middleware_pre_tool_verifier.py to create a boundary-spanning
payload for middleware._analyze_content by inserting a unique marker string that
crosses the _MAX_CONTENT_LENGTH boundary (use _STRIDE to position it), call
middleware._analyze_content(long_content,
function_name=middleware_context.name), then assert mock_llm.ainvoke.call_count
== 2 and that the second call's user message
(mock_llm.ainvoke.call_args_list[1][0][0] -> messages[1]["content"]) contains
the full unique marker while the first call does not, so the test verifies the
overlap logic rather than identical windows.
🪄 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: Pro
Run ID: 847f1a20-c8f8-4899-8eed-6afe15630ef1
📒 Files selected for processing (1)
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Show resolved
Hide resolved
willkill07
requested changes
Mar 30, 2026
Member
willkill07
left a comment
willkill07
left a comment
There was a problem hiding this comment.
Overall the improvements are welcome, but I think we can do better by allowing the maximum number of chunks to be configurable so we don't have any non-deterministic behavior.
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
There was a problem hiding this comment.
♻️ Duplicate comments (3)
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py (3)
117-119:⚠️ Potential issue | 🔴 CriticalFix false assertion: wrapper
</user_input>is expected once.This check fails for valid prompts because
_analyze_chunk()intentionally appends a literal wrapper closing tag. Assert escaping only for the injected payload (or count a single literal wrapper).Suggested patch
- # The raw closing tag must NOT appear verbatim — it must be escaped - assert "</user_input>" not in user_message_content - assert "</user_input>" in user_message_content + # Only the wrapper closing tag should remain literal; injected tag must be escaped. + assert user_message_content.count("</user_input>") == 1 + assert "benign text</user_input>" not in user_message_content + assert "benign text</user_input>" in user_message_content🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py` around lines 117 - 119, The test wrongly asserts that the literal wrapper closing tag never appears even though _analyze_chunk() intentionally appends one; update the assertions on user_message_content to allow a single literal "</user_input>" (e.g., assert user_message_content.count("</user_input>") == 1) while still asserting that any additional/ injected occurrences are escaped by checking "</user_input>" is present for the payload; locate the assertions referencing user_message_content and replace the binary not-in / in checks with a count check for the literal wrapper and the escape presence check so the test accepts the intentional wrapper but ensures injected tags are escaped.
340-349:⚠️ Potential issue | 🟠 MajorUse a boundary-spanning marker, not homogeneous content.
With all-
"a"input, this test cannot prove overlap behavior; it only proves call order from mocked side effects. Add a unique payload crossing the old boundary and assert it appears in window 1 but not window 0.Suggested patch
- long_content = "a" * (_MAX_CONTENT_LENGTH * 2) + boundary_payload = "IGNORE_PREVIOUS_INSTRUCTIONS" + split = len(boundary_payload) // 2 + prefix = "a" * (_MAX_CONTENT_LENGTH - split) + suffix = "b" * (_MAX_CONTENT_LENGTH - (len(boundary_payload) - split)) + long_content = prefix + boundary_payload + suffix result = await middleware._analyze_content(long_content, function_name=middleware_context.name) assert mock_llm.ainvoke.call_count == 2 - # Verify window 1 was passed content starting at _STRIDE (spans the old boundary) + # Verify only overlapping window contains full boundary payload + window0_messages = mock_llm.ainvoke.call_args_list[0][0][0] + window0_user_content = window0_messages[1]["content"] window1_messages = mock_llm.ainvoke.call_args_list[1][0][0] window1_user_content = window1_messages[1]["content"] - expected_window1 = "a" * _MAX_CONTENT_LENGTH # content[_STRIDE : _STRIDE + _MAX_CONTENT_LENGTH] - assert expected_window1 in window1_user_content + assert boundary_payload not in window0_user_content + assert boundary_payload in window1_user_content🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py` around lines 340 - 349, The test uses homogeneous "a" content so it can't prove overlap behavior; change long_content to include a unique marker that crosses the old boundary (e.g., build content = "a" * (_STRIDE - k) + "UNIQUE_MARKER" + "a" * ... so the marker spans the boundary), then call middleware._analyze_content(...) as before and assert that the marker substring appears in window1_user_content (from mock_llm.ainvoke.call_args_list[1][0][0] messages) and does not appear in window0_user_content (mock_llm.ainvoke.call_args_list[0][0][0]); use the exact slice semantics used by the code (content[_STRIDE : _STRIDE + _MAX_CONTENT_LENGTH]) and reference _MAX_CONTENT_LENGTH, _STRIDE, middleware._analyze_content, middleware_context.name, and mock_llm to locate and verify the correct windows.
387-393:⚠️ Potential issue | 🟠 MajorStrengthen over-cap test to validate selection semantics.
Current assertions only verify the call cap. This still passes if selection regresses to a different strategy (e.g., first
_MAX_CHUNKSwindows). Validate which windows are analyzed.Suggested patch
- # (_MAX_CHUNKS * _STRIDE) + 1 chars → _MAX_CHUNKS + 1 windows, exceeding the cap - over_cap_content = "a" * (_MAX_CHUNKS * _STRIDE + 1) + # Build _MAX_CHUNKS + 1 windows with unique stride markers. + markers = [f"[W{i:02d}]" + ("a" * (_STRIDE - 5)) for i in range(_MAX_CHUNKS + 1)] + over_cap_content = "".join(markers) + "x" result = await middleware._analyze_content(over_cap_content, function_name=middleware_context.name) # All selected windows are clean → exactly _MAX_CHUNKS calls, no early exit assert mock_llm.ainvoke.call_count == _MAX_CHUNKS + expected_indices = [int(i * ((_MAX_CHUNKS + 1) / _MAX_CHUNKS)) for i in range(_MAX_CHUNKS)] + for call_idx, window_idx in enumerate(expected_indices): + messages = mock_llm.ainvoke.call_args_list[call_idx][0][0] + user_content = messages[1]["content"] + assert f"[W{window_idx:02d}]" in user_content assert not result.violation_detected assert not result.should_refuse
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In
`@packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py`:
- Around line 117-119: The test wrongly asserts that the literal wrapper closing
tag never appears even though _analyze_chunk() intentionally appends one; update
the assertions on user_message_content to allow a single literal "</user_input>"
(e.g., assert user_message_content.count("</user_input>") == 1) while still
asserting that any additional/ injected occurrences are escaped by checking
"</user_input>" is present for the payload; locate the assertions
referencing user_message_content and replace the binary not-in / in checks with
a count check for the literal wrapper and the escape presence check so the test
accepts the intentional wrapper but ensures injected tags are escaped.
- Around line 340-349: The test uses homogeneous "a" content so it can't prove
overlap behavior; change long_content to include a unique marker that crosses
the old boundary (e.g., build content = "a" * (_STRIDE - k) + "UNIQUE_MARKER" +
"a" * ... so the marker spans the boundary), then call
middleware._analyze_content(...) as before and assert that the marker substring
appears in window1_user_content (from mock_llm.ainvoke.call_args_list[1][0][0]
messages) and does not appear in window0_user_content
(mock_llm.ainvoke.call_args_list[0][0][0]); use the exact slice semantics used
by the code (content[_STRIDE : _STRIDE + _MAX_CONTENT_LENGTH]) and reference
_MAX_CONTENT_LENGTH, _STRIDE, middleware._analyze_content,
middleware_context.name, and mock_llm to locate and verify the correct windows.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 950cd5a3-7d06-474e-8901-ffd19a6f2217
📒 Files selected for processing (2)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.pypackages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
…ix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
willkill07
requested changes
Mar 31, 2026
Member
willkill07
left a comment
willkill07
left a comment
There was a problem hiding this comment.
Minor nit on reasonable lower-bound for max_content_length
Tests should be updated to reflect configuration-driven max content length and max chunks (including tests for varying)
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware_pre_tool_verifier.py
Show resolved
Hide resolved
packages/nvidia_nat_core/tests/nat/middleware/test_defense_middleware_pre_tool_verifier.py
Outdated
Show resolved
Hide resolved
…ix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
…than hardcoded here in tests Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Member
|
/ok to test e11edf2 |
…ix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
…:cparadis-nvidia/NeMo-Agent-Toolkit into fix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
…ix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
Member
|
/ok to test 67d56cc |
…ix-pre-tool-verifier-defense-middleware Signed-off-by: cparadis nvidia <cparadis@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fix content truncation security vulnerability in PreToolVerifierMiddleware
Summary:
The previous _analyze_content implementation truncated oversized inputs by keeping the first and last halves, silently dropping the middle.
An attacker could exploit this by padding benign content around a malicious payload, guaranteeing it lands in the truncated region and bypasses verification.
A subtler variant could also split a directive across the boundary between two disjoint chunks, making it invisible to both.
Replaced truncation with a sliding window approach: content exceeding max_content_length is scanned in overlapping windows of max_content_length chars with a stride of max_content_length // 2 (50% overlap).
Any injection directive up to stride chars long is mathematically guaranteed to appear fully within at least one window.
Windows are analyzed sequentially with early exit on the first refusing result.
Inputs requiring more than max_chunks windows are handled by selecting max_chunks evenly-spaced windows at deterministic intervals, ensuring uniform coverage of the full input.
Both max_content_length (default 32000) and max_chunks (default 16) are now configurable via PreToolVerifierMiddlewareConfig.
sanitized_input is always None for multi-window content since overlapping windows make reconstruction impossible.
Also fixed a secondary vulnerability: chunk content was interpolated verbatim into the LLM prompt inside <user_input> tags, allowing a payload containing </user_input> to break the boundary and inject instructions outside it.
Chunk content is now HTML-escaped before insertion, and the prompt label notes this explicitly so the verifier treats tags as literal text.
Also fixed the test mock helper to serialize LLM response bodies with json.dumps so the positive-path chunking behavior is actually exercised through the full JSON parse path.
Test plan:
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Bug Fixes
Tests