Synced from monorepo
Synced from monorepo Changes: - Release a shell session's resources in one drop - Make the tools blocking-wait cap client-configurable and self-describing - Recognize API "exceeds budget" errors as context overflow - Retry /btw on model overload - Carry running background tasks and subagents across compaction - Require round-trip time for SDK liveness checks - Background-subagent completion reminders with a selectable delivery surface - Make a PTY shell reap itself until it reaches the registry - Recover the OS error code from a TLS-phase connection reset - Consume the attached-client signal and report why idle is withheld - Treat `.grok/sandbox.toml` edits as protected so auto mode prompts before writing - Surface history/search in the Ctrl+. cheatsheet and keep it working in history view - Delete sessions from the dashboard and welcome list - Release a session's activity record when the session ends - Stop charging auth-retry budget for fail-closed 401s; reset it across suspends - Scope skills watches on project vendor roots - Make [stop] cancel in-flight compaction - Make the leader soak measure the leader, not its harness Source-Revision: 8d69c91f02bcacf01e98d5aebbf2f92547c45738
This commit is contained in:
parent
dd04f397b1
commit
a422116582
165 changed files with 15161 additions and 1969 deletions
|
|
@ -92,6 +92,8 @@ urlencoding = "2"
|
|||
xai-fast-worktree = { path = "../xai-fast-worktree", features = ["metadata"] }
|
||||
# tonic: Status/Code mapping for deploy errors in workspace_ops.
|
||||
tonic = { workspace = true }
|
||||
sha1 = { workspace = true, optional = true }
|
||||
zip = { workspace = true, optional = true }
|
||||
|
||||
xai-fsnotify = { path = "../xai-fsnotify" }
|
||||
clap = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ use std::time::Instant;
|
|||
use dashmap::DashMap;
|
||||
use xai_file_utils::events::{Event, EventWriter, ToolCompletedSource, ToolOutcome};
|
||||
use xai_file_utils::queue::UploadQueueStats;
|
||||
use xai_tool_protocol::{ToolServerLifecycleStatus, ToolServerStatusPayload};
|
||||
use xai_tool_protocol::{IdleWithholdReason, ToolServerLifecycleStatus, ToolServerStatusPayload};
|
||||
|
||||
const LIFECYCLE_NONE: u8 = 0;
|
||||
const LIFECYCLE_DRAINING: u8 = 1;
|
||||
|
|
@ -88,10 +88,25 @@ pub struct ActivityTracker {
|
|||
/// Window (ms) recent preview-proxy traffic withholds idle for; defaults to
|
||||
/// [`PREVIEW_ACTIVITY_WINDOW_MS`], overridable via the builder.
|
||||
preview_activity_window_ms: u64,
|
||||
/// Epoch ms of the last scraped preview-proxy activity (`0` = none). Fed by
|
||||
/// the preview-activity scraper (`preview_supervisor`); withholds idle within
|
||||
/// Epoch ms the pane's own status poll was last observed (`0` = none). Fed
|
||||
/// by the preview-activity scraper; withholds idle within
|
||||
/// [`preview_activity_window_ms`](Self::preview_activity_window_ms).
|
||||
last_preview_activity_ms: AtomicU64,
|
||||
///
|
||||
/// The *observation* time, not the proxy's stamp: the two processes are not
|
||||
/// clock-coupled, and only the local clock is comparable with the rest.
|
||||
last_preview_status_ms: AtomicU64,
|
||||
/// Epoch ms real app traffic was last observed (`0` = none).
|
||||
last_preview_routed_ms: AtomicU64,
|
||||
/// Open preview WebSocket (HMR) tunnels as of the last scrape. Nonzero ⇒ a
|
||||
/// client is attached, which no activity stamp would reveal.
|
||||
preview_ws_tunnels_open: AtomicU64,
|
||||
/// In-flight `Routed` preview requests as of the last scrape.
|
||||
preview_routed_in_flight: AtomicU64,
|
||||
/// Epoch ms this process started — the floor of the withhold anchor, so a
|
||||
/// young or freshly-restored workspace is never treated as long-idle.
|
||||
/// Distinct from [`Self::started_at`], a monotonic `Instant` that cannot be
|
||||
/// compared against the epoch stamps around it.
|
||||
started_at_ms: u64,
|
||||
|
||||
sessions: DashMap<String, SessionActivity>,
|
||||
/// call_id → session_id so `tool_call_completed` can decrement
|
||||
|
|
@ -168,7 +183,11 @@ impl ActivityTracker {
|
|||
durability_idle_hold_max_ms,
|
||||
idle_ignores_background: false,
|
||||
preview_activity_window_ms: PREVIEW_ACTIVITY_WINDOW_MS,
|
||||
last_preview_activity_ms: AtomicU64::new(0),
|
||||
last_preview_status_ms: AtomicU64::new(0),
|
||||
last_preview_routed_ms: AtomicU64::new(0),
|
||||
preview_ws_tunnels_open: AtomicU64::new(0),
|
||||
preview_routed_in_flight: AtomicU64::new(0),
|
||||
started_at_ms: now_ms(),
|
||||
sessions: DashMap::new(),
|
||||
call_to_session: DashMap::new(),
|
||||
prune_window_ms: prune_window.as_millis() as u64,
|
||||
|
|
@ -217,15 +236,54 @@ impl ActivityTracker {
|
|||
self.notify.clone()
|
||||
}
|
||||
|
||||
/// Record fresh preview-proxy traffic: withholds `idle_since_ms` for
|
||||
/// Record fresh `Routed` preview traffic. Withholds `idle_since_ms` for
|
||||
/// [`preview_activity_window_ms`](Self::preview_activity_window_ms) and wakes
|
||||
/// the status publisher so the renewed "active" status reaches the server promptly.
|
||||
pub fn note_preview_activity(&self) {
|
||||
self.last_preview_activity_ms
|
||||
pub fn note_preview_routed_activity(&self) {
|
||||
self.last_preview_routed_ms
|
||||
.store(now_ms(), Ordering::Relaxed);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
/// Record a fresh preview status poll. Withholds idle exactly as routed
|
||||
/// traffic does today, but is tracked separately: the poll continues at the
|
||||
/// same cadence whether or not anyone is watching.
|
||||
pub fn note_preview_status_activity(&self) {
|
||||
self.last_preview_status_ms
|
||||
.store(now_ms(), Ordering::Relaxed);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
/// Mirror the proxy's attached-client counters. Absolute values, not edges,
|
||||
/// so a missed scrape self-corrects on the next one.
|
||||
pub fn set_preview_attached(&self, ws_tunnels_open: u64, routed_in_flight: u64) {
|
||||
let was_attached = self.has_preview_client_attached();
|
||||
self.preview_ws_tunnels_open
|
||||
.store(ws_tunnels_open, Ordering::Relaxed);
|
||||
self.preview_routed_in_flight
|
||||
.store(routed_in_flight, Ordering::Relaxed);
|
||||
// The scraper calls this every tick; the common case is 0 → 0.
|
||||
if was_attached != self.has_preview_client_attached() {
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
/// An open WebSocket tunnel or a routed request in flight. Never a poll.
|
||||
fn has_preview_client_attached(&self) -> bool {
|
||||
self.preview_ws_tunnels_open.load(Ordering::Relaxed) > 0
|
||||
|| self.preview_routed_in_flight.load(Ordering::Relaxed) > 0
|
||||
}
|
||||
|
||||
pub fn preview_ws_tunnels_open(&self) -> u64 {
|
||||
self.preview_ws_tunnels_open.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Window recent preview activity withholds idle for — the horizon the
|
||||
/// stamps decay over, and the one the scraper ages an absent proxy against.
|
||||
pub fn preview_activity_window_ms(&self) -> u64 {
|
||||
self.preview_activity_window_ms
|
||||
}
|
||||
|
||||
/// Pending upload-queue items (0 when no queue is coupled).
|
||||
fn upload_queue_pending(&self) -> u64 {
|
||||
self.upload_queue_stats
|
||||
|
|
@ -265,15 +323,33 @@ impl ActivityTracker {
|
|||
let (queue_pending, queue_pending_bytes, queue_inflight, breaker, drain_started) =
|
||||
self.drain_status_fields();
|
||||
let (producers, durability_withhold) = self.durability_gate(queue_pending, breaker);
|
||||
// Withhold idle on durability work OR recent preview traffic, decided here
|
||||
let now = now_ms();
|
||||
let (preview_withhold, preview_reason, preview_anchor) = self.preview_withholds_idle(now);
|
||||
// Withhold idle on durability work OR preview activity, decided here
|
||||
// once so both snapshot paths agree (preview has no hold cap; 12h VM TTL backstops).
|
||||
let withhold_idle = durability_withhold || self.preview_withholds_idle(now_ms());
|
||||
let withhold_idle = durability_withhold || preview_withhold;
|
||||
// Invariants: no reason while genuinely busy (`idle_since == 0` — the
|
||||
// work is the cause, not a concurrent poll); durability outranks
|
||||
// preview; every reason carries a stamp.
|
||||
let (withhold_reason, withhold_since_ms) = if idle_since == 0 {
|
||||
(None, None)
|
||||
} else if durability_withhold {
|
||||
(
|
||||
Some(IdleWithholdReason::Durability),
|
||||
Some(self.durability_busy_since_ms.load(Ordering::Relaxed)),
|
||||
)
|
||||
} else {
|
||||
(preview_reason, preview_reason.map(|_| preview_anchor))
|
||||
};
|
||||
DurabilityPayloadFields {
|
||||
idle_since_ms: if idle_since == 0 || withhold_idle {
|
||||
None
|
||||
} else {
|
||||
Some(idle_since)
|
||||
},
|
||||
withhold_reason,
|
||||
withhold_since_ms,
|
||||
preview_ws_tunnels_open: self.preview_ws_tunnels_open().min(u32::MAX as u64) as u32,
|
||||
upload_queue_pending: queue_pending,
|
||||
upload_queue_pending_bytes: queue_pending_bytes,
|
||||
upload_queue_inflight: queue_inflight,
|
||||
|
|
@ -316,13 +392,41 @@ impl ActivityTracker {
|
|||
(producers, !hold_expired)
|
||||
}
|
||||
|
||||
/// Whether recent preview-proxy traffic should currently withhold idle.
|
||||
fn preview_withholds_idle(&self, now: u64) -> bool {
|
||||
preview_activity_withholds_idle(
|
||||
/// Whether preview activity should withhold idle, and on what grounds.
|
||||
///
|
||||
/// Returns `(withhold, reason, anchor)`, where the anchor is the epoch-ms
|
||||
/// the current hold is measured from. Tiers are checked strongest first;
|
||||
/// all three withhold identically today, only the accounting differs.
|
||||
fn preview_withholds_idle(&self, now: u64) -> (bool, Option<IdleWithholdReason>, u64) {
|
||||
// Including process start means a young or freshly-restored workspace
|
||||
// can never look long-idle — the process restarts on restore, so this
|
||||
// covers revived sessions without a separate minimum-age rule.
|
||||
let anchor = self
|
||||
.last_preview_routed_ms
|
||||
.load(Ordering::Relaxed)
|
||||
.max(self.last_call_completed_ms.load(Ordering::Relaxed))
|
||||
.max(self.started_at_ms);
|
||||
|
||||
if self.has_preview_client_attached() {
|
||||
return (true, Some(IdleWithholdReason::PreviewAttached), anchor);
|
||||
}
|
||||
if preview_activity_withholds_idle(
|
||||
now,
|
||||
self.last_preview_activity_ms.load(Ordering::Relaxed),
|
||||
self.last_preview_routed_ms.load(Ordering::Relaxed),
|
||||
self.preview_activity_window_ms,
|
||||
)
|
||||
) {
|
||||
return (true, Some(IdleWithholdReason::PreviewRouted), anchor);
|
||||
}
|
||||
if preview_activity_withholds_idle(
|
||||
now,
|
||||
self.last_preview_status_ms.load(Ordering::Relaxed),
|
||||
self.preview_activity_window_ms,
|
||||
) {
|
||||
// Holds, but never advances the anchor: a poll must not reset a
|
||||
// clock meant to measure real use.
|
||||
return (true, Some(IdleWithholdReason::PreviewStatusOnly), anchor);
|
||||
}
|
||||
(false, None, anchor)
|
||||
}
|
||||
|
||||
/// Whether any tracked session currently has an active turn (the aggregate
|
||||
|
|
@ -522,12 +626,16 @@ impl ActivityTracker {
|
|||
count
|
||||
}
|
||||
|
||||
/// Mark a session as ended: clear turn-active flag and notify waiters.
|
||||
///
|
||||
/// Called by [`crate::handle::WorkspaceHandle::on_session_ended()`] when
|
||||
/// a `HookEvent::SessionEnded` arrives from the server.
|
||||
/// Releases the session's entry. One with calls in flight keeps it: the
|
||||
/// swap policy reads that count, and the idle prune collects it later.
|
||||
pub fn session_ended(&self, session_id: &str) {
|
||||
if let Some(session) = self.sessions.get(session_id) {
|
||||
let released = self
|
||||
.sessions
|
||||
.remove_if(session_id, |_, s| {
|
||||
s.active_tool_calls.load(Ordering::Acquire) == 0
|
||||
})
|
||||
.is_some();
|
||||
if !released && let Some(session) = self.sessions.get(session_id) {
|
||||
session.turn_active.store(false, Ordering::Release);
|
||||
}
|
||||
self.notify.notify_waiters();
|
||||
|
|
@ -640,6 +748,11 @@ impl ActivityTracker {
|
|||
}
|
||||
}
|
||||
|
||||
/// Resident session records. Does not prune, unlike [`Self::known_sessions`].
|
||||
pub fn session_count(&self) -> usize {
|
||||
self.sessions.len()
|
||||
}
|
||||
|
||||
/// Returns live session IDs. As a side-effect, prunes sessions
|
||||
/// that have been idle longer than the configured prune window.
|
||||
pub fn known_sessions(&self) -> Vec<String> {
|
||||
|
|
@ -731,6 +844,11 @@ impl ActivityTracker {
|
|||
drain_started_ms: d.drain_started_ms,
|
||||
turn_active,
|
||||
idle_ignores_background: self.idle_ignores_background,
|
||||
withhold_reason: d.withhold_reason,
|
||||
withhold_since_ms: d.withhold_since_ms,
|
||||
// No ceilings configured yet, so a hold can never be capped.
|
||||
withhold_capped: false,
|
||||
preview_ws_tunnels_open: d.preview_ws_tunnels_open,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -783,6 +901,11 @@ impl ActivityTracker {
|
|||
drain_started_ms: d.drain_started_ms,
|
||||
turn_active: self.any_turn_active(),
|
||||
idle_ignores_background: self.idle_ignores_background,
|
||||
withhold_reason: d.withhold_reason,
|
||||
withhold_since_ms: d.withhold_since_ms,
|
||||
// No ceilings configured yet, so a hold can never be capped.
|
||||
withhold_capped: false,
|
||||
preview_ws_tunnels_open: d.preview_ws_tunnels_open,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -791,6 +914,9 @@ impl ActivityTracker {
|
|||
/// [`ActivityTracker::durability_payload_fields`].
|
||||
struct DurabilityPayloadFields {
|
||||
idle_since_ms: Option<u64>,
|
||||
withhold_reason: Option<IdleWithholdReason>,
|
||||
withhold_since_ms: Option<u64>,
|
||||
preview_ws_tunnels_open: u32,
|
||||
upload_queue_pending: u32,
|
||||
upload_queue_pending_bytes: u64,
|
||||
upload_queue_inflight: u32,
|
||||
|
|
@ -1287,29 +1413,195 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn note_preview_activity_withholds_then_resumes_idle() {
|
||||
for (label, note) in [
|
||||
(
|
||||
"routed",
|
||||
&ActivityTracker::note_preview_routed_activity as &dyn Fn(&ActivityTracker),
|
||||
),
|
||||
("status", &ActivityTracker::note_preview_status_activity),
|
||||
] {
|
||||
let t = ActivityTracker::new();
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_some(),
|
||||
"{label}: an idle tracker reports idle before any preview activity"
|
||||
);
|
||||
|
||||
note(&t);
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_none(),
|
||||
"{label}: recent preview activity must withhold idle"
|
||||
);
|
||||
assert!(
|
||||
t.snapshot_session("any").idle_since_ms.is_none(),
|
||||
"{label}: the per-session payload must withhold idle too"
|
||||
);
|
||||
|
||||
let stale = now_ms().saturating_sub(PREVIEW_ACTIVITY_WINDOW_MS + 1_000);
|
||||
t.last_preview_routed_ms.store(stale, Ordering::Relaxed);
|
||||
t.last_preview_status_ms.store(stale, Ordering::Relaxed);
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_some(),
|
||||
"{label}: idle must resume once the preview window decays"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn withhold_reason_reports_the_tier_that_is_holding() {
|
||||
let t = ActivityTracker::new();
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_some(),
|
||||
"an idle tracker reports idle before any preview activity"
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
None,
|
||||
"nothing holding ⇒ no reason"
|
||||
);
|
||||
|
||||
t.note_preview_activity();
|
||||
t.note_preview_status_activity();
|
||||
let s = t.snapshot();
|
||||
assert_eq!(
|
||||
s.withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewStatusOnly),
|
||||
"a bare status poll is the weakest tier"
|
||||
);
|
||||
assert!(
|
||||
s.withhold_since_ms.is_some(),
|
||||
"every reason carries a since-stamp so a reader can age the hold"
|
||||
);
|
||||
|
||||
t.note_preview_routed_activity();
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewRouted),
|
||||
"real app traffic outranks the status poll"
|
||||
);
|
||||
|
||||
t.set_preview_attached(1, 0);
|
||||
let s = t.snapshot();
|
||||
assert_eq!(
|
||||
s.withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewAttached),
|
||||
"an open tunnel outranks everything below it"
|
||||
);
|
||||
assert_eq!(s.preview_ws_tunnels_open, 1);
|
||||
assert!(!s.withhold_capped, "no ceilings configured yet");
|
||||
}
|
||||
|
||||
/// An HMR socket writes no activity stamps, so the counter must hold alone
|
||||
/// — which is exactly why the scraper must not clear it on one missed poll.
|
||||
#[test]
|
||||
fn attached_client_withholds_without_any_activity_stamp() {
|
||||
let t = ActivityTracker::new();
|
||||
assert!(t.snapshot().idle_since_ms.is_some());
|
||||
|
||||
t.set_preview_attached(1, 0);
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_none(),
|
||||
"recent preview activity must withhold idle"
|
||||
);
|
||||
assert!(
|
||||
t.snapshot_session("any").idle_since_ms.is_none(),
|
||||
"the per-session payload must withhold idle too"
|
||||
"an open WebSocket tunnel withholds idle by itself"
|
||||
);
|
||||
|
||||
t.last_preview_activity_ms.store(
|
||||
now_ms().saturating_sub(PREVIEW_ACTIVITY_WINDOW_MS + 1_000),
|
||||
Ordering::Relaxed,
|
||||
t.set_preview_attached(0, 1);
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewAttached),
|
||||
"a routed request in flight is equally an attached client"
|
||||
);
|
||||
|
||||
t.set_preview_attached(0, 0);
|
||||
assert!(
|
||||
t.snapshot().idle_since_ms.is_some(),
|
||||
"idle must resume once the preview window decays"
|
||||
"once detached with no stamp in window, idle resumes"
|
||||
);
|
||||
}
|
||||
|
||||
/// A ceiling will be measured from the anchor, so letting the pane's own
|
||||
/// poll reset it would make that ceiling unreachable.
|
||||
#[test]
|
||||
fn status_poll_does_not_advance_the_withhold_anchor() {
|
||||
let t = ActivityTracker::new();
|
||||
t.note_preview_routed_activity();
|
||||
let anchor = t.snapshot().withhold_since_ms.expect("routed holds");
|
||||
|
||||
std::thread::sleep(std::time::Duration::from_millis(5));
|
||||
t.note_preview_status_activity();
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_since_ms,
|
||||
Some(anchor),
|
||||
"a status poll leaves the anchor where the last real use put it"
|
||||
);
|
||||
}
|
||||
|
||||
/// A session busy with real work must not be attributed to the preview just
|
||||
/// because the pane happens to be polling it. Both look like a missing
|
||||
/// `idle_since_ms` on the wire, and conflating them would inflate the
|
||||
/// preview share of exactly the population this field exists to measure.
|
||||
#[test]
|
||||
fn a_genuinely_busy_session_reports_no_withhold_reason() {
|
||||
let t = ActivityTracker::new();
|
||||
t.note_preview_status_activity();
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewStatusOnly),
|
||||
"idle but polled ⇒ the poll is the reason"
|
||||
);
|
||||
|
||||
t.tool_call_started("c1", "read_file", Some("sess-a"));
|
||||
let s = t.snapshot();
|
||||
assert!(
|
||||
s.idle_since_ms.is_none(),
|
||||
"a tool call in flight withholds idle on its own"
|
||||
);
|
||||
assert_eq!(
|
||||
s.withhold_reason, None,
|
||||
"the tool call is the cause, not the concurrent poll"
|
||||
);
|
||||
assert_eq!(s.withhold_since_ms, None);
|
||||
|
||||
t.tool_call_completed("c1", Some("sess-a"), ToolOutcome::Success);
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewStatusOnly),
|
||||
"once the real work finishes, the poll is the reason again"
|
||||
);
|
||||
}
|
||||
|
||||
/// Durability is already bounded by its own cap, so it is the reason worth
|
||||
/// reporting when both hold.
|
||||
#[tokio::test]
|
||||
async fn durability_outranks_preview_in_the_reported_reason() {
|
||||
let t = ActivityTracker::new();
|
||||
let tasks = tokio_util::task::TaskTracker::new();
|
||||
t.set_producer_tasks(tasks.clone());
|
||||
|
||||
t.note_preview_status_activity();
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewStatusOnly),
|
||||
"preview alone reports the preview reason"
|
||||
);
|
||||
|
||||
// An in-flight producer engages the durability gate. Nothing else does
|
||||
// — a background task is not durable work.
|
||||
let gate = Arc::new(tokio::sync::Notify::new());
|
||||
let gate2 = gate.clone();
|
||||
let join = tasks.spawn(async move { gate2.notified().await });
|
||||
|
||||
let s = t.snapshot();
|
||||
assert_eq!(s.artifact_producers_inflight, 1, "the gate is engaged");
|
||||
assert_eq!(
|
||||
s.withhold_reason,
|
||||
Some(IdleWithholdReason::Durability),
|
||||
"durability outranks a concurrent preview hold"
|
||||
);
|
||||
assert!(
|
||||
s.withhold_since_ms.is_some(),
|
||||
"a durability hold reports its own busy-since stamp"
|
||||
);
|
||||
|
||||
gate.notify_one();
|
||||
join.await.expect("producer task must not panic");
|
||||
assert_eq!(
|
||||
t.snapshot().withhold_reason,
|
||||
Some(IdleWithholdReason::PreviewStatusOnly),
|
||||
"once durability clears, the preview hold is reported again"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -1317,13 +1609,13 @@ mod tests {
|
|||
fn configured_preview_window_overrides_default() {
|
||||
let configured = ActivityTracker::new().with_preview_activity_window_ms(500);
|
||||
configured
|
||||
.last_preview_activity_ms
|
||||
.last_preview_status_ms
|
||||
.store(now_ms().saturating_sub(1_000), Ordering::Relaxed);
|
||||
assert!(configured.snapshot().idle_since_ms.is_some());
|
||||
|
||||
let default = ActivityTracker::new();
|
||||
default
|
||||
.last_preview_activity_ms
|
||||
.last_preview_status_ms
|
||||
.store(now_ms().saturating_sub(1_000), Ordering::Relaxed);
|
||||
assert!(default.snapshot().idle_since_ms.is_none());
|
||||
}
|
||||
|
|
@ -1332,7 +1624,7 @@ mod tests {
|
|||
fn preview_activity_does_not_override_active_tool_call() {
|
||||
let t = ActivityTracker::new();
|
||||
t.tool_call_started("c1", "read_file", Some("sess-a"));
|
||||
t.note_preview_activity();
|
||||
t.note_preview_routed_activity();
|
||||
let s = t.snapshot();
|
||||
assert!(s.idle_since_ms.is_none());
|
||||
assert_eq!(
|
||||
|
|
@ -1704,16 +1996,16 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn session_ended_clears_turn_active() {
|
||||
fn session_ended_releases_an_idle_session() {
|
||||
let t = ActivityTracker::new();
|
||||
t.turn_started("sess-a", 3);
|
||||
let session = t.sessions.get("sess-a").expect("session should exist");
|
||||
assert!(session.turn_active.load(Ordering::Acquire));
|
||||
assert!(t.is_turn_active("sess-a"));
|
||||
|
||||
t.session_ended("sess-a");
|
||||
|
||||
assert!(
|
||||
!session.turn_active.load(Ordering::Acquire),
|
||||
"turn_active should be cleared after session_ended"
|
||||
!t.sessions.contains_key("sess-a"),
|
||||
"an ended session must not stay resident"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -1732,10 +2024,9 @@ mod tests {
|
|||
t.tool_call_started("c1", "read_file", Some("sess-a"));
|
||||
t.tool_call_completed("c1", None, ToolOutcome::Success);
|
||||
|
||||
// session_ended should not panic when turn was never active.
|
||||
t.session_ended("sess-a");
|
||||
let session = t.sessions.get("sess-a").expect("session should exist");
|
||||
assert!(!session.turn_active.load(Ordering::Acquire));
|
||||
|
||||
assert!(!t.sessions.contains_key("sess-a"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -1770,8 +2061,11 @@ mod tests {
|
|||
|
||||
t.session_ended("sess-a");
|
||||
|
||||
// turn_active cleared, but the in-flight tool call remains.
|
||||
assert!(!t.is_turn_active("sess-a"));
|
||||
assert!(
|
||||
t.sessions.contains_key("sess-a"),
|
||||
"a session with a call in flight keeps its counters"
|
||||
);
|
||||
assert_eq!(
|
||||
t.snapshot_session("sess-a").active_tool_calls,
|
||||
1,
|
||||
|
|
|
|||
|
|
@ -41,6 +41,10 @@ pub enum WorkspaceError {
|
|||
/// An error from the server connection or tool server.
|
||||
#[error("hub error: {0}")]
|
||||
HubError(String),
|
||||
#[error("unknown workspace method: {0}")]
|
||||
UnknownMethod(String),
|
||||
#[error("workspace archive export failed: {0}")]
|
||||
ExportArchiveLimitExceeded(String),
|
||||
#[error("github export error: {message}")]
|
||||
ExportGithub {
|
||||
kind: xai_grok_workspace_types::rpc::export_github::ExportGithubError,
|
||||
|
|
@ -78,6 +82,8 @@ impl WorkspaceError {
|
|||
Self::InvalidHunkAction(_) => "invalid_hunk_action",
|
||||
Self::HunkActionFailed(_) => "hunk_action_failed",
|
||||
Self::HubError(_) => "hub_error",
|
||||
Self::UnknownMethod(_) => "unknown_method",
|
||||
Self::ExportArchiveLimitExceeded(_) => "export_archive_limit_exceeded",
|
||||
Self::ExportGithub { kind, .. } => kind.wire_code(),
|
||||
Self::ShuttingDown => "shutting_down",
|
||||
Self::ToolsetExternallyOwned(_) => "toolset_externally_owned",
|
||||
|
|
|
|||
|
|
@ -3505,16 +3505,30 @@ impl WorkspaceHandle {
|
|||
let mut any_attempt = false;
|
||||
let mut any_success = false;
|
||||
let session_ids = tracker_for_status.known_sessions();
|
||||
for sid in &session_ids {
|
||||
let mut publish = session_ids.clone();
|
||||
for sid in last_sent.keys().filter_map(|k| k.as_ref()) {
|
||||
if !publish.contains(sid) {
|
||||
publish.push(sid.clone());
|
||||
}
|
||||
}
|
||||
let mut closed: Vec<String> = Vec::new();
|
||||
for sid in &publish {
|
||||
let payload = tracker_for_status.snapshot_session(sid);
|
||||
let key = Some(sid.clone());
|
||||
let ended = !session_ids.iter().any(|s| s == sid);
|
||||
if last_sent.get(&key).map(dedup_key) == Some(dedup_key(&payload)) {
|
||||
if ended {
|
||||
closed.push(sid.clone());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if let Some(ok) = send_status(&server_conn, payload.clone()).await {
|
||||
any_attempt = true;
|
||||
if ok {
|
||||
any_success = true;
|
||||
if ended {
|
||||
closed.push(sid.clone());
|
||||
}
|
||||
last_sent.insert(key, payload);
|
||||
last_successful_send = std::time::Instant::now();
|
||||
}
|
||||
|
|
@ -3522,7 +3536,10 @@ impl WorkspaceHandle {
|
|||
}
|
||||
last_sent.retain(|k, _| match k {
|
||||
None => true,
|
||||
Some(sid) => session_ids.iter().any(|s| s == sid),
|
||||
Some(sid) => {
|
||||
session_ids.iter().any(|s| s == sid)
|
||||
|| (any_success && !closed.contains(sid))
|
||||
}
|
||||
});
|
||||
let payload = tracker_for_status.snapshot();
|
||||
let needs_send = last_sent.get(&None).map(dedup_key) != Some(dedup_key(&payload));
|
||||
|
|
|
|||
|
|
@ -90,17 +90,13 @@ static WORKSPACE_RPC_DURATION_SECONDS: std::sync::LazyLock<HistogramVec> =
|
|||
.unwrap()
|
||||
});
|
||||
const UNKNOWN_METHOD_LABEL: &str = "unknown";
|
||||
/// Prefix of the [`WorkspaceError::HubError`] for an unrecognized method. Shared
|
||||
/// by the dispatch default arm and the metric classifier so the "collapse to
|
||||
/// `unknown`" decision cannot drift from the error it keys on.
|
||||
const UNKNOWN_METHOD_ERR_PREFIX: &str = "unknown workspace method:";
|
||||
/// Zero-init this module's metric families. See [`crate::init_metrics`].
|
||||
pub(crate) fn init_metrics() {
|
||||
WORKSPACE_RPC_REQUESTS_TOTAL
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "error"])
|
||||
.inc_by(0);
|
||||
WORKSPACE_RPC_ERRORS_TOTAL
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"])
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"])
|
||||
.inc_by(0);
|
||||
let _ = WORKSPACE_RPC_DURATION_SECONDS.with_label_values(&[UNKNOWN_METHOD_LABEL]);
|
||||
}
|
||||
|
|
@ -930,9 +926,7 @@ impl WorkspaceRpcHandler {
|
|||
}
|
||||
_ => {
|
||||
tracing::warn!(method, "unknown workspace rpc method");
|
||||
Err(WorkspaceError::HubError(format!(
|
||||
"{UNKNOWN_METHOD_ERR_PREFIX} {method}"
|
||||
)))
|
||||
Err(WorkspaceError::UnknownMethod(method.to_owned()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -989,10 +983,7 @@ impl ToolServerHandler for WorkspaceRpcHandler {
|
|||
bound_session.as_deref().map(|s| s.0.as_str()),
|
||||
)
|
||||
.await;
|
||||
let is_unknown_method = matches!(
|
||||
&result,
|
||||
Err(WorkspaceError::HubError(msg)) if msg.starts_with(UNKNOWN_METHOD_ERR_PREFIX)
|
||||
);
|
||||
let is_unknown_method = matches!(&result, Err(WorkspaceError::UnknownMethod(_)));
|
||||
let method_label = if is_unknown_method {
|
||||
UNKNOWN_METHOD_LABEL
|
||||
} else {
|
||||
|
|
@ -1273,15 +1264,18 @@ mod tests {
|
|||
assert_eq!(reply, turn_hook::HookReply::default());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn dispatch_unknown_method_returns_hub_error() {
|
||||
async fn dispatch_unknown_method_returns_unknown_method_error() {
|
||||
let handle = make_handle();
|
||||
let handler = WorkspaceRpcHandler::new(handle);
|
||||
let result = handler
|
||||
.dispatch("workspace.nonexistent", Value::Null, None)
|
||||
.await;
|
||||
assert!(
|
||||
matches!(result, Err(WorkspaceError::HubError(msg)) if msg.contains("unknown workspace method"))
|
||||
);
|
||||
match result {
|
||||
Err(WorkspaceError::UnknownMethod(method)) => {
|
||||
assert_eq!(method, "workspace.nonexistent");
|
||||
}
|
||||
other => panic!("expected UnknownMethod, got {other:?}"),
|
||||
}
|
||||
}
|
||||
/// A hub evict runs the two-phase drain then settles into terminal
|
||||
/// ShuttingDown (not a lingering Draining) for an evicted workspace.
|
||||
|
|
@ -2360,7 +2354,7 @@ mod tests {
|
|||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "error"])
|
||||
.get();
|
||||
let kind_before = WORKSPACE_RPC_ERRORS_TOTAL
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"])
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"])
|
||||
.get();
|
||||
let mut stream = handler
|
||||
.handle_call(
|
||||
|
|
@ -2378,7 +2372,7 @@ mod tests {
|
|||
);
|
||||
assert!(
|
||||
WORKSPACE_RPC_ERRORS_TOTAL
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "hub_error"])
|
||||
.with_label_values(&[UNKNOWN_METHOD_LABEL, "unknown_method"])
|
||||
.get()
|
||||
> kind_before,
|
||||
"a failed dispatch must also record its error_kind on the errors counter"
|
||||
|
|
|
|||
|
|
@ -183,7 +183,7 @@ mod init_metrics_tests {
|
|||
));
|
||||
assert!(has(
|
||||
"grok_workspace_rpc_errors_total",
|
||||
&[("method", "unknown"), ("error_kind", "hub_error")]
|
||||
&[("method", "unknown"), ("error_kind", "unknown_method")]
|
||||
));
|
||||
for stage in [
|
||||
"startup_recovery",
|
||||
|
|
|
|||
|
|
@ -5123,7 +5123,11 @@ mod tests {
|
|||
let local = tokio::task::LocalSet::new();
|
||||
local
|
||||
.run_until(async {
|
||||
for path in ["/etc/hosts", "/home/user/.grok/hooks/evil.json"] {
|
||||
for path in [
|
||||
"/etc/hosts",
|
||||
"/home/user/.grok/hooks/evil.json",
|
||||
"/home/user/.grok/sandbox.toml",
|
||||
] {
|
||||
let mut auto = crate::permission::types::PermissionConfig::new(vec![]);
|
||||
auto.prompt_policy = PromptPolicy::Auto;
|
||||
let allow =
|
||||
|
|
|
|||
|
|
@ -328,6 +328,7 @@ pub enum ProtectedEditReason {
|
|||
StartupFile,
|
||||
Etc,
|
||||
GrokConfig,
|
||||
GrokSandbox,
|
||||
ClaudeSettings,
|
||||
CursorHooks,
|
||||
/// Fail-closed / unclassified sensitive path; no user copy yet.
|
||||
|
|
@ -343,6 +344,7 @@ impl ProtectedEditReason {
|
|||
Self::StartupFile => "startup_file",
|
||||
Self::Etc => "etc",
|
||||
Self::GrokConfig => "grok_config",
|
||||
Self::GrokSandbox => "grok_sandbox",
|
||||
Self::ClaudeSettings => "claude_settings",
|
||||
Self::CursorHooks => "cursor_hooks",
|
||||
Self::Sensitive => "sensitive",
|
||||
|
|
@ -369,6 +371,9 @@ impl ProtectedEditReason {
|
|||
Self::GrokConfig => Some(
|
||||
"Note: This edit contains changes to Grok config, which can alter permissions, tools, and other behavior in later sessions.",
|
||||
),
|
||||
Self::GrokSandbox => Some(
|
||||
"Note: This edit contains changes to the Grok sandbox config, which can loosen filesystem and network restrictions on commands.",
|
||||
),
|
||||
Self::ClaudeSettings => Some(
|
||||
"Note: This edit contains changes to Claude-compatible settings, which can install hooks or change permission mode without a separate execution approval.",
|
||||
),
|
||||
|
|
@ -471,8 +476,8 @@ fn protected_edit_reason(path: &Path) -> Option<ProtectedEditReason> {
|
|||
if STARTUP_FILES.contains(&file) {
|
||||
return Some(ProtectedEditReason::StartupFile);
|
||||
}
|
||||
if string_components.ends_with(&[".grok", "config.toml"]) {
|
||||
return Some(ProtectedEditReason::GrokConfig);
|
||||
if let Some(reason) = protected_grok_config_file(path, &string_components) {
|
||||
return Some(reason);
|
||||
}
|
||||
if path == Path::new("/etc") || path.starts_with(Path::new("/etc")) {
|
||||
return Some(ProtectedEditReason::Etc);
|
||||
|
|
@ -480,6 +485,52 @@ fn protected_edit_reason(path: &Path) -> Option<ProtectedEditReason> {
|
|||
None
|
||||
}
|
||||
|
||||
/// Grok config files that alter permissions (`config.toml`, the
|
||||
/// `managed_config.toml` defaults tier, the user `requirements.toml` layer) or
|
||||
/// sandbox restrictions (`sandbox.toml`) in the running and later sessions; a
|
||||
/// silent edit would let the agent loosen its own guardrails. Matched directly
|
||||
/// inside any `.grok` dir (user-global default and workspace overlays) and
|
||||
/// directly under a custom `$GROK_HOME`, which the component match cannot see.
|
||||
fn protected_grok_config_file(path: &Path, components: &[&str]) -> Option<ProtectedEditReason> {
|
||||
protected_grok_config_file_with_home(
|
||||
path,
|
||||
components,
|
||||
xai_grok_config::user_grok_home().as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn protected_grok_config_file_with_home(
|
||||
path: &Path,
|
||||
components: &[&str],
|
||||
user_grok_home: Option<&Path>,
|
||||
) -> Option<ProtectedEditReason> {
|
||||
let reason = match components.last().copied() {
|
||||
Some(
|
||||
xai_grok_config::USER_CONFIG_FILENAME
|
||||
| xai_grok_config::MANAGED_CONFIG_FILENAME
|
||||
| xai_grok_config::REQUIREMENTS_FILENAME,
|
||||
) => ProtectedEditReason::GrokConfig,
|
||||
Some("sandbox.toml") => ProtectedEditReason::GrokSandbox,
|
||||
_ => return None,
|
||||
};
|
||||
let in_dot_grok = components.len() >= 2 && components[components.len() - 2] == ".grok";
|
||||
let in_grok_home = || grok_home_matches(user_grok_home, |home| path.parent() == Some(home));
|
||||
(in_dot_grok || in_grok_home()).then_some(reason)
|
||||
}
|
||||
|
||||
/// True when `pred` holds for the user grok home in either its lexical or
|
||||
/// physically-resolved form. Both forms are checked because callers hold a
|
||||
/// lexical and a resolved candidate path, and the home itself may sit behind a
|
||||
/// symlink. The comparison is byte-exact (no case folding), like every other
|
||||
/// resolved-path check in this module.
|
||||
fn grok_home_matches(home: Option<&Path>, pred: impl Fn(&Path) -> bool) -> bool {
|
||||
home.is_some_and(|home| {
|
||||
let lexical = xai_grok_paths::normalize_lexically(home);
|
||||
pred(&lexical)
|
||||
|| resolve_following_symlinks(&lexical, 0).is_some_and(|resolved| pred(&resolved))
|
||||
})
|
||||
}
|
||||
|
||||
fn path_is_under_user_grok_hook_root(path: &Path, grok_home: &Path) -> bool {
|
||||
path.starts_with(grok_home.join("hooks")) || path == grok_home.join("hooks-paths")
|
||||
}
|
||||
|
|
@ -487,12 +538,8 @@ fn path_is_under_user_grok_hook_root(path: &Path, grok_home: &Path) -> bool {
|
|||
fn protected_grok_hook_root(path: &Path, components: &[&str]) -> bool {
|
||||
components.windows(2).any(|pair| pair == [".grok", "hooks"])
|
||||
|| components.ends_with(&[".grok", "hooks-paths"])
|
||||
|| xai_grok_config::user_grok_home().is_some_and(|grok_home| {
|
||||
let lexical_home = xai_grok_paths::normalize_lexically(&grok_home);
|
||||
path_is_under_user_grok_hook_root(path, &lexical_home)
|
||||
|| resolve_following_symlinks(&lexical_home, 0).is_some_and(|resolved_home| {
|
||||
path_is_under_user_grok_hook_root(path, &resolved_home)
|
||||
})
|
||||
|| grok_home_matches(xai_grok_config::user_grok_home().as_deref(), |home| {
|
||||
path_is_under_user_grok_hook_root(path, home)
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -1369,6 +1416,8 @@ mod tests {
|
|||
"/etc",
|
||||
"/etc/grok-test",
|
||||
"/work/subdir/../.git/hooks/pre-commit",
|
||||
"/home/user/.grok/sandbox.toml",
|
||||
"/work/project/.grok/sandbox.toml",
|
||||
] {
|
||||
assert!(
|
||||
edit_target_protection(Path::new(path)).is_some(),
|
||||
|
|
@ -1378,6 +1427,9 @@ mod tests {
|
|||
for path in [
|
||||
"/work/src/main.rs",
|
||||
"/work/project/.grok/config.toml/backup",
|
||||
"/work/project/sandbox.toml",
|
||||
"/work/project/requirements.toml",
|
||||
"/work/project/managed_config.toml",
|
||||
] {
|
||||
assert!(
|
||||
edit_target_protection(Path::new(path)).is_none(),
|
||||
|
|
@ -1428,6 +1480,22 @@ mod tests {
|
|||
"/home/user/.grok/config.toml",
|
||||
ProtectedEditReason::GrokConfig,
|
||||
),
|
||||
(
|
||||
"/home/user/.grok/sandbox.toml",
|
||||
ProtectedEditReason::GrokSandbox,
|
||||
),
|
||||
(
|
||||
"/work/project/.grok/sandbox.toml",
|
||||
ProtectedEditReason::GrokSandbox,
|
||||
),
|
||||
(
|
||||
"/home/user/.grok/managed_config.toml",
|
||||
ProtectedEditReason::GrokConfig,
|
||||
),
|
||||
(
|
||||
"/home/user/.grok/requirements.toml",
|
||||
ProtectedEditReason::GrokConfig,
|
||||
),
|
||||
(
|
||||
"/home/user/.claude/settings.json",
|
||||
ProtectedEditReason::ClaudeSettings,
|
||||
|
|
@ -1551,6 +1619,85 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
/// A custom `$GROK_HOME` has no `.grok` path component, so the live
|
||||
/// `config.toml` / `sandbox.toml` must be caught by the home-prefix branch.
|
||||
#[test]
|
||||
fn grok_config_files_under_custom_grok_home_are_protected() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let home_path = home.path();
|
||||
for (file, reason) in [
|
||||
("config.toml", ProtectedEditReason::GrokConfig),
|
||||
("managed_config.toml", ProtectedEditReason::GrokConfig),
|
||||
("requirements.toml", ProtectedEditReason::GrokConfig),
|
||||
("sandbox.toml", ProtectedEditReason::GrokSandbox),
|
||||
] {
|
||||
let path = home_path.join(file);
|
||||
let components = [file];
|
||||
assert_eq!(
|
||||
protected_grok_config_file_with_home(&path, &components, Some(home_path)),
|
||||
Some(reason),
|
||||
"{file} directly under $GROK_HOME must be protected"
|
||||
);
|
||||
}
|
||||
// Same file names elsewhere (or with no resolvable home) stay ordinary.
|
||||
let elsewhere = home_path.join("sub").join("sandbox.toml");
|
||||
assert_eq!(
|
||||
protected_grok_config_file_with_home(
|
||||
&elsewhere,
|
||||
&["sub", "sandbox.toml"],
|
||||
Some(home_path)
|
||||
),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
protected_grok_config_file_with_home(
|
||||
&home_path.join("sandbox.toml"),
|
||||
&["sandbox.toml"],
|
||||
None
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
/// The resolved-symlink arm of the grok-home match must decide: `$GROK_HOME`
|
||||
/// points at a symlink while the edit targets the physical home directory,
|
||||
/// so the lexical parent-equality arm cannot fire.
|
||||
#[test]
|
||||
#[cfg(unix)]
|
||||
fn grok_config_under_symlinked_grok_home_is_protected() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let real_home = tmp.path().join("real-home");
|
||||
std::fs::create_dir(&real_home).unwrap();
|
||||
let link = tmp.path().join("home-link");
|
||||
symlink(&real_home, &link).unwrap();
|
||||
// tempdir paths can themselves contain symlinks (macOS /var -> /private/var);
|
||||
// compare against the physical home the production resolver will produce.
|
||||
let physical_home = resolve_following_symlinks(&real_home, 0).unwrap();
|
||||
assert_eq!(
|
||||
protected_grok_config_file_with_home(
|
||||
&physical_home.join("sandbox.toml"),
|
||||
&["sandbox.toml"],
|
||||
Some(&link)
|
||||
),
|
||||
Some(ProtectedEditReason::GrokSandbox)
|
||||
);
|
||||
}
|
||||
|
||||
/// `protected_edit_reason` lowercases path components before matching, so
|
||||
/// the canonical filename constants must stay lowercase or the const
|
||||
/// patterns silently stop firing.
|
||||
#[test]
|
||||
fn protected_config_filename_constants_are_lowercase() {
|
||||
for name in [
|
||||
xai_grok_config::USER_CONFIG_FILENAME,
|
||||
xai_grok_config::MANAGED_CONFIG_FILENAME,
|
||||
xai_grok_config::REQUIREMENTS_FILENAME,
|
||||
] {
|
||||
assert_eq!(name, name.to_ascii_lowercase(), "{name}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolved_root_alias_matches_physical_destination() {
|
||||
let resolved_root = resolve_following_symlinks(Path::new("/etc"), 0).unwrap();
|
||||
|
|
|
|||
|
|
@ -400,25 +400,39 @@ fn activity_url(control_port: u16) -> String {
|
|||
)
|
||||
}
|
||||
|
||||
/// One scrape of the proxy's activity endpoint. `last_activity_ms` is the only
|
||||
/// required field; the rest default to zero, so a workspace-server running
|
||||
/// ahead of the proxy binary degrades to the old behaviour rather than failing.
|
||||
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, serde::Deserialize)]
|
||||
struct ActivitySample {
|
||||
last_activity_ms: u64,
|
||||
#[serde(default)]
|
||||
last_routed_ms: u64,
|
||||
#[serde(default)]
|
||||
ws_tunnels_open: u64,
|
||||
#[serde(default)]
|
||||
routed_requests_in_flight: u64,
|
||||
}
|
||||
|
||||
/// Classified result of one scrape, so a missing proxy (quiet no-op) is never
|
||||
/// confused with a genuine error response or with real activity.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
enum ScrapeOutcome {
|
||||
/// The proxy answered with a parseable activity stamp (epoch-ms).
|
||||
Stamp(u64),
|
||||
/// The proxy answered with a parseable activity sample.
|
||||
Stamp(ActivitySample),
|
||||
/// The proxy isn't reachable (connection refused / not up yet): quiet no-op.
|
||||
Absent,
|
||||
/// The proxy answered but the response was unusable (error status / bad body).
|
||||
BadResponse,
|
||||
}
|
||||
|
||||
/// Parse `{ "last_activity_ms": <u64> }`; `None` for a malformed body, a missing
|
||||
/// field, or a non-integer value.
|
||||
fn parse_activity_body(body: &str) -> Option<u64> {
|
||||
serde_json::from_str::<serde_json::Value>(body)
|
||||
.ok()?
|
||||
.get("last_activity_ms")?
|
||||
.as_u64()
|
||||
/// Parse an activity body; `None` for a malformed body or a missing/non-integer
|
||||
/// `last_activity_ms`. Unknown fields are ignored so the proxy can add more.
|
||||
fn parse_activity_body(body: &str) -> Option<ActivitySample> {
|
||||
let value: serde_json::Value = serde_json::from_str(body).ok()?;
|
||||
// Required, and must be an integer.
|
||||
value.get("last_activity_ms")?.as_u64()?;
|
||||
serde_json::from_value(value).ok()
|
||||
}
|
||||
|
||||
/// Classify a completed response by status + body (transport failures are
|
||||
|
|
@ -428,7 +442,7 @@ fn classify_activity_response(status: u16, body: &str) -> ScrapeOutcome {
|
|||
return ScrapeOutcome::BadResponse;
|
||||
}
|
||||
match parse_activity_body(body) {
|
||||
Some(ms) => ScrapeOutcome::Stamp(ms),
|
||||
Some(sample) => ScrapeOutcome::Stamp(sample),
|
||||
None => ScrapeOutcome::BadResponse,
|
||||
}
|
||||
}
|
||||
|
|
@ -457,6 +471,19 @@ async fn scrape_activity(client: &reqwest::Client, url: &str) -> ScrapeOutcome {
|
|||
}
|
||||
}
|
||||
|
||||
/// Clear the mirrored attached-client counters once we have been without
|
||||
/// trustworthy data for `grace`. Starts the clock on the first bad scrape.
|
||||
fn clear_attached_if_stale(
|
||||
tracker: &ActivityTracker,
|
||||
stale_since: &mut Option<Instant>,
|
||||
grace: Duration,
|
||||
) {
|
||||
let since = *stale_since.get_or_insert_with(Instant::now);
|
||||
if since.elapsed() >= grace {
|
||||
tracker.set_preview_attached(0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll the proxy's loopback activity endpoint until `shutdown` flips, feeding
|
||||
/// the tracker on each advance. Spawn after the `ActivityTracker` exists (post
|
||||
/// hub-connect), gated on preview being enabled. `control_port` is the proxy's
|
||||
|
|
@ -507,22 +534,56 @@ async fn scrape_activity_loop(
|
|||
// `None` until the first successful scrape establishes a baseline. Baselining
|
||||
// (rather than starting at 0) avoids a spurious withhold when a workspace-server
|
||||
// restart meets a proxy whose stamp is already non-zero but stale.
|
||||
let mut last_seen: Option<u64> = None;
|
||||
let mut last_seen: Option<ActivitySample> = None;
|
||||
// When we last had trustworthy attached-client data. The counters are the
|
||||
// ONLY hold a WS-only client has — a tunnel writes no stamps — so one bad
|
||||
// scrape must not clear them; but unlike the stamps they do not decay, so a
|
||||
// sustained loss must, or a proxy that dies mid-tunnel holds the sandbox to
|
||||
// the TTL. `Absent` and `BadResponse` both count as loss: one cannot reach
|
||||
// the proxy, the other cannot understand it. The threshold is the stamp
|
||||
// window, so both hold mechanisms expire on the same clock.
|
||||
let attached_grace = Duration::from_millis(tracker.preview_activity_window_ms());
|
||||
let mut attached_stale_since: Option<Instant> = None;
|
||||
loop {
|
||||
if sleep_or_shutdown(interval, &mut shutdown).await {
|
||||
return;
|
||||
}
|
||||
match scrape_activity(&client, &url).await {
|
||||
ScrapeOutcome::Stamp(current) => {
|
||||
if last_seen.is_some_and(|prev| preview_activity_advanced(prev, current)) {
|
||||
tracker.note_preview_activity();
|
||||
if let Some(prev) = last_seen {
|
||||
// The routed stamp is a subset of the generic one, so check
|
||||
// it first and attribute only the remainder to a poll —
|
||||
// otherwise one routed request would report as both.
|
||||
if preview_activity_advanced(prev.last_routed_ms, current.last_routed_ms) {
|
||||
tracker.note_preview_routed_activity();
|
||||
} else if preview_activity_advanced(
|
||||
prev.last_activity_ms,
|
||||
current.last_activity_ms,
|
||||
) {
|
||||
tracker.note_preview_status_activity();
|
||||
}
|
||||
}
|
||||
// Absolute counters, republished every tick: a mirror of the
|
||||
// proxy's state, not an edge.
|
||||
tracker.set_preview_attached(
|
||||
current.ws_tunnels_open,
|
||||
current.routed_requests_in_flight,
|
||||
);
|
||||
last_seen = Some(current);
|
||||
attached_stale_since = None;
|
||||
}
|
||||
// Proxy absent (preview disabled / starting / restarting): no-op.
|
||||
ScrapeOutcome::Absent => {}
|
||||
// Proxy absent (preview disabled / starting / restarting): leave
|
||||
// the stamps, and age the attached counters out. See above.
|
||||
ScrapeOutcome::Absent => {
|
||||
clear_attached_if_stale(&tracker, &mut attached_stale_since, attached_grace);
|
||||
}
|
||||
// Answering, but unusably. Same staleness clock: an error status or
|
||||
// an unparseable body tells us nothing about attached clients, and
|
||||
// leaving them untouched would let a persistently broken proxy hold
|
||||
// the withhold open forever.
|
||||
ScrapeOutcome::BadResponse => {
|
||||
tracing::debug!(%url, "preview-activity scrape returned an unusable response");
|
||||
clear_attached_if_stale(&tracker, &mut attached_stale_since, attached_grace);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -886,15 +947,23 @@ mod tests {
|
|||
.expect("a pre-flipped shutdown must return without scraping");
|
||||
}
|
||||
|
||||
/// What an older proxy binary reports.
|
||||
fn stamp_only(last_activity_ms: u64) -> ActivitySample {
|
||||
ActivitySample {
|
||||
last_activity_ms,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_activity_body_reads_stamp_and_rejects_bad_shapes() {
|
||||
assert_eq!(
|
||||
parse_activity_body(r#"{"last_activity_ms":1234}"#),
|
||||
Some(1234)
|
||||
Some(stamp_only(1234))
|
||||
);
|
||||
assert_eq!(
|
||||
parse_activity_body(r#"{"last_activity_ms":0,"extra":true}"#),
|
||||
Some(0)
|
||||
Some(stamp_only(0))
|
||||
);
|
||||
assert_eq!(parse_activity_body(r#"{"other":1}"#), None);
|
||||
assert_eq!(parse_activity_body(r#"{"last_activity_ms":"7"}"#), None);
|
||||
|
|
@ -903,11 +972,130 @@ mod tests {
|
|||
assert_eq!(parse_activity_body(""), None);
|
||||
}
|
||||
|
||||
/// The binaries are version-pinned per session, but a restore can repin, so
|
||||
/// this skew is real rather than theoretical.
|
||||
#[test]
|
||||
fn parse_activity_body_tolerates_a_proxy_without_the_new_fields() {
|
||||
let old = parse_activity_body(r#"{"last_activity_ms":5,"status_holds_in_use":2}"#)
|
||||
.expect("an old proxy body must still parse");
|
||||
assert_eq!(old.last_routed_ms, 0);
|
||||
assert_eq!(old.ws_tunnels_open, 0);
|
||||
assert_eq!(
|
||||
old.routed_requests_in_flight, 0,
|
||||
"absent fields read as 'nothing attached', never as attached"
|
||||
);
|
||||
}
|
||||
|
||||
/// A WS-only client has no stamps, so the attached counters are its only
|
||||
/// hold. One unreachable scrape — a proxy restart, a loopback hiccup — must
|
||||
/// not drop it and publish idle; a sustained absence still must.
|
||||
#[tokio::test]
|
||||
async fn one_absent_scrape_does_not_drop_an_attached_client() {
|
||||
let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(10_000));
|
||||
tracker.set_preview_attached(1, 0);
|
||||
assert!(tracker.snapshot().idle_since_ms.is_none());
|
||||
|
||||
// Nothing is listening on this port, so every scrape classifies Absent.
|
||||
let port = reserved_closed_port().await;
|
||||
let (tx, rx) = watch::channel(false);
|
||||
let loop_handle = tokio::spawn(scrape_activity_loop(
|
||||
port,
|
||||
Arc::clone(&tracker),
|
||||
Duration::from_millis(10),
|
||||
rx,
|
||||
));
|
||||
|
||||
// Many consecutive absences, all well inside the 10s grace.
|
||||
tokio::time::sleep(Duration::from_millis(300)).await;
|
||||
assert!(
|
||||
tracker.snapshot().idle_since_ms.is_none(),
|
||||
"repeated unreachable scrapes inside the grace must not drop the hold"
|
||||
);
|
||||
assert_eq!(tracker.preview_ws_tunnels_open(), 1);
|
||||
|
||||
let _ = tx.send(true);
|
||||
let _ = loop_handle.await;
|
||||
}
|
||||
|
||||
/// A proxy that answers unusably tells us nothing about attached clients
|
||||
/// either, and the mirrored counters do not decay on their own — so an
|
||||
/// endless run of error statuses must not pin `PreviewAttached` forever.
|
||||
#[tokio::test]
|
||||
async fn a_persistently_broken_proxy_does_not_pin_attached_forever() {
|
||||
let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(50));
|
||||
tracker.set_preview_attached(1, 0);
|
||||
|
||||
// Answers every time, always with a 500 => BadResponse, never Absent.
|
||||
let port = serve_canned("HTTP/1.1 500 Internal Server Error", "boom", true).await;
|
||||
let (tx, rx) = watch::channel(false);
|
||||
let loop_handle = tokio::spawn(scrape_activity_loop(
|
||||
port,
|
||||
Arc::clone(&tracker),
|
||||
Duration::from_millis(10),
|
||||
rx,
|
||||
));
|
||||
|
||||
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
|
||||
while tracker.preview_ws_tunnels_open() != 0 {
|
||||
assert!(
|
||||
tokio::time::Instant::now() < deadline,
|
||||
"a sustained run of unusable responses must age the hold out"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
|
||||
let _ = tx.send(true);
|
||||
let _ = loop_handle.await;
|
||||
}
|
||||
|
||||
/// The other half: a proxy that stays gone must eventually release, or a
|
||||
/// tunnel that died with it would hold the sandbox to the TTL.
|
||||
#[tokio::test]
|
||||
async fn a_sustained_absence_clears_the_attached_client() {
|
||||
let tracker = Arc::new(ActivityTracker::new().with_preview_activity_window_ms(50));
|
||||
tracker.set_preview_attached(1, 0);
|
||||
|
||||
let port = reserved_closed_port().await;
|
||||
let (tx, rx) = watch::channel(false);
|
||||
let loop_handle = tokio::spawn(scrape_activity_loop(
|
||||
port,
|
||||
Arc::clone(&tracker),
|
||||
Duration::from_millis(10),
|
||||
rx,
|
||||
));
|
||||
|
||||
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
|
||||
while tracker.preview_ws_tunnels_open() != 0 {
|
||||
assert!(
|
||||
tokio::time::Instant::now() < deadline,
|
||||
"an absence past the window must clear the attached counters"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
|
||||
let _ = tx.send(true);
|
||||
let _ = loop_handle.await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_activity_body_reads_the_attached_client_fields() {
|
||||
let sample = parse_activity_body(
|
||||
r#"{"last_activity_ms":9,"status_holds_in_use":1,
|
||||
"held_status_aborts_quieted":0,"ws_tunnels_open":3,
|
||||
"routed_requests_in_flight":2,"last_routed_ms":7}"#,
|
||||
)
|
||||
.expect("parse");
|
||||
assert_eq!(sample.last_activity_ms, 9);
|
||||
assert_eq!(sample.last_routed_ms, 7);
|
||||
assert_eq!(sample.ws_tunnels_open, 3);
|
||||
assert_eq!(sample.routed_requests_in_flight, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_activity_response_distinguishes_stamp_from_bad() {
|
||||
assert_eq!(
|
||||
classify_activity_response(200, r#"{"last_activity_ms":42}"#),
|
||||
ScrapeOutcome::Stamp(42)
|
||||
ScrapeOutcome::Stamp(stamp_only(42))
|
||||
);
|
||||
assert_eq!(
|
||||
classify_activity_response(200, "garbage"),
|
||||
|
|
@ -1022,7 +1210,7 @@ mod tests {
|
|||
let port = serve_canned("HTTP/1.1 200 OK", r#"{"last_activity_ms":9876}"#, false).await;
|
||||
assert_eq!(
|
||||
scrape_activity(&scrape_client(), &activity_url(port)).await,
|
||||
ScrapeOutcome::Stamp(9876)
|
||||
ScrapeOutcome::Stamp(stamp_only(9876))
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,8 @@ pub fn error_code(err: &WorkspaceError) -> &'static str {
|
|||
WorkspaceError::InvalidHunkAction(_) => "invalid_hunk_action",
|
||||
WorkspaceError::HunkActionFailed(_) => "hunk_action_failed",
|
||||
WorkspaceError::HubError(_) => "hub_error",
|
||||
WorkspaceError::UnknownMethod(_) => "unknown_method",
|
||||
WorkspaceError::ExportArchiveLimitExceeded(_) => "export_archive_limit_exceeded",
|
||||
WorkspaceError::ExportGithub { kind, .. } => kind.wire_code(),
|
||||
WorkspaceError::ShuttingDown => "shutting_down",
|
||||
WorkspaceError::ToolsetExternallyOwned(_) => "toolset_externally_owned",
|
||||
|
|
@ -80,6 +82,8 @@ pub fn rpc_error_to_workspace(err: RpcError) -> WorkspaceError {
|
|||
"invalid_hunk_action" => WorkspaceError::InvalidHunkAction(err.message),
|
||||
"hunk_action_failed" => WorkspaceError::HunkActionFailed(err.message),
|
||||
"hub_error" => WorkspaceError::HubError(err.message),
|
||||
"unknown_method" => WorkspaceError::UnknownMethod(err.message),
|
||||
"export_archive_limit_exceeded" => WorkspaceError::ExportArchiveLimitExceeded(err.message),
|
||||
"shutting_down" => WorkspaceError::ShuttingDown,
|
||||
"toolset_externally_owned" => WorkspaceError::ToolsetExternallyOwned(err.message),
|
||||
unknown => {
|
||||
|
|
@ -115,6 +119,8 @@ mod tests {
|
|||
WorkspaceError::InvalidHunkAction("h".into()),
|
||||
WorkspaceError::HunkActionFailed("h".into()),
|
||||
WorkspaceError::HubError("hub".into()),
|
||||
WorkspaceError::UnknownMethod("workspace.bogus".into()),
|
||||
WorkspaceError::ExportArchiveLimitExceeded("too big".into()),
|
||||
WorkspaceError::ShuttingDown,
|
||||
WorkspaceError::ToolsetExternallyOwned("s".into()),
|
||||
];
|
||||
|
|
|
|||
Loading…
Reference in a new issue