CET state polishing

frdel committed Sep 25, 2025 at 15:07 UTC 869d5cc1ef5c299c80ce22134715a3c5b961b302
4 files changed +22 -27
prompts/fw.code.running.md
+1 -1
@@ -1 +1 @@
1 -Terminal session {{session}} is still running. Decide wait for more 'output' or 'reset' base on context.
\ No newline at end of file
1 +Terminal session {{session}} is still running. Decide to wait for more 'output', 'reset', or use another session number based on situation.
\ No newline at end of file
python/helpers/shell_local.py
-2
@@ -10,7 +10,6 @@ class LocalInteractiveSession:
10 def __init__(self):
11 self.session: tty_session.TTYSession|None = None
12 self.full_output = ''
13 - self.is_running = False
13
14 async def connect(self):
15 self.session = tty_session.TTYSession("/bin/bash")
@@ -27,7 +26,6 @@ class LocalInteractiveSession:
26 raise Exception("Shell not connected")
27 self.full_output = ""
28 await self.session.sendline(command)
30 - self.is_running = True
29
30 async def read_output(self, timeout: float = 0, reset_full_output: bool = False) -> Tuple[str, Optional[str]]:
31 if not self.session:
python/helpers/shell_ssh.py
-3
@@ -27,7 +27,6 @@ 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
30
31 async def connect(self, keepalive_interval: int = 5):
32 """
@@ -101,8 +100,6 @@ class SSHInteractiveSession:
100 self.trimmed_command_length = 0
101 self.shell.send(self.last_command)
102
104 - self.is_running = True
105 -
103 async def read_output(
104 self, timeout: float = 0, reset_full_output: bool = False
105 ) -> Tuple[str, str]:
python/tools/code_execution_tool.py
+21 -21
@@ -29,10 +29,16 @@ OUTPUT_TIMEOUTS: dict[str, int] = {
29 "dialog_timeout": 5,
30 }
31
32 +@dataclass
33 +class ShellWrap:
34 + id: int
35 + session: LocalInteractiveSession | SSHInteractiveSession
36 + running: bool
37 +
38 @dataclass
39 class State:
40 ssh_enabled: bool
35 - shells: dict[int, LocalInteractiveSession | SSHInteractiveSession]
41 + shells: dict[int, ShellWrap]
42
43
44 class CodeExecution(Tool):
@@ -112,18 +118,18 @@ class CodeExecution(Tool):
118 # always reset state when ssh_enabled changes
119 if not self.state or self.state.ssh_enabled != self.agent.config.code_exec_ssh_enabled:
120 # initialize shells dictionary if not exists
115 - shells: dict[int, LocalInteractiveSession | SSHInteractiveSession] = {}
121 + shells: dict[int, ShellWrap] = {}
122 else:
123 shells = self.state.shells.copy()
124
125 # Only reset the specified session if provided
126 if reset and session is not None and session in shells:
121 - await shells[session].close()
127 + await shells[session].session.close()
128 del shells[session]
129 elif reset and not session:
130 # Close all sessions if full reset requested
131 for s in list(shells.keys()):
126 - await shells[s].close()
132 + await shells[s].session.close()
133 shells = {}
134
135 # initialize local or remote interactive shell interface for session 0 if needed
@@ -144,7 +150,7 @@ class CodeExecution(Tool):
150 else:
151 shell = LocalInteractiveSession()
152
147 - shells[session] = shell
153 + shells[session] = ShellWrap(id=session, session=shell, running=False)
154 await shell.connect()
155
156 self.state = State(shells=shells, ssh_enabled=self.agent.config.code_exec_ssh_enabled)
@@ -186,14 +192,15 @@ class CodeExecution(Tool):
192 for i in range(2):
193 try:
194
189 - await self.state.shells[session].send_command(command)
195 + self.state.shells[session].running = True
196 + await self.state.shells[session].session.send_command(command)
197
198 locl = (
199 " (local)"
193 - if isinstance(self.state.shells[session], LocalInteractiveSession)
200 + if isinstance(self.state.shells[session].session, LocalInteractiveSession)
201 else (
202 " (remote)"
196 - if isinstance(self.state.shells[session], SSHInteractiveSession)
203 + if isinstance(self.state.shells[session].session, SSHInteractiveSession)
204 else " (unknown)"
205 )
206 )
@@ -258,7 +265,7 @@ class CodeExecution(Tool):
265
266 while True:
267 await asyncio.sleep(sleep_time)
261 - full_output, partial_output = await self.state.shells[session].read_output(
268 + full_output, partial_output = await self.state.shells[session].session.read_output(
269 timeout=1, reset_full_output=reset_full_output
270 )
271 reset_full_output = False # only reset once
@@ -367,16 +374,12 @@ class CodeExecution(Tool):
374 reset_full_output=True,
375 prefix=""
376 ):
370 - state = getattr(self, "state", None)
371 - if not state:
377 + if not self.state or session not in self.state.shells:
378 return None
373 - if not (
374 - session in state.shells
375 - and getattr(state.shells[session], "is_running", False)
376 - ):
379 + if not self.state.shells[session].running:
380 return None
381
379 - full_output, _ = await state.shells[session].read_output(
382 + full_output, _ = await self.state.shells[session].session.read_output(
383 timeout=1, reset_full_output=reset_full_output
384 )
385 truncated_output = self.fix_full_output(full_output)
@@ -418,11 +421,8 @@ class CodeExecution(Tool):
421
422 def mark_session_idle(self, session: int = 0):
423 # Mark session as idle - command finished
421 - state = getattr(self, "state", None)
422 - if state and session in state.shells:
423 - shell = state.shells[session]
424 - if hasattr(shell, "is_running"):
425 - shell.is_running = False
424 + if self.state and session in self.state.shells:
425 + self.state.shells[session].running = False
426
427 async def reset_terminal(self, session=0, reason: str | None = None):
428 # Print the reason for the reset to the console if provided