feat: add managed local shell sessions - #9470
Merged
Merged
Conversation
Deploying with
|
| Status | Name | Latest Commit | Preview URL | Updated (UTC) |
|---|---|---|---|---|
| ✅ Deployment successful! View logs |
astrbot-docs | 9acc71e | Commit Preview URL Branch Preview URL |
Jul 30 2026, 02:43 PM |
Soulter
marked this pull request as ready for review
July 30, 2026 15:05
dosubot
Bot
added
size:XL
📄 Knowledge reviewDosu skipped reviewing this PR because your organization has used its
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The session status derivation logic ("running"/"timed_out"/"terminated"/"completed"/"failed") is duplicated in both list_sessions and poll_session; consider extracting a small helper to centralize this to keep future changes consistent.
- poll_session has multiple branches that re-read the same output file and await reader_task more than once; you could simplify this by structuring the read/await flow into a single helper or clearer state machine to reduce redundant I/O and improve readability.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The session status derivation logic ("running"/"timed_out"/"terminated"/"completed"/"failed") is duplicated in both list_sessions and poll_session; consider extracting a small helper to centralize this to keep future changes consistent.
- poll_session has multiple branches that re-read the same output file and await reader_task more than once; you could simplify this by structuring the read/await flow into a single helper or clearer state machine to reduce redundant I/O and improve readability.
## Individual Comments
### Comment 1
<location path="astrbot/core/computer/booters/local.py" line_range="394" />
<code_context>
+ )
+ return {"sessions": items}
+
+ async def poll_session(
+ self,
+ *,
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring the managed shell session logic into small helpers for polling, status, lifecycle, and termination to reduce duplication and make control flow easier to understand.
- `poll_session` mixes output reading, waiting, status computation, and lifecycle in one method, and repeats `wait_task`/`reader_task`/`returncode` checks and `_read_output` calls. This increases cognitive load and makes it harder to reason about edge cases.
You can keep all behavior while reducing complexity by extracting small helpers and letting `poll_session` focus on “read & report”:
```python
def _read_output_once(
session: _LocalShellSession,
cursor: int,
max_output_chars: int,
) -> tuple[bytes, int, int]:
try:
output_size = session.output_path.stat().st_size
except FileNotFoundError:
return b"", cursor, cursor
normalized_cursor = min(cursor, output_size)
with session.output_path.open("rb") as output_file:
output_file.seek(normalized_cursor)
raw_output = output_file.read(max_output_chars)
return raw_output, normalized_cursor + len(raw_output), output_size
```
```python
async def _wait_for_output_or_exit(
session: _LocalShellSession,
timeout_ms: int,
) -> None:
session.output_event.clear()
output_waiter = asyncio.create_task(session.output_event.wait())
try:
done, _ = await asyncio.wait(
{output_waiter, session.wait_task},
timeout=timeout_ms / 1000,
return_when=asyncio.FIRST_COMPLETED,
)
if output_waiter not in done:
output_waiter.cancel()
try:
await output_waiter
except asyncio.CancelledError:
pass
finally:
# If process exited, ensure reader_task is drained
if session.wait_task.done() and not session.reader_task.done():
await session.reader_task
```
Then `poll_session` can be simplified to a single output-read path:
```python
async def poll_session(...):
...
if session.wait_task.done() and not session.reader_task.done():
await session.reader_task
raw_output, next_cursor, output_size = await asyncio.to_thread(
_read_output_once, session, read_cursor, max_output_chars
)
if (
not raw_output
and session.process.returncode is None
and yield_time_ms > 0
):
await _wait_for_output_or_exit(session, yield_time_ms)
raw_output, next_cursor, output_size = await asyncio.to_thread(
_read_output_once, session, read_cursor, max_output_chars
)
...
```
This removes the repeated `returncode` checks and repeated `_read_output` calls while preserving the current semantics of “wait for either output or exit”.
- Status computation is duplicated across `list_sessions` and `poll_session` with nested conditionals. Extracting a helper both reduces duplication and makes the state transitions clearer:
```python
def _session_status(session: _LocalShellSession) -> str:
exit_code = session.process.returncode
if exit_code is None:
return "running"
if session.timed_out:
return "timed_out"
if session.terminated:
return "terminated"
return "completed" if exit_code == 0 else "failed"
```
Use it in both places:
```python
status = _session_status(session)
```
- Session removal logic (`session_closed` computation) is currently embedded in `poll_session`, while termination paths (`terminate_session`, `shutdown_sessions`) also call `_remove_session` separately. Centralizing the “is this finished” decision in a helper keeps lifecycle rules in one place and makes `poll_session` purely about I/O:
```python
def _session_finished(session: _LocalShellSession, has_more_output: bool) -> bool:
exit_code = session.process.returncode
# same semantics as current code: exited and no remaining output
return exit_code is not None and not has_more_output
```
```python
has_more = next_cursor < output_size
session_closed = _session_finished(session, has_more)
if session_closed:
await self._remove_session(session)
```
- `_terminate_process` currently mixes Windows `taskkill` subprocess handling and POSIX signal logic in one function. Splitting OS-specific behavior into helpers keeps the main function short and easier to scan:
```python
async def _terminate_process(self, session: _LocalShellSession) -> None:
if session.process.returncode is not None:
return
if os.name == "nt":
await self._terminate_process_windows(session)
else:
await self._terminate_process_posix(session)
try:
await asyncio.wait_for(asyncio.shield(session.wait_task), timeout=5)
except asyncio.TimeoutError:
if os.name == "nt":
session.process.kill()
else:
try:
os.killpg(session.process.pid, signal.SIGKILL)
except ProcessLookupError:
pass
await session.wait_task
```
```python
async def _terminate_process_windows(self, session: _LocalShellSession) -> None:
try:
taskkill_result = await asyncio.to_thread(
subprocess.run,
["taskkill", "/F", "/T", "/PID", str(session.process.pid)],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
)
except Exception:
session.process.terminate()
else:
if taskkill_result.returncode != 0:
session.process.terminate()
async def _terminate_process_posix(self, session: _LocalShellSession) -> None:
try:
os.killpg(session.process.pid, signal.SIGTERM)
except ProcessLookupError:
pass
```
These focused helpers keep all existing behavior (managed sessions, timeouts, background I/O) but significantly reduce the complexity and duplication in the control flow, especially around `poll_session`, status computation, lifecycle decisions, and termination.
</issue_to_address>
### Comment 2
<location path="astrbot/core/tools/computer_tools/shell.py" line_range="204" />
<code_context>
+ }
+ )
+
+ async def call(
+ self,
+ context: ContextWrapper[AstrAgentContext],
</code_context>
<issue_to_address>
**issue (complexity):** Consider extracting a shared local shell helper and using a dispatch map for session actions to simplify control flow and remove duplication between tools.
You can reduce the new complexity without changing behavior by introducing a shared helper for local shell access and centralizing the session action dispatch.
### 1. Centralize local shell/runtime checks
Right now, both `ExecuteShellTool.call` and `ShellSessionTool.call` repeat:
- `is_local_runtime(context)`
- `get_booter(...)`
- `isinstance(sb.shell, LocalShellComponent)`
- Returning similar error messages
You can move this into a single helper that either returns a `LocalShellComponent` or a user-facing error string, and reuse it in both tools.
```python
async def _get_local_shell(
context: ContextWrapper[AstrAgentContext],
error_prefix: str,
) -> LocalShellComponent | str:
if not is_local_runtime(context):
return f"{error_prefix}: only local runtime is supported."
sb = await get_booter(
context.context.context,
context.context.event.unified_msg_origin,
)
if not isinstance(sb.shell, LocalShellComponent):
return f"{error_prefix}: local shell component is unavailable."
return sb.shell
```
Usage in `ExecuteShellTool.call`:
```python
local_runtime = is_local_runtime(context)
if local_runtime:
shell = await _get_local_shell(
context,
error_prefix="Error executing command",
)
if isinstance(shell, str):
return shell
current_workspace_root = await workspace_root_for_context(context)
current_workspace_root.mkdir(parents=True, exist_ok=True)
cwd = str(current_workspace_root)
env = dict(env or {})
return json.dumps(
await shell.exec_managed(
command,
owner_id=context.context.event.unified_msg_origin,
cwd=cwd,
env=env,
timeout=timeout,
yield_time_ms=0 if background else yield_time_ms,
),
ensure_ascii=False,
)
```
Usage in `ShellSessionTool.call`:
```python
shell = await _get_local_shell(
context,
error_prefix="Error managing shell session",
)
if isinstance(shell, str):
return shell
owner_id = context.context.event.unified_msg_origin
# ... action dispatch below ...
```
This removes duplicated branching and keeps error messaging consistent.
### 2. Replace long `if/elif` chain in `ShellSessionTool.call` with a handler map
You can keep the `session_id` validation but move the action-specific wiring into a small dispatch table. This reduces branching and makes adding new actions more predictable.
```python
owner_id = context.context.event.unified_msg_origin
if action == "list":
result = await shell.list_sessions(owner_id)
else:
if not session_id:
return (
"Error managing shell session: session_id is required "
f"when action={action}."
)
async def _poll():
return await shell.poll_session(
owner_id=owner_id,
session_id=session_id,
cursor=cursor,
yield_time_ms=yield_time_ms,
max_output_chars=max_output_chars,
)
async def _write():
return await shell.write_session(
owner_id=owner_id,
session_id=session_id,
chars=chars,
)
async def _interrupt():
return await shell.interrupt_session(
owner_id=owner_id,
session_id=session_id,
yield_time_ms=yield_time_ms,
max_output_chars=max_output_chars,
)
async def _terminate():
return await shell.terminate_session(
owner_id=owner_id,
session_id=session_id,
max_output_chars=max_output_chars,
)
handlers = {
"poll": _poll,
"write": _write,
"interrupt": _interrupt,
"terminate": _terminate,
}
handler = handlers.get(action)
if handler is None:
return f"Error managing shell session: unsupported action {action}."
result = await handler()
return json.dumps(result, ensure_ascii=False)
```
This keeps the behavior identical while making the control flow more linear and easier to extend/test.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| ) | ||
| return {"sessions": items} | ||
|
|
||
| async def poll_session( |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider refactoring the managed shell session logic into small helpers for polling, status, lifecycle, and termination to reduce duplication and make control flow easier to understand.
-
poll_sessionmixes output reading, waiting, status computation, and lifecycle in one method, and repeatswait_task/reader_task/returncodechecks and_read_outputcalls. This increases cognitive load and makes it harder to reason about edge cases.You can keep all behavior while reducing complexity by extracting small helpers and letting
poll_sessionfocus on “read & report”:def _read_output_once( session: _LocalShellSession, cursor: int, max_output_chars: int, ) -> tuple[bytes, int, int]: try: output_size = session.output_path.stat().st_size except FileNotFoundError: return b"", cursor, cursor normalized_cursor = min(cursor, output_size) with session.output_path.open("rb") as output_file: output_file.seek(normalized_cursor) raw_output = output_file.read(max_output_chars) return raw_output, normalized_cursor + len(raw_output), output_size
async def _wait_for_output_or_exit( session: _LocalShellSession, timeout_ms: int, ) -> None: session.output_event.clear() output_waiter = asyncio.create_task(session.output_event.wait()) try: done, _ = await asyncio.wait( {output_waiter, session.wait_task}, timeout=timeout_ms / 1000, return_when=asyncio.FIRST_COMPLETED, ) if output_waiter not in done: output_waiter.cancel() try: await output_waiter except asyncio.CancelledError: pass finally: # If process exited, ensure reader_task is drained if session.wait_task.done() and not session.reader_task.done(): await session.reader_task
Then
poll_sessioncan be simplified to a single output-read path:async def poll_session(...): ... if session.wait_task.done() and not session.reader_task.done(): await session.reader_task raw_output, next_cursor, output_size = await asyncio.to_thread( _read_output_once, session, read_cursor, max_output_chars ) if ( not raw_output and session.process.returncode is None and yield_time_ms > 0 ): await _wait_for_output_or_exit(session, yield_time_ms) raw_output, next_cursor, output_size = await asyncio.to_thread( _read_output_once, session, read_cursor, max_output_chars ) ...
This removes the repeated
returncodechecks and repeated_read_outputcalls while preserving the current semantics of “wait for either output or exit”. -
Status computation is duplicated across
list_sessionsandpoll_sessionwith nested conditionals. Extracting a helper both reduces duplication and makes the state transitions clearer:def _session_status(session: _LocalShellSession) -> str: exit_code = session.process.returncode if exit_code is None: return "running" if session.timed_out: return "timed_out" if session.terminated: return "terminated" return "completed" if exit_code == 0 else "failed"
Use it in both places:
status = _session_status(session)
-
Session removal logic (
session_closedcomputation) is currently embedded inpoll_session, while termination paths (terminate_session,shutdown_sessions) also call_remove_sessionseparately. Centralizing the “is this finished” decision in a helper keeps lifecycle rules in one place and makespoll_sessionpurely about I/O:def _session_finished(session: _LocalShellSession, has_more_output: bool) -> bool: exit_code = session.process.returncode # same semantics as current code: exited and no remaining output return exit_code is not None and not has_more_output
has_more = next_cursor < output_size session_closed = _session_finished(session, has_more) if session_closed: await self._remove_session(session)
-
_terminate_processcurrently mixes Windowstaskkillsubprocess handling and POSIX signal logic in one function. Splitting OS-specific behavior into helpers keeps the main function short and easier to scan:async def _terminate_process(self, session: _LocalShellSession) -> None: if session.process.returncode is not None: return if os.name == "nt": await self._terminate_process_windows(session) else: await self._terminate_process_posix(session) try: await asyncio.wait_for(asyncio.shield(session.wait_task), timeout=5) except asyncio.TimeoutError: if os.name == "nt": session.process.kill() else: try: os.killpg(session.process.pid, signal.SIGKILL) except ProcessLookupError: pass await session.wait_task
async def _terminate_process_windows(self, session: _LocalShellSession) -> None: try: taskkill_result = await asyncio.to_thread( subprocess.run, ["taskkill", "/F", "/T", "/PID", str(session.process.pid)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5, ) except Exception: session.process.terminate() else: if taskkill_result.returncode != 0: session.process.terminate() async def _terminate_process_posix(self, session: _LocalShellSession) -> None: try: os.killpg(session.process.pid, signal.SIGTERM) except ProcessLookupError: pass
These focused helpers keep all existing behavior (managed sessions, timeouts, background I/O) but significantly reduce the complexity and duplication in the control flow, especially around poll_session, status computation, lifecycle decisions, and termination.
| } | ||
| ) | ||
|
|
||
| async def call( |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider extracting a shared local shell helper and using a dispatch map for session actions to simplify control flow and remove duplication between tools.
You can reduce the new complexity without changing behavior by introducing a shared helper for local shell access and centralizing the session action dispatch.
1. Centralize local shell/runtime checks
Right now, both ExecuteShellTool.call and ShellSessionTool.call repeat:
is_local_runtime(context)get_booter(...)isinstance(sb.shell, LocalShellComponent)- Returning similar error messages
You can move this into a single helper that either returns a LocalShellComponent or a user-facing error string, and reuse it in both tools.
async def _get_local_shell(
context: ContextWrapper[AstrAgentContext],
error_prefix: str,
) -> LocalShellComponent | str:
if not is_local_runtime(context):
return f"{error_prefix}: only local runtime is supported."
sb = await get_booter(
context.context.context,
context.context.event.unified_msg_origin,
)
if not isinstance(sb.shell, LocalShellComponent):
return f"{error_prefix}: local shell component is unavailable."
return sb.shellUsage in ExecuteShellTool.call:
local_runtime = is_local_runtime(context)
if local_runtime:
shell = await _get_local_shell(
context,
error_prefix="Error executing command",
)
if isinstance(shell, str):
return shell
current_workspace_root = await workspace_root_for_context(context)
current_workspace_root.mkdir(parents=True, exist_ok=True)
cwd = str(current_workspace_root)
env = dict(env or {})
return json.dumps(
await shell.exec_managed(
command,
owner_id=context.context.event.unified_msg_origin,
cwd=cwd,
env=env,
timeout=timeout,
yield_time_ms=0 if background else yield_time_ms,
),
ensure_ascii=False,
)Usage in ShellSessionTool.call:
shell = await _get_local_shell(
context,
error_prefix="Error managing shell session",
)
if isinstance(shell, str):
return shell
owner_id = context.context.event.unified_msg_origin
# ... action dispatch below ...This removes duplicated branching and keeps error messaging consistent.
2. Replace long if/elif chain in ShellSessionTool.call with a handler map
You can keep the session_id validation but move the action-specific wiring into a small dispatch table. This reduces branching and makes adding new actions more predictable.
owner_id = context.context.event.unified_msg_origin
if action == "list":
result = await shell.list_sessions(owner_id)
else:
if not session_id:
return (
"Error managing shell session: session_id is required "
f"when action={action}."
)
async def _poll():
return await shell.poll_session(
owner_id=owner_id,
session_id=session_id,
cursor=cursor,
yield_time_ms=yield_time_ms,
max_output_chars=max_output_chars,
)
async def _write():
return await shell.write_session(
owner_id=owner_id,
session_id=session_id,
chars=chars,
)
async def _interrupt():
return await shell.interrupt_session(
owner_id=owner_id,
session_id=session_id,
yield_time_ms=yield_time_ms,
max_output_chars=max_output_chars,
)
async def _terminate():
return await shell.terminate_session(
owner_id=owner_id,
session_id=session_id,
max_output_chars=max_output_chars,
)
handlers = {
"poll": _poll,
"write": _write,
"interrupt": _interrupt,
"terminate": _terminate,
}
handler = handlers.get(action)
if handler is None:
return f"Error managing shell session: unsupported action {action}."
result = await handler()
return json.dumps(result, ensure_ascii=False)This keeps the behavior identical while making the control flow more linear and easier to extend/test.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
background=trueflow with automatic yielding to managed shell sessionsastrbot_shell_sessionactions for listing, polling, writing stdin, interrupting, and terminating Local sessionsMotivation
Local background commands previously returned only a PID and log path, while long foreground commands could outlive the tool call without a reliable control channel. Managed sessions provide a stable way to recover output, send line-oriented stdin, and stop process groups after the initial tool call returns.
Behavior
Local commands wait for
yield_time_msand return normally when they finish within that window. Commands that remain active return a session ID for follow-up operations. Persistent processes can be polled to confirm startup and then interrupted or terminated when no longer needed. This intentionally does not add PTY/TTY emulation.Validation
uv run ruff check .git diff --check origin/master...HEADuv run pytest -q tests/test_local_shell_component.py tests/unit/test_func_tool_manager.py tests/unit/test_computer.py tests/unit/test_astr_main_agent.py tests/unit/test_cua_computer_use.py(236 passed)Summary by Sourcery
Introduce managed local shell sessions and a dedicated session-management tool, replacing background-mode semantics for the local runtime and integrating lifecycle cleanup.
New Features:
Enhancements:
Tests: