ZacharyZhang-NY
6f31415ed6
The PRD's first acceptance gate now holds: grep -RinE '\bx\.ai\b|grok' crates/ --include='*.rs' → 0 matches (exempt: NOTICE and third-party license archives, README provenance, and the required 'Based on Grok Build Open Source' attribution, now sourced from version_attribution.txt). Wire-visible renames (both sides in this repo, changed in lockstep): - Auth method id 'grok.com' → 'kimi-code' (AuthMethodKind::KimiCode). - Every x.ai/* and _x.ai/* ACP ext method and meta key → kigi/* / _kigi/* (~200 names; grokShell → kigiShell). Session-file replay keeps a read-side alias for the legacy '_x.ai/session/update' method so existing updates.jsonl histories load; writes emit only the new name (both directions test-pinned). - Agent types grok-build* → kigi* with a documented legacy-prefix alias at resolution time so persisted sessions keep resolving. - ToolNamespace/BuiltinAgentName GrokBuild* → Kigi* (wire snake_case kigi/kigi_concise/kigi_hashline; schema regenerated); grok_build implementation dirs renamed to kigi*. - x-grok-* headers → x-kigi-*, __GROK_* sentinels → __KIGI_*, themes grokday/groknight → kigiday/kiginight (old persisted values fall back to the default theme), web_fetch allowlist xAI hosts → kimi.com + moonshot platforms, changelog CDN → this repo, grok-build changelog archives deleted. - BYOK default endpoint removed: [endpoints] api_base_url is now truly optional with NO default — consumers fail fast with the flag name when unset (no silent x.ai egress). Mock harnesses inject it explicitly. - System-prompt identity fixed: 'released by xAI' → 'an unofficial community CLI for Kimi' (template + regenerated encrypted form). Also repaired pre-existing grok-era test debt found by the sweep: the stale trace_classify default-model pin, the grok-pager UA label test, pty-harness stale-binary reuse and non-hermetic moonshot routing (a PTY test could previously reach the real api.moonshot.cn), and the outdated oauth fixture scope key. Gates: §9 grep 0; fmt clean; workspace check/clippy 0/0 (-D warnings); FULL cargo test --workspace: 234 suites, 21,961 passed, 0 failed; deny advisories ok.
2411 lines
101 KiB
Rust
2411 lines
101 KiB
Rust
#![allow(
|
|
unused_imports,
|
|
unused_variables,
|
|
unused_mut,
|
|
unreachable_code,
|
|
dead_code
|
|
)]
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
#[global_allocator]
|
|
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
|
|
#[cfg(all(feature = "jemalloc", feature = "release-dist", unix))]
|
|
mod jemalloc_malloc_conf {
|
|
/// jemalloc looks up `extern const char *malloc_conf` — a thin pointer,
|
|
/// not a Rust `&[u8]` fat pointer.
|
|
#[repr(transparent)]
|
|
struct MallocConfPtr(*const u8);
|
|
unsafe impl Sync for MallocConfPtr {}
|
|
static CONF: [u8; 63] = *b"prof:true,prof_active:false,lg_prof_sample:19,prof_final:false\0";
|
|
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
|
|
#[used]
|
|
#[unsafe(export_name = "malloc_conf")]
|
|
static MALLOC_CONF: MallocConfPtr = MallocConfPtr(CONF.as_ptr());
|
|
#[cfg(any(target_os = "macos", target_os = "ios"))]
|
|
#[used]
|
|
#[unsafe(export_name = "_rjem_malloc_conf")]
|
|
static MALLOC_CONF: MallocConfPtr = MallocConfPtr(CONF.as_ptr());
|
|
}
|
|
use anyhow::Result;
|
|
use kigi_shell::agent::app::{run_leader, run_stdio_agent};
|
|
use kigi_shell::agent::config::Config as AgentConfig;
|
|
use kigi_shell::leader::{
|
|
ClientCapabilities, ClientMode, ControlCommand, LeaderCapabilities, LeaderDescriptor,
|
|
LeaderRegistration, LeaderTarget, leader_is_older_than,
|
|
};
|
|
use kigi_shell::leader::{ControlPayload, LeaderClient, connect_or_spawn, default_socket_path};
|
|
use kigi_tui::app::{
|
|
AgentCmd, Command, LeaderMgmtArgs, LeaderMgmtCommand, LeaderTargetArgs, PagerArgs,
|
|
resolve_use_leader,
|
|
};
|
|
use kigi_tui::client_identity::PAGER_CLIENT_VERSION;
|
|
use kigi_update::{UpdateConfig, auto_update, enforce_minimum_version_or_exit};
|
|
use std::env;
|
|
use std::net::SocketAddr;
|
|
use tokio_util::sync::CancellationToken;
|
|
/// Apply global endpoint CLI args to an existing config.
|
|
fn apply_agent_endpoint_args(agent_args: &kigi_tui::app::AgentArgs, config: &mut AgentConfig) {
|
|
if let Some(v) = &agent_args.coding_api_base_url {
|
|
config.endpoints.coding_api_base_url = Some(v.clone());
|
|
}
|
|
if let Some(v) = &agent_args.api_base_url {
|
|
config.endpoints.api_base_url = Some(v.clone());
|
|
}
|
|
}
|
|
/// Resolve --agent-profile path: canonicalize and verify the file exists.
|
|
fn resolve_agent_profile_path(path: &std::path::Path) -> std::path::PathBuf {
|
|
match dunce::canonicalize(path) {
|
|
Ok(abs) if abs.is_file() => abs,
|
|
Ok(abs) => {
|
|
eprintln!(
|
|
"error: --agent-profile path is not a file: {}",
|
|
abs.display()
|
|
);
|
|
std::process::exit(1);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("error: --agent-profile path '{}': {}", path.display(), e);
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
}
|
|
/// Print startup information for the serve command.
|
|
fn print_serve_startup_info(bind_addr: SocketAddr, secret: &str) {
|
|
eprintln!();
|
|
eprintln!(" Kigi agent server starting...");
|
|
eprintln!();
|
|
eprintln!(" Address: {}:{}", bind_addr.ip(), bind_addr.port());
|
|
eprintln!(" Secret: {}", secret);
|
|
eprintln!();
|
|
eprintln!(
|
|
" WebSocket URL: ws://{}/ws?server-key={}",
|
|
bind_addr, secret
|
|
);
|
|
eprintln!();
|
|
}
|
|
/// Entrypoint tag for `kigi -p`; keys the quiet stderr default in `init_tracing_simple`.
|
|
const HEADLESS_ENTRYPOINT: &str = "headless";
|
|
/// Initialize simple tracing for non-TUI agent modes.
|
|
fn init_tracing_simple(app_entrypoint: &'static str) {
|
|
use kigi_log::debug_log::RMCP_SSE_NOISE_TARGET;
|
|
use tracing_subscriber::{EnvFilter, Layer as _, fmt, layer::SubscriberExt as _};
|
|
let default_filter = if app_entrypoint == HEADLESS_ENTRYPOINT {
|
|
"off"
|
|
} else {
|
|
"error"
|
|
};
|
|
let env_filter = match EnvFilter::try_from_default_env() {
|
|
Ok(filter) => filter.add_directive(
|
|
format!("{RMCP_SSE_NOISE_TARGET}=error")
|
|
.parse()
|
|
.expect("static rmcp directive must parse"),
|
|
),
|
|
Err(_) => EnvFilter::new(default_filter),
|
|
};
|
|
let fmt_layer = fmt::layer()
|
|
.with_target(false)
|
|
.with_ansi(true)
|
|
.with_writer(std::io::stderr);
|
|
let registry = tracing_subscriber::registry()
|
|
.with(fmt_layer.with_filter(env_filter))
|
|
.with(kigi_log::sampling_log::layer())
|
|
.with(kigi_log::instrumentation::layer())
|
|
.with(kigi_log::hooks_log::layer());
|
|
kigi_log::debug_log::install_firehose(registry, app_entrypoint);
|
|
}
|
|
/// `kigi setup`: rendering + exit codes only; fetch logic lives in `kigi_shell::managed_config`.
|
|
/// `json` prints the served configuration instead of installing it.
|
|
async fn run_setup_command(json: bool) {
|
|
use kigi_shell::managed_config::{self, SetupOutcome};
|
|
if !managed_config::has_principal() {
|
|
eprintln!("No deployment key or team sign-in found.");
|
|
eprintln!();
|
|
eprintln!("To install managed configuration, sign in with a team using `kigi login`,");
|
|
eprintln!("or set a deployment key:");
|
|
eprintln!();
|
|
if cfg!(unix) {
|
|
eprintln!(" export KIGI_DEPLOYMENT_KEY=<your-key>");
|
|
} else {
|
|
eprintln!(" $env:KIGI_DEPLOYMENT_KEY=\"<your-key>\"");
|
|
}
|
|
eprintln!(" kigi setup");
|
|
eprintln!();
|
|
eprintln!("Or add the key to ~/.kigi/config.toml:");
|
|
eprintln!();
|
|
eprintln!(" [endpoints]");
|
|
eprintln!(" deployment_key = \"<your-key>\"");
|
|
eprintln!();
|
|
eprintln!(
|
|
"If you don't have a deployment key, contact your organization's Kigi administrator."
|
|
);
|
|
std::process::exit(1);
|
|
}
|
|
if json {
|
|
match managed_config::fetch_setup_report().await {
|
|
Ok(report) => {
|
|
let out = serde_json::to_string_pretty(&report)
|
|
.expect("setup report has no non-serializable values");
|
|
println!("{out}");
|
|
if !report.configured {
|
|
eprintln!(
|
|
"Your team doesn't have a managed configuration yet. Ask a team admin to provision one."
|
|
);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
eprintln!("Couldn't fetch managed configuration. {e}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
match managed_config::run_setup().await {
|
|
SetupOutcome::Installed => eprintln!("Applied managed configuration."),
|
|
SetupOutcome::NothingConfigured => {
|
|
eprintln!(
|
|
"Your team doesn't have a managed configuration yet. Ask a team admin to provision one."
|
|
);
|
|
}
|
|
SetupOutcome::Failed(e) => {
|
|
eprintln!("Couldn't apply managed configuration. {e}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
}
|
|
async fn run_leader_mgmt(args: LeaderMgmtArgs) -> Result<()> {
|
|
match args.command {
|
|
LeaderMgmtCommand::Kill => kill_leaders().await,
|
|
LeaderMgmtCommand::List { json } => {
|
|
let leaders = kigi_shell::leader::discover_leaders().await;
|
|
if json {
|
|
let payload: Vec<_> = leaders.iter().map(leader_descriptor_json).collect();
|
|
println!(
|
|
"{}",
|
|
serde_json::to_string(&serde_json::Value::Array(payload))?
|
|
);
|
|
} else if leaders.is_empty() {
|
|
println!("No leader candidates found.");
|
|
} else {
|
|
for d in &leaders {
|
|
print_leader_descriptor(d);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
LeaderMgmtCommand::Info { target, json } => {
|
|
let (descriptor, client) = connect_to_leader(&target).await?;
|
|
let info = match ensure_control_caps(client.registration()) {
|
|
Ok(_) => client
|
|
.send_control(ControlCommand::GetLeaderInfo)
|
|
.await
|
|
.ok()
|
|
.and_then(|r| r.ok()),
|
|
Err(_) => None,
|
|
};
|
|
if json {
|
|
let payload = leader_info_json(&descriptor, client.registration(), info.as_ref())?;
|
|
println!("{}", serde_json::to_string(&payload)?);
|
|
} else if let Some(info) = info {
|
|
println!("{info:#?}");
|
|
} else {
|
|
print_leader_descriptor(&descriptor);
|
|
eprintln!(
|
|
" (detailed info unavailable — leader does not advertise control capabilities)"
|
|
);
|
|
}
|
|
client.cancel();
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
async fn kill_leaders() -> Result<()> {
|
|
let leaders = kigi_shell::leader::discover_leaders().await;
|
|
if leaders.is_empty() {
|
|
eprintln!("No leader candidates found.");
|
|
return Ok(());
|
|
}
|
|
let mut killed = 0u32;
|
|
let mut cleaned = 0u32;
|
|
for d in &leaders {
|
|
let Some(pid) = leader_pid(d) else {
|
|
continue;
|
|
};
|
|
if !kigi_shell::util::is_kigi_process(pid) {
|
|
if let Some(ref lock) = d.lock_path {
|
|
eprintln!(" PID {pid} is not a kigi process, removing stale lock");
|
|
let _ = std::fs::remove_file(lock);
|
|
cleaned += 1;
|
|
}
|
|
if let Some(ref sock) = d.socket_path {
|
|
let _ = std::fs::remove_file(sock);
|
|
}
|
|
continue;
|
|
}
|
|
eprintln!(" Killing leader PID {pid}");
|
|
if let Err(e) = kigi_shell::util::kill_process_by_pid(pid) {
|
|
eprintln!(" warning: failed to terminate PID {pid}: {e}");
|
|
continue;
|
|
}
|
|
killed += 1;
|
|
}
|
|
if killed > 0 {
|
|
eprintln!("Killed {killed} leader process(es).");
|
|
} else if cleaned > 0 {
|
|
eprintln!("No live leader processes found (cleaned up {cleaned} stale lock(s)).");
|
|
} else {
|
|
eprintln!("No live leader processes found.");
|
|
}
|
|
Ok(())
|
|
}
|
|
fn resolve_target(args: &LeaderTargetArgs) -> LeaderTarget {
|
|
match args.pid {
|
|
Some(pid) => LeaderTarget::Pid(pid),
|
|
None => LeaderTarget::Default,
|
|
}
|
|
}
|
|
async fn connect_to_leader(
|
|
args: &LeaderTargetArgs,
|
|
) -> Result<(LeaderDescriptor, kigi_shell::leader::LeaderClient)> {
|
|
let target = resolve_target(args);
|
|
let selection = kigi_shell::leader::resolve_leader_target(target)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("{}", e.message))?;
|
|
let socket_path = selection
|
|
.socket_path()
|
|
.ok_or_else(|| anyhow::anyhow!("resolved leader target did not include a socket path"))?;
|
|
let client = kigi_shell::leader::LeaderClient::connect(
|
|
socket_path.to_path_buf(),
|
|
"kigi-pager-leader-cli",
|
|
ClientMode::Stdio,
|
|
ClientCapabilities::default(),
|
|
)
|
|
.await?;
|
|
Ok((selection.descriptor, client))
|
|
}
|
|
/// Prefer socket-verified live PID over a possibly-recycled lock file PID.
|
|
fn leader_pid(d: &LeaderDescriptor) -> Option<u32> {
|
|
d.live_info.as_ref().map(|li| li.pid).or(d.pid_from_lock)
|
|
}
|
|
fn print_leader_descriptor(d: &LeaderDescriptor) {
|
|
let pid = leader_pid(d)
|
|
.map(|p| p.to_string())
|
|
.unwrap_or_else(|| "?".into());
|
|
let sock = d
|
|
.socket_path
|
|
.as_deref()
|
|
.map(|p| p.display().to_string())
|
|
.unwrap_or_else(|| "?".into());
|
|
let state = format!("{:?}", d.classification);
|
|
eprintln!(" PID {pid} ({state}) -- {sock}");
|
|
}
|
|
fn leader_descriptor_json(d: &LeaderDescriptor) -> serde_json::Value {
|
|
serde_json::json!(
|
|
{ "pid" : leader_pid(d), "pidFromLock" : d.pid_from_lock, "pidLive" : d.live_info
|
|
.as_ref().map(| li | li.pid), "classification" : format!("{:?}", d
|
|
.classification), "socketPath" : d.socket_path.as_deref().map(| p | p.display()
|
|
.to_string()), "lockPath" : d.lock_path.as_deref().map(| p | p.display()
|
|
.to_string()), "socketSuffix" : d.socket_suffix, }
|
|
)
|
|
}
|
|
fn leader_info_json(
|
|
d: &LeaderDescriptor,
|
|
reg: &LeaderRegistration,
|
|
info: Option<&kigi_shell::leader::ControlPayload>,
|
|
) -> Result<serde_json::Value> {
|
|
let mut val = leader_descriptor_json(d);
|
|
val["clientId"] = serde_json::json!(reg.client_id);
|
|
if let Some(info) = info {
|
|
val["info"] = serde_json::to_value(info)?;
|
|
}
|
|
Ok(val)
|
|
}
|
|
fn ensure_control_caps(reg: &LeaderRegistration) -> Result<&LeaderCapabilities> {
|
|
reg.leader_capabilities
|
|
.as_ref()
|
|
.ok_or_else(|| anyhow::anyhow!("Leader does not advertise capabilities (legacy version)"))
|
|
}
|
|
/// How to rebuild one session's `session/load` after a leader reconnect.
|
|
#[derive(Default, Clone)]
|
|
struct CachedSession {
|
|
/// Verbatim `session/load` request JSON (preferred replay form: preserves
|
|
/// the client's exact cwd / mcpServers / meta). `None` when the session
|
|
/// was only ever created via `session/new` — the load is synthesized.
|
|
load_request_json: Option<String>,
|
|
/// `cwd` captured from `session/new` / `session/load` params.
|
|
cwd: Option<String>,
|
|
/// `mcpServers` captured from `session/new` / `session/load` params.
|
|
mcp_servers_json: Option<String>,
|
|
}
|
|
/// ACP state cached from the stdio stream for replay after leader reconnect.
|
|
///
|
|
/// Tracks EVERY session the external client has open (IDE clients drive
|
|
/// multiple sessions over one bridge), not just the most recent one — a
|
|
/// leader crash must restore all of them or the others die with
|
|
/// "unknown session id" on their next prompt.
|
|
#[derive(Default, Clone)]
|
|
struct StdioReplayState {
|
|
initialize_json: Option<String>,
|
|
/// Sessions to restore on reconnect, keyed by session id, in first-seen
|
|
/// order (Vec keeps replay order deterministic).
|
|
sessions: Vec<(String, CachedSession)>,
|
|
/// cwd/mcp from the most recent `session/new` REQUEST whose response has
|
|
/// not been observed yet. Folded into `sessions` when the response
|
|
/// carrying the assigned session id arrives. Never replayed while
|
|
/// unconfirmed (the id is unknown; the client's own request died with the
|
|
/// old leader and is its to retry).
|
|
pending_new: Option<CachedSession>,
|
|
/// Most recently created/loaded session id — reported in
|
|
/// `kigi/leader_reconnected` as the primary restored session.
|
|
last_session_id: Option<String>,
|
|
}
|
|
impl StdioReplayState {
|
|
fn upsert_session(&mut self, sid: &str, cached: CachedSession) {
|
|
if let Some((_, existing)) = self.sessions.iter_mut().find(|(id, _)| id == sid) {
|
|
*existing = cached;
|
|
} else {
|
|
self.sessions.push((sid.to_string(), cached));
|
|
}
|
|
}
|
|
fn remove_session(&mut self, sid: &str) {
|
|
self.sessions.retain(|(id, _)| id != sid);
|
|
if self.last_session_id.as_deref() == Some(sid) {
|
|
self.last_session_id = None;
|
|
}
|
|
}
|
|
}
|
|
fn cache_outgoing_acp_state(msg: &str, state: &std::sync::Mutex<StdioReplayState>) {
|
|
let Ok(json) = serde_json::from_str::<serde_json::Value>(msg) else {
|
|
return;
|
|
};
|
|
let method = json
|
|
.get("method")
|
|
.and_then(|m| m.as_str())
|
|
.unwrap_or_default();
|
|
match method {
|
|
"initialize" => {
|
|
let mut s = state.lock().unwrap_or_else(|e| e.into_inner());
|
|
s.initialize_json = Some(msg.to_string());
|
|
}
|
|
"session/load" => {
|
|
let mut s = state.lock().unwrap_or_else(|e| e.into_inner());
|
|
if let Some(params) = json.get("params") {
|
|
let sid = params
|
|
.get("sessionId")
|
|
.or_else(|| params.get("session_id"))
|
|
.and_then(|v| v.as_str());
|
|
if let Some(sid) = sid {
|
|
let cached = CachedSession {
|
|
load_request_json: Some(msg.to_string()),
|
|
cwd: params
|
|
.get("cwd")
|
|
.and_then(|v| v.as_str())
|
|
.map(str::to_string),
|
|
mcp_servers_json: params
|
|
.get("mcpServers")
|
|
.and_then(|m| serde_json::to_string(m).ok()),
|
|
};
|
|
s.upsert_session(sid, cached);
|
|
s.last_session_id = Some(sid.to_string());
|
|
}
|
|
}
|
|
}
|
|
"session/new" => {
|
|
let mut s = state.lock().unwrap_or_else(|e| e.into_inner());
|
|
let params = json.get("params");
|
|
s.pending_new = Some(CachedSession {
|
|
load_request_json: None,
|
|
cwd: params
|
|
.and_then(|p| p.get("cwd"))
|
|
.and_then(|v| v.as_str())
|
|
.map(str::to_string),
|
|
mcp_servers_json: params
|
|
.and_then(|p| p.get("mcpServers"))
|
|
.and_then(|m| serde_json::to_string(m).ok()),
|
|
});
|
|
}
|
|
"kigi/session/close" | "_kigi/session/close" => {
|
|
if let Some(sid) = json
|
|
.get("params")
|
|
.and_then(|p| p.get("sessionId").or_else(|| p.get("session_id")))
|
|
.and_then(|v| v.as_str())
|
|
{
|
|
let mut s = state.lock().unwrap_or_else(|e| e.into_inner());
|
|
s.remove_session(sid);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
fn cache_incoming_session_id(msg: &str, state: &std::sync::Mutex<StdioReplayState>) {
|
|
let Ok(json) = serde_json::from_str::<serde_json::Value>(msg) else {
|
|
return;
|
|
};
|
|
if let Some(sid) = json
|
|
.get("result")
|
|
.and_then(|r| r.get("sessionId").or_else(|| r.get("session_id")))
|
|
.and_then(|v| v.as_str())
|
|
{
|
|
let mut s = state.lock().unwrap_or_else(|e| e.into_inner());
|
|
if let Some(pending) = s.pending_new.take() {
|
|
s.upsert_session(sid, pending);
|
|
}
|
|
s.last_session_id = Some(sid.to_string());
|
|
}
|
|
}
|
|
/// Synthetic JSON-RPC id for the `session/load` the bridge constructs itself
|
|
/// (when the external client only ever sent `session/new`). A string id can
|
|
/// never collide with a numeric id the external client may have in flight.
|
|
const REPLAY_LOAD_REQUEST_ID: &str = "kigi/leader-replay/session-load";
|
|
/// Max silence between two messages from the leader during a replayed request.
|
|
/// A `session/load` streams replay notifications continuously once it starts,
|
|
/// but the pre-replay phase (MCP resolution, session file reads) can be quiet
|
|
/// for a while on large sessions.
|
|
const REPLAY_RECV_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
|
|
/// Overall deadline for one replayed request's response.
|
|
const REPLAY_RESPONSE_DEADLINE: std::time::Duration = std::time::Duration::from_secs(180);
|
|
/// Outcome of one replayed request (see [`replay_request_until_response`]).
|
|
enum ReplayOutcome {
|
|
/// The response for the replayed request arrived with a `result`.
|
|
ResponseOk,
|
|
/// The response arrived but carried an `error` (e.g. session not found).
|
|
ResponseErr,
|
|
/// The connection closed / timed out / stdout broke before the response.
|
|
Failed,
|
|
}
|
|
/// True when `msg` is the JSON-RPC *response* to the request with `expected_id`
|
|
/// (no `method` key + matching `id`).
|
|
fn parse_replay_response(msg: &str, expected_id: &serde_json::Value) -> Option<ReplayOutcome> {
|
|
let json = serde_json::from_str::<serde_json::Value>(msg).ok()?;
|
|
if json.get("method").is_some() {
|
|
return None;
|
|
}
|
|
if json.get("id") != Some(expected_id) {
|
|
return None;
|
|
}
|
|
if json.get("error").is_some() {
|
|
Some(ReplayOutcome::ResponseErr)
|
|
} else {
|
|
Some(ReplayOutcome::ResponseOk)
|
|
}
|
|
}
|
|
/// Send one replayed request to the (new) leader and pump messages until its
|
|
/// response arrives.
|
|
///
|
|
/// `session/load` emits the full replay stream (session/update notifications)
|
|
/// BEFORE its response, so "wait for the next message" is not "wait for the
|
|
/// response". Everything that is not the response itself is forwarded verbatim
|
|
/// to the external client's stdout — exactly what the pre-reconnect stream
|
|
/// would have carried. Only the response to the replayed request is swallowed
|
|
/// (the external client already received a response for its original send and
|
|
/// must not see a duplicate or unknown-id response).
|
|
///
|
|
/// Returning before the `session/load` response is the root cause of the
|
|
/// "unknown session id" failures after a leader crash: the bridge declared
|
|
/// the reconnect complete while the new leader was still loading the session,
|
|
/// and the client's next `session/prompt` raced (and lost against) the load.
|
|
async fn replay_request_until_response(
|
|
tx: &tokio::sync::mpsc::UnboundedSender<String>,
|
|
rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
|
|
stdout: &mut (impl tokio::io::AsyncWrite + Unpin),
|
|
request_json: &str,
|
|
what: &str,
|
|
) -> ReplayOutcome {
|
|
use tokio::io::AsyncWriteExt as _;
|
|
let expected_id = serde_json::from_str::<serde_json::Value>(request_json)
|
|
.ok()
|
|
.and_then(|v| v.get("id").cloned());
|
|
if tx.send(request_json.to_string()).is_err() {
|
|
tracing::warn!(what, "replay: failed to send request");
|
|
return ReplayOutcome::Failed;
|
|
}
|
|
let Some(expected_id) = expected_id else {
|
|
return ReplayOutcome::ResponseOk;
|
|
};
|
|
let deadline = tokio::time::Instant::now() + REPLAY_RESPONSE_DEADLINE;
|
|
loop {
|
|
let now = tokio::time::Instant::now();
|
|
if now >= deadline {
|
|
tracing::warn!(what, "replay: response deadline exceeded");
|
|
return ReplayOutcome::Failed;
|
|
}
|
|
let per_recv = REPLAY_RECV_TIMEOUT.min(deadline - now);
|
|
match tokio::time::timeout(per_recv, rx.recv()).await {
|
|
Ok(Some(msg)) => {
|
|
if let Some(outcome) = parse_replay_response(&msg, &expected_id) {
|
|
tracing::debug!(
|
|
what,
|
|
ok = matches!(outcome, ReplayOutcome::ResponseOk),
|
|
"replay: response received"
|
|
);
|
|
return outcome;
|
|
}
|
|
if stdout.write_all(msg.as_bytes()).await.is_err()
|
|
|| stdout.write_all(b"\n").await.is_err()
|
|
|| stdout.flush().await.is_err()
|
|
{
|
|
tracing::warn!(what, "replay: stdout closed while forwarding");
|
|
return ReplayOutcome::Failed;
|
|
}
|
|
}
|
|
Ok(None) => {
|
|
tracing::warn!(what, "replay: leader closed before response");
|
|
return ReplayOutcome::Failed;
|
|
}
|
|
Err(_) => {
|
|
tracing::warn!(what, "replay: timed out waiting for response");
|
|
return ReplayOutcome::Failed;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
/// Build the `session/load` JSON to replay for one cached session: the
|
|
/// verbatim client request when available, else a synthesized load from the
|
|
/// captured `session/new` parameters.
|
|
fn replay_load_json(sid: &str, cached: &CachedSession) -> Option<String> {
|
|
if let Some(ref verbatim) = cached.load_request_json {
|
|
return Some(verbatim.clone());
|
|
}
|
|
let cwd = cached.cwd.as_deref()?;
|
|
let mut params = serde_json::json!({ "sessionId" : sid, "cwd" : cwd, });
|
|
if let Some(ref mcp_raw) = cached.mcp_servers_json
|
|
&& let Ok(mcp_val) = serde_json::from_str::<serde_json::Value>(mcp_raw)
|
|
{
|
|
params["mcpServers"] = mcp_val;
|
|
}
|
|
Some(
|
|
serde_json::json!(
|
|
{ "jsonrpc" : "2.0", "id" : REPLAY_LOAD_REQUEST_ID, "method" :
|
|
"session/load", "params" : params, }
|
|
)
|
|
.to_string(),
|
|
)
|
|
}
|
|
/// Replay cached `initialize` + every cached `session/load` to a freshly
|
|
/// (re-)elected leader, blocking until the leader has actually finished
|
|
/// loading EACH session (loads are sent strictly sequentially, each awaiting
|
|
/// its response — the synthesized-id reuse relies on this ordering).
|
|
///
|
|
/// Returns the primary restored session id (the most recently active one,
|
|
/// falling back to any successfully restored session). `None` when there was
|
|
/// nothing to replay or every restore failed — callers emit
|
|
/// `kigi/leader_reconnected` with empty params in that case, signalling the
|
|
/// external client to re-establish state itself.
|
|
async fn replay_acp_state_after_reconnect(
|
|
tx: &tokio::sync::mpsc::UnboundedSender<String>,
|
|
rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
|
|
stdout: &mut (impl tokio::io::AsyncWrite + Unpin),
|
|
state: &StdioReplayState,
|
|
) -> Option<String> {
|
|
if let Some(ref init_json) = state.initialize_json {
|
|
match replay_request_until_response(tx, rx, stdout, init_json, "initialize").await {
|
|
ReplayOutcome::ResponseOk => {}
|
|
ReplayOutcome::ResponseErr => {
|
|
tracing::warn!("replay: initialize was rejected by new leader");
|
|
return None;
|
|
}
|
|
ReplayOutcome::Failed => return None,
|
|
}
|
|
} else {
|
|
tracing::debug!("replay: no cached initialize; skipping ACP replay");
|
|
return None;
|
|
}
|
|
if state.sessions.is_empty() {
|
|
tracing::debug!("replay: no sessions to replay");
|
|
return None;
|
|
}
|
|
let mut restored: Vec<String> = Vec::new();
|
|
for (sid, cached) in &state.sessions {
|
|
let Some(load_json) = replay_load_json(sid, cached) else {
|
|
tracing::warn!(
|
|
session_id = % sid, "replay: no way to rebuild session/load; skipping"
|
|
);
|
|
continue;
|
|
};
|
|
match replay_request_until_response(tx, rx, stdout, &load_json, "session/load").await {
|
|
ReplayOutcome::ResponseOk => {
|
|
tracing::info!(session_id = % sid, "replay: session restored");
|
|
restored.push(sid.clone());
|
|
}
|
|
ReplayOutcome::ResponseErr => {
|
|
tracing::warn!(
|
|
session_id = % sid, "replay: session/load was rejected by new leader"
|
|
);
|
|
}
|
|
ReplayOutcome::Failed => {
|
|
tracing::warn!(
|
|
session_id = % sid,
|
|
"replay: transport failure during session/load; aborting remaining replays"
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
state
|
|
.last_session_id
|
|
.as_ref()
|
|
.filter(|sid| restored.iter().any(|r| r == *sid))
|
|
.cloned()
|
|
.or_else(|| restored.last().cloned())
|
|
}
|
|
/// Flush observability, then exit. Used by the agent/headless signal handler.
|
|
///
|
|
/// Does NOT write terminal escape codes — agent mode never enables TUI modes.
|
|
/// The TUI has its own signal handler (`app::signal_handler`) that does the
|
|
/// full crossterm teardown.
|
|
fn shutdown_and_flush_telemetry(exit_code: i32) -> ! {
|
|
kigi_log::debug_log::flush();
|
|
std::process::exit(exit_code);
|
|
}
|
|
/// Emitted by both leader guards (server mode and leader-connect) so the two sites
|
|
/// can't drift.
|
|
const PLUGIN_DIR_LEADER_WARNING: &str = "kigi: --plugin-dir is ignored in leader mode; run with --no-leader to \
|
|
load per-process plugins";
|
|
/// Run the `agent` subcommand, dispatching to the appropriate mode.
|
|
async fn run_agent_command(
|
|
agent_args: Box<kigi_tui::app::AgentArgs>,
|
|
permission_mode_flag: Option<String>,
|
|
trust: bool,
|
|
no_auto_update: bool,
|
|
disable_web_search: bool,
|
|
update_config: &UpdateConfig,
|
|
) -> Result<()> {
|
|
let _signal_flush = tokio::spawn(async {
|
|
#[cfg(unix)]
|
|
{
|
|
use tokio::signal::unix::{SignalKind, signal};
|
|
let mut term = signal(SignalKind::terminate()).ok();
|
|
let mut hup = signal(SignalKind::hangup()).ok();
|
|
tokio::select! {
|
|
_ = tokio::signal::ctrl_c() => { shutdown_and_flush_telemetry(130); } _ =
|
|
async { if let Some(sig) = term.as_mut() { let _ = sig.recv(). await; }
|
|
else { std::future::pending::< () > (). await; } } => {
|
|
shutdown_and_flush_telemetry(143); } _ = async { if let Some(sig) = hup
|
|
.as_mut() { let _ = sig.recv(). await; } else { std::future::pending::<
|
|
() > (). await; } } => { shutdown_and_flush_telemetry(129); }
|
|
}
|
|
}
|
|
#[cfg(not(unix))]
|
|
{
|
|
if tokio::signal::ctrl_c().await.is_ok() {
|
|
shutdown_and_flush_telemetry(130);
|
|
}
|
|
}
|
|
});
|
|
init_tracing_simple("agent");
|
|
kigi_log::instrumentation::install_panic_hook();
|
|
if trust {
|
|
match std::env::current_dir() {
|
|
Ok(cwd) => kigi_shell::agent::folder_trust::grant_folder_trust(&cwd),
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
error = % e, "--trust: failed to resolve cwd; folder not trusted"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
// Fire-and-forget model-catalog warmup (nothing joins the handle now that
|
|
// the xAI settings fetch it used to carry is gone).
|
|
drop(kigi_shell::agent::models::start_early_prefetch(None));
|
|
kigi_shell::agent::mvp_agent::warm_async_http_client();
|
|
tokio::task::spawn_blocking(|| {});
|
|
let mode = agent_args.mode.clone().unwrap_or(AgentCmd::Stdio);
|
|
let is_stdio = matches!(mode, AgentCmd::Stdio);
|
|
let is_leader = matches!(mode, AgentCmd::Leader(_));
|
|
if !is_stdio && !is_leader {
|
|
eprintln!(
|
|
"Kigi v{}",
|
|
kigi_version::display_version_with_commit(
|
|
env!("VERSION_WITH_COMMIT"),
|
|
kigi_update::channel_label(),
|
|
)
|
|
);
|
|
if should_check_for_updates(no_auto_update) {
|
|
auto_update::run_update_if_available(
|
|
auto_update::UpdateRunMode::NonBlocking,
|
|
false,
|
|
update_config,
|
|
)
|
|
.await
|
|
.ok();
|
|
}
|
|
}
|
|
let raw_config = kigi_shell::config::load_effective_config()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {}", e))?;
|
|
let mut agent_config = AgentConfig::new_from_toml_cfg(&raw_config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {}", e))?;
|
|
agent_config.default_model_override = agent_args.model.clone();
|
|
agent_config.reasoning_effort_override = agent_args
|
|
.reasoning_effort
|
|
.as_deref()
|
|
.and_then(kigi_shell::sampling::types::parse_canonical_effort_token);
|
|
let launch_yolo = kigi_shell::util::config::effective_yolo_for_launch(
|
|
agent_args.yolo,
|
|
permission_mode_flag.as_deref(),
|
|
None,
|
|
);
|
|
if let Some(warning) = launch_yolo.blocked_warning {
|
|
eprintln!("kigi: {warning}");
|
|
}
|
|
agent_config.default_yolo_mode = launch_yolo.yolo;
|
|
agent_config.default_auto_mode = kigi_shell::util::config::effective_auto_for_launch(
|
|
agent_args.yolo,
|
|
permission_mode_flag.as_deref(),
|
|
None,
|
|
);
|
|
agent_config.agent_profile_path = agent_args
|
|
.agent_profile
|
|
.as_deref()
|
|
.map(resolve_agent_profile_path);
|
|
agent_config.client_version = Some(PAGER_CLIENT_VERSION.to_string());
|
|
if is_leader && !agent_args.plugin_dirs.is_empty() {
|
|
eprintln!("{PLUGIN_DIR_LEADER_WARNING}");
|
|
} else {
|
|
agent_config.plugins.cli_plugin_dirs = agent_args.canonical_plugin_dirs();
|
|
}
|
|
apply_agent_endpoint_args(&agent_args, &mut agent_config);
|
|
agent_config.resolve_runtime_fields(&kigi_shell::agent::config::RuntimeResolutionContext {
|
|
raw_config: &raw_config,
|
|
remote_settings: None,
|
|
cwd: None,
|
|
is_headless: !is_leader,
|
|
cli_subagents: None,
|
|
cli_session_summary_model: None,
|
|
cli_experimental_memory: false,
|
|
cli_no_memory: false,
|
|
disable_web_search,
|
|
todo_gate: false,
|
|
laziness_debug_log: None,
|
|
storage_mode: None,
|
|
});
|
|
let agent_memory_config = agent_config.memory_config.clone();
|
|
let leader_eligible = is_stdio;
|
|
let (use_leader, policy_disable_reason) = resolve_use_leader(
|
|
agent_args.leader,
|
|
agent_args.no_leader,
|
|
&raw_config,
|
|
leader_eligible,
|
|
);
|
|
tracing::info!(use_leader, ?policy_disable_reason, "leader mode resolved");
|
|
if stdio_direct_update_eligible(is_stdio, use_leader)
|
|
&& should_check_for_updates(no_auto_update)
|
|
{
|
|
let update_config = update_config.clone();
|
|
tokio::spawn(async move {
|
|
auto_update::run_update_if_available(
|
|
auto_update::UpdateRunMode::NonBlocking,
|
|
false,
|
|
&update_config,
|
|
)
|
|
.await
|
|
.ok();
|
|
});
|
|
}
|
|
if use_leader {
|
|
if !agent_args.plugin_dirs.is_empty() {
|
|
eprintln!("{PLUGIN_DIR_LEADER_WARNING}");
|
|
}
|
|
use kigi_shell::leader::{
|
|
ClientCapabilities, ClientMode, LeaderReconnector, ReconnectPolicy, connect_or_spawn,
|
|
};
|
|
use std::sync::Arc;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio::sync::Mutex as TokioMutex;
|
|
// Only `agent stdio` is leader-eligible (see `leader_eligible` above).
|
|
let mode = ClientMode::Stdio;
|
|
let default_model = agent_config
|
|
.default_model_override
|
|
.clone()
|
|
.or(agent_config.models.default.clone());
|
|
let client_type = std::env::args().collect::<Vec<_>>().join(" ");
|
|
let capabilities = ClientCapabilities {
|
|
yolo_mode: launch_yolo.yolo,
|
|
auto_mode: agent_config.default_auto_mode && !launch_yolo.yolo,
|
|
default_model,
|
|
client_version: Some(PAGER_CLIENT_VERSION.to_string()),
|
|
code_nav_enabled: false,
|
|
terminal: false,
|
|
fs_read: false,
|
|
fs_write: false,
|
|
};
|
|
let conn = connect_or_spawn(&client_type, mode, capabilities.clone()).await?;
|
|
let (tx, rx) = conn.into_channels();
|
|
let (status_tx, _status_rx) = LeaderReconnector::status_channel();
|
|
let reconnector = LeaderReconnector::new(&client_type, mode, capabilities, status_tx);
|
|
let cancel = CancellationToken::new();
|
|
{
|
|
{
|
|
let replay_state = Arc::new(std::sync::Mutex::new(StdioReplayState::default()));
|
|
let leader_tx = Arc::new(TokioMutex::new(tx));
|
|
let leader_tx_stdin = leader_tx.clone();
|
|
let replay_state_stdin = replay_state.clone();
|
|
let cancel_stdin = cancel.clone();
|
|
let stdin_task = tokio::spawn(async move {
|
|
let mut stdin_lines = kigi_acp_lib::spawn_stdin_line_reader();
|
|
loop {
|
|
tokio::select! {
|
|
biased; _ = cancel_stdin.cancelled() => break, maybe_line =
|
|
stdin_lines.recv() => { let Some(line) = maybe_line else {
|
|
break }; let line = String::from_utf8_lossy(& line); let
|
|
trimmed = line.trim_end_matches(['\r', '\n']).to_string(); if
|
|
trimmed.is_empty() { continue; }
|
|
if trimmed
|
|
.contains("\"initialize\"") || trimmed
|
|
.contains("\"session/load\"") || trimmed
|
|
.contains("\"session/new\"") { cache_outgoing_acp_state(&
|
|
trimmed, & replay_state_stdin); } let send_deadline =
|
|
tokio::time::Instant::now() +
|
|
std::time::Duration::from_secs(300); loop { { let tx =
|
|
leader_tx_stdin.lock(). await; if tx.send(trimmed.clone())
|
|
.is_ok() { break; } } if cancel_stdin.is_cancelled() ||
|
|
tokio::time::Instant::now() >= send_deadline {
|
|
tracing::error!("stdio bridge: dropping client message after \
|
|
reconnect retries were exhausted");
|
|
break; }
|
|
tokio::time::sleep(std::time::Duration::from_millis(250)).
|
|
await; } }
|
|
}
|
|
}
|
|
});
|
|
let cancel_stdout = cancel.clone();
|
|
let stdout_task = tokio::spawn(async move {
|
|
let mut stdout = tokio::io::stdout();
|
|
let mut rx = rx;
|
|
loop {
|
|
match rx.recv().await {
|
|
Some(ref msg) => {
|
|
if msg.contains("\"sessionId\"") || msg.contains("\"session_id\"") {
|
|
cache_incoming_session_id(msg, &replay_state);
|
|
}
|
|
if stdout.write_all(msg.as_bytes()).await.is_err()
|
|
|| stdout.write_all(b"\n").await.is_err()
|
|
|| stdout.flush().await.is_err()
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
None => {
|
|
tracing::warn!(
|
|
"Leader disconnected (stdio), attempting reconnect..."
|
|
);
|
|
match reconnector
|
|
.reconnect(ReconnectPolicy::bounded(), &cancel_stdout)
|
|
.await
|
|
{
|
|
Ok((new_tx, mut new_rx, _disconnect_rx)) => {
|
|
tracing::info!("Reconnected to leader (stdio)");
|
|
let replayed_session_id = {
|
|
let mut tx_guard = leader_tx.lock().await;
|
|
*tx_guard = new_tx;
|
|
let state = replay_state
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.clone();
|
|
replay_acp_state_after_reconnect(
|
|
&tx_guard,
|
|
&mut new_rx,
|
|
&mut stdout,
|
|
&state,
|
|
)
|
|
.await
|
|
};
|
|
rx = new_rx;
|
|
reconnector.notify_connected();
|
|
let params = match replayed_session_id {
|
|
Some(ref sid) => {
|
|
serde_json::json!({ "sessionId" : sid }).to_string()
|
|
}
|
|
None => "{}".to_string(),
|
|
};
|
|
let notification = format!(
|
|
r#"{{"jsonrpc":"2.0","method":"kigi/leader_reconnected","params":{params}}}"#
|
|
);
|
|
let _ = stdout.write_all(notification.as_bytes()).await;
|
|
let _ = stdout.write_all(b"\n").await;
|
|
let _ = stdout.flush().await;
|
|
continue;
|
|
}
|
|
Err(e) => {
|
|
tracing::error!(error = % e, "Failed to reconnect (stdio)");
|
|
cancel_stdout.cancel();
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
tokio::select! {
|
|
_ = stdin_task => {} _ = stdout_task => {}
|
|
}
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
match mode {
|
|
AgentCmd::Stdio => run_stdio_agent(&agent_config, None, agent_memory_config).await,
|
|
AgentCmd::Serve(a) => {
|
|
let secret = a.get_secret();
|
|
let server_config = kigi_shell::agent::ServerConfig {
|
|
bind_addr: a.bind,
|
|
secret: secret.clone(),
|
|
};
|
|
print_serve_startup_info(a.bind, &secret);
|
|
kigi_shell::agent::run_agent_server(server_config, agent_config).await
|
|
}
|
|
AgentCmd::Leader(a) => {
|
|
let leader_auto_update = if !should_check_for_updates(
|
|
no_auto_update || a.no_auto_update,
|
|
) {
|
|
tracing::info!("Leader auto-update disabled");
|
|
None
|
|
} else {
|
|
let update_config_for_leader = update_config.clone();
|
|
Some(kigi_shell::agent::app::LeaderAutoUpdateConfig {
|
|
check_interval: std::time::Duration::from_secs(60 * 60),
|
|
check_fn: Box::new(move || {
|
|
let uc = update_config_for_leader.clone();
|
|
Box::pin(async move {
|
|
let current_config = kigi_shell::util::config::load_config().await;
|
|
if current_config.cli.auto_update == Some(false) {
|
|
return false;
|
|
}
|
|
match auto_update::ensure_latest_on_disk(&uc).await {
|
|
Ok(outcome) => {
|
|
if let Some(v) = &outcome.installed {
|
|
if let Err(e) = kigi_shell::managed_config::sync().await {
|
|
tracing::warn!(
|
|
"Leader auto-update: managed config refresh failed: {e}"
|
|
);
|
|
}
|
|
tracing::info!(
|
|
"Leader auto-update: v{v} installed successfully"
|
|
);
|
|
} else if outcome.relaunch_needed {
|
|
tracing::info!(
|
|
"Leader auto-update: newer binary already on disk, \
|
|
relaunching without download"
|
|
);
|
|
}
|
|
outcome.relaunch_needed
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Leader auto-update: check/download failed, \
|
|
staying alive: {e:#}"
|
|
);
|
|
false
|
|
}
|
|
}
|
|
})
|
|
}),
|
|
})
|
|
};
|
|
run_leader(
|
|
&agent_config,
|
|
a.no_exit_on_disconnect,
|
|
leader_auto_update,
|
|
agent_memory_config,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
}
|
|
/// Raise the per-process file descriptor soft limit on macOS.
|
|
///
|
|
/// macOS has a conservative default soft `RLIMIT_NOFILE` (256) that is easily
|
|
/// exceeded by parallel directory walking + file copying in worktree creation,
|
|
/// stdio MCP servers, tool subprocesses, and async runtime sockets.
|
|
///
|
|
/// We raise the soft limit toward the hard limit, capped at 8192 to stay below
|
|
/// `FD_SETSIZE` (1024 on macOS) safety boundaries in any C dependency that may
|
|
/// still use `select(2)` -- Rust std + tokio use `kqueue`, but vendored C code
|
|
/// can corrupt the stack if it select()'s on an fd >= FD_SETSIZE. 8192 also
|
|
/// keeps fork-time fd-table iteration cheap for any child that does
|
|
/// "close all fds up to rlim_cur" on exec.
|
|
///
|
|
/// Best-effort: silently ignores all errors (process limits can be tightened by
|
|
/// containers/cgroups and we should never block startup on a non-essential
|
|
/// optimization).
|
|
#[cfg(target_os = "macos")]
|
|
fn raise_fd_limit() {
|
|
const TARGET: libc::rlim_t = 8192;
|
|
unsafe {
|
|
let mut rlim = libc::rlimit {
|
|
rlim_cur: 0,
|
|
rlim_max: 0,
|
|
};
|
|
if libc::getrlimit(libc::RLIMIT_NOFILE, &mut rlim) != 0 {
|
|
return;
|
|
}
|
|
let new_cur = rlim.rlim_max.min(TARGET);
|
|
if new_cur <= rlim.rlim_cur {
|
|
return;
|
|
}
|
|
let old = rlim.rlim_cur;
|
|
rlim.rlim_cur = new_cur;
|
|
if libc::setrlimit(libc::RLIMIT_NOFILE, &rlim) == 0 {
|
|
tracing::trace!(old, new = new_cur, "raised RLIMIT_NOFILE");
|
|
}
|
|
}
|
|
}
|
|
#[cfg(not(target_os = "macos"))]
|
|
fn raise_fd_limit() {}
|
|
/// Single audit point for the `Command::Dashboard` soft-subcommand.
|
|
/// Sets `KIGI_OPEN_DASHBOARD_AT_STARTUP=1` if the user asked for
|
|
/// `kigi dashboard`, and clears `args.command` so the regular
|
|
/// subcommand match doesn't try to handle it.
|
|
///
|
|
/// The dashboard is independent of leader mode — it renders local
|
|
/// sessions and, when a leader happens to be present, additionally shows
|
|
/// the leader roster — so `kigi dashboard` does NOT force leader mode and
|
|
/// is compatible with `--no-leader`.
|
|
///
|
|
/// The only gate is the feature flag: a disabled dashboard
|
|
/// (`[dashboard].enabled = false` / `KIGI_AGENT_DASHBOARD=0`) is a CLI
|
|
/// error here, before the TUI starts, because the welcome view silently
|
|
/// drops the equivalent runtime toast.
|
|
fn flag_dashboard_at_startup_if_requested(args: &mut PagerArgs) -> Result<()> {
|
|
if !matches!(args.command, Some(Command::Dashboard)) {
|
|
return Ok(());
|
|
}
|
|
if !kigi_tui::views::dashboard::dashboard_enabled() {
|
|
anyhow::bail!(
|
|
"the Agent Dashboard is disabled. Enable it by removing \
|
|
`[dashboard] enabled = false` from ~/.kigi/config.toml and \
|
|
unsetting KIGI_AGENT_DASHBOARD=0."
|
|
);
|
|
}
|
|
args.command = None;
|
|
unsafe { std::env::set_var("KIGI_OPEN_DASHBOARD_AT_STARTUP", "1") };
|
|
Ok(())
|
|
}
|
|
const RUNTIME_SHUTDOWN_GRACE: std::time::Duration = std::time::Duration::from_secs(2);
|
|
/// A plain runtime drop blocks forever on an uncancellable in-flight blocking
|
|
/// task; `shutdown_timeout` abandons it after `grace` so exit can't hang.
|
|
fn run_and_shutdown<F: std::future::Future>(
|
|
runtime: tokio::runtime::Runtime,
|
|
fut: F,
|
|
grace: std::time::Duration,
|
|
) -> F::Output {
|
|
let output = runtime.block_on(fut);
|
|
runtime.shutdown_timeout(grace);
|
|
output
|
|
}
|
|
/// Return freed-but-retained jemalloc pages to the OS.
|
|
///
|
|
/// `arena.<MALLCTL_ARENAS_ALL>.purge` madvises away all dirty/muzzy pages in
|
|
/// every arena. The pager invokes this (via the `memory_release` seam) right
|
|
/// after known memory cliffs — e.g. dropping a session load's replay
|
|
/// transient — so a long-session resume doesn't leave hundreds of MB of dead
|
|
/// pages counted against the process for its lifetime (macOS keeps
|
|
/// `MADV_FREE`d pages in RSS until systemwide pressure).
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
fn purge_jemalloc_retained_pages() {
|
|
static NAME: &[u8] = b"arena.4096.purge\0";
|
|
let ret = unsafe {
|
|
tikv_jemalloc_sys::mallctl(
|
|
NAME.as_ptr().cast(),
|
|
std::ptr::null_mut(),
|
|
std::ptr::null_mut(),
|
|
std::ptr::null_mut(),
|
|
0,
|
|
)
|
|
};
|
|
if ret != 0 {
|
|
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
|
|
WARN_ONCE.call_once(|| {
|
|
tracing::warn!(
|
|
errno = ret,
|
|
"jemalloc arena purge mallctl failed; retained-page release is inert"
|
|
);
|
|
});
|
|
}
|
|
}
|
|
/// Allocator gauges for the memory trace (`memory_trace` seam): advance the
|
|
/// jemalloc epoch so the `stats.*` reads are current, then read each gauge.
|
|
/// Returns `None` if any mallctl fails (trace records the absence). Rides
|
|
/// the `tikv-jemalloc-ctl` raw helpers (introduced by the heap-profile
|
|
/// hooks below) instead of hand-rolled mallctl.
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
fn jemalloc_allocator_stats() -> Option<kigi_tui::memory_trace::AllocatorStats> {
|
|
/// SAFETY: callers pass fixed NUL-terminated `stats.*` size_t ctl names.
|
|
unsafe fn gauge(name: &[u8]) -> Option<u64> {
|
|
unsafe {
|
|
tikv_jemalloc_ctl::raw::read::<usize>(name)
|
|
.ok()
|
|
.map(|v| v as u64)
|
|
}
|
|
}
|
|
unsafe {
|
|
tikv_jemalloc_ctl::raw::write(b"epoch\0", 1u64).ok()?;
|
|
Some(kigi_tui::memory_trace::AllocatorStats {
|
|
allocated: gauge(b"stats.allocated\0")?,
|
|
active: gauge(b"stats.active\0")?,
|
|
resident: gauge(b"stats.resident\0")?,
|
|
mapped: gauge(b"stats.mapped\0")?,
|
|
retained: gauge(b"stats.retained\0")?,
|
|
metadata: gauge(b"stats.metadata\0")?,
|
|
})
|
|
}
|
|
}
|
|
/// Full jemalloc statistics dump for threshold snapshots
|
|
/// (`malloc_stats_print` default human-readable format, arena detail
|
|
/// included) — the artifact the GCS memory-trace upload ships for offline
|
|
/// analysis. Raw `tikv_jemalloc_sys` because jemalloc-ctl has no
|
|
/// callback-form stats_print.
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
fn jemalloc_stats_dump() -> String {
|
|
unsafe extern "C" fn append(opaque: *mut std::ffi::c_void, msg: *const std::ffi::c_char) {
|
|
unsafe {
|
|
let out = &mut *opaque.cast::<String>();
|
|
out.push_str(&std::ffi::CStr::from_ptr(msg).to_string_lossy());
|
|
}
|
|
}
|
|
let mut out = String::new();
|
|
unsafe {
|
|
tikv_jemalloc_sys::malloc_stats_print(
|
|
Some(append),
|
|
(&raw mut out).cast(),
|
|
std::ptr::null(),
|
|
);
|
|
}
|
|
out
|
|
}
|
|
fn main() {
|
|
kigi_pager_minimal::install();
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
kigi_tui::memory_release::install_release_hook(purge_jemalloc_retained_pages);
|
|
#[cfg(all(feature = "jemalloc", unix))]
|
|
{
|
|
kigi_tui::memory_trace::install_allocator_stats_provider(jemalloc_allocator_stats);
|
|
kigi_tui::memory_trace::install_allocator_dump_provider(jemalloc_stats_dump);
|
|
}
|
|
if let Some(code) = kigi_tui::app::mermaid_worker::maybe_run_render_subprocess() {
|
|
std::process::exit(code);
|
|
}
|
|
kigi_tui::memory_trace::start(kigi_shell::util::kigi_home::kigi_home().join("memtrace"));
|
|
raise_fd_limit();
|
|
if let Err(e) = kigi_config::validate_requirements() {
|
|
eprintln!("Couldn't start Kigi: {e}");
|
|
eprintln!();
|
|
eprintln!(
|
|
"Update Kigi to a version the policy allows, or ask your administrator \
|
|
to fix the managed requirements."
|
|
);
|
|
std::process::exit(2);
|
|
}
|
|
kigi_tui::docs::extract_user_guide_docs(&kigi_shell::util::kigi_home::kigi_home());
|
|
kigi_crash_handler::install_terminal_restore_only();
|
|
if kigi_shell::util::config::load_crash_handler_enabled_sync() {
|
|
let crash_dir = kigi_shell::util::kigi_home::kigi_home().join("crash");
|
|
if let Some(report) = kigi_crash_handler::check_previous_crash(&crash_dir) {
|
|
eprintln!("Kigi crashed during your last session.");
|
|
eprintln!(" Signal: {}", report.signal_name);
|
|
eprintln!(" Version: {}", report.app_version);
|
|
eprintln!(" Report: {}", report.report_path.display());
|
|
eprintln!();
|
|
}
|
|
if !kigi_crash_handler::install(kigi_crash_handler::CrashHandlerConfig {
|
|
app_version: env!("VERSION_WITH_COMMIT").to_string(),
|
|
crash_dir: crash_dir.clone(),
|
|
}) {
|
|
eprintln!(
|
|
"warning: crash handler enabled but failed to install (check permissions on {})",
|
|
crash_dir.display()
|
|
);
|
|
}
|
|
}
|
|
let crashed = kigi_shell::active_sessions::collect_crashed().unwrap_or_default();
|
|
if !crashed.is_empty() {
|
|
tracing::info!(
|
|
count = crashed.len(),
|
|
"Found crashed sessions from a previous run"
|
|
);
|
|
}
|
|
let runtime = tokio::runtime::Builder::new_multi_thread()
|
|
.enable_all()
|
|
.build()
|
|
.unwrap_or_else(|e| panic!("failed to start tokio runtime: {e}"));
|
|
let result = run_and_shutdown(runtime, async_main(), RUNTIME_SHUTDOWN_GRACE);
|
|
kigi_log::debug_log::flush();
|
|
if let Err(e) = result {
|
|
kigi_tty_utils::restore_native_stderr();
|
|
eprintln!("Error: {e:#}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
async fn async_main() -> Result<()> {
|
|
let _ = rustls::crypto::ring::default_provider().install_default();
|
|
let mut args = PagerArgs::parse_and_apply_cwd()?;
|
|
if let Some(ref mode) = args.compaction_mode {
|
|
unsafe { std::env::set_var("KIGI_COMPACTION_MODE", mode) };
|
|
}
|
|
if let Some(ref detail) = args.compaction_detail {
|
|
unsafe { std::env::set_var("KIGI_COMPACTION_DETAIL", detail) };
|
|
}
|
|
if args.chat() {
|
|
anyhow::bail!(
|
|
"--chat is no longer supported: the kimi.com chat frontend it drove was \
|
|
removed along with the xAI backend."
|
|
);
|
|
}
|
|
if let Some(ref socket) = args.leader_socket {
|
|
unsafe { std::env::set_var(kigi_shell::leader::LEADER_SOCKET_ENV, socket) };
|
|
}
|
|
if !args.mcp_config_file.is_empty() {
|
|
// Fail fast on unusable files: an explicitly passed config that
|
|
// cannot be read must abort, not silently degrade the session.
|
|
for path in &args.mcp_config_file {
|
|
let content = std::fs::read_to_string(path)
|
|
.map_err(|e| anyhow::anyhow!("--mcp-config-file {}: {e}", path.display()))?;
|
|
serde_json::from_str::<serde_json::Value>(&content).map_err(|e| {
|
|
anyhow::anyhow!("--mcp-config-file {}: invalid JSON: {e}", path.display())
|
|
})?;
|
|
}
|
|
let joined = std::env::join_paths(&args.mcp_config_file).map_err(|e| {
|
|
anyhow::anyhow!("--mcp-config-file paths cannot be joined into the env: {e}")
|
|
})?;
|
|
unsafe { std::env::set_var(kigi_shell::util::config::MCP_CONFIG_FILES_ENV, joined) };
|
|
}
|
|
if let Some(ref path) = args.debug_file {
|
|
unsafe {
|
|
std::env::set_var("KIGI_DEBUG_LOG", path);
|
|
std::env::remove_var("KIGI_LOG_FILE");
|
|
}
|
|
}
|
|
if args.debug || args.debug_file.is_some() {
|
|
let set_if_unset = |k: &str, v: &str| {
|
|
if std::env::var_os(k).is_none() {
|
|
unsafe { std::env::set_var(k, v) };
|
|
}
|
|
};
|
|
set_if_unset("KIGI_DEBUG_LOG", "1");
|
|
set_if_unset("KIGI_HOOKS_LOG", "1");
|
|
}
|
|
if let Some(Command::Completions { shell }) = &args.command {
|
|
kigi_tui::completions_cmd::run(*shell);
|
|
return Ok(());
|
|
}
|
|
if let Some(Command::Wrap(ref wrap_args)) = args.command {
|
|
return kigi_tui::wrap_cmd::run(wrap_args);
|
|
}
|
|
let saved_profile = args.saved_resume_profile();
|
|
let sandbox_profile_arg = match args.startup_sandbox_profile(saved_profile.as_deref()) {
|
|
kigi_tui::app::cli::SandboxStartup::Apply(profile) => profile,
|
|
kigi_tui::app::cli::SandboxStartup::Conflict { requested, saved } => {
|
|
eprintln!(
|
|
"error: cannot resume this session under sandbox profile '{requested}' — \
|
|
it was created with '{saved}'. Omit --sandbox to resume with '{saved}', \
|
|
or start a new session to use '{requested}'."
|
|
);
|
|
std::process::exit(1);
|
|
}
|
|
};
|
|
kigi_shell::config::apply_sandbox(None, sandbox_profile_arg.as_deref(), args.cwd.as_deref());
|
|
flag_dashboard_at_startup_if_requested(&mut args)?;
|
|
let is_interactive = args.command.is_none()
|
|
&& args.single.is_none()
|
|
&& args.prompt_json.is_none()
|
|
&& args.prompt_file.is_none();
|
|
kigi_shell::http::set_client_name(if is_interactive {
|
|
kigi_workspace::permission::ClientType::KigiPager
|
|
} else {
|
|
kigi_workspace::permission::ClientType::Generic
|
|
});
|
|
let update_config = build_update_config();
|
|
if let Some(command) = args.command.take() {
|
|
match command {
|
|
Command::Version { json } => {
|
|
if json {
|
|
let payload = serde_json::json!(
|
|
{ "currentVersion" : env!("VERSION_WITH_COMMIT"), "channel" :
|
|
kigi_update::channel_name().unwrap_or("unknown"), }
|
|
);
|
|
println!("{}", serde_json::to_string(&payload)?);
|
|
} else {
|
|
println!(
|
|
"kigi {}",
|
|
kigi_version::display_version_with_commit(
|
|
env!("VERSION_WITH_COMMIT"),
|
|
kigi_update::channel_label(),
|
|
)
|
|
);
|
|
}
|
|
return Ok(());
|
|
}
|
|
Command::Acp(agent_args) => {
|
|
if let Some(mode) = &agent_args.mode {
|
|
anyhow::bail!(
|
|
"`kigi acp` always runs the stdio server; `{mode:?}` does not apply. \
|
|
Use `kigi agent ...` for other runtime modes."
|
|
);
|
|
}
|
|
enforce_minimum_version_or_exit(&update_config).await;
|
|
return run_agent_command(
|
|
agent_args,
|
|
args.permission_mode_flag.clone(),
|
|
args.trust,
|
|
args.no_auto_update,
|
|
args.disable_web_search,
|
|
&update_config,
|
|
)
|
|
.await;
|
|
}
|
|
Command::Agent(agent_args) => {
|
|
if args.leader || args.no_leader {
|
|
let flag = if args.leader {
|
|
"--leader"
|
|
} else {
|
|
"--no-leader"
|
|
};
|
|
anyhow::bail!(
|
|
"top-level {flag} applies to the pager TUI, not the agent subcommand. \
|
|
Use `kigi-pager agent {flag}` instead."
|
|
);
|
|
}
|
|
enforce_minimum_version_or_exit(&update_config).await;
|
|
return run_agent_command(
|
|
agent_args,
|
|
args.permission_mode_flag.clone(),
|
|
args.trust,
|
|
args.no_auto_update,
|
|
args.disable_web_search,
|
|
&update_config,
|
|
)
|
|
.await;
|
|
}
|
|
Command::Inspect { json } => {
|
|
let cwd = std::env::current_dir().unwrap_or_default();
|
|
kigi_shell::inspect::inspect(&cwd, json).await?;
|
|
return Ok(());
|
|
}
|
|
Command::ImportKimi { dry_run } => {
|
|
init_tracing_simple("cli");
|
|
return kigi_tui::import_kimi_cmd::run(dry_run);
|
|
}
|
|
Command::Setup { json } => {
|
|
init_tracing_simple("cli");
|
|
run_setup_command(json).await;
|
|
return Ok(());
|
|
}
|
|
Command::Mcp(mcp_args) => {
|
|
init_tracing_simple("cli");
|
|
return kigi_tui::mcp_cmd::run(mcp_args).await;
|
|
}
|
|
Command::Plugin(plugin_args) => {
|
|
init_tracing_simple("cli");
|
|
return kigi_tui::plugin_cmd::run(plugin_args).await;
|
|
}
|
|
Command::Models => {
|
|
init_tracing_simple("cli");
|
|
let config = kigi_shell::config::load_effective_config_disk_only()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
|
|
let agent_config = AgentConfig::new_from_toml_cfg(&config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
|
|
return kigi_tui::models::list_available_models(&agent_config).await;
|
|
}
|
|
Command::Leader(leader_args) => {
|
|
init_tracing_simple("cli");
|
|
return run_leader_mgmt(leader_args).await;
|
|
}
|
|
Command::Worktree(worktree_args) => {
|
|
init_tracing_simple("cli");
|
|
let config = kigi_shell::config::load_effective_config_disk_only()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
|
|
let agent_config = AgentConfig::new_from_toml_cfg(&config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
|
|
return kigi_tui::worktree_cmd::run(worktree_args, &agent_config).await;
|
|
}
|
|
Command::Sessions(sessions_args) => {
|
|
init_tracing_simple("cli");
|
|
return kigi_tui::sessions_cmd::run(sessions_args).await;
|
|
}
|
|
Command::Export(export_args) => {
|
|
init_tracing_simple("cli");
|
|
return kigi_tui::export_cmd::run(export_args);
|
|
}
|
|
Command::Trace(trace_args) => {
|
|
init_tracing_simple("cli");
|
|
let config = kigi_shell::config::load_effective_config_disk_only()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
|
|
let agent_config = AgentConfig::new_from_toml_cfg(&config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
|
|
return kigi_tui::trace_cmd::run(trace_args, &agent_config).await;
|
|
}
|
|
Command::Memory(memory_args) => {
|
|
return kigi_tui::memory_cmd::run(memory_args);
|
|
}
|
|
Command::Update {
|
|
check,
|
|
json,
|
|
force_reinstall,
|
|
version,
|
|
alpha,
|
|
stable,
|
|
enterprise,
|
|
} => {
|
|
init_tracing_simple("cli");
|
|
let channel_switch = get_channel_switch(alpha, stable, enterprise);
|
|
return run_update_command(
|
|
check,
|
|
json,
|
|
force_reinstall,
|
|
version,
|
|
channel_switch,
|
|
&update_config,
|
|
)
|
|
.await;
|
|
}
|
|
Command::Login => {
|
|
init_tracing_simple("cli");
|
|
let config = kigi_shell::config::load_effective_config_disk_only()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
|
|
let config = AgentConfig::new_from_toml_cfg(&config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
|
|
kigi_shell::auth::run_cli_login(&config).await?;
|
|
println!();
|
|
kigi_shell::instrumentation::finalize_and_exit(0);
|
|
}
|
|
Command::Logout => {
|
|
init_tracing_simple("cli");
|
|
let config = kigi_shell::config::load_effective_config_disk_only()
|
|
.map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?;
|
|
let config = AgentConfig::new_from_toml_cfg(&config)
|
|
.map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?;
|
|
kigi_shell::auth::run_cli_logout(&config)?;
|
|
kigi_shell::instrumentation::finalize_and_exit(0);
|
|
}
|
|
Command::Wrap(ref wrap_args) => {
|
|
return kigi_tui::wrap_cmd::run(wrap_args);
|
|
}
|
|
Command::Completions { shell } => {
|
|
kigi_tui::completions_cmd::run(shell);
|
|
return Ok(());
|
|
}
|
|
Command::Dashboard => {
|
|
args.command = Some(Command::Dashboard);
|
|
flag_dashboard_at_startup_if_requested(&mut args)?;
|
|
}
|
|
}
|
|
}
|
|
let headless_prompt = kigi_tui::headless::HeadlessPrompt::from_args(
|
|
args.single.as_deref(),
|
|
args.prompt_json.as_deref(),
|
|
args.prompt_file.as_deref(),
|
|
)?;
|
|
if let Some(prompt) = headless_prompt {
|
|
init_tracing_simple(HEADLESS_ENTRYPOINT);
|
|
enforce_minimum_version_or_exit(&update_config).await;
|
|
let launch_yolo = kigi_shell::util::config::effective_yolo_for_launch(
|
|
args.yolo,
|
|
args.permission_mode_flag.as_deref(),
|
|
None,
|
|
);
|
|
if let Some(warning) = launch_yolo.blocked_warning {
|
|
eprintln!("kigi: {warning}");
|
|
}
|
|
let json_schema = args
|
|
.json_schema
|
|
.as_deref()
|
|
.map(kigi_tui::headless::parse_json_schema)
|
|
.transpose()?;
|
|
if json_schema.is_some() {
|
|
if args.output_format == kigi_tui::headless::OutputFormat::Plain {
|
|
args.output_format = kigi_tui::headless::OutputFormat::Json;
|
|
}
|
|
if args.self_verify {
|
|
anyhow::bail!(
|
|
"--json-schema and --self-verify cannot be used together: \
|
|
verification output would corrupt the structured response"
|
|
);
|
|
}
|
|
}
|
|
return kigi_tui::headless::run_single_turn(
|
|
prompt,
|
|
args.verbatim,
|
|
kigi_tui::headless::HeadlessOptions {
|
|
session_id: args.session_id.clone(),
|
|
resume: args.resume_session.or(args.load_session),
|
|
cwd: args.cwd,
|
|
yolo: launch_yolo.yolo,
|
|
trust: args.trust,
|
|
output_format: args.output_format,
|
|
json_schema,
|
|
model: args.model,
|
|
rules: args.rules,
|
|
system_prompt_override: args.system_prompt_override.clone(),
|
|
continue_last_session: args.continue_last_session,
|
|
fork_session: args.fork_session,
|
|
worktree: args.worktree,
|
|
restore_code: args.restore_code,
|
|
agent: args.agent.clone(),
|
|
agents_json: args.agents_json.clone(),
|
|
cli_tools: args.cli_tools.clone(),
|
|
cli_disallowed_tools: args.cli_disallowed_tools.clone(),
|
|
disable_web_search: args.disable_web_search,
|
|
allow_rules: args.allow_rules.clone(),
|
|
deny_rules: args.deny_rules.clone(),
|
|
max_turns: args.max_turns,
|
|
permission_mode_flag: args.permission_mode_flag.clone(),
|
|
reasoning_effort: args.reasoning_effort.clone(),
|
|
self_verify: args.self_verify,
|
|
best_of_n: args.best_of_n,
|
|
wait_for_background: !args.no_wait_for_background,
|
|
background_wait_timeout: std::time::Duration::from_secs(
|
|
args.background_wait_timeout_secs,
|
|
),
|
|
},
|
|
)
|
|
.await;
|
|
}
|
|
enforce_minimum_version_or_exit(&update_config).await;
|
|
type UpdateWaitHandle = tokio::task::JoinHandle<std::io::Result<std::process::ExitStatus>>;
|
|
let bg_update_wait: std::sync::Arc<tokio::sync::Mutex<Option<UpdateWaitHandle>>> =
|
|
std::sync::Arc::new(tokio::sync::Mutex::new(None));
|
|
let bg_update_rx: Option<tokio::sync::oneshot::Receiver<Option<auto_update::UpdateAvailable>>> =
|
|
if should_check_for_updates(args.no_auto_update) {
|
|
let update_config = update_config.clone();
|
|
let wait_slot = bg_update_wait.clone();
|
|
let (tx, rx) = tokio::sync::oneshot::channel();
|
|
tokio::spawn(async move {
|
|
let check = auto_update::check_update_background(&update_config).await;
|
|
if let Some(mut child) = check.download {
|
|
*wait_slot.lock().await = Some(tokio::spawn(async move { child.wait().await }));
|
|
}
|
|
let _ = tx.send(check.update);
|
|
});
|
|
Some(rx)
|
|
} else {
|
|
None
|
|
};
|
|
let result = kigi_tui::app::run(args, bg_update_rx).await;
|
|
kigi_sandbox::flush();
|
|
match result {
|
|
Ok(true) => {
|
|
let adopted = bg_update_wait.lock().await.take();
|
|
if finish_update_on_exit(adopted, &update_config).await {
|
|
eprintln!("Update installed. Run `kigi` to start.");
|
|
} else {
|
|
eprintln!("Update did not complete. Run `kigi update` to retry.");
|
|
}
|
|
Ok(())
|
|
}
|
|
Ok(false) => Ok(()),
|
|
Err(e) => Err(e),
|
|
}
|
|
}
|
|
/// Complete the update after a quit-for-update (Ctrl+U) exit. Returns `true`
|
|
/// when an update path completed without a reported failure.
|
|
///
|
|
/// Prefers awaiting the parked waiter for the background `kigi update` child
|
|
/// spawned at startup — the download is usually already done or in flight.
|
|
/// Only when there is no waiter (spawn failed, or no download was needed
|
|
/// because the target was already on disk) or the child failed does this
|
|
/// fall back to a fresh blocking `kigi update`, which itself resolves to
|
|
/// "Already up to date" without downloading when the disk is current.
|
|
async fn finish_update_on_exit(
|
|
adopted: Option<tokio::task::JoinHandle<std::io::Result<std::process::ExitStatus>>>,
|
|
update_config: &UpdateConfig,
|
|
) -> bool {
|
|
let run_blocking = |reason: Option<String>| async move {
|
|
if let Some(reason) = reason {
|
|
eprintln!("{reason}");
|
|
}
|
|
auto_update::run_update_if_available(
|
|
auto_update::UpdateRunMode::Blocking,
|
|
false,
|
|
update_config,
|
|
)
|
|
.await
|
|
.is_ok()
|
|
};
|
|
match adopted {
|
|
Some(handle) => {
|
|
eprintln!("Waiting for the update download to finish...");
|
|
match handle.await {
|
|
Ok(Ok(status)) if status.success() => true,
|
|
Ok(Ok(status)) => {
|
|
run_blocking(Some(format!(
|
|
"Background update exited with {status}; retrying..."
|
|
)))
|
|
.await
|
|
}
|
|
Ok(Err(e)) => {
|
|
run_blocking(Some(format!(
|
|
"Could not wait for the background update ({e}); retrying..."
|
|
)))
|
|
.await
|
|
}
|
|
Err(join_err) => {
|
|
run_blocking(Some(format!(
|
|
"Background update waiter failed ({join_err}); retrying..."
|
|
)))
|
|
.await
|
|
}
|
|
}
|
|
}
|
|
None => run_blocking(None).await,
|
|
}
|
|
}
|
|
/// Build an [`UpdateConfig`] from the current environment and config files.
|
|
fn build_update_config() -> UpdateConfig {
|
|
let mut config = UpdateConfig::from_environment();
|
|
cryptify::flow_stmt!({
|
|
{
|
|
config.deployment_key =
|
|
kigi_shell::agent::config::EndpointsConfig::default().deployment_key;
|
|
}
|
|
});
|
|
if let Ok(root) = kigi_shell::config::load_effective_config_disk_only()
|
|
&& let Some(ch) = kigi_shell::util::config::channel_from_toml_opt(&root)
|
|
{
|
|
config.channel = ch;
|
|
}
|
|
config
|
|
}
|
|
/// Centralized gate for all auto-update checks. Add new suppression
|
|
/// rules here — not at each call site.
|
|
fn should_check_for_updates(no_auto_update_flag: bool) -> bool {
|
|
if cfg!(debug_assertions) {
|
|
return false;
|
|
}
|
|
if no_auto_update_flag {
|
|
return false;
|
|
}
|
|
if std::env::var_os("KIGI_DISABLE_AUTOUPDATER").is_some() {
|
|
return false;
|
|
}
|
|
true
|
|
}
|
|
/// Mode-gate for the direct stdio agent's background auto-update.
|
|
///
|
|
/// Only the *direct* stdio agent is newly eligible: every other agent mode
|
|
/// already self-updates at the top of `run_agent_command`, and a leader-backed
|
|
/// stdio process is a thin bridge whose updates are owned by the leader
|
|
/// (`LeaderAutoUpdateConfig`). Update suppression (`--no-auto-update`,
|
|
/// `KIGI_DISABLE_AUTOUPDATER`, debug builds) is layered on separately via
|
|
/// [`should_check_for_updates`].
|
|
fn stdio_direct_update_eligible(is_stdio: bool, use_leader: bool) -> bool {
|
|
is_stdio && !use_leader
|
|
}
|
|
/// Map the mutually-exclusive channel flags to a channel name. clap enforces
|
|
/// that at most one is set, so the order is irrelevant.
|
|
fn get_channel_switch(alpha: bool, stable: bool, enterprise: bool) -> Option<&'static str> {
|
|
if alpha {
|
|
Some("alpha")
|
|
} else if stable {
|
|
Some("stable")
|
|
} else if enterprise {
|
|
Some("enterprise")
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
/// Handle `kigi-pager update [--check] [--json] [--force-reinstall] [--version X] [--alpha|--stable|--enterprise]`.
|
|
async fn run_update_command(
|
|
check: bool,
|
|
json: bool,
|
|
force_reinstall: bool,
|
|
version: Option<String>,
|
|
channel_switch: Option<&str>,
|
|
base_update_config: &UpdateConfig,
|
|
) -> Result<()> {
|
|
if json && !check {
|
|
anyhow::bail!("--json requires --check");
|
|
}
|
|
let mut update_config = base_update_config.clone();
|
|
if check {
|
|
if version.is_some() {
|
|
anyhow::bail!("--version cannot be used with --check");
|
|
}
|
|
auto_update::apply_channel_switch(channel_switch, &mut update_config).await;
|
|
let status = auto_update::check_update_status(&update_config).await;
|
|
auto_update::print_update_status(&status, json)?;
|
|
return Ok(());
|
|
}
|
|
if let Some(ref v) = version
|
|
&& semver::Version::parse(v).is_err()
|
|
{
|
|
anyhow::bail!(
|
|
"'{}' is not a valid version. Expected semver like 0.1.150",
|
|
v
|
|
);
|
|
}
|
|
let installed = auto_update::run_update(
|
|
force_reinstall,
|
|
version.as_deref(),
|
|
channel_switch,
|
|
&mut update_config,
|
|
)
|
|
.await?;
|
|
if let Some(installed_version) = installed {
|
|
signal_leaders_to_relaunch(&installed_version).await;
|
|
}
|
|
Ok(())
|
|
}
|
|
/// After a successful `kigi update`, ask any running leader on this machine that
|
|
/// is older than `installed_version` to relaunch onto the new binary (bounded
|
|
/// grace; running sessions close and reconnect via `session/load`).
|
|
///
|
|
/// Best-effort and non-fatal: discovery/connect/control failures are logged and
|
|
/// skipped. The leader re-checks the directional version guard authoritatively;
|
|
/// the pager-side `live_info` check just avoids connecting to newer leaders.
|
|
async fn signal_leaders_to_relaunch(installed_version: &str) {
|
|
for d in kigi_shell::leader::discover_leaders().await {
|
|
if d.classification != kigi_shell::leader::LeaderDiscoveryState::Reachable {
|
|
continue;
|
|
}
|
|
let Some(socket_path) = d.socket_path.clone() else {
|
|
continue;
|
|
};
|
|
if let Some(ref live) = d.live_info
|
|
&& !leader_is_older_than(&live.leader_binary_version, installed_version)
|
|
{
|
|
continue;
|
|
}
|
|
let client = match kigi_shell::leader::LeaderClient::connect(
|
|
socket_path,
|
|
"kigi-pager-update",
|
|
ClientMode::Stdio,
|
|
ClientCapabilities::default(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
tracing::debug!(
|
|
error = % e, "Could not connect to leader to signal relaunch"
|
|
);
|
|
continue;
|
|
}
|
|
};
|
|
if !client.registration().supports_relaunch() {
|
|
client.cancel();
|
|
continue;
|
|
}
|
|
match client
|
|
.send_control(ControlCommand::RelaunchForUpdate {
|
|
to_version: installed_version.to_string(),
|
|
})
|
|
.await
|
|
{
|
|
Ok(Ok(kigi_shell::leader::ControlPayload::Relaunching {
|
|
from_version,
|
|
to_version,
|
|
..
|
|
})) => {
|
|
eprintln!(" ↻ Relaunching shared session (leader {from_version} → {to_version})…");
|
|
}
|
|
Ok(Ok(kigi_shell::leader::ControlPayload::RelaunchDeclined { reason })) => {
|
|
tracing::debug!(% reason, "Leader declined relaunch");
|
|
}
|
|
Ok(Ok(_)) => {}
|
|
Ok(Err(e)) => {
|
|
tracing::debug!(error = % e.message, "Leader relaunch control error");
|
|
}
|
|
Err(e) => {
|
|
tracing::debug!(
|
|
error = % e,
|
|
"Leader relaunch ack not received (leader may be exiting)"
|
|
);
|
|
}
|
|
}
|
|
client.cancel();
|
|
}
|
|
}
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
/// Only the direct (non-leader) stdio agent is newly eligible for the
|
|
/// background auto-update. Leader-backed stdio defers to the leader's own
|
|
/// updater, and non-stdio modes already update at the top of
|
|
/// `run_agent_command`.
|
|
#[test]
|
|
fn stdio_direct_update_eligible_only_for_non_leader_stdio() {
|
|
assert!(
|
|
stdio_direct_update_eligible(true, false),
|
|
"direct stdio agent should be eligible",
|
|
);
|
|
assert!(
|
|
!stdio_direct_update_eligible(true, true),
|
|
"leader-backed stdio defers to the leader's updater",
|
|
);
|
|
assert!(
|
|
!stdio_direct_update_eligible(false, false),
|
|
"non-stdio modes update at the top of run_agent_command",
|
|
);
|
|
assert!(
|
|
!stdio_direct_update_eligible(false, true),
|
|
"non-stdio leader path is not stdio-eligible",
|
|
);
|
|
}
|
|
use clap::Parser as _;
|
|
/// `kigi dashboard` flags the startup hook without forcing leader mode —
|
|
/// the dashboard is independent of leader mode, so the launch keeps
|
|
/// whatever leader setting the user (or config) chose.
|
|
#[serial_test::serial(KIGI_AGENT_DASHBOARD)]
|
|
#[test]
|
|
fn dashboard_subcommand_flags_startup_without_forcing_leader() {
|
|
let mut args = PagerArgs::try_parse_from(["kigi", "dashboard"]).unwrap();
|
|
assert!(!args.leader, "fixture: no explicit --leader");
|
|
flag_dashboard_at_startup_if_requested(&mut args).unwrap();
|
|
assert!(!args.leader, "dashboard must NOT force leader mode");
|
|
assert!(
|
|
args.command.is_none(),
|
|
"soft subcommand must be consumed so the interactive path runs",
|
|
);
|
|
assert_eq!(
|
|
std::env::var("KIGI_OPEN_DASHBOARD_AT_STARTUP").as_deref(),
|
|
Ok("1"),
|
|
"startup hook flag must be set",
|
|
);
|
|
unsafe { std::env::remove_var("KIGI_OPEN_DASHBOARD_AT_STARTUP") };
|
|
}
|
|
/// `kigi dashboard --no-leader` is allowed — the dashboard does not
|
|
/// require a leader, so the combination launches into the dashboard in
|
|
/// non-leader mode.
|
|
#[serial_test::serial(KIGI_AGENT_DASHBOARD)]
|
|
#[test]
|
|
fn dashboard_subcommand_allows_no_leader() {
|
|
let mut args = PagerArgs::try_parse_from(["kigi", "--no-leader", "dashboard"]).unwrap();
|
|
flag_dashboard_at_startup_if_requested(&mut args)
|
|
.expect("--no-leader + dashboard must be allowed");
|
|
assert!(args.no_leader, "--no-leader must be preserved");
|
|
assert!(!args.leader, "dashboard must not force leader mode");
|
|
assert!(
|
|
args.command.is_none(),
|
|
"soft subcommand must be consumed so the interactive path runs",
|
|
);
|
|
assert_eq!(
|
|
std::env::var("KIGI_OPEN_DASHBOARD_AT_STARTUP").as_deref(),
|
|
Ok("1"),
|
|
"startup hook flag must be set",
|
|
);
|
|
unsafe { std::env::remove_var("KIGI_OPEN_DASHBOARD_AT_STARTUP") };
|
|
}
|
|
/// `KIGI_AGENT_DASHBOARD=0` disables the feature — the subcommand
|
|
/// must error visibly before the TUI starts.
|
|
#[serial_test::serial(KIGI_AGENT_DASHBOARD)]
|
|
#[test]
|
|
fn dashboard_subcommand_errors_when_disabled() {
|
|
unsafe { std::env::set_var("KIGI_AGENT_DASHBOARD", "0") };
|
|
let mut args = PagerArgs::try_parse_from(["kigi", "dashboard"]).unwrap();
|
|
let result = flag_dashboard_at_startup_if_requested(&mut args);
|
|
unsafe { std::env::remove_var("KIGI_AGENT_DASHBOARD") };
|
|
let err = result.expect_err("disabled dashboard must error");
|
|
assert!(err.to_string().contains("disabled"), "got: {err}");
|
|
assert!(
|
|
std::env::var("KIGI_OPEN_DASHBOARD_AT_STARTUP").is_err(),
|
|
"failure path must not flag the startup hook",
|
|
);
|
|
}
|
|
fn make_state() -> std::sync::Mutex<StdioReplayState> {
|
|
std::sync::Mutex::new(StdioReplayState::default())
|
|
}
|
|
#[test]
|
|
fn cache_initialize_request() {
|
|
let state = make_state();
|
|
let msg = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
|
|
cache_outgoing_acp_state(msg, &state);
|
|
let s = state.lock().unwrap();
|
|
assert_eq!(s.initialize_json.as_deref(), Some(msg));
|
|
}
|
|
#[test]
|
|
fn cache_session_load_preserves_full_request() {
|
|
let state = make_state();
|
|
let msg = r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"s1","cwd":"/tmp","mcpServers":[]}}"#;
|
|
cache_outgoing_acp_state(msg, &state);
|
|
let s = state.lock().unwrap();
|
|
let (sid, cached) = &s.sessions[0];
|
|
assert_eq!(sid, "s1");
|
|
assert_eq!(cached.load_request_json.as_deref(), Some(msg));
|
|
assert_eq!(cached.cwd.as_deref(), Some("/tmp"));
|
|
assert!(cached.mcp_servers_json.is_some());
|
|
assert_eq!(s.last_session_id.as_deref(), Some("s1"));
|
|
}
|
|
#[test]
|
|
fn cache_session_new_is_pending_until_response_assigns_id() {
|
|
let state = make_state();
|
|
let load = r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"s1","cwd":"/tmp"}}"#;
|
|
cache_outgoing_acp_state(load, &state);
|
|
let new = r#"{"jsonrpc":"2.0","id":3,"method":"session/new","params":{"cwd":"/home"}}"#;
|
|
cache_outgoing_acp_state(new, &state);
|
|
{
|
|
let s = state.lock().unwrap();
|
|
assert_eq!(s.sessions.len(), 1);
|
|
assert_eq!(s.sessions[0].0, "s1");
|
|
assert!(s.pending_new.is_some());
|
|
assert_eq!(
|
|
s.pending_new.as_ref().unwrap().cwd.as_deref(),
|
|
Some("/home")
|
|
);
|
|
}
|
|
cache_incoming_session_id(
|
|
r#"{"jsonrpc":"2.0","id":3,"result":{"sessionId":"s2"}}"#,
|
|
&state,
|
|
);
|
|
let s = state.lock().unwrap();
|
|
assert!(s.pending_new.is_none());
|
|
assert_eq!(s.sessions.len(), 2);
|
|
assert_eq!(s.sessions[1].0, "s2");
|
|
assert_eq!(s.sessions[1].1.cwd.as_deref(), Some("/home"));
|
|
assert_eq!(s.last_session_id.as_deref(), Some("s2"));
|
|
}
|
|
#[test]
|
|
fn cache_session_close_stops_replaying_it() {
|
|
let state = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"s1","cwd":"/tmp"}}"#,
|
|
&state,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":3,"method":"_kigi/session/close","params":{"sessionId":"s1"}}"#,
|
|
&state,
|
|
);
|
|
let s = state.lock().unwrap();
|
|
assert!(s.sessions.is_empty(), "closed session must not be replayed");
|
|
assert!(s.last_session_id.is_none());
|
|
}
|
|
/// An UNCONFIRMED `session/new` (leader died before its response) must not
|
|
/// be replayed — its id was never assigned — but previously loaded
|
|
/// sessions still restore.
|
|
#[tokio::test]
|
|
async fn replay_after_unconfirmed_session_new_restores_prior_sessions() {
|
|
let state = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state,
|
|
);
|
|
let load_a = r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"session-A","cwd":"/old"}}"#;
|
|
cache_outgoing_acp_state(load_a, &state);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":3,"method":"session/new","params":{"cwd":"/new"}}"#,
|
|
&state,
|
|
);
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let load = leader_rx.recv().await.unwrap();
|
|
assert!(load.contains("session-A"), "unexpected replay: {load}");
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":2,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
});
|
|
let s = state.lock().unwrap().clone();
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &s).await;
|
|
assert_eq!(
|
|
result.as_deref(),
|
|
Some("session-A"),
|
|
"prior session must be restored even though the new one was unconfirmed"
|
|
);
|
|
responder.await.unwrap();
|
|
}
|
|
#[test]
|
|
fn fallback_replay_json_escapes_special_chars() {
|
|
let cached = CachedSession {
|
|
load_request_json: None,
|
|
cwd: Some(r#"C:\Users\test path"#.into()),
|
|
mcp_servers_json: None,
|
|
};
|
|
let json = replay_load_json(r#"session"with"quotes"#, &cached)
|
|
.expect("cwd present → load synthesized");
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_str(&json).expect("fallback replay JSON must be valid");
|
|
assert_eq!(
|
|
parsed["params"]["sessionId"].as_str().unwrap(),
|
|
r#"session"with"quotes"#
|
|
);
|
|
assert_eq!(
|
|
parsed["params"]["cwd"].as_str().unwrap(),
|
|
r#"C:\Users\test path"#
|
|
);
|
|
assert_eq!(parsed["id"].as_str(), Some(REPLAY_LOAD_REQUEST_ID));
|
|
}
|
|
#[test]
|
|
fn cache_incoming_session_id_from_response() {
|
|
let state = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"session/new","params":{"cwd":"/tmp"}}"#,
|
|
&state,
|
|
);
|
|
let msg = r#"{"jsonrpc":"2.0","id":1,"result":{"sessionId":"abc123"}}"#;
|
|
cache_incoming_session_id(msg, &state);
|
|
let s = state.lock().unwrap();
|
|
assert_eq!(s.last_session_id.as_deref(), Some("abc123"));
|
|
assert_eq!(s.sessions.len(), 1);
|
|
assert_eq!(s.sessions[0].0, "abc123");
|
|
}
|
|
/// A multi-session client (IDE driving several sessions over one bridge)
|
|
/// gets EVERY session replayed after a reconnect, in first-seen order.
|
|
#[tokio::test]
|
|
async fn replay_restores_all_cached_sessions() {
|
|
let state = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"sess-1","cwd":"/a"}}"#,
|
|
&state,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":3,"method":"session/new","params":{"cwd":"/b"}}"#,
|
|
&state,
|
|
);
|
|
cache_incoming_session_id(
|
|
r#"{"jsonrpc":"2.0","id":3,"result":{"sessionId":"sess-2"}}"#,
|
|
&state,
|
|
);
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let load1 = leader_rx.recv().await.unwrap();
|
|
assert!(load1.contains("sess-1"), "expected sess-1 first: {load1}");
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":2,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let load2 = leader_rx.recv().await.unwrap();
|
|
assert!(load2.contains("sess-2"), "expected sess-2 second: {load2}");
|
|
let load2_json: serde_json::Value = serde_json::from_str(&load2).unwrap();
|
|
assert_eq!(load2_json["id"].as_str(), Some(REPLAY_LOAD_REQUEST_ID));
|
|
response_tx
|
|
.send(
|
|
serde_json::json!(
|
|
{ "jsonrpc" : "2.0", "id" : REPLAY_LOAD_REQUEST_ID, "result" : {}
|
|
}
|
|
)
|
|
.to_string(),
|
|
)
|
|
.unwrap();
|
|
});
|
|
let s = state.lock().unwrap().clone();
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &s).await;
|
|
assert_eq!(result.as_deref(), Some("sess-2"));
|
|
responder.await.unwrap();
|
|
}
|
|
/// One broken session must not doom the rest: a rejected load is skipped
|
|
/// and the remaining sessions still restore.
|
|
#[tokio::test]
|
|
async fn replay_skips_rejected_session_and_restores_the_rest() {
|
|
let state = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"sess-bad","cwd":"/a"}}"#,
|
|
&state,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":3,"method":"session/load","params":{"sessionId":"sess-good","cwd":"/b"}}"#,
|
|
&state,
|
|
);
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let _bad = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(
|
|
r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32602,"message":"Invalid params","data":"unknown session id"}}"#
|
|
.to_string(),
|
|
)
|
|
.unwrap();
|
|
let good = leader_rx.recv().await.unwrap();
|
|
assert!(good.contains("sess-good"));
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":3,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
});
|
|
let s = state.lock().unwrap().clone();
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &s).await;
|
|
assert_eq!(result.as_deref(), Some("sess-good"));
|
|
responder.await.unwrap();
|
|
}
|
|
#[test]
|
|
fn cache_incoming_ignores_non_session_response() {
|
|
let state = make_state();
|
|
let msg = r#"{"jsonrpc":"2.0","id":1,"result":{"models":[]}}"#;
|
|
cache_incoming_session_id(msg, &state);
|
|
let s = state.lock().unwrap();
|
|
assert!(s.last_session_id.is_none());
|
|
assert!(s.sessions.is_empty());
|
|
}
|
|
#[tokio::test]
|
|
async fn replay_with_no_cached_state_returns_none() {
|
|
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let state = StdioReplayState::default();
|
|
let mut sink = Vec::new();
|
|
let result = replay_acp_state_after_reconnect(&tx, &mut rx, &mut sink, &state).await;
|
|
assert!(result.is_none());
|
|
}
|
|
#[tokio::test]
|
|
async fn replay_sends_initialize_and_session_load() {
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let (response_tx, response_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let state_mutex = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state_mutex,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"s1","cwd":"/tmp"}}"#,
|
|
&state_mutex,
|
|
);
|
|
let state = state_mutex.lock().unwrap().clone();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let _load = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":2,"result":{"sessionId":"s1"}}"#.to_string())
|
|
.unwrap();
|
|
});
|
|
let mut rx = response_rx;
|
|
let mut sink = Vec::new();
|
|
let result = replay_acp_state_after_reconnect(&leader_tx, &mut rx, &mut sink, &state).await;
|
|
assert_eq!(result.as_deref(), Some("s1"));
|
|
assert!(
|
|
sink.is_empty(),
|
|
"responses to replayed requests must be swallowed, not forwarded"
|
|
);
|
|
responder.await.unwrap();
|
|
}
|
|
/// Regression test for the post-leader-crash "unknown session id" bug.
|
|
///
|
|
/// `session/load` streams replay notifications BEFORE its response. The
|
|
/// old drain logic consumed exactly one message per replayed request and
|
|
/// returned — declaring the reconnect complete while the new leader was
|
|
/// still loading the session. The replay must instead:
|
|
/// 1. wait for the actual `session/load` RESPONSE (matched by id),
|
|
/// 2. forward interleaved notifications to the client verbatim,
|
|
/// 3. swallow only the responses to the replayed requests.
|
|
#[tokio::test]
|
|
async fn replay_waits_for_load_response_through_notifications() {
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let state_mutex = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":7,"method":"initialize","params":{}}"#,
|
|
&state_mutex,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":8,"method":"session/load","params":{"sessionId":"s9","cwd":"/tmp"}}"#,
|
|
&state_mutex,
|
|
);
|
|
let state = state_mutex.lock().unwrap().clone();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(
|
|
r#"{"jsonrpc":"2.0","method":"kigi/leader/version_mismatch","params":{}}"#
|
|
.to_string(),
|
|
)
|
|
.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":7,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let _load = leader_rx.recv().await.unwrap();
|
|
for i in 0..3 {
|
|
response_tx
|
|
.send(
|
|
format!(
|
|
r#"{{"jsonrpc":"2.0","method":"session/update","params":{{"sessionId":"s9","n":{i}}}}}"#
|
|
),
|
|
)
|
|
.unwrap();
|
|
}
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":8,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
});
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &state).await;
|
|
assert_eq!(
|
|
result.as_deref(),
|
|
Some("s9"),
|
|
"replay must succeed once the load response arrives"
|
|
);
|
|
let forwarded = String::from_utf8(sink).unwrap();
|
|
let lines: Vec<&str> = forwarded.lines().collect();
|
|
assert_eq!(
|
|
lines.len(),
|
|
4,
|
|
"expected exactly the 4 notifications, got: {lines:?}"
|
|
);
|
|
assert!(lines[0].contains("version_mismatch"));
|
|
assert!(lines[1].contains(r#""n":0"#));
|
|
assert!(lines[2].contains(r#""n":1"#));
|
|
assert!(lines[3].contains(r#""n":2"#));
|
|
assert!(!forwarded.contains(r#""id":7"#), "init response leaked");
|
|
assert!(!forwarded.contains(r#""id":8"#), "load response leaked");
|
|
responder.await.unwrap();
|
|
}
|
|
/// A `session/load` rejected by the new leader (error response) must
|
|
/// surface as a failed replay (`None`) so the bridge emits
|
|
/// `kigi/leader_reconnected` with empty params and the external client
|
|
/// knows to re-establish state itself.
|
|
#[tokio::test]
|
|
async fn replay_returns_none_when_load_is_rejected() {
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel();
|
|
let state_mutex = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state_mutex,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/load","params":{"sessionId":"gone","cwd":"/tmp"}}"#,
|
|
&state_mutex,
|
|
);
|
|
let state = state_mutex.lock().unwrap().clone();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let _load = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(
|
|
r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32602,"message":"Invalid params","data":"unknown session id"}}"#
|
|
.to_string(),
|
|
)
|
|
.unwrap();
|
|
});
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &state).await;
|
|
assert!(result.is_none(), "rejected load must not claim success");
|
|
responder.await.unwrap();
|
|
}
|
|
/// The synthetic fallback `session/load` (client only ever sent
|
|
/// `session/new`) uses a string request id that cannot collide with the
|
|
/// external client's numeric ids — and the response matcher honors it.
|
|
#[tokio::test]
|
|
async fn replay_fallback_load_uses_reserved_string_id() {
|
|
let (leader_tx, mut leader_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
|
let state_mutex = make_state();
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
|
|
&state_mutex,
|
|
);
|
|
cache_outgoing_acp_state(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"session/new","params":{"cwd":"/tmp"}}"#,
|
|
&state_mutex,
|
|
);
|
|
cache_incoming_session_id(
|
|
r#"{"jsonrpc":"2.0","id":2,"result":{"sessionId":"s-new"}}"#,
|
|
&state_mutex,
|
|
);
|
|
let state = state_mutex.lock().unwrap().clone();
|
|
let responder = tokio::spawn(async move {
|
|
let _init = leader_rx.recv().await.unwrap();
|
|
response_tx
|
|
.send(r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.to_string())
|
|
.unwrap();
|
|
let load = leader_rx.recv().await.unwrap();
|
|
let load_json: serde_json::Value = serde_json::from_str(&load).unwrap();
|
|
assert_eq!(
|
|
load_json["id"].as_str(),
|
|
Some(REPLAY_LOAD_REQUEST_ID),
|
|
"fallback load must use the reserved string id"
|
|
);
|
|
response_tx
|
|
.send(
|
|
serde_json::json!(
|
|
{ "jsonrpc" : "2.0", "id" : REPLAY_LOAD_REQUEST_ID, "result" : {}
|
|
}
|
|
)
|
|
.to_string(),
|
|
)
|
|
.unwrap();
|
|
});
|
|
let mut sink = Vec::new();
|
|
let result =
|
|
replay_acp_state_after_reconnect(&leader_tx, &mut response_rx, &mut sink, &state).await;
|
|
assert_eq!(result.as_deref(), Some("s-new"));
|
|
responder.await.unwrap();
|
|
}
|
|
fn multi_thread_runtime() -> tokio::runtime::Runtime {
|
|
tokio::runtime::Builder::new_multi_thread()
|
|
.enable_all()
|
|
.build()
|
|
.expect("build runtime")
|
|
}
|
|
#[test]
|
|
fn run_and_shutdown_bounds_teardown_despite_stuck_blocking_task() {
|
|
use std::time::{Duration, Instant};
|
|
let grace = Duration::from_millis(200);
|
|
let ceiling = grace * 8;
|
|
let stuck_sleep = Duration::from_secs(10);
|
|
let runtime = multi_thread_runtime();
|
|
let (started_tx, started_rx) = std::sync::mpsc::channel::<()>();
|
|
runtime.spawn_blocking(move || {
|
|
let _ = started_tx.send(());
|
|
std::thread::sleep(stuck_sleep);
|
|
});
|
|
started_rx.recv().expect("blocking task must start");
|
|
let start = Instant::now();
|
|
let out = run_and_shutdown(runtime, async { 7_u32 }, grace);
|
|
let elapsed = start.elapsed();
|
|
assert_eq!(out, 7, "must return the future's output");
|
|
assert!(
|
|
elapsed >= grace,
|
|
"returned in {elapsed:?}, before the {grace:?} grace — timeout not exercised",
|
|
);
|
|
assert!(
|
|
elapsed < ceiling,
|
|
"teardown took {elapsed:?}; stuck task must be abandoned under {ceiling:?}",
|
|
);
|
|
}
|
|
#[test]
|
|
fn run_and_shutdown_is_fast_without_blocking_work() {
|
|
use std::time::{Duration, Instant};
|
|
let runtime = multi_thread_runtime();
|
|
let grace = Duration::from_secs(5);
|
|
let start = Instant::now();
|
|
let out = run_and_shutdown(runtime, async { 42_u32 }, grace);
|
|
let elapsed = start.elapsed();
|
|
assert_eq!(out, 42, "must pass the future's output through");
|
|
assert!(
|
|
elapsed < grace,
|
|
"clean teardown took {elapsed:?}; grace must be a ceiling, not a floor",
|
|
);
|
|
}
|
|
#[test]
|
|
fn run_and_shutdown_passes_err_output_through() {
|
|
use std::time::Duration;
|
|
let runtime = multi_thread_runtime();
|
|
let out = run_and_shutdown(
|
|
runtime,
|
|
async { Err::<(), String>("boom".to_string()) },
|
|
Duration::from_secs(5),
|
|
);
|
|
assert_eq!(
|
|
out,
|
|
Err("boom".to_string()),
|
|
"Err output must pass through unchanged",
|
|
);
|
|
}
|
|
}
|