From fdbffae4d7cc799b8c974aeddff0f0c385278974 Mon Sep 17 00:00:00 2001
From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com>
Date: Fri, 23 May 2025 14:13:31 +0800
Subject: [PATCH] fix: fix the bug in the regular expression matching for
 stdout (#890)

---
 rdagent/utils/__init__.py | 48 +++++++++++++++++++++++++++++++++------
 rdagent/utils/fmt.py      | 12 +++++++---
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/rdagent/utils/__init__.py b/rdagent/utils/__init__.py
index 904b6dd6..0800d140 100644
--- a/rdagent/utils/__init__.py
+++ b/rdagent/utils/__init__.py
@@ -11,6 +11,7 @@ import importlib
 import json
 import re
 import sys
+from multiprocessing import Process, Queue
 from pathlib import Path
 from types import ModuleType
 from typing import Union
@@ -75,6 +76,42 @@ def remove_ansi_codes(s: str) -> str:
     return ansi_escape.sub("", s)
 
 
+def safe_sub(pattern: str, text: str, queue: Queue) -> None:
+    try:
+        result = re.sub(pattern, "", text)
+        queue.put(result)
+    except Exception as e:
+        queue.put(e)
+
+
+def apply_regex_with_timeout(pattern: str, text: str, timeout: int = 120) -> str:
+    queue: Queue[str | Exception] = Queue()
+    p = Process(target=safe_sub, args=(pattern, text, queue))
+    p.start()
+    p.join(timeout)
+
+    if p.is_alive():
+        p.terminate()
+        p.join()
+        logger.warning(f"Pattern {pattern} timed out after {timeout} seconds, skipping it.")
+        return text
+    else:
+        result = queue.get()
+        if isinstance(result, Exception):
+            logger.warning(f"Pattern {pattern} raised an error: {result}")
+            return text
+        return result
+
+
+def filter_with_time_limit(regex_patterns: Union[str, list[str]], filtered_stdout: str) -> str:
+    if isinstance(regex_patterns, list):
+        for pattern in regex_patterns:
+            filtered_stdout = apply_regex_with_timeout(pattern, filtered_stdout)
+    else:
+        filtered_stdout = re.sub(regex_patterns, "", filtered_stdout)
+    return filtered_stdout
+
+
 def filter_redundant_text(stdout: str) -> str:
     """
     Filter out progress bars from stdout using regex.
@@ -97,11 +134,12 @@ def filter_redundant_text(stdout: str) -> str:
     filtered_stdout = re.sub(r"\s*\n\s*", "\n", filtered_stdout)
 
     needs_sub = True
-    # Attempt further filtering up to 5 times
-    for _ in range(5):
+    # Attempt further filtering up to 3 times
+    for _ in range(3):
         filtered_stdout_shortened = filtered_stdout
         system_prompt = T(".prompts:filter_redundant_text.system").r()
 
+        # Check if all regex patterns have been collected
         for __ in range(10):
             user_prompt = T(".prompts:filter_redundant_text.user").r(
                 stdout=filtered_stdout_shortened,
@@ -131,11 +169,7 @@ def filter_redundant_text(stdout: str) -> str:
         needs_sub = response.get("needs_sub", True)
         regex_patterns = response.get("regex_patterns", [])
         try:
-            if isinstance(regex_patterns, list):
-                for pattern in regex_patterns:
-                    filtered_stdout = re.sub(pattern, "", filtered_stdout)
-            else:
-                filtered_stdout = re.sub(regex_patterns, "", filtered_stdout)
+            filter_with_time_limit(regex_patterns, filtered_stdout_shortened)
 
             if not needs_sub:
                 break
diff --git a/rdagent/utils/fmt.py b/rdagent/utils/fmt.py
index 79af8657..44426e99 100644
--- a/rdagent/utils/fmt.py
+++ b/rdagent/utils/fmt.py
@@ -3,13 +3,18 @@ Tools that support generating better formats.
 """
 
 
-def shrink_text(text: str, context_lines: int = 200, line_len: int = 5000) -> str:
+def shrink_text(
+    text: str, context_lines: int = 200, line_len: int = 5000, *, row_shrink: bool = True, col_shrink: bool = True
+) -> str:
     """
     When the context is too long, hide the part in the middle.
 
     >>> shrink_text("line1\\nline2\\nline3", context_lines=2, line_len=5)
     'line1\\n... (1 lines are hidden) ...\\nline3'
 
+    >>> shrink_text("line1\\nline2\\nline3", context_lines=2, line_len=5, row_shrink=False)
+    'line1\\nline2\\nline3'
+
     >>> shrink_text("short line", context_lines=2, line_len=5)
     'sh... (5 chars are hidden) ...ine'
 
@@ -22,15 +27,16 @@ def shrink_text(text: str, context_lines: int = 200, line_len: int = 5000) -> st
 
     new_lines = []
     for line in lines:
-        if len(line) > line_len:
+        if col_shrink and len(line) > line_len:
             # If any line is longer than line_len, we can't shrink it
             line = f"{line[:line_len // 2]}... ({len(line) - line_len} chars are hidden) ...{line[- line_len + line_len // 2:]}"
         new_lines.append(line)
     lines = new_lines
 
-    if total_lines <= context_lines:
+    if not row_shrink or total_lines <= context_lines:
         return "\n".join(lines)
 
+    # shrink row only when it is enabled and the total number of lines is greater than context_lines
     # Calculate how many lines to show from start and end
     half_lines = context_lines // 2
     start = "\n".join(lines[:half_lines])