Fixing and customizing the enhancer prompt (#865)

In enhancer agent:
1. Teach LLM to write fuzz targets and build scripts in the same way as prototyper. Before this, it writes them within \`\`\`.
2. Make enhancer prompt independent and customizable.
3. More Flexible agent names in logs.

Author: DonggeLiu

agent/enhancer.py

@@ -16,10 +16,9 @@
 """
 import logger
 from agent.prototyper import Prototyper
-#from experiment.workdir import WorkDirs
-from llm_toolkit.prompt_builder import DefaultTemplateBuilder, JvmFixingBuilder
+from llm_toolkit.prompt_builder import EnhancerTemplateBuilder, JvmFixingBuilder
 from llm_toolkit.prompts import Prompt
-from results import AnalysisResult, Result
+from results import AnalysisResult, BuildResult, Result
 
 
 class Enhancer(Prototyper):
@@ -36,20 +35,29 @@ def _initial_prompt(self, results: list[Result]) -> Prompt:
                    trial=self.trial)
       return Prompt()
 
+    last_build_result = None
+    for result in results[::-1]:
+      if isinstance(result, BuildResult):
+        last_build_result = result
+        break
+    if not last_build_result:
+      logger.error('Unable to find the last build result in Enhancer : %s',
+                   results,
+                   trial=self.trial)
+      return Prompt()
+
     if benchmark.language == 'jvm':
       # TODO: Do this in a separate agent for JVM coverage.
       builder = JvmFixingBuilder(self.llm, benchmark,
                                  last_result.run_result.fuzz_target_source, [])
       prompt = builder.build([], None, None)
     else:
       error_desc, errors = last_result.semantic_result.get_error_info()
-      builder = DefaultTemplateBuilder(self.llm)
-      prompt = builder.build_fixer_prompt(benchmark,
-                                          last_result.fuzz_target_source,
-                                          error_desc,
-                                          errors,
-                                          context='',
-                                          instruction='')
+      builder = EnhancerTemplateBuilder(self.llm, benchmark, last_build_result,
+                                        error_desc, errors)
+      prompt = builder.build(example_pair=[],
+                             tool_guides=self.inspect_tool.tutorial(),
+                             project_dir=self.inspect_tool.project_dir)
       # TODO: A different file name/dir.
       prompt.save(self.args.work_dirs.prompt)
 

agent/prototyper.py

@@ -217,15 +217,17 @@ def _generate_prompt_from_build_result(
           'Default /src/build.sh works perfectly, no need for a new '
           'buid script',
           trial=build_result.trial)
-      logger.info('***** Prototyper succeded in %02d rounds *****',
+      logger.info('***** %s succeeded in %02d rounds *****',
+                  self.name,
                   cur_round,
                   trial=build_result.trial)
       return build_result_alt, None
 
     if build_result_ori and build_result_ori.success:
       # Preference 2: New fuzz target + new build.sh can compile, save
       # binary to expected path, and reference function-under-test.
-      logger.info('***** Prototyper succeded in %02d rounds *****',
+      logger.info('***** %s succeeded in %02d rounds *****',
+                  self.name,
                   cur_round,
                   trial=build_result.trial)
       return build_result_ori, None

llm_toolkit/prompt_builder.py

@@ -643,6 +643,59 @@ def build(self,
     return self._prompt
 
 
+class EnhancerTemplateBuilder(PrototyperTemplateBuilder):
+  """Builder specifically targeted C (and excluding C++)."""
+
+  def __init__(self,
+               model: models.LLM,
+               benchmark: Benchmark,
+               build_result: BuildResult,
+               error_desc: str,
+               errors: list[str],
+               template_dir: str = DEFAULT_TEMPLATE_DIR,
+               initial: Any = None):
+    super().__init__(model, benchmark, template_dir, initial)
+    # Load templates.
+    self.priming_template_file = self._find_template(self.agent_templare_dir,
+                                                     'enhancer-priming.txt')
+    self.build_result = build_result
+    self.error_desc = error_desc
+    self.errors = errors
+
+  def build(self,
+            example_pair: list[list[str]],
+            project_example_content: Optional[list[list[str]]] = None,
+            project_context_content: Optional[dict] = None,
+            tool_guides: str = '',
+            project_dir: str = '') -> prompts.Prompt:
+    """Constructs a prompt using the templates in |self| and saves it."""
+    del (example_pair, project_example_content, project_context_content)
+    if not self.benchmark:
+      return self._prompt
+
+    priming = self._get_template(self.priming_template_file)
+    priming = priming.replace('{LANGUAGE}', self.benchmark.file_type.value)
+    priming = priming.replace('{FUNCTION_SIGNATURE}',
+                              self.benchmark.function_signature)
+    # TODO(dongge): Add build script to .
+    priming = priming.replace('{PROJECT_DIR}', project_dir)
+    if self.build_result.build_script_source:
+      build_text = (f'<build script>\n{self.build_result.build_script_source}\n'
+                    '</build script>')
+    else:
+      build_text = 'Build script reuses `/src/build.bk.sh`.'
+    priming = priming.replace('{BUILD_TEXT}', build_text)
+    priming = priming.replace('{TOOL_GUIDES}', tool_guides)
+    priming_weight = self._model.estimate_token_num(priming)
+
+    problem = self._format_fixer_problem(self.build_result.fuzz_target_source,
+                                         self.error_desc, self.errors,
+                                         priming_weight, '', '')
+
+    self._prepare_prompt(priming, problem)
+    return self._prompt
+
+
 class DefaultJvmTemplateBuilder(PromptBuilder):
   """Default builder for JVM projects."""
 

prompts/agent/enhancer-priming.txt

@@ -0,0 +1,94 @@
+<system>
+As a security testing engineer, you must refine the following compilable {LANGUAGE} fuzz target to make it more suitable for fuzzing function {FUNCTION_SIGNATURE}, based on its current runtime errors description.
+Objective: Your task is to understand the runtime error, and refine the fuzz target (and build script if needed) accordingly. Note that the fuzz target can already compile.
+</system>
+
+<steps>
+Follow these steps to refine the fuzz target:
+
+Step 1. Determine the information you need to understand the runtime error of the fuzz target.
+This includes:
+* The existing compilable fuzz target provided below.
+* The existing build script provided below.
+* **Source code** of the function under test.
+* **Custom Types and Dependencies** definitions and implementations.
+* **Initialization and setup** requirements and steps.
+* **Build details** and integration steps.
+* Valid and edge-case input values.
+* Environmental and runtime dependencies.
+
+Step 2. Collect information using the Bash tool.
+Use the bash tool (see <tool> section) and follow its rules to gather the necessary information. You can collect information from:
+* The project source code directory `{PROJECT_DIR}/` cloned from the project repository.
+* Documentation about the project, the function, and the variables/constants involved.
+* Environment variables.
+* Knowledge about OSS-Fuzz's build infrastructure: It will compile your fuzz target in the same way as the exiting human written fuzz target with the build script.
+
+Step 3. Analyze the function and its parameters.
+Understand the function under test by analyzing its source code and documentation:
+* **Purpose and functionality** of the function.
+* **Input processing** and internal logic.
+* **Dependencies** on other functions or global variables.
+* **Error handling** and edge cases.
+
+Step 4. Understand initialization requirements.
+Identify what is needed to properly initialize the function:
+* **Header files** and their relative paths used by include statements in the fuzz target.
+* **Complex input parameters or objects** initialization.
+* **Constructor functions** or initialization routines.
+* **Global state** or configuration needs to be set up.
+* **Mocking** external dependencies if necessary.
+
+Step 5. Understand Constraints and edge cases.
+For each input parameter, understand:
+* Valid ranges and data types.
+* Invalid or edge-case values (e.g., zero, NULL, predefined constants, maximum values).
+* Special values that trigger different code paths.
+
+Step 6: Plan Fuzz Target Implementation.
+Decide how to implement the refined fuzz target:
+* The fuzz target can compile so your can reuse most of the code as a scaffold.
+* Only modify the parts caused the runtime error, no more no less.
+* Prepare to output the FULL new fuzz target, do not leave out any code that is the same as before.
+* **Extract parameters** from the `data` and `size` variable of `LLVMFuzzerTestOneInput(const uint8_t *data, size_t size)`.
+* Handle fixed-size versus variable-size data.
+* **Initialize function's parameters** by appropriately mapping the raw input bytes.
+* Ensure that the fuzz target remains deterministic and avoids side effects.
+* Avoid `goto` statements.
+
+*
+Step 7 (Optional): **Modify** the Build Script.
+Modify the build script only if the existing one in this prompt is insufficient:
+* Decide if you need to modify the build script to successfully build the refined fuzz target.
+* If the build script needs to be modified, prepare to output the FULL new build script, do not leave out any code that is the same as before.
+* Leave it empty if no modification is needed.
+
+Step 9: Providing Your Conclusion:
+* Provide your conclusion on the FULL new fuzz target and build script **ONLY AFTER** you have gathered all necessary information.
+* **DO NOT SEND** any other content (e.g., bash tool commands) in the conclusion message. ALWAYS send other commands individually and ONLY SEND conclusion after collecting all information.
+* Conclusion Format:
+* Overall Description:
+* Summarize the error, the root cause your found, and describe your fuzz target refinement.
+* Wrap this summary within <conclusion> and </conclusion> tags.
+* Modified Fuzz Target:
+* Provide the full code of the refined fuzz target.
+* Wrap the code within <fuzz target> and </fuzz target> tags.
+* Modified Build Script (if applicable):
+* If you need to modify the build script, provide the full code.
+* Wrap it within <build script> and </build script> tags.
+* Format Example:
+<conclusion>
+The fuzz target has runtime error ___, which is caused by ___.
+I will refined it by ___.
+Additionally, the build script requires modification to link against the necessary libraries.
+</conclusion>
+<fuzz target>
+[Your FULL fuzz target code here, do not omit existing code]
+</fuzz target>
+<build script>
+[Your FULL build script code here, do not omit existing code.]
+</build script>
+
+</steps>
+
+{TOOL_GUIDES}