Synced from monorepo
Synced from monorepo Changes: - Workspace server: surface preview-proxy metrics through the hub metric pump - Shell: reclaim a session’s retained state in one entry - Shell: reclaim a session’s resident state in one entry - Pager: withhold key event types from Alacritty builds that double keys - Tools: cancel a session’s subagents when it closes - Pager: keep the whole plan in scrollback and separate reasoning from output in minimal mode - Pager: probe terminal version over DA2 and include it with feedback - SuperGrok Plus: identity, CLI, and analytics tier surfaces - Shell: inherit the session process scope into subagents - Pager: build @-file-search matcher lazily on first use - Tools: fix description and output contradictions in tool definitions - Workspace: degrade @-file-search instead of aborting on thread exhaustion - Tools: reap a session’s LSP servers when it closes - Tools: fix contradictions and defects in tool descriptions, schemas, and harness pools - MCP: reap stdio MCP children on session close - Shell: reuse spawn-time skill discovery for session telemetry - Tools: stop leaking shell-wrapper positional params into sourced scripts (fixes activate_conda under persistent/static shell) - Shell: self-heal corrupt session-search SQLite cache - Workspace: cap workspace-server tokio workers on many-core hosts - Shell: reap a session’s child processes when it closes - Crash handler: capture SIGABRT so panic-aborts leave crash reports - CLI chat proxy: team-scoped Grok Code managed-config admin routes - MCP: add CLI enable/disable for MCP servers - Shell: cap tokio worker threads for startup thread demand - Workspace: harden git_commit and add git_sync_base operation - Circuit breaker: add feature-gated gRPC retry policy Source-Revision: 2a818575225183d8ca915f5632a09b8067b5156a
This commit is contained in:
grokkybara[bot]
parent
02d9359435
commit
5da6962e4a
192 changed files with 10337 additions and 3421 deletions
|
|
@ -1163,7 +1163,7 @@ impl acp::Agent for MvpAgent {
|
|||
.await
|
||||
.map_err(|e| crate::session::persistence::io_error_to_acp(&e))?
|
||||
};
|
||||
self.session_turn_numbers.borrow_mut().insert(session_id.clone(), 0u64);
|
||||
self.set_turn_number(&session_id, 0u64);
|
||||
let chat_history = vec![];
|
||||
let client_code_nav_enabled = arguments
|
||||
.meta
|
||||
|
|
@ -1550,9 +1550,7 @@ impl acp::Agent for MvpAgent {
|
|||
let restored_awaiting_plan_approval = persisted_plan_mode
|
||||
.as_ref()
|
||||
.is_some_and(|s| s.awaiting_plan_approval);
|
||||
self.session_turn_numbers
|
||||
.borrow_mut()
|
||||
.insert(session_id.clone(), summary.next_trace_turn);
|
||||
self.set_turn_number(&session_id, summary.next_trace_turn);
|
||||
tracing::info!(
|
||||
session_id = %session_id.0,
|
||||
next_trace_turn = summary.next_trace_turn,
|
||||
|
|
@ -2454,10 +2452,7 @@ impl acp::Agent for MvpAgent {
|
|||
);
|
||||
}
|
||||
let next_trace_turn = self
|
||||
.session_turn_numbers
|
||||
.borrow()
|
||||
.get(&arguments.session_id)
|
||||
.copied()
|
||||
.session_turn_number(&arguments.session_id)
|
||||
.unwrap_or_else(|| turn_number.saturating_add(1));
|
||||
let _ = handle
|
||||
.cmd_tx
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
use super::*;
|
||||
use crate::auth::PreferredAuthMethod;
|
||||
use xai_grok_tools::implementations::grok_build::task::backend::SubagentBackend;
|
||||
use xai_tty_utils::ProcessScope;
|
||||
/// `preferred` model, else catalog `current`, else first with own credentials.
|
||||
fn byok_from_models(
|
||||
models: &indexmap::IndexMap<String, ModelEntry>,
|
||||
|
|
@ -719,7 +720,7 @@ impl MvpAgent {
|
|||
/// Most recently allocated turn number for `sid`, or `None` if the
|
||||
/// session has not started a turn yet.
|
||||
pub(crate) fn session_turn_number(&self, sid: &acp::SessionId) -> Option<u64> {
|
||||
self.session_turn_numbers.borrow().get(sid).copied()
|
||||
self.retained_resources.borrow().get(sid).and_then(|d| d.turn_number)
|
||||
}
|
||||
/// Return the current GrokAuth credentials, if authenticated and not expired.
|
||||
pub(crate) fn current_auth(&self) -> Option<crate::auth::GrokAuth> {
|
||||
|
|
@ -1877,7 +1878,7 @@ impl MvpAgent {
|
|||
sessions: RefCell::new(HashMap::new()),
|
||||
activity,
|
||||
loading_sessions: RefCell::new(HashMap::new()),
|
||||
dispatch_locks: RefCell::new(HashMap::new()),
|
||||
retained_resources: RefCell::new(HashMap::new()),
|
||||
session_threads: RefCell::new(HashMap::new()),
|
||||
resident_roster_titles: RefCell::new(HashMap::new()),
|
||||
initialize_request: OnceLock::new(),
|
||||
|
|
@ -1925,12 +1926,10 @@ impl MvpAgent {
|
|||
relay_sync_enabled,
|
||||
buffering_settings: RefCell::new(None),
|
||||
background_copy_context: BackgroundCopyContext::new(),
|
||||
session_turn_numbers: RefCell::new(HashMap::new()),
|
||||
permission_event_receivers: RefCell::new(HashMap::new()),
|
||||
codebase_indexes: Arc::new(
|
||||
parking_lot::Mutex::new(CodebaseIndexManager::new()),
|
||||
),
|
||||
session_index_claims: RefCell::new(HashMap::new()),
|
||||
resident_resources: RefCell::new(HashMap::new()),
|
||||
worktree_type,
|
||||
restore_code,
|
||||
session_registry_local,
|
||||
|
|
@ -1952,9 +1951,6 @@ impl MvpAgent {
|
|||
std::sync::atomic::AtomicBool::new(false),
|
||||
),
|
||||
workspace_ops: RefCell::new(None),
|
||||
require_gateway_sessions: Rc::new(
|
||||
RefCell::new(std::collections::HashSet::new()),
|
||||
),
|
||||
session_live_state: RefCell::new(HashMap::new()),
|
||||
supervisor_started: std::cell::Cell::new(false),
|
||||
settings_reapply_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)),
|
||||
|
|
@ -2058,9 +2054,8 @@ impl MvpAgent {
|
|||
continue;
|
||||
}
|
||||
self.request_session_shutdown(&id);
|
||||
if self.sessions.borrow_mut().remove(&id).is_some() {
|
||||
self.session_index_claims.borrow_mut().remove(&id);
|
||||
self.require_gateway_sessions.borrow_mut().remove(&id);
|
||||
if self.take_session(&id).is_some() {
|
||||
self.resident_resources.borrow_mut().remove(&id);
|
||||
self.set_session_live_state(&id, SessionLiveState::Dormant);
|
||||
unloaded += 1;
|
||||
tracing::debug!(session_id = %id.0, "idle session unloaded to disk on disconnect");
|
||||
|
|
@ -2886,27 +2881,20 @@ impl MvpAgent {
|
|||
}
|
||||
/// Read a session's next trace turn number without advancing the counter.
|
||||
fn peek_turn_number(&self, session_id: &acp::SessionId) -> u64 {
|
||||
self.session_turn_numbers.borrow().get(session_id).copied().unwrap_or(0u64)
|
||||
self.session_turn_number(session_id).unwrap_or(0u64)
|
||||
}
|
||||
/// Set a session's next trace turn number. The sole writer of the
|
||||
/// `session_turn_numbers` counter, shared by `allocate_turn_number` and the
|
||||
/// batched harness-sibling allocation so both honor the same storage.
|
||||
fn set_turn_number(&self, session_id: &acp::SessionId, next: u64) {
|
||||
self.session_turn_numbers.borrow_mut().insert(session_id.clone(), next);
|
||||
/// Set a session's next trace turn number.
|
||||
pub(super) fn set_turn_number(&self, session_id: &acp::SessionId, next: u64) {
|
||||
self
|
||||
.retained_resources
|
||||
.borrow_mut()
|
||||
.entry(session_id.clone())
|
||||
.or_default()
|
||||
.turn_number = Some(next);
|
||||
}
|
||||
/// Upload each drained harness trace turn (the goal planner at setup, and
|
||||
/// each verifier skeptic panel) as its OWN sibling `turn_{N}` artifact.
|
||||
///
|
||||
/// These phases run inside the single user-facing goal turn but are
|
||||
/// recorded out-of-band (synthetic `task` pairs in a side buffer), so the
|
||||
/// normal per-round `turn_messages.json` never references them. Giving each
|
||||
/// phase its own monotonic turn number — from the SAME `session_turn_numbers`
|
||||
/// counter the model turns use (see [`Self::allocate_turn_number`]), via
|
||||
/// [`Self::get_trace_context`] + [`upload_turn_messages`] — makes the
|
||||
/// subagents discoverable in remote/web clients
|
||||
/// via the `<subagent_result>` footer each synthetic `task` result carries.
|
||||
/// The advanced counter is persisted via `SetNextTraceTurn` so the siblings
|
||||
/// survive a restart. Best-effort and non-blocking.
|
||||
/// Upload each drained harness trace turn as its own `turn_{N}` artifact,
|
||||
/// numbered from the same counter as model turns so subagents interleave
|
||||
/// correctly in remote clients. Best-effort and non-blocking.
|
||||
pub(super) async fn upload_harness_trace_turns(
|
||||
&self,
|
||||
session_id: &acp::SessionId,
|
||||
|
|
@ -3533,6 +3521,7 @@ impl MvpAgent {
|
|||
session_env,
|
||||
)
|
||||
.with_hunk_tracking_enabled(hunk_tracking_enabled);
|
||||
tool_ctx.process_scope = Some(ProcessScope::new());
|
||||
let workspace_ops = self
|
||||
.resolve_workspace_ops()
|
||||
.map_err(|_| {
|
||||
|
|
@ -3761,11 +3750,12 @@ impl MvpAgent {
|
|||
let mgr = std::sync::Arc::new(
|
||||
tokio::sync::Mutex::new(
|
||||
LspManager::new(
|
||||
servers,
|
||||
tool_ctx.cwd.as_path().to_path_buf(),
|
||||
true,
|
||||
xai_grok_tools::notification::ToolNotificationHandle::noop(),
|
||||
),
|
||||
servers,
|
||||
tool_ctx.cwd.as_path().to_path_buf(),
|
||||
true,
|
||||
xai_grok_tools::notification::ToolNotificationHandle::noop(),
|
||||
)
|
||||
.with_process_scope(tool_ctx.process_scope.clone()),
|
||||
),
|
||||
);
|
||||
let adapter = std::sync::Arc::new(LspBackendAdapter::new(mgr));
|
||||
|
|
@ -4153,9 +4143,12 @@ impl MvpAgent {
|
|||
}
|
||||
});
|
||||
}
|
||||
self.permission_event_receivers
|
||||
self
|
||||
.retained_resources
|
||||
.borrow_mut()
|
||||
.insert(session_info.id.clone(), permission_events_rx);
|
||||
.entry(session_info.id.clone())
|
||||
.or_default()
|
||||
.permission_event_receiver = Some(permission_events_rx);
|
||||
if handle_display_cwd.is_some() {
|
||||
handle.display_cwd = handle_display_cwd;
|
||||
}
|
||||
|
|
@ -4167,7 +4160,14 @@ impl MvpAgent {
|
|||
});
|
||||
self.notify_session_cwd_for_watch(std::path::Path::new(&session_info.cwd));
|
||||
self.activity.register_session(&session_info.id.0, &handle);
|
||||
self.sessions.borrow_mut().insert(session_info.id.clone(), handle);
|
||||
if let Some(old) = self
|
||||
.sessions
|
||||
.borrow_mut()
|
||||
.insert(session_info.id.clone(), handle)
|
||||
&& let Some(scope) = &old.tool_context.process_scope
|
||||
{
|
||||
scope.kill_all();
|
||||
}
|
||||
self.spawn_managed_gateway_tool_catalog_fetch();
|
||||
let cwd_for_maintenance = session_info.cwd.clone();
|
||||
tokio::spawn(async move {
|
||||
|
|
@ -4183,10 +4183,10 @@ impl MvpAgent {
|
|||
session_id: &acp::SessionId,
|
||||
) -> Vec<PermissionEvent> {
|
||||
let mut events = Vec::new();
|
||||
if let Some(rx) = self
|
||||
.permission_event_receivers
|
||||
.borrow_mut()
|
||||
let mut retained = self.retained_resources.borrow_mut();
|
||||
if let Some(rx) = retained
|
||||
.get_mut(session_id)
|
||||
.and_then(|d| d.permission_event_receiver.as_mut())
|
||||
{
|
||||
while let Ok(event) = rx.try_recv() {
|
||||
events.push(event);
|
||||
|
|
|
|||
|
|
@ -34,9 +34,11 @@ impl MvpAgent {
|
|||
// Pin the index to the requesting session so the Weak in
|
||||
// CodebaseIndexManager doesn't orphan it immediately.
|
||||
if let Some(sid) = session_id {
|
||||
self.session_index_claims
|
||||
self.resident_resources
|
||||
.borrow_mut()
|
||||
.insert(sid.clone(), std::sync::Arc::clone(&handle));
|
||||
.entry(sid.clone())
|
||||
.or_default()
|
||||
.codebase_index = Some(std::sync::Arc::clone(&handle));
|
||||
}
|
||||
Some((handle, was_newly_started))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -178,7 +178,7 @@ pub(crate) fn jwt_claim_matches_user_subscription_tier(
|
|||
"XPremiumPlus" => jwt_claim == "x_premium_plus",
|
||||
"SuperGrokPro" => jwt_claim == "supergrok_heavy",
|
||||
"SuperGrokLite" => jwt_claim == "supergrok_lite",
|
||||
_ => false,
|
||||
_ => jwt_claim.parse::<u64>().is_ok_and(|n| n != 0),
|
||||
}
|
||||
}
|
||||
fn parse_session_computer_sessions(_meta: Option<&acp::Meta>) -> Option<Vec<()>> {
|
||||
|
|
@ -609,39 +609,40 @@ const SESSION_SUPERVISOR_TICK: std::time::Duration = std::time::Duration::from_m
|
|||
/// actor is between turns and responsive); on timeout we conservatively treat
|
||||
/// the session as busy and keep it resident.
|
||||
const IDLE_QUERY_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(500);
|
||||
/// Per-session state freed on removal or idle-unload (but kept across a reload
|
||||
/// rebuild); retained state instead survives an unload and is freed only at
|
||||
/// removal.
|
||||
#[derive(Default)]
|
||||
struct ResidentResources {
|
||||
/// Strong ref pinning the code-nav index; the manager holds only a `Weak`.
|
||||
codebase_index: Option<std::sync::Arc<xai_codebase_graph::IndexManagerHandle>>,
|
||||
require_gateway: bool,
|
||||
}
|
||||
/// Per-session state that survives an idle-unload (so the session stays
|
||||
/// resumable); freed only at `remove_session`. See [`ResidentResources`].
|
||||
#[derive(Default)]
|
||||
struct RetainedResources {
|
||||
turn_number: Option<u64>,
|
||||
dispatch_lock: Option<std::rc::Rc<tokio::sync::Mutex<()>>>,
|
||||
permission_event_receiver: Option<
|
||||
tokio::sync::mpsc::UnboundedReceiver<PermissionEvent>,
|
||||
>,
|
||||
}
|
||||
pub struct MvpAgent {
|
||||
/// LEADER-SAFE(per-session): keyed by SessionId. Sessions are created/removed
|
||||
/// per client request; no cross-session iteration except cleanup
|
||||
/// (`remove_session`, `sweep_dead_sessions`).
|
||||
/// LEADER-SAFE(per-session). Removed by `remove_session` / `sweep_dead_sessions`.
|
||||
pub(crate) sessions: RefCell<HashMap<acp::SessionId, SessionHandle>>,
|
||||
/// `Send + Sync` mirror of per-session activity (running turn, pending
|
||||
/// interactions, subagent gauge) shared with the leader's auto-update
|
||||
/// checker, which cannot read the `!Send` maps above. Sessions are
|
||||
/// registered at handle creation and expire when their actor exits — no
|
||||
/// unregister bookkeeping. See [`crate::agent::activity::AgentActivity`].
|
||||
/// LEADER-SAFE(shared): `Send + Sync` mirror of per-session activity for the
|
||||
/// leader's auto-update checker, which cannot read the `!Send` maps. Expires
|
||||
/// when the actor exits. See [`crate::agent::activity::AgentActivity`].
|
||||
pub(crate) activity: crate::agent::activity::AgentActivity,
|
||||
/// Sessions with a `session/load` currently in flight. LEADER-SAFE(per-session).
|
||||
///
|
||||
/// Inserted by [`Self::begin_session_load`] at the top of `load_session`
|
||||
/// and removed when the returned RAII guard drops (any exit path). Lets
|
||||
/// racing session-scoped requests — notably `session/prompt` sent right
|
||||
/// behind a reconnect-replayed `session/load` after a leader restart —
|
||||
/// wait for the load via [`Self::wait_for_in_flight_session_load`]
|
||||
/// instead of failing with "unknown session id". The watch channel closes
|
||||
/// when the guard drops, waking all waiters.
|
||||
/// LEADER-SAFE(per-session): in-flight `session/load` guards. Lets a racing
|
||||
/// `session/prompt` wait via [`Self::wait_for_in_flight_session_load`] instead
|
||||
/// of failing "unknown session id"; the RAII guard's drop wakes waiters.
|
||||
loading_sessions: RefCell<
|
||||
HashMap<acp::SessionId, tokio::sync::watch::Receiver<bool>>,
|
||||
>,
|
||||
/// Per-session lock ordering dispatch onto the actor's mailbox:
|
||||
/// [`Self::prompt`] holds it across its intake preamble and
|
||||
/// [`Self::cancel`] around its `Cancel` send, so prompts land in
|
||||
/// submission order and a cancel cannot overtake the prompt it targets
|
||||
/// (see `cancel_never_overtakes_in_flight_prompt_intake`). Cancels wait
|
||||
/// out preambles held ahead of them — keep preambles lean; bridge cancels
|
||||
/// are unordered. LEADER-SAFE(per-session): mirrors `sessions` lifecycle.
|
||||
dispatch_locks: RefCell<
|
||||
HashMap<acp::SessionId, std::rc::Rc<tokio::sync::Mutex<()>>>,
|
||||
>,
|
||||
/// LEADER-SAFE(per-session): reclaimed at `remove_session`. See [`RetainedResources`].
|
||||
retained_resources: RefCell<HashMap<acp::SessionId, RetainedResources>>,
|
||||
/// LEADER-SAFE(per-session): keyed by SessionId. Mirrors `sessions` lifecycle.
|
||||
session_threads: RefCell<HashMap<acp::SessionId, SessionThread>>,
|
||||
/// Title per resident session id, refreshed each `build_roster`. Lets the
|
||||
|
|
@ -758,24 +759,11 @@ pub struct MvpAgent {
|
|||
buffering_settings: RefCell<Option<update_chunk_merge::BufferingSettings>>,
|
||||
/// Context for managing background copy operations (e.g., copying ignored files)
|
||||
pub(crate) background_copy_context: BackgroundCopyContext,
|
||||
/// LEADER-SAFE(per-session): keyed by SessionId, no cross-session iteration.
|
||||
/// Released by `remove_session`.
|
||||
pub(crate) session_turn_numbers: RefCell<HashMap<acp::SessionId, u64>>,
|
||||
/// LEADER-SAFE(per-session): keyed by SessionId, no cross-session iteration.
|
||||
/// Released by `remove_session`.
|
||||
permission_event_receivers: RefCell<
|
||||
HashMap<acp::SessionId, tokio::sync::mpsc::UnboundedReceiver<PermissionEvent>>,
|
||||
>,
|
||||
/// Agent-level codebase index manager for code navigation.
|
||||
/// Indexes are shared across sessions with the same cwd.
|
||||
/// LEADER-SAFE(shared): keyed internally by cwd. No per-client state.
|
||||
/// LEADER-SAFE(shared): agent-level code-nav index manager, keyed by cwd,
|
||||
/// no per-client state.
|
||||
codebase_indexes: Arc<parking_lot::Mutex<CodebaseIndexManager>>,
|
||||
/// Per-session strong refs that keep the code-nav index alive. The
|
||||
/// CodebaseIndexManager holds only Weak; without these the actor would
|
||||
/// be reaped immediately. Cleaned up in remove_session.
|
||||
session_index_claims: RefCell<
|
||||
HashMap<acp::SessionId, std::sync::Arc<xai_codebase_graph::IndexManagerHandle>>,
|
||||
>,
|
||||
/// LEADER-SAFE(per-session): reclaimed on removal / idle-unload. See [`ResidentResources`].
|
||||
resident_resources: RefCell<HashMap<acp::SessionId, ResidentResources>>,
|
||||
/// Worktree creation type (resolved: local config > remote > default Linked).
|
||||
pub(crate) worktree_type: crate::util::config::WorktreeType,
|
||||
/// Restore codebase state on worktree resume (resolved: local config > remote > default false).
|
||||
|
|
@ -869,10 +857,6 @@ pub struct MvpAgent {
|
|||
/// The agent never opens Computer Hub as a harness/client; remote cloud
|
||||
/// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`).
|
||||
workspace_ops: RefCell<Option<xai_grok_workspace::WorkspaceOps>>,
|
||||
/// Sessions opened with `require_gateway` / chat light-frontend (K13).
|
||||
/// Prompt-time guard consults this when the bridge map entry is missing,
|
||||
/// independent of prompt `_meta` (pager often omits kind on prompt).
|
||||
require_gateway_sessions: Rc<RefCell<std::collections::HashSet<acp::SessionId>>>,
|
||||
/// Per-session coarse lifecycle state (residency + turn-state).
|
||||
/// Updated by `spawn_and_register_session` (→ `IdleResident`) and the
|
||||
/// join-handle supervisor on actor exit (→ `DeadFailed`) / explicit close
|
||||
|
|
|
|||
|
|
@ -28,35 +28,57 @@ impl MvpAgent {
|
|||
});
|
||||
}
|
||||
}
|
||||
/// The funnel for a handle leaving `self.sessions` (the spawn-path insert
|
||||
/// reaps any displaced handle the same way): reaps its child-process tree
|
||||
/// on the agent thread, so even a wedged session's tree is reclaimed.
|
||||
///
|
||||
/// The reap is a synchronous SIGKILL issued before the actor drains
|
||||
/// `Shutdown` — deliberately, including on non-terminal idle-unload:
|
||||
/// enrolled children are resident-session-scoped by contract, so graceful
|
||||
/// child teardown belongs to the enrolling owner while the actor is live,
|
||||
/// not to this funnel.
|
||||
pub(super) fn take_session(&self, id: &acp::SessionId) -> Option<SessionHandle> {
|
||||
let handle = self.sessions.borrow_mut().remove(id);
|
||||
if let Some(handle) = &handle
|
||||
&& let Some(scope) = &handle.tool_context.process_scope
|
||||
{
|
||||
scope.kill_all();
|
||||
}
|
||||
handle
|
||||
}
|
||||
/// Remove a session without finalizing; it stays resumable on disk.
|
||||
pub(crate) fn remove_session(&self, id: &acp::SessionId) {
|
||||
self.sessions.borrow_mut().remove(id);
|
||||
self.dispatch_locks.borrow_mut().remove(id);
|
||||
let _ = self
|
||||
.subagent_event_tx
|
||||
.send(xai_grok_tools::implementations::grok_build::task::types::SubagentEvent::TeardownSession {
|
||||
parent_session_id: id.0.to_string(),
|
||||
});
|
||||
self.take_session(id);
|
||||
self.session_threads.borrow_mut().remove(id);
|
||||
self.session_index_claims.borrow_mut().remove(id);
|
||||
self.require_gateway_sessions.borrow_mut().remove(id);
|
||||
self.resident_resources.borrow_mut().remove(id);
|
||||
self.retained_resources.borrow_mut().remove(id);
|
||||
self.model_unavailable_sessions
|
||||
.borrow_mut()
|
||||
.remove(id.0.as_ref());
|
||||
self.permission_event_receivers.borrow_mut().remove(id);
|
||||
self.session_turn_numbers.borrow_mut().remove(id);
|
||||
self.session_live_state.borrow_mut().remove(id);
|
||||
if let Some(ops) = self.workspace_ops.borrow().as_ref() {
|
||||
ops.end_local_session(id.0.as_ref());
|
||||
}
|
||||
let _ = self
|
||||
.subagent_event_tx
|
||||
.send(xai_grok_tools::implementations::grok_build::task::types::SubagentEvent::DiscardSessionCompletions {
|
||||
parent_session_id: id.0.to_string(),
|
||||
});
|
||||
}
|
||||
/// Get-or-create the per-session dispatch lock (see
|
||||
/// [`Self::dispatch_locks`]). Cheap clone of the shared `Rc`.
|
||||
/// Get-or-create the per-session dispatch lock. `prompt` holds it across
|
||||
/// intake so a cancel cannot overtake the prompt it targets.
|
||||
/// Get-or-create the per-session prompt-intake lock. `prompt` holds it across
|
||||
/// its intake preamble and `cancel` around its `Cancel` send, so prompts land
|
||||
/// in submission order and a cancel cannot overtake the prompt it targets.
|
||||
/// Cancels therefore wait behind an intake preamble: keep preambles lean.
|
||||
/// Bridge cancels take their own path and stay unordered against this lock.
|
||||
pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> std::rc::Rc<tokio::sync::Mutex<()>> {
|
||||
self.dispatch_locks
|
||||
self.retained_resources
|
||||
.borrow_mut()
|
||||
.entry(id.clone())
|
||||
.or_default()
|
||||
.dispatch_lock
|
||||
.get_or_insert_with(Default::default)
|
||||
.clone()
|
||||
}
|
||||
/// Close a session in response to an **explicit** terminal close
|
||||
|
|
@ -416,16 +438,40 @@ impl MvpAgent {
|
|||
)
|
||||
.registry_counts()
|
||||
.await;
|
||||
let (session_index_claims, require_gateway_sessions) = {
|
||||
let resident = self.resident_resources.borrow();
|
||||
(
|
||||
resident
|
||||
.values()
|
||||
.filter(|r| r.codebase_index.is_some())
|
||||
.count(),
|
||||
resident.values().filter(|r| r.require_gateway).count(),
|
||||
)
|
||||
};
|
||||
let retained = self.retained_resources.borrow();
|
||||
let dispatch_locks = retained
|
||||
.values()
|
||||
.filter(|d| d.dispatch_lock.is_some())
|
||||
.count();
|
||||
let session_turn_numbers = retained
|
||||
.values()
|
||||
.filter(|d| d.turn_number.is_some())
|
||||
.count();
|
||||
let permission_event_receivers = retained
|
||||
.values()
|
||||
.filter(|d| d.permission_event_receiver.is_some())
|
||||
.count();
|
||||
drop(retained);
|
||||
RegistrySnapshot {
|
||||
sessions: self.sessions.borrow().len(),
|
||||
session_threads: self.session_threads.borrow().len(),
|
||||
dispatch_locks: self.dispatch_locks.borrow().len(),
|
||||
session_turn_numbers: self.session_turn_numbers.borrow().len(),
|
||||
permission_event_receivers: self.permission_event_receivers.borrow().len(),
|
||||
dispatch_locks,
|
||||
session_turn_numbers,
|
||||
permission_event_receivers,
|
||||
model_unavailable_sessions: self.model_unavailable_sessions.borrow().len(),
|
||||
session_live_state: self.session_live_state.borrow().len(),
|
||||
session_index_claims: self.session_index_claims.borrow().len(),
|
||||
require_gateway_sessions: self.require_gateway_sessions.borrow().len(),
|
||||
session_index_claims,
|
||||
require_gateway_sessions,
|
||||
subagent_pending: subagents.pending,
|
||||
subagent_active: subagents.active,
|
||||
subagent_completed: subagents.completed,
|
||||
|
|
|
|||
|
|
@ -326,11 +326,13 @@ impl MvpAgent {
|
|||
})
|
||||
}?;
|
||||
let available_models = self.models_manager.models();
|
||||
let parent_lsp = {
|
||||
let (parent_lsp, parent_process_scope) = {
|
||||
let sessions = self.sessions.borrow();
|
||||
sessions
|
||||
.get(&parent_sid)
|
||||
.and_then(|h| h.tool_context.lsp.clone())
|
||||
let parent = sessions.get(&parent_sid);
|
||||
(
|
||||
parent.and_then(|h| h.tool_context.lsp.clone()),
|
||||
parent.and_then(|h| h.tool_context.process_scope.clone()),
|
||||
)
|
||||
};
|
||||
let am = self.auth_manager.clone();
|
||||
let inference_idle_timeout_secs = {
|
||||
|
|
@ -407,6 +409,7 @@ impl MvpAgent {
|
|||
};
|
||||
Some(crate::agent::subagent::SubagentSpawnContext {
|
||||
lsp: parent_lsp,
|
||||
process_scope: parent_process_scope,
|
||||
client_hooks: Default::default(),
|
||||
sampling_config: self.sampling_config.borrow().clone(),
|
||||
managed_mcp_proxy_base_url: parent_managed_mcp_proxy_base_url
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ fn jwt_tier_claim_maps_free_and_paid() {
|
|||
jwt_tier_claim(&jwt_with_tier(6)).as_deref(),
|
||||
Some("supergrok_lite")
|
||||
);
|
||||
assert_eq!(jwt_tier_claim(&jwt_with_tier(9)).as_deref(), Some("9"));
|
||||
assert_eq!(jwt_tier_claim(&jwt_with_tier(99)).as_deref(), Some("99"));
|
||||
}
|
||||
fn auth_with_mode(mode: crate::auth::AuthMode, key: &str) -> crate::auth::GrokAuth {
|
||||
|
|
@ -99,6 +100,7 @@ fn jwt_claim_matches_user_subscription_tier_known_pairs() {
|
|||
("x_premium", "XPremium"),
|
||||
("x_premium_plus", "XPremiumPlus"),
|
||||
("supergrok_heavy", "SuperGrokPro"),
|
||||
("9", "EnterpriseMystery"),
|
||||
("supergrok_lite", "SuperGrokLite"),
|
||||
];
|
||||
for (claim, user_tier) in cases {
|
||||
|
|
@ -124,6 +126,10 @@ fn jwt_claim_matches_user_subscription_tier_rejects_stale_and_unknown() {
|
|||
"supergrok_heavy",
|
||||
"EnterpriseMystery"
|
||||
));
|
||||
assert!(!jwt_claim_matches_user_subscription_tier(
|
||||
"0",
|
||||
"EnterpriseMystery"
|
||||
));
|
||||
}
|
||||
/// Single-flight flag must clear on Drop even if the retry task panics /
|
||||
/// aborts mid-backoff (guards against the flag stuck true forever).
|
||||
|
|
@ -390,8 +396,8 @@ fn settings_allow_access_field_absent_is_blocked() {
|
|||
};
|
||||
assert!(!settings_allow_access(Some(&rs)));
|
||||
}
|
||||
/// After allocating a turn number, `session_turn_numbers` holds the next
|
||||
/// value (current + 1). This is the value that must be persisted via
|
||||
/// After allocating a turn number, the retained (in-memory) turn counter holds
|
||||
/// the next value (current + 1). This is the value that must be persisted via
|
||||
/// `SetNextTraceTurn` so the counter survives restarts.
|
||||
#[test]
|
||||
fn allocate_turn_number_advances_counter() {
|
||||
|
|
@ -1809,6 +1815,8 @@ async fn ensure_plugin_registry_lazily_populates_snapshot() {
|
|||
"repeat call must keep the populated snapshot"
|
||||
);
|
||||
}
|
||||
#[cfg(unix)]
|
||||
mod process_scope_reclaim;
|
||||
mod subagent_spawn_context_tests;
|
||||
/// No load in flight and no session → the wait returns immediately
|
||||
/// (the caller then surfaces "unknown session id" exactly as before).
|
||||
|
|
@ -3143,16 +3151,21 @@ async fn remove_session_releases_workspace_binding_and_side_maps() {
|
|||
sid.0.to_string(),
|
||||
acp::ModelId::new(std::sync::Arc::from("gone-model")),
|
||||
);
|
||||
agent
|
||||
.session_turn_numbers
|
||||
.borrow_mut()
|
||||
.insert(sid.clone(), 3);
|
||||
agent.set_turn_number(&sid, 3);
|
||||
let (_permission_tx, permission_rx) =
|
||||
tokio::sync::mpsc::unbounded_channel::<xai_grok_workspace::permission::PermissionEvent>();
|
||||
agent
|
||||
.permission_event_receivers
|
||||
.retained_resources
|
||||
.borrow_mut()
|
||||
.insert(sid.clone(), permission_rx);
|
||||
.entry(sid.clone())
|
||||
.or_default()
|
||||
.permission_event_receiver = Some(permission_rx);
|
||||
agent
|
||||
.resident_resources
|
||||
.borrow_mut()
|
||||
.entry(sid.clone())
|
||||
.or_default()
|
||||
.require_gateway = true;
|
||||
agent.remove_session(&sid);
|
||||
assert!(
|
||||
toolset_weak.upgrade().is_none(),
|
||||
|
|
@ -3164,8 +3177,11 @@ async fn remove_session_releases_workspace_binding_and_side_maps() {
|
|||
.borrow()
|
||||
.contains_key(sid.0.as_ref())
|
||||
);
|
||||
assert!(!agent.session_turn_numbers.borrow().contains_key(&sid));
|
||||
assert!(!agent.permission_event_receivers.borrow().contains_key(&sid));
|
||||
assert!(!agent.resident_resources.borrow().contains_key(&sid));
|
||||
assert!(
|
||||
!agent.retained_resources.borrow().contains_key(&sid),
|
||||
"retained per-session resources must be reclaimed on removal"
|
||||
);
|
||||
}
|
||||
/// Without a bridge, `ext_method` falls through to the unchanged local
|
||||
/// dispatch (`rewind::handle`), which reports the missing session — proving
|
||||
|
|
@ -3213,7 +3229,7 @@ fn cancel_does_not_forward_to_bridge_in_local_mode() {
|
|||
});
|
||||
}
|
||||
/// Regression (post-cancel slot hang, first bad release 0.2.101; see
|
||||
/// `dispatch_locks`). SDK e2e shape:
|
||||
/// `dispatch_lock`). SDK e2e shape:
|
||||
/// `test_cancel_ends_in_flight_turn_and_frees_slot` (grok-agent-sdk).
|
||||
#[test]
|
||||
fn cancel_never_overtakes_in_flight_prompt_intake() {
|
||||
|
|
|
|||
|
|
@ -38,18 +38,16 @@ fn populate_and_evict(agent: &MvpAgent, i: usize) {
|
|||
|
||||
let (_ptx, prx) = tokio::sync::mpsc::unbounded_channel::<PermissionEvent>();
|
||||
agent
|
||||
.permission_event_receivers
|
||||
.retained_resources
|
||||
.borrow_mut()
|
||||
.insert(sid.clone(), prx);
|
||||
agent
|
||||
.session_turn_numbers
|
||||
.borrow_mut()
|
||||
.insert(sid.clone(), i as u64);
|
||||
.entry(sid.clone())
|
||||
.or_default()
|
||||
.permission_event_receiver = Some(prx);
|
||||
agent.set_turn_number(&sid, i as u64);
|
||||
agent.model_unavailable_sessions.borrow_mut().insert(
|
||||
sid.0.to_string(),
|
||||
acp::ModelId::new(std::sync::Arc::from("gone-model")),
|
||||
);
|
||||
|
||||
agent.remove_session(&sid);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,85 @@
|
|||
//! Per-session process-tree reclaim on close, and cross-session isolation.
|
||||
//!
|
||||
//! Closing a session reaps the children enrolled in its scope; the reap runs on
|
||||
//! the agent thread via `take_session`, so it works even if the actor wedged.
|
||||
//! End-to-end reaping of specific subsystems is covered by their own soak tests.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{build_minimal_agent_for_tests, make_test_handle, run_local_for_bridge_test};
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
fn sleeper() -> tokio::process::Command {
|
||||
let mut c = tokio::process::Command::new("sleep");
|
||||
c.arg("600");
|
||||
c
|
||||
}
|
||||
|
||||
async fn died(child: &mut tokio::process::Child) -> bool {
|
||||
tokio::time::timeout(Duration::from_secs(5), child.wait())
|
||||
.await
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
async fn still_running(child: &mut tokio::process::Child) -> bool {
|
||||
tokio::time::timeout(Duration::from_millis(500), child.wait())
|
||||
.await
|
||||
.is_err()
|
||||
}
|
||||
|
||||
/// Register a session handle carrying `scope`, keyed by `sid`.
|
||||
fn insert_session_with_scope(
|
||||
agent: &super::MvpAgent,
|
||||
sid: &acp::SessionId,
|
||||
scope: xai_tty_utils::ProcessScope,
|
||||
) {
|
||||
let mut handle = make_test_handle("test", false, None);
|
||||
handle.info.id = sid.clone();
|
||||
handle.tool_context.process_scope = Some(scope);
|
||||
agent.sessions.borrow_mut().insert(sid.clone(), handle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn close_reaps_enrolled_session_child() {
|
||||
run_local_for_bridge_test(|| async {
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
let sid = acp::SessionId::new("sess-scope-reclaim");
|
||||
let scope = xai_tty_utils::ProcessScope::new();
|
||||
insert_session_with_scope(&agent, &sid, scope.clone());
|
||||
let (mut child, _owner) = scope.spawn(sleeper()).expect("spawn enrolled child");
|
||||
|
||||
agent.remove_session(&sid);
|
||||
|
||||
assert!(
|
||||
!agent.sessions.borrow().contains_key(&sid),
|
||||
"close must remove the session handle"
|
||||
);
|
||||
assert!(died(&mut child).await, "close must reap the enrolled child");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn close_is_isolated_across_sessions() {
|
||||
run_local_for_bridge_test(|| async {
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
let sid_a = acp::SessionId::new("sess-A");
|
||||
let sid_b = acp::SessionId::new("sess-B");
|
||||
let scope_a = xai_tty_utils::ProcessScope::new();
|
||||
let scope_b = xai_tty_utils::ProcessScope::new();
|
||||
insert_session_with_scope(&agent, &sid_a, scope_a.clone());
|
||||
insert_session_with_scope(&agent, &sid_b, scope_b.clone());
|
||||
let (mut child_a, _owner_a) = scope_a.spawn(sleeper()).expect("spawn A child");
|
||||
let (mut child_b, _owner_b) = scope_b.spawn(sleeper()).expect("spawn B child");
|
||||
|
||||
agent.remove_session(&sid_a);
|
||||
|
||||
assert!(died(&mut child_a).await, "closing A must reap A's child");
|
||||
assert!(
|
||||
still_running(&mut child_b).await,
|
||||
"closing A must not touch B's child"
|
||||
);
|
||||
|
||||
agent.remove_session(&sid_b);
|
||||
assert!(died(&mut child_b).await, "closing B must reap B's child");
|
||||
});
|
||||
}
|
||||
|
|
@ -173,3 +173,30 @@ async fn subagent_spawn_context_inherits_parent_configured_cutoff() {
|
|||
"an unbounded parent must not hand a subagent a cutoff"
|
||||
);
|
||||
}
|
||||
|
||||
/// A subagent inherits the parent's `process_scope`, so an owner enrolled through it stays visible via the child.
|
||||
/// End-to-end reaping is covered by the spine's `process_scope_reclaim` tests.
|
||||
#[tokio::test]
|
||||
async fn subagent_spawn_context_inherits_parent_process_scope() {
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
let sid = acp::SessionId::new("parent-process-scope");
|
||||
let mut handle = make_test_handle("test-model", false, None);
|
||||
let parent_scope = xai_tty_utils::ProcessScope::new();
|
||||
handle.tool_context.process_scope = Some(parent_scope.clone());
|
||||
agent.sessions.borrow_mut().insert(sid.clone(), handle);
|
||||
|
||||
// Hold an owner Arc in the parent scope so live_count == 1.
|
||||
let owner = std::sync::Arc::new(xai_tty_utils::ProcessGroup::new().expect("process group"));
|
||||
parent_scope.register(&owner);
|
||||
|
||||
let ctx = agent.build_subagent_spawn_context(sid.0.as_ref());
|
||||
let inherited = ctx
|
||||
.process_scope
|
||||
.expect("subagent context must inherit the parent's process scope");
|
||||
|
||||
assert_eq!(
|
||||
inherited.live_count(),
|
||||
1,
|
||||
"the child sees the owner enrolled through the parent scope"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -720,6 +720,7 @@ pub(crate) async fn run_shell_child(
|
|||
tool_ctx.monitor_event_buffer = Some(MonitorEventBuffer::default());
|
||||
tool_ctx.subagent_depth = child_depth;
|
||||
tool_ctx.lsp = ctx.lsp.clone();
|
||||
tool_ctx.process_scope = ctx.process_scope.clone();
|
||||
let parent_traceparent = xai_file_utils::trace_context::current_traceparent();
|
||||
let tracker_child_cwd = child_session_info.cwd.clone();
|
||||
let tracker_model_id = effective_model_id.0.to_string();
|
||||
|
|
|
|||
|
|
@ -103,6 +103,12 @@ impl AutoCompactThresholdTiers {
|
|||
pub(crate) struct SubagentSpawnContext {
|
||||
/// Parent's LSP runtime — inherited via ToolContext, same as fs/terminal.
|
||||
pub lsp: Option<std::sync::Arc<dyn xai_grok_tools::implementations::lsp::LspBackend>>,
|
||||
/// Root session's process scope, inherited so the subagent's own child
|
||||
/// processes are reaped when the parent session closes. It is the root's
|
||||
/// (not an intermediate parent's) because xai-grok-tools task/coordinator.rs
|
||||
/// `handle_command`'s Spawn arm re-parents nested Spawn requests to the root
|
||||
/// parent, so every subagent resolves back to the root session.
|
||||
pub process_scope: Option<xai_tty_utils::ProcessScope>,
|
||||
/// Parent's client-registered hooks, inherited so the subagent's tool calls hit the
|
||||
/// same PreToolUse gate and its events fire the same observe hooks over the parent's
|
||||
/// connection. Empty when the parent has none. Filled by the coordinator after the
|
||||
|
|
|
|||
|
|
@ -16,17 +16,15 @@ use crate::auth::manager::RefreshReason;
|
|||
use crate::auth::token_type::TokenType;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
/// Subscription tiers that qualify for Grok Build access.
|
||||
/// Any active subscription qualifies -- the access gate in remote settings
|
||||
/// controls which tiers are actually allowed.
|
||||
const QUALIFYING_TIERS: &[&str] = &[
|
||||
"SuperGrokPro",
|
||||
"GrokPro",
|
||||
"SuperGrokLite",
|
||||
"XPremiumPlus",
|
||||
"XPremium",
|
||||
"XBasic",
|
||||
];
|
||||
/// Whether a `/user?include=subscription` tier qualifies for Grok Build
|
||||
/// access. Any active subscription qualifies -- the proxy only returns a
|
||||
/// tier when an active subscription exists (`None` otherwise), and the
|
||||
/// access gate in remote settings controls which tiers are actually
|
||||
/// allowed. The `"Free"` guard is defense-in-depth should the proxy ever
|
||||
/// start stamping free users explicitly.
|
||||
fn is_qualifying_tier(tier: &str) -> bool {
|
||||
!tier.is_empty() && tier != "Free"
|
||||
}
|
||||
/// Successful subscription check result: a confirmed qualifying tier.
|
||||
pub(crate) struct UnblockResult {
|
||||
pub(crate) new_tier: String,
|
||||
|
|
@ -110,7 +108,7 @@ pub(crate) async fn single_check(
|
|||
Some(tier) if !tier.is_empty() => tier.clone(),
|
||||
_ => return None,
|
||||
};
|
||||
if !QUALIFYING_TIERS.contains(&new_tier.as_str()) {
|
||||
if !is_qualifying_tier(&new_tier) {
|
||||
return None;
|
||||
}
|
||||
xai_grok_telemetry::unified_log::info(
|
||||
|
|
@ -146,7 +144,7 @@ pub(crate) async fn single_check(
|
|||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn qualifying_tiers_includes_all_paid_tiers() {
|
||||
fn all_paid_tiers_qualify() {
|
||||
for tier in &[
|
||||
"SuperGrokPro",
|
||||
"GrokPro",
|
||||
|
|
@ -155,26 +153,12 @@ mod tests {
|
|||
"XPremium",
|
||||
"XBasic",
|
||||
] {
|
||||
assert!(
|
||||
QUALIFYING_TIERS.contains(tier),
|
||||
"{tier} must be in QUALIFYING_TIERS"
|
||||
);
|
||||
assert!(is_qualifying_tier(tier), "{tier} must qualify");
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn free_tier_is_not_qualifying() {
|
||||
assert!(!QUALIFYING_TIERS.contains(&"Free"));
|
||||
}
|
||||
#[test]
|
||||
fn empty_tier_is_not_qualifying() {
|
||||
assert!(!QUALIFYING_TIERS.contains(&""));
|
||||
}
|
||||
/// The subscription check only returns `Some` when `/user` reports a
|
||||
/// qualifying tier. Verify the tier matching is exact (no prefix match).
|
||||
#[test]
|
||||
fn partial_tier_name_is_not_qualifying() {
|
||||
assert!(!QUALIFYING_TIERS.contains(&"Super"));
|
||||
assert!(!QUALIFYING_TIERS.contains(&"Grok"));
|
||||
assert!(!QUALIFYING_TIERS.contains(&"XPremium+"));
|
||||
fn free_and_empty_tiers_are_not_qualifying() {
|
||||
assert!(!is_qualifying_tier("Free"));
|
||||
assert!(!is_qualifying_tier(""));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -539,6 +539,7 @@ pub async fn handle(
|
|||
signoff: req.signoff,
|
||||
push: req.push,
|
||||
sync: req.sync,
|
||||
..Default::default()
|
||||
};
|
||||
let commit_result = ops
|
||||
.dispatch(&op, None)
|
||||
|
|
|
|||
|
|
@ -785,18 +785,12 @@ pub async fn init_agent_mcp_pool(mcp_state: &Arc<TokioMutex<McpState>>, cwd: &st
|
|||
}
|
||||
|
||||
let noop = xai_file_utils::events::EventWriter::noop();
|
||||
let results = start_mcp_servers(
|
||||
configs,
|
||||
None,
|
||||
Some(cwd),
|
||||
&Default::default(),
|
||||
&Default::default(),
|
||||
&noop,
|
||||
// Pass Interactive to preserve prior deferred-OAuth behavior. A session-less SDK agent can
|
||||
// reach this non-interactively; threading real non-interactivity here is a deliberate follow-up.
|
||||
crate::session::mcp_servers::OauthInteractivity::Interactive,
|
||||
)
|
||||
.await;
|
||||
// session_less picks Interactive to preserve prior deferred-OAuth behavior. A session-less SDK
|
||||
// agent can reach this non-interactively; threading real non-interactivity is a deliberate follow-up.
|
||||
let ctx = crate::session::mcp_servers::McpSpawnCtx::session_less(&noop);
|
||||
let meta = Default::default();
|
||||
let oauth = Default::default();
|
||||
let results = start_mcp_servers(configs, Some(cwd), &meta, &oauth, &ctx).await;
|
||||
let clients: HashMap<McpServerName, Arc<McpClient>> = results
|
||||
.into_iter()
|
||||
.filter_map(|r| match r {
|
||||
|
|
@ -1731,7 +1725,7 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
|||
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
|
||||
if let Some(connector_id) = gateway_connector_id {
|
||||
if let Err(e) =
|
||||
crate::util::config::save_mcp_server_enabled(&req.server_name, true).await
|
||||
crate::util::config::save_mcp_server_enabled_in(&req.server_name, true, &cwd).await
|
||||
{
|
||||
tracing::warn!(
|
||||
server = req.server_name.as_str(),
|
||||
|
|
@ -1752,7 +1746,9 @@ async fn handle_toggle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
|||
crate::session::managed_mcp::invalidate_cache(agent.managed_mcp_cache()).await;
|
||||
}
|
||||
let managed_configs = agent.get_managed_mcp_configs().await;
|
||||
if let Err(e) = crate::util::config::save_mcp_server_enabled(&req.server_name, true).await {
|
||||
if let Err(e) =
|
||||
crate::util::config::save_mcp_server_enabled_in(&req.server_name, true, &cwd).await
|
||||
{
|
||||
tracing::warn!(
|
||||
server = req.server_name.as_str(),
|
||||
error = %e,
|
||||
|
|
@ -1933,8 +1929,8 @@ async fn handle_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
|||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
// The toggle path spawns a task that adds the server to
|
||||
// `disabled_mcp_servers`. Clean that up since we're deleting entirely.
|
||||
let _ = crate::util::config::save_mcp_server_enabled(&req.server_name, true).await;
|
||||
// `disabled_mcp_servers`. Clear user list only — do not unstick project.
|
||||
let _ = crate::util::config::save_user_mcp_server_enabled(&req.server_name, true).await;
|
||||
|
||||
to_ext_response(Ok(McpToggleResponse { ok: true }))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -363,17 +363,8 @@ async fn check_server_start(
|
|||
) -> Result<(mcp_servers::McpClient, Check), Check> {
|
||||
let start = std::time::Instant::now();
|
||||
let noop = xai_file_utils::events::EventWriter::noop();
|
||||
match mcp_servers::start_mcp_server(
|
||||
acp_server,
|
||||
None,
|
||||
Some(cwd),
|
||||
None,
|
||||
None,
|
||||
&noop,
|
||||
mcp_servers::OauthInteractivity::Interactive,
|
||||
)
|
||||
.await
|
||||
{
|
||||
let ctx = mcp_servers::McpSpawnCtx::session_less(&noop);
|
||||
match mcp_servers::start_mcp_server(acp_server, Some(cwd), None, None, &ctx).await {
|
||||
Ok(client) => {
|
||||
let elapsed = start.elapsed();
|
||||
Ok((
|
||||
|
|
|
|||
|
|
@ -481,15 +481,18 @@ impl SessionActor {
|
|||
crate::util::config::load_mcp_servers_with_oauth(cwd, &self.rebuild_spec.compat);
|
||||
let byo_config = oauth_config_map.get(server_name).cloned();
|
||||
let event_writer = self.events.writer();
|
||||
let mode = crate::session::mcp_servers::OauthInteractivity::Interactive;
|
||||
let ctx = crate::session::mcp_servers::McpSpawnCtx::for_session(
|
||||
session_id,
|
||||
&event_writer,
|
||||
crate::session::mcp_servers::OauthInteractivity::Interactive,
|
||||
self.tool_context.process_scope.as_ref(),
|
||||
);
|
||||
let new_client = crate::session::mcp_servers::start_mcp_server(
|
||||
server_config,
|
||||
Some(session_id),
|
||||
Some(cwd),
|
||||
meta_config.as_ref(),
|
||||
byo_config.as_ref(),
|
||||
&event_writer,
|
||||
mode,
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to prepare OAuth for '{}': {}", server_name, e))?;
|
||||
|
|
@ -973,15 +976,18 @@ impl SessionActor {
|
|||
crate::util::config::load_mcp_servers_with_oauth(cwd, &self.rebuild_spec.compat);
|
||||
let byo_config = oauth_config_map.get(server).cloned();
|
||||
let event_writer = self.events.writer();
|
||||
let mode = OauthInteractivity::from_non_interactive(self.startup_hints.non_interactive);
|
||||
let ctx = crate::session::mcp_servers::McpSpawnCtx::for_session(
|
||||
session_id,
|
||||
&event_writer,
|
||||
OauthInteractivity::from_non_interactive(self.startup_hints.non_interactive),
|
||||
self.tool_context.process_scope.as_ref(),
|
||||
);
|
||||
let new_client = crate::session::mcp_servers::start_mcp_server(
|
||||
server_config.clone(),
|
||||
Some(session_id),
|
||||
Some(cwd),
|
||||
meta_config.as_ref(),
|
||||
byo_config.as_ref(),
|
||||
&event_writer,
|
||||
mode,
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
|
@ -1242,16 +1248,19 @@ impl SessionActor {
|
|||
&toml_mcp_names,
|
||||
);
|
||||
let spawn_writer = self.events.writer();
|
||||
let mode = OauthInteractivity::from_non_interactive(self.startup_hints.non_interactive);
|
||||
let ctx = crate::session::mcp_servers::McpSpawnCtx::for_session(
|
||||
session_id,
|
||||
&spawn_writer,
|
||||
OauthInteractivity::from_non_interactive(self.startup_hints.non_interactive),
|
||||
self.tool_context.process_scope.as_ref(),
|
||||
);
|
||||
let mcp_results = build_pending_clients(
|
||||
&self.mcp_state,
|
||||
configs_to_start,
|
||||
Some(session_id),
|
||||
Some(cwd),
|
||||
&meta_config_map,
|
||||
&oauth_config_map,
|
||||
&spawn_writer,
|
||||
mode,
|
||||
&ctx,
|
||||
)
|
||||
.await;
|
||||
tokio::task::yield_now().await;
|
||||
|
|
|
|||
|
|
@ -1446,11 +1446,16 @@ pub(super) async fn run_session(
|
|||
|
||||
let session_for_mcp = session.clone();
|
||||
let sname = server_name.clone();
|
||||
let session_cwd = session.session_info.cwd.clone();
|
||||
tokio::task::spawn_local(async move {
|
||||
session_for_mcp.ensure_mcp_tools_initialized().await;
|
||||
if let Err(e) = crate::util::config::save_mcp_server_enabled(
|
||||
&sname, enabled,
|
||||
).await {
|
||||
if let Err(e) = crate::util::config::save_mcp_server_enabled_in(
|
||||
&sname,
|
||||
enabled,
|
||||
std::path::Path::new(&session_cwd),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
server = sname.as_str(),
|
||||
error = %e,
|
||||
|
|
|
|||
|
|
@ -730,47 +730,8 @@ pub(crate) async fn spawn_session_actor(
|
|||
};
|
||||
let bridge_state_path =
|
||||
crate::session::persistence::session_dir(&session_info).join("tool_state.json");
|
||||
let initial_agent_type = Some(agent_definition.name.clone());
|
||||
let harness_metrics = if telemetry_enabled || xai_grok_telemetry::external::is_active() {
|
||||
let plugin_names = plugin_registry
|
||||
.as_ref()
|
||||
.map(|reg| {
|
||||
reg.active_plugins()
|
||||
.iter()
|
||||
.map(|p| p.name.clone())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
Some(super::telemetry::SessionHarnessMetrics {
|
||||
session_id: session_info.id.0.to_string(),
|
||||
client_identifier: session_client_identifier.clone(),
|
||||
model_id: session_model_id.0.to_string(),
|
||||
agent_name: agent_definition.name.clone(),
|
||||
permission_mode: if session_yolo_mode {
|
||||
xai_grok_telemetry::enums::PermissionMode::AlwaysApprove
|
||||
} else if session_auto_mode
|
||||
&& crate::util::config::auto_permission_mode_enabled_from_disk()
|
||||
{
|
||||
xai_grok_telemetry::enums::PermissionMode::Auto
|
||||
} else {
|
||||
xai_grok_telemetry::enums::PermissionMode::Ask
|
||||
},
|
||||
mcp_server_names: mcp_servers
|
||||
.iter()
|
||||
.map(|s| mcp_server_name(s).to_owned())
|
||||
.collect(),
|
||||
lsp_server_names: tool_context.lsp_server_names.clone(),
|
||||
memory_enabled: memory_config.is_some(),
|
||||
auto_update,
|
||||
cwd: tool_context.cwd.as_str().to_owned(),
|
||||
skills_config: skills_config.clone(),
|
||||
compat,
|
||||
plugin_registry: plugin_registry.clone(),
|
||||
plugin_names,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let initial_agent_name = agent_definition.name.clone();
|
||||
let initial_agent_type = Some(initial_agent_name.clone());
|
||||
let compaction_policy = xai_grok_agent::CompactionPolicy {
|
||||
auto_compact_threshold_percent: auto_compact_threshold_percent as u32,
|
||||
compact_model: None,
|
||||
|
|
@ -1052,6 +1013,46 @@ pub(crate) async fn spawn_session_actor(
|
|||
.tool_bridge()
|
||||
.update_resource(task_wake_suppressed)
|
||||
.await;
|
||||
let harness_metrics = if telemetry_enabled || xai_grok_telemetry::external::is_active() {
|
||||
let plugin_names = plugin_registry
|
||||
.as_ref()
|
||||
.map(|reg| {
|
||||
reg.active_plugins()
|
||||
.iter()
|
||||
.map(|p| p.name.clone())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
Some(super::telemetry::SessionHarnessMetrics {
|
||||
session_id: session_info.id.0.to_string(),
|
||||
client_identifier: session_client_identifier.clone(),
|
||||
model_id: session_model_id.0.to_string(),
|
||||
agent_name: initial_agent_name,
|
||||
permission_mode: if session_yolo_mode {
|
||||
xai_grok_telemetry::enums::PermissionMode::AlwaysApprove
|
||||
} else if session_auto_mode
|
||||
&& crate::util::config::auto_permission_mode_enabled_from_disk()
|
||||
{
|
||||
xai_grok_telemetry::enums::PermissionMode::Auto
|
||||
} else {
|
||||
xai_grok_telemetry::enums::PermissionMode::Ask
|
||||
},
|
||||
mcp_server_names: mcp_servers
|
||||
.iter()
|
||||
.map(|s| mcp_server_name(s).to_owned())
|
||||
.collect(),
|
||||
lsp_server_names: tool_context.lsp_server_names.clone(),
|
||||
memory_enabled: memory_config.is_some(),
|
||||
auto_update,
|
||||
cwd: tool_context.cwd.as_str().to_owned(),
|
||||
skill_names: agent.tool_bridge().skill_discovery_snapshot_names().await,
|
||||
compat,
|
||||
plugin_registry: plugin_registry.clone(),
|
||||
plugin_names,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let resolved_task_output =
|
||||
xai_grok_tools::reminders::task_completion::resolve_task_output_tool_name(
|
||||
agent.tool_bridge(),
|
||||
|
|
|
|||
|
|
@ -3,9 +3,9 @@
|
|||
pub use xai_grok_mcp::servers::{
|
||||
AcpServerEntry, HttpConfig, MCP_TOOL_NAME_DELIMITER, McpClient, McpClientTimeoutOverrides,
|
||||
McpConfigDiff, McpError, McpInitStrategy, McpMetaConfigMap, McpServerMetaConfig, McpServerName,
|
||||
McpService, McpState, McpTool, McpToolRegistration, OauthInteractivity, SharedMcpPool,
|
||||
mcp_server_name, mcp_target_str, mcp_transport_str, parse_mcp_meta_config, parse_mcp_tool_name,
|
||||
sanitize_descriptor_segment, validate_tool_name,
|
||||
McpService, McpSpawnCtx, McpState, McpTool, McpToolRegistration, OauthInteractivity,
|
||||
SharedMcpPool, mcp_server_name, mcp_target_str, mcp_transport_str, parse_mcp_meta_config,
|
||||
parse_mcp_tool_name, sanitize_descriptor_segment, validate_tool_name,
|
||||
};
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
|
@ -65,24 +65,13 @@ pub fn build_config_resolved_event(
|
|||
|
||||
pub async fn start_mcp_server(
|
||||
mcp_server: acp::McpServer,
|
||||
session_id: Option<&str>,
|
||||
cwd: Option<&Path>,
|
||||
meta_config: Option<&inner::McpServerMetaConfig>,
|
||||
byo_config: Option<&McpOAuthConfig>,
|
||||
event_writer: &xai_file_utils::events::EventWriter,
|
||||
mode: OauthInteractivity,
|
||||
ctx: &inner::McpSpawnCtx<'_>,
|
||||
) -> Result<inner::McpClient, inner::McpError> {
|
||||
let overrides = resolve_overrides(inner::mcp_server_name(&mcp_server), cwd);
|
||||
inner::start_mcp_server(
|
||||
mcp_server,
|
||||
session_id,
|
||||
overrides.as_ref(),
|
||||
meta_config,
|
||||
byo_config,
|
||||
event_writer,
|
||||
mode,
|
||||
)
|
||||
.await
|
||||
inner::start_mcp_server(mcp_server, overrides.as_ref(), meta_config, byo_config, ctx).await
|
||||
}
|
||||
|
||||
/// Build all pending MCP clients for one init pass as a single merged list: config-declared
|
||||
|
|
@ -93,21 +82,17 @@ pub async fn start_mcp_server(
|
|||
pub async fn build_pending_clients(
|
||||
mcp_state: &tokio::sync::Mutex<inner::McpState>,
|
||||
configs_to_start: Vec<acp::McpServer>,
|
||||
session_id: Option<&str>,
|
||||
cwd: Option<&Path>,
|
||||
meta_config_map: &inner::McpMetaConfigMap,
|
||||
oauth_config_map: &McpOAuthConfigMap,
|
||||
event_writer: &xai_file_utils::events::EventWriter,
|
||||
mode: OauthInteractivity,
|
||||
ctx: &inner::McpSpawnCtx<'_>,
|
||||
) -> Vec<Result<inner::McpClient, inner::McpError>> {
|
||||
let mut results = start_mcp_servers(
|
||||
configs_to_start,
|
||||
session_id,
|
||||
cwd,
|
||||
meta_config_map,
|
||||
oauth_config_map,
|
||||
event_writer,
|
||||
mode,
|
||||
ctx,
|
||||
)
|
||||
.await;
|
||||
// Re-resolve SDK (ACP) config.toml overrides for THIS init, matching HTTP/stdio, so a
|
||||
|
|
@ -131,12 +116,10 @@ pub async fn build_pending_clients(
|
|||
|
||||
pub async fn start_mcp_servers(
|
||||
mcp_servers: Vec<acp::McpServer>,
|
||||
session_id: Option<&str>,
|
||||
cwd: Option<&Path>,
|
||||
meta_config_map: &inner::McpMetaConfigMap,
|
||||
oauth_config_map: &McpOAuthConfigMap,
|
||||
event_writer: &xai_file_utils::events::EventWriter,
|
||||
mode: OauthInteractivity,
|
||||
ctx: &inner::McpSpawnCtx<'_>,
|
||||
) -> Vec<Result<inner::McpClient, inner::McpError>> {
|
||||
let overrides_map: HashMap<String, inner::McpClientTimeoutOverrides> = mcp_servers
|
||||
.iter()
|
||||
|
|
@ -147,12 +130,10 @@ pub async fn start_mcp_servers(
|
|||
.collect();
|
||||
inner::start_mcp_servers(
|
||||
mcp_servers,
|
||||
session_id,
|
||||
&overrides_map,
|
||||
meta_config_map,
|
||||
oauth_config_map,
|
||||
event_writer,
|
||||
mode,
|
||||
ctx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ pub mod jsonl;
|
|||
pub(crate) mod relocation;
|
||||
pub mod search;
|
||||
pub mod search_fts;
|
||||
mod search_recovery;
|
||||
pub mod search_remote_sync;
|
||||
pub(crate) mod summary_write;
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@
|
|||
//! and re-runs the full bootstrap when it is missing.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::io::{self, BufRead, Seek};
|
||||
use std::io::{self, BufRead};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
|
|
@ -23,7 +23,8 @@ use std::time::Duration;
|
|||
use tokio::sync::{Semaphore, mpsc};
|
||||
use tokio::time::Instant;
|
||||
|
||||
use super::search_fts::{SessionDoc, SessionSearchIndex, SessionSearchRow};
|
||||
use super::search_fts::{self, SessionDoc, SessionSearchIndex, SessionSearchRow};
|
||||
use super::search_recovery;
|
||||
use super::search_remote_sync;
|
||||
use super::{
|
||||
ContentPeek, PromptExtractEvent, RawLinePeek, RawParamsPeek, StorageAdapter,
|
||||
|
|
@ -39,33 +40,9 @@ const SEARCH_CONTENT_CHAR_LIMIT: usize = 200_000;
|
|||
const BOOTSTRAP_WAIT_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const BOOTSTRAP_POLL_INTERVAL: Duration = Duration::from_millis(50);
|
||||
|
||||
/// Configuration for bootstrap resource limits.
|
||||
///
|
||||
/// Phase 1-3 use hardcoded defaults via `BootstrapConfig::default()`.
|
||||
/// User-configurable overrides via `~/.grok/settings.json` are deferred
|
||||
/// to a follow-up.
|
||||
struct BootstrapConfig {
|
||||
/// Max concurrent sessions being indexed (default: 4).
|
||||
/// Used by the parallel bootstrap pipeline.
|
||||
max_concurrent: usize,
|
||||
/// Per-session timeout (default: 30 seconds).
|
||||
/// Wraps the `spawn_blocking` await — if the timeout fires, the pipeline
|
||||
/// moves on but the blocking task continues to completion.
|
||||
per_session_timeout: Duration,
|
||||
/// Max `updates.jsonl` size to index during bootstrap (default: 30 MB).
|
||||
/// Sessions exceeding this are skipped and indexed incrementally later.
|
||||
max_file_size: u64,
|
||||
}
|
||||
|
||||
impl Default for BootstrapConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_concurrent: 4,
|
||||
per_session_timeout: Duration::from_secs(30),
|
||||
max_file_size: 30 * 1024 * 1024, // 30 MB
|
||||
}
|
||||
}
|
||||
}
|
||||
const BOOTSTRAP_MAX_CONCURRENT: usize = 4;
|
||||
const BOOTSTRAP_PER_SESSION_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const BOOTSTRAP_MAX_FILE_SIZE: u64 = 30 * 1024 * 1024;
|
||||
|
||||
/// Pre-check: skip sessions with excessively large updates files.
|
||||
///
|
||||
|
|
@ -125,8 +102,6 @@ enum SearchManagerCmd {
|
|||
struct SearchManagerState {
|
||||
workers: HashMap<PathBuf, mpsc::UnboundedSender<SearchIndexJob>>,
|
||||
bootstrapped: HashSet<PathBuf>,
|
||||
#[expect(dead_code, reason = "carried for future use by worker tasks")]
|
||||
progress: Arc<BootstrapProgress>,
|
||||
}
|
||||
|
||||
/// Singleton that manages background session indexing.
|
||||
|
|
@ -135,8 +110,9 @@ struct SearchManagerState {
|
|||
///
|
||||
/// TODO: When multiple grok processes run concurrently, they each have
|
||||
/// their own `SearchIndexManager` writing to the same SQLite database.
|
||||
/// WAL mode prevents corruption, but redundant work is done. Consider
|
||||
/// adding reindex claim coordination (like the memory system's
|
||||
/// WAL mode reduces corruption risk and [`search_fts`] self-heals an
|
||||
/// unusable file, but redundant work is still done. Consider adding
|
||||
/// reindex claim coordination (like the memory system's
|
||||
/// `try_claim_reindex()` / `release_claim()` pattern) if this becomes
|
||||
/// a problem.
|
||||
pub struct SearchIndexManager {
|
||||
|
|
@ -176,14 +152,12 @@ pub struct SearchIndexStatus {
|
|||
impl SearchIndexManager {
|
||||
fn start() -> Self {
|
||||
let progress = Arc::new(BootstrapProgress::default());
|
||||
let progress_clone = progress.clone();
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<SearchManagerCmd>();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut state = SearchManagerState {
|
||||
workers: HashMap::new(),
|
||||
bootstrapped: HashSet::new(),
|
||||
progress: progress_clone,
|
||||
};
|
||||
while let Some(cmd) = rx.recv().await {
|
||||
match cmd {
|
||||
|
|
@ -284,6 +258,69 @@ fn sqlite_to_io_error(error: rusqlite::Error) -> io::Error {
|
|||
io::Error::other(format!("sqlite error: {error}"))
|
||||
}
|
||||
|
||||
/// Rate-limits a repetitive log site: the first `cap` events go to `warn`, the
|
||||
/// rest to `debug`. Resets its budget whenever the search cache is healed, so a
|
||||
/// fresh cache starts logging loudly again instead of staying silent forever.
|
||||
struct HealAwareLogCounter {
|
||||
count: AtomicU64,
|
||||
epoch_seen: AtomicU64,
|
||||
cap: u64,
|
||||
}
|
||||
|
||||
impl HealAwareLogCounter {
|
||||
const fn new(cap: u64) -> Self {
|
||||
Self {
|
||||
count: AtomicU64::new(0),
|
||||
epoch_seen: AtomicU64::new(0),
|
||||
cap,
|
||||
}
|
||||
}
|
||||
|
||||
fn should_warn(&self, kind: &str) -> bool {
|
||||
let epoch = search_recovery::current_epoch();
|
||||
if self.epoch_seen.swap(epoch, Ordering::Relaxed) != epoch {
|
||||
self.count.store(0, Ordering::Relaxed);
|
||||
}
|
||||
let n = self.count.fetch_add(1, Ordering::Relaxed);
|
||||
if n < self.cap {
|
||||
return true;
|
||||
}
|
||||
if n == self.cap {
|
||||
tracing::warn!(cap = self.cap, "further {kind} will be logged at debug");
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
static INDEX_FAIL_LOG: HealAwareLogCounter = HealAwareLogCounter::new(8);
|
||||
static BOOTSTRAP_TIMEOUT_LOG: HealAwareLogCounter = HealAwareLogCounter::new(8);
|
||||
|
||||
fn log_session_index_failure(session_id: &str, error: &io::Error, message: &str) {
|
||||
if INDEX_FAIL_LOG.should_warn("index failures") {
|
||||
tracing::warn!(error = %error, session_id = %session_id, "{message}");
|
||||
} else {
|
||||
tracing::debug!(error = %error, session_id = %session_id, "{message}");
|
||||
}
|
||||
}
|
||||
|
||||
fn log_bootstrap_timeout(session_id: &str, timeout_secs: u64) {
|
||||
let msg = "session indexing timed out during bootstrap";
|
||||
if BOOTSTRAP_TIMEOUT_LOG.should_warn("bootstrap timeouts") {
|
||||
tracing::warn!(session_id = %session_id, timeout_secs, "{msg}");
|
||||
} else {
|
||||
tracing::debug!(session_id = %session_id, timeout_secs, "{msg}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Open the index (self-heals unusable files) and run `op`, mapping errors
|
||||
/// to `io::Error` the way the rest of this module expects.
|
||||
fn with_search_index<R>(
|
||||
db_path: &Path,
|
||||
op: impl Fn(&SessionSearchIndex) -> Result<R, rusqlite::Error>,
|
||||
) -> io::Result<R> {
|
||||
search_fts::with_index(db_path, op).map_err(sqlite_to_io_error)
|
||||
}
|
||||
|
||||
/// Execute a session search query.
|
||||
///
|
||||
/// On first call, triggers a background bootstrap that indexes all
|
||||
|
|
@ -306,6 +343,7 @@ pub async fn execute_search(
|
|||
|
||||
SEARCH_INDEX_MANAGER.bootstrap_once(root_dir.to_path_buf());
|
||||
|
||||
let epoch = search_recovery::CacheEpoch::now();
|
||||
let deadline = tokio::time::Instant::now() + BOOTSTRAP_WAIT_TIMEOUT;
|
||||
while SEARCH_INDEX_MANAGER
|
||||
.progress
|
||||
|
|
@ -325,22 +363,27 @@ pub async fn execute_search(
|
|||
let query_owned = query.to_string();
|
||||
|
||||
let qr = tokio::task::spawn_blocking(move || {
|
||||
let index = SessionSearchIndex::open_or_create(&db_path).map_err(sqlite_to_io_error)?;
|
||||
index
|
||||
.query(&query_owned, cwd.as_deref(), limit, offset, include_content)
|
||||
.map_err(sqlite_to_io_error)
|
||||
with_search_index(&db_path, |index| {
|
||||
index.query(&query_owned, cwd.as_deref(), limit, offset, include_content)
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)??;
|
||||
|
||||
let healed = epoch.changed();
|
||||
if healed {
|
||||
SEARCH_INDEX_MANAGER.bootstrap_once(root_dir.to_path_buf());
|
||||
}
|
||||
|
||||
Ok(SessionSearchResponse {
|
||||
results: qr.results,
|
||||
next_offset: qr.next_offset,
|
||||
total_estimate: qr.total_estimate,
|
||||
bootstrapping: SEARCH_INDEX_MANAGER
|
||||
.progress
|
||||
.bootstrapping
|
||||
.load(Ordering::Relaxed),
|
||||
bootstrapping: healed
|
||||
|| SEARCH_INDEX_MANAGER
|
||||
.progress
|
||||
.bootstrapping
|
||||
.load(Ordering::Relaxed),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -466,10 +509,10 @@ async fn flush_ready(
|
|||
for key in ready {
|
||||
pending.remove(&key);
|
||||
if let Err(e) = upsert_by_key(root_dir, storage, &key).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
session_id = %key.session_id,
|
||||
"failed upserting session in search index"
|
||||
log_session_index_failure(
|
||||
&key.session_id,
|
||||
&e,
|
||||
"failed upserting session in search index",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -526,21 +569,20 @@ async fn upsert_session(
|
|||
// Storage backend doesn't expose file paths — no content to index
|
||||
return Ok(UpsertOutcome::NoContent);
|
||||
};
|
||||
let doc = build_session_doc(summary, content, bytes_read);
|
||||
let doc = build_session_doc(summary, content);
|
||||
let db_path = search_db_path(root_dir);
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let index = SessionSearchIndex::open_or_create(&db_path).map_err(sqlite_to_io_error)?;
|
||||
with_search_index(&db_path, |index| {
|
||||
if let Ok(Some(existing_hash)) = index.get_content_hash(&doc.session_id)
|
||||
&& existing_hash == doc.content_hash
|
||||
{
|
||||
return Ok(UpsertOutcome::Unchanged { bytes_read });
|
||||
}
|
||||
|
||||
// Skip if content hasn't changed
|
||||
if let Ok(Some(existing_hash)) = index.get_content_hash(&doc.session_id)
|
||||
&& existing_hash == doc.content_hash
|
||||
{
|
||||
return Ok(UpsertOutcome::Unchanged { bytes_read });
|
||||
}
|
||||
|
||||
index.upsert_doc(&doc).map_err(sqlite_to_io_error)?;
|
||||
Ok(UpsertOutcome::Indexed { bytes_read })
|
||||
index.upsert_doc(&doc)?;
|
||||
Ok(UpsertOutcome::Indexed { bytes_read })
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)?
|
||||
|
|
@ -550,15 +592,14 @@ async fn delete_session(root_dir: &Path, session_id: &str) -> io::Result<()> {
|
|||
let db_path = search_db_path(root_dir);
|
||||
let session_id = session_id.to_string();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let index = SessionSearchIndex::open_or_create(&db_path).map_err(sqlite_to_io_error)?;
|
||||
index.delete_doc(&session_id).map_err(sqlite_to_io_error)
|
||||
with_search_index(&db_path, |index| index.delete_doc(&session_id))
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)?
|
||||
}
|
||||
|
||||
async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Result<()> {
|
||||
let config = BootstrapConfig::default();
|
||||
let epoch = search_recovery::CacheEpoch::now();
|
||||
let progress = &SEARCH_INDEX_MANAGER.progress;
|
||||
|
||||
// Reset progress counters (bootstrapping flag already set by bootstrap_once)
|
||||
|
|
@ -589,7 +630,7 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
let mut skipped_large = 0u64;
|
||||
for (_, path) in &sessions {
|
||||
if let Some(updates_path) = path
|
||||
&& should_skip_session(updates_path, config.max_file_size)
|
||||
&& should_skip_session(updates_path, BOOTSTRAP_MAX_FILE_SIZE)
|
||||
{
|
||||
skipped_large += 1;
|
||||
}
|
||||
|
|
@ -605,7 +646,7 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
// each acquiring a permit before doing the heavy I/O work.
|
||||
// max_concurrent (default 4) limits disk I/O contention and keeps
|
||||
// the tokio blocking thread pool available for other work.
|
||||
let semaphore = Arc::new(Semaphore::new(config.max_concurrent.max(1)));
|
||||
let semaphore = Arc::new(Semaphore::new(BOOTSTRAP_MAX_CONCURRENT.max(1)));
|
||||
let progress_arc = SEARCH_INDEX_MANAGER.progress.clone();
|
||||
let root_owned = root_dir.to_path_buf();
|
||||
|
||||
|
|
@ -615,8 +656,8 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
let sem = semaphore.clone();
|
||||
let progress = progress_arc.clone();
|
||||
let root = root_owned.clone();
|
||||
let timeout_dur = config.per_session_timeout;
|
||||
let max_file_size = config.max_file_size;
|
||||
let timeout_dur = BOOTSTRAP_PER_SESSION_TIMEOUT;
|
||||
let max_file_size = BOOTSTRAP_MAX_FILE_SIZE;
|
||||
|
||||
join_set.spawn(async move {
|
||||
// Acquire semaphore permit — this provides backpressure,
|
||||
|
|
@ -641,19 +682,17 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
);
|
||||
// Insert a title-only placeholder so title search still works;
|
||||
// insert-if-absent so an existing (fuller) row is never touched.
|
||||
let doc = build_session_doc(&summary, String::new(), 0);
|
||||
let doc = build_session_doc(&summary, String::new());
|
||||
let db_path = search_db_path(&root);
|
||||
let title_only = tokio::task::spawn_blocking(move || {
|
||||
SessionSearchIndex::open_or_create(&db_path)
|
||||
.and_then(|index| index.insert_doc_if_absent(&doc))
|
||||
.map_err(sqlite_to_io_error)
|
||||
with_search_index(&db_path, |index| index.insert_doc_if_absent(&doc))
|
||||
})
|
||||
.await;
|
||||
if let Err(e) = title_only.map_err(io::Error::other).and_then(|r| r) {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
session_id = %session_id,
|
||||
"failed to write title-only index row for large session"
|
||||
log_session_index_failure(
|
||||
&session_id,
|
||||
&e,
|
||||
"failed to write title-only index row for large session",
|
||||
);
|
||||
}
|
||||
progress.skipped.fetch_add(1, Ordering::Relaxed);
|
||||
|
|
@ -680,21 +719,21 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
return Ok(UpsertOutcome::NoContent);
|
||||
};
|
||||
|
||||
let doc = build_session_doc(&summary, content, bytes_read);
|
||||
let doc = build_session_doc(&summary, content);
|
||||
let db_path = search_db_path(&root);
|
||||
|
||||
// Each task opens its own SessionSearchIndex connection.
|
||||
// SQLite WAL mode handles concurrent readers + serialized writers.
|
||||
match tokio::task::spawn_blocking(move || {
|
||||
let index =
|
||||
SessionSearchIndex::open_or_create(&db_path).map_err(sqlite_to_io_error)?;
|
||||
if let Ok(Some(existing_hash)) = index.get_content_hash(&doc.session_id)
|
||||
&& existing_hash == doc.content_hash
|
||||
{
|
||||
return Ok(UpsertOutcome::Unchanged { bytes_read });
|
||||
}
|
||||
index.upsert_doc(&doc).map_err(sqlite_to_io_error)?;
|
||||
Ok(UpsertOutcome::Indexed { bytes_read })
|
||||
with_search_index(&db_path, |index| {
|
||||
if let Ok(Some(existing_hash)) = index.get_content_hash(&doc.session_id)
|
||||
&& existing_hash == doc.content_hash
|
||||
{
|
||||
return Ok(UpsertOutcome::Unchanged { bytes_read });
|
||||
}
|
||||
index.upsert_doc(&doc)?;
|
||||
Ok(UpsertOutcome::Indexed { bytes_read })
|
||||
})
|
||||
})
|
||||
.await
|
||||
{
|
||||
|
|
@ -715,10 +754,10 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
UpsertOutcome::NoContent => {}
|
||||
},
|
||||
Ok(Err(e)) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
session_id = %session_id,
|
||||
"failed to index session for search"
|
||||
log_session_index_failure(
|
||||
&session_id,
|
||||
&e,
|
||||
"failed to index session for search",
|
||||
);
|
||||
progress.skipped.fetch_add(1, Ordering::Relaxed);
|
||||
return;
|
||||
|
|
@ -726,11 +765,7 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
Err(_) => {
|
||||
// Timeout expired — the spawn_blocking task continues to
|
||||
// completion but the pipeline moves on to the next session.
|
||||
tracing::warn!(
|
||||
session_id = %session_id,
|
||||
timeout_secs = timeout_dur.as_secs(),
|
||||
"session indexing timed out during bootstrap"
|
||||
);
|
||||
log_bootstrap_timeout(&session_id, timeout_dur.as_secs());
|
||||
progress.skipped.fetch_add(1, Ordering::Relaxed);
|
||||
return;
|
||||
}
|
||||
|
|
@ -749,17 +784,15 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
// Prune orphaned entries
|
||||
let db_path = search_db_path(root_dir);
|
||||
tokio::task::spawn_blocking(move || -> io::Result<()> {
|
||||
let index = SessionSearchIndex::open_or_create(&db_path).map_err(sqlite_to_io_error)?;
|
||||
let indexed_ids = index
|
||||
.all_indexed_session_ids()
|
||||
.map_err(sqlite_to_io_error)?;
|
||||
|
||||
for id in indexed_ids {
|
||||
if !expected_ids.contains(&id) {
|
||||
let _ = index.delete_doc(&id);
|
||||
with_search_index(&db_path, |index| {
|
||||
let indexed_ids = index.all_indexed_session_ids()?;
|
||||
for id in indexed_ids {
|
||||
if !expected_ids.contains(&id) {
|
||||
let _ = index.delete_doc(&id);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)??;
|
||||
|
|
@ -776,11 +809,22 @@ async fn reindex_all(root_dir: &Path, storage: &dyn StorageAdapter) -> io::Resul
|
|||
|
||||
progress.bootstrapping.store(false, Ordering::Release);
|
||||
|
||||
// Record bootstrap completion timestamp in the meta table.
|
||||
// Used by remote sync to determine local index staleness.
|
||||
let db_path_meta = search_db_path(root_dir);
|
||||
if let Err(e) = search_remote_sync::write_last_bootstrap_at(&db_path_meta) {
|
||||
let mut needs_rebootstrap = epoch.changed();
|
||||
if needs_rebootstrap {
|
||||
tracing::warn!("session search cache healed during bootstrap; completion marker withheld");
|
||||
} else if let Err(e) = search_remote_sync::write_last_bootstrap_at(&db_path_meta) {
|
||||
tracing::warn!(error = %e, "failed to write last_bootstrap_at metadata");
|
||||
} else if epoch.changed() {
|
||||
tracing::warn!("session search cache healed while writing completion marker; clearing it");
|
||||
if let Err(e) = search_remote_sync::clear_last_bootstrap_at(&db_path_meta) {
|
||||
tracing::warn!(error = %e, "failed to clear stale completion marker after heal");
|
||||
}
|
||||
needs_rebootstrap = true;
|
||||
}
|
||||
|
||||
if needs_rebootstrap {
|
||||
SEARCH_INDEX_MANAGER.bootstrap_once(root_dir.to_path_buf());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
|
@ -1104,176 +1148,7 @@ fn collect_all_indexable_content_single_pass(updates_path: &Path) -> io::Result<
|
|||
Ok((joined, bytes_read))
|
||||
}
|
||||
|
||||
/// Result of a delta content collection attempt.
|
||||
#[cfg_attr(
|
||||
not(test),
|
||||
expect(
|
||||
dead_code,
|
||||
reason = "Used once delta path is wired into upsert_session"
|
||||
)
|
||||
)]
|
||||
enum DeltaResult {
|
||||
/// New content extracted from the delta window.
|
||||
Content {
|
||||
/// Extracted indexable text from the new bytes.
|
||||
text: String,
|
||||
/// File size at open time — becomes the new `last_indexed_offset`.
|
||||
file_size: u64,
|
||||
},
|
||||
/// The delta window contains a `rewind_marker`, so the caller must fall
|
||||
/// back to a full re-read to rebuild prompt history correctly.
|
||||
NeedsFullReread,
|
||||
}
|
||||
|
||||
/// Collect indexable content from the *new* portion of `updates.jsonl`,
|
||||
/// starting at `offset` bytes.
|
||||
///
|
||||
/// Uses the same selective peek logic as [`collect_all_indexable_content_single_pass`]
|
||||
/// but operates only on the delta window. If a `rewind_marker` is encountered,
|
||||
/// returns [`DeltaResult::NeedsFullReread`] so the caller can fall back to a
|
||||
/// full re-read (rewind affects prompt history, which requires the full file).
|
||||
#[cfg_attr(
|
||||
not(test),
|
||||
expect(
|
||||
dead_code,
|
||||
reason = "Used once delta path is wired into upsert_session"
|
||||
)
|
||||
)]
|
||||
fn collect_delta_content(updates_path: &Path, offset: u64) -> io::Result<DeltaResult> {
|
||||
let mut file = std::fs::File::open(updates_path)?;
|
||||
let file_size = file.metadata()?.len();
|
||||
|
||||
if file_size <= offset {
|
||||
return Ok(DeltaResult::Content {
|
||||
text: String::new(),
|
||||
file_size,
|
||||
});
|
||||
}
|
||||
|
||||
file.seek(io::SeekFrom::Start(offset))?;
|
||||
let reader = io::BufReader::new(file);
|
||||
|
||||
let mut user_texts: Vec<String> = Vec::new();
|
||||
let mut assistant_texts: Vec<String> = Vec::new();
|
||||
let mut current_assistant = String::new();
|
||||
let mut tool_meta: Vec<String> = Vec::new();
|
||||
|
||||
let flush_assistant = |current: &mut String, texts: &mut Vec<String>| {
|
||||
if !current.is_empty() {
|
||||
let t = current.trim().to_string();
|
||||
if !t.is_empty() {
|
||||
texts.push(t);
|
||||
}
|
||||
current.clear();
|
||||
}
|
||||
};
|
||||
|
||||
for line_result in reader.lines() {
|
||||
let line = match line_result {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let (raw_params, is_xai) = if let Ok(env) = serde_json::from_str::<RawLinePeek<'_>>(trimmed)
|
||||
{
|
||||
let raw = env.params.map(|p| p.get()).unwrap_or(trimmed);
|
||||
let xai = env.method == Some(XAI_SESSION_UPDATE_METHOD);
|
||||
(raw, xai)
|
||||
} else {
|
||||
(trimmed, false)
|
||||
};
|
||||
|
||||
let tag = serde_json::from_str::<RawParamsPeek<'_>>(raw_params)
|
||||
.ok()
|
||||
.and_then(|p| p.update)
|
||||
.map(|u| u.session_update);
|
||||
|
||||
match tag {
|
||||
Some(t) if is_xai && t == *REWIND_MARKER => {
|
||||
return Ok(DeltaResult::NeedsFullReread);
|
||||
}
|
||||
Some(t) if !is_xai && t == *USER_MESSAGE_CHUNK => {
|
||||
flush_assistant(&mut current_assistant, &mut assistant_texts);
|
||||
if let Ok(peek) = serde_json::from_str::<UserContentPeek<'_>>(raw_params)
|
||||
&& let Some(content) = peek.update.content
|
||||
&& content.content_type == Some("text")
|
||||
&& let Some(text) = content.text
|
||||
&& content
|
||||
.meta
|
||||
.as_ref()
|
||||
.is_none_or(|m| m.bash_command.is_none())
|
||||
&& peek
|
||||
.update
|
||||
.meta
|
||||
.as_ref()
|
||||
.is_none_or(|m| m.host_turn != Some(true))
|
||||
{
|
||||
user_texts.push(text.into_owned());
|
||||
}
|
||||
}
|
||||
Some("agent_message_chunk") if !is_xai => {
|
||||
if let Ok(peek) = serde_json::from_str::<AgentContentPeek<'_>>(raw_params)
|
||||
&& let Some(content) = peek.update.content
|
||||
&& content.content_type == Some("text")
|
||||
&& let Some(text) = content.text
|
||||
&& !text.is_empty()
|
||||
{
|
||||
if !current_assistant.is_empty() {
|
||||
current_assistant.push(' ');
|
||||
}
|
||||
current_assistant.push_str(&text);
|
||||
}
|
||||
}
|
||||
Some("tool_call") if !is_xai => {
|
||||
flush_assistant(&mut current_assistant, &mut assistant_texts);
|
||||
if let Ok(peek) = serde_json::from_str::<ToolCallPeek<'_>>(raw_params) {
|
||||
if let Some(title) = peek.update.title
|
||||
&& !title.is_empty()
|
||||
{
|
||||
tool_meta.push(title.into_owned());
|
||||
}
|
||||
if let Some(locs) = peek.update.locations {
|
||||
for loc in locs {
|
||||
if let Some(p) = loc.path
|
||||
&& !p.is_empty()
|
||||
{
|
||||
tool_meta.push(p.into_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
flush_assistant(&mut current_assistant, &mut assistant_texts);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
flush_assistant(&mut current_assistant, &mut assistant_texts);
|
||||
|
||||
let user_part = user_texts.join("\n\n");
|
||||
let assistant_part = assistant_texts.join("\n");
|
||||
let tool_part = tool_meta.join("\n");
|
||||
let mut parts: Vec<&str> = Vec::new();
|
||||
if !user_part.is_empty() {
|
||||
parts.push(&user_part);
|
||||
}
|
||||
if !assistant_part.is_empty() {
|
||||
parts.push(&assistant_part);
|
||||
}
|
||||
if !tool_part.is_empty() {
|
||||
parts.push(&tool_part);
|
||||
}
|
||||
let text = parts.join("\n\n");
|
||||
|
||||
Ok(DeltaResult::Content { text, file_size })
|
||||
}
|
||||
|
||||
fn build_session_doc(summary: &Summary, content: String, last_indexed_offset: u64) -> SessionDoc {
|
||||
fn build_session_doc(summary: &Summary, content: String) -> SessionDoc {
|
||||
let title = summary.display_title().to_owned();
|
||||
|
||||
let mut hasher = blake3::Hasher::new();
|
||||
|
|
@ -1289,7 +1164,6 @@ fn build_session_doc(summary: &Summary, content: String, last_indexed_offset: u6
|
|||
title,
|
||||
content,
|
||||
content_hash,
|
||||
last_indexed_offset,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1372,14 +1246,14 @@ mod tests {
|
|||
fn test_build_session_doc_hashes_content() {
|
||||
let summary = test_summary("test-session", "/workspace", "My session title");
|
||||
|
||||
let doc = build_session_doc(&summary, "prompt text".to_string(), 0);
|
||||
let doc = build_session_doc(&summary, "prompt text".to_string());
|
||||
assert_eq!(doc.session_id, "test-session");
|
||||
assert_eq!(doc.title, "My session title");
|
||||
assert_eq!(doc.content, "prompt text");
|
||||
assert!(!doc.content_hash.is_empty());
|
||||
|
||||
// Same content + same title → same hash
|
||||
let doc2 = build_session_doc(&summary, "prompt text".to_string(), 0);
|
||||
let doc2 = build_session_doc(&summary, "prompt text".to_string());
|
||||
assert_eq!(doc.content_hash, doc2.content_hash);
|
||||
}
|
||||
|
||||
|
|
@ -1688,8 +1562,8 @@ mod tests {
|
|||
let new = test_summary("s1", "/workspace", "New title");
|
||||
let content = "same prompt text".to_string();
|
||||
|
||||
let doc_old = build_session_doc(&old, content.clone(), 0);
|
||||
let doc_new = build_session_doc(&new, content, 0);
|
||||
let doc_old = build_session_doc(&old, content.clone());
|
||||
let doc_new = build_session_doc(&new, content);
|
||||
|
||||
assert_ne!(
|
||||
doc_old.content_hash, doc_new.content_hash,
|
||||
|
|
@ -1701,24 +1575,14 @@ mod tests {
|
|||
fn test_build_session_doc_prefers_generated_title() {
|
||||
let mut summary = test_summary("s1", "/workspace", "session summary");
|
||||
summary.generated_title = Some("Generated Title".to_string());
|
||||
let doc = build_session_doc(&summary, "content".to_string(), 0);
|
||||
let doc = build_session_doc(&summary, "content".to_string());
|
||||
assert_eq!(doc.title, "Generated Title");
|
||||
|
||||
summary.generated_title = Some(String::new());
|
||||
let doc2 = build_session_doc(&summary, "content".to_string(), 0);
|
||||
let doc2 = build_session_doc(&summary, "content".to_string());
|
||||
assert_eq!(doc2.title, "session summary");
|
||||
}
|
||||
|
||||
// ── bootstrap config tests ─────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_config_defaults() {
|
||||
let config = BootstrapConfig::default();
|
||||
assert_eq!(config.max_concurrent, 4);
|
||||
assert_eq!(config.per_session_timeout, Duration::from_secs(30));
|
||||
assert_eq!(config.max_file_size, 30 * 1024 * 1024);
|
||||
}
|
||||
|
||||
// ── should_skip_session tests ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
|
|
@ -1808,140 +1672,6 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
// ── delta indexing tests ───────────────────────────────────────────────
|
||||
|
||||
/// Append new lines to an existing updates.jsonl file and return the
|
||||
/// byte offset where the new content starts.
|
||||
fn append_updates_jsonl(path: &Path, lines: &[String]) -> u64 {
|
||||
use std::io::Write as _;
|
||||
let offset = std::fs::metadata(path).unwrap().len();
|
||||
let mut f = std::fs::OpenOptions::new().append(true).open(path).unwrap();
|
||||
for line in lines {
|
||||
writeln!(f, "{line}").unwrap();
|
||||
}
|
||||
offset
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_append_extracts_new_content() {
|
||||
// Write initial content
|
||||
let initial = vec![
|
||||
acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"hello world"}}"#,
|
||||
),
|
||||
acp_update(
|
||||
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hi there"}}"#,
|
||||
),
|
||||
];
|
||||
let f = write_updates_jsonl(&initial);
|
||||
let offset = std::fs::metadata(f.path()).unwrap().len();
|
||||
|
||||
// Append new content
|
||||
let delta = vec![
|
||||
acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"second question"}}"#,
|
||||
),
|
||||
acp_update(
|
||||
r#"{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"second answer"}}"#,
|
||||
),
|
||||
];
|
||||
append_updates_jsonl(f.path(), &delta);
|
||||
|
||||
let result = collect_delta_content(f.path(), offset).unwrap();
|
||||
match result {
|
||||
DeltaResult::Content { text, file_size } => {
|
||||
assert!(
|
||||
text.contains("second question"),
|
||||
"delta should contain new user prompt"
|
||||
);
|
||||
assert!(
|
||||
text.contains("second answer"),
|
||||
"delta should contain new assistant text"
|
||||
);
|
||||
assert!(
|
||||
!text.contains("hello world"),
|
||||
"delta should not contain pre-offset content"
|
||||
);
|
||||
assert!(
|
||||
file_size > offset,
|
||||
"file_size should be larger than the starting offset"
|
||||
);
|
||||
}
|
||||
DeltaResult::NeedsFullReread => {
|
||||
panic!("expected Content, got NeedsFullReread");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_rewind_triggers_full_reread() {
|
||||
let initial = vec![acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"first"}}"#,
|
||||
)];
|
||||
let f = write_updates_jsonl(&initial);
|
||||
let offset = std::fs::metadata(f.path()).unwrap().len();
|
||||
|
||||
// Append a rewind marker in the delta window
|
||||
let delta = vec![
|
||||
xai_update(
|
||||
r#"{"sessionUpdate":"rewind_marker","target_prompt_index":0,"created_at":"2024-01-01"}"#,
|
||||
),
|
||||
acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"replacement"}}"#,
|
||||
),
|
||||
];
|
||||
append_updates_jsonl(f.path(), &delta);
|
||||
|
||||
let result = collect_delta_content(f.path(), offset).unwrap();
|
||||
assert!(
|
||||
matches!(result, DeltaResult::NeedsFullReread),
|
||||
"rewind in delta should trigger NeedsFullReread"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_no_new_bytes() {
|
||||
let lines = vec![acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"hello"}}"#,
|
||||
)];
|
||||
let f = write_updates_jsonl(&lines);
|
||||
let file_size = std::fs::metadata(f.path()).unwrap().len();
|
||||
|
||||
// Offset equals file size — no new bytes
|
||||
let result = collect_delta_content(f.path(), file_size).unwrap();
|
||||
match result {
|
||||
DeltaResult::Content { text, .. } => {
|
||||
assert!(text.is_empty(), "no new bytes should produce empty text");
|
||||
}
|
||||
DeltaResult::NeedsFullReread => {
|
||||
panic!("expected Content with empty text, got NeedsFullReread");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_truncation_detected() {
|
||||
let lines = vec![acp_update(
|
||||
r#"{"sessionUpdate":"user_message_chunk","content":{"type":"text","text":"hello world this is a long message"}}"#,
|
||||
)];
|
||||
let f = write_updates_jsonl(&lines);
|
||||
let file_size = std::fs::metadata(f.path()).unwrap().len();
|
||||
|
||||
// Offset larger than file size — simulates truncation
|
||||
let result = collect_delta_content(f.path(), file_size + 100).unwrap();
|
||||
match result {
|
||||
DeltaResult::Content { text, .. } => {
|
||||
assert!(
|
||||
text.is_empty(),
|
||||
"offset beyond file size should produce empty text"
|
||||
);
|
||||
}
|
||||
DeltaResult::NeedsFullReread => {
|
||||
panic!("expected Content with empty text for truncation");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── bootstrap_once eager flag tests ────────────────────────────────────
|
||||
// NOTE: SEARCH_INDEX_MANAGER is a process-wide singleton, so tests
|
||||
// that depend on the `bootstrapping` flag transitioning to `false`
|
||||
|
|
@ -2057,7 +1787,7 @@ mod tests {
|
|||
let db_path = search_db_path(tmp.path());
|
||||
|
||||
let summary = test_summary("stub", "/ws", "");
|
||||
let stub = build_session_doc(&summary, String::new(), 0);
|
||||
let stub = build_session_doc(&summary, String::new());
|
||||
{
|
||||
let index = SessionSearchIndex::open_or_create(&db_path).unwrap();
|
||||
index.upsert_doc(&stub).unwrap();
|
||||
|
|
@ -2079,9 +1809,4 @@ mod tests {
|
|||
"the upgrade drop must clear stub rows so their stale hashes cannot block re-indexing"
|
||||
);
|
||||
}
|
||||
|
||||
// Note: tests for upsert_session_blocking (delta path, truncation
|
||||
// fallback, rewind fallback, no-new-bytes skip) are deferred until
|
||||
// the delta indexing is wired into upsert_session. The delta content
|
||||
// collection function (collect_delta_content) is tested above.
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,10 +13,17 @@
|
|||
//! triggers so callers never need to touch the FTS table directly.
|
||||
//! The `cwd` column is intentionally excluded from the FTS table — it is a
|
||||
//! filter dimension only, applied via JOIN on `session_docs`.
|
||||
//!
|
||||
//! The index is a rebuildable cache: an unusable file is quarantined and
|
||||
//! recreated once (see [`super::search_recovery`] and [`with_index`]).
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use rusqlite::{Connection, OptionalExtension, params};
|
||||
use xai_sqlite_journal::JournalMode;
|
||||
|
||||
use super::search_recovery;
|
||||
|
||||
/// Bump when making breaking schema changes that require dropping and
|
||||
/// recreating tables, or to force a rebuild of stale index content
|
||||
/// (v3 → v4: messages with JSON escapes were silently dropped at indexing).
|
||||
|
|
@ -33,19 +40,6 @@ pub struct SessionDoc {
|
|||
pub content: String,
|
||||
/// blake3 hash of `content` — used to skip redundant upserts.
|
||||
pub content_hash: String,
|
||||
/// Byte offset in `updates.jsonl` up to which content has been indexed.
|
||||
/// Used for delta indexing: on subsequent updates, only bytes after this
|
||||
/// offset are parsed and merged with existing content.
|
||||
pub last_indexed_offset: u64,
|
||||
}
|
||||
|
||||
/// State of a previously indexed session, returned by
|
||||
/// [`SessionSearchIndex::get_session_index_state`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionIndexState {
|
||||
pub content: String,
|
||||
pub content_hash: String,
|
||||
pub last_indexed_offset: u64,
|
||||
}
|
||||
|
||||
/// A single search result row.
|
||||
|
|
@ -73,6 +67,28 @@ pub struct SessionSearchIndex {
|
|||
db: Connection,
|
||||
}
|
||||
|
||||
pub fn with_index<R>(
|
||||
db_path: &Path,
|
||||
op: impl Fn(&SessionSearchIndex) -> Result<R, rusqlite::Error>,
|
||||
) -> Result<R, rusqlite::Error> {
|
||||
let index = SessionSearchIndex::open_or_create(db_path)?;
|
||||
match op(&index) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(e) if search_recovery::is_unusable_db_error(&e) => {
|
||||
drop(index);
|
||||
search_recovery::heal_unusable(
|
||||
db_path,
|
||||
&e,
|
||||
SessionSearchIndex::probe_usable,
|
||||
SessionSearchIndex::recreate,
|
||||
);
|
||||
let index = SessionSearchIndex::open_or_create(db_path)?;
|
||||
op(&index)
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
impl SessionSearchIndex {
|
||||
/// Open (or create) the FTS index at `db_path`.
|
||||
///
|
||||
|
|
@ -82,25 +98,53 @@ impl SessionSearchIndex {
|
|||
/// and deletes the `last_bootstrap_at` completed-bootstrap marker so the
|
||||
/// wipe is observable to bootstrap/staleness checks.
|
||||
/// A NEWER stored version is tolerated read/write without dropping.
|
||||
pub fn open_or_create(db_path: &std::path::Path) -> Result<Self, rusqlite::Error> {
|
||||
///
|
||||
/// If the existing file is corrupt / not a database, quarantines it and
|
||||
/// opens a fresh empty index (see [`super::search_recovery::heal_unusable`]).
|
||||
pub fn open_or_create(db_path: &Path) -> Result<Self, rusqlite::Error> {
|
||||
if let Some(parent) = db_path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
|
||||
// The mode decision statfs's the parent dir created above.
|
||||
let journal_mode = JournalMode::for_db_path(db_path);
|
||||
|
||||
match Self::open_with_journal_mode(db_path, journal_mode) {
|
||||
Ok(index) => Ok(index),
|
||||
Err(e) if search_recovery::is_unusable_db_error(&e) => {
|
||||
search_recovery::heal_unusable(db_path, &e, Self::probe_usable, Self::recreate);
|
||||
Self::open_with_journal_mode(db_path, journal_mode)
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Open an existing index without the corruption self-heal; returns the
|
||||
/// unusable-DB error instead of quarantining.
|
||||
pub fn open_existing(db_path: &Path) -> Result<Self, rusqlite::Error> {
|
||||
if let Some(parent) = db_path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
Self::open_with_journal_mode(db_path, JournalMode::for_db_path(db_path))
|
||||
}
|
||||
|
||||
/// Open with an explicit journal mode — the seam tests use to exercise
|
||||
/// the network-filesystem decision on a local disk.
|
||||
fn probe_usable(db_path: &Path) -> Result<bool, rusqlite::Error> {
|
||||
let conn = JournalMode::for_db_path(db_path).open_readonly(db_path)?;
|
||||
let mut stmt = conn.prepare("PRAGMA integrity_check")?;
|
||||
let first: Option<String> = stmt.query_map([], |row| row.get(0))?.next().transpose()?;
|
||||
Ok(first.as_deref() == Some("ok"))
|
||||
}
|
||||
|
||||
fn recreate(db_path: &Path) -> Result<(), rusqlite::Error> {
|
||||
Self::open_with_journal_mode(db_path, JournalMode::for_db_path(db_path)).map(|_| ())
|
||||
}
|
||||
|
||||
fn open_with_journal_mode(
|
||||
db_path: &std::path::Path,
|
||||
db_path: &Path,
|
||||
journal_mode: JournalMode,
|
||||
) -> Result<Self, rusqlite::Error> {
|
||||
// busy_timeout + journal pragma live in the helper (see JournalMode::open).
|
||||
let db = journal_mode.open(db_path)?;
|
||||
|
||||
// Check existing schema version
|
||||
let stored_version: Option<String> = db
|
||||
.query_row(
|
||||
"SELECT value FROM meta WHERE key = 'session_search_schema_version'",
|
||||
|
|
@ -155,7 +199,6 @@ impl SessionSearchIndex {
|
|||
);
|
||||
}
|
||||
|
||||
// Create tables + content-synced FTS5 with auto-sync triggers
|
||||
db.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS meta (
|
||||
|
|
@ -202,20 +245,6 @@ impl SessionSearchIndex {
|
|||
",
|
||||
)?;
|
||||
|
||||
// Add last_indexed_offset column (idempotent migration).
|
||||
match db.execute(
|
||||
"ALTER TABLE session_docs ADD COLUMN last_indexed_offset INTEGER NOT NULL DEFAULT 0",
|
||||
[],
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if !msg.contains("duplicate column") {
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Persist schema version — but never regress the row a newer
|
||||
// generation owns (it would re-trigger that binary's upgrade drop).
|
||||
if stored != Some(current) && !owned_by_newer {
|
||||
|
|
@ -235,23 +264,21 @@ impl SessionSearchIndex {
|
|||
/// automatically.
|
||||
pub fn upsert_doc(&self, doc: &SessionDoc) -> Result<(), rusqlite::Error> {
|
||||
self.db.execute(
|
||||
"INSERT INTO session_docs(session_id, cwd, updated_at, title, content, content_hash, last_indexed_offset)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
|
||||
"INSERT INTO session_docs(session_id, cwd, updated_at, title, content, content_hash)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
|
||||
ON CONFLICT(session_id) DO UPDATE SET
|
||||
cwd = excluded.cwd,
|
||||
updated_at = excluded.updated_at,
|
||||
title = excluded.title,
|
||||
content = excluded.content,
|
||||
content_hash = excluded.content_hash,
|
||||
last_indexed_offset = excluded.last_indexed_offset",
|
||||
content_hash = excluded.content_hash",
|
||||
params![
|
||||
doc.session_id,
|
||||
doc.cwd,
|
||||
doc.updated_at_unix,
|
||||
doc.title,
|
||||
doc.content,
|
||||
doc.content_hash,
|
||||
doc.last_indexed_offset as i64
|
||||
doc.content_hash
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
|
|
@ -264,8 +291,8 @@ impl SessionSearchIndex {
|
|||
/// written between the check and the insert.
|
||||
pub fn insert_doc_if_absent(&self, doc: &SessionDoc) -> Result<(), rusqlite::Error> {
|
||||
self.db.execute(
|
||||
"INSERT INTO session_docs(session_id, cwd, updated_at, title, content, content_hash, last_indexed_offset)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
|
||||
"INSERT INTO session_docs(session_id, cwd, updated_at, title, content, content_hash)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6)
|
||||
ON CONFLICT(session_id) DO NOTHING",
|
||||
params![
|
||||
doc.session_id,
|
||||
|
|
@ -273,8 +300,7 @@ impl SessionSearchIndex {
|
|||
doc.updated_at_unix,
|
||||
doc.title,
|
||||
doc.content,
|
||||
doc.content_hash,
|
||||
doc.last_indexed_offset as i64
|
||||
doc.content_hash
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
|
|
@ -302,44 +328,6 @@ impl SessionSearchIndex {
|
|||
.optional()
|
||||
}
|
||||
|
||||
/// Return the full index state for a session: content, content_hash,
|
||||
/// and last_indexed_offset. Used by the delta indexing path.
|
||||
pub fn get_session_index_state(
|
||||
&self,
|
||||
session_id: &str,
|
||||
) -> Result<Option<SessionIndexState>, rusqlite::Error> {
|
||||
self.db
|
||||
.query_row(
|
||||
"SELECT content, content_hash, last_indexed_offset FROM session_docs WHERE session_id = ?1",
|
||||
params![session_id],
|
||||
|row| {
|
||||
let content: String = row.get(0)?;
|
||||
let content_hash: String = row.get(1)?;
|
||||
let offset: i64 = row.get(2)?;
|
||||
Ok(SessionIndexState {
|
||||
content,
|
||||
content_hash,
|
||||
last_indexed_offset: offset as u64,
|
||||
})
|
||||
},
|
||||
)
|
||||
.optional()
|
||||
}
|
||||
|
||||
/// Update only the `last_indexed_offset` for a session without touching
|
||||
/// content or hash (avoids firing FTS triggers when content is unchanged).
|
||||
pub fn update_indexed_offset(
|
||||
&self,
|
||||
session_id: &str,
|
||||
offset: u64,
|
||||
) -> Result<(), rusqlite::Error> {
|
||||
self.db.execute(
|
||||
"UPDATE session_docs SET last_indexed_offset = ?2 WHERE session_id = ?1",
|
||||
params![session_id, offset as i64],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a value from the `meta` key-value table.
|
||||
pub fn get_meta(&self, key: &str) -> Result<Option<String>, rusqlite::Error> {
|
||||
self.db
|
||||
|
|
@ -360,6 +348,13 @@ impl SessionSearchIndex {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a value from the `meta` key-value table.
|
||||
pub fn delete_meta(&self, key: &str) -> Result<(), rusqlite::Error> {
|
||||
self.db
|
||||
.execute("DELETE FROM meta WHERE key = ?1", params![key])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Return all session IDs currently in the index.
|
||||
///
|
||||
/// Used during reindex to detect and prune orphaned entries.
|
||||
|
|
@ -573,7 +568,6 @@ mod tests {
|
|||
title: title.to_string(),
|
||||
content: content.to_string(),
|
||||
content_hash: blake3::hash(content.as_bytes()).to_hex().to_string(),
|
||||
last_indexed_offset: 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -769,6 +763,73 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_malformed_db_file_is_quarantined_and_recreated() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let path = tmp.path().join("session_search.sqlite");
|
||||
// Not a SQLite database — classic "file is not a database" / NOTADB.
|
||||
std::fs::write(&path, b"this is not a sqlite database at all").unwrap();
|
||||
|
||||
// Drive the production entrypoint: it heals on open, then the op runs.
|
||||
with_index(&path, |index| {
|
||||
index.upsert_doc(&test_doc("s1", "after heal", "session search works again"))
|
||||
})
|
||||
.expect("with_index self-heals then upserts");
|
||||
let qr = with_index(&path, |index| index.query("works", None, 10, 0, false))
|
||||
.expect("query after heal");
|
||||
assert_eq!(qr.results[0].session_id, "s1");
|
||||
|
||||
// Original path is a real DB again; a quarantine sibling should exist.
|
||||
assert!(path.is_file());
|
||||
let quarantined: Vec<_> = std::fs::read_dir(tmp.path())
|
||||
.unwrap()
|
||||
.filter_map(|e| e.ok())
|
||||
.map(|e| e.file_name().to_string_lossy().into_owned())
|
||||
.filter(|n| n.contains("corrupt"))
|
||||
.collect();
|
||||
assert!(
|
||||
!quarantined.is_empty(),
|
||||
"expected a quarantined corrupt sibling, got {quarantined:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_index_retries_op_once_after_heal() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let path = tmp.path().join("session_search.sqlite");
|
||||
// A non-db file gives a real "unusable" error to inject into the op.
|
||||
let bogus = tmp.path().join("bogus.sqlite");
|
||||
std::fs::write(&bogus, b"not-sqlite").unwrap();
|
||||
|
||||
// The first op attempt reports the DB unusable, as a mid-op corruption
|
||||
// would; with_index heals (a no-op here, the DB is healthy) and retries
|
||||
// the op exactly once, which then succeeds.
|
||||
let calls = std::cell::Cell::new(0u32);
|
||||
let result = with_index(&path, |index| {
|
||||
let n = calls.get();
|
||||
calls.set(n + 1);
|
||||
if n == 0 {
|
||||
match SessionSearchIndex::open_with_journal_mode(&bogus, JournalMode::Wal) {
|
||||
Ok(_) => unreachable!("a non-sqlite file cannot open as a database"),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
} else {
|
||||
index.upsert_doc(&test_doc("s1", "t", "retried body"))
|
||||
}
|
||||
});
|
||||
|
||||
assert!(result.is_ok(), "op should succeed on the retry: {result:?}");
|
||||
assert_eq!(
|
||||
calls.get(),
|
||||
2,
|
||||
"op runs once, fails unusable, then runs once more"
|
||||
);
|
||||
|
||||
let index = SessionSearchIndex::open_or_create(&path).unwrap();
|
||||
let qr = index.query("retried", None, 10, 0, false).unwrap();
|
||||
assert_eq!(qr.results.len(), 1, "the retried op's write is persisted");
|
||||
}
|
||||
|
||||
/// Repro: the on-disk state left behind by a pre-ratchet binary that
|
||||
/// wiped the shared DB and ran its own bootstrap — a v3-stamped index
|
||||
/// with a *recent* bootstrap marker. Pins that the current binary's open
|
||||
|
|
|
|||
|
|
@ -0,0 +1,232 @@
|
|||
//! Corruption self-heal for the session-search SQLite cache: classify an
|
||||
//! unusable file, then quarantine it under a lock so a fresh empty database
|
||||
//! can be recreated. The index layer drives the retry.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use rusqlite::ErrorCode;
|
||||
use xai_sqlite_journal::JournalMode;
|
||||
|
||||
static HEAL_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
/// Bumped each time the cache is quarantined and recreated, so callers can tell
|
||||
/// they are now looking at a different incarnation of the on-disk index.
|
||||
static CACHE_EPOCH: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
pub(super) fn current_epoch() -> u64 {
|
||||
CACHE_EPOCH.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
/// A snapshot of the [`CACHE_EPOCH`], used to detect whether the cache was
|
||||
/// quarantined and recreated between two points in this process.
|
||||
pub(super) struct CacheEpoch(u64);
|
||||
|
||||
impl CacheEpoch {
|
||||
pub(super) fn now() -> Self {
|
||||
Self(CACHE_EPOCH.load(Ordering::Acquire))
|
||||
}
|
||||
|
||||
pub(super) fn changed(&self) -> bool {
|
||||
CACHE_EPOCH.load(Ordering::Acquire) != self.0
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn is_unusable_db_error(error: &rusqlite::Error) -> bool {
|
||||
match error {
|
||||
rusqlite::Error::SqliteFailure(err, msg) => {
|
||||
if matches!(
|
||||
err.code,
|
||||
ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
msg.as_deref().is_some_and(message_indicates_unusable_db)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Specific phrases only: a bare "malformed" would also match a bad FTS query.
|
||||
fn message_indicates_unusable_db(msg: &str) -> bool {
|
||||
let lower = msg.to_ascii_lowercase();
|
||||
lower.contains("disk image is malformed")
|
||||
|| lower.contains("database schema is malformed")
|
||||
|| lower.contains("database is corrupt")
|
||||
|| lower.contains("file is not a database")
|
||||
|| lower.contains("file is encrypted or is not a database")
|
||||
}
|
||||
|
||||
fn with_suffix(path: &Path, suffix: &str) -> PathBuf {
|
||||
let mut name = path.as_os_str().to_os_string();
|
||||
name.push(suffix);
|
||||
PathBuf::from(name)
|
||||
}
|
||||
|
||||
pub(super) fn quarantine_db_files(db_path: &Path) -> Option<PathBuf> {
|
||||
let ts = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
let corrupt_suffix = format!(".corrupt.{ts}");
|
||||
|
||||
for suffix in ["-wal", "-shm", "-journal"] {
|
||||
let side = with_suffix(db_path, suffix);
|
||||
if !side.exists() {
|
||||
continue;
|
||||
}
|
||||
let dest = with_suffix(&side, &corrupt_suffix);
|
||||
if let Err(e) = std::fs::rename(&side, &dest) {
|
||||
tracing::debug!(
|
||||
error = %e,
|
||||
path = %side.display(),
|
||||
"could not quarantine sqlite sidecar; left in place"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let main_quarantine = with_suffix(db_path, &corrupt_suffix);
|
||||
let renamed = if db_path.exists() {
|
||||
match std::fs::rename(db_path, &main_quarantine) {
|
||||
Ok(()) => true,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
path = %db_path.display(),
|
||||
"failed to quarantine corrupt session search db; left in place"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
renamed.then_some(main_quarantine)
|
||||
}
|
||||
|
||||
pub(super) fn heal_unusable(
|
||||
db_path: &Path,
|
||||
cause: &rusqlite::Error,
|
||||
reprobe: impl FnOnce(&Path) -> Result<bool, rusqlite::Error>,
|
||||
recreate: impl FnOnce(&Path) -> Result<(), rusqlite::Error>,
|
||||
) {
|
||||
use fs2::FileExt as _;
|
||||
|
||||
let _guard = HEAL_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
|
||||
let effective = JournalMode::for_db_path(db_path).effective_db_path(db_path);
|
||||
|
||||
let lock_path = with_suffix(&effective, ".lock");
|
||||
let _lock_file = match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.write(true)
|
||||
.open(&lock_path)
|
||||
{
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
tracing::debug!(error = %e, path = %lock_path.display(), "could not open heal lock; skipping quarantine");
|
||||
return;
|
||||
}
|
||||
};
|
||||
if let Err(e) = _lock_file.lock_exclusive() {
|
||||
tracing::debug!(error = %e, "could not acquire cross-process heal lock; skipping quarantine");
|
||||
return;
|
||||
}
|
||||
|
||||
match reprobe(&effective) {
|
||||
Ok(true) => return,
|
||||
Ok(false) => {}
|
||||
Err(e) if is_unusable_db_error(&e) => {}
|
||||
Err(_) => return,
|
||||
}
|
||||
|
||||
let quarantine = quarantine_db_files(&effective);
|
||||
let recreated = recreate(&effective);
|
||||
if let Err(e) = &recreated {
|
||||
tracing::warn!(error = %e, "failed to recreate session search index after quarantine");
|
||||
}
|
||||
|
||||
if quarantine.is_none() && recreated.is_err() {
|
||||
tracing::warn!(
|
||||
db_path = %effective.display(),
|
||||
error = %cause,
|
||||
"session search index unusable but could not be quarantined or recreated; left in place"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
CACHE_EPOCH.fetch_add(1, Ordering::Release);
|
||||
tracing::warn!(
|
||||
db_path = %effective.display(),
|
||||
quarantine = %quarantine
|
||||
.as_ref()
|
||||
.map(|p| p.display().to_string())
|
||||
.unwrap_or_else(|| "(removed or missing)".into()),
|
||||
error = %cause,
|
||||
"session search index unusable; quarantined and recreated empty cache"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn has_corrupt_sibling(dir: &Path) -> bool {
|
||||
std::fs::read_dir(dir)
|
||||
.unwrap()
|
||||
.filter_map(|e| e.ok())
|
||||
.any(|e| e.file_name().to_string_lossy().contains("corrupt"))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quarantine_moves_main_and_sidecars() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let db = tmp.path().join("session_search.sqlite");
|
||||
std::fs::write(&db, b"main").unwrap();
|
||||
std::fs::write(tmp.path().join("session_search.sqlite-wal"), b"wal").unwrap();
|
||||
std::fs::write(tmp.path().join("session_search.sqlite-shm"), b"shm").unwrap();
|
||||
|
||||
let moved = quarantine_db_files(&db).expect("main db should be quarantined");
|
||||
|
||||
assert!(!db.exists());
|
||||
assert!(moved.exists());
|
||||
assert!(!tmp.path().join("session_search.sqlite-wal").exists());
|
||||
assert!(!tmp.path().join("session_search.sqlite-shm").exists());
|
||||
}
|
||||
|
||||
fn quarantined_after(reprobe: impl FnOnce(&Path) -> Result<bool, rusqlite::Error>) -> bool {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let db = tmp.path().join("session_search.sqlite");
|
||||
std::fs::write(&db, b"looks-like-a-db").unwrap();
|
||||
heal_unusable(&db, &rusqlite::Error::QueryReturnedNoRows, reprobe, |_| {
|
||||
Ok(())
|
||||
});
|
||||
has_corrupt_sibling(tmp.path())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heal_quarantines_only_on_confirmed_corruption() {
|
||||
assert!(!quarantined_after(|_| Ok(true)), "healthy: no quarantine");
|
||||
assert!(
|
||||
!quarantined_after(|_| Err(rusqlite::Error::QueryReturnedNoRows)),
|
||||
"transient failure: no quarantine"
|
||||
);
|
||||
assert!(quarantined_after(|_| Ok(false)), "corrupt: quarantine");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifier_ignores_bad_query_but_catches_corruption() {
|
||||
assert!(message_indicates_unusable_db(
|
||||
"database disk image is malformed"
|
||||
));
|
||||
assert!(message_indicates_unusable_db("file is not a database"));
|
||||
assert!(!message_indicates_unusable_db("malformed MATCH expression"));
|
||||
}
|
||||
}
|
||||
|
|
@ -139,13 +139,22 @@ pub fn try_read_last_bootstrap_at(db_path: &Path) -> Result<Option<i64>, String>
|
|||
/// Write `last_bootstrap_at` into the sqlite meta table.
|
||||
pub fn write_last_bootstrap_at(db_path: &Path) -> io::Result<()> {
|
||||
let index =
|
||||
SessionSearchIndex::open_or_create(db_path).map_err(|e| io::Error::other(e.to_string()))?;
|
||||
SessionSearchIndex::open_existing(db_path).map_err(|e| io::Error::other(e.to_string()))?;
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
index
|
||||
.set_meta(META_KEY_LAST_BOOTSTRAP, &now.to_string())
|
||||
.map_err(|e| io::Error::other(e.to_string()))
|
||||
}
|
||||
|
||||
/// Remove the completed-bootstrap marker.
|
||||
pub fn clear_last_bootstrap_at(db_path: &Path) -> io::Result<()> {
|
||||
let index =
|
||||
SessionSearchIndex::open_existing(db_path).map_err(|e| io::Error::other(e.to_string()))?;
|
||||
index
|
||||
.delete_meta(META_KEY_LAST_BOOTSTRAP)
|
||||
.map_err(|e| io::Error::other(e.to_string()))
|
||||
}
|
||||
|
||||
/// Determine whether the local index is stale enough to warrant downloading
|
||||
/// the remote copy.
|
||||
///
|
||||
|
|
|
|||
|
|
@ -135,8 +135,9 @@ pub(crate) struct SessionHarnessMetrics {
|
|||
pub memory_enabled: bool,
|
||||
pub auto_update: Option<bool>,
|
||||
pub cwd: String,
|
||||
pub skills_config: xai_grok_agent::prompt::skills::SkillsConfig,
|
||||
/// Resolved vendor-compat config, so recorded skill / AGENTS.md names match
|
||||
/// Filled from the built agent's bridge so `into_event` doesn't re-walk the disk.
|
||||
pub skill_names: Vec<String>,
|
||||
/// Resolved vendor-compat config, so recorded AGENTS.md names match
|
||||
/// what the session actually discovers.
|
||||
pub compat: xai_grok_tools::types::compat::CompatConfig,
|
||||
pub plugin_registry: Option<std::sync::Arc<xai_grok_agent::plugins::PluginRegistry>>,
|
||||
|
|
@ -189,16 +190,6 @@ impl SessionHarnessMetrics {
|
|||
.map(|n| n.to_string_lossy().into_owned())
|
||||
})
|
||||
.collect();
|
||||
let skill_names = xai_grok_agent::prompt::skills::list_skills_with_plugins(
|
||||
Some(&self.cwd),
|
||||
&self.skills_config,
|
||||
self.plugin_registry.as_deref(),
|
||||
self.compat,
|
||||
)
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|s| s.name)
|
||||
.collect();
|
||||
SessionHarness {
|
||||
session_id: self.session_id,
|
||||
client_identifier: self.client_identifier,
|
||||
|
|
@ -207,7 +198,7 @@ impl SessionHarnessMetrics {
|
|||
permission_mode: self.permission_mode,
|
||||
mcp_server_names: self.mcp_server_names,
|
||||
plugin_names: self.plugin_names,
|
||||
skill_names,
|
||||
skill_names: self.skill_names,
|
||||
lsp_server_names: self.lsp_server_names,
|
||||
hook_names,
|
||||
agents_md_dir_names,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
|
|||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
SubagentSpawnContext {
|
||||
lsp: None,
|
||||
process_scope: None,
|
||||
parent_max_turns: None,
|
||||
client_hooks: Default::default(),
|
||||
sampling_config: xai_grok_sampler::SamplerConfig {
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ use xai_grok_paths::AbsPathBuf;
|
|||
use xai_grok_workspace::file_system::{AsyncFileSystem, AsyncFsWrapper};
|
||||
use xai_grok_workspace::session::file_state::FileStateHandle;
|
||||
use xai_hunk_tracker::HunkTrackerHandle;
|
||||
use xai_tty_utils::ProcessScope;
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TaskOutputTokenBudget {
|
||||
inner: Arc<parking_lot::Mutex<TaskOutputTokenBudgetState>>,
|
||||
|
|
@ -216,6 +217,10 @@ pub struct ToolContext {
|
|||
pub blocking_wait_depth: Arc<BlockingWaitState>,
|
||||
pub task_output_token_budget: Option<TaskOutputTokenBudget>,
|
||||
pub(crate) sampler_retry_only_before_output: bool,
|
||||
/// This session's child-process reaper, set at session spawn; `None` for
|
||||
/// contexts without one (subagents, defaults). Spawn sites enroll children
|
||||
/// into it; enrolled children are killed when the session closes.
|
||||
pub process_scope: Option<ProcessScope>,
|
||||
}
|
||||
impl ToolContext {
|
||||
pub(crate) fn clamp_task_model_request(
|
||||
|
|
@ -281,6 +286,7 @@ impl ToolContext {
|
|||
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
|
||||
task_output_token_budget: None,
|
||||
sampler_retry_only_before_output: false,
|
||||
process_scope: None,
|
||||
}
|
||||
}
|
||||
pub fn with_preloaded_env(
|
||||
|
|
@ -321,6 +327,7 @@ impl ToolContext {
|
|||
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
|
||||
task_output_token_budget: None,
|
||||
sampler_retry_only_before_output: false,
|
||||
process_scope: None,
|
||||
}
|
||||
}
|
||||
pub fn with_file_state_handle(mut self, handle: FileStateHandle) -> Self {
|
||||
|
|
@ -414,6 +421,7 @@ mod tests {
|
|||
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
|
||||
task_output_token_budget: None,
|
||||
sampler_retry_only_before_output: false,
|
||||
process_scope: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -625,22 +625,182 @@ pub async fn save_mcp_disabled_tools(server_name: &str, disabled_tools: &[String
|
|||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist the enabled/disabled state for a single MCP server.
|
||||
/// Persist enable/disable for one MCP server in user config.
|
||||
///
|
||||
/// Uses the top-level `disabled_mcp_servers` array in `~/.grok/config.toml`.
|
||||
/// For local servers that have a `[mcp_servers.X]` entry, also sets/clears
|
||||
/// the `enabled` field so `to_acp_mcp_server()` respects it at load time.
|
||||
/// Always updates user `disabled_mcp_servers` (and user
|
||||
/// `[mcp_servers.<name>].enabled` when present). On **enable** only, also clears
|
||||
/// sticky `enabled = false` on the **nearest** project definition that defines
|
||||
/// the server (cwd-nearest wins; shadowed ancestors are left alone) — never
|
||||
/// writes project configs on disable.
|
||||
///
|
||||
/// Uses process cwd for project unstick; session callers should prefer
|
||||
/// [`save_mcp_server_enabled_in`] with the session cwd.
|
||||
pub async fn save_mcp_server_enabled(server_name: &str, enabled: bool) -> Result<()> {
|
||||
let path = config_path();
|
||||
let mut root: TomlValue = match tokio::fs::read_to_string(&path).await {
|
||||
Ok(s) => toml::from_str(&s).unwrap_or(TomlValue::Table(TomlMap::new())),
|
||||
Err(_) => TomlValue::Table(TomlMap::new()),
|
||||
let cwd = std::env::current_dir().unwrap_or_default();
|
||||
save_mcp_server_enabled_in(server_name, enabled, &cwd)
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
/// Like [`save_mcp_server_enabled`], with explicit cwd for project config walks.
|
||||
pub async fn save_mcp_server_enabled_in(
|
||||
server_name: &str,
|
||||
enabled: bool,
|
||||
cwd: &std::path::Path,
|
||||
) -> Result<Vec<PathBuf>> {
|
||||
let mut modified = Vec::new();
|
||||
|
||||
let user_path = config_path();
|
||||
if write_toml_table_if_changed(&user_path, |table| {
|
||||
apply_mcp_server_enabled(table, server_name, enabled);
|
||||
})
|
||||
.await?
|
||||
{
|
||||
modified.push(user_path);
|
||||
}
|
||||
|
||||
// Enable-only: unstick the nearest (winning) project def with
|
||||
// `enabled = false` via toml_edit (preserves comments/layout). Disable is
|
||||
// personal (user `disabled_mcp_servers`) and must not dirty shared files.
|
||||
if enabled
|
||||
&& let Some(path) = nearest_project_mcp_definition(cwd, server_name)
|
||||
&& clear_sticky_project_disabled_at(&path, server_name).await?
|
||||
{
|
||||
modified.push(path);
|
||||
}
|
||||
|
||||
Ok(modified)
|
||||
}
|
||||
|
||||
/// User config only — no project unstick.
|
||||
///
|
||||
/// Use after delete (or similar) when the toggle path dirtied
|
||||
/// `disabled_mcp_servers` but shared project configs must stay untouched.
|
||||
pub async fn save_user_mcp_server_enabled(server_name: &str, enabled: bool) -> Result<()> {
|
||||
write_toml_table_if_changed(&config_path(), |table| {
|
||||
apply_mcp_server_enabled(table, server_name, enabled);
|
||||
})
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
/// Nearest project config defining `server_name` (cwd last → reverse; nearest wins).
|
||||
fn nearest_project_mcp_definition(cwd: &std::path::Path, server_name: &str) -> Option<PathBuf> {
|
||||
crate::config::find_project_configs(cwd)
|
||||
.into_iter()
|
||||
.rev()
|
||||
.find(|path| mcp_server_defined_at(path, server_name))
|
||||
}
|
||||
|
||||
/// Apply `f`, write only if the serialized table changed. Returns whether written.
|
||||
///
|
||||
/// Aligns with [`super::persist::save_config`] safety: refuse unparseable
|
||||
/// files (no wipe-to-empty), unique tmp + mode preserve via
|
||||
/// [`super::persist::atomic_write_string`], and the user-config write lock.
|
||||
async fn write_toml_table_if_changed(
|
||||
path: &std::path::Path,
|
||||
f: impl FnOnce(&mut TomlMap<String, TomlValue>),
|
||||
) -> Result<bool> {
|
||||
let is_user = path == config_path().as_path();
|
||||
let _guard = if is_user {
|
||||
Some(super::persist::lock_config_writes().await)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let original = match tokio::fs::read_to_string(path).await {
|
||||
Ok(s) => s,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound && is_user => String::new(),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!("failed to read {}: {e}", path.display()));
|
||||
}
|
||||
};
|
||||
let mut root: TomlValue = if original.is_empty() {
|
||||
TomlValue::Table(TomlMap::new())
|
||||
} else {
|
||||
match toml::from_str(&original) {
|
||||
Ok(v) => v,
|
||||
Err(parse_err) => {
|
||||
return Err(anyhow::anyhow!(
|
||||
"refusing to overwrite unparseable {}: {}; fix the syntax before retrying",
|
||||
path.display(),
|
||||
parse_err
|
||||
));
|
||||
}
|
||||
}
|
||||
};
|
||||
let table = root
|
||||
.as_table_mut()
|
||||
.ok_or_else(|| anyhow::anyhow!("config root is not a table"))?;
|
||||
f(table);
|
||||
let toml_str = toml::to_string_pretty(&root)?;
|
||||
// Normalize empty original so first enable/disable still writes when needed.
|
||||
let before = if original.is_empty() {
|
||||
toml::to_string_pretty(&TomlValue::Table(TomlMap::new()))?
|
||||
} else {
|
||||
// Re-serialize original for stable comparison (ignore formatting noise).
|
||||
match toml::from_str::<TomlValue>(&original) {
|
||||
Ok(v) => toml::to_string_pretty(&v).unwrap_or(original),
|
||||
Err(_) => original,
|
||||
}
|
||||
};
|
||||
if before == toml_str {
|
||||
return Ok(false);
|
||||
}
|
||||
super::persist::atomic_write_string(path, &toml_str)
|
||||
.map_err(|e| anyhow::anyhow!("failed to write {}: {e}", path.display()))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
// Update the `disabled_mcp_servers` list (source of truth for all servers).
|
||||
/// Flip sticky project `enabled = false` → true with toml_edit (comments kept).
|
||||
async fn clear_sticky_project_disabled_at(
|
||||
path: &std::path::Path,
|
||||
server_name: &str,
|
||||
) -> Result<bool> {
|
||||
let original = match tokio::fs::read_to_string(path).await {
|
||||
Ok(s) => s,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
||||
Err(e) => {
|
||||
return Err(anyhow::anyhow!("failed to read {}: {e}", path.display()));
|
||||
}
|
||||
};
|
||||
let mut doc: toml_edit::DocumentMut = original
|
||||
.parse()
|
||||
.map_err(|e| anyhow::anyhow!("refusing to rewrite unparseable {}: {e}", path.display()))?;
|
||||
|
||||
let Some(servers) = doc
|
||||
.get_mut("mcp_servers")
|
||||
.and_then(|item| item.as_table_like_mut())
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
let Some(entry) = servers.get_mut(server_name) else {
|
||||
return Ok(false);
|
||||
};
|
||||
let Some(server_table) = entry.as_table_like_mut() else {
|
||||
return Ok(false);
|
||||
};
|
||||
if server_table.get("enabled").and_then(|v| v.as_bool()) != Some(false) {
|
||||
return Ok(false);
|
||||
}
|
||||
server_table.insert("enabled", toml_edit::value(true));
|
||||
|
||||
let updated = doc.to_string();
|
||||
if updated == original {
|
||||
return Ok(false);
|
||||
}
|
||||
super::persist::atomic_write_string(path, &updated)
|
||||
.map_err(|e| anyhow::anyhow!("failed to write {}: {e}", path.display()))?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// Update user `disabled_mcp_servers` and, if present, per-server `enabled`.
|
||||
fn apply_mcp_server_enabled(
|
||||
table: &mut TomlMap<String, TomlValue>,
|
||||
server_name: &str,
|
||||
enabled: bool,
|
||||
) {
|
||||
let mut disabled_list: Vec<String> = table
|
||||
.get("disabled_mcp_servers")
|
||||
.and_then(|v| v.as_array())
|
||||
|
|
@ -667,14 +827,20 @@ pub async fn save_mcp_server_enabled(server_name: &str, enabled: bool) -> Result
|
|||
table.insert("disabled_mcp_servers".to_string(), TomlValue::Array(arr));
|
||||
}
|
||||
|
||||
let toml_str = toml::to_string_pretty(&root)?;
|
||||
let tmp = path.with_extension("toml.tmp");
|
||||
if let Some(parent) = path.parent() {
|
||||
let _ = tokio::fs::create_dir_all(parent).await;
|
||||
set_mcp_server_enabled_field(table, server_name, enabled);
|
||||
}
|
||||
|
||||
fn set_mcp_server_enabled_field(
|
||||
table: &mut TomlMap<String, TomlValue>,
|
||||
server_name: &str,
|
||||
enabled: bool,
|
||||
) {
|
||||
if let Some(servers) = table.get_mut("mcp_servers").and_then(|v| v.as_table_mut())
|
||||
&& let Some(entry) = servers.get_mut(server_name)
|
||||
&& let Some(server_table) = entry.as_table_mut()
|
||||
{
|
||||
server_table.insert("enabled".to_string(), TomlValue::Boolean(enabled));
|
||||
}
|
||||
tokio::fs::write(&tmp, &toml_str).await?;
|
||||
tokio::fs::rename(&tmp, &path).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Upsert an MCP server entry in `~/.grok/config.toml`.
|
||||
|
|
@ -1453,6 +1619,60 @@ pub fn disabled_mcp_server_names(cwd: &std::path::Path) -> std::collections::Has
|
|||
disabled
|
||||
}
|
||||
|
||||
/// Names `grok mcp enable`/`disable` may target: user/project TOML (including
|
||||
/// setup-required/invalid entries that session merge drops), the user
|
||||
/// `disabled_mcp_servers` list, compat JSON (`.mcp.json`, Claude, Cursor),
|
||||
/// **plugin** MCP servers (same discovery as doctor/`/mcps`), and legacy
|
||||
/// managed `grok_com_*` (special-cased in the CLI).
|
||||
///
|
||||
/// Does **not** include gateway connectors (`managed_gateway:…`); those use
|
||||
/// `disabled_mcp_tools.__managed_gateway_connectors` via the `/mcps` Space.
|
||||
pub fn cli_known_mcp_server_names(cwd: &std::path::Path) -> std::collections::HashSet<String> {
|
||||
let mut names = disabled_mcp_server_names(cwd);
|
||||
// Full TOML key set (list parity) — merge drops setup-required/invalid.
|
||||
names.extend(all_toml_mcp_server_names(cwd));
|
||||
|
||||
// Doctor/Space path: resolved TOML + plugins + Claude + Cursor + `.mcp.json`.
|
||||
let registry = load_cli_plugin_registry(cwd);
|
||||
let compat = CompatConfig::default();
|
||||
for (server, _) in crate::session::managed_mcp::merge_managed_mcp_servers_sourced(
|
||||
cwd,
|
||||
Some(®istry),
|
||||
&compat,
|
||||
) {
|
||||
let name = crate::session::managed_mcp::mcp_server_name(&server);
|
||||
if !name.is_empty() {
|
||||
names.insert(name.to_string());
|
||||
}
|
||||
}
|
||||
names
|
||||
}
|
||||
|
||||
/// Plugin registry for one-shot CLI discovery (matches mcp doctor gating).
|
||||
fn load_cli_plugin_registry(cwd: &std::path::Path) -> xai_grok_agent::plugins::PluginRegistry {
|
||||
let trust_store = xai_grok_agent::plugins::TrustStore::load();
|
||||
let mut plugins_cfg: crate::agent::config::PluginsConfig =
|
||||
crate::config::load_effective_config()
|
||||
.ok()
|
||||
.and_then(|t| t.get("plugins").and_then(|v| v.clone().try_into().ok()))
|
||||
.unwrap_or_default();
|
||||
plugins_cfg.merge_claude_enabled_plugins(Some(cwd));
|
||||
let mut plugin_config = plugins_cfg.to_discovery_config();
|
||||
let project_trusted = crate::agent::folder_trust::resolve_and_record(cwd, None, false);
|
||||
let discovered = xai_grok_agent::plugins::discover_plugins(
|
||||
Some(cwd),
|
||||
&plugin_config,
|
||||
&trust_store,
|
||||
project_trusted,
|
||||
);
|
||||
plugin_config.populate_plugin_lists(&discovered);
|
||||
xai_grok_agent::plugins::PluginRegistry::from_discovered(
|
||||
discovered,
|
||||
&plugin_config.disabled,
|
||||
&plugin_config.enabled,
|
||||
)
|
||||
}
|
||||
|
||||
fn config_path() -> PathBuf {
|
||||
crate::util::grok_home::grok_home().join("config.toml")
|
||||
}
|
||||
|
|
@ -2307,5 +2527,162 @@ enabled = false
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_mcp_server_enabled_updates_array_and_per_server_field() {
|
||||
let mut root: TomlValue = toml::from_str(
|
||||
r#"
|
||||
[mcp_servers.local]
|
||||
command = "npx"
|
||||
enabled = true
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let table = root.as_table_mut().unwrap();
|
||||
|
||||
apply_mcp_server_enabled(table, "local", false);
|
||||
let disabled = table
|
||||
.get("disabled_mcp_servers")
|
||||
.and_then(|v| v.as_array())
|
||||
.expect("disabled_mcp_servers array");
|
||||
assert_eq!(disabled.len(), 1);
|
||||
assert_eq!(disabled[0].as_str(), Some("local"));
|
||||
assert_eq!(
|
||||
table
|
||||
.get("mcp_servers")
|
||||
.and_then(|v| v.as_table())
|
||||
.and_then(|s| s.get("local"))
|
||||
.and_then(|v| v.as_table())
|
||||
.and_then(|s| s.get("enabled"))
|
||||
.and_then(|v| v.as_bool()),
|
||||
Some(false)
|
||||
);
|
||||
|
||||
apply_mcp_server_enabled(table, "local", true);
|
||||
assert!(table.get("disabled_mcp_servers").is_none());
|
||||
assert_eq!(
|
||||
table
|
||||
.get("mcp_servers")
|
||||
.and_then(|v| v.as_table())
|
||||
.and_then(|s| s.get("local"))
|
||||
.and_then(|v| v.as_table())
|
||||
.and_then(|s| s.get("enabled"))
|
||||
.and_then(|v| v.as_bool()),
|
||||
Some(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_mcp_server_enabled_managed_name_only_updates_array() {
|
||||
let mut root: TomlValue = toml::from_str("disabled_mcp_servers = []\n").unwrap();
|
||||
let table = root.as_table_mut().unwrap();
|
||||
|
||||
apply_mcp_server_enabled(table, "grok_com_slack", false);
|
||||
let disabled = table
|
||||
.get("disabled_mcp_servers")
|
||||
.and_then(|v| v.as_array())
|
||||
.expect("disabled_mcp_servers array");
|
||||
assert_eq!(disabled.len(), 1);
|
||||
assert_eq!(disabled[0].as_str(), Some("grok_com_slack"));
|
||||
assert!(table.get("mcp_servers").is_none());
|
||||
|
||||
apply_mcp_server_enabled(table, "grok_com_slack", true);
|
||||
assert!(table.get("disabled_mcp_servers").is_none());
|
||||
assert!(table.get("mcp_servers").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clear_sticky_project_disabled_at_only_flips_false() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = tmp.path().join("config.toml");
|
||||
std::fs::write(
|
||||
&path,
|
||||
r#"
|
||||
[mcp_servers.proj]
|
||||
command = "npx"
|
||||
enabled = false
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(
|
||||
clear_sticky_project_disabled_at(&path, "proj")
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
let body = std::fs::read_to_string(&path).unwrap();
|
||||
assert!(body.contains("enabled = true"), "{body}");
|
||||
// Already true: no write.
|
||||
assert!(
|
||||
!clear_sticky_project_disabled_at(&path, "proj")
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
/// Nested project configs both sticky-disable the same server; enable
|
||||
/// unstick must rewrite cwd-nearest only (not the shadowed ancestor).
|
||||
#[tokio::test]
|
||||
async fn enable_unstick_only_touches_nearest_project_definition() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
git2::Repository::init(tmp.path()).unwrap();
|
||||
let nested = tmp.path().join("pkg");
|
||||
std::fs::create_dir_all(nested.join(".grok")).unwrap();
|
||||
std::fs::create_dir_all(tmp.path().join(".grok")).unwrap();
|
||||
|
||||
let sticky = r#"
|
||||
# keep me
|
||||
[mcp_servers.svc]
|
||||
command = "npx"
|
||||
enabled = false
|
||||
"#;
|
||||
let ancestor = tmp.path().join(".grok").join("config.toml");
|
||||
let nearer = nested.join(".grok").join("config.toml");
|
||||
std::fs::write(&ancestor, sticky).unwrap();
|
||||
std::fs::write(&nearer, sticky).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
nearest_project_mcp_definition(&nested, "svc").as_ref(),
|
||||
Some(&nearer)
|
||||
);
|
||||
|
||||
let path = nearest_project_mcp_definition(&nested, "svc").unwrap();
|
||||
clear_sticky_project_disabled_at(&path, "svc")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let nearer_body = std::fs::read_to_string(&nearer).unwrap();
|
||||
assert!(
|
||||
nearer_body.contains("enabled = true"),
|
||||
"nearest should be unstuck: {nearer_body}"
|
||||
);
|
||||
assert!(
|
||||
nearer_body.contains("# keep me"),
|
||||
"toml_edit must preserve comments: {nearer_body}"
|
||||
);
|
||||
let ancestor_body = std::fs::read_to_string(&ancestor).unwrap();
|
||||
assert!(
|
||||
ancestor_body.contains("enabled = false"),
|
||||
"shadowed ancestor must stay sticky: {ancestor_body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_toml_table_if_changed_refuses_unparseable() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let path = tmp.path().join("broken.toml");
|
||||
std::fs::write(&path, "not = [valid\n").unwrap();
|
||||
let err = write_toml_table_if_changed(&path, |_t| {})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("unparseable"),
|
||||
"expected refuse, got {err}"
|
||||
);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&path).unwrap(),
|
||||
"not = [valid\n",
|
||||
"file must be left intact"
|
||||
);
|
||||
}
|
||||
|
||||
// === merge_section tests ===
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue