//! Per-session and connection-level activity tracking for tool server //! lifecycle status reporting. use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; use std::sync::{Arc, OnceLock}; 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}; const LIFECYCLE_NONE: u8 = 0; const LIFECYCLE_DRAINING: u8 = 1; const LIFECYCLE_SHUTTING_DOWN: u8 = 2; const DEFAULT_SESSION: &str = "__default__"; const SESSION_IDLE_PRUNE_MS: u64 = 5 * 60 * 1000; /// Default cap (ms) on how long pending durability work (artifact producers / /// queued uploads) may withhold `idle_since_ms`. Overridable via /// `GROK_WORKSPACE_DURABILITY_IDLE_HOLD_MAX_MS`. const DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS: u64 = 600_000; /// How long recent preview-proxy traffic withholds `idle_since_ms` — a decaying /// window (not a reset), so a polled preview stays alive but a single stale poll /// can't pin it. Larger than the 5s status poll, smaller than the idle grace. pub(crate) const PREVIEW_ACTIVITY_WINDOW_MS: u64 = 60_000; struct SessionActivity { active_tool_calls: AtomicU32, active_tools: DashMap, last_call_started_ms: AtomicU64, last_call_completed_ms: AtomicU64, idle_since_ms: AtomicU64, /// Current turn number (set by `turn_started`). current_turn: AtomicU64, /// Whether a turn is currently active. turn_active: AtomicBool, } impl SessionActivity { fn new() -> Self { Self { active_tool_calls: AtomicU32::new(0), active_tools: DashMap::new(), last_call_started_ms: AtomicU64::new(0), last_call_completed_ms: AtomicU64::new(0), idle_since_ms: AtomicU64::new(now_ms()), current_turn: AtomicU64::new(0), turn_active: AtomicBool::new(false), } } } /// Tracks in-flight tool calls and background tasks for /// [`ToolServerStatusPayload`] reporting. /// /// All methods are `&self` — share via `Arc` across the tool handler, the /// activity feed, and the status publisher. pub struct ActivityTracker { active_tool_calls: AtomicU32, active_tools: DashMap, background_tasks: AtomicU32, background_ids: DashMap, last_call_started_ms: AtomicU64, last_call_completed_ms: AtomicU64, idle_since_ms: AtomicU64, started_at: Instant, lifecycle: AtomicU8, /// `Arc` so the upload queue (via [`notify_handle`](Self::notify_handle)) /// can wake the same waiter the status publisher blocks on. notify: Arc, /// Coupled upload-queue stats; unset for bare trackers (tests, queue-less mode). upload_queue_stats: OnceLock>, /// Epoch ms a graceful drain began; `0` means "not draining". drain_started_ms: AtomicU64, /// Coupled artifact-producer task tracker; unset for bare trackers. producer_tasks: OnceLock, /// Epoch ms the durability-busy condition started; `0` while clear. Stamped /// lazily at snapshot time. durability_busy_since_ms: AtomicU64, /// Cap (ms) on how long durability work may withhold `idle_since_ms`. durability_idle_hold_max_ms: u64, /// When set, the idle verdict ignores background tasks so `idle_since_ms` /// tracks foreground tool-call activity only. Drain/`status` stay bg-aware. idle_ignores_background: bool, /// 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 /// [`preview_activity_window_ms`](Self::preview_activity_window_ms). last_preview_activity_ms: AtomicU64, sessions: DashMap, /// call_id → session_id so `tool_call_completed` can decrement /// the right session without the caller repeating it. call_to_session: DashMap, /// Idle window (ms) after which an inactive session is pruned by /// [`known_sessions`]. Set once at construction; no locking. prune_window_ms: u64, /// Per-session `events.jsonl` writers, shared (`Arc`) with /// [`WorkspaceShared`](crate::session::WorkspaceShared). `None` until /// [`set_event_writers`](Self::set_event_writers) is called during /// `WorkspaceHandle` construction; bare trackers (the existing unit tests) /// leave it unset so no `Tool*` events are emitted — behaviour-preserving. event_writers: OnceLock>>, /// Per-call start timestamps (epoch ms), keyed by `call_id`. Populated only /// when the call's session has an OPEN `events.jsonl` writer — i.e. the same /// condition under which `ToolStarted` is emitted — so the flag-off / /// no-writer path inserts nothing (the per-session writer map is empty when /// `GROK_WORKSPACE_EVENTS_ENABLED` is off). Consumed by /// [`tool_call_completed`](Self::tool_call_completed) to report a truthful /// `ToolCompleted.duration_ms`. /// /// The stored value pairs the start timestamp with a clone of the session's /// `EventWriter` captured at `ToolStarted` time. Holding the writer handle /// here guarantees a symmetric `ToolStarted`/`ToolCompleted` pair: presence /// of this entry is the single source of truth that a `ToolStarted` was /// emitted, and the captured `Arc`-backed writer keeps the session's /// `events.jsonl` fd alive so the paired `ToolCompleted` is still written /// even if the per-session writer was evicted (e.g. `on_session_ended`) /// mid-call. call_started_ms: DashMap, } impl Default for ActivityTracker { fn default() -> Self { Self::new() } } impl ActivityTracker { /// Construct a tracker with the default 5-minute session-prune window. pub fn new() -> Self { Self::with_prune_window(std::time::Duration::from_millis(SESSION_IDLE_PRUNE_MS)) } /// Construct a tracker with a custom session-prune window. The durability /// idle-hold cap comes from `GROK_WORKSPACE_DURABILITY_IDLE_HOLD_MAX_MS` /// (default [`DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS`]). pub fn with_prune_window(prune_window: std::time::Duration) -> Self { Self::with_prune_window_and_idle_hold(prune_window, durability_idle_hold_max_from_env()) } /// [`with_prune_window`](Self::with_prune_window) with an explicit /// durability idle-hold cap, so tests never race process env. pub fn with_prune_window_and_idle_hold( prune_window: std::time::Duration, durability_idle_hold_max_ms: u64, ) -> Self { Self { active_tool_calls: AtomicU32::new(0), active_tools: DashMap::new(), background_tasks: AtomicU32::new(0), background_ids: DashMap::new(), last_call_started_ms: AtomicU64::new(0), last_call_completed_ms: AtomicU64::new(0), idle_since_ms: AtomicU64::new(now_ms()), started_at: Instant::now(), lifecycle: AtomicU8::new(LIFECYCLE_NONE), notify: Arc::new(tokio::sync::Notify::new()), upload_queue_stats: OnceLock::new(), drain_started_ms: AtomicU64::new(0), producer_tasks: OnceLock::new(), durability_busy_since_ms: AtomicU64::new(0), durability_idle_hold_max_ms, idle_ignores_background: false, preview_activity_window_ms: PREVIEW_ACTIVITY_WINDOW_MS, last_preview_activity_ms: AtomicU64::new(0), sessions: DashMap::new(), call_to_session: DashMap::new(), prune_window_ms: prune_window.as_millis() as u64, event_writers: OnceLock::new(), call_started_ms: DashMap::new(), } } /// Opt into foreground-only idle: background tasks stop withholding /// `idle_since_ms`. pub fn with_idle_ignores_background(mut self, enabled: bool) -> Self { self.idle_ignores_background = enabled; self } /// Override the preview-activity withhold window; the WorkspaceServer sources /// it from `StatusConfig`. pub fn with_preview_activity_window_ms(mut self, window_ms: u64) -> Self { self.preview_activity_window_ms = window_ms; self } /// Wire the shared per-session `events.jsonl` writer map into the tracker so /// [`tool_call_started`](Self::tool_call_started) / /// [`tool_call_completed`](Self::tool_call_completed) can emit `Tool*` /// events. Set once during `WorkspaceHandle` construction; calling it a /// second time is a no-op (the first map wins). pub fn set_event_writers(&self, writers: Arc>) { let _ = self.event_writers.set(writers); } /// Couple upload-queue stats (status reports queue depth; `is_drained` waits /// for the queue). Set once; a second call is a no-op. pub fn set_upload_queue_stats(&self, stats: Arc) { let _ = self.upload_queue_stats.set(stats); } /// Couple the artifact-producer tracker (status counts producers, withholds /// idle while they run). Set once; a second call is a no-op. pub fn set_producer_tasks(&self, tasks: tokio_util::task::TaskTracker) { let _ = self.producer_tasks.set(tasks); } /// Clone of the internal `Notify` for the upload queue to drive republishes. pub fn notify_handle(&self) -> Arc { self.notify.clone() } /// Record fresh preview-proxy 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 .store(now_ms(), Ordering::Relaxed); self.notify.notify_waiters(); } /// Pending upload-queue items (0 when no queue is coupled). fn upload_queue_pending(&self) -> u64 { self.upload_queue_stats .get() .map(|s| s.pending.load(Ordering::Relaxed)) .unwrap_or(0) } /// Queue/drain status fields shared by both snapshot paths; all zero when no /// queue is coupled. Best-effort: independent `Relaxed` loads can transiently /// show `inflight > pending` (cosmetic; the drain reads `pending` alone). fn drain_status_fields(&self) -> (u32, u64, u32, bool, Option) { let (pending, pending_bytes, inflight, breaker) = match self.upload_queue_stats.get() { Some(s) => ( s.pending.load(Ordering::Relaxed) as u32, s.pending_bytes.load(Ordering::Relaxed), s.inflight.load(Ordering::Relaxed) as u32, s.circuit_breaker_active.load(Ordering::Relaxed), ), None => (0, 0, 0, false), }; let drain_started = match self.drain_started_ms.load(Ordering::Relaxed) { 0 => None, ms => Some(ms), }; (pending, pending_bytes, inflight, breaker, drain_started) } /// Whether a drain has been started (`set_draining` ran) in this process. pub fn drain_started(&self) -> bool { self.drain_started_ms.load(Ordering::Relaxed) != 0 } /// Durability tail shared by [`Self::snapshot`] and [`Self::snapshot_session`] /// (one construction site so the two payloads can't drift). fn durability_payload_fields(&self, idle_since: u64) -> DurabilityPayloadFields { 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 // 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()); DurabilityPayloadFields { idle_since_ms: if idle_since == 0 || withhold_idle { None } else { Some(idle_since) }, upload_queue_pending: queue_pending, upload_queue_pending_bytes: queue_pending_bytes, upload_queue_inflight: queue_inflight, upload_queue_circuit_breaker_tripped: breaker, artifact_producers_inflight: producers, drain_started_ms: drain_started, } } /// Durability gate: returns (producers in flight, whether `idle_since_ms` /// must be withheld). Idle is withheld while producers or queued uploads are /// outstanding, except past the bounded hold cap, or — queued items only — /// when the circuit breaker is tripped (queued items survive via disk spill + /// restart recovery; an in-flight producer has nothing on disk yet, so it /// withholds regardless of queue health). The hold resets when it clears. fn durability_gate(&self, queue_pending: u32, breaker_tripped: bool) -> (u32, bool) { let producers = self .producer_tasks .get() .map(|t| t.len() as u32) .unwrap_or(0); let withholding = producers > 0 || (queue_pending > 0 && !breaker_tripped); if !withholding { self.durability_busy_since_ms.store(0, Ordering::Relaxed); return (producers, false); } let now = now_ms(); // First observation of the busy condition wins the stamp. Relaxed: // the stamp guards no other data, it's just a monotonic-enough clock. let since = match self.durability_busy_since_ms.compare_exchange( 0, now, Ordering::Relaxed, Ordering::Relaxed, ) { Ok(_) => now, Err(prev) => prev, }; let hold_expired = now.saturating_sub(since) >= self.durability_idle_hold_max_ms; (producers, !hold_expired) } /// Whether recent preview-proxy traffic should currently withhold idle. fn preview_withholds_idle(&self, now: u64) -> bool { preview_activity_withholds_idle( now, self.last_preview_activity_ms.load(Ordering::Relaxed), self.preview_activity_window_ms, ) } /// Whether any tracked session currently has an active turn (the aggregate /// `turn_active`). fn any_turn_active(&self) -> bool { self.sessions .iter() .any(|s| s.value().turn_active.load(Ordering::Acquire)) } /// Resolve the `events.jsonl` writer for `session_id` — `Some` only when an /// event sink is configured AND the session already has an open writer /// (opened at turn start). Returns `None` (so no event is emitted) for the /// `__default__` / unknown-session cases. fn session_writer(&self, session_id: Option<&str>) -> Option { let session_id = session_id?; let writers = self.event_writers.get()?; writers.get(session_id).map(|w| w.value().clone()) } pub fn tool_call_started(&self, call_id: &str, tool_name: &str, session_id: Option<&str>) { if self.active_tools.contains_key(call_id) { return; } let now = now_ms(); self.active_tool_calls.fetch_add(1, Ordering::Relaxed); self.active_tools .insert(call_id.to_owned(), tool_name.to_owned()); self.last_call_started_ms.store(now, Ordering::Relaxed); self.idle_since_ms.store(0, Ordering::Relaxed); let sid = session_id.unwrap_or(DEFAULT_SESSION); self.call_to_session .insert(call_id.to_owned(), sid.to_owned()); let session = self .sessions .entry(sid.to_owned()) .or_insert_with(SessionActivity::new); session.active_tool_calls.fetch_add(1, Ordering::Relaxed); session .active_tools .insert(call_id.to_owned(), tool_name.to_owned()); session.last_call_started_ms.store(now, Ordering::Relaxed); session.idle_since_ms.store(0, Ordering::Relaxed); self.notify.notify_waiters(); // events.jsonl: only when the session's writer is already open (turn // started, under the events flag) do we record the start time (for a // truthful completion duration) and emit `ToolStarted`. When the writer // is absent — flag-off (empty per-session map) or no turn yet — this // whole block is skipped, so the flag-off path allocates nothing. if let Some(writer) = self.session_writer(session_id) { // Capture the writer handle alongside the start time so the paired // `ToolCompleted` is emitted to the same `events.jsonl` even if the // session writer is evicted mid-call. self.call_started_ms .insert(call_id.to_owned(), (now, writer.clone())); writer.emit(Event::ToolStarted { tool_name: tool_name.to_owned(), }); } } /// Mark an in-flight tool call as completed. /// /// `session_id` identifies the owning session for the caller's bookkeeping; /// the internal session counters key off the recorded `call_id → session` /// mapping, and the [`ToolCompleted`](Event::ToolCompleted) event is written /// via the `EventWriter` handle captured at `ToolStarted` time (so it lands /// in the right `events.jsonl` even if the session writer was evicted /// mid-call). `outcome` is the truthful terminal status from the caller /// (`Success`/`Error` at the tool handler, `Cancelled` from the cancel /// paths). pub fn tool_call_completed( &self, call_id: &str, _session_id: Option<&str>, outcome: ToolOutcome, ) { let Some((_, tool_name)) = self.active_tools.remove(call_id) else { return; }; let now = now_ms(); self.active_tool_calls.fetch_sub(1, Ordering::AcqRel); self.last_call_completed_ms.store(now, Ordering::Relaxed); // Re-read after decrement: a concurrent `tool_call_started` may // have bumped the counter back up between our `fetch_sub` and // this load. Only transition to idle if we're truly at zero. if self.active_tool_calls.load(Ordering::Acquire) == 0 && (self.idle_ignores_background || self.background_tasks.load(Ordering::Acquire) == 0) { self.idle_since_ms.store(now, Ordering::Relaxed); } if let Some((_, sid)) = self.call_to_session.remove(call_id) && let Some(session) = self.sessions.get(&sid) { session.active_tool_calls.fetch_sub(1, Ordering::AcqRel); session.active_tools.remove(call_id); session.last_call_completed_ms.store(now, Ordering::Relaxed); if session.active_tool_calls.load(Ordering::Acquire) == 0 { session.idle_since_ms.store(now, Ordering::Relaxed); } } self.notify.notify_waiters(); // events.jsonl: emit `ToolCompleted` only when a paired `ToolStarted` // was recorded for this call (its `call_started_ms` entry is present). // Gating on that entry — rather than re-checking writer state at // completion — keeps the start/completion pair symmetric: no orphan // zero-duration `ToolCompleted` when a writer opens mid-call, and no // dropped completion when the session writer is evicted mid-call (the // captured writer handle keeps the file open). `duration_ms` is always // truthful because the start time is paired with the writer. if let Some((_, (started_ms, writer))) = self.call_started_ms.remove(call_id) { // Workspace = hub/proxy hop; join package durations on shell rows (no source). writer.emit(Event::ToolCompleted { tool_name, duration_ms: now.saturating_sub(started_ms), outcome, tool_call_id: call_id.to_owned(), source: ToolCompletedSource::Workspace, }); } } pub fn background_task_started(&self, task_id: &str) { if self.background_ids.contains_key(task_id) { return; } self.background_tasks.fetch_add(1, Ordering::Relaxed); self.background_ids.insert(task_id.to_owned(), ()); if self.idle_ignores_background { // An active fg call's store(0) must keep idle withheld. if self.active_tool_calls.load(Ordering::Acquire) == 0 { self.idle_since_ms.store(now_ms(), Ordering::Relaxed); } } else { self.idle_since_ms.store(0, Ordering::Relaxed); } self.notify.notify_waiters(); } pub fn background_task_completed(&self, task_id: &str) { if self.background_ids.remove(task_id).is_none() { return; } let prev = self.background_tasks.fetch_sub(1, Ordering::AcqRel); if !self.idle_ignores_background && prev == 1 && self.active_tool_calls.load(Ordering::Acquire) == 0 { self.idle_since_ms.store(now_ms(), Ordering::Relaxed); } self.notify.notify_waiters(); } pub fn turn_started(&self, session_id: &str, turn_number: u64) { let session = self .sessions .entry(session_id.to_owned()) .or_insert_with(SessionActivity::new); session.current_turn.store(turn_number, Ordering::Release); session.turn_active.store(true, Ordering::Release); self.notify.notify_waiters(); } pub fn turn_completed(&self, session_id: &str, turn_number: u64, _duration_ms: u64) { if let Some(session) = self.sessions.get(session_id) && session.current_turn.load(Ordering::Acquire) == turn_number { session.turn_active.store(false, Ordering::Release); self.notify.notify_waiters(); } } /// Complete all in-flight tool calls for the given session. /// /// Returns the number of calls that were marked as completed. /// Called when a session-wide Cancel hook arrives without a specific /// `call_id` (broadcast cancel from Ctrl+C). pub fn cancel_all_session_calls(&self, session_id: &str) -> usize { let call_ids: Vec = self .call_to_session .iter() .filter(|entry| entry.value() == session_id) .map(|entry| entry.key().clone()) .collect(); let count = call_ids.len(); for call_id in call_ids { self.tool_call_completed(&call_id, Some(session_id), ToolOutcome::Cancelled); } 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. pub fn session_ended(&self, session_id: &str) { if let Some(session) = self.sessions.get(session_id) { session.turn_active.store(false, Ordering::Release); } self.notify.notify_waiters(); } /// Whether a turn is currently active for the given session. pub fn is_turn_active(&self, session_id: &str) -> bool { self.sessions .get(session_id) .is_some_and(|s| s.turn_active.load(Ordering::Acquire)) } /// In-flight tool calls for the given session (`0` when unknown). Only /// the model-facing tool handler ticks the underlying counter, so /// `workspace_rpc` traffic never contributes. pub fn session_active_tool_calls(&self, session_id: &str) -> u32 { self.sessions .get(session_id) .map_or(0, |s| s.active_tool_calls.load(Ordering::Acquire)) } pub fn set_active(&self) { self.lifecycle.store(LIFECYCLE_NONE, Ordering::Release); // Clear the drain stamp symmetrically with `set_draining`: leaving it set // after a resume would make `drain_started_ms` mean "a drain ever began" // rather than "currently draining", and the server's idle gate keys its // unconditional drain escape off that stamp (a stale stamp would let it // report idle while durable work is still outstanding). self.drain_started_ms.store(0, Ordering::Release); self.notify.notify_waiters(); } pub fn set_draining(&self) { self.lifecycle.store(LIFECYCLE_DRAINING, Ordering::Release); // First transition wins, so `drain_started_ms` is stable across calls. let _ = self.drain_started_ms.compare_exchange( 0, now_ms(), Ordering::AcqRel, Ordering::Relaxed, ); self.notify.notify_waiters(); } pub fn set_shutting_down(&self) { self.lifecycle .store(LIFECYCLE_SHUTTING_DOWN, Ordering::Release); self.notify.notify_waiters(); } pub fn is_draining(&self) -> bool { self.lifecycle.load(Ordering::Acquire) >= LIFECYCLE_DRAINING } /// Fully drained: draining, no active tool calls/background tasks, and the /// upload queue emptied (must not exit with artifacts still pending). /// /// In-flight artifact producers are intentionally *not* part of this gate: /// they are withheld from idle by [`Self::durability_gate`] and are awaited /// by the graceful drain's producer phase (`phase 1.5` of /// `WorkspaceHandle::two_phase_drain`) /// *before* the queue flush, so by the time the queue empties their work has /// already been enqueued and is reflected in `upload_queue_pending`. pub fn is_drained(&self) -> bool { self.is_draining() && self.total_active() == 0 && self.upload_queue_pending() == 0 } /// Phase-1 drain condition: all in-flight tool calls and background tasks /// finished, independent of the upload queue. pub fn tools_idle(&self) -> bool { self.total_active() == 0 } pub fn total_active(&self) -> u32 { self.active_tool_calls.load(Ordering::Relaxed) + self.background_tasks.load(Ordering::Relaxed) } pub async fn wait_for_change(&self, timeout: std::time::Duration) { let _ = tokio::time::timeout(timeout, self.notify.notified()).await; } /// Wake the status publisher so it sends a heartbeat immediately. pub fn poke(&self) { self.notify.notify_waiters(); } /// Wait until the tracker is both draining and all active work /// (tool calls + background tasks + the upload queue) has completed. pub async fn wait_until_drained(&self) { loop { if self.is_drained() { return; } // `notify_waiters` stores no permit, so a wake between the check and // this await is missed — safe because every caller is timeout-bounded. self.notify.notified().await; } } /// Wait until all in-flight tool calls and background tasks have finished, /// ignoring the upload queue (phase 1 of the two-phase drain). pub async fn wait_until_tools_idle(&self) { loop { if self.tools_idle() { return; } // Same timeout-bounded missed-wakeup tolerance as `wait_until_drained`. self.notify.notified().await; } } /// 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 { let now = now_ms(); let mut live = Vec::new(); let mut stale = Vec::new(); for entry in self.sessions.iter() { let key = entry.key(); if key == DEFAULT_SESSION { continue; } if entry.value().active_tool_calls.load(Ordering::Relaxed) > 0 { live.push(key.clone()); continue; } let idle_since = entry.value().idle_since_ms.load(Ordering::Relaxed); if idle_since > 0 && now.saturating_sub(idle_since) > self.prune_window_ms { stale.push(key.clone()); } else { live.push(key.clone()); } } for key in &stale { self.sessions.remove(key); self.call_to_session.retain(|_, sid| sid != key); } live } /// Per-session snapshot. `background_task_ids` is the connection aggregate, /// re-published to the session's client. pub fn snapshot_session(&self, session_id: &str) -> ToolServerStatusPayload { let lifecycle = self.lifecycle.load(Ordering::Acquire); let bg = self.background_tasks.load(Ordering::Relaxed); let (active, active_tool_names, last_started, last_completed, idle_since, turn_active) = if let Some(session) = self.sessions.get(session_id) { let a = session.active_tool_calls.load(Ordering::Relaxed); let names: Vec = session .active_tools .iter() .map(|r| r.value().clone()) .collect(); let started = session.last_call_started_ms.load(Ordering::Relaxed); let completed = session.last_call_completed_ms.load(Ordering::Relaxed); let idle = session.idle_since_ms.load(Ordering::Relaxed); let turn = session.turn_active.load(Ordering::Acquire); (a, names, started, completed, idle, turn) } else { (0, vec![], 0, 0, now_ms(), false) }; let status = match lifecycle { LIFECYCLE_SHUTTING_DOWN => ToolServerLifecycleStatus::ShuttingDown, LIFECYCLE_DRAINING => ToolServerLifecycleStatus::Draining, _ if active > 0 => ToolServerLifecycleStatus::Busy, _ => ToolServerLifecycleStatus::Ready, }; let background_task_ids: Vec = self .background_ids .iter() .map(|r| r.key().clone()) .collect(); let d = self.durability_payload_fields(idle_since); ToolServerStatusPayload { status, session_id: xai_tool_protocol::SessionId::new(session_id).ok(), connection_id: None, active_tool_calls: active, active_tool_names, background_tasks: bg, background_task_ids, pending_tool_calls: 0, last_tool_call_started_ms: last_started, last_tool_call_completed_ms: last_completed, uptime_ms: self.started_at.elapsed().as_millis() as u64, idle_since_ms: d.idle_since_ms, upload_queue_pending: d.upload_queue_pending, upload_queue_pending_bytes: d.upload_queue_pending_bytes, upload_queue_inflight: d.upload_queue_inflight, upload_queue_circuit_breaker_tripped: d.upload_queue_circuit_breaker_tripped, artifact_producers_inflight: d.artifact_producers_inflight, drain_started_ms: d.drain_started_ms, turn_active, idle_ignores_background: self.idle_ignores_background, } } /// Aggregate snapshot across all sessions. pub fn snapshot(&self) -> ToolServerStatusPayload { let lifecycle = self.lifecycle.load(Ordering::Acquire); let active = self.active_tool_calls.load(Ordering::Relaxed); let bg = self.background_tasks.load(Ordering::Relaxed); let status = match lifecycle { LIFECYCLE_SHUTTING_DOWN => ToolServerLifecycleStatus::ShuttingDown, LIFECYCLE_DRAINING => ToolServerLifecycleStatus::Draining, _ if active + bg > 0 => ToolServerLifecycleStatus::Busy, _ => ToolServerLifecycleStatus::Ready, }; let active_tool_names: Vec = self .active_tools .iter() .map(|r| r.value().clone()) .collect(); let background_task_ids: Vec = self .background_ids .iter() .map(|r| r.key().clone()) .collect(); let idle_since = self.idle_since_ms.load(Ordering::Relaxed); let d = self.durability_payload_fields(idle_since); ToolServerStatusPayload { status, session_id: None, connection_id: None, active_tool_calls: active, active_tool_names, background_tasks: bg, background_task_ids, pending_tool_calls: 0, last_tool_call_started_ms: self.last_call_started_ms.load(Ordering::Relaxed), last_tool_call_completed_ms: self.last_call_completed_ms.load(Ordering::Relaxed), uptime_ms: self.started_at.elapsed().as_millis() as u64, idle_since_ms: d.idle_since_ms, upload_queue_pending: d.upload_queue_pending, upload_queue_pending_bytes: d.upload_queue_pending_bytes, upload_queue_inflight: d.upload_queue_inflight, upload_queue_circuit_breaker_tripped: d.upload_queue_circuit_breaker_tripped, artifact_producers_inflight: d.artifact_producers_inflight, drain_started_ms: d.drain_started_ms, turn_active: self.any_turn_active(), idle_ignores_background: self.idle_ignores_background, } } } /// Durability tail of a [`ToolServerStatusPayload`], built only by /// [`ActivityTracker::durability_payload_fields`]. struct DurabilityPayloadFields { idle_since_ms: Option, upload_queue_pending: u32, upload_queue_pending_bytes: u64, upload_queue_inflight: u32, upload_queue_circuit_breaker_tripped: bool, artifact_producers_inflight: u32, drain_started_ms: Option, } /// The durability idle-hold cap from `GROK_WORKSPACE_DURABILITY_IDLE_HOLD_MAX_MS`. fn durability_idle_hold_max_from_env() -> u64 { durability_idle_hold_from_raw(std::env::var("GROK_WORKSPACE_DURABILITY_IDLE_HOLD_MAX_MS").ok()) } /// Pure parse of the idle-hold env value: a non-negative integer ms wins (0 /// disables the hold entirely); absent or malformed falls back to the default. fn durability_idle_hold_from_raw(raw: Option) -> u64 { raw.and_then(|s| s.trim().parse::().ok()) .unwrap_or(DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS) } /// Whether a preview-activity stamp still withholds idle at `now`: true while it /// is within `window` ms. A zero stamp (no activity recorded) never withholds, /// and the window is exclusive at the boundary so it decays rather than pins. fn preview_activity_withholds_idle(now: u64, last_activity_ms: u64, window_ms: u64) -> bool { last_activity_ms != 0 && now.saturating_sub(last_activity_ms) < window_ms } fn now_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis() as u64) .unwrap_or(0) } #[cfg(test)] mod tests { use super::*; #[test] fn starts_ready() { let t = ActivityTracker::new(); let s = t.snapshot(); assert_eq!(s.status, ToolServerLifecycleStatus::Ready); assert_eq!(s.active_tool_calls, 0); assert!(s.idle_since_ms.is_some()); assert!(s.session_id.is_none()); } #[test] fn tool_call_transitions_to_busy() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); let s = t.snapshot(); assert_eq!(s.status, ToolServerLifecycleStatus::Busy); assert_eq!(s.active_tool_calls, 1); assert_eq!(s.active_tool_names, vec!["read_file"]); assert!(s.idle_since_ms.is_none()); } #[test] fn tool_call_completion_returns_to_ready() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_completed("c1", None, ToolOutcome::Success); let s = t.snapshot(); assert_eq!(s.status, ToolServerLifecycleStatus::Ready); assert_eq!(s.active_tool_calls, 0); assert!(s.idle_since_ms.is_some()); } #[test] fn background_task_makes_busy() { let t = ActivityTracker::new(); t.background_task_started("t1"); let s = t.snapshot(); assert_eq!(s.status, ToolServerLifecycleStatus::Busy); assert_eq!(s.background_tasks, 1); } #[test] fn background_task_started_dedups_by_id() { let t = ActivityTracker::new(); t.background_task_started("dup"); t.background_task_started("dup"); assert_eq!(t.snapshot().background_tasks, 1); } #[test] fn background_task_completed_unknown_id_does_not_underflow() { let t = ActivityTracker::new(); t.background_task_completed("never-started"); assert_eq!(t.snapshot().background_tasks, 0); t.background_task_started("real"); assert_eq!(t.snapshot().background_tasks, 1); } #[test] fn background_task_decrement_restores_idle_only_at_zero() { let t = ActivityTracker::new(); t.background_task_started("a"); t.background_task_started("b"); assert!(t.snapshot().idle_since_ms.is_none()); t.background_task_completed("a"); assert!( t.snapshot().idle_since_ms.is_none(), "one bg task left → still not idle" ); t.background_task_completed("b"); assert!( t.snapshot().idle_since_ms.is_some(), "idle restored only when the last bg task completes" ); } #[test] fn background_after_calls_complete_pins_idle_only_when_flag_off() { let on = ActivityTracker::new().with_idle_ignores_background(true); on.tool_call_started("c1", "read_file", None); on.tool_call_completed("c1", None, ToolOutcome::Success); on.background_task_started("bg1"); assert!(on.snapshot().idle_since_ms.is_some()); let off = ActivityTracker::new(); off.tool_call_started("c1", "read_file", None); off.tool_call_completed("c1", None, ToolOutcome::Success); off.background_task_started("bg1"); assert!(off.snapshot().idle_since_ms.is_none()); } #[test] fn flag_on_active_call_withholds_idle_until_it_completes_despite_background() { let t = ActivityTracker::new().with_idle_ignores_background(true); t.tool_call_started("c1", "read_file", None); t.background_task_started("bg1"); assert!(t.snapshot().idle_since_ms.is_none()); t.tool_call_completed("c1", None, ToolOutcome::Success); assert!(t.snapshot().idle_since_ms.is_some()); } #[test] fn flag_on_keeps_busy_status_while_reporting_idle() { let t = ActivityTracker::new().with_idle_ignores_background(true); t.background_task_started("bg1"); let s = t.snapshot(); assert_eq!(s.status, ToolServerLifecycleStatus::Busy); assert!(s.idle_since_ms.is_some()); } #[test] fn flag_on_background_completion_does_not_advance_idle() { let t = ActivityTracker::new().with_idle_ignores_background(true); t.background_task_started("bg1"); let after_start = t.snapshot().idle_since_ms; std::thread::sleep(std::time::Duration::from_millis(5)); t.background_task_completed("bg1"); assert_eq!(t.snapshot().idle_since_ms, after_start); } #[test] fn flag_on_background_start_advances_idle_when_no_active_call() { let t = ActivityTracker::new().with_idle_ignores_background(true); t.tool_call_started("c1", "read_file", None); t.tool_call_completed("c1", None, ToolOutcome::Success); let before = t.snapshot().idle_since_ms.unwrap(); std::thread::sleep(std::time::Duration::from_millis(5)); t.background_task_started("bg1"); let after = t.snapshot().idle_since_ms.unwrap(); assert!(after > before); } #[test] fn drain_counts_background_tasks_regardless_of_flag() { for flag in [false, true] { let t = ActivityTracker::new().with_idle_ignores_background(flag); t.background_task_started("bg1"); assert_eq!(t.total_active(), 1); t.set_draining(); assert!(!t.is_drained()); assert!(!t.tools_idle()); t.background_task_completed("bg1"); assert!(t.is_drained()); assert!(t.tools_idle()); } } #[test] fn snapshot_payloads_report_idle_ignores_background_flag() { let on = ActivityTracker::new().with_idle_ignores_background(true); assert!(on.snapshot().idle_ignores_background); assert!(on.snapshot_session("sess-a").idle_ignores_background); let off = ActivityTracker::new(); assert!(!off.snapshot().idle_ignores_background); assert!(!off.snapshot_session("sess-a").idle_ignores_background); } #[test] fn draining_overrides_busy() { let t = ActivityTracker::new(); t.tool_call_started("c1", "grep", None); t.set_draining(); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Draining); } #[test] fn set_active_clears_draining() { let t = ActivityTracker::new(); t.set_draining(); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Draining); t.set_active(); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Ready); assert!(!t.is_draining()); t.tool_call_started("c1", "read_file", Some("sess-a")); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Busy); } #[test] fn shutting_down_overrides_draining() { let t = ActivityTracker::new(); t.set_draining(); t.set_shutting_down(); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::ShuttingDown); } #[test] fn is_drained_requires_zero_active() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", None); t.set_draining(); assert!(!t.is_drained()); t.tool_call_completed("c1", None, ToolOutcome::Success); assert!(t.is_drained()); } // ── upload-queue coupling + drain status fields ─────────────── /// Coupled queue depth surfaces in both aggregate and per-session payloads. #[test] fn snapshot_reports_coupled_upload_queue_stats() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(3, Ordering::Relaxed); stats.pending_bytes.store(9000, Ordering::Relaxed); stats.inflight.store(1, Ordering::Relaxed); stats.circuit_breaker_active.store(true, Ordering::Relaxed); t.set_upload_queue_stats(stats); let agg = t.snapshot(); assert_eq!(agg.upload_queue_pending, 3); assert_eq!(agg.upload_queue_pending_bytes, 9000); assert_eq!(agg.upload_queue_inflight, 1); assert!(agg.upload_queue_circuit_breaker_tripped); t.tool_call_started("c1", "x", Some("sess-a")); let sess = t.snapshot_session("sess-a"); assert_eq!(sess.upload_queue_pending, 3); assert_eq!(sess.upload_queue_inflight, 1); assert!(sess.upload_queue_circuit_breaker_tripped); } /// The queue-less path reports zeroed queue fields (legacy behaviour). #[test] fn snapshot_queue_fields_zero_without_coupled_queue() { let t = ActivityTracker::new(); let s = t.snapshot(); assert_eq!(s.upload_queue_pending, 0); assert_eq!(s.upload_queue_pending_bytes, 0); assert_eq!(s.upload_queue_inflight, 0); assert!(!s.upload_queue_circuit_breaker_tripped); assert_eq!(s.drain_started_ms, None); } /// Draining with no tool calls but a non-empty queue is NOT drained. #[test] fn is_drained_requires_empty_upload_queue() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(2, Ordering::Relaxed); t.set_upload_queue_stats(stats.clone()); t.set_draining(); assert!( !t.is_drained(), "queue still has 2 pending → not drained even with no tool calls" ); stats.pending.store(0, Ordering::Relaxed); assert!(t.is_drained(), "queue emptied → now fully drained"); } /// The phase-1 drain condition must not wait on the queue. #[test] fn tools_idle_ignores_upload_queue() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(5, Ordering::Relaxed); t.set_upload_queue_stats(stats); assert!( t.tools_idle(), "no tool calls → tools idle regardless of queue" ); assert!( !t.is_drained(), "but not fully drained while queue is non-empty" ); } // ── durability-aware idle gating ─────────────────────────────── #[tokio::test] async fn idle_withheld_while_producer_in_flight() { let t = ActivityTracker::new(); let tasks = tokio_util::task::TaskTracker::new(); t.set_producer_tasks(tasks.clone()); let s = t.snapshot(); assert_eq!(s.artifact_producers_inflight, 0); assert!(s.idle_since_ms.is_some(), "no producers → idle reported"); 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); assert!( s.idle_since_ms.is_none(), "an in-flight producer must withhold idle" ); assert!(t.snapshot_session("any").idle_since_ms.is_none()); gate.notify_one(); join.await.expect("producer task must not panic"); let s = t.snapshot(); assert_eq!(s.artifact_producers_inflight, 0); assert!( s.idle_since_ms.is_some(), "idle must be restored once the producer completes" ); } #[test] fn idle_withheld_while_upload_queue_pending() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(1, Ordering::Relaxed); t.set_upload_queue_stats(stats.clone()); assert!( t.snapshot().idle_since_ms.is_none(), "pending uploads must withhold idle" ); stats.pending.store(0, Ordering::Relaxed); assert!( t.snapshot().idle_since_ms.is_some(), "idle must be restored once the queue empties" ); } #[test] fn durability_hold_cap_expiry_allows_idle() { let t = ActivityTracker::with_prune_window_and_idle_hold( std::time::Duration::from_millis(SESSION_IDLE_PRUNE_MS), 1_000, ); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(1, Ordering::Relaxed); t.set_upload_queue_stats(stats); assert!(t.snapshot().idle_since_ms.is_none(), "within the hold cap"); // Backdate the busy stamp past the cap. t.durability_busy_since_ms .store(now_ms() - 2_000, Ordering::Relaxed); let s = t.snapshot(); assert!( s.idle_since_ms.is_some(), "expired hold cap must allow idle despite pending work" ); assert_eq!( s.upload_queue_pending, 1, "the pending depth stays truthfully reported" ); } #[test] fn durability_busy_stamp_resets_when_condition_clears() { let t = ActivityTracker::with_prune_window_and_idle_hold( std::time::Duration::from_millis(SESSION_IDLE_PRUNE_MS), 1_000, ); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(1, Ordering::Relaxed); t.set_upload_queue_stats(stats.clone()); let _ = t.snapshot(); // stamps the busy start t.durability_busy_since_ms .store(now_ms() - 600, Ordering::Relaxed); stats.pending.store(0, Ordering::Relaxed); let _ = t.snapshot(); assert_eq!( t.durability_busy_since_ms.load(Ordering::Relaxed), 0, "a clear condition must reset the stamp" ); // A new busy condition measures its hold from a fresh stamp. stats.pending.store(1, Ordering::Relaxed); assert!(t.snapshot().idle_since_ms.is_none()); } #[test] fn breaker_tripped_allows_idle_despite_pending_work() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.pending.store(3, Ordering::Relaxed); stats.circuit_breaker_active.store(true, Ordering::Relaxed); t.set_upload_queue_stats(stats); let s = t.snapshot(); assert!( s.idle_since_ms.is_some(), "a tripped breaker must allow hibernation despite pending work" ); assert!(s.upload_queue_circuit_breaker_tripped); } /// The breaker escape covers queued items only (they survive via disk /// spill + restart recovery); an in-flight producer has nothing on disk /// yet and must keep withholding idle even with the breaker tripped. #[tokio::test] async fn breaker_tripped_does_not_bypass_producer_hold() { let t = ActivityTracker::new(); let stats = Arc::new(UploadQueueStats::new()); stats.circuit_breaker_active.store(true, Ordering::Relaxed); t.set_upload_queue_stats(stats); let tasks = tokio_util::task::TaskTracker::new(); t.set_producer_tasks(tasks.clone()); let gate = Arc::new(tokio::sync::Notify::new()); let gate2 = gate.clone(); let join = tasks.spawn(async move { gate2.notified().await }); assert!( t.snapshot().idle_since_ms.is_none(), "a producer must withhold idle even with the breaker tripped" ); gate.notify_one(); join.await.expect("producer task must not panic"); assert!( t.snapshot().idle_since_ms.is_some(), "idle restored once the producer completes (breaker alone is no hold)" ); } #[test] fn preview_activity_withholds_idle_window_boundaries() { let now = 10_000_000; assert!( !preview_activity_withholds_idle(now, 0, PREVIEW_ACTIVITY_WINDOW_MS), "a zero stamp (no activity ever) must never withhold" ); assert!( preview_activity_withholds_idle(now, now, PREVIEW_ACTIVITY_WINDOW_MS), "activity right now withholds" ); assert!( preview_activity_withholds_idle( now, now - (PREVIEW_ACTIVITY_WINDOW_MS - 1), PREVIEW_ACTIVITY_WINDOW_MS ), "just inside the window withholds" ); assert!( !preview_activity_withholds_idle( now, now - PREVIEW_ACTIVITY_WINDOW_MS, PREVIEW_ACTIVITY_WINDOW_MS ), "exactly at the window edge no longer withholds (decaying, exclusive)" ); assert!( !preview_activity_withholds_idle( now, now - (PREVIEW_ACTIVITY_WINDOW_MS + 5_000), PREVIEW_ACTIVITY_WINDOW_MS ), "past the window no longer withholds" ); } #[test] fn note_preview_activity_withholds_then_resumes_idle() { let t = ActivityTracker::new(); assert!( t.snapshot().idle_since_ms.is_some(), "an idle tracker reports idle before any preview activity" ); t.note_preview_activity(); 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" ); t.last_preview_activity_ms.store( now_ms().saturating_sub(PREVIEW_ACTIVITY_WINDOW_MS + 1_000), Ordering::Relaxed, ); assert!( t.snapshot().idle_since_ms.is_some(), "idle must resume once the preview window decays" ); } #[test] fn configured_preview_window_overrides_default() { let configured = ActivityTracker::new().with_preview_activity_window_ms(500); configured .last_preview_activity_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 .store(now_ms().saturating_sub(1_000), Ordering::Relaxed); assert!(default.snapshot().idle_since_ms.is_none()); } #[test] 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(); let s = t.snapshot(); assert!(s.idle_since_ms.is_none()); assert_eq!( s.status, ToolServerLifecycleStatus::Busy, "an in-flight tool call stays Busy; preview activity only gates idle reporting" ); } #[test] fn durability_idle_hold_from_raw_parses_and_falls_back() { assert_eq!( durability_idle_hold_from_raw(None), DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS ); assert_eq!(durability_idle_hold_from_raw(Some("1234".into())), 1234); assert_eq!(durability_idle_hold_from_raw(Some(" 99 ".into())), 99); assert_eq!(durability_idle_hold_from_raw(Some("0".into())), 0); assert_eq!( durability_idle_hold_from_raw(Some("nonsense".into())), DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS ); assert_eq!( durability_idle_hold_from_raw(Some("-5".into())), DEFAULT_DURABILITY_IDLE_HOLD_MAX_MS ); } #[test] fn set_draining_stamps_drain_started_ms_once() { let t = ActivityTracker::new(); assert_eq!(t.snapshot().drain_started_ms, None); t.set_draining(); let first = t.snapshot().drain_started_ms.expect("drain_started_ms set"); assert!(first > 0); // A second set_draining must not move the stamp. t.set_draining(); assert_eq!(t.snapshot().drain_started_ms, Some(first)); } #[test] fn set_active_clears_drain_started_ms() { let t = ActivityTracker::new(); t.set_draining(); assert!(t.snapshot().drain_started_ms.is_some()); // Resuming clears the stamp so it tracks "currently draining", not // "ever drained" — the server idle gate's drain escape keys off it. t.set_active(); assert_eq!(t.snapshot().drain_started_ms, None); // A fresh drain after the resume re-stamps with a new value. t.set_draining(); assert!(t.snapshot().drain_started_ms.is_some()); } #[test] fn turn_active_surfaces_in_snapshots() { let t = ActivityTracker::new(); assert!(!t.snapshot().turn_active); t.turn_started("sess-a", 1); assert!(t.snapshot().turn_active, "aggregate: any session active"); assert!(t.snapshot_session("sess-a").turn_active); assert!( !t.snapshot_session("sess-b").turn_active, "a different session is not turn-active" ); t.turn_completed("sess-a", 1, 10); assert!(!t.snapshot().turn_active); } /// The notify handle wakes the same waiter the status publisher blocks on. #[tokio::test] async fn notify_handle_shares_the_publisher_waiter() { let t = Arc::new(ActivityTracker::new()); let notify = t.notify_handle(); let t2 = t.clone(); let waiter = tokio::spawn( async move { t2.wait_for_change(std::time::Duration::from_secs(5)).await }, ); tokio::task::yield_now().await; notify.notify_waiters(); tokio::time::timeout(std::time::Duration::from_secs(2), waiter) .await .expect("publisher waiter must wake via the shared notify handle") .expect("waiter task should not panic"); } #[test] fn multiple_concurrent_calls() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", None); t.tool_call_started("c2", "grep", None); assert_eq!(t.snapshot().active_tool_calls, 2); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 1); t.tool_call_completed("c2", None, ToolOutcome::Success); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Ready); } #[test] fn background_task_keeps_busy_after_calls_complete() { let t = ActivityTracker::new(); t.tool_call_started("c1", "run_terminal_cmd", None); t.background_task_started("bg1"); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Busy); t.background_task_completed("bg1"); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Ready); } #[test] fn per_session_independent_tracking() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_started("c2", "grep", Some("sess-b")); let sa = t.snapshot_session("sess-a"); assert_eq!(sa.active_tool_calls, 1); assert_eq!(sa.session_id.as_ref().map(|s| s.as_str()), Some("sess-a")); let sb = t.snapshot_session("sess-b"); assert_eq!(sb.active_tool_calls, 1); assert_eq!(t.snapshot().active_tool_calls, 2); } #[test] fn per_session_completion_independent() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_started("c2", "grep", Some("sess-b")); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!( t.snapshot_session("sess-a").status, ToolServerLifecycleStatus::Ready ); assert_eq!( t.snapshot_session("sess-b").status, ToolServerLifecycleStatus::Busy ); } #[test] fn unknown_session_returns_ready() { let t = ActivityTracker::new(); assert_eq!( t.snapshot_session("nonexistent").status, ToolServerLifecycleStatus::Ready ); } #[test] fn known_sessions_excludes_default() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", None); t.tool_call_started("c2", "y", Some("sess-a")); assert_eq!(t.known_sessions(), vec!["sess-a"]); } #[test] fn draining_propagates_to_session_snapshot() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("sess-a")); t.set_draining(); assert_eq!( t.snapshot_session("sess-a").status, ToolServerLifecycleStatus::Draining ); } #[test] fn prune_keeps_active_sessions() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("active")); assert_eq!(t.known_sessions(), vec!["active"]); } #[test] fn prune_keeps_recently_idle() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("recent")); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.known_sessions(), vec!["recent"]); } #[test] fn prune_removes_stale() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("stale")); t.tool_call_completed("c1", None, ToolOutcome::Success); if let Some(session) = t.sessions.get("stale") { let old = now_ms() - SESSION_IDLE_PRUNE_MS - 1000; session.idle_since_ms.store(old, Ordering::Relaxed); } assert!(t.known_sessions().is_empty()); assert!(!t.sessions.contains_key("stale")); } #[test] fn small_prune_window_evicts_session_default_window_retains() { // A session idle for ~50ms: pruned under a 10ms window, retained // under the default 300s window. Proves the window is actually used. let idle_ago = 50; let small = ActivityTracker::with_prune_window(std::time::Duration::from_millis(10)); small.tool_call_started("c1", "x", Some("sess")); small.tool_call_completed("c1", None, ToolOutcome::Success); if let Some(session) = small.sessions.get("sess") { session .idle_since_ms .store(now_ms() - idle_ago, Ordering::Relaxed); } assert!( small.known_sessions().is_empty(), "session idle past the small window must be pruned" ); assert!(!small.sessions.contains_key("sess")); let default = ActivityTracker::new(); default.tool_call_started("c1", "x", Some("sess")); default.tool_call_completed("c1", None, ToolOutcome::Success); if let Some(session) = default.sessions.get("sess") { session .idle_since_ms .store(now_ms() - idle_ago, Ordering::Relaxed); } assert_eq!( default.known_sessions(), vec!["sess"], "the same idle time must be retained under the default window" ); } #[test] fn duplicate_start_is_noop() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_started("c1", "grep", Some("sess-b")); // Second start with the same call_id must be ignored. assert_eq!(t.snapshot().active_tool_calls, 1); assert_eq!(t.snapshot_session("sess-a").active_tool_calls, 1); // sess-b should not have been created by the duplicate. assert_eq!(t.snapshot_session("sess-b").active_tool_calls, 0); // Completion still works normally. t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); } #[test] fn duplicate_completion_is_noop() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", None); assert_eq!(t.snapshot().active_tool_calls, 1); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); // Second completion of the same call_id must not underflow. t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); } #[test] fn unknown_call_id_completion_is_noop() { let t = ActivityTracker::new(); // Completing a call_id that was never started must not underflow. t.tool_call_completed("never-started", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); assert_eq!(t.snapshot().status, ToolServerLifecycleStatus::Ready); } #[test] fn duplicate_background_start_is_noop() { let t = ActivityTracker::new(); t.background_task_started("bg1"); t.background_task_started("bg1"); assert_eq!(t.snapshot().background_tasks, 1); t.background_task_completed("bg1"); assert_eq!(t.snapshot().background_tasks, 0); } #[test] fn duplicate_background_completion_is_noop() { let t = ActivityTracker::new(); t.background_task_started("bg1"); assert_eq!(t.snapshot().background_tasks, 1); t.background_task_completed("bg1"); assert_eq!(t.snapshot().background_tasks, 0); // Second completion must not underflow. t.background_task_completed("bg1"); assert_eq!(t.snapshot().background_tasks, 0); } #[test] fn unknown_background_task_completion_is_noop() { let t = ActivityTracker::new(); t.background_task_completed("never-started"); assert_eq!(t.snapshot().background_tasks, 0); } #[test] fn prune_cleans_call_to_session_mapping() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("stale")); t.tool_call_completed("c1", None, ToolOutcome::Success); // Simulate stale idle time. if let Some(session) = t.sessions.get("stale") { let old = now_ms() - SESSION_IDLE_PRUNE_MS - 1000; session.idle_since_ms.store(old, Ordering::Relaxed); } // Prune should remove the session and its call_to_session entries. t.known_sessions(); assert!(!t.sessions.contains_key("stale")); // Verify call_to_session was cleaned (no dangling entries). assert!(!t.call_to_session.contains_key("c1")); } #[test] fn per_session_completion_after_prune_is_safe() { let t = ActivityTracker::new(); t.tool_call_started("c1", "x", Some("sess")); // Force-prune the session while the call is still "active" // (simulates the theoretical race). t.sessions.remove("sess"); // Completing should not panic or underflow — the session // lookup returns None, so only the global counter decrements. t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); } #[test] fn turn_started_sets_current_turn_and_active() { let t = ActivityTracker::new(); t.turn_started("sess-a", 1); let session = t.sessions.get("sess-a").expect("session should exist"); assert_eq!(session.current_turn.load(Ordering::Acquire), 1); assert!(session.turn_active.load(Ordering::Acquire)); } #[test] fn turn_completed_clears_active_when_turn_matches() { let t = ActivityTracker::new(); t.turn_started("sess-a", 1); t.turn_completed("sess-a", 1, 500); let session = t.sessions.get("sess-a").expect("session should exist"); assert!(!session.turn_active.load(Ordering::Acquire)); } #[test] fn turn_completed_stale_turn_does_not_clear_active() { let t = ActivityTracker::new(); t.turn_started("sess-a", 1); t.turn_started("sess-a", 2); // Completing the stale turn 1 must not clear turn_active for turn 2. t.turn_completed("sess-a", 1, 500); let session = t.sessions.get("sess-a").expect("session should exist"); assert!( session.turn_active.load(Ordering::Acquire), "turn_active should still be true for current turn 2" ); assert_eq!(session.current_turn.load(Ordering::Acquire), 2); } #[test] fn turn_completed_unknown_session_is_noop() { let t = ActivityTracker::new(); // Must not panic or create a session entry. t.turn_completed("nonexistent", 1, 100); assert!(!t.sessions.contains_key("nonexistent")); } #[test] fn session_ended_clears_turn_active() { 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)); t.session_ended("sess-a"); assert!( !session.turn_active.load(Ordering::Acquire), "turn_active should be cleared after session_ended" ); } #[test] fn session_ended_unknown_session_is_noop() { let t = ActivityTracker::new(); // Must not panic or create a session entry. t.session_ended("nonexistent"); assert!(!t.sessions.contains_key("nonexistent")); } #[test] fn session_ended_without_active_turn_is_safe() { let t = ActivityTracker::new(); // Create a session via a tool call cycle (turn never started). 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)); } #[tokio::test] async fn session_ended_notifies_waiters() { let t = std::sync::Arc::new(ActivityTracker::new()); t.turn_started("sess-a", 1); let t2 = t.clone(); let waiter = tokio::spawn(async move { t2.wait_for_change(std::time::Duration::from_secs(5)).await; }); // Give the waiter a moment to register. tokio::task::yield_now().await; t.session_ended("sess-a"); // The waiter should complete promptly (not timeout). tokio::time::timeout(std::time::Duration::from_secs(2), waiter) .await .expect("waiter should complete within timeout") .expect("waiter task should not panic"); } #[test] fn session_ended_with_inflight_tool_calls() { let t = ActivityTracker::new(); t.turn_started("sess-a", 1); t.tool_call_started("c1", "read_file", Some("sess-a")); assert!(t.is_turn_active("sess-a")); assert_eq!(t.snapshot_session("sess-a").active_tool_calls, 1); t.session_ended("sess-a"); // turn_active cleared, but the in-flight tool call remains. assert!(!t.is_turn_active("sess-a")); assert_eq!( t.snapshot_session("sess-a").active_tool_calls, 1, "session_ended must not clear active tool calls" ); assert_eq!( t.snapshot_session("sess-a").status, ToolServerLifecycleStatus::Busy, "session should still be busy due to in-flight tool call" ); } #[test] fn cancel_all_session_calls_completes_all_for_session() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_started("c2", "grep", Some("sess-a")); t.tool_call_started("c3", "write", Some("sess-b")); assert_eq!(t.snapshot().active_tool_calls, 3); let cancelled = t.cancel_all_session_calls("sess-a"); assert_eq!(cancelled, 2, "should cancel 2 calls for sess-a"); assert_eq!( t.snapshot().active_tool_calls, 1, "only sess-b call remains" ); assert_eq!(t.snapshot_session("sess-a").active_tool_calls, 0); assert_eq!(t.snapshot_session("sess-b").active_tool_calls, 1); } #[test] fn cancel_all_session_calls_noop_for_empty_session() { let t = ActivityTracker::new(); let cancelled = t.cancel_all_session_calls("nonexistent"); assert_eq!(cancelled, 0); assert_eq!(t.snapshot().active_tool_calls, 0); } #[test] fn cancel_all_session_calls_idempotent() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); assert_eq!(t.cancel_all_session_calls("sess-a"), 1); assert_eq!( t.cancel_all_session_calls("sess-a"), 0, "second cancel_all should find no calls" ); assert_eq!(t.snapshot().active_tool_calls, 0); } #[test] fn double_cancel_via_tool_call_completed_is_idempotent() { let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); assert_eq!(t.snapshot().active_tool_calls, 1); t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!(t.snapshot().active_tool_calls, 0); // Second completion of the same call_id must not underflow. t.tool_call_completed("c1", None, ToolOutcome::Success); assert_eq!( t.snapshot().active_tool_calls, 0, "double cancel must not underflow active_tool_calls" ); assert_eq!( t.snapshot_session("sess-a").active_tool_calls, 0, "double cancel must not underflow session active_tool_calls" ); } // ── events.jsonl emission ───────────────────────────────── /// Build a tracker whose event sink points at a fresh tempdir, with the /// `events.jsonl` writer for `session` pre-opened (as turn-start would do). fn tracker_with_writer(session: &str) -> (ActivityTracker, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let writers: Arc> = Arc::new(DashMap::new()); writers.insert(session.to_owned(), EventWriter::open(dir.path())); let t = ActivityTracker::new(); t.set_event_writers(writers); (t, dir) } fn read_events(dir: &tempfile::TempDir) -> Vec { let text = std::fs::read_to_string(dir.path().join("events.jsonl")).unwrap(); text.trim() .lines() .map(|l| serde_json::from_str(l).unwrap()) .collect() } #[test] fn emits_tool_started_and_completed_with_real_content() { let (t, dir) = tracker_with_writer("sess-a"); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_completed("c1", Some("sess-a"), ToolOutcome::Success); let events = read_events(&dir); assert_eq!(events.len(), 2, "expected ToolStarted + ToolCompleted"); assert_eq!(events[0]["type"], "tool_started"); assert_eq!(events[0]["tool_name"], "read_file"); assert_eq!(events[1]["type"], "tool_completed"); assert_eq!(events[1]["tool_name"], "read_file"); assert_eq!(events[1]["outcome"], "success"); assert_eq!(events[1]["source"], "workspace"); assert!( events[1]["duration_ms"].as_u64().is_some(), "ToolCompleted must carry a duration_ms" ); } #[test] fn tool_completed_outcome_is_truthful_error() { let (t, dir) = tracker_with_writer("sess-a"); t.tool_call_started("c1", "run_terminal_command", Some("sess-a")); t.tool_call_completed("c1", Some("sess-a"), ToolOutcome::Error); let events = read_events(&dir); assert_eq!(events[1]["type"], "tool_completed"); assert_eq!(events[1]["outcome"], "error"); } #[test] fn cancel_all_marks_tool_completed_cancelled() { let (t, dir) = tracker_with_writer("sess-a"); t.tool_call_started("c1", "run_terminal_command", Some("sess-a")); assert_eq!(t.cancel_all_session_calls("sess-a"), 1); let events = read_events(&dir); let last = events.last().unwrap(); assert_eq!(last["type"], "tool_completed"); assert_eq!(last["outcome"], "cancelled"); } #[test] fn no_tool_events_without_event_sink() { // Behaviour preservation: the default tracker (no event sink — the state // for all the legacy tests above) must never touch the filesystem. let t = ActivityTracker::new(); t.tool_call_started("c1", "read_file", Some("sess-a")); t.tool_call_completed("c1", Some("sess-a"), ToolOutcome::Success); // Counters still behave, and nothing was recorded. assert_eq!(t.snapshot().active_tool_calls, 0); assert!(t.call_started_ms.is_empty()); } #[test] fn no_event_when_session_writer_not_open() { // Sink configured but the session's writer was never opened (no // turn-start): no event is emitted and no start time leaks. let dir = tempfile::tempdir().unwrap(); let writers: Arc> = Arc::new(DashMap::new()); writers.insert("other".to_owned(), EventWriter::open(dir.path())); let t = ActivityTracker::new(); t.set_event_writers(writers); t.tool_call_started("c1", "read_file", Some("unopened-sess")); t.tool_call_completed("c1", Some("unopened-sess"), ToolOutcome::Success); // The "other" session's events.jsonl must stay empty. let text = std::fs::read_to_string(dir.path().join("events.jsonl")).unwrap(); assert!(text.trim().is_empty(), "no event should be written"); // No start time was ever recorded: the insert is gated on this session's // writer being open, and "unopened-sess" has none — so the map is empty // (this is exactly the production flag-off shape: sink wired, map empty). assert!(t.call_started_ms.is_empty()); } }