diff --git a/services/claude/chat_with_claude.py b/services/claude/chat_with_claude.py index a7d08c6e..56797f36 100644 --- a/services/claude/chat_with_claude.py +++ b/services/claude/chat_with_claude.py @@ -122,7 +122,7 @@ def chat_with_claude( # Combine system message with user messages for logging system_msg: MessageParam = {"role": "user", "content": system_content} - full_messages = [system_msg] + list(messages) + full_messages = [system_msg, *messages] insert_llm_request( usage_id=usage_id, provider="claude", diff --git a/services/claude/remove_outdated_apply_diff_to_file_attempts_and_results.py b/services/claude/remove_outdated_apply_diff_to_file_attempts_and_results.py index ad5b6a48..e8a0dd4c 100644 --- a/services/claude/remove_outdated_apply_diff_to_file_attempts_and_results.py +++ b/services/claude/remove_outdated_apply_diff_to_file_attempts_and_results.py @@ -1,6 +1,3 @@ -# Standard imports -from copy import deepcopy - # Third party imports from anthropic.types import MessageParam @@ -15,9 +12,6 @@ def remove_outdated_apply_diff_to_file_attempts_and_results( if not messages: return messages - # Make a deep copy to avoid modifying the original - result = deepcopy(messages) - # Track latest diff results by filename - store position and type (failed/successful/input) file_latest_positions = {} @@ -189,6 +183,6 @@ def remove_outdated_apply_diff_to_file_attempts_and_results( new_content.append(new_item) if new_content != msg["content"]: - result[i]["content"] = new_content + messages[i]["content"] = new_content - return result + return messages diff --git a/services/claude/test_remove_outdated_apply_diff_to_file_attempts_and_results.py b/services/claude/test_remove_outdated_apply_diff_to_file_attempts_and_results.py index de3a0f6d..1a940cea 100644 --- a/services/claude/test_remove_outdated_apply_diff_to_file_attempts_and_results.py +++ b/services/claude/test_remove_outdated_apply_diff_to_file_attempts_and_results.py @@ -1161,13 +1161,16 @@ def test_remove_outdated_apply_diff_to_file_attempts_and_results_interleaved_fil def test_handles_exception_gracefully(): - # Test by mocking deepcopy to raise an exception + # Test by mocking enumerate to raise an exception during first pass messages: list[MessageParam] = cast( list[MessageParam], [{"role": "user", "content": []}] ) - with patch("copy.deepcopy", side_effect=RuntimeError("Simulated failure")): + with patch( + "services.claude.remove_outdated_apply_diff_to_file_attempts_and_results.enumerate", + side_effect=RuntimeError("Simulated failure"), + ): result = remove_outdated_apply_diff_to_file_attempts_and_results(messages) # Should return original messages unchanged when exception occurs assert result == messages diff --git a/services/webhook/check_suite_handler.py b/services/webhook/check_suite_handler.py index f0b79b62..5e73b1c7 100644 --- a/services/webhook/check_suite_handler.py +++ b/services/webhook/check_suite_handler.py @@ -69,6 +69,7 @@ from utils.logging.add_log_message import add_log_message from utils.logging.logging_config import logger, set_pr_number, set_trigger from utils.logs.clean_logs import clean_logs +from utils.memory.gc_collect_and_log import gc_collect_and_log from utils.logs.detect_infra_failure import detect_infra_failure from utils.progress_bar.progress_bar import create_progress_bar @@ -690,6 +691,9 @@ async def handle_check_suite( ) break + # Force GC between rounds to free temporary objects (messages, diffs) and reduce Lambda OOM risk + gc_collect_and_log() + # Log if loop exhausted without completion and force verification if not is_completed: logger.warning( diff --git a/services/webhook/new_pr_handler.py b/services/webhook/new_pr_handler.py index fe10d36e..203d4b73 100644 --- a/services/webhook/new_pr_handler.py +++ b/services/webhook/new_pr_handler.py @@ -71,6 +71,7 @@ from utils.images.get_base64 import get_base64 from utils.logging.add_log_message import add_log_message from utils.logging.logging_config import logger, set_pr_number, set_trigger +from utils.memory.gc_collect_and_log import gc_collect_and_log from utils.progress_bar.progress_bar import create_progress_bar from utils.text.text_copy import ( pull_request_completed, @@ -545,6 +546,9 @@ async def handle_new_pr( ) break + # Force GC between rounds to free temporary objects (messages, diffs) and reduce Lambda OOM risk + gc_collect_and_log() + # Log if loop exhausted without completion and force verification if not is_completed: logger.warning( diff --git a/services/webhook/review_run_handler.py b/services/webhook/review_run_handler.py index c0331dd1..5e6ec374 100644 --- a/services/webhook/review_run_handler.py +++ b/services/webhook/review_run_handler.py @@ -44,6 +44,7 @@ from services.webhook.utils.should_bail import should_bail from utils.logging.add_log_message import add_log_message from utils.logging.logging_config import logger, set_pr_number, set_trigger +from utils.memory.gc_collect_and_log import gc_collect_and_log from utils.progress_bar.progress_bar import create_progress_bar @@ -336,6 +337,9 @@ async def handle_review_run( ) break + # Force GC between rounds to free temporary objects (messages, diffs) and reduce Lambda OOM risk + gc_collect_and_log() + # Log if loop exhausted without completion and force verification if not is_completed: logger.warning( diff --git a/utils/memory/gc_collect_and_log.py b/utils/memory/gc_collect_and_log.py new file mode 100644 index 00000000..a48fdcea --- /dev/null +++ b/utils/memory/gc_collect_and_log.py @@ -0,0 +1,18 @@ +import gc + +from utils.error.handle_exceptions import handle_exceptions +from utils.logging.logging_config import logger +from utils.memory.get_rss_mb import get_rss_mb + + +@handle_exceptions(default_return_value=None, raise_on_error=False) +def gc_collect_and_log(): + before_mb = get_rss_mb() + collected = gc.collect() + after_mb = get_rss_mb() + logger.info( + "gc: collected=%d rss_before=%.1fMB rss_after=%.1fMB", + collected, + before_mb, + after_mb, + ) diff --git a/utils/memory/get_rss_mb.py b/utils/memory/get_rss_mb.py new file mode 100644 index 00000000..46b2737a --- /dev/null +++ b/utils/memory/get_rss_mb.py @@ -0,0 +1,15 @@ +import platform +import resource + +from utils.error.handle_exceptions import handle_exceptions + +# macOS returns bytes, Linux returns KB +_IS_MACOS = platform.system() == "Darwin" + + +@handle_exceptions(default_return_value=0.0, raise_on_error=False) +def get_rss_mb(): + ru_maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if _IS_MACOS: + return ru_maxrss / (1024 * 1024) + return ru_maxrss / 1024 diff --git a/utils/memory/is_lambda_oom_approaching.py b/utils/memory/is_lambda_oom_approaching.py index 3061cb3f..f93ab7c1 100644 --- a/utils/memory/is_lambda_oom_approaching.py +++ b/utils/memory/is_lambda_oom_approaching.py @@ -1,23 +1,13 @@ -# Standard imports -import platform -import resource - # Local imports from utils.error.handle_exceptions import handle_exceptions +from utils.memory.get_rss_mb import get_rss_mb # Must match MemorySize in infrastructure/deploy-lambda-with-vpc-efs.yml LAMBDA_MEMORY_MB = 2048 -# macOS returns bytes, Linux returns KB -_IS_MACOS = platform.system() == "Darwin" - @handle_exceptions(default_return_value=(False, 0), raise_on_error=False) def is_lambda_oom_approaching(): - ru_maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - if _IS_MACOS: - used_mb = ru_maxrss / (1024 * 1024) - else: - used_mb = ru_maxrss / 1024 + used_mb = get_rss_mb() threshold_mb = LAMBDA_MEMORY_MB * 0.9 return used_mb > threshold_mb, used_mb diff --git a/utils/memory/test_gc_collect_and_log.py b/utils/memory/test_gc_collect_and_log.py new file mode 100644 index 00000000..4be4933e --- /dev/null +++ b/utils/memory/test_gc_collect_and_log.py @@ -0,0 +1,13 @@ +from unittest.mock import patch + +from utils.memory.gc_collect_and_log import gc_collect_and_log + + +def test_gc_collect_and_log_runs_without_error(): + gc_collect_and_log() + + +@patch("utils.memory.gc_collect_and_log.gc.collect", return_value=42) +def test_gc_collect_and_log_calls_gc(mock_gc): + gc_collect_and_log() + mock_gc.assert_called_once() diff --git a/utils/memory/test_is_lambda_oom_approaching.py b/utils/memory/test_is_lambda_oom_approaching.py index e67892c1..bcb1abde 100644 --- a/utils/memory/test_is_lambda_oom_approaching.py +++ b/utils/memory/test_is_lambda_oom_approaching.py @@ -22,8 +22,8 @@ def _mock_rusage(ru_maxrss: int): class TestIsLambdaOomApproaching: - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", False) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", False) + @patch("utils.memory.get_rss_mb.resource") def test_below_threshold_linux(self, mock_resource): # 1000 MB in KB (Linux units) mock_resource.getrusage.return_value = _mock_rusage(1000 * 1024) @@ -32,8 +32,8 @@ def test_below_threshold_linux(self, mock_resource): assert is_approaching is False assert used_mb == 1000.0 - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", False) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", False) + @patch("utils.memory.get_rss_mb.resource") def test_above_threshold_linux(self, mock_resource): # 1900 MB in KB (above 1843.2 MB threshold) mock_resource.getrusage.return_value = _mock_rusage(1900 * 1024) @@ -42,8 +42,8 @@ def test_above_threshold_linux(self, mock_resource): assert is_approaching is True assert used_mb == 1900.0 - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", True) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", True) + @patch("utils.memory.get_rss_mb.resource") def test_below_threshold_macos(self, mock_resource): # 1000 MB in bytes (macOS units) mock_resource.getrusage.return_value = _mock_rusage(1000 * 1024 * 1024) @@ -52,8 +52,8 @@ def test_below_threshold_macos(self, mock_resource): assert is_approaching is False assert used_mb == 1000.0 - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", True) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", True) + @patch("utils.memory.get_rss_mb.resource") def test_above_threshold_macos(self, mock_resource): # 1900 MB in bytes (macOS units) mock_resource.getrusage.return_value = _mock_rusage(1900 * 1024 * 1024) @@ -62,8 +62,8 @@ def test_above_threshold_macos(self, mock_resource): assert is_approaching is True assert used_mb == 1900.0 - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", False) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", False) + @patch("utils.memory.get_rss_mb.resource") def test_exact_threshold_not_approaching(self, mock_resource): # Exactly at threshold (1843.2 MB) - use 1843 MB, not greater, so False mock_resource.getrusage.return_value = _mock_rusage(1843 * 1024) @@ -72,8 +72,8 @@ def test_exact_threshold_not_approaching(self, mock_resource): assert is_approaching is False assert used_mb == 1843.0 - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", False) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", False) + @patch("utils.memory.get_rss_mb.resource") def test_just_above_threshold(self, mock_resource): # 1844 MB - just above 1843.2 threshold mock_resource.getrusage.return_value = _mock_rusage(1844 * 1024) @@ -94,8 +94,8 @@ def test_just_above_threshold(self, mock_resource): ], ids=["zero", "512mb", "1024mb", "at_threshold", "above_threshold", "at_limit"], ) - @patch("utils.memory.is_lambda_oom_approaching._IS_MACOS", False) - @patch("utils.memory.is_lambda_oom_approaching.resource") + @patch("utils.memory.get_rss_mb._IS_MACOS", False) + @patch("utils.memory.get_rss_mb.resource") def test_parametrized_linux(self, mock_resource, used_kb, expected_approaching): mock_resource.getrusage.return_value = _mock_rusage(used_kb) mock_resource.RUSAGE_SELF = 0