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:
2026-07-17 05:31:01 -04:00
commit d6c20fc13f
2612 changed files with 1353757 additions and 0 deletions
@@ -0,0 +1,64 @@
//! Shared helpers for kigi-shell integration tests.
use kigi_shell::sampling::{ApiBackend, Client, SamplerConfig};
/// Create a sampling client configured for a mock server. Shared by the
/// integration tests so the ~30-field `SamplerConfig` literal lives in one
/// place (`SamplerConfig` has no `Default`).
pub fn create_test_client(base_url: &str, api_backend: ApiBackend) -> Client {
create_test_client_with_extra_headers(base_url, api_backend, &[])
}
/// Like [`create_test_client`] but seeds `SamplerConfig::extra_headers`, so a
/// test can assert that session-injected headers reach the wire.
pub fn create_test_client_with_extra_headers(
base_url: &str,
api_backend: ApiBackend,
extra_headers: &[(&str, &str)],
) -> Client {
Client::new(test_sampler_config(base_url, api_backend, extra_headers)).unwrap()
}
/// The shared mock-server `SamplerConfig`; tests needing a non-default field
/// (e.g. `doom_loop_recovery`) mutate the returned value before building the
/// client themselves.
pub fn test_sampler_config(
base_url: &str,
api_backend: ApiBackend,
extra_headers: &[(&str, &str)],
) -> SamplerConfig {
// Shell `Client` is `kigi_sampler::SamplingClient`, which takes a
// `SamplerConfig` directly. Construct one inline here.
SamplerConfig {
api_key: Some("test-api-key".to_string()),
base_url: base_url.to_string(),
model: "test-model".to_string(),
max_completion_tokens: Some(1000),
temperature: Some(0.7),
top_p: None,
api_backend,
auth_scheme: Default::default(),
extra_headers: extra_headers
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
supports_backend_search: false,
compactions_remaining: None,
compaction_at_tokens: None,
doom_loop_recovery: None,
header_injector: None,
}
}
@@ -0,0 +1,25 @@
{
"name": "clean_completion",
"description": "Assistant ended its turn with no outstanding todos. The gate must NOT fire — this is the only true 'happy path' for end-of-turn.",
"turns": [
{
"turn_index": 0,
"kind": "user",
"content": "<synthetic user prompt — content scrubbed>"
},
{
"turn_index": 1,
"kind": "assistant",
"tool_calls_emitted": [],
"todo_state_after_turn": [
{
"id": "ship-it",
"status": "completed",
"content": "land the PR"
}
],
"backing_task_count": 0,
"expected_gate_decision": "continue"
}
]
}
@@ -0,0 +1,41 @@
{
"name": "pr_babysit_partial_backing",
"description": "Three /pr-babysit todos are in_progress but only one polling subagent is alive — the gate's tightened heuristic (PR2 partition) must classify the two unbacked items and fire with the in-flight reason. Backed items are intentionally NOT listed in the reminder body (the gate already decided not to nudge on them); negative coverage lives in the inline turn_end_guard_tests in acp_session.rs.",
"turns": [
{
"turn_index": 0,
"kind": "user",
"content": "<synthetic user prompt — content scrubbed>"
},
{
"turn_index": 1,
"kind": "assistant",
"tool_calls_emitted": [],
"todo_state_after_turn": [
{
"id": "pr-1",
"status": "in_progress",
"content": "pr-1:ci-green"
},
{
"id": "pr-2",
"status": "in_progress",
"content": "pr-2:ci-green"
},
{
"id": "pr-3",
"status": "in_progress",
"content": "pr-3:ci-green"
}
],
"backing_task_count": 1,
"expected_gate_decision": "nudge",
"expected_reason": "in_flight",
"expected_reminder_contains": [
"In-progress (no backing background task)",
"pr-2:ci-green",
"pr-3:ci-green"
]
}
]
}
@@ -0,0 +1,31 @@
{
"name": "stranded_narration",
"description": "Assistant ended its turn with a content-only message while a pending todo remains and no backing background task exists. Canonical PR1 §B 'stranded narration' failure shape: the gate must fire with the in-flight reason.",
"turns": [
{
"turn_index": 0,
"kind": "user",
"content": "<synthetic user prompt — content scrubbed>"
},
{
"turn_index": 1,
"kind": "assistant",
"tool_calls_emitted": [],
"todo_state_after_turn": [
{
"id": "fix-round-1",
"status": "pending",
"content": "fix open review issues"
}
],
"backing_task_count": 0,
"expected_gate_decision": "nudge",
"expected_reason": "in_flight",
"expected_reminder_contains": [
"ended your turn",
"Pending:",
"fix open review issues"
]
}
]
}
@@ -0,0 +1,484 @@
//! End-to-end repro of the git refresh storm during agent-run rebases.
//!
//! Stands up a real in-process `MvpAgent` over duplex ACP pipes with a
//! hand-rolled `initialize` (the stock test client advertises empty
//! capabilities, so the fs watcher would never spawn) and drives two scripted
//! turns through `MockInferenceServer`:
//!
//! 1. `search_replace` creates two files, so the AgentOnly hunk tracker has
//! tracked paths (defeating its nothing-tracked early return).
//! 2. `run_terminal_command` runs a real multi-pick `git rebase -i` whose
//! picks chain back to back, reproducing the continuous
//! `.git/index.lock` / HEAD-move churn of an agent-run rebase.
//!
//! With the fs-watch machinery on (`x.ai/hunkTracker` + `x.ai/gitHeadChanged`
//! advertised), fsnotify merges rapid lock cycles into one operation and the
//! session defers debounce fires while an op is in flight, so one rebase
//! costs at most a couple of `refresh_all_baselines` scans (each scoped to
//! the tracked paths) instead of one full-worktree scan per inter-pick gap;
//! with capabilities absent the watcher never spawns and no scans run. The
//! test counts real scans via a global tracing layer (sessions run on their
//! own threads, so a thread-scoped subscriber would miss them), prints both
//! runs, and asserts only invariants: zero scans with the machinery off, a
//! small bounded number per rebase with it on.
//!
//! Knobs: KIGI_PERF_GIT_FILES (default 300), KIGI_PERF_GIT_PICKS (default 6).
//! Keep the rebase under ~15s or the terminal tool auto-backgrounds the
//! command and the turn wall time loses meaning.
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use agent_client_protocol::{self as acp, Agent as _};
use kigi_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
LineBufferedRead,
};
use kigi_shell::agent::config::Config as AgentConfig;
use kigi_shell::agent::mvp_agent::MvpAgent;
use kigi_test_support::{MockInferenceServer, ScriptedResponse, SseEvent};
use serde_json::{Value, json};
use tempfile::TempDir;
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
const DUPLEX_BUFFER_BYTES: usize = 8 * 1024 * 1024;
use kigi_test_utils::env::env_usize;
/// Run a git command with deterministic author/committer; assert success.
fn git(dir: &Path, args: &[&str]) -> String {
kigi_test_utils::git::run_git(dir, args)
}
// ── scan counting ─────────────────────────────────────────────────────────
use kigi_hunk_tracker::{REFRESH_SCAN_LOG_PREFIX, REFRESH_SKIP_LOG_PREFIX};
use kigi_test_utils::tracing_capture::MessagePrefixCounter;
/// Counts hunk-tracker scan completions/skips across all threads (the session
/// actor and its consumers run off the test thread). Only real scans log the
/// scan prefix; the unchanged-git-state skip logs the other.
#[derive(Clone)]
struct ScanCounter(MessagePrefixCounter);
impl ScanCounter {
fn scans(&self) -> usize {
self.0.count(REFRESH_SCAN_LOG_PREFIX)
}
fn skips(&self) -> usize {
self.0.count(REFRESH_SKIP_LOG_PREFIX)
}
}
fn install_global_scan_counter() -> ScanCounter {
// KIGI_E2E_LOG=<filter> tees shell logs to stderr for local debugging.
let filter = std::env::var("KIGI_E2E_LOG").ok();
ScanCounter(
kigi_test_utils::tracing_capture::install_prefix_counter_global(
&[REFRESH_SCAN_LOG_PREFIX, REFRESH_SKIP_LOG_PREFIX],
filter.as_deref(),
),
)
}
// ── repo fixture ──────────────────────────────────────────────────────────
/// Committed tree of ~`files` files plus a `feature` branch with `picks`
/// one-file commits and an advanced base branch, `feature` checked out.
/// Returns the repo dir and the base branch name.
fn build_repo(files: usize, picks: usize) -> (TempDir, String) {
let dir = TempDir::new().expect("repo tempdir");
let wd = dir.path();
git(wd, &["init"]);
git(wd, &["config", "user.name", "Test User"]);
git(wd, &["config", "user.email", "test@test.com"]);
kigi_test_utils::git::write_fanout_tree(wd, files, 100);
git(wd, &["add", "."]);
git(wd, &["commit", "-m", "populate tree"]);
let base = kigi_test_utils::git::make_feature_branch(wd, picks);
(dir, base)
}
// ── scripted responses (chat-completions SSE) ────────────────────────────
fn chat_chunk(delta: Value, finish_reason: Value) -> SseEvent {
SseEvent::data(
json!({
"id": "chatcmpl-test",
"object": "chat.completion.chunk",
"created": 1234567890,
"model": "test-model",
"choices": [{ "index": 0, "delta": delta, "finish_reason": finish_reason }]
})
.to_string(),
)
}
fn tool_call_sse(name: &str, arguments: &Value) -> ScriptedResponse {
static CALL_SEQ: AtomicUsize = AtomicUsize::new(0);
let call_id = format!("call_{name}_{}", CALL_SEQ.fetch_add(1, Ordering::Relaxed));
let events = vec![
chat_chunk(
json!({
"role": "assistant",
"content": null,
"tool_calls": [{
"id": call_id,
"type": "function",
"function": { "name": name, "arguments": arguments.to_string() }
}]
}),
Value::Null,
),
chat_chunk(json!({}), json!("tool_calls")),
SseEvent::data("[DONE]"),
];
ScriptedResponse::sse(events)
}
fn text_sse(text: &str) -> ScriptedResponse {
let events = vec![
chat_chunk(json!({ "role": "assistant", "content": text }), Value::Null),
chat_chunk(json!({}), json!("stop")),
SseEvent::data("[DONE]"),
];
ScriptedResponse::sse(events)
}
// ── client ────────────────────────────────────────────────────────────────
/// Auto-approves permissions (AllowOnce preferred) and drops notifications.
struct AutoApproveClient;
#[async_trait::async_trait(?Send)]
impl acp::Client for AutoApproveClient {
async fn request_permission(
&self,
args: acp::RequestPermissionRequest,
) -> acp::Result<acp::RequestPermissionResponse> {
let outcome = args
.options
.iter()
.find(|o| o.kind == acp::PermissionOptionKind::AllowOnce)
.or(args.options.first())
.map(|o| {
acp::RequestPermissionOutcome::Selected(acp::SelectedPermissionOutcome::new(
o.option_id.clone(),
))
})
.unwrap_or(acp::RequestPermissionOutcome::Cancelled);
Ok(acp::RequestPermissionResponse::new(outcome))
}
async fn session_notification(&self, _args: acp::SessionNotification) -> acp::Result<()> {
Ok(())
}
}
// ── one full agent run ────────────────────────────────────────────────────
struct RunStats {
scans: usize,
skips: usize,
edit_turn: Duration,
rebase_turn: Duration,
}
async fn prompt_turn(
client_conn: &acp::ClientSideConnection,
session_id: &acp::SessionId,
text: &str,
label: &str,
) -> Duration {
let started = Instant::now();
let resp = tokio::time::timeout(
Duration::from_secs(180),
client_conn.prompt(acp::PromptRequest::new(
session_id.clone(),
vec![acp::ContentBlock::Text(acp::TextContent::new(
text.to_owned(),
))],
)),
)
.await
.unwrap_or_else(|_| panic!("{label}: prompt timed out"))
.unwrap_or_else(|e| panic!("{label}: prompt failed: {e}"));
assert!(
matches!(resp.stop_reason, acp::StopReason::EndTurn),
"{label}: expected EndTurn, got {:?}",
resp.stop_reason
);
started.elapsed()
}
/// Drive both scripted turns through a fresh in-process agent over a fresh
/// repo. `caps_meta` is the `client_capabilities._meta` advertised on
/// `initialize` (None = fs-watch machinery off).
async fn run_storm(
server: &MockInferenceServer,
caps_meta: Option<Value>,
counter: &ScanCounter,
label: &str,
) -> RunStats {
let files = env_usize("KIGI_PERF_GIT_FILES", 300);
let picks = env_usize("KIGI_PERF_GIT_PICKS", 6);
let (repo, base) = build_repo(files, picks);
eprintln!(
"[perf] {label}: repo ~{files} files, {picks} picks at {:?}",
repo.path()
);
// Turn 1: two file-creating edits, then a final text.
server.enqueue_response(
"/v1/chat/completions",
tool_call_sse(
"search_replace",
&json!({
"file_path": "agent_notes_a.md",
"old_string": "",
"new_string": "agent notes a\n"
}),
),
);
server.enqueue_response(
"/v1/chat/completions",
tool_call_sse(
"search_replace",
&json!({
"file_path": "agent_notes_b.md",
"old_string": "",
"new_string": "agent notes b\n"
}),
),
);
server.enqueue_response("/v1/chat/completions", text_sse("created the notes files"));
// Turn 2: the storm shape — a real multi-pick rebase whose picks chain
// continuously. No --exec: each inserted exec spawns shells that add >1s
// of lock-free idle per pick in this environment, and idle-gapped
// rebases legitimately refresh per gap; continuous lock churn is what
// must merge into a single operation (one refresh), and is what an
// agent-run rebase looks like.
let rebase_cmd = format!("GIT_SEQUENCE_EDITOR=: git rebase -i {base} 2>&1");
server.enqueue_response(
"/v1/chat/completions",
tool_call_sse(
"run_terminal_command",
&json!({ "command": rebase_cmd, "description": "storm repro rebase" }),
),
);
server.enqueue_response("/v1/chat/completions", text_sse("rebase complete"));
let counter = counter.clone();
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
// Blocks this thread on the startup models/settings prefetch —
// served by the mock's dedicated runtime thread, so it completes
// and the catalog contains the mock's chat-completions model.
let agent_config = AgentConfig::default();
let auth_manager = Arc::new(agent_config.create_auth_manager());
let (gw_tx, gw_rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(gw_tx);
let agent =
MvpAgent::new(gateway, &agent_config, auth_manager, None).expect("valid config");
let (c2a_a, c2a_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
let (a2c_a, a2c_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
let agent_incoming = LineBufferedRead::spawn_local(c2a_b.compat());
let (agent_conn, agent_io) =
acp::AgentSideConnection::new(agent, a2c_a.compat_write(), agent_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(
GatewayReceiver::new(gw_rx, agent_conn)
.with_on_meta(kigi_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
tokio::task::spawn_local(agent_io);
let client_incoming = LineBufferedRead::spawn_local(a2c_b.compat());
let (client_conn, client_io) = acp::ClientSideConnection::new(
AutoApproveClient,
c2a_a.compat_write(),
client_incoming,
|fut| {
tokio::task::spawn_local(fut);
},
);
tokio::task::spawn_local(client_io);
let init = tokio::time::timeout(
Duration::from_secs(60),
client_conn.initialize(
acp::InitializeRequest::new(acp::ProtocolVersion::V1)
.client_capabilities(
acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false)
.meta(caps_meta.and_then(|v| v.as_object().cloned())),
)
.meta(
json!({
"startupHints": {
"nonInteractive": true,
"skipGitStatus": true,
"skipProjectLayout": true,
},
"clientType": "git-contention-e2e",
"clientVersion": "0.0-test",
})
.as_object()
.cloned(),
),
),
)
.await
.expect("initialize timed out")
.expect("initialize failed");
// Strict: authenticating is what triggers the remote model fetch,
// and the session must resolve the mock's chat-completions model.
let method = init
.auth_methods
.iter()
.find(|m| &*m.id().0 == "xai.api_key")
.unwrap_or_else(|| panic!("{label}: xai.api_key auth method not advertised"));
client_conn
.authenticate(
acp::AuthenticateRequest::new(method.id().clone())
.meta(json!({ "headless": true }).as_object().cloned()),
)
.await
.unwrap_or_else(|e| panic!("{label}: authenticate failed: {e}"));
let session = tokio::time::timeout(
Duration::from_secs(60),
client_conn.new_session(
acp::NewSessionRequest::new(repo.path().to_path_buf())
.meta(json!({ "modelId": "test-model" }).as_object().cloned()),
),
)
.await
.expect("session/new timed out")
.expect("session/new failed");
let session_id = session.session_id;
let edit_turn =
prompt_turn(&client_conn, &session_id, "create the notes files", label).await;
// The scripted edits must actually have run: the storm depends on
// the hunk tracker holding tracked paths during the rebase.
assert!(
repo.path().join("agent_notes_a.md").exists()
&& repo.path().join("agent_notes_b.md").exists(),
"{label}: scripted search_replace edits did not run\n{}",
server.request_log_summary()
);
// Let the edit turn's fs events settle before windowing the storm.
tokio::time::sleep(Duration::from_millis(500)).await;
let scans_before = counter.scans();
let skips_before = counter.skips();
let rebase_turn = prompt_turn(&client_conn, &session_id, "run the rebase", label).await;
assert!(
repo.path().join("base_advance.txt").exists(),
"{label}: the scripted rebase did not run (HEAD: {})\n{}",
git(repo.path(), &["log", "--oneline", "-1"]),
server.request_log_summary()
);
// Trailing drain: the last debounce window (quiet 500ms, cap 3s)
// plus the spawned refresh itself.
tokio::time::sleep(Duration::from_secs(4)).await;
RunStats {
scans: counter.scans() - scans_before,
skips: counter.skips() - skips_before,
edit_turn,
rebase_turn,
}
})
.await
}
/// Real FS watcher + real git + real timers: too timing-dependent for CI.
#[test]
#[ignore = "perf repro; real FS events; run locally with --ignored --nocapture"]
fn git_rebase_refresh_storm_e2e() {
let _ = rustls::crypto::ring::default_provider().install_default();
let counter = install_global_scan_counter();
// The mock gets its own runtime thread: agent startup blocks the test
// thread on a models/settings prefetch (thread spawn + join), which would
// starve a mock sharing the agent's runtime and time the fetch out.
let mock_rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.expect("mock runtime");
let server = mock_rt
.block_on(MockInferenceServer::start())
.expect("mock server");
let kigi_home = TempDir::new().expect("grok home");
// SAFETY: the only live threads are the mock runtime's workers, which
// serve HTTP and never read the process environment.
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
std::env::set_var("KIGI_TRACE_UPLOAD", "false");
}
let agent_rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("agent runtime");
// Machinery off first so the on-run's trailing refreshes cannot bleed
// into the off-run's counting window.
let off = agent_rt.block_on(run_storm(&server, None, &counter, "machinery-off"));
let on = agent_rt.block_on(run_storm(
&server,
Some(json!({
"x.ai/hunkTracker": { "mode": "agent_only" },
"x.ai/gitHeadChanged": true,
})),
&counter,
"machinery-on",
));
eprintln!("\n[perf] ===== git refresh storm e2e (rebase turn) =====");
eprintln!(
" machinery ON : scans={} skips={} edit_turn={:?} rebase_turn={:?}",
on.scans, on.skips, on.edit_turn, on.rebase_turn
);
eprintln!(
" machinery OFF: scans={} skips={} edit_turn={:?} rebase_turn={:?}",
off.scans, off.skips, off.edit_turn, off.rebase_turn
);
eprintln!("==================================================\n");
assert_eq!(
off.scans, 0,
"without fs-watch capabilities no watcher spawns and no scans run"
);
// Merged lock cycles + in-op deferral: at least the post-op refresh runs,
// and at most one more fire lands in a live window. A regression to
// per-cycle completions or mid-op fires storms this back to one
// full-worktree scan per pick.
assert!(
(1..=2).contains(&on.scans),
"expected the merged operation to cost 1-2 scans, got {}",
on.scans
);
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,22 @@
//! Verify `requirements.toml` can pin the base sandbox `profile`.
use kigi_shell::agent::config::{ConfigSource, SandboxSettingsConfig};
#[test]
fn requirements_pin_profile() {
let config = SandboxSettingsConfig::default();
let resolved = config.resolve_profile(None, Some("strict"));
assert_eq!(resolved.value, "strict");
assert_eq!(resolved.source, ConfigSource::Requirement);
}
#[test]
fn cli_flag_overrides_config_but_not_requirement() {
let config = SandboxSettingsConfig {
profile: Some("workspace".to_string()),
..Default::default()
};
let resolved = config.resolve_profile(Some("read-only"), Some("strict"));
assert_eq!(resolved.value, "strict");
assert_eq!(resolved.source, ConfigSource::Requirement);
}
@@ -0,0 +1,812 @@
//! End-to-end measurement of why resuming a large session is slow — the time
//! spent before the client can render anything.
//!
//! The pager resumes via `session/load` and blocks on the response. The shell
//! answers by (1) `load_light` (chat history; rewind points now load lazily) and
//! (2) `replay_session_updates` — reading `updates.jsonl`, filtering it, typed-
//! parsing every line, and forwarding each as a `session/update`. All of that
//! happens while the client waits; both tests drive the real production code.
//!
//! * [`phase_breakdown_real_functions`] drives the exact load-path functions
//! (`load_session_without_updates`, `load_updates_for_replay_at`) and attributes
//! wall-clock to rewind load, chat+summary load, and updates read+parse+filter,
//! then prints a per-`sessionUpdate`-kind byte breakdown of `updates.jsonl`.
//! * [`full_session_load_e2e`] stands up a real `MvpAgent` over in-process ACP
//! pipes; times `session/load` end-to-end, counts replayed notifications, and
//! dumps the shell's own per-phase `instrumentation_timer!` events.
//!
//! Session data (both tests): a synthetic session mirroring the pathological real
//! one (redundant `available_commands_update` + big rewind snapshots; size knobs
//! via env, see [`GenOpts::from_env`]), or a real session dir via
//! `KIGI_PERF_SESSION_SRC=/path/to/<session-dir>`.
//!
//! Run:
//! cargo test -p kigi-shell --test session_load_perf -- --nocapture
//! cargo test -p kigi-shell --test session_load_perf full_session_load_e2e -- --ignored --nocapture
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use agent_client_protocol::{self as acp};
use tempfile::TempDir;
use kigi_shell::session::info::Info;
use kigi_shell::session::storage::{
JsonlStorageAdapter, StorageAdapter, load_updates_for_replay_at,
};
use kigi_workspace::session::file_state::{FileSnapshot, FlexiblePath, RewindPoint};
// ───────────────────────── size knobs ─────────────────────────
/// Generation parameters. Defaults produce a session large enough that the
/// per-phase costs are clearly measurable (tens of MB) while still finishing
/// in a few seconds. Scale up via env to approach a real heavy session.
struct GenOpts {
turns: usize,
/// `available_commands_update`s persisted per turn. The real session had
/// ~12.5 of these per turn — the slash-command catalog re-advertised on
/// every skill discovery / subagent boundary.
acu_per_turn: usize,
catalog_commands: usize,
catalog_desc_len: usize,
agent_chunks_per_turn: usize,
agent_chunk_len: usize,
rewind_points: usize,
files_per_rewind: usize,
file_content_len: usize,
}
impl GenOpts {
fn from_env() -> Self {
fn g(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
// A single multiplier for quick scaling of the dominant contributors.
let scale = g("KIGI_PERF_SCALE", 1).max(1);
Self {
turns: g("KIGI_PERF_TURNS", 80) * scale,
acu_per_turn: g("KIGI_PERF_ACU_PER_TURN", 15),
catalog_commands: g("KIGI_PERF_CATALOG_COMMANDS", 64),
catalog_desc_len: g("KIGI_PERF_CATALOG_DESC_LEN", 320),
agent_chunks_per_turn: g("KIGI_PERF_AGENT_CHUNKS_PER_TURN", 8),
agent_chunk_len: g("KIGI_PERF_AGENT_CHUNK_LEN", 2000),
rewind_points: g("KIGI_PERF_REWIND_POINTS", 60) * scale,
files_per_rewind: g("KIGI_PERF_FILES_PER_REWIND", 40),
file_content_len: g("KIGI_PERF_FILE_CONTENT_LEN", 8000),
}
}
}
// ───────────────────────── filler ─────────────────────────
/// Deterministic, non-trivially-compressible-ish filler of `n` bytes. Uses a
/// rotating word list so serde has real strings to allocate (not one repeated
/// byte), matching the cost profile of real prose/code content.
fn filler(n: usize) -> String {
const WORDS: &[&str] = &[
"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india",
"juliet", "kilo", "lima", "mike", "november", "oscar", "papa", "quebec", "romeo",
];
let mut s = String::with_capacity(n + 8);
let mut i = 0usize;
while s.len() < n {
s.push_str(WORDS[i % WORDS.len()]);
s.push(' ');
i += 1;
}
s.truncate(n);
s
}
// ───────────────────────── update synthesis ─────────────────────────
fn sid(session_id: &str) -> acp::SessionId {
acp::SessionId::new(session_id.to_string())
}
fn text_chunk(text: String) -> acp::ContentChunk {
acp::ContentChunk::new(acp::ContentBlock::Text(acp::TextContent::new(text)))
}
/// Build one large `AvailableCommandsUpdate` — the redundant catalog that the
/// real session re-persisted thousands of times.
fn available_commands_update(opts: &GenOpts) -> acp::SessionUpdate {
let desc = filler(opts.catalog_desc_len);
let commands: Vec<acp::AvailableCommand> = (0..opts.catalog_commands)
.map(|i| {
acp::AvailableCommand::new(format!("command-number-{i:03}"), desc.clone()).input(Some(
acp::AvailableCommandInput::Unstructured(acp::UnstructuredCommandInput::new(
"[optional arguments here]".to_string(),
)),
))
})
.collect();
acp::SessionUpdate::AvailableCommandsUpdate(acp::AvailableCommandsUpdate::new(commands))
}
/// Serialize one notification into the exact on-disk `updates.jsonl` envelope:
/// `{"timestamp":..,"method":"session/update","params":<SessionNotification>}`.
///
/// Params are plain JSON (not the typed `acp::SessionNotification`) so generation
/// doesn't depend on the acp crate's `_meta` field type; the production replay
/// still parses it back into a typed notification — the cost we're measuring.
fn envelope_line(session_id: &str, update: acp::SessionUpdate) -> String {
let update_val = serde_json::to_value(&update).expect("serialize update");
let params = serde_json::json!({
"sessionId": session_id,
"update": update_val,
});
let envelope = serde_json::json!({
"timestamp": 0u64,
"method": "session/update",
"params": params,
});
serde_json::to_string(&envelope).expect("serialize envelope")
}
/// Per-kind statistics for the generated/loaded updates file.
#[derive(Default)]
struct KindStats {
count: BTreeMap<String, u64>,
bytes: BTreeMap<String, u64>,
}
fn generate_updates_jsonl(path: &Path, session_id: &str, opts: &GenOpts) {
let mut out = String::new();
for turn in 0..opts.turns {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::UserMessageChunk(text_chunk(format!(
"user prompt for turn {turn}"
))),
));
out.push('\n');
for _ in 0..opts.acu_per_turn {
out.push_str(&envelope_line(session_id, available_commands_update(opts)));
out.push('\n');
}
for _ in 0..opts.agent_chunks_per_turn {
out.push_str(&envelope_line(
session_id,
acp::SessionUpdate::AgentMessageChunk(text_chunk(filler(opts.agent_chunk_len))),
));
out.push('\n');
}
}
std::fs::write(path, out).expect("write updates.jsonl");
}
fn generate_rewind_jsonl(path: &Path, opts: &GenOpts) {
let mut out = String::new();
for p in 0..opts.rewind_points {
let mut rp = RewindPoint::new(p);
for f in 0..opts.files_per_rewind {
let fp =
FlexiblePath::Absolute(PathBuf::from(format!("/repo/src/module_{p}/file_{f}.rs")));
rp.add_snapshot(FileSnapshot::new_flexible(
fp.clone(),
Some(filler(opts.file_content_len)),
));
rp.set_after_snapshot(FileSnapshot::new_flexible(
fp,
Some(filler(opts.file_content_len + 64)),
));
}
out.push_str(&serde_json::to_string(&rp).expect("serialize rewind point"));
out.push('\n');
}
std::fs::write(path, out).expect("write rewind_points.jsonl");
}
// ───────────────────────── session setup ─────────────────────────
/// Find `<root>/sessions/<enc-cwd>/<id>` without depending on the (internal)
/// cwd encoder: scan the one level of cwd dirs for a child named `<id>`.
fn locate_session_dir(root: &Path, id: &str) -> PathBuf {
let sessions = root.join("sessions");
for entry in std::fs::read_dir(&sessions)
.expect("read sessions dir")
.flatten()
{
let candidate = entry.path().join(id);
if candidate.is_dir() {
return candidate;
}
}
panic!(
"could not locate session dir for {id} under {}",
sessions.display()
);
}
/// Recursively copy a directory tree.
fn copy_tree(src: &Path, dst: &Path) {
std::fs::create_dir_all(dst).unwrap();
for entry in std::fs::read_dir(src).unwrap().flatten() {
let from = entry.path();
let to = dst.join(entry.file_name());
if from.is_dir() {
copy_tree(&from, &to);
} else {
std::fs::copy(&from, &to).unwrap();
}
}
}
/// Prepare a session on disk under `root` for working dir `cwd`. Returns the
/// `Info` and the session directory path. Uses `KIGI_PERF_SESSION_SRC` if set
/// (copies a real session), otherwise synthesizes one via the production
/// storage adapter (summary) + raw envelope writes (updates/rewind).
async fn prepare_session(root: &Path, cwd: &Path, opts: &GenOpts) -> (Info, PathBuf) {
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
if let Ok(src) = std::env::var("KIGI_PERF_SESSION_SRC") {
// Real session: create a registered session shell to get the encoded
// cwd dir + a valid summary, then overlay the real files on top.
let id = uuid::Uuid::new_v4().to_string();
let info = Info {
id: sid(&id),
cwd: cwd.to_string_lossy().to_string(),
};
adapter
.init_session(&info, acp::ModelId::new("test-model"))
.await
.expect("init_session");
let dir = locate_session_dir(root, &id);
// Copy real session files (updates/rewind/chat/etc.) over the stub,
// but keep our freshly-written summary.json (correct id + cwd + model).
for name in ["updates.jsonl", "rewind_points.jsonl", "chat_history.jsonl"] {
let from = Path::new(&src).join(name);
if from.exists() {
std::fs::copy(&from, dir.join(name)).unwrap();
}
}
// Compaction checkpoints may be referenced by replay; copy if present.
let ckpt = Path::new(&src).join("compaction_checkpoints");
if ckpt.is_dir() {
copy_tree(&ckpt, &dir.join("compaction_checkpoints"));
}
eprintln!("[perf] using REAL session copied from {src}");
return (info, dir);
}
let id = uuid::Uuid::new_v4().to_string();
let info = Info {
id: sid(&id),
cwd: cwd.to_string_lossy().to_string(),
};
adapter
.init_session(&info, acp::ModelId::new("test-model"))
.await
.expect("init_session");
let dir = locate_session_dir(root, &id);
let t = Instant::now();
generate_updates_jsonl(&dir.join("updates.jsonl"), &id, opts);
generate_rewind_jsonl(&dir.join("rewind_points.jsonl"), opts);
eprintln!(
"[perf] generated synthetic session in {} ms (turns={}, acu/turn={})",
t.elapsed().as_millis(),
opts.turns,
opts.acu_per_turn
);
(info, dir)
}
fn file_size_mb(path: &Path) -> f64 {
std::fs::metadata(path).map(|m| m.len()).unwrap_or(0) as f64 / 1e6
}
/// `(len, content_hash)` fingerprint of a file, for asserting it is byte-for-byte
/// unchanged across an operation (zero-data-loss guard). Missing file → `(0, 0)`.
fn file_fingerprint(path: &Path) -> (u64, u64) {
use std::hash::{Hash, Hasher};
let Ok(bytes) = std::fs::read(path) else {
return (0, 0);
};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
bytes.hash(&mut hasher);
(bytes.len() as u64, hasher.finish())
}
/// Per-`sessionUpdate`-kind byte + count breakdown of an `updates.jsonl`.
fn updates_kind_breakdown(path: &Path) -> KindStats {
let mut stats = KindStats::default();
let Ok(contents) = std::fs::read_to_string(path) else {
return stats;
};
for line in contents.lines() {
if line.trim().is_empty() {
continue;
}
let len = line.len() as u64 + 1;
let kind = serde_json::from_str::<serde_json::Value>(line)
.ok()
.and_then(|v| {
v.get("params")
.and_then(|p| p.get("update"))
.and_then(|u| u.get("sessionUpdate"))
.and_then(|s| s.as_str())
.map(String::from)
})
.unwrap_or_else(|| "<unparsed>".to_string());
*stats.count.entry(kind.clone()).or_default() += 1;
*stats.bytes.entry(kind).or_default() += len;
}
stats
}
fn print_kind_breakdown(label: &str, stats: &KindStats) {
let total: u64 = stats.bytes.values().sum();
eprintln!(
"\n[perf] {label}: updates.jsonl composition ({:.1} MB total):",
total as f64 / 1e6
);
eprintln!(
" {:<32} {:>8} {:>10} {:>7}",
"sessionUpdate kind", "count", "MB", "%"
);
let mut rows: Vec<(&String, &u64)> = stats.bytes.iter().collect();
rows.sort_by(|a, b| b.1.cmp(a.1));
for (kind, bytes) in rows {
let count = stats.count.get(kind).copied().unwrap_or(0);
let pct = if total > 0 {
*bytes as f64 / total as f64 * 100.0
} else {
0.0
};
eprintln!(
" {:<32} {:>8} {:>10.1} {:>6.1}%",
kind,
count,
*bytes as f64 / 1e6,
pct
);
}
}
// ───────────────────────── TEST 1: phase breakdown ─────────────────────────
/// Attribute the pre-render load cost to its real phases using the exact
/// production functions, isolating rewind-point load from everything else.
///
/// `#[ignore]`: this is a measurement tool (generates tens of MB, ~3 s), and its
/// only correctness assertion is covered by the unit tests. Run explicitly with
/// `--ignored` (optionally `KIGI_PERF_SESSION_SRC=...`) to get the numbers.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "perf measurement tool; run with --ignored"]
async fn phase_breakdown_real_functions() {
let root = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = GenOpts::from_env();
let (info, dir) = prepare_session(root.path(), cwd.path(), &opts).await;
let updates_path = dir.join("updates.jsonl");
let rewind_path = dir.join("rewind_points.jsonl");
eprintln!(
"\n[perf] session dir: {}\n[perf] updates.jsonl = {:.1} MB\n[perf] rewind_points.jsonl = {:.1} MB",
dir.display(),
file_size_mb(&updates_path),
file_size_mb(&rewind_path),
);
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
// Phase A: load_light core (summary + chat_history) — what mvp_agent's
// `load_light` blocks on before replay.
let t = Instant::now();
let light = adapter
.load_session_without_updates(&info)
.await
.expect("load_session_without_updates");
let full_load_light = t.elapsed();
// load_light no longer reads rewind_points.jsonl (deferred/lazy), so 0 by
// construction — `PersistedDataLight` has no rewind field.
let light_rewind_in_load = 0usize;
drop(light);
// Lazy rewind path (T2): the deferred cost moved here. The picker only needs
// a cheap metadata scan; an actual rewind triggers the full content load.
// Both read the same file that `load_light` no longer touches.
use kigi_workspace::session::file_state::FileStateTracker;
let t = Instant::now();
let lazy_metas = FileStateTracker::with_lazy_source(rewind_path.clone())
.get_rewind_point_metas()
.await;
let lazy_metas_scan = t.elapsed();
let t = Instant::now();
let lazy_points = FileStateTracker::with_lazy_source(rewind_path.clone())
.get_rewind_points()
.await;
let lazy_full_load = t.elapsed();
let num_rewind = lazy_points.len();
assert_eq!(
lazy_metas.len(),
num_rewind,
"picker metadata scan must see every rewind point"
);
// Phase A': isolate rewind cost — delete rewind file and re-measure. The
// delta is the rewind-point deserialization (full file-content snapshots).
std::fs::remove_file(&rewind_path).ok();
let t = Instant::now();
let _light2 = adapter
.load_session_without_updates(&info)
.await
.expect("load_session_without_updates (no rewind)");
let load_light_no_rewind = t.elapsed();
// restore for downstream/manual reruns
generate_or_restore_rewind(&rewind_path, &opts);
let rewind_cost = full_load_light.saturating_sub(load_light_no_rewind);
// Phase B: updates replay parse — production `load_updates_for_replay_at`
// reads the whole file, typed-parses every line, applies rewind filtering.
let t = Instant::now();
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
.expect("load_updates_for_replay_at")
.unwrap_or_default();
let updates_parse = t.elapsed();
let stats = updates_kind_breakdown(&updates_path);
print_kind_breakdown("phase_breakdown", &stats);
eprintln!("\n[perf] ===== PRE-RENDER LOAD PHASE BREAKDOWN (real production fns) =====");
eprintln!(" rewind_points (on disk) : {num_rewind}");
eprintln!(" rewind_points loaded in load : {light_rewind_in_load} (deferred → lazy)");
eprintln!(" updates replayed (acp) : {}", replayed.len());
eprintln!(" ----------------------------------------------------------------");
eprintln!(
" load_light (summary+chat) : {:>8.1} ms",
full_load_light.as_secs_f64() * 1e3
);
eprintln!(
" └─ rewind in load_light (now) : {:>8.1} ms",
rewind_cost.as_secs_f64() * 1e3
);
eprintln!(
" └─ summary + chat only : {:>8.1} ms",
load_light_no_rewind.as_secs_f64() * 1e3
);
eprintln!(
" lazy rewind: picker metas scan : {:>8.1} ms (on /rewind open)",
lazy_metas_scan.as_secs_f64() * 1e3
);
eprintln!(
" lazy rewind: full content load : {:>8.1} ms (on rewind execute)",
lazy_full_load.as_secs_f64() * 1e3
);
eprintln!(
" updates read+parse+filter : {:>8.1} ms",
updates_parse.as_secs_f64() * 1e3
);
eprintln!(" ----------------------------------------------------------------");
eprintln!(
" TOTAL pre-render parse work : {:>8.1} ms",
(full_load_light + updates_parse).as_secs_f64() * 1e3
);
eprintln!("================================================================\n");
assert!(!stats.bytes.is_empty(), "expected a non-empty updates file");
}
/// Re-create the rewind file after the isolation step deletes it (synthetic
/// case). For a real session copy we cannot regenerate; leave it absent.
fn generate_or_restore_rewind(path: &Path, opts: &GenOpts) {
if std::env::var("KIGI_PERF_SESSION_SRC").is_ok() {
return;
}
generate_rewind_jsonl(path, opts);
}
// ───────────────────────── TEST 2: true e2e ─────────────────────────
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use kigi_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
LineBufferedRead,
};
use kigi_shell::agent::config::Config as AgentConfig;
use kigi_shell::agent::mvp_agent::MvpAgent;
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
const DUPLEX_BUFFER_BYTES: usize = 16 * 1024 * 1024;
/// Counts replayed notifications and records first/last receipt timestamps so
/// we can see how long the client streams history before `load` returns.
#[derive(Default)]
struct LoadCounters {
count: u64,
/// `available_commands_update` notifications forwarded during the load. T1
/// skips the (thousands of) historical ones, so this must stay tiny.
acu_count: u64,
first_at: Option<Instant>,
last_at: Option<Instant>,
}
struct CountingClient {
counters: Rc<RefCell<LoadCounters>>,
}
#[async_trait::async_trait(?Send)]
impl acp::Client for CountingClient {
async fn request_permission(
&self,
args: acp::RequestPermissionRequest,
) -> acp::Result<acp::RequestPermissionResponse> {
let outcome = args
.options
.first()
.map(|o| {
acp::RequestPermissionOutcome::Selected(acp::SelectedPermissionOutcome::new(
o.option_id.clone(),
))
})
.unwrap_or(acp::RequestPermissionOutcome::Cancelled);
Ok(acp::RequestPermissionResponse::new(outcome))
}
async fn session_notification(&self, args: acp::SessionNotification) -> acp::Result<()> {
let mut c = self.counters.borrow_mut();
let now = Instant::now();
c.count += 1;
if matches!(args.update, acp::SessionUpdate::AvailableCommandsUpdate(_)) {
c.acu_count += 1;
}
c.first_at.get_or_insert(now);
c.last_at = Some(now);
Ok(())
}
}
/// Parse the production instrumentation JSON log into `(name -> elapsed_ms)`.
fn parse_instrumentation_log(path: &Path) -> Vec<(String, f64)> {
let Ok(contents) = std::fs::read_to_string(path) else {
return Vec::new();
};
let mut out = Vec::new();
for line in contents.lines() {
let Ok(v) = serde_json::from_str::<serde_json::Value>(line) else {
continue;
};
let fields = v.get("fields").unwrap_or(&v);
if fields.get("event").and_then(|e| e.as_str()) != Some("timing") {
continue;
}
let Some(name) = fields.get("name").and_then(|n| n.as_str()) else {
continue;
};
let us = fields
.get("elapsed_us")
.and_then(|u| u.as_u64())
.or_else(|| {
fields
.get("elapsed_ms")
.and_then(|m| m.as_u64())
.map(|m| m * 1000)
})
.unwrap_or(0);
out.push((name.to_string(), us as f64 / 1000.0));
}
out
}
/// True end-to-end: real `MvpAgent` over real ACP pipes. Times `session/load`
/// (what the pager blocks on), counts replayed notifications, and prints the
/// shell's own per-phase instrumentation.
///
/// `#[ignore]` by default because it stands up the full agent; run with
/// `--ignored --nocapture`.
#[tokio::test(flavor = "current_thread")]
#[ignore = "heavy: builds a full MvpAgent and replays a large session; run with --ignored"]
async fn full_session_load_e2e() {
let _ = rustls::crypto::ring::default_provider().install_default();
let server = kigi_test_support::MockInferenceServer::start()
.await
.unwrap();
let kigi_home = TempDir::new().unwrap();
let cwd = TempDir::new().unwrap();
let opts = GenOpts::from_env();
let instr_log = kigi_home.path().join("instr.jsonl");
// SAFETY: single-threaded current-thread runtime; set before any agent code
// reads these process-globals (kigi_home()/instrumentation mode are OnceLock).
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_INSTRUMENTATION", "log");
std::env::set_var("KIGI_INSTRUMENTATION_LOG", &instr_log);
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
std::env::set_var("KIGI_TRACE_UPLOAD", "false");
}
// Install the production instrumentation layer so `instrumentation_timer!`
// events are written to our temp log file.
use tracing_subscriber::Registry;
use tracing_subscriber::prelude::*;
let _ = tracing_subscriber::registry()
.with(kigi_shell::instrumentation::layer::<Registry>())
.try_init();
let (info, dir) = prepare_session(kigi_home.path(), cwd.path(), &opts).await;
let updates_path = dir.join("updates.jsonl");
let rewind_path = dir.join("rewind_points.jsonl");
eprintln!(
"\n[perf] e2e session: updates={:.1} MB rewind={:.1} MB",
file_size_mb(&updates_path),
file_size_mb(&rewind_path)
);
let stats = updates_kind_breakdown(&updates_path);
print_kind_breakdown("e2e", &stats);
// Zero-data-loss guard (C1): a pure load must never rewrite rewind_points.jsonl
// (T2 reads it lazily, never on the load path). Captured here, asserted after.
let rewind_path_guard = rewind_path.clone();
let rewind_fp_before = file_fingerprint(&rewind_path_guard);
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let agent_config = AgentConfig::default();
let auth_manager = Arc::new(agent_config.create_auth_manager());
let (gw_tx, gw_rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(gw_tx);
let agent =
MvpAgent::new(gateway, &agent_config, auth_manager, None).expect("valid config");
let (c2a_a, c2a_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
let (a2c_a, a2c_b) = tokio::io::duplex(DUPLEX_BUFFER_BYTES);
// Agent side.
let agent_incoming = LineBufferedRead::spawn_local(c2a_b.compat());
let (agent_conn, agent_io) =
acp::AgentSideConnection::new(agent, a2c_a.compat_write(), agent_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(
GatewayReceiver::new(gw_rx, agent_conn)
.with_on_meta(kigi_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
tokio::task::spawn_local(agent_io);
// Client side.
let counters = Rc::new(RefCell::new(LoadCounters::default()));
let client = CountingClient {
counters: counters.clone(),
};
let client_incoming = LineBufferedRead::spawn_local(a2c_b.compat());
let (client_conn, client_io) =
acp::ClientSideConnection::new(client, c2a_a.compat_write(), client_incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(client_io);
use acp::Agent as _;
// initialize + authenticate (api-key, like the pager does).
let init = tokio::time::timeout(
Duration::from_secs(60),
client_conn.initialize(acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(acp::ClientCapabilities::new().fs(acp::FileSystemCapabilities::new()).terminal(false)).meta(serde_json::json!({
"startupHints": { "nonInteractive": true, "skipGitStatus": true, "skipProjectLayout": true },
"clientType": "perf-test",
"clientVersion": "0.0-test",
}).as_object().cloned())),
)
.await
.expect("initialize timed out")
.expect("initialize failed");
if let Some(method) = init.auth_methods.iter().find(|m| &*m.id().0 == "xai.api_key") {
let _ = client_conn
.authenticate(acp::AuthenticateRequest::new(method.id().clone()).meta(serde_json::json!({ "headless": true }).as_object().cloned()))
.await;
}
// The measurement: time the full session/load round-trip.
let load_started = Instant::now();
let resp = tokio::time::timeout(
Duration::from_secs(180),
client_conn.load_session(acp::LoadSessionRequest::new(info.id.clone(), cwd.path().to_path_buf())),
)
.await
.expect("session/load timed out (>180s)")
.expect("session/load failed");
let load_elapsed = load_started.elapsed();
let _ = resp;
// Snapshot replay results immediately — BEFORE the post-load
// AdvertiseCommands re-advertise can arrive — so `acu_replayed` is the
// count of ACUs forwarded during history replay (the T1 skip count).
let (replay_count, acu_replayed, ttfn, ttln) = {
let c = counters.borrow();
(
c.count,
c.acu_count,
c.first_at
.map(|t| t.duration_since(load_started).as_secs_f64() * 1e3)
.unwrap_or(0.0),
c.last_at
.map(|t| t.duration_since(load_started).as_secs_f64() * 1e3)
.unwrap_or(0.0),
)
};
// The post-load `AdvertiseCommands` re-advertise (the safety basis for
// dropping historical ACUs on replay) must reach the client. It's
// enqueued at the end of `load_session` and forwarded async, so poll.
// Replay forwards 0 ACUs, so any received ACU is the re-advertise.
let readvertised = tokio::time::timeout(Duration::from_secs(10), async {
while counters.borrow().acu_count == 0 {
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.is_ok();
// Flush the instrumentation writer and read the per-phase log.
let _ = kigi_shell::instrumentation::finalize();
std::thread::sleep(Duration::from_millis(150));
let mut phases = parse_instrumentation_log(&instr_log);
phases.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
// T1 guard: the historical available_commands_update copies (3197 in
// the pathological real session, hundreds in the synthetic one) must
// NOT be replayed.
let acu_persisted = stats.count.get("available_commands_update").copied().unwrap_or(0);
eprintln!("\n[perf] ===== END-TO-END session/load (what the pager waits on) =====");
eprintln!(" total session/load round-trip : {:>9.1} ms", load_elapsed.as_secs_f64() * 1e3);
eprintln!(" notifications replayed : {:>9}", replay_count);
eprintln!(" available_commands_update : {acu_replayed:>9} replayed / {acu_persisted} on disk");
eprintln!(" post-load re-advertise reached : {readvertised:>9}");
eprintln!(" time-to-first notification : {ttfn:>9.1} ms");
eprintln!(" time-to-last notification : {ttln:>9.1} ms");
eprintln!(" ---- shell-side per-phase instrumentation (elapsed) ----");
if phases.is_empty() {
eprintln!(" (no instrumentation events captured)");
} else {
for (name, ms) in &phases {
eprintln!(" {name:<40} {ms:>9.1} ms");
}
}
eprintln!("================================================================\n");
assert!(replay_count > 0, "expected replayed notifications during load");
// C1: the lazy rewind file must be byte-for-byte unchanged by a load.
assert_eq!(
file_fingerprint(&rewind_path_guard),
rewind_fp_before,
"rewind_points.jsonl must be unchanged after a load (zero data loss)"
);
// The thousands of persisted ACUs must be skipped on replay (T1)...
assert!(
acu_persisted > 100,
"fixture should have many persisted ACUs to exercise the skip"
);
assert!(
acu_replayed < 100,
"historical available_commands_update must be skipped on replay \
(replayed {acu_replayed} of {acu_persisted} persisted)"
);
// ...but the catalog IS re-advertised to the client after load.
assert!(
readvertised,
"post-load available_commands_update re-advertise must reach the client"
);
})
.await;
}
@@ -0,0 +1,282 @@
//! Core end-to-end KEYED managed-config tests — verified persist, rejected
//! persist-nothing, and the stripped-sidecar refusal. The harness (and the
//! seam/serial constraints every test here must follow) lives in
//! `signed_managed_config/common.rs`.
//!
//! Placement rule: this binary pins the review-cited security claims
//! (verify-persists / reject-persists-nothing / sidecar-deletion-refuses); new
//! keyed scenarios go in `signed_managed_config_extended.rs` unless they alter
//! one of those three claims.
#[path = "signed_managed_config/common.rs"]
mod common;
use common::{
MANAGED, REQUIREMENTS_FAIL_CLOSED, forged_team_body, install_test_key, reset, signed_team_body,
spawn_mock, team_identity, test_home, write_config, write_team_auth,
};
use kigi_config::signed_policy;
use serial_test::serial;
/// A rejected envelope persists NOTHING: the prior principal's files survive
/// (verify-before-evict), no sidecar appears, and the marker is not rewritten.
#[tokio::test]
#[serial]
async fn rejected_signature_persists_nothing_and_records_no_marker() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
// Prior trusted state: team-b's files + marker (as if synced earlier).
std::fs::write(home.join("managed_config.toml"), "[cli]\nprior = true\n").unwrap();
std::fs::write(home.join("requirements.toml"), "[features]\n").unwrap();
kigi_shell::config::mark_managed_config_synced(kigi_shell::config::SyncMarker {
principal: Some("team-b"),
had_managed_config: true,
had_requirements: true,
key_fingerprint: None,
fail_closed: false,
});
let url = spawn_mock(forged_team_body(&kp, "team-007"));
write_config(&home, &url);
write_team_auth(&home, "team-007");
let wrote = kigi_shell::managed_config::sync()
.await
.expect("a rejected signature is a no-op, not a transport error");
assert!(!wrote, "nothing may be persisted for a rejected envelope");
assert_eq!(
std::fs::read_to_string(home.join("managed_config.toml")).unwrap(),
"[cli]\nprior = true\n",
"verify-before-evict: the prior policy must survive the identity switch"
);
assert!(home.join("requirements.toml").exists());
assert!(
!home.join("managed_config.sig.json").exists(),
"no sidecar may be written for a rejected envelope"
);
let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap();
let v: serde_json::Value = serde_json::from_str(&marker).unwrap();
assert_eq!(
v["principal"].as_str(),
Some("team-b"),
"the marker must not be rewritten for a rejected fetch: {marker}"
);
}
/// A good envelope persists the policy files AND a sidecar that verifies over the
/// exact on-disk bytes; the cache then reads fresh and the gate allows.
#[tokio::test]
#[serial]
async fn verified_envelope_persists_policy_and_sidecar() {
let home = test_home().clone();
reset(&home);
let (kp, pubkey) = install_test_key();
let url = spawn_mock(signed_team_body(
&kp,
"team-007",
Some(MANAGED),
Some(REQUIREMENTS_FAIL_CLOSED),
));
write_config(&home, &url);
write_team_auth(&home, "team-007");
let wrote = kigi_shell::managed_config::sync()
.await
.expect("a verified sync should succeed");
assert!(wrote);
let on_disk_managed = std::fs::read_to_string(home.join("managed_config.toml")).unwrap();
let on_disk_requirements = std::fs::read_to_string(home.join("requirements.toml")).unwrap();
let sidecar: serde_json::Value = serde_json::from_str(
&std::fs::read_to_string(home.join("managed_config.sig.json")).unwrap(),
)
.unwrap();
let payload = signed_policy::verify_signed_payload(
sidecar["signed_payload"].as_str().unwrap(),
sidecar["signature"].as_str().unwrap(),
&[("v1", &pubkey)],
)
.expect("the persisted sidecar must verify");
assert_eq!(
payload.managed_config.as_deref(),
Some(on_disk_managed.as_str()),
"the sidecar covers the exact on-disk managed_config bytes"
);
assert_eq!(
payload.requirements.as_deref(),
Some(on_disk_requirements.as_str()),
"the sidecar covers the exact on-disk requirements bytes"
);
assert!(payload.fail_closed, "the signed opt-in is carried");
assert!(
!kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"a covered cache is not hard-stale"
);
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"an intact verified policy must not be refused"
);
}
/// Deleting the sidecar under a fail-closed marker REFUSES at the gate (stripping it
/// must not downgrade enforcement to the forgeable marker path); the refetch triggers
/// fire so an online start self-heals.
#[tokio::test]
#[serial]
async fn deleted_sidecar_under_fail_closed_marker_refuses_at_gate() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
let url = spawn_mock(signed_team_body(
&kp,
"team-007",
Some(MANAGED),
Some(REQUIREMENTS_FAIL_CLOSED),
));
write_config(&home, &url);
write_team_auth(&home, "team-007");
kigi_shell::managed_config::sync()
.await
.expect("initial sync should succeed");
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"the covered fail-closed policy is allowed"
);
std::fs::remove_file(home.join("managed_config.sig.json")).unwrap();
assert!(
kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"a stripped sidecar must trigger the session-start refetch"
);
assert!(
kigi_shell::config::is_managed_config_stale_for(&team_identity("team-007")),
"the TIMER staleness sibling must fire too (background tick self-heal), even though the marker is timer-fresh"
);
let gate = kigi_shell::managed_config::managed_policy_gate();
assert!(
gate.is_err(),
"a fail-closed policy without its sidecar must refuse offline"
);
assert!(
gate.unwrap_err()
.contains("Managed policy is required for this account"),
"the refusal is the managed-policy gate message"
);
}
/// The keyed availability fix: after a fail_closed team-A install (signed sidecar + marker), an
/// OFFLINE switch to team B previously read Compromised (the authentic sidecar is bound to A) and
/// refused a legitimate switch. The gate's identity-change purge must shed team A's artifacts
/// INCLUDING the sidecar, PERMIT team B, and leave the cache hard-stale so the next online start
/// fetches team B's own policy.
#[tokio::test]
#[serial]
async fn offline_team_switch_purges_sidecar_and_permits_new_team() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
let url = spawn_mock(signed_team_body(
&kp,
"team-a",
Some(MANAGED),
Some(REQUIREMENTS_FAIL_CLOSED),
));
write_config(&home, &url);
write_team_auth(&home, "team-a");
kigi_shell::managed_config::sync()
.await
.expect("team A keyed sync should succeed");
assert!(
home.join("managed_config.sig.json").exists(),
"the keyed sync persists a sidecar"
);
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"team A's verified fail_closed policy must start"
);
// Switch the signed-in team to B; the gate is sync, so no fetch can rebind first.
write_team_auth(&home, "team-b");
// The bug this fixes: without the purge, team B evaluates against team A's
// foreign-bound sidecar → Compromised → a legitimate switch refused startup.
assert!(
kigi_shell::config::managed_policy_compromised_for(&team_identity("team-b")),
"pre-purge, the foreign-bound sidecar must read compromised for team B"
);
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"the gate must purge team A and permit the legitimate offline switch to team B"
);
for f in [
"requirements.toml",
"managed_config.toml",
"managed_config_cache.json",
"managed_config.sig.json",
] {
assert!(
!home.join(f).exists(),
"{f} must be purged on the identity change"
);
}
assert!(
kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-b")),
"the purged cache must read hard-stale so the next online start fetches team B's policy"
);
}
/// A blank `team_id` in `auth.json` (a parse blip) over an authentic team-A-bound fail_closed
/// sidecar: the blank→None filter resolves the identity to None, the marker principal backstops
/// the signed binding (team-a vs team-a → Trusted), so the KEYED gate PERMITS — instead of
/// binding to "" and refusing as Compromised — and nothing is purged.
#[tokio::test]
#[serial]
async fn keyed_blank_team_id_is_not_refused_and_does_not_purge() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
let url = spawn_mock(signed_team_body(
&kp,
"team-a",
Some(MANAGED),
Some(REQUIREMENTS_FAIL_CLOSED),
));
write_config(&home, &url);
write_team_auth(&home, "team-a");
kigi_shell::managed_config::sync()
.await
.expect("team A keyed sync should succeed");
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"team A's verified fail_closed policy must start"
);
// auth.json now carries a team principal with a BLANK team_id.
write_team_auth(&home, "");
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"a blank team_id must read as unknown, not a foreign binding that reads compromised"
);
for f in [
"requirements.toml",
"managed_config.toml",
"managed_config_cache.json",
"managed_config.sig.json",
] {
assert!(
home.join(f).exists(),
"{f} must be retained on a blank team_id (a parse blip is not an identity change)"
);
}
}
@@ -0,0 +1,220 @@
//! Shared harness for the KEYED managed-config integration tests: a test-only
//! signing seam injects a throwaway trusted key so the real
//! sync → verify → persist → gate paths run with verification ACTIVE (the dark
//! behavior is covered by `team_managed_config.rs`).
//!
//! Every test MUST be `#[serial]` and install its own seam keys first: the test
//! binary shares one process-global `KIGI_SHARE_DIR`, process env, and key override.
use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::OnceLock;
use base64::Engine as _;
use kigi_config::signed_policy::{self, SignedPayload};
pub const MANAGED: &str = "[cli]\ntheme = \"dark\"\n";
pub const REQUIREMENTS_FAIL_CLOSED: &str = "fail_closed = true\n[features]\nweb_fetch = false\n";
/// Far-future expiry — envelopes in these tests never expire.
pub const TEST_EXPIRES_AT: u64 = 4_000_000_000;
/// The sole trusted key id: [`install_test_key`] installs it and [`sign_envelope`]
/// signs under it, so the two can't drift.
pub const TEST_KEY_ID: &str = "v1";
/// Shared temp dir used as KIGI_SHARE_DIR for the whole test binary (the kigi_home
/// `OnceLock` only allows one value per process); scrubs the env this suite
/// depends on before any test thread reads it.
pub fn test_home() -> &'static PathBuf {
static HOME: OnceLock<PathBuf> = OnceLock::new();
HOME.get_or_init(|| {
let path = tempfile::TempDir::new().unwrap().keep();
// SAFETY: set once at init before other threads read the vars.
unsafe {
std::env::set_var("KIGI_SHARE_DIR", &path);
for var in [
"KIGI_DEPLOYMENT_KEY",
"KIGI_MANAGED_CONFIG",
"KIGI_DEPLOYMENT_CONFIG_REFRESH_INTERVAL_SECS",
"KIGI_DEPLOYMENT_CONFIG_CACHE_TTL_SECS",
"HTTP_PROXY",
"HTTPS_PROXY",
"ALL_PROXY",
"http_proxy",
"https_proxy",
"all_proxy",
] {
std::env::remove_var(var);
}
std::env::set_var("KIGI_DEPLOYMENT_CONFIG_BACKOFF_MS", "10");
}
path
})
}
pub fn reset(home: &std::path::Path) {
for f in [
"config.toml",
"auth.json",
"managed_config.toml",
"requirements.toml",
"managed_config_cache.json",
"managed_config.lock",
"managed_config.sig.json",
] {
let _ = std::fs::remove_file(home.join(f));
}
}
/// Minimal mock deployment-config server serving `body` to every request.
pub fn spawn_mock(body: String) -> String {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
std::thread::spawn(move || {
for stream in listener.incoming() {
let Ok(mut stream) = stream else { continue };
// Drain the request headers before responding.
let mut reader = BufReader::new(&mut stream);
loop {
let mut line = String::new();
if reader.read_line(&mut line).unwrap_or(0) == 0 || line.trim_end().is_empty() {
break;
}
}
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body.len(),
body
);
let _ = stream.write_all(resp.as_bytes());
let _ = stream.flush();
}
});
format!("http://{addr}/deployment/config")
}
pub fn write_config(home: &std::path::Path, managed_config_url: &str) {
std::fs::write(
home.join("config.toml"),
format!("[endpoints]\nmanaged_config_url = \"{managed_config_url}\"\n"),
)
.unwrap();
}
/// [`write_config`] plus a `deployment_key` (dead-code-allowed: compiled into
/// both binaries, called by one).
#[allow(dead_code)]
pub fn write_dk_config(home: &std::path::Path, managed_config_url: &str, deployment_key: &str) {
std::fs::write(
home.join("config.toml"),
format!(
"[endpoints]\nmanaged_config_url = \"{managed_config_url}\"\ndeployment_key = \"{deployment_key}\"\n"
),
)
.unwrap();
}
pub fn write_team_auth(home: &std::path::Path, team_id: &str) {
let scope = kigi_shell::auth::GrokComConfig::default().auth_scope();
let auth = serde_json::json!({
scope: {
"key": "team-session-token",
"auth_mode": "oidc",
"create_time": "2026-01-01T00:00:00Z",
"expires_at": "2099-01-01T00:00:00Z",
"user_id": "user-1",
"principal_type": "Team",
"team_id": team_id,
}
});
std::fs::write(home.join("auth.json"), auth.to_string()).unwrap();
}
/// A fresh Ed25519 keypair plus its raw public key, installed as the sole trusted
/// key ([`TEST_KEY_ID`]) via the test seam.
pub fn install_test_key() -> (ring::signature::Ed25519KeyPair, Vec<u8>) {
use ring::signature::KeyPair as _;
let rng = ring::rand::SystemRandom::new();
let pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&rng).unwrap();
let kp = ring::signature::Ed25519KeyPair::from_pkcs8(pkcs8.as_ref()).unwrap();
let pubkey = kp.public_key().as_ref().to_vec();
signed_policy::test_seam::set_embedded_keys(&[(TEST_KEY_ID, &pubkey)]);
assert!(
signed_policy::verification_active(),
"the seam must arm verification"
);
(kp, pubkey)
}
/// Serialize → sign → base64: the one `signatures[]` entry for `payload`, signed
/// by `kp` under the payload's own `key_id` (the untrusted outer hint can't drift
/// from the signed one).
pub fn sign_envelope(
kp: &ring::signature::Ed25519KeyPair,
payload: &SignedPayload,
) -> serde_json::Value {
let signed_payload = serde_json::to_string(payload).unwrap();
let signature = base64::engine::general_purpose::STANDARD
.encode(kp.sign(signed_payload.as_bytes()).as_ref());
serde_json::json!({
"signed_payload": signed_payload,
"signature": signature,
"key_id": payload.key_id.as_str(),
})
}
/// A team deployment-config response signed by `kp` under [`TEST_KEY_ID`]. The
/// body's legacy fields mirror the payload exactly (the client rejects a divergence).
pub fn signed_team_body(
kp: &ring::signature::Ed25519KeyPair,
team_id: &str,
managed: Option<&str>,
requirements: Option<&str>,
) -> String {
let payload = SignedPayload {
version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION,
deployment_id: None,
team_id: Some(team_id.to_owned()),
managed_config: managed.map(str::to_owned),
requirements: requirements.map(str::to_owned),
fail_closed: requirements.is_some_and(kigi_config::fail_closed_flag_from_str),
expires_at: TEST_EXPIRES_AT,
key_id: TEST_KEY_ID.into(),
};
serde_json::json!({
"deployment_id": serde_json::Value::Null,
"team_id": team_id,
"managed_config": managed,
"requirements": requirements,
"signatures": [sign_envelope(kp, &payload)],
})
.to_string()
}
/// A [`signed_team_body`] (managed config only) with the signature corrupted —
/// valid base64, wrong bytes — so the verifier must reject the envelope.
pub fn forged_team_body(kp: &ring::signature::Ed25519KeyPair, team_id: &str) -> String {
let mut body: serde_json::Value =
serde_json::from_str(&signed_team_body(kp, team_id, Some(MANAGED), None)).unwrap();
body["signatures"][0]["signature"] = base64::engine::general_purpose::STANDARD
.encode([0u8; 64])
.into();
body.to_string()
}
pub fn team_identity(id: &str) -> kigi_shell::config::ServingIdentity {
kigi_shell::config::ServingIdentity::Team(id.to_owned())
}
/// True when `path` reads despite `chmod 000` (root / DAC bypass): chmod-based
/// tests must then skip LOUDLY — a silent return would pass forever. CI runners
/// are assumed unprivileged; the shared guard keeps skips greppable.
#[cfg(unix)]
#[allow(dead_code)]
pub fn skip_as_root(path: &std::path::Path, test: &str) -> bool {
let skip = std::fs::read_to_string(path).is_ok();
if skip {
eprintln!("{test}: skipping — chmod unreadability not enforced (running as root?)");
}
skip
}
@@ -0,0 +1,296 @@
//! Extended KEYED managed-config scenarios. Harness + seam/serial constraints:
//! `signed_managed_config/common.rs`.
//!
//! Placement rule: new keyed scenarios land HERE; `signed_managed_config.rs`
//! stays fixed to the review-cited security claims (verify-persists /
//! reject-persists-nothing / sidecar-deletion-refuses).
#[path = "signed_managed_config/common.rs"]
mod common;
#[cfg(unix)]
use common::skip_as_root;
use common::{
MANAGED, REQUIREMENTS_FAIL_CLOSED, TEST_EXPIRES_AT, TEST_KEY_ID, forged_team_body,
install_test_key, reset, sign_envelope, signed_team_body, spawn_mock, team_identity, test_home,
write_config, write_dk_config, write_team_auth,
};
use kigi_config::signed_policy::{self, SignedPayload};
use serial_test::serial;
/// The healthy fail-closed starting state the tamper/heal scenarios mutate;
/// the mock keeps serving the same body, so a healing sync can refetch it.
async fn sync_fail_closed_policy(home: &std::path::Path, kp: &ring::signature::Ed25519KeyPair) {
let url = spawn_mock(signed_team_body(
kp,
"team-007",
Some(MANAGED),
Some(REQUIREMENTS_FAIL_CLOSED),
));
write_config(home, &url);
write_team_auth(home, "team-007");
kigi_shell::managed_config::sync()
.await
.expect("initial sync should succeed");
assert!(kigi_shell::managed_config::managed_policy_gate().is_ok());
}
/// The signed-empty deployment response: a `{}` body (no legacy fields) whose
/// envelope binds ABSENCE to `deployment_id` — what the server serves for a
/// provisioned key with no config row.
fn signed_dk_empty_body(kp: &ring::signature::Ed25519KeyPair, deployment_id: &str) -> String {
let payload = SignedPayload {
version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION,
deployment_id: Some(deployment_id.to_owned()),
team_id: None,
managed_config: None,
requirements: None,
fail_closed: false,
expires_at: TEST_EXPIRES_AT,
key_id: TEST_KEY_ID.into(),
};
serde_json::json!({ "signatures": [sign_envelope(kp, &payload)] }).to_string()
}
/// The marker principal for an applied signed-EMPTY dk response comes from the
/// VERIFIED payload's deployment_id (the `{}` body carries none), so the gate's
/// cross-tenant binding holds even on an unprovisioned dk machine.
#[tokio::test]
#[serial]
async fn empty_dk_response_marker_binds_the_verified_deployment_id() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
let url = spawn_mock(signed_dk_empty_body(&kp, "dep-42"));
write_dk_config(&home, &url, "dep-key-1");
// No team auth: the empty dk body is applied (converges), not fallen through.
let wrote = kigi_shell::managed_config::sync()
.await
.expect("signed-empty dk sync should succeed");
assert!(!wrote, "nothing to write for an empty row");
let marker = std::fs::read_to_string(home.join("managed_config_cache.json")).unwrap();
let v: serde_json::Value = serde_json::from_str(&marker).unwrap();
assert_eq!(
v["principal"].as_str(),
Some("dep-42"),
"the marker must bind the VERIFIED deployment id: {marker}"
);
assert!(
home.join("managed_config.sig.json").exists(),
"the absence envelope is persisted"
);
assert!(kigi_shell::managed_config::managed_policy_gate().is_ok());
}
/// A signature-rejected sync surfaces as failure in BOTH `grok setup` and the
/// post-login sync — never as Installed/NoChange while nothing was persisted.
#[tokio::test]
#[serial]
async fn rejected_signature_surfaces_as_setup_and_login_failure() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
let url = spawn_mock(forged_team_body(&kp, "team-007"));
write_config(&home, &url);
write_team_auth(&home, "team-007");
let outcome = kigi_shell::managed_config::run_setup().await;
assert!(
matches!(
outcome,
kigi_shell::managed_config::SetupOutcome::Failed(
kigi_shell::managed_config::ManagedConfigError::SignatureRejected
)
),
"setup must surface the signature rejection, got {outcome:?}"
);
let login = kigi_shell::managed_config::post_login_sync(None).await;
assert_eq!(
login,
kigi_shell::managed_config::ManagedConfigSync::Failed,
"post-login sync must report Failed, not NoChange"
);
}
/// A response that stops serving requirements deletes the on-disk file, and the
/// NEW sidecar (written after the deletion) covers the absence — the converged cache
/// reads fresh and the gate allows.
#[tokio::test]
#[serial]
async fn withdrawn_requirements_is_deleted_and_covered_by_the_new_sidecar() {
let home = test_home().clone();
reset(&home);
let (kp, pubkey) = install_test_key();
sync_fail_closed_policy(&home, &kp).await;
assert!(home.join("requirements.toml").exists());
let url_partial = spawn_mock(signed_team_body(&kp, "team-007", Some(MANAGED), None));
write_config(&home, &url_partial);
let wrote = kigi_shell::managed_config::sync()
.await
.expect("withdrawing sync should succeed");
assert!(wrote, "the deletion is a change");
assert!(
!home.join("requirements.toml").exists(),
"the withdrawn artifact is removed"
);
let sidecar: serde_json::Value = serde_json::from_str(
&std::fs::read_to_string(home.join("managed_config.sig.json")).unwrap(),
)
.unwrap();
let payload = signed_policy::verify_signed_payload(
sidecar["signed_payload"].as_str().unwrap(),
sidecar["signature"].as_str().unwrap(),
&[(TEST_KEY_ID, &pubkey)],
)
.expect("the refreshed sidecar must verify");
assert!(
payload.requirements.is_none(),
"the new sidecar covers the absence"
);
assert!(
!kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"the converged, covered cache is not hard-stale"
);
assert!(kigi_shell::managed_config::managed_policy_gate().is_ok());
}
/// A directory squatting at a signed artifact path reads COMPROMISED at the gate
/// (not lenient-unreadable), and an online sync converges over it — clearing the
/// directory, rewriting the file, and restoring enforcement.
#[tokio::test]
#[serial]
async fn directory_squat_reads_compromised_and_online_sync_heals() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
sync_fail_closed_policy(&home, &kp).await;
// Dir-squat the enforced artifact (with a child, like a real squat).
std::fs::remove_file(home.join("requirements.toml")).unwrap();
std::fs::create_dir(home.join("requirements.toml")).unwrap();
std::fs::write(home.join("requirements.toml").join("junk"), "x").unwrap();
let gate = kigi_shell::managed_config::managed_policy_gate();
assert!(
gate.is_err(),
"a directory squat on a fail-closed policy must refuse offline"
);
// The gate verdict, not an incidental error; classification is unit-pinned
// in signed_policy::directory_squat_is_tamper_not_unreadable.
assert!(
gate.unwrap_err()
.contains("Managed policy is required for this account"),
"the refusal is the managed-policy gate message"
);
assert!(
kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"the squat must trigger the refetch"
);
let wrote = kigi_shell::managed_config::sync()
.await
.expect("healing sync should succeed");
assert!(wrote, "the healing sync must rewrite the squatted artifact");
assert_eq!(
std::fs::read_to_string(home.join("requirements.toml")).unwrap(),
REQUIREMENTS_FAIL_CLOSED,
"the served file replaces the squatting directory"
);
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"enforcement is restored after the heal"
);
}
/// A sidecar read blip (chmod 000) is not tamper: the gate allows while the
/// refetch trigger fires — mirroring the artifact-slot blip semantics.
#[cfg(unix)]
#[tokio::test]
#[serial]
async fn sidecar_read_blip_allows_session_and_triggers_refetch() {
use std::os::unix::fs::PermissionsExt;
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
sync_fail_closed_policy(&home, &kp).await;
let sidecar_path = home.join("managed_config.sig.json");
std::fs::set_permissions(&sidecar_path, std::fs::Permissions::from_mode(0o000)).unwrap();
if skip_as_root(
&sidecar_path,
"sidecar_read_blip_allows_session_and_triggers_refetch",
) {
let _ = std::fs::set_permissions(&sidecar_path, std::fs::Permissions::from_mode(0o600));
return;
}
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"a transient sidecar read blip must not refuse the session"
);
assert!(
kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"the blip must trigger the refetch so the self-heal runs"
);
// Restore so the tempdir (and later tests) stay clean.
std::fs::set_permissions(&sidecar_path, std::fs::Permissions::from_mode(0o600)).unwrap();
}
/// A directory squatting at the SIDECAR path refuses at the gate, and the online
/// sync clears it — a bare rename would error forever.
#[tokio::test]
#[serial]
async fn sidecar_directory_squat_refuses_then_online_sync_heals() {
let home = test_home().clone();
reset(&home);
let (kp, _pubkey) = install_test_key();
sync_fail_closed_policy(&home, &kp).await;
// Dir-squat the sidecar (with a child, like a real squat).
let sidecar_path = home.join("managed_config.sig.json");
std::fs::remove_file(&sidecar_path).unwrap();
std::fs::create_dir(&sidecar_path).unwrap();
std::fs::write(sidecar_path.join("junk"), "x").unwrap();
let gate = kigi_shell::managed_config::managed_policy_gate();
assert!(
gate.is_err(),
"an unreadable (squatted) sidecar under a fail-closed marker must refuse offline"
);
// The gate verdict, not an incidental error; classification is unit-pinned
// in signed_policy::sidecar_directory_squat_is_absence_not_a_blip.
assert!(
gate.unwrap_err()
.contains("Managed policy is required for this account"),
"the refusal is the managed-policy gate message"
);
assert!(
kigi_shell::config::is_managed_config_hard_stale_for(&team_identity("team-007")),
"the squat must trigger the refetch"
);
kigi_shell::managed_config::sync()
.await
.expect("healing sync should succeed");
assert!(
sidecar_path.is_file(),
"the rewrite must replace the squatting directory with a sidecar FILE"
);
// Under a fail-closed marker the gate requires an authentic sidecar, so
// allowing here also pins that the healed sidecar verifies.
assert!(
kigi_shell::managed_config::managed_policy_gate().is_ok(),
"enforcement is restored after the heal"
);
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,40 @@
//! End-to-end smoke test exercising the full active_sessions lifecycle.
use chrono::Utc;
use kigi_shell::active_sessions::*;
use tempfile::TempDir;
fn session(id: &str, pid: u32) -> ActiveSession {
ActiveSession {
session_id: agent_client_protocol::SessionId::new(id),
pid,
cwd: "/tmp/test".into(),
opened_at: Utc::now(),
}
}
#[test]
fn full_lifecycle() {
let dir = TempDir::new().unwrap();
let r = dir.path();
let pid = std::process::id();
let sid = |s: &str| agent_client_protocol::SessionId::new(s);
// Start session, verify listed.
register_in(r, session("s1", pid)).unwrap();
assert_eq!(list_in(r).unwrap().len(), 1);
// Clean exit, verify gone.
unregister_in(r, &sid("s1")).unwrap();
assert!(list_in(r).unwrap().is_empty());
// Simulate crash (dead PID) + live session.
register_in(r, session("crashed", 2_000_000_000)).unwrap();
register_in(r, session("alive", pid)).unwrap();
// Crash detection finds dead PID, keeps live one.
let crashed = collect_crashed_in(r).unwrap();
assert_eq!(crashed.len(), 1);
assert_eq!(&*crashed[0].session_id.0, "crashed");
assert_eq!(list_in(r).unwrap().len(), 1);
}
@@ -0,0 +1,343 @@
//! Agent-type invariant integration tests.
//!
//! These tests exercise the full shell lifecycle via ACP stdio against a mock
//! inference server, verifying that `agent_type = f(model)` holds across:
//!
//! - Session creation with default models
//! - Zero-turn model switching (harness rebuild)
//! - Mid-session model switching (rejection)
//! - Same-type model switching (no rebuild)
//! - Session resume
//!
//! Each test spawns a real `grok agent stdio` process, speaks the full ACP
//! protocol, and asserts on the inference request bodies (system prompt) and
//! stderr tracing output to verify the correct harness was used.
//!
//! Run locally:
//! ```bash
//! cargo test -p kigi-shell --test test_agent_type_invariant -- --ignored
//! ```
use agent_client_protocol::Agent as _;
use kigi_test_support::*;
use std::future::Future;
use std::time::Duration;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Delete the shell's on-disk models cache so the next process is forced to
/// re-fetch from the mock server's `/v1/models` endpoint. Without this, the
/// second spawn in a resume test reads the stale cache written by phase 1
/// and never sees the updated model list.
fn invalidate_models_cache(home: &std::path::Path) {
let cache = home.join(".kigi").join("models_cache.json");
if cache.exists() {
std::fs::remove_file(&cache).expect("failed to delete models_cache.json");
}
}
/// Start a mock server with two models:
/// - `default-model`: no agent_type (→ defaults to "grok-build")
async fn dual_model_server() -> MockInferenceServer {
MockInferenceServer::start_with_models(vec![
MockModelEntry::new("default-model"),
MockModelEntry::with_agent_type("cursor-model", "cursor"),
])
.await
.expect("start mock server")
}
/// Start a mock server with two models that share the same agent_type:
/// - `model-a`: no agent_type (→ "grok-build")
/// - `model-b`: no agent_type (→ "grok-build")
async fn same_type_server() -> MockInferenceServer {
MockInferenceServer::start_with_models(vec![
MockModelEntry::new("model-a"),
MockModelEntry::new("model-b"),
])
.await
.expect("start mock server")
}
/// Session created with a model that has no `agent_type` should use the
/// `grok-build` harness. The system prompt sent to the LLM should contain
/// the grok-build identity string.
#[tokio::test]
#[ignore]
async fn test_default_model_uses_grok_build_harness() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let client = GrokStdioClient::spawn(&server, workdir.path()).await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "prompt failed: {:?}", result.err());
let sys_prompt = server
.last_system_prompt()
.expect("should have at least one inference request");
assert!(
sys_prompt.contains("Grok") || sys_prompt.contains("grok"),
"default model should use grok-build harness\nsystem prompt preview: {}",
&sys_prompt[..sys_prompt.len().min(500)]
);
})
.await;
}
/// Switching between two models with the same agent_type should succeed
/// without a harness rebuild.
#[tokio::test]
#[ignore]
async fn test_same_type_model_switch_no_rebuild() {
with_local_set(|| async {
let server = same_type_server().await;
let workdir = git_workdir();
let client = GrokStdioClient::spawn(&server, workdir.path()).await;
client.initialize_with_timeout().await;
let session_id = client
.create_session_with_model_timeout(workdir.path(), "model-a")
.await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "first prompt failed: {:?}", result.err());
let switch_result = client.set_model_with_timeout(&session_id, "model-b").await;
assert!(
switch_result.is_ok(),
"same-type model switch should succeed\nerror: {:?}\nstderr: {}",
switch_result.err(),
stderr_tail(&client.stderr(), 2000)
);
let result2 = client.prompt_with_timeout(&session_id, "say goodbye").await;
assert!(
result2.is_ok(),
"second prompt after model switch failed: {:?}",
result2.err()
);
})
.await;
}
/// A session created with the default model, persisted, and reloaded should
/// still use the same harness — the system prompt in the resumed session's
/// first inference request should match the original.
#[tokio::test]
#[ignore]
async fn test_session_resume_preserves_harness() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let mut writer = GrokStdioClient::spawn(&server, workdir.path()).await;
writer.initialize_with_timeout().await;
let session_id = writer.create_session_with_timeout(workdir.path()).await;
let result = writer.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "prompt failed: {:?}", result.err());
let original_sys_prompt = server
.last_system_prompt()
.expect("should have captured system prompt");
let shared_home = writer.take_home();
invalidate_models_cache(shared_home.path());
drop(writer);
let reader = GrokStdioClient::spawn_with_home(&server, workdir.path(), shared_home).await;
reader.initialize_with_timeout().await;
let _ = reader
.load_session_with_timeout(&session_id, workdir.path())
.await;
let result2 = reader.prompt_with_timeout(&session_id, "say goodbye").await;
assert!(
result2.is_ok(),
"resumed prompt failed: {:?}",
result2.err()
);
let resumed_sys_prompt = server
.last_system_prompt()
.expect("should have captured resumed system prompt");
let original_has_grok =
original_sys_prompt.contains("Grok") || original_sys_prompt.contains("grok");
let resumed_has_grok =
resumed_sys_prompt.contains("Grok") || resumed_sys_prompt.contains("grok");
assert_eq!(
original_has_grok,
resumed_has_grok,
"resumed session should use the same harness as the original\n\
original identity markers: grok={original_has_grok}\n\
resumed identity markers: grok={resumed_has_grok}\n\
original prompt (first 300): {}\n\
resumed prompt (first 300): {}",
&original_sys_prompt[..original_sys_prompt.len().min(300)],
&resumed_sys_prompt[..resumed_sys_prompt.len().min(300)],
);
})
.await;
}
/// A model that doesn't declare `agent_type` in its metadata should
/// default to `"grok-build"`. This exercises the serde default.
#[tokio::test]
#[ignore]
async fn test_model_without_agent_type_defaults_to_grok_build() {
with_local_set(|| async {
let server = MockInferenceServer::start_with_models(
vec![MockModelEntry::new("no-agent-type-model"),],
)
.await
.expect("start mock server");
let workdir = git_workdir();
let client = GrokStdioClient::spawn(&server, workdir.path()).await;
client.initialize_with_timeout().await;
let session_id = client
.create_session_with_model_timeout(workdir.path(), "no-agent-type-model")
.await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "prompt failed: {:?}", result.err());
let sys_prompt = server
.last_system_prompt()
.expect("should have at least one inference request");
assert!(
sys_prompt.contains("Grok") || sys_prompt.contains("grok"),
"model without agent_type should default to grok-build harness\nsystem prompt preview: {}",
& sys_prompt[..sys_prompt.len().min(500)]
);
})
.await;
}
/// The `KIGI_AGENT` escape hatch should override the model's agent_type.
/// Setting `KIGI_AGENT=grok-build` with an alternate-agent model should use
/// grok-build harness.
#[tokio::test]
#[ignore]
async fn test_grok_agent_env_overrides_model_agent_type() {
with_local_set(|| async {
let server = dual_model_server().await;
let workdir = git_workdir();
let binary = grok_binary();
let home = tempfile::TempDir::new().expect("create temp home");
let mut cmd = tokio::process::Command::new(&binary);
cmd.args(["agent", "stdio"])
.current_dir(workdir.path())
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
kigi_test_support::env::test_env_cmd_tokio(
&mut cmd,
&server.url(),
home.path(),
);
cmd.env("KIGI_AGENT", "grok-build");
let mut child = cmd.spawn().expect("spawn grok");
let outgoing = child.stdin.take().unwrap();
let incoming = child.stdout.take().unwrap();
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
let outgoing = outgoing.compat_write();
let incoming = incoming.compat();
let incoming = kigi_acp_lib::LineBufferedRead::spawn_local(incoming);
use agent_client_protocol as acp;
struct NoopClient;
#[async_trait::async_trait(?Send)]
impl acp::Client for NoopClient {
async fn request_permission(
&self,
args: acp::RequestPermissionRequest,
) -> acp::Result<acp::RequestPermissionResponse> {
let outcome = args
.options
.iter()
.find(|o| o.kind == acp::PermissionOptionKind::AllowOnce)
.or(args.options.first())
.map(|o| acp::RequestPermissionOutcome::Selected(
acp::SelectedPermissionOutcome::new(o.option_id.clone()),
))
.unwrap_or(acp::RequestPermissionOutcome::Cancelled);
Ok(acp::RequestPermissionResponse::new(outcome))
}
async fn session_notification(
&self,
_args: acp::SessionNotification,
) -> acp::Result<()> {
Ok(())
}
}
let (conn, handle_io) = acp::ClientSideConnection::new(
NoopClient,
outgoing,
incoming,
|fut| {
tokio::task::spawn_local(fut);
},
);
tokio::task::spawn_local(handle_io);
let _init = tokio::time::timeout(
Duration::from_secs(20),
conn
.initialize(
acp::InitializeRequest::new(acp::ProtocolVersion::V1)
.client_capabilities(
acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false),
)
.meta(
serde_json::json!(
{ "startupHints" : { "nonInteractive" : true,
"skipGitStatus" : true, "skipProjectLayout" : true },
"clientType" : "test-client", "clientVersion" : "0.0.0-test"
}
)
.as_object()
.cloned(),
),
),
)
.await
.expect("init timed out")
.expect("init failed");
conn.authenticate(
acp::AuthenticateRequest::new(acp::AuthMethodId::new("xai.api_key"))
.meta(
serde_json::json!({ "headless" : true }).as_object().cloned(),
),
)
.await
.expect("auth failed");
let session = tokio::time::timeout(
Duration::from_secs(20),
conn
.new_session(
acp::NewSessionRequest::new(workdir.path().to_path_buf())
.meta(
serde_json::json!({ "modelId" : "cursor-model" })
.as_object()
.cloned(),
),
),
)
.await
.expect("session/new timed out")
.expect("session/new failed");
let _prompt = tokio::time::timeout(
Duration::from_secs(30),
conn
.prompt(
acp::PromptRequest::new(
session.session_id.clone(),
vec![
acp::ContentBlock::Text(acp::TextContent::new("say hello"))
],
),
),
)
.await
.expect("prompt timed out")
.expect("prompt failed");
let sys_prompt = server
.last_system_prompt()
.expect("should have inference request");
assert!(
sys_prompt.contains("Grok") || sys_prompt.contains("grok"),
"KIGI_AGENT=grok-build should override cursor model's agent_type\nsystem prompt preview: {}",
& sys_prompt[..sys_prompt.len().min(500)]
);
})
.await;
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,172 @@
//! Regression test: `update_config` must not leak values from
//! `managed_config.toml` or `requirements.toml` into the user's `config.toml`.
//!
//! Bug: `update_config` used `load_effective_config()` (which merges all config
//! layers) to populate the `Config` struct, then `save_config` wrote that merged
//! result back to the user's `config.toml`. If `requirements.toml` contained
//! `auto_update = false`, any unrelated config write (theme change, model
//! preference, yolo toggle) would permanently poison the user's config.
use std::fs;
use std::path::PathBuf;
use std::sync::OnceLock;
use serial_test::serial;
/// Shared temp directory that lives for the entire test binary.
/// All tests share this as KIGI_SHARE_DIR (the `OnceLock` in kigi-config
/// only allows one value per process).
fn test_home() -> &'static PathBuf {
static HOME: OnceLock<PathBuf> = OnceLock::new();
HOME.get_or_init(|| {
let dir = tempfile::TempDir::new().unwrap();
// Keep so the directory survives the entire test process.
let path = dir.keep();
// SAFETY: called once at init before other threads touch this var.
unsafe { std::env::set_var("KIGI_SHARE_DIR", &path) };
path
})
}
/// Clean up config files between tests.
fn reset_config_files(home: &std::path::Path) {
let _ = fs::remove_file(home.join("config.toml"));
let _ = fs::remove_file(home.join("requirements.toml"));
let _ = fs::remove_file(home.join("managed_config.toml"));
}
#[tokio::test]
#[serial]
async fn update_config_does_not_leak_requirements_into_user_config() {
let home = test_home();
reset_config_files(home);
// --- Arrange ---
// User's config.toml: auto_update = true
fs::write(
home.join("config.toml"),
"[cli]\nauto_update = true\ninstaller = \"internal\"\n",
)
.unwrap();
// Enterprise requirements.toml overrides auto_update to false
fs::write(
home.join("requirements.toml"),
"[cli]\nauto_update = false\n",
)
.unwrap();
// Sanity-check: effective config should show auto_update = false
// (requirements wins over user config).
let effective = kigi_shell::config::load_effective_config().unwrap();
let effective_cfg = kigi_shell::util::config::load_config_from_toml(&effective);
assert_eq!(
effective_cfg.cli.auto_update,
Some(false),
"precondition: effective config should merge requirements (auto_update=false)"
);
// --- Act ---
// Simulate an unrelated config write (e.g. persisting a model preference).
kigi_shell::util::config::update_config(|cfg| {
cfg.models.default = Some("grok-3".to_string());
})
.await
.expect("update_config should succeed");
// --- Assert ---
// Read the user's config.toml back from disk (raw, no merge).
let raw = fs::read_to_string(home.join("config.toml")).unwrap();
let user_toml: toml::Value = toml::from_str(&raw).unwrap();
let user_cfg = kigi_shell::util::config::load_config_from_toml(&user_toml);
assert_eq!(
user_cfg.cli.auto_update,
Some(true),
"BUG REPRODUCED: auto_update in user config.toml was overwritten by \
requirements.toml value. The raw file contents:\n{raw}"
);
// Also verify the unrelated write succeeded.
assert_eq!(user_cfg.models.default.as_deref(), Some("grok-3"));
}
#[tokio::test]
#[serial]
async fn update_config_preserves_none_when_only_requirements_sets_value() {
let home = test_home();
reset_config_files(home);
// User config has no auto_update field at all
fs::write(
home.join("config.toml"),
"[cli]\ninstaller = \"internal\"\n",
)
.unwrap();
// requirements.toml sets auto_update = false
fs::write(
home.join("requirements.toml"),
"[cli]\nauto_update = false\n",
)
.unwrap();
// Write an unrelated field
kigi_shell::util::config::update_config(|cfg| {
cfg.ui.yolo = true;
})
.await
.expect("update_config should succeed");
// Read back
let raw = fs::read_to_string(home.join("config.toml")).unwrap();
let user_toml: toml::Value = toml::from_str(&raw).unwrap();
let user_cfg = kigi_shell::util::config::load_config_from_toml(&user_toml);
assert_eq!(
user_cfg.cli.auto_update, None,
"auto_update should remain absent in user config — requirements.toml \
value must not leak. Raw file:\n{raw}"
);
}
#[tokio::test]
#[serial]
async fn update_config_does_not_leak_managed_config_values() {
let home = test_home();
reset_config_files(home);
// User config has no auto_update — only installer
fs::write(
home.join("config.toml"),
"[cli]\ninstaller = \"internal\"\n",
)
.unwrap();
// managed_config.toml sets auto_update = false and channel = "stable"
fs::write(
home.join("managed_config.toml"),
"[cli]\nauto_update = false\nchannel = \"stable\"\n",
)
.unwrap();
kigi_shell::util::config::update_config(|cfg| {
cfg.models.default = Some("test-model".to_string());
})
.await
.expect("update_config should succeed");
let raw = fs::read_to_string(home.join("config.toml")).unwrap();
let user_toml: toml::Value = toml::from_str(&raw).unwrap();
let user_cfg = kigi_shell::util::config::load_config_from_toml(&user_toml);
assert_eq!(
user_cfg.cli.auto_update, None,
"auto_update from managed_config.toml leaked into user config. Raw:\n{raw}"
);
assert_eq!(
user_cfg.cli.channel, None,
"channel from managed_config.toml leaked into user config. Raw:\n{raw}"
);
}
@@ -0,0 +1,338 @@
//! End-to-end tests for the `--debug` firehose file logging.
//!
//! Runs the built grok binary against the mock inference server with a
//! caller-owned `$KIGI_SHARE_DIR`, then inspects `~/.kigi/debug/`:
//! - the `--debug` FLAG drives the firehose end to end through the master switch:
//! a live `agent` session launched with `--debug` writes a non-empty per-session
//! `~/.kigi/debug/<sessionId>.txt` with first-party content, and does NOT enable
//! sampling/instrumentation. Regression for the master switch having bundled
//! `KIGI_LOG_SAMPLING`/`KIGI_INSTRUMENTATION`, whose global `TargetFilterLayer`
//! suppressed every other target and starved the firehose.
//! - `--debug` (headless) runs cleanly without crashing arg-parsing (smoke).
//! - no `--debug` writes no firehose files.
//! - a live `agent` session (explicit `KIGI_DEBUG_LOG=1`) writes a per-session
//! `~/.kigi/debug/<sessionId>.txt` with real first-party content + `latest.txt`.
//! - `--debug-file <path>` writes one explicit file and bypasses per-session
//! routing entirely (no `~/.kigi/debug/` files).
//! - `KIGI_LOG_FILE=<path>` writes that explicit file (back-compat single file).
//!
//! Per-session content is asserted via the live `agent`, not the headless run:
//! the agent's `run_session` future runs under the `session` span (carrying
//! `session_id`), so its first-party debug events route to `<sessionId>.txt`.
//! This is the same `init_tracing_simple("agent")` path the spawned leader uses,
//! so it covers leader capture deterministically without a flaky detached
//! process. Buffered logs from runs that DO log are not lost: the firehose
//! worker guards are flushed at process exit via `debug_log::flush()` (normal +
//! signal exit paths).
//!
//! `#[ignore]` (they need a built binary). Run locally (auto-builds the pager):
//! ```bash
//! cargo test -p kigi-shell --test test_debug_logging -- --ignored
//! ```
use std::future::Future;
use std::path::{Path, PathBuf};
use std::time::Duration;
use kigi_test_support::*;
use tempfile::TempDir;
/// Run an async body inside a `LocalSet` (required by ACP's `!Send` futures).
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// The per-session firehose directory under a pinned `$KIGI_SHARE_DIR`.
fn debug_dir(home: &Path) -> PathBuf {
home.join(".kigi").join("debug")
}
/// List firehose `*.txt` files under `~/.kigi/debug` (excluding the `latest.txt`
/// symlink). Empty if the dir is missing.
fn firehose_txt_files(home: &Path) -> Vec<PathBuf> {
let Ok(entries) = std::fs::read_dir(debug_dir(home)) else {
return Vec::new();
};
entries
.flatten()
.map(|e| e.path())
.filter(|p| {
p.file_name()
.and_then(|n| n.to_str())
.is_some_and(|n| n.ends_with(".txt") && n != "latest.txt")
})
.collect()
}
/// Build a headless `grok -p` command with a pinned `$KIGI_SHARE_DIR` so the firehose
/// lands under `<home>/.kigi/debug`. Firehose env knobs are cleared so the test
/// is hermetic regardless of the developer's shell.
fn debug_cmd(
server: &MockInferenceServer,
home: &Path,
workdir: &Path,
extra: &[&str],
) -> tokio::process::Command {
let mut cmd = tokio::process::Command::new(grok_binary());
cmd.args(["-p", "say hi", "--yolo", "--output-format", "json"])
.args(extra)
.arg("--cwd")
.arg(workdir)
.current_dir(workdir)
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
kigi_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home);
// Pin the home location and drop inherited firehose toggles for determinism.
cmd.env("KIGI_SHARE_DIR", home.join(".kigi"));
cmd.env_remove("KIGI_DEBUG_LOG");
cmd.env_remove("KIGI_LOG_FILE");
cmd.env_remove("KIGI_LOG_SAMPLING");
cmd.env_remove("KIGI_HOOKS_LOG");
cmd
}
/// Poll up to 50×100ms for the per-session firehose at `path` to become non-empty
/// (its worker flushes asynchronously while the agent process stays alive), then
/// assert it carries first-party (`xai_grok`) content. Panics with the captured
/// stderr tail if it never fills. Shared by the live-agent tests.
async fn read_session_firehose_when_ready(path: &Path, client: &GrokStdioClient) -> String {
let mut content = None;
for _ in 0..50 {
if let Ok(text) = std::fs::read_to_string(path)
&& !text.is_empty()
{
content = Some(text);
break;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
let content = content.unwrap_or_else(|| {
panic!(
"no non-empty per-session firehose {path:?}\nstderr:\n{}",
stderr_tail(&client.stderr(), 800)
)
});
// The firehose filter routes first-party crate logs here; assert that rather
// than a bare non-empty check.
assert!(
content.contains("xai_grok"),
"session firehose {path:?} should contain first-party logs, got {} bytes",
content.len()
);
content
}
/// `--debug` (headless) runs cleanly: arg-parsing + the master switch + tracing
/// init don't crash. Per-session routing + content is proven deterministically by
/// the live `agent` tests (incl. `debug_flag_master_switch_enables_firehose`); a
/// headless `grok -p` client is near-silent, so its lazily-opened firehose may
/// legitimately stay empty here — file existence is intentionally not asserted.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn debug_flag_enables_firehose_without_crashing() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let cmd = debug_cmd(&server, home.path(), workdir.path(), &["--debug"]);
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(&result, "grok --debug headless", Some(&server));
assert_no_crashes(&result.stderr);
}
/// Without `--debug` (and no firehose env), no firehose files are written.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn no_debug_flag_writes_no_debug_dir() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let cmd = debug_cmd(&server, home.path(), workdir.path(), &[]);
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(&result, "grok headless (no --debug)", Some(&server));
assert!(
firehose_txt_files(home.path()).is_empty(),
"no firehose *.txt expected without --debug, found: {:?}",
firehose_txt_files(home.path())
);
}
/// A live `agent` session writes `~/.kigi/debug/<sessionId>.txt` with real
/// first-party content, and points `latest.txt` at it. This is the same
/// `init_tracing_simple("agent")` path the spawned leader uses, so it covers
/// leader capture deterministically without a flaky detached process.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn agent_session_writes_named_session_file() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let kigi_home = home.path().join(".kigi");
let kigi_home_str = kigi_home.to_string_lossy().into_owned();
let client = GrokStdioClient::spawn_with_home_and_env(
&server,
workdir.path(),
home,
&[("KIGI_DEBUG_LOG", "1"), ("KIGI_SHARE_DIR", &kigi_home_str)],
)
.await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
// New session ids are UUID v7 (filesystem-safe), so the firehose file is
// named verbatim `<sessionId>.txt`.
let sid = session_id.0.to_string();
let _ = client.prompt_with_timeout(&session_id, "say hi").await;
let session_file = kigi_home.join("debug").join(format!("{sid}.txt"));
read_session_firehose_when_ready(&session_file, &client).await;
// `latest.txt` is a sibling symlink pointing at the just-opened session
// file, so `tail -f ~/.kigi/debug/latest.txt` follows the live session.
#[cfg(unix)]
{
let link = kigi_home.join("debug").join("latest.txt");
let target = std::fs::read_link(&link)
.unwrap_or_else(|e| panic!("latest.txt should be a symlink ({link:?}): {e}"));
assert_eq!(target, Path::new(&format!("{sid}.txt")));
}
})
.await;
}
/// The `--debug` FLAG (not `KIGI_DEBUG_LOG` directly) drives the firehose end to
/// end through the master switch. Regression: the master switch used to also set
/// `KIGI_LOG_SAMPLING`/`KIGI_INSTRUMENTATION`, whose `TargetFilterLayer` globally
/// suppresses every non-matching target — starving the firehose so `--debug`
/// produced no logs. Drives a real agent session with `--debug` and asserts the
/// per-session file has first-party content (would FAIL pre-fix), and that
/// sampling/instrumentation are NOT enabled by `--debug`.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn debug_flag_master_switch_enables_firehose() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let kigi_home = home.path().join(".kigi");
let kigi_home_str = kigi_home.to_string_lossy().into_owned();
// Drive `grok --debug agent stdio`: the master switch (which runs before
// the agent dispatch) must be what enables the firehose — NOT a direct
// KIGI_DEBUG_LOG env. The spawn helper clears inherited firehose toggles,
// so the `--debug` flag is the only thing that can enable logging here.
let client = GrokStdioClient::spawn_with_home_env_and_args(
&server,
workdir.path(),
home,
&[("KIGI_SHARE_DIR", &kigi_home_str)],
&["--debug"],
)
.await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
let sid = session_id.0.to_string();
let _ = client.prompt_with_timeout(&session_id, "say hi").await;
let session_file = kigi_home.join("debug").join(format!("{sid}.txt"));
read_session_firehose_when_ready(&session_file, &client).await;
// Slimming guard: `--debug` must NOT enable sampling. The agent spawn
// clears KIGI_LOG_SAMPLING (hermetic), so the sampling layer stays off and
// `~/.kigi/logs/sampling.jsonl` is never written — the `--debug`
// set-if-unset must not flip it on (the pre-fix code did, starving the
// firehose). Instrumentation isn't checked: the harness pins
// KIGI_INSTRUMENTATION=disabled, so that assertion would be vacuous.
let sampling = kigi_home.join("logs").join("sampling.jsonl");
let len = std::fs::metadata(&sampling).map(|m| m.len()).unwrap_or(0);
assert_eq!(
len, 0,
"--debug must not enable sampling, found {len} bytes at {sampling:?}"
);
})
.await;
}
/// `--debug-file <path>` writes one explicit file and bypasses per-session
/// routing entirely (no `~/.kigi/debug/` files created).
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn debug_file_flag_writes_single_file_and_bypasses_routing() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let explicit = home.path().join("explicit-firehose.txt");
let explicit_str = explicit.to_string_lossy().into_owned();
let cmd = debug_cmd(
&server,
home.path(),
workdir.path(),
&["--debug-file", &explicit_str],
);
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(&result, "grok --debug-file", Some(&server));
assert_no_crashes(&result.stderr);
assert!(
explicit.exists(),
"explicit --debug-file path not written: {explicit:?}\nstderr tail:\n{}",
stderr_tail(&result.stderr, 800)
);
// Routing bypassed: nothing should land in the per-session debug dir.
assert!(
firehose_txt_files(home.path()).is_empty(),
"--debug-file must bypass per-session routing, found: {:?}",
firehose_txt_files(home.path())
);
}
/// `KIGI_LOG_FILE=<path>` (no `--debug`) writes that exact file (back-compat).
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn grok_log_file_explicit_path_is_written() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = TempDir::new().expect("create temp home");
let custom = home.path().join("custom-log-file.log");
let mut cmd = debug_cmd(&server, home.path(), workdir.path(), &[]);
cmd.env("KIGI_LOG_FILE", &custom);
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(&result, "grok KIGI_LOG_FILE=path", Some(&server));
assert_no_crashes(&result.stderr);
assert!(
custom.exists(),
"explicit KIGI_LOG_FILE path not written: {custom:?}\nstderr tail:\n{}",
stderr_tail(&result.stderr, 800)
);
// Single-file mode bypasses per-session routing.
assert!(
firehose_txt_files(home.path()).is_empty(),
"KIGI_LOG_FILE must bypass per-session routing, found: {:?}",
firehose_txt_files(home.path())
);
}
@@ -0,0 +1,597 @@
//! Mock-HTTP integration suite for the server-side doom-loop check on the
//! Responses API wire: trigger parsing/dedup onto
//! `ConversationResponse.doom_loop_signals`, the recovery/resample
//! contract, and a headless config-to-header lane.
//!
//! Scripts use `MockInferenceServer`'s FIFO `enqueue_response`: request N
//! consumes script N, so "turn 1 doomed, request 2 is the resample" needs no
//! content-keyed dispatch. Parse tests use non-confident triggers (channel
//! `response`, `low_logprob`, or over-threshold) so they stay orthogonal to
//! the recovery, which acts only on confident signals.
mod common;
use common::{create_test_client, test_sampler_config};
use kigi_sampler::RetryPolicy;
use kigi_sampling_types::doom_loop::{DoomLoopSignalKind, SAMPLE_CHECK_EVENT_DATA_CUMULATIVE};
use kigi_shell::sampling::{
ApiBackend, Client, ConversationItem, ConversationRequest, RequestId, SamplerActor,
SamplerHandle,
};
use kigi_test_support::sse::{
responses_api_doom_loop_check_events, responses_api_doom_loop_terminal_only_events,
responses_api_reasoning_and_text_events, responses_api_reasoning_only_events,
responses_api_with_doom_loop_frame,
};
use kigi_test_support::{MockInferenceServer, MockModelEntry, ScriptedResponse};
const MODEL: &str = "test-model";
/// A sampling client with the doom-loop check enabled (default tunables:
/// `max_threshold` 8, `max_retries` 2).
fn doom_loop_client(base_url: &str) -> Client {
let mut config = test_sampler_config(base_url, ApiBackend::Responses, &[]);
config.doom_loop_recovery = Some(Default::default());
Client::new(config).unwrap()
}
/// A sampler actor (the rung that owns retry/recovery) with the given
/// doom-loop policy. Events are fire-and-forget, so the receiver is dropped.
fn spawn_actor(base_url: &str, doom_loop_enabled: bool) -> SamplerHandle {
let mut config = test_sampler_config(base_url, ApiBackend::Responses, &[]);
if doom_loop_enabled {
config.doom_loop_recovery = Some(Default::default());
}
// Small transport budget so a broken spec fails fast instead of spinning.
let retry = RetryPolicy {
max_retries: 2,
rate_limit_retry_threshold: 2,
};
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel();
SamplerActor::spawn(config, retry, event_tx)
}
fn user_request(text: &str) -> ConversationRequest {
ConversationRequest::from_items(vec![ConversationItem::user(text)])
}
fn responses_request_count(server: &MockInferenceServer) -> usize {
server
.requests()
.iter()
.filter(|e| e.method == "POST" && e.path.contains("/responses"))
.count()
}
// ---------------------------------------------------------------------------
// Trigger parsing (live)
// ---------------------------------------------------------------------------
/// Mid-stream check frames populate `doom_loop_signals`, deduplicated across
/// the cumulative re-sends, with the label grammar fully parsed.
#[tokio::test]
async fn mid_stream_check_frames_populate_and_dedupe_signals() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_check_events(
&["tail_repetition:4@response", "tail_repetition:2@response"],
"around and around we go",
MODEL,
)),
);
let client = doom_loop_client(&server.url());
let response = client
.conversation_collect(user_request("hello"))
.await
.expect("a doomed stream still completes");
let signals = &response.doom_loop_signals;
assert_eq!(signals.len(), 2, "cumulative re-sends dedupe by raw label");
assert_eq!(signals[0].kind, DoomLoopSignalKind::TailRepetition(4));
assert_eq!(signals[0].channel, "response");
assert_eq!(signals[0].raw, "tail_repetition:4@response");
assert_eq!(signals[1].kind, DoomLoopSignalKind::TailRepetition(2));
// The doomed turn shape itself is preserved: reasoning-only.
assert!(response.assistant_text().is_empty());
}
/// The inference API's byte-exact cumulative frame parses into both signals
/// through the full HTTP/SSE client path.
#[tokio::test]
async fn byte_exact_cumulative_frame_parses_through_the_client() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_with_doom_loop_frame(
SAMPLE_CHECK_EVENT_DATA_CUMULATIVE,
"thinking",
"the answer",
MODEL,
)),
);
let client = doom_loop_client(&server.url());
let response = client
.conversation_collect(user_request("hello"))
.await
.unwrap();
let raws: Vec<&str> = response
.doom_loop_signals
.iter()
.map(|s| s.raw.as_str())
.collect();
assert_eq!(
raws,
vec!["tail_repetition:4@response", "tail_repetition:2@response"]
);
assert!(response.assistant_text().contains("the answer"));
}
/// The terminal-only copy of the signal (no mid-stream frame) also lands on
/// the response.
#[tokio::test]
async fn terminal_only_field_populates_signals() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["low_logprob@thinking"],
"brief thought",
"an ordinary answer",
MODEL,
)),
);
let client = doom_loop_client(&server.url());
let response = client
.conversation_collect(user_request("hello"))
.await
.unwrap();
assert_eq!(response.doom_loop_signals.len(), 1);
assert_eq!(
response.doom_loop_signals[0].kind,
DoomLoopSignalKind::LowLogprob
);
assert_eq!(response.doom_loop_signals[0].channel, "thinking");
assert!(response.empty_reason().is_none(), "a normal answer turn");
}
/// Malformed check frames are swallowed: the stream completes normally, the
/// answer is intact, and no signal is recorded.
#[tokio::test]
async fn malformed_check_frames_complete_cleanly_without_signals() {
let malformed = [
// triggers as a string
r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":"tail_repetition:8@thinking"}}"#,
// triggers as a number
r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":8}}"#,
// triggers as an array of objects
r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":[{"kind":"tail_repetition"}]}}"#,
// missing doom_loop_check key entirely
r#"{"type":"response.doom_loop_check","sequence_number":9}"#,
// not JSON at all (only the SSE event name identifies it)
"definitely not json",
];
let server = MockInferenceServer::start().await.unwrap();
for payload in malformed {
let events = responses_api_with_doom_loop_frame(payload, "hm", "fine", MODEL);
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
}
let client = doom_loop_client(&server.url());
for payload in malformed {
let response = client
.conversation_collect(user_request("hello"))
.await
.unwrap_or_else(|e| panic!("stream must survive malformed frame {payload}: {e}"));
assert!(
response.doom_loop_signals.is_empty(),
"no signal from malformed frame {payload}"
);
assert!(response.assistant_text().contains("fine"));
}
}
/// Unknown extra keys on a well-formed frame do not impede parsing.
#[tokio::test]
async fn unknown_extra_keys_still_parse() {
let payload = r#"{"sequence_number":7,"type":"response.doom_loop_check","doom_loop_check":{"triggers":["tail_repetition:4@response"]},"future_field":true}"#;
let server = MockInferenceServer::start().await.unwrap();
let events = responses_api_with_doom_loop_frame(payload, "hm", "fine", MODEL);
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
let response = doom_loop_client(&server.url())
.conversation_collect(user_request("hello"))
.await
.unwrap();
assert_eq!(response.doom_loop_signals.len(), 1);
assert_eq!(
response.doom_loop_signals[0].raw,
"tail_repetition:4@response"
);
}
/// No check frame and no terminal field: the signal set stays empty.
#[tokio::test]
async fn absent_field_leaves_signals_empty() {
let server = MockInferenceServer::start().await.unwrap();
let events = responses_api_reasoning_and_text_events("thinking", "hello", MODEL);
server.enqueue_response("/v1/responses", ScriptedResponse::sse(events));
let response = doom_loop_client(&server.url())
.conversation_collect(user_request("hello"))
.await
.unwrap();
assert!(response.doom_loop_signals.is_empty());
}
/// Label kinds this client version does not know are preserved verbatim as
/// `Unknown` (never dropped, never an error).
#[tokio::test]
async fn unknown_label_kinds_preserved_as_unknown() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_check_events(
&["novel_detector:9@thinking"],
"hmmm",
MODEL,
)),
);
let response = doom_loop_client(&server.url())
.conversation_collect(user_request("hello"))
.await
.unwrap();
assert_eq!(response.doom_loop_signals.len(), 1);
assert_eq!(
response.doom_loop_signals[0].kind,
DoomLoopSignalKind::Unknown("novel_detector:9".to_string())
);
assert_eq!(
response.doom_loop_signals[0].raw,
"novel_detector:9@thinking"
);
}
/// With the check disabled, the terminal field is never even parsed — the
/// policy gates all signal work, not just the header.
#[tokio::test]
async fn disabled_policy_leaves_terminal_field_unparsed() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["tail_repetition:8@thinking"],
"thinking",
"an answer",
MODEL,
)),
);
let client = create_test_client(&server.url(), ApiBackend::Responses);
let response = client
.conversation_collect(user_request("hello"))
.await
.unwrap();
assert!(response.doom_loop_signals.is_empty());
assert!(response.assistant_text().contains("an answer"));
}
// ---------------------------------------------------------------------------
// Recovery contract (the acceptance spec for the resample behavior)
// ---------------------------------------------------------------------------
/// A confident signal (`tail_repetition:8@thinking` at the default
/// `max_threshold` 8) on a completed turn is resampled once: two requests,
/// the clean second script is the accepted response, and the resample
/// request body is identical to the first — the poisoned turn's output never
/// enters the conversation.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn confident_signal_resamples_once_and_discards_poisoned_turn() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["tail_repetition:8@thinking"],
"loop loop loop",
"poisoned answer",
MODEL,
)),
);
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_reasoning_and_text_events(
"fresh thought",
"clean answer",
MODEL,
)),
);
let handle = spawn_actor(&server.url(), true);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-confident"), user_request("hello"))
.await
.expect("recovery accepts the clean resample");
assert_eq!(responses_request_count(&server), 2);
assert_eq!(response.assistant_text(), "clean answer");
assert!(
response.doom_loop_signals.is_empty(),
"the accepted response is the clean resample, not the poisoned turn"
);
let bodies = server.request_bodies();
assert_eq!(
bodies[0]["input"], bodies[1]["input"],
"the resample re-sends the same prefix; poisoned output never enters it"
);
}
/// Budget exhaustion: with `max_retries` 2, three consecutively doomed turns
/// consume the budget and the LAST doomed response is accepted as-is — the
/// turn still succeeds.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn budget_exhaustion_accepts_last_doomed_response() {
let server = MockInferenceServer::start().await.unwrap();
for _ in 0..3 {
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["tail_repetition:8@thinking"],
"loop loop loop",
"still looping answer",
MODEL,
)),
);
}
let handle = spawn_actor(&server.url(), true);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-budget"), user_request("hello"))
.await
.expect("an exhausted budget accepts the response instead of erroring");
assert_eq!(
responses_request_count(&server),
3,
"initial attempt + max_retries (2) resamples"
);
assert_eq!(response.assistant_text(), "still looping answer");
assert!(
!response.doom_loop_signals.is_empty(),
"the accepted doomed response keeps its signals (warn-only fallback)"
);
}
/// Non-confident signals never resample: threshold above `max_threshold`,
/// a non-thinking channel, and `low_logprob` are warn-only. The
/// misclassification fence for the recovery's confidence rule.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn not_confident_signals_do_not_resample() {
for trigger in [
"tail_repetition:64@thinking",
"tail_repetition:2@response",
"low_logprob@thinking",
] {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&[trigger],
"some thought",
"kept answer",
MODEL,
)),
);
let handle = spawn_actor(&server.url(), true);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-lax"), user_request("hello"))
.await
.unwrap();
assert_eq!(
responses_request_count(&server),
1,
"{trigger} is not confident and must not resample"
);
assert_eq!(response.assistant_text(), "kept answer");
assert_eq!(response.doom_loop_signals[0].raw, trigger);
}
}
/// A disabled policy ignores even a confident signal end-to-end through the
/// actor: one request, field unparsed.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn disabled_policy_ignores_confident_signal() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["tail_repetition:8@thinking"],
"loop loop loop",
"accepted anyway",
MODEL,
)),
);
let handle = spawn_actor(&server.url(), false);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-disabled"), user_request("hello"))
.await
.unwrap();
assert_eq!(responses_request_count(&server), 1);
assert_eq!(response.assistant_text(), "accepted anyway");
assert!(response.doom_loop_signals.is_empty());
}
/// A confident signal arriving mid-stream aborts the attempt and resamples.
/// The early abort itself is not externally assertable; the contract is two
/// requests and the clean final response. The poisoned script carries a
/// visible answer so the existing empty-response retry cannot mask the
/// doom-loop path.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn mid_stream_signal_aborts_and_resamples() {
let confident_frame = r#"{"type":"response.doom_loop_check","doom_loop_check":{"triggers":["tail_repetition:8@thinking"]}}"#;
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_with_doom_loop_frame(
confident_frame,
"loop loop loop",
"poisoned answer",
MODEL,
)),
);
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_reasoning_and_text_events(
"fresh thought",
"clean answer",
MODEL,
)),
);
let handle = spawn_actor(&server.url(), true);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-midstream"), user_request("hello"))
.await
.unwrap();
assert_eq!(responses_request_count(&server), 2);
assert_eq!(response.assistant_text(), "clean answer");
}
/// The doom-loop budget and the existing empty-response retry class coexist,
/// one debit each: turn 1 is doomed but NON-empty (confident trigger plus a
/// visible answer), so only the doom class can advance past it; turn 2 is
/// reasoning-only without a trigger, so only the empty class fires; turn 3
/// is the clean accept.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn doomed_then_reasoning_only_empty_coexist() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_doom_loop_terminal_only_events(
&["tail_repetition:8@thinking"],
"loop loop loop",
"poisoned answer",
MODEL,
)),
);
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_reasoning_only_events(
"empty but not doomed",
MODEL,
)),
);
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_reasoning_and_text_events(
"fresh thought",
"clean answer",
MODEL,
)),
);
let handle = spawn_actor(&server.url(), true);
let (response, _metrics) = handle
.submit_and_collect(RequestId::from("doom-coexist"), user_request("hello"))
.await
.expect("both retry classes stay within their budgets");
assert_eq!(responses_request_count(&server), 3);
assert_eq!(response.assistant_text(), "clean answer");
assert!(response.doom_loop_signals.is_empty());
}
// ---------------------------------------------------------------------------
// Headless lifecycle lane
// ---------------------------------------------------------------------------
/// `[doom_loop_recovery] enabled = true` in `config.toml` reaches the wire
/// through the real binary: the session TURN request (marked by
/// `x-grok-turn-idx`) carries the opt-in header. Aux side-queries the binary
/// also fires at `/v1/responses` (e.g. session-title generation) must NOT
/// carry it — they collect without the actor's retry loop, so an armed
/// abort there could only fail them, never resample. The recovery behavior
/// itself is covered by the mock-HTTP suite above.
///
/// `#[ignore]` (needs a built binary). Run locally (auto-builds the pager):
/// ```bash
/// cargo test -p kigi-shell --test test_doom_loop_recovery -- --ignored
/// ```
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn headless_config_enables_doom_loop_check_header() {
let models = vec![MockModelEntry::new(MODEL).with_api_backend("responses")];
let server = MockInferenceServer::start_with_models(models)
.await
.expect("start mock server");
let workdir = kigi_test_support::git_workdir();
let home = tempfile::TempDir::new().unwrap();
let kigi_home = home.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).expect("create .kigi home");
std::fs::write(
kigi_home.join("config.toml"),
"[doom_loop_recovery]\nenabled = true\n",
)
.expect("write config.toml");
let mut cmd = tokio::process::Command::new(kigi_test_support::grok_binary());
cmd.args(["-p", "say hi", "--yolo", "--output-format", "json"])
.arg("--cwd")
.arg(workdir.path())
.current_dir(workdir.path())
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
kigi_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home.path());
cmd.env("KIGI_SHARE_DIR", kigi_home);
// Don't attach to a developer's ambient leader; spawn fresh against the mock.
cmd.env_remove("KIGI_LEADER_SOCKET");
let result = kigi_test_support::run_headless_with_cmd(cmd).await;
kigi_test_support::assert_headless_success(&result, "doom-loop header e2e", Some(&server));
let requests = server.requests();
let responses_posts: Vec<_> = requests
.iter()
.filter(|e| e.method == "POST" && e.path.contains("/responses"))
.collect();
// The session turn carries `x-grok-turn-idx`; aux side-queries (session
// title, etc.) do not.
let (turns, aux): (Vec<_>, Vec<_>) = responses_posts
.into_iter()
.partition(|e| e.header("x-grok-turn-idx").is_some());
assert!(
!turns.is_empty(),
"no session turn POST /v1/responses logged; requests:\n{}",
server.request_log_summary()
);
for turn in turns {
assert_eq!(
turn.header("x-grok-doom-loop-check"),
Some("true"),
"[doom_loop_recovery] enabled must reach the turn request header; requests:\n{}",
server.request_log_summary()
);
}
for side_query in aux {
assert_eq!(
side_query.header("x-grok-doom-loop-check"),
None::<&str>,
"the session policy must not leak into aux side-query clients; requests:\n{}",
server.request_log_summary()
);
}
}
@@ -0,0 +1,96 @@
//! Mock-HTTP integration test for reasoning-only detection on the Responses
//! API wire — the trigger that drives the model doomloop.
//!
//! Spawns a `MockInferenceServer` that serves a `/v1/responses` SSE stream
//! carrying only reasoning (reasoning summary deltas, no output text, no tool
//! call) and asserts the shell sampling client classifies the collected
//! response as `EmptyReason::ReasoningOnly`. This is the exact check the
//! sampler's retry loop runs on every completed turn to decide to resample
//! (and accumulate the out-of-band streaming-capture segments verified by the
//! actor-level capture test).
mod common;
use common::create_test_client;
use kigi_sampling_types::EmptyReason;
use kigi_shell::sampling::{ApiBackend, ConversationItem, ConversationRequest};
use kigi_test_support::sse::responses_api_reasoning_only_events;
use kigi_test_support::{MockInferenceServer, ScriptedResponse};
/// A `/v1/responses` stream that streams only reasoning and finishes with no
/// visible content must be collected into a response the client classifies as
/// `EmptyReason::ReasoningOnly` — the detection that makes the shell resample
/// and spin the doomloop. Exercises the real SSE/HTTP path
/// (`conversation_collect` -> `stream_responses` -> `collect_response`).
#[tokio::test]
async fn responses_api_reasoning_only_is_classified_as_reasoning_only() {
let server = MockInferenceServer::start().await.unwrap();
server.enqueue_response(
"/v1/responses",
ScriptedResponse::sse(responses_api_reasoning_only_events(
"let me think carefully about this",
"grok-test",
)),
);
let client = create_test_client(&server.url(), ApiBackend::Responses);
let request = ConversationRequest::from_items(vec![ConversationItem::user(
"Solve this without writing anything",
)]);
// The stream completes normally — it is just empty — so collect succeeds.
let response = client
.conversation_collect(request)
.await
.expect("collect must succeed: the stream completes, the response is empty");
assert_eq!(
response.empty_reason(),
Some(EmptyReason::ReasoningOnly),
"a reasoning-only Responses stream must classify as reasoning_only",
);
assert!(response.is_empty());
// The reasoning sibling survived; the assistant is present but empty.
assert!(
response.reasoning_items().next().is_some(),
"the reasoning item must be collected as a sibling",
);
let assistant = response
.assistant()
.expect("an empty assistant is synthesized for the turn");
assert!(
assistant.content.is_empty(),
"reasoning-only means the assistant carried no visible content",
);
}
/// Negative control for the classifier: a normal text `/v1/responses` stream
/// carrying visible assistant content must NOT be classified empty —
/// `empty_reason()` is `None`, distinguishing real content from the
/// reasoning-only case above. (No recovery/resample loop is exercised here.)
#[tokio::test]
async fn normal_text_response_is_not_classified_reasoning_only() {
let server = MockInferenceServer::start().await.unwrap();
server.set_response("The answer is 42.");
let client = create_test_client(&server.url(), ApiBackend::Responses);
let response = client
.conversation_collect(ConversationRequest::from_items(vec![
ConversationItem::user("What is the answer?"),
]))
.await
.unwrap();
assert!(
response.empty_reason().is_none(),
"a normal text turn must not be classified empty",
);
let assistant = response
.assistant()
.expect("assistant present on a text turn");
assert!(
assistant.content.contains("42"),
"the text turn must carry the model's content, got: {:?}",
assistant.content,
);
}
@@ -0,0 +1,135 @@
//! Integration tests for the fork session flow.
//!
//! These tests verify the complete fork session flow:
//! 1. Fork session data with parent tracking
//! 2. Verify forked session has correct metadata
//! 3. Test worktree creation from worktree types
use agent_client_protocol as acp;
use kigi_shell::sampling::ConversationItem;
use kigi_shell::session::info::Info;
use kigi_shell::session::storage::{JsonlStorageAdapter, StorageAdapter};
use tempfile::TempDir;
/// Helper to create a test session in a temp directory
async fn create_test_session(storage: &JsonlStorageAdapter, session_id: &str, cwd: &str) -> Info {
let info = Info {
id: acp::SessionId::new(session_id),
cwd: cwd.to_string(),
};
let model_id = acp::ModelId::new("grok-code-fast-1");
storage.init_session(&info, model_id).await.unwrap();
// Add some chat messages
let msg = ConversationItem::user("Hello world");
storage.append_chat_message(&info, &msg).await.unwrap();
// Add an update
let notification = acp::SessionNotification::new(
acp::SessionId::new(session_id),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new("Test response".to_string()),
))),
);
storage
.append_update(
&info,
&kigi_shell::session::storage::SessionUpdate::Acp(Box::new(notification)),
)
.await
.unwrap();
info
}
#[tokio::test]
async fn test_fork_session_creates_new_session_with_parent_tracking() {
let temp_dir = TempDir::new().unwrap();
let storage = JsonlStorageAdapter::with_root(temp_dir.path().to_path_buf());
// Create source session
let source_info = create_test_session(&storage, "source-session-123", "/source/path").await;
let target_info = Info {
id: acp::SessionId::new("fork-session-456"),
cwd: "/new/path".to_string(),
};
let options = kigi_shell::session::storage::CopySessionOptions {
parent_session_id: Some("source-session-123".to_string()),
new_model_id: Some("grok-3".to_string()),
target_prompt_index: None,
..Default::default()
};
let result = storage
.copy_session_data(&source_info, &target_info, options)
.await
.unwrap();
// Verify result
assert_eq!(result.chat_messages_copied, 1);
assert_eq!(result.updates_copied, 1);
// Load the forked session and verify metadata
let loaded = storage.load_session(&target_info).await.unwrap();
assert_eq!(loaded.summary.info.id.to_string(), "fork-session-456");
assert_eq!(loaded.summary.info.cwd, "/new/path");
assert_eq!(loaded.summary.current_model_id, acp::ModelId::new("grok-3"));
assert_eq!(
loaded.summary.parent_session_id,
Some("source-session-123".to_string())
);
assert!(loaded.summary.forked_at.is_some());
// Verify chat history was copied
assert_eq!(loaded.chat_history.len(), 1);
// Verify updates were copied with transformed session ID
assert_eq!(loaded.updates.len(), 1);
match &loaded.updates[0] {
kigi_shell::session::storage::SessionUpdate::Acp(notification) => {
assert_eq!(notification.session_id.to_string(), "fork-session-456");
}
_ => panic!("Expected ACP update"),
}
}
#[tokio::test]
async fn test_fork_preserves_session_title() {
let temp_dir = TempDir::new().unwrap();
let storage = JsonlStorageAdapter::with_root(temp_dir.path().to_path_buf());
// Create source session
let source_info = create_test_session(&storage, "titled-session", "/source").await;
// Update source session with a title
storage
.update_session_title(&source_info, "My Important Session".to_string())
.await
.unwrap();
// Fork the session
let target_info = Info {
id: acp::SessionId::new("fork-titled"),
cwd: "/new".to_string(),
};
let options = kigi_shell::session::storage::CopySessionOptions {
parent_session_id: Some("titled-session".to_string()),
new_model_id: None,
target_prompt_index: None,
..Default::default()
};
storage
.copy_session_data(&source_info, &target_info, options)
.await
.unwrap();
// Load and verify title was preserved (generated_title is the LLM title field).
let loaded = storage.load_session(&target_info).await.unwrap();
assert_eq!(loaded.summary.display_title(), "My Important Session");
}
@@ -0,0 +1,83 @@
//! End-to-end test for the global `[models]` defaults.
//!
//! Runs the built grok binary against the mock inference server with a
//! caller-owned `$KIGI_SHARE_DIR` whose `config.toml` sets every global `[models]`
//! default. Asserts the turn succeeds with all of them set and that the
//! wire-observable one — `extra_headers` — reaches the `/v1/chat/completions`
//! request header, for a model with no per-model `[model.<id>]` override.
//!
//! The scalar defaults (temperature, top_p, max_completion_tokens, max_retries,
//! inference_idle_timeout_secs, stream_tool_calls) are exercised here to prove
//! they parse and the turn still completes; their resolution onto the model is
//! covered directly by `config.rs` unit tests. The headless turn does not
//! surface sampling params in the chat-completions body, so they are not
//! wire-asserted here.
//!
//! `#[ignore]` (needs a built binary). Run locally (auto-builds the pager):
//! ```bash
//! cargo test -p kigi-shell --test test_global_extra_headers_e2e -- --ignored
//! ```
use kigi_test_support::*;
/// Every global `[models]` default is accepted, and the wire-observable
/// `extra_headers` reaches the inference request with no per-model block in play.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn global_models_config_reaches_inference_request() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = tempfile::TempDir::new().unwrap();
let kigi_home = home.path().join(".kigi");
std::fs::create_dir_all(&kigi_home).expect("create .kigi home");
std::fs::write(
kigi_home.join("config.toml"),
r#"[models]
extra_headers = { "X-Request-Tags" = "team=example,env=prod" }
temperature = 0.5
top_p = 0.25
max_completion_tokens = 4096
max_retries = 7
inference_idle_timeout_secs = 600
stream_tool_calls = true
"#,
)
.expect("write config.toml");
let mut cmd = tokio::process::Command::new(grok_binary());
cmd.args(["-p", "say hi", "--yolo", "--output-format", "json"])
.arg("--cwd")
.arg(workdir.path())
.current_dir(workdir.path())
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
kigi_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), home.path());
cmd.env("KIGI_SHARE_DIR", kigi_home);
// Don't attach to a developer's ambient leader; spawn fresh against the mock.
cmd.env_remove("KIGI_LEADER_SOCKET");
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(&result, "global models config e2e", Some(&server));
let requests = server.requests();
let chat = requests
.iter()
.find(|e| e.method == "POST" && e.path.contains("chat/completions"))
.unwrap_or_else(|| {
panic!(
"no POST /v1/chat/completions request logged; requests:\n{}",
server.request_log_summary()
)
});
assert_eq!(
chat.header("x-request-tags"),
Some("team=example,env=prod"),
"global [models].extra_headers must reach the request header; requests:\n{}",
server.request_log_summary()
);
}
@@ -0,0 +1,330 @@
//! Repro + regression test: leader process dies → connected clients must
//! re-elect a leader and transparently restore their sessions.
//!
//! Scenario (mirrors the field report "clients see `unknown session id` after
//! the leader dies"):
//!
//! 1. Two stdio clients (`grok agent --leader stdio`) share one leader.
//! 2. Each creates its own session and completes a prompt round-trip.
//! 3. The leader is killed with SIGKILL (crash, no graceful shutdown).
//! 4. Each client's bridge must reconnect (re-electing / spawning a fresh
//! leader), replay `initialize` + `session/load`, and then prompts against
//! the ORIGINAL session IDs must succeed again.
//!
//! Tests are `#[ignore]`d by default — they require a pre-built binary:
//!
//! ```bash
//! cargo test -p kigi-shell --test test_leader_death_repro -- --ignored --nocapture
//! ```
#![cfg(unix)]
use std::time::Duration;
use agent_client_protocol::{self as acp, Agent as _};
use kigi_test_support::leader::{
LeaderStdioClient, leader_log, wait_for_live_leader, wait_for_new_leader,
wait_for_replay_notifications,
};
use kigi_test_support::*;
/// THE repro. Kill the shared leader with SIGKILL while two clients are
/// connected; both must recover their sessions on the re-elected leader.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_leader_sigkill_clients_recover_sessions() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
// ── Phase 1: two clients, one leader, two sessions ────────────
let client_a = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client_a.initialize().await;
let session_a = client_a.create_session(workdir.path()).await;
let r = client_a.prompt(&session_a, "hello from A").await;
assert!(
r.is_ok(),
"pre-crash prompt A failed: {:?}\nstderr:\n{}\nleader log:\n{}",
r.err(),
client_a.stderr_text(),
leader_log(home.path()),
);
let client_b = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client_b.initialize().await;
let session_b = client_b.create_session(workdir.path()).await;
let r = client_b.prompt(&session_b, "hello from B").await;
assert!(
r.is_ok(),
"pre-crash prompt B failed: {:?}\nstderr:\n{}\nleader log:\n{}",
r.err(),
client_b.stderr_text(),
leader_log(home.path()),
);
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(5))
.await
.expect("no live leader PID in lock file");
assert_ne!(leader_pid, client_a.child.id().unwrap_or(0));
assert_ne!(leader_pid, client_b.child.id().unwrap_or(0));
// ── Phase 2: SIGKILL the leader (simulated crash) ─────────────
let base_a = client_a.notification_count();
let base_b = client_b.notification_count();
eprintln!("killing leader pid {leader_pid}");
unsafe {
libc::kill(leader_pid as i32, libc::SIGKILL);
}
// ── Phase 3: clients must re-elect a leader and reconnect ─────
let new_pid = wait_for_new_leader(home.path(), leader_pid, Duration::from_secs(60))
.await
.unwrap_or_else(|| {
panic!(
"no new leader was elected after SIGKILL\n\
client A stderr:\n{}\nclient B stderr:\n{}\nleader log:\n{}",
client_a.stderr_text(),
client_b.stderr_text(),
leader_log(home.path()),
)
});
eprintln!("new leader elected: pid {new_pid}");
let a_reconnected =
wait_for_replay_notifications(&client_a, base_a, Duration::from_secs(60)).await;
let b_reconnected =
wait_for_replay_notifications(&client_b, base_b, Duration::from_secs(60)).await;
eprintln!("replay evidence: A={a_reconnected} B={b_reconnected}");
// ── Phase 4: prompts on the ORIGINAL session IDs must work ────
let res_a = client_a.prompt(&session_a, "after crash A").await;
let res_b = client_b.prompt(&session_b, "after crash B").await;
assert!(
res_a.is_ok(),
"client A prompt after leader crash failed: {:?}\n\
stderr:\n{}\nleader log:\n{}",
res_a.err(),
client_a.stderr_text(),
leader_log(home.path()),
);
assert!(
res_b.is_ok(),
"client B prompt after leader crash failed: {:?}\n\
stderr:\n{}\nleader log:\n{}",
res_b.err(),
client_b.stderr_text(),
leader_log(home.path()),
);
})
.await;
}
/// Single-client variant: kill -9 the leader, the lone client must re-elect
/// and restore. Narrower failure surface than the two-client test.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_leader_sigkill_single_client_recovers() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let client = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client.initialize().await;
let session = client.create_session(workdir.path()).await;
client
.prompt(&session, "hello")
.await
.expect("pre-crash prompt failed");
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(5))
.await
.expect("no live leader PID in lock file");
let base = client.notification_count();
eprintln!("killing leader pid {leader_pid}");
unsafe {
libc::kill(leader_pid as i32, libc::SIGKILL);
}
let new_pid = wait_for_new_leader(home.path(), leader_pid, Duration::from_secs(60))
.await
.unwrap_or_else(|| {
panic!(
"no new leader was elected after SIGKILL\nstderr:\n{}\nleader log:\n{}",
client.stderr_text(),
leader_log(home.path()),
)
});
eprintln!("new leader elected: pid {new_pid}");
let reconnected =
wait_for_replay_notifications(&client, base, Duration::from_secs(60)).await;
eprintln!("replay evidence: {reconnected}");
let res = client.prompt(&session, "after crash").await;
assert!(
res.is_ok(),
"prompt after leader crash failed: {:?}\nstderr:\n{}\nleader log:\n{}",
res.err(),
client.stderr_text(),
leader_log(home.path()),
);
})
.await;
}
/// One client driving TWO sessions over a single stdio bridge (the IDE
/// shape). After a leader SIGKILL, BOTH sessions must be replayed onto the
/// re-elected leader — restoring only the most recent one left the other
/// failing with "unknown session id".
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_leader_sigkill_multi_session_client_recovers_all_sessions() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let client = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client.initialize().await;
let session_one = client.create_session(workdir.path()).await;
client
.prompt(&session_one, "hello one")
.await
.expect("pre-crash prompt on session one failed");
let session_two = client.create_session(workdir.path()).await;
client
.prompt(&session_two, "hello two")
.await
.expect("pre-crash prompt on session two failed");
assert_ne!(session_one.0, session_two.0);
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(5))
.await
.expect("no live leader PID in lock file");
let base = client.notification_count();
eprintln!("killing leader pid {leader_pid}");
unsafe {
libc::kill(leader_pid as i32, libc::SIGKILL);
}
wait_for_new_leader(home.path(), leader_pid, Duration::from_secs(60))
.await
.unwrap_or_else(|| {
panic!(
"no new leader was elected after SIGKILL\nstderr:\n{}\nleader log:\n{}",
client.stderr_text(),
leader_log(home.path()),
)
});
wait_for_replay_notifications(&client, base, Duration::from_secs(60)).await;
// BOTH sessions must work on the new leader.
let res_one = client.prompt(&session_one, "after crash one").await;
let res_two = client.prompt(&session_two, "after crash two").await;
assert!(
res_one.is_ok(),
"session one prompt after crash failed: {:?}\nstderr:\n{}\nleader log:\n{}",
res_one.err(),
client.stderr_text(),
leader_log(home.path()),
);
assert!(
res_two.is_ok(),
"session two prompt after crash failed: {:?}\nstderr:\n{}\nleader log:\n{}",
res_two.err(),
client.stderr_text(),
leader_log(home.path()),
);
})
.await;
}
/// Prompt sent DURING the outage (after the bridge noticed the dead leader
/// but before the new one is ready). The stdio bridge must hold and deliver
/// it once the session is restored — not silently drop it (which left the
/// client's request hanging forever).
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_prompt_sent_during_outage_is_delivered_after_recovery() {
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let client = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client.initialize().await;
let session = client.create_session(workdir.path()).await;
client
.prompt(&session, "hello")
.await
.expect("pre-crash prompt failed");
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(5))
.await
.expect("no live leader PID in lock file");
eprintln!("killing leader pid {leader_pid}");
unsafe {
libc::kill(leader_pid as i32, libc::SIGKILL);
}
// Give the bridge a moment to observe the dead socket (its send
// channel closes), then prompt mid-outage: re-election + session
// restore are still seconds away.
tokio::time::sleep(Duration::from_millis(300)).await;
let res = tokio::time::timeout(
Duration::from_secs(90),
client.conn.prompt(acp::PromptRequest::new(session.clone(), vec![acp::ContentBlock::Text(acp::TextContent::new("sent during outage".to_string()))])),
)
.await
.unwrap_or_else(|_| {
panic!(
"prompt sent during outage never completed (dropped by bridge?)\n\
stderr:\n{}\nleader log:\n{}",
client.stderr_text(),
leader_log(home.path()),
)
});
assert!(
res.is_ok(),
"prompt sent during outage failed: {:?}\nstderr:\n{}\nleader log:\n{}",
res.err(),
client.stderr_text(),
leader_log(home.path()),
);
// A session-scoped request other than prompt (model switch) must
// also survive — same "unknown session id" class.
let set_model = tokio::time::timeout(
Duration::from_secs(30),
client.conn.set_session_model(acp::SetSessionModelRequest::new(session.clone(), acp::ModelId::new("test-model"))),
)
.await
.unwrap_or_else(|_| {
panic!(
"set_session_model after recovery never completed\nstderr:\n{}\nleader log:\n{}",
client.stderr_text(),
leader_log(home.path()),
)
});
assert!(
set_model.is_ok(),
"set_session_model after recovery failed: {:?}\nstderr:\n{}\nleader log:\n{}",
set_model.err(),
client.stderr_text(),
leader_log(home.path()),
);
})
.await;
}
@@ -0,0 +1,400 @@
//! Leader soak: an in-process leader server fronting a REAL `MvpAgent`, hammered
//! by churning `LeaderClient`s until a time budget expires. Asserts the leader
//! neither leaks memory nor accumulates zombie clients, and that no response is
//! ever dropped on a live-client send (`leader.response.send_failed`).
//!
//! Duration is bounded by `LEADER_SOAK_SECS` (default 10s so an ad-hoc
//! `--ignored` run stays quick). RSS growth is bounded by
//! `LEADER_SOAK_MAX_RSS_GROWTH_MB` (default 1024). On-demand today — no CI
//! lane runs it; a real soak is the long form:
//!
//! ```bash
//! LEADER_SOAK_SECS=1200 cargo test -p kigi-shell --test test_leader_soak -- --ignored --nocapture
//! ```
#![cfg(unix)]
use std::sync::Arc;
use std::time::Duration;
use agent_client_protocol as acp;
use kigi_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
LineBufferedRead,
};
use kigi_shell::agent::config::Config as AgentConfig;
use kigi_shell::agent::mvp_agent::MvpAgent;
use kigi_shell::leader::{
ClientCapabilities, ClientMode, LeaderClient, LeaderServerControlState, LeaderServerMetadata,
run_leader_server,
};
use tempfile::TempDir;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
use tokio_util::sync::CancellationToken;
const SIMPLEX_BUF: usize = 8 * 1024 * 1024;
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
/// Resident set size of THIS process (leader server + agent are in-process).
/// Copied from `kigi-codebase-graph/tests/memory_integration.rs`.
fn rss_bytes() -> Option<usize> {
#[cfg(target_os = "linux")]
{
let status = std::fs::read_to_string("/proc/self/status").ok()?;
for line in status.lines() {
if let Some(val) = line.strip_prefix("VmRSS:") {
let kb: usize = val.trim().trim_end_matches(" kB").trim().parse().ok()?;
return Some(kb * 1024);
}
}
None
}
#[cfg(target_os = "macos")]
{
use std::process::Command;
let output = Command::new("ps")
.args(["-o", "rss=", "-p", &std::process::id().to_string()])
.output()
.ok()?;
let kb: usize = String::from_utf8_lossy(&output.stdout)
.trim()
.parse()
.ok()?;
Some(kb * 1024)
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
{
None
}
}
/// `leader.response.send_failed` entries written by THIS process.
fn send_failed_count() -> usize {
let Some(bytes) = kigi_log::unified_log::snapshot_log() else {
return 0;
};
String::from_utf8_lossy(&bytes)
.lines()
.filter(|line| {
serde_json::from_str::<serde_json::Value>(line).is_ok_and(|entry| {
entry["msg"] == "leader.response.send_failed" && entry["pid"] == std::process::id()
})
})
.count()
}
/// Send one JSON-RPC request through a `LeaderClient` and await the response
/// with the matching id, skipping interleaved notifications.
async fn rpc(client: &mut LeaderClient, payload: String, id: u64, what: &str) -> serde_json::Value {
client
.send(payload)
.unwrap_or_else(|e| panic!("{what}: send failed: {e}"));
let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
loop {
let remaining = deadline
.checked_duration_since(tokio::time::Instant::now())
.unwrap_or_else(|| panic!("{what}: timed out waiting for response id {id}"));
let msg = tokio::time::timeout(remaining, client.recv())
.await
.unwrap_or_else(|_| panic!("{what}: timed out waiting for response id {id}"))
.unwrap_or_else(|| panic!("{what}: connection closed awaiting response id {id}"));
let json: serde_json::Value = match serde_json::from_str(&msg) {
Ok(v) => v,
Err(_) => continue,
};
if json["id"] == id && (json.get("result").is_some() || json.get("error").is_some()) {
assert!(
json.get("error").is_none(),
"{what}: error response: {json}"
);
return json;
}
}
}
#[tokio::test(flavor = "current_thread")]
#[ignore = "leader soak; run with --ignored (LEADER_SOAK_SECS bounds the duration)"]
async fn leader_soak_churning_clients_no_leaks_no_zombies() {
let _ = rustls::crypto::ring::default_provider().install_default();
let server = kigi_test_support::MockInferenceServer::start()
.await
.unwrap();
let kigi_home = TempDir::new().unwrap();
let workdir = TempDir::new().unwrap();
// SAFETY: single-threaded current-thread runtime; set before any agent
// code reads these process-globals (same pattern as session_load_perf).
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
std::env::set_var("KIGI_FEEDBACK_ENABLED", "false");
std::env::set_var("KIGI_TRACE_UPLOAD", "false");
}
let sock_path = kigi_home.path().join("leader-soak.sock");
let soak_secs = env_u64("LEADER_SOAK_SECS", 10);
let max_growth_mb = env_u64("LEADER_SOAK_MAX_RSS_GROWTH_MB", 1024);
let send_failed_before = send_failed_count();
let local = tokio::task::LocalSet::new();
local
.run_until(async {
// ── Leader server (survives client churn) ────────────────────
let (acp_tx, mut acp_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
let (response_tx, response_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
let cancel = CancellationToken::new();
let client_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let control_state = LeaderServerControlState::new(LeaderServerMetadata {
pid: std::process::id(),
socket_path: sock_path.clone(),
lock_path: sock_path.with_extension("lock"),
socket_suffix: String::new(),
leader_binary_version: env!("CARGO_PKG_VERSION").to_string(),
});
let cancel_for_server = cancel.clone();
let sock_for_server = sock_path.clone();
let client_count_for_server = client_count.clone();
tokio::task::spawn_local(async move {
let _ = run_leader_server(
sock_for_server,
acp_tx,
response_rx,
cancel_for_server,
true,
client_count_for_server,
Arc::new(std::sync::atomic::AtomicBool::new(false)),
kigi_shell::agent::activity::AgentActivity::default(),
tokio::sync::watch::channel(true).1,
tokio::sync::watch::channel(kigi_shell::leader::ShutdownReason::Manual).0,
None,
control_state,
)
.await;
});
// ── Real agent behind it ──────────────────────────────────────
// Copied from `run_leader`'s agent-spawn + IPC/stdout bridge
// blocks in src/agent/app.rs (inside its LocalSet body); kept as
// a deliberate copy so production stays untouched. Second copy of
// the same wiring: kigi-tui/src/app/leader_cluster/mod.rs
// (`spawn_leader_generation`) — keep the two copies behaviorally
// identical.
let (agent_in_read, agent_in_write) = tokio::io::simplex(SIMPLEX_BUF);
let (agent_out_read, agent_out_write) = tokio::io::simplex(SIMPLEX_BUF);
tokio::task::spawn_local(async move {
let agent_config = AgentConfig::default();
let auth_manager = Arc::new(agent_config.create_auth_manager());
let (gw_tx, gw_rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(gw_tx);
let agent = MvpAgent::new(gateway, &agent_config, auth_manager, None)
.expect("valid agent config");
let incoming = LineBufferedRead::spawn_local(agent_in_read.compat());
let (conn, handle_io) = acp::AgentSideConnection::new(
agent,
agent_out_write.compat_write(),
incoming,
|fut| {
tokio::task::spawn_local(fut);
},
);
tokio::task::spawn_local(
GatewayReceiver::new(gw_rx, conn)
.with_on_meta(kigi_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
let _ = handle_io.await;
});
// Leader → agent stdin.
tokio::task::spawn_local(async move {
let mut agent_in_write = agent_in_write;
while let Some(msg) = acp_rx.recv().await {
if agent_in_write.write_all(msg.as_bytes()).await.is_err()
|| agent_in_write.write_all(b"\n").await.is_err()
{
break;
}
}
});
// Agent stdout → leader responses.
let response_tx_for_agent = response_tx.clone();
tokio::task::spawn_local(async move {
let mut reader = BufReader::new(agent_out_read);
let mut line = String::new();
loop {
line.clear();
match reader.read_line(&mut line).await {
Ok(0) => break,
Ok(_) => {
let msg = line.trim_end_matches(['\r', '\n']).to_string();
if !msg.is_empty() {
let _ = response_tx_for_agent.send(msg);
}
}
Err(_) => break,
}
}
});
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
while !sock_path.exists() && tokio::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(20)).await;
}
assert!(sock_path.exists(), "leader socket never bound");
// ── One-time initialize + authenticate through the leader ────
let mut bootstrap = LeaderClient::connect(
sock_path.clone(),
"soak-bootstrap",
ClientMode::Stdio,
ClientCapabilities::default(),
)
.await
.expect("bootstrap connect");
rpc(
&mut bootstrap,
r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":1,"clientCapabilities":{"fs":{"readTextFile":false,"writeTextFile":false},"terminal":false},"_meta":{"startupHints":{"nonInteractive":true,"skipGitStatus":true,"skipProjectLayout":true},"clientType":"soak","clientVersion":"0.0.0-test"}}}"#.to_string(),
1,
"initialize",
)
.await;
rpc(
&mut bootstrap,
r#"{"jsonrpc":"2.0","id":2,"method":"authenticate","params":{"methodId":"xai.api_key","_meta":{"headless":true}}}"#.to_string(),
2,
"authenticate",
)
.await;
let rss_baseline = rss_bytes();
let soak_deadline = tokio::time::Instant::now() + Duration::from_secs(soak_secs);
let workdir_str = workdir.path().to_string_lossy().to_string();
let mut cycles: u64 = 0;
let mut turns: u64 = 0;
// ── Churn: 10 fresh clients per cycle, 2 sessions each, one
// scripted turn per session, then all disconnect ───────────────
while tokio::time::Instant::now() < soak_deadline {
cycles += 1;
let mut clients = Vec::new();
for i in 0..10u64 {
let client = LeaderClient::connect(
sock_path.clone(),
"soak-client",
ClientMode::Stdio,
ClientCapabilities::default(),
)
.await
.unwrap_or_else(|e| panic!("cycle {cycles} client {i} connect: {e}"));
clients.push(client);
}
for (i, client) in clients.iter_mut().enumerate() {
for s in 0..2u64 {
let new_id = 100 + s;
let resp = rpc(
client,
format!(
r#"{{"jsonrpc":"2.0","id":{new_id},"method":"session/new","params":{{"cwd":"{workdir_str}","mcpServers":[]}}}}"#
),
new_id,
"session/new",
)
.await;
let sid = resp["result"]["sessionId"]
.as_str()
.unwrap_or_else(|| panic!("no sessionId in {resp}"))
.to_string();
let prompt_id = 200 + s;
rpc(
client,
format!(
r#"{{"jsonrpc":"2.0","id":{prompt_id},"method":"session/prompt","params":{{"sessionId":"{sid}","prompt":[{{"type":"text","text":"soak c{i} s{s} cycle {cycles}"}}]}}}}"#
),
prompt_id,
"session/prompt",
)
.await;
turns += 1;
}
}
// Churn: everyone disconnects; the roster must drain fully.
for client in clients {
client.cancel();
}
let drain_deadline = tokio::time::Instant::now() + Duration::from_secs(30);
while client_count.load(std::sync::atomic::Ordering::Relaxed) > 1 {
assert!(
tokio::time::Instant::now() < drain_deadline,
"cycle {cycles}: roster kept {} zombie clients after churn",
client_count.load(std::sync::atomic::Ordering::Relaxed)
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
eprintln!("[soak] {cycles} cycles, {turns} turns in {soak_secs}s budget");
assert!(cycles > 0, "soak budget too small to complete one cycle");
// ── Convergence: only the bootstrap client remains, and the
// leader still serves a healthy round-trip ────────────────────
assert_eq!(
client_count.load(std::sync::atomic::Ordering::Relaxed),
1,
"roster must converge to the bootstrap client after churn"
);
let resp = rpc(
&mut bootstrap,
format!(
r#"{{"jsonrpc":"2.0","id":900,"method":"session/new","params":{{"cwd":"{workdir_str}","mcpServers":[]}}}}"#
),
900,
"post-soak session/new",
)
.await;
assert!(resp["result"]["sessionId"].is_string());
// ── No response was ever dropped on a live-client send ────────
assert_eq!(
send_failed_count(),
send_failed_before,
"leader.response.send_failed must not occur during the soak"
);
// ── RSS bound ─────────────────────────────────────────────────
if let (Some(before), Some(after)) = (rss_baseline, rss_bytes()) {
let growth_mb = after.saturating_sub(before) as f64 / (1024.0 * 1024.0);
eprintln!(
"[soak] rss: {:.1} MB -> {:.1} MB (growth {growth_mb:.1} MB)",
before as f64 / (1024.0 * 1024.0),
after as f64 / (1024.0 * 1024.0),
);
assert!(
growth_mb < max_growth_mb as f64,
"leader RSS grew {growth_mb:.1} MB over the soak (bound {max_growth_mb} MB)"
);
} else {
eprintln!("[soak] rss measurement unavailable on this platform; bound skipped");
}
bootstrap.cancel();
cancel.cancel();
})
.await;
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,408 @@
//! Two-binary version-skew tests: a real OLD released binary and a real NEW
//! binary sharing one leader socket. This is the only harness that exercises
//! cross-version eviction with real processes.
//!
//! Binaries are resolved per role:
//! - `KIGI_BINARY_LEADER` — the binary that elects the initial leader
//! (typically the latest released stable, e.g. fetched from
//! `https://storage.googleapis.com/grok-build-public-artifacts/cli/grok-<ver>-linux-x86_64`).
//! - `KIGI_BINARY_CLIENT` — the second client (typically a freshly built main).
//!
//! All tests are `#[ignore]`d: they need two pre-built binaries and spawn real
//! leader subprocesses. On-demand today — no CI lane runs them; invoke with:
//!
//! ```bash
//! KIGI_BINARY_LEADER=/path/to/grok-old KIGI_BINARY_CLIENT=/path/to/grok-new \
//! cargo test -p kigi-shell --test test_leader_version_skew -- --ignored --nocapture
//! ```
#![cfg(unix)]
use std::path::Path;
use std::time::Duration;
use kigi_shell::leader::{
ClientCapabilities, ClientMode, ControlCommand, ControlPayload, LeaderClient,
};
use kigi_test_support::leader::{
LeaderStdioClient, client_binary, leader_binary, leader_log, pid_alive, read_leader_pid,
wait_for_live_leader, wait_for_new_leader, wait_for_replay_notifications,
};
use kigi_test_support::*;
/// Skew tests are meaningless when both roles resolve to the same binary
/// (e.g. a local `--ignored` run without the env vars): the version floor
/// never trips. Skip loudly instead of failing.
fn skew_binaries() -> Option<(std::path::PathBuf, std::path::PathBuf)> {
let old = leader_binary();
let new = client_binary();
if old == new {
eprintln!(
"SKIP: KIGI_BINARY_LEADER/KIGI_BINARY_CLIENT resolve to the same binary ({})",
old.display()
);
return None;
}
Some((old, new))
}
async fn wait_for_pid_death(pid: u32, timeout: Duration) -> bool {
let deadline = tokio::time::Instant::now() + timeout;
while tokio::time::Instant::now() < deadline {
if !pid_alive(pid) {
return true;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
false
}
fn sandbox_unified_log(home: &Path) -> String {
std::fs::read_to_string(home.join(".kigi").join("logs").join("unified.jsonl"))
.unwrap_or_default()
}
/// End-to-end version-skew: an old leader is running; a newer client connects,
/// evicts it under the version floor, spawns a replacement from its own
/// binary, and the old client's session survives via reconnect + reload.
#[tokio::test]
#[ignore = "two-binary version-skew test; set KIGI_BINARY_LEADER/KIGI_BINARY_CLIENT and run with --ignored"]
async fn new_client_evicts_old_leader_and_sessions_reload() {
let Some((old_bin, new_bin)) = skew_binaries() else {
return;
};
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
// Old binary elects the leader and completes a turn.
let old_client = LeaderStdioClient::spawn_with_binary(
&old_bin,
&server,
workdir.path(),
home.path(),
)
.await;
old_client.initialize().await;
let session = old_client.create_session(workdir.path()).await;
old_client
.prompt(&session, "hello from the old world")
.await
.expect("pre-skew prompt failed");
let old_pid = wait_for_live_leader(home.path(), Duration::from_secs(10))
.await
.expect("no live old leader");
let base = old_client.notification_count();
// New binary connects: version floor → evict → respawn.
let new_client = LeaderStdioClient::spawn_with_binary(
&new_bin,
&server,
workdir.path(),
home.path(),
)
.await;
new_client.initialize().await;
let new_pid = wait_for_new_leader(home.path(), old_pid, Duration::from_secs(60))
.await
.unwrap_or_else(|| {
panic!(
"no replacement leader after version-floor eviction\n\
old client stderr:\n{}\nnew client stderr:\n{}\nleader log:\n{}",
old_client.stderr_text(),
new_client.stderr_text(),
leader_log(home.path()),
)
});
assert_ne!(new_pid, old_pid);
// The evicted leader must actually exit within the evict grace
// (EVICT_WAIT_TIMEOUT is 8s; force-kill covers overruns).
assert!(
wait_for_pid_death(old_pid, Duration::from_secs(30)).await,
"old leader pid {old_pid} still alive after eviction\nleader log:\n{}",
leader_log(home.path()),
);
// The old client reconnects and its original session still works.
wait_for_replay_notifications(&old_client, base, Duration::from_secs(60)).await;
let res = old_client.prompt(&session, "after the eviction").await;
assert!(
res.is_ok(),
"old client prompt after eviction failed: {:?}\nstderr:\n{}\nleader log:\n{}",
res.err(),
old_client.stderr_text(),
leader_log(home.path()),
);
// And the new client works against the leader it spawned.
let new_session = new_client.create_session(workdir.path()).await;
new_client
.prompt(&new_session, "hello from the new world")
.await
.expect("new client prompt failed");
})
.await;
}
/// New leader + old client: the older client adopts the newer leader (the
/// floor is directional — never downgrade), keeps functioning through
/// serde-default compat, and the leader records the version mismatch.
#[tokio::test]
#[ignore = "two-binary version-skew test; set KIGI_BINARY_LEADER/KIGI_BINARY_CLIENT and run with --ignored"]
async fn old_client_adopts_new_leader_and_still_functions() {
let Some((old_bin, new_bin)) = skew_binaries() else {
return;
};
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
// NEW binary elects the leader first.
let new_client = LeaderStdioClient::spawn_with_binary(
&new_bin,
&server,
workdir.path(),
home.path(),
)
.await;
new_client.initialize().await;
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(10))
.await
.expect("no live new leader");
// OLD binary connects: must adopt (no downgrade eviction).
let old_client = LeaderStdioClient::spawn_with_binary(
&old_bin,
&server,
workdir.path(),
home.path(),
)
.await;
old_client.initialize().await;
assert_eq!(
read_leader_pid(home.path()),
Some(leader_pid),
"an older client must never evict a newer leader"
);
// Old client functions across the skew: session + prompt succeed,
// exercising serde-default wire compat in anger.
let session = old_client.create_session(workdir.path()).await;
old_client
.prompt(&session, "old client on new leader")
.await
.expect("old client prompt on new leader failed");
// The leader records the client/leader version mismatch (the
// x.ai/leader/version_mismatch notification's server-side warn).
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut saw_mismatch = false;
while tokio::time::Instant::now() < deadline {
if leader_log(home.path()).contains("Version mismatch") {
saw_mismatch = true;
break;
}
tokio::time::sleep(Duration::from_millis(200)).await;
}
assert!(
saw_mismatch,
"leader never logged the version mismatch\nleader log:\n{}",
leader_log(home.path()),
);
})
.await;
}
/// `grok update`'s relaunch signal against a REAL old leader: connect,
/// require `relaunch_v1`, send `RelaunchForUpdate`, and the leader exits so
/// the surviving client re-elects. Mirrors the private
/// `signal_leaders_to_relaunch` in `kigi-bin/src/main.rs` (which is
/// bin-private, so the per-leader body is replicated here).
#[tokio::test]
#[ignore = "two-binary version-skew test; set KIGI_BINARY_LEADER/KIGI_BINARY_CLIENT and run with --ignored"]
async fn relaunch_for_update_drives_real_old_leader_to_exit() {
let Some((old_bin, _new_bin)) = skew_binaries() else {
return;
};
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let old_client = LeaderStdioClient::spawn_with_binary(
&old_bin,
&server,
workdir.path(),
home.path(),
)
.await;
old_client.initialize().await;
let session = old_client.create_session(workdir.path()).await;
old_client
.prompt(&session, "before relaunch")
.await
.expect("pre-relaunch prompt failed");
let old_pid = wait_for_live_leader(home.path(), Duration::from_secs(10))
.await
.expect("no live old leader");
let base = old_client.notification_count();
// The update-signal body, against the sandboxed socket.
let control = LeaderClient::connect(
home.path().join(".kigi").join("leader.sock"),
"grok-pager-update",
ClientMode::Stdio,
ClientCapabilities::default(),
)
.await
.expect("control connect to old leader failed");
if !control.registration().supports_relaunch() {
// Pre-relaunch_v1 releases degrade to the manual-restart
// message; nothing to drive here.
eprintln!(
"SKIP: old leader {:?} does not advertise relaunch_v1",
control.registration().leader_binary_version
);
control.cancel();
return;
}
let ack = control
.send_control(ControlCommand::RelaunchForUpdate {
to_version: "999.0.0".to_string(),
})
.await;
control.cancel();
match ack {
Ok(Ok(ControlPayload::Relaunching { .. })) => {}
// The leader may exit before the ack flushes — acceptable.
Err(_) => {}
other => panic!("unexpected RelaunchForUpdate reply: {other:?}"),
}
assert!(
wait_for_pid_death(old_pid, Duration::from_secs(30)).await,
"old leader pid {old_pid} did not exit after accepting relaunch\nleader log:\n{}",
leader_log(home.path()),
);
// The surviving client re-elects and restores its session.
wait_for_new_leader(home.path(), old_pid, Duration::from_secs(60))
.await
.unwrap_or_else(|| {
panic!(
"no re-elected leader after relaunch\nstderr:\n{}\nleader log:\n{}",
old_client.stderr_text(),
leader_log(home.path()),
)
});
wait_for_replay_notifications(&old_client, base, Duration::from_secs(60)).await;
old_client
.prompt(&session, "after relaunch")
.await
.expect("prompt after relaunch failed");
})
.await;
}
/// Single-ownership after eviction: exactly one leader remains (old pid dead,
/// lock names the live replacement), the eviction is attributable in the
/// sandbox unified log, and no second writer touched `auth.json` during the
/// swap (API-key auth here, so any write would be a regression).
#[tokio::test]
#[ignore = "two-binary version-skew test; set KIGI_BINARY_LEADER/KIGI_BINARY_CLIENT and run with --ignored"]
async fn eviction_leaves_single_leader_and_single_auth_owner() {
let Some((old_bin, new_bin)) = skew_binaries() else {
return;
};
tokio::task::LocalSet::new()
.run_until(async {
let server = MockInferenceServer::start().await.unwrap();
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let old_client = LeaderStdioClient::spawn_with_binary(
&old_bin,
&server,
workdir.path(),
home.path(),
)
.await;
old_client.initialize().await;
let old_pid = wait_for_live_leader(home.path(), Duration::from_secs(10))
.await
.expect("no live old leader");
let auth_path = home.path().join(".kigi").join("auth.json");
let auth_before = std::fs::metadata(&auth_path)
.ok()
.and_then(|m| m.modified().ok());
let new_client = LeaderStdioClient::spawn_with_binary(
&new_bin,
&server,
workdir.path(),
home.path(),
)
.await;
new_client.initialize().await;
let new_pid = wait_for_new_leader(home.path(), old_pid, Duration::from_secs(60))
.await
.expect("no replacement leader after eviction");
assert!(
wait_for_pid_death(old_pid, Duration::from_secs(30)).await,
"evicted leader must exit"
);
assert!(pid_alive(new_pid), "replacement leader must stay alive");
assert_eq!(
read_leader_pid(home.path()),
Some(new_pid),
"the lock file must name exactly the surviving leader"
);
// Attribution: the evicting client recorded the vacate/replace in
// the sandbox unified log.
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut attributed = false;
while tokio::time::Instant::now() < deadline {
let log = sandbox_unified_log(home.path());
if log.contains("leader.evict.vacate_requested")
|| log.contains("leader.spawn.replacement")
{
attributed = true;
break;
}
tokio::time::sleep(Duration::from_millis(200)).await;
}
assert!(
attributed,
"eviction must be attributable in unified.jsonl\nlog:\n{}",
sandbox_unified_log(home.path()),
);
// API-key sandbox: neither leader generation may write auth.json
// during the swap (single auth ownership; a concurrent refresher
// in the dying leader would show up as a write here).
let auth_after = std::fs::metadata(&auth_path)
.ok()
.and_then(|m| m.modified().ok());
assert_eq!(
auth_before, auth_after,
"auth.json must not be written during an eviction swap"
);
})
.await;
}
@@ -0,0 +1,318 @@
use std::borrow::Cow;
use std::sync::Arc;
use serde::Deserialize;
use serde_json::json;
// rmcp is quarantined in kigi-mcp; see that crate's docs.
use kigi_mcp::rmcp;
use kigi_mcp::rmcp::ServerHandler;
use kigi_mcp::rmcp::model::{
CallToolRequestParams, CallToolResult, ContentBlock, ErrorData as McpError, JsonObject,
ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool,
};
#[derive(Clone)]
struct TestMcpServer {
tools: Arc<Vec<Tool>>,
}
impl TestMcpServer {
fn new() -> Self {
let tools = vec![Self::echo_tool()];
Self {
tools: Arc::new(tools),
}
}
fn echo_tool() -> Tool {
let schema: JsonObject = serde_json::from_value(json!({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "Message to echo back"
}
},
"required": ["message"],
"additionalProperties": false
}))
.unwrap();
Tool::new(
Cow::Borrowed("echo"),
Cow::Borrowed("Echo back the provided message"),
Arc::new(schema),
)
}
}
#[derive(Deserialize)]
struct EchoArgs {
message: String,
}
impl ServerHandler for TestMcpServer {
fn get_info(&self) -> ServerInfo {
let mut info = ServerInfo::default();
info.capabilities = ServerCapabilities::builder().enable_tools().build();
info
}
fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> impl std::future::Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
let tools = self.tools.clone();
async move {
Ok(ListToolsResult {
tools: (*tools).clone(),
next_cursor: None,
meta: None,
})
}
}
async fn call_tool(
&self,
request: CallToolRequestParams,
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> Result<CallToolResult, McpError> {
match request.name.as_ref() {
"echo" => {
let args: EchoArgs = match request.arguments {
Some(arguments) => serde_json::from_value(serde_json::Value::Object(
arguments.into_iter().collect(),
))
.map_err(|err| {
McpError::invalid_params(
format!("'message' is a required property: {}", err),
None,
)
})?,
None => {
return Err(McpError::invalid_params(
"'message' is a required property",
None,
));
}
};
Ok(CallToolResult::success(vec![ContentBlock::text(format!(
"ECHO: {}",
args.message
))]))
}
other => Err(McpError::invalid_params(
format!("unknown tool: {other}"),
None,
)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_schema_json_conversion() {
let schema: JsonObject = serde_json::from_value(json!({
"type": "object",
"required": ["field1"],
"properties": {
"field1": {"type": "string"},
"field2": {"type": "number"}
}
}))
.unwrap();
let arc_schema = Arc::new(schema);
let value = serde_json::to_value(arc_schema.as_ref()).unwrap();
assert_eq!(value["type"], "object");
assert_eq!(value["required"][0], "field1");
assert!(
value["properties"]
.as_object()
.unwrap()
.contains_key("field1")
);
}
#[test]
fn test_echo_tool_schema_has_required_field() {
let server = TestMcpServer::new();
let tools = server.tools.clone();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "echo");
let schema_value = serde_json::to_value(tools[0].input_schema.as_ref()).unwrap();
let required = schema_value["required"].as_array().unwrap();
assert_eq!(required.len(), 1);
assert_eq!(required[0], "message");
let properties = schema_value["properties"].as_object().unwrap();
assert!(properties.contains_key("message"));
}
}
#[cfg(test)]
mod mcp_apps_tests {
use super::*;
/// Build an rmcp Tool with `_meta.ui` like an MCP Apps server would.
fn ui_tool(name: &'static str, resource_uri: &str, visibility: Option<Vec<&str>>) -> Tool {
let schema: JsonObject = serde_json::from_value(json!({
"type": "object",
"properties": { "query": { "type": "string" } }
}))
.unwrap();
let mut ui_meta = json!({ "resourceUri": resource_uri });
if let Some(vis) = visibility {
ui_meta["visibility"] = json!(vis);
}
let mut tool = Tool::new(
Cow::Borrowed(name),
Cow::Borrowed("A UI tool"),
Arc::new(schema),
);
let meta_map: JsonObject = serde_json::from_value(json!({ "ui": ui_meta })).unwrap();
tool.meta = Some(rmcp::model::Meta(meta_map));
tool
}
#[test]
fn test_meta_ui_survives_serialization_roundtrip() {
// rmcp Meta is #[serde(transparent)] over JsonObject.
// Our pipeline does: tool.meta → serde_json::to_value → Option<Value>.
// Verify the ui.resourceUri survives this conversion.
let tool = ui_tool("dashboard", "ui://server/dash", None);
let meta_value: serde_json::Value =
serde_json::to_value(tool.meta.as_ref().unwrap()).unwrap();
assert_eq!(meta_value["ui"]["resourceUri"], "ui://server/dash");
}
#[test]
fn test_visibility_app_only_hides_from_model() {
let tool = ui_tool("refresh", "ui://s/d", Some(vec!["app"]));
let meta_value: serde_json::Value =
serde_json::to_value(tool.meta.as_ref().unwrap()).unwrap();
let model_visible = meta_value
.get("ui")
.and_then(|ui| ui.get("visibility"))
.and_then(|v| v.as_array())
.map(|arr| arr.iter().any(|s| s.as_str() == Some("model")))
.unwrap_or(true);
assert!(!model_visible);
}
}
#[cfg(test)]
mod unit_tests {
#[test]
fn test_mcp_tool_schema_preserves_required() {
use kigi_mcp::rmcp;
use kigi_mcp::rmcp::model::Tool as RmcpTool;
use serde_json::json;
use std::borrow::Cow;
use std::sync::Arc;
// Simulate an MCP tool schema like browser_goto would have
let schema_json = json!({
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "URL to navigate to"
}
},
"required": ["url"]
});
let json_object: rmcp::model::JsonObject =
serde_json::from_value(schema_json.clone()).unwrap();
let rmcp_tool = RmcpTool::new(
Cow::Borrowed("browser_goto"),
Cow::Borrowed("Navigate to a URL"),
Arc::new(json_object),
);
let converted_schema = serde_json::to_value(rmcp_tool.input_schema.as_ref())
.unwrap_or_else(|_| serde_json::json!({}));
// Verify the required field is preserved
assert_eq!(converted_schema["type"], "object");
assert_eq!(converted_schema["required"], json!(["url"]));
assert!(converted_schema["properties"]["url"].is_object());
}
/// Verifies that McpToolRegistration carries the MCP server's actual
/// schema (with "type": "object" patched in), not a generic schemars-
/// derived schema for serde_json::Value. This is the regression test
/// for the bug where register_tool() re-derived the schema via
/// generate_schema::<serde_json::Value>() → `{}`, causing Anthropic Messages
/// and Bedrock backends to reject tool calls with:
/// "tools.N.custom.input_schema.type: Field required"
#[test]
fn test_mcp_registration_carries_server_schema_not_schemars_derived() {
use kigi_mcp::rmcp;
use kigi_mcp::rmcp::model::Tool as RmcpTool;
use serde_json::json;
use std::borrow::Cow;
use std::sync::Arc;
// Build an rmcp Tool with a real schema (properties, required, etc.)
let server_schema = json!({
"type": "object",
"properties": {
"query": { "type": "string" },
"limit": { "type": "integer" }
},
"required": ["query"]
});
let json_object: rmcp::model::JsonObject =
serde_json::from_value(server_schema.clone()).unwrap();
let rmcp_tool = RmcpTool::new(
Cow::Borrowed("search"),
Cow::Borrowed("Search for items"),
Arc::new(json_object),
);
// Simulate the conversion path in McpClient::get_tool_registrations()
let mut schema = serde_json::to_value(rmcp_tool.input_schema.as_ref())
.unwrap_or_else(|_| json!({"type": "object"}));
if let Some(obj) = schema.as_object_mut() {
obj.entry("type").or_insert_with(|| json!("object"));
}
// The schema in the registration must be the MCP server's schema
assert_eq!(schema["type"], "object");
assert_eq!(schema["required"], json!(["query"]));
assert!(schema["properties"]["query"].is_object());
assert!(schema["properties"]["limit"].is_object());
}
/// Verifies that an empty inputSchema `{}` (sent by some MCP servers
/// like VSCode for parameterless tools) gets patched with "type": "object".
#[test]
fn test_empty_mcp_schema_gets_type_object_injected() {
use serde_json::json;
let mut schema = json!({});
if let Some(obj) = schema.as_object_mut() {
obj.entry("type").or_insert_with(|| json!("object"));
}
assert_eq!(schema["type"], "object");
}
}
@@ -0,0 +1,874 @@
//! End-to-end actor tests for the MCP "always allow" persistence path.
//!
//! Spawns a real `spawn_permission_manager` actor with a fake gateway whose
//! `request_permission` returns canned responses. Drives the actor through
//! the request → prompt → grant → re-request flow and verifies that the
//! state file on disk reflects the grant.
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
use agent_client_protocol as acp;
use kigi_acp_lib::{AcpAgentGatewaySender, AcpClientMessage};
use kigi_paths::AbsPathBuf;
use kigi_workspace::permission::types::{
PatternMode, PermissionConfig, PermissionRule, RuleAction, ToolFilter,
};
use kigi_workspace::permission::{
AccessKind, ClientType, Decision, PermissionCommand, PermissionHandle, PermissionState,
spawn_permission_manager, spawn_permission_manager_with_hub,
};
use serial_test::serial;
use tokio::sync::{mpsc, oneshot};
/// Shared `KIGI_SHARE_DIR` for the entire test binary. The `OnceLock` in
/// `kigi-config` only allows one value per process, so all tests share
/// this temp directory and `#[serial]` keeps them from clobbering each
/// other's state files.
fn test_home() -> &'static PathBuf {
static HOME: OnceLock<PathBuf> = OnceLock::new();
HOME.get_or_init(|| {
let dir = tempfile::TempDir::new().unwrap();
let path = dir.keep();
// SAFETY: called once at init before other threads touch this var.
unsafe { std::env::set_var("KIGI_SHARE_DIR", &path) };
path
})
}
fn fresh_cwd() -> AbsPathBuf {
let home = test_home();
let unique = format!(
"test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let cwd = home.join(unique);
std::fs::create_dir_all(&cwd).unwrap();
AbsPathBuf::new(cwd).unwrap()
}
fn permission_state_path(cwd: &AbsPathBuf) -> PathBuf {
test_home()
.join("sessions")
.join(urlencoding::encode(cwd.as_str()).into_owned())
.join("permission.toml")
}
fn load_state(cwd: &AbsPathBuf) -> PermissionState {
let path = permission_state_path(cwd);
let contents = std::fs::read_to_string(&path)
.unwrap_or_else(|_| panic!("permission state not yet written at {:?}", path));
toml::from_str(&contents).unwrap()
}
fn tool_call_update(id: &str, name: &str) -> acp::ToolCallUpdate {
acp::ToolCallUpdate::new(
acp::ToolCallId::new(Arc::from(id)),
acp::ToolCallUpdateFields::new()
.kind(Some(acp::ToolKind::Other))
.title(Some(name.to_owned())),
)
}
fn make_session_id() -> acp::SessionId {
acp::SessionId::new(Arc::from("test-session"))
}
/// Build a fake gateway plus a handle to enqueue scripted responses.
///
/// Each `expect_*` call queues one response; the gateway task pops them in
/// FIFO order as `RequestPermission` messages arrive.
struct FakeGateway {
sender: AcpAgentGatewaySender,
/// Queue of (option_id, optional response meta) pairs.
expected: tokio::sync::mpsc::UnboundedSender<(String, Option<serde_json::Value>)>,
}
fn fake_gateway() -> (FakeGateway, tokio::task::JoinHandle<()>) {
let (gw_tx, mut gw_rx) = mpsc::unbounded_channel::<AcpClientMessage>();
let (script_tx, mut script_rx) =
mpsc::unbounded_channel::<(String, Option<serde_json::Value>)>();
let join = tokio::task::spawn_local(async move {
while let Some(msg) = gw_rx.recv().await {
if let AcpClientMessage::RequestPermission(args) = msg {
let (option_id, meta) = script_rx
.recv()
.await
.expect("test ran out of scripted responses");
let mut response = acp::RequestPermissionResponse::new(
acp::RequestPermissionOutcome::Selected(acp::SelectedPermissionOutcome::new(
acp::PermissionOptionId::new(Arc::from(option_id.as_str())),
)),
);
if let Some(m) = meta.and_then(|v| v.as_object().cloned()) {
response = response.meta(m);
}
let _ = args.response_tx.send(Ok(response));
}
}
});
let sender = AcpAgentGatewaySender::new(gw_tx);
(
FakeGateway {
sender,
expected: script_tx,
},
join,
)
}
impl FakeGateway {
fn expect_allow_always_mcp_tool(&self, tool_name: &str) {
let meta = serde_json::json!({
"kind": "tool",
"tool_name": tool_name,
});
self.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
}
fn expect_allow_always_mcp_server(&self, server: &str) {
let meta = serde_json::json!({
"kind": "server",
"server": server,
});
self.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
}
fn expect_plain_allow_always(&self) {
// Legacy `"always-allow"` option id from `fallback_options`.
self.expected
.send(("always-allow".to_string(), None))
.unwrap();
}
}
async fn request(handle: &PermissionHandle, access: AccessKind, id: &str) -> Decision {
let (tx, rx) = oneshot::channel();
let cmd = PermissionCommand::Request {
access,
tool_call_update: tool_call_update(id, "mcp"),
respond_to: tx,
session_id: None,
subagent_type: None,
subagent_description: None,
};
let PermissionHandle::Actor { cmd_tx, .. } = handle else {
panic!("expected actor handle");
};
cmd_tx.send(cmd).unwrap();
rx.await.unwrap()
}
/// Build an MCP access kind from a tool name; these persistence tests only
/// exercise the name, so args are empty.
fn mcp(name: &str) -> AccessKind {
AccessKind::MCPTool {
name: name.to_string(),
input: serde_json::Value::Null,
}
}
async fn run_actor_test<F, Fut>(client_type: ClientType, body: F)
where
F: FnOnce(PermissionHandle, FakeGateway, AbsPathBuf) -> Fut,
Fut: std::future::Future<Output = ()>,
{
run_actor_test_with_policy(client_type, None, body).await;
}
async fn run_actor_test_with_policy<F, Fut>(
client_type: ClientType,
policy: Option<PermissionConfig>,
body: F,
) where
F: FnOnce(PermissionHandle, FakeGateway, AbsPathBuf) -> Fut,
Fut: std::future::Future<Output = ()>,
{
run_actor_test_full(client_type, policy, false, body).await;
}
async fn run_actor_test_full<F, Fut>(
client_type: ClientType,
policy: Option<PermissionConfig>,
initial_yolo: bool,
body: F,
) where
F: FnOnce(PermissionHandle, FakeGateway, AbsPathBuf) -> Fut,
Fut: std::future::Future<Output = ()>,
{
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let cwd = fresh_cwd();
let (gw, _gw_task) = fake_gateway();
let (handle, _events) = spawn_permission_manager(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
client_type,
policy,
vec![], // deny_read_globs
vec![],
initial_yolo,
None,
);
body(handle, gw, cwd).await;
})
.await;
}
fn rule(action: RuleAction, pattern: &str) -> PermissionRule {
PermissionRule {
action,
tool: ToolFilter::Mcp,
pattern: Some(pattern.to_owned()),
pattern_mode: PatternMode::Glob,
}
}
// --- mcp_pre_decision-style end-to-end ---
#[tokio::test]
#[serial]
async fn mcp_tool_grant_persists_and_short_circuits_next_request() {
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
// First request prompts; user picks tool-scope.
gw.expect_allow_always_mcp_tool("linear__list");
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
// Allow disk write to land.
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(state.allowed_mcp_tools.contains("linear__list"));
assert!(state.allowed_mcp_servers.is_empty());
// Second request for the same tool must NOT prompt — actor returns
// Allow without consuming a scripted response. If it tried to
// prompt, the gateway task would block forever, and the request
// call below would hang; we assert by simply receiving an Allow
// synchronously.
let d = request(&handle, mcp("linear__list"), "2").await;
assert!(matches!(d, Decision::Allow));
// A different tool from the same server still prompts (tool-scope
// is exact). We script a server-scope grant for "linear" next.
gw.expect_allow_always_mcp_server("linear");
let d = request(&handle, mcp("linear__create"), "3").await;
assert!(matches!(d, Decision::Allow));
// Wait for the new write.
for _ in 0..50 {
let s = load_state(&cwd);
if s.allowed_mcp_servers.contains("linear") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(state.allowed_mcp_servers.contains("linear"));
// Now any other linear__* tool short-circuits via server-scope.
let d = request(&handle, mcp("linear__update"), "4").await;
assert!(matches!(d, Decision::Allow));
// A different server still prompts.
gw.expect_allow_always_mcp_tool("notion__fetch");
let d = request(&handle, mcp("notion__fetch"), "5").await;
assert!(matches!(d, Decision::Allow));
})
.await;
}
#[tokio::test]
#[serial]
async fn fallback_client_plain_allow_always_persists_mcp_tool() {
// Regression: Generic / GrokWeb / Extension clients
// submit the legacy `"always-allow"` option id. The prompter maps that
// to plain `PromptOutcome::AllowAlways`, and the manager's plain arm
// must persist tool-scope into `allowed_mcp_tools`.
run_actor_test(ClientType::Generic, |handle, gw, cwd| async move {
gw.expect_plain_allow_always();
let d = request(&handle, mcp("notion__fetch"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(
state.allowed_mcp_tools.contains("notion__fetch"),
"fallback client AllowAlways must persist tool-scope, got state={state:?}"
);
// Re-request the same tool — must short-circuit without prompting.
let d = request(&handle, mcp("notion__fetch"), "2").await;
assert!(matches!(d, Decision::Allow));
})
.await;
}
#[tokio::test]
#[serial]
async fn policy_ask_suppresses_mcp_tool_allowlist() {
// With `remember_tool_approvals` OFF (the default), a policy `Ask` rule on an
// MCP tool overrides a session tool-scope grant: the actor must prompt rather
// than auto-allow. (The gate-ON "grant satisfies ask" path is covered by the
// `mcp_pre_decision` unit tests in `manager.rs`.)
let policy = PermissionConfig::new(vec![rule(RuleAction::Ask, "linear__*")]);
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let cwd = fresh_cwd();
// Pre-seed the state file with a tool-scope grant.
let mut state = PermissionState::default();
state.allowed_mcp_tools.insert("linear__list".to_owned());
let dir = test_home()
.join("sessions")
.join(urlencoding::encode(cwd.as_str()).into_owned());
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
dir.join("permission.toml"),
toml::to_string_pretty(&state).unwrap(),
)
.unwrap();
let (gw, _gw_task) = fake_gateway();
// Gate OFF so the `ask` rule stays a hard floor over the grant.
let (handle, _events) = spawn_permission_manager_with_hub(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
ClientType::GrokPager,
Some(policy),
vec![], // deny_read_globs
vec![],
false,
None,
false, // remember_tool_approvals
None,
);
// Script an outright reject so we can confirm the prompt fires.
gw.expected.send(("reject-once".to_string(), None)).unwrap();
let d = request(&handle, mcp("linear__list"), "1").await;
// If the allowlist had won, the actor would have returned Allow
// without consuming the scripted response; the gateway's reject
// proves the prompt path executed.
assert!(matches!(d, Decision::Reject(_)));
})
.await;
}
#[tokio::test]
#[serial]
async fn policy_ask_suppresses_mcp_server_allowlist() {
// Gate-OFF floor over a server-scope grant (see the tool-scope test above).
let policy = PermissionConfig::new(vec![rule(RuleAction::Ask, "linear__*")]);
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let cwd = fresh_cwd();
let mut state = PermissionState::default();
state.allowed_mcp_servers.insert("linear".to_owned());
let dir = test_home()
.join("sessions")
.join(urlencoding::encode(cwd.as_str()).into_owned());
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
dir.join("permission.toml"),
toml::to_string_pretty(&state).unwrap(),
)
.unwrap();
let (gw, _gw_task) = fake_gateway();
// Gate OFF so the `ask` rule stays a hard floor over the grant.
let (handle, _events) = spawn_permission_manager_with_hub(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
ClientType::GrokPager,
Some(policy),
vec![], // deny_read_globs
vec![],
false,
None,
false, // remember_tool_approvals
None,
);
gw.expected.send(("reject-once".to_string(), None)).unwrap();
let d = request(&handle, mcp("linear__create"), "1").await;
assert!(matches!(d, Decision::Reject(_)));
})
.await;
}
#[tokio::test]
#[serial]
async fn policy_deny_takes_precedence_over_mcp_allowlist() {
let policy = PermissionConfig::new(vec![rule(RuleAction::Deny, "linear__*")]);
let local = tokio::task::LocalSet::new();
local
.run_until(async move {
let cwd = fresh_cwd();
let mut state = PermissionState::default();
state.allowed_mcp_tools.insert("linear__list".to_owned());
let dir = test_home()
.join("sessions")
.join(urlencoding::encode(cwd.as_str()).into_owned());
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(
dir.join("permission.toml"),
toml::to_string_pretty(&state).unwrap(),
)
.unwrap();
let (gw, _gw_task) = fake_gateway();
let (handle, _events) = spawn_permission_manager(
make_session_id(),
gw.sender.clone(),
cwd.clone(),
ClientType::GrokPager,
Some(policy),
vec![], // deny_read_globs
vec![],
false,
None,
);
// Do NOT script a response: a policy Deny must short-circuit
// before the prompt path is reached, so the gateway must never
// be invoked.
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::PolicyDeny(_)));
})
.await;
}
#[tokio::test]
#[serial]
async fn policy_allow_short_circuits_before_mcp_allowlist() {
// Sanity: a policy Allow returns Allow immediately and never touches
// the pre-decision lookup.
let policy = PermissionConfig::new(vec![rule(RuleAction::Allow, "linear__*")]);
run_actor_test_with_policy(
ClientType::GrokPager,
Some(policy),
|handle, _gw, _cwd| async move {
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
},
)
.await;
}
#[tokio::test]
#[serial]
async fn empty_server_prefix_falls_back_to_tool_scope() {
// Defense-in-depth: even if a malformed `McpScopeSelection::Server { server: "" }`
// somehow makes it through, the prompter must downgrade to tool-scope
// and persist via `AllowAlwaysMcpTool` — never write an empty server
// prefix into `allowed_mcp_servers`.
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
let meta = serde_json::json!({
"kind": "server",
"server": "",
});
gw.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(state.allowed_mcp_servers.is_empty());
assert!(state.allowed_mcp_tools.contains("linear__list"));
})
.await;
}
#[tokio::test]
#[serial]
async fn allow_always_mcp_tool_ignores_client_supplied_tool_name() {
// Security regression: the response meta `tool_name` is informational
// only. The manager MUST persist the name from `AccessKind::MCPTool`
// so a buggy or malicious client cannot whitelist a different tool
// than the one the user saw in the prompt.
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
// Request approves `linear__list`, but the response claims a
// different tool name (e.g. `notion__fetch`).
let meta = serde_json::json!({
"kind": "tool",
"tool_name": "notion__fetch",
});
gw.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(
state.allowed_mcp_tools.contains("linear__list"),
"must persist the access-kind name"
);
assert!(
!state.allowed_mcp_tools.contains("notion__fetch"),
"must NOT persist the client-supplied name"
);
})
.await;
}
#[tokio::test]
#[serial]
async fn allow_always_mcp_server_rejects_mismatched_prefix() {
// Security regression: the response meta `server` must match the
// canonical server prefix derived from the access kind. On mismatch,
// the manager downgrades to tool-scope on the access-kind name -- the
// smallest blast radius the user actually approved.
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
// Approve `linear__list` (canonical server prefix is `linear`),
// but the client claims `notion` as the server.
let meta = serde_json::json!({
"kind": "server",
"server": "notion",
});
gw.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(
state.allowed_mcp_servers.is_empty(),
"must NOT persist any server-scope grant on mismatch"
);
assert!(
!state.allowed_mcp_servers.contains("notion"),
"must NOT trust the client-supplied server prefix"
);
assert!(
state.allowed_mcp_tools.contains("linear__list"),
"must downgrade to tool-scope using the access-kind name"
);
})
.await;
}
#[tokio::test]
#[serial]
async fn allow_always_mcp_server_persists_canonical_prefix_on_match() {
// Sanity: a client that supplies the correct canonical prefix
// succeeds (this is the common case post-fix).
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
let meta = serde_json::json!({
"kind": "server",
"server": "linear",
});
gw.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(state.allowed_mcp_servers.contains("linear"));
assert!(state.allowed_mcp_tools.is_empty());
})
.await;
}
#[tokio::test]
#[serial]
async fn allow_always_mcp_server_downgrades_when_access_has_no_separator() {
// Defensive: if the access name itself has no `__` (e.g. via a
// malformed `ToolInput::MCPTool`), the canonical prefix is None and
// server-scope is unreachable. The manager downgrades to tool-scope
// on the raw access name rather than persisting the client prefix.
run_actor_test(ClientType::GrokPager, |handle, gw, cwd| async move {
let meta = serde_json::json!({
"kind": "server",
"server": "linear",
});
gw.expected
.send(("allow-always-mcp".to_string(), Some(meta)))
.unwrap();
let d = request(&handle, mcp("standalone"), "1").await;
assert!(matches!(d, Decision::Allow));
for _ in 0..50 {
if permission_state_path(&cwd).exists() {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let state = load_state(&cwd);
assert!(state.allowed_mcp_servers.is_empty());
assert!(state.allowed_mcp_tools.contains("standalone"));
})
.await;
}
#[tokio::test]
#[serial]
async fn dont_ask_policy_denies_without_prompting() {
use kigi_workspace::permission::types::{PermissionConfig, PromptPolicy};
let mut policy = PermissionConfig::new(vec![]);
policy.prompt_policy = PromptPolicy::Deny;
// No scripted gateway responses: if the manager tried to prompt,
// the gateway would block forever, proving dont_ask short-circuits.
run_actor_test_with_policy(
ClientType::GrokPager,
Some(policy),
|handle, _gw, _cwd| async move {
let d = request(&handle, mcp("linear__list"), "1").await;
assert!(matches!(d, Decision::PolicyDeny(_)));
let d = request(&handle, AccessKind::Bash("npm install".to_string()), "2").await;
assert!(matches!(d, Decision::PolicyDeny(_)));
// Reads still auto-approve (pre-decision, before dont_ask)
let d = request(
&handle,
AccessKind::Read(Some("/tmp/test.txt".to_string())),
"3",
)
.await;
assert!(matches!(d, Decision::Allow));
},
)
.await;
}
// --- deny rules survive YOLO mode ---
#[tokio::test]
#[serial]
async fn deny_rule_enforced_in_yolo_mode_bash() {
let policy = PermissionConfig::new(vec![PermissionRule {
action: RuleAction::Deny,
tool: ToolFilter::Bash,
pattern: Some("rm*".to_owned()),
pattern_mode: PatternMode::Glob,
}]);
run_actor_test_full(
ClientType::GrokPager,
Some(policy),
true,
|handle, _gw, _cwd| async move {
let d = request(
&handle,
AccessKind::Bash("rm -rf /tmp/foo".to_string()),
"1",
)
.await;
assert!(
matches!(d, Decision::PolicyDeny(_)),
"deny rule must block even in YOLO mode, got {d:?}"
);
let d = request(&handle, AccessKind::Bash("cargo test".to_string()), "2").await;
assert!(
matches!(d, Decision::Allow),
"non-denied bash must auto-approve in YOLO mode, got {d:?}"
);
},
)
.await;
}
#[tokio::test]
#[serial]
async fn deny_rule_enforced_in_yolo_mode_mcp() {
let policy = PermissionConfig::new(vec![PermissionRule {
action: RuleAction::Deny,
tool: ToolFilter::Mcp,
pattern: Some("dangerous__*".to_owned()),
pattern_mode: PatternMode::Glob,
}]);
run_actor_test_full(
ClientType::GrokPager,
Some(policy),
true,
|handle, _gw, _cwd| async move {
let d = request(&handle, mcp("dangerous__delete_all"), "1").await;
assert!(
matches!(d, Decision::PolicyDeny(_)),
"deny rule must block MCP tool even in YOLO mode, got {d:?}"
);
let d = request(&handle, mcp("linear__list"), "2").await;
assert!(
matches!(d, Decision::Allow),
"non-denied MCP tool must auto-approve in YOLO mode, got {d:?}"
);
},
)
.await;
}
#[tokio::test]
#[serial]
async fn deny_rule_enforced_in_yolo_mode_edit() {
let policy = PermissionConfig::new(vec![PermissionRule {
action: RuleAction::Deny,
tool: ToolFilter::Edit,
pattern: Some("/etc/**".to_owned()),
pattern_mode: PatternMode::Glob,
}]);
run_actor_test_full(
ClientType::GrokPager,
Some(policy),
true,
|handle, _gw, _cwd| async move {
let d = request(&handle, AccessKind::Edit("/etc/passwd".to_string()), "1").await;
assert!(
matches!(d, Decision::PolicyDeny(_)),
"deny rule must block edits even in YOLO mode, got {d:?}"
);
let d = request(&handle, AccessKind::Edit("src/main.rs".to_string()), "2").await;
assert!(
matches!(d, Decision::Allow),
"non-denied edit must auto-approve in YOLO mode, got {d:?}"
);
},
)
.await;
}
#[tokio::test]
#[serial]
async fn deny_rule_enforced_in_yolo_mode_web_fetch() {
let policy = PermissionConfig::new(vec![PermissionRule {
action: RuleAction::Deny,
tool: ToolFilter::WebFetch,
pattern: Some("evil.com".to_owned()),
pattern_mode: PatternMode::Domain,
}]);
run_actor_test_full(
ClientType::GrokPager,
Some(policy),
true,
|handle, _gw, _cwd| async move {
let d = request(
&handle,
AccessKind::WebFetch("https://evil.com/exfiltrate".to_string()),
"1",
)
.await;
assert!(
matches!(d, Decision::PolicyDeny(_)),
"deny rule must block web_fetch even in YOLO mode, got {d:?}"
);
let d = request(
&handle,
AccessKind::WebFetch("https://docs.rs/tokio".to_string()),
"2",
)
.await;
assert!(
matches!(d, Decision::Allow),
"non-denied web_fetch must auto-approve in YOLO mode, got {d:?}"
);
},
)
.await;
}
#[tokio::test]
#[serial]
async fn yolo_mode_without_deny_rules_approves_everything() {
run_actor_test_full(
ClientType::GrokPager,
None,
true,
|handle, _gw, _cwd| async move {
let d = request(&handle, AccessKind::Bash("rm -rf /".to_string()), "1").await;
assert!(matches!(d, Decision::Allow));
let d = request(&handle, mcp("linear__list"), "2").await;
assert!(matches!(d, Decision::Allow));
let d = request(&handle, AccessKind::Edit("/etc/passwd".to_string()), "3").await;
assert!(matches!(d, Decision::Allow));
let d = request(
&handle,
AccessKind::WebFetch("https://evil.com".to_string()),
"4",
)
.await;
assert!(matches!(d, Decision::Allow));
},
)
.await;
}
@@ -0,0 +1,196 @@
//! Regression tests: a Anthropic Messages API `/v1/messages` stream that terminates with
//! `stop_reason: "refusal"` must complete the turn cleanly with EXACTLY ONE
//! inference request.
//!
//! Previously the unknown `stop_reason` failed the terminal `message_delta`
//! parse (discarding the fully-streamed response) and the resulting
//! serialization error was misclassified as a retryable stream error,
//! producing a ~10-minute retry storm per turn. Covered here end-to-end
//! through both the plain stdio agent and a leader-hosted session.
//!
//! Tests are `#[ignore]`d by default — they require a pre-built binary
//! (auto-built locally when missing):
//!
//! ```bash
//! cargo test -p kigi-shell --test test_refusal_stop_reason -- --ignored
//! ```
use std::future::Future;
use agent_client_protocol as acp;
use kigi_test_support::*;
/// Run an async test body inside a `LocalSet` (required by ACP's `!Send` futures).
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Mock with a single Anthropic-style model whose `/v1/messages` stream ends
/// with `stop_reason: "refusal"`.
async fn refusal_messages_server() -> MockInferenceServer {
let server = MockInferenceServer::start_with_models(vec![
MockModelEntry::new("messages-compatible-model").with_api_backend("messages"),
])
.await
.expect("start mock server");
server.set_messages_stop_reason("refusal");
server
}
/// `/v1/messages` requests belonging to the prompt turn. The session also
/// fires a one-shot title-generation call on the first user message; it is
/// identified (and excluded) by its forced `session_title` tool.
fn turn_messages_request_count(server: &MockInferenceServer) -> usize {
server
.requests()
.iter()
.filter(|e| e.path == "/v1/messages")
.filter(|e| {
!e.body.as_ref().is_some_and(|b| {
b.get("tools")
.and_then(|t| t.as_array())
.is_some_and(|tools| {
tools.iter().any(|t| {
t.get("name").and_then(|n| n.as_str()) == Some("session_title")
})
})
})
})
.count()
}
/// THE regression test: a refusal-terminated `/v1/messages` turn must return
/// a successful prompt response from exactly one inference request.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_refusal_turn_completes_with_single_messages_request() {
with_local_set(|| async {
let server = refusal_messages_server().await;
let workdir = git_workdir();
let client = GrokStdioClient::spawn(&server, workdir.path()).await;
client.initialize_with_timeout().await;
let session_id = client
.create_session_with_model_timeout(workdir.path(), "messages-compatible-model")
.await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
let response = result.unwrap_or_else(|e| {
panic!(
"refusal-terminated turn must complete, got error: {e:?}\nrequest log:\n{}\nstderr:\n{}",
server.request_log_summary(),
stderr_tail(&client.stderr(), 1200)
)
});
assert_eq!(
response.stop_reason,
acp::StopReason::EndTurn,
"refusal must end the turn cleanly"
);
assert!(
client.captured_text().contains("Echo:"),
"streamed response text must be delivered, got: {:?}",
client.captured_text()
);
assert_eq!(
turn_messages_request_count(&server),
1,
"exactly one turn request to /v1/messages (no retry storm)\nrequest log:\n{}",
server.request_log_summary()
);
assert!(
server.messages_request_count() <= 2,
"at most turn + title-generation requests\nrequest log:\n{}",
server.request_log_summary()
);
})
.await;
}
// ============================================================================
// Leader mode: the same refusal scenario through a leader-hosted session
// (client → stdio bridge → leader unix socket → leader-hosted agent).
// ============================================================================
#[cfg(unix)]
mod leader {
use std::time::Duration;
use agent_client_protocol as acp;
use kigi_test_support::leader::{LeaderStdioClient, wait_for_live_leader};
use kigi_test_support::*;
use super::{refusal_messages_server, turn_messages_request_count, with_local_set};
/// Leader-mode variant of the regression: the refusal-terminated turn
/// must complete cleanly (single request, prompt response delivered)
/// when the session is hosted by the leader IPC server.
#[tokio::test]
#[ignore] // requires pre-built binary; run with --ignored
async fn test_leader_refusal_turn_completes_with_single_messages_request() {
with_local_set(|| async {
let server = refusal_messages_server().await;
let workdir = git_workdir();
let home = tempfile::tempdir().unwrap();
std::fs::create_dir_all(home.path().join(".kigi")).unwrap();
let client = LeaderStdioClient::spawn(&server, workdir.path(), home.path()).await;
client.initialize().await;
let session_id = client
.create_session_with_model(workdir.path(), "messages-compatible-model")
.await;
let result = client.prompt(&session_id, "say hello").await;
// Prove the session is leader-hosted: a live leader process,
// distinct from the client subprocess, holds the lock.
let leader_pid = wait_for_live_leader(home.path(), Duration::from_secs(5))
.await
.unwrap_or_else(|| {
panic!(
"no live leader PID in lock file — turn did not run under the leader\nstderr:\n{}",
client.stderr_text()
)
});
assert_ne!(
Some(leader_pid),
client.child.id(),
"leader must be a separate process from the stdio client"
);
let response = result.unwrap_or_else(|e| {
panic!(
"leader-hosted refusal turn must complete, got error: {e:?}\nrequest log:\n{}\nstderr:\n{}",
server.request_log_summary(),
client.stderr_text()
)
});
assert_eq!(
response.stop_reason,
acp::StopReason::EndTurn,
"refusal must end the turn cleanly under the leader"
);
assert!(
client.captured_text().contains("Echo:"),
"streamed response text must reach the client through the leader, got: {:?}",
client.captured_text()
);
assert_eq!(
turn_messages_request_count(&server),
1,
"exactly one turn request to /v1/messages (no retry storm)\nrequest log:\n{}",
server.request_log_summary()
);
assert!(
server.messages_request_count() <= 2,
"at most turn + title-generation requests\nrequest log:\n{}",
server.request_log_summary()
);
})
.await;
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,258 @@
//! Integration test: MockInferenceServer `/v1/settings` endpoint and
//! remote settings settings refresh infrastructure.
//!
//! Tests the mock endpoint directly (no binary needed) and verifies
//! the `fetch_settings_blocking` client round-trips correctly with
//! runtime-mutated mock settings.
//!
//! Run locally:
//! ```bash
//! cargo test -p kigi-shell --test test_settings_refresh
//! ```
use std::future::Future;
use kigi_shell::util::config::RemoteSettings;
use kigi_test_support::*;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Verify the mock `/v1/settings` endpoint returns 404 when no settings
/// are configured (the default). This preserves backward compatibility:
/// existing tests that never call `set_settings` see a 404, and
/// `fetch_settings_blocking` returns `None`.
#[tokio::test]
async fn test_settings_endpoint_returns_404_when_unconfigured() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let resp = reqwest::get(format!("{}/settings", server.url()))
.await
.expect("request failed");
assert_eq!(resp.status(), 404);
})
.await;
}
/// Verify the mock `/v1/settings` endpoint returns configured settings
/// and that `set_settings` runtime mutation is reflected immediately.
#[tokio::test]
async fn test_settings_endpoint_returns_configured_settings() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Configure initial settings
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v1".into()]),
leader_mode: Some(false),
..Default::default()
});
// Fetch and verify
let resp = reqwest::get(format!("{}/settings", server.url()))
.await
.expect("request failed");
assert_eq!(resp.status(), 200);
let settings: RemoteSettings = resp.json().await.expect("parse failed");
assert_eq!(settings.tips, Some(vec!["tip_v1".into()]));
assert_eq!(settings.leader_mode, Some(false));
})
.await;
}
/// Verify that `set_settings` updates are visible to subsequent requests
/// (runtime mutation for multi-session test scenarios).
#[tokio::test]
async fn test_settings_endpoint_reflects_runtime_mutations() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Initial settings
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v1".into()]),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.tips, Some(vec!["tip_v1".into()]));
// Mutate settings (simulating a remote feature flag change)
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v2".into()]),
leader_mode: Some(true),
..Default::default()
});
// Subsequent request sees the updated values
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.tips, Some(vec!["tip_v2".into()]));
assert_eq!(settings.leader_mode, Some(true));
})
.await;
}
/// Verify `fetch_settings_blocking` round-trips through the mock server.
/// This is the actual client function used by `refresh_remote_settings`.
#[tokio::test]
async fn test_fetch_settings_blocking_round_trip() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Without settings configured: returns None (404 from mock)
let auth = kigi_shell::auth::GrokAuth {
key: "test-key".into(),
..Default::default()
};
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None)
})
.await
.unwrap();
assert!(
result.is_none(),
"Expected None when settings not configured"
);
// With settings configured: returns Some(settings)
server.set_settings(RemoteSettings {
tips: Some(vec!["fetched_tip".into()]),
..Default::default()
});
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None)
})
.await
.unwrap();
let settings = result.expect("Expected Some when settings are configured");
assert_eq!(settings.tips, Some(vec!["fetched_tip".into()]));
})
.await;
}
/// Verify the `doom_loop_recovery` settings object survives the
/// `/v1/settings` round-trip, that its absence deserializes to `None` (old
/// servers), and that a partial object keeps its unset fields `None`.
#[tokio::test]
async fn test_doom_loop_recovery_settings_round_trip() {
use kigi_shell::util::config::DoomLoopRecoverySettings;
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Absent from the payload ⇒ None on the client.
server.set_settings(RemoteSettings::default());
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.doom_loop_recovery, None);
server.set_settings(RemoteSettings {
doom_loop_recovery: Some(DoomLoopRecoverySettings {
enabled: Some(true),
max_threshold: Some(16),
max_retries: Some(1),
}),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
let recovery = settings.doom_loop_recovery.expect("object round-trips");
assert_eq!(recovery.enabled, Some(true));
assert_eq!(recovery.max_threshold, Some(16));
assert_eq!(recovery.max_retries, Some(1));
// Partial object: only the set field comes through; the rest stay
// None so the resolver falls through per-field.
server.set_settings(RemoteSettings {
doom_loop_recovery: Some(DoomLoopRecoverySettings {
max_threshold: Some(32),
..Default::default()
}),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
let recovery = settings.doom_loop_recovery.expect("object round-trips");
assert_eq!(recovery.enabled, None);
assert_eq!(recovery.max_threshold, Some(32));
assert_eq!(recovery.max_retries, None);
})
.await;
}
/// Verify that the mock server's request log correctly tracks
/// GET /v1/settings requests for assertion in multi-session tests.
#[tokio::test]
async fn test_settings_requests_appear_in_request_log() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
server.set_settings(RemoteSettings::default());
assert_eq!(server.request_count(), 0);
// First request
let _ = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap();
let settings_reqs: Vec<_> = server
.requests()
.into_iter()
.filter(|r| r.method == "GET" && r.path.contains("/settings"))
.collect();
assert_eq!(settings_reqs.len(), 1, "Expected 1 settings request");
// Second request (simulating /new refresh)
let _ = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap();
let settings_reqs: Vec<_> = server
.requests()
.into_iter()
.filter(|r| r.method == "GET" && r.path.contains("/settings"))
.collect();
assert_eq!(settings_reqs.len(), 2, "Expected 2 settings requests");
})
.await;
}
@@ -0,0 +1,110 @@
//! End-to-end test for subagent orphan reconciliation on session resume.
//!
//! When a process dies mid-subagent, the subagent's `meta.json` is left
//! `status: "running"` with no `SubagentFinished` — so on resume the client
//! shows it Running forever. `MvpAgent::load_session` heals this: it scans the
//! session's `subagents/` dir and flips any stale `running` meta (not tracked by
//! the live coordinator) to `cancelled` (mechanism A, the meta pass).
//!
//! This test spawns a real `grok agent stdio` process, seeds an orphaned
//! `running` meta on disk, resumes the session, and asserts the meta was
//! reconciled to `cancelled`.
//!
//! Run locally (needs a pre-built binary):
//! ```bash
//! cargo test -p kigi-shell --test test_subagent_orphan_reconcile -- --ignored
//! ```
use std::future::Future;
use std::path::{Path, PathBuf};
use kigi_test_support::*;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Find `<home>/sessions/<enc-cwd>/<id>` without depending on the internal cwd
/// encoder: scan the one level of cwd dirs for a child named `<id>`.
fn locate_session_dir(home: &Path, id: &str) -> PathBuf {
let sessions = home.join("sessions");
for entry in std::fs::read_dir(&sessions)
.expect("read sessions dir")
.flatten()
{
let candidate = entry.path().join(id);
if candidate.is_dir() {
return candidate;
}
}
panic!(
"session dir for {id} not found under {}",
sessions.display()
);
}
#[tokio::test]
#[ignore] // requires pre-built binary
async fn resume_reconciles_orphaned_running_subagent() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
// Phase 1: create a real session, then take its home so we can seed it.
let mut writer = GrokStdioClient::spawn(&server, workdir.path()).await;
writer.initialize_with_timeout().await;
let session_id = writer.create_session_with_timeout(workdir.path()).await;
let shared_home = writer.take_home();
drop(writer);
// Simulate a crash: inject a subagent meta left `running` on disk (no
// terminal write, no SubagentFinished) — exactly what a dead process
// leaves behind.
// GrokStdioClient sets HOME=<temp>; the binary uses <HOME>/.kigi as KIGI_SHARE_DIR.
let kigi_home = shared_home.path().join(".kigi");
let session_dir = locate_session_dir(&kigi_home, session_id.0.as_ref());
let sub_id = "sa-orphan";
let meta_path = session_dir.join("subagents").join(sub_id).join("meta.json");
std::fs::create_dir_all(meta_path.parent().unwrap()).unwrap();
std::fs::write(
&meta_path,
serde_json::json!({
"subagent_id": sub_id,
"parent_session_id": session_id.0.as_ref(),
"child_session_id": "child-orphan",
"subagent_type": "general-purpose",
"description": "stuck task",
"prompt": "do work",
"status": "running",
"started_at": chrono::Utc::now().to_rfc3339(),
})
.to_string(),
)
.unwrap();
// Phase 2: resume in a fresh process. `load_session` runs the reconcile.
let reader = GrokStdioClient::spawn_with_home(&server, workdir.path(), shared_home).await;
reader.initialize_with_timeout().await;
let _ = reader
.load_session_with_timeout(&session_id, workdir.path())
.await;
// The orphan's on-disk meta must now be terminal (cancelled), not running.
let reread: serde_json::Value =
serde_json::from_str(&std::fs::read_to_string(&meta_path).expect("read orphan meta"))
.expect("parse orphan meta");
assert_eq!(
reread.get("status").and_then(|s| s.as_str()),
Some("cancelled"),
"resume must reconcile the orphaned running subagent to cancelled\nstderr:\n{}",
stderr_tail(&reader.stderr(), 2000)
);
})
.await;
}
@@ -0,0 +1,116 @@
//! summary.json reasoning-effort persistence tests.
//!
//! Regression tests for "what effort was this session run on?": a fresh
//! session must record its resolved `reasoning_effort` in `summary.json` at
//! creation — not only after an explicit model/effort switch (the old
//! behavior, which left the field absent for sessions that never switched).
//!
//! Each test spawns a real `grok agent stdio` process against a mock
//! inference server and asserts on the persisted `summary.json`.
//!
//! Run locally:
//! ```bash
//! cargo test -p kigi-shell --test test_summary_reasoning_effort -- --ignored
//! ```
use std::future::Future;
use kigi_test_support::*;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Find `summary.json` for `session_id` under `<home>/.kigi/sessions/` and
/// parse it. The sessions tree is `<encoded-cwd>/<session-id>/summary.json`;
/// matching on the directory name avoids re-implementing the cwd encoding.
fn read_summary(home: &std::path::Path, session_id: &str) -> serde_json::Value {
let sessions_root = home.join(".kigi").join("sessions");
let cwd_dirs = std::fs::read_dir(&sessions_root)
.unwrap_or_else(|e| panic!("no sessions dir at {}: {e}", sessions_root.display()));
for cwd_dir in cwd_dirs.flatten() {
let candidate = cwd_dir.path().join(session_id).join("summary.json");
if candidate.is_file() {
let raw = std::fs::read_to_string(&candidate).expect("read summary.json");
return serde_json::from_str(&raw).expect("parse summary.json");
}
}
panic!(
"summary.json for session {session_id} not found under {}",
sessions_root.display()
);
}
/// A fresh session on a model with a configured reasoning effort must persist
/// that effort in `summary.json` without any model/effort switch.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn test_fresh_session_persists_reasoning_effort() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
// Configure the mock catalog's model with an explicit effort via the
// user config override (the same path a remote settings catalog entry or
// `--effort` would populate).
let home = tempfile::TempDir::new().expect("create temp home");
let grok_dir = home.path().join(".kigi");
std::fs::create_dir_all(&grok_dir).expect("create .kigi dir");
std::fs::write(
grok_dir.join("config.toml"),
r#"
[model.test-model]
supports_reasoning_effort = true
reasoning_effort = "high"
"#,
)
.expect("write config.toml");
let client = GrokStdioClient::spawn_with_home(&server, workdir.path(), home).await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "prompt failed: {:?}", result.err());
let summary = read_summary(client.home_path(), &session_id.0);
assert_eq!(
summary.get("reasoning_effort").and_then(|v| v.as_str()),
Some("high"),
"fresh session must record its effort in summary.json; got: {summary}"
);
})
.await;
}
/// A fresh session on a model with no configured effort must not invent one:
/// `summary.json` omits the field (the model uses its server-side default).
#[tokio::test]
#[ignore] // requires pre-built binary
async fn test_fresh_session_without_effort_omits_field() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let client = GrokStdioClient::spawn(&server, workdir.path()).await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
let result = client.prompt_with_timeout(&session_id, "say hello").await;
assert!(result.is_ok(), "prompt failed: {:?}", result.err());
let summary = read_summary(client.home_path(), &session_id.0);
assert_eq!(
summary.get("reasoning_effort"),
None,
"session without a configured effort must omit the field; got: {summary}"
);
})
.await;
}
@@ -0,0 +1,339 @@
//! E2E: trusted local plugin install snapshot refresh on session start.
//!
//! Replicates an enterprise feedback scenario:
//! 1. Install a local plugin (full copy into `installed-plugins/`).
//! 2. Add a new agent only on the **live** source tree.
//! 3. Start a headless session — startup must re-copy trusted/user-home locals.
//! 4. Smoke-validate session JSON under `$KIGI_SHARE_DIR/sessions/` after exit.
//!
//! Requires a built `grok` binary (`KIGI_BINARY` or cargo-built pager) for the
//! ignored headless test.
//!
//! ```bash
//! cargo test -p kigi-shell --test test_trusted_local_plugin_refresh_e2e
//! cargo test -p kigi-shell --test test_trusted_local_plugin_refresh_e2e -- --ignored
//! ```
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use kigi_agent::plugins::SharedPluginRegistryHandle;
use kigi_agent::plugins::discovery::DiscoveryConfig;
use kigi_agent::plugins::git_install::{InstallSource, install_from_source};
use kigi_agent::plugins::install_registry::{
InstallKind, InstallRegistry, InstalledRepo, RepoPlugin,
};
use kigi_test_support::*;
use serial_test::serial;
use tempfile::TempDir;
fn write_minimal_plugin(dir: &Path, name: &str) {
std::fs::create_dir_all(dir).unwrap();
std::fs::write(dir.join("plugin.json"), format!(r#"{{"name":"{name}"}}"#)).unwrap();
}
fn write_agent(dir: &Path, file_stem: &str, name: &str, description: &str) {
std::fs::create_dir_all(dir.join("agents")).unwrap();
let body = format!("---\nname: {name}\ndescription: {description}\n---\n\n# {name}\n");
std::fs::write(dir.join("agents").join(format!("{file_stem}.md")), body).unwrap();
}
fn register_local_install(registry: &mut InstallRegistry, source: &Path) -> InstalledRepo {
let installed = install_from_source(
&InstallSource::Local {
path: source.to_path_buf(),
subdir: None,
},
registry,
)
.expect("install local plugin");
let plugins = installed
.plugins
.iter()
.map(|p| {
(
p.name.clone(),
RepoPlugin {
subdir: p.subdir.clone(),
version: p.version.clone(),
},
)
})
.collect();
let now = chrono::Utc::now().to_rfc3339();
let repo = InstalledRepo {
kind: InstallKind::Local {
source_path: source.to_path_buf(),
subdir: None,
},
installed_at: now.clone(),
updated_at: now,
path: installed.repo_path.clone(),
plugins,
};
registry.insert(installed.repo_key.clone(), repo.clone());
repo
}
/// Library-level e2e in a sandboxed tmp dir (always runs — no external binary).
///
/// Proves the enterprise symptom is fixed: an agent added to the live source after
/// install must surface to discovery (the `/agents` dashboard reads the same
/// `all_subagents_with_plugins` list) without a reinstall, driven through the
/// real session-spawn path (`refresh_and_build_for_cwd`, which refreshes first).
/// RAII: set an env var, restore the prior value (or unset) on drop, so a test
/// never leaves process-global env pointing at a dropped tempdir. Local copy —
/// each test holds two guards at once, and the canonical lock-holding guard
/// deadlocks when nested.
struct EnvVarGuard {
key: &'static str,
prev: Option<std::ffi::OsString>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: impl AsRef<std::ffi::OsStr>) -> Self {
let prev = std::env::var_os(key);
unsafe { std::env::set_var(key, value) };
Self { key, prev }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match self.prev.take() {
Some(v) => unsafe { std::env::set_var(self.key, v) },
None => unsafe { std::env::remove_var(self.key) },
}
}
}
#[test]
#[serial]
fn trusted_local_refresh_surfaces_new_agent_via_discovery() {
// Canonicalize so under-home auto-trust holds where the temp root is a
// symlink (macOS `/var` -> `/private/var`).
let home_tmp = TempDir::new().unwrap();
let home = dunce::canonicalize(home_tmp.path()).unwrap();
let kigi_home = home.join(".kigi");
let _home_guard = EnvVarGuard::set("HOME", &home);
let _grok_guard = EnvVarGuard::set("KIGI_SHARE_DIR", &kigi_home);
// Live source: a user-home local plugin (mirrors a `~/.claude` local tree).
let source = home
.join(".claude")
.join("local-marketplace")
.join("demo-plugin");
write_minimal_plugin(&source, "demo-plugin");
write_agent(&source, "old", "old-agent", "exists at install");
// Install = full snapshot copy into installed-plugins (not a live symlink).
let mut registry = InstallRegistry::empty(kigi_home.join("installed-plugins"));
let installed = register_local_install(&mut registry, &source);
registry.save().expect("save registry");
// New agent added to the live source only — not yet in the snapshot.
write_agent(&source, "new", "new-agent", "added after install");
assert!(!installed.path.join("agents/new.md").exists());
// Session spawn: refresh_and_build_for_cwd re-copies trusted local installs,
// then rediscovers. Mirror the install command auto-enabling the plugin.
let cwd = home.join("workspace");
std::fs::create_dir_all(&cwd).unwrap();
let handle = SharedPluginRegistryHandle::new(None, Vec::new());
let config = DiscoveryConfig {
cli_plugin_dirs: Vec::new(),
config_paths: Vec::new(),
disabled: Vec::new(),
enabled: vec!["demo-plugin".to_string()],
};
let plugin_registry = handle
.refresh_and_build_for_cwd(&cwd, &config, &[], true)
.expect("registry built with installed plugin");
assert!(
installed.path.join("agents/new.md").exists(),
"session-spawn refresh must re-copy the new agent into the snapshot"
);
// The new agent must surface to discovery (the reported symptom).
let agents = kigi_agent::discovery::all_subagents_with_plugins(
&cwd,
&HashMap::new(),
Some(plugin_registry.as_ref()),
);
let names: Vec<&str> = agents.iter().map(|a| a.name.as_str()).collect();
assert!(
names.contains(&"demo-plugin:new-agent"),
"new agent must surface in /agents after session-start refresh; got {names:?}"
);
// Session `_meta.pluginDirs` load. Lives in the same test because
// kigi_home() caches the first KIGI_SHARE_DIR per process; a separate test
// could seed the cache first and break the assertions above.
let plugin_dir = home.join("session-plugin");
write_minimal_plugin(&plugin_dir, "session-plugin");
write_agent(&plugin_dir, "helper", "helper-agent", "session-scoped");
let session_handle = SharedPluginRegistryHandle::new(None, Vec::new());
let session_config = DiscoveryConfig {
cli_plugin_dirs: Vec::new(),
config_paths: Vec::new(),
disabled: Vec::new(),
enabled: Vec::new(),
};
let session_dirs = vec![plugin_dir.clone()];
let registry = session_handle
.build_for_cwd(&cwd, &session_config, &session_dirs, true)
.expect("registry built with session plugin dir");
let plugin = registry
.get("session-plugin")
.expect("session plugin discovered");
assert_eq!(plugin.scope, kigi_agent::plugins::PluginScope::CliOverride);
assert!(plugin.trusted && plugin.enabled);
assert_eq!(registry.session_plugin_dirs(), session_dirs.as_slice());
// A rebuild without the dirs (the shared fan-out shape) must not carry them.
let shared = session_handle.build_for_cwd(&cwd, &session_config, &[], true);
assert!(shared.is_none_or(|r| r.session_plugin_dirs().is_empty()));
}
/// Full binary smoke: session start runs refresh then writes session JSON.
#[tokio::test]
#[ignore = "requires pre-built grok binary; run with --ignored"]
#[serial]
async fn headless_session_refreshes_trusted_local_plugin_and_writes_session_json() {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Canonicalize so under-home auto-trust holds where the temp root is a
// symlink (macOS `/var` -> `/private/var`).
let home_tmp = TempDir::new().unwrap();
let home = dunce::canonicalize(home_tmp.path()).unwrap();
let kigi_home = home.join(".kigi");
std::fs::create_dir_all(&kigi_home).unwrap();
let source = home
.join(".claude")
.join("local-marketplace")
.join("demo-plugin");
write_minimal_plugin(&source, "demo-plugin");
write_agent(&source, "old", "old-agent", "exists at install");
// The spawned binary gets HOME/KIGI_SHARE_DIR via `cmd.env` below; this global env
// is only for the in-process post-run discovery assertion (which resolves the
// registry via kigi_home()). `#[serial]` keeps it from racing other tests.
let _home_guard = EnvVarGuard::set("HOME", &home);
let _grok_guard = EnvVarGuard::set("KIGI_SHARE_DIR", &kigi_home);
let mut registry = InstallRegistry::empty(kigi_home.join("installed-plugins"));
let installed = register_local_install(&mut registry, &source);
registry.save().expect("save registry");
write_agent(&source, "new", "new-agent", "added after install");
assert!(!installed.path.join("agents/new.md").exists());
let workdir = git_workdir();
let mut cmd = tokio::process::Command::new(grok_binary());
cmd.args([
"-p",
"say hello",
"--yolo",
"--output-format",
"json",
"--cwd",
])
.arg(workdir.path())
.current_dir(workdir.path())
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
kigi_test_support::env::test_env_cmd_tokio(&mut cmd, &server.url(), &home);
cmd.env("HOME", &home);
cmd.env("KIGI_SHARE_DIR", &kigi_home);
let result = run_headless_with_cmd(cmd).await;
assert_headless_success(
&result,
"headless session with trusted local plugin refresh",
Some(&server),
);
assert_no_crashes(&result.stderr);
assert!(
installed.path.join("agents/new.md").exists(),
"session startup should have re-copied new agent into installed-plugins; stderr=\n{}",
result.stderr
);
// The real binary's session start refreshed the on-disk snapshot; rebuild the
// registry for the workdir and assert the new agent surfaces in `/agents`
// (same proof as the library test, but through the real binary).
let config = DiscoveryConfig {
cli_plugin_dirs: Vec::new(),
config_paths: Vec::new(),
disabled: Vec::new(),
enabled: vec!["demo-plugin".to_string()],
};
let plugin_registry = SharedPluginRegistryHandle::new(None, Vec::new())
.build_for_cwd(workdir.path(), &config, &[], true)
.expect("registry built from refreshed snapshot");
let agents = kigi_agent::discovery::all_subagents_with_plugins(
workdir.path(),
&HashMap::new(),
Some(plugin_registry.as_ref()),
);
assert!(
agents.iter().any(|a| a.name == "demo-plugin:new-agent"),
"new agent must surface in /agents after the binary's session-start refresh"
);
// Smoke: session storage under KIGI_SHARE_DIR/sessions has JSON artifacts.
let sessions_root = kigi_home.join("sessions");
assert!(
sessions_root.is_dir(),
"expected sessions dir at {}",
sessions_root.display()
);
let mut json_files: Vec<PathBuf> = Vec::new();
collect_json_files(&sessions_root, &mut json_files);
assert!(
!json_files.is_empty(),
"expected session JSON under {}",
sessions_root.display()
);
for path in &json_files {
let text = std::fs::read_to_string(path).unwrap_or_default();
let trimmed = text.trim();
if trimmed.is_empty() {
continue;
}
if path.extension().is_some_and(|e| e == "jsonl") {
for line in trimmed.lines().filter(|l| !l.trim().is_empty()) {
serde_json::from_str::<serde_json::Value>(line).unwrap_or_else(|e| {
panic!("invalid JSONL line in {}: {e}\n{line}", path.display())
});
}
} else {
serde_json::from_str::<serde_json::Value>(trimmed)
.unwrap_or_else(|e| panic!("invalid JSON in {}: {e}", path.display()));
}
}
}
fn collect_json_files(dir: &Path, out: &mut Vec<PathBuf>) {
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
collect_json_files(&path, out);
} else if path
.extension()
.is_some_and(|e| e == "json" || e == "jsonl")
{
out.push(path);
}
}
}
@@ -0,0 +1,366 @@
//! Vendor-compatibility end-to-end tests.
//!
//! Each test builds a fake `$HOME` containing skills/rules/AGENTS.md under the
//! `.kigi`, `.cursor`, and `.claude` vendor dirs, spawns a real `grok agent
//! stdio` process against the mock inference server (toggling the
//! `GROK_<VENDOR>_<SURFACE>_ENABLED` env vars via `cmd.env`), sends one prompt,
//! and asserts on the full inference request bodies:
//!
//! - the Grok-native skill is always present regardless of toggles
//! - each of the 6 (vendor x surface) cells toggles independently
//! - a vendor-shipped default skill (`shell`) under `~/.cursor` is always
//! dropped by the denylist
//! - cross-vendor combos (all-cursor-off, all-claude-off, all-off) work
//!
//! These are `#[ignore]` (they spawn a built binary) like the agent-type
//! invariant suite. Run locally:
//! ```bash
//! cargo test -p kigi-shell --test test_vendor_compat -- --ignored
//! ```
use std::future::Future;
use std::path::Path;
use kigi_test_support::*;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Unique markers placed in skill descriptions / file contents so assertions
/// can't be fooled by incidental occurrences of a bare word like "shell".
const MARKER_GROK_SKILL: &str = "ZZ_GROK_SKILL_MARKER";
const MARKER_CURSOR_SKILL: &str = "ZZ_CURSOR_SKILL_MARKER";
const MARKER_CURSOR_SHELL: &str = "ZZ_CURSOR_SHELL_DENYLISTED_MARKER";
const MARKER_CLAUDE_SKILL: &str = "ZZ_CLAUDE_SKILL_MARKER";
const MARKER_CURSOR_RULE: &str = "ZZ_CURSOR_RULE_MARKER";
const MARKER_CLAUDE_RULE: &str = "ZZ_CLAUDE_RULE_MARKER";
const MARKER_CLAUDE_AGENTS: &str = "ZZ_CLAUDE_AGENTS_MARKER";
const MARKER_CURSOR_AGENTS: &str = "ZZ_CURSOR_AGENTS_MARKER";
fn write_file(path: &Path, contents: &str) {
std::fs::create_dir_all(path.parent().expect("path has parent")).expect("create dirs");
std::fs::write(path, contents).expect("write file");
}
/// Write a `<vendor>/skills/<name>/SKILL.md` with the given description marker.
fn write_skill(home: &Path, vendor_dir: &str, name: &str, marker: &str) {
let p = home
.join(vendor_dir)
.join("skills")
.join(name)
.join("SKILL.md");
write_file(
&p,
&format!("---\nname: {name}\ndescription: {marker}\n---\n\nSkill body.\n"),
);
}
/// Populate a fake `$HOME` + repo cwd with the full vendor-compat fixture set.
fn seed_fixtures(home: &Path, cwd: &Path) {
// Skills (User scope, home-based).
write_skill(home, ".kigi", "grok-skill", MARKER_GROK_SKILL);
write_skill(home, ".cursor", "my-cursor-skill", MARKER_CURSOR_SKILL);
// `shell` is a Cursor vendor-default → must be denylisted under ~/.cursor.
write_skill(home, ".cursor", "shell", MARKER_CURSOR_SHELL);
write_skill(home, ".claude", "my-claude-skill", MARKER_CLAUDE_SKILL);
// Rules: repo-local `.cursor/rules/r.md` and `.claude/rules/c.md`
// (discovered via the cwd→root walk, gated by their respective rules cell).
write_file(
&cwd.join(".cursor").join("rules").join("r.md"),
&format!("# rule\n{MARKER_CURSOR_RULE}\n"),
);
write_file(
&cwd.join(".claude").join("rules").join("c.md"),
&format!("# rule\n{MARKER_CLAUDE_RULE}\n"),
);
// AGENTS.md: `~/.claude/CLAUDE.md` and `~/.cursor/AGENTS.md`
// (discovered via the home compat scan, gated by their respective agents cell).
write_file(
&home.join(".claude").join("CLAUDE.md"),
&format!("# claude instructions\n{MARKER_CLAUDE_AGENTS}\n"),
);
write_file(
&home.join(".cursor").join("AGENTS.md"),
&format!("# cursor instructions\n{MARKER_CURSOR_AGENTS}\n"),
);
}
/// Spawn the agent with the given compat env overrides, send one prompt, and
/// return every inference request body concatenated into one string for
/// substring assertions (system prompt + skill listing + injected reminders).
async fn run_scenario(env: &[(&str, &str)]) -> String {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let workdir = git_workdir();
let home = tempfile::TempDir::new().expect("create temp home");
seed_fixtures(home.path(), workdir.path());
let client = GrokStdioClient::spawn_with_home_and_env(&server, workdir.path(), home, env).await;
client.initialize_with_timeout().await;
let session_id = client.create_session_with_timeout(workdir.path()).await;
let _ = client.prompt_with_timeout(&session_id, "hello").await;
let bodies: Vec<String> = server
.requests()
.iter()
.filter_map(|e| e.body.as_ref().map(|b| b.to_string()))
.collect();
assert!(
!bodies.is_empty(),
"expected at least one inference request; stderr:\n{}",
client.stderr()
);
bodies.join("\n---\n")
}
// ── Skills ──────────────────────────────────────────────────────────────────
/// Defaults (all vendors on): grok + cursor-vendor + claude-vendor skills present; the
/// denylisted vendor builtin `shell` is dropped.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_defaults_include_vendor_skills_but_drop_denylisted() {
with_local_set(|| async {
let body = run_scenario(&[]).await;
assert!(
body.contains(MARKER_GROK_SKILL),
"grok-skill must always be present"
);
assert!(
body.contains(MARKER_CURSOR_SKILL),
"cursor skill present when cursor.skills on (default)"
);
assert!(
body.contains(MARKER_CLAUDE_SKILL),
"claude skill present when claude.skills on (default)"
);
assert!(
!body.contains(MARKER_CURSOR_SHELL),
"denylisted Cursor builtin `shell` must be dropped"
);
})
.await;
}
/// `KIGI_CURSOR_SKILLS_ENABLED=false` drops the cursor-vendor skill; grok stays.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_cursor_skills_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CURSOR_SKILLS_ENABLED", "false")]).await;
assert!(
body.contains(MARKER_GROK_SKILL),
"grok-skill always present"
);
assert!(
!body.contains(MARKER_CURSOR_SKILL),
"cursor skill must be absent when cursor.skills disabled"
);
// Denylist still applies regardless of the toggle.
assert!(!body.contains(MARKER_CURSOR_SHELL));
})
.await;
}
/// `KIGI_CLAUDE_SKILLS_ENABLED=false` drops the claude-vendor skill; grok stays.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_claude_skills_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CLAUDE_SKILLS_ENABLED", "false")]).await;
assert!(
body.contains(MARKER_GROK_SKILL),
"grok-skill always present"
);
assert!(
!body.contains(MARKER_CLAUDE_SKILL),
"claude skill must be absent when claude.skills disabled"
);
})
.await;
}
// ── Rules + AGENTS.md ────────────────────────────────────────────────────────
/// Defaults: all rules and AGENTS.md surfaces are present.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_rules_and_agents_present_by_default() {
with_local_set(|| async {
let body = run_scenario(&[]).await;
assert!(
body.contains(MARKER_CURSOR_RULE),
"cursor rule present when cursor.rules on (default)"
);
assert!(
body.contains(MARKER_CLAUDE_RULE),
"claude rule present when claude.rules on (default)"
);
assert!(
body.contains(MARKER_CLAUDE_AGENTS),
"claude AGENTS.md present when claude.agents on (default)"
);
assert!(
body.contains(MARKER_CURSOR_AGENTS),
"cursor AGENTS.md present when cursor.agents on (default)"
);
})
.await;
}
// ── Per-cell toggles (rules + agents) ────────────────────────────────────────
/// `KIGI_CURSOR_RULES_ENABLED=false` drops cursor-vendor rules; claude-vendor rules stay.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_cursor_rules_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CURSOR_RULES_ENABLED", "false")]).await;
assert!(
!body.contains(MARKER_CURSOR_RULE),
"cursor rule must be absent when cursor.rules disabled"
);
assert!(
body.contains(MARKER_CLAUDE_RULE),
"claude rule unaffected by cursor.rules toggle"
);
})
.await;
}
/// `KIGI_CLAUDE_RULES_ENABLED=false` drops claude-vendor rules; cursor-vendor rules stay.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_claude_rules_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CLAUDE_RULES_ENABLED", "false")]).await;
assert!(
!body.contains(MARKER_CLAUDE_RULE),
"claude rule must be absent when claude.rules disabled"
);
assert!(
body.contains(MARKER_CURSOR_RULE),
"cursor rule unaffected by claude.rules toggle"
);
})
.await;
}
/// `KIGI_CURSOR_AGENTS_ENABLED=false` drops cursor-vendor AGENTS.md; claude-vendor stays.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_cursor_agents_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CURSOR_AGENTS_ENABLED", "false")]).await;
assert!(
!body.contains(MARKER_CURSOR_AGENTS),
"cursor AGENTS.md must be absent when cursor.agents disabled"
);
assert!(
body.contains(MARKER_CLAUDE_AGENTS),
"claude AGENTS.md unaffected by cursor.agents toggle"
);
})
.await;
}
/// `KIGI_CLAUDE_AGENTS_ENABLED=false` drops claude-vendor AGENTS.md; cursor-vendor stays.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_claude_agents_disabled() {
with_local_set(|| async {
let body = run_scenario(&[("KIGI_CLAUDE_AGENTS_ENABLED", "false")]).await;
assert!(
!body.contains(MARKER_CLAUDE_AGENTS),
"claude AGENTS.md must be absent when claude.agents disabled"
);
assert!(
body.contains(MARKER_CURSOR_AGENTS),
"cursor AGENTS.md unaffected by claude.agents toggle"
);
})
.await;
}
// ── Cross-vendor combinations ────────────────────────────────────────────────
/// All cursor-vendor compat OFF: cursor skills, rules, and AGENTS.md all absent;
/// all claude-vendor surfaces unaffected.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_all_cursor_disabled() {
with_local_set(|| async {
let body = run_scenario(&[
("KIGI_CURSOR_SKILLS_ENABLED", "false"),
("KIGI_CURSOR_RULES_ENABLED", "false"),
("KIGI_CURSOR_AGENTS_ENABLED", "false"),
])
.await;
assert!(!body.contains(MARKER_CURSOR_SKILL));
assert!(!body.contains(MARKER_CURSOR_SHELL));
assert!(!body.contains(MARKER_CURSOR_RULE));
assert!(!body.contains(MARKER_CURSOR_AGENTS));
assert!(body.contains(MARKER_GROK_SKILL), "grok always present");
assert!(body.contains(MARKER_CLAUDE_SKILL), "claude unaffected");
assert!(body.contains(MARKER_CLAUDE_RULE), "claude unaffected");
assert!(body.contains(MARKER_CLAUDE_AGENTS), "claude unaffected");
})
.await;
}
/// All claude-vendor compat OFF: claude skills, rules, and AGENTS.md all absent;
/// all cursor-vendor surfaces unaffected.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_all_claude_disabled() {
with_local_set(|| async {
let body = run_scenario(&[
("KIGI_CLAUDE_SKILLS_ENABLED", "false"),
("KIGI_CLAUDE_RULES_ENABLED", "false"),
("KIGI_CLAUDE_AGENTS_ENABLED", "false"),
])
.await;
assert!(!body.contains(MARKER_CLAUDE_SKILL));
assert!(!body.contains(MARKER_CLAUDE_RULE));
assert!(!body.contains(MARKER_CLAUDE_AGENTS));
assert!(body.contains(MARKER_GROK_SKILL), "grok always present");
assert!(body.contains(MARKER_CURSOR_SKILL), "cursor unaffected");
assert!(body.contains(MARKER_CURSOR_RULE), "cursor unaffected");
assert!(body.contains(MARKER_CURSOR_AGENTS), "cursor unaffected");
assert!(!body.contains(MARKER_CURSOR_SHELL), "denylist still active");
})
.await;
}
/// All vendor compat OFF: only grok-native skill survives.
#[tokio::test]
#[ignore] // requires pre-built binary
async fn vendor_compat_all_vendors_disabled() {
with_local_set(|| async {
let body = run_scenario(&[
("KIGI_CURSOR_SKILLS_ENABLED", "false"),
("KIGI_CURSOR_RULES_ENABLED", "false"),
("KIGI_CURSOR_AGENTS_ENABLED", "false"),
("KIGI_CLAUDE_SKILLS_ENABLED", "false"),
("KIGI_CLAUDE_RULES_ENABLED", "false"),
("KIGI_CLAUDE_AGENTS_ENABLED", "false"),
])
.await;
assert!(body.contains(MARKER_GROK_SKILL), "grok always present");
assert!(!body.contains(MARKER_CURSOR_SKILL));
assert!(!body.contains(MARKER_CURSOR_SHELL));
assert!(!body.contains(MARKER_CURSOR_RULE));
assert!(!body.contains(MARKER_CURSOR_AGENTS));
assert!(!body.contains(MARKER_CLAUDE_SKILL));
assert!(!body.contains(MARKER_CLAUDE_RULE));
assert!(!body.contains(MARKER_CLAUDE_AGENTS));
})
.await;
}
@@ -0,0 +1,313 @@
//! Integration test for `_x.ai/session/update` notifications.
//!
//! This test verifies that:
//! 1. xAI session notifications (e.g., diff_review) can be sent via ext_notification
//! 2. The notifications are persisted to storage
//! 3. When a session is loaded, the notifications are replayed with `isReplay: true`
use agent_client_protocol as acp;
use serde_json::json;
use std::path::PathBuf;
use tempfile::TempDir;
use kigi_shell::extensions::notification::{
DiffContent, SessionNotification, SessionUpdate as XaiSessionUpdate,
};
use kigi_shell::session::info::Info as SessionInfo;
use kigi_shell::session::persistence::default_model_id;
use kigi_shell::session::storage::{JsonlStorageAdapter, SessionUpdate, StorageAdapter};
/// Test that xAI session notifications round-trip through storage correctly.
#[tokio::test]
async fn test_xai_session_notification_storage_roundtrip() {
let temp_dir = TempDir::new().unwrap();
let adapter = JsonlStorageAdapter::with_root(temp_dir.path().to_path_buf());
let session_id = acp::SessionId::new("test-session-roundtrip");
let info = SessionInfo {
id: session_id.clone(),
cwd: "/test/workspace".to_string(),
};
// Initialize the session
adapter
.init_session(&info, default_model_id())
.await
.unwrap();
// Create a diff_review notification
let xai_notification = SessionNotification {
session_id: session_id.clone(),
update: XaiSessionUpdate::DiffReview {
content: vec![DiffContent {
diff: acp::Diff::new(PathBuf::from("/test/file.rs"), "fn new() {}".to_string())
.old_text(Some("fn old() {}".to_string())),
}],
},
meta: Some(json!({ "totalTokens": 1234 })),
};
// Persist the notification
adapter
.append_update(
&info,
&SessionUpdate::Xai(Box::new(xai_notification.clone())),
)
.await
.unwrap();
// Also add an ACP notification to verify mixed storage works
let acp_notification = acp::SessionNotification::new(
session_id.clone(),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new("Hello from agent".to_string()),
))),
)
.meta(json!({ "totalTokens": 5678 }).as_object().cloned());
adapter
.append_update(&info, &SessionUpdate::Acp(Box::new(acp_notification)))
.await
.unwrap();
// Load the session and verify both notifications are present
let loaded = adapter.load_session(&info).await.unwrap();
assert_eq!(
loaded.updates.len(),
2,
"Should have 2 updates (1 xAI + 1 ACP)"
);
// Verify xAI notification
match &loaded.updates[0] {
SessionUpdate::Xai(notification) => {
assert_eq!(notification.session_id, session_id);
match &notification.update {
XaiSessionUpdate::DiffReview { content } => {
assert_eq!(content.len(), 1);
assert_eq!(content[0].diff.path, PathBuf::from("/test/file.rs"));
assert_eq!(content[0].diff.old_text, Some("fn old() {}".to_string()));
assert_eq!(content[0].diff.new_text, "fn new() {}");
}
_ => {
panic!("Expected DiffReview, got different update type");
}
}
// Verify meta is preserved
assert_eq!(
notification
.meta
.as_ref()
.and_then(|m| m.get("totalTokens")),
Some(&json!(1234))
);
}
_ => panic!("Expected xAI update as first item"),
}
// Verify ACP notification
match &loaded.updates[1] {
SessionUpdate::Acp(notification) => {
assert_eq!(notification.session_id, session_id);
assert_eq!(
notification
.meta
.as_ref()
.and_then(|m| m.get("totalTokens")),
Some(&json!(5678))
);
}
_ => panic!("Expected ACP update as second item"),
}
}
/// Test that a `TurnCompleted` terminal round-trips through storage — the
/// persistence half of the "stuck on Waiting…" fix, where the durable terminal
/// must survive `updates.jsonl` and reload as a replayable `_x.ai/session/update`.
#[tokio::test]
async fn test_turn_completed_round_trips_through_storage() {
let temp_dir = TempDir::new().unwrap();
let adapter = JsonlStorageAdapter::with_root(temp_dir.path().to_path_buf());
let session_id = acp::SessionId::new("test-session-turn-completed");
let info = SessionInfo {
id: session_id.clone(),
cwd: "/test/workspace".to_string(),
};
adapter
.init_session(&info, default_model_id())
.await
.unwrap();
// Persist a terminal carrying the prompt id + outcome the viewer keys on,
// plus an optional agent result.
let xai_notification = SessionNotification {
session_id: session_id.clone(),
update: XaiSessionUpdate::TurnCompleted {
prompt_id: "prompt-1".to_string(),
stop_reason: "end_turn".to_string(),
agent_result: Some("all done".to_string()),
usage: None,
},
meta: None,
};
adapter
.append_update(&info, &SessionUpdate::Xai(Box::new(xai_notification)))
.await
.unwrap();
// Reload the session (the replay path) and confirm the terminal survives
// with its fields intact.
let loaded = adapter.load_session(&info).await.unwrap();
assert_eq!(
loaded.updates.len(),
1,
"Should have 1 update (the terminal)"
);
match &loaded.updates[0] {
SessionUpdate::Xai(notification) => {
assert_eq!(notification.session_id, session_id);
match &notification.update {
XaiSessionUpdate::TurnCompleted {
prompt_id,
stop_reason,
agent_result,
..
} => {
assert_eq!(prompt_id, "prompt-1");
assert_eq!(stop_reason, "end_turn");
assert_eq!(agent_result.as_deref(), Some("all done"));
}
_ => panic!("Expected TurnCompleted, got different update type"),
}
}
_ => panic!("Expected xAI update"),
}
}
/// Test that totalTokens can be extracted from both ACP and xAI notifications.
#[tokio::test]
async fn test_extract_total_tokens_from_mixed_updates() {
let temp_dir = TempDir::new().unwrap();
let adapter = JsonlStorageAdapter::with_root(temp_dir.path().to_path_buf());
let session_id = acp::SessionId::new("test-session-tokens");
let info = SessionInfo {
id: session_id.clone(),
cwd: "/test/workspace".to_string(),
};
adapter
.init_session(&info, default_model_id())
.await
.unwrap();
// Add ACP notification with totalTokens
let acp_notification = acp::SessionNotification::new(
session_id.clone(),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new("First message".to_string()),
))),
)
.meta(json!({ "totalTokens": 100 }).as_object().cloned());
adapter
.append_update(&info, &SessionUpdate::Acp(Box::new(acp_notification)))
.await
.unwrap();
// Add xAI notification with totalTokens
let xai_notification = SessionNotification {
session_id: session_id.clone(),
update: XaiSessionUpdate::DiffReview { content: vec![] },
meta: Some(json!({ "totalTokens": 200 })),
};
adapter
.append_update(&info, &SessionUpdate::Xai(Box::new(xai_notification)))
.await
.unwrap();
// Add another ACP notification with higher totalTokens
let acp_notification2 = acp::SessionNotification::new(
session_id.clone(),
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::Text(
acp::TextContent::new("Second message".to_string()),
))),
)
.meta(json!({ "totalTokens": 300 }).as_object().cloned());
adapter
.append_update(&info, &SessionUpdate::Acp(Box::new(acp_notification2)))
.await
.unwrap();
// Load and extract totalTokens (simulating what load_session does in mvp_agent)
let loaded = adapter.load_session(&info).await.unwrap();
let last_total_tokens = loaded
.updates
.iter()
.rev()
.find_map(|notification| match notification {
SessionUpdate::Acp(n) => n
.meta
.as_ref()
.and_then(|m| m.get("totalTokens"))
.and_then(|v| v.as_u64()),
SessionUpdate::Xai(n) => n
.meta
.as_ref()
.and_then(|m| m.get("totalTokens"))
.and_then(|v| v.as_u64()),
})
.unwrap_or(0);
assert_eq!(
last_total_tokens, 300,
"Should get the last totalTokens value"
);
}
/// Test the serialization format of SessionNotification for wire compatibility.
#[test]
fn test_xai_session_notification_serialization() {
let notification = SessionNotification {
session_id: acp::SessionId::new("sess-123"),
update: XaiSessionUpdate::DiffReview {
content: vec![DiffContent {
diff: acp::Diff::new(PathBuf::from("src/main.rs"), "new".to_string())
.old_text(Some("old".to_string())),
}],
},
meta: Some(json!({ "isReplay": true })),
};
let json = serde_json::to_value(&notification).unwrap();
// Print actual JSON for debugging
println!(
"Serialized JSON: {}",
serde_json::to_string_pretty(&json).unwrap()
);
// Verify camelCase field names
assert!(json.get("sessionId").is_some(), "Expected sessionId field");
assert!(json.get("_meta").is_some(), "Expected _meta field");
// The update field contains the nested SessionUpdate which has the tag
let update_obj = json.get("update").expect("Expected update field");
let session_update_tag = update_obj
.get("sessionUpdate")
.expect("Expected sessionUpdate tag in update");
assert_eq!(session_update_tag, "diff_review");
// Verify diff content is inside the update
let content = update_obj
.get("content")
.expect("Expected content in update")
.as_array()
.unwrap();
assert_eq!(content.len(), 1);
assert_eq!(content[0]["type"], "diff");
assert_eq!(content[0]["path"], "src/main.rs");
}
@@ -0,0 +1,278 @@
//! Replay-trace verification harness for the turn-end TodoGate.
//!
//! Reads synthetic JSON fixtures from `tests/fixtures/synthetic_*.json`,
//! walks each turn, and at every assistant `end-of-turn` snapshot asserts
//! that `evaluate_todo_gate` returns the decision the fixture declares.
//!
//! Pure-function integration test — no `SessionActor`, no completion
//! stream. Complements unit tests of the pure function with
//! data-driven trace-replay coverage.
//!
//! Data-driven: dropping a new `synthetic_*.json` fixture into
//! `tests/fixtures/` enrolls it in the harness automatically.
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use serde::Deserialize;
use kigi_shell::session::{
CollectedTodoGateInput, TodoGateDecision, TodoGateReason, evaluate_todo_gate,
};
use kigi_shell::tools::todo::TodoStatus;
/// Closed canonical set of shipped fixtures. Adding a fixture
/// here is the one required Rust-side change when a new failure shape
/// is enrolled — `canonical_fixtures_present` asserts set equality
/// against this list so a missing-or-extra fixture fails the harness.
const CANONICAL_FIXTURES: &[&str] = &[
"synthetic_clean_completion.json",
"synthetic_pr_babysit_partial_backing.json",
"synthetic_stranded_narration.json",
];
#[derive(Debug, Deserialize)]
struct Fixture {
name: String,
#[allow(dead_code)] // human-readable; surfaced only on assertion failure
description: String,
turns: Vec<Turn>,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
enum Turn {
#[allow(dead_code)] // user turns are walked but never gate-evaluated
User(UserTurn),
Assistant(AssistantTurn),
}
#[derive(Debug, Deserialize)]
#[allow(dead_code)]
struct UserTurn {
turn_index: usize,
content: String,
}
#[derive(Debug, Deserialize)]
struct AssistantTurn {
turn_index: usize,
#[allow(dead_code)] // present for fixture clarity; the gate does not consult it
tool_calls_emitted: Vec<serde_json::Value>,
todo_state_after_turn: Vec<TodoSnapshot>,
backing_task_count: usize,
expected_gate_decision: ExpectedGateDecision,
#[serde(default)]
expected_reason: Option<ExpectedReason>,
#[serde(default)]
expected_reminder_contains: Vec<String>,
}
#[derive(Debug, Deserialize)]
struct TodoSnapshot {
#[allow(dead_code)] // id is preserved for fixture readability
id: String,
status: TodoStatus,
content: String,
}
/// Typed mirror of the fixture's `expected_gate_decision` field. A
/// closed enum (not `String`) catches typos at deserialize time rather
/// than silently passing the wrong assertion branch.
#[derive(Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
enum ExpectedGateDecision {
Nudge,
Continue,
}
#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
enum ExpectedReason {
InFlight,
}
impl From<TodoGateReason> for ExpectedReason {
fn from(reason: TodoGateReason) -> Self {
match reason {
TodoGateReason::InFlight => Self::InFlight,
}
}
}
fn fixtures_dir() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures")
}
/// Every `synthetic_*.json` in `tests/fixtures/`, sorted for stable
/// output. Directory-iteration IO errors panic with diagnostic context
/// rather than being silently dropped (a permissioned-out fixture must
/// not vanish from the set unnoticed).
fn discover_fixtures() -> Vec<PathBuf> {
let dir = fixtures_dir();
let entries = std::fs::read_dir(&dir)
.unwrap_or_else(|e| panic!("read fixtures dir {}: {e}", dir.display()));
let mut paths: Vec<PathBuf> = entries
.map(|entry| {
entry
.unwrap_or_else(|e| panic!("read entry in {}: {e}", dir.display()))
.path()
})
.filter(|path| {
path.file_name()
.and_then(|n| n.to_str())
.is_some_and(|name| name.starts_with("synthetic_") && name.ends_with(".json"))
})
.collect();
paths.sort();
assert!(
!paths.is_empty(),
"no synthetic_*.json fixtures found in {}",
dir.display()
);
paths
}
fn load_fixture(path: &Path) -> Fixture {
let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {}: {e}", path.display()));
serde_json::from_slice(&bytes).unwrap_or_else(|e| panic!("parse {}: {e}", path.display()))
}
/// Consume a fixture's assistant-turn snapshot into a
/// `CollectedTodoGateInput`. Moves todo `content` strings out of the
/// snapshot — no clones. Calling `as_input()` on the result also
/// re-exercises the production "first N in_progress are backed"
/// partition heuristic.
fn collected_from(
snapshots: Vec<TodoSnapshot>,
backing_task_count: usize,
) -> CollectedTodoGateInput {
let todos = snapshots
.into_iter()
.map(|t| (t.id, t.content, t.status))
.collect();
CollectedTodoGateInput {
todos,
backing_task_count,
}
}
/// Run the gate against one assistant turn. `fixture_name` is only used
/// in panic messages so failures point straight at the offending JSON.
fn check_assistant_turn(fixture_name: &str, turn: AssistantTurn) {
let AssistantTurn {
turn_index,
tool_calls_emitted: _,
todo_state_after_turn,
backing_task_count,
expected_gate_decision,
expected_reason,
expected_reminder_contains,
} = turn;
let collected = collected_from(todo_state_after_turn, backing_task_count);
let input = collected.as_input();
let decision = evaluate_todo_gate(&input);
match (expected_gate_decision, decision) {
(ExpectedGateDecision::Continue, TodoGateDecision::Continue) => {
assert!(
expected_reason.is_none(),
"fixture {fixture_name} turn {turn_index} declares `continue` with \
`expected_reason` a continue decision has no reason",
);
assert!(
expected_reminder_contains.is_empty(),
"fixture {fixture_name} turn {turn_index} declares `continue` with \
`expected_reminder_contains` a continue decision emits no reminder",
);
}
(ExpectedGateDecision::Nudge, TodoGateDecision::Nudge { reminder, reason }) => {
if let Some(expected) = expected_reason {
assert_eq!(
expected,
ExpectedReason::from(reason),
"fixture {fixture_name} turn {turn_index}: gate reason mismatch",
);
}
for needle in &expected_reminder_contains {
assert!(
reminder.contains(needle.as_str()),
"fixture {fixture_name} turn {turn_index}: reminder missing substring \
{needle:?}.\nFull reminder:\n{reminder}",
);
}
}
(expected, TodoGateDecision::Continue) => {
panic!("fixture {fixture_name} turn {turn_index}: expected {expected:?}, got Continue",)
}
(expected, TodoGateDecision::Nudge { reason, .. }) => panic!(
"fixture {fixture_name} turn {turn_index}: expected {expected:?}, got Nudge({reason:?})",
),
}
}
#[test]
fn replay_all_synthetic_fixtures() {
for path in discover_fixtures() {
let fixture = load_fixture(&path);
let Fixture {
name,
description: _,
turns,
} = fixture;
let mut saw_assistant = false;
for turn in turns {
match turn {
Turn::User(_) => {}
Turn::Assistant(at) => {
saw_assistant = true;
check_assistant_turn(&name, at);
}
}
}
assert!(
saw_assistant,
"fixture {name} ({}) has no assistant turns to evaluate",
path.display(),
);
}
}
/// Set-equality guard: the on-disk fixture set must exactly equal
/// [`CANONICAL_FIXTURES`]. Adding a fixture without updating the
/// constant — or losing one — fails the harness. Closes the
/// "open-ended presence check" gap.
#[test]
fn canonical_fixtures_match_disk() {
let actual: BTreeSet<String> = discover_fixtures()
.iter()
.map(|p| {
p.file_name()
.and_then(|n| n.to_str())
.unwrap_or_else(|| panic!("non-UTF8 fixture path: {}", p.display()))
.to_string()
})
.collect();
let expected: BTreeSet<String> = CANONICAL_FIXTURES
.iter()
.map(|s| (*s).to_string())
.collect();
assert_eq!(actual, expected, "fixture set drift vs CANONICAL_FIXTURES");
}
/// Compile-time guard: `ExpectedGateDecision` must stay a closed
/// two-variant enum so a fixture typo fails to load instead of
/// silently passing the wrong branch.
#[test]
fn expected_decision_is_closed() {
let parsed: ExpectedGateDecision = serde_json::from_str(r#""nudge""#).unwrap();
assert_eq!(parsed, ExpectedGateDecision::Nudge);
let parsed: ExpectedGateDecision = serde_json::from_str(r#""continue""#).unwrap();
assert_eq!(parsed, ExpectedGateDecision::Continue);
let err = serde_json::from_str::<ExpectedGateDecision>(r#""maybe""#).unwrap_err();
assert!(
err.to_string().contains("unknown variant"),
"expected unknown-variant error, got: {err}"
);
}