M0: compilable skeleton — Kigi 0.1.0 fork surgery
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.
Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
ptyctl, ptyctl-cli, third_party/ unchanged; proto package
xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
(templates re-encrypted)
Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
module & dc_log, heap-profile uploader, auth-diagnostics uploader,
session-analytics halves of feedback; local zero-egress observability
preserved in new kigi-log crate (unified log, --debug firehose,
subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
shell util
Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted
Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean
Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
(new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
fast-worktree); RSS measurement tests serialized via serial_test
Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
notices sustained; kigi-tools ported-code notices extended; README,
CONTRIBUTING, SECURITY, AGENTS.md rewritten
Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,602 @@
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::{self, Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use fs2::FileExt;
|
||||
use kigi_workspace::util::is_lock_contended;
|
||||
|
||||
use crate::util::kigi_home::kigi_home;
|
||||
|
||||
/// Env var that overrides the leader socket path (and, by extension, the lock
|
||||
/// path — the sibling `.lock`). Set by the `--leader-socket` flag, or exported
|
||||
/// directly. Lets a developer sandbox a leader instance away from the default
|
||||
/// `~/.kigi/leader.sock` — e.g. run a local branch build's leader without
|
||||
/// colliding with an installed stable leader on the same machine. Honored by
|
||||
/// BOTH the client (`connect_or_spawn`) and the leader (`run_leader`), and
|
||||
/// inherited by the spawned leader subprocess, so all parties bind the same
|
||||
/// path.
|
||||
pub const LEADER_SOCKET_ENV: &str = "KIGI_LEADER_SOCKET";
|
||||
|
||||
/// The explicit socket-path override, if [`LEADER_SOCKET_ENV`] is set and
|
||||
/// non-empty.
|
||||
fn leader_socket_override() -> Option<PathBuf> {
|
||||
std::env::var_os(LEADER_SOCKET_ENV)
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(PathBuf::from)
|
||||
}
|
||||
|
||||
/// The lock path paired with a given socket path: the sibling file with a
|
||||
/// `.lock` extension (`/x/leader-foo.sock` → `/x/leader-foo.lock`). Matches the
|
||||
/// default `leader.sock`/`leader.lock` pairing so the two never disagree.
|
||||
fn lock_path_for_socket(socket: &Path) -> PathBuf {
|
||||
socket.with_extension("lock")
|
||||
}
|
||||
|
||||
/// Resolve the socket path: the explicit override wins, else the default
|
||||
/// `leader.sock` under `root`. Pure (the override is passed in) so it is
|
||||
/// unit-testable without touching process env.
|
||||
fn resolve_socket_path(override_socket: Option<PathBuf>, root: &Path) -> PathBuf {
|
||||
override_socket.unwrap_or_else(|| default_socket_path_in(root))
|
||||
}
|
||||
|
||||
/// Resolve the lock path: the sibling `.lock` of the override socket if set,
|
||||
/// else the default `leader.lock` under `root`. Pure (see
|
||||
/// [`resolve_socket_path`]).
|
||||
fn resolve_lock_path(override_socket: Option<PathBuf>, root: &Path) -> PathBuf {
|
||||
match override_socket {
|
||||
Some(socket) => lock_path_for_socket(&socket),
|
||||
None => default_lock_path_in(root),
|
||||
}
|
||||
}
|
||||
|
||||
/// Default leader lock path under `root` (`leader.lock`).
|
||||
pub fn default_lock_path_in(root: &Path) -> PathBuf {
|
||||
root.join("leader.lock")
|
||||
}
|
||||
|
||||
/// Effective leader lock path: the [`LEADER_SOCKET_ENV`] override's sibling
|
||||
/// `.lock` when set, else the default under grok home.
|
||||
pub fn default_lock_path() -> PathBuf {
|
||||
resolve_lock_path(leader_socket_override(), &kigi_home())
|
||||
}
|
||||
|
||||
/// Default leader socket path under `root` (`leader.sock`).
|
||||
pub fn default_socket_path_in(root: &Path) -> PathBuf {
|
||||
root.join("leader.sock")
|
||||
}
|
||||
|
||||
/// Effective leader socket path: the [`LEADER_SOCKET_ENV`] override when set,
|
||||
/// else the default under grok home.
|
||||
pub fn default_socket_path() -> PathBuf {
|
||||
resolve_socket_path(leader_socket_override(), &kigi_home())
|
||||
}
|
||||
|
||||
/// The instance suffix encoded in a lock/socket file-name pair
|
||||
/// (`leader<suffix>.lock` / `leader<suffix>.sock`). Empty suffix = the default
|
||||
/// instance; non-default instances only arise via [`LEADER_SOCKET_ENV`].
|
||||
/// `None` when the two file names disagree.
|
||||
pub fn socket_suffix_from_paths(lock_path: &Path, socket_path: &Path) -> Option<String> {
|
||||
let lock_name = lock_path.file_name()?.to_str()?;
|
||||
let socket_name = socket_path.file_name()?.to_str()?;
|
||||
let lock_suffix = lock_name
|
||||
.strip_prefix("leader")?
|
||||
.strip_suffix(".lock")?
|
||||
.to_string();
|
||||
let socket_suffix = socket_name
|
||||
.strip_prefix("leader")?
|
||||
.strip_suffix(".sock")?
|
||||
.to_string();
|
||||
|
||||
if lock_suffix == socket_suffix {
|
||||
Some(lock_suffix)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum LockError {
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("Lock held by another process")]
|
||||
AlreadyLocked,
|
||||
#[error("Timed out waiting to acquire lock after {0:?}")]
|
||||
Timeout(Duration),
|
||||
}
|
||||
|
||||
/// Lock manager for the leader process using OS-level file locking (flock).
|
||||
///
|
||||
/// The lock file serves two purposes:
|
||||
/// 1. Exclusive lock indicates who is the leader (or who is spawning)
|
||||
/// 2. File contents store the leader's PID for diagnostics
|
||||
///
|
||||
/// Lock semantics:
|
||||
/// - Leader holds exclusive lock for its entire lifetime
|
||||
/// - Clients use try_lock to check if leader exists and coordinate spawning
|
||||
///
|
||||
/// Cleanup behavior:
|
||||
/// - If lock is held when dropped (crash/exit), files are cleaned up
|
||||
/// - If `release()` is called before drop, files are NOT cleaned up (handoff to leader)
|
||||
#[derive(Debug)]
|
||||
pub struct LeaderLock {
|
||||
lock_path: PathBuf,
|
||||
sock_path: PathBuf,
|
||||
lock_file: Option<File>,
|
||||
/// Tracks if we should clean up files on drop.
|
||||
/// Set to true when lock is acquired, set to false when explicitly released.
|
||||
/// This ensures cleanup happens if we crash while holding the lock,
|
||||
/// but NOT if we explicitly hand off to another process via release().
|
||||
was_leader: bool,
|
||||
}
|
||||
|
||||
impl LeaderLock {
|
||||
/// Create a new LeaderLock using the default paths in grok home
|
||||
/// (or the [`LEADER_SOCKET_ENV`] override when set).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
lock_path: default_lock_path(),
|
||||
sock_path: default_socket_path(),
|
||||
lock_file: None,
|
||||
was_leader: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn socket_path(&self) -> &PathBuf {
|
||||
&self.sock_path
|
||||
}
|
||||
|
||||
pub fn lock_path(&self) -> &PathBuf {
|
||||
&self.lock_path
|
||||
}
|
||||
|
||||
/// Open (or create) the lock file for subsequent locking operations.
|
||||
fn open_lock_file(&self) -> Result<File, LockError> {
|
||||
Ok(OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.open(&self.lock_path)?)
|
||||
}
|
||||
|
||||
/// Record a successful lock acquisition in our state.
|
||||
fn mark_acquired(&mut self, file: File) {
|
||||
self.lock_file = Some(file);
|
||||
self.was_leader = true;
|
||||
}
|
||||
|
||||
/// Try to acquire exclusive lock without blocking.
|
||||
///
|
||||
/// Returns `Ok(true)` if lock acquired, `Ok(false)` if already held by another process.
|
||||
/// After acquiring, call `write_pid()` to record the leader's PID.
|
||||
pub fn try_acquire(&mut self) -> Result<bool, LockError> {
|
||||
let file = self.open_lock_file()?;
|
||||
|
||||
match file.try_lock_exclusive() {
|
||||
Ok(()) => {
|
||||
self.mark_acquired(file);
|
||||
Ok(true)
|
||||
}
|
||||
Err(e) if is_lock_contended(&e) => Ok(false),
|
||||
Err(e) => Err(LockError::Io(e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Acquire exclusive lock, blocking until available.
|
||||
///
|
||||
/// Used by the leader process on startup. Blocks until the lock is available.
|
||||
/// After acquiring, call `write_pid()` to record the leader's PID.
|
||||
pub fn acquire_blocking(&mut self) -> Result<(), LockError> {
|
||||
let file = self.open_lock_file()?;
|
||||
|
||||
file.lock_exclusive()?;
|
||||
self.mark_acquired(file);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Try to acquire exclusive lock with a timeout.
|
||||
///
|
||||
/// Polls `try_lock_exclusive()` every 200ms until the lock is acquired or the
|
||||
/// timeout elapses. Returns `LockError::Timeout` if the deadline is exceeded.
|
||||
///
|
||||
/// Used by the leader subprocess in the socket-then-lock startup flow: the
|
||||
/// spawning client holds the lock while the leader binds its IPC socket, then
|
||||
/// releases it. This method waits for that handoff, but gives up after `timeout`
|
||||
/// so a duplicate leader (started while another is already running) exits
|
||||
/// cleanly instead of blocking forever.
|
||||
pub fn try_acquire_timeout(&mut self, timeout: Duration) -> Result<(), LockError> {
|
||||
let file = self.open_lock_file()?;
|
||||
|
||||
let deadline = Instant::now() + timeout;
|
||||
let poll_interval = Duration::from_millis(200);
|
||||
|
||||
loop {
|
||||
match file.try_lock_exclusive() {
|
||||
Ok(()) => {
|
||||
self.mark_acquired(file);
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) if is_lock_contended(&e) => {
|
||||
if Instant::now() >= deadline {
|
||||
return Err(LockError::Timeout(timeout));
|
||||
}
|
||||
std::thread::sleep(poll_interval);
|
||||
}
|
||||
Err(e) => return Err(LockError::Io(e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write our PID to the lock file. Call after acquiring lock.
|
||||
pub fn write_pid(&mut self) -> Result<(), LockError> {
|
||||
if let Some(ref mut file) = self.lock_file {
|
||||
file.set_len(0)?;
|
||||
write!(file, "{}", std::process::id())?;
|
||||
file.sync_all()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read PID from lock file (for diagnostics).
|
||||
pub fn read_pid(&self) -> Option<u32> {
|
||||
Self::read_pid_from_path(&self.lock_path)
|
||||
}
|
||||
|
||||
pub fn read_pid_from_path(path: &Path) -> Option<u32> {
|
||||
let mut content = String::new();
|
||||
File::open(path)
|
||||
.and_then(|mut f| f.read_to_string(&mut content))
|
||||
.ok()?;
|
||||
content.trim().parse().ok()
|
||||
}
|
||||
|
||||
/// Delete the socket file. Call while holding the lock.
|
||||
pub fn cleanup_socket(&self) -> io::Result<()> {
|
||||
match fs::remove_file(&self.sock_path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Release the lock explicitly.
|
||||
///
|
||||
/// This is used by the spawner to release the lock after the leader has bound
|
||||
/// its socket. After calling this, the `Drop` impl will NOT clean up files,
|
||||
/// since we're intentionally handing off to the leader process.
|
||||
pub fn release(&mut self) -> io::Result<()> {
|
||||
if let Some(file) = self.lock_file.take() {
|
||||
file.unlock()?;
|
||||
}
|
||||
// Clear was_leader so Drop doesn't delete files.
|
||||
// The actual leader process will clean up when it exits.
|
||||
self.was_leader = false;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if we currently hold the lock.
|
||||
pub fn is_held(&self) -> bool {
|
||||
self.lock_file.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LeaderLock {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl LeaderLock {
|
||||
/// Bind a lock to explicit paths for tests running outside the default home.
|
||||
pub(crate) fn from_paths(lock_path: PathBuf, sock_path: PathBuf) -> Self {
|
||||
Self {
|
||||
lock_path,
|
||||
sock_path,
|
||||
lock_file: None,
|
||||
was_leader: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for LeaderLock {
|
||||
fn drop(&mut self) {
|
||||
// Lock is automatically released when file is closed.
|
||||
// We only clean up files if was_leader is true, which means:
|
||||
// - We acquired the lock AND
|
||||
// - We did NOT call release() (which clears was_leader)
|
||||
// This ensures the spawner doesn't delete files when handing off to the leader.
|
||||
if self.was_leader {
|
||||
let _ = fs::remove_file(&self.lock_path);
|
||||
let _ = fs::remove_file(&self.sock_path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_lock(temp: &TempDir) -> LeaderLock {
|
||||
LeaderLock::from_paths(
|
||||
temp.path().join("leader.lock"),
|
||||
temp.path().join("leader.sock"),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn override_socket_path_wins_over_default_derivation() {
|
||||
let root = Path::new("/home/u/.kigi");
|
||||
let override_sock = PathBuf::from("/home/u/.kigi/leader-branch.sock");
|
||||
|
||||
// With an override, the path is taken verbatim.
|
||||
assert_eq!(
|
||||
resolve_socket_path(Some(override_sock.clone()), root),
|
||||
override_sock
|
||||
);
|
||||
// The lock is the sibling `.lock`, NOT the default name.
|
||||
assert_eq!(
|
||||
resolve_lock_path(Some(override_sock), root),
|
||||
PathBuf::from("/home/u/.kigi/leader-branch.lock")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_override_falls_back_to_default_paths() {
|
||||
let root = Path::new("/home/u/.kigi");
|
||||
assert_eq!(resolve_socket_path(None, root), root.join("leader.sock"));
|
||||
assert_eq!(resolve_lock_path(None, root), root.join("leader.lock"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lock_path_for_socket_swaps_extension() {
|
||||
assert_eq!(
|
||||
lock_path_for_socket(Path::new("/x/leader-foo.sock")),
|
||||
PathBuf::from("/x/leader-foo.lock")
|
||||
);
|
||||
// A socket path without an extension still gets a `.lock` sibling.
|
||||
assert_eq!(
|
||||
lock_path_for_socket(Path::new("/x/myleader")),
|
||||
PathBuf::from("/x/myleader.lock")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_acquire_succeeds_when_unlocked() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
assert!(lock.try_acquire().unwrap());
|
||||
assert!(lock.is_held());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_acquire_fails_when_held() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock1 = test_lock(&temp);
|
||||
let mut lock2 = test_lock(&temp);
|
||||
|
||||
assert!(lock1.try_acquire().unwrap());
|
||||
assert!(!lock2.try_acquire().unwrap()); // Should return false, not error
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_and_read_pid() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
lock.try_acquire().unwrap();
|
||||
lock.write_pid().unwrap();
|
||||
|
||||
let pid = lock.read_pid().unwrap();
|
||||
assert_eq!(pid, std::process::id());
|
||||
assert_eq!(
|
||||
LeaderLock::read_pid_from_path(lock.lock_path()),
|
||||
Some(std::process::id())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_lock_and_socket_paths_have_empty_suffix() {
|
||||
let root = Path::new("/home/u/.kigi");
|
||||
let lock_path = default_lock_path_in(root);
|
||||
let socket_path = default_socket_path_in(root);
|
||||
|
||||
assert_eq!(
|
||||
socket_suffix_from_paths(&lock_path, &socket_path),
|
||||
Some(String::new())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn override_paths_yield_matching_non_default_suffix() {
|
||||
let override_sock = PathBuf::from("/x/leader-branch.sock");
|
||||
let lock = lock_path_for_socket(&override_sock);
|
||||
assert_eq!(
|
||||
socket_suffix_from_paths(&lock, &override_sock),
|
||||
Some("-branch".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mismatched_suffixes_yield_none() {
|
||||
assert!(
|
||||
socket_suffix_from_paths(Path::new("/x/leader-a.lock"), Path::new("/x/leader-b.sock"))
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cleanup_socket_removes_file() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
// Create socket file
|
||||
fs::write(&lock.sock_path, "").unwrap();
|
||||
assert!(lock.sock_path.exists());
|
||||
|
||||
lock.try_acquire().unwrap();
|
||||
lock.cleanup_socket().unwrap();
|
||||
|
||||
assert!(!lock.sock_path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cleanup_socket_ok_if_missing() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
lock.try_acquire().unwrap();
|
||||
// Should not error even if socket doesn't exist
|
||||
lock.cleanup_socket().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn release_allows_reacquisition() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock1 = test_lock(&temp);
|
||||
let mut lock2 = test_lock(&temp);
|
||||
|
||||
assert!(lock1.try_acquire().unwrap());
|
||||
lock1.release().unwrap();
|
||||
|
||||
assert!(lock2.try_acquire().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_releases_lock() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock2 = test_lock(&temp);
|
||||
|
||||
{
|
||||
let mut lock1 = test_lock(&temp);
|
||||
assert!(lock1.try_acquire().unwrap());
|
||||
// lock1 dropped here
|
||||
}
|
||||
|
||||
// lock2 should be able to acquire now
|
||||
assert!(lock2.try_acquire().unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn release_prevents_file_cleanup_on_drop() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
// Create socket file (simulating leader binding)
|
||||
fs::write(&lock.sock_path, "").unwrap();
|
||||
assert!(lock.sock_path.exists());
|
||||
|
||||
// Acquire and then release (simulating spawner handoff)
|
||||
assert!(lock.try_acquire().unwrap());
|
||||
lock.release().unwrap();
|
||||
|
||||
// Drop should NOT delete the socket file
|
||||
drop(lock);
|
||||
|
||||
// Socket file should still exist (leader would still be using it)
|
||||
assert!(
|
||||
temp.path().join("leader.sock").exists(),
|
||||
"Socket file should NOT be deleted after release()"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_without_release_cleans_up_files() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
|
||||
{
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
// Create socket file
|
||||
fs::write(&lock.sock_path, "").unwrap();
|
||||
assert!(lock.sock_path.exists());
|
||||
|
||||
// Acquire but do NOT release (simulating crash/normal exit)
|
||||
assert!(lock.try_acquire().unwrap());
|
||||
// lock dropped here without release()
|
||||
}
|
||||
|
||||
// Socket file should be deleted
|
||||
assert!(
|
||||
!temp.path().join("leader.sock").exists(),
|
||||
"Socket file SHOULD be deleted when dropped without release()"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_pid_returns_none_for_missing_lock_file() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let lock = test_lock(&temp);
|
||||
// Lock file doesn't exist yet
|
||||
assert!(lock.read_pid().is_none());
|
||||
assert!(LeaderLock::read_pid_from_path(lock.lock_path()).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_pid_returns_none_for_empty_lock_file() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let lock = test_lock(&temp);
|
||||
fs::write(lock.lock_path(), "").unwrap();
|
||||
assert!(lock.read_pid().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_pid_returns_none_for_non_numeric_content() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let lock = test_lock(&temp);
|
||||
fs::write(lock.lock_path(), "not-a-pid").unwrap();
|
||||
assert!(lock.read_pid().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_acquire_timeout_succeeds_when_unlocked() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock = test_lock(&temp);
|
||||
|
||||
lock.try_acquire_timeout(Duration::from_secs(1)).unwrap();
|
||||
assert!(lock.is_held());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_acquire_timeout_returns_timeout_when_held() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock1 = test_lock(&temp);
|
||||
let mut lock2 = test_lock(&temp);
|
||||
|
||||
assert!(lock1.try_acquire().unwrap());
|
||||
|
||||
let result = lock2.try_acquire_timeout(Duration::from_millis(500));
|
||||
assert!(
|
||||
matches!(result, Err(LockError::Timeout(_))),
|
||||
"Expected Timeout error, got {:?}",
|
||||
result
|
||||
);
|
||||
assert!(!lock2.is_held());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_acquire_timeout_succeeds_after_release() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut lock1 = test_lock(&temp);
|
||||
let mut lock2 = test_lock(&temp);
|
||||
|
||||
assert!(lock1.try_acquire().unwrap());
|
||||
|
||||
// Release lock1 in a background thread after a short delay
|
||||
let lock_path = lock1.lock_path.clone();
|
||||
let handle = std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(200));
|
||||
lock1.release().unwrap();
|
||||
lock_path // keep the path for verification, lock1 is consumed
|
||||
});
|
||||
|
||||
// lock2 should acquire within the timeout because lock1 is released after 200ms
|
||||
lock2.try_acquire_timeout(Duration::from_secs(5)).unwrap();
|
||||
assert!(lock2.is_held());
|
||||
|
||||
handle.join().unwrap();
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,734 @@
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
|
||||
use crate::cpu_profile::{ControlError, ProfileArtifactFormat};
|
||||
|
||||
const MAX_MESSAGE_SIZE: u32 = 64 * 1024 * 1024; // 64MB
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ProtocolError {
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("Message too large: {0} bytes (max: {MAX_MESSAGE_SIZE})")]
|
||||
MessageTooLarge(u32),
|
||||
#[error("Invalid JSON: {0}")]
|
||||
InvalidJson(#[from] serde_json::Error),
|
||||
#[error("Connection closed")]
|
||||
ConnectionClosed,
|
||||
}
|
||||
|
||||
pub async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, ProtocolError> {
|
||||
let mut len_buf = [0u8; 4];
|
||||
match reader.read_exact(&mut len_buf).await {
|
||||
Ok(_) => {}
|
||||
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
|
||||
return Err(ProtocolError::ConnectionClosed);
|
||||
}
|
||||
Err(e) => return Err(ProtocolError::Io(e)),
|
||||
}
|
||||
|
||||
let len = u32::from_be_bytes(len_buf);
|
||||
if len > MAX_MESSAGE_SIZE {
|
||||
return Err(ProtocolError::MessageTooLarge(len));
|
||||
}
|
||||
|
||||
let mut buf = vec![0u8; len as usize];
|
||||
reader.read_exact(&mut buf).await?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
pub async fn write_frame<W: AsyncWrite + Unpin>(
|
||||
writer: &mut W,
|
||||
data: &[u8],
|
||||
) -> Result<(), ProtocolError> {
|
||||
let len = data.len() as u32;
|
||||
if len > MAX_MESSAGE_SIZE {
|
||||
return Err(ProtocolError::MessageTooLarge(len));
|
||||
}
|
||||
|
||||
writer.write_all(&len.to_be_bytes()).await?;
|
||||
writer.write_all(data).await?;
|
||||
writer.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn read_message<R, T>(reader: &mut R) -> Result<T, ProtocolError>
|
||||
where
|
||||
R: AsyncRead + Unpin,
|
||||
T: serde::de::DeserializeOwned,
|
||||
{
|
||||
let data = read_frame(reader).await?;
|
||||
Ok(serde_json::from_slice(&data)?)
|
||||
}
|
||||
|
||||
pub async fn write_message<W, T>(writer: &mut W, msg: &T) -> Result<(), ProtocolError>
|
||||
where
|
||||
W: AsyncWrite + Unpin,
|
||||
T: serde::Serialize,
|
||||
{
|
||||
let data = serde_json::to_vec(msg)?;
|
||||
write_frame(writer, &data).await
|
||||
}
|
||||
|
||||
/// Unique identifier for a connected client.
|
||||
///
|
||||
/// Each client gets a unique ID when connecting to the leader server.
|
||||
/// IDs are monotonically increasing and wrap around at u64::MAX.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct ClientId(pub u64);
|
||||
|
||||
impl ClientId {
|
||||
/// Generate a new unique client ID.
|
||||
///
|
||||
/// Uses an atomic counter that wraps around at u64::MAX.
|
||||
/// While collisions are theoretically possible after 2^64 IDs,
|
||||
/// this is practically impossible in real-world usage.
|
||||
pub fn new() -> Self {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(1);
|
||||
// Use wrapping_add to handle overflow gracefully
|
||||
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
Self(if id == 0 {
|
||||
COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||
} else {
|
||||
id
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ClientId {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Client mode determines how the leader handles communication for this client.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ClientMode {
|
||||
/// Headless mode (grok agent, grok agent headless) - uses websocket relay.
|
||||
/// Leader connects to websocket relay once and forwards messages.
|
||||
Headless,
|
||||
/// Stdio mode (grok agent stdio, grok -p) - uses local IPC.
|
||||
/// Client sends/receives ACP messages directly via IPC.
|
||||
Stdio,
|
||||
}
|
||||
|
||||
/// Client capabilities reported during registration.
|
||||
///
|
||||
/// These capabilities are used by the leader to customize behavior for each client,
|
||||
/// such as injecting settings into session requests.
|
||||
pub const LEADER_PROTOCOL_VERSION: u32 = 1;
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ClientCapabilities {
|
||||
/// Auto-approve all tool executions without confirmation (YOLO mode).
|
||||
/// When true, the leader will inject `yoloMode: true` into session/new requests.
|
||||
#[serde(default)]
|
||||
pub yolo_mode: bool,
|
||||
|
||||
/// Classifier permission mode (auto). When true and not yolo, the leader
|
||||
/// injects `autoMode: true` into session/new and session/load `_meta`.
|
||||
#[serde(default)]
|
||||
pub auto_mode: bool,
|
||||
|
||||
/// Default model ID to use for new sessions.
|
||||
/// When set, the leader will inject `modelId` into session/new requests
|
||||
/// (only if the request doesn't already specify a modelId).
|
||||
#[serde(default)]
|
||||
pub default_model: Option<String>,
|
||||
|
||||
/// Client binary version (e.g., "0.1.150").
|
||||
/// Used by the leader to detect version mismatches after client auto-updates.
|
||||
/// If the client version differs from the leader's version, a warning is logged.
|
||||
#[serde(default)]
|
||||
pub client_version: Option<String>,
|
||||
|
||||
/// Whether this client has advertised `x.ai/codeNavigation.enabled`.
|
||||
/// When true, the leader injects `codeNavEnabled: true` into `session/new`
|
||||
/// and `session/load` requests so the agent can gate code-nav startup on a
|
||||
/// per-client basis rather than reading from shared last-initialized state.
|
||||
#[serde(default)]
|
||||
pub code_nav_enabled: bool,
|
||||
|
||||
/// Whether the client handles terminal ACP messages (create, output, kill, etc.).
|
||||
/// When true, the leader injects `clientTerminal: true` into `session/new` and
|
||||
/// `session/load` so the agent routes terminal commands to the client via ACP
|
||||
/// instead of running them locally. Per-client so a TUI (`terminal: false`) and
|
||||
/// a web client (`terminal: true`) sharing the same leader get independent routing.
|
||||
#[serde(default)]
|
||||
pub terminal: bool,
|
||||
|
||||
/// Whether the client handles filesystem ACP read/write messages.
|
||||
/// Same per-client isolation rationale as `terminal`.
|
||||
#[serde(default)]
|
||||
pub fs_read: bool,
|
||||
#[serde(default)]
|
||||
pub fs_write: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub struct LeaderCapabilities {
|
||||
#[serde(default)]
|
||||
pub control_v1: bool,
|
||||
#[serde(default)]
|
||||
pub runtime_cpu_profile: bool,
|
||||
#[serde(default)]
|
||||
pub profile_formats: Vec<ProfileArtifactFormat>,
|
||||
#[serde(default)]
|
||||
pub workspace_exposure: bool,
|
||||
/// Whether the leader supports [`ControlCommand::RelaunchForUpdate`] — a
|
||||
/// disruptive, bounded-grace relaunch onto a freshly-installed binary
|
||||
/// (driven by `grok update`). Old leaders default to `false`, so a new
|
||||
/// client falls back to advising a manual restart (graceful degradation).
|
||||
#[serde(default)]
|
||||
pub relaunch_v1: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ControlCommand {
|
||||
GetLeaderInfo,
|
||||
CpuProfileStatus,
|
||||
StartCpuProfile {
|
||||
#[serde(default)]
|
||||
output: Option<String>,
|
||||
#[serde(default)]
|
||||
frequency_hz: Option<i32>,
|
||||
},
|
||||
StopCpuProfile,
|
||||
WorkspaceStart {
|
||||
#[serde(default)]
|
||||
hub_url: Option<String>,
|
||||
cwd: String,
|
||||
},
|
||||
WorkspacePause,
|
||||
WorkspaceResume,
|
||||
WorkspaceStop,
|
||||
WorkspaceStatus,
|
||||
/// Ask the leader to relaunch onto a freshly-installed binary (driven by
|
||||
/// `grok update`). The leader stops admitting new turns, waits a bounded
|
||||
/// grace period for in-flight turns to finish, flushes session state, then
|
||||
/// exits with [`ShutdownReason::AutoUpdate`] so connected clients reconnect
|
||||
/// onto the new binary and restore their sessions via `session/load`.
|
||||
///
|
||||
/// `to_version` is the version `grok update` just installed; the leader uses
|
||||
/// it to decline if it is already running that version or newer.
|
||||
RelaunchForUpdate {
|
||||
to_version: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ControlPayload {
|
||||
LeaderInfo {
|
||||
pid: u32,
|
||||
socket_path: PathBuf,
|
||||
lock_path: PathBuf,
|
||||
socket_suffix: String,
|
||||
leader_protocol_version: u32,
|
||||
leader_binary_version: String,
|
||||
profiling_supported: bool,
|
||||
profiling_compiled_in: bool,
|
||||
cpu_profile_active: bool,
|
||||
#[serde(default)]
|
||||
cpu_profile_stopping: bool,
|
||||
profile_started_at: Option<String>,
|
||||
profile_formats: Vec<ProfileArtifactFormat>,
|
||||
},
|
||||
CpuProfileStatus {
|
||||
active: bool,
|
||||
#[serde(default)]
|
||||
stopping: bool,
|
||||
started_at: Option<String>,
|
||||
svg_path: Option<PathBuf>,
|
||||
frequency_hz: Option<i32>,
|
||||
},
|
||||
CpuProfileStarted {
|
||||
pid: u32,
|
||||
svg_path: PathBuf,
|
||||
frequency_hz: i32,
|
||||
started_at: String,
|
||||
},
|
||||
CpuProfileStopped {
|
||||
pid: u32,
|
||||
svg_path: PathBuf,
|
||||
started_at: String,
|
||||
stopped_at: String,
|
||||
},
|
||||
WorkspaceStatus {
|
||||
state: String,
|
||||
#[serde(default)]
|
||||
hub_url: Option<String>,
|
||||
#[serde(default)]
|
||||
cwd: Option<String>,
|
||||
uptime_ms: u64,
|
||||
active_tool_calls: u32,
|
||||
#[serde(default)]
|
||||
sessions: Vec<String>,
|
||||
pid: u32,
|
||||
},
|
||||
/// Ack for [`ControlCommand::RelaunchForUpdate`]: the leader accepted the
|
||||
/// request and will exit after a bounded grace period of `grace_ms`.
|
||||
Relaunching {
|
||||
from_version: String,
|
||||
to_version: String,
|
||||
grace_ms: u64,
|
||||
},
|
||||
/// Response to [`ControlCommand::RelaunchForUpdate`] when the leader will not
|
||||
/// relaunch — e.g. it is already running `to_version` or newer, or a relaunch
|
||||
/// is already in progress.
|
||||
RelaunchDeclined { reason: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ClientMessage {
|
||||
Register {
|
||||
client_type: String,
|
||||
/// Client mode determines how leader handles this client's communication
|
||||
mode: ClientMode,
|
||||
#[serde(default)]
|
||||
capabilities: ClientCapabilities,
|
||||
},
|
||||
Acp {
|
||||
payload: String,
|
||||
},
|
||||
Control {
|
||||
request_id: String,
|
||||
command: ControlCommand,
|
||||
},
|
||||
Ping,
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
/// Reason for a planned leader shutdown, sent with [`ServerMessage::ShuttingDown`].
|
||||
///
|
||||
/// ## Runtime status
|
||||
///
|
||||
/// | Variant | Emitted today? | Notes |
|
||||
/// |---------|---------------|-------|
|
||||
/// | `AutoUpdate` | **Yes** — when `run_auto_update_checker` triggers shutdown | |
|
||||
/// | `Manual` | **Yes** — default for SIGTERM, test cancellation, all other paths | |
|
||||
/// | `IdleTimeout` | **No** — reserved for a future idle-timeout feature | |
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ShutdownReason {
|
||||
/// Leader is shutting down to install a downloaded binary auto-update.
|
||||
/// Clients should reconnect immediately via `connect_or_spawn`; the new binary
|
||||
/// will be picked up automatically.
|
||||
AutoUpdate,
|
||||
/// Reserved for a future idle-timeout feature (no active clients for a configurable
|
||||
/// duration). **Not emitted in the current implementation.**
|
||||
IdleTimeout,
|
||||
/// Unspecified or externally-triggered shutdown (SIGTERM, programmatic cancel, etc.).
|
||||
Manual,
|
||||
}
|
||||
|
||||
/// Old leaders that predate `ready` are already initialised, so default to `true`.
|
||||
fn default_ready() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// New fields must use `#[serde(default)]` — the leader and client can run different binary versions.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ServerMessage {
|
||||
/// Registration confirmation.
|
||||
///
|
||||
/// `ready` indicates whether the leader has already completed its startup
|
||||
/// (auth + model prefetch). When `ready = false` the client **must** wait for a
|
||||
/// subsequent [`LeaderReady`](Self::LeaderReady) message before sending any ACP
|
||||
/// traffic — the server will hold the connection open until the leader is ready.
|
||||
Registered {
|
||||
client_id: u64,
|
||||
/// Whether the leader is fully initialised and ready to forward ACP traffic.
|
||||
#[serde(default = "default_ready")]
|
||||
ready: bool,
|
||||
#[serde(default)]
|
||||
leader_protocol_version: Option<u32>,
|
||||
#[serde(default)]
|
||||
leader_binary_version: Option<String>,
|
||||
#[serde(default)]
|
||||
leader_capabilities: Option<LeaderCapabilities>,
|
||||
},
|
||||
Acp {
|
||||
payload: String,
|
||||
},
|
||||
ControlResult {
|
||||
request_id: String,
|
||||
result: Result<ControlPayload, ControlError>,
|
||||
},
|
||||
Pong,
|
||||
Error {
|
||||
code: i32,
|
||||
message: String,
|
||||
},
|
||||
/// Advance notice of a planned shutdown. Sent before [`Shutdown`](Self::Shutdown)
|
||||
/// to give clients time to prepare for reconnection.
|
||||
///
|
||||
/// Clients should treat this as a signal that [`Shutdown`](Self::Shutdown) is
|
||||
/// imminent and pre-arm their reconnection handlers (e.g. show a banner).
|
||||
ShuttingDown {
|
||||
reason: ShutdownReason,
|
||||
/// Milliseconds until the actual [`Shutdown`](Self::Shutdown) message.
|
||||
///
|
||||
/// **Currently always `0`** — the server sends `Shutdown` immediately after
|
||||
/// `ShuttingDown` with no intervening sleep. Clients must not rely on this
|
||||
/// field providing a real grace window in the current implementation; treat
|
||||
/// `ShuttingDown` as equivalent to an imminent `Shutdown` regardless of this
|
||||
/// value.
|
||||
delay_ms: u64,
|
||||
},
|
||||
Shutdown,
|
||||
/// Sent by the server after a `Registered { ready: false }` once the leader
|
||||
/// finishes initialising. The client should treat this as the signal that
|
||||
/// ACP traffic will now be forwarded correctly.
|
||||
LeaderReady,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::io::duplex;
|
||||
|
||||
#[tokio::test]
|
||||
async fn frame_roundtrip() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
let data = b"hello world";
|
||||
|
||||
write_frame(&mut client, data).await.unwrap();
|
||||
let received = read_frame(&mut server).await.unwrap();
|
||||
|
||||
assert_eq!(received, data);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn message_roundtrip() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
let msg = ClientMessage::Register {
|
||||
client_type: "test".into(),
|
||||
mode: ClientMode::Stdio,
|
||||
capabilities: ClientCapabilities::default(),
|
||||
};
|
||||
|
||||
write_message(&mut client, &msg).await.unwrap();
|
||||
let received: ClientMessage = read_message(&mut server).await.unwrap();
|
||||
|
||||
match received {
|
||||
ClientMessage::Register {
|
||||
client_type, mode, ..
|
||||
} => {
|
||||
assert_eq!(client_type, "test");
|
||||
assert_eq!(mode, ClientMode::Stdio);
|
||||
}
|
||||
_ => panic!("wrong message type"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn control_message_roundtrip() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
let msg = ClientMessage::Control {
|
||||
request_id: "req-1".into(),
|
||||
command: ControlCommand::StartCpuProfile {
|
||||
output: Some("/tmp/profile.folded".into()),
|
||||
frequency_hz: Some(250),
|
||||
},
|
||||
};
|
||||
|
||||
write_message(&mut client, &msg).await.unwrap();
|
||||
let received: ClientMessage = read_message(&mut server).await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
received,
|
||||
ClientMessage::Control {
|
||||
request_id,
|
||||
command: ControlCommand::StartCpuProfile {
|
||||
output: Some(output),
|
||||
frequency_hz: Some(250),
|
||||
},
|
||||
} if request_id == "req-1" && output == "/tmp/profile.folded"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_closed_on_eof() {
|
||||
let (client, mut server) = duplex(1024);
|
||||
drop(client);
|
||||
|
||||
match read_frame(&mut server).await {
|
||||
Err(ProtocolError::ConnectionClosed) => {}
|
||||
other => panic!("expected ConnectionClosed, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_oversized_messages() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
|
||||
// Write a length header claiming a huge message
|
||||
client
|
||||
.write_all(&(MAX_MESSAGE_SIZE + 1).to_be_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
match read_frame(&mut server).await {
|
||||
Err(ProtocolError::MessageTooLarge(size)) => {
|
||||
assert_eq!(size, MAX_MESSAGE_SIZE + 1);
|
||||
}
|
||||
other => panic!("expected MessageTooLarge, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiple_frames_in_sequence() {
|
||||
let (mut client, mut server) = duplex(4096);
|
||||
|
||||
for i in 0..10 {
|
||||
let data = format!("message {}", i);
|
||||
write_frame(&mut client, data.as_bytes()).await.unwrap();
|
||||
}
|
||||
drop(client);
|
||||
|
||||
for i in 0..10 {
|
||||
let received = read_frame(&mut server).await.unwrap();
|
||||
assert_eq!(received, format!("message {}", i).as_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registered_serde_compatibility_without_optional_metadata() {
|
||||
let json = r#"{"type":"registered","client_id":7}"#;
|
||||
let msg: ServerMessage = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
msg,
|
||||
ServerMessage::Registered {
|
||||
client_id: 7,
|
||||
// `ready` defaults to `true` via `default_ready()` — old leaders
|
||||
// that predate the field are already initialised.
|
||||
ready: true,
|
||||
leader_protocol_version: None,
|
||||
leader_binary_version: None,
|
||||
leader_capabilities: None,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registered_serde_compatibility_with_all_optional_metadata() {
|
||||
let msg = ServerMessage::Registered {
|
||||
client_id: 7,
|
||||
ready: true,
|
||||
leader_protocol_version: Some(LEADER_PROTOCOL_VERSION),
|
||||
leader_binary_version: Some("1.2.3".into()),
|
||||
leader_capabilities: Some(LeaderCapabilities {
|
||||
control_v1: true,
|
||||
runtime_cpu_profile: true,
|
||||
profile_formats: vec![ProfileArtifactFormat::Svg],
|
||||
workspace_exposure: true,
|
||||
relaunch_v1: true,
|
||||
}),
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&msg).unwrap();
|
||||
let decoded: ServerMessage = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(
|
||||
decoded,
|
||||
ServerMessage::Registered {
|
||||
client_id: 7,
|
||||
ready: true,
|
||||
leader_protocol_version: Some(LEADER_PROTOCOL_VERSION),
|
||||
leader_binary_version: Some(_),
|
||||
leader_capabilities: Some(LeaderCapabilities {
|
||||
control_v1: true,
|
||||
runtime_cpu_profile: true,
|
||||
profile_formats,
|
||||
workspace_exposure: true,
|
||||
relaunch_v1: true,
|
||||
}),
|
||||
} if profile_formats == vec![ProfileArtifactFormat::Svg]
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn profile_artifact_format_serde_names_are_stable() {
|
||||
// Wire compat contract: `svg` must stay decodable (old leaders
|
||||
// advertise it), and `folded` is the name new binaries will start
|
||||
// advertising once the fleet can decode it. Renaming either variant
|
||||
// breaks the Registered handshake across version skew.
|
||||
assert_eq!(
|
||||
serde_json::to_string(&ProfileArtifactFormat::Svg).unwrap(),
|
||||
"\"svg\""
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_string(&ProfileArtifactFormat::Folded).unwrap(),
|
||||
"\"folded\""
|
||||
);
|
||||
let decoded: ProfileArtifactFormat = serde_json::from_str("\"svg\"").unwrap();
|
||||
assert_eq!(decoded, ProfileArtifactFormat::Svg);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn control_payload_serde_defaults_new_stopping_flags() {
|
||||
let leader_info_json = r#"{
|
||||
"type":"leader_info",
|
||||
"pid":123,
|
||||
"socket_path":"/tmp/leader.sock",
|
||||
"lock_path":"/tmp/leader.lock",
|
||||
"socket_suffix":"suffix",
|
||||
"leader_protocol_version":1,
|
||||
"leader_binary_version":"1.2.3",
|
||||
"profiling_supported":true,
|
||||
"profiling_compiled_in":true,
|
||||
"cpu_profile_active":false,
|
||||
"profile_started_at":null,
|
||||
"profile_formats":["svg"]
|
||||
}"#;
|
||||
let status_json = r#"{
|
||||
"type":"cpu_profile_status",
|
||||
"active":false,
|
||||
"started_at":null,
|
||||
"svg_path":null,
|
||||
"frequency_hz":null
|
||||
}"#;
|
||||
|
||||
let leader_info: ControlPayload = serde_json::from_str(leader_info_json).unwrap();
|
||||
let status: ControlPayload = serde_json::from_str(status_json).unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
leader_info,
|
||||
ControlPayload::LeaderInfo {
|
||||
cpu_profile_active: false,
|
||||
cpu_profile_stopping: false,
|
||||
profile_started_at: None,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert!(matches!(
|
||||
status,
|
||||
ControlPayload::CpuProfileStatus {
|
||||
active: false,
|
||||
stopping: false,
|
||||
started_at: None,
|
||||
svg_path: None,
|
||||
frequency_hz: None,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn workspace_control_command_roundtrip() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
let msg = ClientMessage::Control {
|
||||
request_id: "ws-1".into(),
|
||||
command: ControlCommand::WorkspaceStart {
|
||||
hub_url: Some("wss://hub.example/v1/tools".into()),
|
||||
cwd: "/home/u/proj".into(),
|
||||
},
|
||||
};
|
||||
|
||||
write_message(&mut client, &msg).await.unwrap();
|
||||
let received: ClientMessage = read_message(&mut server).await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
received,
|
||||
ClientMessage::Control {
|
||||
request_id,
|
||||
command: ControlCommand::WorkspaceStart { hub_url: Some(url), cwd },
|
||||
} if request_id == "ws-1"
|
||||
&& url == "wss://hub.example/v1/tools"
|
||||
&& cwd == "/home/u/proj"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_status_payload_roundtrip() {
|
||||
let payload = ControlPayload::WorkspaceStatus {
|
||||
state: "running".into(),
|
||||
hub_url: Some("wss://hub.example/v1/tools".into()),
|
||||
cwd: Some("/home/u/proj".into()),
|
||||
uptime_ms: 4200,
|
||||
active_tool_calls: 2,
|
||||
sessions: vec!["grok-a".into(), "grok-b".into()],
|
||||
pid: 4242,
|
||||
};
|
||||
let json = serde_json::to_string(&payload).unwrap();
|
||||
let decoded: ControlPayload = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(decoded, payload);
|
||||
assert!(json.contains("\"type\":\"workspace_status\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_status_payload_defaults_optional_fields() {
|
||||
let json = r#"{"type":"workspace_status","state":"none","uptime_ms":0,"active_tool_calls":0,"pid":1}"#;
|
||||
let decoded: ControlPayload = serde_json::from_str(json).unwrap();
|
||||
assert!(matches!(
|
||||
decoded,
|
||||
ControlPayload::WorkspaceStatus {
|
||||
state,
|
||||
hub_url: None,
|
||||
cwd: None,
|
||||
sessions,
|
||||
..
|
||||
} if state == "none" && sessions.is_empty()
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_exposure_capability_defaults_false() {
|
||||
let json = r#"{"control_v1":true,"runtime_cpu_profile":false,"profile_formats":[]}"#;
|
||||
let caps: LeaderCapabilities = serde_json::from_str(json).unwrap();
|
||||
assert!(!caps.workspace_exposure);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_id_is_unique() {
|
||||
let ids: Vec<_> = (0..100).map(|_| ClientId::new()).collect();
|
||||
let unique: std::collections::HashSet<_> = ids.iter().map(|c| c.0).collect();
|
||||
assert_eq!(unique.len(), 100);
|
||||
}
|
||||
|
||||
// --- ShuttingDown / ShutdownReason tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutting_down_message_roundtrip() {
|
||||
let (mut client, mut server) = duplex(1024);
|
||||
let msg = ServerMessage::ShuttingDown {
|
||||
reason: ShutdownReason::AutoUpdate,
|
||||
delay_ms: 2000,
|
||||
};
|
||||
|
||||
write_message(&mut client, &msg).await.unwrap();
|
||||
let received: ServerMessage = read_message(&mut server).await.unwrap();
|
||||
|
||||
match received {
|
||||
ServerMessage::ShuttingDown { reason, delay_ms } => {
|
||||
assert_eq!(reason, ShutdownReason::AutoUpdate);
|
||||
assert_eq!(delay_ms, 2000);
|
||||
}
|
||||
_ => panic!("Expected ShuttingDown, got {:?}", received),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_reason_variants_serialize_correctly() {
|
||||
let auto = serde_json::to_string(&ShutdownReason::AutoUpdate).unwrap();
|
||||
assert_eq!(auto, "\"auto_update\"");
|
||||
|
||||
let idle = serde_json::to_string(&ShutdownReason::IdleTimeout).unwrap();
|
||||
assert_eq!(idle, "\"idle_timeout\"");
|
||||
|
||||
let manual = serde_json::to_string(&ShutdownReason::Manual).unwrap();
|
||||
assert_eq!(manual, "\"manual\"");
|
||||
|
||||
// Verify deserialization
|
||||
let parsed: ShutdownReason = serde_json::from_str("\"auto_update\"").unwrap();
|
||||
assert_eq!(parsed, ShutdownReason::AutoUpdate);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,174 @@
|
||||
//! In-crate fake leaders for exercising client-side handling of misbehaving
|
||||
//! leaders (hung, half-framed, wrong-versioned) — wire shapes the real
|
||||
//! `spawn_leader_server` can never produce.
|
||||
//!
|
||||
//! All stalls are cancellation-based (`cancel.cancelled().await`), never
|
||||
//! timer-based, so `#[tokio::test(start_paused = true)]` auto-advance jumps
|
||||
//! the client-side timeouts under test without waking the fake.
|
||||
use super::protocol::{
|
||||
ClientMessage, LEADER_PROTOCOL_VERSION, LeaderCapabilities, ServerMessage, read_message,
|
||||
write_message,
|
||||
};
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
/// Version metadata a fake leader reports in `Registered`.
|
||||
pub(crate) struct FakeVersions {
|
||||
pub(crate) protocol_version: Option<u32>,
|
||||
pub(crate) binary_version: Option<String>,
|
||||
}
|
||||
impl FakeVersions {
|
||||
/// The versions a same-build real leader would report (`run_leader` stamps
|
||||
/// `kigi_version::VERSION` into its metadata).
|
||||
pub(crate) fn current() -> Self {
|
||||
Self {
|
||||
protocol_version: Some(LEADER_PROTOCOL_VERSION),
|
||||
binary_version: Some(kigi_version::VERSION.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
/// `LeaderCapabilities` has no `Default` (serde-only defaults), so fakes build
|
||||
/// their capability shape through this helper.
|
||||
pub(crate) fn fake_caps(control_v1: bool, relaunch_v1: bool) -> LeaderCapabilities {
|
||||
LeaderCapabilities {
|
||||
control_v1,
|
||||
runtime_cpu_profile: false,
|
||||
profile_formats: Vec::new(),
|
||||
workspace_exposure: false,
|
||||
relaunch_v1,
|
||||
}
|
||||
}
|
||||
/// Wire behavior of a [`spawn_fake_leader`] instance.
|
||||
pub(crate) enum FakeLeaderBehavior {
|
||||
/// Well-formed: `Registered { ready: true }` with the given metadata, then
|
||||
/// idle until cancelled. Backs the discovery and adopt/evict tests; also
|
||||
/// the composition point for metadata skew (wrong protocol version, stale
|
||||
/// binary version) via an explicit [`FakeVersions`] — no sugar variants.
|
||||
Normal {
|
||||
versions: FakeVersions,
|
||||
caps: LeaderCapabilities,
|
||||
},
|
||||
/// Accepts the connection but never sends anything (hung pre-`Registered`).
|
||||
SilentAfterAccept,
|
||||
/// `Registered { ready: false }`, then never sends `LeaderReady`.
|
||||
ReadyFalseForever,
|
||||
/// Writes only `bytes` (< 4) of the 4-byte length prefix, then stalls.
|
||||
PartialFrame { bytes: usize },
|
||||
/// Valid length prefix followed by a non-JSON body.
|
||||
GarbageFrame,
|
||||
/// Well-formed `Registered { ready: true }`, then closes the connection.
|
||||
CloseAfterRegister,
|
||||
}
|
||||
/// Handle for a running fake leader; cancelling stops the accept loop and any
|
||||
/// held-open connections, and removes the socket.
|
||||
pub(crate) struct FakeLeaderHandle {
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
impl FakeLeaderHandle {
|
||||
pub(crate) fn cancel(&self) {
|
||||
self.cancel.cancel();
|
||||
}
|
||||
}
|
||||
/// Bind a fake leader at `socket_path` behaving per `behavior`.
|
||||
///
|
||||
/// Returns once the listener is bound (readiness signalled via oneshot, no
|
||||
/// fixed startup sleep), so callers can connect immediately. Serves clients
|
||||
/// sequentially: the point of a fake is wire shape, not concurrency.
|
||||
pub(crate) async fn spawn_fake_leader(
|
||||
socket_path: PathBuf,
|
||||
behavior: FakeLeaderBehavior,
|
||||
) -> FakeLeaderHandle {
|
||||
let cancel = CancellationToken::new();
|
||||
let cancel_clone = cancel.clone();
|
||||
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
tokio::spawn(async move {
|
||||
let _ = fs::remove_file(&socket_path);
|
||||
let listener = match super::transport::LeaderListener::bind(&socket_path) {
|
||||
Ok(listener) => listener,
|
||||
Err(_) => return,
|
||||
};
|
||||
let _ = ready_tx.send(());
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = cancel_clone.cancelled() => break, accept_result = listener.accept()
|
||||
=> { let Ok((stream, _)) = accept_result else { break; };
|
||||
serve_client(stream, & behavior, & cancel_clone). await; }
|
||||
}
|
||||
}
|
||||
let _ = fs::remove_file(&socket_path);
|
||||
});
|
||||
let _ = ready_rx.await;
|
||||
FakeLeaderHandle { cancel }
|
||||
}
|
||||
async fn serve_client(
|
||||
stream: super::transport::LeaderStream,
|
||||
behavior: &FakeLeaderBehavior,
|
||||
cancel: &CancellationToken,
|
||||
) {
|
||||
let (mut reader, mut writer) = tokio::io::split(stream);
|
||||
/// A `Registered` with `client_id: 1` and the given shape.
|
||||
fn registered(
|
||||
ready: bool,
|
||||
versions: &FakeVersions,
|
||||
caps: &LeaderCapabilities,
|
||||
) -> ServerMessage {
|
||||
ServerMessage::Registered {
|
||||
client_id: 1,
|
||||
ready,
|
||||
leader_protocol_version: versions.protocol_version,
|
||||
leader_binary_version: versions.binary_version.clone(),
|
||||
leader_capabilities: Some(caps.clone()),
|
||||
}
|
||||
}
|
||||
match behavior {
|
||||
FakeLeaderBehavior::SilentAfterAccept => {
|
||||
cancel.cancelled().await;
|
||||
}
|
||||
FakeLeaderBehavior::PartialFrame { bytes } => {
|
||||
let prefix = 1024u32.to_be_bytes();
|
||||
let n = (*bytes).min(prefix.len());
|
||||
let _ = writer.write_all(&prefix[..n]).await;
|
||||
let _ = writer.flush().await;
|
||||
cancel.cancelled().await;
|
||||
}
|
||||
FakeLeaderBehavior::GarbageFrame => {
|
||||
let body = b"this is not json";
|
||||
let _ = writer.write_all(&(body.len() as u32).to_be_bytes()).await;
|
||||
let _ = writer.write_all(body).await;
|
||||
let _ = writer.flush().await;
|
||||
cancel.cancelled().await;
|
||||
}
|
||||
FakeLeaderBehavior::Normal { versions, caps } => {
|
||||
let register: Result<ClientMessage, _> = read_message(&mut reader).await;
|
||||
if register.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = write_message(&mut writer, ®istered(true, versions, caps)).await;
|
||||
cancel.cancelled().await;
|
||||
}
|
||||
FakeLeaderBehavior::ReadyFalseForever => {
|
||||
let register: Result<ClientMessage, _> = read_message(&mut reader).await;
|
||||
if register.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = write_message(
|
||||
&mut writer,
|
||||
®istered(false, &FakeVersions::current(), &fake_caps(true, false)),
|
||||
)
|
||||
.await;
|
||||
cancel.cancelled().await;
|
||||
}
|
||||
FakeLeaderBehavior::CloseAfterRegister => {
|
||||
let register: Result<ClientMessage, _> = read_message(&mut reader).await;
|
||||
if register.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = write_message(
|
||||
&mut writer,
|
||||
®istered(true, &FakeVersions::current(), &fake_caps(true, false)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
//! Cross-platform IPC transport for leader<->client communication.
|
||||
//!
|
||||
//! - **Unix:** [`LeaderStream`] / [`LeaderListener`] are type aliases for
|
||||
//! `tokio::net::UnixStream` / `UnixListener`. Zero wrapper, no unsafe.
|
||||
//! - **Windows:** wraps `tokio::net::windows::named_pipe::*` (tokio doesn't
|
||||
//! expose AF_UNIX on Windows). The leader's filesystem path is hashed
|
||||
//! into `\\.\pipe\grok-leader-<hash>` so callers keep their path-based API.
|
||||
//!
|
||||
#[cfg(unix)]
|
||||
pub use tokio::net::UnixListener as LeaderListener;
|
||||
#[cfg(unix)]
|
||||
pub use tokio::net::UnixStream as LeaderStream;
|
||||
|
||||
/// Has a leader bound a listener at `path`?
|
||||
///
|
||||
/// - Unix: stats the socket file.
|
||||
/// - Windows: probes the named pipe (Named Pipes don't appear in the
|
||||
/// filesystem, so `path.exists()` doesn't work).
|
||||
pub fn listener_is_ready(path: &std::path::Path) -> bool {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
path.exists()
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
windows_impl::listener_is_ready(path)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
pub use windows_impl::{LeaderListener, LeaderStream};
|
||||
|
||||
#[cfg(windows)]
|
||||
mod windows_impl {
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tracing::debug;
|
||||
|
||||
/// Bidirectional IPC stream wrapping a connected named pipe (server-
|
||||
/// or client-side, depending on how it was created).
|
||||
pub struct LeaderStream {
|
||||
inner: StreamInner,
|
||||
}
|
||||
|
||||
enum StreamInner {
|
||||
Server(tokio::net::windows::named_pipe::NamedPipeServer),
|
||||
Client(tokio::net::windows::named_pipe::NamedPipeClient),
|
||||
}
|
||||
|
||||
impl LeaderStream {
|
||||
/// Connect to a listener at `path`. The path is translated to a
|
||||
/// named-pipe name and `ClientOptions::open` is used.
|
||||
pub async fn connect<P: AsRef<Path>>(path: P) -> io::Result<Self> {
|
||||
use tokio::net::windows::named_pipe::ClientOptions;
|
||||
|
||||
// ClientOptions::open returns ERROR_PIPE_BUSY if all pipe
|
||||
// instances are in use; the caller's CONNECT_TIMEOUT loop
|
||||
// already retries, so we surface the error and let it handle.
|
||||
let pipe_name = path_to_pipe_name(path.as_ref());
|
||||
let inner = ClientOptions::new().open(pipe_name)?;
|
||||
Ok(Self {
|
||||
inner: StreamInner::Client(inner),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// tokio's NamedPipeServer / NamedPipeClient are auto-Unpin (they wrap
|
||||
// PollEvented<mio::windows::NamedPipe>, which is Unpin), so our
|
||||
// wrapping enum and struct are auto-Unpin as well. That means
|
||||
// Pin<&mut Self>::get_mut() is safe — no unsafe needed for the
|
||||
// structural projection into `inner`.
|
||||
impl AsyncRead for LeaderStream {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
match &mut self.get_mut().inner {
|
||||
StreamInner::Server(s) => Pin::new(s).poll_read(cx, buf),
|
||||
StreamInner::Client(c) => Pin::new(c).poll_read(cx, buf),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for LeaderStream {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
match &mut self.get_mut().inner {
|
||||
StreamInner::Server(s) => Pin::new(s).poll_write(cx, buf),
|
||||
StreamInner::Client(c) => Pin::new(c).poll_write(cx, buf),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
match &mut self.get_mut().inner {
|
||||
StreamInner::Server(s) => Pin::new(s).poll_flush(cx),
|
||||
StreamInner::Client(c) => Pin::new(c).poll_flush(cx),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
match &mut self.get_mut().inner {
|
||||
StreamInner::Server(s) => Pin::new(s).poll_shutdown(cx),
|
||||
StreamInner::Client(c) => Pin::new(c).poll_shutdown(cx),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Listener for incoming leader IPC connections. Holds the pipe name
|
||||
/// plus the next pre-created server instance (Windows named pipes
|
||||
/// require pre-creating an instance per pending connection).
|
||||
pub struct LeaderListener {
|
||||
pipe_name: std::ffi::OsString,
|
||||
/// Next pre-created server instance, ready for `connect().await`.
|
||||
/// We rotate: take this one, await its connect, immediately create
|
||||
/// the next one for the following accept(). The first instance is
|
||||
/// created in `bind()` with `first_pipe_instance(true)` to lock
|
||||
/// out other processes from squatting the pipe name.
|
||||
///
|
||||
/// tokio::sync::Mutex (not parking_lot) because accept() holds the
|
||||
/// lock across `server.connect().await`.
|
||||
next_server: tokio::sync::Mutex<Option<tokio::net::windows::named_pipe::NamedPipeServer>>,
|
||||
}
|
||||
|
||||
impl LeaderListener {
|
||||
/// Reserve a named-pipe name (no on-disk file is created).
|
||||
pub fn bind<P: AsRef<Path>>(path: P) -> io::Result<Self> {
|
||||
use tokio::net::windows::named_pipe::ServerOptions;
|
||||
|
||||
let pipe_name = path_to_pipe_name(path.as_ref());
|
||||
let first = ServerOptions::new()
|
||||
.first_pipe_instance(true)
|
||||
.create(&pipe_name)?;
|
||||
Ok(Self {
|
||||
pipe_name,
|
||||
next_server: tokio::sync::Mutex::new(Some(first)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Wait for the next incoming connection. Mirrors
|
||||
/// `UnixListener::accept`, returning a connected stream and a unit
|
||||
/// placeholder where Unix would return the peer address (named
|
||||
/// pipes don't carry one).
|
||||
pub async fn accept(&self) -> io::Result<(LeaderStream, ())> {
|
||||
use tokio::net::windows::named_pipe::ServerOptions;
|
||||
|
||||
// Take the pending instance (or create one), await a client, then
|
||||
// pre-create the next. On connect() error, drop the instance and
|
||||
// retry with a fresh one — returning early would leave the slot
|
||||
// empty and brick the listener. Bounded with a backoff so a
|
||||
// persistently failing connect() can't busy-spin.
|
||||
const MAX_ACCEPT_ATTEMPTS: usize = 10;
|
||||
const RETRY_BACKOFF: Duration = Duration::from_millis(20);
|
||||
|
||||
let mut slot = self.next_server.lock().await;
|
||||
let mut last_err: Option<io::Error> = None;
|
||||
for attempt in 0..MAX_ACCEPT_ATTEMPTS {
|
||||
let server = match slot.take() {
|
||||
Some(server) => server,
|
||||
None => ServerOptions::new().create(&self.pipe_name)?,
|
||||
};
|
||||
match server.connect().await {
|
||||
Ok(()) => {
|
||||
*slot = Some(ServerOptions::new().create(&self.pipe_name)?);
|
||||
return Ok((
|
||||
LeaderStream {
|
||||
inner: StreamInner::Server(server),
|
||||
},
|
||||
(),
|
||||
));
|
||||
}
|
||||
Err(e) => {
|
||||
// Failed `server` drops here, freeing the instance.
|
||||
debug!(attempt, error = %e, "named-pipe accept connect failed; retrying");
|
||||
last_err = Some(e);
|
||||
tokio::time::sleep(RETRY_BACKOFF).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Best-effort re-arm; take-or-create above still recovers if this fails.
|
||||
if let Ok(fresh) = ServerOptions::new().create(&self.pipe_name) {
|
||||
*slot = Some(fresh);
|
||||
}
|
||||
Err(last_err
|
||||
.unwrap_or_else(|| io::Error::other("LeaderListener: accept exhausted retries")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a leader has a pipe bound at `path`.
|
||||
///
|
||||
/// Probes with `WaitNamedPipeW` (non-connecting), not `ClientOptions::open`,
|
||||
/// which would open a real client the leader's `accept()` consumes as a
|
||||
/// phantom session. `ERROR_FILE_NOT_FOUND` means absent; `TRUE` or any other
|
||||
/// error (e.g. `ERROR_SEM_TIMEOUT`: exists but busy) means ready.
|
||||
pub fn listener_is_ready(path: &Path) -> bool {
|
||||
use std::os::windows::ffi::OsStrExt;
|
||||
|
||||
use windows::Win32::Foundation::{ERROR_FILE_NOT_FOUND, GetLastError};
|
||||
use windows::Win32::System::Pipes::WaitNamedPipeW;
|
||||
use windows::core::PCWSTR;
|
||||
|
||||
// 1 ms (a real timeout, not 0 = "server default").
|
||||
const PROBE_TIMEOUT_MS: u32 = 1;
|
||||
|
||||
let pipe_name = path_to_pipe_name(path);
|
||||
let wide: Vec<u16> = pipe_name
|
||||
.as_os_str()
|
||||
.encode_wide()
|
||||
.chain(std::iter::once(0))
|
||||
.collect();
|
||||
|
||||
if unsafe { WaitNamedPipeW(PCWSTR(wide.as_ptr()), PROBE_TIMEOUT_MS) }.as_bool() {
|
||||
return true;
|
||||
}
|
||||
// FALSE: only a missing pipe means not-ready.
|
||||
let err = unsafe { GetLastError() };
|
||||
err != ERROR_FILE_NOT_FOUND
|
||||
}
|
||||
|
||||
/// Full named-pipe path: `\\.\pipe\<leaf>`.
|
||||
fn path_to_pipe_name(path: &Path) -> std::ffi::OsString {
|
||||
let mut name = std::ffi::OsString::from(r"\\.\pipe\");
|
||||
name.push(pipe_leaf_name(path));
|
||||
name
|
||||
}
|
||||
|
||||
/// Deterministic leaf name (`grok-leader-<hash>`) for a filesystem path.
|
||||
///
|
||||
/// Uses SipHash-1-3 with fixed keys so the hash is stable across Rust
|
||||
/// versions (unlike `DefaultHasher`, whose algorithm is unspecified).
|
||||
fn pipe_leaf_name(path: &Path) -> std::ffi::OsString {
|
||||
use siphasher::sip::SipHasher13;
|
||||
use std::hash::{Hash, Hasher};
|
||||
|
||||
// Fixed keys — must never change once shipped.
|
||||
let mut hasher = SipHasher13::new_with_keys(0x67726f6b_6c656164, 0x65725f70_69706521);
|
||||
path.hash(&mut hasher);
|
||||
let hash = hasher.finish();
|
||||
std::ffi::OsString::from(format!("grok-leader-{hash:016x}"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::Path;
|
||||
|
||||
#[test]
|
||||
fn pipe_name_is_deterministic() {
|
||||
let a = path_to_pipe_name(Path::new("/tmp/grok.sock"));
|
||||
let b = path_to_pipe_name(Path::new("/tmp/grok.sock"));
|
||||
assert_eq!(a, b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_paths_produce_different_names() {
|
||||
let a = path_to_pipe_name(Path::new("/tmp/a.sock"));
|
||||
let b = path_to_pipe_name(Path::new("/tmp/b.sock"));
|
||||
assert_ne!(a, b);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pipe_name_has_correct_prefix() {
|
||||
let name = path_to_pipe_name(Path::new("/tmp/test.sock"));
|
||||
let s = name.to_string_lossy();
|
||||
assert!(s.starts_with(r"\\.\pipe\grok-leader-"), "got: {s}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pipe_name_is_bounded() {
|
||||
let long_path = "/".to_owned() + &"a".repeat(500);
|
||||
let name = path_to_pipe_name(Path::new(&long_path));
|
||||
// \\.\pipe\grok-leader- (20 chars) + 16 hex chars = 36 total
|
||||
assert!(name.len() <= 256, "pipe name too long: {}", name.len());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn listener_is_ready_tracks_pipe_lifecycle() {
|
||||
// Unique path per process so parallel test binaries don't collide on
|
||||
// the derived pipe name.
|
||||
let path =
|
||||
std::env::temp_dir().join(format!("grok-ready-probe-{}.sock", std::process::id()));
|
||||
|
||||
// Nothing bound yet -> ERROR_FILE_NOT_FOUND -> not ready.
|
||||
assert!(!listener_is_ready(&path));
|
||||
|
||||
let listener = LeaderListener::bind(&path).unwrap();
|
||||
// Ready as soon as the pipe is bound, before any accept().
|
||||
assert!(listener_is_ready(&path));
|
||||
|
||||
// After the last instance is dropped the pipe name disappears.
|
||||
drop(listener);
|
||||
assert!(!listener_is_ready(&path));
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user