From 7a50b234a62868c6a9bbf414dd2acc2d210533fc Mon Sep 17 00:00:00 2001 From: Riccardo Schirone <562321+ret2libc@users.noreply.github.com> Date: Sun, 22 Jun 2025 16:40:57 +0200 Subject: [PATCH] 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 --- .../multi_pov_confirmed_vulnerability.txt | 2 +- .../patcher/agents/context_retriever.py | 30 +- .../src/buttercup/patcher/agents/leader.py | 6 +- .../src/buttercup/patcher/agents/rootcause.py | 4 +- patcher/src/buttercup/patcher/agents/swe.py | 5 +- patcher/src/buttercup/patcher/patcher.py | 1 + patcher/tests/test_context_retriever.py | 276 ++++++++++++++++++ 7 files changed, 310 insertions(+), 14 deletions(-) diff --git a/common/examples/example-libpng/multi_pov_confirmed_vulnerability.txt b/common/examples/example-libpng/multi_pov_confirmed_vulnerability.txt index 71d02fdb..3b210fa8 100644 --- a/common/examples/example-libpng/multi_pov_confirmed_vulnerability.txt +++ b/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" diff --git a/patcher/src/buttercup/patcher/agents/context_retriever.py b/patcher/src/buttercup/patcher/agents/context_retriever.py index cd6bd1f5..6bbfef01 100644 --- a/patcher/src/buttercup/patcher/agents/context_retriever.py +++ b/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={ diff --git a/patcher/src/buttercup/patcher/agents/leader.py b/patcher/src/buttercup/patcher/agents/leader.py index eb8fceb9..064c684d 100644 --- a/patcher/src/buttercup/patcher/agents/leader.py +++ b/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 diff --git a/patcher/src/buttercup/patcher/agents/rootcause.py b/patcher/src/buttercup/patcher/agents/rootcause.py index 0f6950f8..78c222d3 100644 --- a/patcher/src/buttercup/patcher/agents/rootcause.py +++ b/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( diff --git a/patcher/src/buttercup/patcher/agents/swe.py b/patcher/src/buttercup/patcher/agents/swe.py index ddbfdca7..634a0924 100644 --- a/patcher/src/buttercup/patcher/agents/swe.py +++ b/patcher/src/buttercup/patcher/agents/swe.py @@ -150,7 +150,10 @@ TOOLS: OUTPUT FORMAT (MANDATORY) 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 as needed. diff --git a/patcher/src/buttercup/patcher/patcher.py b/patcher/src/buttercup/patcher/patcher.py index 106398b3..fabe3af4 100644 --- a/patcher/src/buttercup/patcher/patcher.py +++ b/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, diff --git a/patcher/tests/test_context_retriever.py b/patcher/tests/test_context_retriever.py index 5e2c4c75..7b191973 100644 --- a/patcher/tests/test_context_retriever.py +++ b/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),)