Skip to content

Commit aceabdc

Browse files
Report injection_ml on ScanResult when the ML layer did not load
1 parent 056c86f commit aceabdc

2 files changed

Lines changed: 51 additions & 11 deletions

File tree

‎sdk/src/unplug/guard.py‎

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -634,7 +634,7 @@ def scan_output_request(
634634
start_time = time.perf_counter()
635635
violation = self._limits.check_input_length(request.text)
636636
if violation is not None:
637-
return _limit_result(violation, len(request.text))
637+
return self._report_ml_degraded(_limit_result(violation, len(request.text)))
638638
try:
639639
with correlation_scope():
640640
if self._server_client is not None:
@@ -644,11 +644,11 @@ def scan_output_request(
644644
result = self._output_pipeline.run(body, context=ctx)
645645
if not isolated:
646646
self._maybe_mark_session_tainted_from_scan(Source.TOOL_OUTPUT)
647-
return result
647+
return self._report_ml_degraded(result)
648648
except Exception as exc:
649649
_log.error("guard.scan_output_request failed: %s", exc)
650650
latency_ms = (time.perf_counter() - start_time) * 1000
651-
return _fail_closed(exc, latency_ms)
651+
return self._report_ml_degraded(_fail_closed(exc, latency_ms))
652652

653653
def check_tool_call(
654654
self,
@@ -662,16 +662,19 @@ def check_tool_call(
662662
start_time = time.perf_counter()
663663
_ = approved # resume-only via ApprovalProvider; caller flag is ignored
664664
if not self._limits.is_tool_allowed(tool_name):
665-
return _limit_result(
665+
return self._report_ml_degraded(
666+
_limit_result(
666667
LimitViolation(
667668
kind="tool_blocked",
668669
limit=0,
669670
actual=0,
670671
message=f"Tool not allowed: {tool_name}",
671672
),
672673
)
674+
)
673675
if not self._config.tools.is_permitted(tool_name):
674-
return _limit_result(
676+
return self._report_ml_degraded(
677+
_limit_result(
675678
LimitViolation(
676679
kind="tool_profile_blocked",
677680
limit=0,
@@ -681,9 +684,10 @@ def check_tool_call(
681684
),
682685
),
683686
)
687+
)
684688
count_violation = self._limits.check_tool_call_count(len(self._context.tool_calls) + 1)
685689
if count_violation is not None:
686-
return _limit_result(count_violation)
690+
return self._report_ml_degraded(_limit_result(count_violation))
687691
tc = ToolCall(
688692
tool_name=tool_name,
689693
arguments=arguments,
@@ -705,11 +709,11 @@ def check_tool_call(
705709
self._context.add_tool_call(tc)
706710
if self._config.tools.is_taint_source(tool_name):
707711
self.notify_taint_source(tool_name)
708-
return result
712+
return self._report_ml_degraded(result)
709713
except Exception as exc:
710714
_log.error("guard.check_tool_call failed: %s", exc)
711715
latency_ms = (time.perf_counter() - start_time) * 1000
712-
return _fail_closed(exc, latency_ms)
716+
return self._report_ml_degraded(_fail_closed(exc, latency_ms))
713717

714718
def scan_request(
715719
self,
@@ -721,7 +725,7 @@ def scan_request(
721725
start_time = time.perf_counter()
722726
violation = self._limits.check_input_length(request.text)
723727
if violation is not None:
724-
return _limit_result(violation, len(request.text))
728+
return self._report_ml_degraded(_limit_result(violation, len(request.text)))
725729
try:
726730
with correlation_scope():
727731
if self._server_client is not None:
@@ -744,15 +748,15 @@ def scan_request(
744748
result = self._run_input_with_cache(request, ctx)
745749
if not isolated:
746750
self._maybe_mark_session_tainted_from_scan(request.source)
747-
return result
751+
return self._report_ml_degraded(result)
748752
finally:
749753
ctx.allowed_scanners = prev_allowed
750754
except ConfigError:
751755
raise
752756
except Exception as exc:
753757
_log.error("guard.scan_request failed: %s", exc)
754758
latency_ms = (time.perf_counter() - start_time) * 1000
755-
return _fail_closed(exc, latency_ms)
759+
return self._report_ml_degraded(_fail_closed(exc, latency_ms))
756760

757761
@property
758762
def is_server_mode(self) -> bool:
@@ -762,6 +766,23 @@ def is_server_mode(self) -> bool:
762766
def ml_model_loaded(self) -> bool:
763767
return self._ml_provider is not None and self._ml_provider.loaded
764768

769+
770+
def _report_ml_degraded(self, result: ScanResult) -> ScanResult:
771+
"""Copy the private ML-unavailable flag onto the result callers inspect.
772+
773+
``require_ml=True`` already fails closed. The default path keeps scanning
774+
with regex only, but ``ScanResult.degraded`` stayed false because nothing
775+
copied ``_ml_degraded`` onto the result. Server mode reports this itself.
776+
"""
777+
if not self._ml_degraded:
778+
return result
779+
layers = list(result.degraded_layers)
780+
if "injection_ml" not in layers:
781+
layers.append("injection_ml")
782+
if result.degraded and layers == list(result.degraded_layers):
783+
return result
784+
return result.model_copy(update={"degraded": True, "degraded_layers": layers})
785+
765786
@property
766787
def ml_degraded(self) -> bool:
767788
"""True when active_model is configured but ML failed to load."""

‎sdk/tests/unit/test_ml_degraded_warning.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,3 +74,22 @@ def test_no_active_model_does_not_warn(self, caplog: pytest.LogCaptureFixture) -
7474

7575
assert _degraded_records(caplog) == []
7676
assert guard._ml_degraded is False
77+
78+
def test_scan_result_reports_injection_ml_degraded(self) -> None:
79+
"""Callers inspect ScanResult, not the private flag. #107."""
80+
with patch("unplug.guard.prepare_active_model_spec", return_value=None):
81+
guard = Guard(config=GuardConfig(active_model="tiny"))
82+
83+
assert guard.ml_degraded is True
84+
result = guard.scan("some text")
85+
assert result.degraded is True
86+
assert result.degraded_layers == ["injection_ml"]
87+
88+
def test_regex_only_scan_is_not_degraded(self) -> None:
89+
"""No active_model means the ML layer was never configured."""
90+
guard = Guard(config=GuardConfig())
91+
result = guard.scan("some text")
92+
assert guard.ml_degraded is False
93+
assert result.degraded is False
94+
assert result.degraded_layers == []
95+

0 commit comments

Comments
 (0)