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: