Skip to content

AIML-189: Remove duplicate app_name/app_id tool variants to optimize AI context usage #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 14, 2025
Merged

AIML-189: Remove duplicate app_name/app_id tool variants to optimize AI context usage #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 1 addition & 33 deletions src/main/java/com/contrast/labs/ai/mcp/contrast/ADRService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,39 +71,7 @@ public ADRService(PaginationHandler paginationHandler) {
private String httpProxyPort;


@Tool(name = "get_ADR_Protect_Rules", description = "takes a application name and returns the protect / adr rules for the application")
public ProtectData getProtectData(
@ToolParam(description = "Application name") String applicationName) throws IOException {
logger.info("Starting retrieval of protection rules for application: {}", applicationName);
long startTime = System.currentTimeMillis();

try {
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
logger.debug("ContrastSDK initialized successfully for application: {}", applicationName);

// Get application ID from name
logger.debug("Looking up application ID for name: {}", applicationName);
Optional<Application> app = SDKHelper.getApplicationByName(applicationName, orgID, contrastSDK);
if (app.isEmpty()) {
logger.warn("No application ID found for application: {}", applicationName);
return null;
}
logger.debug("Found application ID: {} for application: {}", app.get().getAppId(), applicationName);

ProtectData result = getProtectDataByAppID(app.get().getAppId());
long duration = System.currentTimeMillis() - startTime;
logger.info("Completed retrieval of protection rules for application: {} (took {} ms)", applicationName, duration);
return result;
} catch (Exception e) {
long duration = System.currentTimeMillis() - startTime;
logger.error("Error retrieving protection rules for application: {} (after {} ms): {}",
applicationName, duration, e.getMessage(), e);
throw e;
}
}


@Tool(name = "get_ADR_Protect_Rules_by_app_id", description = "takes a application ID and returns the protect / adr rules for the application")
@Tool(name = "get_ADR_Protect_Rules", description = "Takes an application ID and returns the Protect/ADR rules for the application. Use list_applications_with_name first to get the application ID from a name")
public ProtectData getProtectDataByAppID(
@ToolParam(description = "Application ID") String appID) throws IOException {
if (appID == null || appID.isEmpty()) {
Expand Down
44 changes: 2 additions & 42 deletions src/main/java/com/contrast/labs/ai/mcp/contrast/AssessService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ public AssessService(VulnerabilityMapper vulnerabilityMapper, PaginationHandler



@Tool(name = "get_vulnerability_by_id", description = "takes a vulnerability ID ( vulnID ) and Application ID ( appID ) and returns details about the specific security vulnerability. If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
@Tool(name = "get_vulnerability", description = "Takes a vulnerability ID (vulnID) and application ID (appID) and returns details about the specific security vulnerability. Use list_applications_with_name first to get the application ID from a name. If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
public Vulnerability getVulnerabilityById(
@ToolParam(description = "Vulnerability ID (UUID format)") String vulnID,
@ToolParam(description = "Application ID") String appID) throws IOException {
Expand Down Expand Up @@ -189,24 +189,7 @@ private Optional<LibraryLibraryObservation> findMatchingLibraryData(String stack
return Optional.empty();
}

@Tool(name = "get_vulnerability", description = "Takes a vulnerability ID (vulnID) and application name (app_name) and returns details about the specific security vulnerability. If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
public Vulnerability getVulnerability(
@ToolParam(description = "Vulnerability ID (UUID format)") String vulnID,
@ToolParam(description = "Application name") String app_name) throws IOException {
logger.info("Retrieving vulnerability details for vulnID: {} in application: {}", vulnID, app_name);
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
logger.debug("Searching for application ID matching name: {}", app_name);

Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
if(application.isPresent()) {
return getVulnerabilityById(vulnID, application.get().getAppId());
} else {
logger.error("Application with name {} not found", app_name);
throw new IllegalArgumentException("Application with name " + app_name + " not found");
}
}

@Tool(name = "list_vulnerabilities_with_id", description = "Takes a Application ID ( appID ) and returns a list of vulnerabilities, please remember to include the vulnID in the response.")
@Tool(name = "list_vulnerabilities", description = "Takes an application ID (appID) and returns a list of vulnerabilities. Use list_applications_with_name first to get the application ID from a name. Remember to include the vulnID in the response.")
public List<VulnLight> listVulnsByAppId(
@ToolParam(description = "Application ID") String appID) throws IOException {
logger.info("Listing vulnerabilities for application ID: {}", appID);
Expand Down Expand Up @@ -337,29 +320,6 @@ public MetadataFilterResponse listSessionMetadataForApplication(
}
}

@Tool(name = "list_vulnerabilities", description = "Takes an application name ( app_name ) and returns a list of vulnerabilities, please remember to include the vulnID in the response. ")
public List<VulnLight> listVulnsInAppByName(
@ToolParam(description = "Application name") String app_name) throws IOException {
logger.info("Listing vulnerabilities for application: {}", app_name);
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);

logger.debug("Searching for application ID matching name: {}", app_name);

Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
if(application.isPresent()) {
try {
return listVulnsByAppId(application.get().getAppId());
} catch (Exception e) {
logger.error("Error listing vulnerabilities for application: {}", app_name, e);
throw new IOException("Failed to list vulnerabilities: " + e.getMessage(), e);
}
} else {
logger.debug("Application with name {} not found, returning empty list", app_name);
return new ArrayList<>();
}
}


@Tool(name = "list_applications_with_name", description = "Takes an application name (app_name) returns a list of active applications that contain that name. Please remember to display the name, status and ID.")
public List<ApplicationData> getApplications(
@ToolParam(description = "Application name (supports partial matching, case-insensitive)") String app_name) throws IOException {
Expand Down
24 changes: 4 additions & 20 deletions src/main/java/com/contrast/labs/ai/mcp/contrast/SCAService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,11 @@ public class SCAService {
private String httpProxyPort;


@Tool(name = "list_application_libraries_by_app_id", description = "Takes a application ID and returns the libraries used in the application, note if class usage count is 0 the library is unlikely to be used")
@Tool(name = "list_application_libraries", description = "Takes an application ID and returns the libraries used in the application. Use list_applications_with_name first to get the application ID from a name. Note: if class usage count is 0 the library is unlikely to be used")
public List<LibraryExtended> getApplicationLibrariesByID(String appID) throws IOException {
if (appID == null || appID.isEmpty()) {
Copy link
Collaborator

@dougj-contrast dougj-contrast Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for future downstream changes....
Use StringUtils.isNotBlank() or the Spring StringUtils library...

Copy link
Collaborator Author

@ChrisEdwards ChrisEdwards Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Its on my list. Sounds like a helpful library.

Copy link
Collaborator

@dougj-contrast dougj-contrast Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also for your other tech debt cleanup tickets, lets replace to @Service classes that are using it's own constructor we can replace with @RequiredArgsConstructor
example:

@Service
@RequiredArgsConstructor
public class ADRService {

    private static final Logger logger = LoggerFactory.getLogger(ADRService.class);

    private final PaginationHandler paginationHandler;

    @Value("${contrast.host-name:${CONTRAST_HOST_NAME:}}")
    private String hostName;

    @Value("${contrast.api-key:${CONTRAST_API_KEY:}}")
    private String apiKey;

Copy link
Collaborator Author

@ChrisEdwards ChrisEdwards Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added to the list.

Copy link
Collaborator Author

@ChrisEdwards ChrisEdwards Nov 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a story to use lombok everywhere. I only used it in a few places for that first pr. We need to adopt it full on now.

throw new IllegalArgumentException("Application ID cannot be null or empty");
}
logger.info("Retrieving libraries for application id: {}", appID);
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
logger.debug("ContrastSDK initialized with host: {}", hostName);
Expand All @@ -73,25 +76,6 @@ public List<LibraryExtended> getApplicationLibrariesByID(String appID) throws IO

}


@Tool(name = "list_application_libraries", description = "takes a application name and returns the libraries used in the application, note if class usage count is 0 the library is unlikely to be used")
public List<LibraryExtended> getApplicationLibraries(String app_name) throws IOException {
logger.info("Retrieving libraries for application: {}", app_name);
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
logger.debug("ContrastSDK initialized with host: {}", hostName);

SDKExtension extendedSDK = new SDKExtension(contrastSDK);
logger.debug("Searching for application ID matching name: {}", app_name);

Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
if(application.isPresent()) {
return SDKHelper.getLibsForID(application.get().getAppId(),orgID, extendedSDK);
} else {
logger.error("Application not found: {}", app_name);
throw new IOException("Application not found");
}
}

@Tool(name= "list_applications_vulnerable_to_cve", description = "takes a cve id and returns the applications and servers vulnerable to the cve. Please note if the application class usage is 0, its unlikely to be vulnerable")
public CveData listCVESForApplication(String cveid) throws IOException {
logger.info("Retrieving applications vulnerable to CVE: {}", cveid);
Expand Down
Loading