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

M0: compilable skeleton — Kigi 0.1.0 fork surgery

Browse Source 
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.

Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
  kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
  ptyctl, ptyctl-cli, third_party/ unchanged; proto package
  xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
  KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
  (templates re-encrypted)

Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
  trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
  module & dc_log, heap-profile uploader, auth-diagnostics uploader,
  session-analytics halves of feedback; local zero-egress observability
  preserved in new kigi-log crate (unified log, --debug firehose,
  subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
  direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
  relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
  ~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
  kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
  session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
  shell util

Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
  https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
  https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
  Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted

Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
  workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
  all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
  exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
  insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean

Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
  (new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
  fast-worktree); RSS measurement tests serialized via serial_test

Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
  notices sustained; kigi-tools ported-code notices extended; README,
  CONTRIBUTING, SECURITY, AGENTS.md rewritten

Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
ZacharyZhang-NY
2026-07-17 05:31:01 -04:00
commit d6c20fc13f
2612 changed files with 1353757 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-workspace/src/session/checkpoint.rs
+1021
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/session/checkpoint_store.rs
+725
View File
@@ -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-"));
}
}
crates/codegen/kigi-workspace/src/session/file_state.rs
+1813
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/session/git.rs
+4158
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-workspace/src/session/jj.rs
+226
View File
@@ -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,
})
}
crates/codegen/kigi-workspace/src/session/mod.rs
+971
View File
@@ -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");
}
}
crates/codegen/kigi-workspace/src/session/swap_policy.rs
+779
View File
@@ -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
);
}
}
crates/codegen/kigi-workspace/src/session/tool_config.rs
+1207
View File
File diff suppressed because it is too large Load Diff