M0: compilable skeleton — Kigi 0.1.0 fork surgery
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.
Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
ptyctl, ptyctl-cli, third_party/ unchanged; proto package
xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
(templates re-encrypted)
Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
module & dc_log, heap-profile uploader, auth-diagnostics uploader,
session-analytics halves of feedback; local zero-egress observability
preserved in new kigi-log crate (unified log, --debug firehose,
subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
shell util
Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted
Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean
Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
(new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
fast-worktree); RSS measurement tests serialized via serial_test
Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
notices sustained; kigi-tools ported-code notices extended; README,
CONTRIBUTING, SECURITY, AGENTS.md rewritten
Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
@@ -0,0 +1,118 @@
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
|
||||
use super::runner::{AsyncTerminalRunner, TerminalError, TerminalRunRequest, TerminalRunResult};
|
||||
|
||||
pub struct AcpTerminalRunner {
|
||||
pub gateway: GatewaySender,
|
||||
pub session_id: acp::SessionId,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
// Terminal release on cancel is now handled by kill_and_release_all_for_session()
|
||||
// in cancel_running_task() — see acp_session.rs.
|
||||
impl AsyncTerminalRunner for AcpTerminalRunner {
|
||||
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, TerminalError> {
|
||||
let session_id = self.session_id.clone();
|
||||
// On Windows the ACP client spawns with its own shell; sending the
|
||||
// raw command avoids the /bin/bash dependency.
|
||||
#[cfg(unix)]
|
||||
let command = {
|
||||
let quoted =
|
||||
shlex::try_quote(&request.command).map_err(|_| TerminalError::CommandNotQuoted)?;
|
||||
format!("{} -lc {}", super::default_shell_path(), quoted)
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let command = request.command.clone();
|
||||
let create_res = self
|
||||
.gateway
|
||||
.send(
|
||||
acp::CreateTerminalRequest::new(session_id.clone(), command)
|
||||
.args(vec![])
|
||||
.env(
|
||||
request
|
||||
.env
|
||||
.into_iter()
|
||||
.map(|(name, value)| acp::EnvVariable::new(name, value))
|
||||
.collect::<Vec<_>>(),
|
||||
)
|
||||
.cwd(Some(request.cwd.to_path_buf()))
|
||||
.output_byte_limit(Some(request.output_byte_limit as u64)),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| TerminalError::Other(e.to_string()))?;
|
||||
|
||||
// notify the client about the terminal
|
||||
let notification = acp::SessionNotification::new(
|
||||
session_id.clone(),
|
||||
acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate::new(
|
||||
request.tool_call_id.clone(),
|
||||
acp::ToolCallUpdateFields::new()
|
||||
.status(Some(acp::ToolCallStatus::InProgress))
|
||||
.content(Some(vec![acp::ToolCallContent::Terminal(
|
||||
acp::Terminal::new(create_res.terminal_id.clone()),
|
||||
)])),
|
||||
)),
|
||||
);
|
||||
let _ = self.gateway.send(notification).await;
|
||||
|
||||
let result = tokio::time::timeout(
|
||||
request.timeout,
|
||||
self.gateway.send(acp::WaitForTerminalExitRequest::new(
|
||||
session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
|
||||
let timed_out = match result {
|
||||
Ok(Ok(_)) => false,
|
||||
Ok(Err(e)) => return Err(TerminalError::Other(e.to_string())),
|
||||
Err(_) => {
|
||||
// timeout occurred, need to stop the command
|
||||
let _ = self
|
||||
.gateway
|
||||
.send(acp::KillTerminalRequest::new(
|
||||
session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
))
|
||||
.await;
|
||||
true
|
||||
}
|
||||
};
|
||||
|
||||
let output = self
|
||||
.gateway
|
||||
.send(acp::TerminalOutputRequest::new(
|
||||
session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
))
|
||||
.await
|
||||
.map_err(|e| TerminalError::Other(e.to_string()))?;
|
||||
|
||||
let _ = self
|
||||
.gateway
|
||||
.send(acp::ReleaseTerminalRequest::new(
|
||||
session_id,
|
||||
create_res.terminal_id,
|
||||
))
|
||||
.await;
|
||||
|
||||
let exit_status = output.exit_status.clone();
|
||||
let combined_output = output.output;
|
||||
let truncated = output.truncated;
|
||||
|
||||
let exit_code = exit_status
|
||||
.as_ref()
|
||||
.and_then(|e| e.exit_code.map(|v| v as i32));
|
||||
let signal = exit_status.and_then(|e| e.signal);
|
||||
|
||||
Ok(TerminalRunResult {
|
||||
combined_output,
|
||||
exit_code,
|
||||
truncated,
|
||||
signal,
|
||||
timed_out,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,746 @@
|
||||
//! AcpTerminalAdapter: implements `kigi-tools::TerminalBackend` using ACP gateway calls.
|
||||
//!
|
||||
//! This adapter enables bash tool execution over ACP (remote execution).
|
||||
//! It translates kigi-tools' `TerminalBackend` trait into ACP protocol calls:
|
||||
//! `run()` → create_terminal → wait_for_exit → terminal_output → release_terminal
|
||||
//! `run_background()` → create_terminal + spawn exit watcher
|
||||
//! `get_task()` → terminal_output (merged with tracked metadata)
|
||||
//! `kill_task()` → kill_terminal_command (watcher detects exit)
|
||||
//! `wait_for_completion()` → wait_for_terminal_exit with timeout
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
use kigi_tools::computer::types::{
|
||||
BackgroundHandle, ComputerError, KillOutcome, TaskSnapshot, TerminalBackend,
|
||||
TerminalRunRequest, TerminalRunResult,
|
||||
};
|
||||
use kigi_tools::notification::types::ToolNotificationHandle;
|
||||
|
||||
// ── Tracked task state ───────────────────────────────────────────────
|
||||
|
||||
struct TrackedTask {
|
||||
command: String,
|
||||
display_command: Option<String>,
|
||||
cwd: String,
|
||||
output_file: PathBuf,
|
||||
start_time: std::time::SystemTime,
|
||||
completed: bool,
|
||||
exit_code: Option<i32>,
|
||||
signal: Option<String>,
|
||||
last_output: String,
|
||||
last_truncated: bool,
|
||||
block_waited: bool,
|
||||
explicitly_killed: bool,
|
||||
}
|
||||
|
||||
impl TrackedTask {
|
||||
fn mark_completed(
|
||||
&mut self,
|
||||
exit_code: Option<i32>,
|
||||
signal: Option<String>,
|
||||
output: String,
|
||||
truncated: bool,
|
||||
) {
|
||||
self.completed = true;
|
||||
self.exit_code = exit_code;
|
||||
self.signal = signal;
|
||||
self.last_output = output;
|
||||
self.last_truncated = truncated;
|
||||
}
|
||||
|
||||
fn to_snapshot(
|
||||
&self,
|
||||
task_id: &str,
|
||||
output: String,
|
||||
truncated: bool,
|
||||
exit_code: Option<i32>,
|
||||
signal: Option<String>,
|
||||
) -> TaskSnapshot {
|
||||
let completed = self.completed || exit_code.is_some();
|
||||
TaskSnapshot {
|
||||
task_id: task_id.to_string(),
|
||||
command: self.command.clone(),
|
||||
display_command: self.display_command.clone(),
|
||||
cwd: self.cwd.clone(),
|
||||
start_time: self.start_time,
|
||||
end_time: completed.then(std::time::SystemTime::now),
|
||||
output,
|
||||
output_file: self.output_file.clone(),
|
||||
truncated,
|
||||
exit_code,
|
||||
signal,
|
||||
completed,
|
||||
block_waited: self.block_waited,
|
||||
explicitly_killed: self.explicitly_killed,
|
||||
kind: kigi_tools::computer::types::TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type TaskMap = Arc<Mutex<HashMap<String, TrackedTask>>>;
|
||||
|
||||
// ── Exit watcher ─────────────────────────────────────────────────────
|
||||
|
||||
/// Spawned per background task. Blocks on `WaitForTerminalExitRequest`,
|
||||
/// then fetches final output, emits `TaskCompleted`, and releases the
|
||||
/// terminal.
|
||||
async fn watch_for_exit(
|
||||
gateway: GatewaySender,
|
||||
session_id: acp::SessionId,
|
||||
task_id: String,
|
||||
tasks: TaskMap,
|
||||
notification_handle: ToolNotificationHandle,
|
||||
) {
|
||||
let terminal_id = acp::TerminalId::new(task_id.clone());
|
||||
|
||||
match gateway
|
||||
.send(acp::WaitForTerminalExitRequest::new(
|
||||
session_id.clone(),
|
||||
terminal_id.clone(),
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
task_id,
|
||||
error = %e,
|
||||
"watch_for_exit: gateway error waiting for terminal exit, polling until exit"
|
||||
);
|
||||
if !poll_for_terminal_exit(&gateway, &session_id, &terminal_id, None).await {
|
||||
// Gateway lost — mark the task as completed so it doesn't
|
||||
// remain as a ghost "running" entry forever.
|
||||
let snapshot = {
|
||||
let mut tasks = tasks.lock().unwrap();
|
||||
let Some(task) = tasks.get_mut(&task_id) else {
|
||||
return;
|
||||
};
|
||||
task.mark_completed(None, Some("gateway-lost".into()), String::new(), false);
|
||||
task.to_snapshot(
|
||||
&task_id,
|
||||
String::new(),
|
||||
false,
|
||||
None,
|
||||
Some("gateway-lost".into()),
|
||||
)
|
||||
};
|
||||
notification_handle.send_task_complete(snapshot);
|
||||
let _ = gateway
|
||||
.send(acp::ReleaseTerminalRequest::new(session_id, terminal_id))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (exit_code, signal, output_text, truncated) = match gateway
|
||||
.send(acp::TerminalOutputRequest::new(
|
||||
session_id.clone(),
|
||||
terminal_id.clone(),
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(o) => {
|
||||
let (code, sig) = parse_exit(&o.exit_status);
|
||||
(code, sig, o.output, o.truncated)
|
||||
}
|
||||
Err(_) => (None, None, String::new(), false),
|
||||
};
|
||||
|
||||
let snapshot = {
|
||||
let mut tasks = tasks.lock().unwrap();
|
||||
let Some(task) = tasks.get_mut(&task_id) else {
|
||||
return;
|
||||
};
|
||||
task.mark_completed(exit_code, signal.clone(), output_text.clone(), truncated);
|
||||
task.to_snapshot(&task_id, output_text, truncated, exit_code, signal)
|
||||
};
|
||||
|
||||
notification_handle.send_task_complete(snapshot);
|
||||
|
||||
let _ = gateway
|
||||
.send(acp::ReleaseTerminalRequest::new(session_id, terminal_id))
|
||||
.await;
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Poll `TerminalOutputRequest` at 500ms intervals until `exit_status` is
|
||||
/// present, a deadline is hit, or 60 consecutive gateway errors occur.
|
||||
/// Returns `true` when an exit was detected.
|
||||
async fn poll_for_terminal_exit(
|
||||
gateway: &GatewaySender,
|
||||
session_id: &acp::SessionId,
|
||||
terminal_id: &acp::TerminalId,
|
||||
deadline: Option<tokio::time::Instant>,
|
||||
) -> bool {
|
||||
let mut consecutive_errors = 0u32;
|
||||
loop {
|
||||
if let Some(dl) = deadline
|
||||
&& tokio::time::Instant::now() >= dl
|
||||
{
|
||||
return false;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
match gateway
|
||||
.send(acp::TerminalOutputRequest::new(
|
||||
session_id.clone(),
|
||||
terminal_id.clone(),
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(output) => {
|
||||
consecutive_errors = 0;
|
||||
if output.exit_status.is_some() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
consecutive_errors += 1;
|
||||
if consecutive_errors >= 60 {
|
||||
tracing::error!(
|
||||
terminal_id = %terminal_id.0,
|
||||
error = %e,
|
||||
"gateway unreachable after 60 consecutive poll failures"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn wrap_command(command: &str) -> Result<String, ComputerError> {
|
||||
// On Windows the ACP client (grok-desktop) spawns with `shell: true`
|
||||
// which delegates to cmd.exe. Wrapping in /bin/bash would fail because
|
||||
// that path doesn't exist on Windows. Send the raw command instead.
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = command;
|
||||
Ok(command.to_string())
|
||||
}
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let quoted = shlex::try_quote(command).map_err(|_| ComputerError::CommandNotQuoted)?;
|
||||
Ok(format!(
|
||||
"{} -lc {quoted}",
|
||||
crate::terminal::default_shell_path()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn to_env(env: HashMap<String, String>) -> Vec<acp::EnvVariable> {
|
||||
env.into_iter()
|
||||
.map(|(name, value)| acp::EnvVariable::new(name, value))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_exit(status: &Option<acp::TerminalExitStatus>) -> (Option<i32>, Option<String>) {
|
||||
match status {
|
||||
Some(e) => (e.exit_code.map(|v| v as i32), e.signal.clone()),
|
||||
None => (None, None),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Adapter ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Wraps kigi-shell's ACP gateway to satisfy kigi-tools' TerminalBackend.
|
||||
pub struct AcpTerminalAdapter {
|
||||
gateway: GatewaySender,
|
||||
session_id: acp::SessionId,
|
||||
tasks: TaskMap,
|
||||
}
|
||||
|
||||
impl AcpTerminalAdapter {
|
||||
pub fn new(gateway: GatewaySender, session_id: acp::SessionId) -> Self {
|
||||
Self {
|
||||
gateway,
|
||||
session_id,
|
||||
tasks: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_terminal(
|
||||
&self,
|
||||
command: String,
|
||||
request: &TerminalRunRequest,
|
||||
) -> Result<acp::CreateTerminalResponse, ComputerError> {
|
||||
self.gateway
|
||||
.send(
|
||||
acp::CreateTerminalRequest::new(self.session_id.clone(), command)
|
||||
.args(vec![])
|
||||
.env(to_env(request.env.clone()))
|
||||
.cwd(Some(request.working_directory.clone()))
|
||||
.output_byte_limit(Some(request.output_byte_limit as u64)),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| ComputerError::io(e.to_string()))
|
||||
}
|
||||
|
||||
fn terminal_id(&self, task_id: &str) -> acp::TerminalId {
|
||||
acp::TerminalId::new(task_id)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl TerminalBackend for AcpTerminalAdapter {
|
||||
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, ComputerError> {
|
||||
let command = wrap_command(&request.command)?;
|
||||
let create_res = self.create_terminal(command, &request).await?;
|
||||
|
||||
let timed_out = match tokio::time::timeout(
|
||||
request.timeout,
|
||||
self.gateway.send(acp::WaitForTerminalExitRequest::new(
|
||||
self.session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(_)) => false,
|
||||
Ok(Err(e)) => return Err(ComputerError::io(e.to_string())),
|
||||
Err(_) => {
|
||||
let _ = self
|
||||
.gateway
|
||||
.send(acp::KillTerminalRequest::new(
|
||||
self.session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
))
|
||||
.await;
|
||||
true
|
||||
}
|
||||
};
|
||||
|
||||
let output = self
|
||||
.gateway
|
||||
.send(acp::TerminalOutputRequest::new(
|
||||
self.session_id.clone(),
|
||||
create_res.terminal_id.clone(),
|
||||
))
|
||||
.await
|
||||
.map_err(|e| ComputerError::io(e.to_string()))?;
|
||||
|
||||
let _ = self
|
||||
.gateway
|
||||
.send(acp::ReleaseTerminalRequest::new(
|
||||
self.session_id.clone(),
|
||||
create_res.terminal_id,
|
||||
))
|
||||
.await;
|
||||
|
||||
let (exit_code, signal) = parse_exit(&output.exit_status);
|
||||
let total_bytes = output.output.len();
|
||||
Ok(TerminalRunResult {
|
||||
combined_output: output.output,
|
||||
exit_code,
|
||||
truncated: output.truncated,
|
||||
signal,
|
||||
timed_out,
|
||||
output_file: request.output_file,
|
||||
total_bytes,
|
||||
// ACP gateway does not surface a local PID -- the process
|
||||
// runs on the remote side.
|
||||
pid: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn run_background(
|
||||
&self,
|
||||
request: TerminalRunRequest,
|
||||
) -> Result<BackgroundHandle, ComputerError> {
|
||||
let command = wrap_command(&request.command)?;
|
||||
let notification_handle = request.notification_handle.clone();
|
||||
let display_command = request.display_command.clone();
|
||||
let cwd = request.working_directory.to_string_lossy().to_string();
|
||||
let output_file = request.output_file.clone();
|
||||
|
||||
let create_res = self.create_terminal(command.clone(), &request).await?;
|
||||
let task_id = create_res.terminal_id.0.to_string();
|
||||
|
||||
{
|
||||
let mut tasks = self.tasks.lock().unwrap();
|
||||
tasks.insert(
|
||||
task_id.clone(),
|
||||
TrackedTask {
|
||||
command,
|
||||
display_command,
|
||||
cwd,
|
||||
output_file: output_file.clone(),
|
||||
start_time: std::time::SystemTime::now(),
|
||||
completed: false,
|
||||
exit_code: None,
|
||||
signal: None,
|
||||
last_output: String::new(),
|
||||
last_truncated: false,
|
||||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
tokio::spawn(watch_for_exit(
|
||||
self.gateway.clone(),
|
||||
self.session_id.clone(),
|
||||
task_id.clone(),
|
||||
Arc::clone(&self.tasks),
|
||||
notification_handle,
|
||||
));
|
||||
|
||||
Ok(BackgroundHandle {
|
||||
task_id,
|
||||
output_file,
|
||||
// ACP gateway does not surface a local PID -- the process
|
||||
// runs on the remote side.
|
||||
pid: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_task(&self, task_id: &str) -> Option<TaskSnapshot> {
|
||||
let live = self
|
||||
.gateway
|
||||
.send(acp::TerminalOutputRequest::new(
|
||||
self.session_id.clone(),
|
||||
self.terminal_id(task_id),
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
|
||||
let tasks = self.tasks.lock().unwrap();
|
||||
let tracked = tasks.get(task_id);
|
||||
|
||||
match (live, tracked) {
|
||||
(Some(output), Some(tracked)) => {
|
||||
let (exit_code, signal) = parse_exit(&output.exit_status);
|
||||
Some(tracked.to_snapshot(
|
||||
task_id,
|
||||
output.output,
|
||||
output.truncated,
|
||||
exit_code,
|
||||
signal,
|
||||
))
|
||||
}
|
||||
(Some(output), None) => {
|
||||
let (exit_code, signal) = parse_exit(&output.exit_status);
|
||||
let completed = exit_code.is_some();
|
||||
Some(TaskSnapshot {
|
||||
task_id: task_id.to_string(),
|
||||
command: String::new(),
|
||||
display_command: None,
|
||||
cwd: String::new(),
|
||||
start_time: std::time::SystemTime::now(),
|
||||
end_time: completed.then(std::time::SystemTime::now),
|
||||
output: output.output,
|
||||
output_file: PathBuf::new(),
|
||||
truncated: output.truncated,
|
||||
exit_code,
|
||||
signal,
|
||||
completed,
|
||||
kind: kigi_tools::computer::types::TaskKind::Bash,
|
||||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
})
|
||||
}
|
||||
(None, Some(tracked)) if tracked.completed => Some(tracked.to_snapshot(
|
||||
task_id,
|
||||
tracked.last_output.clone(),
|
||||
tracked.last_truncated,
|
||||
tracked.exit_code,
|
||||
tracked.signal.clone(),
|
||||
)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn kill_task(&self, task_id: &str) -> KillOutcome {
|
||||
// Mark as explicitly killed BEFORE sending the kill request so the
|
||||
// exit watcher's snapshot carries the flag.
|
||||
{
|
||||
let mut tasks = self.tasks.lock().unwrap();
|
||||
if let Some(task) = tasks.get_mut(task_id) {
|
||||
task.explicitly_killed = true;
|
||||
}
|
||||
}
|
||||
|
||||
match self
|
||||
.gateway
|
||||
.send(acp::KillTerminalRequest::new(
|
||||
self.session_id.clone(),
|
||||
self.terminal_id(task_id),
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(_) => KillOutcome::Killed,
|
||||
Err(_) => KillOutcome::NotFound,
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_completion(
|
||||
&self,
|
||||
task_id: &str,
|
||||
timeout: Option<Duration>,
|
||||
) -> Option<TaskSnapshot> {
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(30));
|
||||
|
||||
// Mark BEFORE waiting so watch_for_exit sees the flag in its snapshot.
|
||||
{
|
||||
let mut tasks = self.tasks.lock().unwrap();
|
||||
if let Some(task) = tasks.get_mut(task_id) {
|
||||
task.block_waited = true;
|
||||
}
|
||||
}
|
||||
|
||||
let gateway_result = tokio::time::timeout(
|
||||
timeout,
|
||||
self.gateway.send(acp::WaitForTerminalExitRequest::new(
|
||||
self.session_id.clone(),
|
||||
self.terminal_id(task_id),
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
|
||||
match &gateway_result {
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Err(e)) => {
|
||||
tracing::warn!(task_id, error = %e, "gateway error waiting for terminal exit, falling back to polling");
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
poll_for_terminal_exit(
|
||||
&self.gateway,
|
||||
&self.session_id,
|
||||
&self.terminal_id(task_id),
|
||||
Some(deadline),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!(task_id, "timeout waiting for terminal exit");
|
||||
// The block timed out: the agent did not receive the
|
||||
// completion result, so auto-wake should still fire
|
||||
// when the task eventually completes.
|
||||
let mut tasks = self.tasks.lock().unwrap();
|
||||
if let Some(task) = tasks.get_mut(task_id) {
|
||||
task.block_waited = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.get_task(task_id).await
|
||||
}
|
||||
|
||||
async fn list_tasks(&self) -> Vec<TaskSnapshot> {
|
||||
let task_ids: Vec<String> = {
|
||||
let tasks = self.tasks.lock().unwrap();
|
||||
tasks.keys().cloned().collect()
|
||||
};
|
||||
let mut snapshots = Vec::new();
|
||||
for task_id in task_ids {
|
||||
if let Some(snapshot) = self.get_task(&task_id).await {
|
||||
snapshots.push(snapshot);
|
||||
}
|
||||
}
|
||||
snapshots
|
||||
}
|
||||
|
||||
async fn kill_all_background_tasks(&self) {
|
||||
let task_ids: Vec<String> = {
|
||||
let tasks = self.tasks.lock().unwrap();
|
||||
tasks
|
||||
.iter()
|
||||
.filter(|(_, t)| !t.completed)
|
||||
.map(|(id, _)| id.clone())
|
||||
.collect()
|
||||
};
|
||||
for task_id in task_ids {
|
||||
self.kill_task(&task_id).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn kill_foreground_commands(&self) {
|
||||
let session_id = self.session_id.0.to_string();
|
||||
crate::terminal::kill_and_release_all_for_session(&session_id).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_tracked_task(command: &str) -> TrackedTask {
|
||||
TrackedTask {
|
||||
command: command.to_string(),
|
||||
display_command: None,
|
||||
cwd: "/tmp".to_string(),
|
||||
output_file: PathBuf::from("/tmp/out.log"),
|
||||
start_time: std::time::SystemTime::now(),
|
||||
completed: false,
|
||||
exit_code: None,
|
||||
signal: None,
|
||||
last_output: String::new(),
|
||||
last_truncated: false,
|
||||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrap_command_quotes_shell_metacharacters() {
|
||||
let cmd = wrap_command("echo 'hello world' && ls").unwrap();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
// The resolved bash path may live in any prefix (`/bin`,
|
||||
// `/opt/homebrew/bin`, `/run/current-system/sw/bin`, …), so just
|
||||
// assert the prefix shape: `<resolved-bash> -lc <quoted-cmd>`.
|
||||
let shell = crate::terminal::default_shell_path();
|
||||
assert!(
|
||||
cmd.starts_with(&format!("{shell} -lc")),
|
||||
"expected wrapped cmd to begin with `{shell} -lc`, got: {cmd}"
|
||||
);
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
assert_eq!(cmd, "echo 'hello world' && ls");
|
||||
assert!(cmd.contains("echo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_exit_with_code() {
|
||||
let status = Some(acp::TerminalExitStatus::new().exit_code(Some(42)));
|
||||
let (code, sig) = parse_exit(&status);
|
||||
assert_eq!(code, Some(42));
|
||||
assert_eq!(sig, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_exit_with_signal() {
|
||||
let status = Some(acp::TerminalExitStatus::new().signal(Some("SIGKILL".into())));
|
||||
let (code, sig) = parse_exit(&status);
|
||||
assert_eq!(code, None);
|
||||
assert_eq!(sig, Some("SIGKILL".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_exit_none() {
|
||||
assert_eq!(parse_exit(&None), (None, None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracked_task_mark_completed() {
|
||||
let mut task = make_tracked_task("sleep 10");
|
||||
assert!(!task.completed);
|
||||
assert_eq!(task.exit_code, None);
|
||||
|
||||
task.mark_completed(Some(137), Some("SIGTERM".into()), "output".into(), false);
|
||||
assert!(task.completed);
|
||||
assert_eq!(task.exit_code, Some(137));
|
||||
assert_eq!(task.signal, Some("SIGTERM".into()));
|
||||
assert_eq!(task.last_output, "output");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracked_task_to_snapshot_running() {
|
||||
let task = make_tracked_task("ls -la");
|
||||
let snap = task.to_snapshot("t-1", "file1\nfile2".into(), false, None, None);
|
||||
|
||||
assert_eq!(snap.task_id, "t-1");
|
||||
assert_eq!(snap.command, "ls -la");
|
||||
assert_eq!(snap.cwd, "/tmp");
|
||||
assert_eq!(snap.output, "file1\nfile2");
|
||||
assert!(!snap.completed);
|
||||
assert!(snap.end_time.is_none());
|
||||
assert_eq!(snap.exit_code, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracked_task_to_snapshot_completed() {
|
||||
let mut task = make_tracked_task("echo done");
|
||||
task.mark_completed(Some(0), None, "done\n".into(), false);
|
||||
let snap = task.to_snapshot("t-2", "done\n".into(), false, Some(0), None);
|
||||
|
||||
assert!(snap.completed);
|
||||
assert!(snap.end_time.is_some());
|
||||
assert_eq!(snap.exit_code, Some(0));
|
||||
assert_eq!(snap.signal, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracked_task_to_snapshot_completed_by_exit_code_alone() {
|
||||
let task = make_tracked_task("fast cmd");
|
||||
let snap = task.to_snapshot("t-3", String::new(), false, Some(1), None);
|
||||
assert!(snap.completed);
|
||||
assert!(snap.end_time.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracked_task_to_snapshot_preserves_display_command() {
|
||||
let mut task = make_tracked_task("/bin/bash -lc 'echo hi'");
|
||||
task.display_command = Some("echo hi".into());
|
||||
let snap = task.to_snapshot("t-4", String::new(), false, None, None);
|
||||
assert_eq!(snap.display_command, Some("echo hi".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_map_insert_and_mark_completed() {
|
||||
let tasks: TaskMap = Arc::new(Mutex::new(HashMap::new()));
|
||||
{
|
||||
let mut map = tasks.lock().unwrap();
|
||||
map.insert("t-1".into(), make_tracked_task("sleep 60"));
|
||||
}
|
||||
{
|
||||
let mut map = tasks.lock().unwrap();
|
||||
let task = map.get_mut("t-1").unwrap();
|
||||
task.mark_completed(Some(143), Some("SIGTERM".into()), String::new(), false);
|
||||
assert!(task.completed);
|
||||
}
|
||||
{
|
||||
let map = tasks.lock().unwrap();
|
||||
let task = map.get("t-1").unwrap();
|
||||
assert!(task.completed);
|
||||
assert_eq!(task.exit_code, Some(143));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_map_filter_running() {
|
||||
let tasks: TaskMap = Arc::new(Mutex::new(HashMap::new()));
|
||||
{
|
||||
let mut map = tasks.lock().unwrap();
|
||||
map.insert("running-1".into(), make_tracked_task("sleep 60"));
|
||||
let mut done = make_tracked_task("echo done");
|
||||
done.mark_completed(Some(0), None, String::new(), false);
|
||||
map.insert("done-1".into(), done);
|
||||
map.insert("running-2".into(), make_tracked_task("sleep 120"));
|
||||
}
|
||||
let running: Vec<String> = {
|
||||
let map = tasks.lock().unwrap();
|
||||
map.iter()
|
||||
.filter(|(_, t)| !t.completed)
|
||||
.map(|(id, _)| id.clone())
|
||||
.collect()
|
||||
};
|
||||
assert_eq!(running.len(), 2);
|
||||
assert!(running.contains(&"running-1".into()));
|
||||
assert!(running.contains(&"running-2".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completed_task_snapshot_uses_cached_output() {
|
||||
let mut task = make_tracked_task("echo hello");
|
||||
task.mark_completed(Some(0), None, "hello\n".into(), false);
|
||||
|
||||
let snap = task.to_snapshot(
|
||||
"t-5",
|
||||
task.last_output.clone(),
|
||||
task.last_truncated,
|
||||
task.exit_code,
|
||||
task.signal.clone(),
|
||||
);
|
||||
assert!(snap.completed);
|
||||
assert_eq!(snap.output, "hello\n");
|
||||
assert_eq!(snap.exit_code, Some(0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,772 @@
|
||||
//! Background task registry for tracking long-running commands.
|
||||
//!
|
||||
//! This module provides a per-session registry for background tasks that allows
|
||||
//! the model to query task status and output after launching commands with
|
||||
//! `is_background: true`.
|
||||
//!
|
||||
//! ## Architecture
|
||||
//!
|
||||
//! The registry works alongside the existing terminal infrastructure:
|
||||
//! - `StreamingLocalTerminalRunner` handles process spawning and output streaming
|
||||
//! - `BackgroundTaskRegistry` provides model-facing queries by task_id
|
||||
//!
|
||||
//! ## Output Storage
|
||||
//!
|
||||
//! Output is stored in two places:
|
||||
//! 1. **In memory (`output` field)**: May be truncated if > output_byte_limit
|
||||
//! 2. **On disk (`output_file`)**: Full output written incrementally
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{Mutex, Notify, RwLock};
|
||||
|
||||
/// Task identifier (UUID string)
|
||||
pub type TaskId = String;
|
||||
|
||||
/// Snapshot of a background task's current state.
|
||||
///
|
||||
/// This is a clone-able view of the task that can be returned to callers
|
||||
/// without holding locks.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskSnapshot {
|
||||
/// Unique task ID (UUID) given to the model
|
||||
pub task_id: TaskId,
|
||||
/// Internal tool_call_id for terminal registry lookup
|
||||
pub tool_call_id: String,
|
||||
/// The command that was executed
|
||||
pub command: String,
|
||||
/// Working directory where command was run
|
||||
pub cwd: String,
|
||||
/// Wall-clock start time
|
||||
pub start_time: DateTime<Utc>,
|
||||
/// Wall-clock end time (set when task completes)
|
||||
pub end_time: Option<DateTime<Utc>>,
|
||||
/// In-memory output (may be truncated if > output_byte_limit)
|
||||
pub output: String,
|
||||
/// Path to full output file on disk
|
||||
pub output_file: PathBuf,
|
||||
/// Whether in-memory output was truncated
|
||||
pub truncated: bool,
|
||||
/// Exit code if completed
|
||||
pub exit_code: Option<i32>,
|
||||
/// Signal name if terminated by signal
|
||||
pub signal: Option<String>,
|
||||
/// Whether task has completed (exited or was killed)
|
||||
pub completed: bool,
|
||||
/// Whether a blocking waiter has claimed this task.
|
||||
pub block_waited: bool,
|
||||
/// Whether this task was explicitly killed via the kill tool.
|
||||
pub explicitly_killed: bool,
|
||||
}
|
||||
|
||||
impl TaskSnapshot {
|
||||
/// Calculate duration in seconds.
|
||||
///
|
||||
/// If task is still running, returns time since start.
|
||||
/// If task completed, returns total runtime.
|
||||
pub fn duration_secs(&self) -> f64 {
|
||||
let end = self.end_time.unwrap_or_else(Utc::now);
|
||||
(end - self.start_time).num_milliseconds() as f64 / 1000.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal entry storing task data and completion notification
|
||||
struct TaskEntry {
|
||||
/// The task snapshot (protected by RwLock for concurrent reads)
|
||||
snapshot: RwLock<TaskSnapshot>,
|
||||
/// Notifier for waiters when task completes
|
||||
exit_notify: Arc<Notify>,
|
||||
}
|
||||
|
||||
/// Per-session registry for background tasks.
|
||||
///
|
||||
/// Each session has its own instance via `ToolContext.background_tasks`.
|
||||
/// This ensures:
|
||||
/// - Task IDs only need to be unique within a session
|
||||
/// - Session cleanup automatically cleans up tasks
|
||||
/// - No global state pollution between agents
|
||||
pub struct BackgroundTaskRegistry {
|
||||
/// Map from task_id -> entry
|
||||
tasks: Mutex<HashMap<TaskId, Arc<TaskEntry>>>,
|
||||
/// Maximum number of concurrent tasks
|
||||
max_tasks: usize,
|
||||
}
|
||||
|
||||
/// Default maximum number of concurrent background tasks per session
|
||||
const DEFAULT_MAX_BACKGROUND_TASKS: usize = 10;
|
||||
|
||||
impl Default for BackgroundTaskRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl BackgroundTaskRegistry {
|
||||
/// Create a new registry with default max tasks limit.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tasks: Mutex::new(HashMap::new()),
|
||||
max_tasks: DEFAULT_MAX_BACKGROUND_TASKS,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a registry with custom max tasks limit (for testing).
|
||||
pub fn with_max_tasks(max_tasks: usize) -> Self {
|
||||
Self {
|
||||
tasks: Mutex::new(HashMap::new()),
|
||||
max_tasks,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new background task.
|
||||
///
|
||||
/// If at capacity, completed tasks are cleaned up first.
|
||||
/// Returns error if still at capacity after cleanup.
|
||||
pub async fn register(&self, snapshot: TaskSnapshot) -> Result<(), String> {
|
||||
let mut tasks = self.tasks.lock().await;
|
||||
|
||||
// Cleanup completed tasks if at capacity
|
||||
if tasks.len() >= self.max_tasks {
|
||||
tasks.retain(|_, entry| {
|
||||
// Keep if not completed (check synchronously via try_read)
|
||||
entry
|
||||
.snapshot
|
||||
.try_read()
|
||||
.map(|s| !s.completed)
|
||||
.unwrap_or(true)
|
||||
});
|
||||
|
||||
if tasks.len() >= self.max_tasks {
|
||||
return Err(format!(
|
||||
"Maximum background tasks ({}) reached. Wait for tasks to complete or kill existing tasks.",
|
||||
self.max_tasks
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let task_id = snapshot.task_id.clone();
|
||||
let entry = Arc::new(TaskEntry {
|
||||
snapshot: RwLock::new(snapshot),
|
||||
exit_notify: Arc::new(Notify::new()),
|
||||
});
|
||||
tasks.insert(task_id, entry);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get current snapshot of a task.
|
||||
///
|
||||
/// Returns `None` if task_id is not found.
|
||||
pub async fn get(&self, task_id: &str) -> Option<TaskSnapshot> {
|
||||
let tasks = self.tasks.lock().await;
|
||||
let entry = tasks.get(task_id)?;
|
||||
Some(entry.snapshot.read().await.clone())
|
||||
}
|
||||
|
||||
/// Update task output (called by output collector after completion).
|
||||
pub async fn update_output(&self, task_id: &str, output: String, truncated: bool) {
|
||||
let tasks = self.tasks.lock().await;
|
||||
if let Some(entry) = tasks.get(task_id) {
|
||||
let mut snapshot = entry.snapshot.write().await;
|
||||
snapshot.output = output;
|
||||
snapshot.truncated = truncated;
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark task as completed.
|
||||
///
|
||||
/// This sets the end_time, exit_code/signal, and notifies any waiters.
|
||||
pub async fn mark_completed(
|
||||
&self,
|
||||
task_id: &str,
|
||||
exit_code: Option<i32>,
|
||||
signal: Option<String>,
|
||||
) {
|
||||
let tasks = self.tasks.lock().await;
|
||||
if let Some(entry) = tasks.get(task_id) {
|
||||
{
|
||||
let mut snapshot = entry.snapshot.write().await;
|
||||
snapshot.completed = true;
|
||||
snapshot.exit_code = exit_code;
|
||||
snapshot.signal = signal;
|
||||
snapshot.end_time = Some(Utc::now());
|
||||
}
|
||||
// Notify all waiters that task has completed
|
||||
entry.exit_notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
/// Wait for task completion with optional timeout.
|
||||
///
|
||||
/// Returns the task snapshot after completion or timeout.
|
||||
/// Returns `None` if task_id is not found.
|
||||
pub async fn wait_for_completion(
|
||||
&self,
|
||||
task_id: &str,
|
||||
timeout: Option<std::time::Duration>,
|
||||
) -> Option<TaskSnapshot> {
|
||||
// Get entry without holding the lock during wait
|
||||
let entry = {
|
||||
let tasks = self.tasks.lock().await;
|
||||
tasks.get(task_id)?.clone()
|
||||
};
|
||||
|
||||
// Register notification interest BEFORE checking completion to avoid a
|
||||
// race where mark_completed fires between the check and the wait,
|
||||
// causing notify_waiters() to wake zero futures and the notification
|
||||
// to be permanently lost.
|
||||
let notified = entry.exit_notify.notified();
|
||||
tokio::pin!(notified);
|
||||
notified.as_mut().enable();
|
||||
|
||||
{
|
||||
let snapshot = entry.snapshot.read().await;
|
||||
if snapshot.completed {
|
||||
return Some(snapshot.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(timeout) = timeout {
|
||||
let _ = tokio::time::timeout(timeout, notified).await;
|
||||
} else {
|
||||
notified.await;
|
||||
}
|
||||
|
||||
Some(entry.snapshot.read().await.clone())
|
||||
}
|
||||
|
||||
/// Get a cloneable notification handle for a specific task.
|
||||
///
|
||||
/// Used by multi-wait to select across multiple task exit notifications.
|
||||
/// Returns `None` if the task is not registered.
|
||||
pub async fn get_exit_notify(&self, task_id: &str) -> Option<Arc<Notify>> {
|
||||
let tasks = self.tasks.lock().await;
|
||||
tasks.get(task_id).map(|e| Arc::clone(&e.exit_notify))
|
||||
}
|
||||
|
||||
/// List all tasks in the registry.
|
||||
pub async fn list(&self) -> Vec<TaskSnapshot> {
|
||||
let tasks = self.tasks.lock().await;
|
||||
let mut result = Vec::with_capacity(tasks.len());
|
||||
for entry in tasks.values() {
|
||||
result.push(entry.snapshot.read().await.clone());
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Get number of active (non-completed) tasks.
|
||||
pub async fn active_count(&self) -> usize {
|
||||
let tasks = self.tasks.lock().await;
|
||||
let mut count = 0;
|
||||
for entry in tasks.values() {
|
||||
if !entry.snapshot.read().await.completed {
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
count
|
||||
}
|
||||
}
|
||||
|
||||
/// Get output file path for a background task.
|
||||
///
|
||||
/// Creates the directory structure if it doesn't exist.
|
||||
/// Path format: `~/.kigi/sessions/{session_id}/tasks/{task_id}.log`
|
||||
pub fn get_task_output_path(session_id: &str, task_id: &str) -> PathBuf {
|
||||
use crate::util::kigi_home::kigi_home;
|
||||
|
||||
let tasks_dir = kigi_home().join("sessions").join(session_id).join("tasks");
|
||||
// Create directory (ignore errors - will fail on write if dir creation fails)
|
||||
std::fs::create_dir_all(&tasks_dir).ok();
|
||||
tasks_dir.join(format!("{}.log", task_id))
|
||||
}
|
||||
|
||||
// ── Background task manifest for session resume ──
|
||||
|
||||
const MANIFEST_FILENAME: &str = "background_tasks_manifest.json";
|
||||
|
||||
/// Minimal snapshot of a running background task, persisted on session exit
|
||||
/// so that a resumed session can inform the model about orphaned tasks.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BackgroundTaskManifestEntry {
|
||||
pub task_id: String,
|
||||
pub command: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub display_command: Option<String>,
|
||||
pub output_file: PathBuf,
|
||||
pub start_time: std::time::SystemTime,
|
||||
pub cwd: String,
|
||||
#[serde(default)]
|
||||
pub kind: kigi_tools::computer::types::TaskKind,
|
||||
}
|
||||
|
||||
/// Persist a manifest of running background tasks to the session directory.
|
||||
///
|
||||
/// Only writes a file when `entries` is non-empty. Called during session
|
||||
/// shutdown when background tasks are intentionally left alive.
|
||||
pub fn persist_manifest(session_dir: &Path, entries: Vec<BackgroundTaskManifestEntry>) {
|
||||
if entries.is_empty() {
|
||||
return;
|
||||
}
|
||||
let path = session_dir.join(MANIFEST_FILENAME);
|
||||
match serde_json::to_vec(&entries) {
|
||||
Ok(data) => {
|
||||
if let Err(e) = std::fs::write(&path, data) {
|
||||
tracing::warn!(%e, "failed to write background task manifest");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(%e, "failed to serialize background task manifest");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the background task manifest from the session directory and delete it.
|
||||
///
|
||||
/// Returns an empty vec if the manifest doesn't exist or can't be parsed.
|
||||
pub fn load_and_clear_manifest(session_dir: &Path) -> Vec<BackgroundTaskManifestEntry> {
|
||||
let path = session_dir.join(MANIFEST_FILENAME);
|
||||
let data = match std::fs::read(&path) {
|
||||
Ok(d) => d,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
// Delete regardless of parse success — stale manifests should not accumulate.
|
||||
let _ = std::fs::remove_file(&path);
|
||||
serde_json::from_slice(&data).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Format a system-reminder about background tasks that were running when the
|
||||
/// session was last active.
|
||||
pub fn format_resumed_tasks_reminder(entries: &[BackgroundTaskManifestEntry]) -> String {
|
||||
use std::fmt::Write;
|
||||
|
||||
let now = std::time::SystemTime::now();
|
||||
let mut buf = String::from(
|
||||
"This session was resumed. The following background tasks were running \
|
||||
when the session was last active and may still be in progress:\n",
|
||||
);
|
||||
for entry in entries {
|
||||
let cmd = entry.display_command.as_deref().unwrap_or(&entry.command);
|
||||
let ago = format_duration_ago(now, entry.start_time);
|
||||
let kind_label = match entry.kind {
|
||||
kigi_tools::computer::types::TaskKind::Monitor => " [monitor]",
|
||||
kigi_tools::computer::types::TaskKind::Bash => "",
|
||||
};
|
||||
let _ = writeln!(
|
||||
buf,
|
||||
"- \"{}\"{} (started {}): {}",
|
||||
entry.task_id, kind_label, ago, cmd
|
||||
);
|
||||
let _ = writeln!(buf, " Output log: {}", entry.output_file.display());
|
||||
}
|
||||
buf.push_str(
|
||||
"Check whether each is still running and read its output log to determine \
|
||||
if it completed successfully.",
|
||||
);
|
||||
buf
|
||||
}
|
||||
|
||||
fn format_duration_ago(now: std::time::SystemTime, start: std::time::SystemTime) -> String {
|
||||
let secs = now.duration_since(start).map(|d| d.as_secs()).unwrap_or(0);
|
||||
if secs < 60 {
|
||||
format!("{secs}s ago")
|
||||
} else if secs < 3600 {
|
||||
format!("{}m ago", secs / 60)
|
||||
} else {
|
||||
let hours = secs / 3600;
|
||||
let mins = (secs % 3600) / 60;
|
||||
if mins == 0 {
|
||||
format!("{hours}h ago")
|
||||
} else {
|
||||
format!("{hours}h {mins}m ago")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
fn make_test_snapshot(task_id: &str) -> TaskSnapshot {
|
||||
TaskSnapshot {
|
||||
task_id: task_id.to_string(),
|
||||
tool_call_id: format!("tc-{}", task_id),
|
||||
command: "echo hello".to_string(),
|
||||
cwd: "/tmp".to_string(),
|
||||
start_time: Utc::now(),
|
||||
end_time: None,
|
||||
output: String::new(),
|
||||
output_file: PathBuf::from(format!("/tmp/{}.log", task_id)),
|
||||
truncated: false,
|
||||
exit_code: None,
|
||||
signal: None,
|
||||
completed: false,
|
||||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_and_get() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
let snapshot = make_test_snapshot("test-1");
|
||||
|
||||
registry.register(snapshot).await.unwrap();
|
||||
|
||||
let got = registry.get("test-1").await.unwrap();
|
||||
assert_eq!(got.task_id, "test-1");
|
||||
assert_eq!(got.command, "echo hello");
|
||||
assert!(!got.completed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_not_found() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
assert!(registry.get("nonexistent").await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_update_output() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
registry
|
||||
.register(make_test_snapshot("test-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
registry
|
||||
.update_output("test-1", "hello world".to_string(), false)
|
||||
.await;
|
||||
|
||||
let got = registry.get("test-1").await.unwrap();
|
||||
assert_eq!(got.output, "hello world");
|
||||
assert!(!got.truncated);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mark_completed() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
registry
|
||||
.register(make_test_snapshot("test-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
registry.mark_completed("test-1", Some(0), None).await;
|
||||
|
||||
let got = registry.get("test-1").await.unwrap();
|
||||
assert!(got.completed);
|
||||
assert_eq!(got.exit_code, Some(0));
|
||||
assert!(got.end_time.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wait_for_completion_already_done() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
registry
|
||||
.register(make_test_snapshot("test-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
registry.mark_completed("test-1", Some(0), None).await;
|
||||
|
||||
// Should return immediately since already completed
|
||||
let got = registry
|
||||
.wait_for_completion("test-1", Some(Duration::from_millis(100)))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(got.completed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wait_for_completion_with_timeout() {
|
||||
let registry = Arc::new(BackgroundTaskRegistry::new());
|
||||
registry
|
||||
.register(make_test_snapshot("test-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Start waiting with short timeout
|
||||
let got = registry
|
||||
.wait_for_completion("test-1", Some(Duration::from_millis(50)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should return with incomplete status (timed out)
|
||||
assert!(!got.completed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wait_for_completion_notified() {
|
||||
let registry = Arc::new(BackgroundTaskRegistry::new());
|
||||
registry
|
||||
.register(make_test_snapshot("test-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let registry_clone = registry.clone();
|
||||
|
||||
// Spawn task to complete after short delay
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
registry_clone
|
||||
.mark_completed("test-1", Some(42), None)
|
||||
.await;
|
||||
});
|
||||
|
||||
// Wait for completion
|
||||
let got = registry
|
||||
.wait_for_completion("test-1", Some(Duration::from_secs(5)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(got.completed);
|
||||
assert_eq!(got.exit_code, Some(42));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_max_tasks_limit() {
|
||||
let registry = BackgroundTaskRegistry::with_max_tasks(2);
|
||||
|
||||
// Register up to limit
|
||||
registry
|
||||
.register(make_test_snapshot("task-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.register(make_test_snapshot("task-2"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Third should fail
|
||||
let result = registry.register(make_test_snapshot("task-3")).await;
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("Maximum background tasks"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_max_tasks_cleanup_completed() {
|
||||
let registry = BackgroundTaskRegistry::with_max_tasks(2);
|
||||
|
||||
registry
|
||||
.register(make_test_snapshot("task-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.register(make_test_snapshot("task-2"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Mark first as completed
|
||||
registry.mark_completed("task-1", Some(0), None).await;
|
||||
|
||||
// Now third should succeed (completed task cleaned up)
|
||||
registry
|
||||
.register(make_test_snapshot("task-3"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify task-1 was cleaned up
|
||||
assert!(registry.get("task-1").await.is_none());
|
||||
assert!(registry.get("task-3").await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_tasks() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
registry
|
||||
.register(make_test_snapshot("task-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.register(make_test_snapshot("task-2"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let tasks = registry.list().await;
|
||||
assert_eq!(tasks.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_active_count() {
|
||||
let registry = BackgroundTaskRegistry::new();
|
||||
registry
|
||||
.register(make_test_snapshot("task-1"))
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.register(make_test_snapshot("task-2"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(registry.active_count().await, 2);
|
||||
|
||||
registry.mark_completed("task-1", Some(0), None).await;
|
||||
assert_eq!(registry.active_count().await, 1);
|
||||
}
|
||||
|
||||
// ── Manifest tests ──
|
||||
|
||||
fn make_manifest_entry(task_id: &str, secs_ago: u64) -> BackgroundTaskManifestEntry {
|
||||
BackgroundTaskManifestEntry {
|
||||
task_id: task_id.to_string(),
|
||||
command: format!("rsync -aP src:{task_id} /data/"),
|
||||
display_command: None,
|
||||
output_file: PathBuf::from(format!("/tmp/sessions/tasks/{task_id}.log")),
|
||||
start_time: std::time::SystemTime::now() - Duration::from_secs(secs_ago),
|
||||
cwd: "/home/user".to_string(),
|
||||
kind: kigi_tools::computer::types::TaskKind::Bash,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_roundtrip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let entries = vec![
|
||||
make_manifest_entry("task-a", 3600),
|
||||
make_manifest_entry("task-b", 120),
|
||||
];
|
||||
persist_manifest(dir.path(), entries);
|
||||
|
||||
let loaded = load_and_clear_manifest(dir.path());
|
||||
assert_eq!(loaded.len(), 2);
|
||||
assert_eq!(loaded[0].task_id, "task-a");
|
||||
assert_eq!(loaded[1].task_id, "task-b");
|
||||
assert_eq!(loaded[0].command, "rsync -aP src:task-a /data/");
|
||||
|
||||
// File is deleted after load
|
||||
let again = load_and_clear_manifest(dir.path());
|
||||
assert!(again.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_empty_entries_no_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
persist_manifest(dir.path(), Vec::new());
|
||||
assert!(!dir.path().join(MANIFEST_FILENAME).exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_missing_file_returns_empty() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let loaded = load_and_clear_manifest(dir.path());
|
||||
assert!(loaded.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_malformed_json_returns_empty_and_deletes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join(MANIFEST_FILENAME);
|
||||
std::fs::write(&path, b"not valid json {{{").unwrap();
|
||||
|
||||
let loaded = load_and_clear_manifest(dir.path());
|
||||
assert!(loaded.is_empty());
|
||||
assert!(!path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_partial_json_missing_fields_returns_empty() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join(MANIFEST_FILENAME);
|
||||
// Valid JSON array but missing required fields
|
||||
std::fs::write(&path, br#"[{"task_id": "x"}]"#).unwrap();
|
||||
|
||||
let loaded = load_and_clear_manifest(dir.path());
|
||||
assert!(loaded.is_empty());
|
||||
assert!(!path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_reminder_single_task() {
|
||||
let entries = vec![make_manifest_entry("bg-1", 7200)];
|
||||
let reminder = format_resumed_tasks_reminder(&entries);
|
||||
assert!(reminder.contains("This session was resumed"));
|
||||
assert!(reminder.contains("bg-1"));
|
||||
assert!(reminder.contains("2h ago"));
|
||||
assert!(reminder.contains("rsync -aP src:bg-1 /data/"));
|
||||
assert!(reminder.contains("Output log:"));
|
||||
assert!(reminder.contains("Check whether each is still running"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_reminder_prefers_display_command() {
|
||||
let mut entry = make_manifest_entry("bg-1", 60);
|
||||
entry.display_command = Some("rsync /data".to_string());
|
||||
let reminder = format_resumed_tasks_reminder(&[entry]);
|
||||
assert!(reminder.contains("rsync /data"));
|
||||
assert!(!reminder.contains("rsync -aP"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_reminder_labels_monitor_tasks() {
|
||||
let mut entry = make_manifest_entry("mon-1", 300);
|
||||
entry.kind = kigi_tools::computer::types::TaskKind::Monitor;
|
||||
let reminder = format_resumed_tasks_reminder(&[entry]);
|
||||
assert!(reminder.contains("[monitor]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_reminder_no_label_for_bash_tasks() {
|
||||
let entry = make_manifest_entry("bg-1", 300);
|
||||
let reminder = format_resumed_tasks_reminder(&[entry]);
|
||||
assert!(!reminder.contains("[monitor]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_roundtrip_preserves_kind() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut entry = make_manifest_entry("mon-1", 60);
|
||||
entry.kind = kigi_tools::computer::types::TaskKind::Monitor;
|
||||
persist_manifest(dir.path(), vec![entry]);
|
||||
|
||||
let loaded = load_and_clear_manifest(dir.path());
|
||||
assert_eq!(loaded.len(), 1);
|
||||
assert_eq!(
|
||||
loaded[0].kind,
|
||||
kigi_tools::computer::types::TaskKind::Monitor
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_duration_seconds() {
|
||||
let now = std::time::SystemTime::now();
|
||||
assert_eq!(
|
||||
format_duration_ago(now, now - Duration::from_secs(30)),
|
||||
"30s ago"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_duration_minutes() {
|
||||
let now = std::time::SystemTime::now();
|
||||
assert_eq!(
|
||||
format_duration_ago(now, now - Duration::from_secs(300)),
|
||||
"5m ago"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_duration_hours_and_minutes() {
|
||||
let now = std::time::SystemTime::now();
|
||||
assert_eq!(
|
||||
format_duration_ago(now, now - Duration::from_secs(5400)),
|
||||
"1h 30m ago"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_duration_exact_hours() {
|
||||
let now = std::time::SystemTime::now();
|
||||
assert_eq!(
|
||||
format_duration_ago(now, now - Duration::from_secs(7200)),
|
||||
"2h ago"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_duration_future_start_returns_zero() {
|
||||
let now = std::time::SystemTime::now();
|
||||
assert_eq!(
|
||||
format_duration_ago(now, now + Duration::from_secs(100)),
|
||||
"0s ago"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
// todo: add support for signal handling
|
||||
use std::process::Stdio;
|
||||
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::process::Command;
|
||||
use tokio::time;
|
||||
|
||||
use crate::terminal::runner::{
|
||||
AsyncTerminalRunner, TerminalError, TerminalRunRequest, TerminalRunResult,
|
||||
};
|
||||
|
||||
pub struct LocalTerminalRunner;
|
||||
|
||||
async fn read_stream(mut stream: impl AsyncReadExt + Unpin) -> Vec<u8> {
|
||||
let mut buffer = Vec::new();
|
||||
let _ = stream.read_to_end(&mut buffer).await;
|
||||
buffer
|
||||
}
|
||||
|
||||
/// Truncate buffer to keep only the last `limit` bytes (drops oldest bytes).
|
||||
/// Returns true if truncation occurred.
|
||||
///
|
||||
/// This function ensures we don't split UTF-8 characters when truncating
|
||||
/// by using char_indices to find a valid character boundary.
|
||||
fn truncate_buffer(buf: &mut Vec<u8>, limit: usize) -> bool {
|
||||
if buf.len() > limit {
|
||||
// Convert to string to work with character boundaries
|
||||
let s = String::from_utf8_lossy(buf);
|
||||
let excess = buf.len().saturating_sub(limit);
|
||||
|
||||
// Find the first char boundary at or after `excess` bytes
|
||||
let start_idx = s
|
||||
.char_indices()
|
||||
.find(|(i, _)| *i >= excess)
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(s.len());
|
||||
|
||||
// Slice from that boundary and update buffer
|
||||
*buf = s[start_idx..].as_bytes().to_vec();
|
||||
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncTerminalRunner for LocalTerminalRunner {
|
||||
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, TerminalError> {
|
||||
// Build and spawn the command via the platform shell.
|
||||
#[cfg(unix)]
|
||||
let mut cmd = {
|
||||
let mut c = Command::new(crate::terminal::default_shell_path());
|
||||
c.arg("-lc").arg(&request.command);
|
||||
c
|
||||
};
|
||||
#[cfg(not(unix))]
|
||||
let mut cmd = {
|
||||
let inv = kigi_config::shell::shell_command_argv(&request.command);
|
||||
let mut c = Command::new(inv.program);
|
||||
c.args(&inv.args).envs(inv.env);
|
||||
c
|
||||
};
|
||||
cmd.current_dir(&request.cwd)
|
||||
.envs(&request.env)
|
||||
.envs(crate::terminal::pager_env())
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
|
||||
// Detach from the controlling terminal so child processes
|
||||
// (e.g. GPG pinentry) cannot open /dev/tty and corrupt the TUI.
|
||||
kigi_tools::util::detach_command(&mut cmd);
|
||||
|
||||
let mut child = cmd
|
||||
.spawn()
|
||||
.map_err(|e| TerminalError::Other(format!("Failed to start shell: {e}")))?;
|
||||
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| TerminalError::Other("Failed to capture stdout".into()))?;
|
||||
let stderr = child
|
||||
.stderr
|
||||
.take()
|
||||
.ok_or_else(|| TerminalError::Other("Failed to capture stderr".into()))?;
|
||||
|
||||
let stdout_task = tokio::spawn(read_stream(stdout));
|
||||
let stderr_task = tokio::spawn(read_stream(stderr));
|
||||
|
||||
let mut timed_out = false;
|
||||
|
||||
let wait_result = time::timeout(request.timeout, child.wait()).await;
|
||||
let exit_status = match wait_result {
|
||||
Ok(status_res) => status_res
|
||||
.map_err(|e| TerminalError::Other(format!("Failed to wait for process: {e}")))?,
|
||||
Err(_) => {
|
||||
timed_out = true;
|
||||
if let Err(e) = child.start_kill() {
|
||||
tracing::warn!("Failed to kill timed-out process: {e}");
|
||||
}
|
||||
child.wait().await.map_err(|e| {
|
||||
TerminalError::Other(format!("Failed to wait for killed process: {e}"))
|
||||
})?
|
||||
}
|
||||
};
|
||||
|
||||
let stdout_result = stdout_task.await.unwrap_or_else(|_| Vec::new());
|
||||
let stderr_result = stderr_task.await.unwrap_or_else(|_| Vec::new());
|
||||
|
||||
// Combine stdout and stderr, then truncate if needed
|
||||
let mut combined = stdout_result;
|
||||
combined.extend(stderr_result);
|
||||
let truncated = truncate_buffer(&mut combined, request.output_byte_limit);
|
||||
|
||||
let combined_output = String::from_utf8_lossy(&combined).into_owned();
|
||||
|
||||
Ok(TerminalRunResult {
|
||||
combined_output,
|
||||
exit_code: exit_status.code(),
|
||||
truncated,
|
||||
signal: None,
|
||||
timed_out,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::terminal::DEFAULT_OUTPUT_BYTE_LIMIT;
|
||||
use crate::terminal::runner::TerminalRunRequest;
|
||||
use kigi_paths::AbsPathBuf;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn make_request(command: &str) -> TerminalRunRequest {
|
||||
TerminalRunRequest {
|
||||
tool_call_id: agent_client_protocol::ToolCallId::new("test"),
|
||||
command: command.to_string(),
|
||||
cwd: AbsPathBuf::new(std::env::current_dir().unwrap()).unwrap(),
|
||||
env: HashMap::new(),
|
||||
timeout: std::time::Duration::from_secs(10),
|
||||
output_byte_limit: DEFAULT_OUTPUT_BYTE_LIMIT,
|
||||
stream: false,
|
||||
output_file: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify that `detach_from_tty` prevents child processes from opening
|
||||
/// `/dev/tty`. After setsid(), the child has no controlling terminal.
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn test_child_cannot_open_dev_tty() {
|
||||
// Skip in CI / environments without a controlling terminal.
|
||||
if std::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.open("/dev/tty")
|
||||
.is_err()
|
||||
{
|
||||
eprintln!("skipping: no controlling terminal");
|
||||
return;
|
||||
}
|
||||
|
||||
let result = LocalTerminalRunner
|
||||
.run(make_request(
|
||||
"(exec 3>/dev/tty && echo ATTACHED || echo DETACHED) 2>/dev/null",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.combined_output.trim(),
|
||||
"DETACHED",
|
||||
"child process should not be able to open /dev/tty after detach_from_tty()"
|
||||
);
|
||||
}
|
||||
|
||||
/// Basic regression: commands still produce output and exit normally.
|
||||
#[tokio::test]
|
||||
async fn test_basic_command_output() {
|
||||
let result = LocalTerminalRunner
|
||||
.run(make_request("echo hello"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.combined_output.trim(), "hello");
|
||||
assert_eq!(result.exit_code, Some(0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
pub mod runner;
|
||||
pub use runner::{AsyncTerminalRunner, TerminalError, TerminalRunRequest, TerminalRunResult};
|
||||
|
||||
mod background_task;
|
||||
pub use background_task::{
|
||||
BackgroundTaskManifestEntry, BackgroundTaskRegistry, TaskId, TaskSnapshot,
|
||||
format_resumed_tasks_reminder, get_task_output_path, load_and_clear_manifest, persist_manifest,
|
||||
};
|
||||
|
||||
mod local_terminal;
|
||||
pub use local_terminal::LocalTerminalRunner;
|
||||
|
||||
mod acp_terminal;
|
||||
pub use acp_terminal::AcpTerminalRunner;
|
||||
|
||||
pub mod adapter;
|
||||
pub use adapter::AcpTerminalAdapter;
|
||||
|
||||
pub mod pty_session;
|
||||
|
||||
pub const DEFAULT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
||||
pub const DEFAULT_OUTPUT_BYTE_LIMIT: usize = 30_000; // 30k characters
|
||||
|
||||
/// Resolved absolute path to bash. On Unix uses the `kigi_config` shell
|
||||
/// resolution cascade (`$KIGI_SHELL` > `$SHELL` > `which` > common dirs >
|
||||
/// `/bin/bash`) and is cached process-wide. On non-Unix returns `"/bin/bash"`
|
||||
/// — but every caller in this crate is gated behind `#[cfg(unix)]`, so the
|
||||
/// non-Unix value should not be observed in practice.
|
||||
pub fn default_shell_path() -> &'static str {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
kigi_config::shell::unix_shell_path(kigi_config::shell::UnixShellKind::Bash)
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
"/bin/bash"
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, serde::Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TerminalStatus {
|
||||
Connecting,
|
||||
Connected,
|
||||
Exited,
|
||||
Error,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TerminalInfo {
|
||||
pub terminal_id: String,
|
||||
pub status: TerminalStatus,
|
||||
pub interactive: bool,
|
||||
pub name: Option<String>,
|
||||
pub exit_code: Option<i32>,
|
||||
pub cwd: Option<String>,
|
||||
pub output_offset: u64,
|
||||
pub created_at: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum TerminalExtError {
|
||||
#[error("terminal '{terminal_id}' not found")]
|
||||
NotFound { terminal_id: String },
|
||||
#[error("terminal '{terminal_id}' is not an interactive PTY")]
|
||||
NotInteractive { terminal_id: String },
|
||||
#[error("terminal '{terminal_id}' exited")]
|
||||
Exited { terminal_id: String },
|
||||
#[error("terminal '{terminal_id}' input channel closed")]
|
||||
InputClosed { terminal_id: String },
|
||||
#[error("{0}")]
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
impl TerminalExtError {
|
||||
pub fn code(&self) -> &'static str {
|
||||
match self {
|
||||
Self::NotFound { .. } => "TERMINAL_NOT_FOUND",
|
||||
Self::NotInteractive { .. } => "TERMINAL_NOT_INTERACTIVE",
|
||||
Self::Exited { .. } => "TERMINAL_EXITED",
|
||||
Self::InputClosed { .. } => "TERMINAL_INPUT_CLOSED",
|
||||
Self::Internal(_) => "TERMINAL_INTERNAL_ERROR",
|
||||
}
|
||||
}
|
||||
|
||||
fn terminal_id(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::NotFound { terminal_id }
|
||||
| Self::NotInteractive { terminal_id }
|
||||
| Self::Exited { terminal_id }
|
||||
| Self::InputClosed { terminal_id } => Some(terminal_id),
|
||||
Self::Internal(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: serde::Serialize> From<TerminalExtError> for crate::session::result::ExtMethodResult<T> {
|
||||
fn from(err: TerminalExtError) -> Self {
|
||||
let data = err
|
||||
.terminal_id()
|
||||
.map(|id| serde_json::json!({ "terminalId": id }));
|
||||
Self {
|
||||
result: None,
|
||||
error: serde_json::to_value(crate::session::result::ExtMethodError {
|
||||
code: err.code().to_string(),
|
||||
message: err.to_string(),
|
||||
data,
|
||||
})
|
||||
.ok(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_terminals() -> Vec<TerminalInfo> {
|
||||
let mut terminals = pty_session::list_ptys().await;
|
||||
let mut piped = streaming_local_terminal::list_piped_terminals().await;
|
||||
terminals.append(&mut piped);
|
||||
terminals
|
||||
}
|
||||
|
||||
/// Returns environment variables that prevent CLI tools from launching any blocking/waiting programs.
|
||||
/// Delegates to the canonical implementation in `kigi-tools`.
|
||||
pub use kigi_tools::util::pager_env;
|
||||
|
||||
/// Returns environment variables that encourage CLI tools to emit colored output
|
||||
/// and show progress bars/spinners even when running through pipes (non-TTY).
|
||||
pub fn color_env() -> std::collections::HashMap<String, String> {
|
||||
std::collections::HashMap::from([
|
||||
("TERM".to_string(), "xterm-256color".to_string()),
|
||||
("COLORTERM".to_string(), "truecolor".to_string()),
|
||||
("FORCE_COLOR".to_string(), "1".to_string()),
|
||||
("CLICOLOR_FORCE".to_string(), "1".to_string()),
|
||||
("CLICOLOR".to_string(), "1".to_string()),
|
||||
// Cargo: always show progress bar
|
||||
("CARGO_TERM_PROGRESS_WHEN".to_string(), "always".to_string()),
|
||||
("CARGO_TERM_PROGRESS_WIDTH".to_string(), "80".to_string()),
|
||||
// CI mode - many tools show progress in CI
|
||||
("CI".to_string(), "true".to_string()),
|
||||
// npm/yarn progress
|
||||
("NPM_CONFIG_PROGRESS".to_string(), "true".to_string()),
|
||||
// pip progress
|
||||
("PIP_PROGRESS_BAR".to_string(), "on".to_string()),
|
||||
// gradle
|
||||
(
|
||||
"GRADLE_OPTS".to_string(),
|
||||
"-Dorg.gradle.console=rich".to_string(),
|
||||
),
|
||||
// Maven
|
||||
("MAVEN_OPTS".to_string(), "-Dstyle.color=always".to_string()),
|
||||
])
|
||||
}
|
||||
|
||||
/// Returns environment variables that disable colors and ANSI escape codes in CLI
|
||||
/// tool output. Used when the client sets `x.ai/bashOutputNoColor: true` (e.g.
|
||||
/// kigi-tui which renders its own UI and doesn't need raw ANSI codes).
|
||||
///
|
||||
/// Follows the <https://no-color.org/> convention plus tool-specific overrides.
|
||||
pub fn no_color_env() -> std::collections::HashMap<String, String> {
|
||||
std::collections::HashMap::from([
|
||||
// https://no-color.org/ — respected by many CLI tools
|
||||
("NO_COLOR".to_string(), "1".to_string()),
|
||||
// Override TERM to dumb — prevents cursor movement, color codes
|
||||
("TERM".to_string(), "dumb".to_string()),
|
||||
// Disable forced color in tools that check these
|
||||
("FORCE_COLOR".to_string(), "0".to_string()),
|
||||
("CLICOLOR_FORCE".to_string(), "0".to_string()),
|
||||
("CLICOLOR".to_string(), "0".to_string()),
|
||||
// Cargo: disable color and progress bar
|
||||
("CARGO_TERM_COLOR".to_string(), "never".to_string()),
|
||||
// npm/yarn: disable color
|
||||
("NPM_CONFIG_COLOR".to_string(), "false".to_string()),
|
||||
// pip: disable color and progress
|
||||
("PIP_NO_COLOR".to_string(), "1".to_string()),
|
||||
("PIP_PROGRESS_BAR".to_string(), "off".to_string()),
|
||||
// gradle
|
||||
(
|
||||
"GRADLE_OPTS".to_string(),
|
||||
"-Dorg.gradle.console=plain".to_string(),
|
||||
),
|
||||
// Maven
|
||||
("MAVEN_OPTS".to_string(), "-Dstyle.color=never".to_string()),
|
||||
])
|
||||
}
|
||||
|
||||
mod streaming_local_terminal;
|
||||
pub use streaming_local_terminal::{
|
||||
ExitStatus, GatedNotifier, KillOutcome, OutputSnapshot, SessionNotificationSender,
|
||||
StreamingLocalTerminalRunner, background_terminal, create_terminal, find_terminal_session_id,
|
||||
get_terminal_output, kill_and_release_all_for_session, kill_terminal, release_terminal,
|
||||
wait_for_terminal_exit,
|
||||
};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Terminal runner that routes requests based on the `stream` flag:
|
||||
/// - `stream: true` → StreamingLocalTerminalRunner (updates, killable)
|
||||
/// - `stream: false` → LocalTerminalRunner (silent, fire-and-forget)
|
||||
pub struct TerminalRunner {
|
||||
notifier: Arc<dyn SessionNotificationSender>,
|
||||
session_id: agent_client_protocol::SessionId,
|
||||
}
|
||||
|
||||
impl TerminalRunner {
|
||||
pub fn new(
|
||||
notifier: Arc<dyn SessionNotificationSender>,
|
||||
session_id: agent_client_protocol::SessionId,
|
||||
) -> Self {
|
||||
Self {
|
||||
notifier,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncTerminalRunner for TerminalRunner {
|
||||
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, TerminalError> {
|
||||
if request.stream {
|
||||
StreamingLocalTerminalRunner {
|
||||
notifier: self.notifier.clone(),
|
||||
session_id: self.session_id.clone(),
|
||||
}
|
||||
.run(request)
|
||||
.await
|
||||
} else {
|
||||
LocalTerminalRunner.run(request).await
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,754 @@
|
||||
//! Agent-scoped interactive PTY manager. PTYs are keyed by `terminalId`,
|
||||
//! outlive sessions, and multiplex I/O over the existing ACP WebSocket.
|
||||
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::io::{Read, Write};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
|
||||
use crate::extensions::routing::{TargetClientId, send_routed_notification};
|
||||
use crate::terminal::{TerminalExtError, TerminalInfo, TerminalStatus};
|
||||
|
||||
const NOTIFICATION_METHOD: &str = "x.ai/terminal/pty/notification";
|
||||
const OUTPUT_RING_BUFFER_SIZE: usize = 256 * 1024;
|
||||
const OUTPUT_BATCH_INTERVAL_MS: u64 = 16;
|
||||
const BUSY_POLL_INTERVAL_MS: u64 = 500;
|
||||
const INPUT_CHANNEL_CAPACITY: usize = 256;
|
||||
|
||||
pub struct PtySession {
|
||||
master: Box<dyn MasterPty + Send>,
|
||||
input_tx: mpsc::Sender<Vec<u8>>,
|
||||
output_offset: u64,
|
||||
output_ring: VecDeque<u8>,
|
||||
child: Box<dyn portable_pty::Child + Send + Sync>,
|
||||
cwd: Option<String>,
|
||||
name: Option<String>,
|
||||
created_at: u64,
|
||||
rows: u16,
|
||||
cols: u16,
|
||||
target_client_id: TargetClientId,
|
||||
busy: bool,
|
||||
gateway: GatewaySender,
|
||||
}
|
||||
|
||||
type PtyMap = HashMap<String, Arc<Mutex<PtySession>>>;
|
||||
|
||||
static PTY_REGISTRY: LazyLock<Mutex<PtyMap>> = LazyLock::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
pub async fn get_pty(pty_id: &str) -> Option<Arc<Mutex<PtySession>>> {
|
||||
PTY_REGISTRY.lock().await.get(pty_id).cloned()
|
||||
}
|
||||
|
||||
pub async fn require_pty(terminal_id: &str) -> Result<Arc<Mutex<PtySession>>, TerminalExtError> {
|
||||
get_pty(terminal_id)
|
||||
.await
|
||||
.ok_or_else(|| TerminalExtError::NotInteractive {
|
||||
terminal_id: terminal_id.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn create_pty(
|
||||
shell: Option<&str>,
|
||||
cwd: Option<&str>,
|
||||
env: HashMap<String, String>,
|
||||
rows: u16,
|
||||
cols: u16,
|
||||
name: Option<&str>,
|
||||
gateway: GatewaySender,
|
||||
target_client_id: TargetClientId,
|
||||
) -> Result<String, TerminalExtError> {
|
||||
let pty_id = uuid::Uuid::now_v7().to_string();
|
||||
|
||||
let pty_system = native_pty_system();
|
||||
let size = PtySize {
|
||||
rows,
|
||||
cols,
|
||||
pixel_width: 0,
|
||||
pixel_height: 0,
|
||||
};
|
||||
|
||||
let pair = pty_system
|
||||
.openpty(size)
|
||||
.map_err(|e| TerminalExtError::Internal(format!("failed to open pty: {e}")))?;
|
||||
|
||||
let (shell_path, shell_args) = resolve_pty_shell(shell);
|
||||
|
||||
let mut cmd = CommandBuilder::new(&shell_path);
|
||||
for arg in &shell_args {
|
||||
cmd.arg(arg);
|
||||
}
|
||||
|
||||
if let Some(dir) = cwd {
|
||||
cmd.cwd(dir);
|
||||
} else if let Ok(dir) = std::env::current_dir() {
|
||||
cmd.cwd(dir);
|
||||
}
|
||||
|
||||
for (k, v) in &env {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
cmd.env("TERM", "xterm-256color");
|
||||
cmd.env("COLORTERM", "truecolor");
|
||||
cmd.env("LANG", "en_US.UTF-8");
|
||||
cmd.env("LC_ALL", "en_US.UTF-8");
|
||||
|
||||
let child = pair
|
||||
.slave
|
||||
.spawn_command(cmd)
|
||||
.map_err(|e| TerminalExtError::Internal(format!("failed to spawn shell: {e}")))?;
|
||||
|
||||
let reader = pair
|
||||
.master
|
||||
.try_clone_reader()
|
||||
.map_err(|e| TerminalExtError::Internal(format!("failed to clone pty reader: {e}")))?;
|
||||
let writer = pair
|
||||
.master
|
||||
.take_writer()
|
||||
.map_err(|e| TerminalExtError::Internal(format!("failed to take pty writer: {e}")))?;
|
||||
|
||||
let (input_tx, input_rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY);
|
||||
spawn_pty_input_loop(writer, input_rx);
|
||||
|
||||
let created_at = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
|
||||
let resolved_cwd = cwd.map(|s| s.to_string()).or_else(|| {
|
||||
std::env::current_dir()
|
||||
.ok()
|
||||
.map(|p| p.to_string_lossy().to_string())
|
||||
});
|
||||
|
||||
let resolved_name = name.map(|s| s.to_string()).or_else(|| {
|
||||
// Default to cwd basename, fall back to shell basename.
|
||||
resolved_cwd
|
||||
.as_deref()
|
||||
.and_then(|p| {
|
||||
std::path::Path::new(p)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
})
|
||||
.or_else(|| {
|
||||
std::path::Path::new(&shell_path)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
})
|
||||
});
|
||||
|
||||
let session = PtySession {
|
||||
master: pair.master,
|
||||
input_tx,
|
||||
output_offset: 0,
|
||||
output_ring: VecDeque::with_capacity(OUTPUT_RING_BUFFER_SIZE),
|
||||
child,
|
||||
cwd: resolved_cwd,
|
||||
name: resolved_name,
|
||||
created_at,
|
||||
rows,
|
||||
cols,
|
||||
target_client_id,
|
||||
busy: false,
|
||||
gateway: gateway.clone(),
|
||||
};
|
||||
|
||||
let entry = Arc::new(Mutex::new(session));
|
||||
PTY_REGISTRY
|
||||
.lock()
|
||||
.await
|
||||
.insert(pty_id.clone(), entry.clone());
|
||||
|
||||
let pty_id_clone = pty_id.clone();
|
||||
tokio::task::spawn_local(run_pty_output_loop(reader, entry, pty_id_clone, gateway));
|
||||
|
||||
Ok(pty_id)
|
||||
}
|
||||
|
||||
fn spawn_pty_input_loop(mut writer: Box<dyn Write + Send>, mut input_rx: mpsc::Receiver<Vec<u8>>) {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
while let Some(mut chunk) = input_rx.blocking_recv() {
|
||||
while let Ok(more) = input_rx.try_recv() {
|
||||
chunk.extend_from_slice(&more);
|
||||
}
|
||||
if writer.write_all(&chunk).is_err() {
|
||||
break;
|
||||
}
|
||||
if writer.flush().is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Reads PTY output, batches on a 16ms tick, and sends notifications.
|
||||
/// Also samples the foreground process group on a slower tick and pushes
|
||||
/// `process_started`/`process_ended` on idle↔busy transitions — time-based
|
||||
/// rather than output-driven because a busy process can be silent.
|
||||
async fn run_pty_output_loop(
|
||||
reader: Box<dyn Read + Send>,
|
||||
pty: Arc<Mutex<PtySession>>,
|
||||
pty_id: String,
|
||||
gateway: GatewaySender,
|
||||
) {
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::{Duration, interval};
|
||||
|
||||
let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(64);
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut reader = reader;
|
||||
let mut buf = [0u8; 4096];
|
||||
loop {
|
||||
match reader.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
if data_tx.blocking_send(buf[..n].to_vec()).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut pending = Vec::new();
|
||||
let mut tick = interval(Duration::from_millis(OUTPUT_BATCH_INTERVAL_MS));
|
||||
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
let mut busy_tick = interval(Duration::from_millis(BUSY_POLL_INTERVAL_MS));
|
||||
busy_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
chunk = data_rx.recv() => {
|
||||
match chunk {
|
||||
Some(data) => {
|
||||
{
|
||||
let mut session = pty.lock().await;
|
||||
session.output_ring.extend(data.iter().copied());
|
||||
if session.output_ring.len() > OUTPUT_RING_BUFFER_SIZE {
|
||||
let excess = session.output_ring.len() - OUTPUT_RING_BUFFER_SIZE;
|
||||
session.output_ring.drain(..excess);
|
||||
}
|
||||
session.output_offset += data.len() as u64;
|
||||
}
|
||||
pending.extend_from_slice(&data);
|
||||
}
|
||||
None => {
|
||||
if !pending.is_empty() {
|
||||
flush_output(&pty, &pty_id, &mut pending, &gateway).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ = tick.tick() => {
|
||||
if !pending.is_empty() {
|
||||
flush_output(&pty, &pty_id, &mut pending, &gateway).await;
|
||||
}
|
||||
}
|
||||
|
||||
_ = busy_tick.tick() => {
|
||||
sample_busy_transition(&pty, &pty_id).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Child exited
|
||||
let (exit_code, signal, target_client_id, was_busy) = tokio::task::spawn_blocking({
|
||||
let pty = pty.clone();
|
||||
move || {
|
||||
let mut session = pty.blocking_lock();
|
||||
let target_client_id = session.target_client_id.clone();
|
||||
let was_busy = session.busy;
|
||||
match session.child.wait() {
|
||||
Ok(es) => (
|
||||
Some(es.exit_code() as i32),
|
||||
None::<String>,
|
||||
target_client_id,
|
||||
was_busy,
|
||||
),
|
||||
Err(_) => (None, None, target_client_id, was_busy),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
if was_busy {
|
||||
send_busy_notification(&pty_id, false, &target_client_id, &gateway);
|
||||
}
|
||||
|
||||
send_routed_notification(
|
||||
&gateway,
|
||||
NOTIFICATION_METHOD,
|
||||
serde_json::json!({
|
||||
"terminalId": pty_id,
|
||||
"type": "exit",
|
||||
"exitCode": exit_code,
|
||||
"signal": signal,
|
||||
}),
|
||||
&target_client_id,
|
||||
);
|
||||
}
|
||||
|
||||
/// Whether the PTY's controlling terminal has a foreground process group
|
||||
/// distinct from the shell itself — i.e. a command is actively running
|
||||
/// rather than the shell sitting idle at its prompt.
|
||||
///
|
||||
/// `process_group_leader()` issues `tcgetpgrp` on the master fd; an idle
|
||||
/// shell is its own foreground process group, so it matches the shell
|
||||
/// child's pid. When a command runs in the foreground the kernel reports
|
||||
/// that command's process group instead. Returns false when the value is
|
||||
/// unavailable (the shell exited or runs without job control).
|
||||
///
|
||||
/// Limitation: a shell that `exec`s a program in place keeps the same pid and
|
||||
/// pgid, so `tcgetpgrp` still matches the recorded child pid and the program
|
||||
/// reads as idle. Telling that apart from a real idle prompt needs per-OS
|
||||
/// process inspection, so a command launched the usual way (fork then exec) is
|
||||
/// detected while an `exec`-replaced shell is not.
|
||||
#[cfg(unix)]
|
||||
fn session_has_foreground_process(session: &PtySession) -> bool {
|
||||
let Some(foreground_pgid) = session.master.process_group_leader() else {
|
||||
return false;
|
||||
};
|
||||
match session.child.process_id() {
|
||||
Some(shell_pid) => i64::from(foreground_pgid) != i64::from(shell_pid),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// `tcgetpgrp` has no ConPTY equivalent (`process_group_leader` is
|
||||
/// unix-only in portable-pty), so non-unix PTYs never report a foreground
|
||||
/// process and clients close terminals without confirmation.
|
||||
#[cfg(not(unix))]
|
||||
fn session_has_foreground_process(_session: &PtySession) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Emits `process_started` / `process_ended` only on idle↔busy transitions so
|
||||
/// a steady state never repeats notifications.
|
||||
async fn sample_busy_transition(pty: &Arc<Mutex<PtySession>>, pty_id: &str) {
|
||||
let mut session = pty.lock().await;
|
||||
let now_busy = session_has_foreground_process(&session);
|
||||
if now_busy == session.busy {
|
||||
return;
|
||||
}
|
||||
session.busy = now_busy;
|
||||
send_busy_notification(
|
||||
pty_id,
|
||||
now_busy,
|
||||
&session.target_client_id,
|
||||
&session.gateway,
|
||||
);
|
||||
}
|
||||
|
||||
fn send_busy_notification(
|
||||
pty_id: &str,
|
||||
busy: bool,
|
||||
target_client_id: &TargetClientId,
|
||||
gateway: &GatewaySender,
|
||||
) {
|
||||
send_routed_notification(
|
||||
gateway,
|
||||
NOTIFICATION_METHOD,
|
||||
serde_json::json!({
|
||||
"terminalId": pty_id,
|
||||
"type": if busy { "process_started" } else { "process_ended" },
|
||||
}),
|
||||
target_client_id,
|
||||
);
|
||||
}
|
||||
|
||||
async fn flush_output(
|
||||
pty: &Arc<Mutex<PtySession>>,
|
||||
pty_id: &str,
|
||||
pending: &mut Vec<u8>,
|
||||
gateway: &GatewaySender,
|
||||
) {
|
||||
use base64::Engine as _;
|
||||
|
||||
let (output_offset, target_client_id) = {
|
||||
let session = pty.lock().await;
|
||||
(session.output_offset, session.target_client_id.clone())
|
||||
};
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&*pending);
|
||||
pending.clear();
|
||||
|
||||
send_routed_notification(
|
||||
gateway,
|
||||
NOTIFICATION_METHOD,
|
||||
serde_json::json!({
|
||||
"terminalId": pty_id,
|
||||
"type": "output",
|
||||
"data": b64,
|
||||
"outputOffset": output_offset,
|
||||
}),
|
||||
&target_client_id,
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn write_pty_input(pty_id: &str, data: &[u8]) -> Result<(), TerminalExtError> {
|
||||
let entry = require_pty(pty_id).await?;
|
||||
let input_tx = { entry.lock().await.input_tx.clone() };
|
||||
input_tx
|
||||
.send(data.to_vec())
|
||||
.await
|
||||
.map_err(|_| TerminalExtError::InputClosed {
|
||||
terminal_id: pty_id.into(),
|
||||
})?;
|
||||
sample_busy_transition(&entry, pty_id).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn resize_pty(pty_id: &str, rows: u16, cols: u16) -> Result<(), TerminalExtError> {
|
||||
let entry = require_pty(pty_id).await?;
|
||||
let mut session = entry.lock().await;
|
||||
session
|
||||
.master
|
||||
.resize(PtySize {
|
||||
rows,
|
||||
cols,
|
||||
pixel_width: 0,
|
||||
pixel_height: 0,
|
||||
})
|
||||
.map_err(|e| TerminalExtError::Internal(format!("failed to resize pty: {e}")))?;
|
||||
session.rows = rows;
|
||||
session.cols = cols;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn is_exited(pty_id: &str) -> bool {
|
||||
match get_pty(pty_id).await {
|
||||
Some(entry) => entry.lock().await.child.try_wait().ok().flatten().is_some(),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn close_pty(pty_id: &str) -> Result<(), String> {
|
||||
if let Some(entry) = PTY_REGISTRY.lock().await.remove(pty_id) {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut session = entry.blocking_lock();
|
||||
let _ = session.child.kill();
|
||||
let _ = session.child.wait();
|
||||
})
|
||||
.await
|
||||
.map_err(|e| format!("close task failed: {e}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Called on agent disconnect to clean up all PTYs.
|
||||
pub async fn close_all() {
|
||||
let entries: Vec<Arc<Mutex<PtySession>>> = {
|
||||
let mut reg = PTY_REGISTRY.lock().await;
|
||||
reg.drain().map(|(_, v)| v).collect()
|
||||
};
|
||||
for entry in entries {
|
||||
let _ = tokio::task::spawn_blocking(move || {
|
||||
let mut session = entry.blocking_lock();
|
||||
let _ = session.child.kill();
|
||||
let _ = session.child.wait();
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the shell binary and arguments for an interactive PTY session.
|
||||
///
|
||||
/// Priority: explicit `shell` param > `$SHELL` env > platform default.
|
||||
/// On Windows falls back to the `detect_windows_shell` cascade
|
||||
/// (pwsh > powershell.exe > Git Bash > cmd.exe, overridable via
|
||||
/// `KIGI_SHELL`) since `$SHELL` is absent.
|
||||
fn resolve_pty_shell(shell: Option<&str>) -> (String, Vec<String>) {
|
||||
if let Some(s) = shell {
|
||||
return (s.to_string(), vec![]);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let path = std::env::var("SHELL").unwrap_or_else(|_| "/bin/bash".to_string());
|
||||
(path, vec!["-l".to_string()])
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
use kigi_config::shell::{WindowsShell, detect_windows_shell};
|
||||
match detect_windows_shell() {
|
||||
WindowsShell::GitBash(path) => (path.clone(), vec!["-l".to_string()]),
|
||||
WindowsShell::Pwsh => ("pwsh".to_string(), vec!["-NoLogo".to_string()]),
|
||||
WindowsShell::PowerShell => ("powershell.exe".to_string(), vec!["-NoLogo".to_string()]),
|
||||
WindowsShell::Cmd => ("cmd.exe".to_string(), vec![]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_ptys() -> Vec<TerminalInfo> {
|
||||
let entries: Vec<(String, Arc<Mutex<PtySession>>)> = {
|
||||
let reg = PTY_REGISTRY.lock().await;
|
||||
reg.iter()
|
||||
.map(|(id, entry)| (id.clone(), entry.clone()))
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut result = Vec::with_capacity(entries.len());
|
||||
for (id, entry) in entries {
|
||||
let mut session = entry.lock().await;
|
||||
let (status, exit_code) = match session.child.try_wait() {
|
||||
Ok(Some(es)) => (TerminalStatus::Exited, Some(es.exit_code() as i32)),
|
||||
Ok(None) => (TerminalStatus::Connected, None),
|
||||
Err(_) => (TerminalStatus::Error, None),
|
||||
};
|
||||
result.push(TerminalInfo {
|
||||
terminal_id: id,
|
||||
status,
|
||||
interactive: true,
|
||||
name: session.name.clone(),
|
||||
exit_code,
|
||||
cwd: session.cwd.clone(),
|
||||
output_offset: session.output_offset,
|
||||
created_at: session.created_at,
|
||||
});
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PtyLoadResult {
|
||||
pub terminal_id: String,
|
||||
pub rows: u16,
|
||||
pub cols: u16,
|
||||
pub exited: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub exit_code: Option<i32>,
|
||||
}
|
||||
|
||||
/// Reconnect to a PTY. Replays the full ring buffer (with `isReplay: true`)
|
||||
/// so the client can reset its VTE emulator and feed all bytes from scratch.
|
||||
/// Exited PTYs are still loadable so the client can see final output.
|
||||
///
|
||||
/// Updates the stored `target_client_id` so that subsequent output
|
||||
/// notifications from the output loop are routed to the reconnecting client.
|
||||
pub async fn load(
|
||||
pty_id: &str,
|
||||
gateway: &GatewaySender,
|
||||
target_client_id: TargetClientId,
|
||||
) -> Result<PtyLoadResult, TerminalExtError> {
|
||||
let entry = require_pty(pty_id).await?;
|
||||
let (replay, output_offset, exit_info, rows, cols, busy) = {
|
||||
let mut session = entry.lock().await;
|
||||
session.target_client_id = target_client_id.clone();
|
||||
let exit_info = session
|
||||
.child
|
||||
.try_wait()
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|es| es.exit_code() as i32);
|
||||
let busy = session_has_foreground_process(&session);
|
||||
session.busy = busy;
|
||||
(
|
||||
session.output_ring.iter().copied().collect::<Vec<u8>>(),
|
||||
session.output_offset,
|
||||
exit_info,
|
||||
session.rows,
|
||||
session.cols,
|
||||
busy,
|
||||
)
|
||||
};
|
||||
|
||||
if !replay.is_empty() {
|
||||
use base64::Engine as _;
|
||||
send_routed_notification(
|
||||
gateway,
|
||||
NOTIFICATION_METHOD,
|
||||
serde_json::json!({
|
||||
"terminalId": pty_id,
|
||||
"type": "output",
|
||||
"data": base64::engine::general_purpose::STANDARD.encode(&replay),
|
||||
"outputOffset": output_offset,
|
||||
"isReplay": true,
|
||||
}),
|
||||
&target_client_id,
|
||||
);
|
||||
}
|
||||
|
||||
let exited = exit_info.is_some();
|
||||
if exited {
|
||||
send_routed_notification(
|
||||
gateway,
|
||||
NOTIFICATION_METHOD,
|
||||
serde_json::json!({
|
||||
"terminalId": pty_id,
|
||||
"type": "exit",
|
||||
"exitCode": exit_info,
|
||||
"isReplay": true,
|
||||
}),
|
||||
&target_client_id,
|
||||
);
|
||||
} else {
|
||||
send_busy_notification(pty_id, busy, &target_client_id, gateway);
|
||||
}
|
||||
|
||||
Ok(PtyLoadResult {
|
||||
terminal_id: pty_id.to_string(),
|
||||
rows,
|
||||
cols,
|
||||
exited,
|
||||
exit_code: exit_info,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::cell::RefCell;
|
||||
use std::rc::Rc;
|
||||
use std::time::Duration;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::acp_gateway;
|
||||
|
||||
type RecordedNotifications = Rc<RefCell<Vec<(String, serde_json::Value)>>>;
|
||||
|
||||
struct RecordingClient {
|
||||
notifications: RecordedNotifications,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait(?Send)]
|
||||
impl acp::Client for RecordingClient {
|
||||
async fn request_permission(
|
||||
&self,
|
||||
_: acp::RequestPermissionRequest,
|
||||
) -> acp::Result<acp::RequestPermissionResponse> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn session_notification(&self, _: acp::SessionNotification) -> acp::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ext_notification(&self, args: acp::ExtNotification) -> acp::Result<()> {
|
||||
let params: serde_json::Value =
|
||||
serde_json::from_str(args.params.get()).unwrap_or_default();
|
||||
self.notifications
|
||||
.borrow_mut()
|
||||
.push((args.method.to_string(), params));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn recording_gateway() -> (GatewaySender, RecordedNotifications) {
|
||||
let notifications: RecordedNotifications = Rc::new(RefCell::new(Vec::new()));
|
||||
let (sender, receiver) = acp_gateway::<acp::AgentSide, _>(RecordingClient {
|
||||
notifications: notifications.clone(),
|
||||
});
|
||||
tokio::task::spawn_local(receiver.run());
|
||||
(sender, notifications)
|
||||
}
|
||||
|
||||
async fn create_test_pty(gateway: GatewaySender) -> String {
|
||||
let env = HashMap::from([("ENV".to_string(), String::new())]);
|
||||
create_pty(
|
||||
Some("/bin/sh"),
|
||||
None,
|
||||
env,
|
||||
24,
|
||||
80,
|
||||
None,
|
||||
gateway,
|
||||
TargetClientId::None,
|
||||
)
|
||||
.await
|
||||
.expect("create test pty")
|
||||
}
|
||||
|
||||
fn busy_event_types(notifications: &RecordedNotifications, pty_id: &str) -> Vec<String> {
|
||||
notifications
|
||||
.borrow()
|
||||
.iter()
|
||||
.filter(|(method, params)| {
|
||||
method == NOTIFICATION_METHOD && params["terminalId"] == pty_id
|
||||
})
|
||||
.filter_map(|(_, params)| match params["type"].as_str() {
|
||||
Some(t @ ("process_started" | "process_ended")) => Some(t.to_string()),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn wait_for_busy_events(
|
||||
notifications: &RecordedNotifications,
|
||||
pty_id: &str,
|
||||
expected: &[&str],
|
||||
deadline: Duration,
|
||||
) -> Vec<String> {
|
||||
let started = tokio::time::Instant::now();
|
||||
loop {
|
||||
let events = busy_event_types(notifications, pty_id);
|
||||
if events == expected {
|
||||
return events;
|
||||
}
|
||||
if started.elapsed() > deadline {
|
||||
return events;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn long_running_command_emits_balanced_started_then_ended_pair() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
let (gateway, notifications) = recording_gateway();
|
||||
let pty_id = create_test_pty(gateway).await;
|
||||
|
||||
write_pty_input(&pty_id, b"sleep 2\n")
|
||||
.await
|
||||
.expect("write command");
|
||||
|
||||
let after_start = wait_for_busy_events(
|
||||
¬ifications,
|
||||
&pty_id,
|
||||
&["process_started"],
|
||||
Duration::from_secs(10),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(after_start, vec!["process_started"]);
|
||||
|
||||
let after_end = wait_for_busy_events(
|
||||
¬ifications,
|
||||
&pty_id,
|
||||
&["process_started", "process_ended"],
|
||||
Duration::from_secs(10),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(after_end, vec!["process_started", "process_ended"]);
|
||||
|
||||
close_pty(&pty_id).await.expect("close pty");
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn idle_shell_emits_no_busy_notifications() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
let (gateway, notifications) = recording_gateway();
|
||||
let pty_id = create_test_pty(gateway).await;
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(BUSY_POLL_INTERVAL_MS * 3)).await;
|
||||
|
||||
assert_eq!(
|
||||
busy_event_types(¬ifications, &pty_id),
|
||||
Vec::<String>::new()
|
||||
);
|
||||
|
||||
close_pty(&pty_id).await.expect("close pty");
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_paths::AbsPathBuf;
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum TerminalError {
|
||||
#[error("{0}")]
|
||||
Other(String),
|
||||
#[error("Command could not be quoted")]
|
||||
CommandNotQuoted,
|
||||
}
|
||||
|
||||
pub struct TerminalRunRequest {
|
||||
pub tool_call_id: acp::ToolCallId,
|
||||
pub command: String,
|
||||
pub cwd: AbsPathBuf,
|
||||
pub env: HashMap<String, String>,
|
||||
pub timeout: Duration,
|
||||
pub output_byte_limit: usize,
|
||||
/// Whether to stream output updates and register in the terminal registry.
|
||||
/// - `true`: Agent tool calls (streaming updates, killable via x.ai/terminal/kill)
|
||||
/// - `false`: Extension methods, git helpers (no updates, not killable)
|
||||
pub stream: bool,
|
||||
/// Optional file path to write output incrementally (for background tasks).
|
||||
/// When Some, the streaming loop writes output to this file as it arrives.
|
||||
/// This allows retrieval of full output even after in-memory buffer is truncated.
|
||||
pub output_file: Option<PathBuf>,
|
||||
}
|
||||
|
||||
pub struct TerminalRunResult {
|
||||
pub combined_output: String,
|
||||
pub exit_code: Option<i32>,
|
||||
pub truncated: bool,
|
||||
pub signal: Option<String>,
|
||||
pub timed_out: bool,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub trait AsyncTerminalRunner: Send + Sync {
|
||||
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, TerminalError>;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user