Synced from monorepo
Changes: - Detect the herdr multiplexer - Mark /gboom as non-production code - Bound peak memory when loading a large session - Add a subagent lifecycle soak bounding threads, fds, and heap - Stream inherited replay to bound fork memory - Copy full plan from plan approval with y - Stop armed signature verification from deleting the managed-deny smoke policy - Add source-tagged terminal version telemetry - Show the UI instantly and fetch models and settings in the background - Session test helpers - computer_reason on the ConversationHistoryDone trailer
This commit is contained in:
grokkybara[bot]
parent
47348d13ec
commit
b41c75a578
92 changed files with 9410 additions and 3788 deletions
|
|
@ -0,0 +1,186 @@
|
|||
//! Regression guard for the fork/resume replay OOM: a counting allocator checks
|
||||
//! the streaming load ([`stream_replay_updates_at`]) peaks far below the old
|
||||
//! parse-all load and forwards identical content.
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
use std::alloc::{GlobalAlloc, Layout, System};
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use xai_grok_shell::session::storage::{
|
||||
ReplayEmission, SessionUpdate, UpdatesIterator, filter_rewind_updates,
|
||||
stream_replay_updates_at, strip_context_wrappers,
|
||||
};
|
||||
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
|
||||
|
||||
// `Relaxed` suffices: single-threaded high-water counters read after the
|
||||
// measured section, with no ordering dependency on other memory.
|
||||
struct CountingAlloc;
|
||||
static LIVE: AtomicUsize = AtomicUsize::new(0);
|
||||
static PEAK: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
unsafe impl GlobalAlloc for CountingAlloc {
|
||||
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
||||
let ptr = unsafe { System.alloc(layout) };
|
||||
if !ptr.is_null() {
|
||||
let now = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
|
||||
PEAK.fetch_max(now, Ordering::Relaxed);
|
||||
}
|
||||
ptr
|
||||
}
|
||||
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
||||
LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
|
||||
unsafe { System.dealloc(ptr, layout) }
|
||||
}
|
||||
}
|
||||
|
||||
#[global_allocator]
|
||||
static ALLOC: CountingAlloc = CountingAlloc;
|
||||
|
||||
fn begin_measure() -> usize {
|
||||
let base = LIVE.load(Ordering::Relaxed);
|
||||
PEAK.store(base, Ordering::Relaxed);
|
||||
base
|
||||
}
|
||||
fn peak() -> usize {
|
||||
PEAK.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Heavy ACU catalog, no rewind points: the parity assert covers the
|
||||
/// typed-to-string swap over this transcript, not rewind-marker filtering
|
||||
/// (`synth` emits no `rewind_marker` envelopes on the replay path).
|
||||
fn fork_replay_spec() -> SessionSpec {
|
||||
SessionSpec::from_env_prefixed(
|
||||
"FORK_REPLAY",
|
||||
SessionSpec {
|
||||
turns: 60,
|
||||
acu_per_turn: 8,
|
||||
catalog_commands: 200,
|
||||
catalog_desc_len: 48,
|
||||
agent_chunks_per_turn: 4,
|
||||
agent_chunk_len: 1500,
|
||||
rewind_points: 0,
|
||||
files_per_rewind: 0,
|
||||
file_content_len: 0,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn reference_load_all(updates_path: &Path) -> Vec<acp::SessionUpdate> {
|
||||
let iter = UpdatesIterator::open(updates_path)
|
||||
.expect("open updates")
|
||||
.expect("updates file exists");
|
||||
let all: Vec<SessionUpdate> = iter.filter_map(|r| r.ok()).collect();
|
||||
let filtered = filter_rewind_updates(all);
|
||||
filtered
|
||||
.into_iter()
|
||||
.filter_map(|u| match u {
|
||||
SessionUpdate::Acp(notif) => Some(strip_context_wrappers(notif.update)),
|
||||
SessionUpdate::Xai(_) => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn serialize(u: &acp::SessionUpdate) -> String {
|
||||
serde_json::to_string(u).expect("serialize replayed update")
|
||||
}
|
||||
|
||||
/// Streaming holds one `read_to_string` copy (~1x) plus transient per-update
|
||||
/// structs; the old parse-all path held several multiples. Headroom over 1x.
|
||||
const MAX_STREAM_PEAK_TO_DISK_RATIO: f64 = 2.0;
|
||||
|
||||
// Serial: the process-global counters are only valid when no other test
|
||||
// allocates concurrently.
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn fork_replay_stream_is_bounded_and_faithful() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let opts = fork_replay_spec();
|
||||
|
||||
let (id, updates_path) = {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let out = rt.block_on(async {
|
||||
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
|
||||
(info.id.0.to_string(), dir.join("updates.jsonl"))
|
||||
});
|
||||
drop(rt);
|
||||
out
|
||||
};
|
||||
|
||||
let on_disk = std::fs::metadata(&updates_path).unwrap().len() as usize;
|
||||
// The ratio bound is only meaningful once fixed overhead is dwarfed by
|
||||
// content; guard against a shrunk spec making the assert trivial.
|
||||
assert!(
|
||||
on_disk > 256 * 1024,
|
||||
"fork-replay fixture must be sizeable to bound the peak ratio, got {on_disk} B"
|
||||
);
|
||||
|
||||
let base_old = begin_measure();
|
||||
let reference = reference_load_all(&updates_path);
|
||||
let old_peak = peak() - base_old;
|
||||
let ref_count = reference.len();
|
||||
let ref_serialized: Vec<String> = reference.iter().map(serialize).collect();
|
||||
drop(reference);
|
||||
|
||||
let mut stream_count = 0usize;
|
||||
let base_new = begin_measure();
|
||||
let outcome = stream_replay_updates_at(&id, root.path(), |_update| {
|
||||
stream_count += 1;
|
||||
})
|
||||
.expect("stream_replay_updates_at");
|
||||
let new_peak = peak() - base_new;
|
||||
|
||||
// Serializing during the measured pass would inflate its peak, so replay a
|
||||
// second time purely to collect the parity data.
|
||||
let mut streamed_serialized: Vec<String> = Vec::new();
|
||||
let _ = stream_replay_updates_at(&id, root.path(), |u| {
|
||||
streamed_serialized.push(serialize(&u))
|
||||
})
|
||||
.expect("stream again");
|
||||
|
||||
let new_ratio = new_peak as f64 / on_disk.max(1) as f64;
|
||||
let old_ratio = old_peak as f64 / on_disk.max(1) as f64;
|
||||
eprintln!(
|
||||
"FORK_REPLAY_MEMORY {}",
|
||||
serde_json::json!({
|
||||
"on_disk_mb": on_disk as f64 / 1e6,
|
||||
"old_parse_all_peak_mb": old_peak as f64 / 1e6,
|
||||
"old_ratio": old_ratio,
|
||||
"new_stream_peak_mb": new_peak as f64 / 1e6,
|
||||
"new_ratio": new_ratio,
|
||||
"reduction_x": old_peak as f64 / new_peak.max(1) as f64,
|
||||
"count": stream_count,
|
||||
})
|
||||
);
|
||||
|
||||
assert!(
|
||||
outcome == ReplayEmission::Emitted && stream_count > 0,
|
||||
"expected a non-empty replay"
|
||||
);
|
||||
assert_eq!(
|
||||
streamed_serialized, ref_serialized,
|
||||
"streamed updates must match the typed parse-all path byte-for-byte, in order"
|
||||
);
|
||||
assert_eq!(
|
||||
stream_count, ref_count,
|
||||
"streamed update count must equal the typed path count"
|
||||
);
|
||||
assert!(
|
||||
new_peak < old_peak,
|
||||
"streaming peak ({new_peak} B) must be below the parse-all peak ({old_peak} B)"
|
||||
);
|
||||
assert!(
|
||||
new_ratio < MAX_STREAM_PEAK_TO_DISK_RATIO,
|
||||
"streaming peak {new_ratio:.2}x on-disk must stay near the file size \
|
||||
(< {MAX_STREAM_PEAK_TO_DISK_RATIO}x); the whole transcript is no longer \
|
||||
materialized as typed structs"
|
||||
);
|
||||
}
|
||||
|
|
@ -1,10 +1,10 @@
|
|||
//! End-to-end measurement of why resuming a large session is slow — the time
|
||||
//! 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
|
||||
//! (2) `replay_session_updates`, which reads `updates.jsonl`, filters it, typed
|
||||
//! parses every line, and forwards 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
|
||||
|
|
@ -12,20 +12,23 @@
|
|||
//! 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.
|
||||
//! pipes (via [`load_session_via_agent`]); 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
|
||||
//! `GROK_PERF_SESSION_SRC=/path/to/<session-dir>`.
|
||||
//! Session data (both tests): a synthetic session from the shared
|
||||
//! [`synth`](xai_grok_shell::session::testkit::synth) generator (redundant
|
||||
//! `available_commands_update` + big rewind snapshots; size knobs via
|
||||
//! `GROK_PERF_*`), or a real session dir via `GROK_PERF_SESSION_SRC=<session-dir>`.
|
||||
//!
|
||||
//! Run:
|
||||
//! cargo test -p xai-grok-shell --test session_load_perf -- --nocapture
|
||||
//! cargo test -p xai-grok-shell --test session_load_perf full_session_load_e2e -- --ignored --nocapture
|
||||
//! Run (needs the `test-support` feature; on by default under Bazel):
|
||||
//! cargo test -p xai-grok-shell --features test-support --test session_load_perf -- --nocapture
|
||||
//! cargo test -p xai-grok-shell --features test-support --test session_load_perf full_session_load_e2e -- --ignored --nocapture
|
||||
|
||||
use std::cell::RefCell;
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::rc::Rc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use agent_client_protocol::{self as acp};
|
||||
|
|
@ -35,195 +38,30 @@ use xai_grok_shell::session::info::Info;
|
|||
use xai_grok_shell::session::storage::{
|
||||
JsonlStorageAdapter, StorageAdapter, load_updates_for_replay_at,
|
||||
};
|
||||
use xai_grok_workspace::session::file_state::{FileSnapshot, FlexiblePath, RewindPoint};
|
||||
use xai_grok_shell::session::testkit::e2e::load_session_via_agent;
|
||||
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
|
||||
|
||||
// ───────────────────────── size knobs ─────────────────────────
|
||||
// ───────────────────────── session spec ─────────────────────────
|
||||
|
||||
/// 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("GROK_PERF_SCALE", 1).max(1);
|
||||
Self {
|
||||
turns: g("GROK_PERF_TURNS", 80) * scale,
|
||||
acu_per_turn: g("GROK_PERF_ACU_PER_TURN", 15),
|
||||
catalog_commands: g("GROK_PERF_CATALOG_COMMANDS", 64),
|
||||
catalog_desc_len: g("GROK_PERF_CATALOG_DESC_LEN", 320),
|
||||
agent_chunks_per_turn: g("GROK_PERF_AGENT_CHUNKS_PER_TURN", 8),
|
||||
agent_chunk_len: g("GROK_PERF_AGENT_CHUNK_LEN", 2000),
|
||||
rewind_points: g("GROK_PERF_REWIND_POINTS", 60) * scale,
|
||||
files_per_rewind: g("GROK_PERF_FILES_PER_REWIND", 40),
|
||||
file_content_len: g("GROK_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");
|
||||
/// Perf-tool defaults over the shared [`SessionSpec`], tuned to the pathological
|
||||
/// real session; scale/override via `GROK_PERF_*` (e.g. `GROK_PERF_TURNS`,
|
||||
/// `GROK_PERF_SCALE`), or point `GROK_PERF_SESSION_SRC` at a real session dir.
|
||||
fn perf_spec() -> SessionSpec {
|
||||
SessionSpec::from_env_prefixed(
|
||||
"GROK_PERF",
|
||||
SessionSpec {
|
||||
turns: 80,
|
||||
rewind_points: 60,
|
||||
files_per_rewind: 40,
|
||||
file_content_len: 8000,
|
||||
..SessionSpec::default()
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ───────────────────────── 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.
|
||||
/// Recursively copy a directory tree (real-session overlay only).
|
||||
fn copy_tree(src: &Path, dst: &Path) {
|
||||
std::fs::create_dir_all(dst).unwrap();
|
||||
for entry in std::fs::read_dir(src).unwrap().flatten() {
|
||||
|
|
@ -237,66 +75,58 @@ fn copy_tree(src: &Path, dst: &Path) {
|
|||
}
|
||||
}
|
||||
|
||||
/// Prepare a session on disk under `root` for working dir `cwd`. Returns the
|
||||
/// `Info` and the session directory path. Uses `GROK_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) {
|
||||
/// Prepare a session on disk under `root` for working dir `cwd`. With
|
||||
/// `GROK_PERF_SESSION_SRC` set, copy a real session over a registered stub
|
||||
/// (keeping our `summary.json`); otherwise synthesize one via
|
||||
/// [`synth::prepare_session`].
|
||||
async fn prepare_session(root: &Path, cwd: &Path, spec: &SessionSpec) -> (Info, PathBuf) {
|
||||
let Ok(src) = std::env::var("GROK_PERF_SESSION_SRC") else {
|
||||
return synth::prepare_session(root, cwd, spec).await;
|
||||
};
|
||||
|
||||
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
|
||||
|
||||
if let Ok(src) = std::env::var("GROK_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),
|
||||
id: synth::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
|
||||
);
|
||||
let dir = synth::locate_session_dir(root, &id);
|
||||
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();
|
||||
}
|
||||
}
|
||||
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}");
|
||||
(info, dir)
|
||||
}
|
||||
|
||||
/// 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, spec: &SessionSpec) {
|
||||
if std::env::var("GROK_PERF_SESSION_SRC").is_ok() {
|
||||
return;
|
||||
}
|
||||
synth::write_rewind_jsonl(path, spec);
|
||||
}
|
||||
|
||||
// ───────────────────────── updates.jsonl stats ─────────────────────────
|
||||
|
||||
/// Per-kind statistics for the generated/loaded updates file.
|
||||
#[derive(Default)]
|
||||
struct KindStats {
|
||||
count: BTreeMap<String, u64>,
|
||||
bytes: BTreeMap<String, u64>,
|
||||
}
|
||||
|
||||
fn file_size_mb(path: &Path) -> f64 {
|
||||
std::fs::metadata(path).map(|m| m.len()).unwrap_or(0) as f64 / 1e6
|
||||
}
|
||||
|
|
@ -382,9 +212,9 @@ fn print_kind_breakdown(label: &str, stats: &KindStats) {
|
|||
async fn phase_breakdown_real_functions() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let opts = GenOpts::from_env();
|
||||
let spec = perf_spec();
|
||||
|
||||
let (info, dir) = prepare_session(root.path(), cwd.path(), &opts).await;
|
||||
let (info, dir) = prepare_session(root.path(), cwd.path(), &spec).await;
|
||||
|
||||
let updates_path = dir.join("updates.jsonl");
|
||||
let rewind_path = dir.join("rewind_points.jsonl");
|
||||
|
|
@ -397,7 +227,7 @@ async fn phase_breakdown_real_functions() {
|
|||
|
||||
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
|
||||
|
||||
// Phase A: load_light core (summary + chat_history) — what mvp_agent's
|
||||
// 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
|
||||
|
|
@ -405,12 +235,9 @@ async fn phase_breakdown_real_functions() {
|
|||
.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
|
||||
// Lazy rewind path: 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 xai_grok_workspace::session::file_state::FileStateTracker;
|
||||
|
|
@ -431,8 +258,8 @@ async fn phase_breakdown_real_functions() {
|
|||
"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).
|
||||
// Phase A': isolate rewind cost by deleting the rewind file and re-measuring.
|
||||
// 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
|
||||
|
|
@ -441,12 +268,14 @@ async fn phase_breakdown_real_functions() {
|
|||
.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);
|
||||
generate_or_restore_rewind(&rewind_path, &spec);
|
||||
|
||||
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.
|
||||
// Phase B: updates replay parse. The typed `load_updates_for_replay_at`
|
||||
// reads the whole file, typed-parses every line, and applies rewind
|
||||
// filtering; production now streams via `stream_replay_updates_at`, so this
|
||||
// measures the materialize-all parse cost.
|
||||
let t = Instant::now();
|
||||
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
|
||||
.expect("load_updates_for_replay_at")
|
||||
|
|
@ -458,7 +287,7 @@ async fn phase_breakdown_real_functions() {
|
|||
|
||||
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!(" rewind_points loaded in load : 0 (deferred → lazy)");
|
||||
eprintln!(" updates replayed (acp) : {}", replayed.len());
|
||||
eprintln!(" ----------------------------------------------------------------");
|
||||
eprintln!(
|
||||
|
|
@ -495,38 +324,15 @@ async fn phase_breakdown_real_functions() {
|
|||
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("GROK_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 tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
|
||||
use xai_acp_lib::{
|
||||
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
|
||||
LineBufferedRead,
|
||||
};
|
||||
use xai_grok_shell::agent::config::Config as AgentConfig;
|
||||
use xai_grok_shell::agent::mvp_agent::MvpAgent;
|
||||
|
||||
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.
|
||||
/// `available_commands_update` notifications forwarded during the load.
|
||||
/// History replay skips the (thousands of) historical ones, so this stays tiny.
|
||||
acu_count: u64,
|
||||
first_at: Option<Instant>,
|
||||
last_at: Option<Instant>,
|
||||
|
|
@ -616,7 +422,7 @@ async fn full_session_load_e2e() {
|
|||
|
||||
let grok_home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let opts = GenOpts::from_env();
|
||||
let spec = perf_spec();
|
||||
let instr_log = grok_home.path().join("instr.jsonl");
|
||||
|
||||
// SAFETY: single-threaded current-thread runtime; set before any agent code
|
||||
|
|
@ -641,7 +447,7 @@ async fn full_session_load_e2e() {
|
|||
.with(xai_grok_shell::instrumentation::layer::<Registry>())
|
||||
.try_init();
|
||||
|
||||
let (info, dir) = prepare_session(grok_home.path(), cwd.path(), &opts).await;
|
||||
let (info, dir) = prepare_session(grok_home.path(), cwd.path(), &spec).await;
|
||||
let updates_path = dir.join("updates.jsonl");
|
||||
let rewind_path = dir.join("rewind_points.jsonl");
|
||||
eprintln!(
|
||||
|
|
@ -652,85 +458,33 @@ async fn full_session_load_e2e() {
|
|||
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.
|
||||
// Zero-data-loss guard: a pure load must never rewrite rewind_points.jsonl
|
||||
// (it is read 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(xai_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())),
|
||||
let loaded = load_session_via_agent(
|
||||
client,
|
||||
"perf-test",
|
||||
info.id.clone(),
|
||||
cwd.path().to_path_buf(),
|
||||
)
|
||||
.await
|
||||
.expect("initialize timed out")
|
||||
.expect("initialize failed");
|
||||
.await;
|
||||
let load_started = loaded.load_started;
|
||||
let load_elapsed = loaded.load_elapsed;
|
||||
// Keep the connection alive so the post-load re-advertise still arrives.
|
||||
let _client_conn = loaded.client_conn;
|
||||
|
||||
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).
|
||||
// Snapshot replay results immediately, before the post-load
|
||||
// AdvertiseCommands re-advertise can arrive, so `acu_replayed` counts
|
||||
// the ACUs forwarded during history replay (the skip count).
|
||||
let (replay_count, acu_replayed, ttfn, ttln) = {
|
||||
let c = counters.borrow();
|
||||
(
|
||||
|
|
@ -763,8 +517,8 @@ async fn full_session_load_e2e() {
|
|||
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
|
||||
// Replay-skip guard: the historical available_commands_update copies
|
||||
// (3197 in the pathological real session, hundreds synthetic) must
|
||||
// NOT be replayed.
|
||||
let acu_persisted = stats.count.get("available_commands_update").copied().unwrap_or(0);
|
||||
|
||||
|
|
@ -786,13 +540,13 @@ async fn full_session_load_e2e() {
|
|||
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.
|
||||
// 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)...
|
||||
// The thousands of persisted ACUs must be skipped on replay...
|
||||
assert!(
|
||||
acu_persisted > 100,
|
||||
"fixture should have many persisted ACUs to exercise the skip"
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ use xai_grok_shell::leader::{
|
|||
ClientCapabilities, ClientMode, LeaderClient, LeaderServerControlState, LeaderServerMetadata,
|
||||
run_leader_server,
|
||||
};
|
||||
use xai_grok_test_support::resources::ResourceSnapshot;
|
||||
|
||||
const SIMPLEX_BUF: usize = 8 * 1024 * 1024;
|
||||
|
||||
|
|
@ -42,41 +43,6 @@ fn env_u64(key: &str, default: u64) -> u64 {
|
|||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
/// Resident set size of THIS process (leader server + agent are in-process).
|
||||
/// Copied from `xai-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) = xai_grok_telemetry::unified_log::snapshot_log() else {
|
||||
|
|
@ -281,7 +247,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
|
|||
)
|
||||
.await;
|
||||
|
||||
let rss_baseline = rss_bytes();
|
||||
let rss_before = ResourceSnapshot::capture();
|
||||
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;
|
||||
|
|
@ -379,8 +345,12 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
|
|||
);
|
||||
|
||||
// ── 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);
|
||||
let rss_after = ResourceSnapshot::capture();
|
||||
let growth = rss_after.growth_from(&rss_before);
|
||||
if let (Some(before), Some(after), Some(growth_bytes)) =
|
||||
(rss_before.rss, rss_after.rss, growth.rss)
|
||||
{
|
||||
let growth_mb = growth_bytes 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),
|
||||
|
|
|
|||
183
crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs
Normal file
183
crates/codegen/xai-grok-shell/tests/test_nonblocking_startup.rs
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
//! Non-blocking startup regression tests. `#[ignore]`; requires pre-built binary.
|
||||
//!
|
||||
//! These exercise leader startup through the persistent-leader fixture: the
|
||||
//! leader must bind its socket and become ready without blocking on the remote
|
||||
//! `/settings` + `/v1/models` fetch, and must self-heal its catalog once the
|
||||
//! endpoint recovers.
|
||||
//!
|
||||
//! ```bash
|
||||
//! cargo test -p xai-grok-shell --test test_nonblocking_startup -- --ignored
|
||||
//! ```
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
mod common;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use xai_grok_test_support::leader::LeaderFixture;
|
||||
use xai_grok_test_support::*;
|
||||
|
||||
async fn poll_until(ceiling: Duration, interval: Duration, condition: impl Fn() -> bool) -> bool {
|
||||
let deadline = Instant::now() + ceiling;
|
||||
while Instant::now() < deadline {
|
||||
if condition() {
|
||||
return true;
|
||||
}
|
||||
tokio::time::sleep(interval).await;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// A hanging proxy must not delay leader readiness: the fixture (which waits for
|
||||
/// the leader socket) must come up and a session must be created well within the
|
||||
/// blocking-fetch window.
|
||||
#[tokio::test]
|
||||
#[ignore] // requires pre-built binary; run with --ignored
|
||||
async fn leader_ready_while_proxy_hangs() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
let server = MockInferenceServer::start().await.unwrap();
|
||||
server.set_hang(true);
|
||||
|
||||
let workdir = git_workdir();
|
||||
let sandbox = TestSandbox::new();
|
||||
|
||||
let started = Instant::now();
|
||||
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("leader must become ready while the proxy hangs");
|
||||
let mut clients = Vec::new();
|
||||
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
|
||||
Box::pin(async move {
|
||||
clients.push(
|
||||
fixture
|
||||
.spawn_client(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("spawn leader client"),
|
||||
);
|
||||
clients[0].initialize().await;
|
||||
let _session = clients[0].create_session(workdir.workspace()).await;
|
||||
|
||||
let elapsed = started.elapsed();
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(25),
|
||||
"startup took {elapsed:?} with a hanging proxy; readiness appears \
|
||||
to block on the network fetch\nstderr:\n{}",
|
||||
clients[0].stderr_text(),
|
||||
);
|
||||
})
|
||||
})
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// The background catalog refresh must re-fetch and push `x.ai/models/update`
|
||||
/// once a previously-hanging endpoint recovers.
|
||||
#[tokio::test]
|
||||
#[ignore] // requires pre-built binary; run with --ignored
|
||||
async fn catalog_self_heals_after_endpoint_recovers() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
let server = MockInferenceServer::start().await.unwrap();
|
||||
server.set_hang(true);
|
||||
|
||||
let workdir = git_workdir();
|
||||
let sandbox = TestSandbox::new();
|
||||
|
||||
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("leader must become ready while the proxy hangs");
|
||||
let mut clients = Vec::new();
|
||||
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
|
||||
Box::pin(async move {
|
||||
clients.push(
|
||||
fixture
|
||||
.spawn_client(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("spawn leader client"),
|
||||
);
|
||||
clients[0].initialize().await;
|
||||
|
||||
// Recover the endpoint. The background catalog refresh
|
||||
// (5s-base backoff) re-fetches and pushes `x.ai/models/update`.
|
||||
server.set_hang(false);
|
||||
|
||||
let healed =
|
||||
poll_until(Duration::from_secs(60), Duration::from_millis(500), || {
|
||||
clients[0].models_update_count() > 0
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
healed,
|
||||
"no x.ai/models/update after the endpoint recovered\nstderr:\n{}",
|
||||
clients[0].stderr_text(),
|
||||
);
|
||||
})
|
||||
})
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Custom-backend reality: a user points at their own backend that serves
|
||||
/// `/v1/models` + chat but blocks the cli-chat-proxy `/settings` (404). The
|
||||
/// leader must boot fast, load its catalog from the served models, and run a
|
||||
/// real prompt, even though remote settings never arrive.
|
||||
#[tokio::test]
|
||||
#[ignore] // requires pre-built binary; run with --ignored
|
||||
async fn leader_usable_when_settings_blocked_but_models_served() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
// Models + chat are served; `set_settings` is never called, so
|
||||
// `/v1/settings` 404s, mimicking a blocked proxy settings endpoint.
|
||||
let server = MockInferenceServer::start().await.unwrap();
|
||||
|
||||
let workdir = git_workdir();
|
||||
let sandbox = TestSandbox::new();
|
||||
|
||||
let started = Instant::now();
|
||||
let fixture = LeaderFixture::start(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("leader must become ready with /settings blocked");
|
||||
let mut clients = Vec::new();
|
||||
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
|
||||
Box::pin(async move {
|
||||
clients.push(
|
||||
fixture
|
||||
.spawn_client(&server, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.expect("spawn leader client"),
|
||||
);
|
||||
clients[0].initialize().await;
|
||||
let session = clients[0].create_session(workdir.workspace()).await;
|
||||
|
||||
// The custom backend serves chat, so a prompt round-trips
|
||||
// despite the blocked settings endpoint.
|
||||
clients[0]
|
||||
.prompt(&session, "ping")
|
||||
.await
|
||||
.expect("prompt against the custom backend");
|
||||
assert!(
|
||||
server.has_chat_completion_request() || server.has_responses_request(),
|
||||
"the prompt must reach the served chat endpoint\nstderr:\n{}",
|
||||
clients[0].stderr_text(),
|
||||
);
|
||||
|
||||
assert!(
|
||||
started.elapsed() < Duration::from_secs(25),
|
||||
"startup/usage blocked on the unreachable /settings\nstderr:\n{}",
|
||||
clients[0].stderr_text(),
|
||||
);
|
||||
assert_eq!(
|
||||
clients[0].settings_update_count(),
|
||||
0,
|
||||
"no settings update should land while /settings is blocked",
|
||||
);
|
||||
})
|
||||
})
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
//! Leader boots from local data when the endpoint is fully unreachable
|
||||
//! (connection refused), not merely hanging. `#[ignore]`: needs the built binary.
|
||||
//!
|
||||
//! ```bash
|
||||
//! cargo test -p xai-grok-shell --test test_nonblocking_startup_offline -- --ignored
|
||||
//! ```
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
mod common;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use xai_grok_test_support::leader::LeaderFixture;
|
||||
use xai_grok_test_support::*;
|
||||
|
||||
/// A loopback URL on a closed port (bind, read addr, drop); refuses instantly.
|
||||
fn closed_port_base_url() -> String {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
|
||||
let addr = listener.local_addr().expect("local_addr");
|
||||
drop(listener);
|
||||
format!("http://{addr}/v1")
|
||||
}
|
||||
|
||||
async fn assert_boots_fast(base_url: String, scenario: &'static str) {
|
||||
let workdir = git_workdir();
|
||||
let sandbox = TestSandbox::new();
|
||||
|
||||
let started = Instant::now();
|
||||
let fixture = LeaderFixture::start_with_base_url(&base_url, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.unwrap_or_else(|error| {
|
||||
panic!("[{scenario}] leader never became ready with an unreachable endpoint: {error}")
|
||||
});
|
||||
let mut clients = Vec::new();
|
||||
common::leader::run_with_cleanup(&fixture, &mut clients, |fixture, clients| {
|
||||
Box::pin(async move {
|
||||
clients.push(
|
||||
fixture
|
||||
.spawn_client_with_base_url(&base_url, workdir.workspace(), &sandbox)
|
||||
.await
|
||||
.unwrap_or_else(|error| panic!("[{scenario}] spawn leader client: {error}")),
|
||||
);
|
||||
clients[0].initialize().await;
|
||||
// Catalog resolves offline (built-in/cache), so session creation succeeds.
|
||||
let _session = clients[0].create_session(workdir.workspace()).await;
|
||||
|
||||
let elapsed = started.elapsed();
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(25),
|
||||
"[{scenario}] startup took {elapsed:?} with an unreachable endpoint; \
|
||||
readiness appears to block on the network fetch\nstderr:\n{}",
|
||||
clients[0].stderr_text(),
|
||||
);
|
||||
})
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // needs the built binary
|
||||
async fn leader_ready_with_connection_refused() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
assert_boots_fast(closed_port_base_url(), "connection-refused").await;
|
||||
})
|
||||
.await;
|
||||
}
|
||||
730
crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs
Normal file
730
crates/codegen/xai-grok-shell/tests/test_session_load_memory.rs
Normal file
|
|
@ -0,0 +1,730 @@
|
|||
//! Memory tests for the session-load path: prove the resume peek borrows the
|
||||
//! transcript instead of copying, and bound peak memory with full cleanup.
|
||||
//! Resuming a large session once OOM-killed the process under a cgroup cap.
|
||||
//!
|
||||
//! Run:
|
||||
//! cargo test -p xai-grok-shell --features dhat-heap --test test_session_load_memory \
|
||||
//! session_load_dhat_bounded_and_freed -- --ignored --nocapture
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
#[global_allocator]
|
||||
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;
|
||||
|
||||
use xai_grok_shell::session::storage::{JsonlStorageAdapter, StorageAdapter, prepare_replay_lines};
|
||||
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
use std::path::Path;
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
use xai_grok_shell::session::info::Info;
|
||||
|
||||
use tempfile::TempDir;
|
||||
|
||||
const BYTES_PER_MB: f64 = 1024.0 * 1024.0;
|
||||
const BYTES_PER_MB_U64: u64 = 1024 * 1024;
|
||||
|
||||
fn file_len(path: &std::path::Path) -> u64 {
|
||||
// Fail loud: a silent 0 would collapse the ratio budget instead of
|
||||
// surfacing a missing or unreadable fixture.
|
||||
std::fs::metadata(path).expect("stat updates.jsonl").len()
|
||||
}
|
||||
|
||||
fn env_parse<T: std::str::FromStr>(key: &str, default: T) -> T {
|
||||
let Ok(text) = std::env::var(key) else {
|
||||
return default;
|
||||
};
|
||||
match text.parse() {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
eprintln!("[test_session_load] ignoring unparseable {key}={text:?}; using default");
|
||||
default
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn memory_spec() -> SessionSpec {
|
||||
SessionSpec::from_env_prefixed("SESSION_LOAD", SessionSpec::default())
|
||||
}
|
||||
|
||||
// Replay keeps one user chunk plus the agent chunks per turn and drops the
|
||||
// redundant ACUs, mirroring `synth::prepare_session` and `prepare_replay_lines`.
|
||||
fn expected_replayed_lines(spec: &SessionSpec) -> usize {
|
||||
spec.turns * (1 + spec.agent_chunks_per_turn)
|
||||
}
|
||||
|
||||
/// Non-ignored zero-copy guard: every replay line must borrow from the
|
||||
/// transcript, so an owned-copy regression fails here in CI.
|
||||
#[tokio::test]
|
||||
async fn prepare_replay_lines_borrows_the_transcript() {
|
||||
let spec = SessionSpec {
|
||||
turns: 3,
|
||||
acu_per_turn: 2,
|
||||
catalog_commands: 2,
|
||||
catalog_desc_len: 8,
|
||||
agent_chunks_per_turn: 2,
|
||||
agent_chunk_len: 32,
|
||||
rewind_points: 0,
|
||||
files_per_rewind: 0,
|
||||
file_content_len: 0,
|
||||
};
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let (_info, dir) = synth::prepare_session(root.path(), cwd.path(), &spec).await;
|
||||
let transcript =
|
||||
std::fs::read_to_string(dir.join("updates.jsonl")).expect("read updates.jsonl");
|
||||
|
||||
let prepared = prepare_replay_lines(&transcript, None);
|
||||
assert_eq!(
|
||||
prepared.lines.len(),
|
||||
expected_replayed_lines(&spec),
|
||||
"replay line count regressed"
|
||||
);
|
||||
|
||||
let start = transcript.as_ptr() as usize;
|
||||
let end = start + transcript.len();
|
||||
for line in &prepared.lines {
|
||||
let line_start = line.as_ptr() as usize;
|
||||
assert!(
|
||||
line_start >= start && line_start + line.len() <= end,
|
||||
"replay line must borrow from the transcript (zero-copy), not own a copy"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Let ready tasks drain and timer-driven cleanup run before reading heap
|
||||
/// stats, so a drop's frees show up in the next `curr_bytes`/`curr_blocks`.
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
async fn quiesce() {
|
||||
const YIELDS: usize = 50;
|
||||
const SETTLE: std::time::Duration = std::time::Duration::from_millis(10);
|
||||
for _ in 0..YIELDS {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
tokio::time::sleep(SETTLE).await;
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
async fn run_load_cycle(adapter: &JsonlStorageAdapter, info: &Info, updates_path: &Path) -> usize {
|
||||
let light = adapter
|
||||
.load_session_without_updates(info)
|
||||
.await
|
||||
.expect("load_session_without_updates");
|
||||
let transcript = std::fs::read_to_string(updates_path).expect("read updates.jsonl");
|
||||
let prepared = prepare_replay_lines(&transcript, None);
|
||||
let replayed_lines = prepared.lines.len();
|
||||
drop(prepared);
|
||||
drop(transcript);
|
||||
drop(light);
|
||||
quiesce().await;
|
||||
replayed_lines
|
||||
}
|
||||
|
||||
/// Env-derived gates and cycle counts, separated from the measured results. The
|
||||
/// caller clamps `cycles` to at least one so the per-cycle divisions are safe.
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
struct DhatBudget {
|
||||
warmup: usize,
|
||||
cycles: usize,
|
||||
ratio: f64,
|
||||
abs_budget_mb: u64,
|
||||
max_bytes_per_cycle: i64,
|
||||
max_blocks_per_cycle: i64,
|
||||
}
|
||||
|
||||
/// Heap readings captured across the measured window, named so nothing can
|
||||
/// silently transpose the same-typed counts.
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
struct DhatMeasured {
|
||||
replayed_lines: usize,
|
||||
expected_lines: usize,
|
||||
on_disk_bytes: u64,
|
||||
peak_over_baseline: u64,
|
||||
net_bytes: i64,
|
||||
net_blocks: i64,
|
||||
}
|
||||
|
||||
/// The measured window paired with the budget it is judged against; every gate
|
||||
/// threshold derives from `budget`, so nothing is stored twice.
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
struct DhatOutcome<'a> {
|
||||
budget: &'a DhatBudget,
|
||||
measured: DhatMeasured,
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
impl DhatOutcome<'_> {
|
||||
fn ratio_budget_bytes(&self) -> u64 {
|
||||
(self.budget.ratio * self.measured.on_disk_bytes as f64) as u64
|
||||
}
|
||||
|
||||
fn abs_budget_bytes(&self) -> u64 {
|
||||
self.budget.abs_budget_mb * BYTES_PER_MB_U64
|
||||
}
|
||||
|
||||
fn per_cycle_bytes(&self) -> i64 {
|
||||
self.measured.net_bytes / self.budget.cycles as i64
|
||||
}
|
||||
|
||||
fn per_cycle_blocks(&self) -> i64 {
|
||||
self.measured.net_blocks / self.budget.cycles as i64
|
||||
}
|
||||
|
||||
fn no_spike(&self) -> bool {
|
||||
self.measured.peak_over_baseline < self.ratio_budget_bytes()
|
||||
&& self.measured.peak_over_baseline < self.abs_budget_bytes()
|
||||
}
|
||||
|
||||
fn cleaned_up(&self) -> bool {
|
||||
self.per_cycle_bytes() < self.budget.max_bytes_per_cycle
|
||||
&& self.per_cycle_blocks() < self.budget.max_blocks_per_cycle
|
||||
}
|
||||
|
||||
fn pass(&self) -> bool {
|
||||
self.no_spike()
|
||||
&& self.cleaned_up()
|
||||
&& self.measured.replayed_lines == self.measured.expected_lines
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
#[test]
|
||||
fn dhat_outcome_verdict_arithmetic() {
|
||||
let budget = DhatBudget {
|
||||
warmup: 0,
|
||||
cycles: 4,
|
||||
ratio: 2.0,
|
||||
abs_budget_mb: 1,
|
||||
max_bytes_per_cycle: 100,
|
||||
max_blocks_per_cycle: 10,
|
||||
};
|
||||
|
||||
let ok = DhatOutcome {
|
||||
budget: &budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines: 5,
|
||||
expected_lines: 5,
|
||||
on_disk_bytes: 1024,
|
||||
peak_over_baseline: 1000,
|
||||
net_bytes: 40,
|
||||
net_blocks: 4,
|
||||
},
|
||||
};
|
||||
assert_eq!(ok.per_cycle_bytes(), 10);
|
||||
assert_eq!(ok.per_cycle_blocks(), 1);
|
||||
assert!(ok.no_spike() && ok.cleaned_up() && ok.pass());
|
||||
|
||||
// Peak over the ratio budget (2x the 1024-byte file) trips no_spike.
|
||||
let ratio_spike = DhatOutcome {
|
||||
budget: &budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines: 5,
|
||||
expected_lines: 5,
|
||||
on_disk_bytes: 1024,
|
||||
peak_over_baseline: 4096,
|
||||
net_bytes: 0,
|
||||
net_blocks: 0,
|
||||
},
|
||||
};
|
||||
assert!(!ratio_spike.no_spike() && !ratio_spike.pass());
|
||||
|
||||
// Per-cycle residual over the gate trips cleaned_up.
|
||||
let leak = DhatOutcome {
|
||||
budget: &budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines: 5,
|
||||
expected_lines: 5,
|
||||
on_disk_bytes: 1024,
|
||||
peak_over_baseline: 1000,
|
||||
net_bytes: 4 * 200,
|
||||
net_blocks: 4 * 20,
|
||||
},
|
||||
};
|
||||
assert_eq!(leak.per_cycle_bytes(), 200);
|
||||
assert!(!leak.cleaned_up() && !leak.pass());
|
||||
|
||||
// Clean gates but a mismatched replay count still fails pass.
|
||||
let miscount = DhatOutcome {
|
||||
budget: &budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines: 4,
|
||||
expected_lines: 5,
|
||||
on_disk_bytes: 1024,
|
||||
peak_over_baseline: 1000,
|
||||
net_bytes: 40,
|
||||
net_blocks: 4,
|
||||
},
|
||||
};
|
||||
assert!(miscount.no_spike() && miscount.cleaned_up() && !miscount.pass());
|
||||
|
||||
// A peak under the ratio budget but over the absolute budget trips no_spike
|
||||
// via its other arm.
|
||||
let abs_budget = DhatBudget {
|
||||
ratio: 1.0,
|
||||
..budget
|
||||
};
|
||||
let abs_spike = DhatOutcome {
|
||||
budget: &abs_budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines: 5,
|
||||
expected_lines: 5,
|
||||
on_disk_bytes: 2 * 1024 * 1024,
|
||||
peak_over_baseline: 1024 * 1024 + 1,
|
||||
net_bytes: 0,
|
||||
net_blocks: 0,
|
||||
},
|
||||
};
|
||||
assert!(!abs_spike.no_spike() && !abs_spike.pass());
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
fn report_summary(o: &DhatOutcome<'_>) {
|
||||
eprintln!(
|
||||
"SESSION_LOAD_DHAT_SUMMARY {}",
|
||||
serde_json::json!({
|
||||
"mode": "dhat-heap",
|
||||
"cycles": o.budget.cycles,
|
||||
"warmup": o.budget.warmup,
|
||||
"replayed_lines": o.measured.replayed_lines,
|
||||
"expected_lines": o.measured.expected_lines,
|
||||
"on_disk_updates_bytes": o.measured.on_disk_bytes,
|
||||
"on_disk_updates_mb": o.measured.on_disk_bytes as f64 / BYTES_PER_MB,
|
||||
"peak_over_baseline_bytes": o.measured.peak_over_baseline,
|
||||
"peak_over_baseline_mb": o.measured.peak_over_baseline as f64 / BYTES_PER_MB,
|
||||
"peak_over_on_disk": o.measured.peak_over_baseline as f64 / o.measured.on_disk_bytes.max(1) as f64,
|
||||
"ratio_budget": o.budget.ratio,
|
||||
"ratio_budget_mb": o.ratio_budget_bytes() as f64 / BYTES_PER_MB,
|
||||
"abs_budget_mb": o.budget.abs_budget_mb,
|
||||
"no_spike": o.no_spike(),
|
||||
"per_cycle_residual_bytes": o.per_cycle_bytes(),
|
||||
"per_cycle_residual_blocks": o.per_cycle_blocks(),
|
||||
"max_bytes_per_cycle": o.budget.max_bytes_per_cycle,
|
||||
"max_blocks_per_cycle": o.budget.max_blocks_per_cycle,
|
||||
"net_window_bytes": o.measured.net_bytes,
|
||||
"net_window_blocks": o.measured.net_blocks,
|
||||
"cleaned_up": o.cleaned_up(),
|
||||
"pass": o.pass(),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
fn assert_bounds(o: &DhatOutcome<'_>) {
|
||||
assert_eq!(
|
||||
o.measured.replayed_lines, o.measured.expected_lines,
|
||||
"replayed line count must equal the non-ACU update count"
|
||||
);
|
||||
|
||||
assert!(
|
||||
o.no_spike(),
|
||||
"load peak {:.1} MB over baseline is {:.2}x the {:.1} MB on-disk updates and exceeds a gate \
|
||||
(RATIO {}x = {:.1} MB, ABSOLUTE {} MB); load memory is super-linear in session size",
|
||||
o.measured.peak_over_baseline as f64 / BYTES_PER_MB,
|
||||
o.measured.peak_over_baseline as f64 / o.measured.on_disk_bytes.max(1) as f64,
|
||||
o.measured.on_disk_bytes as f64 / BYTES_PER_MB,
|
||||
o.budget.ratio,
|
||||
o.ratio_budget_bytes() as f64 / BYTES_PER_MB,
|
||||
o.budget.abs_budget_mb,
|
||||
);
|
||||
|
||||
assert!(
|
||||
o.cleaned_up(),
|
||||
"leak: {} bytes/cycle and {} blocks/cycle retained over {} load/drop cycles \
|
||||
({} net bytes, {} net blocks) exceed the {}-byte / {}-block gate; load does not free \
|
||||
everything",
|
||||
o.per_cycle_bytes(),
|
||||
o.per_cycle_blocks(),
|
||||
o.budget.cycles,
|
||||
o.measured.net_bytes,
|
||||
o.measured.net_blocks,
|
||||
o.budget.max_bytes_per_cycle,
|
||||
o.budget.max_blocks_per_cycle,
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[ignore = "memory soak; run with --features dhat-heap --ignored --nocapture"]
|
||||
async fn session_load_dhat_bounded_and_freed() {
|
||||
let opts = memory_spec();
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
|
||||
let updates_path = dir.join("updates.jsonl");
|
||||
let on_disk_bytes = file_len(&updates_path);
|
||||
let expected_lines = expected_replayed_lines(&opts);
|
||||
|
||||
let budget = DhatBudget {
|
||||
warmup: env_parse("SESSION_LOAD_WARMUP", 3usize),
|
||||
cycles: env_parse("SESSION_LOAD_CYCLES", 8usize).max(1),
|
||||
ratio: env_parse("SESSION_LOAD_HEAP_RATIO", 4.0),
|
||||
abs_budget_mb: env_parse("SESSION_LOAD_MAX_PEAK_HEAP_MB", 512u64),
|
||||
max_bytes_per_cycle: env_parse("SESSION_LOAD_MAX_BYTES_PER_CYCLE", 1i64 << 20),
|
||||
max_blocks_per_cycle: env_parse("SESSION_LOAD_MAX_BLOCKS_PER_CYCLE", 128i64),
|
||||
};
|
||||
|
||||
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
|
||||
|
||||
let profiler = dhat::Profiler::builder().testing().build();
|
||||
|
||||
for _ in 0..budget.warmup {
|
||||
let _ = run_load_cycle(&adapter, &info, &updates_path).await;
|
||||
}
|
||||
|
||||
let window_before = dhat::HeapStats::get();
|
||||
let mut replayed_lines = 0usize;
|
||||
for _ in 0..budget.cycles {
|
||||
replayed_lines = run_load_cycle(&adapter, &info, &updates_path).await;
|
||||
}
|
||||
let window_after = dhat::HeapStats::get();
|
||||
drop(profiler);
|
||||
|
||||
// `max_bytes` is a running maximum over the profiler's whole life (warmup
|
||||
// included), so subtracting the post-warmup baseline yields a conservative
|
||||
// upper bound on the load peak, never an underestimate.
|
||||
let peak_over_baseline =
|
||||
(window_after.max_bytes as u64).saturating_sub(window_before.curr_bytes as u64);
|
||||
|
||||
// Net change across the measured window; goes negative if a cycle frees more
|
||||
// than warmup left resident, which still satisfies the leak gate.
|
||||
let net_bytes = window_after.curr_bytes as i64 - window_before.curr_bytes as i64;
|
||||
let net_blocks = window_after.curr_blocks as i64 - window_before.curr_blocks as i64;
|
||||
|
||||
let outcome = DhatOutcome {
|
||||
budget: &budget,
|
||||
measured: DhatMeasured {
|
||||
replayed_lines,
|
||||
expected_lines,
|
||||
on_disk_bytes,
|
||||
peak_over_baseline,
|
||||
net_bytes,
|
||||
net_blocks,
|
||||
},
|
||||
};
|
||||
|
||||
report_summary(&outcome);
|
||||
assert_bounds(&outcome);
|
||||
}
|
||||
|
||||
// dhat replaces the global allocator and perturbs RSS, so the RSS-based forms
|
||||
// only compile without the `dhat-heap` feature.
|
||||
#[cfg(not(feature = "dhat-heap"))]
|
||||
mod rss {
|
||||
use super::*;
|
||||
|
||||
use std::cell::RefCell;
|
||||
use std::path::PathBuf;
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use agent_client_protocol::{self as acp};
|
||||
|
||||
use xai_grok_test_support::resources::ResourceSnapshot;
|
||||
|
||||
const SAMPLE_INTERVAL_MS: u64 = 3;
|
||||
|
||||
/// Poll RSS from a thread; the load path is synchronous, so this is how its
|
||||
/// peak gets captured.
|
||||
fn spawn_rss_sampler(stop: Arc<AtomicBool>) -> std::thread::JoinHandle<usize> {
|
||||
std::thread::spawn(move || {
|
||||
let mut peak = ResourceSnapshot::capture_rss().unwrap_or(0);
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
if let Some(r) = ResourceSnapshot::capture_rss() {
|
||||
peak = peak.max(r);
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(SAMPLE_INTERVAL_MS));
|
||||
}
|
||||
if let Some(r) = ResourceSnapshot::capture_rss() {
|
||||
peak = peak.max(r);
|
||||
}
|
||||
peak
|
||||
})
|
||||
}
|
||||
|
||||
/// Owns the RSS baseline and the background sampler for one measured load.
|
||||
struct RssSampler {
|
||||
baseline: Option<usize>,
|
||||
stop: Arc<AtomicBool>,
|
||||
handle: std::thread::JoinHandle<usize>,
|
||||
}
|
||||
|
||||
impl RssSampler {
|
||||
/// Capture the RSS baseline and start a background sampler.
|
||||
fn start() -> Self {
|
||||
let baseline = ResourceSnapshot::capture_rss();
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let handle = spawn_rss_sampler(stop.clone());
|
||||
Self {
|
||||
baseline,
|
||||
stop,
|
||||
handle,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stop sampling and build the outcome. Takes a final synchronous read
|
||||
/// first as a best-effort backstop: when the caller still holds the
|
||||
/// loaded state it pins a load that peaked and freed between the
|
||||
/// sampler's ticks; otherwise the background sampler is the sole source.
|
||||
fn finish(self, budget_mb: u64) -> RssOutcome {
|
||||
let final_rss = ResourceSnapshot::capture_rss().unwrap_or(0);
|
||||
self.stop.store(true, Ordering::Relaxed);
|
||||
let peak_rss = self.handle.join().expect("sampler thread").max(final_rss);
|
||||
RssOutcome {
|
||||
baseline: self.baseline,
|
||||
peak_rss,
|
||||
budget_mb,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct RssOutcome {
|
||||
baseline: Option<usize>,
|
||||
peak_rss: usize,
|
||||
budget_mb: u64,
|
||||
}
|
||||
|
||||
impl RssOutcome {
|
||||
fn baseline_bytes(&self) -> usize {
|
||||
self.baseline.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn measurable(&self) -> bool {
|
||||
self.baseline.is_some() && self.peak_rss > 0
|
||||
}
|
||||
|
||||
fn peak_growth_bytes(&self) -> usize {
|
||||
self.peak_rss.saturating_sub(self.baseline_bytes())
|
||||
}
|
||||
|
||||
fn budget_bytes(&self) -> u64 {
|
||||
self.budget_mb * BYTES_PER_MB_U64
|
||||
}
|
||||
|
||||
fn within_budget(&self) -> bool {
|
||||
(self.peak_growth_bytes() as u64) < self.budget_bytes()
|
||||
}
|
||||
|
||||
fn pass(&self) -> bool {
|
||||
!self.measurable() || self.within_budget()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rss_outcome_verdict_arithmetic() {
|
||||
let mb = BYTES_PER_MB_U64 as usize;
|
||||
|
||||
let over = RssOutcome {
|
||||
baseline: Some(mb),
|
||||
peak_rss: mb + 3 * mb,
|
||||
budget_mb: 2,
|
||||
};
|
||||
assert_eq!(over.peak_growth_bytes(), 3 * mb);
|
||||
assert!(!over.within_budget());
|
||||
assert!(!over.pass());
|
||||
|
||||
let under = RssOutcome {
|
||||
baseline: Some(mb),
|
||||
peak_rss: mb + mb,
|
||||
budget_mb: 2,
|
||||
};
|
||||
assert!(under.within_budget());
|
||||
assert!(under.pass());
|
||||
|
||||
// An unmeasurable baseline passes vacuously.
|
||||
let unmeasurable = RssOutcome {
|
||||
baseline: None,
|
||||
peak_rss: 0,
|
||||
budget_mb: 1,
|
||||
};
|
||||
assert!(!unmeasurable.measurable());
|
||||
assert!(unmeasurable.pass());
|
||||
}
|
||||
|
||||
fn report_summary(mode: &str, counts: serde_json::Value, on_disk_bytes: u64, o: &RssOutcome) {
|
||||
let mut summary = serde_json::json!({
|
||||
"mode": mode,
|
||||
"on_disk_updates_bytes": on_disk_bytes,
|
||||
"on_disk_updates_mb": on_disk_bytes as f64 / BYTES_PER_MB,
|
||||
"baseline_rss_mb": o.baseline_bytes() as f64 / BYTES_PER_MB,
|
||||
"peak_rss_mb": o.peak_rss as f64 / BYTES_PER_MB,
|
||||
"peak_rss_growth_mb": o.peak_growth_bytes() as f64 / BYTES_PER_MB,
|
||||
"budget_mb": o.budget_mb,
|
||||
"rss_measurable": o.measurable(),
|
||||
"pass": o.pass(),
|
||||
});
|
||||
let obj = summary
|
||||
.as_object_mut()
|
||||
.expect("summary literal is a JSON object");
|
||||
let extra = counts.as_object().expect("counts must be a JSON object");
|
||||
for (k, v) in extra {
|
||||
obj.insert(k.clone(), v.clone());
|
||||
}
|
||||
eprintln!("SESSION_LOAD_MEMORY_SUMMARY {summary}");
|
||||
}
|
||||
|
||||
fn assert_bounds(label: Option<&str>, on_disk_bytes: u64, o: &RssOutcome) {
|
||||
if o.measurable() {
|
||||
let prefix = label.map(|l| format!("{l} ")).unwrap_or_default();
|
||||
assert!(
|
||||
o.within_budget(),
|
||||
"{prefix}peak RSS grew {:.1} MB over baseline while loading a {:.1} MB updates file \
|
||||
(bound {} MB)",
|
||||
o.peak_growth_bytes() as f64 / BYTES_PER_MB,
|
||||
on_disk_bytes as f64 / BYTES_PER_MB,
|
||||
o.budget_mb,
|
||||
);
|
||||
} else {
|
||||
eprintln!("[soak] RSS measurement unavailable on this platform; bound skipped");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[ignore = "peak-memory soak; run with --ignored --nocapture"]
|
||||
async fn session_load_peak_rss_under_budget() {
|
||||
let opts = memory_spec();
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let (info, dir) = synth::prepare_session(root.path(), cwd.path(), &opts).await;
|
||||
let updates_path = dir.join("updates.jsonl");
|
||||
let on_disk_bytes = file_len(&updates_path);
|
||||
let expected_lines = expected_replayed_lines(&opts);
|
||||
|
||||
let budget_mb = env_parse("SESSION_LOAD_MAX_PEAK_MB", 1024u64);
|
||||
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
|
||||
|
||||
let sampler = RssSampler::start();
|
||||
|
||||
let light = adapter
|
||||
.load_session_without_updates(&info)
|
||||
.await
|
||||
.expect("load_session_without_updates");
|
||||
let transcript = std::fs::read_to_string(&updates_path).expect("read updates.jsonl");
|
||||
let prepared = prepare_replay_lines(&transcript, None);
|
||||
let replayed = prepared.lines.len();
|
||||
|
||||
let outcome = sampler.finish(budget_mb);
|
||||
drop(prepared);
|
||||
drop(transcript);
|
||||
drop(light);
|
||||
|
||||
// Report before asserting so a count regression still emits the summary.
|
||||
report_summary(
|
||||
"rss",
|
||||
serde_json::json!({
|
||||
"replayed_lines": replayed,
|
||||
"expected_lines": expected_lines,
|
||||
}),
|
||||
on_disk_bytes,
|
||||
&outcome,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
replayed, expected_lines,
|
||||
"replayed line count must equal the non-ACU update count"
|
||||
);
|
||||
assert_bounds(None, on_disk_bytes, &outcome);
|
||||
}
|
||||
|
||||
struct CountingClient {
|
||||
count: Rc<RefCell<u64>>,
|
||||
}
|
||||
|
||||
#[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<()> {
|
||||
*self.count.borrow_mut() += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn count_replayed_notifications(session_id: acp::SessionId, cwd: PathBuf) -> u64 {
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async move {
|
||||
let count = Rc::new(RefCell::new(0u64));
|
||||
let client = CountingClient {
|
||||
count: count.clone(),
|
||||
};
|
||||
let loaded = xai_grok_shell::session::testkit::e2e::load_session_via_agent(
|
||||
client, "mem-soak", session_id, cwd,
|
||||
)
|
||||
.await;
|
||||
drop(loaded);
|
||||
*count.borrow()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[ignore = "heavy: builds a full MvpAgent and loads a large session; run with --ignored --nocapture"]
|
||||
async fn session_load_e2e_peak_rss() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
|
||||
let server = xai_grok_test_support::MockInferenceServer::start()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let grok_home = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
let opts = memory_spec();
|
||||
let budget_mb = env_parse("SESSION_LOAD_MAX_PEAK_MB", 1024u64);
|
||||
|
||||
// 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("GROK_HOME", grok_home.path());
|
||||
std::env::set_var("GROK_CLI_CHAT_PROXY_BASE_URL", server.url());
|
||||
std::env::set_var("GROK_XAI_API_BASE_URL", server.url());
|
||||
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
|
||||
std::env::set_var("GROK_TELEMETRY_ENABLED", "false");
|
||||
std::env::set_var("GROK_FEEDBACK_ENABLED", "false");
|
||||
std::env::set_var("GROK_TRACE_UPLOAD", "false");
|
||||
}
|
||||
|
||||
let (info, dir) = synth::prepare_session(grok_home.path(), cwd.path(), &opts).await;
|
||||
let on_disk_bytes = file_len(&dir.join("updates.jsonl"));
|
||||
|
||||
let sampler = RssSampler::start();
|
||||
|
||||
let replay_count =
|
||||
count_replayed_notifications(info.id.clone(), cwd.path().to_path_buf()).await;
|
||||
|
||||
// The agent load already dropped the loaded state, so the peak here comes
|
||||
// from the background sampler; the final read is only a backstop.
|
||||
let outcome = sampler.finish(budget_mb);
|
||||
report_summary(
|
||||
"rss-e2e",
|
||||
serde_json::json!({ "replayed_notifications": replay_count }),
|
||||
on_disk_bytes,
|
||||
&outcome,
|
||||
);
|
||||
|
||||
// At least one notification per synthesized turn must replay; a near-empty
|
||||
// replay that still grew memory would otherwise pass silently.
|
||||
assert!(
|
||||
replay_count >= opts.turns as u64,
|
||||
"replayed {replay_count} notifications, expected at least {} (one per turn)",
|
||||
opts.turns,
|
||||
);
|
||||
assert_bounds(Some("e2e"), on_disk_bytes, &outcome);
|
||||
}
|
||||
}
|
||||
|
|
@ -129,7 +129,7 @@ async fn test_fetch_settings_blocking_round_trip() {
|
|||
let result = tokio::task::spawn_blocking({
|
||||
let url = server.url().to_string();
|
||||
let auth = auth.clone();
|
||||
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None)
|
||||
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None).into_option()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
@ -146,7 +146,7 @@ async fn test_fetch_settings_blocking_round_trip() {
|
|||
let result = tokio::task::spawn_blocking({
|
||||
let url = server.url().to_string();
|
||||
let auth = auth.clone();
|
||||
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None)
|
||||
move || xai_grok_shell::remote::fetch_settings_blocking(&url, &auth, None).into_option()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
|
|||
796
crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs
Normal file
796
crates/codegen/xai-grok-shell/tests/test_subagent_soak.rs
Normal file
|
|
@ -0,0 +1,796 @@
|
|||
//! Subagent lifecycle soak: churn spawn/run/completion/eviction and assert
|
||||
//! threads, fds, and heap/RSS reach steady state. A stub `ChildRunner` drives
|
||||
//! the real coordinator/transport.
|
||||
//!
|
||||
//! SUBAGENT_SOAK_CYCLES=20000 cargo test -p xai-grok-shell \
|
||||
//! [--features dhat-heap] --test test_subagent_soak -- --ignored --nocapture
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
#[global_allocator]
|
||||
static DHAT_ALLOC: dhat::Alloc = dhat::Alloc;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::ser::SerializeMap;
|
||||
use serde::{Serialize, Serializer};
|
||||
use strum::{EnumCount, IntoEnumIterator};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use xai_grok_test_support::env::env_parse;
|
||||
use xai_grok_test_support::resources::{ResourceGrowth, ResourceSnapshot};
|
||||
use xai_grok_tools::implementations::grok_build::task::backend::{ChannelBackend, SubagentBackend};
|
||||
use xai_grok_tools::implementations::grok_build::task::coordinator::{
|
||||
ChildCompletion, ChildControl, ChildRunOutput, ChildRunRequest, ChildRunner, CoordinatorConfig,
|
||||
LocalBoxFuture, MAX_COMPLETED_ENTRIES, StartedChild, SubagentCoordinator, SubagentProgress,
|
||||
};
|
||||
use xai_grok_tools::implementations::grok_build::task::types::{
|
||||
SubagentDescribeOutcome, SubagentOwner, SubagentRegistryCounts, SubagentRequest,
|
||||
SubagentResult, SubagentValidateTypeOutcome,
|
||||
};
|
||||
|
||||
const PARENT_SESSION_ID: &str = "subagent-soak-parent";
|
||||
|
||||
#[derive(Clone, Copy, strum::EnumCount, strum::EnumIter)]
|
||||
enum Metric {
|
||||
Rss,
|
||||
Threads,
|
||||
Fds,
|
||||
}
|
||||
|
||||
impl Metric {
|
||||
fn label(self) -> &'static str {
|
||||
match self {
|
||||
Metric::Rss => "rss",
|
||||
Metric::Threads => "threads",
|
||||
Metric::Fds => "fds",
|
||||
}
|
||||
}
|
||||
|
||||
/// RSS reports raw bytes, so its key names the unit.
|
||||
fn summary_key(self) -> &'static str {
|
||||
match self {
|
||||
Metric::Rss => "rss_bytes",
|
||||
Metric::Threads => "threads",
|
||||
Metric::Fds => "fds",
|
||||
}
|
||||
}
|
||||
|
||||
fn unit(self) -> Option<&'static str> {
|
||||
match self {
|
||||
Metric::Rss => Some("MiB"),
|
||||
Metric::Threads | Metric::Fds => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn budget(self, bounds: &Bounds) -> f64 {
|
||||
match self {
|
||||
Metric::Rss => bounds.max_rss_growth_mib as f64,
|
||||
Metric::Threads => bounds.max_thread_growth as f64,
|
||||
Metric::Fds => bounds.max_fd_growth as f64,
|
||||
}
|
||||
}
|
||||
|
||||
/// RSS growth samples are bytes; convert to MiB for the budget comparison.
|
||||
fn growth_in_budget_unit(self, raw: usize) -> f64 {
|
||||
match self {
|
||||
Metric::Rss => bytes_to_mib(raw),
|
||||
Metric::Threads | Metric::Fds => raw as f64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads a metric's field from a snapshot or a growth delta so serialization and
|
||||
/// the gates share one projection instead of repeating it.
|
||||
trait MetricValue {
|
||||
fn value_of(&self, metric: Metric) -> Option<usize>;
|
||||
}
|
||||
|
||||
impl MetricValue for ResourceSnapshot {
|
||||
fn value_of(&self, metric: Metric) -> Option<usize> {
|
||||
// Destructure so a new resource field is a compile error here, not a
|
||||
// silently dropped metric.
|
||||
let ResourceSnapshot { rss, threads, fds } = *self;
|
||||
match metric {
|
||||
Metric::Rss => rss,
|
||||
Metric::Threads => threads,
|
||||
Metric::Fds => fds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MetricValue for ResourceGrowth {
|
||||
fn value_of(&self, metric: Metric) -> Option<usize> {
|
||||
let ResourceGrowth { rss, threads, fds } = *self;
|
||||
match metric {
|
||||
Metric::Rss => rss,
|
||||
Metric::Threads => threads,
|
||||
Metric::Fds => fds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize_metrics<T: MetricValue, S: Serializer>(
|
||||
value: &T,
|
||||
serializer: S,
|
||||
) -> Result<S::Ok, S::Error> {
|
||||
let mut map = serializer.serialize_map(Some(Metric::COUNT))?;
|
||||
for metric in Metric::iter() {
|
||||
map.serialize_entry(metric.summary_key(), &value.value_of(metric))?;
|
||||
}
|
||||
map.end()
|
||||
}
|
||||
|
||||
fn bytes_to_mib(bytes: usize) -> f64 {
|
||||
bytes as f64 / (1024.0 * 1024.0)
|
||||
}
|
||||
|
||||
#[cfg_attr(not(feature = "dhat-heap"), allow(dead_code))]
|
||||
#[derive(Clone, Copy, Serialize)]
|
||||
struct HeapSample {
|
||||
blocks: i64,
|
||||
bytes: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Serialize)]
|
||||
struct HeapMetrics {
|
||||
before: HeapSample,
|
||||
after: HeapSample,
|
||||
blocks_per_cycle: f64,
|
||||
bytes_per_cycle: f64,
|
||||
}
|
||||
|
||||
impl HeapMetrics {
|
||||
fn new(before: HeapSample, after: HeapSample, cycles: u64) -> Self {
|
||||
// `SUBAGENT_SOAK_CYCLES=0` would otherwise divide by zero and feed
|
||||
// NaN/inf into the leak gates.
|
||||
let cycles = cycles.max(1) as f64;
|
||||
Self {
|
||||
before,
|
||||
after,
|
||||
blocks_per_cycle: (after.blocks - before.blocks) as f64 / cycles,
|
||||
bytes_per_cycle: (after.bytes - before.bytes) as f64 / cycles,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Bounds {
|
||||
#[serde(rename = "warmup_cycles")]
|
||||
warmup: u64,
|
||||
#[serde(rename = "measured_cycles")]
|
||||
measure: u64,
|
||||
max_thread_growth: u64,
|
||||
max_fd_growth: u64,
|
||||
max_rss_growth_mib: u64,
|
||||
max_blocks_per_cycle: f64,
|
||||
max_bytes_per_cycle: f64,
|
||||
}
|
||||
|
||||
impl Bounds {
|
||||
fn from_env() -> Self {
|
||||
Self {
|
||||
// Default warmup to the completed-entry cap so the ring is saturated
|
||||
// and the measured window observes steady-state eviction rather than
|
||||
// one-time cache fill.
|
||||
warmup: env_parse("SUBAGENT_SOAK_WARMUP", MAX_COMPLETED_ENTRIES as u64),
|
||||
measure: env_parse("SUBAGENT_SOAK_CYCLES", 512u64),
|
||||
max_thread_growth: env_parse("SUBAGENT_SOAK_MAX_THREAD_GROWTH", 32u64),
|
||||
max_fd_growth: env_parse("SUBAGENT_SOAK_MAX_FD_GROWTH", 64u64),
|
||||
max_rss_growth_mib: env_parse("SUBAGENT_SOAK_MAX_RSS_GROWTH_MIB", 256u64),
|
||||
max_blocks_per_cycle: env_parse("SUBAGENT_SOAK_MAX_BLOCKS_PER_CYCLE", 2.0f64),
|
||||
max_bytes_per_cycle: env_parse("SUBAGENT_SOAK_MAX_BYTES_PER_CYCLE", 4096.0f64),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Measurement {
|
||||
#[serde(serialize_with = "serialize_metrics")]
|
||||
before: ResourceSnapshot,
|
||||
#[serde(serialize_with = "serialize_metrics")]
|
||||
after: ResourceSnapshot,
|
||||
#[serde(serialize_with = "serialize_metrics")]
|
||||
growth: ResourceGrowth,
|
||||
#[serde(serialize_with = "serialize_counts")]
|
||||
counts: SubagentRegistryCounts,
|
||||
heap: Option<HeapMetrics>,
|
||||
quiesced: bool,
|
||||
}
|
||||
|
||||
fn serialize_counts<S: Serializer>(
|
||||
counts: &SubagentRegistryCounts,
|
||||
serializer: S,
|
||||
) -> Result<S::Ok, S::Error> {
|
||||
// Exhaustive destructure so a new count field is a compile error here, not a
|
||||
// silently dropped summary key.
|
||||
let SubagentRegistryCounts {
|
||||
pending,
|
||||
active,
|
||||
completed,
|
||||
} = counts;
|
||||
let mut map = serializer.serialize_map(Some(3))?;
|
||||
map.serialize_entry("pending", pending)?;
|
||||
map.serialize_entry("active", active)?;
|
||||
map.serialize_entry("completed", completed)?;
|
||||
map.end()
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Summary<'a> {
|
||||
#[serde(flatten)]
|
||||
bounds: &'a Bounds,
|
||||
#[serde(flatten)]
|
||||
measurement: &'a Measurement,
|
||||
}
|
||||
|
||||
fn heap_capture() -> Option<HeapSample> {
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
{
|
||||
let stats = dhat::HeapStats::get();
|
||||
Some(HeapSample {
|
||||
blocks: stats.curr_blocks as i64,
|
||||
bytes: stats.curr_bytes as i64,
|
||||
})
|
||||
}
|
||||
#[cfg(not(feature = "dhat-heap"))]
|
||||
{
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
async fn quiesce(backend: &ChannelBackend) -> bool {
|
||||
const MAX_POLLS: usize = 200;
|
||||
const SLEEP: Duration = Duration::from_millis(5);
|
||||
for _ in 0..MAX_POLLS {
|
||||
let counts = backend.registry_counts().await;
|
||||
if counts.pending == 0 && counts.active == 0 {
|
||||
return true;
|
||||
}
|
||||
tokio::time::sleep(SLEEP).await;
|
||||
}
|
||||
let counts = backend.registry_counts().await;
|
||||
eprintln!(
|
||||
"[soak] quiesce budget expired with pending={} active={}; snapshot may be noisy",
|
||||
counts.pending, counts.active
|
||||
);
|
||||
false
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SoakControl {
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
impl ChildControl for SoakControl {
|
||||
type ProgressFuture = std::future::Ready<SubagentProgress>;
|
||||
|
||||
fn progress(&self) -> Self::ProgressFuture {
|
||||
std::future::ready(SubagentProgress::default())
|
||||
}
|
||||
|
||||
fn cancel(&self) {
|
||||
self.cancellation.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
struct SoakRunner;
|
||||
|
||||
impl ChildRunner for SoakRunner {
|
||||
type Control = SoakControl;
|
||||
type CompletionData = ();
|
||||
type RunFuture = LocalBoxFuture<ChildRunOutput<()>>;
|
||||
type ValidateFuture = LocalBoxFuture<SubagentValidateTypeOutcome>;
|
||||
type DescribeFuture = LocalBoxFuture<SubagentDescribeOutcome>;
|
||||
|
||||
fn run(&self, run: ChildRunRequest<Self::Control>) -> Self::RunFuture {
|
||||
Box::pin(async move {
|
||||
let ChildRunRequest {
|
||||
request,
|
||||
cancellation,
|
||||
reporter,
|
||||
} = run;
|
||||
let promoted = reporter
|
||||
.started(StartedChild {
|
||||
child_session_id: request.id.clone(),
|
||||
persona: None,
|
||||
resumed_from: request.resume_from.clone(),
|
||||
child_cwd: request.cwd.clone().unwrap_or_default(),
|
||||
worktree_path: None,
|
||||
effective_model_id: "soak-model".to_owned(),
|
||||
definition_background: false,
|
||||
control: SoakControl {
|
||||
cancellation: cancellation.clone(),
|
||||
},
|
||||
})
|
||||
.await;
|
||||
if !promoted || cancellation.is_cancelled() {
|
||||
return ChildRunOutput {
|
||||
result: SubagentResult {
|
||||
success: false,
|
||||
cancelled: true,
|
||||
error: Some("cancelled before start".to_owned()),
|
||||
subagent_id: request.id.clone(),
|
||||
child_session_id: request.id,
|
||||
..Default::default()
|
||||
},
|
||||
completion_data: (),
|
||||
snapshot_ref: None,
|
||||
};
|
||||
}
|
||||
ChildRunOutput {
|
||||
result: SubagentResult {
|
||||
success: true,
|
||||
output: Arc::from("soak child output"),
|
||||
subagent_id: request.id.clone(),
|
||||
child_session_id: request.id,
|
||||
tool_calls: 1,
|
||||
turns: 1,
|
||||
..Default::default()
|
||||
},
|
||||
completion_data: (),
|
||||
snapshot_ref: None,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_type(&self, _subagent_type: String, _parent: String) -> Self::ValidateFuture {
|
||||
Box::pin(std::future::ready(SubagentValidateTypeOutcome::Ok))
|
||||
}
|
||||
|
||||
fn describe_type(
|
||||
&self,
|
||||
_subagent_type: String,
|
||||
_harness_agent_type: Option<String>,
|
||||
_parent: String,
|
||||
) -> Self::DescribeFuture {
|
||||
Box::pin(std::future::ready(SubagentDescribeOutcome::Unavailable))
|
||||
}
|
||||
|
||||
fn on_completed(&self, _completion: ChildCompletion<Self::CompletionData>) {}
|
||||
}
|
||||
|
||||
fn soak_request(id: String, background: bool) -> SubagentRequest {
|
||||
SubagentRequest {
|
||||
id,
|
||||
prompt: "soak work".to_owned(),
|
||||
description: "soak child".to_owned(),
|
||||
subagent_type: "explore".to_owned(),
|
||||
parent_session_id: PARENT_SESSION_ID.to_owned(),
|
||||
parent_prompt_id: Some("soak-prompt".to_owned()),
|
||||
resume_from: None,
|
||||
cwd: None,
|
||||
runtime_overrides: Default::default(),
|
||||
run_in_background: background,
|
||||
surface_completion: true,
|
||||
await_to_completion: false,
|
||||
fork_context: false,
|
||||
owner: SubagentOwner::Task,
|
||||
cancel_token: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_cycle(backend: &ChannelBackend, i: u64) {
|
||||
let fg = backend
|
||||
.spawn(soak_request(format!("fg-{i}"), false))
|
||||
.await
|
||||
.expect("foreground spawn round-trips through the coordinator");
|
||||
assert!(fg.success, "cycle {i}: foreground child must complete");
|
||||
|
||||
let bg_id = format!("bg-{i}");
|
||||
let bg = backend
|
||||
.spawn(soak_request(bg_id.clone(), true))
|
||||
.await
|
||||
.expect("background spawn round-trips through the coordinator");
|
||||
assert!(bg.success, "cycle {i}: background child must complete");
|
||||
|
||||
let blocking = true;
|
||||
let timeout_ms = Some(5_000);
|
||||
let snapshot = backend.query(&bg_id, blocking, timeout_ms).await;
|
||||
assert!(
|
||||
snapshot.is_some(),
|
||||
"cycle {i}: completed subagent must be queryable"
|
||||
);
|
||||
}
|
||||
|
||||
async fn warmup(backend: &ChannelBackend, cycles: u64) -> bool {
|
||||
for i in 0..cycles {
|
||||
run_cycle(backend, i).await;
|
||||
}
|
||||
quiesce(backend).await
|
||||
}
|
||||
|
||||
async fn measure(backend: &ChannelBackend, bounds: &Bounds, warmup_quiesced: bool) -> Measurement {
|
||||
let heap_before = heap_capture();
|
||||
let before = ResourceSnapshot::capture();
|
||||
|
||||
// Continue ids past the warmup window so measured cycles use fresh entries
|
||||
// and keep exercising eviction instead of colliding with warmup ids.
|
||||
for i in bounds.warmup..(bounds.warmup + bounds.measure) {
|
||||
run_cycle(backend, i).await;
|
||||
}
|
||||
// A warmup that never drained already poisons the `before` baseline, so skip
|
||||
// the measured-window drain and report the window as not quiesced.
|
||||
let quiesced = warmup_quiesced && quiesce(backend).await;
|
||||
|
||||
let heap_after = heap_capture();
|
||||
let after = ResourceSnapshot::capture();
|
||||
let counts = backend.registry_counts().await;
|
||||
|
||||
Measurement {
|
||||
before,
|
||||
after,
|
||||
growth: after.growth_from(&before),
|
||||
counts,
|
||||
heap: heap_before
|
||||
.zip(heap_after)
|
||||
.map(|(before, after)| HeapMetrics::new(before, after, bounds.measure)),
|
||||
quiesced,
|
||||
}
|
||||
}
|
||||
|
||||
fn check_bounds(bounds: &Bounds, m: &Measurement) -> Vec<String> {
|
||||
// Drain first: a non-quiesced window has nonzero counts and noisy growth, so
|
||||
// report the quiesce failure alone; the gates below only mean anything once
|
||||
// drained.
|
||||
if !m.quiesced {
|
||||
return vec![
|
||||
"quiesce budget expired before the measured window drained; soak result is unreliable"
|
||||
.to_owned(),
|
||||
];
|
||||
}
|
||||
|
||||
let mut failures = Vec::new();
|
||||
if m.counts.pending != 0 {
|
||||
failures.push(format!(
|
||||
"no subagent may remain pending, saw {}",
|
||||
m.counts.pending
|
||||
));
|
||||
}
|
||||
if m.counts.active != 0 {
|
||||
failures.push(format!(
|
||||
"no subagent may remain active, saw {}",
|
||||
m.counts.active
|
||||
));
|
||||
}
|
||||
if m.counts.completed > MAX_COMPLETED_ENTRIES {
|
||||
failures.push(format!(
|
||||
"completed retention must stay bounded by its cap, saw {}",
|
||||
m.counts.completed
|
||||
));
|
||||
}
|
||||
|
||||
for metric in Metric::iter() {
|
||||
let Some(raw) = m.growth.value_of(metric) else {
|
||||
continue;
|
||||
};
|
||||
let growth = metric.growth_in_budget_unit(raw);
|
||||
let budget = metric.budget(bounds);
|
||||
if growth > budget {
|
||||
let unit = metric.unit().map(|u| format!(" {u}")).unwrap_or_default();
|
||||
failures.push(format!(
|
||||
"{}: grew {growth:.1}{unit} over the soak (bound {budget:.1}{unit})",
|
||||
metric.label()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(h) = m.heap {
|
||||
let measure = bounds.measure;
|
||||
if h.blocks_per_cycle > bounds.max_blocks_per_cycle {
|
||||
failures.push(format!(
|
||||
"block-count leak: {:.3} blocks/cycle retained ({} over {measure} cycles) \
|
||||
exceeds the {} gate",
|
||||
h.blocks_per_cycle,
|
||||
h.after.blocks - h.before.blocks,
|
||||
bounds.max_blocks_per_cycle
|
||||
));
|
||||
}
|
||||
if h.bytes_per_cycle > bounds.max_bytes_per_cycle {
|
||||
failures.push(format!(
|
||||
"byte leak: {:.1} bytes/cycle retained ({} over {measure} cycles) \
|
||||
exceeds the {} gate",
|
||||
h.bytes_per_cycle,
|
||||
h.after.bytes - h.before.bytes,
|
||||
bounds.max_bytes_per_cycle
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
failures
|
||||
}
|
||||
|
||||
fn assert_bounds(bounds: &Bounds, m: &Measurement) {
|
||||
let failures = check_bounds(bounds, m);
|
||||
assert!(
|
||||
failures.is_empty(),
|
||||
"subagent soak bounds violated:\n - {}",
|
||||
failures.join("\n - ")
|
||||
);
|
||||
}
|
||||
|
||||
/// Keep this the only test in the binary that creates a `dhat::Profiler`.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
#[ignore = "subagent soak; run with --ignored (SUBAGENT_SOAK_CYCLES bounds the measured window)"]
|
||||
async fn subagent_lifecycle_soak_bounds_threads_fds_and_heap() {
|
||||
#[cfg(feature = "dhat-heap")]
|
||||
let _profiler = dhat::Profiler::builder().testing().build();
|
||||
|
||||
let bounds = Bounds::from_env();
|
||||
|
||||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async move {
|
||||
let (command_tx, command_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let config = CoordinatorConfig {
|
||||
foreground_budget: Duration::from_secs(600),
|
||||
..CoordinatorConfig::default()
|
||||
};
|
||||
tokio::task::spawn_local(
|
||||
SubagentCoordinator::new(command_rx, SoakRunner, config).run(),
|
||||
);
|
||||
let backend = ChannelBackend::new(command_tx);
|
||||
|
||||
let warmup_quiesced = warmup(&backend, bounds.warmup).await;
|
||||
let measurement = measure(&backend, &bounds, warmup_quiesced).await;
|
||||
|
||||
let summary = Summary {
|
||||
bounds: &bounds,
|
||||
measurement: &measurement,
|
||||
};
|
||||
eprintln!(
|
||||
"SUBAGENT_SOAK_SUMMARY {}",
|
||||
serde_json::to_string(&summary).expect("summary serializes")
|
||||
);
|
||||
|
||||
assert_bounds(&bounds, &measurement);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn value_of_reads_the_matching_slot_of_snapshot_and_growth() {
|
||||
let snapshot = ResourceSnapshot {
|
||||
rss: Some(11),
|
||||
threads: Some(22),
|
||||
fds: Some(33),
|
||||
};
|
||||
assert_eq!(snapshot.value_of(Metric::Rss), Some(11));
|
||||
assert_eq!(snapshot.value_of(Metric::Threads), Some(22));
|
||||
assert_eq!(snapshot.value_of(Metric::Fds), Some(33));
|
||||
|
||||
let growth = ResourceGrowth {
|
||||
rss: Some(1),
|
||||
threads: None,
|
||||
fds: Some(3),
|
||||
};
|
||||
assert_eq!(growth.value_of(Metric::Rss), Some(1));
|
||||
assert_eq!(growth.value_of(Metric::Threads), None);
|
||||
assert_eq!(growth.value_of(Metric::Fds), Some(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialize_metrics_keys_match_summary_keys_in_order() {
|
||||
#[derive(Serialize)]
|
||||
struct Wrap(#[serde(serialize_with = "serialize_metrics")] ResourceSnapshot);
|
||||
let snapshot = ResourceSnapshot {
|
||||
rss: Some(1),
|
||||
threads: None,
|
||||
fds: Some(3),
|
||||
};
|
||||
let json = serde_json::to_string(&Wrap(snapshot)).expect("snapshot serializes");
|
||||
assert_eq!(json, r#"{"rss_bytes":1,"threads":null,"fds":3}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bytes_to_mib_divides_by_1024_squared() {
|
||||
assert_eq!(bytes_to_mib(0), 0.0);
|
||||
assert_eq!(bytes_to_mib(1024 * 1024), 1.0);
|
||||
assert_eq!(bytes_to_mib(3 * 1024 * 1024), 3.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn growth_in_budget_unit_scales_only_rss() {
|
||||
assert_eq!(Metric::Rss.growth_in_budget_unit(2 * 1024 * 1024), 2.0);
|
||||
assert_eq!(Metric::Threads.growth_in_budget_unit(7), 7.0);
|
||||
assert_eq!(Metric::Fds.growth_in_budget_unit(7), 7.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn budget_reads_per_metric_bound() {
|
||||
let bounds = Bounds {
|
||||
warmup: 0,
|
||||
measure: 0,
|
||||
max_thread_growth: 3,
|
||||
max_fd_growth: 5,
|
||||
max_rss_growth_mib: 7,
|
||||
max_blocks_per_cycle: 1.0,
|
||||
max_bytes_per_cycle: 2.0,
|
||||
};
|
||||
assert_eq!(Metric::Rss.budget(&bounds), 7.0);
|
||||
assert_eq!(Metric::Threads.budget(&bounds), 3.0);
|
||||
assert_eq!(Metric::Fds.budget(&bounds), 5.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heap_metrics_clamps_zero_cycles() {
|
||||
let before = HeapSample {
|
||||
blocks: 10,
|
||||
bytes: 100,
|
||||
};
|
||||
let after = HeapSample {
|
||||
blocks: 20,
|
||||
bytes: 400,
|
||||
};
|
||||
let heap = HeapMetrics::new(before, after, 0);
|
||||
assert!(heap.blocks_per_cycle.is_finite());
|
||||
assert!(heap.bytes_per_cycle.is_finite());
|
||||
assert_eq!(heap.blocks_per_cycle, 10.0);
|
||||
assert_eq!(heap.bytes_per_cycle, 300.0);
|
||||
}
|
||||
|
||||
fn generous_bounds() -> Bounds {
|
||||
Bounds {
|
||||
warmup: 0,
|
||||
measure: 4,
|
||||
max_thread_growth: 100,
|
||||
max_fd_growth: 100,
|
||||
max_rss_growth_mib: 100,
|
||||
max_blocks_per_cycle: 10.0,
|
||||
max_bytes_per_cycle: 10_000.0,
|
||||
}
|
||||
}
|
||||
|
||||
fn drained(growth: ResourceGrowth, heap: Option<HeapMetrics>) -> Measurement {
|
||||
Measurement {
|
||||
before: ResourceSnapshot::default(),
|
||||
after: ResourceSnapshot::default(),
|
||||
growth,
|
||||
counts: SubagentRegistryCounts {
|
||||
pending: 0,
|
||||
active: 0,
|
||||
completed: 0,
|
||||
},
|
||||
heap,
|
||||
quiesced: true,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_passes_a_clean_drained_window() {
|
||||
let m = drained(ResourceGrowth::default(), None);
|
||||
assert!(check_bounds(&generous_bounds(), &m).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_reports_non_quiesce_first_and_alone() {
|
||||
let mut m = drained(ResourceGrowth::default(), None);
|
||||
m.quiesced = false;
|
||||
m.counts.pending = 3;
|
||||
let failures = check_bounds(&generous_bounds(), &m);
|
||||
assert_eq!(failures.len(), 1);
|
||||
assert!(failures[0].contains("quiesce"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_over_budget_growth() {
|
||||
let growth = ResourceGrowth {
|
||||
rss: Some(200 * 1024 * 1024),
|
||||
threads: Some(0),
|
||||
fds: Some(0),
|
||||
};
|
||||
let failures = check_bounds(&generous_bounds(), &drained(growth, None));
|
||||
assert!(
|
||||
failures.iter().any(|f| f.starts_with("rss:")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_treats_the_budget_as_an_inclusive_max() {
|
||||
let growth = ResourceGrowth {
|
||||
rss: Some(100 * 1024 * 1024),
|
||||
threads: Some(100),
|
||||
fds: Some(100),
|
||||
};
|
||||
assert!(check_bounds(&generous_bounds(), &drained(growth, None)).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_nonzero_counts_and_heap_leak() {
|
||||
let mut m = drained(
|
||||
ResourceGrowth::default(),
|
||||
Some(HeapMetrics {
|
||||
before: HeapSample {
|
||||
blocks: 0,
|
||||
bytes: 0,
|
||||
},
|
||||
after: HeapSample {
|
||||
blocks: 0,
|
||||
bytes: 0,
|
||||
},
|
||||
blocks_per_cycle: 0.0,
|
||||
bytes_per_cycle: 1_000_000.0,
|
||||
}),
|
||||
);
|
||||
m.counts.active = 2;
|
||||
let failures = check_bounds(&generous_bounds(), &m);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.contains("active")),
|
||||
"{failures:?}"
|
||||
);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.contains("byte leak")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_pending_while_quiesced() {
|
||||
let mut m = drained(ResourceGrowth::default(), None);
|
||||
m.counts.pending = 3;
|
||||
let failures = check_bounds(&generous_bounds(), &m);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.contains("pending")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_completed_over_cap() {
|
||||
let mut m = drained(ResourceGrowth::default(), None);
|
||||
m.counts.completed = MAX_COMPLETED_ENTRIES + 1;
|
||||
let failures = check_bounds(&generous_bounds(), &m);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.contains("completed retention")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_thread_and_fd_over_budget() {
|
||||
let growth = ResourceGrowth {
|
||||
rss: Some(0),
|
||||
threads: Some(200),
|
||||
fds: Some(200),
|
||||
};
|
||||
let failures = check_bounds(&generous_bounds(), &drained(growth, None));
|
||||
assert!(
|
||||
failures.iter().any(|f| f.starts_with("threads:")),
|
||||
"{failures:?}"
|
||||
);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.starts_with("fds:")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_bounds_flags_block_count_leak() {
|
||||
let m = drained(
|
||||
ResourceGrowth::default(),
|
||||
Some(HeapMetrics {
|
||||
before: HeapSample {
|
||||
blocks: 0,
|
||||
bytes: 0,
|
||||
},
|
||||
after: HeapSample {
|
||||
blocks: 0,
|
||||
bytes: 0,
|
||||
},
|
||||
blocks_per_cycle: 50.0,
|
||||
bytes_per_cycle: 0.0,
|
||||
}),
|
||||
);
|
||||
let failures = check_bounds(&generous_bounds(), &m);
|
||||
assert!(
|
||||
failures.iter().any(|f| f.contains("block-count leak")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
//! Non-ignored guard for the testkit's core path: synthesize a small session
|
||||
//! with [`synth::prepare_session`], then confirm the production replay reader
|
||||
//! parses every persisted update back with the right per-kind counts.
|
||||
//!
|
||||
//! `load_updates_for_replay_at` is the typed reader and keeps every update
|
||||
//! (only `Xai` updates are dropped); the redundant-ACU skip is a later
|
||||
//! line-based step in the client replay path, not asserted here.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use xai_grok_shell::session::storage::{
|
||||
JsonlStorageAdapter, StorageAdapter, load_updates_for_replay_at,
|
||||
};
|
||||
use xai_grok_shell::session::testkit::synth::{self, SessionSpec};
|
||||
|
||||
#[tokio::test]
|
||||
async fn synth_replay_roundtrip_parses_every_persisted_update() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let cwd = TempDir::new().unwrap();
|
||||
// Distinct per-turn counts so a miscount of any kind is unambiguous.
|
||||
let spec = SessionSpec {
|
||||
turns: 3,
|
||||
acu_per_turn: 2,
|
||||
catalog_commands: 2,
|
||||
catalog_desc_len: 8,
|
||||
agent_chunks_per_turn: 4,
|
||||
agent_chunk_len: 16,
|
||||
rewind_points: 0,
|
||||
files_per_rewind: 0,
|
||||
file_content_len: 0,
|
||||
};
|
||||
|
||||
let (info, _dir) = synth::prepare_session(root.path(), cwd.path(), &spec).await;
|
||||
|
||||
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
|
||||
.expect("load_updates_for_replay_at")
|
||||
.unwrap_or_default();
|
||||
|
||||
let count = |pred: fn(&acp::SessionUpdate) -> bool| replayed.iter().filter(|u| pred(u)).count();
|
||||
let users = count(|u| matches!(u, acp::SessionUpdate::UserMessageChunk(_)));
|
||||
let acus = count(|u| matches!(u, acp::SessionUpdate::AvailableCommandsUpdate(_)));
|
||||
let agents = count(|u| matches!(u, acp::SessionUpdate::AgentMessageChunk(_)));
|
||||
|
||||
assert_eq!(users, spec.turns, "one user chunk per turn");
|
||||
assert_eq!(
|
||||
acus,
|
||||
spec.turns * spec.acu_per_turn,
|
||||
"every ACU is preserved"
|
||||
);
|
||||
assert_eq!(
|
||||
agents,
|
||||
spec.turns * spec.agent_chunks_per_turn,
|
||||
"every agent chunk is preserved"
|
||||
);
|
||||
assert_eq!(
|
||||
replayed.len(),
|
||||
spec.turns * (1 + spec.acu_per_turn + spec.agent_chunks_per_turn),
|
||||
"no update is dropped or duplicated by the typed replay reader"
|
||||
);
|
||||
}
|
||||
|
||||
/// The adapter-driven bench generator reaches its byte target and emits updates
|
||||
/// the production reader parses. Synchronous because `synthesize_to_target_bytes`
|
||||
/// drives the adapter on its own runtime.
|
||||
#[test]
|
||||
fn synthesize_to_target_bytes_reaches_target_and_parses() {
|
||||
let root = TempDir::new().unwrap();
|
||||
let target: u64 = 8 * 1024;
|
||||
|
||||
let info = synth::synthesize_to_target_bytes(root.path(), target);
|
||||
|
||||
let adapter = JsonlStorageAdapter::with_root(root.path().to_path_buf());
|
||||
let updates_path = adapter.updates_file_path(&info).expect("updates path");
|
||||
let len = std::fs::metadata(&updates_path)
|
||||
.expect("stat updates.jsonl")
|
||||
.len();
|
||||
assert!(
|
||||
len >= target,
|
||||
"updates.jsonl ({len} B) reached the target ({target} B)"
|
||||
);
|
||||
|
||||
let replayed = load_updates_for_replay_at(info.id.0.as_ref(), root.path())
|
||||
.expect("load_updates_for_replay_at")
|
||||
.unwrap_or_default();
|
||||
assert!(!replayed.is_empty(), "the emitted updates parse back");
|
||||
}
|
||||
Loading…
Reference in a new issue