Kigi-CLI/crates/codegen/kigi-workspace/src/activity.rs

//! 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 kigi_file_utils::events::{Event, EventWriter, ToolOutcome};
use kigi_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;

/// 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<String, String>,
    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<String, String>,
    background_tasks: AtomicU32,
    background_ids: DashMap<String, ()>,
    last_call_started_ms: AtomicU64,
    last_call_completed_ms: AtomicU64,
    idle_since_ms: AtomicU64,
    started_at: Instant,
    lifecycle: AtomicU8,
    /// `Arc` so external wakers (via [`notify_handle`](Self::notify_handle))
    /// can wake the same waiter the status publisher blocks on.
    notify: Arc<tokio::sync::Notify>,
    /// Epoch ms a graceful drain began; `0` means "not draining".
    drain_started_ms: AtomicU64,
    /// 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<String, SessionActivity>,
    /// call_id → session_id so `tool_call_completed` can decrement
    /// the right session without the caller repeating it.
    call_to_session: DashMap<String, String>,
    /// 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<Arc<DashMap<String, EventWriter>>>,
    /// 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
    /// `KIGI_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<String, (u64, EventWriter)>,
}

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.
    pub fn with_prune_window(prune_window: std::time::Duration) -> 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()),
            drain_started_ms: AtomicU64::new(0),
            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<DashMap<String, EventWriter>>) {
        let _ = self.event_writers.set(writers);
    }

    /// Clone of the internal `Notify` for external wakers to drive republishes.
    pub fn notify_handle(&self) -> Arc<tokio::sync::Notify> {
        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();
    }

    /// 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
    }

    /// Idle/drain tail shared by [`Self::snapshot`] and [`Self::snapshot_session`]
    /// (one construction site so the two payloads can't drift). Returns
    /// `(idle_since_ms, drain_started_ms)`; recent preview traffic withholds idle.
    fn idle_and_drain_fields(&self, idle_since: u64) -> (Option<u64>, Option<u64>) {
        let drain_started = match self.drain_started_ms.load(Ordering::Relaxed) {
            0 => None,
            ms => Some(ms),
        };
        let withhold_idle = self.preview_withholds_idle(now_ms());
        let idle_since_ms = if idle_since == 0 || withhold_idle {
            None
        } else {
            Some(idle_since)
        };
        (idle_since_ms, drain_started)
    }

    /// 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<EventWriter> {
        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) {
            let duration_ms = now.saturating_sub(started_ms);
            writer.emit(Event::ToolCompleted {
                tool_name,
                duration_ms,
                outcome,
            });
        }
    }

    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<String> = 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 and no active tool calls/background tasks.
    pub fn is_drained(&self) -> bool {
        self.is_draining() && self.total_active() == 0
    }

    /// Drain condition: all in-flight tool calls and background tasks finished.
    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) 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.
    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<String> {
        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<String> = 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<String> = self
            .background_ids
            .iter()
            .map(|r| r.key().clone())
            .collect();

        let (idle_since_ms, drain_started_ms) = self.idle_and_drain_fields(idle_since);

        ToolServerStatusPayload {
            status,
            session_id: kigi_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,
            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<String> = self
            .active_tools
            .iter()
            .map(|r| r.value().clone())
            .collect();
        let background_task_ids: Vec<String> = self
            .background_ids
            .iter()
            .map(|r| r.key().clone())
            .collect();

        let idle_since = self.idle_since_ms.load(Ordering::Relaxed);

        let (idle_since_ms, drain_started_ms) = self.idle_and_drain_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,
            drain_started_ms,
            turn_active: self.any_turn_active(),
            idle_ignores_background: self.idle_ignores_background,
        }
    }
}

/// 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());
    }

    #[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 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<DashMap<String, EventWriter>> = 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<serde_json::Value> {
        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!(
            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<DashMap<String, EventWriter>> = 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());
    }
}