feat: add session state tracking for code execution
linuztx committed
Aug 7, 2025 at 10:41 UTC
a246c14c5033059f67db6f5d1f3185a8f2bb06a3
3 files changed
+85
-16
prompts/fw.code.running.md
new
+1
@@ -0,0 +1 @@
1
+Terminal session {{session}} is still running. Decide wait for more 'output' or 'reset' base on context.
\ No newline at end of file
python/helpers/shell_ssh.py
+3
@@ -27,6 +27,7 @@ class SSHInteractiveSession:
27
self.full_output = b""
28
self.last_command = b""
29
self.trimmed_command_length = 0 # Initialize trimmed_command_length
30
+ self.is_running = False
31
32
async def connect(self):
33
# try 3 times with wait and then except
@@ -82,6 +83,8 @@ class SSHInteractiveSession:
83
self.last_command = command.encode()
84
self.trimmed_command_length = 0
85
self.shell.send(self.last_command)
86
+
87
+ self.is_running = True
88
89
async def read_output(
90
self, timeout: float = 0, reset_full_output: bool = False
python/tools/code_execution_tool.py
+81
-16
@@ -21,6 +21,20 @@ class State:
21
22
class CodeExecution(Tool):
23
24
+ # Common shell prompt regex patterns (add more as needed)
25
+ prompt_patterns = [
26
+ re.compile(r"\\(venv\\).+[$#] ?$"), # (venv) ...$ or (venv) ...#
27
+ re.compile(r"root@[^:]+:[^#]+# ?$"), # root@container:~#
28
+ re.compile(r"[a-zA-Z0-9_.-]+@[^:]+:[^$#]+[$#] ?$"), # user@host:~$
29
+ ]
30
+ # potential dialog detection
31
+ dialog_patterns = [
32
+ re.compile(r"Y/N", re.IGNORECASE), # Y/N anywhere in line
33
+ re.compile(r"yes/no", re.IGNORECASE), # yes/no anywhere in line
34
+ re.compile(r":\s*$"), # line ending with colon
35
+ re.compile(r"\?\s*$"), # line ending with question mark
36
+ ]
37
+
38
async def execute(self, **kwargs):
39
40
await self.agent.handle_intervention() # wait for intervention and handle it, if paused
@@ -31,6 +45,7 @@ class CodeExecution(Tool):
45
46
runtime = self.args.get("runtime", "").lower().strip()
47
session = int(self.args.get("session", 0))
48
+ self.allow_running = bool(self.args.get("allow_running", False))
49
50
if runtime == "python":
51
response = await self.execute_python_code(
@@ -157,6 +172,12 @@ class CodeExecution(Tool):
172
):
173
174
await self.agent.handle_intervention() # wait for intervention and handle it, if paused
175
+
176
+ # Check if session is running and handle it
177
+ if not self.allow_running:
178
+ if response := await self.handle_running_session(session):
179
+ return response
180
+
181
# try again on lost connection
182
for i in range(2):
183
try:
@@ -221,20 +242,6 @@ class CodeExecution(Tool):
242
sleep_time=0.1,
243
prefix="",
244
):
224
- # Common shell prompt regex patterns (add more as needed)
225
- prompt_patterns = [
226
- re.compile(r"\\(venv\\).+[$#] ?$"), # (venv) ...$ or (venv) ...#
227
- re.compile(r"root@[^:]+:[^#]+# ?$"), # root@container:~#
228
- re.compile(r"[a-zA-Z0-9_.-]+@[^:]+:[^$#]+[$#] ?$"), # user@host:~$
229
- ]
230
-
231
- # potential dialog detection
232
- dialog_patterns = [
233
- re.compile(r"Y/N", re.IGNORECASE), # Y/N anywhere in line
234
- re.compile(r"yes/no", re.IGNORECASE), # yes/no anywhere in line
235
- re.compile(r":\s*$"), # line ending with colon
236
- re.compile(r"\?\s*$"), # line ending with question mark
237
- ]
245
246
start_time = time.time()
247
last_output_time = start_time
@@ -271,7 +278,7 @@ class CodeExecution(Tool):
278
)
279
last_lines.reverse()
280
for idx, line in enumerate(last_lines):
274
- for pat in prompt_patterns:
281
+ for pat in self.prompt_patterns:
282
if pat.search(line.strip()):
283
PrintStyle.info(
284
"Detected shell prompt, returning output early."
@@ -281,6 +288,7 @@ class CodeExecution(Tool):
288
"\n".join(last_lines), idx + 1, True
289
)
290
self.log.update(heading=heading)
291
+ self.mark_session_idle(session)
292
return truncated_output
293
294
# Check for max execution time
@@ -327,7 +335,7 @@ class CodeExecution(Tool):
335
truncated_output.splitlines()[-2:] if truncated_output else []
336
)
337
for line in last_lines:
330
- for pat in dialog_patterns:
338
+ for pat in self.dialog_patterns:
339
if pat.search(line.strip()):
340
PrintStyle.info(
341
"Detected dialog prompt, returning output early."
@@ -350,6 +358,63 @@ class CodeExecution(Tool):
358
)
359
return response
360
361
+ async def handle_running_session(
362
+ self,
363
+ session=0,
364
+ reset_full_output=False,
365
+ prefix=""
366
+ ):
367
+ if not (
368
+ session in self.state.shells
369
+ and getattr(self.state.shells[session], "is_running", False)
370
+ ):
371
+ return None
372
+
373
+ full_output, _ = await self.state.shells[session].read_output(
374
+ timeout=1, reset_full_output=reset_full_output
375
+ )
376
+ truncated_output = self.fix_full_output(full_output)
377
+ heading = self.get_heading_from_output(truncated_output, 0)
378
+
379
+ last_lines = (
380
+ truncated_output.splitlines()[-3:] if truncated_output else []
381
+ )
382
+ last_lines.reverse()
383
+ for idx, line in enumerate(last_lines):
384
+ for pat in self.prompt_patterns:
385
+ if pat.search(line.strip()):
386
+ PrintStyle.info(
387
+ "Detected shell prompt, returning output early."
388
+ )
389
+ self.mark_session_idle(session)
390
+ return None
391
+
392
+ has_dialog = False
393
+ for line in last_lines:
394
+ for pat in self.dialog_patterns:
395
+ if pat.search(line.strip()):
396
+ has_dialog = True
397
+ break
398
+ if has_dialog:
399
+ break
400
+
401
+ if has_dialog:
402
+ sys_info = self.agent.read_prompt("fw.code.pause_dialog.md", timeout=1)
403
+ else:
404
+ sys_info = self.agent.read_prompt("fw.code.running.md", session=session)
405
+
406
+ response = self.agent.read_prompt("fw.code.info.md", info=sys_info)
407
+ if truncated_output:
408
+ response = truncated_output + "\n\n" + response
409
+ PrintStyle(font_color="#FFA500", bold=True).print(response)
410
+ self.log.update(content=prefix + response, heading=heading)
411
+ return response
412
+
413
+ def mark_session_idle(self, session: int = 0):
414
+ # Mark session as idle - command finished
415
+ if session in self.state.shells:
416
+ self.state.shells[session].is_running = False
417
+
418
async def reset_terminal(self, session=0, reason: str | None = None):
419
# Print the reason for the reset to the console if provided
420
if reason: