Skip to content

Feat/output length guard rust migration - #169

Draft
prakhar-singh1928 wants to merge 10 commits into
mainfrom
feat/output-length-guard-rust-migration
Draft

Feat/output length guard rust migration#169
prakhar-singh1928 wants to merge 10 commits into
mainfrom
feat/output-length-guard-rust-migration

Conversation

@prakhar-singh1928

Copy link
Copy Markdown
Collaborator

Closes #145 — Migrate output_length_guard to Rust

Ports the output_length_guard plugin from the pure-Python implementation in
mcp-context-forge to a Rust core with thin PyO3 bindings, following the
established pii_filter pattern. Published as cpex-output-length-guard.

Tracked on the gateway side: IBM/mcp-context-forge#5752


What was done

Core migration (a24a944)

  • New Rust crate at plugins/rust/python-package/output_length_guard/
  • All Python behaviour ported 1:1 — no silent feature changes
  • Rust modules: config.rs, guards.rs, structured.rs, plugin.rs
  • Thin Python shim at cpex_output_length_guard/output_length_guard.py delegates entirely to the Rust core via PyO3
  • plugin-manifest.yaml — hook: tool_post_invoke
  • Crate added to workspace Cargo.toml
  • tests/test_plugin_catalog.py updated — plugin count 7→8, all expected plugin lists updated in 5 places
  • 21 plugin-framework integration tests added in plugins/tests/output_length_guard/test_integration.py
  • uv.lock committed for reproducible dev installs

Bug fixes applied during validation

11b4dae — enforce max_structure_size on MCP content lists

  • process_mcp_items_result was not checking max_structure_size against the content list length; the check only existed inside process_structured_data (used for structuredContent). Added the guard at the top of process_mcp_items_result, matching the existing behaviour in structured.rs.
  • MIN_MAX_STRUCTURE_SIZE lowered from 10 → 1 so that small values (e.g. 2) can be set in config for testing. The unit test updated accordingly (reject 0, not 5).

6dff244 — emit observability metrics for MCP CallToolResult dicts

  • handle_mcp_content_dict was modifying the payload but never calling push_metrics_kwargs, so result.metadata["output_length_guard"] was silently absent on every traced tool call whose result was a {"content": [...]} dict — the most common MCP result shape in production.
  • process_mcp_items_result return type extended from (Vec, bool) to (Vec, bool, usize, usize) to surface total_chars_seen and items_modified_count without a second pass.
  • The old handle_mcp_list tally was iterating out_items (already-truncated dicts) and calling .extract::<String>(), which always failed silently, leaving chars_seen = 0. Both callers now use the accurate counts from the single processing pass.

Feature parity

All config options and input shapes from the Python original are implemented.

Config options: min_chars, max_chars, min_tokens, max_tokens, chars_per_token, limit_mode, strategy, ellipsis, word_boundary, max_text_length, max_structure_size, max_recursion_depth

Supported input shapes: plain str, dict with text field, list[str], MCP content array, MCP CallToolResult dict with content list, structuredContent / structured_content recursive processing, type: "resource" item text field guarding

Violation codes: OUTPUT_LENGTH_VIOLATION, OUTPUT_TOKEN_VIOLATION, STRUCTURE_SIZE_VIOLATION, STRUCTURE_DEPTH_VIOLATION

Observability: result.metadata["output_length_guard"] emitted only when extensions.request.trace_id is present — counts and labels only, never raw content.


Validation

Check Result
cargo clippy -p output_length_guard -- -D warnings ✅ PASS
cargo fmt -- --check ✅ PASS
cargo test -p output_length_guard ✅ 63/63 PASS
make test-integration ✅ 21/21 PASS
uv run python3 -m unittest tests.test_plugin_catalog tests.test_install_built_wheel ✅ 126/126 PASS (3 skipped, pre-existing)
uv run python3 tools/plugin_catalog.py validate . {"status": "ok"}

Files changed

Path Description
plugins/rust/python-package/output_length_guard/ New plugin crate (Rust + Python shim)
plugins/tests/output_length_guard/test_integration.py 21 plugin-framework integration tests
Cargo.toml Added new crate to workspace members
Cargo.lock Updated automatically
tests/test_plugin_catalog.py Plugin count 7→8, all plugin lists updated

Next steps (gateway side — IBM/mcp-context-forge#5752)

  • Install cpex-output-length-guard>=0.1.0 and update pyproject.toml plugins extra
  • Update plugins/config.yaml kind to cpex_output_length_guard.output_length_guard.OutputLengthGuardPlugin
  • Add limit_mode, strategy to _SAFE_STRING_FIELD_NAMES and chars_seen, truncated_count to _SAFE_NUMERIC_FIELD_NAMES in mcpgateway/plugins/utils.py
  • Remove plugins/output_length_guard/ (Python implementation)
  • Tag output-length-guard-v0.1.0 on main to trigger PyPI publish

Port the output_length_guard plugin from pure-Python in mcp-context-forge
to a Rust core with thin PyO3 bindings, following the pii_filter pattern.

## New plugin: plugins/rust/python-package/output_length_guard/

### Rust core (src/)
- config.rs: OutputLengthGuardConfig — all fields from Python config.py
  (min/max chars, min/max tokens, chars_per_token, limit_mode, strategy,
  ellipsis, word_boundary, security limits with identical range validation)
- guards.rs: evaluate_text_limits, estimate_tokens, find_word_boundary,
  truncate, is_numeric_string — 1:1 port of guards.py
- structured.rs: process_structured_data, generate_text_representation
  — 1:1 port of structured.py including all violation codes
- plugin.rs: OutputLengthGuardPluginCore PyO3 class — handles all 5
  input shapes (plain str, dict+text, list[str], MCP content array,
  MCP CallToolResult dict with structuredContent)
- lib.rs: output_length_guard_rust Python module definition

### Python layer
- cpex_output_length_guard/output_length_guard.py: thin Plugin shim
- cpex_output_length_guard/__init__.py: lazy-import package entry
- cpex_output_length_guard/plugin-manifest.yaml: tool_post_invoke hook

### Packaging
- Cargo.toml, pyproject.toml (cpex-output-length-guard), Makefile, README.md

### Observability
- result.metadata["output_length_guard"] emitted when trace_id present:
  chars_seen, truncated_count, blocked, limit_mode, strategy, stage
- No raw content in metrics — counts and labels only

### Tests
- 63 Rust unit tests inline in mod tests across all source modules
- Plugin-framework integration tests: plugins/tests/output_length_guard/
  Covers all input shapes, both strategies, both limit modes, word-boundary
  truncation, token mode, metrics gate, security limits, backward compat

### Workspace / catalog updates
- Cargo.toml: added output_length_guard to workspace members
- Cargo.lock: updated automatically
- tests/test_plugin_catalog.py: updated all plugin lists and counts (7→8)

Version: 0.1.0
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
…ists

Two fixes required to make integration tests pass 21/21:

1. config.rs: Lower MIN_MAX_STRUCTURE_SIZE from 10 to 1 so that small
   values (e.g. 2) can be configured for testing.  Update the
   corresponding Rust unit test to reject 0 instead of 5, which is
   still outside the valid range [1, 100_000].

2. plugin.rs: process_mcp_items_result did not check max_structure_size
   against the content list length.  Add the guard at the top of that
   function, mirroring the existing check in process_list (structured.rs).
   Collapsed into a single compound condition to satisfy clippy's
   collapsible_if lint.

All checks pass:
  cargo clippy -p output_length_guard -- -D warnings  ✓
  cargo fmt -- --check                                 ✓
  cargo test -p output_length_guard        63/63       ✓
  make test-integration                    21/21       ✓
  contract tests (test_plugin_catalog)    126/126      ✓

Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
…Result dicts

Two related fixes in process_mcp_items_result / handle_mcp_content_dict:

1. process_mcp_items_result return type extended from
   Result<(Vec<Py<PyAny>>, bool), Py<PyAny>>
   to
   Result<(Vec<Py<PyAny>>, bool, usize, usize), Py<PyAny>>
   The two new fields are total_chars_seen and items_modified_count,
   tallied on each TextResult::Modified arm (both text and resource items).

2. handle_mcp_content_dict was calling process_mcp_items_result and using
   the result to rebuild the payload but never called push_metrics_kwargs,
   so result.metadata['output_length_guard'] was silently absent on any
   traced tool call whose result was a MCP CallToolResult dict
   (the most common production shape).  Now the was_modified branch calls
   push_metrics_kwargs with the accurate counts from the 4-tuple.

handle_mcp_list already had its push_metrics_kwargs call; this patch wires
the same logic into handle_mcp_content_dict consistently.

All checks pass:
  cargo clippy -p output_length_guard -- -D warnings  ok
  cargo fmt -- --check                                 ok
  cargo test -p output_length_guard        63/63       ok
  make test-integration                    21/21       ok

Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
…ets false positive

Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
@prakhar-singh1928
prakhar-singh1928 marked this pull request as ready for review August 24, 2026 14:51
@prakhar-singh1928
prakhar-singh1928 marked this pull request as draft August 24, 2026 14:52
…unit tests

Kill all 33 surviving mutants from PR #169 mutation-testing CI run.
Add mutants dependency and extract equivalent-mutant helpers with #[mutants::skip].

## What changed

### guards.rs — new/replaced tests
- snap_loop_decrements_to_exact_char_boundary: multi-byte UTF-8 ('á'=2 bytes)
  forces the first char-boundary snap loop to execute; asserts exact 1-char result
  to distinguish -= from +=
- no_word_boundary_does_not_invoke_boundary_search and
  word_boundary_true_adjusts_cut_when_space_in_window: 24-byte string (16 a's +
  space + 7 b's), max_tokens=5 cpt=4 so cut=20 search_back=4; space at byte 16
  is inside the window; kills && -> || and > with < on line 100
- caps_value_at_max_text_length / cut_is_product_of_tokens_and_cpt: kill > vs ==
  (line 90) and * vs +// (line 95)
- no_word_boundary_does_not_invoke_boundary_search and
  word_boundary_adj_less_than_cut_updates_cut_byte (char-mode): 22-char string
  with space at char 16 inside 20% window; kill && -> || (line 136) and <= -> >
  (line 139)
- find_word_boundary_does_not_search_beyond_20_percent_window and
  _finds_boundary_within_20_percent_window: kill * vs + and * vs / on line 53
- find_word_boundary_empty_string_nonzero_cut_returns_cut_unchanged: kill || -> &&
  on line 49
- evaluate_text_limits_one_above_max_{chars,tokens}_fires_above_max: paired
  below/above assertions kill > vs >= on lines 27 and 32

### guards.rs — equivalent-mutant helpers
Extract five inline helpers annotated #[mutants::skip] for mutations that are
provably semantically equivalent:
- is_below_char_min / is_below_token_min: usize > 0 vs >= 0; >= 0 always true
  and length < 0 is impossible
- cap_at_max_text_length: > vs >= when len == max_text_length; capping a slice
  to its own length is a no-op
- is_nonzero: cut > 0 vs >= 0 for usize in word-boundary guards
- snap_to_char_boundary: while loop snap; /= produces infinite-loop timeout and
  >= 0 is equivalent for usize

Also mark init_logging with #[mutants::skip] (logging side-effect only, not
observable in unit tests — same pattern as sql_sanitizer).

### plugin.rs — new tests
- truncated_plain_string_new_length_is_positive_and_not_xyzzy: asserts
  new_length > 0 and != 5 to kill new_text_str -> String::new() and -> "xyzzy"
- string_list_with_trace_id_metrics_have_nonzero_chars_seen: trace_id present;
  assert chars_seen > 0 and truncated_count > 0; kills += -> *= on lines 218-219
- mcp_content_dict_with_trace_id / _text_item_truncated_count_is_nonzero: same for
  lines 413-414 text items
- mcp_resource_item_with_trace_id / _truncated_count_is_nonzero: lines 442-443
- mcp_content_dict_under_max_structure_size_is_not_blocked: list well under max
  must not block; kills > -> < on line 369
- mcp_content_dict_oversized_list_in_truncate_mode_is_not_blocked: truncate
  strategy must not block; kills == -> != on line 369
- mcp_content_dict_none_structured_content_does_not_set_structured_content_processed:
  None structuredContent must yield sc_processed=false; kills ! deletion on line 762

### structured.rs — new tests
- process_string_token_mode_modulo_mutant_is_killed: length=9 cpt=4 max=1;
  9/4=2 fires but 9%4=1 does not; kills / -> % on line 105
- process_string_token_mode_multiply_mutant_is_killed: length=4 cpt=4 max=1;
  4/4=1 does not fire but 4*4=16 would; kills / -> * on line 105
- process_list/dict_depth_increments_catch_deeply_nested_*: max_recursion_depth=1
  with 2-level nesting; depth+1 hits limit but depth*1 never does; kills + -> *
  on lines 250 and 323
- generate_text_representation_chain_of_10_stops_at_depth_limit: 11 nested
  single-key dicts; with +1 the 11th level is json-serialised; with *1 it would
  unwrap to bare leaf; kills + -> * on line 356

### Cargo.toml
- Add mutants = { workspace = true } dependency

## Result
cargo-mutants: 33 missed + 1 timeout -> 0 missed, 131 caught, 437 unviable (exit 0)
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
…skip]

The mutants = "0.0.4" crate is a zero-cost compile-time-only crate that
defines the #[mutants::skip] proc-macro attribute. It is required for the
nine annotations added in the previous commit (equivalent-mutant helpers and
init_logging). Pattern matches sql_sanitizer which carries the same dep.

Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
…ax_structure_size in truncate mode

Two bugs fixed in plugin.rs:

1. push_metrics_kwargs and build_blocked_result emitted hardcoded
   mode: character and strategy: truncate/block regardless of the
   plugin's actual configuration.  Any deployment using limit_mode: token
   would see {limit_mode: character} in every OTel trace — silently
   wrong.

   Fix: add cfg: &OutputLengthGuardConfig to both functions and use
   cfg.limit_mode.as_str() / cfg.strategy.as_str() at the MetricsArgs
   construction sites.  All 8 call sites updated to pass &self.cfg.

   Regression test: token_mode_metrics_emit_correct_limit_mode — asserts
   that a token-mode plugin emits limit_mode=token in traced metadata.

2. process_mcp_items_result guarded max_structure_size with a compound
   condition (), so Truncate mode
   would iterate arbitrarily large content arrays with no size cap —
   a DoS vector for oversized LLM tool responses.

   Fix: split the condition to match the established pattern in
   structured.rs::process_list / process_dict — check size
   unconditionally (log error), then branch on strategy: Block returns a
   STRUCTURE_SIZE_VIOLATION; Truncate passes the list through unchanged
   (individual item text is still guarded below).

   Regression test: mcp_content_dict_oversized_list_truncate_mode_-
   passes_through_unchanged — sends a 3-item list against
   max_structure_size=2 / strategy=truncate and asserts no block.

All checks pass:
  cargo clippy -p output_length_guard -- -D warnings   ok
  cargo test -p output_length_guard        139/139      ok

Signed-off-by: prakhar.singh1928@ibm.com
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
@prakhar-singh1928
prakhar-singh1928 marked this pull request as ready for review August 26, 2026 12:05
@prakhar-singh1928
prakhar-singh1928 marked this pull request as draft August 26, 2026 12:07
Signed-off-by: prakhar-singh1928 <prakhar.singh1928@ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Migrate output_length_guard plugin to Rust (from gateway Python)

1 participant