Professional

Best practices for writing AI extensions

• Last updated: February 13, 2025

• Read time: 3 Minutes

The Montoya API enables you to integrate AI-powered functionality into your Burp Suite extensions. Follow these best practices to ensure your extension is efficient and user-friendly.

Declare AI support and check availability

For your extension to use AI features, you must:

• Explicitly declare AI support for the extension. To do so, override the enhancedCapabilities() method of the BurpExtension interface and return the EnhancedCapability.AI_FEATURES flag.

• Check AI functionality is available in the current instance of Burp Suite using the ai.isEnabled() method.

This ensures that the Use AI checkbox appears for users and that your extension handles scenarios where AI features are unavailable.

Minimize data sent to the AI

Sending unnecessary data increases response time and consumes more AI credits.

To mitigate this:

• Only send essential data. For example, avoid including full HTTP requests if only headers or parameters are needed, and only send traffic that is in-scope for the application you are testing.

• Consider implementing a cache to optimize credit usage and improve response times. For example, you can hash the temperature, system message, and prompt to check whether a request has already been made with the same parameters, and serve the cached response instead of sending a new prompt.

Use effective prompts

A well-structured prompt ensures that the AI provides relevant and high-quality responses. We recommend that you:

• Provide background information to improve accuracy.

• Define the AI's role at the start of the conversation.

• Clearly specify what you want the AI to do.

Use lower temperatures for better accuracy

The temperature setting controls the balance between determinism and creativity in AI responses. In general, lower values produce focused and accurate results. The default is 0.5. Conversely, higher values may cause the AI to go off on tangents, potentially making security-related responses unreliable.

Handle exceptions gracefully

AI calls can fail for various reasons, such as the user not having enough credits or service downtime. Handle exceptions gracefully to ensure your extension is able to continue functioning as expected for your users.

The PromptException class represents errors that may occur during AI prompt execution. It is thrown if there is an issue with the AI request. Wrap AI calls in a try-catch block to handle errors appropriately.

Use an executor service where necessary to avoid blocking threads

AI requests can take longer than typical Burp API calls. Running them in the Swing Event Dispatch Thread can cause Burp to appear unresponsive, as the whole GUI must wait until the slow operation completes.

To keep Burp's UI responsive, execute AI calls asynchronously using an executor service. This prevents long-running operations from locking the Swing thread.