feat: add channel-based AI tools and multi-round tool chaining for SSH sessions#152
Conversation
Move display responsibility from aish-pty to aish-shell callback. SSH session AI now uses the same blue dots spinner, green separator, 🤖 prefix, and markdown rendering as local aish.
- Add mark_prompt_ready() on select timeout to recover lost line-start state after several commands in SSH session - Replace \r\x1b[2K (erase input line) with \r\n (preserve input line) - Show bash_exec-style command display before injecting into remote shell
- Backspace/Delete erases last character from buffer and screen - Ctrl+C / Escape cancels AI input, returns to passthrough - Ctrl+U clears the input line content (keeps ; prefix)
…H sessions Introduce ChannelBashTool and ChannelAskUserTool that communicate via channels instead of direct terminal control, enabling AI tool use in SSH/telnet sessions. Implement followup callback mechanism for multi-round LLM interactions where tool output feeds back into the model for further analysis.
📝 WalkthroughWalkthroughThis PR integrates AI-assisted remote command execution into PTY sessions through input interception and channel-based tool forwarding. A new state machine detects AI-trigger prefixes in user input, buffers PTY output for context, executes AI callbacks in background threads, and forwards bash and ask-user tool events back through the PTY session loop. ChangesRemote Session AI Integration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Thanks for the pull request. A maintainer will review it when available. Please keep the PR focused, explain the why in the description, and make sure local checks pass before requesting review. Contribution guide: https://github.com/AI-Shell-Team/aish/blob/main/CONTRIBUTING.md |
|
This pull request description looks incomplete. Please update the missing sections below before review. Missing items:
|
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/aish-pty/src/persistent.rs (1)
495-546:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftIdle-for-50ms is not a reliable “command finished” signal.
This branch treats one quiet poll interval as “the remote shell is back at a prompt” and immediately runs the followup. Over SSH, silent commands and network gaps routinely exceed 50ms, so the next AI-generated command can be injected while the previous one is still running.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aish-pty/src/persistent.rs` around lines 495 - 546, The code treats sel == 0 (a single 50ms idle poll) as a definitive “command finished” event and invokes followup logic (involving followup_capturing, looks_like_continuation_prompt, pending_followup, pending_response, etc.), which is unreliable over SSH; change the logic so a single timeout does not trigger followup or mark_prompt_ready. Specifically, in the sel == 0 branch (within the draining/is_session/followup_capturing handling), require a stronger prompt confirmation before treating the shell as ready: either detect a real prompt via an existing prompt-detection routine (e.g. extend or reuse looks_like_continuation_prompt/strip_ansi_and_prompt to detect a normal PS1 prompt) or require N consecutive empty polls/timeouts before invoking interceptor.mark_prompt_ready() and executing pending_followup; only send the Ctrl+C or '\r' and clear followup_capturing after that confirmed prompt/threshold is reached. Update any related state (skip_leading_newline, pending_followup, followup_captured) accordingly so followups are fired only on confirmed prompt presence rather than a single 50ms idle.
🧹 Nitpick comments (1)
crates/aish-pty/src/output_buffer.rs (1)
22-30: ⚡ Quick win
appenditerates byte-by-byte — consider a slice-based implementation for hot PTY output.PTY output can be large and frequent. The current loop touches one byte per iteration, which has poor throughput. The writable range can be split into at most two contiguous slices, each copied with
copy_from_slice/extend_from_slice.♻️ Proposed slice-based append
- pub fn append(&mut self, input: &[u8]) { - for &byte in input { - self.data[self.write_pos] = byte; - self.write_pos = (self.write_pos + 1) % self.capacity; - if self.len < self.capacity { - self.len += 1; - } - } - } + pub fn append(&mut self, input: &[u8]) { + if input.is_empty() { + return; + } + // If input is larger than capacity, only keep the tail. + let input = if input.len() > self.capacity { + &input[input.len() - self.capacity..] + } else { + input + }; + let n = input.len(); + let first_chunk = (self.capacity - self.write_pos).min(n); + self.data[self.write_pos..self.write_pos + first_chunk] + .copy_from_slice(&input[..first_chunk]); + if first_chunk < n { + let second_chunk = n - first_chunk; + self.data[..second_chunk].copy_from_slice(&input[first_chunk..]); + } + self.write_pos = (self.write_pos + n) % self.capacity; + self.len = self.capacity.min(self.len + n); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aish-pty/src/output_buffer.rs` around lines 22 - 30, The append method currently writes bytes one-by-one; refactor aish-pty::output_buffer::append to copy the input slice in at most two contiguous chunks into self.data using slice-based operations (e.g., copy_from_slice or extend_from_slice) instead of per-byte assignment: compute the first writable span from write_pos to end-of-buffer, copy min(input.len(), span.len()) bytes, update write_pos (wrap with modulo capacity) and len (capped at capacity), then if remaining input exists copy the rest into the beginning span and update write_pos and len again; make sure to handle empty input and correctly wrap when input length > capacity.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aish-pty/src/output_buffer.rs`:
- Around line 12-19: The constructor new currently allows capacity == 0 which
later causes index/OOB and divide-by-zero in append; guard by normalizing the
input capacity to at least 1 (e.g. let capacity = capacity.max(1)) before
creating data and storing capacity/write_pos/len so data: vec![0u8; capacity]
and self.capacity use the non-zero value; update the new function (symbols: new,
data, capacity, write_pos, len, append) accordingly to avoid the crash.
In `@crates/aish-pty/src/persistent.rs`:
- Around line 1627-1630: The user-facing string for min-length validation is
hardcoded in Chinese (constructed into msg using request.min_length) and
bypasses the i18n layer; update the code that builds msg (and the similar block
at lines ~1847-1860) to use the aish_i18n localization API (e.g., call the
appropriate aish_i18n translation function/key) and format the translated
template with request.min_length so the banner/validation text is localized per
session locale rather than hardcoded Chinese.
- Around line 592-623: The confirmation read only consumes one byte (ans) so a
following newline stays on stdin and is forwarded later; after the single-byte
read in the block that sets approved (the libc::read on stdin_fd producing ans),
add a small consume loop that reads and discards remaining bytes from stdin_fd
(via libc::read) until you hit a '\n' or '\r' or EOF, so the trailing newline
from "y⏎" is removed; keep the existing echo behavior (write the chosen echo
once) and apply the identical fix to the other confirmation site around the code
referenced at lines 921-953 (same ans/stdin_fd/approved pattern).
- Around line 1899-1903: The branch in handle_ask_user that matches
Ok(crate::AiEvent::BashExec { command, output_sender }) currently fakes a tool
result by sending "(command queued: ...)" and breaking, which prevents the
actual bash command from being scheduled; replace that fake-response behavior by
forwarding the BashExec event into the real bash execution path (i.e., enqueue
or call the existing bash handler / scheduler used elsewhere) so the remote
command is actually executed and the original output_sender is preserved for
results; keep the debug log, remove the fake send, and ensure you don't consume
or drop the BashExec event so the normal executor (the function that handles
AiEvent::BashExec) receives it.
In `@crates/aish-shell/src/app.rs`:
- Around line 2417-2425: The specialized system prompt chosen earlier for the
error-correction branch (the (context, system_msg_t) selection that builds the
cmd_error prompt) is being overwritten by later unconditional clones
(system_msg_t = system_msg.clone() and system_msg_th = system_msg.clone());
remove or guard those later assignments so the previously-selected
system_msg_t/system_msg_th are preserved for the error-correction path (i.e.,
only clone/assign system_msg when no branch-specific prompt was chosen), and
apply the same fix for the analogous assignments around the other occurrence
noted (the block at the 2571-2573 equivalent).
In `@crates/aish-tools/src/channel_ask_user.rs`:
- Around line 105-123: The JSON schema in parameters() advertises "placeholder"
and "required" but execute() discards them and AskUserRequest lacks fields to
store them; either remove those keys from the schema or fully implement them:
add placeholder: Option<String> and required: Option<bool> (with default true)
to AskUserRequest, update parameters() to reflect the true defaults, and modify
execute() to read AskUserRequest.placeholder and AskUserRequest.required to
change prompt rendering and validation (e.g., enforce required vs optional input
and respect min_length/allow_cancel). Apply the same fix to the other schema
blocks noted (around the other parameter() definitions at the ranges mentioned).
In `@crates/aish-tools/src/channel_bash.rs`:
- Around line 43-46: The schema defines a configurable "timeout" (default 120)
but the code uses a hard-coded 120s wait; update the execution logic to read the
supplied "timeout" parameter and use it for the actual wait/timeout instead of
the constant. Locate uses of the literal 120 in this module (the "timeout" field
in the schema and the execution/wait logic around it — the two places noted
around lines 43–46 and 72–79) and replace the constant with the parsed timeout
value from the tool parameters (validate/parse the "timeout" param as an integer
and apply it to the timeout/wait call, e.g., passing it into whatever function
implements the command execution timeout).
---
Outside diff comments:
In `@crates/aish-pty/src/persistent.rs`:
- Around line 495-546: The code treats sel == 0 (a single 50ms idle poll) as a
definitive “command finished” event and invokes followup logic (involving
followup_capturing, looks_like_continuation_prompt, pending_followup,
pending_response, etc.), which is unreliable over SSH; change the logic so a
single timeout does not trigger followup or mark_prompt_ready. Specifically, in
the sel == 0 branch (within the draining/is_session/followup_capturing
handling), require a stronger prompt confirmation before treating the shell as
ready: either detect a real prompt via an existing prompt-detection routine
(e.g. extend or reuse looks_like_continuation_prompt/strip_ansi_and_prompt to
detect a normal PS1 prompt) or require N consecutive empty polls/timeouts before
invoking interceptor.mark_prompt_ready() and executing pending_followup; only
send the Ctrl+C or '\r' and clear followup_capturing after that confirmed
prompt/threshold is reached. Update any related state (skip_leading_newline,
pending_followup, followup_captured) accordingly so followups are fired only on
confirmed prompt presence rather than a single 50ms idle.
---
Nitpick comments:
In `@crates/aish-pty/src/output_buffer.rs`:
- Around line 22-30: The append method currently writes bytes one-by-one;
refactor aish-pty::output_buffer::append to copy the input slice in at most two
contiguous chunks into self.data using slice-based operations (e.g.,
copy_from_slice or extend_from_slice) instead of per-byte assignment: compute
the first writable span from write_pos to end-of-buffer, copy min(input.len(),
span.len()) bytes, update write_pos (wrap with modulo capacity) and len (capped
at capacity), then if remaining input exists copy the rest into the beginning
span and update write_pos and len again; make sure to handle empty input and
correctly wrap when input length > capacity.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 203596f1-2251-4a71-a079-fd28e649af7b
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (13)
crates/aish-i18n/locales/en-US.yamlcrates/aish-i18n/locales/zh-CN.yamlcrates/aish-pty/Cargo.tomlcrates/aish-pty/src/lib.rscrates/aish-pty/src/output_buffer.rscrates/aish-pty/src/persistent.rscrates/aish-pty/src/session_interceptor.rscrates/aish-shell/src/ai_handler.rscrates/aish-shell/src/app.rscrates/aish-shell/src/renderer.rscrates/aish-tools/src/channel_ask_user.rscrates/aish-tools/src/channel_bash.rscrates/aish-tools/src/lib.rs
- Add zero-capacity guard in OutputBuffer::new() - Honor timeout parameter in ChannelBashTool instead of hardcoded 120s - Remove unimplemented placeholder/required from ChannelAskUserTool schema - Drain trailing stdin bytes after confirmation read to prevent leaking - Replace hardcoded Chinese strings with i18n in ask_user and min-length - Actually execute BashExec commands in handle_ask_user instead of faking - Preserve branch-selected system prompt in session AI callback - Require N consecutive idle polls before treating shell as idle
…SH sessions - Preserve event_receiver via Arc<Mutex<Option<Receiver>>> to keep the LLM channel alive across multiple tool calls (make_chain_followup) - Call followup with cancellation message when user rejects a command, preventing "Channel closed" errors in both pending_response handlers - Suppress remote shell echo of injected commands
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
crates/aish-shell/src/app.rs (3)
2214-2221: 💤 Low valueUse
*const c_void(not*mut) forlibc::writebuf pointer.The buffers passed to
nix::libc::writehere are read-only*const u8frommsg.as_ptr(), but get cast to*mut nix::libc::c_void. The same file casts correctly to*const nix::libc::c_voidat lines 1939 and 2072 — the inconsistency is misleading even though the implicit*mut → *constcoercion makes it compile.♻️ Suggested fix
- unsafe { - nix::libc::write( - nix::libc::STDOUT_FILENO, - msg.as_ptr() as *mut nix::libc::c_void, - msg.len(), - ); - } + unsafe { + nix::libc::write( + nix::libc::STDOUT_FILENO, + msg.as_ptr() as *const nix::libc::c_void, + msg.len(), + ); + }(apply at both call sites — line 2217 and line 2234)
Also applies to: 2231-2237
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aish-shell/src/app.rs` around lines 2214 - 2221, The write calls that use msg.as_ptr() currently cast the buffer to *mut nix::libc::c_void; change those casts to *const nix::libc::c_void (e.g., use msg.as_ptr() as *const nix::libc::c_void) so the read-only byte slice is passed as a const pointer — update the write invocation(s) that reference msg (the same places where earlier code at the other write sites used *const) to keep the API usage consistent.
2118-2132: ⚡ Quick winUse
is_some()instead ofif let Some(_).
if let Some(_) = next_cmdtriggersclippy::redundant_pattern_matchingand obscures intent —next_cmdis only consumed inside the branch.♻️ Suggested fix
-let next_cmd = text.as_ref().and_then(|t| extract_bash_command(t)); -let ai_response = if let Some(_) = next_cmd { - let next_followup = Self::build_followup_closure( - &api_base_th, &api_key_th, &model_th, temperature, max_tokens, - &system_msg_th, &question_th, &anim_th, &conversation_history_th, - ); - Some(aish_pty::AiResponse { - command: next_cmd, - display_text: String::new(), - followup: Some(next_followup), - ask_user: None, - }) -} else { - None -}; +let ai_response = text + .as_deref() + .and_then(extract_bash_command) + .map(|cmd| { + let next_followup = Self::build_followup_closure( + &api_base_th, &api_key_th, &model_th, temperature, max_tokens, + &system_msg_th, &question_th, &anim_th, &conversation_history_th, + ); + aish_pty::AiResponse { + command: Some(cmd), + display_text: String::new(), + followup: Some(next_followup), + ask_user: None, + } + });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aish-shell/src/app.rs` around lines 2118 - 2132, Replace the pattern match "if let Some(_) = next_cmd" with a simple boolean check "if next_cmd.is_some()" in the ai_response construction so you avoid clippy::redundant_pattern_matching; keep the rest of the block (the next_cmd variable, the aish_pty::AiResponse creation, and the call to Self::build_followup_closure) unchanged so next_cmd is still available to be moved into the AiResponse.command field.
1865-2259: 🏗️ Heavy liftSignificant duplication between
build_followup_closureandbuild_session_ai_callback.The two functions share large near-identical blocks: the LLM event callback that draws the reasoning overlay (≈100 lines, including spinner/elapsed/i18n header, line truncation, and reasoning-line cleanup), the channel + cancellation token plumbing, the stdin select loop for Ctrl+C detection, and the post-loop residual reasoning cleanup. Maintaining two copies in sync (e.g., the i18n keys, the timeout value, the
as *mut/*const c_voidcast, the spinner state) is fragile — any future bugfix has to be applied twice and is easy to miss.Consider extracting:
- a
SessionLlmEventCallbackstruct/builder that owns the reasoning state and produces theArc<dyn Fn(LlmEvent) -> _>for both call sites,- a small
wait_for_callback_eventhelper for thetry_recv+select(STDIN_FILENO)Ctrl+C polling loop,- and a shared
clear_reasoning_overlayfn for the post-loop cleanup.This is a deferrable refactor, not a blocker, but the symmetry will only get harder to preserve as the SSH flow evolves.
Also applies to: 2261-2843
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aish-shell/src/app.rs` around lines 1865 - 2259, build_followup_closure duplicates large blocks (LLM event callback, reasoning overlay rendering, cancellation/token/channel plumbing, stdin select loop, and residual cleanup) that are also present in build_session_ai_callback; extract the shared logic into reusable pieces: create a SessionLlmEventCallback struct (or builder) that owns reasoning state and exposes an Arc<dyn Fn(aish_core::LlmEvent) -> Option<_>> for use in both build_followup_closure and build_session_ai_callback, factor the try_recv+select(STDIN_FILENO) polling loop into a wait_for_callback_event(...) helper that returns AiEvent/None and accepts the event_rx, token_rx, anim handle, cancelled flag and timeout, and move the terminal cleanup code into a clear_reasoning_overlay(...) utility used from the callback and after the wait loop so both call sites share the same spinner/elapsed/i18n/truncation behavior and timeouts.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aish-shell/src/app.rs`:
- Around line 1555-1561: The hot path problem is that build_session_ai_callback
(called from execute_external_command) eagerly constructs a PromptManager, calls
PromptManager::default_dir() and pm.load_all(), and renders the
oracle/system_msg even when remote_host is None; change
build_session_ai_callback to short-circuit and return a no-op callback
immediately when remote_host.is_none(), or alternatively add a cached
PromptManager and cached system_msg on AishShell (e.g., OnceLock fields) and
have build_session_ai_callback reuse them instead of calling pm.load_all() per
invocation; also remove the duplicated pm.load_all()/pm.render(...) in the
error-correction branch inside the callback so both paths share the same cached
PromptManager/system_msg usage.
- Around line 1934-1942: Replace all hardcoded English error/timeout literals in
the session-AI/followup callbacks with calls to aish_i18n::t_with_args (or t)
using new translation keys (e.g., "shell.session.followup_error",
"shell.session.llm_error", "shell.session.llm_timeout") and pass the error
string as the "{error}" arg and the timeout seconds as a "{seconds}" arg so the
message is localized and the timeout value is parameterized; update occurrences
around the AiEvent::Done send/format block (refer to the format! call that
builds "\r\n\x1b[31mFollowup error: {}\x1b[0m\r\n"), the followup LLM error
block (the other format! with "Followup LLM error"), the raw byte timeout
(b"...LLM timeout (60s)..."), and the eprintln! sites so they all use
t_with_args and interpolate error/seconds instead of embedding English literals.
---
Nitpick comments:
In `@crates/aish-shell/src/app.rs`:
- Around line 2214-2221: The write calls that use msg.as_ptr() currently cast
the buffer to *mut nix::libc::c_void; change those casts to *const
nix::libc::c_void (e.g., use msg.as_ptr() as *const nix::libc::c_void) so the
read-only byte slice is passed as a const pointer — update the write
invocation(s) that reference msg (the same places where earlier code at the
other write sites used *const) to keep the API usage consistent.
- Around line 2118-2132: Replace the pattern match "if let Some(_) = next_cmd"
with a simple boolean check "if next_cmd.is_some()" in the ai_response
construction so you avoid clippy::redundant_pattern_matching; keep the rest of
the block (the next_cmd variable, the aish_pty::AiResponse creation, and the
call to Self::build_followup_closure) unchanged so next_cmd is still available
to be moved into the AiResponse.command field.
- Around line 1865-2259: build_followup_closure duplicates large blocks (LLM
event callback, reasoning overlay rendering, cancellation/token/channel
plumbing, stdin select loop, and residual cleanup) that are also present in
build_session_ai_callback; extract the shared logic into reusable pieces: create
a SessionLlmEventCallback struct (or builder) that owns reasoning state and
exposes an Arc<dyn Fn(aish_core::LlmEvent) -> Option<_>> for use in both
build_followup_closure and build_session_ai_callback, factor the
try_recv+select(STDIN_FILENO) polling loop into a wait_for_callback_event(...)
helper that returns AiEvent/None and accepts the event_rx, token_rx, anim
handle, cancelled flag and timeout, and move the terminal cleanup code into a
clear_reasoning_overlay(...) utility used from the callback and after the wait
loop so both call sites share the same spinner/elapsed/i18n/truncation behavior
and timeouts.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 87235526-8d6d-4c67-ad15-af49c98f43e0
📒 Files selected for processing (7)
crates/aish-i18n/locales/en-US.yamlcrates/aish-i18n/locales/zh-CN.yamlcrates/aish-pty/src/output_buffer.rscrates/aish-pty/src/persistent.rscrates/aish-shell/src/app.rscrates/aish-tools/src/channel_ask_user.rscrates/aish-tools/src/channel_bash.rs
✅ Files skipped from review due to trivial changes (3)
- crates/aish-i18n/locales/en-US.yaml
- crates/aish-pty/src/output_buffer.rs
- crates/aish-tools/src/channel_bash.rs
🚧 Files skipped from review as they are similar to previous changes (3)
- crates/aish-i18n/locales/zh-CN.yaml
- crates/aish-tools/src/channel_ask_user.rs
- crates/aish-pty/src/persistent.rs
Summary
ChannelBashToolandChannelAskUserToolthat communicate via channels instead of direct terminal control, enabling AI tool use in SSH/telnet sessionspersistent.rsfor multi-round LLM interactions where tool output feeds back into the model for further analysisapp.rsfor session commandsTest plan
Summary by CodeRabbit
New Features
Improvements