ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

M0: compilable skeleton — Kigi 0.1.0 fork surgery

Browse Source 
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:
ZacharyZhang-NY
2026-07-17 05:31:01 -04:00
commit d6c20fc13f
2612 changed files with 1353757 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-tools/src/computer/local/cgroup.rs
+538
View File
@@ -0,0 +1,538 @@
//! Cgroup v2 memory-high monitor for graceful OOM handling.
//!
//! Approach:
//!
//! 1. On startup we create a child cgroup under the current process's cgroup and
//! configure:
//! - `memory.high` = soft limit (the "desired" ceiling)
//! - `memory.max` = soft limit + headroom (hard OOM kill boundary)
//!
//! 2. Before each spawned command, we write the child PID into `cgroup.procs` to
//! move it (and its entire process group) into the cgroup.
//!
//! 3. A background `MemoryHighMonitor` task watches `memory.events` via inotify.
//! When the kernel increments the `high` counter (meaning a process touched
//! `memory.high`), the monitor reads `memory.current` and — if RSS is still
//! above 90 % of `memory.high` — sends a `MemoryHighEvent` through a
//! `tokio::sync::watch` channel.
//!
//! 4. The terminal actor polls `monitor.try_recv()` on every tick. When it
//! receives an event it kills the offending process group with SIGKILL and
//! reports exit-code **137** (128 + SIGKILL) with signal `"oom"`.
//!
//! The kigi-tools process itself is **never** inside this cgroup — only spawned
//! child commands are. After the child exits the cgroup is empty until the next
//! command.
//!
//! ## Platform
//!
//! Everything compiles on all platforms, but the actual cgroup + inotify logic is
//! gated behind `#[cfg(target_os = "linux")]`. On macOS / Windows the public
//! constructors return `None` / no-op stubs so callers don't need `#[cfg]`.
/// Exit code for processes killed due to memory pressure.
/// Matches the POSIX convention: 128 + signal-number (SIGKILL = 9).
pub const PROCESS_OOM_EXIT_CODE: i32 = 137;
// ============================================================================
// Public types (cross-platform)
// ============================================================================
/// Event emitted when `memory.high` is breached and RSS is still above
/// the 90 % buffer threshold.
#[derive(Debug, Clone)]
pub struct MemoryHighEvent {
/// Current cgroup memory usage in bytes when the event fired.
pub memory_current: u64,
/// The configured `memory.high` threshold in bytes.
pub memory_high_threshold: u64,
}
/// Configuration for cgroup memory limits.
#[derive(Debug, Clone)]
pub struct CgroupMemoryConfig {
/// Soft memory limit (`memory.high`). When a process inside the cgroup
/// exceeds this, the monitor fires.
pub memory_high_bytes: u64,
/// Extra headroom above `memory.high` before the kernel hard-kills
/// (`memory.max = memory_high_bytes + headroom_bytes`).
/// A reasonable default is 256 MiB.
pub headroom_bytes: u64,
}
impl CgroupMemoryConfig {
/// memory.max = memory.high + headroom
#[cfg(target_os = "linux")]
fn memory_max(&self) -> u64 {
self.memory_high_bytes.saturating_add(self.headroom_bytes)
}
}
// ============================================================================
// Linux implementation
// ============================================================================
#[cfg(target_os = "linux")]
mod linux {
use super::*;
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::os::unix::ffi::OsStrExt;
use std::path::PathBuf;
use tokio::io::Interest;
use tokio::io::unix::AsyncFd;
use tokio::sync::watch;
// ── inotify FFI ──────────────────────────────────────────────────────
unsafe fn inotify_init1(flags: libc::c_int) -> std::io::Result<i32> {
#[allow(clippy::cast_possible_truncation)]
let ret = unsafe { libc::syscall(libc::SYS_inotify_init1, flags) as libc::c_int };
if ret < 0 {
return Err(std::io::Error::last_os_error());
}
Ok(ret)
}
unsafe fn inotify_add_watch(
fd: libc::c_int,
pathname: *const libc::c_char,
mask: u32,
) -> std::io::Result<i32> {
#[allow(clippy::cast_possible_truncation)]
let ret = unsafe {
libc::syscall(libc::SYS_inotify_add_watch, fd, pathname, mask) as libc::c_int
};
if ret < 0 {
return Err(std::io::Error::last_os_error());
}
Ok(ret)
}
const IN_MODIFY: u32 = 0x0000_0002;
// ── Inotify wrapper ──────────────────────────────────────────────────
struct Inotify {
fd: AsyncFd<OwnedFd>,
}
impl Inotify {
fn new() -> std::io::Result<Self> {
let raw = unsafe { inotify_init1(libc::O_NONBLOCK | libc::O_CLOEXEC) }?;
let owned = unsafe { OwnedFd::from_raw_fd(raw) };
let fd = AsyncFd::with_interest(owned, Interest::READABLE)?;
Ok(Inotify { fd })
}
fn add_watch(&self, path: &std::path::Path) -> std::io::Result<i32> {
let mut bytes = path.as_os_str().as_bytes().to_vec();
bytes.push(0); // NUL-terminate
let wd = unsafe {
inotify_add_watch(
self.fd.get_ref().as_raw_fd(),
bytes.as_ptr().cast(),
IN_MODIFY,
)
}?;
Ok(wd)
}
async fn wait_and_drain(&self) -> std::io::Result<()> {
let mut guard = self.fd.readable().await?;
// Drain all pending inotify events
let mut buf = [0u8; 4096];
loop {
let n = unsafe {
libc::read(
guard.get_inner().as_raw_fd(),
buf.as_mut_ptr().cast(),
buf.len(),
)
};
if n > 0 {
continue;
}
break;
}
guard.clear_ready();
Ok(())
}
}
// ── Cgroup helpers ───────────────────────────────────────────────────
/// Read `/proc/self/cgroup` to find our own cgroup path (cgroup v2 unified).
fn read_self_cgroup() -> std::io::Result<String> {
let contents = std::fs::read_to_string("/proc/self/cgroup")?;
// In cgroupv2 unified hierarchy, the line is "0::<path>"
for line in contents.lines() {
if let Some(rest) = line.strip_prefix("0::") {
return Ok(rest.to_owned());
}
}
Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
"Could not find cgroupv2 entry in /proc/self/cgroup",
))
}
/// Parse the `high <N>` counter from `memory.events` contents.
fn parse_memory_events_high(contents: &str) -> Option<u64> {
for line in contents.lines() {
if let Some(value) = line.strip_prefix("high ") {
return value.trim().parse::<u64>().ok();
}
}
None
}
// ── CgroupHandle ─────────────────────────────────────────────────────
/// Owns the lifecycle of a child cgroup directory.
pub(crate) struct CgroupHandle {
fs_path: PathBuf,
}
impl CgroupHandle {
/// Create a child cgroup under the current process's cgroup and
/// configure memory limits.
pub(crate) async fn create(config: &CgroupMemoryConfig) -> std::io::Result<Self> {
let self_cgroup = read_self_cgroup()?;
let name = format!("kigi-tools-{}", uuid::Uuid::now_v7());
let fs_path = PathBuf::from(format!("/sys/fs/cgroup{}/{}", self_cgroup, name));
tokio::fs::create_dir_all(&fs_path).await?;
// Enable memory + cpu controllers in the child cgroup's parent
// (the parent's subtree_control must list the controllers).
let parent = fs_path.parent().unwrap();
let subtree_ctl = parent.join("cgroup.subtree_control");
// Best-effort; may already be enabled or not permitted.
let _ = tokio::fs::write(&subtree_ctl, "+memory +cpu").await;
// Configure memory.high (soft limit)
let memory_high_path = fs_path.join("memory.high");
tokio::fs::write(&memory_high_path, config.memory_high_bytes.to_string()).await?;
// Configure memory.max (hard limit = high + headroom)
let memory_max_path = fs_path.join("memory.max");
tokio::fs::write(&memory_max_path, config.memory_max().to_string()).await?;
tracing::info!(
cgroup = %fs_path.display(),
memory_high = config.memory_high_bytes,
memory_max = config.memory_max(),
"Created cgroup with memory limits"
);
Ok(CgroupHandle { fs_path })
}
/// Move a process (by PID) into this cgroup.
pub(crate) async fn add_process(&self, pid: u32) -> std::io::Result<()> {
let procs_path = self.fs_path.join("cgroup.procs");
tokio::fs::write(&procs_path, pid.to_string()).await
}
/// Read `memory.current` from this cgroup.
#[allow(dead_code)]
pub(crate) async fn memory_current(&self) -> std::io::Result<u64> {
let s: String = tokio::fs::read_to_string(self.fs_path.join("memory.current")).await?;
s.trim()
.parse::<u64>()
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
}
/// Filesystem path to this cgroup.
pub(crate) fn path(&self) -> &std::path::Path {
&self.fs_path
}
}
impl Drop for CgroupHandle {
fn drop(&mut self) {
let path = self.fs_path.clone();
// Always use tokio::spawn: the cleanup future is Send and Drop
// can fire after the LocalSet has shut down, making spawn_local
// unsafe here.
tokio::spawn(async move {
let kill_path = path.join("cgroup.kill");
let _ = tokio::fs::write(&kill_path, "1").await;
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
if let Err(e) = tokio::fs::remove_dir(&path).await {
tracing::debug!(
cgroup = %path.display(),
"Failed to remove cgroup dir on drop (may already be gone): {e}"
);
}
});
}
}
// ── MemoryHighMonitor ────────────────────────────────────────────────
/// Watches `memory.events` via inotify and signals when the `high`
/// counter increments while RSS is still above 90% of the threshold.
pub(crate) struct MemoryHighMonitor {
rx: watch::Receiver<Option<MemoryHighEvent>>,
/// Dropping this aborts the background task.
_handle: tokio::task::JoinHandle<()>,
}
impl MemoryHighMonitor {
/// Start monitoring the given cgroup for memory.high events.
pub(crate) async fn start(
cgroup_path: PathBuf,
memory_high_threshold: u64,
use_spawn_local: bool,
) -> std::io::Result<Self> {
let (tx, rx) = watch::channel(None);
let inotify = Inotify::new()?;
let events_path = cgroup_path.join("memory.events");
inotify.add_watch(&events_path)?;
let monitor_fut = Self::monitor_loop(inotify, cgroup_path, memory_high_threshold, tx);
let handle = if use_spawn_local {
tokio::task::spawn_local(monitor_fut)
} else {
tokio::spawn(monitor_fut)
};
Ok(MemoryHighMonitor {
rx,
_handle: handle,
})
}
/// Non-blocking check: returns `Some(event)` if memory.high was
/// breached since the last call, `None` otherwise.
pub(crate) fn try_recv(&mut self) -> Option<MemoryHighEvent> {
if self.rx.has_changed().unwrap_or(false) {
self.rx.borrow_and_update().clone()
} else {
None
}
}
async fn monitor_loop(
inotify: Inotify,
cgroup_path: PathBuf,
memory_high_threshold: u64,
tx: watch::Sender<Option<MemoryHighEvent>>,
) {
let events_path = cgroup_path.join("memory.events");
let current_path = cgroup_path.join("memory.current");
// Read baseline high counter
let mut last_high_count = Self::read_high_counter(&events_path).await.unwrap_or(0);
loop {
// Block until inotify fires (memory.events was modified)
if inotify.wait_and_drain().await.is_err() {
break;
}
// Read new high counter
let current_high = Self::read_high_counter(&events_path).await.unwrap_or(0);
if current_high <= last_high_count {
continue;
}
last_high_count = current_high;
// Read current memory usage
let memory_current = match tokio::fs::read_to_string(&current_path).await {
Ok(s) => {
let s: String = s;
s.trim().parse::<u64>().unwrap_or(0)
}
Err(_) => continue,
};
// Only fire if still above 90 % of threshold (avoids false
// positives from transient spikes the kernel already handled).
let buffer_threshold = memory_high_threshold.saturating_mul(9) / 10;
if memory_current >= buffer_threshold {
let event = MemoryHighEvent {
memory_current,
memory_high_threshold,
};
if tx.send(Some(event)).is_err() {
break; // receiver dropped
}
}
}
}
async fn read_high_counter(events_path: &std::path::Path) -> Option<u64> {
let contents = tokio::fs::read_to_string(events_path).await.ok()?;
parse_memory_events_high(&contents)
}
}
impl Drop for MemoryHighMonitor {
fn drop(&mut self) {
self._handle.abort();
}
}
}
// ============================================================================
// Cross-platform re-exports
// ============================================================================
/// Cgroup handle — owns the child cgroup's lifecycle.
///
/// On Linux, this creates a real cgroupv2 directory with memory limits.
/// On other platforms, this is a no-op.
pub struct CgroupGuard {
#[cfg(target_os = "linux")]
inner: Option<linux::CgroupHandle>,
}
impl CgroupGuard {
/// Try to create a cgroup with the given memory config.
/// Returns a guard that cleans up the cgroup on drop.
///
/// On non-Linux platforms this always returns a no-op guard.
/// On Linux, if cgroup creation fails (e.g., not running as root,
/// cgroupv2 not available), it logs a warning and returns a no-op guard.
pub async fn try_create(config: &CgroupMemoryConfig) -> Self {
#[cfg(target_os = "linux")]
{
match linux::CgroupHandle::create(config).await {
Ok(handle) => CgroupGuard {
inner: Some(handle),
},
Err(e) => {
tracing::warn!("Failed to create cgroup (falling back to no limits): {e}");
CgroupGuard { inner: None }
}
}
}
#[cfg(not(target_os = "linux"))]
{
let _ = config;
CgroupGuard {}
}
}
/// No-op guard with no backing cgroup.
pub fn noop() -> Self {
#[cfg(target_os = "linux")]
{
CgroupGuard { inner: None }
}
#[cfg(not(target_os = "linux"))]
{
CgroupGuard {}
}
}
/// Move a process into this cgroup by PID.
/// No-op if cgroup was not created.
pub async fn add_process(&self, _pid: u32) -> std::io::Result<()> {
#[cfg(target_os = "linux")]
{
if let Some(ref handle) = self.inner {
return handle.add_process(_pid).await;
}
}
Ok(())
}
/// Returns the cgroup filesystem path, if available.
#[allow(dead_code)]
pub fn path(&self) -> Option<&std::path::Path> {
#[cfg(target_os = "linux")]
{
if let Some(ref handle) = self.inner {
return Some(handle.path());
}
}
None
}
/// Returns true if this guard has a real cgroup backing it.
pub fn is_active(&self) -> bool {
#[cfg(target_os = "linux")]
{
self.inner.is_some()
}
#[cfg(not(target_os = "linux"))]
{
false
}
}
}
/// Memory-high monitor — watches for memory pressure events.
///
/// On Linux, uses inotify on `memory.events`.
/// On other platforms, this is a no-op that never fires.
pub struct MemoryMonitor {
#[cfg(target_os = "linux")]
inner: Option<linux::MemoryHighMonitor>,
}
impl MemoryMonitor {
/// Start monitoring the given cgroup guard for memory.high events.
/// Returns a no-op monitor if the guard has no backing cgroup.
pub async fn start(
guard: &CgroupGuard,
config: &CgroupMemoryConfig,
use_spawn_local: bool,
) -> Self {
#[cfg(target_os = "linux")]
{
if let Some(ref handle) = guard.inner {
match linux::MemoryHighMonitor::start(
handle.path().to_path_buf(),
config.memory_high_bytes,
use_spawn_local,
)
.await
{
Ok(monitor) => {
return MemoryMonitor {
inner: Some(monitor),
};
}
Err(e) => {
tracing::warn!("Failed to start memory monitor: {e}");
}
}
}
MemoryMonitor { inner: None }
}
#[cfg(not(target_os = "linux"))]
{
let _ = (guard, config, use_spawn_local);
MemoryMonitor {}
}
}
/// No-op monitor that never fires.
pub fn noop() -> Self {
#[cfg(target_os = "linux")]
{
MemoryMonitor { inner: None }
}
#[cfg(not(target_os = "linux"))]
{
MemoryMonitor {}
}
}
/// Non-blocking poll: returns `Some(event)` if memory.high was breached
/// since the last call, `None` otherwise.
pub fn try_recv(&mut self) -> Option<MemoryHighEvent> {
#[cfg(target_os = "linux")]
{
if let Some(ref mut monitor) = self.inner {
return monitor.try_recv();
}
}
None
}
}
crates/codegen/kigi-tools/src/computer/local/embedded_search_tools.rs
+805
View File
@@ -0,0 +1,805 @@
//! Shadow `find`→`bfs` and `grep`→`ugrep` when those binaries resolve.
//!
//! Per-tool enable state (default on) is resolved by the host via the shared
//! config helper `kigi-shell::util::config::resolve_search_tools_enabled`
//! (requirements > env `KIGI_TOOLS_FIND_BFS` / `KIGI_TOOLS_GREP_UGREP` (+
//! `KIGI_FIND_BFS` / `KIGI_GREP_UGREP` aliases, `DISABLE_EMBEDDED_SEARCH_TOOLS`
//! master) > `[toolset.bash]` config.toml > managed > default), baked into the
//! `LocalTerminalBackend` as a [`SearchShadowConfig`] and passed to
//! [`search_injection`] per command. The enable state lives on the backend (not
//! a process-global): a subagent that reuses the parent's backend inherits the
//! parent's shadows instead of clobbering a shared static. This module no longer
//! parses the flags itself.
//!
//! Resolve (host side, memoized): env override if a regular file → bundled binary
//! (release builds, self-extracted to `~/.kigi/vendor/<name>-<ver>-<target>`) →
//! `~/.kigi/vendor/{name}` if a regular file → `which` on the agent `$PATH`.
//! Env/vendor only require `is_file()` as a lenient hint (no `--version` probe).
//! This memoized path is only a *hint*: the injected shadow re-resolves at
//! **call time** — it uses the hint when it's still *executable* (`[ -x ]`), else
//! `command -v {bin}` on the live shell `PATH` (which includes login/rc additions
//! the agent process may lack), else falls back to the OS `{name}`. So a removed
//! or non-executable binary self-heals to OS `find`/`grep`, and a binary
//! reachable only through the login shell is still found.
//!
//! Inject is **always** non-empty on Unix callers: either install a shadow
//! function (which tags itself with a `__grok_shadow_{name}` marker) or a
//! marker-gated `unalias`+`unset -f` that drops *only* a prior harness shadow —
//! never a user-defined `find`/`grep` function replayed from the snapshot.
use super::SearchShadowConfig;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
/// BRE + ignore-files; recursive flags are no-ops on stdin (computer parity).
const UGREP_DEFAULT_ARGS: &[&str] = &[
"-G",
"--ignore-files",
"--hidden",
"-I",
"--exclude-dir=.git",
"--exclude-dir=.svn",
"--exclude-dir=.hg",
"--exclude-dir=.bzr",
"--exclude-dir=.jj",
"--exclude-dir=.sl",
];
// Binaries embedded by build.rs when `KIGI_TOOLS_BUNDLE_{BFS,UGREP}_PATH` is set
// (release pipeline). Self-extracted to `~/.kigi/vendor` on first use, mirroring
// the ripgrep bundling in `grok_build::grep::ripgrep`.
#[cfg(bundle_bfs)]
const BFS_BYTES: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/bundle-bfs/bfs-",
env!("KIGI_TOOLS_BFS_VER"),
"-",
env!("KIGI_TOOLS_BFS_TARGET"),
".bin"
));
#[cfg(bundle_ugrep)]
const UGREP_BYTES: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
"/bundle-ugrep/ugrep-",
env!("KIGI_TOOLS_UGREP_VER"),
"-",
env!("KIGI_TOOLS_UGREP_TARGET"),
".bin"
));
/// Oneline inject for shell wrappers; always ends with `"; "`.
///
/// `cfg` is the backend's resolved per-tool enable state (see module docs); it
/// is passed in per command rather than read from a process-global so subagents
/// sharing a backend can't clobber each other's shadows.
pub fn search_injection(cfg: SearchShadowConfig) -> String {
build_injection(cfg.find_bfs, cfg.grep_ugrep, resolved_tools())
}
/// Compose the inject from per-tool enable flags + resolved binaries. An enabled
/// tool installs a self-resolving shadow (the memoized `tools` path is only a
/// fast-path hint; the shadow re-resolves at call time and falls back to the OS
/// binary — see [`shell_function`]). A disabled tool emits a marker-gated
/// `restore` that drops only a prior harness shadow. Kept pure (flags/tools
/// passed in) so tests need no process-global env mutation — that is UB against
/// the `shell_state` integration tests that read env / spawn children
/// concurrently.
fn build_injection(find_on: bool, grep_on: bool, tools: &ResolvedTools) -> String {
let find = if find_on {
shell_function("find", "bfs", tools.bfs.as_deref(), &[])
} else {
restore_command("find")
};
let grep = if grep_on {
shell_function("grep", "ugrep", tools.ugrep.as_deref(), UGREP_DEFAULT_ARGS)
} else {
restore_command("grep")
};
format!("{find}; {grep}; ")
}
/// Drop a *previously installed harness* shadow so command-word `{name}` uses the
/// OS binary again. Gated on the `__grok_shadow_{name}` marker that
/// [`shell_function`] sets, so a user-defined `{name}` function replayed from the
/// shell snapshot is left intact — only the harness's own shadow is removed.
/// `set -u`/`set -e` safe and idempotent (`unset -f` is bash + zsh).
fn restore_command(name: &str) -> String {
format!(
"if [ -n \"${{__grok_shadow_{name}-}}\" ]; then \
unalias {name} 2>/dev/null || true; \
unset -f {name} 2>/dev/null || true; \
unset __grok_shadow_{name} 2>/dev/null || true; \
fi"
)
}
struct ResolvedTools {
bfs: Option<PathBuf>,
ugrep: Option<PathBuf>,
}
fn resolved_tools() -> &'static ResolvedTools {
static TOOLS: OnceLock<ResolvedTools> = OnceLock::new();
TOOLS.get_or_init(|| ResolvedTools {
bfs: resolve_tool("bfs", "KIGI_TOOLS_BFS_PATH", bundled_bfs()),
ugrep: resolve_tool("ugrep", "KIGI_TOOLS_UGREP_PATH", bundled_ugrep()),
})
}
/// Write embedded `bytes` to `~/.kigi/vendor/<versioned_name>` (chmod 755) on
/// first use and return the path; reused on later runs. Versioned so bumping the
/// bundled version writes a fresh file instead of reusing a stale one.
#[cfg(any(bundle_bfs, bundle_ugrep))]
fn extract_bundled(versioned_name: &str, bytes: &[u8]) -> std::io::Result<PathBuf> {
use std::os::unix::fs::PermissionsExt;
let dir = crate::util::kigi_home().join("vendor");
let dest = dir.join(versioned_name);
if !dest.exists() {
std::fs::create_dir_all(&dir)?;
// Write to a unique temp then atomically rename, so a concurrent first
// use (or an interrupted write) can't leave a half-written binary that
// gets cached and exec'd.
let tmp = dir.join(format!(
"{versioned_name}.tmp.{}.{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0)
));
std::fs::write(&tmp, bytes)?;
let mut perms = std::fs::metadata(&tmp)?.permissions();
perms.set_mode(0o755);
std::fs::set_permissions(&tmp, perms)?;
// Rename is atomic on the same filesystem. If another process won the
// race, `dest` already exists and is correct — drop our temp copy.
if let Err(e) = std::fs::rename(&tmp, &dest) {
let _ = std::fs::remove_file(&tmp);
if !dest.exists() {
return Err(e);
}
}
}
Ok(dest)
}
/// Path to the bundled `bfs` (extracted on first use), or `None` when not bundled.
fn bundled_bfs() -> Option<PathBuf> {
#[cfg(bundle_bfs)]
{
extract_bundled(
concat!(
"bfs-",
env!("KIGI_TOOLS_BFS_VER"),
"-",
env!("KIGI_TOOLS_BFS_TARGET")
),
BFS_BYTES,
)
.ok()
}
#[cfg(not(bundle_bfs))]
{
None
}
}
/// Path to the bundled `ugrep` (extracted on first use), or `None` when not bundled.
fn bundled_ugrep() -> Option<PathBuf> {
#[cfg(bundle_ugrep)]
{
extract_bundled(
concat!(
"ugrep-",
env!("KIGI_TOOLS_UGREP_VER"),
"-",
env!("KIGI_TOOLS_UGREP_TARGET")
),
UGREP_BYTES,
)
.ok()
}
#[cfg(not(bundle_ugrep))]
{
None
}
}
fn resolve_tool(bin_name: &str, env_override: &str, bundled: Option<PathBuf>) -> Option<PathBuf> {
resolve_tool_from(
std::env::var_os(env_override).map(PathBuf::from),
bundled,
crate::util::kigi_home().join("vendor").join(bin_name),
bin_name,
)
}
/// Resolution order: explicit env path → bundled (self-extracted) →
/// `~/.kigi/vendor/<bin>` → `which`. Env and vendor only require `is_file()` here
/// (a lenient hint, no `+x` probe) so an odd-permission copy still resolves; the
/// injected shadow gates on `[ -x ]` at call time and falls back to the OS binary
/// if the hint isn't executable, so a non-exec path can't hard-fail `find`/`grep`.
fn resolve_tool_from(
env_path: Option<PathBuf>,
bundled: Option<PathBuf>,
vendor: PathBuf,
bin_name: &str,
) -> Option<PathBuf> {
if let Some(path) = env_path
&& path.is_file()
{
return Some(path);
}
if let Some(path) = bundled {
return Some(path);
}
if vendor.is_file() {
return Some(vendor);
}
which::which(bin_name).ok()
}
fn bash_safe_quote(s: &str) -> String {
if s.chars().all(|c| {
c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '/' | '=' | ':' | '@' | '+' | '-')
}) {
return s.to_string();
}
format!("'{}'", s.replace('\'', "'\\''"))
}
/// Oneline `name() { … }` for `-c` inject — a *self-resolving* shadow.
///
/// At call time it picks the binary: the host-resolved `preferred` path
/// (bundled/env/vendor/which) when that file still exists, else `command -v
/// {bin_name}` on the live shell `PATH` (which carries login/rc additions the
/// agent process may not have), else it falls back to the OS `{name}`. This keeps
/// the fast hard-coded path for the common case while self-healing when the
/// binary was removed (revalidation) or is only reachable through the shell's
/// richer `PATH`.
///
/// `exec -a` runs inside a subshell so a top-level call can't replace the wrapper
/// shell (it must survive to dump state); a call already inside a subshell
/// (`BASH_SUBSHELL > 0`, bash) execs directly to skip a fork. `${ZSH_VERSION-}`
/// keeps the probe `set -u`-safe (a bare `$ZSH_VERSION` aborts bash under
/// nounset). `exec -a` gives the binary the `find`/`grep` argv0 (ps display +
/// ugrep grep-personality) in both bash and zsh. The trailing
/// `__grok_shadow_{name}=1` marks this as a harness shadow so `restore_command`
/// only ever removes our own function — never a user's.
fn shell_function(
name: &str,
bin_name: &str,
preferred: Option<&Path>,
prepend_args: &[&str],
) -> String {
let qpref = preferred
.map(|p| bash_safe_quote(&p.to_string_lossy()))
.unwrap_or_else(|| "''".to_string());
let prepend = {
let qargs: Vec<String> = prepend_args.iter().map(|a| bash_safe_quote(a)).collect();
if qargs.is_empty() {
String::new()
} else {
format!("{} ", qargs.join(" "))
}
};
// `local __grok_bin` is re-resolved every call. The host hint is trusted
// only when it's *executable* (`[ -x ]`, not just `[ -f ]`): the resolver
// accepts any regular file as a hint, but `exec` needs `+x`, so a non-exec
// hint must fall through rather than hard-fail with no OS fallback. Then
// `command -v` on the live shell PATH (returns an executable), else the OS
// binary. `|| __grok_bin=''` keeps the lookup `set -e`-safe (a failed
// `command -v` would otherwise abort the function under errexit). The OS
// fallback uses `command {name}` to bypass this function. `{prepend}` is
// empty for find, the ugrep default flags for grep (and is omitted from the
// OS fallback, which gets the original args).
format!(
"unalias {name} 2>/dev/null || true; \
{name}() {{ \
local __grok_bin={qpref}; \
[ -x \"$__grok_bin\" ] || __grok_bin=$(command -v {bin_name} 2>/dev/null) || __grok_bin=''; \
if [ -z \"$__grok_bin\" ]; then command {name} \"$@\"; return; fi; \
if [[ -z ${{ZSH_VERSION-}} ]] && (( BASH_SUBSHELL > 0 )); then \
exec -a {name} \"$__grok_bin\" {prepend}\"$@\"; \
else \
(exec -a {name} \"$__grok_bin\" {prepend}\"$@\"); \
fi; \
}}; \
__grok_shadow_{name}=1"
)
}
#[cfg(test)]
mod tests {
use super::*;
/// Both binaries resolved, for `build_injection` shape tests.
fn both_tools() -> ResolvedTools {
ResolvedTools {
bfs: Some(PathBuf::from("/tmp/bfs")),
ugrep: Some(PathBuf::from("/tmp/ugrep")),
}
}
#[test]
fn shell_function_shape() {
let fn_body = shell_function("find", "bfs", Some(Path::new("/tmp/bfs")), &[]);
assert!(fn_body.contains("unalias find"));
// Preferred path is the fast-path hint; the shadow execs `$__grok_bin`.
assert!(fn_body.contains("local __grok_bin=/tmp/bfs"));
assert!(fn_body.contains("exec -a find \"$__grok_bin\" \"$@\""));
// Hint is trusted only when executable (`[ -x ]`, not `[ -f ]`), so a
// non-exec hint falls through instead of hard-failing exec.
assert!(fn_body.contains("[ -x \"$__grok_bin\" ]"));
assert!(!fn_body.contains("[ -f \"$__grok_bin\" ]"));
// Self-heal: live-PATH lookup + OS fallback.
assert!(fn_body.contains("command -v bfs"));
assert!(fn_body.contains("command find \"$@\""));
assert!(fn_body.contains("BASH_SUBSHELL > 0"));
assert!(fn_body.contains("(exec -a find"));
// Marker so `restore_command` only removes our own shadow.
assert!(fn_body.contains("__grok_shadow_find=1"));
// set -u-safe zsh probe (a bare $ZSH_VERSION aborts bash under nounset).
assert!(fn_body.contains("${ZSH_VERSION-}"));
assert!(!fn_body.contains("[[ -n $ZSH_VERSION ]]"));
}
#[test]
fn shell_function_unresolved_uses_empty_hint() {
// No host-resolved path → empty hint, relies on live-PATH `command -v`.
let fn_body = shell_function("find", "bfs", None, &[]);
assert!(fn_body.contains("local __grok_bin=''"));
assert!(fn_body.contains("command -v bfs"));
assert!(fn_body.contains("command find \"$@\""));
}
#[test]
fn grep_prepends_ugrep_defaults() {
let fn_body = shell_function(
"grep",
"ugrep",
Some(Path::new("/tmp/ugrep")),
UGREP_DEFAULT_ARGS,
);
assert!(fn_body.contains("local __grok_bin=/tmp/ugrep"));
assert!(fn_body.contains("\"$__grok_bin\" -G --ignore-files --hidden -I"));
assert!(fn_body.contains("--exclude-dir=.git"));
assert!(fn_body.contains("command -v ugrep"));
}
#[test]
fn bash_safe_quote_escapes_metacharacters() {
assert_eq!(bash_safe_quote("/usr/bin/bfs"), "/usr/bin/bfs");
assert_eq!(bash_safe_quote("/tmp/my bfs"), "'/tmp/my bfs'");
let body = shell_function("find", "bfs", Some(Path::new("/tmp/evil$(id)")), &[]);
assert!(body.contains("local __grok_bin='/tmp/evil$(id)'"), "{body}");
assert!(!body.contains("=/tmp/evil$(id)"));
}
#[test]
fn restore_command_is_marker_gated() {
let r = restore_command("find");
// Only removes the harness shadow when our marker is set.
assert!(r.contains("if [ -n \"${__grok_shadow_find-}\" ]"));
assert!(r.contains("unalias find"));
assert!(r.contains("unset -f find"));
assert!(r.contains("unset __grok_shadow_find"));
}
#[test]
fn config_default_is_both_on() {
// Standalone/no-host backends default to shadowing both tools.
let cfg = SearchShadowConfig::default();
assert!(cfg.find_bfs);
assert!(cfg.grep_ugrep);
}
#[test]
fn build_injection_off_emits_marker_gated_restore_not_function() {
// Disabled tools emit a marker-gated restore so a stale harness shadow
// from a prior snapshot is dropped, but a user function is left intact.
let inject = build_injection(false, false, &both_tools());
assert!(inject.ends_with("; "));
assert!(inject.contains("if [ -n \"${__grok_shadow_find-}\" ]"));
assert!(inject.contains("if [ -n \"${__grok_shadow_grep-}\" ]"));
assert!(inject.contains("unset -f find"));
assert!(inject.contains("unset -f grep"));
assert!(!inject.contains("find()"));
assert!(!inject.contains("grep()"));
}
#[test]
fn build_injection_on_shadows_resolved_tools() {
let inject = build_injection(true, true, &both_tools());
assert!(inject.contains("find()"));
assert!(inject.contains("grep()"));
assert!(inject.contains("-G --ignore-files"));
assert!(inject.contains("__grok_shadow_find=1"));
}
#[test]
fn build_injection_enabled_unresolved_still_self_heals() {
// Enabled but no host-resolved path → still install a self-resolving
// shadow (live-PATH `command -v` + OS fallback), never a bare restore.
let tools = ResolvedTools {
bfs: None,
ugrep: None,
};
let inject = build_injection(true, true, &tools);
assert!(inject.contains("find()"));
assert!(inject.contains("grep()"));
assert!(inject.contains("command -v bfs"));
assert!(inject.contains("command -v ugrep"));
// OS fallback present; not a marker-gated restore.
assert!(inject.contains("command find \"$@\""));
assert!(!inject.contains("if [ -n \"${__grok_shadow_find-}\" ]"));
}
#[test]
fn injection_always_nonempty_and_trailing_sep() {
// Structural invariants that hold regardless of host flags/binaries:
// never empty (so a stale snapshot shadow is always overwritten) and a
// trailing `"; "` separator before the user command.
for cfg in [
SearchShadowConfig {
find_bfs: true,
grep_ugrep: true,
},
SearchShadowConfig {
find_bfs: false,
grep_ugrep: false,
},
] {
let inject = search_injection(cfg);
assert!(!inject.is_empty());
assert!(inject.ends_with("; "));
assert!(!inject.contains('\n'));
}
}
#[test]
fn env_override_accepts_regular_file_without_exec_bit() {
let bin = std::env::temp_dir().join(format!(
"grok-bfs-noexec-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::write(&bin, b"#!/bin/sh\necho ok\n").unwrap();
// Mode 0o644 — no execute bit; Nix-like / restricted copies.
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mut perms = std::fs::metadata(&bin).unwrap().permissions();
perms.set_mode(0o644);
std::fs::set_permissions(&bin, perms).unwrap();
}
// Bundled + vendor intentionally absent so the env override is what wins.
let got = resolve_tool_from(
Some(bin.clone()),
None,
PathBuf::from("/nonexistent/vendor/bfs"),
"bfs",
);
let _ = std::fs::remove_file(&bin);
assert_eq!(got.as_deref(), Some(bin.as_path()));
}
#[test]
fn resolve_tool_precedence() {
// Real temp files so the is_file() checks pass.
let dir = std::env::temp_dir().join(format!("grok-resolve-{}-{:?}", std::process::id(), {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
}));
std::fs::create_dir_all(&dir).unwrap();
let envp = dir.join("env");
let bundled = dir.join("bundled");
let vendor = dir.join("vendor");
for p in [&envp, &bundled, &vendor] {
std::fs::write(p, b"x").unwrap();
}
// env override beats everything.
assert_eq!(
resolve_tool_from(
Some(envp.clone()),
Some(bundled.clone()),
vendor.clone(),
"bfs"
)
.as_deref(),
Some(envp.as_path())
);
// bundled beats the manual vendor copy.
assert_eq!(
resolve_tool_from(None, Some(bundled.clone()), vendor.clone(), "bfs").as_deref(),
Some(bundled.as_path())
);
// vendor used when neither env nor bundled is present.
assert_eq!(
resolve_tool_from(None, None, vendor.clone(), "bfs").as_deref(),
Some(vendor.as_path())
);
// A non-file env override is ignored and falls through to vendor.
assert_eq!(
resolve_tool_from(Some(dir.join("missing")), None, vendor.clone(), "bfs").as_deref(),
Some(vendor.as_path())
);
let _ = std::fs::remove_dir_all(&dir);
}
/// Only compiled when the binaries are actually bundled (release pipeline, or
/// `KIGI_TOOLS_BUNDLE_{BFS,UGREP}_PATH` at build time). Verifies the embedded
/// bytes self-extract under `~/.kigi/vendor` and the extracted `bfs` runs.
#[cfg(all(bundle_bfs, bundle_ugrep))]
#[test]
fn bundled_binaries_extract_and_run() {
let vendor = crate::util::kigi_home().join("vendor");
let bfs = bundled_bfs().expect("bfs should be bundled");
let ugrep = bundled_ugrep().expect("ugrep should be bundled");
assert!(bfs.is_file() && bfs.starts_with(&vendor), "bfs at {bfs:?}");
assert!(
ugrep.is_file() && ugrep.starts_with(&vendor),
"ugrep at {ugrep:?}"
);
let v = std::process::Command::new(&bfs)
.arg("--version")
.output()
.unwrap();
assert!(
String::from_utf8_lossy(&v.stdout)
.to_lowercase()
.contains("bfs"),
"bfs --version: {}",
String::from_utf8_lossy(&v.stdout)
);
}
/// Regression: a bare `$ZSH_VERSION` probe aborted the shadow under `set -u`
/// ("unbound variable"), killing find/grep and dropping the state dump. Runs
/// the generated function with `/bin/echo` standing in for the binary.
#[test]
fn shadow_runs_under_nounset_bash() {
let Ok(bash) = which::which("bash") else {
return;
};
let inject = shell_function("find", "bfs", Some(Path::new("/bin/echo")), &[]);
let script = format!("set -euo pipefail; {inject}; find hello world");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"find shadow aborted under set -u: {:?}\nstderr: {}",
out.status,
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hello world");
}
/// Command substitution runs with `BASH_SUBSHELL > 0`, taking the direct
/// `exec -a` branch; it must still forward args and succeed under `set -u`,
/// and the outer shell must survive (so its state dump can run).
#[test]
fn shadow_execs_directly_in_subshell_bash() {
let Ok(bash) = which::which("bash") else {
return;
};
let inject = shell_function("find", "bfs", Some(Path::new("/bin/echo")), &[]);
let script =
format!("set -euo pipefail; {inject}; printf '[%s]' \"$(find sub shell)\"; echo ALIVE");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"stderr: {}",
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout), "[sub shell]ALIVE\n");
}
/// zsh takes the subshell `exec -a` branch (consistent argv0 with bash); it
/// must run under nounset without referencing the bash-only `BASH_SUBSHELL`.
#[test]
fn shadow_runs_under_nounset_zsh() {
let Ok(zsh) = which::which("zsh") else {
return;
};
let inject = shell_function("grep", "ugrep", Some(Path::new("/bin/echo")), &[]);
let script = format!("setopt nounset errexit pipefail; {inject}; grep hi there");
let out = std::process::Command::new(&zsh)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"grep shadow aborted in zsh: {:?}\nstderr: {}",
out.status,
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "hi there");
}
/// #3 regression: a user-defined `find` must survive a disabled-tool restore.
/// Without the marker gate, the unconditional `unset -f find` dropped it.
#[test]
fn restore_preserves_user_function_without_marker() {
let Ok(bash) = which::which("bash") else {
return;
};
let restore = restore_command("find");
let script =
format!("set -euo pipefail; find() {{ echo USERFIND; }}; {restore}; find anything");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"stderr: {}",
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "USERFIND");
}
/// The marker-gated restore *does* drop a prior harness shadow (which set the
/// marker), so a flipped-off tool falls back to the OS binary.
#[test]
fn restore_removes_harness_shadow_with_marker() {
let Ok(bash) = which::which("bash") else {
return;
};
let shadow = shell_function("find", "bfs", Some(Path::new("/bin/echo")), &[]);
let restore = restore_command("find");
let script = format!("set -euo pipefail; {shadow}; {restore}; type -t find || echo NOFUNC");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"stderr: {}",
String::from_utf8_lossy(&out.stderr)
);
let stdout = String::from_utf8_lossy(&out.stdout);
assert!(
!stdout.lines().any(|l| l.trim() == "function"),
"find should no longer be a function after restore: {stdout:?}"
);
}
/// #1 + #4: the host hint points at a missing file, but the binary is on the
/// live shell `PATH` — the shadow re-resolves via `command -v` at call time.
#[test]
fn shadow_self_heals_via_path_lookup() {
let Ok(bash) = which::which("bash") else {
return;
};
let dir = std::env::temp_dir().join(format!(
"grok-selfheal-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).unwrap();
let bin = dir.join("myfind");
std::fs::write(&bin, "#!/bin/sh\necho SELFHEAL \"$@\"\n").unwrap();
{
use std::os::unix::fs::PermissionsExt;
let mut perms = std::fs::metadata(&bin).unwrap().permissions();
perms.set_mode(0o755);
std::fs::set_permissions(&bin, perms).unwrap();
}
// Hint is a nonexistent path; `myfind` is only reachable via PATH.
let shadow = shell_function("find", "myfind", Some(Path::new("/nonexistent/bfs")), &[]);
let script = format!(
"set -euo pipefail; export PATH={}:\"$PATH\"; {shadow}; find X",
bash_safe_quote(&dir.to_string_lossy())
);
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
let _ = std::fs::remove_dir_all(&dir);
assert!(
out.status.success(),
"stderr: {}",
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "SELFHEAL X");
}
/// #1 fallback: neither the hint nor `command -v` resolves → the OS binary
/// runs (via `command find`), so the shadow never breaks `find`.
#[test]
fn shadow_falls_back_to_os_when_binary_absent() {
let Ok(bash) = which::which("bash") else {
return;
};
let shadow = shell_function(
"find",
"grok_no_such_search_bin_xyz",
Some(Path::new("/nonexistent/bfs")),
&[],
);
// `find /dev/null` prints the path on every find implementation.
let script = format!("set -euo pipefail; {shadow}; find /dev/null");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
assert!(
out.status.success(),
"stderr: {}",
String::from_utf8_lossy(&out.stderr)
);
assert_eq!(String::from_utf8_lossy(&out.stdout).trim(), "/dev/null");
}
/// #1 regression: a host hint that exists but is **not executable** (e.g. a
/// mode-0644 `KIGI_TOOLS_*_PATH` / vendor copy) must fall through to the OS
/// binary rather than hard-fail `exec` with EACCES. The `[ -x ]` guard (not
/// `[ -f ]`) is what makes this work.
#[test]
fn shadow_falls_back_when_hint_not_executable() {
let Ok(bash) = which::which("bash") else {
return;
};
let dir = std::env::temp_dir().join(format!(
"grok-noexec-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).unwrap();
let hint = dir.join("bfs");
std::fs::write(&hint, "#!/bin/sh\necho SHOULD_NOT_RUN\n").unwrap();
{
// Mode 0644 — exists but not executable by anyone.
use std::os::unix::fs::PermissionsExt;
let mut perms = std::fs::metadata(&hint).unwrap().permissions();
perms.set_mode(0o644);
std::fs::set_permissions(&hint, perms).unwrap();
}
// Hint is the non-exec file; bin_name isn't on PATH → must reach OS find.
let shadow = shell_function(
"find",
"grok_no_such_search_bin_xyz",
Some(hint.as_path()),
&[],
);
let script = format!("set -euo pipefail; {shadow}; find /dev/null");
let out = std::process::Command::new(&bash)
.args(["-c", &script])
.output()
.unwrap();
let _ = std::fs::remove_dir_all(&dir);
assert!(
out.status.success(),
"non-exec hint should fall back to OS find, not fail: {:?}\nstderr: {}",
out.status,
String::from_utf8_lossy(&out.stderr)
);
let stdout = String::from_utf8_lossy(&out.stdout);
assert_eq!(stdout.trim(), "/dev/null");
assert!(
!stdout.contains("SHOULD_NOT_RUN"),
"non-exec hint was run: {stdout:?}"
);
}
}
crates/codegen/kigi-tools/src/computer/local/file_system.rs
+344
View File
@@ -0,0 +1,344 @@
use std::{future::Future, io, path::Path, time::Duration};
use tokio::{fs, time::sleep};
use crate::computer::types::{AsyncFileSystem, ComputerError};
/// Creates a local FS access which allows writing and reading from the local files
pub struct LocalFs;
// Keep the window short: these retries absorb brief Windows editor/indexer/AV
// races without hiding persistent locks, ACL failures, or sandbox denials.
const WRITE_RETRY_DELAYS: &[Duration] = &[
Duration::from_millis(25),
Duration::from_millis(50),
Duration::from_millis(100),
Duration::from_millis(200),
Duration::from_millis(400),
];
#[cfg(any(windows, test))]
const WINDOWS_ERROR_SHARING_VIOLATION: i32 = 32;
#[cfg(any(windows, test))]
const WINDOWS_ERROR_LOCK_VIOLATION: i32 = 33;
/// Check if an IO error is a permission denial (EACCES or EPERM),
/// which indicates a sandbox violation.
fn is_permission_error(e: &io::Error) -> bool {
matches!(e.kind(), io::ErrorKind::PermissionDenied)
}
#[cfg(any(windows, test))]
fn is_windows_transient_write_lock_raw_os_error(raw_os_error: Option<i32>) -> bool {
matches!(
raw_os_error,
Some(WINDOWS_ERROR_SHARING_VIOLATION | WINDOWS_ERROR_LOCK_VIOLATION)
)
}
fn is_transient_write_lock_error(e: &io::Error) -> bool {
#[cfg(windows)]
{
is_windows_transient_write_lock_raw_os_error(e.raw_os_error())
}
#[cfg(not(windows))]
{
let _ = e;
false
}
}
#[cfg(test)]
fn is_test_transient_write_lock_error(e: &io::Error) -> bool {
is_windows_transient_write_lock_raw_os_error(e.raw_os_error())
}
async fn write_file_with_transient_lock_retries(path: &Path, data: &[u8]) -> io::Result<()> {
write_file_with_retry_hooks(
|| fs::write(path, data),
|delay| sleep(delay),
|retry_count| {
tracing::debug!(
path = %path.display(),
retry_count,
"file write succeeded after transient lock retries"
);
},
|error, retry_count, delay| {
tracing::debug!(
path = %path.display(),
error = %error,
retry_count,
delay_ms = delay.as_millis(),
"file write hit transient lock; retrying"
);
},
|error, retry_count| {
tracing::debug!(
path = %path.display(),
error = %error,
retry_count,
"file write exhausted transient lock retries"
);
},
is_transient_write_lock_error,
)
.await
}
async fn write_file_with_retry_hooks<W, WFut, S, SFut, Success, Retry, Exhausted, IsRetryable>(
mut write: W,
mut sleep_for: S,
mut on_retry_success: Success,
mut on_retry: Retry,
mut on_exhausted: Exhausted,
is_retryable: IsRetryable,
) -> io::Result<()>
where
W: FnMut() -> WFut,
WFut: Future<Output = io::Result<()>>,
S: FnMut(Duration) -> SFut,
SFut: Future<Output = ()>,
Success: FnMut(usize),
Retry: FnMut(&io::Error, usize, Duration),
Exhausted: FnMut(&io::Error, usize),
IsRetryable: Fn(&io::Error) -> bool,
{
let mut retry_count = 0usize;
loop {
match write().await {
Ok(()) => {
if retry_count > 0 {
on_retry_success(retry_count);
}
return Ok(());
}
Err(e) if is_retryable(&e) => {
if retry_count >= WRITE_RETRY_DELAYS.len() {
on_exhausted(&e, retry_count);
return Err(e);
}
let delay = WRITE_RETRY_DELAYS[retry_count];
retry_count += 1;
on_retry(&e, retry_count, delay);
sleep_for(delay).await;
}
Err(e) => return Err(e),
}
}
}
#[async_trait::async_trait]
impl AsyncFileSystem for LocalFs {
#[tracing::instrument(name = "fs.read_file", skip_all)]
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, ComputerError> {
match fs::read(path).await {
Ok(data) => Ok(data),
Err(e) => {
if is_permission_error(&e) {
kigi_sandbox::log_violation(&path.display().to_string(), "read");
}
Err(e.into())
}
}
}
#[tracing::instrument(name = "fs.write_file", skip_all)]
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), ComputerError> {
// implicitly creates the missing directories if any
if let Some(dir) = path.parent()
&& let Err(e) = fs::create_dir_all(dir).await
{
if is_permission_error(&e) {
kigi_sandbox::log_violation(&dir.display().to_string(), "mkdir");
}
return Err(e.into());
}
if let Err(e) = write_file_with_transient_lock_retries(path, data).await {
if is_permission_error(&e) {
kigi_sandbox::log_violation(&path.display().to_string(), "write");
}
return Err(e.into());
}
Ok(())
}
#[tracing::instrument(name = "fs.delete_file", skip_all)]
async fn delete_file(&self, path: &Path) -> Result<(), ComputerError> {
if let Err(e) = fs::remove_file(path).await {
if is_permission_error(&e) {
kigi_sandbox::log_violation(&path.display().to_string(), "delete");
}
return Err(e.into());
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
use std::rc::Rc;
#[test]
fn classifies_windows_transient_write_lock_errors() {
assert!(is_windows_transient_write_lock_raw_os_error(Some(32)));
assert!(is_windows_transient_write_lock_raw_os_error(Some(33)));
assert!(!is_windows_transient_write_lock_raw_os_error(Some(5)));
assert!(!is_windows_transient_write_lock_raw_os_error(None));
}
#[cfg(windows)]
#[test]
fn classifies_windows_transient_write_lock_io_errors() {
assert!(is_transient_write_lock_error(
&io::Error::from_raw_os_error(WINDOWS_ERROR_SHARING_VIOLATION,)
));
assert!(is_transient_write_lock_error(
&io::Error::from_raw_os_error(WINDOWS_ERROR_LOCK_VIOLATION,)
));
assert!(!is_transient_write_lock_error(
&io::Error::from_raw_os_error(5)
));
}
#[cfg(not(windows))]
#[test]
fn non_windows_does_not_retry_windows_raw_error_numbers() {
assert!(!is_transient_write_lock_error(
&io::Error::from_raw_os_error(WINDOWS_ERROR_SHARING_VIOLATION,)
));
}
#[tokio::test]
async fn first_attempt_success_does_not_fire_retry_callbacks() {
let result = write_file_with_retry_hooks(
|| async { Ok(()) },
|_| async { panic!("should not sleep") },
|_| panic!("should not fire on_retry_success"),
|_, _, _| panic!("should not fire on_retry"),
|_, _| panic!("should not fire on_exhausted"),
is_test_transient_write_lock_error,
)
.await;
result.unwrap();
}
#[tokio::test]
async fn transient_lock_errors_are_retried_until_success() {
let attempts = Rc::new(Cell::new(0usize));
let sleeps = Rc::new(Cell::new(0usize));
let retry_success_count = Rc::new(Cell::new(0usize));
let retry_log_count = Rc::new(Cell::new(0usize));
let result = write_file_with_retry_hooks(
{
let attempts = Rc::clone(&attempts);
move || {
let attempts = Rc::clone(&attempts);
async move {
let next = attempts.get() + 1;
attempts.set(next);
if next <= 2 {
Err(io::Error::from_raw_os_error(
WINDOWS_ERROR_SHARING_VIOLATION,
))
} else {
Ok(())
}
}
}
},
{
let sleeps = Rc::clone(&sleeps);
move |_| {
let sleeps = Rc::clone(&sleeps);
async move {
sleeps.set(sleeps.get() + 1);
}
}
},
|count| retry_success_count.set(count),
|_, _, _| retry_log_count.set(retry_log_count.get() + 1),
|_, _| panic!("retry budget should not be exhausted"),
is_test_transient_write_lock_error,
)
.await;
result.unwrap();
assert_eq!(attempts.get(), 3);
assert_eq!(sleeps.get(), 2);
assert_eq!(retry_log_count.get(), 2);
assert_eq!(retry_success_count.get(), 2);
}
#[tokio::test]
async fn non_transient_errors_are_not_retried() {
let attempts = Rc::new(Cell::new(0usize));
let result = write_file_with_retry_hooks(
{
let attempts = Rc::clone(&attempts);
move || {
let attempts = Rc::clone(&attempts);
async move {
attempts.set(attempts.get() + 1);
Err(io::Error::new(io::ErrorKind::NotFound, "missing"))
}
}
},
|_| async {},
|_| panic!("write did not succeed"),
|_, _, _| panic!("non-transient errors must not be retried"),
|_, _| panic!("non-transient errors must not exhaust retry budget"),
is_test_transient_write_lock_error,
)
.await;
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::NotFound);
assert_eq!(attempts.get(), 1);
}
#[tokio::test]
async fn persistent_transient_lock_exhausts_retry_budget() {
let attempts = Rc::new(Cell::new(0usize));
let sleeps = Rc::new(Cell::new(0usize));
let exhausted_count = Rc::new(Cell::new(0usize));
let result = write_file_with_retry_hooks(
{
let attempts = Rc::clone(&attempts);
move || {
let attempts = Rc::clone(&attempts);
async move {
attempts.set(attempts.get() + 1);
Err(io::Error::from_raw_os_error(WINDOWS_ERROR_LOCK_VIOLATION))
}
}
},
{
let sleeps = Rc::clone(&sleeps);
move |_| {
let sleeps = Rc::clone(&sleeps);
async move {
sleeps.set(sleeps.get() + 1);
}
}
},
|_| panic!("write did not succeed"),
|_, _, _| {},
|_, count| exhausted_count.set(count),
is_test_transient_write_lock_error,
)
.await;
assert_eq!(
result.unwrap_err().raw_os_error(),
Some(WINDOWS_ERROR_LOCK_VIOLATION)
);
assert_eq!(attempts.get(), WRITE_RETRY_DELAYS.len() + 1);
assert_eq!(sleeps.get(), WRITE_RETRY_DELAYS.len());
assert_eq!(exhausted_count.get(), WRITE_RETRY_DELAYS.len());
}
}
crates/codegen/kigi-tools/src/computer/local/mock_fs.rs
+122
View File
@@ -0,0 +1,122 @@
//! Mock file system implementation for testing.
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::computer::types::{AsyncFileSystem, ComputerError};
/// In-memory file system for testing.
/// Thread-safe and async-compatible.
pub struct MockFs {
files: Arc<RwLock<HashMap<PathBuf, Vec<u8>>>>,
}
impl Default for MockFs {
fn default() -> Self {
Self::new()
}
}
impl MockFs {
/// Create a new empty mock file system.
pub fn new() -> Self {
Self {
files: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Set a file's contents directly (for test setup).
pub async fn set_file(&self, path: impl AsRef<Path>, content: &[u8]) {
self.files
.write()
.await
.insert(path.as_ref().to_path_buf(), content.to_vec());
}
/// Get a file's contents directly (for test assertions).
pub async fn get_file(&self, path: impl AsRef<Path>) -> Option<Vec<u8>> {
self.files.read().await.get(path.as_ref()).cloned()
}
/// Check if a file exists.
pub async fn exists(&self, path: impl AsRef<Path>) -> bool {
self.files.read().await.contains_key(path.as_ref())
}
/// List all files in the mock filesystem.
pub async fn list_files(&self) -> Vec<PathBuf> {
self.files.read().await.keys().cloned().collect()
}
}
#[async_trait::async_trait]
impl AsyncFileSystem for MockFs {
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, ComputerError> {
self.files.read().await.get(path).cloned().ok_or_else(|| {
ComputerError::IOError(
format!("File not found: {}", path.display()),
Some(std::io::ErrorKind::NotFound),
)
})
}
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), ComputerError> {
self.files
.write()
.await
.insert(path.to_path_buf(), data.to_vec());
Ok(())
}
async fn delete_file(&self, path: &Path) -> Result<(), ComputerError> {
self.files.write().await.remove(path);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_mock_fs_read_write() {
let fs = MockFs::new();
// File doesn't exist initially
assert!(fs.read_file(Path::new("/test.txt")).await.is_err());
// Write a file
fs.write_file(Path::new("/test.txt"), b"hello world")
.await
.unwrap();
// Read it back
let content = fs.read_file(Path::new("/test.txt")).await.unwrap();
assert_eq!(content, b"hello world");
}
#[tokio::test]
async fn test_mock_fs_delete() {
let fs = MockFs::new();
fs.write_file(Path::new("/test.txt"), b"hello")
.await
.unwrap();
assert!(fs.exists(Path::new("/test.txt")).await);
fs.delete_file(Path::new("/test.txt")).await.unwrap();
assert!(!fs.exists(Path::new("/test.txt")).await);
}
#[tokio::test]
async fn test_mock_fs_set_file() {
let fs = MockFs::new();
fs.set_file("/preset.txt", b"preset content").await;
let content = fs.read_file(Path::new("/preset.txt")).await.unwrap();
assert_eq!(content, b"preset content");
}
}
crates/codegen/kigi-tools/src/computer/local/mod.rs
+37
View File
@@ -0,0 +1,37 @@
pub mod cgroup;
#[cfg(unix)]
pub mod embedded_search_tools;
pub mod file_system;
pub mod mock_fs;
#[cfg(unix)]
pub mod shell_state;
pub mod terminal;
pub use cgroup::{CgroupMemoryConfig, PROCESS_OOM_EXIT_CODE};
pub use file_system::LocalFs;
pub use mock_fs::MockFs;
pub use terminal::{ExitStatus, LocalTerminalBackend};
/// Per-backend enable state for the bash-harness `find`→`bfs` / `grep`→`ugrep`
/// shadows.
///
/// Resolved once by the host (config.toml `[toolset.bash]` / env / requirements)
/// and baked into a [`LocalTerminalBackend`] at creation. Keeping it on the
/// backend instead of a process-global means a subagent that reuses the parent's
/// `LocalTerminalBackend` inherits the parent's shadows — it can't overwrite the
/// enable state for bash that later runs on the shared backend. Defaults to
/// both-on for standalone backends with no host wiring.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SearchShadowConfig {
pub find_bfs: bool,
pub grep_ugrep: bool,
}
impl Default for SearchShadowConfig {
fn default() -> Self {
Self {
find_bfs: true,
grep_ugrep: true,
}
}
}
crates/codegen/kigi-tools/src/computer/local/shell_state.rs
+1179
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-tools/src/computer/local/terminal.rs
+4774
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-tools/src/computer/mod.rs
+3
View File
@@ -0,0 +1,3 @@
pub mod local;
/// Contains the computer implementation
pub mod types;
crates/codegen/kigi-tools/src/computer/types.rs
+422
View File
@@ -0,0 +1,422 @@
use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};
use crate::notification::types::ToolNotificationHandle;
// ============================================================================
// Error types
// ============================================================================
#[derive(thiserror::Error, Debug, Clone)]
pub enum ComputerError {
#[error("IO Error: {0}")]
IOError(String, Option<std::io::ErrorKind>),
#[error("UnQuoted command")]
CommandNotQuoted,
}
impl ComputerError {
/// Create an IO error from a message string with no preserved error kind.
pub fn io(msg: impl Into<String>) -> Self {
Self::IOError(msg.into(), None)
}
pub fn io_with_kind(msg: impl Into<String>, kind: std::io::ErrorKind) -> Self {
Self::IOError(msg.into(), Some(kind))
}
/// Returns the underlying `io::ErrorKind` or `None`
pub fn io_error_kind(&self) -> Option<std::io::ErrorKind> {
match self {
Self::IOError(_, kind) => *kind,
_ => None,
}
}
}
impl From<std::io::Error> for ComputerError {
fn from(err: std::io::Error) -> Self {
Self::IOError(err.to_string(), Some(err.kind()))
}
}
// ============================================================================
// File system trait
// ============================================================================
#[async_trait::async_trait]
pub trait AsyncFileSystem: Send + Sync {
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, ComputerError>;
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), ComputerError>;
async fn delete_file(&self, path: &Path) -> Result<(), ComputerError>;
}
// ============================================================================
// Terminal types
// ============================================================================
pub struct TerminalRunRequest {
pub command: String,
pub working_directory: PathBuf,
pub env: HashMap<String, String>,
pub timeout: Duration,
pub output_byte_limit: usize,
/// File path to write output incrementally as it arrives.
/// This ensures full output is always available even after in-memory buffer is truncated.
/// For background tasks, this allows retrieval of output after the agent has moved on.
pub output_file: PathBuf,
/// Notification handle for streaming output chunks during execution.
/// The backend sends `BashOutputChunk` notifications every ~100ms.
/// Callers that don't need streaming pass `ToolNotificationHandle::noop()`
/// — messages are silently dropped. No `Option` wrapper needed.
pub notification_handle: ToolNotificationHandle,
/// Tool call ID for correlating notifications with the tool invocation.
/// Flows from `ToolContext::tool_call_id()` through the actor to
/// `BashOutputChunk.base.tool_call_id`.
pub tool_call_id: String,
/// Original user command before isolation wrapping.
///
/// When set, the terminal actor stores this on the `ProcessEntry` so
/// `get_task()` returns it in `TaskSnapshot.display_command`. This
/// ensures model-facing `get_task_output` shows the user's command
/// instead of the `unshare`/mount wrapper.
pub display_command: Option<String>,
/// Auto-background on timeout instead of killing (default `false`).
pub auto_background_on_timeout: bool,
/// When [`Self::auto_background_on_timeout`] is true, maximum time the
/// command may block the turn before being moved to the background (process
/// keeps running). Independent of [`Self::timeout`].
///
/// - `None` → use the terminal backend default (typically 15s).
/// - `Some(Duration::MAX)` → no short budget; auto-bg only when `timeout` elapses.
/// - `Some(d)` → auto-bg after `d` if still running.
pub foreground_block_budget: Option<Duration>,
/// Task kind for distinguishing monitor tasks from regular bash tasks.
pub kind: TaskKind,
/// Session that owns this process. Used to scope kill operations so
/// `kill_all_background_tasks_by_owner` only targets the requesting
/// session's processes — not the parent's or sibling's.
pub owner_session_id: Option<String>,
}
/// Distinguishes different types of background tasks.
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
Default,
serde::Serialize,
serde::Deserialize,
schemars::JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum TaskKind {
/// Regular bash command.
#[default]
Bash,
/// Monitor tool — streams stdout events with rate limiting.
Monitor,
}
#[derive(Clone)]
pub struct TerminalRunResult {
pub combined_output: String,
pub exit_code: Option<i32>,
pub truncated: bool,
pub signal: Option<String>,
pub timed_out: bool,
/// Path to the output file where full output is stored.
/// Use read_file tool to retrieve full output when truncated.
pub output_file: PathBuf,
/// Total bytes of output (before truncation).
/// When truncated, combined_output contains the first and last portions up to output_byte_limit chars.
pub total_bytes: usize,
/// PID of the spawned shell process, when available. Set by the
/// local terminal backend at spawn time. Useful for foreground
/// commands that auto-background on timeout: the resulting
/// `BackgroundTaskStarted` can carry the real PID instead of a
/// placeholder. `None` for backends that cannot surface a local
/// PID (e.g. ACP/remote terminals) or when the process exited
/// before `child.id()` could be queried.
pub pid: Option<u32>,
}
/// Returned by `TerminalBackend::run_background` — gives the caller the task_id
/// to use for subsequent queries via `get_task`, `kill_task`, `wait_for_completion`.
pub struct BackgroundHandle {
pub task_id: String,
pub output_file: PathBuf,
/// PID of the spawned shell process, when available. `None` for
/// backends that do not surface a local PID (e.g. ACP/remote
/// gateways) or when the process exited before the PID could be
/// captured.
pub pid: Option<u32>,
}
/// Full snapshot of a task's state.
/// Used by both local and ACP backends.
#[derive(
Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize, schemars::JsonSchema,
)]
pub struct TaskSnapshot {
pub task_id: String,
/// The actual command that was executed (may be isolation-wrapped).
pub command: String,
/// The original user command before isolation wrapping.
///
/// When set, model/user-facing output should prefer this over `command`
/// to avoid exposing internal isolation mechanics (unshare/mount wrapper).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_command: Option<String>,
pub cwd: String,
pub start_time: std::time::SystemTime,
pub end_time: Option<std::time::SystemTime>,
pub output: String,
pub output_file: PathBuf,
pub truncated: bool,
pub exit_code: Option<i32>,
pub signal: Option<String>,
pub completed: bool,
/// Task kind: bash (default) or monitor.
#[serde(default)]
pub kind: TaskKind,
/// Whether a block-waiter (`block=true`) consumed this task's result.
/// When set, the notification bridge skips auto-wake synthetic prompts
/// because the blocking caller already received the result directly.
#[serde(default)]
pub block_waited: bool,
/// Whether this task was explicitly killed via the `kill_command_or_subagent` tool.
/// When set, auto-wake synthetic prompts are suppressed because the model
/// already received the kill result via `KillTaskResult`.
/// Also set during `kill_all_background_tasks` teardown (e.g. subagent
/// cleanup), where auto-wake suppression is irrelevant since the session
/// is shutting down.
#[serde(default)]
pub explicitly_killed: bool,
/// Session that owns this task. Used for scoped kill operations so
/// subagent teardown only kills the subagent's own tasks, not
/// the parent's or sibling's.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub owner_session_id: Option<String>,
}
impl TaskSnapshot {
/// Calculate duration in seconds.
/// If task is still running, returns time since start.
pub fn duration_secs(&self) -> f64 {
let end = self.end_time.unwrap_or_else(std::time::SystemTime::now);
end.duration_since(self.start_time)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
/// True iff the task has NOT yet completed — covers bash AND
/// monitor task kinds (the `kind` field doesn't change this
/// predicate; the runtime turn-end TodoGate counts both as
/// backing work).
pub fn is_outstanding(&self) -> bool {
!self.completed
}
}
/// Result of killing a terminal task.
///
/// Serialized over the wire in the `x.ai/task/kill` ext response
/// (`kigi-shell::extensions::task::KillTaskResponse`) and deserialized
/// by clients (kigi-tui), so it derives both serde directions.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum KillOutcome {
Killed,
AlreadyExited,
NotFound,
}
// ============================================================================
// TerminalBackend trait
// ============================================================================
/// The single abstraction over terminal execution backends.
///
/// Implemented by:
/// - `LocalTerminalBackend` (in kigi-tools, spawns processes)
/// - `AcpTerminalBackend` (in kigi-shell, calls ACP protocol)
#[async_trait::async_trait]
pub trait TerminalBackend: Send + Sync {
/// Run a command. Blocks until completion or timeout.
async fn run(&self, request: TerminalRunRequest) -> Result<TerminalRunResult, ComputerError>;
/// Start a command in the background. Returns immediately with a handle ID.
/// The process continues running; use get_task/kill_task/wait_for_completion to manage it.
async fn run_background(
&self,
request: TerminalRunRequest,
) -> Result<BackgroundHandle, ComputerError>;
/// Get current snapshot of a background task.
async fn get_task(&self, task_id: &str) -> Option<TaskSnapshot>;
/// Kill a background task.
async fn kill_task(&self, task_id: &str) -> KillOutcome;
/// Kill all running foreground processes.
async fn kill_foreground_commands(&self) {}
/// Kill all running foreground processes owned by a specific session.
/// Used on a shared terminal backend so a subagent's cancel doesn't
/// kill the parent's foreground commands.
async fn kill_foreground_commands_by_owner(&self, _owner_session_id: &str) {}
/// Kill all running background tasks.
/// Used during subagent teardown to clean up orphaned processes.
async fn kill_all_background_tasks(&self) {}
/// Kill all running background tasks owned by a specific session.
/// Used during subagent teardown on a shared terminal backend so
/// only the subagent's own tasks are killed — not the parent's.
async fn kill_all_background_tasks_by_owner(&self, _owner_session_id: &str) {}
/// Fire-and-forget prewarm of the persistent login shell; default no-op for
/// backends without one (ACP/remote, non-persistent).
async fn warm_persistent_shell(&self, _cwd: &std::path::Path) {}
/// Reparent notification handles for all tasks owned by `old_owner_session_id`.
/// Swaps the dead child session's notification handle with the parent's
/// live handle so events from surviving processes route correctly.
/// Also re-spawns monitor pipelines on the caller's runtime so monitor
/// events continue streaming to the parent.
///
/// `backend_weak` is a [`Weak`](std::sync::Weak) to *this* backend (anchored
/// by the parent session's `Arc`); it drives re-spawned monitor pipelines
/// without keeping the backend alive. See `run_monitor_pipeline`.
async fn reparent_notifications(
&self,
_old_owner_session_id: &str,
_new_owner_session_id: &str,
_new_handle: crate::notification::types::ToolNotificationHandle,
_backend_weak: std::sync::Weak<dyn TerminalBackend>,
) {
}
/// Move a foreground command to background by tool_call_id.
/// The process keeps running but the foreground waiter is unblocked.
/// Returns `true` if a matching foreground process was found.
async fn background_foreground_command(&self, _tool_call_id: &str) -> bool {
false
}
/// Wait for a background task to complete, with optional timeout.
async fn wait_for_completion(
&self,
task_id: &str,
timeout: Option<Duration>,
) -> Option<TaskSnapshot>;
/// List all known background tasks (running and completed).
/// Used for context compaction to include task state in summaries.
async fn list_tasks(&self) -> Vec<TaskSnapshot>;
/// Return the persistent shell's current working directory, if persistent
/// shell state is enabled. Returns `None` when persistence is off or the
/// backend doesn't support it (e.g. ACP/remote).
async fn get_shell_cwd(&self) -> Option<std::path::PathBuf> {
None
}
}
// ============================================================================
// Computer struct
// ============================================================================
/// Contains the computer struct which provides access to both the terminal and the fs
pub struct Computer {
pub terminal: Arc<dyn TerminalBackend>,
pub file_system: Arc<dyn AsyncFileSystem>,
}
impl Computer {
pub fn new(terminal: Arc<dyn TerminalBackend>, file_system: Arc<dyn AsyncFileSystem>) -> Self {
Self {
terminal,
file_system,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn io_error_kind_preserved_through_from() {
let io_err = std::io::Error::new(std::io::ErrorKind::NotFound, "gone");
let ce = ComputerError::from(io_err);
assert_eq!(ce.io_error_kind(), Some(std::io::ErrorKind::NotFound));
}
#[test]
fn io_error_kind_is_a_directory() {
let io_err = std::io::Error::new(std::io::ErrorKind::IsADirectory, "it's a dir");
let ce = ComputerError::from(io_err);
assert_eq!(ce.io_error_kind(), Some(std::io::ErrorKind::IsADirectory));
}
#[test]
fn io_error_kind_none_for_string_constructor() {
let ce = ComputerError::io("something broke");
assert_eq!(ce.io_error_kind(), None);
}
#[test]
fn io_error_kind_none_for_non_io_variant() {
let ce = ComputerError::CommandNotQuoted;
assert_eq!(ce.io_error_kind(), None);
}
#[test]
fn io_with_kind_preserves_not_found() {
let ce = ComputerError::io_with_kind("Resource not found", std::io::ErrorKind::NotFound);
assert_eq!(ce.io_error_kind(), Some(std::io::ErrorKind::NotFound));
}
#[test]
fn io_with_kind_preserves_permission_denied() {
let ce = ComputerError::io_with_kind("access denied", std::io::ErrorKind::PermissionDenied);
assert_eq!(
ce.io_error_kind(),
Some(std::io::ErrorKind::PermissionDenied)
);
}
#[test]
fn io_with_kind_matches_local_fs_dispatch_for_not_found() {
// Simulate what LocalFs produces for a missing file
let local_err = ComputerError::from(std::io::Error::new(
std::io::ErrorKind::NotFound,
"No such file or directory (os error 2)",
));
// Simulate what AcpFsAdapter now produces for RESOURCE_NOT_FOUND
let acp_err =
ComputerError::io_with_kind("Resource not found", std::io::ErrorKind::NotFound);
// Both must produce the same io_error_kind so read_file dispatches
// to FileNotFound("Error: {path} does not exist.") in both cases.
assert_eq!(local_err.io_error_kind(), acp_err.io_error_kind());
}
}