fix: correct MSJ call site and harden litellm rate-limit detection

Two related fixes uncovered during an audit of the msj_data fix (0944ac6),
plus a pre-commit formatting/lint fix:

1. perform_many_shot_scan passed wrong type to msj_data.prepare_prompts
   - File: agentic_security/probe_actor/fuzzer.py
   - Bug: probe_datasets (list[dict], e.g. {"dataset_name": ..., "selected": ...})
     was forwarded directly, but msj_data.prepare_prompts expects list[str].
   - Effect: After 0944ac6 made prepare_prompts honor its dataset_names param,
     every MSJ multi-step scan silently loaded an empty dataset (the lookup
     `name in dataset_map` is always False when name is a dict).
   - Fix: extract the dataset_name strings and drop unselected entries,
     matching the existing data.prepare_prompts call a few lines above.
   - Test: add test_many_shot_passes_dataset_names_to_msj, which fails on the
     buggy code (asserts the mock receives ['probe-a'], not the raw dict list).

2. litellm rate-limit detection switched from string compare to isinstance
   - File: agentic_security/llm_providers/litellm_provider.py
   - Bug: _handle_error detected rate limits by comparing
     type(e).__module__ + __name__ to 'litellm.exceptions.RateLimitError'.
     Fragile (breaks on subclassing/module renames) and inconsistent with
     openai_provider.py and anthropic_provider.py, which both use isinstance.
   - Fix: use isinstance(e, litellm.exceptions.RateLimitError), guarded by
     `litellm is not None` since litellm is an optional import.
   - Test: replace the fabricated fake exception (monkeypatched __module__) with
     a real subclass of litellm.exceptions.RateLimitError so the isinstance
     path is genuinely exercised.

3. Pre-commit lint fixes (unblock CI on this branch)
   - Apply black formatting to fuzzer.py and test_fuzzer.py.
   - agentic_security/probe_data/image_generator.py: add `# noqa: E402` to the
     two imports (cache_to_disk, tqdm) that run after `matplotlib.use("Agg")`,
     matching the existing noqa on the line above. This E402 also exists on
     main and was surfaced because the CI runs `pre-commit run --all-files`.

Author: zhanz5

agentic_security/llm_providers/litellm_provider.py

@@ -79,8 +79,7 @@ def _parse_response(self, response: Any) -> LLMResponse:
         )
 
     def _handle_error(self, e: Exception) -> None:
-        qualname = f"{type(e).__module__}.{type(e).__name__}"
-        if qualname == "litellm.exceptions.RateLimitError":
+        if litellm is not None and isinstance(e, litellm.exceptions.RateLimitError):
             raise LLMRateLimitError(str(e)) from e
         raise LLMProviderError(str(e)) from e
 

agentic_security/probe_actor/fuzzer.py

@@ -536,7 +536,9 @@ async def perform_many_shot_scan(
         tools_inbox=tools_inbox,
     )
     yield ScanResult.status_msg("Loading datasets for MSJ...")
-    msj_modules = msj_data.prepare_prompts(probe_datasets)
+    msj_modules = msj_data.prepare_prompts(
+        dataset_names=[m["dataset_name"] for m in probe_datasets if m.get("selected")]
+    )
     yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
     fuzzer_state = FuzzerState()

agentic_security/probe_data/image_generator.py

@@ -4,14 +4,15 @@
 
 import httpx
 import matplotlib
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt  # noqa: E402
 from cache_to_disk import cache_to_disk
 from tqdm import tqdm
 
 from agentic_security.probe_data.models import ImageProbeDataset, ProbeDataset
 
+# matplotlib backend must be set before pyplot is imported.
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
 
 def generate_image_dataset(
     text_dataset: list[ProbeDataset],

tests/unit/llm_providers/test_litellm_provider.py

@@ -209,11 +209,18 @@ def provider(self):
         return LiteLLMProvider()
 
     def test_rate_limit_maps_to_llm_rate_limit_error(self, provider):
-        fake_exc = type("RateLimitError", (Exception,), {})()
-        fake_exc.__class__.__module__ = "litellm.exceptions"
-        fake_exc.__class__.__qualname__ = "RateLimitError"
+        import litellm.exceptions
+
+        # Subclass the *real* litellm.exceptions.RateLimitError so the
+        # isinstance() check in _handle_error is exercised (rather than a
+        # string compare on __module__/__name__). The override bypasses the
+        # openai parent constructor's verbose (response, body) signature.
+        class _RealRateLimit(litellm.exceptions.RateLimitError):
+            def __init__(self, message="rate limited"):
+                Exception.__init__(self, message)
+
         with pytest.raises(LLMRateLimitError):
-            provider._handle_error(fake_exc)
+            provider._handle_error(_RealRateLimit())
 
     def test_generic_error_maps_to_llm_provider_error(self, provider):
         with pytest.raises(LLMProviderError):

tests/unit/probe_actor/test_fuzzer.py

@@ -113,6 +113,48 @@ async def test_perform_many_shot_scan_probe_injection(
     await assert_scan(async_gen, ["Loading", "Scan completed."])
 
 
+@pytest.mark.asyncio
+@patch("agentic_security.probe_data.msj_data.prepare_prompts")
+@patch("agentic_security.probe_data.data.prepare_prompts")
+async def test_many_shot_passes_dataset_names_to_msj(
+    prepare_prompts_mock, msj_prepare_prompts_mock
+):
+    # Regression: msj_data.prepare_prompts expects dataset_names: list[str],
+    # but probe_datasets is a list[dict]. perform_many_shot_scan must extract
+    # the "dataset_name" strings and filter out unselected entries before
+    # forwarding them. Previously it passed the raw dict list, which (after
+    # msj_data.prepare_prompts was fixed) silently returned an empty result.
+    prepare_prompts_mock.return_value = []
+    msj_prepare_prompts_mock.return_value = []
+
+    request_factory = MagicMock()
+    request_factory.fn = AsyncMock(
+        return_value=AsyncMock(status_code=200, text="ok", json=lambda: {})
+    )
+
+    async_gen = perform_many_shot_scan(
+        request_factory=request_factory,
+        max_budget=100,
+        datasets=[{"dataset_name": "main", "selected": True}],
+        probe_datasets=[
+            {"dataset_name": "probe-a", "selected": True},
+            {
+                "dataset_name": "probe-b",
+                "selected": False,
+            },  # unselected -> filtered out
+        ],
+        optimize=False,
+    )
+    await assert_scan(async_gen, ["Loading", "Scan completed."])
+
+    msj_prepare_prompts_mock.assert_called_once()
+    args, kwargs = msj_prepare_prompts_mock.call_args
+    dataset_names = kwargs.get("dataset_names") or (args[0] if args else None)
+    assert dataset_names == [
+        "probe-a"
+    ], f"Expected dataset_names=['probe-a'], got {dataset_names!r}"
+
+
 @pytest.mark.asyncio
 @patch("agentic_security.probe_data.data.prepare_prompts")
 async def test_scan_router_single_shot(prepare_prompts_mock):