Update code_execution_tool.py

frdel committed May 22, 2025 at 10:06 UTC 92bdb208047f56f42cfef095fb02ba4c12801f53
1 file changed +56 -14
python/tools/code_execution_tool.py
+56 -14
@@ -8,7 +8,6 @@ from python.helpers.print_style import PrintStyle
8 from python.helpers.shell_local import LocalInteractiveSession
9 from python.helpers.shell_ssh import SSHInteractiveSession
10 from python.helpers.docker import DockerContainerManager
11 -from python.helpers.messages import truncate_text
11 import re
12
13
@@ -54,13 +53,8 @@ class CodeExecution(Tool):
53 "fw.code_runtime_wrong.md", runtime=runtime
54 )
55
57 - # if response contains only whitespace, clear it
58 - if isinstance(response, str) and response.strip() == "":
59 - response = None
60 -
56 if not response:
57 response = self.agent.read_prompt("fw.code_no_output.md")
63 - self.log.update(content=response)
58 return Response(message=response, break_loop=False)
59
60 # async def before_execution(self, **kwargs):
@@ -237,11 +231,15 @@ class CodeExecution(Tool):
231 start_time = time.time()
232 last_output_time = start_time
233 full_output = ""
240 -
241 - while max_exec_time <= 0 or time.time() - start_time < max_exec_time:
242 - await asyncio.sleep(SLEEP_TIME) # Wait for some output to be generated
243 - full_output, partial_output = await self.state.shells[session].read_output(
244 - timeout=1, reset_full_output=reset_full_output
234 + got_output = False
235 + # --- Configuration for repetition detection ---
236 + chunk_size = 128 # Size of the block to check for repetition
237 + repeat_threshold = 5 # How many times the block must repeat consecutively
238 +
239 + while True:
240 + await asyncio.sleep(sleep_time)
241 + current_full_output, partial_output = await self.state.shells[session].read_output(
242 + timeout=between_output_timeout, reset_full_output=reset_full_output
243 )
244 reset_full_output = False
245
@@ -250,9 +248,53 @@ class CodeExecution(Tool):
248 now = time.time()
249 if partial_output:
250 PrintStyle(font_color="#85C1E9").stream(partial_output)
253 - truncated_output = truncate_text(self.agent, full_output, 10_000)
254 - self.log.update(content=truncated_output)
255 - idle = 0
251 + full_output += partial_output # Append new output
252 + self.log.update(content=full_output)
253 + last_output_time = now
254 + got_output = True
255 +
256 + # --- Check for repeating output pattern ---
257 + required_len = chunk_size * repeat_threshold
258 + if len(full_output) >= required_len:
259 + check_segment = full_output[-required_len:]
260 + last_chunk = full_output[-chunk_size:]
261 + expected_segment = last_chunk * repeat_threshold
262 + if check_segment == expected_segment:
263 + loop_detected_msg = f"Detected repeating output pattern (last {chunk_size} chars repeated {repeat_threshold} times), likely an infinite loop."
264 + PrintStyle.error(f"{loop_detected_msg} Resetting session {session}.")
265 +
266 + # --- Truncate output for feedback ---
267 + max_output_length = 2048
268 + truncated_output = full_output[-max_output_length:]
269 + if len(full_output) > max_output_length:
270 + summary_prefix = f"[... Output truncated to last {max_output_length} characters ...]\n"
271 + truncated_output = summary_prefix + truncated_output
272 + else:
273 + summary_prefix = ""
274 + # --- End of truncation ---
275 +
276 + self.log.update(content=truncated_output + f"\n--- LOOP DETECTED --- Session {session} reset.")
277 + # Automatically reset the problematic session
278 + reset_reason = f"Detected repeating output pattern (last {chunk_size} chars repeated {repeat_threshold} times)"
279 + await self.reset_terminal(session=session, reason=reset_reason)
280 + # Return informative message to the agent with truncated output and pattern
281 + return f"{truncated_output}\n--- LOOP DETECTED & SESSION RESET ---\n{loop_detected_msg}\nRepeating pattern identified: '{last_chunk[:128]}{'...' if len(last_chunk)>128 else ''}'\nThe terminal session {session} was automatically reset.\nPlease review the code or command that caused the loop to prevent recurrence."
282 + # --- End of repetition check ---
283 +
284 +
285 + # Check for shell prompt at the end of output
286 + last_lines = full_output.splitlines()[-3:] if full_output else []
287 + for line in last_lines:
288 + for pat in prompt_patterns:
289 + if pat.search(line.strip()):
290 + PrintStyle(font_color="#229954").print("Detected shell prompt, returning output early.")
291 + return full_output
292 +
293 + if not got_output:
294 + # Waiting for first output
295 + if now - start_time > first_output_timeout:
296 + PrintStyle.error(f"No output for {first_output_timeout}s after start, returning.")
297 + break
298 else:
299 # Waiting for more output after first output
300 if now - last_output_time > between_output_timeout: