akto-api-security · notshivansh · Nov 7, 2025
diff --git a/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java b/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java
@@ -189,7 +189,7 @@ public YamlTestResult execute(ExecutorNode node, RawApi rawApi, Map<String, Obje
                 TestResult res = null;
                 if (AgentClient.isRawApiValidForAgenticTest(testReq)) {
                     // execute agentic test here
-                    res = agentClient.executeAgenticTest(testReq);
+                    res = agentClient.executeAgenticTest(testReq, apiInfoKey.getApiCollectionId());
                 }else{
                     List<String> contentType = origRawApi.getRequest().getHeaders().getOrDefault("content-type", new ArrayList<>());
                     String contentTypeString = "";
@@ -655,7 +655,8 @@ public ExecutorSingleOperationResp runOperation(String operationType, RawApi raw
                 }
             case "conversations_list":
                 List<String> conversationsList = (List<String>) value;
-                return Operations.addHeader(rawApi, "x-agent-conversations", String.join(",", conversationsList));
+                rawApi.setConversationsList(conversationsList);
+                return Operations.addHeader(rawApi, Constants.AKTO_AGENT_CONVERSATIONS , "0");
             case "attach_file":
                 return Operations.addHeader(rawApi, Constants.AKTO_ATTACH_FILE , key.toString());
             case "add_body_param":

diff --git a/apps/mini-testing/src/main/java/com/akto/testing/Main.java b/apps/mini-testing/src/main/java/com/akto/testing/Main.java
@@ -120,7 +120,6 @@ private static void handleTestEditorPlayground(TestingRunPlayground testingRunPl
         ApiInfo.ApiInfoKey infoKey = testingRunPlayground.getApiInfoKey();
 
         List<String> sampleData = testingRunPlayground.getSamples(); // get sample data from DB
-        AgenticUtils.checkAndInitializeAgent(Collections.singleton(infoKey.getApiCollectionId()), true, testConfig.getApiSelectionFilters().getNode());
 
         List<TestingRunResult.TestLog> testLogs = new ArrayList<>();
         Map<ApiInfo.ApiInfoKey, List<String>> sampleDataMap = new HashMap<>();

diff --git a/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java b/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java
@@ -277,7 +277,6 @@ public void apiWiseInit(TestingRun testingRun, ObjectId summaryId, boolean debug
 
             // create count down latch to know when inserting kafka records are completed.
             CountDownLatch latch = new CountDownLatch(apiInfoKeyList.size());
-            AgenticUtils.checkAndInitializeAgent(Main.extractApiCollectionIds(apiInfoKeyList), false, null);
             int tempRunTime = 10 * 60;
             if(!Constants.IS_NEW_TESTING_ENABLED){
                 tempRunTime = testingRun.getTestRunTime() <= 0 ? 30*60 : testingRun.getTestRunTime();

diff --git a/.../mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java b/.../mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java
@@ -157,7 +157,7 @@ public NodeResult processNode(Node node, Map<String, Object> varMap, Boolean all
                 TestResult res = null;
                 if (AgentClient.isRawApiValidForAgenticTest(testReq)) {
                     // execute agentic test here
-                    res = agentClient.executeAgenticTest(testReq);
+                    res = agentClient.executeAgenticTest(testReq, yamlNodeDetails.getApiCollectionId());
                 }else{
                     tsBeforeReq = Context.nowInMillis();
                     testResponse = ApiExecutor.sendRequest(testReq.getRequest(), followRedirect, testingRunConfig, debug, testLogs, Main.SKIP_SSRF_CHECK);

diff --git a/libs/dao/src/main/java/com/akto/dto/ApiCollection.java b/libs/dao/src/main/java/com/akto/dto/ApiCollection.java
@@ -389,4 +389,11 @@ public boolean isGenAICollection() {
         return false;
     }
 
+    public boolean isGuardRailCollection() {
+        if (!CollectionUtils.isEmpty(this.getTagsList())) {
+            return this.getTagsList().stream().anyMatch(t -> Constants.AKTO_GUARD_RAIL_TAG.equals(t.getKeyName()));
+        }
+        return false;
+    }
+
 }
diff --git a/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java b/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java
@@ -2,19 +2,18 @@
 
 import com.akto.dto.testing.TLSAuthParam;
 import com.akto.dto.type.RequestTemplate;
+import com.akto.util.Constants;
 import com.akto.util.HttpRequestResponseUtils;
 import com.alibaba.fastjson2.JSON;
 import com.google.gson.Gson;
 import com.mongodb.BasicDBObject;
 
-import java.util.stream.Collectors;
 import lombok.Getter;
 import lombok.Setter;
 import okhttp3.Headers;
 import okhttp3.HttpUrl;
 import java.net.URI;
 import java.util.*;
-import org.apache.commons.lang3.StringUtils;
 
 public class OriginalHttpRequest {
 
@@ -412,6 +411,32 @@ public Headers toOkHttpHeaders() {
         return builder.build();
     }
 
+    public String fetchHeadersJsonString() {
+        Map<String, List<String>> headersMap = this.getHeaders();
+        List<String> forbiddenHeaders = Arrays.asList("content-length", "accept-encoding");
+        if (headersMap == null)
+            headersMap = new HashMap<>();
+        headersMap.put(Constants.AKTO_IGNORE_FLAG, Collections.singletonList("0"));
+        Map<String, String> filteredHeaders = new HashMap<>();
+        for (String headerName : headersMap.keySet()) {
+            if (forbiddenHeaders.contains(headerName))
+                continue;
+            if (headerName.contains(" "))
+                continue;
+            if(headerName.startsWith(":")) continue;
+            List<String> headerValueList = headersMap.get(headerName);
+            if (headerValueList == null || headerValueList.isEmpty())
+                continue;
+            for (String headerValue : headerValueList) {
+                if (headerValue == null)
+                    continue;
+                filteredHeaders.put(headerName, headerValue);
+                break;
+            }
+        }
+        return gson.toJson(filteredHeaders);
+    }
+
     @Override
     public String toString() {
         return "OriginalHttpRequest{" +

diff --git a/libs/dao/src/main/java/com/akto/dto/RawApi.java b/libs/dao/src/main/java/com/akto/dto/RawApi.java
@@ -10,14 +10,18 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.mongodb.BasicDBObject;
 
-import javax.naming.LimitExceededException;
+import lombok.Getter;
+import lombok.Setter;
 
 public class RawApi {
 
     private OriginalHttpRequest request;
     private OriginalHttpResponse response;
     private String originalMessage;
     static ObjectMapper om = new ObjectMapper();
+    @Getter
+    @Setter
+    private List<String> conversationsList;
 
     public RawApi(OriginalHttpRequest request, OriginalHttpResponse response, String originalMessage) {
         this.request = request;

diff --git a/libs/dao/src/main/java/com/akto/util/Constants.java b/libs/dao/src/main/java/com/akto/util/Constants.java
@@ -24,6 +24,7 @@ private Constants() {}
     public static final String AKTO_TOKEN_KEY = "x-akto-key";
     public static final String AKTO_NODE_ID = "x-akto-node";
     public static final String AKTO_REMOVE_AUTH= "x-akto-remove-auth";
+    public static final String AKTO_AGENT_CONVERSATIONS= "x-agent-conversations";
 
     public static final String LOCAL_KAFKA_BROKER_URL = System.getenv("KAFKA_BROKER_URL") != null ? System.getenv("KAFKA_BROKER_URL") : "localhost:29092"; // run kafka process with name kafka1 in docker
     public static final String TEST_RESULTS_TOPIC_NAME = "akto.test.messages";
@@ -45,6 +46,7 @@ private Constants() {}
     public final static String _AKTO = "AKTO";
     public static final String AKTO_MCP_SERVER_TAG = "mcp-server";
     public static final String AKTO_GEN_AI_TAG = "gen-ai";
+    public static final String AKTO_GUARD_RAIL_TAG = "guard-rail";
     public static final String AKTO_MCP_TOOLS_TAG = "mcp-tool";
     public static final String AKTO_MCP_RESOURCES_TAG = "mcp-resource";
     public static final String AKTO_MCP_PROMPTS_TAG = "mcp-prompt";

diff --git a/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java b/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java
@@ -21,6 +21,9 @@ public String getName() {
 
     /* Category of tests perfomred */
     public enum TestCategory {
+
+        // whenever adding a new test category, please ensure that if mcp or LLM is used, update the list of categories in TestTemplateUtils.java which will be used to fetch templates in fetchAllSubCategories
+
         BOLA("BOLA", Severity.HIGH, "Broken Object Level Authorization (BOLA)", "BOLA"),
         NO_AUTH("NO_AUTH", Severity.HIGH, "Broken User Authentication (BUA)", "Broken Authentication"),
         BFLA("BFLA", Severity.HIGH, "Broken Function Level Authorization (BFLA)", "Broken Function Level Authorization"),
@@ -44,7 +47,35 @@ public enum TestCategory {
         LFI("LFI", Severity.HIGH, "Local File Inclusion (LFI)", "Local File Inclusion"),
         XSS("XSS", Severity.HIGH, "Cross-site scripting (XSS)", "Cross-site scripting"),
         IIM("IIM", Severity.HIGH, "Improper Inventory Management (IIM)", "Improper Inventory Management"),
-        LLM("LLM",Severity.HIGH,"LLM (Large Language Models) Top 10","LLM");
+        LLM("LLM",Severity.HIGH,"LLM (Large Language Models) Top 10","LLM"),
+        INJECT("INJECT", Severity.MEDIUM, "Injection Attacks (INJECT)", "Injection Attacks"),
+        INPUT("INPUT", Severity.MEDIUM, "Input Validation (INPUT)", "Input Validation"),
+        PROMPT_INJECTION("PROMPT_INJECTION", Severity.HIGH, "Prompt Injection", "PromptInjection"),
+        SENSITIVE_INFORMATION_DISCLOSURE("SENSITIVE_INFORMATION_DISCLOSURE", Severity.HIGH, "Sensitive Information Disclosure", "SensitiveInformationDisclosure"),
+        SUPPLY_CHAIN("SUPPLY_CHAIN", Severity.HIGH, "Supply Chain", "SupplyChain"),
+        DATA_AND_MODEL_POISONING("DATA_AND_MODEL_POISONING", Severity.HIGH, "DataAndModelPoisoning", "DataAndModelPoisoning"),
+        IMPROPER_OUTPUT_HANDLING("IMPROPER_OUTPUT_HANDLING", Severity.HIGH, "ImproperOutputHandling", "ImproperOutputHandling"),
+        EXCESSIVE_AGENCY("EXCESSIVE_AGENCY", Severity.HIGH, "Excessive Agency", "ExcessiveAgency"),
+        SYSTEM_PROMPT_LEAKAGE("SYSTEM_PROMPT_LEAKAGE", Severity.HIGH, "System Prompt Leakage", "SystemPromptLeakage"),
+        VECTOR_AND_EMBEDDING_WEAKNESSES("VECTOR_AND_EMBEDDING_WEAKNESSES", Severity.HIGH, "Vector and Embedding Weaknesses", "VectorAndEmbeddingWeaknesses"),
+        MISINFORMATION("MISINFORMATION", Severity.HIGH, "Misinformation", "Misinformation"),
+        UNBOUNDED_CONSUMPTION("UNBOUNDED_CONSUMPTION", Severity.HIGH, "Unbounded Consumption", "UnboundedConsumption"),
+        MCP_PROMPT_INJECTION("MCP_PROMPT_INJECTION", Severity.HIGH, "MCP - Prompt Injection", "MCP_PROMPT_INJECTION"),
+        MCP_TOOL_POISONING("MCP_TOOL_POISONING", Severity.HIGH, "MCP - Tool Poisoning", "MCP_TOOL_POISONING"),
+        MCP_PRIVILEGE_ABUSE("MCP_PRIVILEGE_ABUSE", Severity.HIGH, "MCP - Privilege Abuse", "MCP_PRIVILEGE_ABUSE"),
+        MCP_INDIRECT_PROMPT_INJECTION("MCP_INDIRECT_PROMPT_INJECTION", Severity.HIGH, "MCP - Indirect Prompt Injection", "MCP_INDIRECT_PROMPT_INJECTION"),
+        MCP_SENSITIVE_DATA_LEAKAGE("MCP_SENSITIVE_DATA_LEAKAGE", Severity.HIGH, "MCP - Data Leak", "MCP_SENSITIVE_DATA_LEAKAGE"),
+        MCP_DOS("MCP_DOS", Severity.HIGH, "MCP - Denial of Service", "MCP_DOS"),
+        MCP_AUTH("MCP_AUTH", Severity.HIGH, "MCP - Broken Authentication", "MCP_AUTH"),
+        MCP_MALICIOUS_CODE_EXECUTION("MCP_MALICIOUS_CODE_EXECUTION", Severity.HIGH, "MCP - Malicious Code Execution", "MCP_MALICIOUS_CODE_EXECUTION"),
+        MCP("MCP", Severity.HIGH, "Model Context Protocol (MCP) Security", "MCP"),
+        MCP_INPUT_VALIDATION("MCP_INPUT_VALIDATION", Severity.HIGH, "MCP - Input Validation", "MCP_INPUT_VALIDATION"),
+        MCP_FUNCTION_MANIPULATION("MCP_FUNCTION_MANIPULATION", Severity.CRITICAL, "Model Context Protocol (MCP) Security - Function Call Manipulation", "MCP_FUNC_MANIP"),
+        MCP_SECURITY("MCP_SECURITY", Severity.HIGH, "Model Context Protocol (MCP) Security", "MCP_SEC"),
+        AGENTIC_BUSINESS_ALIGNMENT("AGENTIC_BUSINESS_ALIGNMENT", Severity.HIGH, "Agent Business Alignment", "AGENTIC_BUSINESS_ALIGNMENT"),
+        AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS("AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS", Severity.HIGH, "Agent Hallucination and Trustworthiness", "AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS"),
+        AGENTIC_SAFETY("AGENTIC_SAFETY", Severity.HIGH, "Agent Safety", "AGENTIC_SAFETY"),
+        AGENTIC_SECURITY("AGENTIC_SECURITY", Severity.HIGH, "Agent Security", "AGENTIC_SECURITY");
 
         private final String name;
         private final Severity severity;
@@ -108,6 +139,30 @@ public enum YamlTemplateSource {
         CUSTOM
     }
 
+    public enum TemplatePlan {
+        FREE, STANDARD, PRO, ENTERPRISE
+    }
+    public enum TemplateFeatureAccess {
+        PRO_TESTS, ENTERPRISE_TESTS
+    }
+    public enum TemplateNature {
+        INTRUSIVE, NON_INTRUSIVE
+    }
+    public enum TemplateDuration {
+        SLOW, FAST
+    }
+
+    public enum ENCODING_TYPE{
+        BASE_64_ENCODED, JWT
+    }
+
+    public enum TicketSource {
+        JIRA, AZURE_BOARDS, SERVICENOW
+    }
+
+    public enum CONTEXT_SOURCE {
+        API, MCP, GEN_AI, AGENTIC
+    }
 
     /* ********************************************************************** */
 }
diff --git a/libs/utils/src/main/java/com/akto/agent/AgentClient.java b/libs/utils/src/main/java/com/akto/agent/AgentClient.java
@@ -42,19 +42,28 @@ public AgentClient(String agentBaseUrl) {
         this.agentBaseUrl = agentBaseUrl.endsWith("/") ? agentBaseUrl.substring(0, agentBaseUrl.length() - 1) : agentBaseUrl;
     }
 
-    public TestResult executeAgenticTest(RawApi rawApi) throws Exception {
+    public TestResult executeAgenticTest(RawApi rawApi, int apiCollectionId) throws Exception {
         String conversationId = UUID.randomUUID().toString();
-        String prompts = rawApi.getRequest().getHeaders().get("x-agent-conversations").get(0);
-        List<String> promptsList = Arrays.asList(prompts.split(","));
+        List<String> promptsList = rawApi.getConversationsList();
         String testMode = getTestModeFromRole();
 
+        try {
+            /*
+             * the rawApi already has been modified by the testRole, 
+             * and should have the updated auth request headers
+             */
+            AgenticUtils.checkAndInitializeAgent(conversationId, rawApi, apiCollectionId);
+        } catch(Exception e){
+        }
+
         try {
             List<AgentConversationResult> conversationResults = processConversations(promptsList, conversationId, testMode);
 
             boolean isVulnerable = conversationResults.get(conversationResults.size() - 1).isValidation();
             List<String> errors = new ArrayList<>();
 
             TestResult testResult = new TestResult();
+            // TODO: Fill in message field
             testResult.setMessage(null);
             testResult.setConversationId(conversationId);
             testResult.setResultTypeAgentic(true);
@@ -84,7 +93,7 @@ public TestResult executeAgenticTest(RawApi rawApi) throws Exception {
         }
     }
 
-    public List<AgentConversationResult> processConversations(List<String> prompts, String conversationId, String testMode) throws Exception {
+    private List<AgentConversationResult> processConversations(List<String> prompts, String conversationId, String testMode) throws Exception {
         List<AgentConversationResult> results = new ArrayList<>();
         int index = 0;
         int totalRequests = prompts.size();
@@ -102,7 +111,7 @@ public List<AgentConversationResult> processConversations(List<String> prompts,
         return results;
     }
 
-    public AgentConversationResult sendChatRequest(String prompt, String conversationId, String testMode, boolean isLastRequest) throws Exception {
+    private AgentConversationResult sendChatRequest(String prompt, String conversationId, String testMode, boolean isLastRequest) throws Exception {
         Request request = buildOkHttpChatRequest(prompt, conversationId, isLastRequest);
 
         try (Response response = agentHttpClient.newCall(request).execute()) {
@@ -172,12 +181,13 @@ private void storeConversationResults(List<AgentConversationResult> conversation
         try {
             dataActor.storeConversationResults(conversationResults);
         } catch (Exception e) {
-            loggerMaker.error("Error storing conversation results: " + e.getMessage());
+            loggerMaker.errorAndAddToDb("Error storing conversation results: " + e.getMessage());
         }
     }
 
     public static boolean isRawApiValidForAgenticTest(RawApi rawApi) {
-        return rawApi.getRequest().getHeaders().containsKey("x-agent-conversations");
+        List<String> temp = rawApi.getConversationsList();
+        return (temp != null && !temp.isEmpty());
     }
 
     public boolean performHealthCheck() {
@@ -192,8 +202,24 @@ public boolean performHealthCheck() {
         }
     }
     public void initializeAgent(String sseUrl, String authorizationToken) {
+        Map<String, Object> requestBody = new HashMap<>();
+        requestBody.put("sseUrl", sseUrl);
+        requestBody.put("authorization", authorizationToken);
+        initializeAgent(requestBody);
+    }
+
+    public void initializeAgent(String sessionUrl, String requestHeaders, String apiRequestBody, String conversationId) {
+        Map<String, Object> requestBody = new HashMap<>();
+        requestBody.put("sessionUrl", sessionUrl);
+        requestBody.put("requestHeaders", requestHeaders);
+        requestBody.put("requestBody", apiRequestBody);
+        requestBody.put("conversationId", conversationId);
+        initializeAgent(requestBody);
+    }
+
+    public void initializeAgent(Map<String, Object> requestBody) {
         try {
-            Request initRequest = buildOkHttpInitializeRequest(sseUrl, authorizationToken);
+            Request initRequest = buildOkHttpInitializeRequest(requestBody);
             try (Response response = agentHttpClient.newCall(initRequest).execute()) {
                 if (!response.isSuccessful()) {
                     loggerMaker.errorAndAddToDb("Agent initialization failed with status: " + response.code());
@@ -212,18 +238,15 @@ private Request buildOkHttpHealthCheckRequest() {
                 .build();
     }
 
-
-    private Request buildOkHttpInitializeRequest(String sseUrl, String authorizationToken) {
-        Map<String, Object> requestBody = new HashMap<>();
-        requestBody.put("sseUrl", sseUrl);
-        requestBody.put("authorization", authorizationToken);
-
-        String body;
+    private Request buildOkHttpInitializeRequest(Map<String, Object> requestBody) {
+
+        String body = "";
         try {
             body = objectMapper.writeValueAsString(requestBody);
         } catch (Exception e) {
             loggerMaker.errorAndAddToDb("Error serializing initialize request body: " + e.getMessage());
-            body = "{\"sseUrl\":\"" + sseUrl.replace("\"", "\\\"") + "\"}";
+            // TODO: Fix at sse URL.
+            // body = "{\"sseUrl\":\"" + sseUrl.replace("\"", "\\\"") + "\"}";
         }
 
         RequestBody requestBodyObj = RequestBody.create(body, MediaType.parse("application/json"));