patcher: small improvements to find_test logic (#912)

* patcher: fix type hints

* patcher: save custom test instructions in redis

* patcher: small improvement to patch_strategy prompt

common/examples/example-libpng/multi_pov_confirmed_vulnerability.txt

@@ -32,4 +32,4 @@ crashes: {
       }
       tracer_stacktrace: "+ FUZZER=libpng_read_fuzzer\n+ shift\n+ '[' '!' -v TESTCASE ']'\n+ TESTCASE=/testcase\n+ '[' '!' -f /testcase ']'\n+ export RUN_FUZZER_MODE=interactive\n+ RUN_FUZZER_MODE=interactive\n+ export FUZZING_ENGINE=libfuzzer\n+ FUZZING_ENGINE=libfuzzer\n+ export SKIP_SEED_CORPUS=1\n+ SKIP_SEED_CORPUS=1\n+ run_fuzzer libpng_read_fuzzer -runs=100 /testcase\nvm.mmap_rnd_bits = 28\n/out/libpng_read_fuzzer -rss_limit_mb=2560 -timeout=25 -runs=100 /testcase -dict=png.dict < /dev/null\nDictionary: 28 entries\nINFO: Running with entropic power schedule (0xFF, 100).\nINFO: Seed: 4229423848\nINFO: Loaded 1 modules   (6849 inline 8-bit counters): 6849 [0x55a9f2c1bfe8, 0x55a9f2c1daa9), \nINFO: Loaded 1 PC tables (6849 PCs): 6849 [0x55a9f2c1dab0,0x55a9f2c386c0), \n/out/libpng_read_fuzzer: Running 1 inputs 100 time(s) each.\nRunning: /testcase\n\x1b[1mpngrutil.c:1447:10:\x1b[1m\x1b[31m runtime error: \x1b[1m\x1b[0m\x1b[1mindex 41 out of bounds for type 'wpng_byte[max_keyword_wbytes]' (aka 'unsigned short[max_keyword_wbytes]')\x1b[1m\x1b[0m\n    #0 0x55a9f2b8bc17 in OSS_FUZZ_png_handle_iCCP /src/libpng/pngrutil.c:1447:10\n    #1 0x55a9f2b721fd in OSS_FUZZ_png_read_info /src/libpng/pngread.c:229:10\n    #2 0x55a9f2ad069f in LLVMFuzzerTestOneInput /src/libpng/contrib/oss-fuzz/libpng_read_fuzzer.cc:156:3\n    #3 0x55a9f2aede10 in fuzzer::Fuzzer::ExecuteCallback(unsigned char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerLoop.cpp:614:13\n    #4 0x55a9f2ad9085 in fuzzer::RunOneTest(fuzzer::Fuzzer*, char const*, unsigned long) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:327:6\n    #5 0x55a9f2adeb1f in fuzzer::FuzzerDriver(int*, char***, int (*)(unsigned char const*, unsigned long)) /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerDriver.cpp:862:9\n    #6 0x55a9f2b09dc2 in main /src/llvm-project/compiler-rt/lib/fuzzer/FuzzerMain.cpp:20:10\n    #7 0x7f34f05c1082 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24082) (BuildId: 0323ab4806bee6f846d9ad4bccfc29afdca49a58)\n    #8 0x55a9f2aa5c7d in _start (/out/libpng_read_fuzzer+0x54c7d)\n\nDEDUP_TOKEN: OSS_FUZZ_png_handle_iCCP--OSS_FUZZ_png_read_info--LLVMFuzzerTestOneInput\nSUMMARY: UndefinedBehaviorSanitizer: undefined-behavior pngrutil.c:1447:10 \nsubprocess command returned a non-zero exit status: 1\n"
   }
-submission_index: "my-vuln-id"
+internal_patch_id: "my-vuln-id"

patcher/src/buttercup/patcher/agents/context_retriever.py

@@ -15,6 +15,7 @@ from langchain_openai.chat_models.base import BaseChatOpenAI
 from concurrent.futures import as_completed, TimeoutError
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
+from redis import Redis
 from buttercup.common.stack_parsing import parse_stacktrace
 from langgraph.checkpoint.memory import InMemorySaver
 from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
@@ -31,7 +32,7 @@ from buttercup.common.challenge_task import ChallengeTask
 from buttercup.common.stack_parsing import CrashInfo
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.config import get_executor_for_config
-from buttercup.patcher.utils import truncate_output, get_challenge, TruncatePosition, get_codequery
+from buttercup.patcher.utils import truncate_output, get_challenge, TruncatePosition
 from buttercup.patcher.agents.common import (
     PatcherAgentBase,
     ContextRetrieverState,
@@ -63,6 +64,8 @@ from buttercup.patcher.agents.tools import (
 
 logger = logging.getLogger(__name__)
 
+CUSTOM_TEST_MAP_NAME = "custom_test_map"
+
 SYSTEM_TMPL = """You are an agent - please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved.
 If you are not sure about file content or codebase structure pertaining to the user's request, use your tools to read files and gather the relevant information: do NOT guess or make up an answer.
 You MUST plan extensively before each function call, and reflect extensively on the outcomes of the previous function calls. DO NOT do this entire process by making function calls only, as this can impair your ability to solve the problem and think insightfully.
@@ -687,6 +690,7 @@ You CANNOT stop here, you MUST fix the test instructions and call this tool agai
 class ContextRetrieverAgent(PatcherAgentBase):
     """Agent that retrieves code snippets from the project."""
 
+    redis: Redis | None = None
     agent: Runnable = field(init=False)
     llm: BaseChatOpenAI = field(init=False)
     cheap_llm: BaseChatOpenAI = field(init=False)
@@ -1180,13 +1184,20 @@ class ContextRetrieverAgent(PatcherAgentBase):
             goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
         )
 
-    def _get_custom_test_path(self, configuration: PatcherConfig) -> Path:
-        codequery = get_codequery(self.challenge.task_dir, configuration.work_dir)
-        return codequery.challenge.task_dir.joinpath("custom_test.sh")  # type: ignore[no-any-return]
+    def _get_custom_test_instructions(self) -> str | None:
+        if self.redis is None:
+            return None
+
+        return self.redis.hget(CUSTOM_TEST_MAP_NAME, self.challenge.task_meta.task_id)  # type: ignore[return-value]
+
+    def _save_custom_test_instructions(self, instructions: str) -> None:
+        if self.redis is None:
+            return
+        self.redis.hset(CUSTOM_TEST_MAP_NAME, self.challenge.task_meta.task_id, instructions)
 
     def find_tests_node(
         self, state: PatcherAgentState, config: RunnableConfig
-    ) -> Command[Literal[PatcherAgentName.ROOT_CAUSE_ANALYSIS.value, PatcherAgentName.REFLECTION.value]]:  # type: ignore[name-defined]
+    ) -> Command[Literal[PatcherAgentName.ROOT_CAUSE_ANALYSIS.value]]:  # type: ignore[name-defined]
         """Determine instructions to run tests in the challenge task."""
         configuration = PatcherConfig.from_configurable(config)
         logger.info(
@@ -1204,11 +1215,11 @@ class ContextRetrieverAgent(PatcherAgentBase):
                 goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
             )
 
-        custom_test_path = self._get_custom_test_path(configuration)
-        if custom_test_path and custom_test_path.exists() and custom_test_path.is_file():
+        custom_test_instructions = self._get_custom_test_instructions()
+        if custom_test_instructions:
             return Command(
                 update={
-                    "tests_instructions": custom_test_path.read_text(),
+                    "tests_instructions": custom_test_instructions,
                 },
                 goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
             )
@@ -1283,8 +1294,7 @@ class ContextRetrieverAgent(PatcherAgentBase):
                             self.challenge.task_meta.task_id,
                             self.challenge.name,
                         )
-                        if not custom_test_path.exists() or not custom_test_path.is_file():
-                            custom_test_path.write_text(agent_state.tests_instructions)
+                        self._save_custom_test_instructions(agent_state.tests_instructions)
 
                         return Command(
                             update={

patcher/src/buttercup/patcher/agents/leader.py

@@ -20,6 +20,7 @@ from buttercup.patcher.agents.reflection import ReflectionAgent
 from buttercup.patcher.agents.input_processing import InputProcessingAgent
 from buttercup.patcher.utils import PatchOutput
 from buttercup.common.llm import get_langfuse_callbacks
+from redis import Redis
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +31,7 @@ RECURSION_LIMIT = 200
 class PatcherLeaderAgent(PatcherAgentBase):
     """LLM-based Patcher Agent."""
 
+    redis: Redis | None
     work_dir: Path
     tasks_storage: Path
     model_name: str | None = None
@@ -38,7 +40,9 @@ class PatcherLeaderAgent(PatcherAgentBase):
         rootcause_agent = RootCauseAgent(self.challenge, self.input, chain_call=self.chain_call)
         swe_agent = SWEAgent(self.challenge, self.input, chain_call=self.chain_call)
         qe_agent = QEAgent(self.challenge, self.input, chain_call=self.chain_call)
-        context_retriever_agent = ContextRetrieverAgent(self.challenge, self.input, chain_call=self.chain_call)
+        context_retriever_agent = ContextRetrieverAgent(
+            self.challenge, self.input, chain_call=self.chain_call, redis=self.redis
+        )
         reflection_agent = ReflectionAgent(self.challenge, self.input, chain_call=self.chain_call)
         input_processing_agent = InputProcessingAgent(self.challenge, self.input, chain_call=self.chain_call)
         self.model_name = swe_agent.default_llm.model_name

patcher/src/buttercup/patcher/agents/rootcause.py

@@ -222,7 +222,9 @@ class RootCauseAgent(PatcherAgentBase):
             requests.append(CodeSnippetRequest(request=match.strip()))
         return requests
 
-    def analyze_vulnerability(self, state: PatcherAgentState) -> Command[Literal[PatcherAgentName.CREATE_PATCH.value]]:  # type: ignore[name-defined]
+    def analyze_vulnerability(
+        self, state: PatcherAgentState
+    ) -> Command[Literal[PatcherAgentName.PATCH_STRATEGY.value, PatcherAgentName.REFLECTION.value]]:  # type: ignore[name-defined]
         """Analyze the diff analysis and the code to understand the
         vulnerability in the current code."""
         logger.info(

patcher/src/buttercup/patcher/agents/swe.py

@@ -150,7 +150,10 @@ TOOLS:
 OUTPUT FORMAT (MANDATORY)
 <patch_development_process>
    a. List 2-4 alternative mitigation ideas, each with pros/cons.
-   b. Identify your selected approach and justify why it is the best trade-off. Prefer the most minimalistic approach, unless instructed otherwise.
+   b. Identify your selected approach and justify why it is the best trade-off. \
+Prefer the simplest (but sound) approaches first, if not instructed \
+otherwise. For example, you could try reverting the patch that introduced the \
+issue, if available, or validate the inputs before using them.
    c. Reference line numbers / function names from <code_snippets> as needed.
 </patch_development_process>
 <full_description>

patcher/src/buttercup/patcher/patcher.py

@@ -71,6 +71,7 @@ class Patcher:
         patcher_agent = PatcherLeaderAgent(
             ro_task,
             input,
+            redis=self.redis,
             chain_call=self._chain_call,
             work_dir=self.scratch_dir,
             tasks_storage=self.task_storage_dir,

patcher/tests/test_context_retriever.py

@@ -34,6 +34,8 @@ from buttercup.common.challenge_task import ChallengeTask, CommandResult
 from buttercup.common.task_meta import TaskMeta
 from langchain_core.messages import AIMessage
 from langchain_core.messages.tool import ToolCall
+from redis import Redis
+from buttercup.patcher.agents.context_retriever import CUSTOM_TEST_MAP_NAME
 
 
 original_subprocess_run = subprocess.run
@@ -2375,3 +2377,277 @@ def test_find_tests_parallel(
         assert result.update is not None
         assert result.update["tests_instructions"] == "#!/bin/bash\ncd /src2\nmake test2\n"
         assert result.goto == PatcherAgentName.ROOT_CAUSE_ANALYSIS.value
+
+
+def test_find_tests_redis_save_and_retrieve(
+    mock_agent: ContextRetrieverAgent,
+    mock_runnable_config: dict,
+    mock_challenge: ChallengeTask,
+) -> None:
+    """Test that find_tests_node correctly saves and retrieves test instructions from Redis."""
+    from unittest.mock import MagicMock
+
+    # Create a mock Redis instance
+    mock_redis = MagicMock(spec=Redis)
+    mock_agent.redis = mock_redis
+
+    # Test data
+    task_id = "test-task-id"
+    test_instructions = "#!/bin/bash\ncd /src\nmake test\n"
+
+    # Update the challenge task ID to match our test data
+    mock_challenge.task_meta.task_id = task_id
+
+    state = PatcherAgentState(
+        context=PatchInput(
+            challenge_task_dir=Path("/test/dir"),
+            task_id=task_id,
+            internal_patch_id="1",
+            povs=[
+                PatchInputPoV(
+                    challenge_task_dir=Path("/test/dir"),
+                    sanitizer="address",
+                    pov=Path("test.pov"),
+                    pov_token="test-token",
+                    sanitizer_output="test output",
+                    engine="libfuzzer",
+                    harness_name="test-harness",
+                )
+            ],
+        ),
+        relevant_code_snippets=set(),
+        execution_info={},
+    )
+
+    # Test 1: Verify that instructions are saved to Redis when found
+    # Mock the find_tests_agent to return successful test instructions
+    mock_find_tests_agent = MagicMock()
+    mock_find_tests_agent.invoke.return_value = AIMessage(
+        content="I found test instructions in the README.",
+        tool_calls=[
+            ToolCall(
+                id="test_instructions_call_1",
+                name="test_instructions",
+                args={
+                    "instructions": [
+                        "cd /src",
+                        "make test",
+                    ],
+                },
+            )
+        ],
+    )
+    # Mock the state with all required fields
+    mock_find_tests_agent.get_state.return_value.values = {
+        "tests_instructions": test_instructions,
+        "messages": [],  # Required field
+        "challenge_task_dir": Path("/test/dir"),  # Required field
+        "work_dir": mock_runnable_config["configurable"]["work_dir"],  # Required field
+    }
+    mock_agent.find_tests_agent = mock_find_tests_agent
+
+    # Mock the docker command execution for test_instructions
+    with (
+        patch("buttercup.common.challenge_task.ChallengeTask.exec_docker_cmd") as mock_exec,
+        patch("buttercup.common.challenge_task.ChallengeTask.get_clean_task") as mock_clean_task,
+        patch("buttercup.common.challenge_task.ChallengeTask.apply_patch_diff") as mock_apply_patch_diff,
+        patch("buttercup.common.challenge_task.ChallengeTask.get_oss_fuzz_path") as mock_get_oss_fuzz_path,
+        patch(
+            "buttercup.patcher.agents.context_retriever._are_test_instructions_valid"
+        ) as mock_are_test_instructions_valid,
+    ):
+
+        @contextmanager
+        def yield_challenge(*args, **kwargs):
+            yield mock_challenge
+
+        mock_clean_task.return_value = mock_challenge
+        mock_apply_patch_diff.return_value = True
+        mock_challenge.apply_patch_diff = MagicMock(return_value=True)
+        mock_challenge.get_rw_copy = MagicMock(side_effect=yield_challenge)
+        mock_get_oss_fuzz_path.return_value = Path("/test/oss-fuzz")
+        mock_exec.return_value = CommandResult(
+            success=True,
+            returncode=0,
+            output=b"Tests passed",
+            error=b"",
+        )
+        mock_are_test_instructions_valid.return_value = True
+
+        # Initially, Redis should not have any instructions for this task
+        mock_redis.hget.return_value = None
+
+        result = mock_agent.find_tests_node(state, mock_runnable_config)
+
+        # Verify that the agent found and validated test instructions
+        assert result.update["tests_instructions"] == test_instructions
+        assert result.goto == PatcherAgentName.ROOT_CAUSE_ANALYSIS.value
+
+        # Verify that the instructions were saved to Redis
+        mock_redis.hset.assert_called_once_with(CUSTOM_TEST_MAP_NAME, task_id, test_instructions)
+
+    # Test 2: Verify that instructions are retrieved from Redis when available
+    # Reset the mock to clear previous calls
+    mock_redis.reset_mock()
+
+    # Mock Redis to return existing instructions
+    mock_redis.hget.return_value = test_instructions
+
+    # Mock that test.sh file doesn't exist (so it goes to Redis lookup)
+    with (
+        patch("pathlib.Path.exists", return_value=False),
+        patch("buttercup.common.challenge_task.ChallengeTask.get_oss_fuzz_path") as mock_get_oss_fuzz_path,
+    ):
+        mock_get_oss_fuzz_path.return_value = Path("/test/oss-fuzz")
+        result = mock_agent.find_tests_node(state, mock_runnable_config)
+
+        # Verify that the agent retrieved instructions from Redis
+        assert result.update["tests_instructions"] == test_instructions
+        assert result.goto == PatcherAgentName.ROOT_CAUSE_ANALYSIS.value
+
+        # Verify that Redis was queried for the instructions
+        mock_redis.hget.assert_called_once_with(CUSTOM_TEST_MAP_NAME, task_id)
+
+        # Verify that hset was NOT called (since we retrieved from Redis)
+        mock_redis.hset.assert_not_called()
+
+    # Test 3: Verify behavior when Redis is None
+    mock_agent.redis = None
+
+    # Mock that test.sh file doesn't exist
+    with (
+        patch("pathlib.Path.exists", return_value=False),
+        patch("buttercup.common.challenge_task.ChallengeTask.get_oss_fuzz_path") as mock_get_oss_fuzz_path,
+        patch("buttercup.common.challenge_task.ChallengeTask.get_clean_task") as mock_get_clean_task,
+    ):
+        mock_get_oss_fuzz_path.return_value = Path("/test/oss-fuzz")
+        mock_get_clean_task.return_value = mock_challenge
+
+        # Mock the find_tests_agent to avoid actual execution
+        mock_find_tests_agent = MagicMock()
+        mock_find_tests_agent.invoke.return_value = AIMessage(content="Test")
+        mock_find_tests_agent.get_state.return_value.values = {
+            "tests_instructions": None,
+            "messages": [],
+            "challenge_task_dir": Path("/test/dir"),
+            "work_dir": mock_runnable_config["configurable"]["work_dir"],
+        }
+        mock_agent.find_tests_agent = mock_find_tests_agent
+
+        result = mock_agent.find_tests_node(state, mock_runnable_config)
+
+        # Should proceed to find tests agent since Redis is None
+        # The result should be the same as when Redis was available but empty
+        assert result.goto == PatcherAgentName.ROOT_CAUSE_ANALYSIS.value
+
+    # Test 4: Verify that _get_custom_test_instructions returns None when Redis is None
+    assert mock_agent._get_custom_test_instructions() is None
+
+    # Test 5: Verify that _save_custom_test_instructions does nothing when Redis is None
+    # This should not raise any exceptions
+    mock_agent._save_custom_test_instructions("some instructions")
+
+    # Test 6: Verify Redis operations with actual Redis instance
+    mock_agent.redis = mock_redis
+
+    # Test _get_custom_test_instructions
+    mock_redis.hget.return_value = "cached_instructions"
+    assert mock_agent._get_custom_test_instructions() == "cached_instructions"
+    mock_redis.hget.assert_called_with(CUSTOM_TEST_MAP_NAME, task_id)
+
+    # Test _save_custom_test_instructions
+    mock_agent._save_custom_test_instructions("new_instructions")
+    mock_redis.hset.assert_called_with(CUSTOM_TEST_MAP_NAME, task_id, "new_instructions")
+
+
+def test_find_tests_redis_multiple_tasks(
+    mock_agent: ContextRetrieverAgent,
+    mock_runnable_config: dict,
+    mock_challenge: ChallengeTask,
+) -> None:
+    """Test that Redis correctly handles multiple tasks with different test instructions."""
+    from unittest.mock import MagicMock
+
+    # Create a mock Redis instance
+    mock_redis = MagicMock(spec=Redis)
+    mock_agent.redis = mock_redis
+
+    # Test data for multiple tasks
+    task_id_1 = "test-task-1"
+    task_id_2 = "test-task-2"
+    instructions_1 = "#!/bin/bash\ncd /src\nmake test\n"
+    instructions_2 = "#!/bin/bash\ncd /src\n./run_tests.sh\n"
+
+    # Test that different tasks get different instructions
+    mock_redis.hget.side_effect = lambda map_name, task_id: {
+        task_id_1: instructions_1,
+        task_id_2: instructions_2,
+    }.get(task_id, None)
+
+    # Test task 1
+    mock_challenge.task_meta.task_id = task_id_1
+    state_1 = PatcherAgentState(
+        context=PatchInput(
+            challenge_task_dir=Path("/test/dir"),
+            task_id=task_id_1,
+            internal_patch_id="1",
+            povs=[
+                PatchInputPoV(
+                    challenge_task_dir=Path("/test/dir"),
+                    sanitizer="address",
+                    pov=Path("test.pov"),
+                    pov_token="test-token",
+                    sanitizer_output="test output",
+                    engine="libfuzzer",
+                    harness_name="test-harness",
+                )
+            ],
+        ),
+        relevant_code_snippets=set(),
+        execution_info={},
+    )
+
+    with (
+        patch("pathlib.Path.exists", return_value=False),
+        patch("buttercup.common.challenge_task.ChallengeTask.get_oss_fuzz_path") as mock_get_oss_fuzz_path,
+    ):
+        mock_get_oss_fuzz_path.return_value = Path("/test/oss-fuzz")
+        result_1 = mock_agent.find_tests_node(state_1, mock_runnable_config)
+        assert result_1.update["tests_instructions"] == instructions_1
+
+    # Test task 2
+    mock_challenge.task_meta.task_id = task_id_2
+    state_2 = PatcherAgentState(
+        context=PatchInput(
+            challenge_task_dir=Path("/test/dir"),
+            task_id=task_id_2,
+            internal_patch_id="1",
+            povs=[
+                PatchInputPoV(
+                    challenge_task_dir=Path("/test/dir"),
+                    sanitizer="address",
+                    pov=Path("test.pov"),
+                    pov_token="test-token",
+                    sanitizer_output="test output",
+                    engine="libfuzzer",
+                    harness_name="test-harness",
+                )
+            ],
+        ),
+        relevant_code_snippets=set(),
+        execution_info={},
+    )
+
+    with (
+        patch("pathlib.Path.exists", return_value=False),
+        patch("buttercup.common.challenge_task.ChallengeTask.get_oss_fuzz_path") as mock_get_oss_fuzz_path,
+    ):
+        mock_get_oss_fuzz_path.return_value = Path("/test/oss-fuzz")
+        result_2 = mock_agent.find_tests_node(state_2, mock_runnable_config)
+        assert result_2.update["tests_instructions"] == instructions_2
+
+    # Verify that Redis was called with the correct task IDs
+    assert mock_redis.hget.call_count == 2
+    calls = mock_redis.hget.call_args_list
+    assert calls[0] == ((CUSTOM_TEST_MAP_NAME, task_id_1),)
+    assert calls[1] == ((CUSTOM_TEST_MAP_NAME, task_id_2),)