Publish harness and TUI open-source
initial sync from the monorepo
This commit is contained in:
grokkybara[bot]
commit
c68e39f604
2734 changed files with 1437016 additions and 0 deletions
354
crates/codegen/xai-grok-test-support/src/leader.rs
Normal file
354
crates/codegen/xai-grok-test-support/src/leader.rs
Normal file
|
|
@ -0,0 +1,354 @@
|
|||
//! Leader-mode (`grok agent --leader stdio`) test harness.
|
||||
//!
|
||||
//! Spawns the real binary as a stdio client whose bridge elects a leader
|
||||
//! subprocess hosting the actual sessions, speaks ACP over pipes, and
|
||||
//! exposes lock-file helpers for leader-lifecycle assertions. Unix-only:
|
||||
//! the leader transport is a unix socket.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use agent_client_protocol::{self as acp, Agent as _};
|
||||
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
|
||||
use xai_acp_lib::LineBufferedRead;
|
||||
|
||||
use crate::env::grok_binary;
|
||||
use crate::mock_server::MockInferenceServer;
|
||||
use crate::process::spawn_piped_with_stderr_capture;
|
||||
|
||||
/// Env var naming the binary that elects/hosts the leader in a two-binary
|
||||
/// (version-skew) test. Falls back to [`grok_binary`]'s resolution.
|
||||
pub const LEADER_BINARY_ENV: &str = "GROK_BINARY_LEADER";
|
||||
|
||||
/// Env var naming the binary for the second (usually newer) client in a
|
||||
/// two-binary test. Falls back to [`grok_binary`]'s resolution.
|
||||
pub const CLIENT_BINARY_ENV: &str = "GROK_BINARY_CLIENT";
|
||||
|
||||
fn role_binary(env_key: &str) -> PathBuf {
|
||||
if let Ok(path) = std::env::var(env_key) {
|
||||
let p = PathBuf::from(path);
|
||||
assert!(p.exists(), "{env_key} does not exist: {}", p.display());
|
||||
return p;
|
||||
}
|
||||
grok_binary()
|
||||
}
|
||||
|
||||
/// Binary for the leader-electing side of a version-skew test
|
||||
/// (`GROK_BINARY_LEADER`, else the shared [`grok_binary`] resolution).
|
||||
pub fn leader_binary() -> PathBuf {
|
||||
role_binary(LEADER_BINARY_ENV)
|
||||
}
|
||||
|
||||
/// Binary for the client side of a version-skew test (`GROK_BINARY_CLIENT`,
|
||||
/// else the shared [`grok_binary`] resolution).
|
||||
pub fn client_binary() -> PathBuf {
|
||||
role_binary(CLIENT_BINARY_ENV)
|
||||
}
|
||||
|
||||
/// Capture for notifications + reconnect signals.
|
||||
#[derive(Default)]
|
||||
pub struct Capture {
|
||||
chunks: std::sync::Mutex<Vec<String>>,
|
||||
notification_count: AtomicU32,
|
||||
reconnected_count: AtomicU32,
|
||||
}
|
||||
|
||||
struct LeaderAcpClient {
|
||||
capture: Arc<Capture>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait(?Send)]
|
||||
impl acp::Client for LeaderAcpClient {
|
||||
async fn request_permission(
|
||||
&self,
|
||||
args: acp::RequestPermissionRequest,
|
||||
) -> acp::Result<acp::RequestPermissionResponse> {
|
||||
let outcome = args
|
||||
.options
|
||||
.iter()
|
||||
.find(|o| o.kind == acp::PermissionOptionKind::AllowOnce)
|
||||
.or(args.options.first())
|
||||
.map(|o| {
|
||||
acp::RequestPermissionOutcome::Selected(acp::SelectedPermissionOutcome::new(
|
||||
o.option_id.clone(),
|
||||
))
|
||||
})
|
||||
.unwrap_or(acp::RequestPermissionOutcome::Cancelled);
|
||||
Ok(acp::RequestPermissionResponse::new(outcome))
|
||||
}
|
||||
|
||||
async fn session_notification(&self, args: acp::SessionNotification) -> acp::Result<()> {
|
||||
self.capture
|
||||
.notification_count
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
if let acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk { content, .. }) =
|
||||
args.update
|
||||
&& let acp::ContentBlock::Text(t) = content
|
||||
{
|
||||
self.capture.chunks.lock().unwrap().push(t.text);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ext_notification(&self, args: acp::ExtNotification) -> acp::Result<()> {
|
||||
if &*args.method == "x.ai/leader_reconnected" {
|
||||
self.capture
|
||||
.reconnected_count
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// A `grok agent --leader stdio` client subprocess speaking ACP over pipes.
|
||||
/// The leader subprocess it elects hosts the actual sessions.
|
||||
pub struct LeaderStdioClient {
|
||||
pub conn: acp::ClientSideConnection,
|
||||
// Exposed for PID assertions.
|
||||
pub child: tokio::process::Child,
|
||||
capture: Arc<Capture>,
|
||||
stderr: Arc<std::sync::Mutex<Vec<u8>>>,
|
||||
}
|
||||
|
||||
impl LeaderStdioClient {
|
||||
pub async fn spawn(server: &MockInferenceServer, cwd: &Path, home: &Path) -> Self {
|
||||
Self::spawn_with_binary(&grok_binary(), server, cwd, home).await
|
||||
}
|
||||
|
||||
/// [`Self::spawn`] with an explicit binary, for two-binary version-skew
|
||||
/// tests (pair with [`leader_binary`] / [`client_binary`]).
|
||||
pub async fn spawn_with_binary(
|
||||
binary: &Path,
|
||||
server: &MockInferenceServer,
|
||||
cwd: &Path,
|
||||
home: &Path,
|
||||
) -> Self {
|
||||
let mut cmd = tokio::process::Command::new(binary);
|
||||
cmd.args(["agent", "--leader", "stdio"])
|
||||
.current_dir(cwd)
|
||||
// Hermetic env: the developer's shell may export GROK_* vars
|
||||
// (e.g. GROK_LEADER_SOCKET pointing at a REAL leader on this
|
||||
// machine). env_clear + explicit allowlist guarantees the test
|
||||
// can never touch a leader outside its sandbox home.
|
||||
.env_clear()
|
||||
.env("PATH", std::env::var("PATH").unwrap_or_default())
|
||||
.env("HOME", home)
|
||||
.env("GROK_HOME", home.join(".grok"))
|
||||
// Pin the socket inside the sandbox. The lock file is the
|
||||
// sibling `.lock` (leader.sock -> leader.lock), and the spawned
|
||||
// leader subprocess inherits/forwards this env var, so every
|
||||
// (re-)elected leader binds the same sandboxed path.
|
||||
.env("GROK_LEADER_SOCKET", home.join(".grok").join("leader.sock"))
|
||||
.env("GROK_CLI_CHAT_PROXY_BASE_URL", server.url())
|
||||
.env("GROK_XAI_API_BASE_URL", server.url())
|
||||
.env("XAI_API_KEY", "test-key-for-ci")
|
||||
.env("GROK_TELEMETRY_ENABLED", "false")
|
||||
.env("GROK_FEEDBACK_ENABLED", "false")
|
||||
.env("GROK_TRACE_UPLOAD", "false")
|
||||
.env("GROK_INSTRUMENTATION", "disabled")
|
||||
// Inherited by the spawned leader, whose stderr goes to
|
||||
// ~/.grok/leader.log — keep it chatty for diagnosis.
|
||||
.env("RUST_LOG", "xai_grok_shell=debug");
|
||||
|
||||
let (mut child, stderr) = spawn_piped_with_stderr_capture(cmd);
|
||||
|
||||
let outgoing = child.stdin.take().unwrap().compat_write();
|
||||
let incoming = child.stdout.take().unwrap().compat();
|
||||
|
||||
let capture = Arc::new(Capture::default());
|
||||
let client = LeaderAcpClient {
|
||||
capture: capture.clone(),
|
||||
};
|
||||
let incoming = LineBufferedRead::spawn_local(incoming);
|
||||
let (conn, handle_io) = acp::ClientSideConnection::new(client, outgoing, incoming, |fut| {
|
||||
tokio::task::spawn_local(fut);
|
||||
});
|
||||
tokio::task::spawn_local(handle_io);
|
||||
|
||||
Self {
|
||||
conn,
|
||||
child,
|
||||
capture,
|
||||
stderr,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stderr_text(&self) -> String {
|
||||
String::from_utf8_lossy(&self.stderr.lock().unwrap()).into_owned()
|
||||
}
|
||||
|
||||
pub fn captured_text(&self) -> String {
|
||||
self.capture.chunks.lock().unwrap().join("")
|
||||
}
|
||||
|
||||
pub async fn initialize(&self) -> acp::InitializeResponse {
|
||||
let init = tokio::time::timeout(
|
||||
Duration::from_secs(60),
|
||||
self.conn.initialize(
|
||||
acp::InitializeRequest::new(acp::ProtocolVersion::V1)
|
||||
.client_capabilities(
|
||||
acp::ClientCapabilities::new()
|
||||
.fs(acp::FileSystemCapabilities::new())
|
||||
.terminal(false),
|
||||
)
|
||||
.meta(
|
||||
serde_json::json!({
|
||||
"startupHints": {
|
||||
"nonInteractive": true,
|
||||
"skipGitStatus": true,
|
||||
"skipProjectLayout": true
|
||||
},
|
||||
"clientType": "test-client",
|
||||
"clientVersion": "0.0.0-test"
|
||||
})
|
||||
.as_object()
|
||||
.cloned(),
|
||||
),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|_| panic!("initialize timed out\nstderr:\n{}", self.stderr_text()))
|
||||
.expect("initialize failed");
|
||||
|
||||
let api_key_method = init
|
||||
.auth_methods
|
||||
.iter()
|
||||
.find(|m| &*m.id().0 == "xai.api_key")
|
||||
.expect("xai.api_key auth method");
|
||||
self.conn
|
||||
.authenticate(
|
||||
acp::AuthenticateRequest::new(api_key_method.id().clone())
|
||||
.meta(serde_json::json!({"headless": true}).as_object().cloned()),
|
||||
)
|
||||
.await
|
||||
.expect("authenticate failed");
|
||||
init
|
||||
}
|
||||
|
||||
pub async fn create_session(&self, cwd: &Path) -> acp::SessionId {
|
||||
self.create_session_inner(cwd, None).await
|
||||
}
|
||||
|
||||
pub async fn create_session_with_model(&self, cwd: &Path, model_id: &str) -> acp::SessionId {
|
||||
self.create_session_inner(
|
||||
cwd,
|
||||
serde_json::json!({ "modelId": model_id })
|
||||
.as_object()
|
||||
.cloned(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_session_inner(&self, cwd: &Path, meta: Option<acp::Meta>) -> acp::SessionId {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(30),
|
||||
self.conn.new_session(
|
||||
acp::NewSessionRequest::new(cwd.to_path_buf())
|
||||
.mcp_servers(vec![])
|
||||
.meta(meta),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|_| panic!("session/new timed out\nstderr:\n{}", self.stderr_text()))
|
||||
.expect("session/new failed")
|
||||
.session_id
|
||||
}
|
||||
|
||||
pub async fn prompt(
|
||||
&self,
|
||||
session_id: &acp::SessionId,
|
||||
text: &str,
|
||||
) -> acp::Result<acp::PromptResponse> {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(30),
|
||||
self.conn.prompt(acp::PromptRequest::new(
|
||||
session_id.clone(),
|
||||
vec![acp::ContentBlock::Text(acp::TextContent::new(
|
||||
text.to_string(),
|
||||
))],
|
||||
)),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|_| panic!("prompt timed out\nstderr:\n{}", self.stderr_text()))
|
||||
}
|
||||
|
||||
pub fn reconnected_count(&self) -> u32 {
|
||||
self.capture.reconnected_count.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
pub fn notification_count(&self) -> u32 {
|
||||
self.capture.notification_count.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn leader_lock_path(home: &Path) -> PathBuf {
|
||||
home.join(".grok").join("leader.lock")
|
||||
}
|
||||
|
||||
pub fn read_leader_pid(home: &Path) -> Option<u32> {
|
||||
std::fs::read_to_string(leader_lock_path(home))
|
||||
.ok()?
|
||||
.trim()
|
||||
.parse()
|
||||
.ok()
|
||||
}
|
||||
|
||||
pub fn pid_alive(pid: u32) -> bool {
|
||||
unsafe { libc::kill(pid as i32, 0) == 0 }
|
||||
}
|
||||
|
||||
/// Wait until the leader lock file contains a live PID, return it.
|
||||
pub async fn wait_for_live_leader(home: &Path, timeout: Duration) -> Option<u32> {
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
if let Some(pid) = read_leader_pid(home)
|
||||
&& pid_alive(pid)
|
||||
{
|
||||
return Some(pid);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Wait until the leader lock file contains a live PID *different* from `old_pid`.
|
||||
pub async fn wait_for_new_leader(home: &Path, old_pid: u32, timeout: Duration) -> Option<u32> {
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
if let Some(pid) = read_leader_pid(home)
|
||||
&& pid != old_pid
|
||||
&& pid_alive(pid)
|
||||
{
|
||||
return Some(pid);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Wait for evidence that the bridge finished its reconnect replay.
|
||||
///
|
||||
/// The `x.ai/leader_reconnected` ext notification is dropped by the typed
|
||||
/// `ClientSideConnection` (bare `x.ai/*` methods are rejected by the ACP
|
||||
/// decoder), so we wait for the replayed `session/load` to emit session
|
||||
/// notifications instead: the notification count rises above `baseline`.
|
||||
pub async fn wait_for_replay_notifications(
|
||||
client: &LeaderStdioClient,
|
||||
baseline: u32,
|
||||
timeout: Duration,
|
||||
) -> bool {
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
if client.reconnected_count() > 0 || client.notification_count() > baseline {
|
||||
return true;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
pub fn leader_log(home: &Path) -> String {
|
||||
std::fs::read_to_string(home.join(".grok").join("leader.log")).unwrap_or_default()
|
||||
}
|
||||
Loading…
Reference in a new issue