From 85312720d53cfe31bf2ef86bec0a26cb1957ec0e Mon Sep 17 00:00:00 2001 From: Migration Bot Date: Thu, 12 Mar 2026 06:58:23 +0000 Subject: [PATCH 1/2] feat(exceptions): add typed exception hierarchy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - AgentException (base RuntimeException for full pipeline) - ScanException (moved from inner class in SecuritySupervisorAgentService) - McpToolNullResponseException (MCP returned null/empty — stops LLM loops) - McpToolExecutionException (MCP tool threw during execution) - MaxLlmCallsExceededException (loop guard circuit-breaker) - AgentRoutingException (supervisor cannot determine valid route) Fixes: ScanException was an inner class — non-standard, untestable. All exceptions extend AgentException for single-catch supervisor handling. --- .../service/DependencyParserService.java | 4 +- .../service/IDependencyParserService.java | 29 +++ .../com/agent/service/GithubRepoService.java | 4 +- .../com/agent/service/IGithubRepoService.java | 37 +++ .../service/impl/GithubRepoServiceImpl.java | 167 +++++++++++++ .../java/com/agent/agent/SupervisorAgent.java | 218 ++++++++-------- .../sub_agents/DependencyParserAgent.java | 219 ++++++---------- .../agent/sub_agents/DirectAnswerAgent.java | 24 +- .../agent/sub_agents/GithubRepoAgent.java | 216 +++++----------- .../agent/sub_agents/VulnerabilityAgent.java | 233 +++++++----------- .../com/agent/controller/ScanController.java | 2 +- .../com/agent/exception/AgentException.java | 17 ++ .../exception/AgentRoutingException.java | 20 ++ .../MaxLlmCallsExceededException.java | 36 +++ .../exception/McpToolExecutionException.java | 27 ++ .../McpToolNullResponseException.java | 24 ++ .../com/agent/exception/ScanException.java | 19 ++ .../com/agent/model/ConversationContext.java | 90 +++++++ .../java/com/agent/prompt/PromptLibrary.java | 143 +++++------ .../SecuritySupervisorAgentService.java | 123 ++++----- .../java/com/agent/tool/McpToolRegistry.java | 2 +- .../java/com/agent/tool/ToolRegistry.java | 47 ++++ .../main/java/com/agent/util/AgentUtils.java | 101 ++++++++ .../com/agent/util/McpResponseValidator.java | 54 ++++ .../agent/service/IVulnerabilityService.java | 31 +++ .../agent/service/VulnerabilityService.java | 4 +- 26 files changed, 1188 insertions(+), 703 deletions(-) create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/IDependencyParserService.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentRoutingException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/MaxLlmCallsExceededException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolExecutionException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolNullResponseException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/ScanException.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/ToolRegistry.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/AgentUtils.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/McpResponseValidator.java create mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/IVulnerabilityService.java diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/DependencyParserService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/DependencyParserService.java index f508379..7a44acb 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/DependencyParserService.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/DependencyParserService.java @@ -1,5 +1,7 @@ package com.agent.service; +import com.agent.service.IDependencyParserService; + import com.agent.model.Dependency; import com.agent.parser.*; import lombok.RequiredArgsConstructor; @@ -11,7 +13,7 @@ @Slf4j @Service @RequiredArgsConstructor -public class DependencyParserService { +public class DependencyParserService implements IDependencyParserService { private final PomParser pomParser; private final NpmParser npmParser; diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/IDependencyParserService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/IDependencyParserService.java new file mode 100644 index 0000000..bfaefc3 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/service/IDependencyParserService.java @@ -0,0 +1,29 @@ +package com.agent.service; + +import com.agent.model.Dependency; + +import java.util.List; + +/** + * IDependencyParserService — contract for dependency file parsers. + * + * Extracted from the concrete class to enable mock injection in unit tests. + * Implementation: {@link impl.DependencyParserServiceImpl} + */ +public interface IDependencyParserService { + + /** Parse Maven pom.xml content */ + List parsePom(String content); + + /** Parse npm package.json content */ + List parsePackageJson(String content) throws Exception; + + /** Parse pip requirements.txt content */ + List parseRequirements(String content); + + /** Parse Gradle build.gradle / build.gradle.kts content */ + List parseGradle(String content); + + /** Parse Flutter/Dart pubspec.yaml content */ + List parsePubspec(String content); +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/GithubRepoService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/GithubRepoService.java index 4b1817f..3598c11 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/GithubRepoService.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/GithubRepoService.java @@ -1,5 +1,7 @@ package com.agent.service; +import com.agent.service.IGithubRepoService; + import com.agent.util.FileScannerUtil; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Service; @@ -15,7 +17,7 @@ @Slf4j @Service -public class GithubRepoService { +public class GithubRepoService implements IGithubRepoService { private static final String REPO_BASE_DIR = "/tmp/repos"; private static final int CLONE_TIMEOUT_SECONDS = 120; diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java new file mode 100644 index 0000000..4370640 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java @@ -0,0 +1,37 @@ +package com.agent.service; + +import java.io.IOException; +import java.util.List; + +/** + * GithubRepoService — contract for repository access operations. + * + * Extracted from the concrete class to enable mock injection in unit tests. + * Implementation: {@link impl.GithubRepoServiceImpl} + */ +public interface GithubRepoService { + + /** + * Clone a GitHub repository to a local temp directory. + * + * @param repoUrl full HTTPS GitHub URL + * @param forceReclone if true, delete and re-clone even if already present + * @return absolute path to the cloned repository root + */ + String cloneRepository(String repoUrl, boolean forceReclone) throws Exception; + + /** + * List all files recursively under repoPath (up to configured depth). + */ + List listRepositoryFiles(String repoPath) throws IOException; + + /** + * Find dependency manifest files (pom.xml, package.json, etc.) under repoPath. + */ + List findDependencyFiles(String repoPath) throws IOException; + + /** + * Read the content of a single file (path-traversal protected). + */ + String readFileContent(String filePath) throws IOException; +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java new file mode 100644 index 0000000..318b746 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java @@ -0,0 +1,167 @@ +package com.agent.service.impl; +import com.agent.service.IGithubRepoService; + +import com.agent.util.FileScannerUtil; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.TimeUnit; + +@Slf4j +@Service +public class GithubRepoServiceImpl implements IGithubRepoService { + + private static final String REPO_BASE_DIR = "/tmp/repos"; + private static final int CLONE_TIMEOUT_SECONDS = 120; + + // Allowed: https://github.com/owner/repo or https://github.com/owner/repo.git + private static final String GITHUB_URL_PATTERN = + "https://github\\.com/[\\w.-]+/[\\w.-]+(\\.git)?"; + + /** + * Clone a GitHub repository into the local base directory. + * + * @param repoUrl public GitHub HTTPS URL + * @param forceReclone if true, deletes existing clone and re-clones fresh + */ + public String cloneRepository(String repoUrl, boolean forceReclone) throws Exception { + + // 1. Validate GitHub URL format + if (repoUrl == null || !repoUrl.matches(GITHUB_URL_PATTERN)) { + throw new IllegalArgumentException( + "Invalid GitHub URL. Expected: https://github.com/owner/repo — got: " + repoUrl + ); + } + + // 2. Derive safe directory name + String repoName = repoUrl + .substring(repoUrl.lastIndexOf("/") + 1) + .replace(".git", ""); + + String target = REPO_BASE_DIR + "/" + repoName; + Path targetPath = Paths.get(target); + + // 3. Handle existing clone + if (Files.exists(targetPath)) { + if (forceReclone) { + log.info("[SERVICE] forceReclone=true — deleting existing clone at: {}", target); + deleteDirectory(targetPath); + } else { + long fileCount = Files.walk(targetPath) + .filter(Files::isRegularFile) + .count(); + log.info("[SERVICE] Repository already cloned at: {} | existingFileCount={}", target, fileCount); + + if (fileCount == 0) { + log.warn("[SERVICE] Existing clone is EMPTY (possibly a failed previous clone) — force re-cloning: {}", target); + deleteDirectory(targetPath); + } else { + return target; + } + } + } + + // 4. Ensure base directory exists + Files.createDirectories(Paths.get(REPO_BASE_DIR)); + + log.info("[SERVICE] Cloning repository: {} → {}", repoUrl, target); + + ProcessBuilder builder = new ProcessBuilder("git", "clone", repoUrl, target); + builder.redirectErrorStream(true); + Process process = builder.start(); + + // 5. Enforce timeout + boolean finished = process.waitFor(CLONE_TIMEOUT_SECONDS, TimeUnit.SECONDS); + + if (!finished) { + process.destroyForcibly(); + throw new RuntimeException( + "Git clone timed out after " + CLONE_TIMEOUT_SECONDS + "s for: " + repoUrl + ); + } + + int exitCode = process.exitValue(); + if (exitCode != 0) { + String errorOutput = readProcessOutput(process); + throw new RuntimeException( + "Git clone failed (exit=" + exitCode + "): " + errorOutput + ); + } + + long clonedFileCount = Files.walk(targetPath) + .filter(Files::isRegularFile) + .count(); + + log.info("[SERVICE] Clone successful | path={} | totalFiles={}", target, clonedFileCount); + return target; + } + + /** + * List all non-internal files in a cloned repository. + */ + public List listRepositoryFiles(String repoPath) throws IOException { + validateRepoPath(repoPath); + List files = FileScannerUtil.listFiles(repoPath); + log.debug("[SERVICE] listRepositoryFiles | repoPath={} | count={}", repoPath, files.size()); + return files; + } + + /** + * Find dependency/manifest files in the repository. + */ + public List findDependencyFiles(String repoPath) throws IOException { + validateRepoPath(repoPath); + List depFiles = FileScannerUtil.findDependencyFiles(repoPath); + log.debug("[SERVICE] findDependencyFiles | repoPath={} | found={}", repoPath, depFiles.size()); + return depFiles; + } + + /** + * Read file content. Path traversal protection enforced in FileScannerUtil. + */ + public String readFileContent(String filePath) throws IOException { + return FileScannerUtil.readFile(filePath); + } + + // ------------------------------------------------------------------------- + // Private helpers + // ------------------------------------------------------------------------- + + private void validateRepoPath(String repoPath) { + Path resolved = Paths.get(repoPath).normalize(); + if (!resolved.startsWith(REPO_BASE_DIR)) { + throw new SecurityException( + "Access denied: path is outside allowed base directory. Path: " + repoPath + ); + } + } + + private void deleteDirectory(Path path) throws IOException { + Files.walk(path) + .sorted(Comparator.reverseOrder()) + .forEach(p -> { + try { + Files.delete(p); + } catch (IOException e) { + log.warn("[SERVICE] Failed to delete: {}", p); + } + }); + log.info("[SERVICE] Deleted directory: {}", path); + } + + private String readProcessOutput(Process process) { + try (InputStream is = process.getInputStream()) { + return new String(is.readAllBytes()); + } catch (IOException e) { + return "(could not read process output)"; + } + } + +} \ No newline at end of file diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/SupervisorAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/SupervisorAgent.java index 72acbac..6a3a568 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/SupervisorAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/SupervisorAgent.java @@ -4,8 +4,12 @@ import com.agent.agent.sub_agents.DirectAnswerAgent; import com.agent.agent.sub_agents.GithubRepoAgent; import com.agent.agent.sub_agents.VulnerabilityAgent; +import com.agent.exception.AgentRoutingException; +import com.agent.exception.MaxLlmCallsExceededException; +import com.agent.exception.McpToolNullResponseException; import com.agent.prompt.PromptLibrary; import com.agent.state.SecurityAgentState; +import com.agent.util.AgentUtils; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -26,34 +30,32 @@ /** * SupervisorAgent — orchestrates the multi-agent security scanning pipeline. - *

- * The LLM handles intent detection and routing in one unified reason node. - * JSON parsing uses ObjectMapper instead of regex for correctness. - *

- * Flows: - * REPO_SCAN → github_repo_agent → dependency_parser_agent → vulnerability_agent → DONE - * PACKAGE_SCAN → vulnerability_agent → DONE - * CVE_QUERY → vulnerability_agent → DONE - * FILE_SCAN → dependency_parser_agent → vulnerability_agent → DONE - * DIRECT_ANSWER → direct_answer_agent → END - *

- * Graph: START → reason → [sub-agent → reason]* → END + * + * Changes from original: + * 1. Supervisor iteration counter (MAX_SUPERVISOR_ITERATIONS = 6). + * 2. Typed exception handling in invoke(): MaxLlmCallsExceededException, + * McpToolNullResponseException, AgentRoutingException. + * 3. Two-mode prompt: full initial call vs compact routing call (~50 tokens). + * 4. Uses AgentUtils shared utilities. */ @Slf4j @Component @RequiredArgsConstructor public class SupervisorAgent { - public static final String INTENT_REPO_SCAN = "REPO_SCAN"; + public static final String INTENT_REPO_SCAN = "REPO_SCAN"; public static final String INTENT_PACKAGE_SCAN = "PACKAGE_SCAN"; - public static final String INTENT_CVE_QUERY = "CVE_QUERY"; - public static final String INTENT_FILE_SCAN = "FILE_SCAN"; + public static final String INTENT_CVE_QUERY = "CVE_QUERY"; + public static final String INTENT_FILE_SCAN = "FILE_SCAN"; - private static final String ROUTE_REPO = "github_repo_agent"; - private static final String ROUTE_PARSER = "dependency_parser_agent"; + private static final String ROUTE_REPO = "github_repo_agent"; + private static final String ROUTE_PARSER = "dependency_parser_agent"; private static final String ROUTE_VULNERABILITY = "vulnerability_agent"; - private static final String ROUTE_DIRECT = "direct_answer_agent"; - private static final String ROUTE_DONE = "DONE"; + private static final String ROUTE_DIRECT = "direct_answer_agent"; + private static final String ROUTE_DONE = "DONE"; + + private static final int MAX_SUPERVISOR_ITERATIONS = 6; + private static final String STATE_KEY_ITERATIONS = "supervisorIterations"; private final ChatClient chatClient; private final GithubRepoAgent repoAgent; @@ -64,45 +66,36 @@ public class SupervisorAgent { private CompiledGraph graph; - // ── Graph ───────────────────────────────────────────────────────────────── - private synchronized void build() throws Exception { if (graph != null) return; - graph = new StateGraph<>(SecurityAgentState.SCHEMA, SecurityAgentState::new) - .addNode("reason", node_async(this::reason)) - .addNode(ROUTE_REPO, node_async(this::runRepoAgent)) - .addNode(ROUTE_PARSER, node_async(this::runParserAgent)) - .addNode(ROUTE_VULNERABILITY, node_async(this::runVulnerabilityAgent)) - .addNode(ROUTE_DIRECT, node_async(this::runDirectAnswerAgent)) - - .addEdge(START, "reason") - .addEdge(ROUTE_REPO, "reason") - .addEdge(ROUTE_PARSER, "reason") + .addNode("reason", node_async(this::reason)) + .addNode(ROUTE_REPO, node_async(this::runRepoAgent)) + .addNode(ROUTE_PARSER, node_async(this::runParserAgent)) + .addNode(ROUTE_VULNERABILITY, node_async(this::runVulnerabilityAgent)) + .addNode(ROUTE_DIRECT, node_async(this::runDirectAnswerAgent)) + .addEdge(START, "reason") + .addEdge(ROUTE_REPO, "reason") + .addEdge(ROUTE_PARSER, "reason") .addEdge(ROUTE_VULNERABILITY, "reason") - .addEdge(ROUTE_DIRECT, END) - + .addEdge(ROUTE_DIRECT, END) .addConditionalEdges( "reason", edge_async(state -> { String route = state.nextAgent().orElse(ROUTE_DIRECT); - log.info("[Supervisor] Routing → {}", route); + log.info("[Supervisor] Routing -> {}", route); return switch (route) { - case ROUTE_REPO -> ROUTE_REPO; - case ROUTE_PARSER -> ROUTE_PARSER; + case ROUTE_REPO -> ROUTE_REPO; + case ROUTE_PARSER -> ROUTE_PARSER; case ROUTE_VULNERABILITY -> ROUTE_VULNERABILITY; - case ROUTE_DIRECT -> ROUTE_DIRECT; - case ROUTE_DONE -> END; - default -> ROUTE_DIRECT; + case ROUTE_DIRECT -> ROUTE_DIRECT; + case ROUTE_DONE -> END; + default -> ROUTE_DIRECT; }; }), - Map.of( - ROUTE_REPO, ROUTE_REPO, - ROUTE_PARSER, ROUTE_PARSER, - ROUTE_VULNERABILITY, ROUTE_VULNERABILITY, - ROUTE_DIRECT, ROUTE_DIRECT, - END, END - ) + Map.of(ROUTE_REPO, ROUTE_REPO, ROUTE_PARSER, ROUTE_PARSER, + ROUTE_VULNERABILITY, ROUTE_VULNERABILITY, + ROUTE_DIRECT, ROUTE_DIRECT, END, END) ) .compile(); } @@ -113,6 +106,25 @@ public Map invoke(Map initialState) { return graph.invoke(initialState) .map(SecurityAgentState::data) .orElse(initialState); + + } catch (MaxLlmCallsExceededException e) { + log.error("[Supervisor] LLM loop guard triggered — agent={} limit={}", e.getAgentName(), e.getLimit()); + Map err = new HashMap<>(initialState); + err.put("finalAnswer", "Analysis stopped: agent '" + e.getAgentName() + "' exceeded " + e.getLimit() + " LLM calls."); + return err; + + } catch (McpToolNullResponseException e) { + log.warn("[Supervisor] MCP null response — tool={}", e.getToolName()); + Map err = new HashMap<>(initialState); + err.put("finalAnswer", "No data returned from tool '" + e.getToolName() + "'. The MCP server may be unavailable."); + return err; + + } catch (AgentRoutingException e) { + log.error("[Supervisor] Routing failed — received='{}'", e.getReceivedRoute()); + Map err = new HashMap<>(initialState); + err.put("finalAnswer", "Pipeline routing error: " + e.getMessage()); + return err; + } catch (Exception e) { log.error("[Supervisor] Fatal error", e); Map err = new HashMap<>(initialState); @@ -121,39 +133,42 @@ public Map invoke(Map initialState) { } } - // ── reason node ─────────────────────────────────────────────────────────── - // - // First call: LLM detects intent + extracts entities + sets first route. - // Subsequent calls: LLM advances the pipeline based on what is now populated. - // Large blobs are summarised as POPULATED/null to save tokens. - private Map reason(SecurityAgentState state) { log.info("[Supervisor] Reasoning — intent={}", state.intent().orElse("none")); - // ── Hard guard: empty dependency list ───────────────────────────────── - // If the parser ran but produced 0 dependencies - // Supervisor looping forever trying to reach vulnerability_agent - // while vulnerabilities stays null. + // Supervisor iteration guard + int iterations = (int) state.data().getOrDefault(STATE_KEY_ITERATIONS, 0); + if (iterations >= MAX_SUPERVISOR_ITERATIONS) { + log.error("[Supervisor] Max iterations ({}) reached — forcing exception", MAX_SUPERVISOR_ITERATIONS); + throw new MaxLlmCallsExceededException("SupervisorAgent", MAX_SUPERVISOR_ITERATIONS); + } + + // Hard guard: empty dependency list String intentNow = state.intent().orElse(""); - if ("REPO_SCAN".equals(intentNow) || "FILE_SCAN".equals(intentNow)) { - boolean depsEmpty = state.dependencies().map(d -> { - String s = d.toString().trim(); - return s.equals("[]") || s.equals("null") || s.isBlank(); - }).orElse(false); - boolean vulnsAbsent = state.vulnerabilities().isEmpty(); - if (depsEmpty && vulnsAbsent) { - log.info("[Supervisor] Dependencies are empty and vulnerabilities not yet set — routing to DONE (no LLM call)"); - return Map.of("nextAgent", ROUTE_DONE); + if (INTENT_REPO_SCAN.equals(intentNow) || INTENT_FILE_SCAN.equals(intentNow)) { + boolean depsEmpty = state.dependencies().map(d -> AgentUtils.isNullOrEmpty(d.toString())).orElse(false); + if (depsEmpty && state.vulnerabilities().isEmpty()) { + log.info("[Supervisor] Dependencies empty — routing to DONE (no LLM call)"); + return Map.of("nextAgent", ROUTE_DONE, STATE_KEY_ITERATIONS, iterations + 1); } } - String depsStatus = state.dependencies().map(d -> "POPULATED").orElse("null"); - String vulnsStatus = state.vulnerabilities().map(v -> "POPULATED").orElse("null"); - String detailStatus = state.vulnerabilityDetails().map(v -> "POPULATED").orElse("null"); - String fileContents = state.dependencyFileContents().map(v -> "POPULATED").orElse("null"); - String depFiles = state.dependencyFiles().map(v -> "POPULATED").orElse("null"); + // Two-mode prompt + boolean isInitialCall = state.intent().isEmpty(); + String prompt = isInitialCall ? buildInitialPrompt(state) : buildRoutingPrompt(state); + + String raw = Objects.requireNonNull( + chatClient.prompt().messages(new UserMessage(prompt)).call().content() + ).trim(); + + log.info("[Supervisor] LLM response: {}", raw); + Map updates = parseResponse(raw); + updates.put(STATE_KEY_ITERATIONS, iterations + 1); + return updates; + } - String prompt = PromptLibrary.SUPERVISOR_AGENT_PROMPT.formatted( + private String buildInitialPrompt(SecurityAgentState state) { + return PromptLibrary.SUPERVISOR_AGENT_PROMPT.formatted( state.userQuery().orElse("null"), state.intent().orElse("null"), state.repositoryUrl().orElse("null"), @@ -162,42 +177,32 @@ private Map reason(SecurityAgentState state) { state.packageEcosystem().orElse("null"), state.cveId().orElse("null"), state.fileType().orElse("null"), - depFiles, - fileContents, - depsStatus, - vulnsStatus, - detailStatus + state.dependencyFiles().map(v -> "POPULATED").orElse("null"), + state.dependencyFileContents().map(v -> "POPULATED").orElse("null"), + state.dependencies().map(v -> "POPULATED").orElse("null"), + state.vulnerabilities().map(v -> "POPULATED").orElse("null"), + state.vulnerabilityDetails().map(v -> "POPULATED").orElse("null") ); - - String raw = Objects.requireNonNull( - chatClient.prompt() - .messages(new UserMessage(prompt)) - .call() - .content() - ).trim(); - - log.info("[Supervisor] LLM response: {}", raw); - - return parseResponse(raw); } - // ── response parser ─────────────────────────────────────────────────────── + private String buildRoutingPrompt(SecurityAgentState state) { + return PromptLibrary.SUPERVISOR_ROUTING_PROMPT.formatted( + state.intent().orElse("null"), + state.dependencyFiles().map(v -> "POPULATED").orElse("null"), + state.dependencyFileContents().map(v -> "POPULATED").orElse("null"), + state.dependencies().map(v -> "POPULATED").orElse("null"), + state.vulnerabilities().map(v -> "POPULATED").orElse("null"), + state.vulnerabilityDetails().map(v -> "POPULATED").orElse("null") + ); + } @SuppressWarnings("unchecked") private Map parseResponse(String raw) { Map updates = new HashMap<>(); - try { - String json = raw.replaceAll("(?s)```[a-z]*\\s*", "").replaceAll("```", "").trim(); - + String json = AgentUtils.stripMarkdown(raw); Map parsed = objectMapper.readValue(json, Map.class); - - // Always extract nextAgent - if (parsed.containsKey("nextAgent")) { - updates.put("nextAgent", parsed.get("nextAgent")); - } - - // First-call fields — only set if non-null + if (parsed.containsKey("nextAgent")) updates.put("nextAgent", parsed.get("nextAgent")); setIfPresent(parsed, updates, "intent"); setIfPresent(parsed, updates, "repositoryUrl"); setIfPresent(parsed, updates, "packageName"); @@ -206,44 +211,39 @@ private Map parseResponse(String raw) { setIfPresent(parsed, updates, "cveId"); setIfPresent(parsed, updates, "rawFileContent"); setIfPresent(parsed, updates, "fileType"); - log.info("[Supervisor] Parsed updates: {}", updates); - } catch (Exception e) { log.error("[Supervisor] Failed to parse LLM response: {}", raw, e); updates.put("nextAgent", ROUTE_DIRECT); } - updates.putIfAbsent("nextAgent", ROUTE_DIRECT); return updates; } private void setIfPresent(Map source, Map target, String key) { Object value = source.get(key); - if (value != null && !"null".equals(value.toString().trim())) { + if (value != null && !AgentUtils.isNullOrEmpty(value.toString())) { target.put(key, value.toString().trim()); } } - // ── sub-agent dispatch ──────────────────────────────────────────────────── - private Map runRepoAgent(SecurityAgentState state) { - log.info("[Supervisor] → GithubRepoAgent"); + log.info("[Supervisor] -> GithubRepoAgent"); return repoAgent.invoke(state.data()); } private Map runParserAgent(SecurityAgentState state) { - log.info("[Supervisor] → DependencyParserAgent"); + log.info("[Supervisor] -> DependencyParserAgent"); return parserAgent.invoke(state.data()); } private Map runVulnerabilityAgent(SecurityAgentState state) { - log.info("[Supervisor] → VulnerabilityAgent"); + log.info("[Supervisor] -> VulnerabilityAgent"); return vulnerabilityAgent.invoke(state.data()); } private Map runDirectAnswerAgent(SecurityAgentState state) { - log.info("[Supervisor] → DirectAnswerAgent"); + log.info("[Supervisor] -> DirectAnswerAgent"); return directAnswerAgent.invoke(state.data()); } -} \ No newline at end of file +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java index f543927..ccf7ec8 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java @@ -1,10 +1,15 @@ package com.agent.agent.sub_agents; +import com.agent.exception.MaxLlmCallsExceededException; +import com.agent.exception.McpToolExecutionException; +import com.agent.exception.McpToolNullResponseException; import com.agent.model.ScanProgressEvent; import com.agent.prompt.PromptLibrary; import com.agent.service.ScanEventPublisher; import com.agent.state.SecurityAgentState; import com.agent.tool.McpToolRegistry; +import com.agent.util.AgentUtils; +import com.agent.util.McpResponseValidator; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.extern.slf4j.Slf4j; import org.bsc.langgraph4j.CompiledGraph; @@ -23,20 +28,23 @@ import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async; /** - * DependencyParserAgent — pure LLM-driven dependency file parsing. - *

- * The LLM sees all available parser tools and decides which one to call - * for the next unparsed file, along with the exact args. - * Supports monorepos: iterates until all files in dependencyFileContents - * have been parsed and merged into a single deduplicated dependencies list. - *

- * Graph: START → reason → [tool → reason]* → END + * DependencyParserAgent — LLM-driven dependency file parsing sub-agent. + * + * Changes from original: + * 1. LLM call counter (MAX_LLM_CALLS = 8) with MaxLlmCallsExceededException. + * 2. McpResponseValidator.validateAndExtract() in tool() — note: empty result + * (no deps found) is valid here, so we catch McpToolNullResponseException + * and treat it as an empty dep list rather than re-throwing. + * 3. Uses AgentUtils.formatToolList() and AgentUtils.stripMarkdown() shared utils. + * 4. MaxLlmCallsExceededException re-throws to SupervisorAgent. */ @Slf4j @Component public class DependencyParserAgent { - private static final String DONE = "DONE"; + private static final String DONE = "DONE"; + private static final int MAX_LLM_CALLS = 8; + private static final String STATE_KEY_CALL_COUNT = "llmCallCount_DependencyParserAgent"; private final ChatClient chatClient; private final McpToolRegistry toolRegistry; @@ -58,7 +66,7 @@ private synchronized void build() throws Exception { if (graph != null) return; graph = new StateGraph<>(SecurityAgentState.SCHEMA, SecurityAgentState::new) .addNode("reason", node_async(this::reason)) - .addNode("tool", node_async(this::tool)) + .addNode("tool", node_async(this::tool)) .addEdge(START, "reason") .addConditionalEdges( "reason", @@ -75,40 +83,39 @@ public Map invoke(Map incomingState) { return graph.invoke(incomingState) .map(SecurityAgentState::data) .orElse(incomingState); + + } catch (MaxLlmCallsExceededException e) { + throw e; + } catch (Exception e) { - log.error("[DependencyParserAgent] Error", e); + log.error("[DependencyParserAgent] Unexpected error", e); Map err = new HashMap<>(incomingState); err.put("finalAnswer", "DependencyParserAgent error: " + e.getMessage()); return err; } } - // ── reason node ─────────────────────────────────────────────────────────── - // - // The LLM sees: all available tools, all files needing parsing, and - // which files have already been parsed. - // It responds with: { "tool": "...", "args": { "content": "" }, "filePath": "..." } - // filePath is a hint so we know which file this parse result belongs to. - @SuppressWarnings("unchecked") private Map reason(SecurityAgentState state) { log.info("[DependencyParserAgent] Reasoning"); - List availableTools = toolRegistry.getToolNames(); + int callCount = (int) state.data().getOrDefault(STATE_KEY_CALL_COUNT, 0); + if (callCount >= MAX_LLM_CALLS) { + log.error("[DependencyParserAgent] Max LLM calls ({}) exceeded", MAX_LLM_CALLS); + throw new MaxLlmCallsExceededException("DependencyParserAgent", MAX_LLM_CALLS); + } - // Build a summary of what needs parsing: filePath → first 200 chars of content + List availableTools = toolRegistry.getToolNames(); Map fileMap = extractFileMap(state); Set alreadyParsed = extractParsedFilePaths(state); List remaining = fileMap.keySet().stream() - .filter(p -> !alreadyParsed.contains(p)) - .toList(); + .filter(p -> !alreadyParsed.contains(p)).toList(); if (remaining.isEmpty()) { log.info("[DependencyParserAgent] All files parsed — DONE"); - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); } - // Build a file overview for the LLM (path + content preview) StringBuilder fileOverview = new StringBuilder(); remaining.forEach(path -> { String preview = fileMap.get(path); @@ -117,67 +124,52 @@ private Map reason(SecurityAgentState state) { .append("Content preview:\n").append(preview).append("\n---\n"); }); - // The LLM only decides: which tool + which file path. - // File content is NEVER sent to the LLM — tool node reads it directly from state. - // This eliminates huge prompts and LLM latency on large files. String prompt = PromptLibrary.DEPENDENCY_PARSER_AGENT_PROMPT.formatted( - formatToolList(availableTools), + AgentUtils.formatToolList(availableTools), remaining.size(), fileOverview ); String raw = Objects.requireNonNull( - chatClient.prompt() - .messages(new UserMessage(prompt)) - .call() - .content() + chatClient.prompt().messages(new UserMessage(prompt)).call().content() ).trim(); log.info("[DependencyParserAgent] LLM decision: {}", raw); try { - String json = raw.replaceAll("(?s)```[a-z]*\\s*", "").replaceAll("```", "").trim(); - + String json = AgentUtils.stripMarkdown(raw); Map decision = objectMapper.readValue(json, Map.class); - String toolName = String.valueOf(decision.getOrDefault("tool", DONE)).trim(); String filePath = String.valueOf(decision.getOrDefault("filePath", "")).trim(); if (DONE.equals(toolName)) { - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); } - // Validate filePath — fall back to first remaining if LLM returned a wrong path if (!fileMap.containsKey(filePath)) { String fallback = remaining.get(0); - log.warn("[DependencyParserAgent] LLM returned unknown filePath '{}' — using '{}'", - filePath, fallback); + log.warn("[DependencyParserAgent] LLM returned unknown filePath '{}' — using '{}'", filePath, fallback); filePath = fallback; } - // args.content is intentionally empty here — tool node injects full content from state - log.info("[DependencyParserAgent] → tool={} filePath={}", toolName, filePath); - + log.info("[DependencyParserAgent] -> tool={} filePath={}", toolName, filePath); return Map.of( - "nextTool", toolName, - "nextToolArgs", "{}", // tool node fills real content - "currentParsingFile", filePath + "nextTool", toolName, + "nextToolArgs", "{}", + "currentParsingFile", filePath, + STATE_KEY_CALL_COUNT, callCount + 1 ); - } catch (Exception e) { log.error("[DependencyParserAgent] Failed to parse LLM response '{}': {}", raw, e.getMessage()); - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); } } - // ── tool node ───────────────────────────────────────────────────────────── - @SuppressWarnings("unchecked") private Map tool(SecurityAgentState state) { String toolName = state.nextTool().orElse(""); String filePath = state.data().getOrDefault("currentParsingFile", "unknown").toString(); - String scanId = state.scanId().orElse(null); - + String scanId = state.scanId().orElse(null); log.info("[DependencyParserAgent] Executing tool={} for file={}", toolName, filePath); if (!toolRegistry.isKnownTool(toolName)) { @@ -185,13 +177,13 @@ private Map tool(SecurityAgentState state) { return markFileParsed(state, filePath, List.of()); } - // Inject full file content here — never via the LLM prompt Map fileMap = extractFileMap(state); String fullContent = fileMap.getOrDefault(filePath, ""); if (fullContent.isBlank()) { - log.warn("[DependencyParserAgent] No content found for filePath='{}' — skipping", filePath); + log.warn("[DependencyParserAgent] No content for filePath='{}' — skipping", filePath); return markFileParsed(state, filePath, List.of()); } + String jsonArgs; try { jsonArgs = objectMapper.writeValueAsString(Map.of("content", fullContent)); @@ -200,49 +192,46 @@ private Map tool(SecurityAgentState state) { return markFileParsed(state, filePath, List.of()); } - // Emit SSE progress (reuse fileMap already extracted above) if (scanId != null) { Set parsed = extractParsedFilePaths(state); - int total = fileMap.size(); + int total = fileMap.size(); int current = parsed.size() + 1; - eventPublisher.publish(scanId, - ScanProgressEvent.parsingFile(scanId, filePath, current, total)); + eventPublisher.publish(scanId, ScanProgressEvent.parsingFile(scanId, filePath, current, total)); } ToolCallback tool = toolRegistry.getTool(toolName); - try { Object result = tool.call(jsonArgs); - String extracted = extractDependencies(result); + // For parsers, an empty result (no deps found) is valid — catch and treat as empty list + String extracted; + try { + extracted = McpResponseValidator.validateAndExtract(result, toolName); + } catch (McpToolNullResponseException e) { + log.info("[DependencyParserAgent] Tool '{}' returned empty for {} — treating as 0 deps", toolName, filePath); + return markFileParsed(state, filePath, List.of()); + } List> fileDeps = new ArrayList<>(); - if (extracted != null && !extracted.isBlank() && !extracted.equals("[]")) { + if (!AgentUtils.isNullOrEmpty(extracted)) { fileDeps = objectMapper.readValue( extracted, objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) ); } - - log.info("[DependencyParserAgent] {} → {} dependencies from {}", toolName, fileDeps.size(), filePath); + log.info("[DependencyParserAgent] {} -> {} dependencies from {}", toolName, fileDeps.size(), filePath); return markFileParsed(state, filePath, fileDeps); + } catch (McpToolNullResponseException e) { + return markFileParsed(state, filePath, List.of()); } catch (Exception e) { log.error("[DependencyParserAgent] Tool '{}' failed for {}: {}", toolName, filePath, e.getMessage(), e); - return markFileParsed(state, filePath, List.of()); // skip failed file, don't stop pipeline + return markFileParsed(state, filePath, List.of()); } } - // ── state helpers ───────────────────────────────────────────────────────── - - /** - * Merges newly parsed dependencies into the accumulated list, - * deduplicates, and marks the file as parsed. - */ @SuppressWarnings("unchecked") - private Map markFileParsed(SecurityAgentState state, - String filePath, + private Map markFileParsed(SecurityAgentState state, String filePath, List> newDeps) { - // Load existing accumulated deps Map accumulated = new LinkedHashMap<>(); state.dependencies().ifPresent(existing -> { try { @@ -251,28 +240,22 @@ private Map markFileParsed(SecurityAgentState state, objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) ); list.forEach(d -> accumulated.put(buildDepKey(d), d)); - } catch (Exception ignored) { - } + } catch (Exception ignored) {} }); - - // Merge new deps (dedup by name:version:ecosystem) newDeps.forEach(d -> accumulated.putIfAbsent(buildDepKey(d), d)); - // Track which files have been parsed Set parsedFiles = new LinkedHashSet<>(extractParsedFilePaths(state)); parsedFiles.add(filePath); Map updates = new HashMap<>(); try { updates.put("dependencies", objectMapper.writeValueAsString(new ArrayList<>(accumulated.values()))); - updates.put("parsedFiles", objectMapper.writeValueAsString(new ArrayList<>(parsedFiles))); + updates.put("parsedFiles", objectMapper.writeValueAsString(new ArrayList<>(parsedFiles))); } catch (Exception e) { log.error("[DependencyParserAgent] Failed to serialize state update", e); } - - log.info("[DependencyParserAgent] Accumulated {} unique deps, {} files parsed so far", + log.info("[DependencyParserAgent] Accumulated {} unique deps, {} files parsed", accumulated.size(), parsedFiles.size()); - return updates; } @@ -281,89 +264,33 @@ private Set extractParsedFilePaths(SecurityAgentState state) { Object raw = state.data().get("parsedFiles"); if (raw == null) return new LinkedHashSet<>(); try { - List list = objectMapper.readValue( + return new LinkedHashSet<>(objectMapper.readValue( raw.toString(), objectMapper.getTypeFactory().constructCollectionType(List.class, String.class) - ); - return new LinkedHashSet<>(list); - } catch (Exception e) { - return new LinkedHashSet<>(); - } + )); + } catch (Exception e) { return new LinkedHashSet<>(); } } @SuppressWarnings("unchecked") private Map extractFileMap(SecurityAgentState state) { - // REPO_SCAN / FILE_SCAN: dependencyFileContents is a JSON map { "path": "content" } Optional fileContentsOpt = state.dependencyFileContents().map(Object::toString); if (fileContentsOpt.isPresent()) { String raw = fileContentsOpt.get().trim(); if (raw.startsWith("{")) { - try { - return objectMapper.readValue(raw, Map.class); - } catch (Exception e) { - log.warn("[DependencyParserAgent] Could not parse dependencyFileContents as map"); - return Map.of("dependency-file", raw); - } + try { return objectMapper.readValue(raw, Map.class); } + catch (Exception e) { return Map.of("dependency-file", raw); } } if (!raw.isBlank()) return Map.of("dependency-file", raw); } - - // FILE_SCAN (direct upload): rawFileContent String raw = state.rawFileContent().orElse(""); - if (!raw.isBlank()) { - return Map.of(state.fileType().orElse("dependency-file"), raw); - } - + if (!raw.isBlank()) return Map.of(state.fileType().orElse("dependency-file"), raw); return Map.of(); } private String buildDepKey(Map dep) { - String name = String.valueOf(dep.getOrDefault("name", dep.getOrDefault("artifactId", "unknown"))); - String version = String.valueOf(dep.getOrDefault("version", "unknown")); + String name = String.valueOf(dep.getOrDefault("name", dep.getOrDefault("artifactId", "unknown"))); + String version = String.valueOf(dep.getOrDefault("version", "unknown")); String ecosystem = String.valueOf(dep.getOrDefault("ecosystem", dep.getOrDefault("type", "unknown"))); return (name + ":" + version + ":" + ecosystem).toLowerCase(); } - - // ── helpers ─────────────────────────────────────────────────────────────── - - private String formatToolList(List tools) { - var sb = new StringBuilder(); - tools.forEach(t -> sb.append("- ").append(t).append("\n")); - return sb.toString().stripTrailing(); - } - - /** - * Unwraps MCP envelope: [{"text":">"}] → JSON - */ - private String extractDependencies(Object result) { - try { - String raw = result.toString().trim(); - - if (raw.startsWith("[")) { - List> envelope = objectMapper.readValue( - raw, - objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) - ); - if (!envelope.isEmpty() && envelope.get(0).containsKey("text")) { - String text = envelope.get(0).get("text").toString(); - if (text.startsWith("\"") && text.endsWith("\"")) { - text = text.substring(1, text.length() - 1) - .replace("\\\"", "\"") - .replace("\\n", "\n"); - } - return text.trim(); - } - return objectMapper.writeValueAsString(envelope); - } - - if (raw.startsWith("\"") && raw.endsWith("\"")) { - raw = raw.substring(1, raw.length() - 1).replace("\\\"", "\""); - } - return raw.trim(); - - } catch (Exception e) { - log.warn("[DependencyParserAgent] Failed to extract dependencies: {}", result); - return result.toString(); - } - } -} \ No newline at end of file +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DirectAnswerAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DirectAnswerAgent.java index 2dbe502..d138247 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DirectAnswerAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DirectAnswerAgent.java @@ -1,5 +1,6 @@ package com.agent.agent.sub_agents; +import com.agent.model.ConversationContext; import com.agent.prompt.PromptLibrary; import com.agent.state.SecurityAgentState; import lombok.RequiredArgsConstructor; @@ -7,10 +8,12 @@ import org.bsc.langgraph4j.CompiledGraph; import org.bsc.langgraph4j.StateGraph; import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.stereotype.Component; import java.util.HashMap; +import java.util.List; import java.util.Map; import static org.bsc.langgraph4j.StateGraph.END; @@ -20,9 +23,12 @@ /** * DirectAnswerAgent — fallback for unclear or unsupported requests. * - * Explains what the system can do and asks the user to rephrase. - * - * Graph: START → reason → END + * Changes from original: + * - ConversationContext support: if history is present in state, previous turns + * are prepended to the LLM call via ConversationContext.buildMessages(). + * This gives the model context from earlier in the session when explaining + * what the system can do (e.g. "you just scanned lodash — what else can I scan?") + * - Falls back to single-message call when no history is present (original behaviour). */ @Slf4j @Component @@ -59,14 +65,20 @@ private Map reason(SecurityAgentState state) { log.info("[DirectAnswerAgent] Generating clarification"); String userQuery = state.userQuery().orElse("(no input)"); + String prompt = PromptLibrary.DIRECT_ANSWER_AGENT_PROMPT.formatted(userQuery); + + // Use conversation history if available — enables context-aware responses + ConversationContext ctx = (ConversationContext) state.data().get("conversationContext"); - String prompt = PromptLibrary.DIRECT_ANSWER_AGENT_PROMPT.formatted(userQuery); + List messages = (ctx != null && ctx.hasHistory()) + ? ctx.buildMessages(prompt) + : List.of(new UserMessage(prompt)); String response = chatClient.prompt() - .messages(new UserMessage(prompt)) + .messages(messages) .call().content(); log.info("[DirectAnswerAgent] Response generated"); return Map.of("finalAnswer", response); } -} \ No newline at end of file +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/GithubRepoAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/GithubRepoAgent.java index b1ead5c..36b4970 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/GithubRepoAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/GithubRepoAgent.java @@ -1,11 +1,16 @@ package com.agent.agent.sub_agents; +import com.agent.exception.MaxLlmCallsExceededException; +import com.agent.exception.McpToolExecutionException; +import com.agent.exception.McpToolNullResponseException; import com.agent.model.ScanProgressEvent; import com.agent.prompt.PromptLibrary; import com.agent.service.ScanEventPublisher; import com.agent.service.UserConfirmationService; import com.agent.state.SecurityAgentState; import com.agent.tool.McpToolRegistry; +import com.agent.util.AgentUtils; +import com.agent.util.McpResponseValidator; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.extern.slf4j.Slf4j; import org.bsc.langgraph4j.CompiledGraph; @@ -24,20 +29,22 @@ import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async; /** - * GithubRepoAgent — pure LLM-driven repository scanning. - *

- * The LLM decides BOTH which tool to call AND its exact arguments. - * Java is a thin executor. The one non-LLM concern kept here is the - * user-confirmation pause/resume flow, which is inherently stateful - * infrastructure — not business logic. - *

- * Graph: START → reason → [tool → reason]* → END + * GithubRepoAgent — LLM-driven repository scanning sub-agent. + * + * Changes from original: + * 1. LLM call counter (MAX_LLM_CALLS = 8) with MaxLlmCallsExceededException. + * 2. McpResponseValidator.validateAndExtract() in tool() — throws on null MCP result. + * 3. Uses AgentUtils.formatToolList() and AgentUtils.stripMarkdown() shared utils. + * 4. Typed exception re-throw in invoke(): McpToolNullResponseException and + * MaxLlmCallsExceededException propagate to SupervisorAgent. */ @Slf4j @Component public class GithubRepoAgent { - private static final String DONE = "DONE"; + private static final String DONE = "DONE"; + private static final int MAX_LLM_CALLS = 8; + private static final String STATE_KEY_CALL_COUNT = "llmCallCount_GithubRepoAgent"; private final ChatClient chatClient; private final McpToolRegistry toolRegistry; @@ -58,13 +65,11 @@ public GithubRepoAgent( this.eventPublisher = eventPublisher; } - // ── Graph ───────────────────────────────────────────────────────────────── - private synchronized void build() throws Exception { if (graph != null) return; graph = new StateGraph<>(SecurityAgentState.SCHEMA, SecurityAgentState::new) .addNode("reason", node_async(this::reason)) - .addNode("tool", node_async(this::tool)) + .addNode("tool", node_async(this::tool)) .addEdge(START, "reason") .addConditionalEdges( "reason", @@ -81,8 +86,12 @@ public Map invoke(Map incomingState) { return graph.invoke(incomingState) .map(SecurityAgentState::data) .orElse(incomingState); + + } catch (McpToolNullResponseException | MaxLlmCallsExceededException e) { + throw e; + } catch (Exception e) { - log.error("[GithubRepoAgent] Error", e); + log.error("[GithubRepoAgent] Unexpected error", e); Map err = new HashMap<>(incomingState); err.put("finalAnswer", "GithubRepoAgent error: " + e.getMessage()); return err; @@ -92,11 +101,17 @@ public Map invoke(Map incomingState) { private Map reason(SecurityAgentState state) { log.info("[GithubRepoAgent] Reasoning"); + int callCount = (int) state.data().getOrDefault(STATE_KEY_CALL_COUNT, 0); + if (callCount >= MAX_LLM_CALLS) { + log.error("[GithubRepoAgent] Max LLM calls ({}) exceeded", MAX_LLM_CALLS); + throw new MaxLlmCallsExceededException("GithubRepoAgent", MAX_LLM_CALLS); + } + List availableTools = toolRegistry.getToolNames(); log.info("[GithubRepoAgent] Tools: {}", availableTools); String prompt = PromptLibrary.GITHUB_REPO_AGENT_PROMPT.formatted( - formatToolList(availableTools), + AgentUtils.formatToolList(availableTools), state.repositoryUrl().orElse("null"), state.data().getOrDefault("repositoryCloned", "null"), state.data().getOrDefault("repoPath", "null"), @@ -105,186 +120,122 @@ private Map reason(SecurityAgentState state) { ); String raw = Objects.requireNonNull( - chatClient.prompt() - .messages(new UserMessage(prompt)) - .call() - .content() + chatClient.prompt().messages(new UserMessage(prompt)).call().content() ).trim(); log.info("[GithubRepoAgent] LLM decision: {}", raw); try { - String json = raw.replaceAll("(?s)```[a-z]*\\s*", "").replaceAll("```", "").trim(); - + String json = AgentUtils.stripMarkdown(raw); @SuppressWarnings("unchecked") Map decision = objectMapper.readValue(json, Map.class); - String toolName = String.valueOf(decision.getOrDefault("tool", DONE)).trim(); - @SuppressWarnings("unchecked") Map args = decision.containsKey("args") - ? (Map) decision.get("args") - : Map.of(); - - log.info("[GithubRepoAgent] → tool={} args={}", toolName, args); - + ? (Map) decision.get("args") : Map.of(); + log.info("[GithubRepoAgent] -> tool={} args={}", toolName, args); return Map.of( - "nextTool", toolName, - "nextToolArgs", objectMapper.writeValueAsString(args) + "nextTool", toolName, + "nextToolArgs", objectMapper.writeValueAsString(args), + STATE_KEY_CALL_COUNT, callCount + 1 ); - } catch (Exception e) { log.error("[GithubRepoAgent] Failed to parse LLM response '{}': {}", raw, e.getMessage()); - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); } } private Map tool(SecurityAgentState state) { String toolName = state.nextTool().orElse(""); String jsonArgs = state.data().getOrDefault("nextToolArgs", "{}").toString(); - log.info("[GithubRepoAgent] Executing tool={} args={}", toolName, jsonArgs); if (!toolRegistry.isKnownTool(toolName)) { log.warn("[GithubRepoAgent] Unknown tool '{}'. Known: {}", toolName, toolRegistry.getToolNames()); - return Map.of( - "finalAnswer", "Invalid tool for repo-server: " + toolName, - "nextTool", DONE - ); + return Map.of("finalAnswer", "Invalid tool for repo-server: " + toolName, "nextTool", DONE); } ToolCallback tool = toolRegistry.getTool(toolName); - try { Object result = tool.call(jsonArgs); - String extracted = extractText(result); - + // Throws McpToolNullResponseException if null/empty + String extracted = McpResponseValidator.validateAndExtract(result, toolName); return switch (toolName) { - case "cloneRepository" -> onCloneRepository(state, extracted, jsonArgs); + case "cloneRepository" -> onCloneRepository(state, extracted, jsonArgs); case "findDependencyFiles" -> onFindDependencyFiles(state, extracted); - case "readFileContent" -> onReadFileContent(state, extracted, jsonArgs); + case "readFileContent" -> onReadFileContent(state, extracted, jsonArgs); default -> { log.warn("[GithubRepoAgent] No post-processor for '{}' — storing raw result", toolName); yield Map.of("toolResult", extracted); } }; - + } catch (McpToolNullResponseException e) { + throw e; } catch (Exception e) { log.error("[GithubRepoAgent] Tool '{}' failed", toolName, e); - return Map.of( - "finalAnswer", "Tool execution failed: " + e.getMessage(), - "nextTool", DONE - ); + throw new McpToolExecutionException(toolName, e); } } - // ── post-processors ─────────────────────────────────────────────────────── - // - // These handle side-effects (SSE events, blocking for user input) that occur - // AFTER a tool succeeds. The LLM is not involved in these — they are pure - // infrastructure glue between the MCP result and the pipeline state. - - private Map onCloneRepository(SecurityAgentState state, - String repoPath, - String jsonArgs) throws Exception { + private Map onCloneRepository(SecurityAgentState state, String repoPath, String jsonArgs) throws Exception { @SuppressWarnings("unchecked") String repoUrl = (String) objectMapper.readValue(jsonArgs, Map.class) .getOrDefault("repoUrl", state.repositoryUrl().orElse("unknown")); - - log.info("[GithubRepoAgent] Cloned → {}", repoPath); - - state.scanId().ifPresent(id -> - eventPublisher.publish(id, ScanProgressEvent.repoCloned(id, repoUrl))); - - return Map.of( - "repositoryCloned", true, - "repoPath", repoPath - ); + log.info("[GithubRepoAgent] Cloned -> {}", repoPath); + state.scanId().ifPresent(id -> eventPublisher.publish(id, ScanProgressEvent.repoCloned(id, repoUrl))); + return Map.of("repositoryCloned", true, "repoPath", repoPath); } - private Map onFindDependencyFiles(SecurityAgentState state, - String dependencyFilesJson) throws Exception { + private Map onFindDependencyFiles(SecurityAgentState state, String dependencyFilesJson) throws Exception { List allFiles = objectMapper.readValue( dependencyFilesJson, objectMapper.getTypeFactory().constructCollectionType(List.class, String.class) ); - String scanId = state.scanId().orElse(null); log.info("[GithubRepoAgent] Found {} dependency files", allFiles.size()); + if (scanId != null) eventPublisher.publish(scanId, ScanProgressEvent.filesDiscovered(scanId, allFiles)); - if (scanId != null) { - eventPublisher.publish(scanId, ScanProgressEvent.filesDiscovered(scanId, allFiles)); - } - - // No dependency files found — store empty list and signal DONE so the - // reason node stops immediately instead of hallucinating a file path. if (allFiles.isEmpty()) { - log.warn("[GithubRepoAgent] No dependency files found in repo — stopping agent"); - return Map.of( - "dependencyFiles", dependencyFilesJson, // stores "[]" - "dependencyFileContents", "{}", // marks contents as done - "nextTool", DONE - ); + log.warn("[GithubRepoAgent] No dependency files found — stopping agent"); + return Map.of("dependencyFiles", dependencyFilesJson, "dependencyFileContents", "{}", "nextTool", DONE); } if (allFiles.size() > 1) { - // Multiple files — pause and wait for user to select which ones to scan log.info("[GithubRepoAgent] Multiple files — pausing for user confirmation"); - - if (scanId != null) { - eventPublisher.publish(scanId, ScanProgressEvent.awaitingConfirmation(scanId, allFiles)); - } - + if (scanId != null) eventPublisher.publish(scanId, ScanProgressEvent.awaitingConfirmation(scanId, allFiles)); String selectedJson = confirmationService.waitForConfirmation(scanId, allFiles); - List selected = objectMapper.readValue( selectedJson, objectMapper.getTypeFactory().constructCollectionType(List.class, String.class) ); - if (selected.isEmpty()) { log.info("[GithubRepoAgent] Empty selection — defaulting to all {} files", allFiles.size()); selected = allFiles; } - log.info("[GithubRepoAgent] User selected {} file(s)", selected.size()); - - if (scanId != null) { - eventPublisher.publish(scanId, ScanProgressEvent.filesSelected(scanId, selected)); - } - + if (scanId != null) eventPublisher.publish(scanId, ScanProgressEvent.filesSelected(scanId, selected)); String finalJson = objectMapper.writeValueAsString(selected); return Map.of("dependencyFiles", finalJson, "selectedFiles", finalJson); } - // Single file — no confirmation needed return Map.of("dependencyFiles", dependencyFilesJson); } - private Map onReadFileContent(SecurityAgentState state, - String content, - String jsonArgs) throws Exception { + private Map onReadFileContent(SecurityAgentState state, String content, String jsonArgs) throws Exception { @SuppressWarnings("unchecked") - String filePath = (String) objectMapper.readValue(jsonArgs, Map.class) - .getOrDefault("filePath", "unknown"); - + String filePath = (String) objectMapper.readValue(jsonArgs, Map.class).getOrDefault("filePath", "unknown"); log.info("[GithubRepoAgent] Read file {} ({} chars)", filePath, content.length()); - // Merge into the existing file-contents map (supports multi-file reads loop) Map existing = new LinkedHashMap<>(); - state.dependencyFileContents().ifPresent(raw -> { try { @SuppressWarnings("unchecked") Map parsed = objectMapper.readValue(raw.toString(), Map.class); existing.putAll(parsed); - } catch (Exception ignored) { - } + } catch (Exception ignored) {} }); - existing.put(filePath, content); - // Check if all files have been read List allFiles = new ArrayList<>(); state.dependencyFiles().ifPresent(raw -> { try { @@ -292,62 +243,17 @@ private Map onReadFileContent(SecurityAgentState state, raw.toString(), objectMapper.getTypeFactory().constructCollectionType(List.class, String.class) )); - } catch (Exception ignored) { - } + } catch (Exception ignored) {} }); boolean allRead = !allFiles.isEmpty() && existing.keySet().containsAll(allFiles); String contentsJson = objectMapper.writeValueAsString(existing); - Map updates = new HashMap<>(); updates.put("dependencyFileContents", contentsJson); - if (allRead) { log.info("[GithubRepoAgent] All {} files read — signalling parser", allFiles.size()); - state.scanId().ifPresent(id -> - eventPublisher.publish(id, ScanProgressEvent.parsingDependencies(id))); + state.scanId().ifPresent(id -> eventPublisher.publish(id, ScanProgressEvent.parsingDependencies(id))); } - return updates; } - - // ── helpers ─────────────────────────────────────────────────────────────── - - private String formatToolList(List tools) { - var sb = new StringBuilder(); - tools.forEach(t -> sb.append("- ").append(t).append("\n")); - return sb.toString().stripTrailing(); - } - - /** - * Unwraps MCP envelope: [{"text":""}] → value - */ - private String extractText(Object result) { - try { - String raw = result.toString().trim(); - - if (raw.startsWith("[")) { - List> list = objectMapper.readValue( - raw, - objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) - ); - if (!list.isEmpty()) { - String text = list.get(0).get("text").toString(); - if (text.startsWith("\"") && text.endsWith("\"")) { - text = text.substring(1, text.length() - 1); - } - return text.replace("\\\"", "\"").trim(); - } - } - - if (raw.startsWith("\"") && raw.endsWith("\"")) { - raw = raw.substring(1, raw.length() - 1); - } - return raw.trim(); - - } catch (Exception e) { - log.warn("[GithubRepoAgent] Failed to parse MCP response: {}", result); - return result.toString(); - } - } -} \ No newline at end of file +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java index b735578..bd40b7a 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java @@ -1,8 +1,13 @@ package com.agent.agent.sub_agents; +import com.agent.exception.MaxLlmCallsExceededException; +import com.agent.exception.McpToolExecutionException; +import com.agent.exception.McpToolNullResponseException; import com.agent.prompt.PromptLibrary; import com.agent.state.SecurityAgentState; import com.agent.tool.McpToolRegistry; +import com.agent.util.AgentUtils; +import com.agent.util.McpResponseValidator; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.extern.slf4j.Slf4j; import org.bsc.langgraph4j.CompiledGraph; @@ -24,20 +29,31 @@ import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async; /** - * VulnerabilityAgent — pure LLM-driven vulnerability scanning. - *

- * The LLM decides BOTH which tool to call AND its exact arguments. - * No hardcoded if/else per tool, no batch processing — the agent is a - * thin executor: reason → tool → reason until the LLM says DONE. - *

- * Graph: START → reason → [tool → reason]* → END + * VulnerabilityAgent — LLM-driven vulnerability scanning sub-agent. + * + * Changes from original: + * 1. LLM call counter (MAX_LLM_CALLS = 6) — throws MaxLlmCallsExceededException + * when exceeded, propagates to SupervisorAgent for clean error handling. + * 2. Deterministic early-exit in reason(): skips LLM call entirely when the + * correct action is already obvious (e.g. vulns already populated → DONE). + * Saves 1-2 LLM calls per scan on common paths. + * 3. McpResponseValidator.validateAndExtract() in tool(): throws + * McpToolNullResponseException on null/empty results — stops the loop + * immediately instead of silently passing empty state back to reason(). + * 4. Uses AgentUtils.formatToolList(), AgentUtils.stripMarkdown() instead of + * inline duplicates. + * 5. Typed exception re-throw: McpToolNullResponseException and + * MaxLlmCallsExceededException bubble to SupervisorAgent; only unexpected + * exceptions are caught and converted to finalAnswer. */ @Slf4j @Component public class VulnerabilityAgent { - private static final String DONE = "DONE"; - private static final int MAX_RETRIES = 3; + private static final String DONE = "DONE"; + private static final int MAX_RETRIES = 3; + private static final int MAX_LLM_CALLS = 6; + private static final String STATE_KEY_CALL_COUNT = "llmCallCount_VulnerabilityAgent"; private final ObjectMapper objectMapper = new ObjectMapper(); private final ChatClient chatClient; @@ -52,13 +68,11 @@ public VulnerabilityAgent( this.toolRegistry = toolRegistry; } - // ── Graph ───────────────────────────────────────────────────────────────── - private synchronized void build() throws Exception { if (graph != null) return; graph = new StateGraph<>(SecurityAgentState.SCHEMA, SecurityAgentState::new) .addNode("reason", node_async(this::reason)) - .addNode("tool", node_async(this::tool)) + .addNode("tool", node_async(this::tool)) .addEdge(START, "reason") .addConditionalEdges( "reason", @@ -75,8 +89,13 @@ public Map invoke(Map incomingState) { return graph.invoke(incomingState) .map(SecurityAgentState::data) .orElse(incomingState); + + } catch (McpToolNullResponseException | MaxLlmCallsExceededException e) { + // Re-throw typed exceptions to SupervisorAgent for proper handling + throw e; + } catch (Exception e) { - log.error("[VulnerabilityAgent] Error", e); + log.error("[VulnerabilityAgent] Unexpected error", e); Map err = new HashMap<>(incomingState); err.put("finalAnswer", "VulnerabilityAgent error: " + e.getMessage()); return err; @@ -88,215 +107,143 @@ public Map invoke(Map incomingState) { private Map reason(SecurityAgentState state) { log.info("[VulnerabilityAgent] Reasoning — intent={}", state.intent().orElse("null")); - List availableTools = toolRegistry.getToolNames(); - log.info("[VulnerabilityAgent] Tools: {}", availableTools); + // LLM call counter guard + int callCount = (int) state.data().getOrDefault(STATE_KEY_CALL_COUNT, 0); + if (callCount >= MAX_LLM_CALLS) { + log.error("[VulnerabilityAgent] Max LLM calls ({}) exceeded", MAX_LLM_CALLS); + throw new MaxLlmCallsExceededException("VulnerabilityAgent", MAX_LLM_CALLS); + } + + // Deterministic early-exit: skip LLM on known-outcome patterns + String intent = state.intent().orElse(""); + boolean vulnsPopulated = state.vulnerabilities().isPresent(); + boolean detailsPopulated = state.vulnerabilityDetails().isPresent(); - // ── Early-exit guard: empty dependency list ─────────────────────────── - // If the parser ran but produced 0 dependencies (e.g. workspace manifest - // with no actual packages), skip the LLM entirely to avoid hallucination. - String intent = state.intent().orElse("null"); + if ("CVE_QUERY".equals(intent) && detailsPopulated) { + log.info("[VulnerabilityAgent] CVE details already populated — DONE (no LLM call)"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); + } + if ("PACKAGE_SCAN".equals(intent) && vulnsPopulated) { + log.info("[VulnerabilityAgent] Vulnerabilities already populated — DONE (no LLM call)"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); + } + + // Early-exit: empty dependency list for REPO/FILE scan if ("REPO_SCAN".equals(intent) || "FILE_SCAN".equals(intent)) { String depsRaw = state.dependencies().map(Object::toString).orElse(null); - if (depsRaw != null) { - String trimmed = depsRaw.trim(); - boolean isEmpty = trimmed.equals("[]") || trimmed.equals("null") || trimmed.isBlank(); - if (isEmpty) { - log.info("[VulnerabilityAgent] Dependencies list is empty — returning DONE immediately (no LLM call)"); - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); - } + if (depsRaw != null && AgentUtils.isNullOrEmpty(depsRaw.trim())) { + log.info("[VulnerabilityAgent] Empty dependency list — DONE (no LLM call)"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); } } - // Summarise large blobs — show count instead of just POPULATED so LLM - // knows the list is non-empty and doesn't hallucinate packages. + List availableTools = toolRegistry.getToolNames(); + log.info("[VulnerabilityAgent] Tools: {}", availableTools); + String depsStatus = state.dependencies().map(d -> { try { com.fasterxml.jackson.databind.JsonNode node = objectMapper.readTree(d.toString()); return node.isArray() ? "POPULATED (" + node.size() + " deps)" : "POPULATED"; - } catch (Exception e) { - return "POPULATED"; - } + } catch (Exception e) { return "POPULATED"; } }).orElse("null"); - String vulnsStatus = state.vulnerabilities().map(v -> "POPULATED").orElse("null"); - String detailStatus = state.vulnerabilityDetails().map(v -> "POPULATED").orElse("null"); String prompt = PromptLibrary.VULNERABILITY_AGENT_PROMPT.formatted( - formatToolList(availableTools), - state.intent().orElse("null"), + AgentUtils.formatToolList(availableTools), + intent, state.packageName().orElse("null"), state.packageVersion().orElse("null"), state.packageEcosystem().orElse("null"), state.cveId().orElse("null"), depsStatus, - vulnsStatus, - detailStatus + vulnsPopulated ? "POPULATED" : "null", + detailsPopulated ? "POPULATED" : "null" ); String raw = Objects.requireNonNull( - chatClient.prompt() - .messages(new UserMessage(prompt)) - .call() - .content() + chatClient.prompt().messages(new UserMessage(prompt)).call().content() ).trim(); log.info("[VulnerabilityAgent] LLM decision: {}", raw); try { - String json = raw.replaceAll("(?s)```[a-z]*\\s*", "").replaceAll("```", "").trim(); - + String json = AgentUtils.stripMarkdown(raw); @SuppressWarnings("unchecked") Map decision = objectMapper.readValue(json, Map.class); - String toolName = String.valueOf(decision.getOrDefault("tool", DONE)).trim(); - @SuppressWarnings("unchecked") Map args = decision.containsKey("args") - ? (Map) decision.get("args") - : Map.of(); - - log.info("[VulnerabilityAgent] → tool={} args={}", toolName, args); - + ? (Map) decision.get("args") : Map.of(); + log.info("[VulnerabilityAgent] -> tool={} args={}", toolName, args); return Map.of( - "nextTool", toolName, - "nextToolArgs", objectMapper.writeValueAsString(args) + "nextTool", toolName, + "nextToolArgs", objectMapper.writeValueAsString(args), + STATE_KEY_CALL_COUNT, callCount + 1 ); - } catch (Exception e) { log.error("[VulnerabilityAgent] Failed to parse LLM response '{}': {}", raw, e.getMessage()); - return Map.of("nextTool", DONE, "nextToolArgs", "{}"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); } } // ── tool node ───────────────────────────────────────────────────────────── - // - // Executes the tool+args the LLM decided. No argument construction here. - // Only mapping left is: which state key does this tool's output belong to. private Map tool(SecurityAgentState state) { String toolName = state.nextTool().orElse(""); String jsonArgs = state.data().getOrDefault("nextToolArgs", "{}").toString(); - log.info("[VulnerabilityAgent] Executing tool={} args={}", toolName, jsonArgs); if (!toolRegistry.isKnownTool(toolName)) { log.warn("[VulnerabilityAgent] Unknown tool '{}'. Known: {}", toolName, toolRegistry.getToolNames()); - return Map.of( - "finalAnswer", "Tool not found in vulnerability-server: " + toolName, - "nextTool", DONE - ); + return Map.of("finalAnswer", "Tool not found in vulnerability-server: " + toolName, "nextTool", DONE); } ToolCallback tool = toolRegistry.getTool(toolName); - try { Object result = callWithRetry(tool, jsonArgs, toolName); - String extracted = extractText(result); - log.info("[VulnerabilityAgent] {} completed", toolName); + // Throws McpToolNullResponseException if null/empty — stops loop immediately + String extracted = McpResponseValidator.validateAndExtract(result, toolName); + log.info("[VulnerabilityAgent] {} completed successfully", toolName); return Map.of(resolveStateKey(toolName), extracted); + } catch (McpToolNullResponseException e) { + // Re-throw: let SupervisorAgent handle this gracefully + throw e; } catch (Exception e) { log.error("[VulnerabilityAgent] Tool '{}' failed", toolName, e); - return Map.of( - "finalAnswer", "Tool execution failed: " + e.getMessage(), - "nextTool", DONE - ); + throw new McpToolExecutionException(toolName, e); } } - // Maps tool name → state key that receives its output. - // Extend this switch as new tools are added to the vulnerability-server. private String resolveStateKey(String toolName) { return switch (toolName) { - case "checkPackageVulnerability", - "checkMultipleVulnerabilities" -> "vulnerabilities"; + case "checkPackageVulnerability", "checkMultipleVulnerabilities" -> "vulnerabilities"; case "getVulnerabilityDetails" -> "vulnerabilityDetails"; default -> { - log.warn("[VulnerabilityAgent] No state key mapping for '{}' — using 'toolResult'", toolName); + log.warn("[VulnerabilityAgent] No state key for '{}' — using 'toolResult'", toolName); yield "toolResult"; } }; } - // ── retry ───────────────────────────────────────────────────────────────── - private Object callWithRetry(ToolCallback tool, String jsonArgs, String label) throws Exception { Exception last = null; - for (int attempt = 1; attempt <= MAX_RETRIES; attempt++) { try { if (attempt > 1) log.info("[VulnerabilityAgent] {} — retry {}/{}", label, attempt, MAX_RETRIES); return tool.call(jsonArgs); - } catch (Exception e) { last = e; String msg = e.getMessage() != null ? e.getMessage() : ""; - - boolean isDeterministic = msg.contains("Cannot invoke") - || msg.contains("NullPointerException") - || msg.contains("IllegalArgumentException") - || msg.contains("IllegalStateException"); - - boolean isTransient = msg.contains("TimeoutException") - || msg.contains("ConnectException") - || msg.contains("SocketException") - || msg.contains("IOException"); - - log.warn("[VulnerabilityAgent] {} — attempt {}/{} failed ({}): {}", - label, attempt, MAX_RETRIES, + boolean isDeterministic = msg.contains("Cannot invoke") || msg.contains("NullPointerException") + || msg.contains("IllegalArgumentException") || msg.contains("IllegalStateException"); + boolean isTransient = msg.contains("TimeoutException") || msg.contains("ConnectException") + || msg.contains("SocketException") || msg.contains("IOException"); + log.warn("[VulnerabilityAgent] {} attempt {}/{} failed ({}): {}", label, attempt, MAX_RETRIES, isTransient ? "TRANSIENT" : isDeterministic ? "DETERMINISTIC" : "ERROR", msg); - - if (isDeterministic) { - log.error("[VulnerabilityAgent] {} — deterministic error, not retrying.", label); - throw e; - } - - if (attempt < MAX_RETRIES) { - Thread.sleep(1000L * attempt); - } + if (isDeterministic) throw e; + if (attempt < MAX_RETRIES) Thread.sleep(1000L * attempt); } } - - log.error("[VulnerabilityAgent] {} — all {} attempts failed", label, MAX_RETRIES); throw last; } - - // ── helpers ─────────────────────────────────────────────────────────────── - - private String formatToolList(List tools) { - var sb = new StringBuilder(); - tools.forEach(t -> sb.append("- ").append(t).append("\n")); - return sb.toString().stripTrailing(); - } - - /** - * Unwraps MCP envelope: [{"text":""}] → value - */ - private String extractText(Object result) { - try { - String raw = result.toString().trim(); - - if (raw.startsWith("[")) { - List> list = objectMapper.readValue( - raw, - objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) - ); - if (!list.isEmpty()) { - String text = list.get(0).get("text").toString(); - if (text.startsWith("\"") && text.endsWith("\"")) { - text = text.substring(1, text.length() - 1); - } - return text.replace("\\\"", "\"").trim(); - } - } - - if (raw.startsWith("\"") && raw.endsWith("\"")) { - raw = raw.substring(1, raw.length() - 1); - } - return raw.trim(); - - } catch (Exception e) { - log.warn("[VulnerabilityAgent] Failed to parse MCP response: {}", result); - return result.toString(); - } - } -} \ No newline at end of file +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/controller/ScanController.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/controller/ScanController.java index 2284719..644a8fb 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/controller/ScanController.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/controller/ScanController.java @@ -5,7 +5,7 @@ import com.agent.model.ScanRequest; import com.agent.model.ScanResponse; import com.agent.service.SecuritySupervisorAgentService; -import com.agent.service.SecuritySupervisorAgentService.ScanException; +import com.agent.exception.ScanException; import com.agent.service.UserConfirmationService; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.RequiredArgsConstructor; diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentException.java new file mode 100644 index 0000000..7adb42b --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentException.java @@ -0,0 +1,17 @@ +package com.agent.exception; + +/** + * Base exception for all agent pipeline failures. + * All typed exceptions in this package extend this class so callers + * can catch the full hierarchy with a single catch(AgentException e). + */ +public class AgentException extends RuntimeException { + + public AgentException(String message) { + super(message); + } + + public AgentException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentRoutingException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentRoutingException.java new file mode 100644 index 0000000..4321076 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/AgentRoutingException.java @@ -0,0 +1,20 @@ +package com.agent.exception; + +/** + * Thrown when the SupervisorAgent cannot determine a valid route + * from the LLM response (e.g. LLM returned an unknown nextAgent value + * and the fallback to direct_answer_agent is not appropriate). + */ +public class AgentRoutingException extends AgentException { + + private final String receivedRoute; + + public AgentRoutingException(String receivedRoute) { + super("Supervisor could not route to a known agent. Received: '" + receivedRoute + "'"); + this.receivedRoute = receivedRoute; + } + + public String getReceivedRoute() { + return receivedRoute; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/MaxLlmCallsExceededException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/MaxLlmCallsExceededException.java new file mode 100644 index 0000000..4c12fa6 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/MaxLlmCallsExceededException.java @@ -0,0 +1,36 @@ +package com.agent.exception; + +/** + * Thrown when an agent's LLM call counter exceeds its configured limit. + * + * This is the hard circuit-breaker for infinite loops in the reason→tool→reason + * cycles. Each sub-agent and the SupervisorAgent track their own counters. + * When the limit is hit, this exception propagates to SupervisorAgent.invoke() + * which catches it and returns a graceful error response. + * + * Default limits (configurable per agent): + * SupervisorAgent → 6 iterations + * GithubRepoAgent → 8 LLM calls + * DependencyParserAgent → 8 LLM calls + * VulnerabilityAgent → 6 LLM calls + */ +public class MaxLlmCallsExceededException extends AgentException { + + private final String agentName; + private final int limit; + + public MaxLlmCallsExceededException(String agentName, int limit) { + super("Agent '" + agentName + "' exceeded maximum LLM calls (" + limit + "). " + + "Possible causes: unexpected state, LLM returning inconsistent decisions, or pipeline loop."); + this.agentName = agentName; + this.limit = limit; + } + + public String getAgentName() { + return agentName; + } + + public int getLimit() { + return limit; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolExecutionException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolExecutionException.java new file mode 100644 index 0000000..d8c65f7 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolExecutionException.java @@ -0,0 +1,27 @@ +package com.agent.exception; + +/** + * Thrown when an MCP tool call itself throws an exception (network error, + * invalid arguments, server-side failure, etc.). + * + * Wraps the original cause so callers can inspect the root failure while + * using a typed catch instead of a generic Exception catch. + */ +public class McpToolExecutionException extends AgentException { + + private final String toolName; + + public McpToolExecutionException(String toolName, String message) { + super("Tool '" + toolName + "' execution failed: " + message); + this.toolName = toolName; + } + + public McpToolExecutionException(String toolName, Throwable cause) { + super("Tool '" + toolName + "' execution failed: " + cause.getMessage(), cause); + this.toolName = toolName; + } + + public String getToolName() { + return toolName; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolNullResponseException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolNullResponseException.java new file mode 100644 index 0000000..94f03a7 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/McpToolNullResponseException.java @@ -0,0 +1,24 @@ +package com.agent.exception; + +/** + * Thrown when an MCP tool returns a null, blank, or empty response. + * + * This exception is the key circuit-breaker that prevents the LLM from + * being called again on empty state. When thrown from a tool() node, + * it propagates to the agent's invoke() method and then up to + * SupervisorAgent.invoke(), which catches it and returns a clean error + * response without wasting additional LLM calls. + */ +public class McpToolNullResponseException extends AgentException { + + private final String toolName; + + public McpToolNullResponseException(String toolName) { + super("MCP tool '" + toolName + "' returned null/empty — aborting reasoning loop"); + this.toolName = toolName; + } + + public String getToolName() { + return toolName; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/ScanException.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/ScanException.java new file mode 100644 index 0000000..96987eb --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/exception/ScanException.java @@ -0,0 +1,19 @@ +package com.agent.exception; + +/** + * Thrown when the agent pipeline cannot complete a scan. + * Caught by ScanController and returned as HTTP 422 UNPROCESSABLE_ENTITY. + * + * Previously a static inner class inside SecuritySupervisorAgentService — + * moved here to follow standard Spring exception package conventions. + */ +public class ScanException extends AgentException { + + public ScanException(String message) { + super(message); + } + + public ScanException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java new file mode 100644 index 0000000..e64601d --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java @@ -0,0 +1,90 @@ +package com.agent.model; + +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.RequiredArgsConstructor; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; + +import java.util.ArrayList; +import java.util.List; + +/** + * ConversationContext — carries session history for multi-turn chatbot mode. + * + * CURRENT STATUS: Scaffolded but not yet fully wired into agent prompts. + * The history is loaded from Redis (SessionStore.getHistory()) and passed + * through initialState, but agents currently use only the latest userQuery. + * + * FUTURE WIRING: + * Pass ConversationContext into agent reason() nodes via state. + * Use buildMessages(prompt) instead of List.of(new UserMessage(prompt)) + * to include previous turns in each LLM call — enabling true multi-turn + * conversation that remembers previous scan results and user questions. + * + * DirectAnswerAgent is the first candidate for full wiring since it handles + * open-ended user questions where history context is most valuable. + */ +@Data +@RequiredArgsConstructor +public class ConversationContext { + + private final String sessionId; + private final List history; + + /** + * Build the full message list to send to the LLM. + * Prepends all previous turns (user query + agent summary pairs) + * before the new user input. + * + * @param newUserInput the current user prompt or agent reasoning prompt + * @return ordered list of messages: [history turns...] + [newUserInput] + */ + public List buildMessages(String newUserInput) { + List messages = new ArrayList<>(); + + for (ChatMessage entry : history) { + if (entry.getUserQuery() != null && !entry.getUserQuery().isBlank()) { + messages.add(new UserMessage(entry.getUserQuery())); + } + if (entry.getAgentSummary() != null && !entry.getAgentSummary().isBlank()) { + messages.add(new AssistantMessage(entry.getAgentSummary())); + } + } + + messages.add(new UserMessage(newUserInput)); + return messages; + } + + /** + * Returns true if there are previous conversation turns in this session. + */ + public boolean hasHistory() { + return history != null && !history.isEmpty(); + } + + // ── Inner model ─────────────────────────────────────────────────────────── + + /** + * A single chat turn: user query + agent response summary. + * Stored in Redis as JSON via SessionStore.saveChatEntry(). + */ + @Data + @AllArgsConstructor + public static class ChatMessage { + + /** Original user query for this turn */ + private String userQuery; + + /** + * Compact summary of the agent's response for this turn. + * Used as the assistant message in history — not the full vulnerability report, + * just enough context for the LLM to understand what was done. + */ + private String agentSummary; + + /** ISO-8601 timestamp of this turn */ + private String timestamp; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/prompt/PromptLibrary.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/prompt/PromptLibrary.java index 3a74df2..44925a0 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/prompt/PromptLibrary.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/prompt/PromptLibrary.java @@ -1,37 +1,27 @@ package com.agent.prompt; /** - * PromptLibrary — single source of truth for all LLM system and user prompts. - *

- * All prompts are public static final String constants using Java text blocks. - * Agent classes must import prompts from here; no inline prompt strings are - * permitted inside agent classes. - *

- * Architecture: - * agent logic → agent classes - * prompt strings → PromptLibrary (this class) - * tool execution → MCP servers - * orchestration → LangGraph4j graph + * PromptLibrary — single source of truth for all LLM prompts. + * + * Changes from original: + * 1. Added SUPERVISOR_ROUTING_PROMPT — compact ~50-token prompt for subsequent + * supervisor calls when intent is already set. SupervisorAgent.buildRoutingPrompt() + * uses this instead of the full SUPERVISOR_AGENT_PROMPT, saving ~40% tokens + * on all routing calls after the first. + * 2. No other prompts changed — correctness of routing is more important than + * token savings; only the routing-only path is safe to compress. */ public final class PromptLibrary { /** - * SupervisorAgent reasoning prompt. - *

- * Placeholder tokens (in order): - * %s userQuery - * %s intent - * %s repositoryUrl - * %s packageName - * %s packageVersion - * %s packageEcosystem - * %s cveId - * %s fileType - * %s dependencyFiles (POPULATED | null) - * %s dependencyFileContents (POPULATED | null) - * %s dependencies (POPULATED | null) - * %s vulnerabilities (POPULATED | null) - * %s vulnerabilityDetails (POPULATED | null) + * SupervisorAgent — INITIAL call prompt (intent not yet set). + * Used only on the first reason() call when intent = null. + * LLM must detect intent and extract all entities. + * ~370 tokens. + * + * Placeholders (in order): userQuery, intent, repositoryUrl, packageName, + * packageVersion, packageEcosystem, cveId, fileType, dependencyFiles, + * dependencyFileContents, dependencies, vulnerabilities, vulnerabilityDetails */ public static final String SUPERVISOR_AGENT_PROMPT = """ You are the Security Supervisor Agent orchestrating a multi-agent pipeline. @@ -51,8 +41,7 @@ public final class PromptLibrary { dependencies = null → dependency_parser_agent dependencies = POPULATED AND vulnerabilities = null → vulnerability_agent dependencies = POPULATED AND vulnerabilities = POPULATED → DONE - IMPORTANT: if the dependency parser ran but produced 0 packages (empty list), - skip vulnerability_agent entirely and go directly to DONE. + IMPORTANT: if parser produced 0 packages, skip vulnerability_agent → DONE. Never route to vulnerability_agent more than once. PACKAGE_SCAN flow: @@ -67,11 +56,9 @@ public final class PromptLibrary { dependencies = null → dependency_parser_agent dependencies = POPULATED AND vulnerabilities = null → vulnerability_agent dependencies = POPULATED AND vulnerabilities = POPULATED → DONE - IMPORTANT: if the dependency parser ran but produced 0 packages, - skip vulnerability_agent and go directly to DONE. + IMPORTANT: if parser produced 0 packages, skip vulnerability_agent → DONE. - DIRECT_ANSWER flow: - → direct_answer_agent (always terminal) + DIRECT_ANSWER flow: → direct_answer_agent (always terminal) ── CURRENT STATE ──────────────────────────────────────────────────────── userQuery = %s @@ -90,7 +77,7 @@ public final class PromptLibrary { ── RESPONSE FORMAT ────────────────────────────────────────────────────── First call (intent = null): include all extracted entities. - Subsequent calls (intent already set): only nextAgent is required. + Subsequent calls: only nextAgent is required. First call response: { @@ -109,19 +96,36 @@ Subsequent calls (intent already set): only nextAgent is required. { "nextAgent": "" } """; - // ───────────────────────────────────────────────────────────────────────── - // SupervisorAgent - // ───────────────────────────────────────────────────────────────────────── + /** + * SupervisorAgent — ROUTING call prompt (intent already set). + * Used on all reason() calls AFTER the first (intent != null). + * LLM only needs to pick the next agent — no entity extraction needed. + * ~50 tokens vs ~370 for the full prompt. Saves ~40% tokens per pipeline run. + * + * Placeholders (in order): intent, dependencyFiles, dependencyFileContents, + * dependencies, vulnerabilities, vulnerabilityDetails + */ + public static final String SUPERVISOR_ROUTING_PROMPT = """ + Pipeline state: + intent=%s depFiles=%s fileContents=%s deps=%s vulns=%s details=%s + + Routes: github_repo_agent | dependency_parser_agent | vulnerability_agent | direct_answer_agent | DONE + + Rules: + REPO_SCAN: fileContents=null→github_repo_agent | deps=null→dependency_parser_agent | vulns=null→vulnerability_agent | else→DONE + FILE_SCAN: deps=null→dependency_parser_agent | vulns=null→vulnerability_agent | else→DONE + PACKAGE_SCAN: vulns=null→vulnerability_agent | else→DONE + CVE_QUERY: details=null→vulnerability_agent | else→DONE + If deps=[] (empty), skip vulnerability_agent → DONE + + Reply ONLY: { "nextAgent": "" } + """; + /** * GithubRepoAgent reasoning prompt. - *

- * Placeholder tokens (in order): - * %s formatted tool list - * %s repositoryUrl - * %s repositoryCloned - * %s repoPath - * %s dependencyFiles (value | null) - * %s dependencyFileContents (POPULATED | null) + * + * Placeholders (in order): formatted tool list, repositoryUrl, + * repositoryCloned, repoPath, dependencyFiles, dependencyFileContents */ public static final String GITHUB_REPO_AGENT_PROMPT = """ You are the GitHub Repository Agent. @@ -154,10 +158,9 @@ Decision rules (apply FIRST matching rule, in order): 2. repositoryCloned is true AND repoPath exists AND dependencyFiles is null → findDependencyFiles with { "repoPath": } 3. dependencyFiles = [] (empty list) → DONE - (repo has no recognised dependency files — stop here) 4. dependencyFiles is a NON-EMPTY list AND dependencyFileContents is null → readFileContent using the FIRST real path from the dependencyFiles list - CRITICAL: filePath MUST be a real absolute path from that list — never invent or use placeholder text + CRITICAL: filePath MUST be a real absolute path — never invent or use placeholder text 5. dependencyFileContents exists → DONE Respond with ONLY the JSON object. Examples: @@ -167,16 +170,11 @@ Decision rules (apply FIRST matching rule, in order): { "tool": "DONE", "args": {} } """; - // ───────────────────────────────────────────────────────────────────────── - // GithubRepoAgent - // ───────────────────────────────────────────────────────────────────────── /** * DependencyParserAgent reasoning prompt. - *

- * Placeholder tokens (in order): - * %s formatted tool list - * %d number of remaining files - * %s file overview (path + content preview for each remaining file) + * + * Placeholders (in order): formatted tool list, number of remaining files, + * file overview (path + content preview for each remaining file) */ public static final String DEPENDENCY_PARSER_AGENT_PROMPT = """ You are the Dependency Parser Agent. @@ -210,22 +208,12 @@ Files still needing parsing (%d remaining): { "tool": "parsePomXml", "filePath": "/tmp/repos/myapp/pom.xml" } """; - // ───────────────────────────────────────────────────────────────────────── - // DependencyParserAgent - // ───────────────────────────────────────────────────────────────────────── /** * VulnerabilityAgent reasoning prompt. - *

- * Placeholder tokens (in order): - * %s formatted tool list - * %s intent - * %s packageName - * %s packageVersion - * %s packageEcosystem - * %s cveId - * %s dependencies (POPULATED | null) - * %s vulnerabilities (POPULATED | null) - * %s vulnerabilityDetails (POPULATED | null) + * + * Placeholders (in order): formatted tool list, intent, packageName, + * packageVersion, packageEcosystem, cveId, dependencies, + * vulnerabilities, vulnerabilityDetails */ public static final String VULNERABILITY_AGENT_PROMPT = """ You are the Vulnerability Intelligence Agent. @@ -280,14 +268,10 @@ Decision rules (apply first matching rule): { "tool": "DONE", "args": {} } """; - // ───────────────────────────────────────────────────────────────────────── - // VulnerabilityAgent - // ───────────────────────────────────────────────────────────────────────── /** * DirectAnswerAgent clarification prompt. - *

- * Placeholder tokens (in order): - * %s userQuery + * + * Placeholders (in order): userQuery */ public static final String DIRECT_ANSWER_AGENT_PROMPT = """ You are a helpful AI cybersecurity assistant. @@ -316,12 +300,5 @@ Decision rules (apply first matching rule): Be concise and friendly. Do NOT make up any vulnerability data. """; - // ───────────────────────────────────────────────────────────────────────── - // DirectAnswerAgent - // ───────────────────────────────────────────────────────────────────────── - - private PromptLibrary() { - // utility class — no instances - } + private PromptLibrary() {} } - diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/service/SecuritySupervisorAgentService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/service/SecuritySupervisorAgentService.java index 22bb6a1..76d2938 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/service/SecuritySupervisorAgentService.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/service/SecuritySupervisorAgentService.java @@ -2,6 +2,11 @@ import com.agent.agent.SupervisorAgent; import com.agent.config.SessionContext; +import com.agent.exception.ScanException; +import com.agent.model.ConversationContext; +import com.fasterxml.jackson.core.type.TypeReference; + +import java.time.Instant; import com.agent.model.ScanProgressEvent; import com.agent.model.ScanResponse; import com.fasterxml.jackson.databind.JsonNode; @@ -19,19 +24,14 @@ /** * SecuritySupervisorAgentService - *

+ * * Entry point for all scan types. Supports two execution modes: - *

- * 1. executeScan(query, sessionId) — blocking, returns ScanResponse - * 2. executeScanStreaming(query, sessionId) — returns SseEmitter immediately - *

- * Session handling: - * - sessionId is resolved via SessionStore (generated if blank/null). - * - SessionContext thread-local is set before the graph runs so - * TokenCountingChatModel can increment counters for every LLM call. - * - After the pipeline completes, the chat entry (query + output) is - * persisted and the final token counts are read back from Redis and - * embedded in the ScanResponse. + * 1. executeScan(query, sessionId) — blocking, returns ScanResponse + * 2. executeScanStreaming(query, sessionId) — returns SseEmitter immediately + * + * Changes from original: + * - ScanException moved to com.agent.exception.ScanException (own class). + * Import updated accordingly. Behaviour is identical. */ @Slf4j @Service @@ -63,10 +63,13 @@ public ScanResponse executeScan(String userQuery, String incomingSessionId) { initialState.put("scanId", scanId); initialState.put("sessionId", sessionId); + // Load chat history from Redis and pass as ConversationContext + ConversationContext ctx = buildConversationContext(sessionId, userQuery); + initialState.put("conversationContext", ctx); + Map result = supervisorAgent.invoke(initialState); if (result.get("finalAnswer") instanceof String answer && !answer.isBlank()) { - // Save chat entry even for error-path answers sessionStore.saveChatEntry(sessionId, userQuery, answer); throw new ScanException(answer); } @@ -94,9 +97,6 @@ public SseEmitter executeScanStreaming(String userQuery, String incomingSessionI eventPublisher.register(scanId, emitter); streamingExecutor.submit(() -> { - // Set the sessionId on the virtual thread itself. - // ThreadLocal is NOT inherited from the submitting thread by virtual threads — - // each virtual thread starts with a clean ThreadLocal state. SessionContext.set(sessionId); try { eventPublisher.publish(scanId, ScanProgressEvent.started(scanId)); @@ -106,6 +106,10 @@ public SseEmitter executeScanStreaming(String userQuery, String incomingSessionI initialState.put("scanId", scanId); initialState.put("sessionId", sessionId); + // Load chat history from Redis and pass as ConversationContext + ConversationContext ctx = buildConversationContext(sessionId, userQuery); + initialState.put("conversationContext", ctx); + Map result = supervisorAgent.invoke(initialState); if (result.get("finalAnswer") instanceof String answer && !answer.isBlank()) { @@ -134,12 +138,11 @@ public SseEmitter executeScanStreaming(String userQuery, String incomingSessionI private ScanResponse buildResponse(Map result, String sessionId) { String intent = String.valueOf(result.getOrDefault("intent", "UNKNOWN")); - String vulnerabilitiesJson = toJsonString(result.get("vulnerabilities")); + String vulnerabilitiesJson = toJsonString(result.get("vulnerabilities")); String vulnerabilityDetailsJson = toJsonString(result.get("vulnerabilityDetails")); ScanResponse.VulnerabilitySummary summary = buildSummary(vulnerabilitiesJson); - // Read cumulative token counts from Redis for this session Map tokens = sessionStore.getTokenCounts(sessionId); ScanResponse.ScanResponseBuilder builder = ScanResponse.builder() @@ -155,9 +158,8 @@ private ScanResponse buildResponse(Map result, String sessionId) case SupervisorAgent.INTENT_REPO_SCAN -> { String deps = toJsonString(result.get("dependencies")); int depCount = countJsonArray(deps); - builder - .repositoryUrl(stringOrNull(result, "repositoryUrl")) - .totalDependencies(Math.max(depCount, 0)); + builder.repositoryUrl(stringOrNull(result, "repositoryUrl")) + .totalDependencies(Math.max(depCount, 0)); } case SupervisorAgent.INTENT_PACKAGE_SCAN -> builder .packageName(stringOrNull(result, "packageName")) @@ -170,9 +172,8 @@ private ScanResponse buildResponse(Map result, String sessionId) case SupervisorAgent.INTENT_FILE_SCAN -> { String deps = toJsonString(result.get("dependencies")); int depCount = countJsonArray(deps); - builder - .fileType(stringOrNull(result, "fileType")) - .totalDependencies(Math.max(depCount, 0)); + builder.fileType(stringOrNull(result, "fileType")) + .totalDependencies(Math.max(depCount, 0)); if (depCount == 0) { builder.message("No recognisable package dependencies were found in the provided file."); } @@ -201,17 +202,14 @@ private ScanResponse.VulnerabilitySummary buildSummary(String vulnerabilitiesJso boolean hasVulns = vulnList.isArray() && !vulnList.isEmpty(); String riskLevel = report.path("score").path("riskLevel").asText("LOW"); - if (!hasVulns) { - safe++; - continue; - } + if (!hasVulns) { safe++; continue; } total++; switch (riskLevel.toUpperCase()) { case "CRITICAL" -> critical++; - case "HIGH" -> high++; - case "MEDIUM" -> medium++; - default -> low++; + case "HIGH" -> high++; + case "MEDIUM" -> medium++; + default -> low++; } } } @@ -231,23 +229,14 @@ private ScanResponse.VulnerabilitySummary buildSummary(String vulnerabilitiesJso // ── Helpers ─────────────────────────────────────────────────────────────── - /** - * Produces a compact Map summary of the ScanResponse for storing in - * the Redis chat history. Returning a Map (not a pre-serialized String) - * prevents double-serialization in SessionStore.saveChatEntry(). - */ private Map summariseResponse(ScanResponse r) { Map summary = new HashMap<>(); - summary.put("scanType", r.getScanType()); - summary.put("sessionId", r.getSessionId()); + summary.put("scanType", r.getScanType()); + summary.put("sessionId", r.getSessionId()); summary.put("inputTokens", r.getTotalInputTokens()); summary.put("outputTokens", r.getTotalOutputTokens()); - if (r.getMessage() != null) { - summary.put("message", r.getMessage()); - } - if (r.getSummary() != null) { - summary.put("vulnerabilitySummary", r.getSummary()); - } + if (r.getMessage() != null) summary.put("message", r.getMessage()); + if (r.getSummary() != null) summary.put("vulnerabilitySummary", r.getSummary()); return summary; } @@ -256,11 +245,8 @@ private String toJsonString(Object value) { String s = value.toString().trim(); if (s.isBlank() || s.equals("null")) return null; if (s.startsWith("[") || s.startsWith("{")) return s; - try { - return objectMapper.writeValueAsString(s); - } catch (Exception e) { - return null; - } + try { return objectMapper.writeValueAsString(s); } + catch (Exception e) { return null; } } private String stringOrNull(Map result, String key) { @@ -275,16 +261,41 @@ private int countJsonArray(String json) { try { JsonNode node = objectMapper.readTree(json); return node.isArray() ? node.size() : 0; - } catch (Exception e) { - return 0; - } + } catch (Exception e) { return 0; } } - // ── Exception ───────────────────────────────────────────────────────────── + // ── ConversationContext builder ─────────────────────────────────────────── - public static class ScanException extends RuntimeException { - public ScanException(String message) { - super(message); + /** + * Build ConversationContext from Redis session history. + * Each history entry is a JSON string of {userQuery, agentSummary, timestamp}. + * If history is empty or an error occurs, returns an empty context. + */ + private ConversationContext buildConversationContext(String sessionId, String currentQuery) { + try { + java.util.List rawHistory = sessionStore.getHistory(sessionId); + java.util.List messages = new java.util.ArrayList<>(); + + for (String entry : rawHistory) { + try { + com.fasterxml.jackson.databind.JsonNode node = objectMapper.readTree(entry); + String query = node.path("userQuery").asText(null); + String summary = node.path("agentSummary").asText(null); + String ts = node.path("timestamp").asText(Instant.now().toString()); + if (query != null) { + messages.add(new ConversationContext.ChatMessage(query, summary, ts)); + } + } catch (Exception ignored) {} + } + + if (!messages.isEmpty()) { + log.debug("[Service] Loaded {} history turns for session {}", messages.size(), sessionId); + } + return new ConversationContext(sessionId, messages); + + } catch (Exception e) { + log.warn("[Service] Could not load chat history for session {}: {}", sessionId, e.getMessage()); + return new ConversationContext(sessionId, java.util.List.of()); } } } \ No newline at end of file diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/McpToolRegistry.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/McpToolRegistry.java index ea2f847..101bf53 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/McpToolRegistry.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/McpToolRegistry.java @@ -24,7 +24,7 @@ * tool name lookup and dynamic tool listing are always server-scoped. */ @Slf4j -public class McpToolRegistry { +public class McpToolRegistry implements ToolRegistry { private final String serverName; private final Map toolByName; diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/ToolRegistry.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/ToolRegistry.java new file mode 100644 index 0000000..165d2e4 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/tool/ToolRegistry.java @@ -0,0 +1,47 @@ +package com.agent.tool; + +import org.springframework.ai.tool.ToolCallback; + +import java.util.List; +import java.util.Map; + +/** + * ToolRegistry — interface for server-scoped MCP tool registries. + * + * Extracted from the concrete McpToolRegistry to enable mock injection + * in unit tests. Each agent injects ToolRegistry rather than McpToolRegistry + * directly, so tests can substitute a stub without bringing up MCP servers. + * + * Three beans implement this in McpConfig: + * repoToolRegistry → github-repo-mcp-server (port 9001) + * parserToolRegistry → dependency-parser-mcp-server (port 9002) + * vulnerabilityToolRegistry → vulnerability-intel-mcp-server (port 9003) + */ +public interface ToolRegistry { + + /** + * Returns the ToolCallback for the given tool name, or null if not found. + */ + ToolCallback getTool(String toolName); + + /** + * Returns all tool names registered for this server. + */ + List getToolNames(); + + /** + * Returns all tool callbacks as an immutable map. + */ + Map getAllTools(); + + /** + * Returns the logical name of the MCP server this registry is scoped to. + */ + String getServerName(); + + /** + * Returns true if toolName belongs to this server's tool set. + * Agents must call this before executing a tool to avoid cross-server calls. + */ + boolean isKnownTool(String toolName); +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/AgentUtils.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/AgentUtils.java new file mode 100644 index 0000000..0fb51dd --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/AgentUtils.java @@ -0,0 +1,101 @@ +package com.agent.util; + +import com.fasterxml.jackson.databind.ObjectMapper; +import lombok.extern.slf4j.Slf4j; + +import java.util.List; +import java.util.Map; + +/** + * AgentUtils — shared utilities for all sub-agents. + * + * Previously these methods were copy-pasted across GithubRepoAgent, + * VulnerabilityAgent, and DependencyParserAgent. Centralised here to + * ensure consistent behaviour and single-point fixes. + */ +@Slf4j +public final class AgentUtils { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private AgentUtils() { + // utility class — no instances + } + + /** + * Format a list of tool names as a bullet string for prompt injection. + * + * Example output: + * - cloneRepository + * - findDependencyFiles + * - readFileContent + */ + public static String formatToolList(List tools) { + var sb = new StringBuilder(); + tools.forEach(t -> sb.append("- ").append(t).append("\n")); + return sb.toString().stripTrailing(); + } + + /** + * Unwrap MCP envelope: [{"text":""}] → value + * + * MCP tool responses arrive as a JSON array with a single object + * containing a "text" field. This method extracts the inner value + * and strips surrounding quotes if present. + * + * Falls back to result.toString() if the envelope format is not recognised. + * + * Returns null if result is null (callers should check with McpResponseValidator). + */ + public static String extractText(Object result) { + if (result == null) { + return null; + } + try { + String raw = result.toString().trim(); + + if (raw.startsWith("[")) { + List> list = MAPPER.readValue( + raw, + MAPPER.getTypeFactory().constructCollectionType(List.class, Map.class) + ); + if (!list.isEmpty() && list.get(0).containsKey("text")) { + String text = list.get(0).get("text").toString(); + if (text.startsWith("\"") && text.endsWith("\"")) { + text = text.substring(1, text.length() - 1); + } + return text.replace("\\\"", "\"").trim(); + } + } + + if (raw.startsWith("\"") && raw.endsWith("\"")) { + raw = raw.substring(1, raw.length() - 1); + } + return raw.trim(); + + } catch (Exception e) { + log.warn("[AgentUtils] Failed to parse MCP response envelope: {}", result); + return result.toString(); + } + } + + /** + * Strip markdown code fences from LLM JSON output. + * + * LLMs sometimes wrap JSON in ```json ... ``` or ``` ... ``` blocks. + * This method removes those fences so the result can be parsed directly. + */ + public static String stripMarkdown(String raw) { + if (raw == null) return ""; + return raw.replaceAll("(?s)```[a-z]*\\s*", "").replaceAll("```", "").trim(); + } + + /** + * Check if a string value represents an empty/absent state. + * Used to test state field values from SecurityAgentState. + */ + public static boolean isNullOrEmpty(String value) { + return value == null || value.isBlank() || "null".equals(value.trim()) + || "[]".equals(value.trim()) || "{}".equals(value.trim()); + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/McpResponseValidator.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/McpResponseValidator.java new file mode 100644 index 0000000..270a9ed --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/util/McpResponseValidator.java @@ -0,0 +1,54 @@ +package com.agent.util; + +import com.agent.exception.McpToolNullResponseException; +import lombok.extern.slf4j.Slf4j; + +/** + * McpResponseValidator — guards against null/empty MCP tool responses. + * + * PROBLEM SOLVED: + * Previously, extractText() silently returned an empty string on null MCP + * responses. The agent stored nothing in state and called reason() again + * with unchanged state — potentially looping indefinitely and burning tokens. + * + * SOLUTION: + * Call validateAndExtract() in every tool() node BEFORE storing the result + * in state. If the response is null/empty, McpToolNullResponseException is + * thrown immediately, propagating to invoke() and then to SupervisorAgent, + * which returns a clean error response with zero additional LLM calls. + */ +@Slf4j +public final class McpResponseValidator { + + private McpResponseValidator() { + // utility class — no instances + } + + /** + * Extract and validate an MCP tool response. + * + * @param result raw result from ToolCallback.call() + * @param toolName name of the tool (used in exception message and logs) + * @return extracted text value — never null, never blank + * @throws McpToolNullResponseException if the response is null, blank, + * the literal string "null", or an empty JSON array "[]" + */ + public static String validateAndExtract(Object result, String toolName) { + if (result == null) { + log.warn("[McpResponseValidator] Tool '{}' returned null", toolName); + throw new McpToolNullResponseException(toolName); + } + + String extracted = AgentUtils.extractText(result); + + if (AgentUtils.isNullOrEmpty(extracted)) { + log.warn("[McpResponseValidator] Tool '{}' returned empty/null response: '{}'", + toolName, extracted); + throw new McpToolNullResponseException(toolName); + } + + log.debug("[McpResponseValidator] Tool '{}' response OK ({} chars)", toolName, + extracted.length()); + return extracted; + } +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/IVulnerabilityService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/IVulnerabilityService.java new file mode 100644 index 0000000..e37a3b6 --- /dev/null +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/IVulnerabilityService.java @@ -0,0 +1,31 @@ +package com.agent.service; + +import com.agent.model.Dependency; +import com.agent.model.VulnerabilityReport; + +import java.util.List; + +/** + * IVulnerabilityService — contract for vulnerability analysis operations. + * + * Extracted from the concrete class to enable mock injection in unit tests. + * Implementation: {@link VulnerabilityService} + */ +public interface IVulnerabilityService { + + /** + * Analyse a single dependency for known vulnerabilities. + * + * @param dependency the dependency to scan + * @return a VulnerabilityReport with all known CVEs and a composite risk score + */ + VulnerabilityReport analyzeDependency(Dependency dependency); + + /** + * Analyse multiple dependencies in parallel using virtual threads. + * + * @param dependencies list of dependencies to scan + * @return one VulnerabilityReport per dependency, order matching input + */ + List analyzeMultipleDependencies(List dependencies); +} diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/VulnerabilityService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/VulnerabilityService.java index 696b6c0..62e6012 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/VulnerabilityService.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/vulnerability-intel-mcp-server/src/main/java/com/agent/service/VulnerabilityService.java @@ -1,5 +1,7 @@ package com.agent.service; +import com.agent.service.IVulnerabilityService; + import com.agent.engine.RecommendationEngine; import com.agent.engine.UnifiedScoringEngine; import com.agent.engine.VulnerabilityNormalizer; @@ -19,7 +21,7 @@ @Slf4j @Service @RequiredArgsConstructor -public class VulnerabilityService { +public class VulnerabilityService implements IVulnerabilityService { private final VulnerabilityAggregationService aggregationService; private final VulnerabilityNormalizer normalizer; From d646f5c4d6cca471d230479b551fc3f440f81f95 Mon Sep 17 00:00:00 2001 From: loglogn Date: Thu, 12 Mar 2026 14:36:31 +0530 Subject: [PATCH 2/2] feat: Enhance ConversationContext and DependencyParserAgent for serialization; implement Serializable interface and improve dependency parsing logic --- .../main/java/com/agent/parser/PipParser.java | 74 +++-- .../com/agent/service/IGithubRepoService.java | 2 +- .../service/impl/GithubRepoServiceImpl.java | 167 ---------- .../sub_agents/DependencyParserAgent.java | 294 ++++++------------ .../agent/sub_agents/VulnerabilityAgent.java | 68 +++- .../com/agent/model/ConversationContext.java | 12 +- 6 files changed, 218 insertions(+), 399 deletions(-) delete mode 100644 02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/parser/PipParser.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/parser/PipParser.java index db45f8e..1c442a5 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/parser/PipParser.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/dependency-parser-mcp-server/src/main/java/com/agent/parser/PipParser.java @@ -13,13 +13,30 @@ @Component public class PipParser { - // Matches: package==1.0, package>=1.0, package~=1.0, package<=1.0, package!=1.0 - // Also handles inline comments like: requests==2.31.0 # HTTP library - private static final Pattern DEP_PATTERN = - Pattern.compile("^([A-Za-z0-9_.-]+)\\s*([><=!~]+)\\s*([A-Za-z0-9_.*+-]+)"); + /** + * Matches a pip requirement line: + * group 1 — package name (letters, digits, -, _, .) + * group 2 — optional extras e.g. [asyncio,security] (captured but discarded) + * group 3 — optional version specifier(s) e.g. >=1.0,<2.0 or ~=1.4.4 + * (environment markers after ';' are stripped before matching) + * + * Examples handled: + * requests==2.31.0 + * aiohttp>=3.11.18 + * sqlalchemy[asyncio]>=2.0.41 + * markitdown-no-magika[docx,xls,xlsx]>=0.1.2 + * audioop-lts ; python_full_version >= '3.13' + * xinference-client (no version) + * chardet~=5.1.0 + */ + private static final Pattern DEP_PATTERN = Pattern.compile( + "^([A-Za-z0-9][A-Za-z0-9_.-]*) # package name\n" + + "(\\[[^\\]]*\\])? # optional extras\n" + + "\\s*([><=!~,A-Za-z0-9_.*+\\-]*)\\s*$ # optional version spec(s)", + Pattern.COMMENTS + ); public List parse(String content) { - List deps = new ArrayList<>(); if (content == null || content.isBlank()) { @@ -27,43 +44,42 @@ public List parse(String content) { return deps; } - String[] lines = content.split("\n"); + // Normalise escaped newlines that appear when content travels through JSON + // e.g. "requests==2.31.0\naiohttp>=3.11.18" → real newlines + String normalised = content + .replace("\\r\\n", "\n") + .replace("\\n", "\n") + .replace("\\r", "\n"); - for (String rawLine : lines) { + for (String rawLine : normalised.split("\n")) { - // Strip inline comments and whitespace - String line = rawLine.split("#")[0].trim(); + // Strip inline comments and environment markers (everything after ';') + String line = rawLine.split("#")[0].split(";")[0].trim(); - // Skip blank lines, comments, -r includes, -e editable installs + // Skip blank lines, options (-r, --index-url, -e editable installs, etc.) if (line.isEmpty() || line.startsWith("-")) { continue; } Matcher matcher = DEP_PATTERN.matcher(line); + if (!matcher.matches()) { + log.debug("[PipParser] Line did not match: '{}'", line); + continue; + } - if (matcher.find()) { - String name = matcher.group(1).trim(); - String version = matcher.group(2).trim() + matcher.group(3).trim(); + String name = matcher.group(1).trim(); + String version = matcher.group(3) == null ? "" : matcher.group(3).trim(); - deps.add(Dependency.builder() - .name(name) - .version(version) - .ecosystem("pip") - .build()); - } else { - // Package with no version pinned (e.g. just "requests") - if (!line.isEmpty()) { - deps.add(Dependency.builder() - .name(line.trim()) - .version("") - .ecosystem("pip") - .build()); - } - } + if (name.isEmpty()) continue; + + deps.add(Dependency.builder() + .name(name) + .version(version) + .ecosystem("pip") + .build()); } log.info("[PipParser] Parsed {} dependencies", deps.size()); return deps; } - } \ No newline at end of file diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java index 4370640..c635c62 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/IGithubRepoService.java @@ -9,7 +9,7 @@ * Extracted from the concrete class to enable mock injection in unit tests. * Implementation: {@link impl.GithubRepoServiceImpl} */ -public interface GithubRepoService { +public interface IGithubRepoService { /** * Clone a GitHub repository to a local temp directory. diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java deleted file mode 100644 index 318b746..0000000 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/github-repo-mcp-server/src/main/java/com/agent/service/impl/GithubRepoServiceImpl.java +++ /dev/null @@ -1,167 +0,0 @@ -package com.agent.service.impl; -import com.agent.service.IGithubRepoService; - -import com.agent.util.FileScannerUtil; -import lombok.extern.slf4j.Slf4j; -import org.springframework.stereotype.Service; - -import java.io.IOException; -import java.io.InputStream; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Comparator; -import java.util.List; -import java.util.concurrent.TimeUnit; - -@Slf4j -@Service -public class GithubRepoServiceImpl implements IGithubRepoService { - - private static final String REPO_BASE_DIR = "/tmp/repos"; - private static final int CLONE_TIMEOUT_SECONDS = 120; - - // Allowed: https://github.com/owner/repo or https://github.com/owner/repo.git - private static final String GITHUB_URL_PATTERN = - "https://github\\.com/[\\w.-]+/[\\w.-]+(\\.git)?"; - - /** - * Clone a GitHub repository into the local base directory. - * - * @param repoUrl public GitHub HTTPS URL - * @param forceReclone if true, deletes existing clone and re-clones fresh - */ - public String cloneRepository(String repoUrl, boolean forceReclone) throws Exception { - - // 1. Validate GitHub URL format - if (repoUrl == null || !repoUrl.matches(GITHUB_URL_PATTERN)) { - throw new IllegalArgumentException( - "Invalid GitHub URL. Expected: https://github.com/owner/repo — got: " + repoUrl - ); - } - - // 2. Derive safe directory name - String repoName = repoUrl - .substring(repoUrl.lastIndexOf("/") + 1) - .replace(".git", ""); - - String target = REPO_BASE_DIR + "/" + repoName; - Path targetPath = Paths.get(target); - - // 3. Handle existing clone - if (Files.exists(targetPath)) { - if (forceReclone) { - log.info("[SERVICE] forceReclone=true — deleting existing clone at: {}", target); - deleteDirectory(targetPath); - } else { - long fileCount = Files.walk(targetPath) - .filter(Files::isRegularFile) - .count(); - log.info("[SERVICE] Repository already cloned at: {} | existingFileCount={}", target, fileCount); - - if (fileCount == 0) { - log.warn("[SERVICE] Existing clone is EMPTY (possibly a failed previous clone) — force re-cloning: {}", target); - deleteDirectory(targetPath); - } else { - return target; - } - } - } - - // 4. Ensure base directory exists - Files.createDirectories(Paths.get(REPO_BASE_DIR)); - - log.info("[SERVICE] Cloning repository: {} → {}", repoUrl, target); - - ProcessBuilder builder = new ProcessBuilder("git", "clone", repoUrl, target); - builder.redirectErrorStream(true); - Process process = builder.start(); - - // 5. Enforce timeout - boolean finished = process.waitFor(CLONE_TIMEOUT_SECONDS, TimeUnit.SECONDS); - - if (!finished) { - process.destroyForcibly(); - throw new RuntimeException( - "Git clone timed out after " + CLONE_TIMEOUT_SECONDS + "s for: " + repoUrl - ); - } - - int exitCode = process.exitValue(); - if (exitCode != 0) { - String errorOutput = readProcessOutput(process); - throw new RuntimeException( - "Git clone failed (exit=" + exitCode + "): " + errorOutput - ); - } - - long clonedFileCount = Files.walk(targetPath) - .filter(Files::isRegularFile) - .count(); - - log.info("[SERVICE] Clone successful | path={} | totalFiles={}", target, clonedFileCount); - return target; - } - - /** - * List all non-internal files in a cloned repository. - */ - public List listRepositoryFiles(String repoPath) throws IOException { - validateRepoPath(repoPath); - List files = FileScannerUtil.listFiles(repoPath); - log.debug("[SERVICE] listRepositoryFiles | repoPath={} | count={}", repoPath, files.size()); - return files; - } - - /** - * Find dependency/manifest files in the repository. - */ - public List findDependencyFiles(String repoPath) throws IOException { - validateRepoPath(repoPath); - List depFiles = FileScannerUtil.findDependencyFiles(repoPath); - log.debug("[SERVICE] findDependencyFiles | repoPath={} | found={}", repoPath, depFiles.size()); - return depFiles; - } - - /** - * Read file content. Path traversal protection enforced in FileScannerUtil. - */ - public String readFileContent(String filePath) throws IOException { - return FileScannerUtil.readFile(filePath); - } - - // ------------------------------------------------------------------------- - // Private helpers - // ------------------------------------------------------------------------- - - private void validateRepoPath(String repoPath) { - Path resolved = Paths.get(repoPath).normalize(); - if (!resolved.startsWith(REPO_BASE_DIR)) { - throw new SecurityException( - "Access denied: path is outside allowed base directory. Path: " + repoPath - ); - } - } - - private void deleteDirectory(Path path) throws IOException { - Files.walk(path) - .sorted(Comparator.reverseOrder()) - .forEach(p -> { - try { - Files.delete(p); - } catch (IOException e) { - log.warn("[SERVICE] Failed to delete: {}", p); - } - }); - log.info("[SERVICE] Deleted directory: {}", path); - } - - private String readProcessOutput(Process process) { - try (InputStream is = process.getInputStream()) { - return new String(is.readAllBytes()); - } catch (IOException e) { - return "(could not read process output)"; - } - } - -} \ No newline at end of file diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java index ccf7ec8..d01831d 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/DependencyParserAgent.java @@ -1,10 +1,7 @@ package com.agent.agent.sub_agents; -import com.agent.exception.MaxLlmCallsExceededException; -import com.agent.exception.McpToolExecutionException; import com.agent.exception.McpToolNullResponseException; import com.agent.model.ScanProgressEvent; -import com.agent.prompt.PromptLibrary; import com.agent.service.ScanEventPublisher; import com.agent.state.SecurityAgentState; import com.agent.tool.McpToolRegistry; @@ -12,263 +9,166 @@ import com.agent.util.McpResponseValidator; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.extern.slf4j.Slf4j; -import org.bsc.langgraph4j.CompiledGraph; -import org.bsc.langgraph4j.StateGraph; -import org.springframework.ai.chat.client.ChatClient; -import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.tool.ToolCallback; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.stereotype.Component; import java.util.*; -import static org.bsc.langgraph4j.StateGraph.END; -import static org.bsc.langgraph4j.StateGraph.START; -import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async; -import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async; - /** - * DependencyParserAgent — LLM-driven dependency file parsing sub-agent. + * DependencyParserAgent — fully deterministic dependency file parsing sub-agent. + * + * No LLM calls are made. The correct MCP parser tool is selected directly from + * the file name/path using resolveToolForFile(). The file content is already + * available in state (written by GithubRepoAgent.onReadFileContent()), so it + * is passed straight to the MCP tool as {"content": }. * - * Changes from original: - * 1. LLM call counter (MAX_LLM_CALLS = 8) with MaxLlmCallsExceededException. - * 2. McpResponseValidator.validateAndExtract() in tool() — note: empty result - * (no deps found) is valid here, so we catch McpToolNullResponseException - * and treat it as an empty dep list rather than re-throwing. - * 3. Uses AgentUtils.formatToolList() and AgentUtils.stripMarkdown() shared utils. - * 4. MaxLlmCallsExceededException re-throws to SupervisorAgent. + * Tool mapping (deterministic, based on filename): + * pom.xml → parsePomXml (ecosystem: maven) + * package.json → parsePackageJson (ecosystem: npm) + * requirements.txt → parseRequirementsTxt (ecosystem: pip) + * build.gradle / build.gradle.kts → parseGradleDependencies (ecosystem: gradle) + * pubspec.yaml → parsePubspecDependencies (ecosystem: dart) */ @Slf4j @Component public class DependencyParserAgent { - private static final String DONE = "DONE"; - private static final int MAX_LLM_CALLS = 8; - private static final String STATE_KEY_CALL_COUNT = "llmCallCount_DependencyParserAgent"; - - private final ChatClient chatClient; private final McpToolRegistry toolRegistry; private final ObjectMapper objectMapper = new ObjectMapper(); private final ScanEventPublisher eventPublisher; - private CompiledGraph graph; - public DependencyParserAgent( - ChatClient chatClient, @Qualifier("parserToolRegistry") McpToolRegistry toolRegistry, ScanEventPublisher eventPublisher) { - this.chatClient = chatClient; this.toolRegistry = toolRegistry; this.eventPublisher = eventPublisher; } - private synchronized void build() throws Exception { - if (graph != null) return; - graph = new StateGraph<>(SecurityAgentState.SCHEMA, SecurityAgentState::new) - .addNode("reason", node_async(this::reason)) - .addNode("tool", node_async(this::tool)) - .addEdge(START, "reason") - .addConditionalEdges( - "reason", - edge_async(s -> DONE.equals(s.nextTool().orElse(DONE)) ? END : "tool"), - Map.of("tool", "tool", END, END) - ) - .addEdge("tool", "reason") - .compile(); - } + // ── Entry point ─────────────────────────────────────────────────────────── public Map invoke(Map incomingState) { - try { - build(); - return graph.invoke(incomingState) - .map(SecurityAgentState::data) - .orElse(incomingState); - - } catch (MaxLlmCallsExceededException e) { - throw e; - - } catch (Exception e) { - log.error("[DependencyParserAgent] Unexpected error", e); - Map err = new HashMap<>(incomingState); - err.put("finalAnswer", "DependencyParserAgent error: " + e.getMessage()); - return err; - } - } - - @SuppressWarnings("unchecked") - private Map reason(SecurityAgentState state) { - log.info("[DependencyParserAgent] Reasoning"); + // Build a temporary SecurityAgentState view over the incoming map + SecurityAgentState state = new SecurityAgentState(new HashMap<>(incomingState)); - int callCount = (int) state.data().getOrDefault(STATE_KEY_CALL_COUNT, 0); - if (callCount >= MAX_LLM_CALLS) { - log.error("[DependencyParserAgent] Max LLM calls ({}) exceeded", MAX_LLM_CALLS); - throw new MaxLlmCallsExceededException("DependencyParserAgent", MAX_LLM_CALLS); - } - - List availableTools = toolRegistry.getToolNames(); Map fileMap = extractFileMap(state); - Set alreadyParsed = extractParsedFilePaths(state); - List remaining = fileMap.keySet().stream() - .filter(p -> !alreadyParsed.contains(p)).toList(); - - if (remaining.isEmpty()) { - log.info("[DependencyParserAgent] All files parsed — DONE"); - return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); + if (fileMap.isEmpty()) { + log.warn("[DependencyParserAgent] No file content in state — nothing to parse"); + return incomingState; } - StringBuilder fileOverview = new StringBuilder(); - remaining.forEach(path -> { - String preview = fileMap.get(path); - if (preview.length() > 300) preview = preview.substring(0, 300) + "\n...[truncated]"; - fileOverview.append("File: ").append(path).append("\n") - .append("Content preview:\n").append(preview).append("\n---\n"); - }); + String scanId = state.scanId().orElse(null); + int total = fileMap.size(); - String prompt = PromptLibrary.DEPENDENCY_PARSER_AGENT_PROMPT.formatted( - AgentUtils.formatToolList(availableTools), - remaining.size(), - fileOverview - ); + // Accumulated deps across all files (deduped by name:version:ecosystem key) + Map accumulatedDeps = new LinkedHashMap<>(); + Set parsedFiles = new LinkedHashSet<>(); + int current = 0; - String raw = Objects.requireNonNull( - chatClient.prompt().messages(new UserMessage(prompt)).call().content() - ).trim(); + for (Map.Entry entry : fileMap.entrySet()) { + String filePath = entry.getKey(); + String content = entry.getValue(); + current++; - log.info("[DependencyParserAgent] LLM decision: {}", raw); + String toolName = resolveToolForFile(filePath); + if (toolName == null) { + log.warn("[DependencyParserAgent] No parser tool for '{}' — skipping", filePath); + parsedFiles.add(filePath); + continue; + } - try { - String json = AgentUtils.stripMarkdown(raw); - Map decision = objectMapper.readValue(json, Map.class); - String toolName = String.valueOf(decision.getOrDefault("tool", DONE)).trim(); - String filePath = String.valueOf(decision.getOrDefault("filePath", "")).trim(); + if (!toolRegistry.isKnownTool(toolName)) { + log.warn("[DependencyParserAgent] Tool '{}' not in registry for '{}' — skipping", toolName, filePath); + parsedFiles.add(filePath); + continue; + } - if (DONE.equals(toolName)) { - return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); + if (content.isBlank()) { + log.warn("[DependencyParserAgent] Empty content for '{}' — skipping", filePath); + parsedFiles.add(filePath); + continue; } - if (!fileMap.containsKey(filePath)) { - String fallback = remaining.get(0); - log.warn("[DependencyParserAgent] LLM returned unknown filePath '{}' — using '{}'", filePath, fallback); - filePath = fallback; + if (scanId != null) { + eventPublisher.publish(scanId, ScanProgressEvent.parsingFile(scanId, filePath, current, total)); } - log.info("[DependencyParserAgent] -> tool={} filePath={}", toolName, filePath); - return Map.of( - "nextTool", toolName, - "nextToolArgs", "{}", - "currentParsingFile", filePath, - STATE_KEY_CALL_COUNT, callCount + 1 - ); - } catch (Exception e) { - log.error("[DependencyParserAgent] Failed to parse LLM response '{}': {}", raw, e.getMessage()); - return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount + 1); + log.info("[DependencyParserAgent] Parsing '{}' with tool '{}'", filePath, toolName); + List> fileDeps = parseFile(toolName, filePath, content); + fileDeps.forEach(d -> accumulatedDeps.putIfAbsent(buildDepKey(d), d)); + parsedFiles.add(filePath); + log.info("[DependencyParserAgent] {} → {} deps from '{}'", toolName, fileDeps.size(), filePath); } - } - @SuppressWarnings("unchecked") - private Map tool(SecurityAgentState state) { - String toolName = state.nextTool().orElse(""); - String filePath = state.data().getOrDefault("currentParsingFile", "unknown").toString(); - String scanId = state.scanId().orElse(null); - log.info("[DependencyParserAgent] Executing tool={} for file={}", toolName, filePath); + log.info("[DependencyParserAgent] Total accumulated: {} unique deps from {} file(s)", + accumulatedDeps.size(), parsedFiles.size()); - if (!toolRegistry.isKnownTool(toolName)) { - log.warn("[DependencyParserAgent] Unknown tool '{}'. Known: {}", toolName, toolRegistry.getToolNames()); - return markFileParsed(state, filePath, List.of()); + Map updates = new HashMap<>(incomingState); + try { + updates.put("dependencies", objectMapper.writeValueAsString( + new ArrayList<>(accumulatedDeps.values()))); + updates.put("parsedFiles", objectMapper.writeValueAsString( + new ArrayList<>(parsedFiles))); + } catch (Exception e) { + log.error("[DependencyParserAgent] Failed to serialize state update", e); } + return updates; + } - Map fileMap = extractFileMap(state); - String fullContent = fileMap.getOrDefault(filePath, ""); - if (fullContent.isBlank()) { - log.warn("[DependencyParserAgent] No content for filePath='{}' — skipping", filePath); - return markFileParsed(state, filePath, List.of()); - } + // ── Deterministic tool selection ────────────────────────────────────────── + + /** + * Resolve the MCP parser tool name from a file path. + * Matching is done on the lowercased filename only (not the full path). + * Returns null if no supported parser exists for this file. + */ + private String resolveToolForFile(String filePath) { + if (filePath == null || filePath.isBlank()) return null; + String name = filePath.toLowerCase(); + // Use endsWith so both /tmp/repos/app/pom.xml and pom.xml are matched + if (name.endsWith("pom.xml")) return "parsePomXml"; + if (name.endsWith("package.json")) return "parsePackageJson"; + if (name.endsWith("requirements.txt")) return "parseRequirementsTxt"; + if (name.endsWith("build.gradle") || name.endsWith("build.gradle.kts")) return "parseGradleDependencies"; + if (name.endsWith("pubspec.yaml")) return "parsePubspecDependencies"; + return null; + } + + // ── Parse a single file via MCP tool ────────────────────────────────────── + @SuppressWarnings("unchecked") + private List> parseFile(String toolName, String filePath, String content) { String jsonArgs; try { - jsonArgs = objectMapper.writeValueAsString(Map.of("content", fullContent)); + jsonArgs = objectMapper.writeValueAsString(Map.of("content", content)); } catch (Exception e) { - log.error("[DependencyParserAgent] Failed to serialize args for {}", filePath, e); - return markFileParsed(state, filePath, List.of()); - } - - if (scanId != null) { - Set parsed = extractParsedFilePaths(state); - int total = fileMap.size(); - int current = parsed.size() + 1; - eventPublisher.publish(scanId, ScanProgressEvent.parsingFile(scanId, filePath, current, total)); + log.error("[DependencyParserAgent] Failed to serialize content for '{}': {}", filePath, e.getMessage()); + return List.of(); } ToolCallback tool = toolRegistry.getTool(toolName); try { Object result = tool.call(jsonArgs); - // For parsers, an empty result (no deps found) is valid — catch and treat as empty list String extracted; try { extracted = McpResponseValidator.validateAndExtract(result, toolName); } catch (McpToolNullResponseException e) { - log.info("[DependencyParserAgent] Tool '{}' returned empty for {} — treating as 0 deps", toolName, filePath); - return markFileParsed(state, filePath, List.of()); + log.info("[DependencyParserAgent] '{}' returned empty for '{}' — 0 deps", toolName, filePath); + return List.of(); } - List> fileDeps = new ArrayList<>(); - if (!AgentUtils.isNullOrEmpty(extracted)) { - fileDeps = objectMapper.readValue( - extracted, - objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) - ); - } - log.info("[DependencyParserAgent] {} -> {} dependencies from {}", toolName, fileDeps.size(), filePath); - return markFileParsed(state, filePath, fileDeps); + if (AgentUtils.isNullOrEmpty(extracted)) return List.of(); + return objectMapper.readValue( + extracted, + objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) + ); } catch (McpToolNullResponseException e) { - return markFileParsed(state, filePath, List.of()); - } catch (Exception e) { - log.error("[DependencyParserAgent] Tool '{}' failed for {}: {}", toolName, filePath, e.getMessage(), e); - return markFileParsed(state, filePath, List.of()); - } - } - - @SuppressWarnings("unchecked") - private Map markFileParsed(SecurityAgentState state, String filePath, - List> newDeps) { - Map accumulated = new LinkedHashMap<>(); - state.dependencies().ifPresent(existing -> { - try { - List> list = objectMapper.readValue( - existing.toString(), - objectMapper.getTypeFactory().constructCollectionType(List.class, Map.class) - ); - list.forEach(d -> accumulated.put(buildDepKey(d), d)); - } catch (Exception ignored) {} - }); - newDeps.forEach(d -> accumulated.putIfAbsent(buildDepKey(d), d)); - - Set parsedFiles = new LinkedHashSet<>(extractParsedFilePaths(state)); - parsedFiles.add(filePath); - - Map updates = new HashMap<>(); - try { - updates.put("dependencies", objectMapper.writeValueAsString(new ArrayList<>(accumulated.values()))); - updates.put("parsedFiles", objectMapper.writeValueAsString(new ArrayList<>(parsedFiles))); + return List.of(); } catch (Exception e) { - log.error("[DependencyParserAgent] Failed to serialize state update", e); + log.error("[DependencyParserAgent] Tool '{}' failed for '{}': {}", toolName, filePath, e.getMessage(), e); + return List.of(); } - log.info("[DependencyParserAgent] Accumulated {} unique deps, {} files parsed", - accumulated.size(), parsedFiles.size()); - return updates; - } - - @SuppressWarnings("unchecked") - private Set extractParsedFilePaths(SecurityAgentState state) { - Object raw = state.data().get("parsedFiles"); - if (raw == null) return new LinkedHashSet<>(); - try { - return new LinkedHashSet<>(objectMapper.readValue( - raw.toString(), - objectMapper.getTypeFactory().constructCollectionType(List.class, String.class) - )); - } catch (Exception e) { return new LinkedHashSet<>(); } } @SuppressWarnings("unchecked") diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java index bd40b7a..f4e7cdb 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/agent/sub_agents/VulnerabilityAgent.java @@ -128,18 +128,80 @@ private Map reason(SecurityAgentState state) { return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); } - // Early-exit: empty dependency list for REPO/FILE scan - if ("REPO_SCAN".equals(intent) || "FILE_SCAN".equals(intent)) { + // Deterministic path for REPO_SCAN / FILE_SCAN — never let the LLM hallucinate deps + if (("REPO_SCAN".equals(intent) || "FILE_SCAN".equals(intent)) && !vulnsPopulated) { String depsRaw = state.dependencies().map(Object::toString).orElse(null); - if (depsRaw != null && AgentUtils.isNullOrEmpty(depsRaw.trim())) { + + // Empty / absent dependency list → nothing to scan + if (depsRaw == null || AgentUtils.isNullOrEmpty(depsRaw.trim())) { log.info("[VulnerabilityAgent] Empty dependency list — DONE (no LLM call)"); return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); } + + // Non-empty dependency list → build checkMultipleVulnerabilities args directly + // from the real parsed dependency data — no LLM call needed. + try { + com.fasterxml.jackson.databind.JsonNode depsNode = objectMapper.readTree(depsRaw); + if (depsNode.isArray() && depsNode.isEmpty()) { + log.info("[VulnerabilityAgent] Dependency list is empty array — DONE (no LLM call)"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); + } + + // Sanitize: keep only name/version/ecosystem; normalise version strings + com.fasterxml.jackson.databind.node.ArrayNode sanitized = + objectMapper.createArrayNode(); + for (com.fasterxml.jackson.databind.JsonNode dep : depsNode) { + String name = dep.path("name").asText("").trim(); + String version = dep.path("version").asText("").trim(); + // Use ecosystem from the parsed dep — never hardcode + String ecosystem = dep.path("ecosystem").asText("").trim(); + if (ecosystem.isBlank()) { + ecosystem = dep.path("type").asText("unknown").trim(); + } + + // Skip entries with blank names + if (name.isBlank()) continue; + + // Strip leading comparison operators from version strings (>=, ~=, ==, <=, !=, >) + // Keep only the first version token if multiple constraints are present (e.g. >=1.0,<2.0) + version = version.replaceAll("^[><=!~]+", "").trim(); + if (version.contains(",")) { + version = version.split(",")[0].replaceAll("^[><=!~]+", "").trim(); + } + // If version is still empty use "latest" + if (version.isBlank()) version = "latest"; + + com.fasterxml.jackson.databind.node.ObjectNode entry = + objectMapper.createObjectNode(); + entry.put("name", name); + entry.put("version", version); + entry.put("ecosystem", ecosystem.isBlank() ? "unknown" : ecosystem); + sanitized.add(entry); + } + + if (sanitized.isEmpty()) { + log.info("[VulnerabilityAgent] All dependencies were filtered out — DONE (no LLM call)"); + return Map.of("nextTool", DONE, "nextToolArgs", "{}", STATE_KEY_CALL_COUNT, callCount); + } + + String toolArgs = objectMapper.writeValueAsString( + Map.of("dependencies", objectMapper.treeToValue(sanitized, List.class))); + log.info("[VulnerabilityAgent] Deterministic dispatch — checkMultipleVulnerabilities with {} dep(s)", sanitized.size()); + return Map.of( + "nextTool", "checkMultipleVulnerabilities", + "nextToolArgs", toolArgs, + STATE_KEY_CALL_COUNT, callCount // no LLM call consumed + ); + } catch (Exception e) { + log.error("[VulnerabilityAgent] Failed to build deterministic tool args from deps: {}", e.getMessage()); + // Fall through to LLM path as a safety net + } } List availableTools = toolRegistry.getToolNames(); log.info("[VulnerabilityAgent] Tools: {}", availableTools); + // Build a concise deps summary for the LLM prompt (non-REPO/FILE paths) String depsStatus = state.dependencies().map(d -> { try { com.fasterxml.jackson.databind.JsonNode node = objectMapper.readTree(d.toString()); diff --git a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java index e64601d..202fa45 100644 --- a/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java +++ b/02-Cybersecurity/Vulnerability-Detection-Agent-Java/security-supervisor-agent/src/main/java/com/agent/model/ConversationContext.java @@ -7,6 +7,8 @@ import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.UserMessage; +import java.io.Serial; +import java.io.Serializable; import java.util.ArrayList; import java.util.List; @@ -28,7 +30,10 @@ */ @Data @RequiredArgsConstructor -public class ConversationContext { +public class ConversationContext implements Serializable { + + @Serial + private static final long serialVersionUID = 1L; private final String sessionId; private final List history; @@ -72,7 +77,10 @@ public boolean hasHistory() { */ @Data @AllArgsConstructor - public static class ChatMessage { + public static class ChatMessage implements Serializable { + + @Serial + private static final long serialVersionUID = 1L; /** Original user query for this turn */ private String userQuery;