diff --git a/agentic_security/app.py b/agentic_security/app.py
index cbb8c1d..1cf0ef3 100644
--- a/agentic_security/app.py
+++ b/agentic_security/app.py
@@ -33,6 +33,9 @@ app.add_middleware(
)
tools_inbox = Queue()
+# Global stop event for cancelling scans
+stop_event = Event() # Added stop_event to cancel the scan
+
FEATURE_PROXY = False
@@ -99,6 +102,7 @@ def streaming_response_generator(scan_parameters: Scan):
datasets=scan_parameters.datasets,
tools_inbox=tools_inbox,
optimize=scan_parameters.optimize,
+ stop_event=stop_event, # Pass the stop_event to the generator
):
yield scan_result + "\n" # Adding a newline for separation
@@ -238,6 +242,12 @@ config.dictConfig(
)
+@app.post("/stop")
+async def stop_scan():
+ stop_event.set() # Set the stop event to cancel the scan
+ return {"status": "Scan stopped"}
+
+
class LogNon200ResponsesMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
try:
diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py
index b4fddbb..88b8235 100644
--- a/agentic_security/probe_actor/fuzzer.py
+++ b/agentic_security/probe_actor/fuzzer.py
@@ -1,3 +1,4 @@
+import asyncio
import os
from typing import AsyncGenerator
@@ -49,6 +50,7 @@ async def perform_scan(
datasets: list[dict[str, str]] = [],
tools_inbox=None,
optimize=False,
+ stop_event: asyncio.Event = None,
) -> AsyncGenerator[str, None]:
if IS_VERCEL:
yield ScanResult.status_msg(
@@ -81,6 +83,12 @@ async def perform_scan(
)
should_stop_early = False
async for prompt in prompt_iter(module.prompts):
+ if stop_event and stop_event.is_set(): # Check if stop_event is set
+ stop_event.clear() # Clear the event for the next scan
+ logger.info("Scan stopped by user.")
+ yield ScanResult.status_msg("Scan stopped by user.")
+ return # Exit the scan gracefully
+
processed_prompts += 1
progress = 100 * processed_prompts / total_prompts if total_prompts else 0
diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html
index fc8d325..04c4627 100644
--- a/agentic_security/static/index.html
+++ b/agentic_security/static/index.html
@@ -335,6 +335,7 @@
+
diff --git a/agentic_security/static/main.js b/agentic_security/static/main.js
index 3b02fd0..abbea52 100644
--- a/agentic_security/static/main.js
+++ b/agentic_security/static/main.js
@@ -341,6 +341,15 @@ var app = new Vue({
}
this.budget = value;
},
+ stopScan: async function () {
+ this.scanRunning = false;
+ const response = await fetch(`${URL}/stop`, {
+ method: 'POST',
+ headers: {
+ 'Content-Type': 'application/json',
+ },
+ });
+ },
startScan: async function () {
this.showLLMSpec = false;
let payload = {
@@ -358,6 +367,7 @@ var app = new Vue({
});
this.okMsg = 'Scan started';
this.mainTable = [];
+ this.scanRunning = true;
const reader = response.body.getReader();
let receivedLength = 0; // received that many bytes at the moment
let chunks = []; // array of received binary chunks (comprises the body)
diff --git a/pyproject.toml b/pyproject.toml
index 6e5b17e..81dc2e7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "agentic_security"
-version = "0.2.1"
+version = "0.2.2"
description = "Agentic LLM vulnerability scanner"
authors = ["Alexander Miasoiedov "]
maintainers = ["Alexander Miasoiedov "]